Planning with Diffusion for Flexible Behavior Synthesis

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.

1 Introduction

Planning with a learned model is a conceptually simple framework for reinforcement learning and data-driven decision-making. Its appeal comes from employing learning techniques only where they are the most mature and effective: for the approximation of unknown environment dynamics in what amounts to a supervised learning problem. Afterwards, the learned model may be plugged into classical trajectory optimization routines (Tassa et al., 2012; Posa et al., 2014; Kelly, 2017), which are similarly well-understood in their original context.

Figure 1: Diffuser is a diffusion probabilistic model that plans by iteratively refining trajectories.

However, this combination rarely works as described. Because powerful trajectory optimizers exploit learned models, plans generated by this procedure often look more like adversarial examples than optimal trajectories (Talvitie, 2014; Ke et al., 2018). As a result, contemporary model-based reinforcement learning algorithms often inherit more from model-free methods, such as value functions and policy gradients (Wang et al., 2019), than from the trajectory optimization toolbox. Those methods that do rely on online planning tend to use simple gradient-free trajectory optimization routines like random shooting (Nagabandi et al., 2018) or the cross-entropy method (Botev et al., 2013; Chua et al., 2018) to avoid the aforementioned issues.

In this work, we propose an alternative approach to data-driven trajectory optimization. The core idea is to train a model that is directly amenable to trajectory optimization, in the sense that sampling from the model and planning with it become nearly identical. This goal requires a shift in how the model is designed. Because learned dynamics models are normally meant to be proxies for environment dynamics, improvements are often achieved by structuring the model according to the underlying causal process (Bapst et al., 2019). Instead, we consider how to design a model in line with the planning problem in which it will be used. For example, because the model will ultimately be used for planning, action distributions are just as important as state dynamics and long-horizon accuracy is more important than single-step error. On the other hand, the model should remain agnostic to reward function so that it may be used in multiple tasks, including those unseen during training. Finally, the model should be designed so that its plans, and not just its predictions, improve with experience and are resistant to the myopic failure modes of standard shooting-based planning algorithms.

Figure 2: Diffuser samples plans by iteratively denoising two-dimensional arrays consisting of a variable number of state-action pairs. A small receptive field constrains the model to only enforce local consistency during a single denoising step. By composing many denoising steps together, local consistency can drive global coherence of a sampled plan. An optional guide function $𝒥$ can be used to bias plans toward those optimizing a test-time objective or satisfying a set of constraints.

We instantiate this idea as a trajectory-level diffusion probabilistic model (Sohl-Dickstein et al., 2015; Ho et al., 2020) called Diffuser, visualized in Figure 2. Whereas standard model-based planning techniques predict forward in time autoregressively, Diffuser predicts all timesteps of a plan simultaneously. The iterative sampling process of diffusion models leads to flexible conditioning, allowing for auxiliary guides to modify the sampling procedure to recover trajectories with high return or satisfying a set of constraints. This formulation of data-driven trajectory optimization has several appealing properties:

Long-horizon scalability   Diffuser is trained for the accuracy of its generated trajectories rather than its single-step error, so it does not suffer from the compounding rollout errors of single-step dynamics models and scales more gracefully with respect to long planning horizon.

Task compositionality   Reward functions provide auxiliary gradients to be used while sampling a plan, allowing for a straightforward way of planning by composing multiple rewards simultaneously by adding together their gradients.

Temporal compositionality   Diffuser generates globally coherent trajectories by iteratively improving local consistency, allowing it to generalize to novel trajectories by stitching together in-distribution subsequences.

Effective non-greedy planning   By blurring the line between model and planner, the training procedure that improves the model’s predictions also has the effect of improving its planning capabilities. This design yields a learned planner that can solve the types of long-horizon, sparse-reward problems that prove difficult for many conventional planning methods.

The core contribution of this work is a denoising diffusion model designed for trajectory data and an associated probabilistic framework for behavior synthesis. While unconventional compared to the types of models routinely used in deep model-based reinforcement learning, we demonstrate that Diffuser has a number of useful properties and is particularly effective in offline control settings that require long-horizon reasoning and test-time flexibility.

2 Background

Our approach to planning is a learning-based analogue of past work in behavioral synthesis using trajectory optimization (Witkin & Kass, 1988; Tassa et al., 2012). In this section, we provide a brief background on the problem setting considered by trajectory optimization and the class of generative models we employ for that problem.

2.1 Problem Setting

Consider a system governed by the discrete-time dynamics ${𝐬}_{t+1}=𝒇({𝐬}_t,{𝐚}_t)$ at state ${𝐬}_t$ given an action ${𝐚}_t$. Trajectory optimization refers to finding a sequence of actions ${𝐚}_{0:T}^{\ast }$ that maximizes (or minimizes) an objective $𝒥$ factorized over per-timestep rewards (or costs) $r({𝐬}_t,{𝐚}_t)$:

$${𝐚}_{0:T}^{\ast }=\underset{{𝐚}_{0:T}}{\arg \max }𝒥({𝐬}_0,{𝐚}_{0:T})=\underset{{𝐚}_{0:T}}{\arg \max }\sum _{t=0}^{T}r({𝐬}_t,{𝐚}_t)$$

