summaries: The paper proposes a novel generalisation measure, i.e., measurement that indicates how well the network generalises, based on pruning. The idea is to measure the fraction of the weights that can be pruned (either randomly, or based on the norms) without hurting the training loss of the model. The paper provides thorough discussion of the related methods and motivates the measure in multiple ways. Further, the authors show empirical evidence for the correlation of the pruning robustness to the generalisation ability of networks, based on the paper by Jiang et al., 2019 and dataset (updated with additional models) provided in the paper.

summaries: In order to understand why deep networks generalize well, this paper proposes "prunability" as an empirical measure that can be predictive of the generalization. Prunability is roughly the smallest _fraction_ (i.e., $\in [0,1]$) of parameters that can be retained, while zeroing out everything else, without increasing the model's training loss by too much. The authors experimentally demonstrate the predictive ability of this measure in three ways.

summaries: In the present work, the authors tackle the highly debated (and sometimes confusing) problem of finding a good simplicity/complexity measure able to predict generalization performance of deep networks. A novel measure called 'prunability' is introduced and compared with some of the many alternatives in the literature. This property measures how networks are able to retain low training loss when a fraction of the weights is set to zero, and is clearly related to common training practices (e.g. dropout) that seems to yield better generalization performance in practice. The experimental settings and the evaluation methods for this new metric are inspired by recent extensive studies on deep networks performance. The authors are able to show that prunability is in fact associated with good generalization and seems able to capture some non-trivial phenomena (double-descent), but they also find it to be inferior to pre-existing (margin based) measures. Moreover, the close relationship to perturbation robustness and flatness measures is investigated, but the results are not fully conclusive.

abstractText: Why over-parameterized neural networks generalize as well as they do is a central concern of theoretical analysis in machine learning today. Following Occam’s razor, it has long been suggested that simpler networks generalize better than more complex ones. Successfully quantifying this principle has proved difficult given that many measures of simplicity, such as parameter norms, grow with the size of the network and thus fail to capture the observation that larger networks tend to generalize better in practice. In this paper, we introduce a new, theoretically motivated measure of a network’s simplicity: the smallest fraction of the network’s parameters that can be kept while pruning without adversely affecting its training loss. We show that this measure is highly predictive of a model’s generalization performance across a large set of convolutional networks trained on CIFAR-10. Lastly, we study the mutual information between the predictions of our new measure and strong existing measures based on models’ margin, flatness of minima and optimization speed. We show that our new measure is similar to – but more predictive than – existing flatness-based measures.

authors:
Zeyuan Allen-Zhu
Yuanzhi Li
Yingyu Liang

title: Learning and generalization in overparameterized neural networks, going beyond two layers

venue: In Advances in Neural Information Processing Systems,

authors:
Sanjeev Arora
Rong Ge
Behnam Neyshabur
Yi Zhang

title: Stronger generalization bounds for deep nets via a compression approach

venue: Proceedings of Machine Learning Research,

authors:
Peter L Bartlett
Dylan J Foster
Matus J Telgarsky

title: Spectrally-normalized margin bounds for neural networks

authors:
Mikhail Belkin
Daniel Hsu
Siyuan Ma
Soumik Mandal

title: Reconciling modern machinelearning practice and the classical bias–variance trade-off

venue: Proceedings of the National Academy of Sciences, 116:201903070,

authors:
Davis Blalock
Jose Javier Gonzalez Ortiz
Jonathan Frankle
John Guttag

title: What is the state of neural network pruning

venue: In Proceedings of 3rd MLSys Conference,

authors:
Gintare Karolina Dziugaite
Daniel M Roy

title: Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data

venue: In Proceedings of the Thirty-Third Conference on Uncertainty in Artificial Intelligence,

authors:
Jonathan Frankle
Michael Carbin

title: The lottery ticket hypothesis: Finding sparse, trainable neural networks

venue: In International Conference on Learning Representations,

authors:
Trevor Gale
Erich Elsen
Sara Hooker

title: The state of sparsity in deep neural networks

venue: arXiv preprint arXiv:1902.09574,

authors:
Song Han
Jeff Pool
John Tran
William Dally

title: Learning both weights and connections for efficient neural network

venue: In Advances in neural information processing systems,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Identity mappings in deep residual networks

venue: In European conference on computer vision,

authors:
Geoffrey Hinton
Drew Van Camp

title: Keeping neural networks simple by minimizing the description length of the weights

venue: In in Proc. of the 6th Ann. ACM Conf. on Computational Learning Theory. Citeseer,

authors:
Pavel Izmailov
Dmitrii Podoprikhin
Timur Garipov
Dmitry Vetrov
Andrew Gordon Wilson

title: Averaging weights leads to wider optima and better generalization

venue: In Proceedings of Conference on Uncertainty in Artificial Intelligence,

authors:
Yiding Jiang
Dilip Krishnan
Hossein Mobahi
Samy Bengio

title: Predicting the generalization gap in deep networks with margin distributions

authors:
Yiding Jiang
Behnam Neyshabur
Hossein Mobahi
Dilip Krishnan
Samy Bengio

title: Fantastic generalization measures and where to find them

authors:
Alex Krizhevsky
Vinod Nair
Geoffrey Hinton

title: Cifar-10 and cifar-100 datasets

venue: URl: https://www. cs. toronto. edu/kriz/cifar. html,

authors:
Namhoon Lee
Thalaiyasingam Ajanthan
Philip Torr

title: SNIP: SINGLE-SHOT NETWORK PRUNING BASED ON CONNECTION SENSITIVITY

authors:
Wesley J. Maddox
Gregory Benton
Andrew Gordon Wilson

title: Rethinking parameter counting in deep models: Effective dimensionality revisited, 2020

title: A pac-bayesian tutorial with a dropout

title: Pac-bayesian model averaging

venue: In Proceedings of the twelfth annual conference on Computational learning theory, pp

authors:
Ari S Morcos
David GT Barrett
Neil C Rabinowitz
Matthew Botvinick

title: On the importance of single directions for generalization

authors:
Preetum Nakkiran
Gal Kaplun
Yamini Bansal
Tristan Yang
Boaz Barak
Ilya Sutskever

title: Deep double descent: Where bigger models and more data hurt

authors:
Behnam Neyshabur
Ryota Tomioka
Nathan Srebro

title: Norm-based capacity control in neural networks. volume

venue: Proceedings of Machine Learning Research, pp. 1376–1401,

authors:
Behnam Neyshabur
Srinadh Bhojanapalli
David McAllester
Nati Srebro

title: Exploring generalization in deep learning

authors:
Alex Renda
Jonathan Frankle
Michael Carbin

title: Comparing rewinding and fine-tuning in neural network pruning

title: Stochastic complexity and modeling

venue: The annals of statistics,

authors:
Binxin Ru
Clare Lyle
Lisa Schut
Mark van der Wilk
Yarin Gal

title: Revisiting the train loss: an efficient performance estimator for neural architecture search

venue: arXiv preprint arXiv:2006.04492,

authors:
Nitish Srivastava
Geoffrey Hinton
Alex Krizhevsky
Ilya Sutskever
Ruslan Salakhutdinov

title: Dropout: a simple way to prevent neural networks from overfitting

venue: The journal of machine learning research,

title: Improving generalization of neural networks through pruning

venue: International Journal of Neural Systems,

authors:
Thomas Verma
Judea Pearl

title: Equivalence and synthesis of causal models

venue: UCLA, Computer Science Department,

authors:
Chaoqi Wang
Guodong Zhang
Roger Grosse

title: Picking winning tickets before training by preserving gradient flow

authors:
Chiyuan Zhang
Samy Bengio
Moritz Hardt
Benjamin Recht
Oriol Vinyals

title: Understanding deep learning requires rethinking generalization

authors:
Wenda Zhou
Victor Veitch
Morgane Austern
Ryan P Adams
Peter Orbanz

title: Non-vacuous generalization bounds at the imagenet scale: a pac-bayesian compression approach

authors:
Michael H. Zhu
Suyog Gupta

title: To prune, or not to prune: Exploring the efficacy of pruning for model compression

venue: In International Conference on Learning Representations (Workshop),

venue: We extend the DEMOGEN data

