Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning
Authors: Zhendong Wang
Jonathan J Hunt
Mingyuan Zhou
Venue: ICLR 2023 poster (2023)
Keywords: offline RL
diffusion models
behavior cloning
policy regularization
Q-learning

Abstract: Offline reinforcement learning (RL), which aims to learn an optimal policy using a previously collected static dataset, is an important paradigm of RL. Standard RL methods often perform poorly in this regime due to the function approximation errors on out-of-distribution actions. While a variety of regularization methods have been proposed to mitigate this issue, they are often constrained by policy classes with limited expressiveness that can lead to highly suboptimal solutions. In this paper, we propose representing the policy as a diffusion model, a recent class of highly-expressive deep generative models. We introduce Diffusion Q-learning (Diffusion-QL) that utilizes a conditional diffusion model to represent the policy. In our approach, we learn an action-value function and we add a term maximizing action-values into the training loss of the conditional diffusion model, which results in a loss that seeks optimal actions that are near the behavior policy. We show the expressiveness of the diffusion model-based policy, and the coupling of the behavior cloning and policy improvement under the diffusion model both contribute to the outstanding performance of Diffusion-QL. We illustrate the superiority of our method compared to prior works in a simple 2D bandit example with a multimodal behavior policy. We then show that our method can achieve state-of-the-art performance on the majority of the D4RL benchmark tasks.