where $T$ is the planning horizon. We use the abbreviation $𝝉=({𝐬}_0,{𝐚}_0,{𝐬}_1,{𝐚}_1,\mathrm{\dots },{𝐬}_T,{𝐚}_T)$ to refer to a trajectory of interleaved states and actions and $𝒥(𝝉)$ to denote the objective value of that trajectory.

2.2 Diffusion Probabilistic Models

Diffusion probabilistic models (Sohl-Dickstein et al., 2015; Ho et al., 2020) pose the data-generating process as an iterative denoising procedure $p_{\theta }({𝝉}^{i-1}\mid {𝝉}^i)$. This denoising is the reverse of a forward diffusion process $q({𝝉}^i\mid {𝝉}^{i-1})$ that slowly corrupts the structure in data by adding noise. The data distribution induced by the model is given by:

$$p_{\theta }({𝝉}^0)=\int p({𝝉}^N)\prod _{i=1}^N{p}_{\theta }({𝝉}^{i-1}\mid {𝝉}^i)\mathrm{d}{𝝉}^{1:N}$$

where $p({𝝉}^N)$ is a standard Gaussian prior and ${𝝉}^0$ denotes (noiseless) data. Parameters $\theta$ are optimized by minimizing a variational bound on the negative log likelihood of the reverse process: ${\theta }^{\ast }={\arg \min }_{\theta }-{𝔼}_{{𝝉}^0}\left[\log p_{\theta }({𝝉}^0)\right].$ The reverse process is often parameterized as Gaussian with fixed timestep-dependent covariances:

$$p_{\theta }({𝝉}^{i-1}\mid {𝝉}^i)=𝒩({𝝉}^{i-1}\mid {\mu }_{\theta }({𝝉}^i,i),{\mathrm{\Sigma }}^i).$$

The forward process $q({𝝉}^i\mid {𝝉}^{i-1})$ is typically prespecified.

Notation.   There are two “times” at play in this work: that of the diffusion process and that of the planning problem. We use superscripts ($i$ when unspecified) to denote diffusion timestep and subscripts ($t$ when unspecified) to denote planning timestep. For example, ${𝐬}_t^0$ refers to the $t^{\text{th}}$ state in a noiseless trajectory. When it is unambiguous from context, superscripts of noiseless quantities are omitted: $𝝉={𝝉}^0$. We overload notation slightly by referring to the $t^{\text{th}}$ state (or action) in a trajectory $𝝉$ as ${𝝉}_{{𝐬}_t}$ (or ${𝝉}_{{𝐚}_t}$).

3 Planning with Diffusion

A major obstacle to using trajectory optimization techniques is that they require knowledge of the environment dynamics $𝒇$. Most learning-based methods attempt to overcome this obstacle by training an approximate dynamics model and plugging it in to a conventional planning routine. However, learned models are often poorly suited to the types of planning algorithms designed with ground-truth models in mind, leading to planners that exploit learned models by finding adversarial examples.

We propose a tighter coupling between modeling and planning. Instead of using a learned model in the context of a classical planner, we subsume as much of the planning process as possible into the generative modeling framework, such that planning becomes nearly identical to sampling. We do this using a diffusion model of trajectories, $p_{\theta }(𝝉)$. The iterative denoising process of a diffusion model lends itself to flexible conditioning by way of sampling from perturbed distributions of the form:

$${\tilde{p}}_{\theta }(𝝉)\propto p_{\theta }(𝝉)h(𝝉).$$ (1)

The function $h(𝝉)$ can contain information about prior evidence (such as an observation history), desired outcomes (such as a goal to reach), or general functions to optimize (such as rewards or costs). Performing inference in this perturbed distribution can be seen as a probabilistic analogue to the trajectory optimization problem posed in Section 2.1, as it requires finding trajectories that are both physically realistic under $p_{\theta }(𝝉)$ and high-reward (or constraint-satisfying) under $h(𝝉)$. Because the dynamics information is separated from the perturbation distribution $h(𝝉)$, a single diffusion model $p_{\theta }(𝝉)$ may be reused for multiple tasks in the same environment.

In this section, we describe Diffuser, a diffusion model designed for learned trajectory optimization. We then discuss two specific instantiations of planning with Diffuser, realized as reinforcement learning counterparts to classifier-guided sampling and image inpainting.

3.1 A Generative Model for Trajectory Planning

Temporal ordering.   Blurring the line between sampling from a trajectory model and planning with it yields an unusual constraint: we can no longer predict states autoregressively in temporal order. Consider the goal-conditioned inference $p({𝐬}_1\mid {𝐬}_0,{𝐬}_T)$; the next state ${𝐬}_1$ depends on a future state as well as a prior one. This example is an instance of a more general principle: while dynamics prediction is causal, in the sense that the present is determined by the past, decision-making and control can be anti-causal, in the sense that decisions in the present are conditional on the future. Because we cannot use a temporal autoregressive ordering, we design Diffuser to predict all timesteps of a plan concurrently.

