PaperlayerView exact revision on arXiv

Source-linked AI summary

Planning with Diffusion for Flexible Behavior Synthesis

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

arXiv:2205.09991v2cs.LGcs.AI

TL;DR

The paper addresses the mismatch between learned dynamics models and classical trajectory optimizers by integrating planning into a trajectory-level diffusion model. Diffuser plans through iterative denoising with flexible auxiliary guidance, and the paper reports effectiveness for long-horizon, sparse-reward, and test-time-flexible control. Its main scope limitation is slow individual plan generation, which motivates warm-starting subsequent plans.

  • Problem

    Classical planners can exploit learned dynamics models to produce adversarial examples, leaving learned models poorly suited to standard trajectory optimization.

  • Method

    Diffuser is a trajectory-level diffusion model that plans by iteratively denoising complete trajectories while using auxiliary guides for rewards, goals, or constraints.

  • Results

    Diffuser demonstrates useful planning properties, including long-horizon scalability, reward composition, temporal compositionality, and effectiveness in offline control requiring long-horizon reasoning and test-time flexibility.

  • Takeaways & Limitations

    Planning can be made nearly identical to sampling, allowing new rewards and constraints to guide behavior without retraining and enabling coherent trajectories from in-distribution subsequences.

  • Takeaways & Limitations

    Individual Diffuser plans are slow to generate because generation is iterative, so open-loop execution may require regenerating a plan at each step.

Abstract

from arXiv · show

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.

1 Introduction

The paper proposes Diffuser, a trajectory-level diffusion model that makes sampling and planning nearly identical, addressing shortcomings of classical optimization over learned dynamics. It emphasizes long-horizon scalability, flexible test-time conditioning, reward composition, and temporal compositionality.

  • Motivation: Classical trajectory optimization over learned dynamics can produce adversarial-looking plans, motivating a model designed directly for planning.The paper argues that standard optimizers exploit learned models in ways that differ from their intended use.
  • Approach: Diffuser is a trajectory-level diffusion model whose iterative sampling process predicts all plan timesteps simultaneously.Unlike autoregressive model-based planning, diffusion sampling denoises complete trajectory representations.
  • Approach: Auxiliary guides can steer sampled trajectories toward high return or constraint satisfaction, while remaining separate from the learned dynamics information.This supports test-time flexibility without requiring the model to encode a particular reward function.
  • Properties: Diffuser is trained for generated-trajectory accuracy rather than single-step error, enabling more graceful scaling to long planning horizons.The design targets long-horizon reasoning instead of relying on single-step dynamics accuracy.
  • Properties: The framework supports composing multiple rewards during sampling by adding their gradients, extending planning to task combinations.The paper presents reward composition as a direct use of auxiliary sampling guidance.
  • Properties: Temporal compositionality lets Diffuser generate globally coherent novel trajectories by stitching together in-distribution subsequences.Iterative improvement of local consistency supports generalization to trajectories not seen during training.

2 Background

Trajectory optimization seeks action sequences that optimize a reward or cost under system dynamics, while diffusion models generate trajectories through iterative denoising from noise.

  • 2.1 Problem Setting: Trajectory optimization finds an action sequence that maximizes or minimizes an objective accumulated over per-timestep rewards or costs.The trajectory interleaves states and actions across a planning horizon T.
  • 2.1 Problem Setting: A trajectory τ denotes interleaved states and actions, and J(τ) denotes its objective value.The planning horizon is T, with τ = (s0, a0, s1, a1, ..., sT, aT).
  • 2.2 Diffusion Probabilistic Models: Diffusion probabilistic models define data generation as iterative denoising that reverses a forward process adding noise to trajectories.The reverse transition is written pθ(τ_i−1 | τ_i), while the forward process is q(τ_i | τ_i−1).
  • 2.2 Diffusion Probabilistic Models: The diffusion model begins from a standard Gaussian prior and produces noiseless data through learned reverse-process transitions.The reverse process is commonly parameterized as Gaussian with fixed timestep-dependent covariances.
  • 2.2 Diffusion Probabilistic Models: Diffusion training minimizes a variational bound on the negative log likelihood of the reverse process.The supplied passage identifies this objective as the optimization target for θ.
  • Notation: The notation distinguishes diffusion timesteps, written as superscripts, from planning timesteps, written as subscripts.For example, s^0_t denotes the tth state in a noiseless trajectory.

3 Planning with Diffusion

