Efficient Diffusion Policies For Offline Reinforcement Learning
Authors: Bingyi Kang
Xiao Ma
Chao Du
Tianyu Pang
Shuicheng YAN
Venue: NeurIPS 2023 poster (2023)
Keywords: Offline Reinforcement Learning
Diffusion Models
TL;DR: A faster, better and more general diffusion policy for offline reinforcement learning. Code is attached.

Abstract: Offline reinforcement learning (RL) aims to learn optimal policies from offline datasets, where the parameterization of policies is crucial but often overlooked. Recently, Diffsuion-QL significantly boosts the performance of offline RL by representing a policy with a diffusion model, whose success relies on a parametrized Markov Chain with hundreds of steps for sampling. However, Diffusion-QL suffers from two critical limitations. 1) It is computationally inefficient to forward and backward through the whole Markov chain during training. 2) It is incompatible with maximum likelihood-based RL algorithms (e.g., policy gradient methods) as the likelihood of diffusion models is intractable. Therefore, we propose efficient diffusion policy (EDP) to overcome these two challenges. EDP approximately constructs actions from corrupted ones at training to avoid running the sampling chain. We conduct extensive experiments on the D4RL benchmark. The results show that EDP can reduce the diffusion policy training time from 5 days to 5 hours on gym-locomotion tasks. Moreover, we show that EDP is compatible with various offline RL algorithms (TD3, CRR, and IQL) and achieves new state-of-the-art on D4RL by large margins over previous methods.

Efficient Diffusion Policies for Offline Reinforcement Learning
Bingyi Kang"
Xiao Ma"
Chao Du
Tianyu Pang
Shuicheng Yan
Sea AI Lab
[bingykang, baingykang..usuumaaaae yusufma555, 0aii com
Abstract
Offline reinforcement learning (RL) aims to learn optimal policies from offline datasets, where the parameterization of policies is crucial but often overlooked. Recently, Diftsuion-QLi [37] significantly boosts the performance of offline RL by representing a policy with diffusion model, whose success relies on parametrized Markov Chain with hundreds of steps for sampling. However, Diffusion-QLI suffers from two critical limitations. 1) It is computationally inefficient to forward and backward through the whole Markov chain during training. 2) It is incompatible with maximum likelihood-brat based RL algorithms (e.g., policy gradient methods) the likelihood of diffusion models intractable. Therefore, we propose efficient diffusion policy (EDP) overcome these challenges. EDP approximately constructs actions from corrupted ones at training to avoid running the sampling chain. We conduct extensive experiments the D4RL benchmark. The results show that EDP can reduce the diffusion policy training time from 5 days to 5 hours on gym-locomotion tasks. Moreover, show that EDP is compatible with various offline RL algorithms (TD3, CRR, and IQL) and achieves new state-of- the-art D4RL by large margins over previous methods. Our code available https///gttuub.comsiiii.
Introduction
Offline reinforcement learning (RL) is much desired in real-world applications as it can ex- tract knowledge from previous experiences, thus avoiding costly or risky online interactions. Ex-
Teaining Time
Nornalized Soure
Days
Pre SUTTA
71.3
73
tending online RL algorithms the offline do-
631
59.4
main faces the distributional shift [9] problem.
56.9
521.3
Existing methods mainly focus on addressing this issue by constraining policy to stay close the data-collecting policy [6. 39]. making con-
Hours
D-OL EDP
Lacomotion
Kitchon
Adroit
Antmaze
servative updates for Q-networks [17, 15, 40], or
Figure 1: Efficiency and Generality. D-QL is Diffusion-QL. Left: The training time on the loco- motion tasks in D4RL. Right: the performance of EDP and previous SOTA on each domain in D4RL. EDP trained with TD3 on locomotion and IQL on the other three domains. (Best viewed in color.)
combining these two strategies [21, 12]. How- ever, Offline RI. can also be viewed as a state- conditional generative modeling problem ac- tions, where the parameterization of the pol- icy network is important but largely overlooked. Most offline RL works follow the convention of parameterizing the policy as a diagonal Gaussian distribution with the learned mean and variance. This scheme might become inferior when the data distribution is complex, especially when offline
equal contribution
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

data are collected from various sources and present strong multiimoodlittis [32]. Therefore, more expressive models for the policy are strongly desired.
Recently, Diffusion-QL-1333 made a successful attempt by replacing the diagonal Gaussian policy with diffusion model, significantly boosting the performance of the TD3+BC [6] algorithm. Diffusion models [34, 10] have achieved the new state-of-the (SOTA) in image generation tasks [26, 5], demonstrating superior ability capture complex data distributions.
Despite the impressive improvement that Diffusion-QL has achieved, it has two critical drawbacks preventing it from practical applications. First, training diffusion policy with offline RL is computa- tionally inefficient. Consider a parameterized diffusion policy (a(aas), Diffusion LLL optimizes it by maximizing the Q value Q(s, ag) of a state given a volicy-generater action ag TT s). However, sampling from a diffusion model relies on a long parameterized Markov chain (e.g. 1,000 steps), whose forward inference and gradient backpropagation are unaffordably expensive. Second, diffusion policy is not a generic policy class as it is restricted to TD3-style algorithms. As computing the sample likelihood (a s) is intractable in diffusion models [35], diffusion policy is incompati- ble with a large family of policy gradient algorithms (e.g., V-Trace [22], AWR [28], IQL [15]), which require a tractable and differentiable log- likelihood log m for policy improvement.
In this work, we propose efficient diffusion policy (EDP) to address the above two limitations of diffusion policies. Specifically, we base EDP on the denoising diffusion probabilistic model (DDPM) [10], which learns a moise-predicctee network to predict the noise used to corrupt an example. In the forward diffusion process, corrupted sample follows predefined Gaussian distribution when the clean example and timestep are given. In turn, given a corrupted sample and predicted noise, we can approximate its clean version by leveraging the reparametrization trick. Based on this observation, to avoid the tedious sampling process, we propose action approximation to build an action from a corrupted one, which can be easily constructed from the dataset. In this way, each training step only needs to pass through the noise-prediction network once, thus substantially reducing the training time. As experimented, by simply adding action approximation, we obtain 2x speed-up without performance loss. Moreover, we apply DPM- IPM-Solver [20], a faster DDE-based sampler, to further accelerate both the training and sampling process. Finally, to support likelihood- we leverage the evidence lower bound for the likelihood developed in DDPM and approximate the policy likelihood from constructed Gaussian distribution with variance fixed and mean obtained from action approximation.
RL algorithms,
We evaluate the efficiency and generality of our method the popular D4RL benchmarking, as shown in Fig. 1. We first benchmark the efficiency of EDP on gym-locomotion tasks. By replacing the diffusion policy in Diffusion-QL with our EDP, the training time of Diffusion-QL is reduced substantially from five days to five hours (compared to their official code). Meanwhile, we observe slight to clear performance improvements on different tasks as the improved efficiency enables training DDPM with more timesteps than before. Moreover, we plug EDP into three different offline algorithms (including TD3+BC, CRR, and IQL), and the results justify its superiority over standard diagonal Gaussian policies. As a result, EDP set up new stte-offthhearr on all four domains in D4RL..
2 Related Work
Offline RL Distributional shift between the learned and behavior policies is offline RL's biggest challenge. Existing research mitigates this problem by making modifications to policy evaluation [24,
17, 15, 21] or policy improvement [9, 39, 6, 33, 38, 42, 41]. For example, conservative Q-learning (CQL) [17] penalizes dustof-distribution actions for having higher Q-values, proving that this is
equivalent to optimizing lower bound of Q-values. Onestep RL [1] conducts policy evaluation on in-distribution data to avoid querying unseen actions. IQL [15] introduces expectile regression [14] to approximate dynamic programming with the Bellman optimality function. TD3+BC explicitly constrains the learned policy by adding behavior cloning loss to mimic the behavior policy. Instead, CRR and AWR impose an implicit policy regularization by performing policy gradient-style policy updates. Despite their effectiveness, they ignore that the capacity of policy representation plays a vital role in fitting the data distribution. This paper instead focuses on an orthogonal aspect (i.e., policy parameterization) that all the above methods can benefit. Another line of work tries to cast offline RL as a sequence-to-sequeee -equencee model [3, 11], which is beyond the scope of this work.
Policy Parametrization Different RL algorithms may pose different requirements for parameteriz- ing a policy distribution. There are mainly two categories of requirements: 1) The sampling process
2

---

