Planning from Pixels using Inverse Dynamics Models
Authors: Keiran Paster
Sheila A. McIlraith
Jimmy Ba
Venue: ICLR 2021 Poster (2021)
Keywords: model based reinforcement learning
deep reinforcement learning
multi-task learning
deep learning
goal-conditioned reinforcement learning

Abstract: Learning dynamics models in high-dimensional observation spaces can be challenging for model-based RL agents. We propose a novel way to learn models in a latent space by learning to predict sequences of future actions conditioned on task completion. These models track task-relevant environment dynamics over a distribution of tasks, while simultaneously serving as an effective heuristic for planning with sparse rewards. We evaluate our method on challenging visual goal completion tasks and show a substantial increase in performance compared to prior model-free approaches.

Published as a conference paper at ICLR 2021
PLANNING FROM PIXELS USING INVERSE DYNAMICS MODELS
Keiran Paster
Sheila A. Mcllraith & Jimmy Ba
Department of Computer Science University of Toronto, Vector Institute ke rpêcs. toronto....o... edd
Department of Computer Science University of Toronto, Vector Institute {sheila, jba}ecs. jbafeas.toronter wowonto..ot...ooooo
ABSTRACT
Learning task-agnostic dynamics models in high-dimensional observation spaces can be challenging for model-based RL agents. We propose a novel way to learn latent world models by learning predict sequences of future actions conditioned on task completion. These task-conditioned models adaptively focus modeling capacity on task-relevant dynamics, while simultancously serving an effective heuristic for planning with sparse rewards. We evaluate our method on challeng- ing visual goal completion tasks and show a substantial increase in performance compared to prior model-free approaches.
INTRODUCTION
Deep reinforcement learning has proven to be a powerful and effective framework for solving a di- versity of challenging decision-making problems (Silver al. 2017a Berner et al. 2019). However these algorithms are typically trained to maximize a single reward function, ignoring information that is not directly relevant the associated task hand. This way of learning in stark contrast to how humans learn (Tenenbaum) 2018). Without being prompted by specific task, humans can still explore their environment, practice achieving imaginary goals, and so doing learn about the dy- namics of the environment. When subsequently presented with novel task, humans can utilize this learned knowledge to bootstrap learning property we would like our artificial agents to have. In this work, we investigate one way to bridge this gap by learning world models (Ha & Schmidhuber) 2018) that enable the realization of previously unseen tasks.
By modeling the task-agnostic dynamics of an environment, an agent can make predictions about how its own actions may affect the environment state without the need for additional samples from the environment. Prior work has shown that by using powerful function approximators model en- vironment dynamics, training an agent entirely within its own world models can result large gains in sample efficiency (Ha & Schmidhuber 2018). However, learning world models that are both curate and general has largely remained elusive, with these models experiencing many performance issues in the multi-task setting.
The main reason for poor performance is the so-called planning horizon dilemma (Wang et al., 2019): accurately modeling dynamics over a long horizon necessary accurately estimate re- wards, but performance is often poor when planning over long sequences due to the accumulation of errors. These modeling errors are especially prevalent in high-dimensional observation spaces where loss functions that operate on pixels may focus model capacity on insk-irrelevant features (Kaiser et al., (020). Recent work (Hainer al., 2020 Schrittiiesrr al.] 2019) has attempted to side-step these issues by learning world model in latent space and propagating gradients over multiple time-steps. While these methods are able to learn accurate world models and achieve high performance benchmark tasks, their representations are usually trained with task-specific information such as rewards, encouraging the model to focus on tracking task-relevant features but compromising their ability to generalize to new tasks.
In this work, we propose to learn powerful, latent world models that can predict environment dy- namics when planning for distribution of tasks. The main contributions of our paper are three-fold: we propose to learn a latent world model conditioned on a goal; we train our latent representation to model inverse dynamics sequences of actions that take the agent from state to another, rather

---

Published as a conference paper at ICLR 2021
Move Down
Tur Loft
Down
Linear
Linear
Linear
Linear
Linear
State
Linear
Linear
Linear
Goal
The network architecture for the inverse dynamics model used in GLAMOR. ResNets are used to encode state features and an LSTM predicts the action sequence.
Figure I:
than training it to capture information about reward; and we show that by combining our inverse dynamics model and a prior over action sequences, we can quickly construct plans that maximize the probability of reaching a goal state. We evaluate our world model on a diverse distribution of challenging visual goals in Atari games and the Deepmind Control Suite (Tassa et al.) 2018) to as- sess both its accuracy and sample efficiency. We find that when planning in our latent world model, our agent outperforms prior, model-free methods across most tasks, while providing an order of magnitude better sample efficiency some tasks.
RELATED WORK
Model-based RL has typically focused on learning powerful forward dynamics models, which are trained to predict the next state given the current state and action. In works such as (Kaiser et al.) 2020), these models are trained to predict the next state in observation space often by minimizing L.2 distance. While the performance of these algorithms in the low data regime is often strong, they can struggle reach the asymptotic performance model-free methods (Hafner al. 2020). An alternative approach is to learn forward model in latent space, which may be able to avoid modeling irrelevant features and better optimize for long-term consistency. These latent spaces can be trained to maximize mutual information with the observations (Hafner et al. 2020 2019) task-specific quantities like the reward, value, or policy (Schrittwieser et al. 2019). Using learned forward model, there several ways that an agent could create policy.
While forward dynamics models map a state and action to the next state, an inverse dynamics model maps subsequent states action. Inverse dynamics models have been used in various ways in sequential decision making. In exploration, inverse dynamics serves as a way to learn represen- tations the controllable aspects of the state (Pathak et al. 2017). In imitation learning, inverse dynamics models can be used to map sequence of states to the actions needed to imitate the trajec- tory (Pavse et al. 2019). Christiano al. (2016) use inverse dynamics models to translate actions taken in simulated environment to the real world.
Recently, there has been emergence of work (e.g.. Ghosh al. 2020 Schmidhuber 2019 Srivas-
tava et al. 2019) highlighting the relationship between imitation learning and reinforcement learn- ing. Specifically, rather than learn to map states and actions to reward, learning, Srivastava Srriiint,Srivvstavvv al. (2019) train model to predict actions given state and outcome, which could be the amount of reward the agent is to collect within certain amount of time, Ghosh et al. (2020) use similar idea, predicting actions conditioned on initial state, goal state, and the amount time left to achieve the goal. As explored Appendix A.I these methods are perhaps the nearest neighbors to algorithm.
typical in reinforcement
In our paper, we tackle a visual goal-completion task due to its generality and ability to generate tasks with domain knowledge. Reinforcement learning with multiple goals has been studied
2

---