Diffuser folds trajectory planning into a diffusion model whose samples are iteratively denoised, jointly represent states and actions, and can be flexibly guided by rewards, goals, and constraints.

  • 3 Planning with Diffusion: Diffuser replaces a classical planner paired with learned dynamics by a trajectory generative model in which planning is nearly identical to sampling.The model represents trajectories with pθ(τ), while perturbation functions encode evidence, desired outcomes, rewards, or costs.
  • 3.1 A Generative Model for Trajectory Planning: Diffuser predicts all timesteps concurrently because control decisions can condition on future states rather than following a causal autoregressive ordering.The paper characterizes this as anti-causal decision-making and contrasts it with causal dynamics prediction.
  • 3.1 A Generative Model for Trajectory Planning: Temporally local denoising steps compose into globally coherent trajectories by repeatedly improving local consistency.The model uses repeated temporal convolutional residual blocks with local receptive fields.
  • 3.1 A Generative Model for Trajectory Planning: States and actions are predicted jointly in a two-dimensional trajectory array with one column per planning timestep.Actions are treated as additional dimensions of the state for prediction.
  • 3.1 A Generative Model for Trajectory Planning: Training uses a learned denoising gradient ϵθ(τ_i, i), with noisy trajectories formed by corrupting τ^0 using Gaussian noise at a sampled diffusion timestep.Reverse-process covariances follow a cosine schedule.
  • 3.2 Reinforcement Learning as Guided Sampling: Reward-guided sampling converts reinforcement learning into conditional sampling by modifying reverse-process means with gradients of predicted cumulative reward.A diffusion model is trained on trajectory data, then a separate model Jφ supplies guidance gradients.
  • 3.3 Goal-Conditioned RL as Inpainting: Goal and constraint satisfaction can be formulated as inpainting, where observed state or action values are fixed and unobserved trajectory entries are filled by diffusion.The perturbation function uses a Dirac delta for observed values and is constant elsewhere.
  • 3.3 Goal-Conditioned RL as Inpainting: The framework exhibits learned long-horizon planning, temporal compositionality, variable-length plans, and task compositionality.Figure 3 associates these properties with sparse-reward planning, trajectory stitching, input-noise-controlled horizons, and new reward functions.

4 Properties of Diffusion Planners

Diffuser ties planning to diffusion sampling, enabling long-horizon prediction, compositionality, variable-length plans, and reward-independent task adaptation.

  • Learned long-horizon planning: Diffuser’s planning routine is nearly identical to sampling, so effective long-horizon prediction translates directly into long-horizon planning.Auxiliary perturbation functions guide samples toward desired trajectories.
  • Temporal compositionality: Diffuser composes familiar trajectory subsequences into novel globally coherent behaviors, including v-shaped paths learned from straight-line trajectories.Local consistency is iteratively composed into global coherence.
  • Variable-length plans: Diffuser supports variable-length plans because its fully convolutional horizon dimension is determined by the input-noise size rather than fixed architectural choices.The planning horizon is set by τ_N ∼ N(0, I).
  • Task compositionality: Diffuser can adapt to new reward functions without retraining because lightweight perturbation functions guide a reward-independent prior over possible futures.Multiple perturbations can be composed for different rewards.
  • Experimental evidence: In Maze2D, Diffuser substantially outperforms prior approaches in both single-task and multi-task long-horizon settings.Multi2D resamples goal locations at the beginning of every episode.

5 Experimental Evaluation

Experiments test Diffuser on sparse-reward long-horizon planning, changing test-time goals, heterogeneous offline data, and planning-speed trade-offs; it performs strongly but generates plans slowly.

  • 5.1 Long Horizon Multi-Task Planning: Diffuser achieves scores over 100 in all Maze2D sizes, outperforming a reference expert policy under sparse rewards.It conditions trajectories on start and goal locations through inpainting and executes them open loop.
  • 5.1 Long Horizon Multi-Task Planning: Diffuser retains single-task performance when goal locations are randomized at test time, without retraining, whereas multi-task adaptation substantially reduces IQL performance.MPPI with ground-truth dynamics performs poorly, underscoring the difficulty of long-horizon planning.
  • 5.2 Test-time Flexibility: Diffuser substantially outperforms BCQ and CQL on block-stacking tasks, especially in conditional settings requiring flexible behavior generation.The evaluated tasks include unconditional stacking, conditional block ordering, and rearrangement.
  • 5.3 Offline Reinforcement Learning: Diffuser is evaluated on D4RL locomotion by guiding trajectories toward high-reward regions and conditioning them on the current state using inpainting.The reward predictor is trained on the same trajectories as the diffusion model.
  • 5.3 Offline Reinforcement Learning: Diffuser’s effectiveness depends on coupled modeling and planning: using it as a dynamics model inside conventional trajectory optimizers performs no better than random.This result contrasts with treating improved open-loop prediction as sufficient for planning.
  • 5.4 Warm-Starting Diffusion for Faster Planning: Individual plans are slow to generate because diffusion uses iterative generation, requiring regeneration at each open-loop execution step.Warm-starting from previous plans markedly reduces the planning budget with only a modest performance drop.

6 Related Work