is differentiable. such as the deterministic policy in DDPG [18] and TD3 [7]. 2) The log-likelihood of samples is tractable. For example, policy gradient methods [30, 31, 38, 28] optimize a policy based on maximum likelihood estimation (MI.E.. Therefore, most works represent policy with a diagonal Gaussian distribution with mean and variance parameterized with a multi-layer perceptron (MIP). On the other hand, BCQ [9] and BEAR [16] choose to model policy with conditional variational autoencoder (CVAE). Recently, Diffusion QL [37] introduced diffusion models into offline RL and demonstrated that diffiusion models are superior at modeling complex action distributions than CVAE and diagonal Gaussian. However, it takes tens to hundreds more time to train a diffusion policy than diagonal Gaussian one. Moreover, diffusion policy only satisfies the first requirement, which means many other offline RL algorithms can not use it, including the current SOTA IQL.
Our method is motivated to solve the about two limitations. We first propose more efficient way to train diffusion policies, which reduces training time to the level of a Gaussian policy. Then, we generalize the diffusion policy to be compatible with MLE-based RL methods.
3
Preliminaries
3.1 Offline Reinforcement Learning
A decision-makern problem in reinforcement learning is usually represented by Markov De- cision Process (MDP): M = {S,A, R. 7). & and A are the state and action spaces re- spectively, P(s'(s, a) measures the transition probability from state 8 to state s' after tak- ing action a while R(s, a, s') gives the reward for the corresponding transition, Y E [0, 1) is the discount factor. A policy iTaas)? describes how an agent interacts with the environ- ment. The optimal policy T (a(a(s) is the one achieves maximal cumulative discounted returns: n² -arg max Er 12r= 7'r(s;, a.)]. Reinforcement learning algorithms frequently rely on the definition of value function V(s) = E, /E 7'r(s;, a = s], and action value (Q) function Q(s, a) = Ex I (', az)(s0 = 8, ag = a). which represents the expected cumulative dis- counted return of a policy given the initial state or state-action pair a).
In the offline RL setting, instead of learning from interactions with the environment, agents focus learning optimal policy from previously collected dataset of transitions: D {(s,, as, rt)). Offline RL algorithms for continuous control are usually based on an actor-critic framework that alternates between policy evaluation and policy improvement. Dur- ing policy evaluation, a parameterized Q network Qs,s, a) is optimized based on approxi- mate dynamic programming to minimize the following temporal difference (TD) error (m))
where Q3 a) denotes target net- )',, (r(s,a) + maxa Qg(s', aa) Q6(8,
work. Then the policy improvement step, knowledge in the Q network distilled into the policy network in various ways. Offline RL methods address the distributional shift [9] problem induced by the offline dataset D by either modifying the policy evaluation step to regularize Q learning or constraining the policy improvement directly. In the following, will show that our diffusion policy design is compatible with any offline algorithms and can speed up policy evaluation and improvement.
3.2 Diffusion Models
Consider a real data distribution ((x) and a sample x0 ~ ((x) drawn from it. The (forward) diffusson process fixed to Markov chain gradually adds Gaussian noise to the sample in K steps, producing sequence of noisy samples x .... x Note that we use superscript & to denote diffusion timestep to avoid conflicting with the RL timestep. The noise is controlled by a variance schedule "......,8 3k
q(00 -1) -No";; v1 3ha0 3 I),
((x ))
II k=1
(1)
When K - oo, K distributes as an isotropic Gaussian distribution. Diffusion models learn a conditional distribution po "x=")[=") and generate new samples by reversing the above process:
K
po (kkk = x(x*)
II
1x*) = -xxx* 10("" k), Eo(x", k)),
(2)
km
AA policy could either be deterministic stochastic, we use its stochastic form without loss of generality.
3

---

where p(**) = (N I) under the condition that IT k= (1 384) = O. The training is performed by
maximizing the evidence lower bound (ELBO):Ex, [loglolopss po(x))) 2 Egg log qi
pa( (
Efficient Diffusion Policy
In this section, detail the design of our efficient diffusion policy (EDP). First, we formulate RL policy with a diffusion model. Second, we present a novel algorithm that can train a diffusion policy efficiently, termed Reinforcement- Guiddddd Diffusion Policy Learning (RGDPL). Then, generalize the diffusion policy to work with arbitrary offline RL algorithms and compare our EDP with Diffusion- QL to highlight its superiority in efficiency and generality. Finally, we discuss several methods to sample from the diffusion policy during evaluation.
4.1 Diffusion Policy
Following [37], we use the reverse process of conditional diffusion model parametric policy:
K
To To(aas) pa (a =p(a*)
II
where a I). We choose to parameterize vo based on Denoising Diffusion Prob- abilistic Models (DDPM) [10], which sets Eo(a", k); B fixed time-dependent constants, and constructs the mean Ho from noise prediction model po(a", k; 8)
Van a -i co k; k: where at = g a* I and cg is parametric model.
To obtain an action from DDPM, we need draw samples from K different Gaussian distributions sequentially, as illustrated in Eqn. (2)-(3). The sampling process can be reformulated as
1
at
a ca k):s) + Bkc.
Vat
VI
with the reparametrization trick, where e ~ No0,, 1), k is the reverse timestep from K to 0.
Similar to DDPM, plugging in the conditional Gaussian distributions, the ELBO in Sec. 3.2 can be simplified to the following training objective (f)
||c-co (Va-a V1- che, k; $ s)1?
(5)
where & follows a uniform distribution over the discrete set {l..... "k....,K,, K}. It means the expectation taken over all diffusion steps from clean action to pure noise. Moreover, E (O, I), and (a", 8) E 'D are state-action pairs drawn from the offline dataset. Given a dataset, we can easily and efficiently train a diffusion policy in behavior-cloning manner backward through the network once each iteration. As shown in Diffusion- QL [37], diffusion policies greatly boost the performance when trained with TD3-based Q learning. However, it still faces two main drawbacks that limit its real-world application: 1) It is inefficient in sampling and training: 2) It is not generalizable to other strong offline reinforcement learning algorithms.
only need to forward and
4.2 Reinforcement--uuide Diffusion Policy Learning
To understand how parametric policy Ty is trained with offline RL algorithms, we start with typical Q-learning actor-critic framework for continuous control, which iterates between policy evaluation and policy improvement. Policy evaluation learns a Q network by minimizing the TD error LıD(ő):
((,,,'') ( "aoddaa'soss,dl Qi(s,ai) 27
Q0(s
(6)
where the next action a' ~ (T(-)")) The policy optimized to maximize the expected Q values
max Eg (.. (8,
0

---

It is straightforward to optimize this objective when a Gaussian policy is used, but things get much more difficult when a diffusion policy is considered due to its complicated sampling process. Instead, we propose to view the offline RI. problem from the perspective of generative modeling, where diffusion policy can be easily learned in a supervised manner from a given dataset. However, unlike in computer vision, where the training data are usually perfect, offline RL datasets often contain suboptimal state-action pairs. Suppose we have a well-trained Q network Qs. the question becomes how we can efficiently use Qo to guide diffusion policy training procedure. We now show that this can be achieved without sampling actions from diffusion policies.
Let's revisit the forward diffusion process in Eqn. 1. A notable property of it is that the distribution of noisy action a at any step l can be written in closed form: q(a*|a") = NN(a;; vā"a", aa, (1 (1-a*)r)
Using the reparametrization trick, we are able to connect a", a and c by:
a V v - a*e c~,
(8)
Recall that our diffusion policy is parameterized to predict with ca(a", 8). By relacing e with Es(a", k; s) and rearranging Eqn. (8), we obtain the approximiated action:
1 à a-=-1 a
vI a va
<0(a", k;;))
In this way, instead of rumning the reverse diffusion process to sample an action a", we can cheaply construct â from a state acttin pair (s, a) in the dataset by first corrupting the action a to a then performing one-step denoising to it. We will refer to this technique action approximation in the following. Accordingly, the policy improvement for diffusion policies modified as follows:
(+) a [Qo(s, a")]
(10)
To improve the efficiency of policy evaluation, we propose to replace the DDPM sampling in Eqn. (4) with DPM-Solver [20], which is an ODE-based sampler. The algorithm is defered to the appendix.
4.3 Generalization to Various RL algorithms
There are mainly two types of approaches realize the objective in Eqn. 7 for policy improvement.
Direct policy optimization: It maximizes Q values and directly backpropagate the gradients from Q
network to policy network, i.e., VoT. (0) aQ.(s.a) da This is only applicable to cases where 09 is tractable, e.g., when a deterministic policy a - me (s) is used or when the sampling process can be reparameterized. Sample algorithms belonging to this category include TD3 [7]. TD3+BC [6], and CQL [17]. One can easily verify that both the expensive DDPM sampling in Eqn. (4) and our efficient approximation in Eqn. (9) can be used for direct policy optimization.
da 86
Likelihood based policy optimization: It tries distill the knowledge from the Q network into the policy network indirectly by performing weighted regression or weighted maximum likelihood:
max E( (f(Q (s, a)) log (a)as))]
(11)
where f(Qu(s, a)) is a monotonically increasing function that assigns weight to each state-action pair in the dataset. This objective requires the log- likelihood of the policy to be tractable and differentiable. AWR [28], CRR [38], and IQL [15] fall into this category but each has a unique design in terms of the weighting function f. Since the likelihood of samples in Diffusion models is intractable, we propose the following two variants for realizing Eqn. 11.
First, instead of computing the likelihood, turn lower bound for log me (aa(s))) introduced in DDPM [10]. By discarding the constant term that does not depend on 0, we can have the objective:
8h [(Qa(s, a)) 2aa - 1x-1) ||e E, (at, k; s)]]
ER (x,e((.)))
(12)
Second, instead of directlying optimizing log To (a)s), we propose to replace it with an approximated policy a (a|s) & Aaa",, I), where â" is from Eqn. (9). Then, we get the following objective:
Ex,e((..)) [(Qo(s,a)) l|a à a²1]"]
(13)
5

---

Empirically, we find these two choices perform similarly, but the latter is easier to implement. So we will report results mainly based on the second realization. In our experiments. we consider two offline RI. algorithms under this category, i.e., CRR. and IQL.. They use two weighting schemes: SCRR = exp [Q a) (al)) (ss a')) /rCRR] and fio = exp [Qo(s, a) - V₂(()) /qq.). where refers to the temperature parameter and ( is an additional value network parameterized by v.. We defer the details of these two algorithms to Appendix A.
4.4 Comparison to Diffusion-QL
Now we are ready to compare our method and Diffusion-QL comprehensively. Though our EDP shares the same policy parametrization as Diffusion-QL. it differs from Diffusion-QL significantly in the training algorithm. As result, the computational efficiency and generality of diffusion policies have been improved substantially.
Efficiency The diffusion policy affects both policy evaluation (Eqn. (6)) and policy improvement (Eqn. (7)). First, calculating (o) in policy evaluation requires drawing the next action from it. Diffusion-QL. uses DDPM sampling while EDP employs a DPM-Solver, which can reduce the sampling steps from 1000 15, thus accelerating the training. Second, in policy improvement, Diffusion-QL again applies DDPM for sampling. Then, it calculates the loss function based on sampled actions and backpropagates through the sampling process for network update. This means it needs to forward and backward neural network for K times each training iteration. As a result, Diffusion QL can only work with small K, e.g., 5 100. In comparison, our training scheme only passes through the network once an iteration, no matter how big K is. This enables EDP to use larger K (1000 in our experiments) results in Tab. 1 also show a larger K can give better performance.
train diffusion policy on the more fine- ane-grained scale. The
Generality Diffusion-QL. can only work with direct policy optimization, which contains only small portion of algorithms. Moreover, thanks to their flexibility and high performance, the likelihood- based algorithms are preferred for some tasks g., Antmaze). Our method successfully makes diffusion trainable with any RL algorithm.
4.5 Controlled Sampling from Diffusion Policies
Traditionally, a continuous policy is represented with a state-conditional Gaussian distribution. During evaluation time, policy executes deterministically to reduce variance by outputting the distribution mean as an action. However, with diffusion policies, we can only randomly draw a sample from the underlying distribution without access to its statistics. As result, the sampling process noisy, and the evaluation is of high variance. We consider the following method to reduce variance.
Energy- -ased Action Selection (EAS) Recall that the goal of (offline) RL is learn policy that can maximize the cumulative return or values. Though the policy no is stochastic, the learned Qo provides deterministic critic for action evaluation. We can sample few actions randomly, then use Qø for selection among them to eliminate randomness. EAS first samples N actions from mo by using any samplers (i.e., DPM-Solver), then sample one of them with weights proportional to Q(s.a) This procedure can be understood as sampling from an improved policy p(a s) x e eQ(e,a) All results will be reported based EAS. See Appendix. C.4 for the other two methods.
(eeea))((())) (a((a) s)..
Experiments
We conduct extensive experiments on the D4RL benchmark [2] to verify the following assumptions: I) Our diffusion policy is much more efficient than the previous one regarding training and evaluation costs. 2) Our diffusion policy a generic policy class that can be learned through direct and likelihood-basse policy learning methods. We also provide various ablation studies on the critical components for better understanding.
Baselines We evaluate our method on four domains in D4RL, including Gym-locomotion, AntMaze, Adroit, and Kitchen. For each domain, we consider extensive baselines to provide thorough evaluation. The simplest method the classific behavior cloning (BC) baseline and 10% BC that performs behavior cloning the best 10% data. TD3+BC [6] combines off-policy reinforcement

---

Experimental Setup We keep the backbone network architecture the same for all tasks and algorithms, which is 3-layer MLP (hidden size 256) with Mish [23] activation function following Diffusion-QI. QL. [37]. For the noise prediction network ep ca(a", k; 8) diffusion policy, we first encode timestep with sinusoidal embedding [36], then concatenate with the noisy action at and the conditional state 8, We use the Adam [13] to optimize both diffusion policy and the Q networks. The models are trained for 2000 epochs on Gym-locomotion and 1000 epochs on the other three domains. Each epoch consists of 1000 iterations of policy updates with batch size 256. For DPM-Solver [20], we use the third-order version and set the model call steps to 15. We reimplement DQL strictly following the official PyTorch code [27] for fair comparison and we refer to DQL (JAX) for all sample efficiency comparisons. We defer the complete list of all hyperparameters the appendix learning algorithms with BC. OneStepRL [1] first conducts policy evaluation to obtain the Q-value of the behavior policy from the offline dataset, then use for policy improvement. AWAC [25], AWR [28], and CRR [38] improve policy improvement by adding advantage-based weights policy loss functions. CQI.(17) and IQL [15] constrain the policy evaluation process by making conservative Q updates or replacing the max operator with expectile regression. We also consider the Decision Transformer (DT) [4] baseline that maps offline RL as a sequence-to-sequence translation problem.
due to space limits. Throughout this paper, the results are reported by averaging 5 random seeds.
Evaluation Protocal We consider two evaluation metrics in this paper. First, online model selection (OMS), proposed by Diffusion-QU [37], selects the best- best-perfir performinf model throughout the whole training process. However, though OMS can reflect an algorithm's 's capacity, it is cheating, especially when the training procedure is volatile on some of the tasks. Therefore, we propose another metric to focus the training stability and quality, which is running average al training (RAT). RAT calculates the running average of evaluation performance for ten consecutive checkpoints during training and reports the last score as the final performance.
5.1 Efficiency and Reproducibility
Sperdup on Gym lacomotion
In this section, we focus on the training and evaluation efficiency of our efficient diffusion policy. We choose the OMS evaluation met-
DGL.
DQL((sx)
EDP w/o DPM
EDP w/o AP
EDP
ric to make fair comparison with the baseline
method Diffusion-QU [37]. We consider four
variants of EDP to understand how each com- ponent contributes to the high efficiency of our method. 1) EDP is the complete version of our method. It uses the action approximation tech- nique in Eqn. (9) for policy training and uses DPM-Solver for sampling. 2) EDP w/o DPM modifies EDP by replacing DPM-Solver with the original DDPM sampling method in Eqn. 4. 3) EDP w/o AP removes the action approxima- tion technique. 4) DQL (JAX) is our Jax implementation of Diffusion-QL.
Training
Sampling
Figure 2: Training and evaluation speed compari- son. The training IPS are: 4.66, 22.30, 50.94, 38.4, and 116.21. The sampling SPS are: 18.67, 123.70, 123.06, 411.0, 411.79.
We first benchmark and compare the training/evaluation speed of the above three variants and Diffusion-QL. We choose sullkerdd-mmiium-xxppettt algorithm for 10,000 iterations of policy updates and calculate the corresponding iterations-per-sccond (IPS). Similarly, we sample 10,000 transitions by interacting with the environment and calculate the corresponding peps-per--econd (SPS) for evaluation speed. Based the visualization in Fig. 2, by taking DQL (JAX) as the baseline, we are able to attribute the performance boost to specific techniques proposed. Specifically,
as the testbed. For training speed, we run each
observe that action approximation makes 2.3x training and 3.3x sampling faster, while using the DPM-Solver adds an additional 2.3x training speedup. We observe that DQL (JAX) is 5x faster than Diffusion-QLI QL., which means our Jax implementation is more computationally efficient than Diffusion-QU QL's PoTohc code. This demonstrates that both the action approximation technique and DPM-Solver play a critical role in making the training of diffusion policy efficient. However, this technique does not affect the sampling procedure; thus, EDP w/o DPM and DQL (JAX) are on par with each other regarding sampling speed.
7