venue: We train the networks on CIFAR-10 (Krizhevsky et al.,

title: DOUBLE DESCENT For our double descent experiments we use the models and baselines used by Maddox et al. (2020) which are available on: https://github.com/g-benton/hessian-eff-dim

venue: In particular,

heading: 1 INTRODUCTION

text:
The gap between learning-theoretic generalization bounds for highly overparameterized neural networks and their empirical generalization performance remains a fundamental mystery to the field (Zhang et al., 2016; Jiang et al., 2020; Allen-Zhu et al., 2019). While these models are already being successfully used in many applications, improving our understanding of how neural networks perform on unseen data is crucial for safety-critical use cases. By understanding which factors drive generalization in neural networks we may further be able to develop more efficient and performant network architectures and training methods.
Numerous theoretically and empirically motivated attempts have been made to identify generalization measures, that is, properties of the trained model, training procedure and training data that distinguish models that generalize well from those that do not (Jiang et al., 2020). A number of generalization measures have attempted to quantify Occam’s razor, i.e. the principle that simpler models generalize better than complex ones (Neyshabur et al., 2015; Bartlett et al., 2017). This has proven to be non-trivial, as many measures, particularly norm-based measures, grow with the size of the model and thus incorrectly predict that larger networks generalize worse than smaller networks. Other approaches have tried to establish a connection between model compression and generalization (Arora et al., 2018; Zhou et al., 2019). While both of these approaches are theoretically elegant and yield tighter bounds than bounds that are based on the size of uncompressed networks, they nonetheless grow with the size of the original network. Recent empirical studies (Jiang et al., 2020; 2019), on the other hand, identify three classes of generalization measures that do seem predictive of generalization: measures that estimate the flatness of local minima, the speed of the optimization, and the margin of training samples to decision boundaries. While these measures are correlated with generalization, their failure to fully explain the test performance of the model demonstrate a need for other notions of model simplicity to explain generalization in neural networks.
In this paper, we leverage the empirical observation that large fractions of trained neural networks’ parameters can be pruned – that is, set to 0 – without hurting the models’ performance (Gale et al., 2019; Zhu & Gupta, 2018; Han et al., 2015). Based on this insight, we introduce a new measure of
a model’s simplicity which we call prunability: the smallest fraction of weights we can keep while pruning a network without hurting its training loss.
In a range of empirical studies, we demonstrate that a model’s prunability is indeed highly informative of its generalization. In particular, we find that the larger the fraction of parameters that can be pruned without hurting a model’s training loss, the better the model will generalize. Overall, we show that the smaller the fraction of parameters a model actually “uses” - the simpler it is - the better the network’s generalization performance.
In summary, our main contributions are thus the following:
1. We introduce a new generalization measure called prunability that captures a model’s simplicity (Sec. 4) and show that across a large set of models this measure is highly informative of a network’s generalization performance (Sec. 5.4.1).
2. We show that even in a particularly challenging setting in which we observe a test loss double descent (Nakkiran et al., 2020; He et al., 2016), prunability is informative of models’ test performance and is competitive with an existing strong generalization measures (Sec. 5.4.2).
3. Lastly, we investigate whether the success of prunability can be explained by its relationship to flat local minima. We find that while prunability makes similar predictions to some existing measures that estimate the flatness of minima, it differs from them in important ways, in particular exhibiting a stronger correlation with generalization performance than these flatness-based measures (Sec. 5.4.1 and 5.4.3).

heading: 2 RELATED WORK

text:
Zhang et al. (2016) demonstrate that neural networks can perfectly memorize randomly labeled data while still attaining good test set performance on correctly labeled data, a phenomenon that complexity measures such as the VC-Dimension and Rademacher complexity fail to explain. More recently, increasing the parameter count of models has been shown to improve their generalization performance, yielding ‘double descent’ curves in which the test loss first decreases, then increases and then decreases again as the parameter count is increased (He et al., 2016; Nakkiran et al., 2020; Belkin et al., 2019). In contrast, many existing generalization measures grow monotonically with the model’s size and thus fail to capture double descent, making it a particularly interesting setting to study new generalization measures (Maddox et al., 2020).
While generalization measures have been studied for a long time, Jiang et al. (2020) recently brought new empirical rigor to the field. They perform a large-scale empirical study by generating a large set of trained neural networks with a wide range of generalization gaps. Proposing a range of new evaluation criteria, they test how informative of generalization previously proposed generalization measures actually are. In this paper, we evaluate our new generalization measure in the same framework and use the strongest measures as baselines for comparison against our proposed measure, the neural network’s prunability.
A common theme across many generalization measures is that they try to formalize Occam’s Razor, the idea that simpler models generalize better than more complex models. Rissanen (1986) formalizes this principle as a model’s Minimum Description length which was later applied to small neural networks by Hinton & Van Camp (1993). Similarly, Neyshabur et al. (2015; 2017) and numerous other approaches suggest that networks with smaller parameter norms generalize better. We are aware of two existing generalization measures that are based on compression or pruning. Arora et al. (2018) derive a PAC-Bayesian generalization bound based on the size of a network following a particular compression method based on noise stability. Zhou et al. (2019), on the other hand, formulate a generalization bound that is non-vacuous even for large network architectures based on the description length of a compressed network which in turn also grows with the size of the original network. While these approaches are theoretically elegant, all such measures grow with the number of parameters in a given network and thus do not capture the phenomenon that more highly overparameterized networks tend to generalize better.
Three classes of generalization measures seem to be particularly predictive of generalization. First, measures that estimate the flatness of local minima, either based on random perturbations of the
weights (Neyshabur et al., 2017; Dziugaite & Roy, 2017; Jiang et al., 2020; McAllester, 1999) or directly based on the curvature of the loss surface (Maddox et al., 2020). In a purely empirical study, Morcos et al. (2018) find that a model’s robustness to clamping the activations of individual units to fixed values is predictive of its generalization performance. Second, margin-based measures that use the distance of training samples to decision boundaries (Jiang et al., 2019). Lastly, Jiang et al. (2020); Ru et al. (2020) have proposed measures that estimate the speed of optimization.
A large fraction of a network’s parameters can be removed without impairing performance (Gale et al., 2019; Blalock et al., 2020); in fact, such pruning often improves generalization performance (Thodberg, 1991). There are many approaches to selecting parameters for removal; we focus on randomly selected parameters (random pruning) and removing the smallest-magnitude parameters (magnitude pruning), though other approaches exist in the literature (Han et al., 2015; Renda et al., 2020; Lee et al., 2019; Wang et al., 2020; Blalock et al., 2020). Random pruning is applied by the regularization method dropout (Srivastava et al., 2014), which is used to improve network generalization performance. Magnitude pruning can be used in an iterative pruning and retraining procedure (Frankle & Carbin, 2019) to identify sparse subnetworks that, for some sparsity levels, generalize better than their parent network. Both approaches are therefore natural candidates for a notion of network simplicity, which we will leverage in our proposed generalization measure. While research on dropout and the Lottery Ticket Hypothesis suggest that pruning during training and generalization are connected, we study the distinct question of whether pruning after training and generalization are connected.

heading: 3 NOTATION

text:
Following the notation of Jiang et al. (2020), we denote by D the distribution over inputs X and their labels y and let κ denote the number of classes. By S we denote a given data set containing m i.i.d. tuples {(X1, y1), . . . , (Xm, ym)} drawn from D where Xi ∈ X is the input data and yi ∈ {1, . . . , κ} the class labels. We denote a feedforward neural network by fw : X → Z, its weight parameters by w and the number of its parameters by ω. We denote the weight tensor of the ith layer of the network by Wi and we then have that w = vec(W1, . . . ,Wd), where d is the depth of the network and vec denotes the vectorization operator.
Let L be the 1-0 classification error under the data distribution D : L(fw) := E(X,y)∼D[fw(X) 6= y], and let L̂ be the empirical 1-0 error on S: L̂(fw) := 1m ∑m i=1[fw(Xi) 6= yi]. We define the generalization gap as g(fw) := L(fw) − L̂(fw). When training a network, we denote each hyperparameter by θi ∈ Θi for i = 1, . . . n where n is the total number of hyperparameters. We denote a particular hyperparameter configuration by θ := (θ1, θ2, . . . , θn) ∈ Θ, where Θ := Θ1 × Θ2 × · · · ×Θn.

heading: 4 PRUNABILITY

text:
A generalization measure should satisfy three straightforward desiderata: 1) It should correlate with the generalization gap across a broad range of models. 2) It should not suffer trivially from Goodhart’s law, that is, incorporating it into the optimization procedure of a model should lead to improved generalization. 3) It should ideally have a plausible theoretical justification to allow for further study of its properties.
A measure based on magnitude- and random pruning presents a promising direction for satisfying all three desiderata. First, it is well-established that we can remove large fractions of a neural network’s parameters without negatively impacting its train performance, while also improving its test performance (Blalock et al., 2020; Thodberg, 1991). Second, training models using dropout, i.e. random pruning, is widely known to improve generalization (Srivastava et al., 2014), as does the iterative magnitude pruning procedure of Frankle & Carbin (2019) (up to certain sparsity levels). Finally, pruning has a theoretical grounding in the minimum description length principle and in PAC-Bayesian bounds.
The connection between dropout and generalization has been studied in the PAC-Bayesian framework (McAllester, 1999), which allows us to quantify the worst-case generalization performance of random predictors. PAC-Bayesian bounds define model complexity as the KL divergence be-
tween some fixed prior distribution P , and the probability distribution over functions induced by the random predictor (often called the posterior Q, although this may not correspond to a Bayesian posterior) to obtain generalization bounds of the following form. Theorem 1 (McAllester (1999)). For any δ > 0, data distributionD, prior P , with probability 1−δ over the training set, for any posterior Q the following bound holds:
Ev∼Q[L(fv)] ≤ Ev∼Q[L̂(fv)] +
√ KL(Q||P ) + log(mδ )
2(m− 1)
McAllester (2013) shows that in the case where both the prior and posterior are distributions over networks to which dropout with a given dropout probability α is applied, the KL term becomes: KL(Q||P ) = 1−α2 ||w||
2. We can directly use this formalization to derive a generalization bound for networks that are randomly pruned to the largest α, s.t. Ev∼Q[L̂(fv)] < L̂(fw)× (1 + β), i.e., such that the training loss does not increase more than a relative tolerance β. We have to adjust the bound to account for the search for α. We make use of the fact that we search over a fixed number c of possible values for α and use union bound which will change the log term in the bound to log( cmδ ). We thus obtain the following bound:
Ev∼Q[L(fv)] ≤ Ev∼Q[L̂(fv)] +
√ 1−α
2 ||w|| 2 2 + log( m δ ) + 5
2(m− 1) . (1)
While this bound will be vacuous for many large neural networks, it does provide us with some intuition for why we should expect a model’s generalization to be linked to its robustness to pruning. In particular, this bound predicts the relative generalization performance of models from which different fractions of parameters can be removed without adversely affecting their training loss. It is this prediction that we evaluate with the experiments described below. We provide the details on the derivation in Appendix C.
Given a trained model and a pruning algorithm, we define the model’s prunability as the smallest fraction of the model’s parameters that we can keep while pruning without increasing the model’s training loss by more than a given relative tolerance β. We choose the fraction of parameters pruned rather than the number of parameters remaining based on the empirical observation that naive parameter counting is a poor predictor of generalization performance.
Definition 1 (Prunability). Let L̂ denote the cross-entropy loss. For a given model fw, a training data set Strain, a relative bound on the change in train loss β and a pruning method φ, we thus define prunability as
µprunability(fw) = arg min α φ(fw, α)
s.t. L̂(φ(fw, α),Strain) ≤ (1 + β)× L̂(fw,Strain) , (2)
where φ(fw, α) sets the fraction 1− α of the weights of fw to zero.
In our experiments we evaluate robustness to pruning both with regards to magnitude pruning, in which the parameters with the smallest absolute magnitude are set to 0, and random pruning, where random parameters are set to 0. Magnitude pruning can remove a greater number of model weights without affecting performance than random pruning (Blalock et al., 2020), and can therefore be considered a more effective pruning strategy. We hypothesize that using a more effective pruning method might yield a measure of prunability that is more predictive of generalization. Per default, we refer to a model’s prunability with regards to magnitude pruning in the following. Note that while iterative pruning is usually used when the aim is to achieve the best compression rate, we are only concerned with the generalization of the model at the end of training and thus apply one-shot pruning. In our experiments we use a tolerance β of 0.1 and note that the particular choice of β did not seem to have a large impact on the predictiveness of prunability. See Appendix A.6.1 for pseudo-code of how prunability is computed for a given model.