Temporal locality.   Despite not being autoregressive or Markovian, Diffuser features a relaxed form of temporal locality. In Figure 2, we depict a dependency graph for a diffusion model consisting of a single temporal convolution. The receptive field of a given prediction only consists of nearby timesteps, both in the past and the future. As a result, each step of the denoising process can only make predictions based on local consistency of the trajectory. By composing many of these denoising steps together, however, local consistency can drive global coherence.

Trajectory representation.   Diffuser is a model of trajectories designed for planning, meaning that the effectiveness of the controller derived from the model is just as important as the quality of the state predictions. As a result, states and actions in a trajectory are predicted jointly; for the purposes of prediction the actions are simply additional dimensions of the state. Specifically, we represent inputs (and outputs) of Diffuser as a two-dimensional array:

(2)

with one column per timestep of the planning horizon.

Architecture.   We now have the ingredients needed to specify a Diffuser architecture: (1) an entire trajectory should be predicted non-autoregressively, (2) each step of the denoising process should be temporally local, and (3) the trajectory representation should allow for equivariance along one dimension (the planning horizon) but not the other (the state and action features). We satisfy these criteria with a model consisting of repeated (temporal) convolutional residual blocks. The overall architecture resembles the types of U-Nets that have found success in image-based diffusion models, but with two-dimensional spatial convolutions replaced by one-dimensional temporal convolutions (Figure A1). Because the model is fully convolutional, the horizon of the predictions is determined not by the model architecture, but by the input dimensionality; it can change dynamically during planning if desired.

Training.   We use Diffuser to parameterize a learned gradient ${ϵ}_{\theta }({𝝉}^i,i)$ of the trajectory denoising process, from which the mean ${\mu }_{\theta }$ can be solved in closed form (Ho et al., 2020). We use the simplified objective for training the $ϵ$-model, given by:

$$ℒ(\theta )={𝔼}_{i,ϵ,{𝝉}^0}\left[{\parallel ϵ-{ϵ}_{\theta }({𝝉}^i,i)\parallel }^2\right],$$

in which $i\sim 𝒰\{1,2,\mathrm{\dots },N\}$ is the diffusion timestep, $ϵ\sim 𝒩(\mathrm{𝟎},𝑰)$ is the noise target, and ${𝝉}^i$ is the trajectory ${𝝉}^0$ corrupted with noise $ϵ$. Reverse process covariances ${\mathrm{\Sigma }}^i$ follow the cosine schedule of Nichol & Dhariwal (2021).

Algorithm 1 Guided Diffusion Planning

1:  Require Diffuser ${\mu }_{\theta }$, guide $𝒥$, scale $\alpha$, covariances ${\mathrm{\Sigma }}^i$

2:  while not done do

3:     Observe state $𝐬$; initialize plan ${𝝉}^N\sim 𝒩(\mathrm{𝟎},𝑰)$

4:     for $i=N,\mathrm{\dots },1$ do

5:        // parameters of reverse transition

6:        $\mu \leftarrow {\mu }_{\theta }\left({𝝉}^i\right)$

7:        // guide using gradients of return

8:        ${𝝉}^{i-1}\sim 𝒩(\mu +\alpha \mathrm{\Sigma }\nabla 𝒥\left(\mu \right),{\mathrm{\Sigma }}^i)$

9:        // constrain first state of plan

10:        ${𝝉}_{{𝒔}_0}^{i-1}\leftarrow 𝒔$   

11:     end for

12:     Execute first action of plan ${𝝉}_{{𝐚}_0}^0$

13:  end while

a          

denoising  

b     

data                  plan

c    $\to$    $\to$

d     

reward $-$ $+$        plan

Figure 3: (Properties of diffusion planners) (a) Learned long-horizon planning: Diffuser’s learned planning procedure does not suffer from the myopic failure modes common to shooting algorithms and is able to plan over long horizons with sparse reward. (b) Temporal compositionality: Even though the model is not Markovian, it generates trajectories via iterated refinements to local consistency. As a result, it exhibits the types of generalization usually associated with Markovian models, with the ability to stitch together snippets of trajectories from the training data to generate novel plan. (c) Variable-length plans: Despite being a trajectory-level model, Diffuser’s planning horizon is not determined by its architecture. The horizon can be updated after training by changing the dimensionality of the input noise. (d) Task compositionality: Diffuser can be composed with new reward functions to plan for tasks unseen during training. In all subfigures, denotes a starting state and denotes a goal state.

3.2 Reinforcement Learning as Guided Sampling

In order to solve reinforcement learning problems with Diffuser, we must introduce a notion of reward. We appeal to the control-as-inference graphical model (Levine, 2018) to do so. Let ${𝒪}_t$ be a binary random variable denoting the optimality of timestep $t$ of a trajectory, with $p({𝒪}_t=1)=\exp (r({𝐬}_t,{𝐚}_t))$. We can sample from the set of optimal trajectories by setting $h(𝝉)=p({𝒪}_{1:T}\mid 𝝉)$ in Equation 1:

$${\tilde{p}}_{\theta }(𝝉)=p(𝝉\mid {𝒪}_{1:T}=1)\propto p(𝝉)p({𝒪}_{1:T}=1\mid 𝝉).$$