---

Table 2: Average normalized parameterized with feed-forwad neural networks, while EDP is our efficient diffusion policy. Results of baselines taken directly from [15]. For the missing results of 10%, AWAC, and OneStep RL, we re-implement the baselines and report the results. Our results are reported with the RAT metric.
the D4RL benchmark. FF denotes the original policies
TD3-BC
CRR
IQL
Dataset
BC 10%BC
AWAC OneStep CQL
halfcheelah medium v2
42.6
42.5
42.6
43.5
48.4
44.0
48.3
52.1
44.0
49.2
47.4
48.1
hopper-medium--2
52.9
56.9
67.6
57.0
59.6
58.5
59.3
81.9
58.5
78.7
66.3
63.1
walkor2d-msdium-22
75.3
75.0
74.0
72.4
81.8
72.5
83.7
86.9
72.5
82.5
78.3
85.4
holfcheetah- 2epla
38.1
44.6
45.5
44.2
hopper-medle searr--eediiinpppee reply-y--yy-2222
18.1
75.9
82.7
37.2
97.5
95.0
00.9
101.0
95.0
99.0
94.77
99.1
26.0
62.5
66.6
27.0
49.5
77.2
81.8
94,9
77.2
63.3
73.9
84.0
hulfcheetah medium mediuum-expert-xy expert-v2
55.2
92.9
86.8
42.8
93.4
91.6
90.7
95.5
91.6
85.6
86.7
86.7
hopper medium expert-v2
52.5
110.9
107.6
55.8
103.3
105.4
98.0
97.4
105.4
92.9
91.5
99.6
molkkerdd-edium--ette average
107.5
109.0
108,1
74.5
113.0
108.8
110.1
110.2
108.8
110.1
109.6
109.0
51.9
74.0
74.7
50.1
76.1
77.6
75.3
85.5
77.6
78.3
77.0
79.9
kilchen-complete-u
65.0
7.2
39.3
57.0
43.8
2.2
61.3
43.8
73.9
62.5
75.5
kitchen partial vD
38.0
66.8
36.6
53.1
49.8
0.7
52.8
49.8
40.0
46.3
46.3
sitchen-mixed-xO)
51.5
50.9
22.0
47.6
51.0
60.8
51.0
46.
51.0
56.5
avcrage
51.5
41.6
32.6
52.6
48.2
1.0
58.4
18.2
53.3
53.3
59.4
pen-buman-vv
63.9
-2.0
15.6
71.8
37.5
5.9
18.2
37.5
70.2
71.5
72.7
pen con-cloned-vd
37.0
0.0
24.7
60.0
39.2
17.2
15.9
39.2
54.0
37.3
70.0
average
50.5
-1.0
20.2
65.9
38.4
11.6
32.
38.4
62.1
54.4
71.3
antmaze-umazo-v0
54.6
62.8
64.3
40.2
95.9
87.5
94.2
caunzz--mazz--ivveee
45.6
50.2
53.0
49.3
60.7
84.0
58.0
69.5
11.9
15.9
62.2
79.0
antinaze medium play-vO
0.0
5.4
0.0
0.0
0.3
61.2
0.2
0.0
0.0
33.5
71.2
81.8
antmaae medium diverse vO
0.0
9.8
0.7
0.0
53.7
6.4
32.7
70.0
82.3
antmaze-largee large-play-vO
0.0
0.0
0.0
0.0
0.0
15.8
0.0
1-6
0.0
26.0
39.6
42.3
antmazo-largen lanauz--argo-ivvess--00
6.0
0.0
17.5
60.6
average
16.7 22.4
18.7
18.0
20.9
50.6
16.4
29.8
7.0
43.8
63.0
73.4
To show that EDP does not hurt performance, we compare the normalized scores of all tasks in Tab. 1. The results for Diffusion-QI. are directly copied from the paper, where each task is carefully tuned with its own hyperparameters. Instead, DQL (JAX) and EDP use the same hyperparameters for tasks belonging the same domain. Moreover, since EDP is more computationally efficient, we can train better score model with large K 1000, which is larger than the one (5~100) used in Diffusion QL. Note that training diffusion policies with large K is impossible in Diffusion-Q1 as it needs to forward and backward the neural network K times. Finally, we can observe that EDP and DQL (JAX) are comparable with Diffusion-QU on gym-locomotion tasks but much better on the other three domains. Hence, efficient diffusion policy can boost sample efficiency and improve performance by enabling diffusion training in line-graine noise prediction.
The performance Diffusion-QU QL with efficient diffusion policy. The results for Diffusion- QL are directly quoted from [37]. All results are reported based on the OMS metric.
Table 1:
Diffusion QI.
DQL (IAX))X))
EDP
locomotion adroit
89.4
89.4
90.3
6
77.9
68.3
66.5
89.1
kitchen
80.5
5.2 Generality and Overall Results
This section aims evaluate whether EDP is a general policy class. To this end, we train it with direct policy optimization (TD3) and likelihood-hased policy optimization (CRR, IQL.) methods.
Then, compare them with their feed-forward counterparts and other baseline methods in Table 2. All scores for diffusion policies are reported using the RAT metric, while other scores are directly quoted from their paper. It shows that EDP can beat the standard Gaussian policy parameterized with an MLP on all domains and for all the three algorithms considered. On the gym-locomotion domain, EDP+ + TD3 gives the best performance (average score 85.5), while likelihood-basee policy learning methods are slightly worse. However, on the other three domains (Kitchen, Adroit, and Antmaze), EDP IQL beats all the other methods by a large margin (more than 10 average scores). Therefore, we conclude that EDP can serve as plug-in policy class for different RI. methods.
5.3 Ablation Study
8

---