heading: 5 EXPERIMENTAL RESULTS AND ANALYSIS

text:
The goals of our experiments are fourfold: 1) We verify that prunability predicts generalization across a set of models trained with different hyperparameters and generalization performance. To
this end, we use a similar experimental framework as Jiang et al. (2020). 2) We evaluate how prunability performs in a particularly challenging double descent setting. 3) To understand what notion of simplicity prunability quantifies, we study the conditional mutual information between it and baselines that represent existing classes of different strong generalization measures. 4) Given that prunability and random perturbation robustness both measure the sensitivity of a model’s training loss to a perturbation of its weights, we directly compare how pruning and random weight perturbations impact models’ losses.

heading: 5.1 BASELINES

text:
We compare our generalization measure to a number of baseline measures. Having somewhat tighter computational constraints than Jiang et al. (2020), we select one strong baseline per class of generalization measures as described by Jiang et al. (2020) and Jiang et al. (2019).
From the class of generalization measures that aim to capture a model’s simplicity using parameternorms, we use the Sum of Squared Frobenius Norms (Jiang et al., 2020):
µsum of fro(fw) = d · ( d∏ i=1 ||Wi||2F )1/d,
where d is the number of layers of the model. From the class of measures based on the idea that models with flatter local minima generalize better, we use magnitude aware random perturbation robustness (Jiang et al., 2020):
µpac bayes mag flat(fw) = 1
σ2 ,
where ui ∼ N (0, σ2|wi| + ) and σ2 is the largest number such that Eu[L̂(fw+u)] ≤ (1 + β) × L̂(fw), where β is an arbitrarily selected relative tolerance, in our experiments 0.1. From the class of measures based on the idea that models with wider margins between the training samples and decision boundaries generalize better, we use the Normalized Margins measure (Jiang et al., 2019):
µnorm margins(fw) = a Tφ,
wherein φ consists of five statistics of the normalized margins distribution for each layer of the model. That is, at each hidden layer of the network, the normalized distances of all samples to the decision boundary of the second most likely class are computed. Five statistics of each of these distance distributions are used as input for the linear regression. Lastly, from the class of measures based on the idea that the speed of optimization is predictive of generalization, we use the negative Sum of Train Losses (Ru et al., 2020) throughout training:
µSoTL(fw) = − T∑ t=1 L̂(fwt),
where T is the number of training epochs, wt the parameters at epoch t and L̂ the training crossentropy loss. In Appendix A.2 we provide further details on these generalization measures and additional baselines.

heading: 5.2 MODELS

text: To generate a large set of trained models, we build upon the DEMOGEN set of trained models introduced by Jiang et al. (2019). Our data set consists of 324 models that resemble the Networkin-Network architecture by (Lin et al., 2014) trained on CIFAR10. Jiang et al. (2019) propose this particular architecture because it is parameter-efficient and achieves relatively competitive performance on standard image classification benchmarks. We achieve a wide range of generalization gaps by varying a number of hyperparameters such as the width, depth or L2-regularization of the models. We modify the original set of models, DEMOGEN, used by Jiang et al. (2019) to obtain an even stronger experimental set up: noticing that the depth of neural networks has a large impact on generalization performance, we add models of different depths to the original DEMOGEN data set. Following Jiang et al. (2020) we only use models that apply batch normalization. See Appendix A.3 for further details of our experimental set-up.

heading: 5.3 EVALUATION METRICS

text:
To study our generalization measures comprehensively, we consider 1) their Kendall’s Rank correlation with the generalization gap, 2) how predictive they are in terms of the adjusted R2 and 3) apply the conditional independence test of Jiang et al. (2020), in an attempt to understand the causal relations between the complexity measures and generalization. We briefly summarize the metrics in this section and provide a more detailed explanation in Appendix A.1
We compute Kendall’s rank correlation coefficient which quantifies to what extent the ranking of models according to a given generalization measure corresponds to the ranking according to the generalization gaps. For a given generalization measure µ, we consider: T := ∪θ∈Θ{(µ(θ), g(θ))} where g(θ) is the generalization gap of the model trained with hyperparameters θ. Kendall’s rank correlation coefficient is then defined as the fraction of pairs of tuples that are correctly ranked according to the generalization measure: τ(T ) :=
1 |T |(|T |−1) ∑ (µ1,g1)∈T ∑ (µ2,g2)∈T \(µ1,g1) sign(µ1 − µ2)sign(g1 − g2). While Kendall’s rank correlation coefficient is generally a useful metric, it does not adequately reflect whether a generalization measure’s performance is consistent across all hyperparameter axes. Thus, we also consider the granulated Kendall’s rank correlation coefficient which is essentially the average of Kendall’s rank coefficients ψi obtained by only varying one hyperparameter dimension at a time: Ψ := 1n ∑n i=1 ψi
To go beyond correlation, Jiang et al. (2020) use a conditional independence test inspired by Verma & Pearl (1991). The main goal here is to understand whether there exists an edge in a causal graph between a given generalization measure µ and the generalization gap g: this tells us whether models generalize well because of the value of the generalization measure, or whether the two values are only correlated. This is achieved by estimating the conditional mutual information between the generalization measure and the generalization gap conditioned on different hyperparameters that are observed. Formally, we have that for any function φ : Θ → R, let Vφ : Θ1 × Θ2 → {+1,−1} be as Vφ(θ1, θ2) := sign(φ(θ1) − φ(θ2)). Let US be a random variable that corresponds to the values of the hyperparameters in S. Intuitively, the higher this metric is, the more likely it is that there is indeed an edge in the causal graph between the generalization gap and the generalization measure: K(µ) = minUSs.t.|S|≤2 Î(Vµ, Vg|US), where Î(Vµ, Vg|US) is the normalized conditional mutual information between the generalization measure and the generalization gap conditioned on observing the hyperparameter set S. To evaluate how well we can predict a model’s generalization using a linear function of a given generalization measure and following Jiang et al. (2019), we additionally evaluate the generalization measures’ predictiveness in terms of their adjusted R2, i.e. the proportion of the variance of the generalization gap g(θ) across hyperparameters θ ∈ Θ that can be explained by the generalization measure µ(θ).

heading: 5.4 RESULTS

heading: 5.4.1 PERFORMANCE OF PRUNABILITY ON EXTENDED DEMOGEN DATA SET

text:
Table 1 holds a summary of our experimental results on the extended DEMOGEN data set. We find that prunability is highly informative of generalization across all of our evaluation metrics. In particular, it outperforms random perturbation robustness, the training loss itself and the Frobenius norm measures in terms of Kendall’s rank correlation coefficient.
We find that while prunability based on magnitude pruning outperforms prunability based on random pruning, the difference in performance is relatively small.
We find that the Normalized Margins baseline clearly outperforms prunability while the Negative Sum of Train Losses and the Best Margin Variable seem similarly informative of generalization as prunability. Note that and the margin-based measures achieve much higher conditional mutual information scores than the other baselines, providing further evidence for the causal connection between wide margins and good generalization performance. The conditional mutual information between prunability and generalization is relatively low, suggesting that it seems somewhat less likely that there is a direct causal connection between a model’s prunability and its generalization performance.

heading: Generalization Measure Kendall’s τ Adjusted R2 Cond. Mutual Information

text: In Table 2, we compare the granulated and regular Kendall’s coefficients of some additional baselines. The results reconfirm the weak performance of norm and parameter count-based complexity measures both for pruned and unpruned networks. While this a well known issue, we point out that this is also problematic for recently proposed generalization measures based on the size of the compressed network such as Zhou et al. (2019). We also note that the Frobenius norm measure is strongly negatively correlated with generalization rather than positively correlated as one would expect based on Jiang et al. (2020). It seems that this result is largely due to a strong negative correlation for models with different dropout rates and relatively weak correlation coefficients along other hyperparameter axes.

heading: 5.4.2 DOUBLE DESCENT

text:
In addition to our experiments on the DEMOGEN data set, we also evaluate prunability on a set of models that exhibit a test loss double descent (Nakkiran et al., 2020). This phenomenon is particularly interesting for the study of generalization measures, since many measures that aim to capture a model’s complexity grow monotonically as we increase a network’s width and thus do not display a double descent. We are thus including this experiment to evaluate whether prunability also succumbs to this phenomenon. We use the experimental setup proposed by Maddox et al. (2020), and compare our new measure to their curvature-based generalization measure which also aims to capture double descent. Note that up to now we studied the generalization gap in terms of the 0− 1 loss. In this section, however, we focus on another aspect of generalization and study the relation between prunability and models’ cross-entropy test loss.
Maddox et al. (2020) introduce a generalization measure called Effective Dimensionality which is defined as: Neff(H, z) = ∑k i=1 λi λi+z
where λi are the eigenvalues of the Hessian of the training loss H ∈ Rk×k, and z > 0 is a regularization constant. Intuitively, the Effective Dimensionality measures the flatness of the local minimum of a network in terms of the number of eigenvalues of the Hessian of the training loss that are ”large”. To get a measure on the same scale as our
new generalization measure, we use a normalized version of Neff : µeffective dimensionality = Neffk . Additionally, we also evaluate a number of the strongest baselines from Section 5.4.1 in the double descent setting. We refer to appendix A.4 for further details on the experimental setup.
We find that prunability does indeed display a double descent and is more strongly rank correlated with the test loss than the baseline measures (Kendall’s τ of 0.3613, p < 0.0005).
In Appendix B.1 we study an additional double descent setting and obtain results that are consistent with the results presented in this section.