We have exchanged the reinforcement learning problem for one of conditional sampling. Thankfully, there has been much prior work on conditional sampling with diffusion models. While it is intractable to sample from this distribution exactly, when $p({𝒪}_{1:T}\mid {𝝉}^i)$ is sufficiently smooth, the reverse diffusion process transitions can be approximated as Gaussian (Sohl-Dickstein et al., 2015):

$$p_{\theta }({𝝉}^{i-1}\mid {𝝉}^i,{𝒪}_{1:T})\approx 𝒩({𝝉}^{i-1};\mu +\mathrm{\Sigma }g,\mathrm{\Sigma })$$ (3)

where $\mu ,\mathrm{\Sigma }$ are the parameters of the original reverse process transition $p_{\theta }({𝝉}^{i-1}\mid {𝝉}^i)$ and

$g$ $={{\nabla }_{𝝉}\log p({𝒪}_{1:T}\mid 𝝉)|}_{𝝉=\mu }$

$={{\displaystyle \sum _{t=0}^T}{\nabla }_{{𝐬}_t,{𝐚}_t}r({𝐬}_t,{𝐚}_t)|}_{({𝐬}_t,{𝐚}_t)={\mu }_t}=\nabla 𝒥(\mu ).$

This relation provides a straightforward translation between classifier-guided sampling, used to generate class-conditional images (Dhariwal & Nichol, 2021), and the reinforcement learning problem setting. We first train a diffusion model $p_{\theta }(𝝉)$ on the states and actions of all available trajectory data. We then train a separate model ${𝒥}_{\varphi }$ to predict the cumulative rewards of trajectory samples ${𝝉}^i$. The gradients of ${𝒥}_{\varphi }$ are used to guide the trajectory sampling procedure by modifying the means $\mu$ of the reverse process according to Equation 3. The first action of a sampled trajectory $𝝉\sim p(𝝉\mid {𝒪}_{1:T}=1)$ may be executed in the environment, after which the planning procedure begins again in a standard receding-horizon control loop. Pseudocode for the guided planning method is given in Algorithm 1.

3.3 Goal-Conditioned RL as Inpainting

Some planning problems are more naturally posed as constraint satisfaction than reward maximization. In these settings, the objective is to produce any feasible trajectory that satisfies a set of constraints, such as terminating at a goal location. Appealing to the two-dimensional array representation of trajectories described by Equation 2, this setting can be translated into an inpainting problem, in which state and action constraints act analogously to observed pixels in an image (Sohl-Dickstein et al., 2015). All unobserved locations in the array must be filled in by the diffusion model in a manner consistent with the observed constraints.

The perturbation function required for this task is a Dirac delta for observed values and constant elsewhere. Concretely, if ${𝐜}_t$ is state constraint at timestep $t$, then

The definition for action constraints is identical. In practice, this may be implemented by sampling from the unperturbed reverse process ${𝝉}^{i-1}\sim p_{\theta }({𝝉}^{i-1}\mid {𝝉}^i)$ and replacing the sampled values with conditioning values ${𝐜}_t$ after all diffusion timesteps $i\in \{0,1,\mathrm{\dots },N\}$.

Even reward maximization problems require conditioning-by-inpainting because all sampled trajectories should begin at the current state. This conditioning is described by line 10 in Algorithm 1.

4 Properties of Diffusion Planners

We discuss a number of Diffuser’s important properties, focusing on those that are are either distinct from standard dynamics models or unusual for non-autoregressive trajectory prediction.

Learned long-horizon planning.   Single-step models are typically used as proxies for ground-truth environment dynamics $𝒇$, and as such are not tied to any planning algorithm in particular. In contrast, the planning routine in Algorithm 1 is closely tied to the specific affordances of diffusion models. Because our planning method is nearly identical to sampling (with the only difference being guidance by a perturbation function $h(𝝉)$), Diffuser’s effectiveness as a long-horizon predictor directly translates to effective long-horizon planning. We demonstrate the benefits of learned planning in a goal-reaching setting in Figure 3a, showing that Diffuser is able to generate feasible trajectories in the types of sparse reward settings where shooting-based approaches are known to struggle. We explore a more quantitative version of this problem setting in Section 5.1.

Temporal compositionality.   Single-step models are often motivated using the Markov property, allowing them to compose in-distribution transitions to generalize to out-of-distribution trajectories. Because Diffuser generates globally coherent trajectories by iteratively improving local consistency (Section 3.1), it can also stitch together familiar subsequences in novel ways. In Figure 3b, we train Diffuser on trajectories that only travel in a straight line, and show that it can generalize to v-shaped trajectories by composing trajectories at their point of intersection.

Variable-length plans.   Because our model is fully convolutional in the horizon dimension of its prediction, its planning horizon is not specified by architectural choices. Instead, it is determined by the size of the input noise ${𝝉}^N\sim 𝒩(\mathrm{𝟎},𝐈)$ that initializes the denoising process, allowing for variable-length plans (Figure 3c).

Task compositionality.   While Diffuser contains information about both environment dynamics and behaviors, it is independent of reward function. Because the model acts as a prior over possible futures, planning can be guided by comparatively lightweight perturbation functions $h(𝝉)$ (or even combinations of multiple perturbations) corresponding to different rewards. We demonstrate this by planning for a new reward function unseen during training of the diffusion model (Figure 3d).

