World Model as a Graph: Learning Latent Landmarks for Planning
Authors: Lunjun Zhang
Ge Yang
Bradly C Stadie
Venue:  (2021)
Keywords: Reinforcement Learning
Planning

Abstract: Planning, the ability to analyze the structure of a problem in the large and decompose it into interrelated subproblems, is a hallmark of human intelligence. While deep reinforcement learning (RL) has shown great promise for solving relatively straightforward control tasks, it remains an open problem how to best incorporate planning into existing deep RL paradigms to handle increasingly complex environments. One prominent framework, Model-Based RL, learns a world model and plans using step-by-step virtual rollouts. This type of world model quickly diverges from reality when the planning horizon increases, thus struggling at long-horizon planning. How can we learn world models that endow agents with the ability to do temporally extended reasoning? In this work, we propose to learn graph-structured world models composed of sparse, multi-step transitions. We devise a novel algorithm to learn latent landmarks that are scattered (in terms of reachability) across the goal space as the nodes on the graph. In this same graph, the edges are the reachability estimates distilled from Q-functions. On a variety of high-dimensional continuous control tasks ranging from robotic manipulation to navigation, we demonstrate that our method, named L^{3}P, significantly outperforms prior work, and is oftentimes the only method capable of leveraging both the robustness of model-free RL and generalization of graph-search algorithms. We believe our work is an important step towards scalable planning in reinforcement learning. 

Under review as a conference paper at ICLR 2021
WORLD MODEL AS A GRAPH: LEARNING LATENT LANDMARKS FOR PLANNING
Anonymous authors Paper under double-blind review
ABSTRACT
Planning. the ability to analyze the structure of a problem in the large and decom- pose it into interrelated subproblems, is a hallmark of human intelligence. While deep reinforcement learning (RI.) has shown great promise for solving relatively straightforward control tasks, it remains an open problem how to best incorporate planning into existing deep RL paradigms to handle increasingly complex environ- ments. One prominent framework, Model-Based RL, learns world model and plans using step-by-step virtual rollouts. This type world model quickly diverges from reality when the planning horizon increases, thus struggling at long horizon planning. How can we learn world models that endow agents with the ability to do temporally extended reasoning? In this work, we propose to learn graph-structured world models composed of sparse, multi-step transitions. We devise a novel algo- rithm to learn latent landmarks that are scattered (in terms of reachability) across the goal space as the nodes on the graph. In this same graph, the edges are the reachability estimates distilled from Q-functions. On variety of high-dimensional continuous control tasks ranging from robotic manipulation to navigation, we demonstrate that our method, named P significantly outperforms prior work, and is oftentimes the only method capable of leveraging both the robustness of model-frea RL. and generalization of graph graph-search algorithms. We believe our work is important step towards scalable planning reinforcement learning.
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 [966 Hart
let 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) algorithmss- 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 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 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 probabilistic and unpredictable world, and the action space for robots is oftentimes continuous. As a result, planning 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-prediccive 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 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 deliver the promise of planning.

---

