Matt Strong

Matt Strong

CS PhD at Stanford University

← Paper Digest

Planning with Diffusion for Flexible Behavior Synthesis

Michael Janner, Yilun Du, Joshua B. Tenenbaum, Sergey Levine

Seminal May 2022

Diffuser controlling a simulated Kuka arm in the block-stacking task.

TL;DR Diffuser trains a diffusion model over whole state-action trajectories and plans by denoising them, steered by a learned return gradient or by inpainting the start and goal. On sparse-reward Maze2D it averages a score of 119.5, against 47.0 for the best model-free baseline (IQL).

What they do

Model-based RL usually learns a dynamics model and hands it to a classical trajectory optimizer, which tends to exploit the model’s errors. Diffuser instead folds planning into the generative model: a diffusion model over whole trajectories, where planning is sampling with an added guide. One trained model serves several tasks by swapping only the guide.

How it works

  • Trajectory model. A trajectory is a 2D array of states and actions, one column per timestep. A U-Net of 1D temporal convolutions predicts all timesteps at once; each denoising step only sees nearby timesteps, so global coherence comes from many local refinements. The horizon is set by the size of the input noise.
  • Rewards as guidance. A separate network predicts a trajectory’s return; its gradient shifts the mean of each reverse step, as in classifier-guided image sampling. The first action is executed and the plan is regenerated.
  • Constraints as inpainting. Start and goal states are written into the array after every denoising step, so the model fills in the rest.
  • Composing guides. For block stacking, a final-state classifier and a contact constraint are added to the same model.

Results

  • Maze2D (sparse reward): average score 119.5 vs 47.0 for IQL and 16.2 for MPPI with ground-truth dynamics. With the goal randomized each episode (Multi2D), Diffuser averages 129.4 without retraining; goal-conditioned IQL drops to 16.9.
  • Kuka block stacking: 58.7 (unconditional), 45.6 (conditional) and 58.9 (rearrangement), vs at most 24.4 for CQL and 0 for BCQ.
  • D4RL locomotion: average 77.5, close to CQL (77.6) and IQL (77.0), above DT (74.7) and MOReL (72.9).
  • Warm-starting each plan from the previous one with 10x fewer diffusion steps costs little performance on Walker2d Medium-Expert.

Limitations

Plans are slow to generate because sampling is iterative. In single-task offline RL it trails the best specialized methods, and used as a plain dynamics model inside MPPI it does no better than random. All experiments are in simulation.

Videos

Noise denoised into start-to-goal plans in three Maze2D mazes.
Diffuser turning a noise array into a locomotion plan.

Figures

From Figure 1: Diffuser is a diffusion model that plans by iteratively refining trajectories; the forward process add...
From Figure 1: Diffuser is a diffusion model that plans by iteratively refining trajectories; the forward process adds noise and the learned reverse process denoises.
From Figure 2: Diffuser denoises a 2D array of state-action pairs along the planning horizon. Each step has a small l...
From Figure 2: Diffuser denoises a 2D array of state-action pairs along the planning horizon. Each step has a small local receptive field, so many steps build global coherence; an optional guide J biases plans toward a test-time objective or constraints.
From Figure 4: Planning as inpainting in Maze2D. The start and goal are fixed and denoising fills in the states betwe...
From Figure 4: Planning as inpainting in Maze2D. The start and goal are fixed and denoising fills in the states between them, from scattered noise (left) to a feasible path (right).
From Figure 7: Score versus planning time on Walker2d Medium-Expert as the number of warm-start diffusion steps varie...
From Figure 7: Score versus planning time on Walker2d Medium-Expert as the number of warm-start diffusion steps varies. Using one-tenth of the steps costs little performance when each plan starts from the previous one.
Abstract

Model-based reinforcement learning methods often use learning only for the purpose of estimating an approximate dynamics model, offloading the rest of the decision-making work to classical trajectory optimizers. While conceptually simple, this combination has a number of empirical shortcomings, suggesting that learned models may not be well-suited to standard trajectory optimization. In this paper, we consider what it would look like to fold as much of the trajectory optimization pipeline as possible into the modeling problem, such that sampling from the model and planning with it become nearly identical. The core of our technical approach lies in a diffusion probabilistic model that plans by iteratively denoising trajectories. We show how classifier-guided sampling and image inpainting can be reinterpreted as coherent planning strategies, explore the unusual and useful properties of diffusion-based planning methods, and demonstrate the effectiveness of our framework in control settings that emphasize long-horizon decision-making and test-time flexibility.

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.