Environment	MPPI	CQL	IQL	Diffuser
Maze2D    U-Maze	$33.2$	$5.7$	$47.4$	113.9 $\pm 3.1$
Maze2D    Medium	$10.2$	$5.0$	$34.9$	121.5 $\pm 2.7$
Maze2D    Large	$5.1$	$12.5$	$58.6$	123.0 $\pm 6.4$
Single-task Average	16.2	7.7	47.0	119.5
Multi2D    U-Maze	$41.2$	-	$24.8$	128.9 $\pm 1.8$
Multi2D    Medium	$15.4$	-	$12.1$	127.2 $\pm 3.4$
Multi2D    Large	$8.0$	-	$13.9$	132.1 $\pm 5.8$
Multi-task Average	21.5	-	16.9	129.4

Table 1: (Long-horizon planning) The performance of Diffuser and prior model-free algorithms in the Maze2D environment, which tests long-horizon planning due to its sparse reward structure. The Multi2D setting refers to a multi-task variant with goal locations resampled at the beginning of every episode. Diffuser substantially outperforms prior approaches in both settings. Appendix A details the sources for the scores of the baseline algorithms.

5 Experimental Evaluation

The focus of our experiments is to evaluate Diffuser on the capabilities we would like from a data-driven planner. In particular, we evaluate (1) the ability to plan over long horizons without manual reward shaping, (2) the ability to generalize to new configurations of goals unseen during training, and (3) the ability to recover an effective controller from heterogeneous data of varying quality. We conclude by studying practical runtime considerations of diffusion-based planning, including the most effective ways of speeding up the planning procedure while suffering minimally in terms of performance.

5.1 Long Horizon Multi-Task Planning

We evaluate long-horizon planning in the Maze2D environments (Fu et al., 2020), which require traversing to a goal location where a reward of 1 is given. No reward shaping is provided at any other location. Because it can take hundreds of steps to reach the goal location, even the best model-free algorithms struggle to adequately perform credit assignment and reliably reach the goal (Table 1).

We plan with Diffuser using the inpainting strategy to condition on a start and goal location. (The goal location is also available to the model-free methods; it is identifiable by being the only state in the dataset with non-zero reward.) We then use the sampled trajectory as an open-loop plan. Diffuser achieves scores over 100 in all maze sizes, indicating that it outperforms a reference expert policy. We visualize the reverse diffusion process generating Diffuser’s plans in Figure 4.

While the training data in Maze2D is undirected – consisting of a controller navigating to and from randomly selected locations – the evaluation is single-task in that the goal is always the same. In order to test multi-task flexibility, we modify the environment to randomize the goal location at the beginning of each episode. This setting is denoted as Multi2D in Table 1. Diffuser is naturally a multi-task planner; we do not need to retrain the model from the single-task experiments and simply change the conditioning goal. As a result, Diffuser performs as well in the multi-task setting as in the single-task setting. In contrast, there is a substantial performance drop of the best model-free algorithm in the single-task setting (IQL; Kostrikov et al. 2022) when adapted to the multi-task setting. Details of our multi-task IQL with hindsight experience relabeling (Andrychowicz et al., 2017) are provided in Appendix A. MPPI uses the ground-truth dynamics; its poor performance compared to the learned planning algorithm of Diffuser highlights the difficulty posed by long-horizon planning even when there are no prediction inaccuracies.

U-Maze

Medium

Large

denoising  

Figure 4: (Planning as inpainting) Plans are generated in the Maze2D environment by sampling trajectories consistent with a specified start and goal condition. The remaining states are “inpainted” by the denoising process.

Dataset	Environment	BC	CQL	IQL	DT	TT	MOPO	MOReL	MBOP	Diffuser
Medium-Expert	HalfCheetah	$55.2$	$91.6$	$86.7$	$86.8$	$95.0$	$63.3$	$53.3$	105.9	$88.9$ $\pm 0.3$
Medium-Expert	Hopper	$52.5$	105.4	$91.5$	107.6	110.0	$23.7$	108.7	$55.1$	$103.3$ $\pm 1.3$
Medium-Expert	Walker2d	107.5	108.8	109.6	108.1	$101.9$	$44.6$	$95.6$	$70.2$	106.9 $\pm 0.2$
Medium	HalfCheetah	$42.6$	$44.0$	47.4	$42.6$	46.9	$42.3$	$42.1$	$44.6$	$42.8$ $\pm 0.3$
Medium	Hopper	$52.9$	$58.5$	$66.3$	$67.6$	$61.1$	$28.0$	95.4	$48.8$	$74.3$ $\pm 1.4$
Medium	Walker2d	$75.3$	$72.5$	78.3	$74.0$	79.0	$17.8$	77.8	$41.0$	79.6 $\pm 0.55$
Medium-Replay	HalfCheetah	$36.6$	$45.5$	$44.2$	$36.6$	$41.9$	53.1	$40.2$	$42.3$	$37.7$ $\pm 0.5$
Medium-Replay	Hopper	$18.1$	95.0	94.7	$82.7$	91.5	$67.5$	93.6	$12.4$	93.6 $\pm 0.4$
Medium-Replay	Walker2d	$26.0$	$77.2$	$73.9$	$66.6$	82.6	$39.0$	$49.8$	$9.7$	$70.6$ $\pm 1.6$
Average		51.9	77.6	77.0	74.7	78.9	42.1	72.9	47.8	77.5