Published as a conference paper at ICLR 2021
since Kaelbling (1993). Most agents that trained to achieve multiple goals trained with off- policy reinforcement learning combined with a form of hindsight relabeling (Andrychowicz et al. 2017). where trajectories that do not achieve the desired goal are relabeled
successful trajectory that achieves the goal that was actually reached. (Andrychowiez et al. (2017) uses value-based rein- forcement learning with reward based the euclidean distance between physical objects, which is only possible with access to an object-oriented representation of the state. In environments with high-dimensional observation spaces, goal-achievement rewards are more difficult design. Nair et al. (2018) use a VAE (Kingma & Welling 2014) trained on observations to construct a latent space and uses distances in the latent space for reward. These distances, however, may contain features that are uncontrollable or irrelevant. Warrde-Farley al. (2019) attempt to solve this issue by fram- ing the goal-achievement task as maximizing the mutual information between the goal and achieved state [sgg ST). Our method differs from these approaches since we aim simply to maximize an indicator reward (ST = 8y) and do not explicitly learn a value or Q-function.
METHOD
3.1 PROBLEM FORMULATION
Reinforcement learning is a framework in which an agent acts in an unknown environment and adapts based on experience. We model the problem with an MDP, defined as the tuple (S, A, R, 7). S is a set of states; A is a set of actions; the transition probabilities T S x A x S [0, 1] defines the probability of the environment transitioning from state s to s' given that the agent acts with action a; the reward function R S x A x 5 + R maps state-action transition to a real number, and 0 S. S the discount factor, which controls how much an agent should prefer rewards sooner rather than later. An agent acts in the MDP with a policy S x A - [0, 1]. which determines the probability the agent taking action while in state 8.
The expected return of policy is denoted:
R(st, az s1) st)
Σ E(
JRn = Ey ETNP(r)))
(1)
that is the averaged discounted future rewards for trajectories = {(s,, |T of states and actions sampled from the policy. A reinforcement learning agent's objective is to find the optimal policy o² = arg max, (T) that maximizes the expected return.
In goal-conditioned reinforcement learning, an agent's objective is to find a policy that maximizes this return over the distribution of goals conditioned on g. In our work, g and we consider goal achievement rewards the form (s) 1(s g). Additionally, we consider trajectory to be complete when any Ry(s;) 1 and denote this time-step t = T. With these rewards, an optimal goal- achieviig agent maximizes:
p(g) when acting with policy that is now also
((T) =
= Ryspte [Eer m.) (7 Rg (sT)]]-
(2)
Note that unlike prior works, we consider both the probability of goal achievement as well as the length of the trajectory T in our objective.
3.2 PLANNING
We consider the problem of finding an optimal action sequence aj...., -4-1 to maximize expected return J(s, 9. 01, . -1-1):
J(s, g, a1, ax- 1) Spl (s,ar...... "rr = p(3; = g|s, a....., ax-1)
(3)
Thus, the optimal action sequence is found by solving the following optimization problem:
0 ii 0i....,4h._ 46_1
arg max 7 * p(5k g 11, $1,0,,1 .., a4-1)
01....Rt-1
Even with access to a perfect model of p(s: = g ai.... 1)-1), solving this optimization may be difficult. In many environments, the number of action sequences that reach the goal are vastly outnumbered by the action sequences that do not. Without heuristic or reward- wwwrd--happng, there is little hope of solving this problem in reasonable amount of time.
3

---

Published as a conference paper at ICLR 2021
3.3 GLAMOR: GOAL-CONDITIONED LATENT ACTION MODELS FOR RL
Inspired by sequence modeling in NLP, propose rewrite Eqquation 4 in way that permits factoring across the actions in the action sequence. By factoring, planning in our model can use the heuristic search algorithms that enable sampling high quality language sequences that are hundreds of tokens long. First, note tha
k-1 p(sk g/ aci, p( = g| acs)
p(s = g/s1. 0I........... ax-1)
(5)
Let 2( aci, a;) A p1s
p(ex=gle1,@cc) Intuitively, these terms are equal to the relative gain in probability of reaching state conditioned on taking action a, versus the marginal probability of reaching the goal without conditioning that action. These terms provide useful information can guide search towards high scoring action sequences when constructing a plan.
To learn the values of the z($1,9 g: a-ci, aci;;;)) az), we use Bayes' rule to show that we can equivalenty learn two auto-regressive behavioral models:
p( = g/81, ag:)
p(as/80 = 9: 81, agi 1)pss = 81, (,, ag:-1)
(6)
p(sk 0181,03i-1,
plsi, asilp(ss 9, 81, agi-1)
plas|11, 85 = 9,05i-1)
pla,1s1,
We refer p(a,,-..,,al$5 8k 9) the inverse dynamics model and p(a1,..., (x1s)) the action prior. Using these models, we can find an optimal plan by optimizing the following objective:
oz arg max 7 pP((a...... 01
,( 02-1
(8)
p(a ag a
3.4 LEARNING INVERSE DYNAMICS MODEL AND ACTION PRIOR
We parameterize both the inverse dynamics model and the action prior model with LSTMs (Hochre- iter & Schmidhuber 1997) parameterized by 0 and Φ. Please refer to section 4.2 for more informa- tion about architecture.
From the training data we construct a new dataset D of ($1, C1...., ax-1, Sx) tuples for all possible combinations of 81 and 8g. Note that this similar to hindsight relabeling (Andrychowicz al., 2017). In training, we find parameters to maximize the likelihood of the training data using Adam (Loshchilov & Hutter] 2019) optimize the following loss, where a controls the relative weight of the losses when the two models share parameters:
L(O, 6) =IE
- (1.. log pa (g)((a)... (#1,..., 00-11 [1, sk)]
+ a (s (1.4)....--1)) -l log Ps (s(a),........ 02 [s1]].
While there are many heuristic search algorithms that could be used to optimize Equation opt to use random shooting (RS) (Richards) 2005) due to ease parallelization and strong empirical performance. RS samples N candidate action sequences, evaluates the value of the sequences using the learned models, and takes the first action of the winning action sequence. Rather than sampling action sequences uniformly, sample them auto-regressively using Boltzmann distribution of the scores z.
Herre we assume that all goals have a non-zero probability of heing achieved under the training distribution.
"The notation azi denotes 01.-.., Ai-1-
'By upper-bounding each term any individual factor being small bounds the rest of the score and allows for early pruning of had action sequences. An upper-bound can he achieved by assuming some support over all actions the training policy, which common requirement for the convergence of many RT. algorithms.

---

Published as a conference paper at ICLR 2021
In order to be able to re- re-plan every step, model must be able to model action sequences with variable length. By predicting variable length action sequences using snd-tokens, our model is able to plan to both maximize goal-achievement probability and find shortest path.
Algorithm 1: GLAMOR
Initialize inverse dynamics model pa ps(a....... ak- $1, sk); Initialize action prior model pa (......... .ktii........11-11....- 18)1
Initialize dataset (((s, a)): for training iteration i- 1, 2,3...... do
Sample trajectories with training policy Ne and goals sampled from p(g);
Add tuples (51, 02,.... 4g-1, 8*) to dataset D;
Update 8, by descending the gradient of [Equatinn 9
end
3.5 DATA COLLECTION
While training data need not be collected on-policy, we note that using an arbitrary data- cala-collection policy may result in a causally our open-loop planner to generate a sequence of actions given just the initial state 81 and a sampled goal g and follow this sequence for the rest of the episode. For discussion of the types of errors that may occur with causally incorrect model, see Appendix [A.2
model (Rezende et al. 2020). In our experiments,
For exploration and in order to ensure support over all actions, we use an e-greedy exploration policy where the ith action in the plan is taken with probability 1 and an action is randomly sampled with probability e. We decay to a small final value over time. We store trajectories in a circular replay buffer, meaning dataset D contains trajectories from multiple past versions of the planner.
3.6 COMPARISON TO PRIOR METHODS
Prior methods for goal-conditioned RL include DISCERN (Ward--Fatde et al. 2019) and Goal- Conditioned Supervised Learning (GCSL) (Ghosh et al. 2020). DISCERN learns goal discrimi- nator and Q-network and rewards the agent for ending up in states where its discriminator can easily guess the desired goal. GCSL learns policy with supervised learning by iteratively imitating its own relabeled trajectories. While GLAMOR outputs a sequence of actions using an LSTM, these prior works only output single action for the current time-step.
While GLAMOR shows strong empirical performance in benchmark tasks, we believe that there are several other key advantages over prior goal-conditioned RL algorithms:
GLAMOR learns world model, which more flexible than the Q-networks or policies that prior methods use. For example, despite only being trained to find shortest paths to goal state, the latent world model learned in GLAMOR contains information about many different paths to the goal. One implication is that during planning, the agent can modified to select action sequences that reach the goal at a specific time-step. This is experimentally verified in section 4.7
Q-learning is well known (Fujimoto et al. 2019) to perform poorly with purely off ooflcy data and GCSL relies on many iterations training and data-collection converge to optimal policy. As shown in Appendix A.4.1 GLAMOR performs well even with completely random training policy, often still achieving SOTA performance. Strong off- policy performance may enable the use of offline datasets or more sophisticated exploration policies.
GCSL. does not estimate the prior probability of taking an action sequence and therefore fails to converge to an optimal policy when action sequences that are optimal for one goal may land the agent in another goal state. By using the action prior to disentangle the probability of reaching goal from the probability of taking action sequence, GLAMOR avoids this issue. For more detailed comparison of GLAMOR and GCSL, see section Appendix A.
5

---

Published as a conference paper at ICLR 2021
Control
0.6
0.5
Goab. Adlieved
Goals Achieved
0.4
0.4
0.3
0.3
0.2
0.2
0.1
0.1
0.0
0.0
2
3
Agent Steps (Millions)
Agett Steps (Millions)
algo
DISCERN
GCSL
GLAMOR (Oun)
Figure 2: Both in Atari and on tasks from the Deepmind Control Suite, GLAMOR outperforms prior methods. The goal achievement rate is averaged over all games / control tasks and over three seeds. See Figure Figure 10in the appendix for more detailed training curves.
4 EXPERIMENTS
In our experiments, we aim to answer the following questions:
How accurate our world model? Can the agent plan using its model to achieve a diverse set of goals even in high-dimensional domains?
How sample efficient are our agents? In the low sample regime, is it better to use GLAMOR or use model-free reinforcement learning?
How effective is planning procedure? Can the planner discover how to achieve goals with shortest paths or paths with a specific length?
4.1 ENVIRONMENTS
We evaluate our method on two types of environments: Atari games and control tasks the Deep- mind Control Suite (Tassa al. 2018). In both environments, visual observations converted to grayscale and down-scaled to (80, 104) pixels. In order to incorporate historical information, we opt concatenate the four most recent frames form observation, introduced by MMih al. (001)). However, we opt to use only a single frame for specifying visual goals. We calculate goal achievement both environments using extracted low dimensional features.
ARCADE LEARNING ENVIRONMENT (ALE). We run experiments on a subset of the available ALE games, chosen by the availability of labeling methods and the suitability of the games the goal- achieving tasks. In all games we use a frame-skip of 4, and all Atari environments used a random number initial noops and sticky actions introduce stochasticity.
DEEPMIND CONTROL SUITE. We choose to use the subset of the Deepmind Control Suite Tassa al. 2018) chosen by Warde-Farley et al. (2019). We also adopt the same method of dis- cretization, discretizing most of the A-dimensional continuous control tasks into 3 actions. For manipule ator, we adopt DISCERN's diagonal discretization. In point mass we apply frame-skip of frames, with no frame-skip for any other control environment.
Table I: Goal Achievement Rates at 500k Agent Steps in Atari
Atari
GLAMOR (Ours)
DISCERN
GCSL
DM Control
GLAMOR (Ours)
DISCERN
GCSL
Bowling
0.14
0.11 (0.09)
0.21 (0.11)
ball_in.cup
0.20 (0.08)
0.01 (0.02)
0.03 (0.04)
Boxing
0.06 (0.04)
0.01 (0.01)
0.01 (0.02)
caripole
0.97 (0.03)
0.01 (0.01)
0.09 (0.02)
Breakout
(0.05)
(0.02)
(0.03)
finger
(0.01)
(0.02)
Frosthite
0.35 (0.06)
0.04 (0.04)
0.20 (0.05)
manipulator
0.00 (0.00)
0.00 (0.00)
0.00 (0.00)
Montezuma
0.13 (0.07)
0.03 (0.02)
0.30 (0.07)
pendulum
(0.09)
0.20 (0.05)
(0.05)
MsPacman
0.28 (0.07)
0.04 (0.02)
0.30 (0.05)
point
0.76 (0.16)
0.03 (0.02)
0.82 (0.05)
Pitfall
0.40 (0.08)
0.07 (0.04)
0.26 (0.07)
reacher
0.27 (0.08)
0.08 (0.06)
0.07 (0.04)
Pong
0.44
(0.08)
0.28
(b)
PrivateEEye
0.26 (0.07)
0.04 (0.04)
0.16 (0.06)
Qhert
(0.03)
(0.04)
(0.17)
Riverraid
0,50 (0.05)
(0.03)
0.79 (0.09)
Seaquest
0.57 (0.06)
(0.04)))4
0.48 (0.08)
Skiing
(0.09)
0.02 (0.02)
0,38 (0.08)
Tennis
0.10 (0.04)
0.00 (0.00)
0.05 (0.03)
(a)

---

Published as a conference paper at ICLR 2021
carigale balance
Feecheed
1.0
0.8
Coak Adieve]
0.8
0.6
0.6
0.4
trials
5 trials
10 trials
50 trials
0.4
Shottes patts
0.2
0.2
0.0
Agen Stepe (Millions)
GLAMOR (Ows)
plan-cud.
(b) Goal achievement rates for different termination
(a) Goal achievement rates for a 7x7 grid-world.
strategies.
Figure 3: (a) The agent starts in the center and must travel to the goal tile. Top shows the rate at which the agent eventually achieved the goal and bottom shows the rate at which the agent achieved the goal with the shortest available path. The amount of compute used for planning is shown the x-axis. As the planning budget increases, both the number of successfully reached goals and the number of goals achieved optimally improves substantially. Brighter means higher achievement rate. (b) In "naive-end", the agent greedily tries to take a shortest path to the goal for T timesteps and is evaluated at the end. In "plan-end", the agent explicitely constructs plan achieve the goal state at the end of its trajectory. GLAMOR (Ours) can choose to terminate its episode early.
4.2 IMPLEMENTATION DETAILS
We parameterize our models with two convolutional encoders, similar to the large model in Espe- holt et al. (2018), and an LSTT model (Hochreiter & Schmidhuber 1997) for action prediction. See Figure 1 for visual description of our architecture. We opt share parameters only in the encoders for the inverse dynamics model and action prior. In DISCERN and GCSL, we represent the time remaining with periodic representation (sin cos(2tl/T)) Detailed training hyperparam-
eters are available at Appendix A.3
The code for training agents on both Atari and DM Control Suite along with evaluation code can be
found at //tttss://gitaaae com www.cow...iceceeeeocco wwwkp/iipp/... glmor
4.3 BASELINES
We evaluate our method against both GCSL and DISCERN. We chose DISCERN due to its reported high performance on our benchmark tasks, and GCSL. due to its similarity to our own algorithm. GCSL was implemented within our code-base, ensuring that any differences are due algorithmic differences rather than implementation details. Our implementation was checked against publicly available code. For DISCERN, due to the lack of available source code, we made a best effort at- tempt to reproduce the algorithm. We built our implementation on top of rlpyt's (Stooke & Abbeel, 2019) Rainbow implementation (Hessel et al. 2018). ) All algorithms use the same encoder architec- ture. and hyperparameters were fixed between GCSL and GLAMOR.
4.4 EVALUATION
To evaluate the degree to which agent achieved its desired goal, we extract the positions of the various entities in the scene. For Atari, we extract state entity labels using code from Anand et al. (219). In the Deepmind Control Suite environments, we use the position information from the non- visual representation. Information from both of these agents kept private from the agents both train and evaluation time. As in DISCERN, goal is considered to be achieved if the positions of the entities are within 10% their feasible range. Positions are evaluated at the end the episode (T = 50 in Atari and T = 100 in control tasks) when the agent decides to terminate the episode. In order to inform GLAMOR of the time remaining until evaluation, we limit the length of the plan in the planner after steps to T - t.
Agents are evaluated on set of 30 fixed goals per environment. These goals are generated similarly to the diverse goal buffer described in Warde-Farde al. (2019) goals repeatedly sampled and
7

---

Published as a conference paper at ICLR 2021
goal is only added to the buffer if it is farther away from the closest goal in feature space than the goal that it is replacing is. We found this procedure to generate a good coverage of possible goals for evaluation.
4.3 ACHIEVING SSSUALL-SPECIIEES GOALS
In order to test our model's 's accuracy, we evaluate our agent's performance in achieving visually- specified goals in all of our test environments. Figure 2] shows that by planning its latent world model, our agent learns to achieve goals with at least as much accuracy as DISCERN and GCSL in out of 21 tasks, often achieving as much as twice as many goals. See Figure and Figure 10 in the appendix for individual training curves for each environment. Despite not being explicitly trained to only focus on controllable features as in WWadd--Farley et al. (2019), Figure 12 shows our agent learning to control the position of its finger even when it can't I exactly match the orientation of the spinner. Overall, we conclude that since planning under our agent's world model results in a high gote-ennhev rate relative to previous state-of- ofte--oftthe-ar algorithms, our model is sufficiently accurate. Videos of GLAMOR agents achieving goals in all environments are available at hhtps: 7/sites. .elas.goollee le www.wi/m/viiiiie amor paoo-ppper.
4.6 SAMPLE EFFICIENCY
Model-based reinforcement learning is known for having better sample efficiency than model-free algorithms. We found this to be true for GLAMOR as well. In Table we show the performance of our algorithm trained with only 500k agent steps. GLAMOR achieves decidedly more goals than DISCERN and GCSL. Of particular note that DISCERN sample inefficien learning to control almost none of the control tasks at 500k steps while GLAMOR has already converged.
4.7 PLANNING
We ran several experiments to test properties of our planning procedure.
4.7.1 PLANNING COMPUTE
One benefit of model-based reinforcement learning is that the agent can improve its performance or even adapt its policy at test-time simply by changing its planning procedure. In order to test the performance of GLAMOR with various levels of compute used construct plan, we used an empty grid-world environment where shortest-paths are known. We evaluate the percentage of trajectories that reach the goal in an optimal amount of time for each goal. Figure 3al shows a heatmap of the optimal goal-completion rates for different amounts of compute used during planning. As the amount of samples in the planner is increased, both the rate of achieving goals and the rate of achieving goals with optimal path increase.
4.7.2 EARLY TERMINATION
We also test the flexibility of the learned model by slightly changing the task. While GLAMOR trained to find shortest paths to the goal state, we ran an additional experiment evaluating its performance if the agent is evaluated on its ability to get to the goal state in exactly T steps. For this experiment, we use the cartpole /bal Aballnce ance environment, since terminating early when the cart is in the correct position is significanly easier than manipulating it to reach that position at specific time. In Figure 36. we show that while early termination achieves the highest coal-achhevement rate, explicitly planning achieve the goal at the end of the trajectory outperforms naively taking shortest paths towards the goal state for T time-steps. All three experiments used the exact same learned model and the only difference is in the planning procedure, where "plan-end" does not sample the termination token until the last time-step. Videos the different termination strategies also available on the website.
Warde-Farley al. 2019) trained DISCERN for 200M steps in their paper.
8

---

Published as a conference paper at ICLR 2021
Files
Pitfal
Poug
PrivateEye
Achireed
Goals Achievel
Goals Achievel
Goals Achieved
0.4
0.2
0.0
0.00
Agent Steps (Millions)
Agent Stegu (Millions)
Agemt Steps (Millions)
Agent Steps (Millions)
Qbert
Seaquest
Skiing
Tennis
Goals Achiewe
Gonls Achieved
Con's Achieved
Gonks Achieved
Agent Steps (Millions)
Agent Stegs (Millions)
Aguut Stepe (Millions)
Ageut Steps (Millions)
GLAMOR (Ours)
Sparse Planner
Figure 4: Using intermediate information to guide the planning process helps GLAMOR achieve more goals than when it only looks at the estimated probability of reaching goal the end of the episode.
4.7.3 SPARSE PLANNING
Prior works such MuZero (Schrittwieser et al., 2019) and the Predictron (Silver et al., 2017b) also learn latent dynamics models and are similar to GLAMOR in that they predict rewards given starting state and action sequence. The reward (probability of reaching the goal state) predicted in GLAMOR however is in the form of the inverse dynamics model and action prior, as shown in Equation 8, and the intermediate probabilities predicted by the auto-regressive models are used to guide the search for good action sequence. To test the contribution of this guidance, we compare GLAMOR version where candidate action sequences are generated randomly and the highest scoring one under |Equation 8 is selected. FFiuure 4 shows that performance is significantly worse in the tested environments.
5 CONCLUSION
We have presented novel way to learn latent world models by modeling inverse dynamics. These models learn to track task-relevant dynamics for a diverse distribution of tasks and provide strong heuristic that enables efficient planning. We also demonstrate strong performance in both the low and high sample regime on challenging visual benchmark tasks.
Geal-acchevement tasks already have significant practical value. A next step is to extend GLAMOR to general reward functions. GLAMOR also learns a non-reactive policy. While combining non- reactive planning with Model Predictive Control has proven to be sufficient in many benchmark tasks, a natural future direction is to account for these types of stochastic environments.
ACKNOWLEDGEMENT
We gratefully acknowledge funding from the Natural Sciences and Engineering Research Council of Canada (NSERC), the Canada CIFAR AI Chairs Program, and Microsoft Research. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute for Artificial Intelligence www. www..eectorinstiitte...o wtt.::ttt....aaaat..eee.e.eeee
REFERENCES
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Mare- Mare-Alexandre Côté, and R. Devon Hjelm. Unsupervised state representation learning in atari. In Advances in Neural Information

---

Published as a conference paper at ICLR 2021
Processing Systems 32: Annual Conference Neural Information Processing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Canada, pp. 8766-8779, 2019.
Marein Andrychowicz, Dwight Crow, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech 7aremba. Hindsight experience replay. In Ad- vances Neural Information Processing Systems 30: Anmual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pp. 5048-5058, 2017.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Christopher Hesse, Rafal Józefowicz, Scott Gray, Catherine Olsson, Jakub Pachocki, Michael Petrov, Henrique Pondé de Oliveira Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, Ilya Sutskever, Jie Tang, Filip Wolski, and Susan Zhang. Dota with large scale deep reinforcement learning. arXiv preprint arXiv: 1912.06680, 2019.
Paul Christiano, Zain Shah, Igor Mordatch, Jonas Schneider, Trevor Blackwell, Joshua Tobin, Pieter Abbeel, and Wojciech Zaremba. Transfer from simulation to real world through learning deep inverse dy dynamics model. arXiv preprint arXiv: 1610.03518, 2016.
Lasse Espeholt, Hubert Soyer, Rémi Munos, Karen Simonyan, Volodymyr Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, Shane Legg, and Koray Kavukcuoglu. IMPALA: scalable distributed deep-rl with importance weighted actor-learner architectures. In Proceed- ings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüsson, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1415. PMLR, 2018.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML, 2019, 9-15 June 2019, Long Beach, Cali- fornia, USA, volume 97 of Proceedings of Machine Learning Research, pp. 2052-2062. PMLR, 2019. bttp::/tt////eeeeiiiiio
mlr. www.ppeesssv...... / www.futtlloooooooooooooooooooott 9a html
Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Devin, Benjamin Eysenbach, and Sergey Levine. Learning to reach goals via iterated supervised learning. arXiv preprint arXiv: 1912.06088, 2020.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Ro- man Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, 3-8 December 2018, Montréal, Canada, pp. 2455-24611 2018.
Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, Cali- fornia, USA, volume 97 of Proceedings of Machine Learning Research, pp. 2555 2565. PMLR, 2019.
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learn- ing behaviors by latent imagination. In 8th International Conference on Learning Representa- tions, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, 2020.
Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Gheshlaghi Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the Thirty-Second AAAI Con- ference on Artificial Intelligence, (AAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pp.
AAAI Press, 2018
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9:1735- 80, 12 1997.
10

---

Published as a conference paper at ICLR 2021
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In 8th International Conference on Learning Representations, ICLR 2020. Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL
net/ ii rygggggygvvv
Leslie Pack Kaelbling. Learning to achieve goals. In Proceedings of the 13th International Joint Conference on Artificial Intelligence (ICCAI-93,, pp. 1094-1098, 1993.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020.
Diederik Kingma and Max Welling, Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, 2014.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv: 1711.05101, 2019.
Volodymyr Mnih, Koray Kavukcuoglu. David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Toannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nat., 518(7540))55295533 2015.
Ashvin Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual re- inforcement learning with imagined goals. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Informa- tion Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, 3-8 December 2018, Montréal, Canada, pp. 9209-9220, 2018
Advances
Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. Curiosity- driven exploration by self- sule-suue prediction. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 2778-2787, 2017.
Brahma Pavse, Faraz Torabi, Josiah P. Hanna, Garrett Warnell, and Peter Stone. RIDM: reinforced inverse dynamics modeling for learning from single observed demonstration. arXiv preprint arXiv: 1906.07372, 2019.
Danilo J. Rezende, Ivo Danihelka, George Papamakarios, Nan Rosemary Ke, Ray Jiang, Theophane Weber, Karol Gregor, Hamza Merzic, Fabio Viola, Jane Wang, Jovana Mitrovie, Frederic Besse, Toannis Antonoglou, and Lars Buesing. Causally correct partial models for reinforcement learn- ing. arXiv preprint arXiv: 2002.02836, 2020.
A. Richards. Robust constrained model predictive control. PhD thesis, Massachusetts Institute of Technology, 2005.
Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards - just map them to actions. arXiv preprint arXiv: 1912.02875, 2019.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with a learned model. arXiv preprint arXiv: 191 1. 1.08265, 2019.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy P. Lillicrap, Karen Simonyan, and Demis Hassabis. Mastering chess and shogi by self-play with a general reinforce- ment learning algorithm. arXiv preprint arXiv: 1712.01815. 2017a.
11

---

Published as a conference paper at ICLR 2021
David Silver, Hado van Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac-Amold, David P. Reichert, Neil C. Rabinowitz, André Barreto, and Thomas Degris. The predictron: End-to-end learning and planning. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, ICMI. 2017, Syd- ney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Re- search, pp. 3191-3199. PMLR, 2017b. URL http://proceecings mlr pti.sppees/voc/
si silveri7a... marer1a..hmm html
Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojciech Jaskowski, and Jürgen Schmid- huber. Training agents using upside-down reinforcement learning. arXiv preprint arXiv: 1912.02877, 2019.
Adam Stooke and Pieter Abbeel. ripyt: A research code base for deep reinforcement learning in pytorch. arXiv preprint arXiv: 000,00550,, 2019.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Frez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy P. Lillierap, and Martin A. Riedmiller. Deepmind control suite. arXiv preprint arXiv: 1801.00690, 2018.
Josh Tenenbaum. Building machines that learn and think like people. In Elisabeth André, Sven Koenig, Mchdi Dastani, and Gita Sukthankar (eds.), Proceedings of the 17th International Con- ference Autonomous Agents and MultiAgent Systems, AAMAS 2018, Stockholm, Sweden, July 10-15, pp. 5. International Foundation for Autonomous Agents and Multiagent Systems Rich- land, SC, USA, 2018.
Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforce- ment learning. arXiv preprint arXiv: 1907.02057, 2019.
David Warde-Farley, Tom Van de Wiele. Tejas D. Kulkarni, Catalin Ionescu, Steven Hansen, and Volodymyr Mnih. Unsupervised control through non-parametric discriminative rewards. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019.
A APPENDIX
A.1
COMPARISON TO GOAL CONDITIONAL SUPERVISED LEARNING
Ghosh et al. 2020) propose a similar algorithm where a policy is trained by imitation learning with
a maximum- likelihood loss:
L(T) Ep log (a g,dls,g,... g, h)]
(10)
Here h is the time remaining until evaluation. GCSL iterates between training a policy to clone the behavior of the previous policy conditioned on achicving goal and gathering new training data.
The key differences between GCSL and GLAMOR are: (i) To act, GLAMOR solves optimiza- tion problem while GCSL samples directly from the policy; (ii) GLAMOR learns both an inverse dynamics model and the action prior while GCSL only learns a policy T(s, a, g, h); and (iii) Data for GCSL is collected with a reactive policy while GLAMOR uses an open-loop plan.
In experiments, we found that GCSL performed surprisingly well visual goal-achievement tasks despite the original implementation by Ghosh et al. (200)) only being tested on control tasks with low-dimensional observations.
Ghosh et al. 2020) provide a theoretical analysis of GCSL, and show that their behavioral cloning objective is bound on the RL objective, with looseness induced by both how off-policy the training distribution is and the relabeling step. To illustrate the sub-optimal behavior of GCSL caused by relabeling, consider simplified setting of one-step MDP and on-policy training data.
Proposition A.1. Let (g) be action that maximizes the probability p(s gl.). If there exist two goals g. g' such that p(g) > O, p(g') >> O. and p(s = gl)) > p(s = (g)) > 0, then one-step GCSL does not converge to an optimal policy.
12

