[Paper Review] Planning with Sequence Models through Iterative Energy Minimization

Hongyi Chen, Yilun Du+3 more|arXiv (Cornell University)|Mar 28, 2023

Reinforcement Learning in Robotics4 citations

TL;DR

This paper proposes LEAP, a planner that uses iterative energy minimization over a masked language model to refine trajectories for reinforcement learning. By treating trajectory planning as energy minimization with bidirectional context, LEAP improves performance on BabyAI and Atari tasks, enabling better generalization, test-time adaptation, and plan composition than autoregressive baselines.

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 of actions to reach some goal, is less straightforward. The typical autoregressive generation procedures of sequence 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 sequence 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 website: https://hychen-naza.github.io/projects/LEAP

Motivation & Objective

• To address the limitation of autoregressive sequence models in planning, which cannot refine earlier actions once generated.
• To enable trajectory planning with sequence models by formulating it as iterative energy minimization over a bidirectional masked language model.
• To improve performance on long-horizon decision-making tasks such as those in BabyAI and Atari environments.
• To enable generalization to unseen environments and test-time constraints, including dynamic obstacles and complex mazes.
• To support compositional planning by combining multiple models for multi-objective tasks.

Proposed method

• Train a bidirectional masked language model (MLM) to estimate an implicit energy function over action trajectories using negative pseudo-likelihood.
• Define the energy of a trajectory as the negative log-likelihood of the MLM, enabling energy-based optimization.
• Perform iterative refinement by replacing actions at each timestep with the most likely action given the current context and future actions.
• Use a bidirectional context to allow future information to influence earlier action choices during refinement.
• Apply iterative minimization to converge on low-energy (high-probability) trajectories that satisfy the goal.
• Integrate the energy minimization process into a planner that can adapt to new constraints and compose plans from multiple models.

Experimental results

Research questions

• RQ1 Can iterative energy minimization over a sequence model improve trajectory planning in reinforcement learning compared to standard autoregressive generation?
• RQ2 Does the bidirectional context in the sequence model enable better refinement of earlier actions in a trajectory?
• RQ3 Can the energy minimization framework generalize to unseen environments with new obstacles or complex layouts?
• RQ4 Can the planner adapt to test-time constraints such as dynamic or newly introduced obstacles?
• RQ5 Can multiple learned models be composed together using this framework to achieve multi-objective planning?

Key findings

• LEAP outperforms autoregressive baselines such as DT and IQL in online adaptation to test-time obstacles, achieving 92.0% success on medium difficulty lava environments compared to 68.0% for DT.
• In unseen maze environments, LEAP generalizes effectively, achieving 77.5% success on medium and 61.0% on hard cases, while DT drops to 36.0% and 24.5% respectively.
• On composite task composition, LEAP achieves 83.5% success in easy cases and 43.0% in medium cases, significantly outperforming IQL (42.5% and 11.5%) and DT (58.0% and 15.5%).
• LEAP demonstrates robustness in complex mazes and dynamic obstacle scenarios, successfully bypassing lava grids and unperceivable obstacles through iterative refinement.
• The iterative energy minimization process enables better long-horizon planning, as evidenced by consistent performance gains across all evaluated benchmarks.
• The framework supports compositional planning, allowing multiple models to contribute to a single plan even when each only observes partial information.