Table 2: (Offline reinforcement learning) The performance of Diffuser and a variety of prior algorithms on the D4RL locomotion benchmark (Fu et al., 2020). Results for Diffuser correspond to the mean and standard error over 150 planning seeds. We detail the sources for the performance of prior methods in Appendix A.3. Following Kostrikov et al. (2022), we emphasize in bold scores within 5 percent of the maximum per task ($\ge 0.95\cdot \text{max}$).

5.2 Test-time Flexibility

In order to evaluate the ability to generalize to new test-time goals, we construct a suite of block stacking tasks with three settings: (1) Unconditional Stacking, for which the task is to build a block tower as tall as possible; (2) Conditional Stacking, for which the task is to construct a block tower with a specified order of blocks, and (3) Rearrangement, for which the task is to match a set of reference blocks’ locations in a novel arrangement. We train all methods on 10000 trajectories from demonstrations generated by PDDLStream (Garrett et al., 2020); rewards are equal to one upon successful stack placements and zero otherwise. These block stacking are challenging diagnostics of test-time flexibility; in the course of executing a partial stack for a randomized goal, a controller will venture into novel states not included in the training configuration.

We use one trained Diffuser for all block-stacking tasks, only modifying the perturbation function $h(𝝉)$ between settings. In the Unconditional Stacking task, we directly sample from the unperturbed denoising process $p_{\theta }(𝝉)$ to emulate the PDDLStream controller. In the Conditional Stacking and Rearrangement tasks, we compose two perturbation functions $h(𝝉)$ to bias the sampled trajectories: the first maximizes the likelihood of the trajectory’s final state matching the goal configuration, and the second enforces a contact constraint between the end effector and a cube during stacking motions. (See Appendix B for details.)

Environment	BCQ	CQL	Diffuser
Unconditional Stacking	$0.0$	$24.4$	58.7 $\pm 2.5$
Conditional Stacking	0.0	0.0	45.6 $\pm 3.1$
Rearrangement	0.0	0.0	58.9 $\pm 3.4$
Average	0.0	8.1	54.4

Table 3: (Test-time flexibility) Performance of BCQ, CQL, and Diffuser on block stacking tasks. A score of 100 corresponds to a perfectly executed stack; 0 is that of a random policy.

We compare with two prior model-free offline reinforcement learning algorithms: BCQ (Fujimoto et al., 2019) and CQL (Kumar et al., 2020), training standard variants for Unconditional Stacking and goal-conditioned variants for Conditional Stacking and Rearrangement. (Baseline details are provided in Appendix A.) Quantitative results are given in Table 3, in which a score of 100 corresponds to a perfect execution of the task. Diffuser substantially outperforms both prior methods, with the conditional settings requiring flexible behavior generation proving especially difficult for the model-free algorithms. A visual depiction of an execution by Diffuser is provided in Figure 5.

5.3 Offline Reinforcement Learning

$\to$   

Figure 5: (Block stacking) A block stacking sequence executed by Diffuser. This task is best illustrated by videos viewable at diffusion-planning.github.io.

Finally, we evaluate the capacity to recover an effective single-task controller from heterogeneous data of varying quality using the D4RL offline locomotion suite (Fu et al., 2020). We guide the trajectories generated by Diffuser toward high-reward regions using the sampling procedure described in Section 3.2 and condition the trajectories on the current state using the inpainting procedure described in Section 3.3. The reward predictor ${𝒥}_{\varphi }$ is trained on the same trajectories as the diffusion model.

We compare to a variety of prior algorithms spanning other approaches to data-driven control, including the model-free reinforcement learning algorithms CQL (Kumar et al., 2020) and IQL (Kostrikov et al., 2022); return-conditioning approaches like Decision Transformer (DT; Chen et al. 2021b); and model-based reinforcement learning approaches including Trajectory Transformer (TT; Janner et al. 2021), MOPO (Yu et al., 2020), MOReL (Kidambi et al., 2020), and MBOP (Argenson & Dulac-Arnold, 2021). In the single-task setting, Diffuser performance comparably to prior algorithms: better than the model-based MOReL and MBOP and return-conditioning DT, but worse than the best offline techniques designed specifically for single-task performance. We also investigated a variant using Diffuser as a dynamics model in conventional trajectory optimizers such as MPPI (Williams et al., 2015), but found that this combination performed no better than random, suggesting that the effectiveness of Diffuser stems from coupled modeling and planning, and not from improved open-loop predictive accuracy.

denoising

planning horizon  

Figure 6: (Guided sampling) Diffuser generates all timesteps of a plan concurrently, instead of autoregressively, through the denoising process.

5.4 Warm-Starting Diffusion for Faster Planning

A limitation of Diffuser is that individual plans are slow to generate (due to iterative generation). Naïvely, as we execute plans open loop, a new plan must be regenerated at each step of execution. To improve execution speed of Diffuser, we may further reuse previously generated plans to warm-start generations of subsequent plans.

To warm-start planning, we may run a limited number of forward diffusion steps from a previously generated plan and then run a corresponding number of denoising steps from this partially noised trajectory to regenerate an updated plan. In Figure 7, we illustrate the trade-off between performance and runtime budget as we vary the underlying number of denoising steps used to regenerate each a new plan from 2 to 100. We find that we may reduce the planning budget of our approach markedly with only modest drop in performance.

6 Related Work

Figure 7: (Fast Planning) Performance of Diffuser on Walker2d Medium-Expert when varying the number of diffusion steps to warm-start planning. Performance suffers only minimally even when using one-tenth the number of diffusion steps, as long as plans are initialized from the previous timestep’s plan.

Advances in deep generative modeling have recently made inroads into model-based reinforcement learning, with multiple lines of work exploring dynamics models parameterized as convolutional U-networks (Kaiser et al., 2020), stochastic recurrent networks (Ke et al., 2018; Hafner et al., 2021a; Ha & Schmidhuber, 2018), vector-quantized autoencoders (Hafner et al., 2021b; Ozair et al., 2021), neural ODEs (Du et al., 2020a), normalizing flows (Rhinehart et al., 2020; Janner et al., 2020), generative adversarial networks (Eysenbach et al., 2021), energy-based models (EBMs; Du et al. 2019), graph neural networks (Sanchez-Gonzalez et al., 2018), neural radiance fields (Li et al., 2021), and Transformers (Janner et al., 2021; Chen et al., 2021a). Further, Lambert et al. 2020 have studied non-autoregressive trajectory-level dynamics models for long-horizon prediction. These investigations generally assume an abstraction barrier between the model and planner. Specifically, the role of learning is relegated to approximating environment dynamics; once learning is complete the model may be inserted into any of a variety of planning (Botev et al., 2013; Williams et al., 2015) or policy optimization (Sutton, 1990; Wang et al., 2019) algorithms because the form of the planner does not depend strongly on the form of the model. Our goal is to break this abstraction barrier by designing a model and planning algorithm that are trained alongside one another, resulting in a non-autoregressive trajectory-level model for which sampling and planning are nearly identical.

A number of parallel lines of work have studied how to break the abstraction barrier between model learning and planning in different ways. Approaches include training an autoregressive latent-space model for reward prediction (Tamar et al., 2016; Oh et al., 2017; Schrittwieser et al., 2019); weighing model training objectives by state values (Farahmand et al., 2017); and applying collocation techniques to learned single-step energies (Du et al., 2019; Rybkin et al., 2021). In contrast, our method plans by modeling and generating all timesteps of a trajectory concurrently, instead of autoregressively, and conditioning the sampled trajectories with auxiliary guidance functions.

Diffusion models have emerged as a promising class of generative model that formulates the data-generating process as an iterative denoising procedure (Sohl-Dickstein et al., 2015; Ho et al., 2020). The denoising procedure can be seen as parameterizing the gradients of the data distribution (Song & Ermon, 2019), connecting diffusion models to score matching (Hyvärinen, 2005) and EBMs (LeCun et al., 2006; Du & Mordatch, 2019; Nijkamp et al., 2019; Grathwohl et al., 2020). Iterative, gradient-based sampling lends itself towards flexible conditioning (Dhariwal & Nichol, 2021) and compositionality (Du et al., 2020b), which we use to recover effective behaviors from heterogeneous datasets and plan for reward functions unseen during training. While diffusion models have been developed for the generation of images (Song et al., 2021), waveforms (Chen et al., 2021c), 3D shapes (Zhou et al., 2021), and text (Austin et al., 2021), to the best of our knowledge they have not previously been used in the context of reinforcement learning or decision-making.

7 Conclusion

We have presented Diffuser, a denoising diffusion model for trajectory data. Planning with Diffuser is almost identical to sampling from it, differing only in the addition of auxiliary perturbation functions that serve to guide samples. The learned diffusion-based planning procedure has a number of useful properties, including graceful handling of sparse rewards, the ability to plan for new rewards without retraining, and a temporal compositionality that allows it to produce out-of-distribution trajectories by stitching together in-distribution subsequences. Our results point to a new class of diffusion-based planning procedures for deep model-based reinforcement learning.

Code References

We used the following open-source libraries for this work: NumPy (Harris et al., 2020), PyTorch (Paszke et al., 2019), and Diffusion Models in PyTorch (Wang, 2020).

Acknowledgements

We thank Ajay Jain for feedback on an early draft and Leslie Kaelbling, Tomás Lozano-Pérez, Jascha Sohl-Dickstein, Ben Eysenbach, Amy Zhang, Colin Li, and Toru Lin for helpful discussions. This work was partially supported by computational resource donations from Microsoft. M.J. is supported by fellowships from the National Science Foundation and the Open Philanthropy Project. Y.D. is supported by a fellowship from the National Science Foundation.

Appendix A Baseline details and sources

In this section, we provide details about baselines we ran ourselves. For scores of baselines previously evaluated on standardized tasks, we provide the source of the listed score.

A.1 Maze2D experiments

Single-task.

The performance of CQL and IQL on the standard Maze2D environments is reported in the D4RL whitepaper (Fu et al., 2020) in Table 2.