heading: 5.4.3 PRUNING AND FLATNESS-BASED MEASURES

text:
We now investigate whether a network’s robustness to pruning can be explained by other established generalization measures, or whether it is measuring something entirely distinct from previous notions of model simplicity. To better understand the relation between prunability and the other generalization measures we perform the conditional independence test described in section A.1.3 between them. Recall that, up to now, we used this test to study the relation between the generalization measure and the generalization gap and that the higher the κ between two variables, the likelier it is that they neighbors in the causal DAG. See Table 3 for a summary of the results.
First, we observe that the conditional mutual information between prunability and random perturbation robustness is the highest across the baselines we are evaluating. Because pruning is a special case of perturbation, some connection between the two methods is expected. However, prunability is clearly measuring something subtly different from flatness of the loss surface, as it outperforms random perturbation robustness in our evaluations. Thus, we investigate whether this distinction may arise from the differential effects of pruning and random perturbations in function space.
To evaluate this, we directly compare the impact of both pruning and random perturbations on some of the models in our double descent experimental setting by performing pruning and random perturbations of the same magnitude and computing their respective impact on the loss. We find that pruning has a larger negative impact on a model’s training loss than a random perturbation of the same magnitude (see Figure 2). Pruning can lead to an improvement of the test loss which we do not observe for random perturbations; we conjecture that this relationship may be responsible for the difference between the two methods, though disentangling this phenomenon is left for future work. We also note that prunability achieves much higher Kendall’s τ , AdjustedR2 and Conditional Mutual Information than random perturbation robustness.
In conclusion, pruning affects models differently from randomly perturbing their weights. Because pruning can in certain settings improve test performance (and hence seems to be more aligned with network simplicity) we suggest that this difference leads to the stronger performance of prunability as compared to random perturbation robustness. We compare the impact of pruning and random perturbations on a model’s training and test loss for an additional architecture in Appendix B.2, and find that the results are consistent with the results presented in this section.

heading: Pruning versus Random Weight Perturbations

heading: 6 CONCLUSION

text: In this paper, we show that a network’s prunability, the smallest fraction of parameters we can keep while pruning without hurting the model’s training loss, is highly informative of the network’s generalization performance. We also show that prunability seems to capture something distinct from previously suggested measures based on robustness to random perturbations or the curvature of the loss landscape. We see a number of promising directions for further research. For instance, one could study whether the direction of pruning lies in the span of the main eigenvectors of the Hessian studied by Maddox et al. (2020). Given the strong empirical performance of prunability – and given that the theory behind its success is not yet well understood – this measure of model complexity may be of use to the construction of future generalization bounds. We include an attempt of this in Appendix C, but more thorough investigation will likely prove illuminating.

heading: A EXPERIMENTAL SETUP

heading: A.1 EVALUATION METRICS

text: In this section we describe the metrics we use to evaluate the generalization measures. To study our generalization measures comprehensively, we consider the generalization measures’ Kendall’s Rank correlation with the generalization gap, how predictive they are in terms of Adjusted R2 and Jiang et al. (2020).’s conditional independence test in an attempt to understand the causal relations between the complexity measures and generalization.

heading: A.1.1 KENDALL’S RANK CORRELATION COEFFICIENT

text:
This criterion evaluates the quality of a given generalization in terms of how well it ranks models with different generalization gaps. Intuitively, a generalization measure ranks two models that were trained with hyperparameters θ1 and θ2 correctly iff µ(θ1) > µ(θ2)→ g(θ1) > g(θ2) is satisified. Kendall’s rank coefficient then measures for what fraction of models in our data set this holds. Kendall T takes values between 1 (perfect agreement) and -1 (rankings are exactly reversed) where a Kendall’s t of 0 means that the complexity measure and generalization gap are independent.
Formally, we consider the set of tuples of generalization measures and generalization gaps:
T := ∪θ∈Θ{(µ(θ), g(θ))}
And Kendall’s rank correlation coefficient is defined as:
τ(T ) := 1 |T |(|T | − 1) ∑ (µ1,g1)∈T ∑ (µ2,g2)∈T \(µ1,g1) sign(µ1 − µ2)sign(g1 − g2)

heading: A.1.2 GRANULATED KENDALL’S RANK COEFFICIENT

text:
While Kendall’s rank correlation coefficient is generally a useful metric, it does not adequately reflect whether a generalization measure’s performance is consistent across all hyperparameter axes. To account for this Jiang et al. (2020) propose the granulated Kendall’s rank-correlation coefficient which is essentially the average of Kendall’s rank coefficients obtained by only varying one hyperparameter dimension at a time.
Formally, we have:
mi := |Θ1 × · · · ×Θi−1 ×Θi+1 × ...Θn|
and the Kendall’s rank correlation coefficient is defined as:
ψi := 1
mi ∑ θ1∈Θ1 · · · ∑ θi−1∈Θi−1 ∑ θi+1∈Θi+1 · · · ∑ θn∈Θn τ(∪θi∈Θi(µ(θ), g(θ)))
Averaging all of these values, we obtain the granulated Kendall’s rank correlation coefficient:
Ψ := 1
n n∑ i=1 ψi
Given these definitions, we obtain the following standard errors for the granulated Kendall’s coefficients in our experiments. For the width, dropout rate, weight decay and depth measurements we hsve, sτ = 0.00056 and for the data augmentation measurements we have sτ0.000484. Lastly, for ψ, we have sτ = 0.00011

heading: A.1.3 CONDITIONAL MUTUAL INFORMATION

text:
To go beyond measures of correlation, Jiang et al. (2020) use a conditional independence test inspired by Verma & Pearl (1991). The main goal here is to understand whether there exists an edge in a causal graph between a given generalization measure µ and the generalization gap g. This is achieved by estimating the mutual information between the generalization measure and the generalization gap conditioned on different hyperparameters being observed. The higher this metric is, the more likely it is that there is indeed an edge in the causal DAG between the generalization gap and the generalization measure.
Formally, we have that for any function φ : Θ → R, let Vφ : Θ1 × Θ2 → {+1,−1} be as Vφ(θ1, θ2) := sign(φ(θ1)− φ(θ2)). Let US be a random variable that corresponds to the values of the hyperparameters in S. The conditional mutual information can then be computed as:
I(Vµ, Vg|US) = ∑ US p(US) ∑ Vµ∈{±1} ∑ Vg∈{±1} p(Vµ, Vg|US) log ( p(Vµ, Vg|US) p(Vµ|US)p(Vg|US) )
By dividing the term above by the conditional entropy of generalization: H(Vg|US) = − ∑ US p(US) ∑ Vµ∈{±1} p(Vg|US) log(p(Vg|US))
we obtain the normalized criterion:
Î(Vµ, Vg|US) = I(Vµ, Vg|US) H(Vg|US)
In the original Inductive Causation algorithm (Verma & Pearl, 1991), an edge in the causal graph is kept if there is no subset of hyperparameters S such that the two two nodes are independent, that is Î(Vµ, Vg|US) = 0. Due to computational limitations, and following Jiang et al. (2020), we only consider subsets of hyperparameters of size at most 2:
K(µ) = min USs.t.|S|≤2 Î(Vµ, Vg|US)
Intuitively, the higher K(µ), the more likely it is that there exists an edge between µ and generalization in the causal graph.

heading: A.1.4 ADJUSTED R2

text:
Lastly, we consider a metric to evaluate how predictive a generalization measure is of generalization based on a simple linear regression: ĝ = aTµ+ b. We find a and b by minimizing the mean squared error: a∗, b∗ = arg mina,b ∑ i(a Tφ(µi) + b− gi)2
On a test set of models we then compute the coefficient of determination R2 which measures what fraction of the variance of the data ca be explained by the linear model:
R2 = 1− ∑n j=1(ĝj − gj)2∑n
j=1(gj − 1 n ∑n j=1 gj) 2
We compute the R2 using a 10-fold cross-validation.
Alternatively, one can use the adjusted R2 score (R̄2) which measures how well a linear model fits a training set and which accounts for the number of variables used in the regression:
R̄2 = 1− (1−R2) n− 1 n− dim(µ)− 1

heading: A.2 GENERALIZATION MEASURES

text:
Additional generalization measures we consider are:
• µbest margins variable: To account for the fact that the normalized margins variable is based on a large number of input variables, we select the single independent variable that has the largest standardized coefficient in the Normalized Margins model and directly use it as a generalization measure.
• µtrain loss: L̂(fw), the cross entropy loss of the model fw on the training data set after training is finished.
• µparameter count: The number of parameters ω of the model. • µsum two norms:
µsum of two norms(fw) = ( d∑ i=1 ||Wi||22)
• µsum two norms pruned: Essentially, the same as µsum two norms, but applied to fw that we obtain when we compute µprunability . Formally, we have:
µsum two norms pruned(fw) = µsum two norms(min α φ(fw, α))
s.t. L̂(φ(fw, α),Strain) ≤ (1 + β)× L̂(fw,Strain)
• Analogously, µpruned parameter count is the parameter count ω of the pruned model we obtain while computing µprunability

heading: A.3 DEMOGEN

text:
We extend the DEMOGEN data set in Jiang et al. (2019) with regards to two aspects. First, we follow Jiang et al. (2020) and only consider models that are trained using batch normalization. Second, Jiang et al. (2019) do not consider networks of different depths. We consider this to be a crucial hyperparameter and a extend the original DEMOGEN data set by training additional models with different numbers of NiN-blocks than the models in the original data set.
We use the network architecture proposed by Jiang et al. (2019) in the DEMOGEN data set which is very similar to the architecture used in Jiang et al. (2020). The architecture is very similar to the Network in Network architecture in Lin et al. (2014) with the max pooling and dropout layer removed. We train the networks on CIFAR-10 (Krizhevsky et al., 2009). The architecture we use is thus:
1. The input layer 2. 2, 3 or 4 NiN-blocks as described in Table 4. 3. A single convolution layer with kernel size 4x4 and stride 1.
In this architecture we vary the following hyperparameters:
1. Use channel sizes of 192, 288, and 384 2. Apply dropout after NiN-blocks with p = 0.0, 0.2, 0.5 3. Apply l2 regularization with λ = 0.0, 0.001, 0.005. 4. Train with and without data augumentation, that is random cropping, flipping and shifting.
5. We consider two random initializations per hyperparameter configuration. 6. Models with 2, 3 or 4 NiN blocks as described above.
In total we are thus working with a data set of 324 trained models. We train the models with SGD with momentum (α = 0.9) with a batch size of 128 and an initial learning rate of 0.01. The networks are trained for 380 epochs with a learning rate decay of 10x every 100 epochs. We base our implementation on the original code base: https://github.com/google-research/ google-research/tree/master/demogen.

heading: A.4 DOUBLE DESCENT

text:
For our double descent experiments we use the models and baselines used by Maddox et al. (2020) which are available on: https://github.com/g-benton/hessian-eff-dim. In particular, the CNN we use is the architecture used by Nakkiran et al. (2020) which is availbable on: https://gitlab.com/harvard-machine-learning/double-descent/-/ blob/master/models/mcnn.py Likewise, the ResNet18 architecture can be found here: https://github.com/g-benton/hessian-eff-dim/blob/temp/hess/nets/ resnet.py.
The other baselines correspond to the ones from the DEMOGEN experiments described in Appendix A.3. The margins baseline used in this setting is based on the one used in Jiang et al. (2019): a linear model is trained to predict the test loss based on the statistics of the margin distribution at the output layer of a given model.
We train a set of 32 networks of widths 2, 4, 6, ..., 64 on CIFAR100 (Krizhevsky et al., 2009) using SGD with a learning rate of 10−2, momentum of 0.9, weight decay of 10−4 for 200 epochs with a batch size of 128. The learning rate decays to 10−4 on a piecewise constant learning rate schedule (Izmailov et al., 2018), beginning to decay on epoch 100. Random cropping and flipping are used for data augmentation which is turned off for the computation of eigenvalues. Our plots contain the results across 5 random seeds. The shaded areas indicate the standard error. All experiments were run on Nvidia GeForce RTX 2080 Ti GPUs.

heading: A.5 RANDOM PERTURBATIONS VERSUS PRUNING

text: We use the ResNet18 implementation from https://gitlab.com/ harvard-machine-learning/double-descent/-/blob/master/models/ resnet18k.py and use the same hyperparameters for training as in the double descent settings above. For every magnitude pruned model we also compute the randomly perturbed model in which we perturb the same set of weights by a random vector of the same size.

heading: A.6 ALGORITHMS

heading: A.6.1 ALGORITHM TO DETERMINE PRUNABILITY

text:
Algorithm 1: Maximal Lossless Magnitude Pruning Input: Original model fw, train data, tolerance in change of loss β, step size for search step size Result: Pruned Model fwpruned
1 original train loss = loss(h, train data); 2 fraction of weights to remove = 0.98; 3 wpruned = prune(w, fraction of weights to remove); 4 while loss(fwpruned , train data) > (1 + β)× original train loss do 5 wpruned = prune(w, fraction of weights to remove); 6 fraction of weights to remove -= step size; 7 end 8 return fwpruned
Where prune(w, sparisty) is an algorithm that sets sparsity % of w to 0 and then returns the new vector. In our experiments we evaluate both magnitude pruning, where the parameters with the
smallest absolute magnitude are set to 0, and random pruning, where random parameters are set to 0.
Empirically, we find that the particular choice for the tolerance β does not make a big difference with regards to predictiveness of prunability. In our experiments we use β = 0.1.
Note that to determine a model’s prunability we use the differentiable cross-entropy loss rather than the 1-0 loss used elsewhere in the paper.

heading: B ADDITIONAL EXPERIMENTS

heading: B.1 DOUBLE DESCENT

text: In Section 5.4.2, we conduct an experiment that suggests that prunability captures the particularly challenging double descent phenomenon. In this section, we provide an extended version of Figure 1, containing all baselines (see Figure 3a). Furthermore, we extend that experiment to an additional architecture, ResNet18 trained on CIFAR100 (see Figure 4). The results of this additional experiment are consistent with the one in the main part of the body: prunability is strongly Kendall rank correlated with the test loss and the test error. In this additional setting, however, the margins-based method outperforms prunability. We describe the details of the architecture and training of this experiment in Appendix A.5.

heading: B.2 PRUNING VERSUS RANDOM PERTURBATIONS

text: Likewise, we study the impact of pruning and random perturbations on the training and test loss in an additional experimental setting. As in the experiment in Section 5.4.3, we find that pruning and randomly perturbing weights (by a vector of the same magnitude), have very different impacts on the training and test loss (see Figure 5). In this experiment, we study the CNN used in our double descent experiment of Section 5.4.2. In particular, we find that pruning generally has a larger negative impact on the training loss than a random weight perturbation of the same magnitude. With regards to the test loss, we find that pruning has a smaller negative impact than random perturbations for moderate amounts of pruning, and a larger negative impact for larger amounts of pruning/perturbation. This matches the behavior of the first experiment. In this additional experiment we do not observe an improvement of the test loss through pruning which was not expected.

heading: C PAC-BAYES BOUND

text:
We use the following, simplified formulation of McAllester (1999)’s dropout bounds which is also used in Jiang et al. (2020):
Theorem 1 For any δ > 0, distribution D, prior P, with probability 1 − δ over the training set, for any posterior Q the following bound holds:
Ev∼Q[L(fv)] ≤ Ev∼Q[L̂(fv)] +
√ KL(Q||P ) + log(mδ )
2(m− 1)
For the reader’s convenience we briefly reproduce the derivation of the dropout bound in McAllester (2013).
To draw a sample ω from Q means first drawing a sparsity pattern S ∼ Sα, where α is the dropout probability, and a noise vector ∼ N (0, 1)d. Then, we construct ω as s ◦ (w + ), ◦ being the Hadamard product: (s◦w)i := siwi. Similarly, a sample from P is s◦ . We denote by Sα the distribution on the d-cube B of sparsity patterns. We then have Ev∼Q[fv] = Es∼Sα, ∼N (0,1)d [fs◦(w+ )]. We derive KL(Q||P ) as follows:
KL(Q||P ) = Es∼Sα, ∼N (0,1)d [ ln Sα(s)e
− 12 ||s◦ || 2
Sα(s)e− 1 2 ||s◦(w+ )||2
]
= Es∼Sα [ 1
2 ||s ◦w||2]
= 1− α
2 ||w||2
We can directly obtain a bound for prunability from this, by searching for the largest α, s.t.Ev∼Q[L̂(fv)] < L̂(fw) × (1 + β), where w are the weights learned during training of the network. We have to account for the fact that we search for α in our bound. We make use of the fact that we’re searching over a fixed number c of αs and use union bound in the bound which will change the log term in the bound to log( cmδ ). Gives us the following bound:
Ev∼Q[L(fv)] ≤ Ev∼Q[L̂(fv)] +
√ 1−α
2 ||w|| 2 2 + log( m δ ) + 5
2(m− 1)


summaries: This paper tackles the problem of long horizon visual planning, with the aim of of being able to plan actions to reach distant goals. This is a well studied problem, and like prior work this method considers the setting where the agent is given an offline dataset of interaction, which it learns from to be able to reach new goals (specified by a goal image). The method first does unsupervised representation learning, where it learns a discrete representation of images (using contrastive predictive coding (CPC) with a discrete latent variable using Gumbel Softmax). Using a set of discrete latent variable, it then builds a set of graphs which connects these discrete states based on the collected experience, and derives a planning procedure to find a path in the graph which reaches the goal.

summaries: This paper presents a method that combines learning discrete representations together with planning using graph search to solve long horizon tasks from vision. The approach works by generating data via random exploration and trains a representation encoder on this data. This network extracts objects from the observations and passes them through a CNN and shared encoder to generate one-hot encodings; these encodings are concatenated to generate the discrete representation. The encoder is trained via a contrastive learning objective with a similarity matrix that encourages nearby states to share similar encodings, thereby encouraging spatial and temporal abstractions. Next, these representations are combined together with an abstract planner to generate a sequence of waypoints to the goal. This is done by creating a graph of transitions from the collected exploratory data and search within this graph — this is executed for each encoding at a time with the assumption that the task can be solved by moving each object independently of the others. Finally, a low-level controller is used for reaching the waypoints and final goal, this is done via MPC on an action-conditional predictive model that generates future observations given the current state and goal. The approach is tested on two simple planar planning tasks where the agent has to solve k-object arrangement and open a room with a key respectively.

summaries: This work presents Discrete Object-factorized Representation Planning (DORP), which learns a discrete representation from videos with an enforced temporal consistency. This representation can then be planned over through a sequence of small alterations to the discrete embedding, which are then executed via MPC. DORB is demonstrated to solve long-horizon tasks and learn representations that consider objects and their properties.

abstractText: Discrete representations have been key in enabling robots to plan at more abstract levels and solve temporally-extended tasks more efficiently for decades. However, they typically require expert specifications. On the other hand, deep reinforcement learning aims to learn to solve tasks end-to-end, but struggles with long-horizon tasks. In this work, we propose Discrete Object-factorized Representation Planning (DORP), which learns temporally-abstracted discrete representations from exploratory video data in an unsupervised fashion via a mutual information maximization objective. DORP plans a sequence of abstract states for a low-level model-predictive controller to follow. In our experiments, we show that DORP robustly solves unseen long-horizon tasks. Interestingly, it discovers independent representations per object and binary properties such as a key-and-door.

authors:
Pulkit Agrawal
Ashvin V Nair
Pieter Abbeel
Jitendra Malik
Sergey Levine

title: Learning to poke by poking: Experiential learning of intuitive physics

venue: In Advances in Neural Information Processing Systems,