Under review as a conference paper at ICLR 2021
Learning
Planning
MBRL
Haplay
P(a'l is a)
Step-by-step Rollouts
Reachabiltyy
Edges
World Model
as Graph
Repley
Nodes
Sparse Multi-step Transitions
Figure 1: MBRL versus L (World Model as a Graph). MBRL does step-by-step virtual rollouts with the world model and quickly diverges from reality when the planning horizon increases. 1, p models the world as a graph of sparse multi-step transitions, where the nodes are learned latent landmarks and the edges reachability estimates. L succeeds at temporally extended reasoning.
Another general class of methods, Hierarchical RL (HRL), introduces higher-level learner to address the problem of planning (Dayan & Hinton 1993 Vezhnevets et al. [2117] Nachuum et al.) (0018). In this scenario, goal-based RL agent serves the worker, and manager learns what sequences of goals it must set for the worker to achieve a complex task. While this is apparently a sound solution the problem planning, hierarchical learners neither explicitly learn 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 learn graph-structured world models composed of sparse multi-step transitions. To model the world as graph, we borrow a concept from the navigation literature the idea of landmarks (Wang 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 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 perform clustering in the latent space group together goals that 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 this paper is a new method for planning over learned latent landmarks for nigh-dimensiooal continuous control domains, which we name Learning Latent Landmarks for Planning (L,
The idea of reducing planning in RL. graph search problem has enjoyed some attention recently (Savinov et al. 20 Sal Eysenbach et al. 2019) Huang et al. 2019: Liu et al. 2019 Yang et al. 2020 Laskin al. 2020). A key difference between those works and L is that our use of latent landmarks allows us to substantially reduce the size of the search space. What's ''s more, we make improvements to the graph search module and the online planning algorithm
improve the robustness and sample efficiency of our method. As result of those decisions, our algorithm is able to achieve superior performance on variety robotics domains involving both navigation and manipulation. In addition to the results presented in Section 5 videos of our algorithm's 'ss performance, and an analysis of the sub-tasks discovered by the latent landmarks, may be found at www wetme/viewwiiiie wwww/wateettttt.....aaa landmarks/
RELATED WORKS
The problem of learning landmarks aid in robotics problems has 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, the RRT algorithm Valle 1998) for explicit path planning. The A algorithm (Hart
2

---

Under review as a conference paper at ICLR 2021
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 Hafnnr et al.] 2019: Wang & Ba 2019 Janner al. 2019), MBRL able to achieve some level of success. [Coneil 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 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. [00 1bb) Tamar et al., 2016; Farquhar et al. 2017 Racanière et al.) 2017 Lee et al. (018/Srimivas et al. (0018). The focus of this line of work is to improve the low-level policy and is thus orthogonal
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 Recennty,,uuggenoon et al. (2020); Pertsch 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. 201 Sa, Eysenbach et al. 2019 Huang et al. 2019 Liu et al. 2019 Yang et al. 2020 Laskin al. (2220). Several of these articles will be discussed throughout this paper. LEAP Nasiriaay et al. (219)) also considers the problem of proposing sub-goals for goal-conditioned agent: 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, concept previously explored in Savinov et al. (2018b), and uses latent clustering and graph search rather than sampling-baas 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 parameter- ized by (S, A, P, G, V, R, po). S and A are the state and action space. The probability distribution of the initial states is given by Po(s), and P(s'|s,, a) is the transition probability. v ; S + G is a mapping from the state space to the goal space, which assumes that every state s can be mapped sponding achieved goal g. The reward function R can be defined as R(s, a's g) = 1 (s')) * g). We further assume that each episode has a fixed horizon? T.
corre-
The goal-conditioned policy is a probability distribution rise to trajectory samples of the form t {So, an, 9, 81, ... 8T 8T The purpose of the policy v 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 y E (0, 1). the objective is to maximize [TT) = (gg-,)) E y² 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, 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 V S + G in the multi-goal MDP setting. We can write the the joint objective of multi-goal Q-learning with HER as minimizing:
S G A IR+ The policy gives
(Q(s) R(ss, at S41,9) + + Q( Q(s .0)) (1)
min E Tw (0..TT-- (se.0e your (t-(++1 (.. -V a' ((9-1()))-1110))
3

---

Under review as a conference paper at ICLR 2021
lunzionn
Reachabilty
edpes
Decoder
Planner
Oools
Enooder
Latent Clustering
Latent Landmarks
World Model
Figure 2: An overview of P, which learns a small number of latent landmarks for planning. The main components of our method are: learning reachability estimates (via Q-learning and regression), learning latent space (via an auto-encoder with reachability constraints), learning latent landmarks (via clustering in the latent space), graph search on the world model and online planning.
THE L ALGORITHM
Our overall objective in this section is derive an algorithm that learns 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 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 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 accomplish its high-level goal. The proposed online planning algorithm is simple, scalable, and robust.
the constrained latent space, and decodes the learned
4.1 LEARNING A LATENT SPACE
Let us consider the following question: where the metric reflects reachability?" Suppose we have an auto-encoder (AE) in the agent's 's goal space, with deterministic encoder fE and decoder fn- As usual, the reconstruction loss given by
"How should we go about learning a latent space of goals
(g) = fD - =l|. We want to make sure that the distance between two latent codes would roughly correspond to the number of steps would take the policy to go from one goal to another. Concretely, for any pair of goals (g1, 92), optimize the following loss:
2
/ 1 (V 92) ((92.91)) Liatent (01, (01,92) [IRR(()) -
2
Where V: G RR+ is mapping that estimates how many steps it would take the policy go from one goal to another goal on average. By adding this constraint and solving a joint optimization Cree A- Liatent, the encoding-decoding mapping can no longer be arbitrary, giving more structure to the latent space. Goals that are close by in terms of reachability will be naturally clustered in the latent space, and interpolations between latent codes will lead to meaningful results.
Of course, the constraint in Equation 2 quite meaningless if we do not have a way to estimate the mapping V. We will proceed towards this objective by noting the following interesting connection between multi-goa Q-functions and reachability. In the multi-goal RI. framework considered in the background section,
reward is binary nature. The agent receives reward of until reaches the goal, and then () when it reaches the desired goal. In this setting, the Q- Q-function is implicitly estimating the number of steps takes to reach the goal g from the current state after the action a is taken. Denote this quantity D(s, a, g), the Q-function can be ce-written as:
D(s.a.g)-1
T-1
1 1 D(sa.a)
Σ
(-1)+
Σ D(.0..))))
(3)
1
Choosing to parameterize Q-functions in this way disentangles the effect of 7 on multi-gon -oal Q- learning. It also provides us with access the direct distance estimation function D(s a, g). We note
4

---

Under review as a conference paper at ICLR 2021
that this distance is not mathematical distance in the sense of metric. Instead, we use the word distance to refer to the number of steps the policy т needs to take in the environment.
Given our tractable estimate of D. it is now straightforward matter to estimate the desired quantity V. which approximates how many steps it takes the policy to transition between goals. To get the desired estimate, we regress V towards D follows
D(s as v((s)) - V (V ((V(S4+)) )
2
min EN B, B..- V
(4)
k-(++1..TT)
where V is given by the environment to map the states to the goal space. One erucial detail is the use of (s++1) rather than ((s) in the inputs to V. This is due to the fact that D S x A x G R outputs the number of steps to go after an action is taken, when the state has transitioned into 8-+1- The objective above provides unbiased estimate of the average number of steps between two goals.
The estimates D and V will prove useful beyond helping to optimize the auto-encoder in Equation 2 They will prove essential in weighting and planning over latent landmark nodes in Section 4.3.
4.2 LEARNING LATENT LANDMARKS
Planning on a graph be expensive. the number of edges can grow quadratically with the number of nodes. To battle this issue in scalability, we use the constrained latent space to learn a sparse set of landmarks. A landmark can be thought of as waypoint that the agent can pass through enroute to achieve a desired goal. Ideally, goals that are easily reachable from one another should be grouped to form one single landmark. Since our latent representation captures the temporal reachability between goals, this can be achieved by doing clustering in the latent space. The cluster centroids, when decoded from the decoder, will be precisely the latent landmarks we are seeking.
Clustering proceeds as follows. For N clusters to be learned, we define a mixture of Gaussians in the latent space with N trainable latent centroids, {ci... CN J. and shared trainable variance vector o. We maximize the evidence lower bound (ELBO) with uniform prior p(c):
log = fE(S)) 2 BB((e=)) logp(z
c)] DxL a( 2) l1 p(e))
(5)
Ideally, we would like each batch of data given to the latent clustering model to be representative of the whole replay buffer, such that the centroids will quickly learn to scatter out. To this end, we propose to use the Greedy Latent Sparsification (GLS) algorithm (Algorithm 2 in the Appendix) on each batch of data sampled from the replay before taking gradient step with the batch. GLS is inspired by kmeans++ (Arthur & Vassilvitskii 2007). with several key differences: this sparsification process is used for both training and initialization, it uses neural metric for determining the distance between data points, and that it is compatible with mini-batch-style gradient-bast training.
4.3 PLANNING WITH LATENT LANDMARKS
Having derived a latent encoding algorithm and an algorithm for learning latent landmarks, we at last turn our attention planning. While prior works simply solve for the shortest path, employ soft version of the Floyd algorithm, where the soft relaxation operations can be seen as a soft value iteration procedure (see Equation 7 in the Appendix).
To construct a weight matrix that at first provides a raw distance estimate between latent land- marks, we begin by decoding the learned centroids in the latent space into the nodes in the graph (fD(cr).... f((x)))), To build the graph, we add two edges directed in reverse orders for every pair of latent landmarks. For instance, for an edge going from (D(c, to fD(c,), the weight on that edge is -V (f,)),,,, fn) (c,)). Notice that the distances are negated to be negative. At the start of an episode, the agent receives a goal g, and we construct the following matrix of size (N + 1) x (N + 1):
((fD(C)).fo(eM)))
-V(fo(ct).))
W = - -V(folew), fo(ci)
(6)
-V(fo(eN).9)
0
5

---

Under review as a conference paper at ICLR 2021
TRAIN
TEST
Figure 3: For both Point and Ant, during training, the initialization state distribution and the goal proposal distribution are uniform around the maze. During test time, the agent asked to traverse the longest path in the maze. The success environment demonstrates L Ps ability to generalize to longer horizon goals during test time.
the test environment is reported in Figure 4 This
For online planning, when the agent receives goal at the start of an episode, we use graph search to solve for de-+g (which is fixed through- out episode). For an observation state the algorithm calculates d.-tc:
Algorithm 1 Online Planning in L
Given: Environment env, initial state s, goal g.
1: Cnt = 0. SubG = None.
2: Solve for de-+y with graph search.
3: for tort=1 to 7 do
b One episode
4: if Cnt 2 1.0 then
- ((, ((s fo(ee))))))
5:
Cnt = Cnt -1
6: else We do not re-plan at every step
ds-sc=
D( (s T ((s.:(s..o((cx)
7:
Calculate dg-re
-D(s, 9),
8:
t do-t + de-ig if subG * None then
9:
The chosen landmark is subgoal To further provide tem- poral abstraction and robustness, the agent will be asked to consistently pursue subgoa for ds-eesubboall] number of steps, which is how many steps it thinks it will need. The pro- posed goal does not change during this period.
10:
d[subG] + -0
arg
11: end if
Remove the immediate
previous landmark
SubG, Cnt + arg max d, max d end if
13:
(Ts, SubG): env. env. step (a). end for
15:
The algorithm makes sure that the agent does not re-plan at every step, and this mechanism for temporal abstraction is crucial its robustness. After this many steps, agent will decide on the next landmark to pursue by re-calculating d.-c> but the immediate previous landmark will not be considered as a candidate landmark. The reason that, the agent has failed reach a self-proposed landmark within the reachability limit it has set for itself, then the agent should try something new for the immediate next goal rather than stick to the immediate previous landmark for another round. We have found that this simple algorithm helps the agent avoid getting stuck and improves the overall robustness of the agent.
In summary, have derived algorithm that learns sparse set of latent landmarks scattered across goal space in terms of reachability, and uses those learned landmarks for robust planning.
EXPERIMENTS AND EVALUATION
We investigate the impact of p in variety of robotic manipulation and navigation envi- ronments. These include standard benchmarks such as Fetch-PickAndPlace,, and more diffi- cult environments such AntMaze- and Prace-Inside-Bor that have been engineered to require test-time generalization. Videos of our algorithm in action are available here: https? 7/sites. www.looooe....oooooooooo. wtt.:///w.oooiee///aeeeettt lannttlladmarrs/////
5.1 BASELINES
We compare method with variety of baselines. HER (Andrychowiez et al. 2017) model-free RL algorithm. SORB (Eysenbach et al. 2019) is a method that combines RL and graph search by using the entire replay buffer. Mapping State Space (MSS Huang et al. 2019) reduces the number

---

Under review as a conference paper at ICLR 2021
PointMaze Hard
AntMaze Hard
Fetch PickAnd Place
Test Success Palo
Tox-aside-PickAndPleee
Place-Inside-Bos
L
MSS
SORB
Figure 4: Test time success rate vs. total number of timesteps, on maze and robotic manipulation environments. During test time, new more difficult goals are selected. L shows more robust generalization much more quickly than other methods. For every environment except PointmMaze, L the only algorithm that consistently solves task.
Figure 5: Visualizing planning AntMaze test time. Read images from upper left bottom right. The blue dots are the learned latent landmarks decoded from the latent centroids. The orange dot represents the ant's present location in the maze. The red dot is the final goal that the agent needs reach. At each step, the blue star indicates the landmark chosen by our planning algorithm. Whereas MSS and SORB sample 400 and hundreds of thousands of landmarks (respectively), our method obtains lean graph that only contain 50 landmarks. L is the only method capable of achieving over 80% test success rate on AntMaze-Hann within 3M timesteps.
Figure 6: We consider two environments involving a fetch robot, a block, and a box. In Box-aside- PinkAnd,,ace, the fetch must learn to pick and place the block while avoiding collision with the box. In Place-Inside-Box. the fetch must pick the block and place it inside the box. We visualize the fetch states corresponding to learned landmarks in the second row of images.
7

---

The first question we try to understand is whether P is robust to the number of latent landmarks. In contrast to prior methods, L P Under review as a conference paper at ICLR 2021
of vertices by sub-sampling the replay buffer. LAP, SORB, and MSS all use the same hindsight relabelling strategy proposed in HER. All of the domains are continuous control tasks, so we adopt DDPG (Lillicrap et al. 2013) the learning algorithm for the low-level actor.
5.2 GENERALIZATION TO LONGER HORIZONS
The PointMaze- Hard and AntMaze- Haadd environments introduced in Figure are designed to test an agent's ability generalize longer horizons. While PointMaze and AntMaze have been previously used in Duan et al. 2016); Huang et al. 20 Pitis et al. (2020), we make slight changes to those environments in order to increase their difficulty. We use short, 200-timestep time horizon during training and a Po that is uniform in the maze. At test time, we always initialize the agent on one end of the maze, and set the goal on the other end. The horizon of the test environment is 500 steps. Crucially, no prior knowledge on the shape of the maze is given to the agent. We also set much stricter threshold for determining whether an agent has reached the goal. In Figure we see L is the only algorithm capable of solving AntMaze-Hand consistently.
are
We observe an interesting trend where the success rates for other graph search methods crash and then slowly recover after making some initial progress. We postulate this occurs because methods that are based on using the entire replay or sub-sampling the replay for landmark selection will struggle as the buffer size increases. In contrast to these methods, L P does not exhibit such undesirable instability. The online planning algorithm in L which effectively leverages temporal abstraction to improve robustness, also contributes to the asymptotic success rate. The result convincingly shows that, at least on the navigation tasks considered, 1,3 P is most effective at taking advantage of the problem's 's inherent graph structure, and that learning latent landmarks is significantly more sample efficient and scalable than directly using or sub-sampling the replay buffer to build the graph.
5.3 ROBOTIC MANIPULATION TASKS
We also benchmark challenging robotic manipulations tasks with Fetch robot introduced in |Plappert iet al. a (2018); Andrychowicz et al. (2017). In Figure we introduce two pick and place tasks involving a box on a table. In the Iaccee-nsiie--Bee Box environment, we design simple curriculum cope with the difficulty of the task. During training, the goal distribution has 80% regular pick-and-place goals, enabling the agent to first learn how to fetch in general. Meanwhile, only 20% of the goals are inside the box, which is the harder part of the task. During testing, we evaluate the ability of the agent to pick the object from the table and place inside the box. Our method achieves dominant performance in both learning speed and test-time generalization. We note that on those manipulation tasks considered, many prior planning methods hurt the performance the model-free agent. Our method is the only one that is able to help the model-free agent learn faster and generalize better.
5.4 UNDERSTANDING MODEL CHOICES IN L
We investigate 3.3 P.' sensitivity to different peper-parameters and design choices via a set of ablation studies. More specifically, we study how the following factors affect the performance of 3 P: number of latent landmarks, the choice of (online) planning algorithms, the choice of graph search algorithms, and edge cutoff threshold in graph search (a key syper-paramete in the search module).
Ablation: Number of Landmarks
AntMaze-Hard
is able to learn the landmarks used for graph search from the agent's 's own experience. We vary the number of learned landmarks in the challenging AntMaze-Hart environment, and find that I,3Pis is robust against a decreasing number of landmarks. This is expected, because the landmarks in the latent space of L will
N=100
Test Success Rate
N-50
N-25
try to be equally scattered across the goal space according to the reachability metric. As the
Millions of timesteps
number of landmarks decreases, the learning procedure will automatically push the landmarks to be further away from one another.
8

---

Under review as a conference paper at ICLR 2021
A key component in L is the online planning algorithm described in Algorithm [I We find this algorithm to bear special importance to the good performance of T,a P. Our planning algo- rithm in I,3 P can take advantage of the tempo- ral abstraction provided by the graph-structured world model. It does not re-plan every step, but instead uses the reachability estimates to dynamically decide when to re-plan, striking a balance between adaptability and consistency in planning. This planner also more tolerant of errors: it removes the immediate previous landmark when it re-plans, so that the agent will be less prone to getting stuck. A naive planner, on the other hand, simply re-calculates the shortest path at every step. The curve the left shows that this planning algorithm is crucial to the success of L
Ablation: Choice Planners
AntMaze-Hard
Planneen Naive Planner:
Succoss
timesteps
The particular choice of graph search seems to have small effect the stability of learning. As explained Section 4.3 and Appendix [A.2] find that while employing the Floyd algo- rithm for graph search, a soft operation for relax- ation leads to better stability during training. On the right, we show that a hard version of relax- ation helps the agent take off faster but suffers from greater instability during policy improve- ment. The likely reason that neural distance estimates are not entirely accurate, and in the presence of occasional bad edges, softmax in Equation 7 improves robustness. We therefore use soft relaxation in graph search module.
Ablation: Graph Search Algorithms PointMaze-Haa
Rlydd
(soft Floycl)
Teet Succoce
Millions of timesteps
One of the most important hyper-parameters when combining RL with graph search is d_max, the clipping threshold for the edges the graph (Savinov et al. 2018a Eysenbach et al. 2019] Huang et al. [ Laskin et al., 2020). The motivation for introducing this com- monly used hyper-parameter is two-fold. Firstly, we only trust distance estimates when they are local. Secondly, we want the graph search mod- ule to produce sub-goals that are nearby. The d_mar value determines the maximum distance for each edge on the graph and masks out longer edges. One weakness of our current approach is that it is still quite sensitive to this hyper-parameter a small change to d_max can have considerable impacts on learning. As this weakness is common to this class of approaches, we believe that further research is required to find other ways of encouraging search to be local. See Appendix A.2 A..2 for more details on implementing this clipping threshold.
Ablation: d max in Graph Search AntMaze-Hard
4. max=20 1_max=15
Test Success Rate
4. d_max=10 10
Millions of timesteps
6 CLOSING REMARKS
In this work, we introduce a way of learning graph-structured world models that endow agents with the ability to do temporally extended reasoning. The algorithm, I P, learns of latent landmarks scattered across the goal space to enable scalable planning. We demonstrate that L achieves significantly better sample efficiency, higher asymptotic performance, and better generalization on range of challenging robotic navigation and manipulation tasks. We hope that this work inspires more research in the direction of combining deep RL with classical planning.
9

---

Under review as a conference paper at ICLR 2021
REFERENCES
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay, In Advances in neural information processing systems, pp. 5048-5058, 2017.
David Arthur and Sergei Vassilvitskii. k-means++:
The advantages of careful seeding. Proceedings of the eighteenth anmual ACM-SIAM symposium on Discrete algorithms, 2007.
Dane Corneil, Wulfram Gerstner, and Johanni Brea. Efficient model-based deep reinforcement learning with variational state tabulation. arXiv preprint wrXiv:/802.04322 2018
Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In International conference on computers and games, pp. 72-83. Springer, 2006.
Peter Dayan and Geoffrey Hinton. Feudal reinforcement learning. In Advances in neural information processing systems, pp. 27 1-278, 1993.
Edsger W Dijkstra et al. A note on two problems in connexion with graphs. Numerische mathematik, 1(1):269-22 271, 1959.
James E Doran and Donald Michie. Experiments with the graph traverser program. Proceedings of the Royal Society of London. Series A. Mathematical and Physical Sciences, 294(1437):235-559, 1966.
Yan Duan, Xi Chen, Rein Houthooft, John Schulman, and Pieter Abbeel. Benchmarking deep reinforcement learning for continuous control. In International Conference on Machine Learning, pp. 1339,1338, 2016.
Hugh Durrant-Wiie Whyte and Tim Bailey. Simultaneous localization and mapping: part i. IEEE robotics & automation magazine, 13(2):99-110, 2006.
Ben Eysenbach, Russ Salakhutdinov, and Sergey Levine. Search on the replay buffer: Bridging planning and reinforcement learning. In Advances in Neural Information Processing Systems, pp.
15246-15257, 2019.
Gregory Farquhar, Tim Rocktäschel, Maximilian Igl, and Shimon Whiteson. Treeqn and atreec: Differentiable tree-structured models for deep reinforcement learning. October 2017. URL
htttp://arxi>:
11417
Sabine Gillner and Hanspeter A Mallot. Navigation and acquisition of spatial knowledge in a virtual maze. Journal of cognitive neuroscience, 10(4):445 4633 1998.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems, pp. 2450-2462, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555-2565. PMLR, 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi. and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02 193, 2020.
Peter E Hart, Nils Nilsson, and Bertram Raphael. A formal basis for the heuristic determination of minimum cost paths. IEEE transactions on Syslems Science and Cybernetics, 1220:100-107 1968.
Zhiao Huang, Fangchen Liu, and Hao Su. Mapping state space using landmarks for universal goal reaching. In Advances in Neural Information Processing Systems, pp. 1942-1952, 2019.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems, pp. 12519-12538 2019.
Tom Jurgenson, Or Avner, Edward Groshev, and Aviv Tamar. Sub -oal trees-a framework for goal-based reinforcement learning, arXiv preprint arXiv:2002. /2361, 2020.
10

---

Under review as a conference paper at ICLR 2021
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv:1312.61 2013.
arXiv preprint
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. arXiv preprint arXiv: 1802. 2018.
Michael Laskin, Scott Emmons, Ajay Jain, Thanard Kurutach, Picter Abbeel, and Deepak Pathak. Sparse graphical memory for robust planning. arXiv preprint arXiv: 064417, 2020.
Steven M Valle. Rapidly-exploring random trees: new tool for path planning. 1998.
Lisa Lee, Emilio Parisotto, Devendra Singh Chaplot, Eric Xing, and Ruslan Salakhutdinov. Gated path planning networks. June 2018. httptttp///axxi http:/awwi...oo..oo/.a/oa wtwg/aa////aaoooooooooo/eee 806 06408
Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint (wX.::150.009777000 2015.
Kara
Liu, Thanard Kurutach, Christine Kit-Ching Tung, Pieter Abbeel, and Aviv Tamar. Hallucinative topological memory for zero-shot visual plan-
ning. ArXiv, 2019. URL https: //www www. sema nt
paper/
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic
framework for model-based deep reinforcement learning with theoretical guarantees. arXiv preprint arXiv:/807.03838, 2018.
Ofir Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforcement learning. In Advances in Neural Information Processing Systems, pp. 3303-3313, 2018.
Anusha Nagabandi, Gregory Kahn, Ronald Fearing, and Sergey Levine. Neural network dynam- ics for model-based deep reinforcement learning with model-free fine-tuning. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pp. 7559-7566. IEEE, 2018
Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. In Advances Neural Information Processing Systems, pp. 9191-9200, 2018.
Soroush Nasiriany, Vitchyr Pong. Steven Lin, and Sergey Levine. Planning with goal-conditioned policies. In Advances in Neural Information Processing Systems, pp. 14843-14854, 2019.
Karl Pertsch, Oleh Rybkin, Frederik Ebert, Chelsea Finn, Dinesh Jayaraman, and Sergey Levine. Long-horizon visual planning with goal-conditioned hierarchical predictors. arXiv preprint arXiv:2006. 13205, 2020.
Silviu Pitis, Harris Chan, Stephen Zhao, Bradly Stadie, and Jimmy Ba. Maximum entropy gain exploration for long horizon multi-goal reinforcement learning, arXiv preprint arXhv:2007.02832. 2020.
Matthias Plappert, Marcin Andrychowicz, Alex Ray. Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, et Multi-goal reinforcement learn- ing: Challenging robotics environments and request for research. arXiv preprint arXiv: 1802.0640 2018.
Vitchyr Pong, Shixiang Gu, Murtaza Dalal, and Sergey Levine. Temporal difference models: Model- free deep for model-based control. arXiv preprint 1rXiv:1802.0008. 2018.
Vitchyr H Pong. Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine. Skew-fit: State-covering sepe-super reinforcement learning. arXiv preprint arXiv: a9Xiv: 2019.
11

---

Under review as a conference paper at ICLR 2021
Sébastien Racanière, Theophane Weber, David Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adrià Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, Razvan Pas- canu, Peter Battaglia, Demis Hassabis, David Silver, and Daan Wierstra. Imagination-augmennee augmented agents for deep reinforcement learning. In Guyon, U V Luxburg. S Bengio, H Wallach, R Fergus, S Vishwanathan, and R Garnett (eds.), Advances in Neural Information Processing Systems 30, pp. 5690-5701. Curran Associates, Inc., 2017. URL http: http::/papprrs...... www.::iwss...pppppopoeocopp
7152 imagination augmented agent s for deep reinforcement learning pdf
Nikolay Savinov, Alexey Dosovitskiy, and Vladlen Koltun. Semi-parametric topological memory for navigation. arXiv preprint arXiv: 2018a.
Nikolay Savinov, Anton Raichuk, Raphaël Marinier, Damien Vincent, Mare Pollefeys, Timothy Lilli- crap, and Sylvain Gelly. Episodic curiosity through reachability. arXiv preprint arXiv:/810.022: 2018b.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Joannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 201
David Silver, Hado van Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac- Amold, David Reichert, Neil Rabinowitz, Andre Barreto, and Thomas Degris. The predictron: End-to-end learning and planning. December 2016b. URL httt:://arxii: wttp://wwww.gg.////a.aaaa/
08810
Aravind Srinivas, Allan Jabri, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Universal planning networks. April 2018. URL http:::/aaxxv...rr.. abs/ 1804. 00645
Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4): 1991.
Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. February 2016. URL
Alexander Sasha Vezhnevets, Simon Osindero, Tom Schaul, Nicolas Heess, Max Jaderberg, David Silver, and Koray Kavukcuoglu. Feudal networks for hierarchical reinforcement learning, arXiv preprint UrXiv:1703.01166 2017.
Ranxiao Frances Wang and Flizabeth S Spelke. Human spatial representation: Insights from animals. Trends in cognitive sciences, 6(9):376-322, 2002.
Tingwu Wang and Jimmy Ba. Exploring model-based planning with policy networks. arXiv preprint arXiv:1906.0866 2019.
Yang Wang, David Mulvaney, Ian Sillitoe, and Erick Swere. Robot navigation by waypoints. Journal of Intelligent and Robotic Systems, 5752):, 2008
Ge Yang, Amy Zhang, Ari S. Morcos, Joelle Pineau, Pieter Abbeel, and Roberto Calandra. Plan2vec: Unsupervised representation learning by latent plans. In Proceedings The 2nd Annual Conference on Learning for Dynamics and Control, volume 120 of Proceedings of Machine Learning Research, pp. 1-12, 2020. Stt:22005.036453
Rui Zhao, Xudong Sun, and Volker Tresp. Maximum entropy-regulariized multi-goal reinforcement learning. arXiv preprint arXiv: 2019.
12

---

Under review as a conference paper at ICLR 2021
A
APPENDIX
A.1 GREEDY LATENT SPARSIFICATION
Algorithm 2 Greedy Latent Sparsification (GLS) for Latent Cluster Training
Given: Replay Bufter B, Encoder fE. (}-
Initialize: Latent
Set of embeddings selected.
1: Sample K achieved goals from B.
2: Sample lc ~ {0... K 1}.
3: dist = (g)) fE (2)),, .M I[fE(gk) (9x))
4: for fori=1 to M do
Sub-sampling
5: < arg max dist
6: Add fE (gx) LatentEmeeds.
7: NEWdist = [ll/ fr())
"..,)o))) - [x(
8: dist s.
9:
end for
Optimize equation on Latent
The Greedy Latent Sparsification (GLS) algorithm sub-samples a large batch by sparsification. GLS first randomly selects a latent embedding from the batch, and then greedily chooses the next embedding that is furthest away from already selected embeddings. After collecting some warm-up trajectories before planning starts (see Table 2) during training, we first use GLS to initialize the latent centroids, and then continue use to sample the batches used to train the latent clusters. As mentioned in Section 4.2 GLS is strongly inspired by Arthur & Vassilvitskij (2007), and this type of approach is known to improve clustering.
A.2 GRAPH SEARCH VIA SOFT VALUE ITERATIONS
In this paper, we employ a soft version of Floyd algorithm, which we find empirically work well. Rather than simply using the min operation to do relaxation, the soft value iteration procedure uses soft ft min operation when doing an update (note that, since we negated the distances to be negative in the weight matrix of the graph, which is Equutioo 6, the operations softmax). The reason is that neural distances can be inconsistent and inaccurate at times, and using a soft operation makes the whole procedure more robust. More concretely, repeat the following update on the weight matrix for S steps with temperature f::
are actually max and
1 exp (wi, + wki) B
N+1
Wi.j
Wi,k + wk.)
N+1 kl exp (wir + ww)))
Following the practice in Eysenbach et al. (2019); Huang et al. (2019), we do the following initial- ization to the matrix in Equation 6) for entries smaller than the negative of d_max, we penalize the entry by adding oo to it (in this paper, we use 106 the -co value). The essential idea is that we only trust neural estimate when is local, and we rely on graph search to solve for global, longer-horizon distances. The -c penalty effectively masks out those entries with large negative values in the softmax operation above. If we replace softmax with a hard max, we recover the original update in Floyd algorithm; can interpolate between hard Floyd and soft Floyd by tuning the temperature 3.
A.3 HYPER-PAN PARAMETERS
Table lists the common
across all environments. Table lists the
that differ across the environments.
13

---

Under review as a conference paper at ICLR 2021
Table Hyyee-parametern in Common
Parameter
Value
DDPG
optimizer
Adam Kingma & Ba 2014
number of hidden layers (all networks)
3
number hidden units per layer
256
nonlinearity
ReLU
polyak for target network (r)
0.995
target update interval
10
ratio between
optimization steps
2
Random action probability
0.2
Initial random trajs per worker
100
Hindsight relabelling ratio
0.85
Latent Landmarks Auto-encoder
number of hidden layers
2
number of hidden units per layer
128
nonlinearity
ReLU
embedding size for reachability constraint loss
16
1.0
learning rate
3e-4
Graph Search
probability of using search during train
0.5
S (number of soft value iterations)
20
3 (temperature)
1.1
Table Hyper-parametes for Each Environment
Point-Maze
Ant-Maze
Fetch tasks
DDPG
Learning rate
2e-4
2e-4
1e-3
Number of workers
1
3
12
Batch size
512
1024
1024
Action L.
0.5
0.05
0.01
Gamma
0.98
0.98
0.99
Action noise
0.2
0.2
0.1
Hindsight relabelling range
80
100
50
Latent Landmarks & Auto-encoder
Number of latent landmarks
50
50
80
Number warm-up trajectories
500
500
6000
Batch size
256
256
150
Graph Search
damar (clipping threshold for distances)
20.0
20.0
15.0
Random landmarks added during train
150
150
20
14

---

Under review as a conference paper at ICLR 2021
We find that having centralized replay for all parallel workers is significantly more sample efficient than having separate replays for each worker and simply averaging the gradients across workers.
For Ant-Maze environment, we do grad norm clipping by value of 15.0 for all networks. For Fetch tasks, we normalize the inputs by running means and standard deviations per input dimensions.
Since L is able to decompose a long morgzoorizon goal into many short-horizon goals, we shorten the range of future steps where we do hindsight relabelling: as a result, the agent can focus its optimization effort on more immediate goals. This corresponds to the hyper- parameter: Hindsight relabelling range.
During training, we collect 50% of the data without the planning module, and the other 50% of the data with planning. This corresponds to the hyper-parameter: probability of using search during train.
At train time, to encourage exploration during planning, temporarily add small number of random landmarks from GLS (Algorithm{) to the existing latent landmarks, A new set of random landmarks selected for each episode before graph search starts (Algorithm [T]
This corresponds to the arper-paramete: Random landmarks added during train.
We find that collecting a certain number of warm-up trajectories for every worker before the planning procedure starts (during training) and before GLS (Algorithm ( is used for initialization to help improve the planning results. This corresponds to the hyper-parameter: Number of warm-up trajectories.
15


Reviews:
[{"review_id": "wmP0mIFNJNY", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper approaches long horizon planning by learning a sparse graphical representation. The proposed algorithm, L3P, proceeds by learning a latent space which enforces a distance measure, where this distance is learned to mimic the number of steps between states via a goal conditioned Q-function. A clustering algorithm is then used to represent this latent space through only a small, efficient set of latent landmarks. These landmarks are then connected if nearby via the distance, the estimate of which is refined via soft value iterations over the graph. L3P is demonstrated on a number of environments to be both data efficient and high performing compared to baselines.\n\nThe paper is clear and well written. The topic of graphical representations of state spaces to plan long horizons over via local predictions holds significant promise, as does sparsifying this representation. The latent space construction and sparsification procedures are reasonable, particularly tying the distance to the policy Q-function as in SoRB. A few notes on clarity:\n- The algorithm should be moved to the main body of the paper. (to reduce space, potentially remove a row of images in Fig. 4, stretch figure 1 to use the full width)\n- A video would be helpful of the full process, including the soft value iterations.\n- What algorithmic parameters (e.g., latent space dimensionality, network sizes, training data and number of states from which the latent space was constructed) were used?\n- How are L_rec and L_latent traded off?\n- The tasks, e.g., Fig 2 and 5, should be shown before the results in Fig. 3.\n- More intuition on the soft value iteration would be helpful. \n\nThe results on the shown tasks show clear benefits of L3P, particularly on the tasks where HER has most difficulty (likely the more long range tasks) L3P outperforms substantially. However, the results are somewhat lacking (1) on ablation, (2) demonstrating new environment generalization, and (3) showing significantly complex tasks.\n1) The authors show ablation for number of landmarks and d_max, but not for algorithmic changes like the soft value iteration. How much does this procedure help versus the use of latent marks to sparsify the space? It would also be interesting (though not fully necessary) to see comparison to Savinov 2018a (SPTM), which uses a learned distance predictor rather than the Q-function as in SoRB and L3P.\n2) It is not clear how this method would generalize to new environments, e.g., Section 5.5 in SoRB. This to me is key (particularly with generalization in the title, which I believe currently only refers to task length) and may be a challenge for L3P due to finding landmarks from limited coverage of a new scene.\n3) L3P should be shown on higher dimensional and longer horizon problems, as the maze is fairly short and simple , as is the box task (as can be seen by HER’s performance). These do not push the boundaries of the algorithm. Higher dimensional problems may particularly challenge the landmark learning (e.g., if the landmarks must be learned too in a higher dimensional latent space)\n\n\n__________\n\nAfter author response:\nI appreciate the author’s response. Previous topics:\n1) The author’s add ablations on the hard vs. soft min during the graph search, the additional results are informative, but not conclusive. Given that the overall performance is similar, the authors need to demonstrate the soft-min’s benefits for each experiment and over more training seeds. \n2) For generalization, the author’s confirm that this method is unable to generalize to new environments, though clearly it has other benefits in terms of data efficiency and robustness of solutions. I believe these are still important benefits, though it would be useful to discuss how generalization may be achieved.\n3) For harder experiments, the authors note that the baselines perform poorly there and thus these tasks were not considered. Though reasonable that the baselines are unable to perform in such cases, harder experiments would show the limit of the proposed algorithm. It would be useful to see for instance how well it scales with dimensionality, how quickly the success rate falls off.\n\nNew comments:\na) I believe the title change away from Generalization is an improvement, though the algorithm name \"WORLD MODEL AS A GRAPH\" seems to not capture the novel aspects of this work. This name I believe would be more readily applied to search on the replay buffer or semi-parametric topological memory.\nb) R2's point that much of the robustness may be a factor of choosing states further from the wall is an interesting one. It would be interesting to examine exactly *why* the method is robust.\nOverall, I believe the paper is interesting and proposes some novel ideas that have benefit, it requires more thorough analysis, and thus I am leaving my score unchanged.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "fwPi_pZyptm", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nIn general, I found the paper interesting and the key ideas intuitive and sensible. The paper is reasonably well-written, but there are some important points that are unclear (see below). The experimental section compares with very recent algorithms for planning over learned value functions (SORB and MSS), and the results on five continuous control task show substantial improvement over these methods.\n\nIn terms of clarification, I have a few questions for the authors:\n- With regards to the embedding used and $\\Psi$:\n  - The embedding operates over goals. Are these the goals of the problem? For example, in the AntMaze problem (Fig. 4), is $G = \\{ \\textit{red-square} \\} $? If so, doesn't this mean that  $\\forall s \\Psi(s) = \\textit{red-square}$, and thus $V(\\Psi(s), \\Psi(s')) = 0$ for all pairs of states?\n  - More reasonably, $V$ could be defined over states, so that $V(s_1, s_2)$ indicates the number of steps between $s_1$ and $s_2$. This can be achieved under the current formulation with $G = S$ and $\\Psi(s) = s$. Then $D(s_t, a, \\Psi(s_t))$ represents the number of steps between $s_{t+1}$ and $s_t$. Is this the case for most experiments? No details of the $\\Psi$ function used in the experiments was given, so I find this point somewhat confusing. \n  - If this is not the case, then I think it's important to explain where does this mapping function comes from, since it essentially provides some reward shaping for the problem. \n\n- Can you explain what is the motivation for using Algorithm 2? A common approach is to use Djikstra on the graph representation. Is there a reason why this can't be done (or is worse) in this case? \n\n- I don't fully understand Algorithm 1 to create the batch for Eq. (5). Can you explain this algorithm in more details? For example, the definition of $\\texttt{dist}$ is ambiguous; is it pairwise distance between all sampled goals? distance between $g_1$ and all the others? something else? In general, I'm having quite a hard time figuring out what this is doing.\n\nIn light of above, I think the clarity of the paper can be improved in many places. But, overall, I'm positive towards this work and I think it is a nice contribution, particularly since I'm not aware of other work creating explicit low-dimensional landmarks to be used for search. That being said, one exception that is missing from the literature review is [1], which creates a discrete latent representation of the environment and solves it using prioritized sweeping. \n\n[1] Corneil, Dane, Wulfram Gerstner, and Johanni Brea. \"Efficient Model-Based Deep Reinforcement Learning with Variational State Tabulation.\" International Conference on Machine Learning. 2018.", "extra_scores": {}, "extra_text": {}}, {"review_id": "_WeL21aim2-", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper generalizes the previous graph-based planning algorithm for goal-conditioned RL algorithms by learning a latent metric space and building the graph by clustering. The experiments show that their approach outperforms both the HER and previous MSS/SORB baselines.\n\nStrong points:\n1. The paper achieves good performance on the challenging AntMaze and PickAndPlace environments.\n2. The clustering methods to find latent landmarks are novel in this setting.\n\nWeak points:\n1. The novelty of the paper is limited. Planning on the graph or learning a latent distance embedding for planning are all existing ideas. The combination is straightforward. The robustness of soft-min is well known. \n2. I am confused about why the authors' approach is better than MSS or SORB. It's not obvious that a distance-based clustering will perform significantly better than sampled landmarks (like farthest point sampling). I guess the reason is that clustering helps to avoid the landmarks near the wall (which often appears in the FP sampled methods). Those states are challenging as it has a higher possibility for the ant to collide with the wall. Can the author give some words about this? Moreover, the paper is not clear about the contributions of each component in the improvements. For example, no table or figure shows the benefits of the soft value iteration over the hard version. The author shall provide more ablation studies and explanations.\n3. Although the paper presents the overall idea well, there is still a lot of room for improvement in writing. For example, I suggest that the author replaces the hyphen with the comma in the abstract.\n\n~~~Based on the current presentation, the lack of ablation study, and the limited novelty, I think this paper is not good enough to be accepted. ~~~\n\nHere are some questions:\n1. Can we build graphs by sampling goal states and cluster them based on the metric space defined by the Q networks? What's the performance of this approach? Will it be worse? Why?\n2. Can a sample-based approach work in non-navigation environments? Motion planning algorithms can solve the high-level part of both pick&place or ant maze problem once the geometric model is known. If so, why do we need the expensive RL algorithm?\n\n\n---\n\nAfter author response：\n\nThe authors have improved the presentation and added the necessary ablation studies. I appreciate the authors' effort. I am glad to raise the score to reflect the changes. I hope this work will inspire future research on hierarchical planning algorithms.", "extra_scores": {}, "extra_text": {}}, {"review_id": "DJsCqwuDwTl", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#### Summary:\nThis 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.\n\n##### Pros:\n- The results seem strong.\n- The method is straightforward and nice, and it seems to work more efficiently than previous methods.\n\n##### Cons:\n- Overall, the paper is incomplete. It’s very difficult to understand the whole system, how training proceeds, what the hyperparameters are, etc. There is very little detail about how this all works. The specific components of the latent landmark learning and the planning are mostly fleshed out (albeit with some missing details still) but the rest of the setup is mostly not described. How are the losses combined? Are there multiple learners or does the DDPG agent also do the latent space learning? How is training done? What are the hyperparameters used? \n- The generalization claims are a bit oversold as the generalization achieved is weak. In particular, the method requires that the test goal distribution have sufficient overlap with the train goal distribution for latent landmarks to be created near the desired test goals. \n- Are the experiment domains created here or existing work? If the former, more details are required. If the latter, citations are required.\n- Given the multiple changes with respect to prior work, additional ablations and experiments are needed to better understand why this method is performing well, and what contributions are important.\n\n#### Decision:\nOverall, I find it too hard to know exactly what is happening in this work. It would be very difficult, if not impossible, to reproduce this work from this paper alone. Without sufficient details about the system, it’s hard to evaluate, which means that I default to a reject, especially in conjunction with the other weaknesses listed above.\n\n#### Questions:\n1. What is the goal space G? Is it the state space? This should be defined.\n2. What exactly is \\Psi? Is it learned or is it given? This should be explained better.\n3. How is the autoencoder for the latent space learned within the overall system?\n4. The connection between D and Q is poorly explained. From the text, it reads as if you get Q from D. However, it makes more sense if the DDPG agent estimates Q and then you use eq(3) to compute D from Q. Which is it?\n5. Are both V and D really necessary? It seems like they encode the same thing, essentially. \n6. Why do you use soft value iteration? Is it a good choice? Is there a reason for this choice? Would other choices be better or worse? It seems likely that the need for the dist_max penalty is due to the use of the soft value iteration with a poorly set soft value iteration temperature. Is this true?\n7. Are the domains from Duan et al 2016 or some other existing work? Or did you create them from scratch (unlikely). Why are they labeled “-hard”? What is hard about them? These need to be cited and described. You have tons of space in the appendix to use freely. \n8. In Figure 3, is the x axis training timesteps? \n9. Figure 5 caption says the agent has learned landmarks that help avoid collision with the box. Where are these? Is there a figure? Any sort of information on these? Wouldn’t any landmark away from the box accomplish this?\n10. For Figure 6, why is choosing the maximum distance bound important? Why would the planning algorithm choose subgoals that are near the max distance bound given that the planner is finding the argmin subgoal? It seems likely that this hacky distance bound is unnecessary.\n11. The “planning algorithm” seems trivial. Am I missing something? Given the latent landmarks and the estimates for $d_{c \\rightarrow g}$, the planner just gets the next estimated subgoal and the distance to it, executes the goal-conditioned policy for that subgoal for that many steps, crosses out that subgoal, and then goes to the next subgoal. Is it non-trivial because $d_{s \\rightarrow c}$ gets updated after each step (not included in the pseudocode)?\n12. Is the GLS algorithm a contribution of this work or something from a previous work? Again, it's neither described sufficiently nor cited sufficiently.\n\n#### Comments:\n- The introduction uses unnecessary hyperbole and insufficient citations to make its points. It comes off as less compelling as a result.\n- Section 3 is called Background but really it contains the preliminaries and definitions of this work. It should not be called Background.\n- Should the variance vector of the centroids appear in equation (5)? It is mentioned once and then never discussed again. What values of it are learned? Is it important? \n- Similarly, the temperature for the soft value iteration is mentioned once and then never discussed again. What role does it play? Is it important? Would tuning this better remove the need for the dist_max penalty?\n- The notation $f_D(c_i)$ is unnecessarily complicated. These values correspond to states, no? So define them as $s_{c_i}$ or something similar to make the math and text more clear. As it is, equation (6) is unnecessarily cluttered.\n- I assume that the argmax in the subgoal computation at the end of section 4 should be an argmax over c? \n- There is a typo in the Figure 3 caption (PointmMaze). Further, the text in the subfigures is illegibly small.\n- Break lines 11 and 13 of algorithm 2 into two lines each.\n\n\n*****************\nAfter author response:\n\nI appreciate the updates you've made to the paper to better flesh it out, including the diagrams, pseudocode, and additional ablations. I've increased my score accordingly. Regarding generalization, I fully agree that it can be difficult to show. That said, when it's advertised in the title of the paper, I expect it to be clearly shown in the paper itself. It seems that the language has been greatly toned down in the updated version. However, without entirely re-reviewing the paper, I am unable to fully recommend acceptance.\n\nAside: I would have appreciated responses to my questions directly so that I don't need to dig through the rewritten paper to find the answers to my questions.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a model-based RL algorithm which, instead of simply fitting a parameterized transition model and uses rollout for planning, learns latent landmarks via distance-based clustering and conducts planning on the learned graph. Although some of these ideas themselves have appeared in literatures, the overall approach is very nice, novel and sophisticated. The experimental results appear strong and interesting. Most reviewers feel positive about the contributions of the paper, but there remain concerns that need to be addressed.

The proposed approach is highly nontrivial, and more ablation, generalization and environments need to be studied to fully justify what's going on. The authors agree to expand the paper and add the needed results, which would require substantial work thus reviewers recommend that the paper be submitted again to a future conference and receive another round of review. Showing the generalization is nontrivial, and it would be make the paper stronger if the authors put more thoughts into this issue, although it is not a must.

Minor: Another technical comment is that the approach seems heavily rely the choice of embedding distance. Learning the best embedding with meaningful embedding distance has been considered in other scenarios, see eg https://arxiv.org/abs/1906.00302. It would be interesting to try out and compare difference choices of the embedding distance.


Decision: Reject