We ran IQL using the offical implementation from the authors:

github.com/ikostrikov/implicit_q_learning.

We tuned over two hyperparameters:

1. temperature $\in [3,10]$

2. expectile $\in [0.65,0.95]$

Multi-task.

We only evaluated IQL on the Multi2D environments because it is the strongest baseline in the single-task Maze2D environments by a sizeable margin. To adapt IQL to the multi-task setting, we modified the $Q$-functions, value function, and policy to be goal-conditioned. To select goals during training, we employed a strategy based on hindsight experience replay, in which we sampled a goal from among those states encountered in the future of a trajectory. For a training backup $({𝐬}_t,{𝐚}_t,{𝐬}_{t+1})$, we sampled goals according to a geometric distribution over the future

$$\mathrm{\Delta }\sim \text{Geom}(1-\gamma )𝐠={𝐬}_{t+\mathrm{\Delta }},$$

recalculated rewards based on the sampled goal, and conditioned all relevant models on the goal during updating. During testing, we conditioned the policy on the ground-truth goal.

We tuned over the same IQL parameters as in the single-task setting.

A.2 Block stacking experiments

Single-task.

We ran CQL using the following implementation

https://github.com/young-geng/cql.

and used default hyperparameters in the code. We ran BCQ using the author’s original implementation

https://github.com/sfujim/BCQ.

For BCQ, we tuned over two hyperparameters:

1. discount factor $\in [0.9,0.999]$

2. tau $\in [0.001,0.01]$

Multi-task.

To evaluate BCQ and CQL in the multi-task setting, we modified the $Q$-functions, value function and policy to be goal-conditioned. We trained using goal relabeling as in the Multi2D environments. We tuned over the same hyperparameters described in the single-task block stacking experiments.

A.3 Offline Locomotion

The scores for BC, CQL, IQL, and AWAC are from Table 1 in Kostrikov et al. (2022). The scores for DT are from Table 2 in Chen et al. (2021b). The scores for TT are from Table 1 in Janner et al. (2021). The scores for MOReL are from Table 2 in Kidambi et al. (2020). The scores for MBOP are from Table 1 in Argenson & Dulac-Arnold (2021).

Figure A1: Diffuser has a U-Net architecture with residual blocks consisting of temporal convolutions, group normalization, and Mish nonlinearities.

Appendix B Test-time Flexibility

To guide Diffuser to stack blocks in specified configurations, we used two separate perturbation functions $h(𝝉)$ to specify a given stack of block A on top of block B, which we detail below.

Final State Matching  To enforce a final state consisting of block A on top of block B, we trained a perturbation function $h_{\text{match}}(𝝉)$ as a per-timestep classifier determining whether a a state $𝒔$ exhibits a stack of block A on top of block B. We train the classifier on the demonstration data as the diffusion model.

Contact Constraint  To guide the Kuka arm to stack block A on top of block B, we construct a perturbation function $h_{\text{contact}}(𝝉)={\sum }_{i=0}^{64}-1\ast {\Vert {𝝉}_{c_i}-1\Vert }^2$, where ${𝝉}_{c_i}$ corresponds to the underlying dimension in state ${𝝉}_{s_i}$ that specifies the presence or absence of contact between the Kuka arm and block A. We apply the contact constraint between the Kuka arm and block A for the first 64 timesteps in a trajectory, corresponding to initial contact with block A in a plan.

Appendix C Implementation Details

In this section we describe the architecture and record hyperparameters.

1. The architecture of Diffuser (Figure A1) consists of a U-Net structure with 6 repeated residual blocks. Each block consisted of two temporal convolutions, each followed by group norm (Wu & He, 2018), and a final Mish nonlinearity (Misra, 2019). Timestep embeddings are produced by a single fully-connected layer and added to the activations of the first temporal convolution within each block.

2. We train the model using the Adam optimizer (Kingma & Ba, 2015) with a learning rate of $4\mathrm{e}-05$ and batch size of $32$. We train the models for 500k steps.

3. The return predictor $𝒥$ has the structure of the first half of the U-Net used for the diffusion model, with a final linear layer to produce a scalar output.

4. We use a planning horizon $T$ of 32 in all locomotion tasks, $128$ for block-stacking, $128$ in Maze2D / Multi2D U-Maze, 265 in Maze2D / Multi2D Medium, and 384 in Maze2D / Multi2D Large.

5. We found that we could reduce the planning horizon for many tasks, but that the guide scale would need to be lowered (e.g., to 0.001 for a horizon of $4$ in the halfcheetah tasks) to accommodate. The configuration file in the open-source code demonstrates how to run with a modified scale and horizon.

6. We use $N=20$ diffusion steps for locomotion tasks and $N=100$ for block-stacking.

7. We use a guide scale of $\alpha =0.1$ for all tasks except hopper-medium-expert, in which we use a smaller scale of $0.0001$.

8. We used a discount factor of $0.997$ for the return prediction ${𝒥}_{\varphi }$, though found that above $\gamma =0.99$ planning was fairly insensitive to changes in discount factor.

9. We found that control performance was not substantially affected by the choice of predicting noise $ϵ$ versus uncorrupted data ${𝝉}^0$ with the diffusion model.