authors:
Ankesh Anand
Evan Racah
Sherjil Ozair
Yoshua Bengio
Marc-Alexandre Côté
R Devon Hjelm

title: Unsupervised state representation learning in atari

authors:
Masataro Asai
Alex Fukunaga

title: Classical planning in deep latent space: Bridging the subsymbolic-symbolic boundary

venue: arXiv preprint arXiv:1705.00154,

authors:
Mohammad Babaeizadeh
Chelsea Finn
Dumitru Erhan
Roy H Campbell
Sergey Levine

title: Stochastic variational video prediction

venue: arXiv preprint arXiv:1710.11252,

title: Active perception

venue: Proceedings of the IEEE,

authors:
Ruzena Bajcsy
Yiannis Aloimonos
John K Tsotsos

title: Revisiting active perception

venue: Autonomous Robots,

authors:
Christopher P Burgess
Loic Matthey
Nicholas Watters
Rishabh Kabra
Irina Higgins
Matt Botvinick
Alexander Lerchner. Monet

title: Unsupervised scene decomposition and representation

authors:
Maxime Chevalier-Boisvert
Lucas Willems
Suman Pal

title: Minimalistic gridworld environment for openai gym

authors:
Frederik Ebert
Chelsea Finn
Sudeep Dasari
Annie Xie
Alex Lee
Sergey Levine

title: Visual foresight: Model-based deep reinforcement learning for vision-based robotic control

venue: arXiv preprint arXiv:1812.00568,

authors:
Ben Eysenbach
Russ R Salakhutdinov
Sergey Levine

title: Search on the replay buffer: Bridging planning and reinforcement learning

authors:
Chelsea Finn
Ian Goodfellow
Sergey Levine

title: Unsupervised learning for physical interaction through video prediction

title: First contact: an active vision approach to segmentation

venue: In Proceedings 2003 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS 2003)(Cat. No. 03CH37453),

title: The ecological approach to visual perception: classic edition

authors:
Klaus Greff
Raphaël Lopez Kaufman
Rishabh Kabra
Nick Watters
Chris Burgess
Daniel Zoran
Loic Matthey
Matthew Botvinick
Alexander Lerchner

title: Multi-object representation learning with iterative variational inference

title: Les valeurs extrêmes des distributions statistiques

venue: In Annales de l’institut Henri Poincaré,

authors:
Danijar Hafner
Timothy Lillicrap
Ian Fischer
Ruben Villegas
David Ha
Honglak Lee
James Davidson

title: Learning latent dynamics for planning from pixels

venue: In International Conference on Machine Learning,

authors:
Karol Hausman
Dejan Pangercic
Zoltán-Csaba Márton
Ferenc Bálint-Benczédi
Christian Bersch
Megha Gupta
Gaurav Sukhatme
Michael Beetz

title: Interactive segmentation of textured and textureless objects

venue: In Handling Uncertainty and Networked Structure in Robot Control,

authors:
Kaiming He
Haoqi Fan
Yuxin Wu
Saining Xie
Ross Girshick

title: Momentum contrast for unsupervised visual representation learning

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Eric Jang
Shixiang Gu
Ben Poole

title: Categorical reparameterization with gumbel-softmax

venue: arXiv preprint arXiv:1611.01144,

authors:
Leslie Pack Kaelbling
Tomás Lozano-Pérez

title: Hierarchical task and motion planning in the now

venue: In 2011 IEEE International Conference on Robotics and Automation,

authors:
Ken Kansky
Tom Silver
David A Mély
Mohamed Eldawy
Miguel Lázaro-Gredilla
Xinghua Lou
Nimrod Dorfman
Szymon Sidor
Scott Phoenix
Dileep George

title: Schema networks: Zero-shot transfer with a generative causal model of intuitive physics

venue: arXiv preprint arXiv:1706.04317,

authors:
Jacqueline Kenney
Thomas Buckley
Oliver Brock

title: Interactive segmentation for manipulation in unstructured environments

venue: In 2009 IEEE International Conference on Robotics and Automation,

authors:
Thomas Kipf
Elise van der Pol
Max Welling

title: Contrastive learning of structured world models

venue: arXiv preprint arXiv:1911.12247,

authors:
Thanard Kurutach
Aviv Tamar
Ge Yang
Stuart J Russell
Pieter Abbeel

title: Learning plannable representations with causal infogan

authors:
Suha Kwak
Minsu Cho
Ivan Laptev
Jean Ponce
Cordelia Schmid

title: Unsupervised object discovery and tracking in video collections

venue: In Proceedings of the IEEE international conference on computer vision,

authors:
Michael Laskin
Scott Emmons
Ajay Jain
Thanard Kurutach
Pieter Abbeel
Deepak Pathak

title: Sparse graphical memory for robust planning

venue: arXiv preprint arXiv:2003.06417,

authors:
Kara Liu
Thanard Kurutach
Christine Tung
Pieter Abbeel
Aviv Tamar

title: Hallucinative topological memory for zero-shot visual planning

venue: arXiv preprint arXiv:2002.12336,

authors:
Corey Lynch
Mohi Khansari
Ted Xiao
Vikash Kumar
Jonathan Tompson
Sergey Levine
Pierre Sermanet

title: Learning latent plans from play

venue: In Conference on Robot Learning,

authors:
Xiao Ma
Siwei Chen
David Hsu
Wee Sun Lee

title: Contrastive variational model-based reinforcement learning for complex observations

venue: arXiv preprint arXiv:2008.02430,

authors:
Chris J Maddison
Andriy Mnih
Yee Whye Teh

title: The concrete distribution: A continuous relaxation of discrete random variables

venue: arXiv preprint arXiv:1611.00712,

authors:
Mehdi Mirza
Andrew Jaegle
Jonathan J Hunt
Arthur Guez
Saran Tunyasuvunakool
Alistair Muldal
Théophane Weber
Peter Karkus
Sébastien Racanière
Lars Buesing

title: Physically embedded planning problems: New challenges for reinforcement learning

venue: arXiv preprint arXiv:2009.05524,

authors:
Suraj Nair
Chelsea Finn

title: Hierarchical foresight: Self-supervised learning of long-horizon tasks via visual subgoal generation

venue: arXiv preprint arXiv:1909.05829,

authors:
Soroush Nasiriany
Vitchyr Pong
Steven Lin
Sergey Levine

title: Planning with goal-conditioned policies

authors:
Andrew Y Ng
Daishi Harada
Stuart Russell

title: Policy invariance under reward transformations: Theory and application to reward shaping

venue: In ICML,

authors:
Masashi Okada
Tadahiro Taniguchi

title: Dreaming: Model-based reinforcement learning by latent imagination without reconstruction

venue: arXiv preprint arXiv:2007.14535,

authors:
Aaron van den Oord
Yazhe Li
Oriol Vinyals

title: Representation learning with contrastive predictive coding

venue: arXiv preprint arXiv:1807.03748,

authors:
Karl Pertsch
Oleh Rybkin
Frederik Ebert
Chelsea Finn
Dinesh Jayaraman
Sergey Levine

title: Long-horizon visual planning with goal-conditioned hierarchical predictors

venue: arXiv preprint arXiv:2006.13205,

authors:
Ali Razavi
Aaron van den Oord
Oriol Vinyals

title: Generating diverse high-fidelity images with vq-vae-2

authors:
Sebastian Risi
Kenneth O Stanley

title: Deep neuroevolution of recurrent and discrete world models

venue: In Proceedings of the Genetic and Evolutionary Computation Conference,

authors:
Nikolay Savinov
Alexey Dosovitskiy
Vladlen Koltun

title: Semi-parametric topological memory for navigation

venue: arXiv preprint arXiv:1803.00653,

authors:
Aravind Srinivas
Michael Laskin
Pieter Abbeel

title: Curl: Contrastive unsupervised representations for reinforcement learning

venue: arXiv preprint arXiv:2004.04136,

authors:
Siddharth Srivastava
Eugene Fang
Lorenzo Riano
Rohan Chitnis
Stuart Russell
Pieter Abbeel

title: Combined task and motion planning through an extensible planner-independent interface layer

venue: IEEE international conference on robotics and automation (ICRA),

authors:
Karl Stratos
Sam Wiseman

title: Learning discrete structured representations by adversarially maximizing mutual information

venue: arXiv preprint arXiv:2004.03991,

title: Logic-geometric programming: An optimization-based approach to combined task and motion planning

authors:
Rishi Veerapaneni
John D Co-Reyes
Michael Chang
Michael Janner
Chelsea Finn
Jiajun Wu
Joshua Tenenbaum
Sergey Levine

title: Entity abstraction in visual model-based reinforcement learning

authors:
Ning Wang
Yibing Song
Chao Ma
Wengang Zhou
Wei Liu
Houqiang Li

title: Unsupervised deep tracking

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Zi Wang
Caelan Reed Garrett
Leslie Pack Kaelbling
Tomas Lozano-Perez

title: Active model learning and diverse action sampling for task and motion planning

venue: In International Conference on Intelligent Robots and Systems (IROS),

authors:
Manuel Watter
Jost Springenberg
Joschka Boedecker
Martin Riedmiller

title: Embed to control: A locally linear latent dynamics model for control from raw images

authors:
Nicholas Watters
Loic Matthey
Matko Bosnjak
Christopher P Burgess
Alexander Lerchner

title: Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven exploration

authors:
Fanyi Xiao
Yong Jae Lee

title: Track and segment: An iterative unsupervised approach for video object proposals

authors:
Wilson Yan
Ashwin Vangipuram
Pieter Abbeel
Lerrel Pinto

title: Learning predictive representations for deformable objects using contrastive estimation

venue: arXiv preprint arXiv:2003.05436,

heading: 1 INTRODUCTION

text:
In future, we hope that robots will be able to operate in unstructured environments such as homes and hospitals, and endowed with long-horizon planning ability. Despite successes in deep reinforcement learning (RL) from raw observations, much progress relies on the availability of shaped reward to guide the learning (Ng et al., 1999; Mirza et al., 2020). On the other hand, over past decades, task and motion planning has been shown to solve much longer-horizon goal-directed tasks such as making a cup of coffee from torque control (Kaelbling & Lozano-Pérez, 2011; Srivastava et al., 2014; Toussaint, 2015; Wang et al., 2018). However, these methods often require pre-specified discrete abstract states, task representations and transition models, e.g., whether the robot is holding a cup and what actions (or perturbations) change such an abstract state. In this paper, we aim to learn discrete representations for high-level abstract planning from video interaction data, combined with a learned short-horizon controller.
Learning discrete representations from unsupervised data for planning is challenging for two reasons. First, the relationship between the optimization objective and the true task objective is not well-defined. Second, optimizing a model with a discrete layer is difficult with standard deep learning techniques. Recent methods approach the first problem using reconstruction or constrastive objectives (Watter et al., 2015; Anand et al., 2019; Ha & Schmidhuber, 2018; Kurutach et al., 2018; Hafner et al., 2019; Srinivas et al., 2020); however, the learned representations are in continuous latent space which is unstructured and difficult to combine with high-level abstract planning. While other methods show promise in learning discrete representations, they have not been applied to temporally-extended RL tasks (Oord et al., 2018; Razavi et al., 2019; Risi & Stanley, 2019; Asai & Fukunaga, 2017; Stratos & Wiseman, 2020).
In this work, we propose Discrete Object-factorized Representations for Planning (DORP) – a novel framework for visual planning and control by learning discrete representations and a low-level controller. DORP learns discrete representations from images that change slowly overtime, such as whether or not the agent holds a key or which room the agent is in, along with a low-level predictive model for control. These slow features enable the agent to plan at a low frequency in longer-horizon tasks. More specifically, DORP represents an abstract state as a set of one-hot vectors, and optimizes its encoder by maximizing a mutual information lower bound between the current representations to future observations (Oord et al., 2018). In order to train through the discrete layer, we apply the Gumbel-Softmax reparametrization trick (Jang et al., 2016; Maddison et al., 2016). Using abstract states as nodes, we build an approximate feasibility graph based on observed transition data. When provided with new start and goal images, the agent plan the shortest abstract path. Using the next
abstraction as waypoint, model-predictive control maximizes the objective that is 1 if it reaches the target abstraction and 0 otherwise with a trained video prediction model. Unlike other subgoal planning work (Savinov et al., 2018; Nasiriany et al., 2019; Laskin et al., 2020; Liu et al., 2020), by following abstract waypoints, DORP avoids unneccessary steps to match exact waypoint states.
In a set of experiments, we demonstrate that DORP learns temporally-consistent and objectfactorized representations suitable for planning. We show that these representations enable DORP to handle unseen long-horizon tasks more successfully compared to the states-of-the-arts in visual planning. Interestingly, we observe that latent representations show object-level factorization such as key-and-door.

heading: 2 PRELIMINARIES

text:
We present background material on unsupervised representation learning and discrete optimization that our method builds on, and our problem formulation.
Contrastive Predictive Coding (CPC) (Oord et al., 2018) learns low-dimensional representations that are most predictive of the future future high-dimensional sequential data. A non-linear encoder qθ : O → Rl, parametrized by θ, encodes the observation ot ∈ O to a latent l-dimensional vector representation zt. Let’s define a similarity score fk(zt, ot+k) = exp(ztψqθ(ot+k)) where ψ is a trainable l-by-l similarity matrix and ot+k is a future observation k steps ahead of ot. Given the query observation ot, we aim to classify the key observations – ot+k as positive and other sample õ from the dataset as negative. Formally, we optimize the loss function LCPC = −Eot,ot+k [ log fk(zt, ot+k)− log ∑ oj∈X fk(zt, oj) ] with respect to θ and ψ. This also corresponds to maximizing a lowerbound of the mutual information between the latent representation zt and the future observation ot+k.
Gumbel-Softmax (GS) (Jang et al., 2016) is a discrete optimization technique to compute the gradient of through samples from a categorical distribution πi. GS first applies a reparametrization trick to rewrite samples zi as onehot(argmaxj(gj + logπi)) where gj are i.i.d samples from Gumbel(0, 1) (Gumbel, 1935). GS approximates argmax with softmax making the discrete stochastic layer differentiable, i.e., zi = softmax((gi + log πi)/τ). When the temperature parameter τ approaches 0, the samples converge to the true categorical distribution. Empirically, the temperature τ starts high and is annealed closer to 0 by a certain schedule.

heading: 2.1 PROBLEM STATEMENT: GOAL-DIRECTED VISUAL PLANNING AND CONTROL

text:
We define an unknown, fully-observable, stochastic dynamical system f which maps input observation ot ∈ O and action at ∈ A to the next observation ot+1. Under this dynamical system, we assume a simple exploration policy πrand which can collect data that characterizes the dynamics of the system. This is known as self-supervised data or play data (Agrawal et al., 2016; Lynch et al., 2020). We consider high-dimensional observation such as images.
Our goal is to learn discrete abstraction zt ∈ Z given observation ot, an abstract feasibility model, and a low-level local controller. At test time, given start os and goal og observations, we can plan a sequence of abstract states zs, z1, ..., zg where zs and zg are the representation of os and og , and apply the low-level controller to reach these abstract waypoints and finally to og .

heading: 3 DISCRETE OBJECT-FACTORIZED REPRESENTATION FOR PLANNING

text:
Our objective is to derive representation properties such as object-factorization and temporalconsistency from an unsupervised learning objective to facilitate long-horizon planning. When object representations are factorized, the agent can escape the need of combinatorial data configurations and can quickly generalize to unseen tasks. Also, the connectivity graph can also be memory-efficiently represented and combined with a more powerful planning algorithm. Another important property for a representation is temporal consistency, meaning any two state observations in the same abstraction should be reachable from one-another by a short sequence of actions. When this property holds in the latent space, a high-level plan can be successfully executed by a low-level controller.
We propose Discrete Object-factorized Representation for Planning (DORP) to tackle goal-directed visual planning and control problem. DORP learns latent discrete representations, a connectivity graph for abstract planning, and a low-level model-predictive controller. We describe each subcomponent in Section 3.1 - 3.3.

heading: 3.1 DISCRETE REPRESENTATIONS

text:
Despite successes in applying its unsupervisedly learned representations for downstream tasks, CPC is not endowed with an inductive bias for its representations to facilitate planning (Oord et al., 2018). Additionally, it is unclear how a discrete representations can be extracted and optimized. Toward this goal, we propose an object-factorized architecture with asymmetric Gumbel softmax for discrete optimization, and a diagonally dominating and diagonally positive similarity matrix to encourage temporal consistency.
We implement our DORP architecture to encourage object-factorization and capture discrete representations. The extractor passes an anchor observation through a convolutional neural network and outputs a c-channel feature map. Each feature map is inputted to a shared encoder followed by a Gumbel softmax. For a positive and negative pair of images, however, we opt for a continuous embedding by swapping Gumbel softmax for softmax. Empirically, without this asymmetry trick, the optimization tends to converge to a poor local optima. Together these one-hots are flattened into long vectors, and their bilinear product is the similarity score between the query-key pair (see Figure 2).
We propose an inductive bias in the similarity score function to encourage temporal consistency. Instead of a fully trainable matrix, we parameterize the similarity matrix ψ to be diagonally dominant and diagonally positive (DDDP). This property biases the similarity score to be high when the key embedding z̃t+k is close to the query embedding zt and maximized when z̃t+k = zt. In other words, this incentivizes the representations of positive pairs to share
weight matrix as eψ0(−σ(ψ1) + αIl×l) where Il×l is an identity matrix, α is a positive constant, σ is a sigmoid function, ψ0 is a trainable scalar, and ψ1 is a trainable l-by-l matrix.

heading: 3.2 ABSTRACT PLANNING

text:
A critical component in planning is a connectivity graph that decides whether two discrete representations are connected. One naive solution is to build such a graph from data – all observed representations as nodes and all observed transitions as edges. However as the number of independent entities increases, the amount of data required to generate the nodes and edges in order to cover the state space increases exponentially. To solve this, we propose to reduce the planning problem by exploiting factorization (see Figure 3).
1. First, embed the current image os and goal og as zs = qθ(os) and zg = qθ(og).
2. Pick a random one-hot index i s.t. the current state’s one-hot zis differs from the goal’s one-hot zig .
3. Build a graph based on how this one-hot transits in the data (ignoring all other one-hots). 4. Plan a path zi1, ..., z i g from z i s to z i g . If this does not exists, then we fail the task.
5. Execute the low-level controller following the next target zik, k = 1, · · · , g, while keeping other zj the same.
6. If it doesn’t reach zik, remove the edge from the graph and redo the planning in step 4.
7. If it succeeds, follow the next target until it reaches zig . 8. Repeat step 2 until all the one-hots match. 9. Finally, execute the low-level controller to the goal image og .1
Our proposed method assumes that each latent one-hot can be manipulated independently to solve the task. This assumption can hold various settings. For example, imagine a task that requires reorganizing furniture in a spacious room, and each objects is represented as a separate one-hot. Here, we can manipulate each object separately, and since there is plenty of empty space in this scenario, we can easily plan around objects to avoid collisions.
This assumption does not always hold; along the same scenario, when the space is more cluttered, objects are more likely to collide, and the proposed fully-factorized planning may fail even though a solution exists (incomplete). We extend this planning algorithm by building a graph on a random set of one-hot indices at a time in step 2. When can choose this set to contain all the one-
1An implicit assumption of temporal consistency in the learned discrete representations is required.
hots the method is equivalent to the full graph search (complete but expensive). Thus, our extended planning method can trade between completeness and efficiency.

heading: 3.3 LOW-LEVEL CONTROL

