Planning with Sequence Models through Iterative Energy Minimization
Authors: Hongyi Chen
Yilun Du
Yiye Chen
Joshua B. Tenenbaum
Patricio A. Vela
Venue: ICLR 2023 poster (2023)
Keywords: Reinforcement Learning
Planning
Language Model
Decision Transformer

Abstract: Recent works have shown that language modeling can be effectively used to train reinforcement learning (RL) policies. However, the success of applying existing language models to planning, in which we wish to obtain a trajectory of actions to reach some goal, is less straightforward. The typical autoregressive generation procedures of language models preclude sequential refinement of earlier steps, which limits the effectiveness of a predicted plan. In this paper, we suggest an approach towards integrating planning with language models based on the idea of iterative energy minimization, and illustrate how such a procedure leads to improved RL performance across different tasks. We train a masked language model to capture an implicit energy function over trajectories of actions, and formulate planning as finding a trajectory of actions with minimum energy. We illustrate how this procedure enables improved performance over recent approaches across BabyAI and Atari environments. We further demonstrate unique benefits of our iterative optimization procedure, involving new task generalization, test-time constraints adaptation, and the ability to compose plans together. Project webpage: https://hychen-naza.github.io/projects/LEAP/index.html

Published as a conference paper at ICLR 2023
PLANNING WITH SEQUENCE MODELS THROUGH ITER- ATIVE ENERGY MINIMIZATION
Hongyi Chen" Yilun Du * Yiye Chen",, Joshua Tenenbaum', Patricio Vela' Georgia Institute of Technology
Massachusetts Institute of Technology
ABSTRACT
Recent works have shown that sequence modeling can be effectively used to train reinforcement learning (RL) policies. However, the success of applying existing sequence models to planning, in which we wish to obtain a trajectory actions to reach some goal, is less straightforward. The typical autoregressive generation procedures of sequence models preclude sequential refinement of carlier steps, which limits the effectiveness of predicted plan. In this paper, suggest approach towards integrating planning with sequence models based on the idea of iterative energy minimization, and illustrate how such procedure leads im- proved RL performance across different tasks. We train masked language model capture implicit energy function over trajectories of actions, and formulate planning as finding a trajectory of actions with minimum energy. We illustrate how this procedure enables improved performance over recent approaches across BabyAl and Atari environments. We further demonstrate unique benefits of our iterative optimization procedure, involving new task generalization, test-time con- straints adaptation, and the ability to compose plans together. Project website: https://hychen-neee naza, maw.aaigtttuu../o/oocsLLLEEEEE
INTRODUCTION
Sequence modeling has emerged as unified paradigm to study numerous domains such as language (Brown al., 2020; Radford et al., 2018) and vision (Yu et al., 2022; Dosovitskiy al., 2020). Recently, (Chen et al., 2021; Jan- ner et al., 2021) have shown how similar ap- proach can be effectively applied to decision making, by predicting the next action to take. However, in many decision making domains, is sub-optimal to simply predict the next action to execute as such an action may only lo- cally optimal and Iead to global dead-end. In- stead, it is more desirable to plan sequence of actions towards final goal, and choose the action most optimal for the final overall goal.
Ea(ti-1)
Ea(+.)
Z
High
Low
Energy
Figure 1: Plan Generation through Iteratively En- ergy Minimization. LEAP plans trajectory goal (specified by the yellow star) by iteratively sampling and minimizing trajectory energy function estimated using language model Es.
Unlike greedily picking the next action to execute, effectively constructing wards a given goal requires careful, iterative procedure, where we need to assess and refine in- termediate actions in plan ensure we reach the final goal. To refine action particular timestep in a plan, we must reconsider both actions both before and after the chosen action. Directly applying this procedure standard language generation difficult, as the standard autoregressive decoding procedure prevents regeneration of previous actions based of future ones. For example, if the first five predicted actions places an agent at a location too far reach a given goal, there is no manner we may change the early portions of plan.
action sequence to-
In this paper, we propose an approach to iteratively generate plans using sequence models. Our approach, Multistep Energy-Minimizattion Planner (LEAP), formulates planning as an iterative op-
denotes equal contribution. Correspondence hchen Bggatech, edu,, yi y undugnit eddo, yychen281 99gttechh.eed

---

Published as a conference paper at ICLR 2023
timization procedure energy function over trajectories defined implicitly by sequence model (illustrated in Figure 1). To define an energy function across trajectories, we train a bidirectional sequence model using maasked-language modeling (MLM) objective (Devlin et al., 2019). We de- fine the energy of a trajectory as the negative pseudo-likelihodd (PI..) of this ML.M (Salazar et al,, 2019) and sequentially minimize this energy value by replacing actions at different timepoints in the trajectory with the marginal estimates given by the MLM. Since our MLM is bi-directional in nature, the choice of new action at given time-step is generated based on both future and past actions.
By iteratively generating actions through planning, we illustrate how our proposed framework out- performs prior methods in both BabyAI Schevalier-Boiseeer et al., 2019) and Atari (Bellemare et al,, 2013) tasks. Furthermore, by formulating the action generation process as an iterative energy mini- mization procedure, we illustrate how this enables us to generalize to environments with new sets of test-time constraints as well as more complex planning problems. Finally, we demonstrate how such an energy minimization procedure enables us to compose planning procedures in different models together, enabling the construction of plan which achieves multiple objectives.
Concretely, in this paper, we contribute the following: First, we introduce LEAP, a framework through which we may iteratively plan with sequence models. Second, we illustrate how such planning framework can be beneficial on both BabyAl and Atari domains. Finally, we illustrate how iteratively planning through energy minimization gives set of unique properties, enabling better test time performance on more complex environments and environments with new test-time obstacles, and the ability to compose multiple learned models together, to jointly generate plans that satisfy multiple sets of goals.
RELATED WORK
Sequence Models and Reinforcement Learning. Sequence modeling with deep networks, from gequence-to-squeeee models (Hochreiter & Schmidhuber, 1997; Sutskever et al., 2014) to BERT (Devlin et al., 2019) and XLnet (Yang et al., 2019), have shown promising results in a series of lan- guage modeling problems (Dai et al., 2019b; Sutskever et al., 2014; Liu & Lapata, 2019; Dehghani et al., 2018). With these advances, people start applying sequence models to represent components in standard RL such as policies, value functions, and models to improved performance (Espeholt et al., 2018; Parisotto et al., 2020; Kapturowski et al., 2018). While the sequence models provide memory information to make the agent predictions temporally and spatially coherent, they still rely on standard RL algorithm to fit value functions or compute policy gradients. Furthermore, recent works replace much of the RL pipeline possible with sequence modeling to leverage its scala- bility, flexible representations and causally reasoning (Janner et al., 2021; Chen et al., 2021; Furuta et al., 2021; Zheng et al., 2022; Emmons al., 2021; Li al., 2022). However, those methods adopt autoregressive modeling objectives and the predicted trajectory sequences have no easy way to be optimized, which will inevitably lower the long- nonghoorizon accuracy. Recent studies point out that using sequence models (Chen et al., 2021; Emmons et al., 202 1) rather than typical value-based approaches have difficulty converging in stochastic environments (Paster et al., 2020; 2022).
Planning in Reinforcement Learning. Planning has been explored extensively in model-based RL, which Iearns how the environment respond to actions (Sutton, 1991). The learned world dy- namic model is exploited to predict the conditional distribution over the immediate next state or autoregressively reason the long-term future (Chiappa et al., 2017; Ke et al., 2018). However, due to error compounding, plans generated by this procedure often look more like adversarial examples than optimal trajectories when the planning horizon is extended (Bengio et al., 2015; Talvitie, 2014; Asadi et al., 2018). To avoid the aforementioned issues, simple gradient-free method like Monte Carlo tree search (Coulom, 2007), random shooting (Nagabandi et al., 2018) and beam search (Sun et al., 2022) are explored. Another line of works studied how to break the barrier between model learning and planning, and plan with an imperfect model, include training an autoregressive latent- space model to predict values for abstract future states (Tamar et al., 2016; Oh et al., 2017; Schrit- twieser et al., 2020; Sun et al., 2022); energy-based models of policies for model-free reinforcement learning (Haarnoja et al., 2017); improve the offline policies by planning with learned models for model-based reinforcement learning (Yu et al., 2020; Schrittwieser et al., 2021); directly applying collocation techniques for direct trajectory optimization (Erez & Todorov, 2012; Du et al., 2019); and folding planning into the generative modeling process (Janner et al., 2021). In contrast to these works, we explore having planning directly integrated in language modeling framework.
2