Published as a conference paper at ICLR 2023
DIFFUSION POLICIES AS AN EXPRESSIVE POLICY CLASS FOR OFFLINE REINFORCEMENT LEARNING
Zhendong Wang".", Jonathan J Hunta.t, Mingyuan Zhou 'Te University of Texas at Austin, 2 Twitter
zhendong. wang@ut exas. exaas. mingyuan. zhou thow@mccombs utexas. edd
edu, juunt wowuntttwittte..omm com
ABSTRACT
Offline reinforcement learning (RL), which aims to learn an optimal policy using a previously collected static dataset, is an important paradigm of RL. Standard RL methods often perform poorly in this regime due to the function approximation errors on out-of-distribution actions. While variety of regularization methods have been proposed to mitigate this issue, they are often constrained by policy classes with limited expressiveness that can lead to highly suboptimal solutions. In this paper, we propose representing the policy as a diffusion model, a recent class of highly-expressive deep generative models. We introduce Diffusion Q- learning (Diffusion-QLL) that utilizes a conditional diffusion model to represent the policy. In our approach, we learn an action-value function and we add term maximizing action-values into the training loss of the conditional diffusion model, which results in a loss that seeks optimal actions that are near the behavior policy. We show the expressiveness of the diffusion model-based policy, and the coupling of the behavior cloning and policy improvement under the diffusion model both contribute to the outstanding performance of Diffusion-QI. We illustrate the su- periority of our method compared to prior works in a simple 2D bandit example with multimodal behavior policy. We then show that method can achieve state-of of-te---the-rr performance on the majority of the D4RL benchmark tasks.
INTRODUCTION
Offline reinforcement learning (RL), also known as batch RL, aims at learning effective policies entirely from previously collected data without interacting with the environment (Lange et al.) 2012] Fujimoto et al., 2019). Eliminating the need for online interaction with the environment makes offline RL attractive for wide array real-world applications, such as autonomous driving and patient treatment planning, where souldwordd exploration with an untrained policy is risky, expensive, or consumin..... Instead of relying on world exploration, offline RL emphasizes the use of prior data, such human demonstration, that is often available interactions. However, relying only on previously collected data makes offline RL a challenging task. Applying standard policy improvement approaches to an offline dataset typically leads to relying on evaluating actions that have not been seen in the dataset, and therefore their values are unlikely to be estimated accurately. For this reason, naive approaches offline RL typically learn poor policies that prefer out-of distriibbttion actions whose values have been overestimated, resulting unsatisfactory performance (Fujimoto al. 2019).
much lower cost than online
Previous work on offline RL generally addressed this problem in one of four ways: 1) regularizing how far the policy can deviate from the behavior policy (Fujimoto et al., 2019 Fujimoto & Gu 2021
Kumar et al.) 2019 Wu et al., 2019 Nair et al. 2020 [yu et al. (022); 2) constraining the learned value function to assign low values to out-of-distribution actions (Kostrikov et al.) [22 Kumar et al., 2020); 3) introducing model-based methods, which learn model of the environment dy- namics and perform pessimistic planning in the learned Markov decision process (MDP) (Kidambi
"The work was done in part during a summer internship at Twitter.
Joint senior authors; order determined by flipping coin.

---

Published as a conference paper at ICLR 2023
et al. 2020; Yu et al., 2021 4) treating offline RL as a problem of sequence prediction with return guidance (Chen et al. 2021 1 Janner et al. 2021:000 2022). Our approach falls into the first category.
Empirically, the performance of policy-regularized offline RL methods is typically slightly worse than that of other approaches, and here we show that this is largely because the policy regularization methods perform poorly due to their limited ability to accurately represent the behavior policy. This results in the regularization adversely affecting the policy improvement. For example, the policy regularization may limit the exploration space of the agent to a small region with only suboptimal actions and then the Q-learning will be induced to converge towards a suboptimal policy.
The inaccrrate policy regularization occurs for main reasons: 1) policy classes are not ex- pressive enough: 2) the regularization methods are improper. In most prior work, the policy is a Gaussian distribution with mean and diagonal covariance specified by the output of neural net- work. However, as offline datasets are often collected by mixture of policies, the true behavior policy may exhibit strong nulti-modalities, skewness, or dependencies between different action di- mensions, which cannot be well modeled by diagonal Gaussian policies (Shafiullah et al. (2022). In a particularly extreme, but not uncommon example, a Gaussian policy is used to fit bimodal training data by minimizing the Kullback-Leibler (KL) divergence from the data distribution to the policy distribution. This will result in the policy exhibiting modc-covering behavior and placing high den- sity in the middle area of the two modes, which is actually the low-density region of the training data. In such cases, regularizing a new policy towards the behavior-cloned policy is likely to make the policy learning substantially worse. Second, the regularization, such the KL divergence and maximum mean discrepancy (MMD) (Kumar et al. 2019). is often not well suited for offline RL. The KL divergence needs access to explicit density values and MMD needs multiple action sam- ples at each state for optimization. These methods require an extra step by first learning a behavior cloned policy to provide density values for KL optimization or random action samples for MMD optimization. Regularizing the current policy towards the behavior cloned policy can further induce approximation errors, since the cloned behavior policy may not model the true behavior policy well, due to limitations in the expressiveness of the policy class. We conduct a simple bandit experiment in Section 4 which illustrates these issues can occur even on a simple bandit task.
In this work, we propose method to perform policy regularization using diffusion (or score-based) models (Sohl-Dickstein et al. 2015 Song & Ermon 2019 Ho et al. 2020). Specifically, we use multilayer perceptron (MLP) based denoising diffusion probabilistic model (DDPM) (Ho et al., 2020) as our policy. We construct an objective for the diffusion loss which contains two terms: I) behavior-cloning term that encourages the diffusion model to sample actions in the same distribu- tion as the training set, and 2) a policy improvement term that attempts to sample high-value actions (according a learned Q-value). Our diffusion model is conditional model with states as the con- dition and actions as the outputs. Applying diffusion model here has several appealing properties. First, diffusion models are very expressive and can well capture multi-modal distributions. Second, the diffusion model loss constitutes a strong distribution matching technique and hence it could be seen as powerful sample-basee based policy regularization method without the need for extra behavior cloning. Third, diffusion models perform generation via iterative refinement, and the guidance from maximizing the Q-value function can be added at each reverse diffusion step.
In summary, our contribution is Diffusion-Ql. sion models to do precise policy regularization and successfully injects the Q-learning guidance into the reverse diffusion chain to seek optimal actions. We test Diffusion-QL on the D4RL benchmark tasks for offline RL and show this method outperforms prior methods on the ma- jority of tasks. We also visualize the method outperform prior methods. Code is available at httpps: /tttt::/qqiiee /ttp///tttu///t/ com/ 2hendong-Nang/ Diffus - Prolicien sellicies-se ies-ssssss wwr.iiieesffooffl-ooee Ofr--fffine-RR
new offline RL algorithm that leverages diffu-
simple bandit task to illustrate why it can
2 PRELIMINARIES AND RELATED WORK
Offline RL. The environment in RL is typically defined by Markov decision process (MDP): M = P. do), with state space S, action space A, environment dynamics P(s') s, a) : 1xSxA-L [0. 1]. reward function R S x A + R, discount factor y E [0, 1), and initial state distribution dn (Sutton& & Barto 2018). The goal is to learn policy To (a) s). parameterized by 8 that maximizes the cumulative discounted reward E. EEe ( a:)). The action-value Q-value of
2

---

Published as a conference paper at ICLR 2023
Ei= r'r(s at). In the offline setting (Fu {(s, a.r. s')),, collected by a behavior policy Th. is provided. Offline RL algorithms learn policy entirely from this static offline dataset D, without online interactions with the environment.
policy is defined as Q* ai) Ba 11.0 12.2..
al. 2020), instead of the environment, a static dataset D -
Diffusion Model. Diffusion-basel generative models (Ho et al. 2020 Sohl-Dickstein et al. 2015 Song & Ermon 2019) assume po(2n) : po (ato:r) where Ii: #T are latent variables of the same dimensionality the data Zo ~ p(æo). A forward diffussion chain gradually adds noise to the data 20 ~ ((x)) in T steps with a pre-defined variance schedule 3, expressed as
q(21T za)= = Il= g( xx-1),,
q(x; x1-1) == NN(x; VI
A reverse diffusion chain, constructed as (o(or))) := W(x 0, IT (0t xx), is then optimized by maximizing the evidence lower bound defined as E2y a(mao ma
(o(
Jordan et al.
1999 Blei et al. 2017). After training. sampling from the diffusion model consists of sampling TT p( FT and rumning the reverse diffusion chain to go from t = T to 0. Diffusion models can be straightforwardly extended to conditional models by conditioning po (at--1 x2,
Related Work: Policy Regularization. Most prior methods for offline RL in the class of regular- ized policies rely on behavior cloning for policy regularization: BCQ (Fujimoto et al. 2019) con- structs the policy as a learnable and Conditional-Val VAE (CVAE, |ohnn et al. (2015)) behavior-cloning model; BEAR (Kumar et al. 2019) adds weighted behavior-clonin loss via minimizing MMD into the policy improvement step; TD3+BC (Fujimoto & Gu 2021) applies the same trick as BEAR via maximum likelihood estima- tion (MLE); BRAC (Wu et 2019)) evaluates multiple methods for behavior-cloning regularization, such as the KI. divergence, MMD, and Wasserstein dual form; IQI. (Kostrikov et al. (02 is an advantage weighted behavior-cloning method with "in-sample" learned Q-value functions, Goo & Nickum (2022) emphasize the necessity of conducting explicit behavioral cloning in offline RL, while Ajay et al. (2022) admit the power of conditional generative models for decision making.
deviation from a separately learned
Related Work: Diffusion Models in RL. Pearce et al. (2023) propose to better imitate human behaviors via diffusion models which are expressive and stable. Diffuser (Janner et al. 2022) applies diffusion model as trajectory generator. The full trajectory of state-action pairs form single sample for the diffusion model. A separate return model is learned to predict the cumulative rewards of each trajectory sample. The guidance of the return model is then injected into the reverse sampling stage. This approach is similar to Decision Transformer (Chen et al.] 2021). which also learns a trajectory generator through GPT2 (Radford al.) 2019) with the help of the true trajectory returns. When used online, sequence models can no longer predict actions from states autoregressively (since the states are an outcome of the environment). Thus, in the evaluation stage, whole trajectory is predicted for each state while only the first action is applied, which incurs a large computational cost. Our approach employs diffusion models for RL distinct manner.
We apply the diffusion model to the action space and we form it as conditional diffusion model with states as the condition. This approach model-free and the diffusion model is sampling a single action at time. Further, our Q-value function guidance is injected during training, which provides good empirical performance in our case. While both Diffuser (Janner et al. (022)) and our
work apply diffusion models in Offline RL, Diffuser is from the model-based trajectory-planning perspective while our method is from the offline model-free policy- optimizatiiiiiio perspective.
DIFFUSION Q-LEARNING
Below we explain how we apply a conditional diffusion model as an expressive policy for behavior cloning. Then, we introduce how we add Q-learning guidance into the learning of our diffusion model the training stage with the behavior cloning term acting as a form of policy regularization.
3.1 DIFFUSION POLICY
Notation:
Since there two different types of timesteps in this work, one for the diffusion process and one for reinforcement learning we use superscripts E [1.....N] to denote diffusion timestep and subscripts / (1,..., .,. T} to denote trajectory timestep.
3

---

Published as a conference paper at ICLR 2023
We represent our RL policy via the reverse process of a conditional diffusion model Ta(a| 8) Po(a 8) ; II [ pa la,s)
where the end sample of the reverse chain, a", is the action used for RL evaluation. Generally, p aa., 8) could be modeled as Gaussian distribution WNaa" pg (a', 8,1), Eg(a', 8, s,))). We follow Ho et al. (2020) to parameterize po po(a a', s) as noise prediction model with the covariance matrix fixed as Eo(a', s, B,I and mean constructed
Mo(a' 8, va;
1 (a' -ad--Pe vi-a; co(a', s,)))
We first sample a a'-1 11'1aa a' Following DDPM (Ho al. 2020), when 1, is set 0 to improve the sampling quality. We mimic the simplified objective proposed by Ho et al. (020) to train our conditional e-model via La(0)= =
(0, 1) and then from the reverse diffusion chain parameterized by 0 as
3:
VBec, (O,, I), for i ......
Vas Vo.(1-a.)
|||e EØ( Va:a + vI - 0,,,, (,,]1
(2)
where U is a uniform distribution over the discrete set as (1.....N] and D denotes the offline dataset, collected by behavior policy m.. This diffusion model loss C0()) is behavior-cloning loss, which aims to learn the behavior policy (a 8) (i.e. seeks to sample actions from the same distribution the training data). Note the marginal of the reverse diffusion chain provides an im- plicit, expressive distribution that can capture complex distribution properties, such as skewness and multi-modality, exhibited by the offline datasets. addition, the regularization is sampling- -aased that only requires taking random samples from both 'D and the current policy (i.e. this method does not require us to know the behavior policy, which may be infeasible when the dataset is collected by human demonstrations), Different from the usual two-step strategy, our strategy provides a clean and effective way of applying regularization on a flexible policy.
Ca(O) can efficiently be optimized by sampling single diffusion step for each data point, but the reverse sampling in Equation ( which requires iteratively computing Eg networks times, can become a bottleneck for the rumning time. Thus may want to limit N to relatively small value. To work with small N, with Amin = 0.1 and Finx = 10.0, we follow Xiao et al. 2021 to define
B -- = - e Smti ) 0. - -min))
2/-1
N
which is a noise schedule obtained under the variance preserving SDE of [Son et al. (021).
3.2 Q-LEARNING
The policy-regularrizaion loss L0( is a behavior-cloning term, but would not result in learning policy that can outperform the behavior policy that generated the training data. To improve the policy, we inject Q-value function guidance into the reverse diffusion chain in the training stage in order learn to preferentially sample actions with high values. The final policy-learning objective is a linear combination of policy regularization and policy improvement:
arg min (()) La(O) + (0) = La( a EgnDaa (Qo(s.a)).
(3)
Note that a is reparameterized by Equation 1 and hence the gradient of the Q-value function with respect to the action is backpropagated through the whole diffusion chain.
As the scale of the Q-value function varies in different offline datasets, normalize it, follow Fujimoto& & Gu (2021) to set a as a = (inya)=llllst.aar) where is hyperparameter that balances the two loss terms and the Q in the denominator is for normalization only and not differentiated over.
The Q-value function itself is learned in a conventional way, minimizing the Bellman operator (Li-- licrap et al.] 2015 Fujimoto al. 2019) with the double Q-learning trick (Hasselt 2010). We built two Q-networks, Q Qøy and target networks Qoy Qsy and Tar. We then optimize Φ; for i = [1,2) by minimizing the objective,
[[((sss) ar) + m Qo! (8 aft1) Qo. (s, au)l
(4)
We conduct extensive experiments in Sections and 5 and show that Ca and L, work together to achieve the best performance. We summarize our implementation in Algorithm
4

---

Published as a conference paper at ICLR 2023
Algorithm Diffusion Q-learning
Initialize policy network TT, critic networks Qo and Qø and target networks Tor, Qot and Qat
for each iteration do
Sample transition mini-batch B = {(8t, ae Tt, 11+1)} D. # Q-value function learning
Sample a,+1 ~ Te. (ac+1 sc+1) by Equation
Update Qo, and QΦ by Equation (max Q backup by Kumar et al. (2020) could be added)
# Policy learning
Sample a¹ ~ To (a, st) by Equation 1
Update policy by minimizing Equation 3
# Update target networks
8' pff + (1 p)0, 0 = pot (1 p)o, for = {1.2].
end for
POLICY REGULARIZATION
In this section, we illustrate how the previous policy regularization methods work compared to our conditional diffusion-bes based approach a simple bandit task with 2D continuous action space. Below we first provide a brief review of prior methods.
BC-MLE, The policy ( 8) is modeled by Gaussian distribution N(a; (st)) [s)), where usually μ4 and Σg are parameterized by multi-layer perceptrons (MLPs) and for simplicity Eg is assumed to be diagonal matrix. The policy is optimized by maximizing (,,, [log ma(a 8)]- TD3+BC (Fujimoto & Gu 2021) directly add the behavior-cloning (BC) loss as an additional term in policy learning, while IQL (Kostrikov al., 2021b) proposes using "in-sample" learned advantage functions to reweigh the log term inside the expectation.
BC-CVAE.
The policy (a s) is modeled by CVAE model with an encoder network (e s, a) and a decoder network Pa 8, =). The two networks are optimized by maximizing the evidence lower bound a (E [Rz-g(t1 s.a) [log [logp(a |(a| =s, (L(g(2 s, a a))pp=)))] where p(z) prior distribution that is usually set as standard Gaussian. BCQ (Fujimoto et al. (019) trains CVAE model as an approximation of the behavior policy and trains another deviation model to guide the actions drawn from the CVAB approximation towards the regions with high learned Q-values.
BC-MMD. BEAR (Kumar et al., 2019) also mimies the behavior policy via a CVAE model and proposes to limit the current policy (a s) to be close to the cloned behavior policy via MMD minimization. A Tanh-Gaussian policy used, which is a Gaussian network with a Tanh activation function at the output layer.
BCQ and BEAR can be seen as a two-step regularization: First, an approximation of the behavior policy is learned (behavior cloning), and then the policy learned by policy improvement is regular- ized towards the cloned behavior policy. However, such a two-step approach means that the efficacy of the second-step policy regularization heavily depends the cloning quality, and inaccurate regularization could misguide the subsequent policy improvement step. We illustrate this weakness in 2D continuous action space bandit example.
Example. We consider a simple bandit task with real-valued actions in a 2D space, a E [-1, We construct an offline dataset D {{a;)} with M 10000 action examples, where the actions are collected from an equal mixture of four Gaussian distributions with centers fe e {00.0, 0.8), (0.8, 0.0), (0.0, -..8). ( (-0.8, 0.0)) and standard deviations σ (0.05, 0.05), as de- picted in the first panel of Figure [ Note that this example exhibits strong multi-modality in the behavior policy distribution, which is often the case when the dataset is collected by different poli- cies. For example, if multiple humans are involved, some may be experts and choose actions from different mode from amateur demonstrators.
To evaluate the strength of prior regularization methods, we first compare them to our diffusion- based approach on a behavior-cloning task, where the goal is to just clone the behavior policy that generated the data, not improve it. As shown in the first row of Figure diffusion model captures all the four density modes of the behavior policy. The policy of BC-MLE is limited single mode and hence exhibits strong mode-covering behavior. It fits the four den-
observe that the
5

---

Published as a conference paper at ICLR 2023
Ground Truth
BCALEE
SC-VAE
BCHMD
BC asien
Add
Diffuvion-
Figure 1: Offline RL experiments on a simple bandit task. The first row shows the comparison of behavior cloning between method (BC-Diffusion, . 50) and prior methods. Prior methods struggle to capture the multi- molli- behavior policy (ground truth). The second row shows the comparison results when policy improvement added (first figure shows the rewards). The policy regularization prior methods results in a poor policy, since the behavior-cloming step fails to capture the multi-modal hchavior policy, while our method (Diffusion-QL, Q correctly identifies the high reward behavior mode.
BC 2
BC Diffusionin-S1
BC 801
BC OrtusionN-201
BC DiffusionN-
OMtusion-UUN-22
Outusion-OUNNi00
OUN:000
DiMtusion OLINo
Figure 2: Experiment examining the effect of varying the number of diffusion steps N on the simple bandit talks. The first shows the ablation study of N for BC-Diffusion. The second row shows the ablation study of N. for Diffusion-QI.. Large N results in better fit the data distribution, but leads higher computational cost. By combining the behavior cloning and Q-value losses during training, we are able to learn near-optimal policies with fewer diffusion steps.
sity modes by one Gaussian distribution with large standard deviation, whose high-density regions are actually the low-density regions of the true behavior policy. The CVAE model is shown to ex- hibit mode-covering behavior even though itself an implicit model with enough expressiveness: We see that CVAE captures the four modes but high densities are assigned between them to cover all the modes. Note we also observe the CVAE model sometimes fails to capture all four modes with different random seeds. The Tanh-Gaussin densities near the boundary lines and, due to its limited policy expressiveness, fails to capture the true distribution. These observed failures behavior cloning in this simple example illustrate the limitations of prior regularization methods in the common case of multi-modal behavior policies.
policy optimized under MMD learns to align the
Next, investigate how the policy improvement will be impacted by the corresponding policy regularization. We assign each data point a reward sampled from a Gaussian distribution, whose mean is determined by the data center and standard deviation is fixed as 0.5, as shown in the second row of Figure [] Note here we mimic the offline RL setting, under which the underlying reward function is unknown and needs to be learned. We compare Diffusion Q-learning (QL) with prior methods, including TD3+BC, BCQ. and BEAR-MMD, We train all methods with 1000 epochs to ensure convergence. Due to the strong policy constraint applied by each method, we observe that the policy improvement of prior methods is constrained to suboptimal or even wrong exploration regions, induced by the corresponding behavior-cloning regularization. TD3+BC cannot converge to the optimal mode since the behavior policy places most density in the region where no offline data exists. BCQ learns to place major actions on the four diagonal corners discovered by its CVAE-based behavior cloning. The policy of BEAR- AMAD learns place actions randomly since the exploration region constrained by inaccurate policy regularization. We observe that the prior regularizations typically push the policy to converge to sub-optimal solutions, such BCQ, prevent the policy

---

Published as a conference paper at ICLR 2023
from being concentrated on the optimal corner, such as TD3+BC and BEAR-MMD By contrast, the policy of Diffusion- QL successfully converges to the optimal bottom corner. This is because 1) diffusion policy is expressive enough to recover the behavior policy, which covers all modes for further exploration; 2) the Q-learning guidance, directly through linearly combined loss functions in Equation (3). helps diffusion policy seek optimal actions in the region. The two components working together to produce good performance.
Diffusions steps. We further investigated how the diffusion policy performs as the number of dif- fusion timesteps N is varied. As expected, the first row of Figure shows diffusion model becomes more expressive and learns more details about the underlying data distri- bution. When N is increased to 50, the true data distribution is accurately recovered. The second row shows that with Q-learning applied, a moderately small N able to deliver good performance due to our loss coupling defined in Equation (3). However, we can see with larger N the policy regularization imposed by the diffusion model-based cloning becomes stronger. For example, when N 2, there are still few action points sampled near regions with no training data, whereas when N 50, the policy is constrained in the correct data region for further exploration. The num- ber of timesteps N serves Diffusion-QU QL. We found N = 5 performs well on D4RL (Fu et al., (2020) datasets, which is also a small enough value for cost-effective training and deployment.
shows that as N increases, the
trade-off between policy expressiveness and computational cost for
5 EXPERIMENTS
We evaluate our method on the popular D4RL (Fu et al., 2020) benchmark. Further, we conduct empirical studies on the number of timesteps required for our diffusion model and also perform an ablation study for analyzing the contribution of the two main components of Diffusion-QL.
Datasets. We consider four different domains of tasks in benchmark: Gym, AntMaze, Adroit, and Kitchen. The Gym-MuloCo locomotion tasks are the most commonly used standard tasks for evaluation and are relatively easy, since they usually include a significant fraction of near-optimal trajectories in the dataset and the reward function is quite smooth. AntMaze consists of more chal- lenging tasks, which have sparse rewards and explicitly need the agent to stitch various sub-optimal trajectories to find path towards the goal of the maze (Fu et al.] (020). Adroit datasets are mostly collected by human behavior and the state-action region reflected by the offline data is often very narrow, strong policy regularization is needed to ensure that the agent stays in the expected re- gion. The Kitchen environment requires the agent to complete 4 target subtasks in order to reach a desired state configuration, and hence long-term value optimization is important for it.
Baselines. We consider different classes of baselines that perform well in each domain of tasks. For policy regularization-based methods, we include the classic BC, BEAR (Kumar et al. 2019), BRAC et al. 2019), BCQ (Fujimoto et al., 2019). TD3+BC (Fujimoto & Gul (021). AWR (Peng et al,, 2019). AWAC (Nair et al.. (0020). and IQL (Kostrikov et al.) (021b), along with the Onestep RL [Brandfonbrener et al.] 2021). which is based on single-step improvement. For Q-value constraint methods, we include REM (Agarwal et al.) 2020) and CQL (Kumar et al., (020). For model-based offline RL, we consider MoRel (Kidambi et al. 2020). For sequence modelling approaches, we compare with Decision Transformer (DT) (Chen et al. 202T) and Diffuser (Janner et al. (022). We
report the performance of baseline methods either using the best results reported from their own paper, Fu al, (2020) or Kostrikov et al. (0011b).
Experimental details. We train for 1000 epochs (2000 for Gym tasks). Each epoch consists of 1000 gradient steps with batch size 256. The training is usually quite stable as shown in Figure except that we observe the training for AntMaze tasks has variations due to its sparse reward setting and lack of optimal trajectories in the offline datasets. Hence, we save multiple model checkpoints during training and use a completely offline method, described in Appendix [ to select the best checkpoint for performance evaluation. Using La loss as a lagging indicator of online performance, perform early stopping and select the checkpoint with the second third lowest La value. The results in our main paper are all based on the offline model selection. When a small amount of online interaction can be used for model selection, can achieve even better results (Appendix Table
Effect of hyperparameter N. We conduct an empirical study on the effect of the number of timesteps N of our diffusion model real tasks. As shown in Figure 3 empirically we find as
7

---

Published as conference paper at ICLR 2023
hopper-eedum--xpertv
wallerzbb-eeium-exppere
1.0
1.25
Normal maddzed Score
1.0
0.8
Normalized
0.6-
0.8
0.75
0.4
0.4-
0.25
0.2
0.00
20
20
40
30
10
20
Training Progress 50k steps)
Training Progress steps)
Training Progress (# 50k steps)
Figure 3: Ablation study of N on selected Gym tasks. We consider N grid [2, 5, 10, 20] and we find that
N = 5 is good enough for the selected tasks.
Table 1: The performance of Diffusion-QL. and SOTA baselines on D4RL D4RL.Gym, AntMaze, Adroit, and Kitchen tasks. Results for Diffusion-QI. correspond to the mean standard errors normalized scores over 50 random rollouts (5 independently trained models and 10 trajectories per model) for Gym tasks, which generally exhibit low variance in performance, and over 500 random rollouts (5 independently trained models and 100 trjecttories per model) for the other tasks. Note the standard error of AntMaze is usually large since the return success while
trajectories binomial clear margin on all domain, even on the challenging AntMaze, for which behavior cloning methods would fail and some form of policy improvement is essential.
failure). Our method outperforms all prior methods by a
Gym Tusks
BC
AWAC
Diffuser
MoRel
Onestep RI.
TD3+BC
CQL
IQL.
Diffusion-O0 QL
halfcheetab- m2
42.6
43.5
44.2
42.1
48.4
48.3
42.6
44.0
47.4
$1.1 0.5
hopper- medium-v2
52.9
57.0
58.5
95.4
59.6
59.3
67.6
58.5
66.3
90.5 4.6
walker2d medium-v2
halfcheetah: medium--eree hea-m---ellyaaaea
36.6
40.5
42.2
40.2
38.1
44.6
36.6
45.5
44.2
47.8 0.3
hopper-ne medium pepper-meiim-eeppaayy.
18.1
37.2
96.8
93.5
97.5
60.9
82.7
95.0
94.7
101.3 0.6
walker2d mendiiu-rrpply----
26.0
27.0
61.2
49.5
81.8
66.6
77.2
73.9
95.5
wwdimm-esperede..e. -2
hopper sugper-mndiium-xxxertt22
52.5
55.8
107.2
108.7
103.3
98.0
107.6
105.4
91.5
111.1 1: 1.3
w
107.5
74.5
108.4
95.6
113.0
110.1
108.1
108.8
109.6
110.1 0.3
Average
51.9
50.1
75.3
72.9
76.1
75.3
74.7
77.6
77.0
88.0
AntMaze Tasks
BC
AWAC
BCQ
BEAR
Onestep RI.
TD3+BC
DT
CQL
IQL.
Diffusion
aotonzzs-mmzeev0
54.6
56.7
78.9
73.0
64.3
78.6
59.3
74.0
87.5
3.4
sotmaaa-omazz-diveeee
45.6
49.3
55.0
61.0
60.7
71.4
53.0
84.0
62.2
66.2 8.6
-atmmaa-meddaum-paa---00
0.0
0.0
0.0
0.0
0.3
10.6
0.0
61.2
71.2
76.6 10.8
catmaao-mcdiun-diiree--00
78.6
aatmmaa-lagg-ppaay-v0
0.0
6.7
0.0
0.0
0.2
15.8
39.6
46.4 8.3
0.0
1.0
2.2
0.0
0.0
0.0
0.0
14.9
47.5
56.6 + 7.6
Average
16.7
18.0
23.8
23.7
20.9
27.3
18.7
50.6
63.0
69.6
Adroit Tusks
BC
SAC
BCQ
BEAR
BRAC-p
BRAC-
REM
CQL
1QL.
Diffusion-
pen
pen pen-cloned-v1
38.3
.0.8
44.0
26.5
1.6
-25
-10
27.2
37.3
57.3 + 11.9
Average
32.1
1.8
56.5
12.8
4.9
-1.0
2.2
31.2
$9.4
65.1
Kitchen Tasks
BC
SAC
BCQ
BEAR
BRAC-p
BRAC-V
AWR
CQL
IQL.
Diffusion-Qi QL
kinchen-somplete--i)
15.0
kitchen partial-a
33.8
0.0
18.9
13.1
0.0
0.0
15.4
49.8
46.3
60.5 6.9
kitchen mixed-xi
47.5
2.5
8.1
47.2
0.0
0.0
10.6
51.0
51.0
62.6 5.1
Averuge
38.4
5.8
11.7
20.1
0.0
0.0
48.2
53.3
69.0
the N increases, the model converges faster and the performance becomes more stable. In the fol- lowing D4RL tasks, we set moderate value, N = 5. to balance the performance and computational cost. With 5, the training time of our method is similar to that CQL (Kumar et al. 2020). The other hyperparameters, such as the learning rate and are provided in Appendix E
5.1 COMPARISON TO OTHER METHOOSS
We compare our Diffusion- QL with the baselines on four domains of tasks and report the results in We give the analysis based each specific domain.
Table I
Results for Gym Domain. We can see while most baselines already work well on the Gym tasks, Diffusion- QL can often further improve their performance by clear margin, especially in 'medium"
8

---

Published as a conference paper at ICLR 2023
Table 2: Ablation study. We conduct ablation study compare our diffusion model with CVAE models, and our Q-learning method with BCQ policy improvement.
Gym Tasks
BC-CVAE
BC-Diflusion
BOQ-CVAE
BCQ-Dilfusion
CVAE-QL
DiffusionQQ..
67.3
76.6 7.0
96.1 0.5.
94.6 1.0
70.3 5.5
97.2 i:
heppos-mssfoum-expettt
69.9 8.6
78.0 1 8.9
108.5 0.6
109,3 0:
109.2 1: 4.6
112.3 0.8
walker wulkkrr2dmmmdimmm-cxeett22222222
102.5 4.4
103.1 4.44
110.7 1: 0.2
114.0 0.5
50.5 38.2
111.2 0.9
Average
79,9
85.9
105.1
106.0
106.9
and 'medium-replay" tasks. Note the 'medium' datasets include the trajectories collected by an online SAC (Haarmoja et al., (0118) agent trained to approximately 1/3 the performance of the expert. Hence, the Tanh-Gaussian policy at that time is usually exploratory and not concentrated, which makes the collected data distribution hard learn. As shown Section the diffusion model has the expressivity to mimic the behavior policy even in complicated cases and then the policy improvement term will guide the policy to converge to the optimal actions in the subset of explored action space. These two components are key to the good empirical performance of Diffusion-QL.
Results for AntMaze Domain. The sparse reward and large amount of sub-optimal trajectories make the AntMaze tasks especially difficult. Strong and stable Q-learning is required to achieve good performance. For example, BC-based methods could easily fail on 'medium' and 'large' tasks. We show that the proposed Q-learning guidance added during the training of conditional diffusion model is stable and effective. Empirically, with a proper 1/, we find Diffusion-QL outperforms the prior methods by a clear margin, especially in harder tasks, such 'large-diverse.
Results for Adroit and Kitchen Domain. We find the Adroit domain needs strong policy reg- ularization to overcome the extrapolation error (Fujimoto et al., 2019) in offline RL, due to the narrowness of the human demonstrations. With smaller 1), Diffusion-QL. easily beats the other baselines by its reverse diffusion- based policy, which has high expressiveness and better policy reg- ularization. Moreover, long-term value optimization is required for the Kitchen tasks, and we find Diffusion QL also performs very well in this domain.
5.2 ABLATION STUDY
In this section, we analyze why Diffusion-QL. outperforms the other policy-constraint based meth- ods quantitatively on D4RL tasks. We conduct an ablation study on the two main components of Diffusion-QL: use of a diffusion model as an expressive policy and Q-learning guidance. For the policy part, we compare our diffusion model with the popular CVAE model for behavior cloning. For the Q-learning component, we compare with the BCQ approach policy improvement.
As shown in Table 2 BC-Diffusion model outperforms BC-CVAE, validating that our diffusion- based policy is more expressive and better at capturing the data distributions (as would expect from the results in l.. Under the same BCQ framework, which explicilly limits how far the ac- tion samples from polices could deviate from the cloned action samples, BCQ-Diffusion still works better than BCQ-CVAE. The hard and physical value constraints on the deviation by BCQ actually limits the policy improvement, as we can see that, Diffusion- -L further boosts the performance. Note Diffusion-QL applies the diffusion model learning itself as a soft policy regularization and guides the policy optimization via additive Q-learning. The weak expressiveness and poor cloning behavior of CVAE makes it fail when coupling with free Q-learning guidance, as shown by the re- sults of CVAE-QL. In a nutshell, the ablation study shows that the two components of Diffusion-QL. are working together to produce good performance.
CONCLUSION
In this work, we present Diffusion-QL. which is conditional bassed offline RL algorithm. First, its policy is built by the reverse chain of a conditional diffusion model, which allows for a highly expressive policy class and whose learning itself acts a strong policy regularization method. Second, Q-learning guidance through jointly learned Q-value function is injected in the learning of the diffusion policy, which guides the denoising sampling towards the optimal region in its exploration area. The two key components contribute to its state-of of the-art performance across all tasks in the D4RL benchmark.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
Z. Wang and M. Zhou acknowledge the support of NSF-IS 2212418 and the Texas Advanced Com- puting Center (TACC) for providing HPC resources that have contributed the research results reported within this paper.
REFERENCES
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning, pp. 104-114. PMLR, 2020.
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:221 I. 15657, 2022.
Christopher M Bishop. Mixture density networks. 1994.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: review for statisti- cians. Journal of the American statistical Association, 112(518):859-877 2017.
David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline RI without off- policy evaluation. Advances in Neural Information Processing Systems, 34:4933-4946, 2021.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34, 2021.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint 20X0::
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference Machine Learning, pp. 2052-2062. PMLR, 2019.
Wonjoon Goo and Scott Niekum. Know your boundaries: The necessity of explicit behavioral
cloning in offline rl. arXiv preprint arXiv:2206.00695, 2022.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. PMLR, 2018
Hado Hasselt. Double Q-learning. Advances in neural information processing systems, 23, 2010.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34, 2021.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.0999/ 2022.
Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):183-233 1999.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. Advances in neural information processing systems, 33: 21810-21823, 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.66 1412.6980 2014.
10

---

Published as a conference paper at ICLR 2023
Iya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In International Conference on Machine Learning, pp. 5774--7.7... PMLR, 2021a.
Iya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit Q- learning. arXiv preprint arXiv:2110.06169 2021b.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off off-po Q- learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 1991 2020.
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforce- ment learning, pp. 45-73. Springer, 2012.
Timothy P Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv: arXiv:/509.000
2015.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprint arXiv:2206.0003 00927, 2022.
Jiafei Lyu, Xiaoteng Ma, Xiu Li, Zongqing Lu. Mildly conservative Q-learning for offline reinforcement learning. arXiv preprint arXiv:2206. 220604200 2022.
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online rein- forcement learning with offline datasets. arXiv preprint 0e::20006.09359, 2020.
Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Ser- gio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, al. Imitating human behaviour with diffusion models. arXiv preprint arXiv:2301.10677 2023.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint 1rx00:1910000177,, 2019.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00:: 2022.
Nur Muhammad Mahi Shafiullah, Zichen Jeff Cui, Ariuntuya Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning & modes with one stone. arXiv preprint arXiv:2206. / 1251, 2022.
Jascha Gohl-Dicksteing Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. ArXiv, abs/1503.03588 2015.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28, 2015.
Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances Neural Information Processing Systems, pp. 11918-11938 2019.
Yang Song, Jascha Sohl- Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna tional Conference on Learning Representations, 2021. URL htttps:/oppenreview.. net/

---

Published as a conference paper at ICLR 2023
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Diffusion- GAN: Training gans with diffusion. arXiv preprint arXiv:2206.02262, 2022.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv: /91i:://11.. I.
1361, 2019.
Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. arXiv preprint arXhv:2112.07804, 2021.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems, 34:28954-28967, 2021.
Huangjie Zheng. Pengcheng He, Weizhu Chen, and Mingyuan Zhou Truncated diffusion proba- bilistic models and diffusion-basee adversarial auto-encoders, arXiv preprint arXiv:220..06671, 2022.
12

---

Published as a conference paper at ICLR 2023
Appendix
A MORE TOY EXPERIMENTS
Here we describe an additional toy experiment on bandit task. Actions are again in a real-valued 2D space, a E []. The offline data D lected by sampling actions equally from four Gaussian distributions with centers μ E {(-0.8, 0.8), (0.8, 0.8), -0.8), -0.8) and standard deviations σa (0.05, 0.05). as depicted the first panel of Figure We conduct the same experiments as the ones our main paper (see figure D and show the performance in Figure The only difference in this experiment is that the samples are now in the corners of the ation space.
([a=)]=m are col-
For behavior-cloning experiments, inal data distribution while the prior regularization methods fail in some way. For example, CVAE could only capture the two diagonal modes and place density between them, while MMD tends to align density around the boundaries because of its Tanh-Gaussian policy. For Q-learning experi- ments, we observe that the prior regularization methods typically push the policy converge to sub-optimal solutions in BCQ and BEAR while preventing the policy of TD3+BC from being con- centrated on the right corner. However, the policy of Diffusion- QL successfully converges to the optimal bottom corner. The ablation study experiments are consistent with our conclusion in the main paper.
observe that only our diffusion model could recover the orig-
BC-iffuusion
Add Reward
TD3
BCQ
BEARAMM
Diftusion-QL
N=2
Ne5
N=10
N==
N=50
Leammed Rewardd Function
Diffasion-UU/ee)))
AffuusionQQU/ii1100
Difiision Diffieion-LLIN-222 QLIN=20
Diffusion-UINN-SS
Figure 4: Bandit experiment with strong multi-modal behavior policy. The first row shows the comparison of behavior-cloming results between our method and prior methods. The second row shows the comparison results with Qelariigg involved. The third row shows the ablation study of N for BC-Diffusion The fourth row shows the learned reward function and the ablation study of N for Diffusion-QU
13

---

Published as a conference paper at ICLR 2023
B IMPLEMENTATION DETAILS
Diffusion policy. We build our policy as an MLP-based conditional diffusion model. Following the parameterization of Ho et al. (2020). the model itself is residual model, cø(a", s. i), where the i is the last timestep and is the state condition. We model Eg 3-layer MLPs with Mish activations and we use 256 hidden units for all networks. The input of en is the concatenation of the last step action vector, the current state vector, and the sinusoidal positional embedding of timestep i. The output of is the predicted residual at diffusion timestep
Q networks. We build two Q networks with the same MLP setting as our diffusion policy, which has 3-yayer MLPs with Mish activations and 256 hidden units for all networks.
We use the Adam (Kingma & Ba 2014) optimizer for the training both Diffusion policy and Q networks.
C EXPERIMENTAL DETAILS
We train our algorithm with 2000 epochs for Gym tasks and 1000 epochs for the other tasks, where each epoch consists of 1000 gradient steps. For the Gym locomotion tasks, average mean returns over 6 independently trained models and 10 trajectories per mode. For the other tasks, we average over 6 independently trained models and 100 evaluation trajectories.
We investigate two model selection methods, online and offline model selection. For online model selection, following the convention of supervised learning, the best-performed model is saved during training and used for final evaluation. For offline case, we select the model based on our lagging indicator La and more details are in Appendix
We use the original task rewards from MuJoCo Gym and Kitchen tasks. We standardize Adroit task rewards for training stability. We modify the rewards according the suggestion of CQL (Kumar et al. 2020 for the AntMaze datasets.
D OFFLINE MODEL SELECTION
For reducing the training cost and picking the best model during training without any interaction with the real environment, we provide a way to properly conduct early stopping for Diffusion- QL. Empirically, we found that La loss the behavior cloning loss of our diffusion model-based policy and measures how close the current policy to the behavior policy. For offline RL, we want La to be small to get rid of the OOD issue but we also don't expect La to be the smallest since our goal is policy learning while not behavior cloning. Hence, we monitor the La and save multiple model checkpoints during training. We stop the training whenever La increases in evaluation stage for early stopping. When the training stopped, we picked the checkpoint according to our metric, the second or third lowest La value. Note our model selection part is totally offline and only based on La without any access to the environment. The selection of the La is not very sensitive. Always using the 2nd smallest checkpoints doesn't 't impact the performance much. For example, with 2nd checkpoints selected, for the Gym domain, the average score is 87.6, and for the AntMaze domain, the average score is 1.
lagging indicator of online performance. Note La
E HYPERPARAMETERS
For Diffusion-QL. we consider three hyperparameters in total: learning rate, Q-learning weight n, and whether to use max Q backup from CQL (Kumar et al. 2020). For the learning rate of Adam, we consider values in the grid {1 x 10-3 3 3,3 x 10 x 10 " for the policy, while we use fixed learning rate, 3 x 10-4, for Q-networks. For 7), consider values according to the characteristics of different domains, as we mentioned in the description of datasets that the Adroit and Kitchen tasks require more policy regularization and the AntMaze tasks require more Q-learning. For max Q backup, we only apply it on the AntMaze tasks. Based on these considerations, we provide our hyperparameter setting in Table
14

---

Published as a conference paper at ICLR 2023
Table 3: Hyperparameter settings of all selected tasks.
Tasks
learning rate
Q backup
salfcheetal-mediiu---
3 x 10
1.0
False
hupper-medium--2
3 10
1.0
False
walker2d-medium-22
madicheetaah-meiim---a replayy
8 10
1.0
False
hopper-nediim--eplaayy22
3 10
1.0
False
wuller22-mmddiu-xxeppyy--22
3 x 10
1.0
False
8 x 10-4
hogger-nedduum-xxper-222
3 10
Falsc
wullerrd-mmeddum-ecpptt---
3 10
1.0
False
antmaze- uumazzo-umaz-vv
x
0.5
False
unimaze-umaze-dine diversee-v
3 10
2.0
True
antmaze medium medium-play--O play-vO
1x
True
30
3.0
True
cutmaze-large-play-e
3x
4,5
True
antmaze-large diveeress-
3 10
3.5
True
pen-human-vl
3 10-5
0.15
False
pen-cloned-vl
3 10
False
slichen-complete-xO)
3 10
0.005
False
kitchen kitchen-partial-ve partial-vD
3 3x110 10
0.005
False
kitchen-mixed
3 x 10
0.005
False
F. OPTIMAL RESULTS
If a small amount of online experience is provided during the evaluation stage for model selection, we can pick the best models during training via online evaluations (similar early stopping in supervised learning). This regime provides further boost in the performance of Diffusion- -L as shown in Table
G LIMITATIONS AND FUTURE WORK
Diffusion policies highly expressive and hence they can capture multi-modal distributions well. We have shown this results in learning better policies in offlineRL. However, at the inference time, the reverse sampling defined Equation ( requires iteratively computing eœ networks N times, and this can become a bottleneck for the running time. In our case, we couple the learning of diffusion policies with Q-learning, and achieve good performance with small N with N = 5 could be four to five times slower in action inference compared to previous one-step feedforward policies, and hence the inference cost could prevent the approach from deployment in some real-world scenarios, where fast response is necessary.
Diffusion policies
This motivates potential future works. There have been many recent works focusing on improving the sampling speed of diffusion models (Lu al. 2022 Wang et al.) 2022 Xiao et al., 2021 Sal- Jimans& & Ho. 2022 Song et al. 2020) [Zenng et al. (022). which could be applied to improve the sampling efficiency of Diffusion- QL. For example, the diffusion policy may be able to be distilled into a simpler feedforward policy after training.
H GAUSSIAN MIXTURE POLICY
test a Gaussian Mixture policy as a classic policy class that can capture multi-modal distributions as additional baseline. We modify TD3+BC (Fujimoto & Gu. 2021 by replacing the original deterministic actor with a mixture density network (Bishop) 1994), where each mixture component a Gaussian. Since Gaussian mixture policy is applied, TD3+BC) between predicted actions and real actions, with maximizing the likelihood estimate of Gaussian mixtures real state-action pairs. We keep all the other parts the same as TD3+BC.
replaced minimizing the L2 loss (from
15

---

Published as a conference paper at ICLR 2023
Table 4: Performance comparison with online model selection and offline model selection.
Gym Tasks
Diffusion-O QL (Offline)
Diffusion-QL. (Online)
halfchectah- halfchectab-modiiu--22
51.1 0.5
$1.5 0.3
hopper-medium-v2
90.5
96.6
walker2d-mcdium-v2
87.0
87.3 0.5
hallcheelah
malicheetab-mmccdu---reyyyy-
47.8 0.3
48.3 0,2
hopper-meciiumm-epppyyyy2
101.3 + 0.6
102.0 + 0.4
waller22-mmddiuu-zzelayy-22
95.5
98.0
hallchhetaa-mmediiu--eertrt---
96.8
97.2
hopper-mecduum--sseett.
111.1 = 1.3
112.3 + 0.8
wallecr2--mdduum-xxpettt--
110.1 = 0.3
111.2 + 0.9
Average
88.0
89.3
AntMaze Tasks
Diffusion-QU (Offlime)
Diffusion- QL (Online)
antmaze- umazz-v0
93.4 3.4
96.0 :.0033 3.3
mntmazzeuumazeeeiivneee--
66.2 8.6
84.0 + 10.1
anttmaze-medduum-layy-000
76.6 + 10.8
79.8 8.7
antmaze-mmdiunn-divee--00
78.6 + 10.3
82.0
untmaze-lange--llyy-v0
46.4 8.3
49.0
latmazz-lago--ivere--000
56.6 7.6
61.7 8.2
Average
69.6
75.4
Adroit Tasks
Diffusion- QL (Offline)
Diffusion-O QL (Online)
pen-human-vl
72.8
75.7
pen-cloned-vl
57.3
60.8
Average
65.1
68.3
Kitchen Tasks
Diffusion-O0 QL (Oftline)
Diffusion-QU (Online)
kiichen-complete--00
84.0 7.4
84.5 6,1
kitchen-partial-v0
60.5
63.7
kitchen-mixed-vv
62.6 5.1
66.6 3.3
Average
69.0
71.6
Table Performance comparison for Gaussian Mixture ablation study.
Gym Tasks
TD3+BC
TD3+BC-GM
Diffusion-QL
halfchherah-mmdiuu--xxeeeerr
90.4
49.0
96.8
hopper-medium-experr--2
98.0
40.0
111.1
walker2d-meeium-exppe
110.1
66.5
110.1
We evaluated the model (we call it TD3+BC-GM) both our bandit toy experiment and selected D4RL tasks. As shown Figure 5 with a properly selected number of mixtures, the Gaussian mixture could capture the multi-modis distribution in behavior cloning experiment. However, in the Q-learning experiment, it fails to converge to the optimal target location, and always places some density the suboptimal modes, resulting in suboptimal policy with multi-modes. For D4RL experiments, we set the number of mixtures to be 3, and seen in Table B we observed that TD3+BC-GM does not perform well on the three D4RL tasks, which consistent with our previous observation that TD3+BC-GM is prone to have suboptimal actions.
One reason for the poor performance could be that Gaussian mixture models are challenging fit well, particularly in higher dimensional spaces. We are also forced to choose the number of mixture components, which for the D4RL experiments we don't know a priori how many components there are. Moreover, each mixture component is often parameterized with diagonal Gaussian that has limited ability in capturing the dependence between different dimensions.
16

---

Published as a conference paper at ICLR 2023
Ground Truth
BD3+BC-GM (K-2)
TD3+ (K-3)
TD3+ BC-GM (K-4)
-39
Add Reward
TD3+BC-GM GM (K=2)
TD3+ GM (K=3)
TD3+BC GM (K=4)
7-MB 031
P-811 051
F-815
55%
150 30
35 C8
100
X30 000 130 35
130 500
430
350
00
35
921
13
Figure 5: Gaussian Mixture ablation study. The first row shows the behavior cloning experiment and the second row shows the Qelerning experiment, which are the same of experiments described in Section 4 Here, K the numher of Gaussian mixtures.
17


Reviews:
[{"review_id": "x-ELsanSkU", "reviewer": "Reviewer_JRTd", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a new offline RL algorithm where the distribution of actions (conditionally to the state) is represented by a diffusion model. At the same time, existing methods usually use a simple gaussian. The approach is then quite simple: the diffusion can be used with a classical behavioral cloning loss (diffusion BC), and coupled with a Q function when reward is available (Diffusion QL). When associated with a Q function, the authors use a strategy that is similar to the TD3+BC method proposed in previous work. The method is then evaluated both on a toy bandit environment to better understand the differences between this model and the existing ones, and also on standard offline-RL benchmarks (D4RL). The diffusion algorithm is obtaining very good results and is very competitive w.r.t state-of-the-art.", "summary_of_the_review": "This is a paper that is solid on both the idea and the execution of the idea. I do not see clear reasons to refuse such a paper.", "strength_and_weaknesses": "The paper is well-written and easy to follow. The proposed model is (very) simple -- which is a good point in my mind -- and I am sure that this paper is bringing interesting information to the community. If the idea is almost obvious (using the best-known and trendy generative model instead of simple gaussian models), I think that such a contribution provides important insights and researchers will be able to build on it. What the paper suggests is that the way we are used to capturing the distribution of actions (gaussian conditioned on the state) is not the right one since it is not able to capture multiple modes (while datasets, and particularly medium ones in D4RL are multi-modal since they have been built by using different acquisition policies). On that point, it seems particularly true for BC approaches where two experts may have taken two different actions in a similar state (and the average action would not be an expert action). What is less clear to me is that this also holds when the reward is available (particularly dense reward) where we can expect one action to 'dominate' all the other actions. For instance, in the toy problem (Figure 1), I don't understand why TD3+BC is not able to capture the same distribution as the one captured by Diffusion-QL, even using a simple gaussian. I would be interested in having more explanations on this point.\n\nOne question I also have is what would happen when using classical state-of-the-art methods with a mixture of Gaussians instead of a simple gaussian. Indeed, all the comparisons are made with a simple one-mode gaussian, and reimplementing a BC for instance with a mixture of 2 or 3 gaussians would give insight into the need to use diffusion models while simple approaches could work better. I would be very happy to see this kind of experiment in the ablation study for instance, but also in the toy environment which seems to be solvable through a mixture of Gaussians. Author could also compare their approach to 'implicit behavioral cloning' which proposes a solution to the problem of capturing multiple modes in BC.  \n\nMinor remark: figures are too small, and axis are not readable. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "tLfgm9ZQ__", "reviewer": "Reviewer_jsYy", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes using a conditional diffusion policy for offline RL. The proposed method consists of a diffusion model as a behavior cloning term for policy regularization, and a standard policy improvement term via a Q function. The authors demonstrated the effectiveness of their method on a series of toy examples and the standard D4RL, AntMaze, Adroit, and the kitchen benchmarks.", "summary_of_the_review": "I think this work is well-presented and adds an important contribution to the community, I recommend its acceptance at this venue.", "strength_and_weaknesses": "I liked this paper a lot, the authors presented a simple and straightforward method that yielded incredibly powerful results. Given recent advances in diffusion models, the idea of using a diffusion model for behavior cloning as a regularization term for offline RL seems like quite an obvious approach but is to the best of my knowledge quite novel. The paper itself is very well-written with a clear logical flow and was a joy to read. The authors also did a good job of placing their work in the context of other literature and comparing their work to other similar approaches. Experiments are very thorough and shows clear evidence of improvement compared to other baselines\n\nI especially liked the toy example introduced in Section 5 and the Appendix and how the authors separated the behavior cloning part and the policy improvement part. The example gives clear visualization on how the proposed method is advantageous in dealing with multimodal data compared to other approaches. Though one minor suggestion I have is that I think the toy examples perhaps deserve a separate section in the main text.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "YBKix-r_KPz", "reviewer": "Reviewer_Qge1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes to use a (conditional) diffusion model to represent a policy for the RL agent in the offline RL setting. The objective (3) consists of a Q-learning term and a behavior regularization term. Synthetic experiments on a bandit problem show the advantages of using diffusion models over other generative models, and experiments on common offline RL benchmarks show that the proposed diffusion-QL method can outperform alternative methods when tuned right.", "summary_of_the_review": "The paper proposes to use a diffusion model to parametrize the policy and claim that it can achieve better performance in the offline RL setting. However, the application is straightforward. It does not really show why such models are suitable for the task (the synthetic experiments are questionable) and important experiment details are missing. Hence the score.", "strength_and_weaknesses": "Strengths\n- Writing is clear\n- Good empirical performance\n\nWeakness\n- Limited novelty: the better performance mainly comes from a different parametrization of the policy borrowed from the generative model literature. It is a straightforward application.\n- Some of the comparisons in the synthetic experiment are questionable. It is surprising to see that CVAE and MMD methods compared in Sec.5 fail so miserably for such a simple task. It is very likely that these methods are not tuned properly for the bandit problem (e.g., it is not clear why CVAE produces a diamond shape policy) No enough implementation detail is provided here so it is challenging to see what may have been done incorrectly.\n- Some of the experiment details are missing and require further clarification. (a) The trade-off hyper-parameter is chosen on a dataset-to-dataset basis, which makes the applicability questionable. It is mentioned in Sec.6.1 and the appendix that eta (also see the discussion after Eq.(3)) is chosen differently for different benchmark problems. The actual values of eta are also not provided in the paper. (b) Page 7 mentioned that L_d is a “lagging indicator” and 2nd or 3rd checkpoints are chosen. This is not convincing at all and additional explanation is needed. As a result, the better performance becomes less convincing since it can be due to over-tuning.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "AbQP4oaNLA", "reviewer": "Reviewer_J2jU", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents Diffusion-QL, an offline RL method that can be seen as an actor-critic method, where the actor is a generative diffusion models, and the critic is a standard Q-function. The diffusion model is trained with supervised learning to imitate the data distribution. The main benefit of diffusion models is that they are able to accurately model the data distribution, even when it is multimodal, which is usually a challenge for previous methods. The Q-function is then used to guide the gradient during the training of the diffusion model towards actions with higher Q value. This combined approach prevents the agent from taking out-of-sample actions, while ensuring convergence towards the optimal actions among the explored regions of the state-action set. The authors introduce their approach as a policy regularisation method as it is constrained to a subset of the support of the data distribution.\n\nFirst, the method is illustrated on a toy bandit problem showing how it is the only one able to accurately model the data distribution and the only one to reliably converge towards the optimal action among the baselines. Finally, further experiments show Diffusion-QL achieves competitive or state of the art results on the D4RL benchmark.", "summary_of_the_review": "This is an interesting paper that tackles an important and difficult problem, offline-RL, in a novel, elegant and effective manner. I believe the paper presents a proof of concept that can be extended in several ways and will likely inspire further research.", "strength_and_weaknesses": " ### Strength\n- The paper presents a fresh idea for offline-RL that fixes multiple issues from previous approaches.\n- Identifying the main problem with previous policy regularisation approaches is the inaccurate approximation of the data distribution is insightful and leveraging diffusion models for directly modelling multimodal behaviour policies is an effective and sensible approach.\n- Experimental results on the D4RL benchmark are consistent and very promising.\n- The paper is well written and easy to read.\n\n### Weaknesses\n- Some important baselines are missing in the experiment sections that were designed to deal with multimodal data distributions, like the Behaviour Transformers (Shafiullah et al., 2022).\n- The paper barely discusses the limitations of the proposed approach, like the high computational cost of sampling actions, which might prevent the approach from deployment in some real-world scenarios.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster