I S C ONDITIONAL G ENERATIVE M ODELING ALL YOU NEED FOR D ECISION -M AKING ?
A BSTRACT
Recent improvements in conditional generative modeling have made it possible to generate high-quality images from language descriptions alone.
We investigate whether these methods can directly address the problem of sequential decisionmaking.
We view decision-making not through the lens of reinforcement learning (RL), but rather through conditional generative modeling.
To our surprise, we find that our formulation leads to policies that can outperform existing offline
RL approaches across standard benchmarks.
By modeling a policy as a returnconditional generative model, we avoid the need for dynamic programming and subsequently eliminate many of the complexities that come with traditional offline
RL.
We further demonstrate the advantages of modeling policies as conditional generative models by considering two other conditioning variables: constraints and skills.
Conditioning on a single constraint or skill during training leads to behaviors at test-time that can satisfy several constraints together or demonstrate a composition of skills.
Our results illustrate that conditional generative modeling is a powerful tool for decision-making.
I NTRODUCTION
Over the last few years, there have been incredible advances in the field of conditional generative modeling (Ramesh et al., 2022; Jumper et al., 2021; Brown et al., 2020; Saharia et al., 2022;
Lewkowycz et al., 2022).
Image models like DALL-E (Ramesh et al., 2022) and ImageGen (Saharia et al., 2022) can generate accurate high-resolution images from text descriptions alone.
Language models like GPT (Brown et al., 2020) can generate entire paragraphs and stories from short text prompts.
More recent models like Minerva (Lewkowycz et al., 2022) can generate complete, step-bystep solutions to high-school math questions.
The success of generative models in countless domains motivates us to apply them to decision-making.
Directly applying generative modeling to decision-making is not so straightforward because data is readily available from the internet in the former, but must be actively gathered from the environment in the latter.
Conveniently, there exists a wide body of work that has looked into recovering highperforming policies from data logged by already operational systems (Kostrikov et al., 2022; Kumar et al., 2020;
Walke et al., 2022).
This is particularly useful in real-world settings where interacting with the environment is not always possible, and exploratory decisions can have fatal consequences (DulacArnold et al., 2021).
With access to such offline datasets, the problem of decision-making becomes similar to those where generative models have already found success.
Prior work in offline decision-making has primarily relied on tools from reinforcement learning (RL) (Kumar et al., 2020;
Kostrikov et al., 2022; Wu et al., 2019; Kostrikov et al., 2021; Dadashi et al., 2021; Ajay et al., 2020; Ghosh et al., 2022).
Although these ideas have found some success (Yala et al., 2022; Jang et al., 2022; Mirhoseini et al., 2020), they are not without their flaws.
Since they operate in the same spirit as temporal difference (TD) learning, they are prone to many instabilities that come with using function approximation, off-policy learning, and bootstrapping together, otherwise known as the deadly triad (Sutton & Barto, 2018; Van Hasselt et al., 2018).
Since they must also operate with a limited amount of data, they rely on tricks and heuristics to keep the policy within the distribution of the dataset.
These challenges make it difficult to scale existing offline RL algorithms.
τ τ τ τ τ τ τ τ τ τ τ τ τ τ τ τ τ τ τ τ skills constraints reward low x 2 + y 2 ≥ r 2 x 2 + y 2 ≤ Rrunjump high composeskills satisfyconstraints maximizereward r 2 ≤ x 2 + y 2 ≤ Rrun and jump optimal
[Equation]
[Figure]
In this paper, we formulate the problem of offline decision-making from the perspective of conditional generative modeling.
Given a fixed dataset of reward-labeled trajectories, we adapt diffusion models (Sohl-Dickstein et al., 2015) to learn a return-conditional model of the trajectory.
During inference, we use classifier-free guidance with low-temperature sampling to capture the best behaviors in the dataset and glean return maximizing trajectories.
Our straightforward conditional generative modeling formulation outperforms existing approaches on standard D4RL tasks (Fu et al., 2020).
Viewing offline decision-making through the lens of conditional generative modeling allows going beyond conditioning on returns (Figure 1).
Consider an example (detailed in Appendix A) where a robot with linear dynamics navigates an environment containing two concentric circles (Figure 2).
We are given a dataset of state-action trajectories of the robot, each satisfying one of two constraints: (i) the final position of the robot is within the larger circle, and (ii) the final position of the robot is outside the smaller circle.
With conditional diffusion modeling, we can use the datasets to learn a constraintconditioned model that can generate trajectories satisfying any set of constraints.
During inference, the learned trajectory model can merge constraints from the dataset and generate trajectories that satisfy the combined constraint.
Figure 2 shows that the constraint-conditioned model can generate trajectories such that the final position of the robot lies between the concentric circles.
[Figure]
Here, we demonstrate the benefits of modeling policies as conditional generative models.
First, conditioning on constraints allows policies to not only generate behaviors satisfying individual constraints but also generate novel behaviors by flexibly combining constraints at test time.
Further, conditioning on skills allows policies to not only imitate individual skills but also generate novel behaviors by composing those skills.
We instantiate this idea with a state-sequence based diffusion probabilistic model (Ho et al., 2020) called Decision Diffuser , visualized in Figure 1.
In summary, our contributions include (i) illustrating conditional generative modeling as an effective tool in offline decision making, (ii) using classifier-free guidance with low-temperature sampling, instead of dynamic programming, to get return-maximizing trajectories and, (iii) leveraging the framework of conditional generative modeling to combine constraints and compose skills during inference flexibly.
B ACKGROUND 2. R EINFORCEMENT L EARNING
We formulate the sequential decision-making problem as a discounted Markov Decision Process (MDP) defined by the tuple ⟨ ρ 0 , S , A , T , R , γ ⟩ , where ρ 0 is the initial state distribution, S and A are state and action spaces, T : S ×
A → S is the transition function, R : S ×
A × S → R gives the reward at any transition and γ ∈ [0 , 1) is a discount factor (Puterman, 1995).
The agent acts with a stochastic policy π : S → ∆ A , generating a sequence of state-action-reward transitions or trajectory τ :=
( s k , a k , r k ) k ≥ 0 with probability p π ( τ ) and return R ( τ ) :
= (cid:80) k ≥
γ k r k .
The standard objective in RL is to find a return-maximizing policy π ∗ = arg max π E τ ∼ p
π
[ R ( τ )] .
Temporal Difference Learning TD methods (Fujimoto et al., 2018; Lillicrap et al., 2015) estimate
Q ∗ ( s, a ) := E τ ∼ p
π ∗ [ R ( τ ) | s
=
s, a 0 = a ] , the return achieved under the optimal policy π ∗ when starting in state s and taking action a , with a parameterized Q -function.
This requires minimizing the following TD loss:
[Equation]
Continuous action spaces further require learning a parametric policy π ϕ ( a | s ) that plays the role of the maximizing action in equation 1.
This results in a policy objective that must be maximized:
[Equation]
Here, the dataset of transitions D evolves as the agent interacts with the environment and both Q θ and π ϕ are trained together.
These methods make use of function approximation, off-policy learning, and bootstrapping, leading to several instabilities in practice (Sutton, 1988; Van Hasselt et al., 2018).
Offline RL
In this setting, we must find a return-maximizing policy from a fixed dataset of transitions collected by an unknown behavior policy µ (Levine et al., 2020).
Using TD-learning naively causes the state visitation distribution d π ϕ ( s ) to move away from the distribution of the dataset d µ ( s ) .
In turn, the policy π ϕ begins to take actions that are substantially different from those already seen in the data.
Offline RL algorithms resolve this distribution-shift by imposing a constraint of the form D ( d π ϕ || d µ ) , where D is some divergence metric, directly in the TD-learning procedure.
The constrained optimization problem now demands additional hyper-parameter tuning and implementation heuristics to achieve any reasonable performance (Kumar et al., 2021).
The
Decision Diffuser, in comparison, does not have any of these disadvantages.
It does not require estimating any kind of Q -function, thereby sidestepping TD methods altogether.
It also does not face the risk of distribution-shift because the generative model is trained with maximum-likelihood estimation.
2. D IFFUSION P ROBABILISTIC M ODELS
Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) are a specific type of generative model that learn the data distribution q ( x ) from a dataset D :=
{ x i } 0 ≤ i<M .
They have been used most notably for synthesizing high-quality images from text descriptions (Saharia et al., 2022; Nichol et al., 2021).
Here, the data-generating procedure is modelled with a predefined forward noising process q ( x t +1 | x t ) :
=
N ( x t +1 ; √ α t x t , (1 − α t ) I ) and a trainable reverse process p θ ( x t − 1 | x t ) :=
N ( x t − 1 | µ θ ( x t , t ) , Σ t ) , where N ( µ, Σ) denotes a Gaussian distribution with mean µ and variance Σ , α t ∈ R determines the variance schedule, x 0 :
= x is a sample, x 1 , x 2 , ..., x N − 1 are the latents, and x N ∼ N ( 0 , I ) for carefully chosen α t and long enough T .
Starting with Gaussian noise, samples are then iteratively generated through a series of ”denoising” steps.
Although a tractable variational lower-bound on log p θ can be optimized to train diffusion models,
Ho et al.
(2020) propose a simplified surrogate loss:
[Equation]
The predicted noise ϵ θ ( x t , t ) , parameterized with a deep neural network, estimates the noise ϵ ∼
N (0 , I ) added to the dataset sample x 0 to produce noisy x t .
This is equivalent to predicting the mean of p θ ( x t − 1 | x t ) since µ θ ( x t , t ) can be calculated as a function of ϵ θ ( x t , t )
(Ho et al., 2020).
[Figure]
returns
[Figure]
Guided Diffusion
Modelling the conditional data distribution q ( x | y ) makes it possible to generate samples with attributes of the label y .
The equivalence between diffusion models and scorematching (Song et al., 2021), which shows ϵ θ ( x t , t ) ∝ ∇ x t log p ( x t ) , naturally leads to two kinds of methods for conditioning: classifier-guided
(Nichol & Dhariwal, 2021) and classifier-free (Ho &
Salimans, 2022).
The former requires training an additional classifier p ϕ
( y | x t ) on noisy data so that samples may be generated at test-time with the perturbed noise ϵ θ ( x t , t ) − s √ 1 − ¯ α t ∇ x t log p ( y | x t ) , where s is referred to as the guidance scale.
The latter does not separately train a classifier but modifies the original training setup to learn both a conditional ϵ θ ( x t , y , t ) and an unconditional ϵ θ ( x t , t ) model for the noise.
The unconditional noise is represented, in practice, as the conditional noise ϵ θ ( x t , Ø , t ) where a dummy value Ø takes the place of y .
The perturbed noise ϵ θ ( x t , t )+ s ( ϵ θ ( x t , y , t ) − ϵ θ ( x t , t )) is used to later generate samples.
G ENERATIVE M ODELING WITH THE D ECISION D IFFUSER
It is useful to solve RL from offline data, both without relying on TD-learning and without risking distribution-shift.
To this end, we formulate sequential decision-making as the standard problem of conditional generative modeling:
[Equation]
Our goal is to estimate the conditional data distribution with p θ
so we can later generate portions of a trajectory x 0 ( τ ) from information
y ( τ ) about it.
Examples of y could include the return under the trajectory, the constraints satisfied by the trajectory, or the skill demonstrated in the trajectory.
We construct our generative model according to the conditional diffusion process:
[Equation]
As usual, q represents the forward noising process while p θ the reverse denoising process.
In the following, we discuss how we may use diffusion for decision making.
First, we discuss the modeling choices fof diffusion in Section 3.1.
Next, we discuss how we may utilize classifier-free guidance to capture the best aspects of trajectories in Section 3.2.
We then discuss the different behaviors that may be implemented with conditional diffusion models in Section 3.3.
Finally, we discuss practical training details of our approach in Section 3.4.
3. D IFFUSING OVER S TATES
In images, the diffusion process is applied across all pixel values in an image.
Na¨ıvely, it would therefore be natural to apply a similar process to model the state and actions of a trajectory.
However, in the reinforcement learning setting, directly modeling actions using a diffusion process has several practical issues.
First, while states are typically continuous in nature in RL, actions are more varied, and are often discrete in nature.
Furthermore, sequences over actions, which are often represented as joint torques, tend to be more high-frequency and less smooth, making them much harder to predict and model (Tedrake, 2022).
Due to these practical issues, we choose to diffuse only over states, as defined below:
[Equation]
Here, t denotes the timestep in the forward process and k denotes the time at which a state was visited in trajectory τ .
Moving forward, we will view x t ( τ ) as a noisy sequence of states from a trajectory of length H .
We represent x t ( τ ) as a two-dimensional array with one column for each timestep of the sequence.
Acting with Inverse-Dynamics.
Sampling states from a diffusion model is not enough for defining a controller.
A policy can, however, be inferred from estimating the action a k that led the state s k to s k +1 for any timestep k in x 0 ( τ ) .
Given two consecutive states, we generate an action according to the inverse dynamics model (Agrawal et al., 2016; Pathak et al., 2018):
[Equation]
Note that the same offline data used to train the reverse process p θ can also be used to learn f ϕ .
We illustrate in Table 2 how the design choice of directly diffusing state distributions, with an inverse dynamics model to predict action, significantly improves performance over diffusing across both states and actions jointly.
3. P LANNING WITH C LASSIFIER -F REE G UIDANCE
Given a diffusion model representing the different trajectories in a dataset, we next discuss how we may utilize the diffusion model for planning.
To use the model for planning, it is necessary to additionally condition the diffusion process on characteristics y
( τ ) .
One approach could be to train a classifier p ϕ ( y ( τ ) | x t ( τ )) to predict y ( τ ) from noisy trajectories x t ( τ ) .
In the case that y ( τ ) represents the return under a trajectory, this would require estimating a Q -function in the same way as Janner et al.
(2022), which requires a separate, complex dynamic programming procedure.
One approach to avoid dynamic programming is to directly train a conditional diffusion model conditioned on the returns
y ( τ ) in the offline dataset.
However, as our dataset consists of a set of sub-optimal trajectories, the conditional diffusion model will be polluted by such sub-optimal behaviors.
To circumvent this issue, we utilize classifier-free guidance (Ho & Salimans, 2022) with low-temperature sampling, to extract high-likelihood trajectories in the dataset.
We find that such trajectories correspond to the best set of behaviors in the dataset.
Formally, to implement classifier free guidance, a x 0 ( τ ) is sampled by starting with Gaussian noise x T ( τ ) and refining x t ( τ ) into x t − 1 ( τ ) at each intermediate timestep with the perturbed noise:
[Equation]
where the scalar s applied to ( ϵ θ ( x t ( τ ) , y ( τ ) , t ) − ϵ θ ( x t ( τ ) , Ø , t )) seeks to augment and extract the best portions of trajectories in the dataset that exhibit y ( τ ) .
With these ingredients, sampling from the Decision Diffuser becomes similar to planning in RL.
First, we observe a state in the environment.
Next, we sample states later into the horizon with our diffusion process conditioned on y and history of last K states observed.
Finally, we identify the action that should be taken to reach the most immediate predicted state with our inverse dynamics model.
This procedure repeats in a standard receding-horizon control loop described in Algorithmand visualized in Figure 3.
3. C ONDITIONING BEYOND R ETURNS
So far we have not explicitly defined the conditioning variable y ( τ ) .
Though we have mentioned that it can be the return under a trajectory, we may also consider guiding our diffusion process towards sequences of states that satisfy relevant constraints or demonstrate specific behavior.
We describe these scenarios in greater detail below.
Maximizing Returns
To generate trajectories that maximize return, we condition the noise model on the return of a trajectory so ϵ θ ( x t ( τ ) , y ( τ ) , t ) :
=
ϵ θ ( x t ( τ ) , R ( τ ) , t ) .
These returns are normalized
1: Input: Noise model ϵ θ , inverse dynamics f ϕ , guidance scale s , history length K , condition y 2: Initialize h ← Queue ( length = K ) , k ←// Maintain a history of length K 3: while not done do 4:
Observe state s ; h. insert ( s ) ; Initialize x T ( τ ) ∼ N (0 , αI ) 5: for t = T . . .
1 do 6: x t ( τ )[: length ( h )]
← h //
Constrain plan to be consistent with history 7:
ˆ ϵ ← ϵ θ ( x t ( τ ) , t )
+ s ( ϵ θ ( x t ( τ ) , y , t ) − ϵ θ ( x t ( τ ) , i )) //
Classifier-free guidance 8: ( µ t − 1 , Σ t − 1 ) ← Denoise ( x t ( τ ) , ˆ ϵ ) 9: x t − 1 ∼ N ( µ t − 1 , α Σ t − 1 ) 10: end for 11:
Extract ( s k , s k +1 ) from x 0 ( τ ) 12:
Execute a k = f ϕ
( s t , s k +1 ) ; k ← k + 1 13: end while to keep R ( τ ) ∈ [0 , 1] . Sampling a high return trajectory amounts to conditioning on R ( τ ) = 1 . Note that we do not make use of any Q -values, which would then require dynamic programming. Satisfying Constraints
Trajectories may satisfy a variety of constraints, each represented by the set C i , such as reaching a specific goal, visiting states in a particular order, or avoiding parts of the state space.
To generate trajectories satisfying a given constraint C i , we condition the noise model on a one-hot encoding so that ϵ θ ( x t ( τ ) , y ( τ ) , t ) :
= ϵ θ ( x t ( τ ) , 1 ( τ ∈ C i ) , t ) . Although we train with an offline dataset in which trajectories satisfy only one of the available constraints, at inference we can satisfy several constraints together. Composing Skills
A skill i can be specified from a set of demonstrations B i .
To generate trajectories that demonstrate a given skill, we condition the noise model on a one-hot encoding so that ϵ θ ( x t ( τ ) , y ( τ ) , t ) :
= ϵ θ ( x t ( τ ) , 1 ( τ ∈ B i ) , t ) .
Although we train with individual skills, we may further compose these skills together during inference.
Assuming we have learned the data distributions q ( x 0 ( τ ) | y 1 ( τ )) , . .
, q ( x 0 ( τ ) | y n ( τ )) for n different conditioning variables, we can sample from the composed data distribution q ( x 0 ( τ ) | y 1 ( τ ) , . . .
y n ( τ )) using the perturbed noise (Liu et al., 2022):
[Equation]
We use this property to compose more than one constraint or skill together at test-time.
3. T RAINING THE D ECISION D IFFUSER
The Decision Diffuser, our conditional generative model for decision-making, is trained in a supervised manner.
Given a dataset D of trajectories, each labeled with the return it achieves, the constraint that it satisfies, or the skill that it demonstrates, we simultaneously train the reverse diffusion process p θ , parameterized through the noise model ϵ θ , and the inverse dynamics model f ϕ with the following loss:
[Equation]
For each trajectory τ , we first sample noise ϵ ∼ N ( 0 , I ) and a timestep t ∼ U{ 1 , . .
, T } .
Then, we construct a noisy array of states x t ( τ ) and finally predict the noise as ˆ ϵ θ :=
ϵ θ ( x t ( τ ) , y ( τ ) , t ) .
Note that with probability p we ignore the conditioning information and the inverse dynamics is trained with individual transitions rather than trajectories.
Architecture
We parameterize ϵ θ with a temporal U-Net architecture, a neural network consisting of repeated convolutional residual blocks (Janner et al., 2022).
This effectively treats a sequence of states x t ( τ ) as an image where the height represents the dimension of a single state and the width denotes the length of the trajectory.
We encode the conditioning information y
( τ ) as either a scalar or a one-hot vector and project it into a latent variable z ∈ R h with a multi-layer perceptron (MLP).
When y ( τ ) =
Ø , we zero out the entries of z .
We also parameterize the inverse dynamics f
ϕ with an
MLP.
For implementation details, please refer to the Appendix B.
[Figure]
Low-temperature Sampling
In the denoising step of Algorithm 1, we compute µ t − 1 and Σ t −from a noisy sequence of states and a predicted noise.
We find that sampling x t − 1 ∼ N ( µ t − 1 , α Σ t − 1 ) where the variance is scaled by α ∈ [0 , 1) leads to better quality sequences (corresponding to sampling lower temperature samples).
For a proper ablation study, please refer to Appendix C.
E XPERIMENTS
In our experiments section, we explore the efficacy of the Decision Diffuser on a variety of different decision making tasks (performance illustrated in Figure 4).
In particular, we evaluate (1) the ability to recover effective RL policies from offline data, (2) the ability to generate behavior that satisfies multiple sets of constraints, (3) the ability compose multiple different skills together.
4. O FFLINE R EINFORCEMENT L EARNING Setup
We first test whether the Decision Diffuser can generate return-maximizing trajectories.
To test this, we train a state diffusion process and inverse dynamics model on publicly available
D4RL datasets (Fu et al., 2020).
We compare with existing offline RL methods, including modelfree algorithms like CQL (Kumar et al., 2020) and
IQL (Kostrikov et al., 2022), and model-based algorithms such as trajectory transformer ( TT , Janner et al. (2021)) and MoReL (Kidambi et al., 2020).
We also compare with sequence-models like the Decision Transformer ( DT ) (Chen et al.
and diffusion models like Diffuser (Janner et al., 2022)
[Table]
Results
Across a broad suite of different offline reinforcement learning tasks, we find that the
Decision Diffuser is either competitive or outperforms many of our offline RL baselines (Table 1). It [Table] [Table]
also outperforms Diffuser and sequence modeling approaches, such as Decision Transformer and
Trajectory Transformer.
The difference between Decision Diffuser and other methods becomes even more significant on harder D4RL Kitchen tasks which require long-term credit assignment.
To convey the importance of classifier-free guidance [Table] we also compare with the baseline
CondDiffuser , which diffuses over both state and action sequences as in Diffuser without classifierguidance.
In Table 2, we observe that CondDiffuser improves over Diffuser in 2 out ofenvironments.
Decision Diffuser further improves over CondDiffuser , performing better across all 3 environments.
We conclude that learning the inverse dynamics is a good alternative to diffusing over actions.
Finally, we compare against CondMLPDiffuser , a policy where the current action is denoised according to a diffusion process conditioned on both the state and return.
We see that
CondMLPDiffuser performs the worst amongst diffusion models.
4. C ONSTRAINT S ATISFACTION
Setup We next evaluate how well we can generate trajectories that satisfy a set of constraints using the Kuka Block Stacking environment (Janner et al., 2022) visualized in Figure 5.
In this domain, there are four blocks which can be stacked as a single tower or rearranged into several towers.
A constraint like BlockHeight ( i ) >
BlockHeight ( j ) requires that block i be placed above block j .
We train the Decision Diffuser from 10 , 000 expert demonstrations each satisfying one of these constraints.
We randomize the positions of these blocks and consider two tasks at inference: sampling trajectories that satisfy a single constraint seen before in the dataset or satisfy a group of constraints for which demonstrations were never provided.
In the latter, we ask the Decision Diffuser to generate trajectories
so BlockHeight ( i ) >
BlockHeight ( j ) >
BlockHeight ( k ) for three of the four blocks i, j, k .
For more details, please refer to Appendix D.
Results
In both the stacking and rearrangement settings, Decision Diffuser satisfies single constraints with greater success rate than Diffuser (Table 3).
We also compare with BCQ (Fujimoto et al., 2019) and CQL (Kumar et al., 2020), but they consistently fail to stack or rearrange the blocks and get a success rate of 0 . 0 .
Unlike these baselines, our method can just as effectively satisfy several constraints together according to Equation 9.
For a visualization of these generated trajectories, please see the website https://sites.google.com/view/decisiondiffuser/. [Figure] [Figure]
S KILL C OMPOSITION Setup
Finally, we explore how will we may compose different skills together.
We consider the
Unitree-go-running environment (Margolis & Agrawal, 2022), where a quadruped robot can be found running with various gaits, like bounding, pacing, and trotting.
We explore if it is possible to generate trajectories that transition between these gaits after only training on individual gaits separately.
For each gait, we collect a dataset of 2500 demonstrations.
Then, we train the Decision Diffuser to imitate each of these these trajectories.
Results
During testing, we compose the noise model of our reverse diffusion process according to equation 9.
This allows us to sample trajectories of the quadruped robot with entirely new running behavior.
Figure 6 shows a trajectory that begins with trotting but ends with pacing.
Please see the
Appendix E for more visualizations of composing running gaits together.
R ELATED W ORK Diffusion Models
Diffusion Models have shown great promise in learning generative models of image and text data (Saharia et al., 2022; Nichol et al., 2021; Nichol & Dhariwal, 2021).
It formulates the data sampling process as an iterative denoising procedure (Sohl-Dickstein et al., 2015; Ho et al., 2020).
The denoising procedure can be alternatively interpreted as parameterizing the gradients of the data distribution (Song et al., 2021) optimizing the score matching objective (Hyv ¨ arinen, 2005) and Energy-Based Models (Du & Mordatch, 2019;
Nijkamp et al., 2019; Grathwohl et al., 2020).
To generate data samples (eg: images) conditioned on some additional information (eg:text), prior works (Nichol & Dhariwal, 2021) have learned a classifier to facilitate the conditional sampling.
More recent works (Ho & Salimans, 2022) have argued to leverage gradients of an implicit classifier, formed by the difference in score functions of a conditional and an unconditional model, to facilitate conditional sampling.
The resulting classifier-free guidance has been shown to generate better conditional samples than classifier-based guidance.
All these above mentioned works have mostly focused on generation of text or images.
Janner et al.
(2022) used diffusion models to learn a generative model of trajectories (consisting of states and actions) for decision-making.
However, Janner et al.
(2022) relies on learning a Q -function to guide the trajectory model towards high-return trajectories, thereby still retaining the instabilities and the complexities of offline RL pipeline.
In contrast, Decision Diffuser learns a conditional diffusion model over state sequences and uses classifier-free guidance for sampling high-return state sequences, thereby side-stepping learning of Q -function and avoiding any complexities of offline RL.
Reward Conditioned Policies
Prior works (Kumar et al., 2019; Schmidhuber, 2019; Srivastava et al., 2019; Emmons et al., 2021; Chen et al., 2021) have studied learning of reward conditioned policies via reward conditioned behavioral cloning.
Chen et al.
(2021) used a transformer (Vaswani et al., 2017) to model the reward conditioned policies and obtained a performance competitive with offline RL approaches.
Emmons et al.
(2021) obtained similar performance as Chen et al.
(2021) without using a transformer policy but relied on careful capacity tuning of MLP policy.
In contrast to these works, in addition to modeling returns, Decision Diffuser can also model constraints or skills and generate novel behaviors by flexibly combining multiple constraints or skills during test time.
D ISCUSSION
We propose Decision Diffuser, a conditional generative model for sequential decision making.
It frames offline sequential decision making as conditional generative modeling and sidesteps the need of reinforcement learning, thereby making the decision making pipeline simpler.
By sampling for high returns, it is able to capture the best behaviors in the dataset and outperforms existing offline RL approaches on standard benchmarks (such as D4RL).
In addition to returns, it can also be conditioned on constraints or skills and can generate novel behaviors by flexibly combining constraints or composing skills during test time.
In this work we focused on offline sequential decision making, thus circumventing the need for exploration.
Using ideas from Zheng et al.
(2022), future works could look into online fine-tuning of
Decision Diffuser by leveraging entropy of the state-sequence model for exploration.
While our work focused on state based environments, it can be extended to image based environments by performing the diffusion in latent space, rather than observation space, as done in Rombach et al.
Appendix
In this appendix, we discuss details of the illustrative example in Section A. Next, we discuss hyperparameters and architectural details in Section
B. We analyze the importance of low temperature sampling in Section C.
Finally, we provide details of the Kuka Block Stacking environment in
Section D and the Unitree environment in Section E.
A I LLUSTRATIVE E XAMPLE Setup
In linear system robot navigation, Decision Diffuser is trained on 1000 expert trajectories either satisfying the constraint
∥
s T ∥ ≤ R ( R = 1 ) or the constraint
∥
s T ∥ ≥ r ( r = 0 . 7 ).
Here, s T =
[ x T , y T ] represents the final robot state in a trajectory, specifying its final 2d position.
The maximum trajectory length is 50 .
During test time, Decision Diffuser is asked to generate trajectories satisfying ∥ s T ∥ ≤ R and ∥ s T
∥ ≥ r to test its ability to satisfy single constraints.
Furthermore,
Decision Diffuser is also asked to generate trajectories satisfying r ≤ ∥ s T ∥ ≤ R to test its ability to satisfy combined constraints.
Results
Figure 2 shows that Decision Diffuser learns to generate trajectories perfectly (i.e. with 100% success rate) satisfying single constraints in linear system robot navigation.
Furthermore, it learns to generate trajectories satisfying the composed constraint in linear system robot navigation with 91 .
3%( ± 2 .
6%) accuracy where the standard error is calculated over 5 random seeds.
B H YPERPARAMETER AND A RCHITECTURAL DETAILS
In this section, we describe various architectural and hyperparameter details:
• We represent the noise model ϵ θ with a temporal U-Net (Janner et al., 2022), consisting 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 and condition embeddings, both 128 -dimensional vectors, are produced by separate 2-layered MLP (with 256 hidden units and Mish nonlinearity) and are concatenated together before getting added to the activations of the first temporal convolution within each block.
We borrow the code for temporal U-Net from https://github.com/jannerm/diffuser.
• We represent the inverse dynamics f
ϕ with a 2-layered MLP with 512 hidden units and ReLU activations.
• We train ϵ θ and f ϕ using the Adam optimizer (Kingma & Ba, 2015) with a learning rate of 2 e −and batch size of 32 for 1 e 6 train steps.
• We choose the probability p of removing the conditioning information to be 0 .
• We use T = 100 diffusion steps.
• We use a planning horizon H of 100 in all the D4RL locomotion tasks, 56 in D4RL kitchen tasks, 128 in Kuka block stacking, 56 in unitree-go-running tasks and 50 in the illustrative example.
• We use a guidance scale s ∈ { 1 . 2 , 1 . 4 , 1 . 6 , 1 . 8 } but the exact choice varies by task.
• We choose α = 0 .
5 for low temperature sampling.
• We choose context length K = 20 .
C I MPORTANCE OF LOW TEMPERATURE SAMPLING
In Algorithm 1, we compute µ t − 1 and Σ t − 1 from a noisy sequence of states and predicted noise.
We find that sampling x t − 1 ∼ N ( µ t − 1 , α Σ t − 1 ) (where α ∈ [0 , 1) ) with a reduced variance produces high-likelihood state sequences.
We refer to this as low-temperature sampling.
To empirically show its importance, we compare performances of Decision Diffuser with different values of α (Table A1).
We show that low temperature sampling ( α = 0 . 5) gives the best average returns.
However, reducing the α to 0 eliminates the entropy in sampling and leads to lower returns.
On the other hand, α = 1 .leads to a higher variance in terms of returns of the trajectories.
[Table]
D K UKA B LOCK S TACKING
In the Kuka blocking stacking environment, the underlying goal is to stack a set of blocks on top of each other.
Models have trained on a set of demonstration data, where a set of 4 blocks are sequentially stacked on top of each other to form a block tower.
We construct state-space plans of length 128.
Following (Janner et al., 2022), we utilize a close-loop controller to generate actions for each state in our state-space plan (controlling the 7 degrees of freedom in joints).
The total maximum trajectory length plan in Kuka block stacking is 384.
We detail differences between the two consider conditional stacking environments below:
• Stacking
In the stacking environment, at test time we wish to again construct a tower of four blocks.
Rearrangement
In the rearrangement environment, at test time wish to stack blocks in a configuration where a set of blocks are above a second set.
This set of stack-place relations may not precisely correspond to a single block tower (can instead construct two block towers), making this environment an out-of-distribution challenge.
In addition to Diffuser (Janner et al., 2022), we used goal-conditioned variants of CQL (Kumar et al., 2020) and BCQ (Fujimoto et al., 2019) as baselines for the block stacking and rearrangement with single constraint.
However, they get a success rate of 0 . 0 .
E U NITREE G O R UNNING
We consider Unitree-go-running environment (Margolis & Agrawal, 2022) where a quadruped robot runs in 3 different gaits: bounding, pacing, and trotting.
The state space is 56 dimensional, the action space is 12 dimensional, and the maximum trajectory length is 250 .
As described in Section 4.3, we train Decision Diffuser on expert trajectories demonstrating individual gaits.
Figure A1 shows the ability of Decision Diffuser to imitate bounding, trotting and pacing and their combinations.
[Figure]