Evaluation Metrics To reveal that evaluation metric is impor-
Walker2d
Hopper
Animaze
Normalizzd Return
1.0
0.8
tant, we train EDP with TD3 0. algorithms on three selected environ- ments: walker2d-medium-ene
6.6
64
0.4
0.4
02
Averago
hopper-medium--xxt -xxerttttuxxeee antiaaze-medimmddvveeeee
and
Best
0.0
0.0
We
Training Steps
then compare the scores for OMS (best) and RAT (average) by plotting the training curves in Fig. 3. On
Figure 3: Training curves for EDP +TD3 on three representative environments. Average represents RAT, Best represents OMS.
walker2d, the training is stable; thus, both OMS and RAT scores steadily grow and result in close final scores. A similar trend can be observed on the hopper but with a more significant gap between these two metrics. However, these two metrics diverge significantly when the training succeeds and then crashes on antmaze. Therefore, OMS is misleading and can not give a reliable evaluation of algorithms, which explains the necessity of using RAT in Sec. 5.2.
Energy-Based Action Selection
halfcheetah
hopper
walker2d
We notice that energy-based action selection (EAS) is a general method and can also be used for arbitrary policies. We apply EAS to normal TD3+BC and find no improvement, which shows EAS is only necessary for diffusion sampling. The results are deferred to the Appendix. Moreover, set the number of actions used in EAS from 1 to 200, and report the performance on gym-locomotions tasks in Fig. 4. It shows the normalized score monotonically grows as the number of actions increases on out of 9 tasks. In our main experiments, we set the number of actions to 10 by trading off the performance and computation efficiency.
Normalized pormallznn Retamn
1.0
1.10
0.9
0.
meduum
0.9
1.05
1.00
0.95
medium-ergert
0.7
0.90
0.85
04
10
100
10
100
0.80
10
100
Number Samples
Figure 4: Performance of different number of actions used in EAS.
The experiments are conducted on the nine locomotion tasks.
DPM-Solver We are using DPM-
balfcheetah
hopper
walker2d
Norpalized Roturn
1.0
Solver to speed up the sampling pro- cess of diffusion policies. The number of models in DPM-Solver is an impor- tant hyper- byper-parameter that affects sam- pling efficiency and quality. We vary this number from 3 to 30 and compare the performance on gym-locomotion tasks in Fig. 5. We can observe that the performance increases as more steps of model calls are used. The performance gradually plateaus after 15 model calls. There- fore, we use 15 in our main experiments.
0.9
1.0
0.8
0.8
mcdduum
0,
medtumm replay
0.6
8
export
0.4
6.4
20
25
Number of Model Calls
10
20
30
10
15
20
25
Figure 5: Performance of DPM-Solver with varying steps. The experiments are conducted on the nine locomotion tasks.
6 Conclusion
Diffusion policy has emerged as an expressive policy class for offline reinforcement learning. Despite its effectiveness, diffusion policy is limited by two drawbacks, hindering it from wider applications. First, training a diffusion policy requires to forward and backward through a long parameterized Markov chain, which is computationally expensive. Second, the diffusion policy is restricted policy class that can not work with likelihood-basee RL algorithms, which are preferred in many scenarios. We propose efficient diffusion policy (EDP) to address these limitations and make diffusion policies faster, better, and more general. EDP relies on an action approximation to construct actions from corrupted ones, thus avoiding running the Markov chain for action sampling at training. Our benchmarking shows that EDP achieves 25x speedup over Diffusion- QL at training time on the sym-locomotion tasks in D4RL. We conducted extensive experiments by training EDP with various offline RL algorithms, including TD3, CRR, and IQL, the results clearly justify the superiority of diffusion policies over Gaussian policies. As a result, EDP set new of-e-of--hhe-n on all four domains in D4RL.

---

References
[1] David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline without off-policy evaluation. Advances in Neural Information Processing Systems, 34:4933-4946, 2021.
[2] David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline rl without off-policy evaluation. Advances in Neural Information Processing Systems, 34:4933-44 4946, 2021.
[3] 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: 14:14484-55097,
2021.
[4] 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:150 10084-550997
2021.
[5] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794, 2021.
[6] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach offline reinforcement learning. Advances in neural information processing systems, 14-20132-20144 2021.
[7] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference machine learning, pages 1587-1596. PMLR, 2018
[8] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning, pages 1587-1596. PMLR, 2018
[9] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, 2019.
[10] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840 01680066511,., 2020.
[11] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34: 12:3- 1286, 2021.
[12] Bingyi Kang, Xiao Ma, Yirui Wang, Yang Yue, and Shuicheng Yan. Improving and benchmark- ing offline reinforcement learning algorithms. arXiv preprint arXiv:2306.00972. 2023.
[13] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 2014.
[14] Roger Koenker and Kevin F Hallock. Quantile regression. Journal of economic perspectives, ((()):1431566 2001.
[15] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning arXiv preprint arXiv:21 2021.
[16] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
[17] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative gelaarning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 179- 1191, 2020.
[18] 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: 02971, 2015.
[19] 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: htt.::119900999.. 2015.
10

---

[20] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilisticc model sampling in around 10 steps. arXiv preprint arXiv:2206.0008
2022.
[21] Xiao Ma, Bingyi Kang, Zhongwen Xu, Min Lin, and Shuicheng Yan. Mutual information regularized offline reinforcement learning. arXiv preprint arXiv:2210.07484, 2022.
[22] Michael Mathieu, Sherjil Ozair, Srivatsan Srinivasan, Caglar Gulcehre, Shangtong Zhang. Ray Jiang, Tom Le Paine, Konrad Zolna, Richard Powell, Julian Schrittwieser, et al. Starcraft ii unplugged: Large scale offiin reinforcement learning. In Deep RL Workshop NeurlPS 2021, 2021.
[23] Diganta Misra. Mish: A self regularized non-monotonic neural activation function. arXiv preprint arXiv: 08681, (2)::0-48550, 2019.
[24] Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. Advances in neural information processing systems, 2017.
[25] Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.0933: 2020.
[26]
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[27] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetioo, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019.
[28] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weigghed regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177 2019.
[29] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weiighted regresssoonoioo Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv: 1 2019.
[30] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pages 1889-1897. PMLR, 2015.
[31] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint /rXiv:1707.000 0077,06377,2 2017.
[32] Nur Muhammad Mahi Shafiullah, Zichen Jeff Cui, Ariuntuya Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning k modes with one stone. arXiv preprint arXiv:2206.1125) 2022.
[33] Noah Siegel, Jost Tobias Springenberg, Felix Berkenkamp. Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavior modelling priors for offline reinforcement learning. In International Conference on Learning Representations, 2020.
[34] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. International Conference on Machine Learning, pages 2256-2266 2265. PMLR, 2015.
[35] Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in Neural Information Processing Systems, 34:1415- 1428, 2021.
[36] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Fukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[37] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusson policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06199 061999908.06........ 2022.
learning
[38] Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. Advances in Neural Information Processing Systems, 33:7768-7778, 2020.
11

---

[39] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11380 2019.
[40] 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, 01-28954-28966 2021.
[41] Yang Yue, Bingyi Kang, Xiao Ma. Gao Huang, Shiji Song, and Shuicheng Yan. Offline prioritized experience replay. arXiv preprint arXiv:2306.05412 2023.
[42] Yang Yue, Bingyi Kang, Xiao Ma, Zhongwen Xu, Gao Huang, and Shuicheng Yan, Boosting offline reinforcement learning via data rebalancing, arXiv preprint arXiv:2210.092441 2022.
12

---

Broader Impact
In this paper, we propose an efficient yet powerful policy class for offline reinforcement learning. We show that this method is superior to most existing methods on simulated robotic tasks. However, some war robots or weapon robots might employ our EDP to learn strategic agents considering the generalization ability of EDP. Depending the specific application scenarios, might be harmful to domestic privacy and safety.
A Reinforcement Learning Algorithms
In this section, we introduce the details of the RI. algorithms experimented.
A.1
TD3+BC
TD3 [8] is a popular off-po RL algorithm for continuous control. TD3 improves DDPG [19] by addressing the value overestimation issue. Specifically, TD3 adopts double Q-learning paradigm that computes the TD-target Qss a)
(s. = r(s. a) + ( a'), Q2(s',a')), a'= x=s')s),
(14)
where x (s'') is deterministic policy, Q1 and Q2 are two independent value networks. Specifically, TD3 takes only Qı for policy improvement
arg max E [Q (8, a)), à (8).
(15)
Built on top TD3, TD3+BC simply adds an additional behavior cloning term for its policy improvement
arg max Eaa-D [Q(s, à) a(a a) à ===(s),
(16)
where a is hhpee--pppae parameer that balances these two terms. However, as the scale of Q-values are different from the behavior cloning loss, TD3+BC normalizes it with & Σ [Q(s;, a,)l for numerical stability.
In the context EDP, we observe that 1) the behavior cloning term can be naturally achieved as diffusion loss as defined in Eqn. 5; 2) the sampled action à is replaced by action approximation as in Eqn. 9 for efficient policy improvement. Different from the original TD3+BC that uses only Q1 for policy improvement, we sample from Q1 and Q2 with equal probability for each policy improvement step.
A.2 Critic Regularized Regression
CRR follows Advantage Weighted Regression (AWR) [29], which is simple yet effective off-policy RL method, for policy improvement. Specifically, AWR considers a constrained policy improvement step
arg max (du()) / n(a a))aads s.A. J. [.dulo)Dai
ml 8)] e, (17)
where A(s, a) is the advantage function, μ is behavior policy that is used to generate trajectories during off-policy RL, (s) is the state distribution induced by μ4, DRI. is the KL divergence, and e is threshold parameter. This constrained optimization problem can solved in closed form, which gives an optimal policy of
1 Z(s) uuaa |s) exp (=)(s,a))
aa 8)
(18)
with (s)) being the partition function and a hyper-parameter. For policy improvement, AWR simply distills the one-step improved optimal policy to the learning policy æ by minimizing the
13

---

