Planning from Pixels in Environments with Combinatorially Hard Search Spaces

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 . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 1 Introduction . Decision problems with an underlying combinatorial structure pose a significant challenge for a 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 a simulator or a 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 as the state space of the game grows at a 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 , heuristics 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 , 19 ] or applying heuristic search algorithms [ 38 ] , it also prevents a 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 is trained to 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 a 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 graph 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 to unseen variations of the environment , or equivalently across a distribution of levels [ 13 ] . This is in contrast with traditional benchmarks [ 7 ] , 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 a weakness of modern reinforcement learning approaches , second , we introduce a simple but principled world model architecture that can accurately learn the latent dynamics of a complex system from high dimensional observations ; third , we show how a planning module can simultaneously estimate the latent graph for previously unseen environments and deploy a breadth first search in the latent space to retrieve a 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-conditioned Markov Decision Processes with discrete actions , that is the 6-tuples { ( S , A , T , G , R , γ ) i } 1 ... n where Si is the state set , Ai is the action set , Ti is a transition function Ti : Si × Ai → Si , Gi is the goal set and Ri is a reward function Ri : Si ×Gi → R and γi is the discount factor . We remark that each environment can also be modeled as a BlockMDP [ 14 ] in which the context space X corresponds to the state set Si we introduced . 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 |Si| = |Sj | and |Ai| = |Aj | for all 0 ≤ i , j ≤ n. In our work the reward simplifies to an indicator function for goal achievement R ( s , g ) = 1s=g with G ⊆ S. Given a goal distribution p ( g ) , the objective is that of finding a goal-conditioned policy πg that maximizes the return Jπ = E g∼p ( g ) [ E τ∼p ( τ |πg ) ∑ t γtR ( st , g ) ] ( 1 ) where τ ∼ p ( τ |πg ) is a trajectory ( st , at ) Tt=1 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 . In order to master them , an 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 for non-trivial time horizons . These factors contribute to making such environments very challenging for existing methods . Our method is designed in light of these necessities ; it has two integral parts , the world model and the planner , which we now introduce . 2.1 World Model . The world model relies solely on three jointly trained function approximators : an encoder , a forward model and an 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 a low-dimensional state space Z has several benefits [ 17 , 18 ] . 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 . For this purpose , our method relies on an encoder hθ , that is a neural function approximator mapping each observed state s ∈ S and a low-dimensional representation z ∈ Z ( embedding ) . While there are many suitable choices for the structure of the latent space Z , we choose to map observations to points on an d-dimensional hypersphere taking inspiration from Liu et al . [ 29 ] . 2.1.2 Forward Model . In order to plan ahead in the environment , it is crucial for an agent to estimate the transition function T . In fact , if a mapping to a low-dimensional latent space Z is available , learning directly the projected transition function TZ : Z × A → Z can be largely beneficial [ 17 , 18 ] . The deterministic latent transition function TZ can be learned by a neural function approximator fφ so that if T ( st , at ) = st+1 , then fφ ( hθ ( st ) , at ) : = fφ ( zt , at ) = hθ ( st+1 ) . We refer to this component as forward model . Intuitively , it can be trained to retrieve the representation of the state of the MDP at time t+ 1 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 zt and the action sequence ( at , . . . , at+k ) are sufficient to to predict the latent state at time t + k , as long as zt successfully captures all relevant information from the observed state st . The amount of information to be embedded in zt and to be retained in autoregressive predictions is , however , in most cases , prohibitive . Take for example the case of a simple maze : zt would have to 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 on local information , we adopt an hybrid forward model which can recover the invariant structures in the environment from previous observations . The function that the forward model seeks to approximate can then include an additional input : fφ ( zt , at , sc ) = zt+1 , where sc ∈ 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 , as we show in the ablations in Suppl . A . Given a trajectory ( st , at ) Tt=1 , the forward model can be trained to minimize some distance measure between state embeddings ( zt+1 ) 1 ... T−1 = ( hθ ( st+1 ) ) 1 ... T−1 and one-step predictions ( fφ ( hθ ( st ) , at , sc ) ) 1 ... T−1 . In practice , we choose to minimize a Monte Carlo estimate of the expected Euclidean distance over a finite time horizon , a set of trajectories and a set of levels . When training on a distribution of levels p ( l ) , we extract K trajectories of length H from each level with a uniform random policy π and we minimize LFW = E l∼p ( l ) [ 1 H − 1 H−1∑ h=1 E ah∼π [ ‖fφ ( zlh , ah , slc ) − zlh+1‖22 ] ] ( 2 ) 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 to a single point in the latent space . As embeddings of any pair of states could not be distinguished in this case , this is not a desirable solution . We remark that this is a known problem in metric learning and image retrieval [ 8 ] , for which solutions ranging from siamese networks [ 9 ] to using a triplet loss [ 22 ] have been proposed . The context of latent world models offers a natural solution that isn ’ t available in the general embedding problem , which consists in additionally training a probabilistic inverse model pω ( at | zt , zt+1 ) such that if TZ ( zt , at ) = zt+1 , then pω ( at | zt , zt+1 ) > pω ( ak | zt , zt+1 ) ∀ak 6= at ∈ A . The inverse model , parameterized by ω , can be trained to predict the action at that causes the latent transition between two embeddings zt , zt+1 by minimizing multi-class cross entropy . LCE = E l∼p ( l ) [ 1 H − 1 H−1∑ h=1 E ah∼π [ − log pω ( ah | zlh , zlh+1 ) ] ] . ( 3 ) 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 , it is empirically observed to enforce a 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 a similar role to the reconstruction loss in Hafner et al . [ 18 ] . However , LCE does not force the encoder network to embed information that helps with reconstructing irrelevant parts of the observation , unlike training methods relying on image reconstruction [ 11 , 17–20 ] . 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 . A straightforward way to enforce this is by introducing a margin ε , representing a desirable minimum distance between embeddings of non-bisimilar states [ 43 ] . A third and final loss term can then be introduced to encourage margins in the latent space : Lmargin = E l∼p ( l ) [ 1 H − 1 H−1∑ h=1 max ( 0 , 1− ‖zlh+1 − zlh‖22 ε2 ) ] . ( 4 ) We then propose to reidentify two embeddings as representing the same true state if their Euclidean distance is less than ε2 . Adopting a latent margin effectively constrains the number of margin-separated states that can be represented on an hyperspherical latent space . However , this quantity is lower-bounded by the kissing number [ 41 ] , 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 a combination of the three loss functions : L = αLFW + βLCE + Lmargin . ( 5 ) To summarize , the three components are respectively encouraging accurate dynamics predictions , regularizing latent representations and enforcing a discrete structure for state reidentification .

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.