Related work spans generative dynamics models and methods that integrate model learning with planning; Diffuser instead generates whole trajectories concurrently and guides them with auxiliary functions.

  • Generative models for model-based RL: Recent model-based reinforcement learning work has explored dynamics models based on convolutional U-networks, recurrent networks, autoencoders, flows, GANs, EBMs, graph networks, neural radiance fields, and Transformers.These approaches represent a broad range of generative-model parameterizations.
  • Practical planning speed: Warm-starting previously generated plans is studied as a practical way to reduce diffusion-planning computation while limiting performance loss.Performance suffers only minimally with one-tenth the diffusion steps when plans are initialized from the previous timestep’s plan.
  • Integrating learning and planning: Other approaches reduce the separation between learning and planning through latent-space reward models, value-weighted objectives, or collocation techniques.These methods address the model-planning interface through different training or optimization strategies.
  • Diffuser’s distinction: In contrast, Diffuser models and generates all trajectory timesteps concurrently rather than autoregressively, while conditioning samples with auxiliary guidance functions.This design makes sampling and planning closely related operations.
  • Diffusion foundations: Diffusion models formulate generation as iterative denoising and connect this process to score matching and energy-based models.The denoising procedure parameterizes gradients of the data distribution.

7 Conclusion

The paper concludes that Diffuser makes planning nearly identical to sampling and supports sparse rewards, new rewards, and temporally compositional out-of-distribution behavior.

  • Conclusion: Diffuser is a denoising diffusion model for trajectory data whose planning differs from sampling only through auxiliary perturbation functions.These perturbations guide generated samples.
  • Conclusion: The framework handles sparse rewards, plans for new rewards without retraining, and uses temporal compositionality to produce out-of-distribution behaviors.These are presented as useful properties of diffusion-based planning.
  • Implementation: The work relies on open-source NumPy, PyTorch, and Diffusion Models in PyTorch libraries.These implementation dependencies are listed in the paper’s supplementary material.

A.1 Maze2D experiments

The Maze2D experiments evaluate IQL and CQL implementations, including a goal-conditioned IQL adaptation for multi-task environments.

  • A.1 Maze2D experiments: The reported single-task CQL and IQL Maze2D scores came from the D4RL whitepaper.
  • A.1 Maze2D experiments: IQL was evaluated on Multi2D after being selected as the strongest single-task Maze2D baseline by a sizeable margin.
  • A.1 Maze2D experiments: The multi-task IQL adaptation made its Q-functions, value function, and policy goal-conditioned.
  • A.1 Maze2D experiments: During testing, IQL was conditioned on the ground-truth goal.

A.2 Block stacking experiments

The block-stacking experiments evaluate CQL and BCQ alongside a U-Net-based Diffuser architecture, adapting both baselines to goal-conditioned multi-task control.

  • A.2 Block stacking experiments: The block-stacking comparisons included CQL and BCQ, with CQL using one implementation and BCQ using the authors’ original implementation.
  • A.2 Block stacking experiments: BCQ was tuned over two hyperparameters.
  • A.2 Block stacking experiments: The Diffuser architecture uses a U-Net with residual blocks containing temporal convolutions, group normalization, and Mish nonlinearities.
  • A.2 Block stacking experiments: For multi-task evaluation, BCQ and CQL used goal-conditioned Q-functions, value functions, and policies.

A.3 Offline Locomotion

The offline locomotion section compares Diffuser with model-free, return-conditioned, and model-based control algorithms, while using perturbation functions to guide block-stacking behavior.

  • A.3 Offline Locomotion: The comparison spans CQL, IQL, Decision Transformer, Trajectory Transformer, MOPO, MOReL, and MBOP.
  • A.3 Offline Locomotion: Diffuser was guided toward specified block configurations using separate perturbation functions.
  • A.3 Offline Locomotion: A final-state classifier determined at each timestep whether a state exhibited block A stacked on block B.
  • A.3 Offline Locomotion: The contact constraint penalized the distance between the planned and desired contact indicators for the first 64 trajectory timesteps.

Appendix C Implementation Details

The implementation details specify Diffuser’s architecture, optimization settings, planning horizons, diffusion steps, guidance scales, and return-prediction discounting.

  • Appendix C Implementation Details: Diffuser uses six repeated residual blocks, each with two temporal convolutions, group normalization, Mish nonlinearities, and timestep embeddings.
  • Appendix C Implementation Details: The model is trained with Adam at learning rate 4e−05, batch size 32, for 500k steps.
  • Appendix C Implementation Details: The return predictor J uses the first half of the diffusion model’s U-Net followed by a linear scalar-output layer.
  • Appendix C Implementation Details: Planning horizons range from 32 in locomotion to 384 in Large Maze2D and Multi2D tasks.
  • Appendix C Implementation Details: The authors found that shorter planning horizons require lowering guide scale, such as 0.001 for horizon 4 in halfcheetah tasks.
  • Appendix C Implementation Details: Diffusion uses 20 steps for locomotion and 100 for block-stacking, while the return predictor uses discount factor 0.997.
  • Appendix C Implementation Details: Control performance was not substantially affected by predicting noise ϵ rather than uncorrupted data τ 0.
Loading 2205.09991v2…