World Model as a Graph: Learning Latent Landmarks for Planning

1 INTRODUCTION . An intelligent agent should be able to solve difficult problems by breaking them down into sequences of simpler problems . Classically , planning algorithms have been the tool of choice for endowing AI agents with the ability to reason over complex long-horizon problems ( Doran & Michie , 1966 ; Hart et al. , 1968 ) . Recent years have seen an uptick in monographs examining the intersection of classical planning techniques – which excel at temporal abstraction – with deep reinforcement learning ( RL ) algorithms – which excel at state abstraction . Perhaps the ripest fruit born of this relationship is the AlphaGo algorithm , wherein a model free policy is combined with a MCTS ( Coulom , 2006 ) planning algorithm to achieve superhuman performance on the game of Go ( Silver et al. , 2016a ) . In the field of robotics , progress on combining planning and reinforcement learning has been somewhat less rapid , although still resolute . Indeed , the laws of physics in the real world are infinitely more complex than the simple rules of Go . Unlike board games such as chess and Go , which have deterministic and known dynamics and discrete action space , robots have to deal with a probabilistic and unpredictable world , and the action space for robots is oftentimes continuous . As a result , planning in robotics presents a much harder problem . One general class of methods ( Sutton , 1991 ) seeks to combine model-based planning and deep RL . These methods can be thought of as an extension of model-predictive control ( MPC ) algorithms , with the key difference being that the agent is trained over hypothetical experience in addition to the actually collected experience . The primary shortcoming of this class of methods is that , like MCTS in AlphaGo , they resort to planning with action sequences – forcing the robot to plan for each action at every hundred milliseconds . Planning on the level of action sequences is fundamentally bottlenecked by the accuracy of the learned dynamics model and the horizon of a task , as the learned world model quickly diverges over a long horizon . This limitation shows that world models in the traditional Model-based RL ( MBRL ) setting often fail to deliver the promise of planning . Another general class of methods , Hierarchical RL ( HRL ) , introduces a higher-level learner to address the problem of planning ( Dayan & Hinton , 1993 ; Vezhnevets et al. , 2017 ; Nachum et al. , 2018 ) . In this scenario , a goal-based RL agent serves as the worker , and a manager learns what sequences of goals it must set for the worker to achieve a complex task . While this is apparently a sound solution to the problem of planning , hierarchical learners neither explicitly learn a higher-level model of the world nor take advantage of the graph structure inherent to the problem of search . To better combine classical planning and reinforcement learning , we propose to learn graph-structured world models composed of sparse multi-step transitions . To model the world as a graph , we borrow a concept from the navigation literature – the idea of landmarks ( Wang et al. , 2008 ) . Landmarks are essentially states that an agent can navigate between in order to complete tasks . However , rather than simply using previously seen states as landmarks , as is traditionally done , we will instead develop a novel algorithm to learn the landmarks used for planning . Our key insight is that by mapping previously achieved goals into a latent space that captures the temporal distance between goals , we can perform clustering in the latent space to group together goals that are easily reachable from one another . Subsequently , we can then decode the latent centroids to obtain a set of goals scattered ( in terms of reachability ) across the goal space . Since our learned landmarks are obtained from latent clustering , we call them latent landmarks . The chief algorithmic contribution of this paper is a new method for planning over learned latent landmarks for high-dimensional continuous control domains , which we name Learning Latent Landmarks for Planning ( L3P ) . The idea of reducing planning in RL to a graph search problem has enjoyed some attention recently ( Savinov et al. , 2018a ; Eysenbach et al. , 2019 ; Huang et al. , 2019 ; Liu et al. , 2019 ; Yang et al. , 2020 ; Laskin et al. , 2020 ) . A key difference between those works and L3P is that our use of latent landmarks allows us to substantially reduce the size of the search space . What ’ s more , we make improvements to the graph search module and the online planning algorithm to improve the robustness and sample efficiency of our method . As a result of those decisions , our algorithm is able to achieve superior performance on a variety of robotics domains involving both navigation and manipulation . In addition to the results presented in Section 5 , videos of our algorithm ’ s performance , and an analysis of the sub-tasks discovered by the latent landmarks , may be found at https : //sites.google.com/view/latent-landmarks/ . 2 RELATED WORKS . The problem of learning landmarks to aid in robotics problems has a long and rich history ( Gillner & Mallot , 1998 ; Wang & Spelke , 2002 ; Wang et al. , 2008 ) . Prior art has been deeply rooted in the classical planning literature . For example , traditional methods would utilize Dijkstra et al . ( 1959 ) to plan over generated waypoints , SLAM ( Durrant-Whyte & Bailey , 2006 ) to simultaneously integrate mapping , or the RRT algorithm ( LaValle , 1998 ) for explicit path planning . The A * algorithm ( Hart et al. , 1968 ) further improved the computational efficiency of Dijkstra . Those types of methods often heavily rely on a hand-crafted configuration space that provides prior knowledge . Planning is intimately related to model-based RL ( MBRL ) , as the core ideas underlying learned models and planners can enjoy considerable overlap . Perhaps the most clear instance of this overlap is Model Predictive Control ( MPC ) , and the related Dyna algorithm ( Sutton , 1991 ) . When combined with modern techniques ( Kurutach et al. , 2018 ; Luo et al. , 2018 ; Nagabandi et al. , 2018 ; Ha & Schmidhuber , 2018 ; Hafner et al. , 2019 ; Wang & Ba , 2019 ; Janner et al. , 2019 ) , MBRL is able to achieve some level of success . Corneil et al . ( 2018 ) and Hafner et al . ( 2020 ) also learn a discrete latent representation of the environment in the MBRL framework . As discussed in the introduction , planning on action sequences will fundamentally struggle to scale in robotics . Our method will make extensive use of a parametric goal-based RL agent to accomplish low-level navigation between states . This area has seen rapid progress recently , largely stemming from the success of Hindsight Experience Replay ( HER ) ( Andrychowicz et al. , 2017 ) . Several improvements to HER augment the goal relabeling and sampling strategies to improve performance ( Nair et al. , 2018 ; Pong et al. , 2018 ; 2019 ; Zhao et al. , 2019 ; Pitis et al. , 2020 ) . There have also been attempts at incorporating search as inductive biases within the value function ( Silver et al. , 2016b ; Tamar et al. , 2016 ; Farquhar et al. , 2017 ; Racanière et al. , 2017 ; Lee et al. , 2018 ; Srinivas et al. , 2018 ) . The focus of this line of work is to improve the low-level policy and is thus orthogonal to our work . Recent work in Hierarchical RL ( HRL ) builds upon goal-based RL by learning a high-level parametric manager that feeds goals to the low-level goal-based agent ( Dayan & Hinton , 1993 ; Vezhnevets et al. , 2017 ; Nachum et al. , 2018 ) . This can be viewed as a parametric alternative to classical planning , as discussed in the introduction . Recently , Jurgenson et al . ( 2020 ) ; Pertsch et al . ( 2020 ) have derived HRL methods that are intimately tied to tree search algorithms . These papers are further connected to a recent trend in the literature wherein classical search methods are combined with parametric control ( Savinov et al. , 2018a ; Eysenbach et al. , 2019 ; Huang et al. , 2019 ; Liu et al. , 2019 ; Yang et al. , 2020 ; Laskin et al. , 2020 ) . Several of these articles will be discussed throughout this paper . LEAP ( Nasiriany et al. , 2019 ) also considers the problem of proposing sub-goals for a goal-conditioned agent : it uses a VAE ( Kingma & Welling , 2013 ) and does CEM on the prior distribution to form the landmarks . Our method constrains the latent space with temporal reachability between goals , a concept previously explored in Savinov et al . ( 2018b ) , and uses latent clustering and graph search rather than sampling-based methods to learn and propose sub-goals . 3 BACKGROUND . We consider the problem of Multi-Goal RL under a Markov Decision Process ( MDP ) that is parameterized by ( S , A , P , G , Ψ , R , ρ0 ) . S and A are the state and action space . The probability distribution of the initial states is given by ρ0 ( s ) , andP ( s′|s , a ) is the transition probability . Ψ : S 7→ G is a mapping from the state space to the goal space , which assumes that every state s can be mapped to a corresponding achieved goal g. The reward function R can be defined as R ( s , a , s′ , g ) = −1 { Ψ ( s′ ) 6= g } . We further assume that each episode has a fixed horizon T . The goal-conditioned policy is a probability distribution π : S ×G× A→ R+ . The policy gives rise to trajectory samples of the form τ = { s0 , a0 , g , s1 , · · · sT } . The purpose of the policy π is to learn how to reach the goals drawn from the goal distribution pg , which means maximizing the cumulative rewards . Together with a discount factor γ ∈ ( 0 , 1 ) , the objective is to maximize J ( π ) = Eg∼pg , τ∼π ( g ) [ ∑T−1 t=0 γ t · R ( st , at , st+1 , g ) ] . Q-learning provides a sample-efficient way to optimize the above objective by utilizing off-policy data stored in a replay buffer B. Q ( s , a , g ) estimates the reward-to-go under the current policy π conditioned upon the given goal . An additional technique , called Hindsight Experience Replay , or HER ( Andrychowicz et al. , 2017 ) , uses hindsight relabelling to drastically speed up training . This relabeling crucially relies upon the mapping Ψ : S 7→ G in the multi-goal MDP setting . We can write the the joint objective of multi-goal Q-learning with HER as minimizing : min Q E τ∼B , t∼ { 0···T−1 } ( st , at , st+1 ) ∼τ k∼ { t+1···T } , g=Ψ ( sk ) a′∼π ( ·|st+1 , g ) ( Q ( st , at , g ) − ( R ( st , at , st+1 , g ) + γ ·Q ( st+1 , a′ , g ) ) ) 2 ( 1 ) 4 THE L3P ALGORITHM Our overall objective in this section is to derive an algorithm that learns a small number of landmarks scattered across goal space in terms of reachability and use those learned landmarks for planning . There are three chief difficulties we must overcome when considering such an algorithm . First , how can we group together goals that are easily reachable from one another ? The answer is to embed goals into a latent space , where the latent representation captures some notion of temporal distance between goals – in the sense that goals that would take many timesteps to navigate between are further apart in latent space . Second , we need to find a way to learn a sparse set of landmarks used for planning . Our method performs clustering on the constrained latent space , and decodes the learned centroids as the landmarks we seek . Finally , we need to develop a non-parametric planning algorithm responsible for selecting sequences of landmarks the agent must traverse to accomplish its high-level goal . The proposed online planning algorithm is simple , scalable , and robust .

Summary: This paper proposes an approach for automatically learning state abstraction on a RL problem, which can then be used for temporally extended planning using a search algorithm. The main contribution is introducing the concept of latent *landmarks*, a clustering of low dimensional state embeddings. Landmarks are defined on a latent space wherein the distance between two latents code is small if their corresponding high-dimensional states can reach each other in few environment steps. The paper proposes to cluster latent states, so that each cluster contains states that are easy to reach from each other; the landmarks correspond to the centers of these clusters. An algorithm for automatically learning this clustering is introduced in the paper.  With landmarks in hand, the paper proposes a soft-value iteration method to compute shortest path distances to the problem's goal. The approach is evaluated in a variety of domains, and compares favorably with recent state-of-the-art methods.