Planning from Pixels in Environments with Combinatorially Hard Search Spaces
Authors: Marco Bagatella
Miroslav Olšák
Michal Rolinek
Georg Martius
Venue: NeurIPS 2021 Poster (2021)
Keywords: World Model
Representation Learning
Model-Based Reinforcement Learning

Abstract: The ability to form complex plans based on raw visual input is a litmus test for current capabilities of artificial intelligence, as it requires a seamless combination of visual processing and abstract algorithmic execution, two traditionally separate areas of computer science. A recent surge of interest in this field brought advances that yield good performance in tasks ranging from arcade games to continuous control; these methods however do not come without significant issues, such as limited generalization capabilities and difficulties when dealing with combinatorially hard planning instances. Our contribution is two-fold: (i) we present a method that learns to represent its environment as a latent graph and leverages state reidentification to reduce the complexity of finding a good policy from exponential to linear (ii) we introduce a set of lightweight environments with an underlying discrete combinatorial structure in which planning is challenging even for humans. Moreover, we show that our methods achieves strong empirical generalization to variations in the environment, even across highly disadvantaged regimes, such as “one-shot” planning, or in an offline RL paradigm which only provides low-quality trajectories.

Planning from Pixels in Environments with Combinatorially Hard Search Spaces
Marco Bagatella
Mirek Olšák Computer Science Department University Innsbruck, Austria mirekfoles sak net
Max Planck Institute for Intelligent Systems
Tübingen, Germany
mbagatellad wawellll@uteee.. mpg
mpg de
Michal Rolínek
Georg Martius
Max Planck Institute for Intelligent Systems michal,
Max Planck Institute for Intelligent Systems Tübingen, Germany
Tübingen, Germany
rol rolaalrrllii
neke tue. mpg. de
georg. wpw. de
Abstract
The ability to form complex plans based on raw visual input is a litmus test for current capabilities of artificial intelligence, requires seamless combination of visual processing and abstract algorithmic execution, two traditionally separate areas of computer science. A recent surge of interest in this field brought advances that yield good performance in tasks ranging from arcade games to continuous control; these methods however do not come without significant issues, such limited generalization capabilities and difficulties when dealing with combina- torially hard planning instances. Our contribution is two-fold: method that learns to represent its environment as a latent graph and leverages state reidentification reduce the complexity of finding a good policy from exponential to linear (ii) we introduce set of lightweight environments with an underlying discrete combinatorial structure in which planning is challenging even for humans. Moreover, we show that our methods achieves strong empirical generalization to variations in the environment, even across highly disadvantaged regimes, such as "one-shot" planning, or in an offline RL paradigm which only provides low-quality trajectories.
we present a
latent graph space
1.0
success rate
network
18
0.6
margin
0.2
breadth first search (with reidentification)
IceSlider
Figure 1: Planning from Pixels with Graph Search. Our method leverages learned latent dynamics efficiently build and search a graph representation of the environment. Resulting policies show unrivaled performance across a distribution of hard combinatorial tasks.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Introduction
Decision problems with an underlying combinatorial structure pose a significant challenge for learning agent, as they require both the ability to infer the true low-dimensional state of the environment and the application of abstract reasoning to master it. A traditional approach for common logic games, given that simulator or model of the game are available, consists in applying a graph search algorithm to the state diagram, effectively simulating several trajectories to find the optimal one. As long the state space of the game grows polynomial rate with respect to the planning horizon, the solver is able to efficiently find the optimal solution to the problem. Of course, when this is not the case, heuristies can be introduced at the expense of optimality of solutions.
Learned world models 17 18 can learn to map complex observations to a lower-dimensional latent space and retrieve an approximate simulator of an environment. However, while the continuous structure of the latent space is suitable for training reinforcement learning agents [12 [9] or applying
heuristic search algorithms [38] it also prevents straightforward application of simpler graph search
techniques that rely on identifying and marking visited states.
Our work follows naturally from the following insight: a simple graph search might be sufficient for solving visually complex environments, as long as a world model trained realize a suitable structure in the latent space. Moreover, the complexity of the search can be reduced from exponential to linear by reidentifying visited latent states.
The method we propose is located at the intersection between classical planning, representation learning and model-based reinforcement learning. It relies on novel low-dimensional world model trained through a combination of opposing losses without reconstructing observations. We show how learned latent representations allow a dynamics model to be trained to high accuracy, and how the dynamics model can then be used to reconstruct a latent gruph representing environment states as vertices and transitions as edges. The resulting latent space structure enables powerful graph search algorithms to be deployed for planning with minimal modifications, solving challenging combinatorial environments from pixels. We name our method PPGS as it Plans from Pixels through Graph Search.
We design PPGS to be capable of generalizing unseen variations of the environment, or equivalently across a distribution of levels [13]. This is in contrast with traditional benchmarks 77, which require the agent to be trained and tested on the same fixed environment.
We can describe the main contributions of this paper as follows: first, we introduce a suite of environments that highlights weakness of modern reinforcement learning approaches, second, we introduce a simple but principled world model architecture that can accurately Iearn the latent dynamics of complex system from high dimensional observations; third, we show how planning module can simultaneously estimate the latent graph for previously unseen environments and deploy breadth first search in the latent space to retrieve competitive policy; fourth, we show how combining our insights leads to unrivaled performance and generalization on a challenging class of environments.
2 Method
For the purpose of this paper, each environment can be modeled as a family of fully-observable deterministic goal-conditions {(S, A, T, G. R2 7):).... where S; is the state set, A; is the action set, T; is transition function : x A, + S., G2 is the goal set and R2 is a reward function R; S x G; - IR and % is the discount factor. We remark that each environment can also be modeled as a BlockMDP 44] in which the context space X' corresponds to the state set S, we introduced.
Markov Decision Processes with discrete actions, that is the 6-tuples
In particular, we deal with families of procedurally generated environments. We refer to each of the n elements of a family as a level and omit the index i when dealing with a generic level. We assume that state spaces and action spaces share the same dimensionality across all levels, that is |Sal = |S and |A,| = | for all 0 < i.j S n.
In our work the reward simplifies to an indicator function for goal achievement R(s, = 1, with G S S. Given goal distribution p(g), the objective is that of finding a goal-conditioned policy mg that maximizes the return

---

Encoder hg
Zt
Forward Model
B
f
Inverse Model Pw
Amargin
Lce
at
Cr
Encoder hg
2t-11
LFW
Figure 2: Architecture of the world model. A convolutional encoder extracts latent state representa- tions from observations, while forward model and an inverse model reconstruct latent dynamics by predicting state transitions and actions that cause them. The notation is introduced in Sec. 2.1
Jx
Ee)R(s
(1)
gr-((a) Trpp(T|=s)
where t ~ pg)[ is a trajectory o+) sampled from the policy.
Our environments of interest should challenge both perceptive and reasoning capabilities of an agent. In principle, they should be solvable through extensive search in hard combinatorial spaces. order to master them, agent should therefore be able to (i) identify pairs of bisimilar states 43 (ii) keep track of and reidentify states it has visited in the past and (iii) produce highly accurate predictions non-trivial time horizons. These factors contribute to making such environments very challenging for existing methods. Our method designed in light these necessities; has two integral parts, the world model and the planner, which now introduce.
2.1 World Model
The world model relies solely on three jointly trained function approximators: an encoder, a forward model and inverse model. Their overall orchestration is depicted in Fig. 2 and described in the following.
2.1.1 Encoder
Mapping highly redundant observations from an environment to now-dimensional state space Z has several benefits [8). Ideally, the projection should extract the compressed "true state" of the environment and ignore irrelevant visual cues, discarding all information that is useless for planning. encoder he, that neural function approximator mapping
For this purpose, method relies each observed state s and a oow-dimensional representation 2 E (embedding). While there many suitable choices for the structure of the latent space Z, choose map observations points on an d-dimensional hypersphere taking inspiration from Liu et al. 29
2.1.2 Forward Model
In order to plan ahead the environment, it is crucial for an agent to estimate the transition function T. In fact, if mapping low- oow-dinensional latent space Z is available. learning directly the projected transition function Tz Z A + can be largely beneficial [17 [8]. The deterministic latent transition function T' can be learned by neural function approximator f so that if a.) 8t+1. then fo (s), a.) := fo (2,, az) = hg ($++1). We refer to this component as forward model. Intuitively, it can be trained retrieve the representation of the state of the MDP time + given the representation of the state and the action taken at the previous time step t.
Due to the Markov property of the environment, an initial state embedding and the action sequence (ot,..., ((+k) are sufficient to to predict the latent state time + k; as long as 2, successfully captures all relevant information from the observed state st. The amount of information to be

---

embedded in and to be retained in autoregressive predictions is, however, in most cases, prohibitive. Take for example the case of simple maze: 2, would have encode not only the position of the agent, but, as the predictive horizon increases, most of the structure of the maze.
Invariant Structure Recovery To allow the encoder to only focus local information, adopt an hybrid forward model which can recover the invariant structures in the environment from previous observations. The function that the forward model seeks approximate can then include additional input: fa(z,, Oz, 8e) = 2, where 8, E S is a generic observation from the same environment and level. Through this context input the forward model can retrieve information that is constant across time steps (e.g. the location of walls in static mazes), In practice, we can use randomly sampled observation from the same level during training and use the latest observation during evaluation. This choice allows for more accurate and structure-aware predictions, we show in the ablations in Suppl.
Given a trajectory (st, az){== the forward model can be trained to minimize some distance measure between state embeddings (2.41)1...7T1 (Jo(ha(s)), (t, 8..)....T-I In practice, we choose to minimize a Monte Carlo estimate of the ex- pected Euclidean distance over finite time horizon, set of trajectories and set of levels. When training on a distribution of levels p((), we extract K trajectories of length H from each level with uniform random policy æ and we minimize
(... and one-step predictions
Lew R (xp()
Σ
E
[1,o(Ehs
an, an, +
h=1
where the superscript indicates the level from which the embeddings are extracted.
2.1.3 Inverse Model and Collapse Prevention
Unfortunately, the loss landscape of Equation 2 presents a trivial minimum in case the encoder collapses all embeddings single point in the latent space. As embeddings of any pair of states could not be distinguished in this case, this is not desirable solution. We remark that this is known problem in metric learning and image retrieval [8] for which solutions ranging from siamese networks [9] to using triplet loss (22) have been proposed.
The context of latent world models offers natural solution that isn't available in the general embedding problem, which consists in additionally training a probabilistic inverse model pw 2t, 2(+1) ) such that if T2 4z) 2++1, then (w 2t, 2+11) > (a 2t, 2441)aa # 0g A. The inverse model, parameterized by w, can be trained to predict the action at that causes the latent transition between two embeddings 24, 2++1 by minimizing multi-class cross entropy.
H
log (a 122 sh+
LCE
E Imp(1)
E OA
(3)
H-
Intuitively, LcE increases as embeddings collapse, since it becomes harder for the inverse model to recover the actions responsible for latent transitions. For this reason, it mitigates unwanted local minima. Moreover, is empirically observed to enforce regular structure in the latent space that eases the training procedure, as argued in Sec. A of the Appendix. We note that this loss plays similar role to the reconstruction loss Hafner et [8] However, Log does not force the encoder network embed information that helps with reconstructing irrelevant parts of the observation, unlike training methods relying on image reconstruction [1 17
While Lce is sufficient for preventing collapse of the latent space, a discrete structure needs to be recovered in order to deploy graph search in the latent space. In particular, it is still necessary to define a criterion to reidentify nodes during the search procedure, or to establish whether two embeddings (directly encoded from observations or imagined) represent the same true low-dimensional state.
straightforward way to enforce this is by introducing a margin E, representing a desirable minimum distance between embeddings of non-bisimilar states [43] A third and final loss term then be
introduced to encourage margins in the latent space:
Emargin =
E Imp(()
1
H-1
Σ hml
max (0,
||
2 11
(4)
H
82

---

on
he
Figure 3: Overview of latent-space planning. One-shot planning is possible by (i) embedding the current observation and goal to the latent space and (ii) iteratively growing latent graph until vertex is reidentified with the goal.
We then propose to reidentify two embeddings representing the same true state if their Euclidean distance is less than 5-
Adopting latent margin effectively constrains the number of -eergin-epprateed states that can be represented on an hyperspherical latent space. However, this quantity is lower-bounded by the kissing number [11, that is the number of non-overlapping unit-spheres that can be tightly packed around one d dimensional sphere. The kissing number grows exponentially with the dimensionality d. Thus, the capacity of our d-dimensional unit sphere latent space (d 16 in our case with margin 0.1) is not overly restricted.
The world model can be trained jointly and end-to-end by simply minimizing combination of the three loss functions:
L = aCFw + BCCE + mmargin.
To summarize, the three components are respectively encouraging accurate dynamics predictions, regularizing latent representations and enforeing discrete structure for state reidentification.
2.2 Planning Regimes
deterministic environment can be represented directed graph G whose vertices V represent states E S and whose edges F encode state transitions. An edge from vertex representing a state present if and only if T(s, a) s' for some action
E S vertex representing state s' E A, where 7 is the state transition function of the environment. This edge can then be labelled by action a. Our planning module relies on reconstructing the latent graph, which is projection of graph G to the latent state Z.
In this section we describe how latent graph can be build from the predictions of the world model and efficiently searched to recover a plan, as illustrated in Fig. 3 This method can be used as a one-shot planner, which only needs access visual goal and the initial observation from level. When iterated and augmented with online error correction, this procedure results in a powerful ap- proach, which we refer to as full planner, or simply PPGS.
1012
WRhout reddontiication
new vertices 14,
101:
Gatoff
10a
10"
10a
One-shot Planner Breadth First Search (BFS) is graph search algorithm that relies on a LIFO queue and marking visited states to find an optimal path O(V E) steps. Its simplicity makes it an ideal candidate for solving combinatorial games by exploring their latent graph. If the number of reachable states in the environment grows polynomially, the size of the graph to search will increase at modest rate and the method can be applied efficiently.
10
15
20
Figure 4: Number of leaf vertices when planning in ProcgenMaze, averaged over 100 levels, with 90% confidence inter- vals.
5

---

We propose to execute BFS-like algorithm on the latent graph, which is recovered by autoregressively simulating all transitions from visited states. As depicted in Fig. []] at each step, the new set of leaves I, is retrieved by feeding the leaves from the previous iteration through the forward model fø. The efficiency of the search process can be improved as shown
Fig. 4 by exploiting the
margin e enforced by equation to reidentify states and identify loops in the latent graph. We now provide a simplified description of the planning method in Algorithm [
while details can be found in
Suppl. C,2
Alg Algorithm 1 Simplified one-shot PPGS
Input: Initial observed state 81, visual goal g. model parameters 8, Φ
1: 21; 2g hg hg(s)),ho(9)
> project to latent space Z > sets of leaves and visited vertices
2: L, V {21}
3: for TMAX steps do
4:
L= (ff(2,a,81) E L, E A}
Di grow graph
5:
=f E L. can be reidentified with 2g then
6:
return action sequence from 21 2²
7:
end if
8:
L= VV
reidentify and discard visited vertices (details in Suppl. C.2 D update visited vertices
9:
V VUVUL
10: end for
Full Planner The one-shot variant of PPGS largely relies on highly accurate autoregressive predictions, which learned model cannot usually guarantee. We mitigate this issue by adopting a model predictive control-like approach 15. PPGS recovers an initial guess on the best policy (a:)1..... simply by applying one-shot PPGS described the previous paragraph and Algorithm 2] It then applies the policy step by step and projects new observations to the latent space. When new observations do not match with the latent trajectory, the policy recomputed by applying one-shot PPGS from the latest observation. This happens when the autoregressive prediction of the current embedding (conditioned on the action sequence since the last planning iteration) can not be reidentified with the embedding of the current observation. Moreover, the algorithm stores all observed latent transitions lookup table and, when replanning, only trusts the forward model on previously unseen observation/action pairs. A detailed description can be found in Suppl. C.2
3 Environments
In order to benchmark both perception and abstract reasoning, we empirically show the feasibility of our method three challenging procedurally generated environments. These include the Maze environment from the procgen suite [3), as well as DigitJump and IceSlider, two combinatorially hard environments which stress the reasoning capabilities of learning agent, or even of an human player. In the context of our work, the term "combinatorial hardness" is used loosely. We refer to an environment "combinatorially hard" if only very few of the exponentially many trajectories actually lead to the goal, while deviating from them often results in failure (e.g. Digitlump or IceSlider). Hence, some "intelligent" search algorithm required. In this way, the process of retrieving successful policy resembles that of a graph-traversing algorithm. The last two environments are made available in public repository M. where they can also be tested interactively. More details on their implementation are included in Suppl. D
ProcGenMaze The ProcgenMaze environment consists of a family of procedurally generated 2D mazes. The agent starts in the bottom left corner of the grid and needs to reach position marked by piece of cheese. For each level, an unique shortest solution exists, and its length is usually distributed roughly between 1 and 40 steps. This environment presents significant intra-level variability, with different sizes, textures, and maze structures. While retrieving the optimal solution this environment is already a non-trivial task, its dynamics are uniform and actions only cause local changes in the observations. Moreover, ProcgenMaze is forgiving environment in which errors can always be recovered from. In the real world, many operations are irreversible, for instance, cutting/breaking objects, gluing parts, mixing liquids, etc. Environments containing remote controls, for example, show non-local effects. We use these insights to choose the additional environments.

---

ProegenMaze
DigitJump
IceSlider
5 5/523/62
3 552442 a 2 /
45 156341253 3
13645/ 4
1 1 /
53654/32
535
Figure 5: Environments. Initial observations and one-shot PPGS's solution (arrows) of a random level of each of the three environments. ProegenMaze is from [3] DigitJump and IceSlider are proposed by and can be accessed at [1]
IceSlider IceSlider is in principle similar to ProcgenMaze, since it also consists of procedurally generated mazes. However, each action propels the agent in direction until an obstacle (a rock or the borders the environments) is met. We generate solvable but unforgiving levels that feature irreversible transitions, that, once taken, prevent the agent from reaching the goal.
DigitJump DigitJump features a distribution of randomly generated levels which consist of a 2D 8x8 grid of handwritten digits from 6. The agent needs go from the top left corner to the bottom right corner. The 4 directional actions are available, but each of them causes the agent to move in that directions by the number steps expressed by the digit on the starting cell. Therefore, single action can easily transport the player across the board. This makes navigating the environment very challenging, despite the reduced cardinality of the state space. Moreover, the game presents many cells in which the agent can get irreversibly stuck.
4 Related Work
World Models and Reinforcement Learning The idea of learning to model an environment has been widely explored in recent years. Work by Oh et al. [32] and Chiappa al. [II] has argued that modern machine learning architectures are capable of learning to model the dynamics a generic environment reasonably well for non-trivial time horizons. The seminal work by Ha and Schmidhuber [17] built upon this by learning a world model a low-dimensional latent space instead of conditioning predictions on observations. They achieved this by training VAE on reconstructing observations and a recurrent network for sampling latent trajectories conditioned on an action sequence. Moreover, they showed how sample efficiency could be addressed by recovering a simple controller acting directly on latent representations through an evolutionary approach.
This initial idea was iteratively improved along two main directions. On one hand, some subsequent works focused on learning objectives and suggested to jointly train encoding and dynamics compo- nents. Hafner et al. [18] introduced multi-step variational inference objective to encourage latent representations to be predictive of the future and propagate information through both deterministic and stochastic paths. On the other hand, authors proposed to learn act in the latent space by using zero-order methods [18] such as CEM (36] or policy gradient techniques
19 20 These
improvements gradually led to strong model-based RL agents capable of achieving very competitive
performance in continuous control tasks 19 and on the Atari Learning Environment [ 10 20
Relying on image reconstruction limitation Okada and Taniguchi [33] and Zhang et al. 45 forgo the decoder network, while the latter proposes to rely the notion of bisimilarity to learn meaningful representations. Similarly, Gelada et al. [6] only learn to predict rewards and action-ccn ceniionconditionall state distributions, but only study this task additional loss model-free reinforcement learning methods. Another relevant approach that of [44) who propose to learn a discrete graph representation of the environment, but their final goal is that of recovering a series of subgoals for model-free RL.
however lead vulnerability to visual noise: to overcome this
strong example of how world models can be coupled with elassical planners is given by MuZero 38 MuZero trains recurrent world model to guide Monte Carlo tree search by encouraging
7

---

hidden states to be predictive of future states and a sparse reward signal. While we adopt a similar framework, we focus on recovering a discrete structure in the latent space order to reidentify states and lower the complexity of the search procedure. Moreover, we do not rely on reward signals, but only focus on learning the dynamics of the environment.
Neuro-algorithmic Planning In recent years, several other authors have explored the intersection between representation learning and classical algorithms. This is the case, for instance, of Ichter and Pavone [3]] Kumar et al. (26). Kuo et al. [27] who rely sequence models or VAEs to propose planners. Within planning research, Yonetani et al. 42] introduce a differentiable version of the A search algorithm that can learn suitable representations from images with supervision. The most relevant line of work to us is perhaps the one that attempts to learn representations that are suitable an input for classical solvers. Within this area, Asai and Fukunaga 4. Asai and Muise [ show how symbolic representations can be extracted from complex tasks in end-to-end fashion and directly fed into off-the-sher solvers. More recently, Vlastelica et al. 40 frames MDPs as shortest-pait problems and trains a convolutional neural network to retrieve the weights of fixed graph structure. The extracted graph representation can be solved with combinatorial solver and trained end-to-end by leveraging the blackbox differentiation method [35]
trajectories for sampling-base
Visual Goals A further direction of relevant research is that of planning to achieve multiple goals [30]. While the most common approaches involve learning
condiiondi policy with experience relabeling []. the recently proposed GLAMOR (34) relies on learning inverse dynamies and retrieves policies through a recurrent network. By doing so, it can achieve visual goals without explicitly modeling reward function, an approach that is sensibly closer to ours and can serve as relevant comparison. Another method that sharing similar setting to ours is LEAP [11] which also attempts to fuse reinforcement learning and planning: however, its approach is fundamentally different and designed for dense rewards and continuous control. Similarly, SPTM (37 pursues similar direction, but requires exploratory traversals in the current environment, which would be particularly hard to obtain due to procedural generation.
5 Experiments
The purpose of the experimental section is to empirically verify the following claims: (i) PPGS is able to solve challenging environments with an underlying combinatorial structure and (ii) PPGS is able to generalize to unseen variations of the environments, even when trained on few levels. We aim to demonstrate that forming complex plans in these simple-looking environments is beyond the reach of the best suited state-of-the-arr methods. Our approach, on the other hand, achieves non-trivial performance. With this in mind, we did not insist on perfect fairnsss of all comparisons, as the different methods have different type of access to the data and the environment. However, the largest disadvantage is arguably given to our own method.
While visual goals could be drawn from a distribution P(g), we evaluate a single goal for each test level matching the environment solution (or the only state that would give positive reward in sparse reinforcement working framework). This represents a very challenging task with respect to common visual goal achievement benchmarks [34], while also allowing comparisons with reward- based approaches such as PPO 33 We mainly evaluate the success rate, which is computed as the proportion of solved levels in set of 100 unseen levels. A level is considered to be solved when the agent achieves the visual goal (or receives a non-zero reward) within 256 steps.
Choice of Baselines Our method learns to achieve visual goals by planning with world model learned on a distribution of levels. To the best of our knowledge, no other method in the literature shares these exact settings. For this reason, we select three diverse and strong baselines and we make our best efforts for a fair comparison within our computational limits.
PPO 39 is strong and scalable policy optimization method that has been applied in procedurally generated environments [33 While PPGS requires a visual goal to given, PPO relies on a (sparse) reward signal specializing on unique goal per level. DreamerV2 20 is moseel-based RL approach that also relies on a reward signal, while GLAMMRR (34 is more aligned with PPGS as is also designed to reach visual goals in absence of reward.
8

---

While we restrict PPGS to only access offline dataset of low-quate low-quality random trajectories, all baselines are allowed collect data on policy for a much larger number environment steps. More considerations on these baselines and on the fairness of our comparison can be found in Suppl. Furthermore, we also consider non-learning naive search algorithm (GS ON IMAGES) thoroughly described in c
comprehensive ablation study of PPGS can be found in Section of the Appendix.
5.1 Comparison of Success Rates
Our first claim is supported by Figure 6 PPGS perform its baselines across the three environments. The gap with baselines smaller in ProcgenMaze, a forgiving environment for which accurate plans are not necessary. On the other hand, ProcgenMaze in- volves long-horizon planning, which can be seen as
Dan 1993
0.8
Disament
Success rate
0.6
a limitation to one-shot PPGS. As the combinatorial nature of the environment becomes more important, the gap with all baselines increases drastically.
0.4
0.2
PPO performs fairly well with simple dynamics and long-term planning, but struggles more when com- binatorial reasoning is necessary. GLAMOR and DreamerV2 struggle across the three environments, as they likely fail to generalize across a distribution of levels. The fact that GS ON IMAGES manages to rival other baselines is testament to the harshness of the environments.
ProcgenMaze
IceSlider
Digitjup
Figure 6: Success rates across the three envi- ronments. One-shot planning is competitive with the full method on shorter time horizons.
5.2 Analysis of Generalization
The inductive biases represented by the planning algorithm and our training procedure ensure good generalization from minimal number of training levels. In Fig. 7 compare solution rates between PPGS and PPO as the number levels available for training increases. The same metric for larger training level sets is additionally available in Table Our method generally outperforms its baselines across all environments. In ProcgenMaze, PPGS achieves better success rates than PPO after only
IceSlider
Digitjump
1.00
Success rate
0.75
Success rale
1.0
0.50
0.5
0.25
0.00
0.0
10
10t
102
10a
10a
10a
# training levels
training levels
ProcgenMaze
0.8
SUcCess rate
0.6
PPGS
0.4
PPO
GS Images
0.2
10a
10a
training levels
Figure 7: Solution rates of PPGS and PPO as function of the cardinality of the set of training levels.

---

seeing two orders of magnitude less level, e.g. 10 levels instead of 1000. Note that PPGS uses only 400k samples from a random policy whereas PPO uses 50M on-policy samples. Due to the harshness of the remaining environments, PPO struggles to find a good policy and its solution rate on unseen levels improves slowly as the number of training levels increases. IceSlider, PPGS well above PPO for any size of the training set and outperforms GS ON IMAGES when only having access to 2 training levels. While having a comparable performance to PPO on small training sets method severely outperforms it once approximately 200 levels are available. On the other hand, PPO's ability to generalize plateaus. These results show that PPGS quickly learns to extract meaningful representations that generalize well to unseen scenarios.
in DigitJump,
6 Discussion
Limitations The main limitations of method regard the assumptions that characterize the class of environments we focus on, namely slowly expanding state space and discrete actions. In general, due to the complexity of the search algorithms, scaling to very large action sets becomes challenging. Moreover, a single expansion of the search tree requires a forward pass of the dynamics network, which takes a non-negligible amount of time. Finally, the world model is fundamental component and the accuracy of the forward model is vital to the planner. Training an accurate forward model can be hard when dealing with exceedingly complex observations: very large grid sizes in the environments are a significant obstacle. On the other hand, improvements in the world model would directly benefit the whole pipeline.
Conclusion Hard search from pixels is largely unexplored and unsolved, yet fundamental for future AL In this paper presented how powerful graph planners can be combined with learned perception modules to solve challenging environment with a hidden combinatorial nature. In particular, our training procedure and planning algorithm achieve this by (i) leveraging state reidentification to reduce planning complexity and (ii) overcoming the limitation posed by information-dent through hybrid forward model. We validated our proposed method, PPGS, across three challenging environments in which we found state-of-the- methods to struggle. We believe that our results represent sensible argument in support of the integration of learning-base -basd approaches and classical solvers.
observations
Acknowledgments and Disclosure of Funding
We acknowledge the support from the German Federal Ministry of Education and Research (BMBF) through the Tübingen Al Center (FK7: 011S (ISS1839B). Georg Martius is member of the Machine Learning Cluster of Excellence, funded by the Deutsche Forschungsgemeinschatt (DFG, German Research Foundation) under Germany's 's Excellence Strategy - EXC number 2064/1 Project number 390727645.
References
[1]
bttpes//Etttaaccooaaaariia/a/l/p 2021.
[2]
2021.
[3] M. Andrychowiez, F. Wolski, A. Ray, J. Schneider, R. Fong, P. Welinder, B. McGrew, J. Tobin,
O. Pieter Abbeel, and W. Zaremba. Hindsight experience replay. In Advances in Neural Information Processing Systems, 2017.
[4] M. Asai and A. Fukunaga. Classical planning in deep latent space: Bridging the subsymbolic- symbolic boundary. In S. A. Mcllraith and K. Q. Weinberger, editors, Proceedings of the T'hirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (1AAI-18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pages 6094-6101, 2018
[5] M. Asai and C. Muise. Learning neural-symbolic descriptive planning models via cube-space priors: The voyage home (to STRIPS). C. Bessiere, editor, Proceedings of the Twenty-Ninth International Joint Conference Artificial Intelligence, IJCAI 2020, pages 2676-2682, 2020.
10

---

[6] J. L. Ba, R. Kiros, and G. E. Hinton. Layer normalization, 2016.
[7] M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling, The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
[8] A. Bellet, A. Habrard, and M. Sebhan. A survey on metric learning for feature vectors and structured data, arXiv preprint arXiv: 2013.
[9] Bromley, I. Guyon, Y. LeCun, E. Säckinger, and R. Shah. Signature verification using a" siamese" time delay neural network. Advances in neural information processing systems, 6: 737-744, 1993.
[10] T. Chen, S. Korblith, M. Norouzi, and G. Hinton. simple framework for contrastive learning of visual representations. In International conference machine learning. 2020.
[11] S. Chiappa, S. Racanière, D. Wierstra, and S. Mohamed. Recurrent environment simulators, 2017.
[12] K. Chua, R. Calandra, R. McAllister, and S. Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing (NeurIPS), 2018.
[13] K. Cobbe, C. Hesse, J. Hilton, and J. Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International conference on machine learning, pages 2048-2056, 2020.
[14] S. Du, A. Krishnamurthy, N. Jiang, A. Agarwal, M. Dudik, and J. Langford. Provably efficient with rich observations via latent state decoding. In International Conference on Machine Learning. PMLR, 2019.
[15] C. E. Garcia, D. M. Prett, and M. Morari. Model predictive control: Theory and practice survey. Automatica, 25(3):335-348, 1989.
[16] C. Gelada, S. Kumar, J. Buckman, O. Nachum, and M. G. Bellemare. DeepMDP: Learning con- tinuous latent space models for representation learning. In Proceedings of the 36th International Conference on Machine Learning, volume 97, pages 2170-2179, 2019.
[17] D. Ha and J. Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems, 2018.
[18] D. Hafner, T. Lillicrap, Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. Proceedings of the 36th International Conference on Machine Learning, volume 97, pages 2555-2565, 2019.
[19] D. Hafner, T. Lillicrap, Ba, and M. Norouzi. Dream control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
[20] D. Hafner, T. P. Lillicrap, M. Norouzi, and J. Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021.
[21] K. He, X. Zhang, S. Ren, and Sun. Deep residual learning for image recognition. CoRR. abs/15 151.0058,,,3 2015.
[22] E. Hoffer and N. Ailon. Deep metric learning using triplet network. In International workshop on similarity-bor based pattern recognition, pages 84-92. 2015.
[23] Ichter and M. Pavone. Robot motion planning in learned latent spaces. IEEE Robotics and Automation Letters, 4(3):2407-2414, 2019.
[24] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: /ttp://42.6690 2.990, 2014.
[25] Kipf, F. van der Pol, and M. Welling. Contrastive learning of structured world models. In International Conference Learning Representations, 2020.
11

---

[26] R. Kumar, A. Mandalika, S. Choudbury, and S. Srinivasa. Lego: Leveraging experience in roadmap generation for sampling-bast planning. In 2019 IEEE/RS/ International Conference on Intelligent Robots and Systems (IROS), pages 1488-1495, 2019.
[27] Y.-L. Kuo, A. Barbu, and B. Katz. Deep sequential models for sampling-based planning. In 2018 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages IEEE, 2018.
[28] Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. Gradient- brase learning applied document recognition. Proceedings of the IEEE, 86(11):2278-2324, 1998.
[29] W. Liu, Y. Wen, Z. Yu, M. Li, B. Raj, and L. Song. Sphereface: Deep hypersphere embedding for face recognition. In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 6738 6746, 2017.
[30] A. V. Nair, V. Pong, M. Dalal, S. Bahl, S. Lin, and S. Levine. Visual reinforcement learning with imagined goals. Advances Neural Information Processing Systems, volume 31, 2018.
[31] S. Nasiriany, V. Pong. S. Lin, and S. Levine. Planning with goal-conditioned policies. In Advances in Neural Information Processing Systems, volume 32, 2019.
[32] J. Oh, X. Guo, H. Lee, R. L. Lewis, and S. Singh. Action-conditional video prediction using deep networks in atari games. In Advances in Neural Information Processing Systems, 2015.
[33] M. Okada and T. Taniguchi. Dreaming: Model-based reinforcement learning by latent imagina- tion without reconstruction, 2020.
[34] K. Paster, S. Mcllraith, and J. Ba. Planning from pixels using inverse dynamics models. In International Conference on Learning Representations, 2021.
[35] M. V. Pogančié, A. Paulus, V. Musil, G. Martius, and M. Rolinek. Differentiation of blackbox combinatorial solvers. In International Conference on Learning Representations, 2020.
[36] R. Y. Rubinstein. Optimization of computer simulation models with rare events. European Journal of Operational Research, 99(1)899112,, 1997.
[37] N. Savinov, A. Dosovitskiy, and V. Koltun. Semi-parametric topological memory for navigation. In International Conference on Learning Representations, 2018.
[38] J. Schritwieser,, I. Antonoglou, T. Hubert, K. Simonyan, L. Sifre, S. Schmitt, A. Guez, E. Lock- hart, D. Hassabis, T. Graepel, al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839))004-699, 2020.
[39] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. CoRR, abs/ 707.70377777777 2017.
[40] M. Vlastelica, M. Rolínek, and G. Martius. Neuro-algorithmic policies enable fast combinatorial generalization, 2021.
[41] Wikipedia contributors. Kissing number - Wikipedia, the free encyclopedia, 2021. URL https::/e.. www.awwwiiiiiee.. wwwipedia........mmm..m......iai or.gorr////o w/ indexx.xxxx..//// phpptiitee=kissing_mmnae Bactig__nmberrree oidid=1020394961 [Online; accessed 9-M1y-2021].
[42] R. Yonetani, T. Taniai, M. Barekatain, M. Nishimura, and A. Kanezaki. Path planning using neural a search, 2020.
[43] A. Zhang, R. T. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant repre- sentations for reinforcement learning without reconstruction. In International Conference on Learning Representations, 2021.
[44] I. Z.annn G. Yang, and B. C. Stadie. World model as graph: Learning latent landmarks for planning, 2020.
12


Reviews:
[{"review_id": "lMwYgecTdM", "reviewer": "Reviewer_5kFP", "summary": "This paper introduces a planning method for solving combinatorial environments from pixels called PPGS. PPGS learns a latent space using a ResNet encoder which is trained such that a forward model and inverse model can be learned using the representation. Finally, it is constrained to encourage margins in the latent space. During goal-conditioned planning, visual inputs are mapped into the latent space, discretized, and a BFS-like search algorithm is performed. PPGS achieves strong empirical results on the Procgen maze environment and two 2D Atari-like environments which the authors newly introduce in the paper. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: The choices of components in the objective used to learn the latent space are well-motivated and seem sensible. While the idea of planning from pixels using a learned latent space is not novel, to my knowledge the training objective used in this paper contains new elements of combining forward and inverse model training, as well as a discretization step to allow planning on a graph. \n\nQuality: The method is not compared to other model-based planning methods which are goal-conditioned and do not require a reward signal, for example, SPTM (Savinov 2018). These may be expected to perform better than methods which require rewards due to generalization, better than GLAMOR by performing explicit planning, and better than “GS on images” due to being goal-directed.\n\nIn motivating the development of new environments, it’s stated “its dynamics are uniform and actions only cause local changes in the observations. Moreover, ProcgenMaze is a forgiving environments in which errors can always be recovered from. We use these insights to choose the additional environments”. While this seems true, the choice of making new evaluation environments more difficult in these axes (nonlocal changes from actions and irrecoverable errors) could be better justified for example in terms of downstream application.\n\nThe relatively weak performance of Dreamerv2 and GLAMOR is explained in the text by lack of generalization across a variety of levels. Can this be verified experimentally, e.g. by looking at the performance of these algorithms on fixed levels?\n\nHow sensitive is the method to the choice of $\\epsilon$ in Equation 4, and how important is adding the margin loss to performance? Was this value tuned prior to experiments? The ablations do not include an analysis of this aspect.\n\nClarity: It seems like PPGS can only work if planning towards states which have already been seen in the training data or which have neighboring states which have been seen in the training data. This is confusing to me because it seems that the environments are quite difficult, yet this means the random policy must achieve a reasonable success rate on them to collect data in that range. What is the success rate of the uniformly random policy on each of these three environments? Can state coverage be guaranteed for other environments?\n\nNitpicks: \nIn figure 2, there is a L_hinge which is not described in the text -- did you mean L_margin?.\n\nTypos: L202: “wepsite” -> “website”, L340: “state-of-the-art method” -> “methods”, L210 “envrionments” -> “environment”\n\nConsider citing: \n“Semi-parametric Topological Memory for Navigation”, Savinov et. al 2018\n\nThe code for the environments as well as the algorithm are provided and the authors have made efforts to provide details for reproducibility. \n\nSignificance: The empirical results are quite impressive -- the environments in the experiments are quite difficult to solve even for humans, and PPGS achieves an impressive success rate on these tasks. Furthermore, the method is demonstrated only using offline trajectories and does not use any further online rollouts. The experimental environments are a nice contribution and the demo/presentation is quite helpful for understanding their complexity.\n\nRecommendation:\n\nI believe that this work is a nice contribution towards planning from visual observations. While I have concerns about necessary conditions for environments it can perform well on, the generalization performance and state representation learning component of the work seems straightforward and effective.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I was unable to find much discussion in the paper specifically about the limitations of PPGS -- it would be great if the authors could add some discussion. However, I believe that one major limitation as mentioned earlier is that the success rate of a random policy must reach a minimum threshold to be able to build a dataset to learn world models which can be used for planning without any additional data collection.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "j9cZ7gjgoGB", "reviewer": "Reviewer_XKCd", "summary": "In this work, the authors present a method to perform planning from pixels in environments that exhibit a hard combinatorial search space. They consider procedurally generated environments that can be modeled by a goal-conditioned MDP.  They present a method that first learns to represent the environment as a latent graph, in which nodes represent latent states and the edges represent actions, and then use a graph search algorithm to plan over this graph. They learn a state encoder model, parameterized by a ResNet, to project a pair of environment's states and actions to a d-dimensional hypersphere. They also learn a transition model over the latent space as well as an inverse model that is trained to prevent the loss of the forward model to collapse to a trivial solution. The authors also introduce an extra loss that enforces an epsilon margin between encodings of non-bisimilar states to ensure that the latent exhibits a discrete structure to enable planning. Then, they use the breadth-first search algorithm to find in the latent graph, for a given episode, the sequence of actions that lead from the initial state to the goal. They also use a Model Predictive Control approach to mitigate the prediction error propagation during planning. Finally, the authors introduce three visual environments that exhibit a combinatorial hard search space. They show that their method outperforms strong baselines on these environments and obtain good generalization.", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "First I would like to thank the authors for their work. I believe that learning to plan from pixels to solve combinatorial problems is very important and that this paper steps forward in that direction. My comment on the soundness of the claims, significance and novelty of the contribution, and relevance to the NeurIPS community are below:\n\nThe paper is very clear, well written and overall easy to follow. Notably,  I really enjoyed sections 1, 2 and 3 that are well motivated and didactic.\n\nThe method is sound and the design choices are clearly explained and motivated.\n\nThe experimental section is well-structured. The environments are clearly described.\n\nThe choice of the baselines makes sense and the supplementary material justifies this choice clearly.\n\nThe method seems to work really well and these environments compared to the baselines.\n\nLimitations of this work:\n\nIn the related work, I would recommend the authors to additionally cite and compare their method to LEAP (Nasiriany et al.) that is really close to this method. LEAP would also be a good candidate for additional baselines.\n\nWhile the way the model is trained and how the planning is performed are clearly described, I think this paper lacks a section to discuss how both are intertwined. Notably, if my understanding is correct, the model is first trained on randomly sampled transitions, then frozen and then the authors perform planning on this frozen model for all their experiments. This setting is interesting as it validates the representation capabilities of the model however it is limited. For instance, in some environments, some states and thus some transitions can't be visited by a random policy. Therefore, I think that it is important to also consider settings where the model is improved iteratively on transitions found during the search.\n\nWhile these environments are complex along some axes, they remain simple along others. Notably, the dynamics seems simple enough to be learned from a few randomly sampled transitions. Therefore, I would recommend the authors test their method on other environments to strengthen their claims. For instance, the authors might consider the environments used in GLAMOR and assess if their method still outperforms GLAMOR on these.\n\nThe paper mentions several times the issue of generalization. While the authors addressed generalization on levels, it would be very interesting to address as well generalization on goals as the method allows for it. In this form, it is surprising when starting to read the experimental sections, to learn that the goals are not sampled but fixed. It would be great if the authors run additional experiments to assess generalization over goals. In this case, to ensure a FAIR comparison, they should also add the goal as an input to PPO's policy and DREAMER's model and use HER (Andrychowicz et al.) in addition.\n\nTo conclude, I think that this work is very promising and I really liked the paper. In this form, the paper is marginally above the acceptance threshold. However, I would be happy to increase my score if the authors address the limitations I mentioned.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This paper doesn't add limitations or potential negative societal impact to existing reinforcement learning and classical planning methods.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "IxhHPSFL0Re", "reviewer": "Reviewer_Ft4Z", "summary": "This paper proposes a model-based planning method for pixel-based control in difficult-to-plan environments in which the full environment is observed through a high-dimensional image. The method proceeds by proposing a particular learning objective for learning image representations that can be used for localization (\"reidentification\") and planning. After this encoder is learned, the representations are used inside a graph-based planner that observes the current image and target goal image and uses receding-horizon control (\"Full Planner\") to control the agent. Experiments on an existing set of difficult environments, as well as two newly-proposed environments that satisfy the desired properties (difficulty and observability) compare the method to model-based baselines, a model-free baseline, and a simple learning-free planning baseline. These experiments show the method outperforming the baselines in all key settings, and exhibiting better performance relative to a model-free approach in low-sample regimes.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The setting is an interesting and highly-relevant: model-based control from pixels in a particular class of difficult environments. The proposed method is straightforward (fairly intuitive) and performs well. The proposed evaluation does a reasonable job of highlighting the strengths of the method and deficiencies of existing methods in the setting. I have a few concerns about the problem setup, importance of training objective components, and experiments that I think are important (hence, my initial rating of 6 instead of 7), but I believe should be straightforward to address.\n\n**Problem setup:**\n- L60 I think the set of MDPs need the further assumptions that each $|s_i| = |s_j| \\forall i\\neq j$, and similarly for actions, otherwise the model will need to be able to handle states and actions of different dimensionalities. Furthermore, I think that the problem assumption might match that of Block MDPs [A] -- which assume that each observation contains enough information to uniquely identify the state and implies the existing of a perfect decoder. If true, this would make the problem setup in this paper clearer, in which each $z_t=f(o_t)$ is meant to approximate a true underlying $s_t$. Particularly, it's peculiar to think of the high-dimensional images as Markovian states; the block-MDP assumption fixes this peculiarity.\n- The notion of \"combinatorial hardness\" is important to the setting in which the algorithm is meant to function well, yet it is never precisely defined, which also means that the environments employed are not conclusively shown to possess combinatorial hardness. I think this needs to be formalized in order for the paper to make its claims precise enough.\n\n**Evaluation concerns:**\n- Does the margin loss actually contribute to learning a better latent space? Discretization for the graph search could still be performed using a model trained without the margin loss. It would be good to have evidence to show that the margin loss results in improved planning performance. This evaluation could be added to Table 1 in the appendix, and referenced in the main text when the margin loss is discussed. If the following ends up being true, that could look something like \"we found the margin loss to result in representations better-suited to our planning method, see Appendix Table 1.\" Similar ablations are missing for the $\\mathcal L_{\\text{FW}}$ and $\\mathcal L_{\\text{CE}}$ losses.\n- Success rate fails to disentangle the efficiency with which approaches achieve the goals. I think the paper would have a more robust evaluation if it included the success weighted by path length metric -- see e.g. Eq. 1 from [B]. For instance, it could be true that PPGS finds very efficient paths relative to the baselines, and the baselines merely do a reasonable job of efficiently exploring the space, rather than exhibiting particularly efficient directed behavior.\n- L318 \"GS on Images manages to rival other baselines is a testament to the harshness of the environments\" I'm not sure there's enough evidence to conclude this. I'd say that GS on Images is performing reasonably well (>40% success in all environments). Given that GS on Images is a \"simple\" baseline that doesn't employ learning, this suggests to me that the environments are not particularly harsh (or they are well-suited to GS on Images), and that the other baselines are merely underperforming / are not well-suited to the task.\n- If computationally feasible, it would be good to extend the X axes of Fig 7 IceSlider and ProcgenMaze (and DigitJump), in order to determine if/when the performance of PPGS and PPO overlap. The trends suggest that PPO may reach PPGS performance at ~$10^4$ training levels. Presumably they would eventually overlap. This evaluation would strengthen our understanding of the performance of the approaches as a function of the available training levels.\n- L325 \"Note that PPGS uses only 400k samples from a random policy whereas PPO uses 50M on-policy samples\" This point is too easy to miss. It's not otherwise clear that the amount of training data available to each approach is different. Isn't the 400k samples only true at one of the x-values in Fig7? The number of training samples available to each approach should be made clear in Fig 7. For example, having stacked x-axes for each method in Fig 7 could work, if the number of training examples available to PPGS and PPO are linear functions of the number of levels (current x-axis).\n\n**Minor comments:**\n- L60 There's an extra comma in the set; also, (pedantry to follow) include the tuple parens: $\\{(S, A, T, G, R, \\gamma)_i\\}_{i=1}^n$.\n- L280 \"the largest disadvantage is arguably given to our own method.\" More explanation is needed. Is this due to the reward-free learning?\n- typo: Eq 6 has an extra parenthesis after $z_h^l$.\n- L202 typo \"wepsite\"\n- Alg 1 is missing the pretrained encoder as input.\n\n[A] Provably efficient RL with Rich Observations via Latent State Decoding; ICML 2019 http://proceedings.mlr.press/v97/du19b/du19b.pdf\n\n[B] On Evaluation of Embodied Navigation Agents; https://arxiv.org/pdf/1807.06757.pdf\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper didn't really discuss broader limitations or broadly characterize failure cases -- the Checklist 1(b) answer points us to sections that narrowly reference performance issues. One of the main limitations is the observability assumption, furthermore, determinism in the dynamics seems quite limiting. These should be discussed somewhere. An inclusion of potential broader impacts was added in Checklist response 1(c), but it should be included in the appendix (?)", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Learning dynamics models over pixels for model-based RL remains a largely open problem. There is consensus among the reviewers that the empirical results in this paper are impressive, and the approach constitutes a valuable contribution. The reviewers make several suggestions for improvements, e.g., analyzing generalization to goals and clarifying necessary conditions for environments it can perform well on, that can be addressed in a revision.
Decision: Accept (Poster)