KL-divergence
arg min (g) (a) (RL (+ (.1 () |l x(- [)]]]
log T(a 8) exxpp( 1 A(s, A a))
arg max
(19)
CRR performs policy improvement the same way as AWR, which simply replaces the sampling distribution with fixed dataset
arg max aaDD log ((a (a exp( A(s, a))
(20)
As result, this naturally imposes implicit constraint its policy improvement step.
However, as computing the log-likelihood log (a 8) intractable in diffusion models, practical implementations, we use Eqn. 13 instead. In addition, compute the advantage by
N
Σ A(s, a) = min(Q1 a) Q2(s, N 2=1
min(Q1(s, az), Q2(s,
(21)
where å; ~ Naaa,, o)) is sampled action with the mean of approximated action and an additional standard deviation. In our experiment, we found using fixed standard deviation, an identity matrix of o I, 3 1, and sample size N = 10 generally produce a good performance.
A.3 Implicit Q Learning
Similar to CRR, IQL also adopts the AWR-style policy improvement that naturally imposes constraint to encourage the learning policy to stay close to the behavior policy that generates the dataset. Different from CRR query novel and potentially out-of-distribution (OOD) actions when computing advantages, IQL aims to completely stay in-distribution with only dataset actions, while maintaining the ability to perform effecttve multi-step dynamic programming during policy evaluation. IQL achieves this by introducing an additional value function V(s) and performing expectile regression. Specifically, the policy evaluation in IQL implemented
min F. (Q (s) [s))
min L. ((r(s. - +V ((s a1
(22)
where I, the expectile regression loss defined as (x) - 1(x < 0)]s² with hyper-paramter E (0, 1). The intuition behind IQL is that with larger we will be able to better approximate the max operator. As a result, IQL. approximates the Bellman's optimality equation without querying OOD actions.
In practical implementations, we also adopt double learning for IQL, where we replace the Qs,, a) in Eqn. 22 with min(Q1 a), Q2(8, a)), and then use the updated value network to train both Q value networks. For IQL, we follow the policy improvement steps of CRR, as de- scribed in Eqn. 20 and Eqn. 21. The key difference is that instead of sampling actions to compute
N
N min(Qi a.), Q2(s, a.)), IQL replaces it directly with the learned (ss). As for hyper-
parameters, we use temperature = 1, a fixed standard deviation a I, for expectile ratio we use 0.9 for antmaze-v0 environments and 0.7 for other tasks.
B Reinforcement Guided Diffusion Policy Details
The overall algorithm for our Reinforcement Guided Diffusion Policy Learning is given in Alg. 1. The detailed algorithm for energy-based action selection is given in Alg.
C Environmental Details
C.1 ParaaeParrmeters
We detail our hyperparameters in Tab. 3.
14

---

Algorithm 1: Reinforcement Guided Diffusion Policy
Input: 7, 1) D. co(a", k; 8), (8)
Output: 0,6
for !,...,d / do
// Sample batch of data
{(s. a, s') s D
1/ Sample next actions with DPM-Solven
a' (T(+(s)
④ + Φ - V (VLDD())
11 Action approximation
Val a Vi-at
à
Van E0(a", k;
0 Vo (Lam (0)
Algorithm 2: Energy-based Action Selection
Input: number of actions: N, To (a) Qo(s, a) Output:
a; ~ no (a i i= l,.,..,,
a = categorical__cmmlee@aa ( (elo(and)))
C.2 More Results
We first expand Tab. 1 by providing detailed numbers for each of the tasks used in Tab. 4.
We report the performance of EDP trained with TD3, CRR, and IQL in Tab. 5, where we directly compare the scores of different evaluation metrics, i.e., OMS and RAT. We can observe that there are huge gaps between OMS and RAT for all domains and all algorithms. However, IQL and CRR are relatively more stable than T33. For example, on the antmaze domain, TD3 achieves a best score of 80.5, while the average score is just 29.8. In comparison, the best and average scores of IQL are 89.2 and 73.4, respectively.
C.3 More Results on EAS
We compare normal TD3+BC, TD3+BC with EAS for evaluation and TD3 EDP Tab. 6.
C.4 More Experiments Controlled Sampling
As in Sec. 4.5, discussed reducing variance with EAS. We now detail another two methods experimented as below.
Policy scaling Instead of sampling from the policy directly, we can sample from a sharper policy Tī ())s), where t > 1. The scaled policy zī shares the same distribution modes as mh but with reduced variance. Since diffusion policies are modeling the scores log no (a,)s) policy scaling can be easily achieved by scaling the output experiments on the gym-locomotion tasks in D4RL by varying 7 from 0.5 to 2.0, as shown in Fig. 6, the results show the best performance is achieved when T = 1.0. It means sampling from a scaled policy does not work.
noise-prediction network by the factor т. We conduct
Deterministic Sampling This method is based on the observation that the sampling process of the DPM-Solver is deterministic except for the first step. The first step is to sample from an isotropic Gaussian distribution. We modify it to use the mean, thus avoiding stochasticity. Consider initial noise a K (0, I). we rescale a K by noise scale factor. We show how this factor affects the final performance by varying from 0.0 to 1.0. As illustrated Fig. 7, the best performance achieved at zero noise in most cases, and a normal a performs worst. This means reducing the variance of the initial noise is able improve the performance of diffusion policy. However, the best performance achieved in this way still falls behind EAS.
15

---

Table Hyperparameters used by EDP.
Task
Learning rate
Gradient norm clipping
Loss weight
Fpochs
Batch size
Locomation
0.0003
1.0
2000
256
Antmaze
0.0003
1000
256
Adroit
0,000033
0.1
1000
256
Kitchen
0.0003
0.005
1000
256
Table The performance of Diffusion- QL with efficient diffusion policy. The results for Diffusion- QL are directly quoted from [37]. EDP is our method. DQL (JAX) is a variant that uses the exact same configurations Diffusion QL. All results are reported based on the OMS metric.
Dataset
Diffusion-QI.
DQL (JAX)
EDP
hallchettah-medimm-t medium x2
halfcheetah
51.5
52.3
52.8
hopper-medium--?
96.6
95.3
walker2d-mediun--2
87.3
86.9
89.6
halfcheetah medium- meteaeee-liiimmm-pelaaa22
48.3
50.3
50.4
hoperr-eddiun-cplge
102.0
101.8
102.7
walker2d-me mriiiddmmmdummoppa-----
98.0
96.3
97.7
hulfcheetah medium expert-v2
97.2
hopper medium expert-v2
112.3
113.1
112.0
111.2
111.5
112.0
average
89.4
89.4
90.3
umazeevv
96.0
antmaze umuze divense v0
84.0
74.0
66.0
antmaaze wwd-imeddi-mm---- play-vO
79.8
96.0
82.0
82.0
96.0
antmaae- wwwgllllggg-y---eeeeeee play-v0 aatmano-lage-divene--t average
49.0
66.0
60.0
61.7
60.0
64.0
75.4
78.6
77.9
pen-humaa--vl
76.7
74.0
98.3
pen-cloned-vl
60.8
59.1
79.9
average
68.3
66.5
89.1
Aucchen-prtial-Vi
63.7
73.0
71.5
kitchen Muichen-mixed-ae
66.6
76.0
73.0
average
71.6
83.0
80.5
C.5 Computational Cost Comparison between EDP and Feed-Forward Policy networks
benchmark the training speed of TD3+BC with EDP on walker2d- -eerr-2d-mdduumm-eprt--2222 by training each agent for 10,000 iterations of policy updates. The training speed for TD3+BC is 689 iterations- per-second (IPS), while EDP is 4 412 IPS. other words, it takes around 3 hours to train an agent with the feed-for forward network, while EDP needs around 5 hours. We also conducted experiments by double the network used in TD3+BC; unfortunately, there was no performance gain on the locomotion tasks (75.3 75.6). Moreover, both feed-forward policy and diffusion policy utilize 3-layer MLP (hidden size 256) as the backbone network. Therefore, the network capacity should be comparable.
halfcheetah
hopper
walker2d
Normalized Return
0.8
0.6
0.6
0.4
medium
0.2
0.2
replay
0.0
expert
0.0
0.0
1.0
Figure 6: Performance of EDP + TD3 on gym-locomotion tasks with varying
16

---

Table 5: Average normalized score on the D4RL benchmark EDP trained with different algorithms. "Best" represents the online model selection metric, while "Average" is our runing average at training metric.
EDP + TDS
EDP CRR
EDP + IQL
Average
Best
Average
Best
Average
Best
balfcheetah medium
52.1
52.8
49.2
50.2
48.1
48.7
hopper-medium--2 2alke22d-mmdiun-12
81.9
98.6
78.7
95.0
63.
97.3
86.9
$2.5
85.8
85.4
88.7
49.4
50.4
43.5
47.8
43.8
45.5
hopper medium replay
102.7
101.7
99.1
walker 2d medium replay-v2
94.9
97.7
63.3
89.8
84.0
93.4
hulfcheetah medium expert-v2
95.3
97.1
85.6
93.3
86.7
80.9
hopper mediam- expert-v2
97.4
112.0
92.9
109.4
99.6
95.7
dalker2d-mddimmoxpprryyy
110.2
112.0
110.1
112.3
109.0
111.5
average
87.3
79.9
kilchen-complete-0)
61.3
93.4
73.9
95.8
75.5
95.0
kitchen partial
52.8
66.0
40.0
56.7
46.3
72.5
kitchen-mixed-tt
60.8
88.0
46.1
59.2
56.5
70.0
average
58.1
53.3
70.6
59.4
79.2
pen-human-vD)
48.2
60.0
70.2
127.8
72.7
130.3
clo-cllnedxd x0)
15.9
64.0
54.0
106.0
70.0
138.2
average
32.1
77.9
62.1
116.9
71.3
134.3
animaze umnzze v0
96.6
98.3
95.9
98.0
94.2
98.0
diverss--o
69.5
79.9
15.9
80.0
79.0
90.0
antiaase medium play
0.0
89.1
33.5
82.0
81.8
89.0
antmaze mwdimm--dmmmm- diverse vO
6.4
97.0
32.7
72.0
82.3
88.0
playe--
1.6
71.5
26.0
57.0
42.3
52.0
samazu-largeddvveseee average
4.1
73.0
58.5
71.0
60.6
68.0
76.7
73.4
Table 6: Energy-based Action Selection + Normal TD3
Dataset
TD3+Diff
TD3+FAS
TD3
halfcheettal-mecrun-22
52.1
48.7
47.8
hopper-mediumm-
81.9
30.8
54.0
walker2d medium v2
86.9
T7.7
53.4
halfcheetah, medium- replay-v2
49.4
44.4
44.1
hupper medium--rpplae
walker2d medium replay-v2 v2
94.9
74.8
71.6
haffeheetah-medine 2xperuu--xee--222222
95.5
85.9
copper-mdiunn-xppet-222
97.1
71.9
95.1
110.2
108.4
110.0
gym locumotion v2 (avg)
85.5
69.2
69.7
C.6 The effect of action approximation
We compare with and without action approximation on the following three environments by using OMS metric (Tab. 4). In Tab. 7, the DDPM column will forward and backward a policy network 100 times at training time, but action approximation only needs once. We can observe that action approximation will slightly harm the performance, when the same number of diffusion steps is used. However, supports training diffusion policies with larger K (e.g. 1000), while Diffusion- -LL does
halfcheetah
hopper
walker2d
1.1
1.10
6.9
Rentrrn
1.0
1.05
0.8
0.9
medium replay
medium
0.55
medium export
0.6
0.85
0.5
0.4
0.00
0.75
1.00
0.75
Figure 7: Performance of EDP + TD3 on gym-locomotion tasks with varying initial noise scale.
17

---

Table 7: The effect of action approximation.
OMS K=100
DDPM
Action Approx
walker2d-mediumm-22
86.9
85.5
walker2d- medium-reely--??
96.3
93.3
malker2d-medium medium-expert-v2 expert-v2
111.5
111.1
not. Increasing K is able to avoid performance drop as evidenced by the last column of Tab. 2 and Tab.
D Negative Societal Impacts
In this paper, propose efficient yet powerful policy class for offline reinforcement learning. We show that this method is superior to most existing methods on simulated robotic tasks. However, some robots weapon robots might employ our EDP to learn strategic agents considering the generalization ability of EDP. Depending on the specific application scenarios, might be harmful to domestic privacy and safety.
18


Reviews:
[{"review_id": "Mv6ck2YvMB", "reviewer": "Reviewer_tRDF", "summary": "This paper proposed EDP to address the existing limitation of diffusion policy in offline RL. EDP relies on an action approximation to construct actions from corrupted ones, thus avoiding running the Markov chain for action sampling at training. shows that EDP achieves 25+ speedup over Diffusion-QL at training time on the gym-locomotion tasks in D4RL. Extensive experiments by training EDP with various offline RL algorithms, including TD3, CRR, and IQL justify the superiority of the proposed method.\n", "questions": "* It seems there exist mismatch of results between pre-sota  methods from Figure 1 and Table 2?\n* Line 227 points out that Tab 1 reflects the benefits of increasing K. However, it seems Tab 1 doesn't include this information?\n* The GPU memory needed by Diffusion-QL is quite large. I'm curious whether EDP has the same problem?\n", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* The proposed method EDP focuses on improving the Diffusion-QL baseline from two aspects including computational cost and limited applicability to different RL algorithms, which is the major contribution of this paper. It provides a new way to apply a complicated generative model in RL.\n* Extensive evaluation of datasets of different domains is performed. EDP can outperform several widely-used algorithms by replacing the originally-used Gaussian policy with the refined diffusion policy. And both the training and sampling speed of EDP are way faster than that of Diffusion-QL. Besides, it seems EDP is easier to tune compared with Diffusion-QL as it uses the same hyperparameters for tasks belonging to the same domain and could achieve satisfactory performance.\n* The unreasonableness of evaluation protocol OMS is pointed out and RAT is adopted to justify the performance of the proposed method and baselines.\n", "weaknesses": "* The compared baselines  in Table 2 and Figure 1 are a little bit weak, which cannot demonstrate the superiority of EDP.  As the authors claimed that the proposed method outperforms pre-sota methods, more powerful model-free baselines should be included such as EDAC[1], RORL[2] on Locomotion tasks, X-QL[3], InAC[4], BPPO[5] on other domains. Otherwise, I believed the superiority of the proposed method is overclaimed to some extent.\n\n  [1] An, Gaon, et al. \"Uncertainty-based offline reinforcement learning with diversified q-ensemble.\" *Advances in neural information processing systems* 34 (2021): 7436-7447.\n\n  [2] Yang, Rui, et al. \"Rorl: Robust offline reinforcement learning via conservative smoothing.\" *Advances in Neural Information Processing Systems* 35 (2022): 23851-23866.\n\n  [3] Garg, Divyansh, et al. \"Extreme Q-Learning: MaxEnt RL without Entropy.\" *The Eleventh International Conference on Learning Representations*. 2022.\n\n  [4] Xiao, Chenjun, et al. \"The In-Sample Softmax for Offline Reinforcement Learning.\" *The Eleventh International Conference on Learning Representations*. 2022.\n\n  [5] Zhuang, Zifeng, et al. \"Behavior Proximal Policy Optimization.\" *The Eleventh International Conference on Learning Representations*. 2022.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Q1. The baselines are a little bit weak. The superiority might be overclaimed.\n\nA1. Thank you for reminding us of these important literatures.  We’d like to clarify that our paper focuses on policy representation in offline RL, which is orthogonal to the algorithmic developments by these related works. It means that our EDP can be further integrated into their methods for better policy modeling. Moreover, we have tried our best to make a fair comparison with these methods during rebuttal period. Unfortunately, we failed to do so due to the mismatch of evaluation protocols, network architectures and hyper-parameters. Taking BPPO as an example, in their Figure 3 (c), the running average score for  hopper-medium-expert-v2 is around 90-100, while in Table 1, the OMS score is reported as 112.8. \n\nIn light of the aforementioned challenges and your insightful suggestion, we are happy to adjust our claim to make it more precise. We will restrict the sota claim to the policy representation scope. \n\n> Q2. Mismatch of pre-sota methods from Figure 1 and Table 2. \n\nA2. Thank you so much for pointing out this typo. We will update Figure 1 to make the pre-sota of Adroit be 65.9 instead of the current 56.9.\n\n> Q3. The benefits of increasing K\n\nA3. Sorry for the confusion.  We thank you for the valuable question and will definitely improve this in our revision. In table 1, Diffusion-QL are directly copied from the original paper, which is obtained by setting K to 5. DQL(JAX) is our reimplementation of Diffusion-QL but with K equals to 100. EDP is our final version whose training speed is constant to the number of K (thanks to action approximation). Therefore, EDP is using 1000 for K. On all 4 domains, EDP surpasses Diffusion-QL, which means larger K brings better performance. \n\n> Q4. The GPU memory needed by Diffusion-QL is quite large, how about EDP?\n\nA4. EDP does not suffer from this problem. The reason Diffusion-QL needs large GPU memory is that it needs to forward and backward the policy network K times to sample an action at each training iteration. However, EDP avoids this problem by introducing the action approximation technique, which only forwards and backwards the network once, thus greatly saving GPU memory usage and training time. \n\n"}, {"review_id": "uUxNTpPvH0", "reviewer": "Reviewer_TjdE", "summary": "The focus of this paper is to enhance the diffusion policies introduced in Diffusion-QL for offline reinforcement learning. The authors address the challenges of training and sampling efficiency by incorporating action approximation and employing an advanced ODE solver for diffusion policies. They conducted extensive experiments to demonstrate the effectiveness of their proposed approach, known as Efficient Diffusion Policy (EDP).", "questions": "See weakness. ", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThis paper is well structured and easy to follow.  \n2.\tThe conducted experiments are of enough quantity and quality. \n\n", "weaknesses": "> I think some claims need to be more careful or accurate. \n\n“… reduce the diffusion policy training time from 5 days …”\n“However, it takes tens to hundreds more time to train a diffusion policy than a diagonal Gaussian one”\t\n\nI tried Diffusion-QL before and I don’t think it takes about 5 days for training in their default setting (K=5). In my case, it is only about 20 hours, so I am not sure what setting is inferred here to be 5 days. \n\nMoreover, “In comparison, our training scheme only passes through the network once an iteration, no matter how big K is.” This statement is in correct. According to Algorithm 1 and the description in the paper, “To improve the efficiency of policy evaluation, we propose to replace the DDPM sampling in Eqn. (4) with DPM-Solver”, the policy evaluation part still needs the action samples with the full reverse process, which takes a number of steps even with DPM-Solver. \n\n> The novelty is kind of limited. In short words, EDP tried two things: 1. replace the real action samples generated by the full reverse process in policy improvement step with one-time real action prediction. 2. replace the original DDPM solver with DPM-solver. \n\n> The Energy-based Action Selection (EAS) is not a new thing and has been studied by other works, such as “Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling (ICLR 2023)”.\n\n> Following diffusion model conventions, I use $x$ to represent $a$ here. This paper uses $\\hat{x_{0, t}}$ to replace $x_0$ in policy improvement step. Note using $\\hat{x_{0, t}}$ to replace $x_0$ is not a common thing in diffusion models, since the $\\hat{x_{0, t}}$ has different meaning compared to $x_0$. $\\hat{x_{0, t}}$ usually represents the mean of all $\\hat{x_0}$ that could goes back to the true $x_0$ distribution from time step $t$. Hence, $\\hat{x_{0, t}}$ could be with much noise and averaged when $t$ is large. Not sure why here the $\\hat{x_{0, t}}$ could replace the $x_0$ in Eq (7). If so, why this procedure could not be applied onto Eq (6) as well? \n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Q1. About the training time of Diffusion-QL. \n\nA1. Thank you for bringing this valuable problem into discussion. We noticed that in Diffusion-QL’s official code repo, they default the number of diffusion steps to 100 (K=100). Please refer to their official `diffusion-rl` repo belonging to organization `Twitter` for details. The hyper-parameter is located at line 187 of file `run_offline.py`. Therefore, to make sure we can reproduce their results, we stick to this setting throughout the experiments. It indeed takes 5 days to finish one experiment. \n\n> Q2. Incorrect statement about “our training scheme only passes through the network once an iteration, no matter how big K is”. \n\nA2. Sorry for the confusion caused. In line 218-220, we are discussing the policy evaluation efficiency. It says DPM-Solver reduces the number of steps to 15. However, in line 221-225 (where the quoted sentence appears), we are talking about the training efficiency. For training, EDP just needs to forward and backward through the network once per iteration, thanks to the proposed action approximation technique. Please see Sec. 4.2 for more details. \n\n> Q3. limited novelty, EDP tried two things: 1) action approximation, 2) DPM-solver\n\nA3. Thank you for opening this discussion on the novelty. We do agree that our proposed methods are quite simple. We’d like to emphasize that the contribution/novelty of our paper is two-fold. 1) A class of diffusion policy with superior training efficiency, which relies on two techniques, **i.e.** , action approximation and DPM-solver.  2) A more general class of diffusion policy that is compatible with both DDPG-style and likelihood-based RL methods (See Sec.4.3). \n\n> Q4.  EAS has been studied by “Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling (ICLR 2023)”. \n\nA4. Thank the reviewer for reminding us of this important work. We will discuss this paper in the future version. We’d like to emphasize that our main contribution is proposing an efficient and general diffusion policy, as tested with three different RL algorithms. EAS is a minor contribution for reducing the variance of the learned policy. We agree that the formulation of EAS is not novel, as similar forms have been explored for policy improvement in MPO[R1,R2]. However, EAS is developed for a very specific and different purpose. Moreover, we found EAS differs from SfBC in the following aspects:\n- Motivation: EAS is developed to reduce the variance of the learned diffusion policy. SfBC aims to model a diverse policy. \n- Usage time: EAS is used at evaluation time only. SfBC is used during training time. \n- Role: EAS is for policy execution. SfBC is for value estimation.\n\n \n[R1] Abdolmaleki, Abbas, et al. \"Relative entropy regularized policy iteration.\" arXiv preprint arXiv:1812.02256 (2018).  \n[R2] Abdolmaleki, Abbas, et al. \"Maximum a posteriori policy optimisation.\" arXiv preprint arXiv:1806.06920 (2018).\n\n> Q5. Why $\\hat{x}_{0,t}$ can be used in policy improvement, but not policy evaluation. \n\nA5. The intuitive explanation is that the action plays a different role in these two steps. Policy evaluation aims to estimate the Q-value (cumulative reward) of a given action starting from a state. Therefore, a correctly paired action-Q data point is important. However, in policy improvement, we have two targets: 1) behavior cloning by diffusion objective, 2) Q guidance by RL objective (Eq. (7)). The rl objective is used to guide the policy optimization towards actions with high return. Though $\\hat{x}_{0,t}$ is not the precise action, it can still provide valuable guidance for policy optimization. \n\nMoreover, we can draw connections to classifier guidance in diffusion models. The Q-network here is a bit like the classifier used in diffusion models. Instead of providing guidance on noisy intermediate steps at inference time, Q-network and action approximation gives guidance at training time. "}, {"review_id": "cOuvFXC3Xb", "reviewer": "Reviewer_44Ux", "summary": "The paper proposes to learn a policy for several offline RL tasks in the D4RL benchmark by parameterizing a policy with a diffusion model. The authors claim that their approach is computationally efficient and more compatible with several other RL approaches like maximum likelihood based approaches when compared to Diffusion-QL which is the primary baseline in their experimental studies.", "questions": "Please refer to some of the questions mentioned in the weakness section;\n\nMinor comments:\n\n1) Is the speed only reason for using the DPM-Solver? Are there other alternatives that were considered\n\n2) Please rephrase line 274 (Do not like phrasing the evaluation protocol of another peer-reviewed paper to be labelled as cheating)", "limitations": "I do not see any discussion around limitations of the approach or other assumptions made. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The paper is well written and easy to read\n2) The authors give a clear and concise preface to the method and results in the introduction part of the paper which sets the flow for the rest of the paper.\n3) The authors compare their approach with the Diffusion-QL method which introduced the idea of formulating policies as a diffusion model. They explain the issues with the Diffusion-QL approach and try to overcome those shortcomings through EDP\n4) Their approach is compatible with max-likelihood based approaches such as IQL (which is one of the SoTA in Offline RL)\n5) The run extensive experiments and ablations on the D4RL benchmark and the results reported are consistently better than Diffusion-QL\n", "weaknesses": "1) Line 206: The authors claim that both policy approximations are empirically similar when calculating the objective though there is no evidence provided for the same. Would be good to know their intuition on that as well.\n2) There is no discussion around why EDP + IQL works significantly better in some tasks like kitchen, adroit whereas it under-performs in the locomotion tasks (where EDP + TD3) performs better. Would be good to know why one works better than another in these cases. Is there any investigation done in that direction?\n3) Barely no discussion around limitations of the approach. \n4) Would the approach be robust to data from multi-modalities? There is a mention of why that would be an issue with gaussian policies but there is no details on that in the latter parts of the paper.\n5) Would the method work on more complex offline-RL task like CARLA? (why was that excluded from the investigation? I understand that diffusion-QL does not test their method on this task but was it ever tried out?)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Q1. Eqn. (12) and Eqn. (13) are empirically similar to each other. No evidence\n\nA1. We ran experiments with Eqn. (13) on three environments with TD3 as the base algorithm. These two approximations are compared in the following table. We can observe that they indeed perform similarly. \n\n|    Base algorithm: TD3    | Eqn. (12) | Eqn. (13) |\n|:-------------------------:|:---------:|:---------:|\n|     walker2d-medium-v2    |    86.9   |    86.9   |\n|  walker2d-medium-reply-v2 |    94.8   |    94.9   |\n| walker2d-medium-expert-v2 |   110.3   |   110.2   |\n\n> Q2. Why EDP-IQL out-performs in kitchen, adroit and antmaze, but underperforms in locomotion. \n\nA2. Thank you for raising this valuable question. EDP-IQL actually shares the same observation as IQL. But, why does this phenomenon happen to IQL? Our hypothesis is that compared to locomotion tasks, the other three types of tasks have sparser reward. As a result, the out-of-distribution issue is amplified in these tasks. IQL performs in-distribution value estimation without querying any out-of-distribution actions, making it less prone to overestimation in such sparse tasks, resulting in better performance. \n\n> Q3. Barely no discussion around limitations of the approach.\n\nA3. Thank you very much for pointing out this problem. We will add the following discussion in the revision. \n\nThough EDP is a much more efficient class of diffusion policies, it is still inefficient compared to a feedforward policy network. For example, the training time on walker2d for a feedforward policy network is just 2 hours, while EDP takes 5 hours. Moreover, The implementation of diffusion policies is much more complicated than its feedforward counterparts. \n\n> Q4. Would the approach be robust to data from multi-modalities?\n\nA4. Sure, diffusion models/policies are superior in modeling multimodal data. We do not provide detailed experimental analysis on this as Diffusion-QL has already conducted a comprehensive study. If you are interested, please refer to Figure 1 of Diffusion-QL for more details. \n\n> Q5. Would the method work on more complex offline-RL task like CARLA?\n\nA5. Sorry, we just tested our methods following conversions. Extending diffusion policies to visual/pixel-based environments itself is a research topic. \n\nThank you for the valuable suggestion! We agree that CARLA is definitely an interesting domain for testing the capability of EDP on more complex tasks. However, in this paper, we only focus on prototyping a general diffusion policy for offline RL, handling the vision inputs of CARLA and imbalanced distribution in self-driving scenarios is out of the scope of this paper. We will leave it for future study.\n\n> Q6. Is the speed only reason for using the DPM-Solver? Are there other alternatives that were considered\n\nA6. Yes, that is the only reason. Sure, we actually tried DDIM, it performs similar to DPM-Solver, but a bit slower. Therefore, we only report DPM-Solver in our paper. \n\n> Q7. Please rephrase line 274 (Do not like phrasing the evaluation protocol of another peer-reviewed paper to be labelled as cheating)\n\nA7. Definitely, we will remove the “cheating” statement and focus on the difference between evaluation protocols only. \n\n"}, {"review_id": "49VMy7ql5j", "reviewer": "Reviewer_SXfY", "summary": "The authors propose a method to efficiently train diffusion based policies in the offline-RL setting. The authors suggest three main tricks to enable this: 1) Removing the need to backpropagate through the diffusion sampling chain to update the policy by using what the authors call action approximation; 2) Replacement of intractable policy log likelihoods in reinforcement learning objectives with the ELBO, which allows using diffusion policies generally with many popular offline RL algorithms; 3) Using a fast DPM-Solver for action sampling during policy execution. Altogether, these tricks allow for relatively fast training of expressive multimodal policies which achieve state of the art results across the D4RL benchmark suite.", "questions": "* Generally, diffusion models generate higher quality samples with larger diffusion timesteps K, and the authors note in Tab. (1) that training EDP with K=1000 matched or beat Diffusion-QL which was trained with much smaller K~5 to 100. The authors cite this ability to train diffusion models with higher K without having to backprop through the chain as one of the main advantages of EDP. I am curious how much of an impact K actually has in these offline RL benchmark tasks, given that Diffusion-QL has similar scores. I am interested to see an ablation study (at least in some environments) where EDP is trained with different values of K starting at K=1, to see if generally increasing K actually significantly improves policies in these domains.\n* Energy-based Action Selection (EAS) seems necessary to come close to the performance of DQL using the proposed EDP policies. The appendix shows that other ways to sample actions from the diffusion model result in much worse performance. DQL does not to my knowledge require EAS to produce good action samples from its diffusion policy. Is this due to the effect of action approximation resulting in training a worse diffusion model?\n* As an extension to the above point, have the authors considered using sampling guidance techniques like Classifer-Free Guidance (CFG) to reduce variance of the final policy?\n* The authors allow for training diffusion policies with popular offline RL algorithms which require policy log likelihoods by replacing it with the ELBO. Could this be done with lower bounds of simpler deep generative models, such as the ELBO of a Variational Autoencoder as well? I am not aware of prior work that has done this. It could be that diffusion policies are overkill for the expressiveness needed to learn high value policies in the tasks under consideration, in which case these same tricks could be applied to train simpler VAE policies for example.", "limitations": "The authors have not addressed technical limitations of their work. They have briefly addressed societal impact of better reinforcement learning methods in the appendix.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* The paper is generally well written and easy to follow.\n* The use of action approximation facilitates training diffusion policy models with much larger diffusion noising time K, with the authors using K=1000. The RL community has recently taken great interest in the prospect of using the expressiveness of diffusion models for policies, and this simple trick seems to make diffusion policy training practical without any noticeable drawbacks.\n* The authors also demonstrate that we can use diffusion policies with other common offline RL algorithms such as TD3+BC, CRR and IQL by simply replacing the log likelihood with the diffusion ELBO. This is a valuable contribution to the field, as it could open the door to use diffusion policies to work with many general RL algorithms, a strict improvement over simple gaussian policies.\n* The authors show strong empirical performance across the D4RL benchmark suite. EDP policies are especially strong compared to the FF counterparts on the more challenging antmaze and kitchen suite, which involve high multimodality due to undirected demonstrations.", "weaknesses": "* The main proposal to efficiently train diffusion policies hinges on action approximation, which uses a much higher variance estimate of $\\hat{a}^0$. There are no ablations to show how this affects the final diffusion policy as compared to one trained without action approximation. There is a comparison against Diffusion-QL in the paper, but with a very large difference in diffusion timesteps K, so the effect of action approximation itself is difficult to gauge. I have listed some questions related to this point below, which if clarified would be helpful.\n* The use of DPM-Solver is highlighted as a major contribution is speeding up the sampling process compared to DDPM sampling in Diffusion-QL. It is a well known fact that faster diffusion sampling methods than DDPM exist which can be used for sampling, and Diffusion-QL could have used this as well. While it is useful that the authors show the time savings for using DPM-Solver over DDPM sampling, this is not a novel contribution.\n* There isn’t discussion about important hyperparameters related to the method. A very important hyperparameter that could be discussed more the diffusion timesteps K (I have noted this as a question below as well). Are EDPs more brittle to train than their FF counterparts in algorithms like IQL with regards to changing hyperparameters?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Q1. The effect of action approximation\n\nA1. We compare with and without action approximation on the following three environments by using OMS metric (Table. 4). In the following table, the DDPM column will forward and backward a policy network 100 times at training time, but action approximation only needs once. We can observe that action approximation will slightly harm the performance, when the same number of diffusion steps is used. However, it supports training diffusion policies with larger K (e.g. 1000), while Diffusion-QL does not. Increasing K is able to avoid performance drop as evidenced by the last column of Table 1 and Table 4. \n\n\n|         OMS K=100         |  DDPM | Action Approx |\n|:-------------------------:|:-----:|:-------------:|\n|     walker2d-medium-v2    |  86.9 |      85.5     |\n|  walker2d-medium-reply-v2 |  96.3 |      93.3     |\n| walker2d-medium-expert-v2 | 111.5 |     111.1     |\n\n> Q2. Using DPM-Solver is not a novel contribution. \n\nA2. We totally agree that simply replacing the DDPM sampler with DPM-Solver is not a novel contribution. We will modify the statement accordingly to highlight that DPM-Solver is able to improve efficiency but not novelly developed by us. A diffusion policy that can be efficiently trained and a generation diffusion policy compatible with likelihood-based methods are the main contributions. \n\n> Q3. Are EDPs more brittle to train than their FF counterparts in algorithms like IQL with regards to changing hyperparameters?\n\nA3. In our experiments, we do not tune the IQL parameters when it is combined with EDP. For the hyper-parameters for the diffusion part, we just adopt the one used in TD3+EDP. Therefore, we believe that EDPs are quite stable to train with IQL. \n\n> Q4.  if generally increasing K actually significantly improves policies in these domains.\n\nA4. As discussed in Q1, when K is small, action approximation generally performs worse than its full-chain counterparts. Moreover, we analyze the effect of K with EDP on three gym-locomotion tasks. We can conclude that increasing K is able to improve the performance steadily. EDP often fails when K is smaller than 10, but becomes steady when K is larger than 50. \n\n\n|           EDP  K          |   1  |   2  |   3  |   5   |   10  |   50   |   100  |\n|:-------------------------:|:----:|:----:|:----:|:-----:|:-----:|:------:|:------:|\n|     walker2d-medium-v2    | 1.08 | 4.65 | 5.78 | 10.05 | 81.22 |  84.08 |  85.50 |\n|  walker2d-medium-reply-v2 | 0.45 | 1.76 | 3.10 |  4.01 | 18.89 |  89.18 |  93.03 |\n| walker2d-medium-expert-v2 | 0.02 | 2.00 | 4.73 |  8.34 | 80.92 | 109.90 | 111.00 |\n\n> Q5. Is the usage of EAS due to the effect of action approximation resulting in training a worse diffusion model?\n\nA5. This is a really great question. Actually Diffusion-QL suffers the same high-variance problem. And their code actually involves a similar technique like EAS, which is not discussed in their paper. Please refer to their official `diffusion-rl` repo belonging to organization `Twitter` for details. The code snippet is located at line 169-176 of file `agents/ql_diffusion.py`.\n\n> Q6. Using Classifier-free guidance (CFG) to reduce variance of the final policy. \n\nA6. That is a really great suggestion. Unfortunately, we did not try this before mainly due to the following reason. First, It is hard to decide what conditional information to use, probably we can use normalized return as a condition. Second, we need to retrain a diffusion policy such that it can take a conditional or an idle signal as input. Instead, EAS can serve as a plug-in component for evaluation without modifying the training procedure. \n\n> Q7. Can simpler deep generative models be used as policies? \n\nA7. Sure, the answer is yes. BCQ [R1] has successfully trained a conditional-VAE as the diffusion policy. Diffusion-QL also studies different forms of generative models as policies and compares them with diffusion models.  The conclusion is that Diffusion policies show superior ability at capturing data with strong multi-modalities. So, when the data distribution is not so complicated, simpler generative models are definitely a promising choice. \n\n[R1] Fujimoto, Scott, David Meger, and Doina Precup. \"Off-policy deep reinforcement learning without exploration.\" International conference on machine learning. PMLR, 2019.\n\n"}, {"review_id": "9Bb2MlJ92u", "reviewer": "Reviewer_UWuP", "summary": "This paper focuses on the improvement of computation efficiency of Diffusion-QL by adopting the property of marginal distribution in the diffusion model and the variance control scheme proposed by DPM-Solver. Besides, this paper extends the scope of compatibility with other offline RL methods, from value-based to policy gradient methods. This is achieved by directly approximating the clean examples from corrupted examples at arbitrary diffusion time steps. ", "questions": "See my detailed comments below.\n\n==Major concerns==\n\n- In lines 233-235 of Section 4.5: according to Equation (9), the action approximation is the mean value of action Gaussian distribution, so why adopting $\\hat{a}^0$ can not reduce high variance?\n- In Section 4.4, the improvement of sample efficiency (from $ K=1000$ to $K\\approx 15$) comes from DPM-Solver rather than the action approximation. Firstly, if the author adopts action approximation, whether the $a_K$ comes from normal Gaussian distribution? If so, the noise prediction function $\\epsilon_theta (a^K,K;s)$ should predict the action based on the state and the Gaussian noise $a^K$. Intuitively, generating proper actions based on states seems more difficult, i.e., that classical policy function.\n- As for the log-likelihood, why can we not obtain the exact log-likelihood of $\\pi_{\\theta}(a|s)$ through $log~p(a^K)+\\sum_{k=1}^{K}log~p_{\\theta}(a^{k-1}| a^{k}, s)$?\n- Does the action approximation improve the computation efficiency of diffusion policy? \n- I hope the authors explain how they use the action approximation during training.\n- During the sampling stage, the DPM-Solver speeds up generation rather than action approximation. If action approximation just works on Equation (10), I think it is very similar to the “Target Policy Smoothing Regularization” technique in Section 5.3 of the TD3 paper. I want to know why the authors use the predicted $\\hat{a}^0$ rather than $a^0+\\epsilon$.\n\n\n==Minor concerns==\n- How does the author avoid out-of-distribution actions' effects on the Q function?\n- DPM-Solver directly makes me understand the sample efficiency of EDP, but I can not understand why the training efficiency can also be improved? So I hope more description of the reason.\n- In line 318-320 of Section 5.1: during training the diffusion model, we sample a step $t\\sim [1, K]$, then train the diffusion model according to Equation (5). If the $K$ is large, we can still train the diffusion model. So I hope the authors explain this claim more clearly.\n\n\n==Typos==\n- Line 188: Eqn. 7 -> Eqn. (7).\n- Line 176: Eqn. 1 -> Eqn. (1).\n- Line 201: Eqn. 11 -> Eqn. (11).\n- Line 292: Eqn. 4 -> Eqn. (4).\n- Line 518 of Appendix: Eqn. 5 -> Eqn. (5).\n- Line 519 of Appendix: Eqn. 9 -> Eqn. (9).\nSimilar typos can also be found in many places in Appendix.\n- Line 565 of Appendix: Tab. ??.\n- Line 575 of Appendix: Tab. ??.\n", "limitations": "The authors do not discuss the limitations and discuss broader impacts.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper is well-written, and the proposed method is easy to understand.\n- The evaluation of value at risk seems interesting.\n- The technique seems sound, and the results seem strong.\n", "weaknesses": "- The discussion of limitation seems to be neglected.\n- More clear discussion is needed to highlight the contribution of this paper and show the novelty.\n- The proposed algorithm seems to significantly increase the inference computation complexity by introducing 10000 iterations for each action.\n- More experiments are needed to support the contributions sufficiently.\n- Some typos are needed to fix through further proofreading.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Q1. In lines 233-235 of Section 4.5, why adopting $\\hat{a}^0$ can not reduce high variance?\n\nA1. Here is an intuitive explanation. Given an actual action $a$, action approximation $\\hat{a}^0$ represents the mean of the action Gaussian distribution that can be denoised from $a^k$. Therefore, it is related to a specific example drawn from a multi-modal distribution.   $\\hat{a}^0$ can be used to guide the optimization gradient for sample $a$, but can not represent the whole data distribution.  \n\n> Q2. The improvement of sample efficiency comes from DPM-Solver rather than action approximation. \n\nA2. Yes, the sample or inference efficiency solely comes from DPM-Solver. However, action approximation can greatly boost the training or policy improvement efficiency, as discussed in Sec. 4.2 and Sec. 4.4. This is because one needs to backpropagate the whole sampling chain of an action to optimize the policy. Instead of sampling an action through multiple steps of denoising diffusion, action approximation just needs to forward the network once. \n\n> Q3. Whether the $a^K$ comes from a normal Gaussian distribution? \n\nA3. Yes, when $k=K$, $a^k$ is drawn from a normal gaussian distribution. However, $k$ is uniformly sampled from the discrete set ${1,2,..,K}$, which means $k=K$ is only a small portion. \n\n> Q4. As for the log-likelihood, why can we not obtain the exact log-likelihood of \n through  $\\log p(a^K)+\\sum_{k=1}^{K} \\log p_{\\theta}(a^{k-1}| a^{k}, s) $\n\nA4. This is because $\\pi_\\theta$ is actually written as the following form, \n$$\\pi_\\theta(a|s) = \\int_{a^K}\\int_{a^{K-1}}\\dots\\int_{a^1} p(a^K|s) p_\\theta (a^{K-1}|a^K, s) \\dots p_\\theta(a^0|a^1,s) d_{a^K} d_{a^{K-1}}\\dots d_{a^1}$$\nwhere the integrations are intractable. \n\n> Q5. Does the action approximation improve the computation efficiency of diffusion policy?  I hope the authors explain how they use the action approximation during training.\n\nA5. Sure, action approximation can boost the training efficiency greatly. Recall the policy optimization procedure in TD3. One needs to first draw a sample from the policy, then feed it into the Q-network to get a Q estimation. Then, the policy is optimized to maximize the Q estimation. This procedure requires backpropagation through the action sampling process. Without action approximation, It takes K times to query the diffusion network for action sampling, thus K times of backpropagation. However, with action propagation, we just need to forward the denoising diffusion network once, which greatly saves GPU memory and optimization time. \n\n> Q6. I want to know why the authors use the predicted $\\hat{a}^0$, rather than $a^0 + \\epsilon$. \n\nA6. This is because the Q-network takes a clean example as input. Instead, $a^0 + \\epsilon$ is a noisy one, feeding it into the Q-network can not give a proper estimation of the corresponding Q value. \n\n> Q7. How does the author avoid out-of-distribution actions' effects on the Q function?\n\nA7. Apart from the RL objective, we also have a behavior cloning term (i.e., the diffusion objective), which restricts the policy to be similar with the behavior policy. \n\n> Q8.  why the training efficiency can also be improved?\n\nA8. Please refer to A5. More discussion on this problem is welcomed. \n\n> Q9. Why can we not train Diffusion-QL when K is large, but train diffusion model? \n\nA9. The difference is that Diffusion-QL needs to sample an example and backpropogate through its sampling chain at training time, but normal diffusion models do not. When K is large, the sampling chain will involve thousands of neural networks, resulting in a super huge computation graph. Based on our empirical experience, it is even computational intractable to compile such a big graph with JAX. Feel free to try with our code by setting K=1000, and disable action approximation. \n\n\n"}]

Author rebuttal:
We sincerely thank all the reviewers for recognizing the novelty and contributions of our work, as well as for providing valuable questions for discussion and constructive suggestions.

As there are limited shared questions from the reviewers, we address them individually in the corresponding responses. Specifically, as requested by Reviewer 44Ux and Reviewer SXfY, we added additional experiments that we'd like to attract your attention, including

- ablating the policy approximation by comparing Eqn. 12 and Eqn. 13 in likelihood-based policy optimization (requested by Reviewer 44Ux). 
- ablating the effect of action approximation (requested by Reviewer SXfY)
- ablating the effect of the diffusion timesteps K (requested by Reviewer SXfY)

We have directly appended the updated tables to our responses for the respective reviewers. We value your feedback and please kindly let us know if you have any further questions or concerns.

Sincerely,

Authors

Meta-review: This work showcases a new implementation of a diffusion policy that is algorithm-independent. The reviewers agree this paper shows a contribution across algorithms. The use of DPM-Solver is highlighted as a major contribution in speeding up the sampling process compared to DDPM sampling in Diffusion-QL. There are other algorithms for improving the sampling speed for DPM solvers. The proposed method works well for RL problems. I also tend to agree the breadth of experiments shows a general improvement.
Decision: Accept (poster)