---

Published as a conference paper at ICLR 2021
As example, consider game with state space S {1, 3, 4, 6} and action space {roll fair die, roll loaded die). Rolling the fair die will result in the state being uniformly set to the number that the die lands on, while the loaded die will always land on 1. The optimal goal-achieving policy evident: the agent should roll the loaded die with probability one only when the goal is to land on 1. Otherwise the agent should always roll the fair die. In this example, GCSL will fail converge but GLAMOR will find the correct policy.
Proof. We consider a one-step MDP with goal distribution p(g). In this MDP, we assume a deter- ministic initial state and denote the state the agent transitions to after taking action a as
For simplicity, we assume the policy Ty powerful enough to perfectly match its training distri- bution when trained with maximum-likellhonn Therefore, GCSL consists of the following iterated steps:
Gather trajectories with πT. This gives empirical distribution of goal-conditioned actions Pr41(als = g) p(a=g (a)
p(==g)
Distill Pr+ (als g) into x441
The update to the policy at each iteration depends both on the relative likelihood of transitioning to the goal with the environment dynamics and the probability of taking action a, pc(a). Clearly, Pe(a) is function of the previous policy иe and the goal distribution p(g):
pr(a)= Σ ((ll)))
(11)
Assuming the distillation is exact, the policy evolves like:
Ta/s = =9] p(s -=s==10 gla)
Σ ((dd))i(aal))
(12)
p(s g)
We then analyze the ratio of the probability of any sub-optimal action a to the probability of the optimal action a": me(alg)
(a)aa"
If there is goal interference, goals g, exist such that plg )p(s gla. (g))) and p(s (aa(gg)) > p(s = (g)) > 0.
Assume - is optimal. Then,
Te Te(a Te (a)(()))))
p(s (g')) Pld") (a. p(s gll,(g)) Eg ( (g) (g)
(13)
p(s g/a. (g (od')))'d 1)p(g') p(s = g a. (g)
(14)
(15)
Therefore, if nt-1 optimal, ne will again be sub-optimal and the policy will never converge.
A.2 CAUSALLY CORRECT MODELS
Rezende et al. (2000) explore the connection between model-based reinforcement learning and causality. model is causally correct learned model (z) As p(z) with respect to a set of interventions. In model-based RL, model is trained to predict some aspect of an environment. In order to use the learned model to predict the affect of a new policy in the environment, the model must be causally correct with respect to changes in the policy. Rezende et al. (020) show that some partial models, including MuZero, are not causally correct with respect to action sequence interventions.
13

---

Published as a conference paper at ICLR 2021
Hlyper-paramett
value
optimizer
AdamW
weight-decay
0.01
normalization
GroupNorm
learning-rate
5e-4
replay-ratio
eps-steps
3e5
eps-final
0.1
min-steps-learn
5e4
buffer size
le6
policy trials
50
state size
512
clip-p-actions
-3.15
lstm-hidden-dim
64
lstm-layers
1
train tasks
1000
Figure 5: Hyperparameters used to train GLAMOR.
As an example, consider training a model to predict whether a sequence of actions wins in the game of Simon Says. The model is trained to predict p( p(win|so, (1,.--,,., ag) data produced by some training policy. If this training policy is good at the game, the conditional probability p(win|ss so, a1,..., a&) for all action sequences, even though using the action sequence blindly in the real game would result in a much lower win-rate. This happens because the true data-generating process is actually dependent on intermediate states $1,..... .k......k. Sk. which the training policy has ac- cess to. By conditioning on some action sequence 8o, AI...., without modeling the intermediate states, these states become confounding variables. In order to predict the affect of taking a certain action sequence on the environment, what we really want to do is find p( win so, do(a)........ ...; (a:)).
To learn causally correct models with GLAMOR, we opt to simplify the data-generating process by using training policies that are independent of intermediate states. While this may hurt training some stochastic environments, we find that in our multi-task setting with relabeling, using non- reactive exploration has negligible effect.
A.3 EXPERIMENTAL DETAILS
Figure 5] shows the hyperparameters that were used train our method. While our method is substantially more simple than value-based methods like DISCERN, there are still few important hyperparameters. We found that tuning the replay ratio is important to balance between sample efficiency and wert-fitting in models. We also find that GLAMOR works best with large replay buffer and large model. We also found that avoiding selecting action sequences which have too low probability, similar to tricks used in beam search in NLP (Holtzman et al. 2020), increases the performance of our planner. To achieve this, we introduce a hyperparameter c slip-p-act i ions and only expand action sequence if the immediate log-probability under the inverse dynamics model is over this value.
A.4 ABLATIONS
In order to find which parts of GLAMOR contribute to its performance, we ran ablation study.
A.4.1 TRAINING POLICY
In Figure 6) we vary the way we construct the open-loop action sequence that is followed to collect training data. We vary it in two ways: the amount of compute used to create the plan and whether we consider the action prior. While the performance is poor when using only one planning sample, GLAMOR seems to work well Pong with as little as 5 samples. More interestingly, disabling the action prior during training seems to increase the variance of GLAMOR significantly. Without the action prior, the planning procedure will select action sequences that may have shown up more in
14

---

Published as a conference paper at ICLR 2021
Planning Trial
5 Planning Trials
10 Planning Trials
50 Planning Trials
0.6
0.6
0.6
0.6
0.5
0.5
0.5
0.5
Goal: Adhieved
Goab Achievel
Goals Achieved
Goals Acsieved
0.2
0.2
0.2
0.1
0.1
0.1
0.1
0.0
0.0
0.0
0.0
0
2
Agent Steps (Millione)
Agent Steps (Millions)
Agend Steps (Millions)
Agent. Steps (Millioons)
No Action Price
With Action Priot
Figure 6: In this experiment, we test how changing the training policy affects performance in Pong. As the amount of compute used in the planner during training increases, so does the performance of the evaluated agent. Using the action prior decreases the variance of the agent's 's performance.
Pitfall
Pong
Qbert
Seaqueest
1.0
1.0
1.0
1.0
Achioved
Goals Achiteved
Goals Achirvd
Goals Achieved
0.6
0.0
0.6
0.6
0.2
0.2
0.2
0.2
Agent. Steps (Millioee
Agenl Stepe (Millious)
Agenl Stepe (Millions)
Agent. Steps (Millions)
GLAMOR (Curs)
Figure 7: When training agents with off- off-policy data collected with random policy, GLAMOR outperfoms GCSL and can achieve most goals.
the training data simply due the training policy. We hypothesize that this effect can significantly hurt exploration and performance.
We also test the performance of both GLAMOR and GCSL when using a uniform training policy in Figure 7 We found that GLAMOR performs very well even when trained on completely off-policy data, while GCSL struggles.
A.4.2 PLANNER
We also evaluate whether the planner is necessary at test-time achieve strong goal-achievement performance. To test this, we ran an experiment where the planner simply takes trajectory sample and takes the first action Figure shows that while the agent still surprisingly achieves many goals in this setting, using the planner results in stronger policy. We interperate this result as showing that the heuristic search guided by the factored inverse dynamics and action prior strong and additional compute simply chooses plan among already good options.
A.5 ADDITIONAL. FIGURES
In Figure 9 and Figure 10, we plot the learning curves of GLAMOR, GCSL, and DISCERN trained for 5M agent steps. GLAMOR learns quickly compared the other algorithms and performs better asymptotically with the exception of a few environments (MsPacman, point-mass). Figure 11 and Figure 12] show the states achieved by GLAMOR attemping reach a specific goal. As noted by Warde-Farley et al. (2019), the manipul ator environment is difficult and no algorithm learned to
Note that this is exactly equivalent lowering the planning compute in implementation.
15

---

Published as a conference paper at ICLR 2021
Atari
Control
0.6
0.6
0.5
Goal: Achievel
Goals Achieved
0.1
0.1
0.0
0.0
3
0
Agent Stege (Millions)
Ageut Steps (Millions)
Search Trial
50 Search Tridls
Figure 8: Goal achievement rates for DM Control and Atari. Searching for high scoring action sequence results in more goals achieved. However, even when using compute equivalent model- free agent, GLAMOR performs remarkably well.
Bowling
Boxing
Breakout
Frostbite
Achieeed
1.0
Achioved
1.0
Achioved
1.0
Goals Achirved
1.0
0.5
0.5
0.5
0.5
0.0
0.0
0
Agent Steps (Millions)
Agent Steps (Millions)
Agent Steps (Millions)
Agentt Steps (Millions)
MssPacman
Pitfall
Pong
Goals Adhieved
1.0
Goals Adhieved
Goals Achiowe
Goals Arhieved
0.5
0.5
0.5
0.5
Agent Steps (Millione)
Agent Steps (Millious)
Agent Stepe (Millions)
Agent Steps (Millions)
PrivadeFye
Qbert
Riverraid
Seaquess
Achieeed
1.0
Achieved
1.0
Goals Achoved
3.0
Gouls Achieved
1.0
0.0
0.0
0.0
0.0
2
Agent Steps (Millions)
Agrnt Stegi (Millions)
Agent Stepe (Millions)
Agent Stepa (Millions)
Skingg
Tennis
Achioved
Adhteved
0.5
0.5
Agent Steps (Millions)
Agent Steps (Millions)
DISCERN
GCSL.
GLAMOR (Ours)
Figure 9: Training Curves for Atari Tasks. GLAMOR achieves more goals (often with many fewer steps) than both GCSL and DISCERN.
achieve goals within 5M steps. The low achievement rate on the finger task is due to the agent's inability to reliably control the angle of the spinner.
16

---

Published as a conference paper at ICLR 2021
ball cupp/attth
cartpole/balanee
finger /spin
manipnlator hariip/ntttto/brrrggggg hall
Achioved
1.0
Goals Achieved
1.0
Conls Achicved
1.0
Achieeed
1.0
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
3
2
3
Agros Stope (Millions)
Agent Steps (Millions)
Agontt Stops (Millions)
Agros Strpe (Millions)
poncuulu/sswingg
poiint_masss/oe /eausmasss//ey/
reacher/hor /hard
Goals Achieved
Goals Achieved
Goals Achieved
1.0
0.5
Ageus Stepe (Millionss
Agent Steps (Millions)
Agent Stegs (Million))
(Ours)
Figure 10: Training Curves for Control Tasks. GLAMOR achieves more goals (often with many
fewer steps) than both GCSL and DISCERN.
0.
19:15
A&A
d.
19:57
Figure 11: Goal states (above) and states achieved by the fully trained GLAMOR agent (below) averaged over 5 trials for each Atari game tested. Variance comes from environment and planning stochasticity. Note that most games, GLAMOR learns control the positions of both directly and indirectly controllable objects in the frame.
17

---

Published as a conference paper at ICLR 2021
Figure 12: Goal states (above) and states achieved by the fully trained agent (below) averaged over trials for each control task tested. Variance comes from planning stochasticity since the environment dynamics are deterministic. In most environments, GLAMOR learns to control the agent's state to match the visually specified goal. Note that in the finger environment, GLAMOR learns to control position of the finger despite not often being able to control the angle of the spinner.
18


Reviews:
[{"review_id": "YNM9ZNrfayX", "reviewer": "AnonReviewer3", "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": "The paper proposes a model-based reinforcement learning method. The method builds a partial model of the environment through learning inverse dynamics, which is the distribution of action sequences that would bring one state to another state. Through training the model with an iterative relabeling scheme, the model is able to learn to reach goals in a subset of DM Control and Atari domains. \n\nComments:\n+ This MBRL formulation is conceptually simple and does not rely on learning the full model of an environment, potentially enabling long-horizon planning while being able to high-dimensional input.\n+ The relationship and differences to the closely related work Ghosh et al., 2020 is clearly stated and elaborated.\n\n- My main concern about the paper is that although the difference between GCSL and the proposed method is exposited in theory (Proposition 3.1), there is no direct empirical experiment explaining how this difference might play out in practice. Specifically, there clearly are performance gaps between the proposed method and GCSL, but I cannot seem to grasp what EXACTLY is causing this gap. Is it ONLY because of the difference stated in proposition 3.1? If so, I'd like to see a minimum toy experiment clearly showing the qualitative differences. If there are other differences, I'd like to see them addressed in greater details. None of the reasons (i)-(iii) stated in section 3.6 are substantiated experimentally. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "X7sBnBgmSIa", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes to learn task-agnostic world models in latent space, using a goal-conditioned inverse dynamics formulation and an action priors model, that learn to predict action sequences. The approach is compared against a model-free unsupervised control method and an imitation learning method on the subset of Atari games and a continuous control benchmark, and showed to achieve superior performance in most tasks based on goal completion ratio and sample efficiency. \n\nThe proposed approach learns two networks, the inverse dynamics model and the action prior model, and uses this to plan a sequence of control actions, rather than learning the policy directly as done by the baselines. Although this has shown to achieve more goals in the 2 benchmarks, the authors do not discuss the learning complexity of training two networks compared to one policy network as done in baselines. The authors conduct an ablation study on the computational complexity of planning at test-time for their approach. I believe it would help to evaluate the method better if they also show the number of trials used when comparing to baselines. \n\nCan the authors provide any real-time performance guarantees at test time?\n\nIt's not clear to me how the exploration and exploitation is traded off in the Data Collection section. The planner is used to generate a sequence of actions but an e-greedy policy is used for exploration. This is confusing to me.\n\nThe paper claims that the inverse dynamics formulation learns to represent controllabe aspects of the state. Has this been theoretically proven before? And have any empirical studies been done to show that this actually helps model-based RL to perform better than when using forward models? It would be interesting to compare against model-based RL techniques based on forward models, if avaialble, for this work.", "extra_scores": {}, "extra_text": {}}, {"review_id": "SaAICBJEUS4", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: The author proposes Goal-Conditioned Latent Action Models for RL (GLAMOR) a novel approach to learn latent world models by modeling inverse dynamics. The proposed approach learns to track task-relevant dynamics for a diverse distribution of tasks and provide a strong heuristic that enables efficient planning. GLAMOR demonstrates good performance against its baselines in terms of achieving accurately goals, sample efficiency and effective planning.\n\n\nReason for the score:\nPaper address an interesting question related to Reinforcement Learning and provides strong results to back the proposed method.\n\n\nPros:\n-good flow of the paper\n\n-Correct situation of the problem with respect to current research\n\n-A good problem to address in Reinforcement Learning area\n\n-Strong experimental results\n\n\nCons:\n-Can the authors talk about how the proposed method perform in a real-robotic task?\n\n-In the equation 5, a_{<i} means?\n\n-in page 3 last para is not clear to me. Can you rewrite or provide explanation of how it is possible?\n\n-In para 3 of Section 3.4, why you sample from Boltzmann distribution and not uniform distribution?\n\n-Have you considered other exploration methods than of epsilon-greedy?\n\n\nQuestions for the rebuttal:\nPlease address and clarify the cons above \n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "-jzVQ1YrQbl", "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": "The authors adopted an interesting idea for learning inverse dynamics which is different from common approaches where oftentimes a state-transition model is learned for planning. Though in their scenario this does not make any guarantees on the improvement on generalization, conditioning on the rich semantics could be contained in actions, it might be a justification for learning the inverse dynamics model for tasks that test generalization capabilities. \n\nHowever, the idea of using relabeling to facilitate the learning of goal-conditioned policies and world models is not new as discussed by the authors. This makes the learning framework somewhat less significant in novelty. More specifically, the authors put quite a lot of effort into distinguishing their method from a recent method GCSL. The differences actually lead to a question on the action prior. I'm curious if this prior distribution is different from a uniform distribution over all available paths starting from $s_1$, it is not well justified why we need to parameterize this with LSTM. \n\nThere are also several, I believe, typos that might need fixing. The first term in Equation (9), should it be $\\mathbb{E}[-\\log p_{\\theta}(a_1, \\cdots, a_k | s_1, s_k)]? In the experiment section, the model was trained with 500k agent steps while in the figure it was 5M, this leads to the problem of sample efficiency (section 4.6).\n\nOverall, I feel that though the idea of adopting inverse dynamics for learning action model is interesting, it is still somewhat incremental to prior works with a similar framework. The missing justification and details also contribute to the decision of putting a decision of borderline during pre-rebuttal period.\n\n================================================================================================================\n\nThe authors has addressed my concerns well in the discussion period, I have increased the score to 6.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: While many work in the literature (PlaNet (2018), Dreamer (2020), SimPLe (2019), etc.) learn world models to perform well on a particular task at hand, the motivation behind this work is that dynamics models benefit if they are task-agnostic, hence would be able to perform a wider range of tasks, as opposed to just doing one task really well. In order to do this, they propose to learn a latent representation that models inverse dynamics of the system / environment rather than capturing information about the task-specific rewards, and incorporate a planning for solving specific tasks in which they can measure performance.

To show broad applicability of their method, the authors tested their approach on Atari and DM Control Suite (from pixels), and also simple grid worlds to illustrate the concepts, and demonstrated strong performance over SOTA model-free algorithms (even the ones that do not have open-source implementations). Reviewers and myself agree that the paper is well written, easy to follow, and the approach is well-motivated.

After the review period, the authors have done work to improve the draft, particularly including ablation studies with and without planning, addition comparisons, and improved visualizations, after taking in the comments and feedback from the reviewers after the initial reviews, which satisfied some of the reviewers. One reviewer asked for a real robotic task, but I feel that while it will help the paper, many existing works focus purely on DM control from pixels, and this work has performed experiments on both DM Control and Atari, two reasonably different domains, and IMO makes up for the lack of real-world robotics experiment. That being said, a discussion on how the proposed method would work in a real-robotic task, as suggested by R4 would be good to have.

I believe the work in its current state is ready for acceptance for ICLR 2021, and should be a fine contribution to the visual model-based RL works. I'm excited to see this work presented to the community, and I'm going to recommend acceptance (Poster).
Decision: Accept (Poster)