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 presents a method for learning a sparse set of latent subgoal states during training. Using a goal-conditioned policy and the latent states, a simple planning algorithm that performs soft value iteration between the subgoal states is proposed to facilitate within-dataset generalization. The proposed method outperforms competing approaches on multiple simulated navigation and robotics tasks.