text:
The low-level controller is the key component in order to follow abstract plans and reach the goal. For this reason, the controller must be flexible in achieving multi-objectives. Modelpredictive controller (MPC) achieves this by modeling the dynamics of the world from past data and minimizing with the total predicted cost of horizon T at run-time: a∗t , ..., a ∗ t+T−1 =
argminat,...,at+T−1 E [∑T i=1 γ tĉt+i ] . The predicted cost is computed by applying a known cost function on the predicted outcome of a T -step action sequence from the current observation ot.
MPC requires the world model and the costs. For the world model, we implement a stochastic video prediction model (Babaeizadeh et al., 2017). In our framework, we define two cost functions for reaching abstract goals and for reaching the final goal observation. The first cost function is defined to be 1 if the current embedding exactly matches the goal embedding (all the one-hots match) and 0 otherwise: clat(ot, z′) = 1[qθ(ot) == z′]. The second cost function is defined by its L2 loss in the observation space to the goal: cobs(ot, og) = ||ot − og||. This cost is versatile for reaching nearby observations particularly to reach the goal observation once it has reached the goal code.

heading: 4 RELATED WORK

text:
Object Discovery Recent work has studied unsupervised object segmentations from visual inputs (Burgess et al., 2019; Greff et al., 2019), and entity-factorized representations and models for predictions (Kipf et al., 2019). However, these studies have not been demonstrated to directly solve the task. In this work, we take a step further to evaluate the representations in downstream tasks. While other work (Veerapaneni et al., 2020; Watters et al., 2019) shows how the MPC agents benefit from object-factorized models, they require shaped rewards for the tasks. In contrast, we consider a longhorizon task in which the agent only receives its reward when it has reached the goal. A large body of work in computer vision has studied unsupervised video object segmentation (Aloimonos et al., 1988; Bajcsy, 1988; Bajcsy et al., 2018; Gibson, 2014; Fitzpatrick, 2003; Hausman et al., 2015; Kenney et al., 2009) and unsupervised object detection (Kwak et al., 2015; Wang et al., 2019; Xiao & Jae Lee, 2016). However, semantic segmentation and object detection might not be the correct representation for the end-task. In this work, we learn representations that are more ready to use for planning without explicit segmentation or detection.
World models Much previous work has applied generative models of the world to visual control tasks (Watter et al., 2015; Ha & Schmidhuber, 2018; Hafner et al., 2019). Other work (Ma et al., 2020; Okada & Taniguchi, 2020; Yan et al., 2020) leverages a contrastive objective to learn a latent world model. These methods however require dense reward signals for most tasks. Visual foresight methods (Finn et al., 2016; Ebert et al., 2018) demonstrate impressive results on real robots using unshaped reward such as pixel distance to the goal, but it still remains limited to short-horizon object pushing tasks. In our work, we borrow these methods for low-level controllers. Asai & Fukunaga (2017) learn discrete abstraction of the system such as an 8-piece puzzle, but do not consider temporal abstraction. Kansky et al. (2017) demonstrate that discrete object-factorized representations can be used to learn logic-based transitions. In combination with powerful planning, it improves generalization and data efficiency. However, it assumes supervised ground-truth labels for representation learning. In this work, we aim to learn this in an unsupervised manner.
Hierarchial RL Recent work has tried to approach long-horizon visual planning by breaking down tasks using skills (Lynch et al., 2020; Kipf et al., 2019). Our approach is orthogonal and may deploy such action abstractions as a low-level controller. Other methods propose to plan subgoals and attempt to follow them either in the latent space (Kurutach et al., 2018; Nasiriany et al., 2019; Nair & Finn, 2019) or in the visual space (Savinov et al., 2018; Liu et al., 2020; Eysenbach et al., 2019). However, such methods require labor-intensive engineering to tune the threshold on when to move on to pursue the next subgoal (Savinov et al., 2018; Liu et al., 2020) because the observations or the latent states can never exactly match. Imagine bringing a chair from one office to another. It would be time consuming to match all the positions and orientations of the chair along the way. Rather we should care about rough area the chair has to go through in order to reach the goal. In our work, by learning the discrete codes, our method do not require such threshold as it exactly knows when
it has reached the target discrete representation. Additionally, it avoids aiming to match a specific waypoint observation. Instead of planning to the goal, Pertsch et al. (2020) predict intermediate images iteratively to construct a subgoal tree. However in order to train such prediction model a long-sequential training data are required. Instead our work can be applied to short-horizon or long-horizon trajectories.

heading: 5 EXPERIMENTS

text: We perform experiments aimed at answering the following questions: (1) Are DORP representations temporally-consistent for high-level planning? (3) Can DORP representations factorize objects that translate or change in appearances over time? (3) How does DORP compare to the visual planning SOTA in solving unseen long-horizon goal-directed tasks?

heading: 5.1 ENVIRONMENTS

text:
We evaluate DORP on two main environments – object-rearrangement and key-room. All the observations are provided as color images of size 16× 16× 3. k-object-rearrangement Multiple (k) 2-by-4 blocks of different colors can each be manipulated independently by an agent. The tasks require the agent to manipulate these objects from a start configuration to a goal configuration. In each step, it can manipulate each block by one unit in one of the 4 directions. The agent collects data by randomly interacting with one of the objects at each timestep in a purely exploratory manner, without any goal in mind. In this environment our trajectory has length one that is the configuration is randomly reset after every step.
key-room Two variations of key-room depend on the number of rooms – key-corridor (6 rooms and a corridor) and key-wall (2 rooms. A key and agent are represented by 1 pixel. A key is placed in a fixed location in a room. If the agent steps on the key a door will be removed allowing the agent to enter a locked room. Each step the agent can move in one of the four directions by 1 unit. The agent collect data by a long random walk – 1500 steps for key-corridor and 1000 steps for key-wall. After each reset the agent will be placed randomly in one of the unlocked rooms. This environment is inspired by MiniGrid (Chevalier-Boisvert et al., 2018).

heading: 5.2 BASELINES

text:
We choose the state-of-the-arts visual foresight (VF) (Ebert et al., 2018) implemented using SV2P architecture (Babaeizadeh et al., 2017) as our baseline for two reasons. First, VF only requires selfsupervised data collection and is applicable to unseen tasks – thus aligning with goal-directed visual planning problems. Second, Visual Foresight is a low-level controller of DORP. By sharing the same video model as baseline, we show a direct improvement of abstract planning for temporallyextended tasks. Two optimization variations of VF – VF with random shooting (VF-RS) and VF with cross-entropy method (VF-CEM) – both share the same video prediction model. Their objective is to minimize its L2 distance from the current image to the goal image. While DORP deploys VFRS for its low-level control, VF-CEM comparision are provided for an improved baseline. VF-RS randomizes 1000 trajectories and take the full action sequence that achieves the minimum cost. VF-CEM randomizes 1000 samples per iteration. We pick the top 2% to refit a distribution over sequence and repeat for 3 iterations.
To understand DORP representation learning improvement, we evaluate DORP when replacing its DDDP similarity matrix by an arbitrary weight (Oord et al., 2018) or an identity matrix (He et al., 2020) which have been used extensively in unsupervised representation learning. To understand the benefits of factorized planning, we replace it with the full graph planning with a maximum limit on the number of steps allowed.

heading: 5.3 QUANTITATIVE RESULTS

text:
Temporal consistency We investigate DORP representations by color-coding the discrete embeddings on the configuration map of the environments. For visualization simplicity we choose 1-
3With the exceptions for 5 objects in which we evaluate VF-RS on 10 tasks, VF-CEM on 11 tasks, and DORP on 23 tasks due to time and compute limits.
object-rearrangement to visualize the learned code map. In this environment we have a single one-hot code for representing the object. We demonstrate that DORP learns temporally consistent representations in which two states from the same discrete code are connected by a short-horizon controller (see Figure 5(a)). Similarly, we observe the same property in key-wall environment (see Figure 6).
Object factorization We evaluate the behavior of the one-hot codes by changing one property of the agent at a time such as positions or whether the agent has a key. In Figure 4(c) we observe that by moving one object at a time only at most one one-hot code changes its value k-object-rearrangement. In key-room, we formulate the discrete representation as follows: of the two output one-hot latents, the first is set to a size z0 and the other is set to size 2. Our aim is for one latent one-hot to encode the agent’s abstract state and for the other binary latent to encode the key’s abstract state. In Figure 6, we demonstrate that the learned abstract representations are factorized as desired by observing (1) the larger one-hot changes when the agent moves with no key interactions and (2) the binary one-hot changes when removing and adding the key while maintaining the agent position.
Weight Ablation We study the effect of the proposed similarity matrix against commonly-used similarity matrices (Oord et al., 2018; He et al., 2020) in unsupervised representation learning. In Figure 5, we demonstrate that DORP is able to exploit more available latent code and therefore helping the latent space to be more temporal consistent.

heading: 5.4 QUANTITATIVE RESULTS

text: Long-horizon planning We test the agent by its ability to solve unseen goal-directed tasks. In k-object-rearrangement, we can study the effect of the increase of the length and complexity of the tasks by increasing the number of objects. In Table 1, we demonstrate that as the number of objects increases DORP is able to succeed in most of the tasks. While other methods’ performances degrade quickly. Note that when faced with 5 objects DORP employs the extended planning version by grouping the 5 one-hots into two groups of 2 and 3 one-hots. By considering more that one object at a time, the agent is able to perform non-myopic planning and achieve 87% success rate. We also evaluate DORP on two variations of key-room. In our evaluations, we find that DORP solves keywall with a single wall and door successfully in 9 of 10 tasks. When using a ground-truth dynamics model in place of an SV2P model, DORP solves Key-Room successfully in 37 of 40 tasks, and Key-Corridor in 33 of 40 tasks.

heading: 6 CONCLUSION AND FUTURE DIRECTIONS

text: In this paper, we propose an unsupervised discrete representation learning method for long term planning, which can extract high level abstract states that are good for planning in an purely unsupervised manner. We demonstrate DORP’s effectiveness over other methods on challenging long horizon tasks. We note that our method generate approximate plans more tractably and hence trades off optimality for efficiency, like other state abstraction works (Savinov et al., 2018; Liu et al., 2020; Eysenbach et al., 2019). Those approximate solutions can be used to initialize a model-free policy which can be later fine-tuned to reach optimality. We leave this as our future work.