---

Published as a conference paper at ICLR 2023
Energy Based Learning. Energy-Basee Models (EBMs) capture dependencies between variables by associating a scalar energy to each configuration of the variables, and provide a unified theoret- ical framework for many probabilistic and non-probabilistii approaches to learning (LeCun et al., 2006). Prior works have explored MBMs for policy training in model-free RI. (Haarnoja et al., 2017), modeling the environment dynamics in model-based RL (Du et al., 2019; Janner et al., 2022) and natural images (Du & Mordatch, 2019; Dai et al., 2019a), as well as energy values over text (Goyal al., 2021). Most similar our work, Du al. (2019) illustrates how energy optimization in EBMs naturally support planning given start and goal state distributions. However, the underlying training relies on a constrastive divergence, which difficult to train (Du et al., 2020). On the other hand, our training approach relies on more stable masked language modeling objective.
METHOD
In this section, we describe our framework, Multistep Peergy-Miinmizzaiion Planner (LEAP), which formulates planning as a energy minimization procedure. Given a set of trajectories in a discrete ac- tion space, with each trajectory containing state and action sequences ($), a S2, ......, SN ax), our goal is learn a planning model which can predict a sequence of actions a1:T, given the tra- jectory context Tatz containing the past K steps states and actions, that maximizes the long-term uusk-dependent objective J:
a arg max I (Teta; aI:T)
a1.T
where N denotes the length of the entire trajectory and T is the planning horizon. We use the abbreviation [TT,, where T =: (Tetz; a1:7), to denote the objective value of that trajectory. To formulate this planning procedure, we learn an energy function Ea(T), which maps each trajectory T to a scalar valued energy so that
a air = arg min ET(Teta> aI.T).
airT
3.1 LEARNING TRAJECTORY LEVEL ENERGY FUNCTIONS
We wish to construct an energy function (T a1:r) such that minimal energy is assigned to optimal set actions ai.y- To train our energy function, we assume access to dataset of M near optimal set of demonstrations in the environment, and train our energy function to have low energy across demonstrations. Below, we introduce masked language models, and then discuss how we may get such a desired energy function from masked language modeling.
Masked Language Models. Given a trajectory of the form (s1, a1 S2 ag..... a,), we train transformer language model to model the marginal likelihood pe po Tata. a.r), where we utilize a as shorthand for actions a1.7 except the action at timestep 1. To train this masked language model (MLM), we utilize the standard BERT training objective (Devlin et al., 2019), where we minimize the loss function
LMLM Ey. |- log (a Teta a-r]],
where we mask out and predict the marginal likelihood of percentage of the actions in trajectory (details on masking in the Section A.4).
Constructing Trajectory Level Energy Functions. Given a trained MLM, we define an energy function for trajectory as the sum of negative marginal likelihood of each action in sequence
Ea(T) = Σ log po (as; Teter a)).
Such energy function, also known as the pseudolikelihodd of the MLM, has been used exten- sively in prior work in NLP (Goyal et al., 2021; Salazar et al., 2020), and has been demonstrated effectively score the quality natural language text (outperforming direct autoregressive scoring) (Salazar et al., 2020). In our planning context, this translates to effectively assigning low energy to optimal planned actions, which our desired goal for (T). We illustrate the energy computation process in Figure 2.
3.2 PLANNING AS ENERGY MINIMIZATION
Given a learned energy function (T(T), which assigns low energy to optimal trajectories ai.7- we wish to plan sequence of actions in test-time minimize our energy function. To implement this
3

---

Published as a conference paper at ICLR 2023
Trajectory Energy
Energy Minimization
i+
Σ log
X
Sampler
log
Enengy Estimation
MLM
s
0i
ai
of
T
Horizon
Figure 2: Energy Minimization LEAP. LEAP generates plans via Gibbs sampling different actions based on a learned trajectory energy model Es (T). In each iteration, Masked Language Model (MLM) predicts the energy of alternative actions selected timesteps in trajectory. A trajectory is generated using Gibbs sampler, with individual actions sampled based on the energy distribution. By repeating above steps iteratively, I.EAP generates the trajectory with low energy value.
planning procedure, utilize Gibbs sampling of individual actions each timestep the energy function (T) for low energy plans, which we detail below.
For test-time plan generation, initialize masked trajectory of length T, with small context length of past states and actions, which we illustrate in Equation 3. At each step of Gibb's sampling, randomly mask out to estimate their energy distribution conditioned on trajectory context Ti which is the masked
multiple action tokens in padded locations and perform forward pass
outcome from previous iteration on sampled timesteps collected in index set Z. Then, action is sampled using Gibb's is sampling based the locally normalized energy score a, pa(ar; Teta a-1). The process is illustrated in Figure 2, where the actions with low energy values (in blue) are sampled minimize (T 0) in each iteration. To sample effectively from the underlying energy distri- bution, repeat this procedure for multiple timesteps, which we illustrate in Algorithm 1, The computational time of Algorithm increases linearly with iteration numbers.
81
82
Sp
Sy PAD
Sn PAD
Sn [PD
(3)
plan
Algorithm Iterative Planning through Energy Minimization (for discrete actions)
Require trained energy model hy, context trajectory
Pad the states and actions with length T into context trajectory
3: for I,...,N do
Sample
1e [1,2....,
Bal DeLimeLe the energy distributions maaked Lokena
f(h(m((:9))
8:
// Sample the action tokens based on energy value
E
10:
Update actions a ir. 7 at manked tokans 7/.
11:
12:
end
13: Execute all planned actions a the first planned action in puddedd trujectory
Note that our resultant algorithm has several important differences compared to sequential model based approaches such Decision Transformer (DT) (Chen et al., 2021). First, actions gen- crated using an energy function defined globally across actions, enabling us to choose each action factoring in the entire trajectory context. Second, our action generation procedure is iterative in nature, allowing us to leverage computational time to find better solutions towards final goals.
4 PROPERTIES OF MULTISTEP ENERGY-MINIMIZATI
PLANNER
In LEAP, we formulate planning as an optimization procedure arg min, ET(T). By formulating planning in such manner, we illustrate how approach enables online adaptation new test-time

---

Published as a conference paper at ICLR 2023
(b)
Plan
Online
Train
Test
(c)
Initial
Invisible
Ostacle
Goal State
Visible
O bsttale
Run- time
Trijectoy
bstacle
Plarner
P Paumer
Composite Fkn
Figure 3: Properties Plamning as Energy Minimization. By formulating planning as energy minimization, LEAP enables the following properties: (a): Online adaptation; (b): Generalization; (c): Task composition.
constraints, generalization harder planning problems. and plan composition achieve multiple set of goals (illustrated in Figure 3).
Online adaptation. In LEAP, recall that plans are generated by minimizing energy function ET(T) across trajectories. At test time, if we have new external constraints, we maybe correspond- ingly define energy function Exomstrain (T) (T) encode these constraints. For instance, if state becomes dangerous at test time (illustrated as red grid in Figure (a)) we may directly define an energy function which assigns 0 energy plans which do not utilize this state and high energy to plans which utilize such state. We may then generate plans which satisfies this constraint by simply minimizing the summed energy function
arg min( (T) Econstrintt (T)
and then utilize Gibb's sampling to obtain plan T" from Eo(T) + Exunstraim (T). While online adaptation may also integrated with other sampling based planners using rejection sampling, our approach directly integrates trajectory generation with online constraints.
Novel Environment Generalization. In LEAP, we leverage many steps of sampling to recover an optimal trajectory r" which minimizes our learned trajectory energy function EB(T). In settings at test time when the environment is more complex than those seen training time (illustrated in Fig- ure 3 (b)) the underlying energy function necessary to compute plan feasibility may remain simple (i.e. measure actions enter obstacles), but the underlying planning problem becomes much more difficult. In these settings, long the learned function (T(T) generalizes, and may simply leverage more steps of sampling to recover a successful plan in this more complex environment.
Task compesitionality. Given two different instances of LEAP, E (T) and E3(T) encoding sep- arate tasks for planning. we may generate trajectory which accomplishes the tasks encoded by both models by simply minimizing composed energy function (assuming task independence)
T" arg min Ei! (T) + (T)).
(5)
An simple instance of such a setting is illustrated Figure (c), where the first LEAP model (T) encodes two obstacles in environment, and second LEAP model (T(7) encodes four other obstacles. By jointly optimizing both energy functions (through Gibbs sampling), we may successfully construct plan which avoids all obstacles in both models.
EXPERIMENTS
In this section, we evaluate the planning performance of LEAP in BabyAl and Atari environments. We compare with variety of different offline reinforcement learning approaches, and summarize the main results in Figure

---

Published as a conference paper at ICLR 2023
Ours
Decision Trmmeommer
Performmec
BabyAl
Atai
Genoralization
Figure 4: Quantitative Results of LEAP of Different Domains. Results comparing LEAP to Decision Transformer and TD learning (IQL BabyAl and Generalization tests and CQL in Atari) across BabyAI, Atari, and Generalization Tests. On a diverse set tasks, LEAP performs better than prior approaches.
5.1 BABYAI
Setup The BabyAI comprises extensible suite of tasks in environments with different sizes and shapes, where the reward is given only when the task successfully finished. We evaluate models in trajectory planning that requires the agent move the goal object through turn left, right and move forward actions, and general instruction completion tasks in which the agent needs to take extra pickup, drop and open actions complete the instruction. For more detailed experimental settings and parameters, please refer to Table in Appendix.
Baselines We compare our approach with set of different baselines: Behavior Cloning algo- rithm (BC); model free reinforcement learning (RL) algorithms BatchhConsstrineed deep Q-Learning (BCQ) (Fujimoto et al., 2019), and Implicit Q-Learning (IQL) (Kostrikov al., 2021); model based RL algorithm Model-based Offline Policy (MOPO) (Yu et al., 2020); return-conditioning approach Decision Transformer (DDT) (Chen al., 2021); )); model based Planning Transformer (PlaTe) (Sun et al., 2022). In our experiments, we use, as model inputs, the full observation the environment, the instruction, the agent's current location and the goal object location (if available).
Results Across all environments, LEAP achieves highest success rate, with the margin magnified on larger, harder tasks, see Table 1. In particular, LEAP could solve the easy tasks like GoToo caIS8N7 with nearly 90% success rate, and has huge advantages over baselines in the large maze worlds (GozeO)jMazee7) and complex tasks (GoToSegS5R2) which require going to sequence of objects in correct order. In contrast, baselines perform poorly solving these difficult tasks.
Next, visualize the underlying iterative planing and execution procedure in GoToLoocllS8y and GoToSeqS5R2 tasks. On the left side of Figure 5, we present how the trajectory is optimized and its energy is minimized single time step. Through iterative refinement, the final blue trajectory is closer to the optimal solution than the original red one, which follows the correct direction with higher efficiency and perform the actions like open the door correct situation. On the right side, we present the entire task completion process through many time steps. LEAP successfully plans efficient trajectory visit the objects and opens the door when blocked. We also explore the model performance in the stochastic settings, please refer to Appendix B.
Task
Env
BCQ
IQL
PlaTe
MOPO
LEAP
GoToLoealS7NS
71.5%
73.0%
42.5%
Trajectory Planning
GoToLocalS8N7
61.5%
63.0%
71.5%
63.5%
45.0%
81.5%
95.0%
Go1oObiMazeS4
24.0%
23.0%
52.5%
46.5%
35.5%
60.0%
62.5%
GaToohhjMare7
18.0%
10.5%
29.0%
22.0%
27.5%
30.5%
42.5%
PickUpLoc
57.5%
58.5%
41.0%
59.5%
7.5%
43.5%
67.0%
Instruction
GoToSegS5R2
13.5%
10.0%
28.5%
26.5%
25.0%
30.0%
33.0%
Completion
24.0%
22.5%
31.5%
48.5%
32.5%
36.0%
55.0%
rate of LEAP and variety of prior algorithms BabyAl env. Models trained with 500 optimal trajectory demos each environment, and results craged over random seeds. The SX and NY environment name means its size and the number of obstacles. Effect of Iterative Refinement. We investigate the effect of iterative refinement by testing the success rate of our approach under different sample iteration number 5 environ- ment. From the left side of Figure 6, the task success rate continues to improve as we increase the number sample iteration.
Table 1: Ba Quantitative Performance. The task
Energy Landscape Verification. We further verify our approach by visualizing the energy as- signment on various trajectories in the same environment above. More specifically, com-

---

Published as a conference paper at ICLR 2023
Low
Plan
Execute
Energy
Go to the biu bull
(b)
Open
ball go groen key
High
N=10
Figure 5: Qualitative Visualization of Planning and Execution Procedure BabyAL. Left depicts the planning through iterative energy minimization where N is the sample iteration number. Right shows the execution of the concatenate action sequences. Two task settings are illustrated: (a):
Trajectory planning, where the task is to solely plan trajectory leading to the goal location. (b): Instruction completion, where sequence of tasks are commanded, and an additional "Open" involved to get through the doors. Target locations are marked with
Lock
Seaent Level 19%
92
Teaun Leve
Sucoes Rale (K)
Loake'
70
Sample Number
Figure 6: Analysis of LEAP in the BabyAI Environments. Left: Success rate increases with more sampling steps, suggesting the importance of iterative refinement in LEAP, Right: LEAP captures the correct energy landscape. assigns low energy the optimal trajectory (noise level=0%% and high energy noisy paths.
pare the estimated energy of the labeled optimal trajectory with the noisiness of trajectories, pro- duced by randomizing a percentage of steps in the optimal action sequence. The right Figure depicts the energy assignment to trajectories with various corruption levels as a function of train- ing time. With the progress of training, LEAP learns to (a) reduce the energy value assigned the optimal trajectory; (b) increase energy value assigned to the corrupted trajectories. This result justifies the performance of LEAP.
Optimal
Suboptimal25%)
Effeet of Training Data. In BabyAl, we utilize set of demonstrations generated using oracle planner. We further investigate the performance of LEAP when the training data is not optimal. To achieve it, randomly swap the decisions in the optimal demonstration with an arbitrary action with the probability of 25%. We compare against DT, the autoregressive sequential planner. Despite small performance drop Table LEAP still substan- tially outperforms DT, indicating LEAP works well with non-perfect data.
Roy
DT
LEAP
DT
LEAP
GoTooRRjMane8
46.5%
63.0%
GoToohjjMazs?
GolkSegS3R2
26.5%
38.0%
29.5%
37.5%
Table 2: Performance on Suboptimal Data The task success rate of LEAP and DT, using optimal trajectories and subopti- mal trajectories containing 25% random ac- tions respectively as training data.
5.2 ATARI
Setup We further evaluate our approach on the dynamically changing Atari environment, with higher-dimensionall visual state. Due to above features we train and test our model without the goal state, and update the plan after each step of execution to avoid unexpected changes in the world. We compare our model to BC, DT (Chen et al., 2021). CQL (Kumar et al., 2020), REM (Agarwal et al., 2020), and QR-DQN (Dabney al., 2018). Following Chen et al. (2021), the evaluation is conducted on four Atari games (Breakout, Qbert, Pong, and Seaquest), where 1% of data is used for training. Human normalized score utilized for the performance evaluation.

---

Published as a conference paper at ICLR 2023
Breakout
Pong
R 0.246
R² 0.516
Game
LEAP
DT
CQL QR-DON REM
BC
Nomalleg Rerands
Breakout
400.9:1-5..0
267.5-456.3
211.1
17.1
138.9
Qbert
19.5+1.6
16.4=6.6
104.2
0.0
0.0
17.3
Nomalied
Pong
108.9+1.6
106.1:4.7
111.9
18.0
85.2
Seaquest
1.3.1.22
2.5.10.2
1.7
0.4
0.7
2.1
Avg
127.2
97.9 107.2
2.5 00.9
00
Table Quantitative Comparison on Atari. Gamer- normalized scores for the 1% DQN-replay Atari dataset (Agarwal et al., 2020). We report the mean and standard error score across 5 sceds. LEAP achieves best aver- aged scores over games and performs comparably DT and CQL over all games.
Nomralied
Nomma Szed Eren
Figure 7: Energy vs. Reward on Atari En- ergies and rewards are normalized to [0, 1]. We demonstrate negative correspondence between the achieved rewards and estimated energy by LEAP, which justifies our method.
Results Table 3 shows the result comparison. LEAP achieves the best average performance. Specifically, achieves better or comparable result in games out of 4, whereas baselines typi- cally perform poorly in more than one games.
Energy Landscape Verification In Atari environment, the training trajectories are generated by an online DQN agent during training, whose accumulated rewards are varied a lot. LEAP is trained to estimate the energy values of trajectories depending on their rewards. In Figure 7, we visualize the estimated energy of different training trajectories and their corresponding rewards in Breakout and Pong games. We notice that the underlying energy value estimated to a trajectory is well cor- related with its reward, with low energy value assigned to high reward trajectory. This justifies the correctness of our trained model and further gives a natural objective to assess the relative of quality planned trajectory. In Qbert and Seaquest games that LEAP gets low scores, this negative correlation is not obvious showing that the model is not well-trained.
PROPERTIES OF MULTISTEP ENERGY- -WN-MMIZIIIZOOO PLANNER
Next, we analyze the unique properties enabled by LEAP described in $4 in customized Baby environments. For each environment, design at most three settings with increasing difficulty levels to gradually challenge the planning model. As before, the target reaching success rate is measured as the evaluation criteria. The performances are compared against Implicit Q-Learning (IQL) (Kostrikov et al., 2021) and Decision Transformer (DT) (Chen et al., 2021).
6.1 ONLINE ADAPTATION
Setup To examine LEAP's adaptation ability to test-time constraints, we construct an BabyAl environment with multiple lethal lava zones at the test time as depicted in Figure 8 (a). The planner Eg generates the trajectory without the awareness of the lava zones. Once planned, the energy prediction corrected by the constraint energy function Econstrinn (T), which assigns large energy to the immediate action leading to lava, and zero otherwise. The agent traverses under the guidance of the updated energy estimation. To make the benchmark comparison fair, we also remove the immediate action that will lead to a lava for all baselines. The difficulty levels are characterized by the amount of lava grids added and the way they are added, where easy, medium correspond to adding most 2 and 5 lava grids respectively on the way the goal object in 8x8 grids world. The third case is hard due to the unstructured maze world in which the narrow paths can be easily blocked by lava grids and requires the agents to plan trajectory bypass.
The quantitative comparison is collected in the Table 4, Left. Although drops with harder challenges, the performance of our model still exceeds both baselines under settings. Visual illustration of medium example results can be seen in Figure 8 (a) that the agent goes up first to bypass the lava grids and then drives to the goal object.
Results
6.2
NOVEL ENVIRONMENT GENERALIZATION
Setup To evaluate LEAP's generalization ability in unseen testing environments, we train the model in easier environments but test them in more challenging environments. In easy case, the model is trained in 8x8 world without any obstacles but tested in the world with 14 obstacles as distractors. In medium and hard cases, the model is trained in single-room world but tested in maze
8

---

Published as a conference paper at ICLR 2023
Online Adaptation
Generalization DT
Task Composition
Test
LEAP
IQL
LEAP
LEAP
DT
IQL
92.0%
68.0%
90.3%
77.5%
36.0%
60.5%
83.5%
58 0%
42.5%
Medium
64.5%
20.0%
32.0%
64.0%
37.5%
57.5%
43.0%
15.5%
11.5%
44%
65.5%
N/A
Table 4: Property Test on Modified BabyAI Environments. Three properties performance of LEAP and prior algorithms BabyAl tasks with different difficulty. Left: Online Adaptation; Middle: Generalization; Right: Task Composition.
(a)
(b)
CAS SAS 2012
Train
Test
Train
Test
(c)
x
LEAP
LEAP
Composite LEAP
Figure 8: Qualitative Visualization of Generalization Tests. (a): Online adaptation (medium), trained in plane world and tested world with lavas; (b): Generalization (hard), trained in plane world and tested in maze world: (c): Task composition (easy), each model only perceive half number of obstacles. Target locations and unperceivable obstacles marked with and 8. respectively.
world containing multi-rooms connected by narrow paths (Figure 8 (b)). The maze size and the number rooms hard case are 10x 10 and 9, which are larger than 7x7 and medium case.
Results Our model achieves best averaged performance across three cases, but slightly worse than IQL in hard case, see Table 4, Middle. In contrast, sequential RL. model DT has significantly lower performance when moved to in unseen maze environments. LEAP trained plane world could still plan decent trajectory in unseen maze environment after blocked by walls, see Figure 8 (b).
6.3 TASK COMPOSITIONALITY
Setup We design composite trajectory planning and instruction completion tasks for easy and medium cases respectively. In easy case, all obstacles are equally separated into two subsets, each observable by one of the two planners, see Figure 8 (c). As result, the planning needs add up model's predictions using two partial observations to successfully avoid the obstacles. In medium case. two separate models trained for different tasks: one for planning trajectory in 10x 10 maze world and the other for object pickup in single-room world. The composite task is to complete the object pickup in 10x maze world.
Results Our model significantly outperforms the baselines in both two testing cases, while IQL and DT suffer great success rate drop indicating they can not be applied to composite tasks directly, see Table 4. Right. This proves that LEAP can be easily combined with other models responsible for different tasks, making it more applicable and general for wide-range tasks. In Figure (c), the composite LEAP could reach the goal by avoiding all obstacles even though the first LEAP planner is blocked by unperceivable obstacle.
CONCLUSION
This work proposes and evaluates LEAP, an sequence model that plans and refines a trajectory through energy minimization. The energy- morgyymiinmizzatonn is done iteratively where actions quentially along a trajectory. Our current approach is limited to discrete spaces relaxing this using approaches such discrete binning (Janner et al., 2021) would be interesting.

---

Published as a conference paper at ICLR 2023
REFERENCES
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning, 2020. 7, 8
Kavosh Asadi, Dipendra Misra, and Michael Littman. Lipschitz continuity in model-based rein- forcement learning. In International Conference on Machine Learning, pp. 264-273. PMLR, 2018. 2
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013. 2
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Voss, Gretchen Krueger. Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffiey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement Jearnnn via sequence modeling. Advances in neural information processing systems, 34: 14:11584-15097, 2021. 1, 2, 4, 6, 7, 8, 16
Maxime Chevalier-Boisvertt Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. AabyAl: First steps towards grounded language learning with a human in the loop. In International Conference on Learning Representations, 2019. 2
Silvia Chiappa, Sébastien Racaniere, Daan Wierstra, and Shakir Mohamed. Recurrent environment simulators. arXiv preprint /704. 02254, 2017.
Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In Computers and Games, pp. 72-83, 2007.
Will Dabney, Mark Rowland, Mare Bellemare, and Rémi Munos. Distributional reinforeement learning with quantile regression. In Conference Artificial Intelligence, 2018. 7
Bo Dai, Zhen Liu, Hanjun Dai, Niao He, Arthur Gretton, Le Song, and Dale Schuurmans. Expo- nential family estimation via adversarial dynamics embedding. Advances in Neural Information Processing Systems, 32, 2019a.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv: 0tt.//911.200.. 2019b.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals. Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint urXiv:1807.0388 2018.
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long and Short Papers), pp. 4171-4186. 2019. 2,3
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 1929, 2020.
10

---

Published as a conference paper at ICLR 2023
Yilum Du and Igor Mordatch. Implicit generation and generalization in energy- baseed models. arXiv preprint arXiv:I903.088 08689, 2019. 3
Yilun Du, Toru Lin, and Igor Mordatch. Model based planning with energy based models. In
Conference Robot Learning, 2019. 2,
Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy based models. arXiv preprint arXiv:2012.01316, 2020.
Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint arXiv:2112.10751, 2021.
Tom Erez and Emanuel Todorov. Trajectory optimization for domains with contacts using inverse dynamics. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. IEEE, 2012.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with im- portance weighted actor-learner architectures. In International conference machine learning, pp. 1407-1416. PMLR, 2018. 2
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference Machine Learning, pp.
2019. 6
Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:2111.10366 2021. 2
Kartik Goyal, Chris Dyer, and Taylor Kerg-Kirkpatrick. Exposing the implicit energy networks be- hind masked language models via metropolis hastings. arXiv preprint arXiv:2106.02736, 2021.
Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference machine learning, pp. 1352-1361. PMLR, 2017. 2,3
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997. 2
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning big sequence modeling problem. Advances in neural information processing systems, 34:1273-122 2021. 1, 2.9
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.0999. 2022.
Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabuey. Recurrent ex- perience replay in distributed reinforcement learning. In International conference learning representations, 2018.
Nan Rosemary Ke, Amanpreet Singh, Ahmed Touati, Anirudh Goyal, Yoshua Bengio, Devi Parikh, and Dhruv Batra. Modeling the long term future in model-based reinforcement learning. In International Conference on Learning Representations, 2018. 2
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-
learning. arXiv preprint arXiv:21 10.06110, 2021. 6,
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems, 2020. 7
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Huang. tutorial on energy-based learning. Predicting structured data, 100), 2006. 3
Shuang Li. Xavier Puig, Yilun Du, Clinton Wang, Ekin Akyurek, Antonio Torralba, Jacob Andreas, and Igor Mordatch. Pre-trained language models for interactive decision-making arXiv preprint arXiv:2202.01771, 2022. 2
11

---

Published as a conference paper at ICLR 2023
Yang Liu and Mirella Lapata. Text summarization with pretrained encoders. arXiv preprint arXiv:1908.08344 08345, 2019. 2
Anusha Nagabandi, Gregory Kahn, Ronald S Fearing, and Sergey Levine, Neural network dynamics for model-basee deep reinforcement learning with model-free fine-tuning In 2018 JEEE Interna- tional Conference on Robotics and Automation (ICRA), pp. 7559-7566. IEEE, 2018.
Junhyuk Oh, Satinder Singh, and Honglak Lee.
Value prediction network. Advances in neural
information processing systems, 30, 2017. 2
Emilio Parisotto, Francis Song, Jack Rac, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury. et al. Stabilizing transformers for reinforcement learning. In International conference on machine learning, pp. 7487-7498. PMLR, 2020. 2
Keiran Paster, Sheila A Mcllraith, and Jimmy Ba. Planning from pixels using inverse dynamics models. arXiv preprini arXiv:2012.02419, 2020.
Keiran Paster, Sheila Mcllraith, and Jimmy Ba. You count on luck: Why decision transformers fail in stochastic environments. arXiv preprint WrXiv:2205./5967 2022. 2, 16
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language under- standing by generative pre-training. 2018. 1
Julian Salazar, Davis Liang, Toan Q Nguyen, and Katrin Kirchhoff. Masked language model scoring. arXiv preprint 1910. 14659, 2019. 2
Julian Salazar, Davis Liang Toan Q. Nguyen, and Katrin Kirchhoff. Masked language model scor- ing. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 2699-2712, 2020. 3
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-600 609, 2020.
Julian Schrittweeser, Thomas Hubert, Amol Mandhane, Mohammadamin Barckatain, Ioannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with learned model. Advances in Neural Information Processing Systems, 34:27580-27591, 2021.
Jiankai Sun, De-An Huang, Bo Lu, Yun-Hui Liu, Bolei Zhou, and Animesh Garg. Plate: Visually- grounded planning with transformers in procedural tasks. IEFE Robotics and Automation Letters, 7(2):4924-4930, 2022, 2, 6
Ilya Sutskever, Oriol Vinyals, and Quoe V Le. Sequence to sequence learning with neural networks. Advances in neural information processing systems, 27, 2014. 2
Richard S Sutton. Dyna, integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 164) 1991. 2
Erik Talvitie. Model regularization for stable sample rollouts. In UAI, pp. 780-789, 2014. 2
Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. Advances in neural information processing systems, 29, 2016. 2
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. XInet: Generalized autoregressive pretraining for language understanding. In Advances Neural Information Processing Systems, volume 32, 2019. 2
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang. Vijay Vasudevan, Alexander Ku, Yinfci Yang, Bureu Karagol Ayan, et al. Scaling autoregressive models for content- rich text-to-image generation. arXiv preprint arXiv:2206. 10789, 2022.1
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model based offline policy optimization. Advances in Neural Information 2020. 2,
Processing Systems,
12

---

Published as a conference paper at ICLR 2023
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. arXiv preprint arXiv:2202.05607,, 2022. 2
13

---

Published as a conference paper at ICLR 2023
Appendix
A EXPERIMENTAL DETAILS
A.1 BABYAI ENVIRONMENT DETAILS
We categorize the environments tested in the trajectory planning and instruction completion into the single-room plane world and multi-room maze world which are connect by doors.
1. Trajectory planning:
Plane world: GoTolooc (7x7), GoToloocaIS8N7 (8x8)
Maze world: GoToObjMazeS4 (10x 10), GoToobjjazeS7 (19x19)
2. Instruction completion:
Plane world: PickUpLoc (8x8)
Maze world: GoTToOhjMazeS4R22Coss (7x7), GoTaSeqS5R2 (9x9)
The Table 5 presents the detailed Baby Al environment settings including the environment size, the number of rooms, the number of obstacles, the status of doors and one example instruction in that environment.
Env
Size
Room
Obs
Door
Inst
GoToLoca
x
5
gu the green key
7
to the blue hox
GoToObjMazeS4
10 x 10
9
1
9 Open
to the blue key
GaToshiiMaze77
7 19
Open
go to grey hall
Go MorroOhiaaeSSRRACCosee
4
1
3
Closed
to the blue ball
PickUpLoc
8
pick up the yellow
the green door and go the green door. then go to red door and go to the green door
GoToSeg55R2
4
Closed
Table 5: Baby/ll environment setting details and example instruction.
A.2 NETWORK DETAILS
We build our LEAP implementation based on Decision Transformer (https://githul) com/ model (https:/
and exploit the instruction encoder from the BabyAl agent / //github. wwt..ooommmmmoiooo walll-iiiaabaae btaabyai/blbb/cl1rtt
/babyaaaaaaa//aaa hatboyyiimmoell... /model. py). In detail, we use the Gated Recurrent Units (GRU) encoder to process the instruction and then apply ExpertControllerFii. abs 709... 078777777 7tto fuse the instruction embedding with state embedding. For all our experi- ments we use bidirectional mask in transformer attention layer, except for Atari where we found casual attention perform better. The full list of hyperparameters be found in Table and Table 7, most of the hyperparameters are taken from Decision Transformer and BabyAl agent model.
(inspired by FiL.MeBllok from https:/ /tapps///xiiiv .org/
A.3 BASELINE MODELS
BabyAl Baseline Models We ran BCQ and IQL based on the following implementation
https://gittuu..com///
For BC and DT, we use the author's 's original implementation
https: https:///iithtte httggggthhhg... http:///a..aaaaao//////aaaaaaaaa
For PlaTe, the author's original implementation
14

---

Published as a conference paper at ICLR 2023
Hyperparameter
Value
Number of layers
3
Number of attention hcads
4
Embedding dimension
128
Batch size
64
Image Encoder
an.Conv2d
Image Encoder channels
128, 128
Image Encoder filter sizes
2 2.3 x
Image Encoder maxpeol strides
2, (Image Encoder may vary littl
depending on environment size) nn.GRU
Instruction Encoder
Instruction Encoder channels
128
State Encoder
nn.Linear
State Encoder channels
256, 256, 128
Max epochs
200
Dropout
0.1
Learning rate
6 10
Adam betas
(0.9, 0.95)
Grad norm clip
1.0
Weight decay
0.1
Learning rate decay
Lincar warmup and cosine decay (see code for details)
Table 6: Hyperparameters of I.EAP for BabyAl experiments.
Hyperparameter
Value
Number of layers
Number attention heads
Embedding dimension
128
Batch size
64 Breakout, Qbert
128 Seaquest
256 Pong
Image Encoder
nn. an.Conv2d
Image Encoder channels
32, 64, 64
Image Encoder filter sizes
x 8, x 4,3 x
Image Encoder strides
4,2,
Max epochs
10
Dropout
Learning rate
6 10
Adam betas
(0.9, 0.95)
Grad norm clip
1.0
Weight decay
Learning rate decay
Linear warmup and cosine decay (see code for details)
Table 7: Hyperparameters of LEAP for Atari experiments.
For MOPO, the author's original implementation of dynamic model training and policy learning. For RL policy, we adopt the IQL discussed above.
https://github.co wwt.://ae///ttttttt//eamm /opo.
The actor network and policy network of BCQ and IQL use the transformer architecture which is the same architecture in model, details above. The original DT and BC already use the transformer architecture so we didn't 't change. For all baselines, we add the same instruction encoder and image encoder described above process instruction and image observations.
15

---

Published as a conference paper at ICLR 2023
Atari Baseline Models in Chen et al. (2021).
The scores for DT, BC, CQL, QR-DQN, and REM in Table can be found
A.4 EXPERIMENT DETAILS
BabyAI For LEAP, the larger size environment requires longer horizon and correspondingly more sampling iterations N. After N iteration, all 7 planned actions will be executed. For DT model, it's beneficial of using a longer context length in more complex environments as shown in its original paper (Chen et al., 2021). We list out these parameters for LEAP and DT models in Table 8. We didn't use context information in LEAP in most Baby AI environments as we expect the iterative planning could generate a correct trajectory based solely on the current state observation. While the Go GoToSeqS5R2 environment requires go to sequence of objects in correct order and LEAP needs to remember what objects have been visited from the context. During training, we randomly select and mask one action in trajectory.
LEAP
Eny
context
plan
sample iteration
context
Lo GoToLocalS8N7
10
10
GoToObjMazeS4
10
30
10
GoToObjMazeS7
50
10
PickUpLoc
10
GoIoSegS3R2
10
Table 8: BabyAl environment experiment details for LEAP and DT.
The input to DT model includes the instruction, state context sequence, action context sequence and return-to-ga sequence in which the target reward is set to I initially. The input to other baseline mod- els are the same except they normal reward sequence instead of return-to-go sequence. While LEAP only use the instruction, state context sequence and action context sequence. Inside state sequence, each state s, contains the [x:, y, dér, 9zz 9yl meaning the agent's position, y position, direction and goal object's x position, position (if the goal location available).
In dynamically changing Atari environment, LEAP use context information in all four games and only execute the first planned action avoid the unexpected changes in the world, details in Table 9. During training, we randomly sample and mask one action in trajectory.
Atari
Env
context
plan
sample iteration
Env
w/o return
w return
Breakout
25
5
10
Breakout
182.0
378.9
Qbert
25
5
10
Qbert
41.0
19.6
Pong
25
5
10
Pong
100.7
108.9
Seaquest
25
10
30
Seaquest 0.5
1.3
Table 9: Atari environment experiment de- tails for LEAP.
Table 10: LEAP performance in Atari envi- ronment with and w/o return input.
Note that our approach can easily be conditioned total reward, by simply concatenating the reward as input the sequence model. One hypothesis that when demonstration set contained trajectories of varying quality, taking reward input following will enable the model to recognize the quality of training trajectories and potentially improve the performance. To further validate the importance test the LEAP with and without return-to-go inputs, which sum of future rewards Chen et al. (2021). The results show that the performance degradation without the return-to-go inputs, which is shown in Table 10.
of the rewards,
B STOCHASTIC ENVIRONMENT TESTING
In this section demonstrate the possibility of extending method into stochastic settings. Although Paster al. (2022) reveals that planning by sampling from the learnt policy conditioned on desired reward could lead to suboptimal outcome due the existence of stochastic factors,
16

---

Published as a conference paper at ICLR 2023
model circumvents the problem by formulating the planning optimization problem we the Gibbs sampling method to find the trajectory with the lowest energy evaluated by the trained model. Assuming that the frequency of successful actions dominates in the dataset, our model is trained to assign lower energy to trajectories with higher likelihood of reaching the goal. Consequently, in the stochastic environments, LEAP constructs sequence of actions that has the best opportunity accomplish the target. When executing this plan in a stochastic environment, we may also choose to replan our sequence of actions after each actual action in the environment (to deal with stochasticity of the next state given action). Note then that this sequence of actions will be optimal in the stochastic environment, as we always choose the action that has the maximum likelihood of reaching the final state. Also note that multi-step planning can potentially provide advantage over simple policy to predict the next action stochastic environments, as such policy simply assigns probability distribution to the immediate next step without the awareness of future step adjustments facing stochastic factors.
To verify the assumptions, constructed a stochastic testing in Al environment. The test created by adopting a stochastic dynamic model, where the agent fails to execute the turning actions nurn left/right with 20% chance, and instead performs the remaining actions, including turn right/left, forward, pickup, drop, and open, with uniform probability. The remaining settings follow Ba experiments detailed in Appendix. A.4, except that we train models using demonstrations generated with the above dynamic model. Those training data are noisy in the sense that the actions taken are not optimal, and corrections are required from future actions. We believe LEAP, as multi-step planner, can learn the above correlations between the consecutive actions. We compare LEAP with the baseline DT, the results of which is collected in the Table. 11. can be observed that LEAP has superior performance compared to DT on both tested environments, which indicates both the possibility applying our approach in the stochastic settings, and the advantage of multi- step planning when facing stochastic factors.
Env
LEAP
DT
GoZeObiMazeS4
57.5% 33.3%
30.8%
GoToObjMazeS7
28.3%
Table 11: Comparison of LEAP and DT on stochastic settings
C ABLATION ON ENERGY MODEL AND OPTIMIZATION METHOD
We further ablate on our design choices, including the energy model and sampling methods. We consider Masked Language Model (MLM) and sequence model classifier as the energy model, and random-shooting, Cross-Entropy Method (CEM), and Gibbs Sampling the optimization approach. All combinations are tested, for which the results are collected in Table 12. We observe that the Gibbs sampling gives the best outcome with MLM model and that defining energy value using sequence model classifier doesn't 't work well in all settings.
1. Sampling methods:
Random-shie :andom-shooting:: randomly generated 30 action sequences and pick up the one with lowest estimated energy value.
Cross- Cooss-Entropy method: randomly generated 30 action sequences, keep the three best se- quences with lowest estimated energy values in each iteration. Then we randomly update one action token in the elite sequences to get 30 new sequences for next iteration. The sequence with lowest energy is selected in the final iteration. Gibbs sampling: discussed in the main text.
2. Energy models:
Sequence model classifier: LSTM sequence model predicts the scalar energy value given the entire trajectory 7, and train the loss between ground truth trajectory energy and estimated energy. The optimal trajectories in Babyai are assigned with lowest energy value and the generates suboptimal trajectories are assigned with higher values depending on the degree of randomness.
MLM: discussed in the main text.
17

---

Published as a conference paper at ICLR 2023
Energy model
Random-shooting
CEM
Gibbs sampling
MLM Classifier
23.3% 25.0%
57.5% 12.5%
62.5 %
15.5%
Table 12: Comparison of different combination between energy models and sampling methods
18


Reviews:
[{"review_id": "tCvuy-0r7r", "reviewer": "Reviewer_MKux", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "The paper looks at the use of masked-language models trained on a given set of trajectories to generate action sequences that minimize a specified energy function. If the energy function captures the objectives of a sequential decision-making problem, then the procedure could be utilized to generate the course of action to be performed. The proposed system is then tested in various domains including, BabyAI and ATARI and shows a number of desirable properties such as online adaptation,  novel environment generalization and task compositionality.", "summary_of_the_review": "While I think the paper presents an interesting direction and formulation, I don’t believe the paper in its current form is ready to be published yet.", "strength_and_weaknesses": "One of the main strengths of the approach is the fact that it presents a straightforward application of existing learning methods. \n\nIn terms of weaknesses, one of the central problems is the clarity of exposition. In fact, the current description of the method presented in the paper is not sufficient to clearly evaluate the properties of the method. First off, the method is referred to as planning. In the common usage of the term, planning usually involves reasoning with some form of task model. It isn’t completely clear if there is a task model being maintained in this case. Is it referred to as planning because you are using an energy function or is the probabilistic model meant to stand for the world model? \n\nAlso assuming access to a set of near optimal training data brings the setting closer to a behavior cloning set up than traditional RL set up, where the agent has to learn through interaction with the environment. The fact that for the experiments, you had to initially run an RL agent (at least in atari) kind of speaks against the usefulness of these methods on their own. Also the energy function itself makes no reference to total reward associated with each trajectory. If the original trajectory set contained trajectories of varying quality how does the model differentiate them during the testing time? Would it just have to assume the better action showed up with higher frequency? \n\nIt is also unclear to me if the method was only designed for deterministic settings? For, stochastic settings solutions can’t take the form of a sequence of actions, rather it needs to take the form of policies. Which means after each action, it needs to check what state it leads to and decide the next action based on this information. While in theory the model can support such execution, it is unclear if one really benefits from the use of such a sequential model in stochastic settings. After all the sequence predicted by the model, even if it is the highest likelihood one, may end up not being followed.  Instead at each step, one needs to consider the expected value associated with following a given action. However as discussed above, the model doesn’t really seem capable of associating any kind of value with a specific action.\n\nFinally on a smaller note, while the learning methods itself uses tools from NLP that were used to learn language models, they are not used to learn language models in this context. So I might advise the authors from referring to the method as planning using language models and maybe just refer to the fact that you are using sequential models.\n\nIn terms of experiments, as discussed the generation of action trajectories are a bit unrealistic, but apart from that the experiments presented seem like a reasonable starting point. One point I would like to know is whether the model presented here can perform better than the traces that were present originally. In other words, if it was presented a state and context that correspond to an example trace part of the initial set, could it have come up with a better course of action by leveraging other traces it saw.\n\nPost Rebuttal Comment: After the new set of experiments and the detailed discussion with the authors. I am happy to recommend the paper for acceptance.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "XApn69idWe", "reviewer": "Reviewer_VX64", "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": "This paper proposed a novel approach (LEAP) for modelling RL trajectories and planning with the model. Unlike GPT-2/Trajectory Transformer style of autoregressive model, it applies BERT style modelling where the actions can be masked in any position of the trajectory. With such a model, planning does not have to follow temporal order. The empirical results show LEAP can do flexible planning and surpass model-free offline RL methods in terms of success rate or return.", "summary_of_the_review": "I think the idea of the paper is interesting and the empirical studies also show great potential for this work. However, the write-up of the paper still needs to be improved to improve clarity. Current I still have a lot of questions about the method, which may affect my final recommendation but in general, I'm leaning to accept the paper.", "strength_and_weaknesses": "Strength:\n* The proposed method is rather novel for RL planning as it goes beyond step-by-step planning which gives extra flexibility and might lead to a new thread of future work.\n* Empirical results show LEAP can fit multiple scenarios of planning, eg. goal-oriented, reward maximization and online adaptation by a constraint function.\n\nWeakness:\n* Empirical evaluations did not include any other planning-based (actually any model-based) methods. It's then hard to compare the proposed method with conventional planning.\n* some of the notations are confusing. \n  * The trajectory length is first defined as $n$ and then it redefined as $T$.\n  * I'm also not sure what does eq (3) mean. Why all the padded actions correspond to state $s_n$?\n  * eq(2) defines energy function as pseudolikelihood but then in the Atari experiment, it seems like a different function is used which is associated with reward. Also, if the optimization objective can be anything, why it's called energy?\n\nQuestions:\n1. Does $\\pmb s$ in eq (1) and eq (2) mean $\\pmb{s}_{1}$? So the \"model\" is not predicting states?\n2. Following question 1, if that's the case, how can you define constraint based on future states in eq (4)? Are you using the simulator for planning?\n3. How did you get the value and reward estimation for Atari experiments?\n4. What did you get the goal state for BABYAI tasks?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "K3NGxGyOUE6", "reviewer": "Reviewer_QzJV", "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": "This paper presents an approach to leveraging a language modeling objective to train an energy function which can be used together with a sampling procedure to generate trajectories for discrete-control tasks. The approach uses a masked language modeling loss to train a bi-directional language model to fit expert demonstrations; the likelihood of this model is used to choose actions (given a context of states and actions) via an iterative Gibbs sampling procedure. The method is tested on a subset of BabyAI and Atari tasks and compared to several baselines; performance is significantly better than baselines on the BabyAI tasks and about the same as baselines on Atari tasks. Extensions of the approach to enable solutions to compositional problems are also propose and show good results on the tested tasks.", "summary_of_the_review": "Overall, this paper presents an interesting new approach to utilizing a language modeling loss to estimate an energy for planning action sequences and is tested on a small set of tasks in two different domains. The results are promising and show improvements compared to the presented baselines. I would suggest a borderline accept.", "strength_and_weaknesses": "Strengths:\n1. This paper is very well written. The central message of the paper is clearly conveyed and the figures nicely illustrate the different potential applications of the proposed approach.\n2. The idea of using a bi-directional language model for energy estimation and using a sampling based approach for computing actions leveraging this energy is interesting and shows good results compared to baselines on a subset of the tested tasks. \n3. The approach is extensible and can easily lend itself for compositional objectives which prior methods seem to struggle on.\n\nWeaknesses:\n1. A key point mentioned in the motivation of the approach is that \"auto-regressive\" modeling approaches can struggle to provide good action selections for long-horizon problems due to their uni-directional nature. This was used as a motivation for bi-directional reasoning but it is not clear if this makes a difference as there are no baselines that ablate this selection (I presume a different planning method would also be needed for this ablation).\n2. The key novelty (in my opinion) of the presented approach is the choice of masked language modeling style losses for modeling the expert data distribution, and subsequent use of this likelihood function for planning. Which of these two matter more? Would any generic planning approach (e.g. random shooting vs Gibbs sampling vs cross-entropy style planning) work with the proposed energy function? And vice-versa would other ways of approximating the energy work with the proposed planning approach? This is unclear from the current results, and an ablation for this would be very illuminating.\n3. The results on Atari are quite weak; in fact, CQL performs better than the proposed approach on 3/4 tasks, and only on a single task the approach gets a high final score which makes the overall score be higher than CQL. Additionally, the results in Fig 7 indicate that while there is some correlation between the achieved rewards and the estimated energy there are a significant number of outliers. Why is this the case?\n4. There are three specific extensions of the proposed approach mentioned: online adaptation, novel environment generalization and task compositionality. It is not clear if these are unique to the presented approach. Specifically, the proposed online adaptation schema can work for any sampling based planning approach if rejection sampling is used, and the task compositionality only works if the energy functions can be linearly combined (and even then there can be significant local optimas from the planning perspective). These should be highlighted in the presentation.\n5. A potential limitation of the proposed approach is its applicability to continuous action spaces; several, if not all of the presented baselines can directly be applied to continuous action spaces whereas the proposed approach is limited to discrete tokens. This can potentially be relaxed by binning/discretizing continuous actions but nonetheless it would be useful to discuss this in the main text.\n6. It would useful to understand the computational/time complexity of the proposed planning approach as it requires several sequential iterations with a potentially large model for successful planning.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JZd61ziULAP", "reviewer": "Reviewer_gPMg", "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": "This paper studies planning with language models (LM) using iterative energy minimization. The authors utilize a bidirectional LM with the masked language model (MLM) objective. The model is trained by optimizing for MLM using expert trajectories. For inference, they iteratively mask actions in a trajectory, including future actions, and use Gibbs sampling to iteratively predict and refine actions. The authors use a decision transformer (DT) for the encoding of a trajectory and explain benefits of this iterative planning approach. Empirical results on BabyAI and Atari games show that LEAP outperforms previous offline RL methods on BabyAI and it is in par with DT on Atari. ", "summary_of_the_review": "I think bidirectional LM based planning with Gibbs sampling is interesting. I found some baselines missing, especially relevant model-based and planning baselines.", "strength_and_weaknesses": "**Strengths** The paper is written well and easy to follow through. I found the iterative planning with bidirectional LMs interesting and broadly applicable to other settings including multi-agent learning (not mentioned in the paper).\n\n**Weaknesses**\n\n1. There are no model-based or autoregressive planning baselines. For example, *PlaTe* and *MuZero Unplugged* are two baselines that need to be cited and compared against.\n\n2. Baselines for the Lava experiment are used as is while a large energy is added for those states in LEAP; which a bit unfair in my opinion. Could you add baselines where you just mask an action for DT/IQL if the action is leading to a lava?\n\n3. At the beginning of section 3, in both argmax and argmin, you don't have $s_{2:T}$. Was this intentional or just a mistake in writing? If it was intentional, could you explain why you only condition on the first state in your formulation? You use only the first action in Atari, so I am assuming this doesn't apply to it.\n\n4. Please add confidence intervals to Table-3. DT results have wide intervals and it is difficult to compare LEAP to DT without error bounds.\n\n5. Can you discuss if LEAP would work for non-deterministic dynamics? \n\n6. In Appendix A.2. you mention causal attention while in main text you mention bidirectional attention with future. Please clarify if you use bidirectional or causal attention.\n\n7. You mention that $\\tau^*$ is derived by minimizing $E_\\theta(\\tau)$ but this is not true as your model doesn't predict states. Please clarify.\n\n8. \"baseslines\" --> \"baselines\",\n\"adaptaiton\" --> \"adaptation\"\n\n\nPlaTe: Visually-Grounded Planning with Transformers in Procedural Tasks. Jiankai Sun, De-An Huang, Bo Lu, Yun-Hui Liu, Bolei Zhou, Animesh Garg.\n\nOnline and Offline Reinforcement Learning by Planning with a Learned Model. Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, Ioannis Antonoglou, David Silver.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster