DIME: Diffusion-Based Maximum Entropy Reinforcement Learning
Authors: Onur Celik
Zechu Li
Denis Blessing
Ge Li
Daniel Palenicek
Jan Peters
Georgia Chalvatzaki
Gerhard Neumann
Venue: ICML 2025 poster (2025)
Keywords: Reinforcement Learning
Diffusion Models
Diffusion Based Reinforcement Learning
Maximum Entropy Reinforcement Learning
Primary area: reinforcement_learning->deep_rl
TL;DR: We provide DIME, a method for training diffusion-based policies in the maximum entropy reinforcement learning framework.

Abstract: Maximum entropy reinforcement learning (MaxEnt-RL) has become the standard approach to RL due to its beneficial exploration properties. Traditionally, policies are parameterized using Gaussian distributions, which significantly limits their representational capacity. Diffusion-based policies offer a more expressive alternative, yet integrating them into MaxEnt-RL poses challenges—primarily due to the intractability of computing their marginal entropy. 
To overcome this, we propose Diffusion-Based Maximum Entropy RL (DIME). DIME leverages recent advances in approximate inference with diffusion models to derive a lower bound on the maximum entropy objective. 
Additionally, we propose a policy iteration scheme that provably converges to the optimal diffusion policy. Our method enables the use of expressive diffusion-based policies while retaining the principled exploration benefits of MaxEnt-RL, significantly outperforming other diffusion-based methods on challenging high-dimensional control benchmarks. It is also competitive with state-of-the-art non-diffusion based RL methods while requiring fewer algorithmic design choices and smaller update-to-data ratios, reducing computational complexity.

DIME: Diffusion-Based Maximum Entropy Reinforcement Learning
Onur Celik
Zechu Li
Denis Blessing'
Ge Li Daniel Palenicek Jan Peters 3456 Georgia Chalvatzaki
Gerhard Neumann"
Abstract
several favorable properties among which improved explo- ration (Ziebart, 2010; Haarnoja et al., 2017) is crucial for RL. Recent successful model-free RL algorithms leverage these favorable properties and build upon this framework (Bhatt et al., 2024; Nauman et al., 2024) improving sam- ple efficiency and leading to remarkable results. However, the policies are traditionally parameterized using Gaussian distributions, significantly limiting their representational ca- pacity. On the other hand, diffusion models (Sohl-Dickstein et al., 2015; Ho al., 2020; Song et al., 2021; Karras et al., 2022) are highly expressive generative models and have proven beneficial in representing complex behavior policies (Reuss et al., 2023; Chi et al., 2023). However, important metrics such as the marginal entropy are intractable to com- pute (Zhou et al., 2024) which restricts their usage in RL. Because of this shortcoming. recent methods propose dif- ferent ways train diffusion-based methods in off-policy RL. While these methods are discussed in more detail in the related work section, most of them require additional techniques to add artificial (in most cases Gaussian) noise to the generated actions to induce exploration in the behav- ior generation process. Hence, they do not leverage the diffusion model to generate potentially non-Gaussian explo- ration patterss but fall back to mainly Gaussian exploration. Nonetheless, there have been significant advances in training diffusion-basse models for approximate inference (Berner et al.; Richter & Berner). Since the policy improvement in MaxEnt-RL can also be cast as an approximate inference problem to the energy- policy (Haarnoja et al., 2017), it a natural step to explore these parallels.
Maximum entropy reinforcement learning (MaxEnt-RL) has become the standard approach to RL due to its beneficial exploration properties. Traditionally, policies parameterized using Gaussian distributions, which significantly limits their representational capacity. Diffusion-basee policies offer a more expressive alternative. yet integrating them into MaxEnt-RL poses challenges- primmall due to the intractability of computing their marginal entropy. To overcome this, propose Diffuuson- -Based Maximum Entropy RL (DIME). DIME leverages recent advances in approximate inference with diffusion models to derive a lower bound on the maximum entropy objective. Additionally, we propose a policy iteration scheme that provably converges the optimal diffusson policy. Our method enables the of expressive diffusion- blafus policies while retaining the principled exploration benefits of MaxEnt-RL, significantly outperforming other diffusion- based methods on challenging high-dimennt ligh-dimenssooal control benchmarks. also competitive with state-of-the-ar non-diffusion based RL methods while requiring fewer algorith- mic design choices and smaller podate-to-date ratios, reducing computational complexity!.
1. Introduction
We propose Diffusion-Basse Maximum Entropy Reinforce- ment Learning (DIME). DIME leverages recent advances in approximate inference with diffusion models (Richter & Berner) to derive a lower bound on the MaxEnt objective. We propose a policy iteration framework with monotonic policy improvement that converges to the optimal diffusion policy. Additionally, building on recent off off-policy RI. algo- rithms such as Cross-Q (Bhatt al., 2024) and distributional RL (Bellemare et al., 2017). we propose practical version of DIME that can be used for training diffusion- based RL policies. On 13 challenging continuous high-dimensional control benchmarks, we empirically validate that DIME sig- nificantly outperforms other diffusion- based baselines on
The maximum entropy reinforcement learning (MaxEnt-RL.) objective augments the task reward in each time step with the entropy of the policy (Ziebart et al., 2008; Toussaint, 2009; Haarnoja et al., 2017; 201 8b). This objective has
Autonomous Learning Robots, KIT "Interactive Robot Per- ception & Learning, Darmstadt 'Intelligent Autonomous Sys- tems, TU Darmstadt 'Hessian.Al German Research Center for Al "Centre for Cognitivie Science, TU Darmstadt. Correspondence to: Onur Celik @
Proceedings of the 4gnd International Conference on Machine I Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Phtpps:/// //al /tapp::://arrb...../ rhub. witub...ttuu...... io/dime io/dime-websitee

---

DIME: Diffusion- Billusion-Based Maximum Entropy Reinforcement Learning
all environments and consistently outperforms other state- of-the-art RL methods based on Gaussian policy on 10 out of 13 environments, while being computationally more efficient and requiring less algorithmic design choices as the current state of the art baseline BRO (Nauman et al., 2024).
updated using the gradients of the Q-function (oQ(o, a). DIPO has been extended to develop methods for learning multi-modal behaviors(Li et al., 2024) by leveraging hierar- chical clustering to isolate different behavior modes. DIPO relies on the stochasticity inherent to the diffusion model for exploration and does not explicitly control it via an objec- tive. QSM (Psenka et al., 2024) directly matches the policy's score with the gradient of the Q-function Va (s,, a). While their objective avoids differentiating through the whole dif- fusion chain, the proposed objective disregards the entropy of the policy and, therefore, exploration. Consequently, QSM needs to add noise to the final action of the diffusion chain. More recently, DACER (Wang et al., 2024) proposed using the data generaenng process as the policy representa- tion and backpropagating the gradients through the diffusion chain. However, they do not consider a backward process as we do, and their objective for updating the diffusion model is based on the expected Q-values only. To incentivize the exploration, DACER adds diagonal Gaussian noise to the sampled actions, where the variance of this noise is con- trolled by a scaling term that is updated automatically using an approximation of the marginal entropy by extracting a Gaussian Mixture Model from the diffusion policy. Con- currently, QVPO (Ding al., 2024) proposed weighting their diffusion loss with their respective Q-values after ap- plying transformations. However, QVPO relies on sampling actions from a uniform distribution to enforce exploration.
2. Related Work
Maximum Entropy RL. The maximum entropy RL frame- work uses the entropy of the policy at each time step an additional objective, providing a principled way of in- ducing exploration in the RL policy. It is different from entropy regularized RL (Neu et al., 2017), where the en- tropy of the policy is maximized only for the current time step. Haarnoja et al. (2017) proposed Soft-Q Learning, where amortized Stein variational gradient descent (Wang & Liu, 2016) (SVGD) is used train a parameterized sam- pler that can sample from the energy- -aased policy. SAC (Haarnoja et al., 2018b) proposes an actor-critic RL method but frames the policy update as an approximate inference problem to the energy-based policy using a Gaussian policy parameterization. SAC has been extended to energy-based policies using SVGD in (Messaoud et al.), where the au- thors also propose new method to estimate the entropy in closed form. While SVGD is powerful method for learning an energy based policy, it is harder to scale these approaches to high-dimensional control problems. For im- proving exploration, LSAC (Ishfaq et al., 2025) proposes leveraging Langevin Monte Carlo (Welling & Teh, 2011) in conjunction with a distributed critic objective to sample a state-action value. Haarnoja et al. (2018a) proposes learning a latent variable model as policy representation, but relies on the change of variable formula to express the density of the policy by calculating the Jacobian of the transformations. Recent advances of SAC also define the satte-of-the-are in off-policy RL many domains, such as CrossQ (Bhatt et al., 2024) and BRO (Nauman al., 2024). CrossQ proposed removing the target network by leveraging batch renormal- ization and BRO scales to large networks in RL by using several methods such as optimistic exploration (Nauman & Cygan, 2023), network resets (Nikishin et al., 2022), weight decay, and high update-to- to-lat--to-ataa data ratios.
DIME distinguishes from prior works in that we use the maximum entropy RL framework for training the diffusion policy, which was not considered before. This allows direct control of the auploratioo-eeploitaaaii urally through this objective without the need for additional approximations. DIME is leveraging the diffusion model to generate non-Gaussian exploration actions which is in contrast to most other diffusion RL approaches that still require including Gaussian or uniform exploration noise.
trade-off arising nat-
Approximate Inference with Diffusion Models. Early works on approximate inference with diffusion models were formalized as stochastic optimal control problem using Schrödinger-Fölie Föllmer diffusions (Dai Pra, 1991; Tzen & Ra- ginsky, 2019; Huang et al., 2021) and only recently realized with deep-learning based approaches (Vargas et al., 2023; Zhang & Chen, 2021). Vargas et al.; Berner et al. later extended these results to denoising diffusion models. A more general framework where both forward and backward processes of the diffusion model are learnable was concur- rently proposed by Richter & Berner; Nusken et al. (2024). Recently, many extensions have been proposed, see e.g. (Akhound -adnnh et al., 2024; Noble et al., 2024; Geffner & Domke, 2023; Zhang et al., 2023; Chen et al., 2024; Blessing et al., 2025b;a; Chen et al., 2025). Our work can be seen as instance of the sampler presented in (Berner et al.). However, our formulation allows using different
Diffusion-Baan Based Policies in RL. Early works have re- searched diffusion models in offline RL (Lange al., 2012: Levine et al., 2020) as trajectory generators (Janner al., 2022) or as expressive policy representations (Wang al., 2023; Kang et al., 2023; Hansen-Estruch et al., 2023; Chen et al., 2023; Ding & Jin, 2024; Mao et al., 2024; Fang et al., 2025; Lu et al., 2023). More recently. diffusion models in online RL have become more popular. DIPO (Yang et al., 2023) proposes training a fnffuusion-based policy using a behavior cloning loss. The actions in the replay buffer serve as target actions for the policy improvement step and are
2

---

DIME: Diffusion-Basee Maximum Entropy Reinforcement Learning
diffusion samplers such those presented in (Richter & Berner; Blessing et al., 2025a), while restrict ourselves in this work to the sampler presented in (Berner et al.).
optimal policy with respect to the maximum entropy ob- jective. The two steps include policy evaluation and policy improvement. Given policy policy evaluation aims to evaluate the value of т. To that end, (Haarnoja et al., 2018b) showed that repeated application of the Bellman backup operator T Q^ with
3. Preliminaries
3.1. Maximum Entropy Reinforcement Learning
T"*(ss,, a1) a r + (Q[ 04+1) + H(a++1/s44
Notation We consider the task of learning a policy π : S x A + Rt, where S and A denote continuous state and action space, respectively using reinforcement learning (RL). We formalize the RL problem using infinite horizon Markov decision process consisting of the tuple (S, T, P):7), with bounded reward fune- tion r S x A + [Tminn: Fmux] and transition density R+ which denotes the likelihood for
(5)
converges to Q as k + oc, starting from any Q. To update the policy, that is, to perform the policy improvement step. the Q-function of the previous evaluation step, QTa is used to obtain a new policy according to
ST new arg min DKL (T(eel))) exp (SL) (Taa(st)
(6)
p : S x A transitioning into a state s E S when being in & E and executing an action E A. We follow (Haarnoja et al., 2018b) and slightly overload Prr which denotes the state and state-action marginals induced by policy Moreover, E [0, 1) denotes the discount factor. For brevity, we use A r(st, dz). Lastly, we denote objective functions that we aim to maximize as J and minimize as L.
Tell
where II is set of policies such as a family of parame- terized distributions. Note that (*(H(ss) is not required for optimization as it is independent of т. Haarnoja et al. (2018b) showed that for all state-action pairs (s,a) E S A it holds that QTTT (s, a) 2 QTM"s, a) ensuring that policy iteration converges to the optimal policy n' in the limit of infinite repetitions of policy evaluation and improvement.
Control as inference. The goal of maximum entropy rein- forcement learning (MaxEnt-RL) is to jointly maximize the sum of expected rewards and entropies of policy
3.2. Denoising Diffusion Policies
For a given state S, consider stochastic process the time-interval [0, I1 given by an -mlteen-Uhlenbeeck
J(T)= = Σ 7t EEp Ir. + a((T(()(s)))
(1)
(OU) process (Särkkä & Solin, 2019)
t=1
da, -3, + ny 23,d ao #-(())), (7)
where H(T(als)) - s(a]s) log ((assdda is the differ- ential entropy, and R+ controls the exploration ploitation trade-off (Haarnoja et al., 2017). To keep the notation uncluttered we absorb a into the reward function via r + r/a. Defining the Q-function of a policy π as
with diffusion coefficient/ B [0, + Rt, standard Brow- nian motion (Bt))e(o,7). and some target policy uo. For 1, I E [0, 7). we denote the marginal density of Eq. 7 at as ग. and the conditional density time given I as नय.. Eq. 7 is commonly referred to as forward or noising process since, for suitable choice of B, it holds that कं s NN(0, I))... Denoising diffusion models leverage the fact, that the time- reversed process of Eq. 7 is given by
Q= (8),0)))= = + y Rp, [ret + H ( (a((et))(e)1))
(2)
with Q= 8 x A + R. the MaxEnt objective can be cast as an approximate inference problem of the form
da,- = (-B - 2n1 Be V log (arls)) + ny 23,dB., (8) starting from #y = मे. R NO 17²1) and running back- wards time (Nelson, 2020; Anderson, 1982; Haussmann & Pardoux, 1986). For the backward, generative or de- noising process (Eq. 8), we denote the density as æ. Here, time-reversal means that the marginal densities align, i.e., म. = F. for all t E [0, T.. Hence, starting from ar ~ NN(0, 1)2 one can sample from the target policy Ho by simulating Eq. 8. However, for most densities Mo.
[(*) DKL. (T(arlss) exp (81,8))))
(=t(s)
in a sense that max, (TT) min, [(T). Here, DKL de- notes the Kullback-Leibler divergence and
2 (s) / exp ex Q2 (s, (s, a)da
(4)
is the state-dependent normalization constant.
'leasee note, for clarity, we slightly abuse notation by using to denote the time in the stochastic process. This should not be confused with the time step in RL. The distinction becomes clear when we discretize the processes.
Policy iteration a two-step iterative update scheme that is, under certain assumptions, guaranteed to converge to the
3

---

DIME: Diffusion- Coluusion-Based Maximum Entropy Reinforcement Learning
()
N1 7)
(b)
(c) a>>
Figure 1. The effect of the reward scaling parameter a. The figures in (a)-(b) show diffusion processes for different a values starting at prior distribution (0, I) and going backward in time approximate the target distribution exp (Q" /a)/'"" Small values for (a) lead to concentrated target distributions with less noise in the diffusion trajectories especially at the last time steps. The higher a becomes (b) and (c), the more the target distribution is smoothed and the distribution of the samples the last time steps becomes more noisy. Therefore, the parameter a directly controls the exploration by enforcing noisier samples the higher a becomes.
the scores (V log (a, ()) s) 101,2 are intractable, requiring numerical approximations. To address this, denoising score- matching objectives are commonly employed, that is,
for diffuusion model is difficult the marginal entropy M())(aas)) of the denoising process in Eq. 8 is intractable. Please note that we use superscripts for the actions to in- dicate the diffusion step to avoid collisions with the time step used in RL. Moreover, we will again absorb a into the reward and use r. å ((s),a)). To overcome this in- tractability, we propose to maximize a lower bound. We start by discretizing the stochastic processes introduced in Section 3.2 and use the results as a foundation to derive this lower bound. Note that while similar results can be derived from a continuous-time perspective (see e.g., Berner et al.; Richter & Berner; Nusken et al. (2024)), such derivation would require a background in stochastic calculus, making less accessible a broader audience.
LsM (0) E [8ell (ar, 8) - V log (ajo(a (ao: s)]]?] (9) where is sampled on [0, T) and 1 denotes parameterized score network (Hyvärinen & Dayan, 2005; Vincent, 201 I).
For OU processes, the conditional densities V log 1ojo are explicitly computable, making the objective tractable for optimizing (Song et al., 2021). Once trained, the score network f can be used to simulate the denoising process
da,- (-3,ard 2nj² B. (at, s)) + / B,, (10)
to obtain samples ag F that are approximately dis- tributed according to #o. Here, of denotes the marginal distribution of Eq. 10 at time t. While score-matching tech- niques work well in practice, they cannot be applied to max- imum entropy reinforcement learning. This is because the expectation in Eq. requires samples dg ~ mn xx exp Q" which are not available. However, in the next section, we build recent advances in approximate inference opti- mize diffusion models without requiring samples from dg. relying instead on evaluations of Q".
The Euler-Maruyama (EM) discretization (Särkkä & Solin, 2019) of the noising (Eq. 7) and denoising (Eq. 8) process is given by
a+1 - -0,aaaa t and
(1))
a-1 = a² + (3na" + 20 Bn V log न. (a*)s)) 5 + En. (12)
respectively, with En ~ NO 277 82))). Here, 5 de- notes a constant discretization step size such that N = 7/8 is an integer. To simplify notation, we write a", instead of n6 Under the EM discretization, the noising and denoising process admit the following joint distributions
4. Diffusion-Based Maximum Entropy RL
Here, we explain how diffusion models can be used within a maximum entropy RL framework. To that end, we express the maximum entropy objective as approximate inference problem for diffusion models. We then use these results to introduce a policy iteration scheme that provably converges to the optimal policy. Lastly, we propose a practical algo- rithm for optimizing diffusion models.
No (a (NN)ss #o(a"ls)
II
मे 1p
taa la),s), (13)
N
(a (a
II ==1
1n (a
(am= s)),
(1 (14)
4.1. Control as Inference for Diffusion Policies
Directly maximizing the maximum entropy objective
in a sense that To.N and Wo.N converge to the law of (a)) in Eq. 7 and 8, as 8 + 0, respectively (Doucet et al., 2022). Here, I and - are Gaussian transi- tion densities that directly follow from Eq. 11 and 12.
((H) = Σ Ep, [re(se, a.) aH((t((a(ss))
t=?

---

DIME: Diffusion- Billudion-Based Maximum Entropy Reinforcement Learning
To obtain a maximum entropy objective for diffusion mod- els, make use of the following lower bound on the marginal entropy, that is, H(To((aols) 2 (a s), where
4.2. Diffusion-based Policy Iteration
We propose a policy iteration scheme for learning an optimal maximum entropy policy, similar (Haarnoja et al., 2018b) However, here we restrict the family of stochastic actors to diffusion policies i E II C II. Throughout this section, we assume finite action spaces to enable theoretical analysis, but relax this assumption in Section 4.3. All proofs of this section are deferred to Appendix A.
TENiola'-Naa E (a s) 5) -ENN log #o.v (auN (8)
(15)
Please note that similar bounds have been used, c.g., in (Agakov & Barber, 2004; Tran et al., 2015; Ranganath et al., 2016; Maalge et al., 2016; Arenz et al., 2018), or, more generally, follow from the data processing inequality (Cover, 1999). A derivation can be found in Appendix A. From Eq. 15, it directly follows that
For policy evaluation, we aim compute the value of policy i.. We define the Bellman backup operator
T Q(st( of) _ r [Q(5+11 a.+1) + la (a2
(21)
Note that Eq. 21 contains the entropy-lower bound L.. By applying standard convergence results for policy evaluation (Sutton & Barto, 1999) we can obtain the value of a policy by repeatedly applying T as established in Proposition 4.1, Proposition 4.1 (Policy Evaluation). Let T be the Bellman backup operator for a diffussion policy t as defined in Eq. 21. Further, let Q" : S x A + R and Q = =Ta Q² Then, it holds that litnk-ooo Q² = Q" where Qa is the Q value of न.
(( 2 J(
7 Ep [Te + (a(())
(16)
Next, we cast Eq. 16 as an approximate inference problem to make the objective more interpretable. To that end, let us define the Q-function of denoising policy with respect to the maximum entropy objective
Q (st, a rt + Σ TT [rete + (=(al+ st+i]
For the policy improvement step, we seek to improve the current policy based on its value using the Q-funtiion. For- mally, we need to solve the approximate inference problem
(17) with Q² : S A R. With Eq. 17 we identify the corresponding approximate inference problem as finding के
which minimizes (please see Appendix A for derivation)
π
arg min DKL. (#o:v (N [s)] F (a 0:N (s)) #cñ
C(#) = DKL (To.N (N [s) (c
(18)
(22)
for all 8 E S, where To.N (a-N is as in Eq. 13 with marginal density
where the target policy, i.e., the marginal of the noising process in Eq. 13 given by the exponentiated Q-function of the diffusion policy
no old (a. (s))s) exp Quiu (s.a))
(23)
oo (o(a"ls) exp (s,a") (=(s)
Indeed, solving Eq. 22 results in a policy with higher value as established below.
(19)
Proposition 4.2 (Policy Improvement). Let Rolds Fnew E II be defined as in Eq. 23 and 22, respectively. Then for all (s, a) E S x A it holds that @ a) (s, as.
Recall from Section 3.2 that aim time-reverse the nois- ing process, that is, to ensure for all states E S, it holds that No.N No Please note that this is precisely what Eq. 18 is trying to accomplish, i.e., we aim to learn a diffusion model i., such that the denoising process time-reverses the noising process, and, in particular, has a marginal distri- bution given by Ta = exp Q" / Lastly, from the data processing inequality, it directly follows that
Combining these results leads to the policy iteration method which alternates between policy evaluation (Proposition 4. and policy improvement (Proposition 4.2) and provably converges to the optimal policy in IT (Proposition 4.3).
Proposition 4.3 (Policy Iteration). Let 7 =i+1 , F E II. Further, let #1+1 be the policy obtained from म² after a policy evaluation and improvement step. Then, for any starting policy =0 holds that lim,--oc 5² = # with = such that for all क Il and (s,a) S x A holds that Q* (s 2 Q (s.
(s, exp (*((s )
Dri. [o]
(s)))
< DRL. (#(a 0:N ONN (aN))()))
(20)
which shows the approximate inference problem in Eq. 18 indeed optimizes the same inference problem stated in Eq. 3.
However, performing policy iteration until convergence is in practice often intractable, particularly for continuous control tasks. As such, we will introduce practical algorithm next.
Next, we will use these results to develop a policy iteration scheme for diffusion models.
5

---

DIME: Diffusion-Bare Maximum Entropy Reinforcement Learning
10-01
10-02
10-03
10 10-01
10
10
10
10-12
10M Rotimm
600
Stoum
600
400
400
400
200
200
200
DIML
DHNE
0
0.5
1.5
2
2.5
3
500
600
700
800
0.5
1.5
2
255
0.5
1.3 2
25
Number Lay Interactions -10P
IQM Retur
Number Eny Internctions -10f
Number Err Interactions
-10
(a)
(b)
(c)
(d)
Figure Reward Scaling Sensitivity (a)-(b). The a parameter controls the exploration-exploittiio trade-off. (a) shows the learning curves for varying values on DMC's dog-run task. Too high values (a 0.1) do not incentivize learning whereas too small values (a S 10 10- converge to suboptimal behavior. (b) shows the aggregated end performance for each learning curve in (a). For increasing a values, the end performance increases until it reaches optimum
10 after which the performance starts dropping. Diffusion Policy Benefit (c) and (d), We compare DIME to a Gaussian policy with the same implementation details as DIME on the (a) humanoid-run and (h) dog-run tasks. The baaded policy reaches higher return (a) and converges faster.
4.3. DIME: A Practical Diffusion RL Algorithm
diffusion models, stochastic optimization of [()) does not need access to samples ap exp Qa/Zg, instead relying on stochastic gradients obtained via reparameterizzaioo trick (Kingma, 2013) using samples from the diffusion model n!".
To obtain a practical algorithm. we use a parameterized function approximation for the Q-function and the policy, that is, Qo and 7", with parameters & and 0, respectively. Here, 88 is represented by parameterized score network, see Eq. 10. To perform approximate policy evaluation, we can minimize the Bellman residual,
4.4. Implementation Details
Autotuning Temperature. We follow implementations like SAC (Haarnoja et al., 2018c) where the reward scaling parameter (also see Fig. 1) is not absorbed into the reward but scales the entropy term. Choosing a depends on the reward ranges and the dimensionality of the action space, which requires tuning it per environment. We instead follow prior works (Haarnoja et al., 2018c) for auto-tuning a by optimizing
(o(6) = _ (Qo(se,a)) Quaga (s,, a))) (24)
using stochastic gradients with respect to Φ. We provide implementation details in Section 4.4. Moreover, the expec- tation is computed using state-action pairs collected from environment interactions and saved in replay buffer. For policy improvement, we solve the approximate inference problem
J(a)= = a (Hurgat (x))
(28)
(())= = DKL ( (s) N
(25)
where Hurget is a target value for the mismatch between the noising and denoising processes measured by the log ratio.
where the target policy, i.e. the marginal of the noising process in Eq. 13 is given by the approximate Q-function
Autotuning Diffusion Coefficient. Please note that the objective function in Eq. 27 is fully differentiable with respect to parameters of the diffusion process. As such, additionally treat the diffusion coefficient / as a learnable parameter that is optimized end-to-end, further reducing the need for manual hyperparameter tuning. Further details on the parameterization can be found in Appendices D and G.
exp (s).aa))) (s(s)
(a 1s)
(26)
where states are again sampled from replay buffer. Further expanding C(0) yields
Q-function. Following Bhatt et al. (2024) we adopt the CrossQ algorithm, i.c., we use Batch Renormalization the Q-function and avoid a target network for calculating Quarget- When updating the Q-function, the values for the current and next state-action pairs are queried in parallel. The next Q-values are used as target values where the gradients are stopped. Additionally, we employ distributional Q learning as proposed by (Bellemare et al., 2017). The details are described in Appendix D.
((0)= =B,0 log (a (s) a))
(27)
n (a" 1- 1aa0
log
log (8),
Tn (a a" 11
showing that 2, is not needed to minimize Eq. 27 as it is independent of 0. Moreover, contrary to the score-matching objective (see Eq. 9) that is commonly used to optimize
6

---

DIME: Diffusion- Biflusion-Baaed Maximum Entropy Reinforcement Learning
32
16
OHnikeg-A DIO
0020
600
Mon Ronvine in Ulous
6,000
10M Maan Rottens (x.1M)
1DM Meas Rotan
10-
1QM Betum
400
4,000
200
2,000
0.5
1
1.5
2
2.5
15
0.2
0.4
0.6
0.8
0
0.2
0.4
0.0
0.5
Nanber Internations 110"
Number Diffasien Steps
Number Inseactions
-10"
Numbee Eas
100
(a) Varying the Diffusion Steps
(b) Runtime for 1M Steps
(c) Ant-v3
(d) Humanoid-v2 v3
Figure 3. Varying the Number of diffusion steps (a)-(b), The number of diffusion steps might affect the performance and the computation
time. (a) shows DIME's learning curves for varying diffusion steps. Two diffusion steps perform badly, whereas four and eight diffusion steps perform similar but still worse than 16 and 32 diffusion steps which perform similarly. (b) shows the computation time for IMIO steps of the corresponding learning curves. The smaller the diffusion steps, the less computation time is required. Learning Curves on Gym Benchmark Suite (c)-(d). We compare DIME against various diffusion baselines and CrossQ on the (c) Ant-v3 and (d) Humanoid-v3 from the Gym suite. While all diffusion-based methods are outperformed by DIME, DIME performs par with CrossQ the Ant environment. DIME performs favorably on the high-dimensional Humanoid-v3 environment, where it also outperforms CrossQ.
5. Experiments
5.1. Ablation Studies
We analyze DIME's algorithmic features with an intensive ablation study where we clarify the role of the reward scaling parameter a, the effect of varying diffusion steps, and the gained performance boost when using diffusion policy representation over a Gaussian representation. Additional analysis on employing distributional Q learning is discussed in the Appendix G.
Exploration Control. The parameter a balances the sxploration-explotataii trade-off by scaling the reward sig- nal. We analyze the effect of this parameter by comparing DIME's learning curves with different a values on the dog- run task from the DMC (see Fig. 2a). Additionally, we show the performance of the last return measurements for each learning curve in Fig. 2b. Too high a values (a = 0.1 do not incentivize maximizing the task's return, leading to no learning at all, whereas small values (a < 10 10-5) lead to suboptimal performance because the policy does not explore sufficiently. We can also see a clear trend that starting from a = 10-12, the performance gradually increases until the best performance is reached for a 10-3
In broad range of 13 sophisticated learning environments from different benchmark suits, ranging from mujoco gym (Brockman et al., 2016), deepmind control suit (DMC) (Tun- yasuvunakool et al., 2020), and myo suite (Caggiano et al., 2022), we compare DIME's performance against state-of- the-art RL baselines that employ diffusion and Gaussian policy parameterizations. The considered environments are challenging locomotion and manipulation learning tasks with up to 39-dimensional action and 223-dimensional ob- servation spaces.
Diffusion Policy Benefit. We aim analyze the perfor- mance benefits of the iiffusioonpprammeteriee policy com- pared to a Gaussian parameterization in the same setup by only exchanging the policy and the corresponding policy update. This comparison ensures that the Gaussian policy is trained with the identical implementation details from DIME described in Sec. 4.4 and showcases the perfor- mance benefits of diffusion- based policy. Fig. 2c and 2d show the learning curves of both versions on DMC's humanoid-run and dog-run environments. The diffusion policy's 's expressivity leads to a higher aggregated return in the humanoid run and to significantly faster convergence in the high-dimensional dog-run task. We attribute this performance benefit to an improved exploration behavior.
We consider QSM (Psenka et al., 2024), Diffusion-QI. (Wang et al., 2023), Consistency-AC (Ding & Jin, 2024), DIPO (Yang et al., 2023), QVPO (Ding et al., 2024), and DACER (Wang et al., 2024) as baselines for diffusion-based policy representations.
Additionally, we compare against the state-of-the-art RL methods CrossQ (Bhatt et al., 2024) and BRO (Nauman et al., 2024), where we have used the provided learning curves from the latter. Both methods use a Gaussian param- eterized policy and have shown remarkable results.
Number of Diffusion Steps. The number of diffusion steps determines how accurately the stochastic differential equa- tions are simulated and is a hyperparameter that affects the performance. Usually, the higher the number of diffusion steps the better the model performs at the burden of higher computational costs. In Fig. 3a we plot DIME's perfor-
We have run the learning curves for 10 seeds using the offi- cial code releases and report the interquartile mean (IQM) with a 95% stratified bootstrap confidence interval as sug- gested by Agarwal et al. (2021).
7

---

DIME: Diffusion- Crndudio--Basdd Maximum Entropy Reinforcement Learning
(00
1000
1,000
1,000
10M Mea Returs
800
RBM Moaa MeaM Rasuan
600
10M Mean
6.0
0.00
600
200
400
300
10M
400
200
200
0
0
0
0.2
0.2
0.4
10.6
0.8
0
0.2
0.4
10.5
0.88
0.2
0.44
0.5
0.8
Number Em Internctions
Number Emn Interactions
-10
Number Ene Interactions
Number Env Interactions
(a) Dog Run
(b) Dog Trot
(c) Dog Walk
(d) Dog Stand
1,000
1,000
DIME (ours)
10M Meun Retum
10M Mean Ketum
8000
1QM Meun-ettu Retum
800
BRO
200
600
600
BRO (Fast)
400
400
CrossQ
100
QSM
200
200
Dift-QI.
Consistency-AG
0.2
0.4
0.5
0.8
0.9
0.44
0.6
0.8
Number Ent Internations
-10
Number Env Interactions -10F
Number Env Interactions -10
DIPO
(e) Humanoid Run
(1) Humanoid Walk
(g) Humanood Stand
R0M Nocans Rae
0.8
RgM1 Success Rate
0.8
XQM Socaras Rate
0.8
XQM Success Rac
0.8
0.6
0..
10.6-
00.0
0.4-
0.4
0.2
0.2
0.2-
0.2
0-
0
0.2
0-4
0.6
0.8
0.2
0.4
0.6
0.8
0
0.2
0-4
0.6
0.8
0
0.2
0.6
0.8
Number Interactious
-10P
Number Inseractions -11f
Number Eive Enteractions -10P
Number Eir
.11
(h) Object Hold Hard
(i) Reach Hard
i Key Turn Hard
(k) Pen Twirl Hard
Figure 4. Training curves on DMC's dog, humanoid tasks, and the hand environments from the MYO Suite. DIME performs favorably on the nigh-dimensional dog tasks, where it significantly outperforms all baselines (dog-run) converges faster to the final performance. On the humanoid tasks, DIME outperforms all diffusion- based baselines, CrossQ and BRO Fast, and performs on par with BRO the humanoid-stane task and slightly worse on the humanoid-rum and humanoid walk tasks. In the MYO SUITE environments, DIME performs consistently on all tasks, either outperforming the baselines or performing on par.
mance for varying diffusion steps on DMC's humanoid-rum environment and report the corresponding runtimes for Mio environment steps in Fig. 3b on an Nvidia A 100 GPU machine. With an increasing number of diffusion steps, the performance and runtime increases. However, from 16 diffusion steps on, the performance stays the same.
Dog and Humanoid Tasks (Fig. 4). We bench- mark on DMC suit's 's challenging dog and humanoid envi- ronments, where we additionally consider BRO and BRO policy baseline. BRO Fast is iden-
DMC:
Fast as a Gaussian tical to BRO but differs only in the update-to-data (UTD) ratio of two as DIME and CrossQ. Please note that we used the online available learning curves provided by the official implementation for BRO. DIME outperforms all baselines significantly on the dog-run environment and converges faster to the same end performance on the remaining dog environments (see Fig. 4a - 4d). BRO has slightly higher average performance on the humanoid- run and humanoid- walk (see Fig. 4f-444 4c)) tasks indicating that DIME performs favorably on more nigh-dimensiooal tasks like the dog envi- ronments and tasks from the myo suite. However, DIME's asymptotic behavior in the humanoid -rnn achieves slightly higher aggregated performance than BRO, where we have run both algorithms for 3M steps (Fig. 6c). However, BRO requires full parameter resets leading to performance drops
5.2. Performance Comparisons
We consider environments with high- nigh-dimensional observa- tion and action spaces from three benchmark suits for robust performance assessment (please Appendix C).
Gym Environments. Fig 3c and Fig. 3d show the learning curves for the An-/v3 and Humanoid-v3 tasks respectively. While the diffusion-based baselines perform reasonably well on the Ant-v3 task with DIPO outperforming the rest, they are all outperformed by DIME and CrossQ which perform comparably. On the Humanoid-v3 DIME achieves signifi- cantly higher return than all baselines.
8

---

DIME: Diffusion- Diflusion-Baaed Maximum Entropy Reinforcement Learning
during training and is run with UTD ratio of 10 which is 5 times higher than DIME. This leads to longer training times. As reported in their paper (Nauman et al., 2024), BRO needs an average training time of 8.5h, whereas DIME trains in approximately 4.5h with 16 diffusion steps on the humanoid- run with the same hardware (Nvidia A A100).
DFG Collaborative Research Center 1574, Circular Factory for the Perpetual Product, and by the pilot program Core Informatics of the Helmholtz Association (HGF). This re- search has been additionally supported by the DFG Emmy Noether project CH 2676/1-1. This research was also sup- ported by the research cluster "Third Wave of AI", funded by the excellence program of the Hessian Ministry of Higher Education, Science, Research and the Arts, hessian
MYO Suite (Fig. 4). Except for pen twirl hard (Fig. 4k), DIME consistently outperforms BRO and BRO Fast in that it converges to a higher or faster end success rate. DIME also consistently outperforms CrossQ in terms of the achieved success rates on all the tasks except for the object hold hard task 4h, where DIME converges faster.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
6. Conclusion and Future Work
In this work, we introduced DIME, method for learning diffusion models for maximum entropy reinforcement learn- ing by leveraging connections to approximate inference. We view this work as starting point for exciting future research. Specifically, we explored denoising diffusion models, where the forward process follows an Ornstein-Uhhee Uhlenbeck process. However, approximate inference with diffusion models is an active and rapidly evolving field, with numerous recent advancements that consider alternative stochastic processes. For example, Richter& & Berner proposed learning both the forward and backward processes, while Nusken et al. (2024) further enhanced exploration by incorporating the gradient of the target density into the diffusion process. Additionally, Chen et al. (2024) combined learned diffusion models with Sequential Monte Carlo (Del Moral et al., 2006), resulting in highly effective inference method. These approaches hold significant promise for further improving ceffusiio-bbsedd policies in RL. We have conducted preliminary experiments on the framework from Richter & Berner and provide them in Appendix F. Finally, we note that the loss function used in this work (see Eq. 25) based on the Hellback-Leiblle divergence. However, in principle, any divergence could be used. For instance, the log- variance divergence (Richter & Berner) has shown promising results in optimizing diffusion models for approximate inference (Chen et al., 2024; Noble et al., 2024). Exploring alternative objectives could lead to additional performance improvements. Another interest- ing future research lies in investigating the effects of using more sophisticated critic structures, such as transformers, as proposed by Li et al. (2025).
References
Agakov, F. V. and Barber, D. An auxiliary variational method. In Neural Information Processing: 11th In- ternational Conference, ICONIP 2004, Calcutta, India, November 22-25, 2004. Proceedings 11, pp. 561-566. Springer, 2004.
Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A. C., and Bellemare, M. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural in- formation processing systems, 34:29304-29320, 2021.
Akhound-Sadegh, T., Rector-Brooks, J., Bose, A. J., Mittal, S., Lemos, P., Liu, C.-H., Sendera, M., Ravanbakhsh, S., Gidel, G., Bengio, Y., et al. Iterated denoising energy matching for sampling from boltzmann densities. arXiv preprint arXiv:2402.06111, 2024.
Anderson, B. D. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313- 326, 1982.
Arenz, O., Neumann, G., and Zhong, M. Efficient gradient- free variational inference using policy search. In Inter- national conference on machine learning, pp. 234 PMLR, 2018
Bellemare, M. G., Dabney, W., and Munos, R. A distribu- tional perspective on reinforcement learning. In Inter- national conference on machine learning, pp. 449-458. PMLR, 2017.
Acknowledgements
Berner, J., Richter, L., and Ullrich, K. An optimal con- trol perspective on diffusion-based generative modeling. Transactions on Machine Learning Research.
The authors acknowledge support from the state of Baden- Württemberg through the HoreKa supercomputer funded by the Ministry of Science, Research and the Arts Baden- Württemberg and by the German Federal Ministry of Edu- cation and Research. This work has been supported by the
Bhatt, A., Palenicek, D.,, Belousov, B., Argus, M., Ami- ranashvili, A., Brox, T., and Peters, J. Crossq: Batch normalization in deep reinforcement learning for greater

---

DIME: Diffusion- Balled Maximum Entropy Reinforcement Learning
sample efficiency and simplicity. In The Twelfth Interna- tional Conference on Learning Representations, 2024.
Ding, Z. and Jin, C. Consistency models as rich and ef- ficient policy class for reinforcement learning. In The Twelfth International Conference on Learning Represen- tations, 2024.
Blessing, D., Berner, J., Richter, L., and Neumann, G. Un- derdamped diffusion bridges with applications to sam- pling. In The Thirteenth International Conference Learning Representations, 2025a.
Doucet, A.,. Grathwohl, W., Matthews, A. G., and Strath- mann, H. Score-based diffusion meets annealed impor- tance sampling. Advances in Neural Information Process- ing Systems, 35:21482-2141 2022.
Blessing, D., Jia, X., and Neumann, G. End-to-end learning of gaussian mixture priors for diffusion sampler. In The Thirteenth International Conference Learning Rep- resentations, 2025b. URL https: /ttp:///pppeeee.eeeemmmm bib
Fang, L., Liu, R., Zhang, J., Wang, W., and Jing, B. Diffu- sion actor-critic: Formulating constrained policy iteration as diffusion noise regression for offline reinforcement learning. In The Thirteenth International Conference on Learning Representations, 2025.
Brockman, G., Cheung, V., Pettersson, L., Schneider, J.,
Schulman, J., Tang, J., and Zaremba, W. Openai gym, 2016.
Geffner, T. and Domke, J. Langevin diffusion variational inference. In International Conference on Artificial Intel- ligence and Statistics, pp. 576-593. PMLR, 2023.
Caggiano, V., Wang, H., Durandau, G., Sartori, M., and Kumar, V. Myosuite- a contact-rich simulation suite for musculoskeletal motor control, 2022. arXiv preprint arXiv:2205.00S: 000588..
Haarojja, T., Tang, H., Abbeel, P., and Levine, S. Rein- forcement learning with deep energy- policies. In International conference on machine learning, pp. 1352- 1361. PMLR, 2017.
Chen, H., Lu, C., Ying, C., Su, H., and Zhu, J. Offline rein- forcement learning via high-fidelity generative behavior modeling. In The Eleventh International Conference on Learning Representations, 2023.
Haarnoja, T., Hartikainen, K., Abbeel, P., and Levine, S. Latent space policies for hierarchical reinforcement learn- ing. In International Conference on Machine Learning, pp. 1851-1860. PMLR, 2018a.
Chen, J., Richter, I.., Berner, J., Blessing, D., Neumann, G., and Anandkumar, A. Sequential controlled langevin diffusions. arXiv preprint arXiv:2412.0708) 2024.
Chen, J., Richter, L., Berner, J., Blessing, D., Neumann, G., and Anandkumar, A. Sequential controlled langevin diffusions. In The Thirteenth International Conference on Learning Representations, 2025. URL https:// openreview net /amee/foornni ? id=dImD2sgy8s
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforce- ment learning with a stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018b.
Chi, C., Feng, S., Du, Y., Xu, Z., Cousineau, E., Burchfiel, B., and Song, S. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
Haarnoja, T., Zhou, A., Hartikainen, K., Tucker, G., Ha, S., Tan, J., Kumar, V., Zhu, H., Gupta, A., Abbeel, P., et al. Soft actor-critic algorithms and applications. arXiv preprint 1 1812.0595, 2018c
Cover, T. M. Elements of information theory. John Wiley & Sons, 1999.
Hansen-Estruce P., Kostrikov, I., Janner, M., Kuba, J. G., and Levine, S. Idql: Implicit q-learning as actor- critic method with diffusion policies. arXiv preprint arXiv:2304.105T 2023.
Dai Pra, P. A stochastic control approach to reciprocal dif- fusion processes. Applied mathematics and Optimization, 23(1):313-32, 1991.
Haussmann, U. G. and Pardoux, E. Time reversal of diffu- sions. The Annals of Probability, pp. 1 1188-1205, 1986.
Del Moral, P., Doucet, A., and Jasra, A. Sequential monte carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):411 1-436, 2006.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840-6833 6851, 2020.
Ding, S., Hu, K., Zhang, 7., Ren, K., Zhang, W., Yu, J., Wang, J., and Shi, Y. Diffusion-basee reinforcement learning via q weighted variational policy optimization. In The Thirty-eighth Annual Conference on Neural Infor- mation Processing Systems, 2024.
Huang, J., Jiao, Y., Kang, L., Liao, X., Liu, J., and Liu, Y. Schrödinger-föllle föllmer sampler: sampling without ergodic- ity. arXiv preprint arXiv:2 106. 10880, 1, 2021.
10

---

DIME: Diffusion-Basse Maximum Entropy Reinforcement Learning
Hyvärinen, A. and Dayan, P. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Mao, L., Xu, H., Zhan, X., Zhang, W., and Zhang, A. Diffusion-dice: In-sample diffusion guidance for offline reinforcement learning. In The Thirty-eighth Annual Con- ference on Neural Information Processing Systems, 2024.
Ishfaq, H., Wang, G., Islam, S. N., and Precup. D. Langevin soft actor-critic: Efficient exploration through uncertainty-driven critic learning. In The Thirteenth International Conference on Learning Representations, 2025, URL httnttps://opeereeee http///pppeeeeeeeeeeeeee/ee net////wweeeeeeeeeeeeeeeeeee forum? id-FvOsk3lal7.
Messaoud, S., Mokeddem, B., Xue, Z., Pang, L.,, An, B., Chen, H., and Chawla, S. S 2 ac: Energy-based reinforce- ment learning with stein soft actor critic. In The Twelfth International Conference on Learning Representations.
Nauman, M. and Cygan, M. On the theory of risk-aware agents: Bridging actor-critic and economics. In ICML 2024 Workshop: Aligning Reinforcement Learning Exper- imentalists and Theorists, 2023.
Janner, M., Du, Y., Tenenbaum, J., and Levine, S. Planning with diffusion for flexible behavior synthesis. In Interna- tional Conference on Machine Learning, pp. 9902-9915. PMLR, 2022.
Nauman, M., Ostaszewski, M., Jankowski, K., Mitoš, P., and Cygan, M. Bigger, regularized, optimistic: scaling for compute and sample efficient continuous control. The Thirty-eight Annual Conference on Neural Information Processing Systems, 2024.
Kang, B., Ma, X., Du, C.,, Pang, T., and Yan, S. Efficient diffusion policies for offline reinforcement learning. Ad- vances in Neural Information Processing Systems, 36, 2023.
Nelson, E. Dynamical theories of Brownian motion, volume 101. Princeton university press, 2020.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35: 26565-26577, 2022.
Neu, G., Jonsson, A., and Gómez, A unified view of entropy-regularizer markov decision processes. arXiv preprint arXiv: 1705.077, 2017.
Kingma, D. P. Auto-encoding variational bayes. arXiv preprint arXiv:/312.66. 14, 2013.
Nikishin, E., Schwarzer, M., D'Oro, P., Bacon, P.-L., and Courville, A. The primacy bias in deep reinforcement learning. In International conference on machine learn- ing. pp. 16828-168477 PMLR, 2022.
Lange, S., Gabel, T., and Riedmiller, M. Batch reinforce- ment learning. In Reinforcement learning: oaat-oo-the- art, pp. 45-73. Springer, 2012.
Noble, M., Grenioux, L., Gabrié, M., and Durmus, A. O. Learned reference-based diffusion sampling for multi- modal distributions. arXiv preprint arXiv:2410.19449, 2024.
Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline rein- forcement learning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.01643, 2020.
Li, G., Tian, D., Zhou. H., Jiang, X., Lioutikov, R., and Neumann, G. TOP-ERL: Transformer-basee off-policy episodic reinforcement learning. The Thirteenth Inter- national Conference on Learning Representations, 2025.
Nusken, N., Vargas, F., Padhy, S., and Blessing, D. Trans- port meets variational inference: Controlled monte carlo diffusions. In The Twelfth International Conference on Learning Representations: ICLR 2024, 2024.
Psenka, M., Escontrela, A., Abbeel, P., and Ma, Y. Learning a diffusion model policy from rewards via q-score match- ing. In Forty-first International Conference on Machine Learning, 2024.
Li, Z.,, Krohn, R., Chen, T,, Ajay, A., Agrawal, P., and Chalvatzaki, G. Learning multimodal behaviors from scratch with diffusion policy gradient. arXiv preprint arXiv:2406.00681, 2024.
Ranganath, R., Tran, D., and Blei, D. Hierarchical varia- tional models. In International conference on machine learning, pp. 324-333, PMLR, 2016.
Lu, C., Chen, H., Chen, J., Su, H., Li, C., and Zhu, J. Con- trastive energy prediction for exact energy-guided diffu- sion sampling offline reinforcement learning. In Inter- national Conference on Machine Learning, pp. 22825- 22855, PMLR, 2023.
Reuss, M., Li, M., Jia, X., and Lioutikov, R. Goal con- ditioned imitation learning using score-based diffusion policies. Robotics: Science and Systems, 2023.
Maaløe, L.. Sønderby, C. K., Sønderby, S. K., and Winther, O. Auxiliary deep generative models. In International conference on machine learning, pp. 1445-1453. PMLR, 2016.
Richter, L. and Berner, J. Improved sampling via learned diffusions. In The Twelfth International Conference on Learning Representations.

---

DIME: Diffusion-Base Maximum Entropy Reinforcement Learning
Särkkä, S. and Solin, A. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Wang, Z., Hunt, J. and Zhou, M. Diffusion policies an expressive policy class for offline reinforcement learning. International Conference on Learning Representations, 2023.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International conference on machine learning. pp. 2256 2256-2265, PMLR, 2015.
Welling, M. and Teh, Y. W. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference machine learning (ICML-11), pp. 681-688. Citeseer, 2011.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A.,, Er- mon, S., and Poole, B. Score- basee generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
Yang. L., Huang, Z., Lei, F. h., Zhong, Y., Yang, Y., Fang, C.,
Wen, S., Zhou, B., and Lin, Z. Policy representation via diffusion probability model for reinforcement learning. arXiv preprint arXiv:2305. 13/22, 2023.
Sutton, R. S. and Barto, A. G. Reinforcement learning: An introduction. Robotica, 17(2):22--33 235, 1999.
Zhang, D.,, Chen, R. T., Liu, C.-H., Courville, A., and Ben- gio, Y. Diffusion generative flow samplers: Improving learning signals through partial trajectory optimization. arXiv preprint arXiv:2330.00000 10.02679, 2023.
Toussaint, M. Robot trajectory optimization using approxi- mate inference. In Proceedings of the 26th annual inter- national conference on machine learning, pp. 2009.
Zhang, Q. and Chen, Y. Path integral sampler: stochas- tic control approach for sampling. arXiv preprint arXi:: 2021.
Tran, D., Ranganath, R., and Blei, D. M. The varia- tional gaussian process. arXiv preprint 1XXi::511.00699: 2015.
Zhou, H., Blessing, D., Li, G., Celik, O., Jia, X., Neumann, G., and Lioutikov, Variational distillation of diffusion policies into mixture of experts. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL https: opwnreview.eee..cooooe httteeeiieeeeeieeeeeeeee forum? iidiiiiiiin iddiiddddgHHww....
Tunyasuvunakool, S., Muldal, A., Doron, Y., Liu, S., Bohez, S., Merel, J., Erez, T. Lillicrap, T., Heess, N., and Tassa, Y. dm.control: Software and tasks for continuous control. Software Impacts, 6:1 6:100022, 2020. ISSN 2665-9638.
Tzen, B. and Raginsky, M. Theoretical guarantees for sam- pling and inference in generative models with latent diffu- sions. In Conference on Learning Theory, pp. 3084-3114. PMLR, 2019.
Ziebart, B. D. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. Carnegie Mel- lon University, 2010.
Vargas, F., Grathwohl, W. S., and Doucet, A. Denoising diffusion samplers. In The Eleventh International Con- ference on Learning Representations.
Ziebart, B. D., Maas, A. L., Bagnell, J. A., Dey, A. K., et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pp. 1433-1438. Chicago, II., USA, 2008.
Vargas, F., Ovsianas, A., Fernandes, D., Girolami, M., Lawrence, N. D., and Nüsken, N. Bayesian learning via neural schrödinger-föölmer flows. Statistics and Com- puting, 33(1):3, 2023.
Vincent, A connection between score matching and de- noising autoencoders. Neural computation. 23(7):1661- 1674, 2011.
Wang, D. and Liu, Q. Learning to draw samples: With application to amortized mle for generative adversarial learning. arXiv preprint
2016.
Wang, Y., Wang, L., Jiang, Y., Zou, W., Liu, T., Song, X., Wang, W., Xiao, I.,, WU, J., Duan, J., and Li, S. E. Dif- fusion actor-critic with entropy regulator. In The Thirty- eighth Annual Conference on Neural Information Pro- cessing Systems, 2024. URL https:: openreview. net/ wtr.fformmmiceee d=10011 j4QvTq.
12

---

DIME: Diffusion-Bare Maximum Entropy Reinforcement Learning
A. Derivations
Lower-Bound Derivation. H(To((ols)) 2 (a
To:N (a n:N ss)
H(To(uolss) -EROON
log
(29)
T1:N1o( 8, (,
To:N (a::N (S)
a)
-EoO.N
log
#1:NIU( (al:N ( 1: 18, a9)
# 1I,Njo I:N (s, log
=B
log
+ ENO.
(30)
To:N (s)
(a I:N (a a0)
TI:Njo(a (s
log
ET [DLL (#INN o o |s,a"|) an)l FiiNoolaaii |s, ss,a
(31)
To:N(a
#I:vooaaiin (8,a%) No.N (a0:N11)
2 EFOIN
log
(32)
where we have used the relation
Fo.N(a.:M))
To(a) #1:N1o 1aN 1 a))
(33)
and the fact that the KI. divergence is always non-negative Approximate Inference Formulation. Recall the definition of the Q-function
Qa (st, ag) +
Σ 1=1
/ Ep,, [rees + Ex S(4)))
(34)
and
#I:Njo(a 1:N la",
EFO.N
log
(35)
#o:N (a0:N (s)
We start reformulating the objective
0
J(H) 2 J(H)
t 1- Ep, [re+ + (x(a).
(36)
7 "''Epp Ir, E=
(a(aafss)) st] + Ep- ( + E= st)
(37)
t=l+1
Ep' [Q (st, [ + E (af s1]]
(38)
Ep [Q* (st, a.) (a [
(39)
(a 1:N (a" s) E.p,, Fo.N **($0, ab) + log #o. (ao:N 1s)
(40)
Ep [DRI. (#(a":N\b))) 7(au:N
U.N (s))
log 2 (s)]
(41)
where we used
cxp Q (s,
no (a)
(42)
z*(s)
in the last step. When minimizing, the negative sign in front the KL vanishes. Please note that the expectation over the marginal state distribution was ommited in the main text to avoid cluttered notation.
13

---

DIME: Diffusion-Basse Maximum Entropy Reinforcement Learning
B. Proofs
Proof of Proposition 4.1 (Policy Evaluation). Let's define the entropy-augmented reward a diffusion policy as
#1.vioaa::1o Ja", To.N (a4:N(s)
(sc, ag _ ri (st. + BaaN log
(43)
and the update rule for the Q-function
Q(s a.) t (.) + Be [(+))1
(44)
This formulation allows to apply the standard convergence results for policy evaluation stated in (Sutton & Barto, 1999).
Proof of Proposition 4.2 (Policy Improvement). It holds that
(a = exp exp 1 (s,a
7100 (a
(45)
Moreover, using the fact that the KL divergence is always non-negative, we obtain
0 DKL. (# U:N
(1 1s)
DKI.
(46)
Rewriting the KI. divergences yields
log x(++1) (a0:0 (8)
log 7(0) (a (s)
E B((+1)
< Eycn
(47)
(1+1) (1):
f(++1) (au:N1s)
E 11))
log # (i+1) (1):N(s)
[s)
E 11G1
log म (i+t) 0:N (s)
(48)
< Egto
log #x (1) (a (s)
- Eggt) [log = (a::)ss)
log exp Qm (s, a
E ()0+
log (1++1) (a 0:N 1s)]
- Ez(++1
0:N 1 11a
(49)
Zall (s)
log exp expQ" QT (s,a~))
< Ege
log (0 (a:N(s))))) - Eggo
7(1)
0:N 1a s)
Z=
log 7(1) (a.. ,8)
Bgi4 [Qmc (0)0((88a)) (x + Eg(ce1)
(50)
#(++1) (auN
log #0 IN (8))
(s, ( + Eyco)
2 Egto
(
To keep the notation uncluttered
+(1) (a (a ss) and (s, Egco
7(0) (a -
| s)
(s, =Ey0+1) log #(i+1) (a.:011)))))))))
log # (n0: Is)
(51)
14

---

DIME: Diffusion- Benuusion-Based Maximum Entropy Reinforcement Learning
Figure 5. Considered environments. The Humanoid-v3 and the Ant-v3 are environments from the mujoco gym benehmark (Brockman al., 2016). The three anvironmentshumanid www.aanooddd -wakk and humanoid-smne -ttand are from the deepmind control suite (DMC) benchmark (Tunyasuvunakool al., 2020). The dog enviromments consist of dog-run, dog-walk. dog-stond dog-trot and are also from the DMC sutie benchmark. Finally, the myo suite hand environments object-hold hholhhhaddddeee reach- key-turn-han hard. pen- pen-twirl- ter--wirll-haa hard are from the myo suite (Caggiano al., 2022).
@n (s, ) = ro + E 7 (200) (s1, af
Reo [ort (s1,
(52)
S To + E 7 ( ($), ((ttt) ai 1 (Q=Th) (s1, afnnaa))))
(53)
(d(tt)) aj t + (d2) a²) + Egg [ (82, a afaii]]
= oo + E
(54)
S ro + E
7 (s)
+ 2²
(1itt) a²² + Eg(ctD) (Q (s2, ag² }]]
(55)
(56)
< ro so EE
Σ
(dlit)) at +m))
Q (s,a*)
(57)
Since Q improves monotonically, we eventually reach fixed point Q(++1) = -Q(1 Q
Proof of Proposition 4.3 (Policy Iteration). From Proposition 4.2 it follows that 2*th
(s,a) 2 Q
(s.a). If for
lim # # then it must hold that (=* 2 (*(ssa) for all म fl which is guaranteed by Proposition 4.2.
C. Environments
All environments are visualized in Fig. 5. We consider the Ant-v3 and the Humanoid-v3 environments from mujoco gym (Brockman et al., 2016). The humanoid-stanee humanoid-walk humanoid-run, dog-stand, dog-walk, dog-tre -rrgttrt and dog-run environments from the deepmind control suite (DMC) (Tunyasuvunakool al., 2020). The hand environments from myo suite the Mhjjetthhoddrandon wrando,,reachrraeee reach-random, key-turn-random and ren-tworl-random environments (Caggiano et al., 2022). The action and observation spaces of the respective environments are shown in Table 1.
D. Implementation Details
We consider score network with 3 layers and a 256 dimensional hidden layer with gelu activation function. We use Fourier features to encode the timestep and scale the embedding using a feed-forward neural network with two layers, with hidden dimension of 256. For the diffusion coefficient, we use cosine schedule and additionally optimize a scaling parameter for the diffusion coefficient per dimension end-to-end (i.e.. learn the parameter 3 (please Appendix G).
We employ distributional Q following (Bellemare et al., 2017), where the Q-model outputs probabilities y over bins. Using the bellman backup operator for diffiusion models from Eq. 21 and the bin values we follow (Bellemare et al., 2017) and calculate the target probabilities Starget- Using the entropy-regularinn
cross-entropy loss L(ő) = -- Starget log 9e 0.005 Σ 9e log qs we update the parameters ó of the f-function. Please note that the entropy regularization was not proposed in the original paper from (Bellemare et al., 2017), however, we noticed that small regularization helps improve
15

---

DIME: Diffusion-Basse Maximum Entropy Reinforcement Learning
Training Environment
Observation Space Dim.
Action Space Dim.
Ant-v3
111
Humanoid-v3
376
17
dog-run
223
38
dog-walk
223
38
dog-trot
223
38
dog-stand
223
38
humanoid-run
67
24
humanoid-wall
67
24
humanoid-stann
67
24
myoHamd(OOHHolldaaddoonn
91
39
myoHaandReaciRaaddon-v0
115
39
93
39
myoHaadPea/Tiiikaaddmmo00
83
39
Table Observation and Action Space Dimensions for Various Training Environments
the performance in the early learning stages but does not change the asymptotic performance. Additionally, follow (Nauman et al., 2024) and use the mean of the two Q-values instead of the min it has usually been used in RL so far.
The expected Q-values for updating the actor can be easily calculated using the expectation Q(s,a,) = Σ; (s af)6;
Action Scaling. Practical applications have a bounded action space that can usually scaled to fixed range. However, the action range of the diffusion policy unbounded. Therefore, we follow recent works (Haarnoja et al., 201 and propose applying the change of variables with tanh squashing function the last diffusion step n O. For the backward process squash the action a tanh uD such that a (-1, 1) and its
To.N (N (s) with unbounded action space R density is given by
d du
No.N (a O.N (aN To.N (u det
(58)
with the corresponding log-like ligliike
N
D
0.N log Fo.N (a (s) = log 7N (u + >
log Tn 1a--(uu=11pp lum;,)
- log (1 - tanh² (u ))
(59)
n=1
This means that the Gaussian kernels of the diffusion chain have the same log probabilities except for the correction term of the last step at 0.
Algorithm DIME: Diffusion Befdeeioo-BBased Maximum Entropy Reinforcement Learning
Input: Initialized parameters 6, &, a, learningrates À
1: for M do
2: i % UTD then
3: G.T The (NN (st)
4:
St+1 p(s la st)
5:
D - Dst ab re, 81+1}
6:
end if'
7:
Φ + δ - A V Jo ( (Eq. 24)
8:
if POLICYDELAY then
9:
0 0 V (() (Eq. 25) a X, (a) (Eq. 28)
11: end if
12: end for
16

---

DIME: Diffusion- Biflusion-Based Maximum Entropy Reinforcement Learning
Algorithm shows the learning procedure of DIME. Note that policy delay refers the number of delayed updates of the policy compared to the critic. UTD is the update to data ratio.
E. List of Hyperparameters
Please note that we have used the official code releases of the respective baseline methods for training. For BRO and BRO Fast we used the provided learning curves
DIME
QSM
Diff-QL
Consistency-AC
DIPO
DACER
QVPO
Update-to-data ratio
2
Discount
0.99
0.99
0.99
0.99
0.99
0.99
0.99
batch size
256
256
256
256
256
256
256
Buffer size
lc6
le6
1e5
le5
1e6
1c6
le6
Htarget
4dim((A)
N/A
N/A
N/A
N/A
-0.9dimA
N/A
Critic hidden depth
2
2
2
3
3
3
2
Critic hidden size
2048
2048
256
256
256
256
256
Actor/Score depth
3
3
4
3
Actor/Score size
256
256
256
256
256
256
256
Num. Bins/Quantiles
100
N/A
N/A
N/A
N/A
2
N/A
Temp. Learn. Rate
1e-3
N/A
N/A
N/A
N/A
3e-2
N/A
Learn. Rate Critic
3e-4
3e-4
3e-4
3e-4
3e-4
3e-4
3e-4
Learn Rate Actor/Score
3e-4
3e-4
le-5
le-5
3e-4
3e-4
3e-4
Optimizer
Adam
Adam
Adam
Adam
Adam
Adam
Adam
Diffusion Steps
16
15
5
N/A
100
20
20
Prior Distr.
(0, 2.5)
(O, 1)
N/A
N/A
N/A
N( 1)
N0, 1)
Exploration Steps
5000
le4
le4
le4
lc4
lc4
le4
Score-Q Qaree-Que align. factor
N/A
50
N/A
N/A
N/A
N/A
N/A
Table Hyperparameters of DIME and all diffusion-based algorithms for all benchmark suits. Varying hyperparameters for different henchmark suits descrihed in the text.
DIME
BRO
BRO Fast
Cross
Polyak weight
N/A
0.005
0.005
N/A
Update-to-data ratio
Discount
0.99
0.99
0.99
0.99
batch size
256
128
128
256
Buffer size
le6
1e6
1e6
1e6
Hearget
4dim(A) (A)
dim(AA)2
dim(A)/2
dim(A)
Critic hidden depth
2
BRONET
BRONET
2
Critic hidden size
2048
512
512
2048
Actor/Score depth
3
BRONET
BRONET
3
Actor/Score size
256
256
256
256
Num. Bins/Quantiles
100
100
100
N/A
Temp. Leamn Rate
1e-3
3e-4
3e-4
3e-4
Learn. Rate Critic
3e-4
3e-4
3e-4
7e-4
Learn. Rate Actor/Score
3e-4
3e-4
3e-4
7e-4
Optimizer
Adam
Adam W
AdamW W
Adam
Diffusion Steps
N/A
N/A
N/A
Prior Distr. Exploration Steps
(0, 2.5)
N/A
N/A
N/A
5000
2500
2500
5000
Table 3. Hyperparameters of DIME and Gaussian-bass based algorithms for all benchmark suits. Varying hyperparameters for different benchmark suits are described the text.
17

---

DIME: Diffusion-Basse Maximum Entropy Reinforcement Learning
DIME. For DIME, we use distributional Q. where the maximum and minimum values for the bins have been chosen per benchmark suite. We have used Umin 1600 and Urar 1600 for the gym environments, Unin 200 and Umaa 200 for the DMC suite and Umin 3600 and Umaz = 3600 for the myo suite.
QSM. In certain environments, we observed that QSM with default hyperparameters performed poorly, particularly in several DMC tasks and the Gym Ant-v3 tasks. To address this, we fine-tuned the hyperparameters for QSM in each of these underperforming tasks. For the DMC tasks, we found that QSM often requires an a valre--rerssnnttnng the alignment factor between the score and the Q-function (Psenka al., 2024) in the range of 100-200, rather than the default value of 50 reported in QSM's original implementation. In the A3t-v3 task, we determined that a needs to be to In the original implementation, the number of diffusion steps set to be 5, however, we found using more steps, such 10 and 15, can significantly improve the performance in these under performed tasks.
CrossQ. We used the hyperparameters from the original paper (Bhatt et al., 2024) for the gym benchmark suite. However, we used a different set of hyperparameters for the DMC and MYO suites for improved performance. More precisely, we increased the policy size to 3 layers with 256 hidden size. Additionally, we reduced the learning rate to 7e-4.
F. General Diffusion Policies
DIME's maximum entropy reinforcement learning framework for training diffusion policies is not specifically restricted to denoising diffusion policies but can be extended to general diffusion policies. This can be realized using the General Bridges framework as presented in (Richter & Berner). In this case, we can write the forward and backward process as
da; = [f(a, t) + Bu(a,, t)] dt + 2B
ag #o(-ss),
(60)
da = [f t) Bula,, s, t)] dt + /23,dB.
oT ~ 1),,
(61)
: with the drift and control functions f, u x [0, T| + R.d the diffusion coefficient 8 [0, + R standard Brownian
motion ( (B, (Bo),t and some target policy up- Again we denote the marginal density of the forward process as म. and the conditional density at time given / as Fel for t, E [0, Tj. The backward process starts from Fr = का ~ NN I) and runs
backward in time where we denote its density as #.
The respective discretization using the Fuler Maruyama (EM) (Särkkä & Solin, 2019) method are given by
a a+ [f(a" n) + Bu(a", s, s,n) 5 + Cas
(62)
- -===-"~\f(0".n Bu(a", .taa""ss,m0 n)] 8 En.
(63)
where En, En ~ (o, 2331),, with the constant discretization step size 8 such that N = TT/ is an integer. We have used the simplified notation where we write a" instead of n5 The discretizations admit the joint distributions
N-1
To:N (n0:Ns)) (a(0)())) II
के a In+iin(a" (a"t"(a", 8),
(64)
n=0
N
No.N (1) #Nv(a") [8)
II
1/m #
(a 1- (a"=1("",,0) s)),
(65)
with Gaussian kernels
"a"" a", = ~Vaa*+1aa" + [f(a", n) + Bu(a", s, n, 2361) - (a--1 aa, _N(a"='la" [f(a", n) Bu(a", s, n)) 6, 2.37))
क 1 (a n
(66)
(67)
Following the same framework presented in the main text, we now optimize the controls and using the same objective
C(u, v) DxL (a 0:N [s) Ho.N (a
(68)
where the target policy at time step is given as
exp Q2 (= (s,a²) dd)
To (s)
(69)
2*
18

---

DIME: Diffusion- Biflusion-Based Maximum Entropy Reinforcement Learning
600
300
600
IQM Mean Retur
IQM Mean Return
400
200
IQM Retum
400
100
200
BRO
DIME
0-
0
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.5
Number Env Interactions
-10
Number Env Interactions -100
Number Env Interactions -10"
(a) DIME and GB Dog Run
(b) DIME and GB Humanoid Run
(c) DIME and BRO on Humanoid Run
Figure 6. Preliminary results for the GB sampler on the dog run (a) and humanoid run (b) environments from DMC. Comparison
BRO on the humanoid run for million steps.
1.8-
18-
1.5
1.6
1.6
1.
1.5
de
1QM Bo
10M
1QM
IQM
1.4
1.4
1.4
1.4
1.2
1.2
1.2
1.2
0.2
0.4
0.6
0.8
0
0.2
0.4
0.6
0.5
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.5
Number Eay Interactions 30
Number Internetions -10
Number Interactions 30
Number
-10
(a) 8o on the Dog Run
(b) 810 on the Dog Run
(c) 920 on the Dog Run
(d) 330 on the Dog Run
Figure Learned parameters. DIME's policy improvement objective (Eq. 27) allows to train various parameters end-to-end, such the scaling for the diffusion coefficient B. More concretely, we train a scaling parameter S& per dimension k that scales the cosine schedule. We visualize adaptation of the parameter for the dimension = 0, 10, 20,30 30 over the training, averaged over 10 seeds for the dog-run task. Clearly, DIME first increases the parameter the beginning of the training phase. Depending the dimension, either converges to rather high value (k- 20 and k = 30), or keeps being reduced for other dimensions = 0 and t 10.
In practice, optimize the control functions and using parameterized neural networks. We have run preliminary results using the general bridge framework within the maximum entropy objective suggested in work. The learning curves can be seen in Fig. 6.
G. Additional Experiments
End-To-End Learning of B. We visualize the adaptation of the scaling for the diffusion coefficient in Fig. 7 during learning DMC's dog-run environment.
Extended Analysis Distributional Q Learning. DIME employs distributional Q Learning (Bellemare et al., 2017) to represent the Q-function as a distribution over bins. We compare DIME to baselines when using distributional Q Learning and when using the well-known Bellman residual (see Eq. 24) for updating the parameters of the Q-function.
We start by comparing DIME with distributional Q learning against diffusion-based baselines that employ distributional Q learning. Fig. 8a and Fig. show the learning curves on the Ant-v3 and Humanoid-v3, respectively, where we compare against DACER, distributional Q variant Diff- QL, and Consistency-AC. DIME converges faster to the same performance DACER the Ant-y3 task and outperforms the baselines the Humanoid-yā task. In the setting without distributional Q Learning, i.e., when updating the parameters using the residual Bellman function, DIME performs similarly to DIPO and QVPO the Ant-13 task and outperforms all baselines on the higher-dimensionall Humanoid-v3 task (Fig. 8c and Fig. 8d).
Additionally, we compare DIME with and without distributional Q Learning on the four dog environments from the DMC suite (Fig. 8), where we concentrate on the strong baselines BRO (Nauman et al., 2024) and CrossQ (Bhatt et al., 2024). BRO employs quantile distributional Q learning, whereas CrossQ uses the Bellman residual loss function for updating the
19

---

DIME: Diffusion- Sotuudion-Basdd Maximum Entropy Reinforcement Learning
(ouns)
DIME DistrQ (ouns)
QSM
DHIT-QI
Condisenny-ACC
QVea
RQM Mean Setum (xc10")
IQM Mean Retum (xtnn)
RQM Mean Remun
6,000
1.0
R0M Mean
4,000
4,000
0.5
2,000
2,000
0.2
0.4
10.5
0.8
0
0.2
0.4
0.6
0.8
0.2
0.6
11.6
0.8
0.4
0.6
0.8
Number
-10"
Number Elov Interactions -10P
Number
-10*
Number Livr Interactions
-10"
(a) Ant-v3
(b) Humanoid-v3
(c) Ant-v3 wo DistrQ
(d) Humanoid-v3- w/o DistrQ
Figure Comparison to Diffusion Baselines with (a)-b) and without Distributional Q (c)-d)) on Ant-v3 and Humanoid-v3 tasks. We provide the learning curves for distributional versions for Diff-QL and Consistency-AC alongside DACER, which employs distributional Q by default on the Ant-v3 (a) and Humanoid-v3 (b) tasks. DIME converges faster on Ant-v3 (a) task to the same performance achieved by DACER and outperforms all baselines on the more high-dimensional Humanoid-v3 (b) task. Additionally, we compare DIME without distributional Q against the diffusion baselines without distributional Q on the Ant-v3 (c) and Humanoid-v3 (d) tasks. DIME without distributional Q performs par with the baselines DIPO and QVPO the Ant-v3 (c) and outperforms all baselines on the Humanoid-v3 (d).
DIME (ours)
DIME w/o DistrQ
BRO
BRO (Fast)
CrossQ
600
1,000
1,000
R0M Moan Mean Rouum
TOM Mean Return
10M Moan Retum
800
800
400
600
600
600
200
IQM
400
200
0
0.2
0.2
0.4
0.5
0-
0..5
0.6
0.8
0-
Number Eny Interactions -10
Number Env Interactions
Number Eav Interactions -10
Number Lane Interactions
Dog Run
(b) Dog Trot
(c) Dog Walk
(d) Dog Stand
Figure 9. Ablation Distributional Q. Comparison of DIME and DIME without employing distributional Q (dashed line). While there a small improvement when using distributional Q, DIME w/o Distributional Q still performs on par, or better than BRO, which employs quantile distributional RL. DIME w/o DistrQ outperforms CrossQ and BRO (Fast).
Q-function's 's parameters. In the main text, we have already observed that DIME with distributional Q performs favorably over the baselines. Fig. shows small improvement when using distributional Q However, DIME without distributional Q (dashed line) still performs on par, better than BRO and consistently performs better than BRO (Fast) and CrossQ. Please note that BRO and BRO (Fast) employ quantile distributional RI. (Nauman et al., 2024).
20


Reviews:
[{"review_id": "FkVRljZLkW", "reviewer": "Reviewer_Etdt", "summary": "Authors have introduced the idea of using diffusion-based policy in maximum entropy RL. To tackle the issue of intractable entropy calculation of using a diffusion-based policy, the authors derive a computationally tractable lower bound on the maximum entropy objective. Through extensive experiments, the proposed algorithm Diffusion-based Maximum Entropy RL (DIME) is shown to outperform existing diffusion-based methods and state-of-the-art non-diffusion based methods.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised.\n \n---\n\n> Some ablations is reported on just one environment. It would be better to include all. For example, for hyper-parameter alpha and diffusion steps, are the optimal values the same across different environments? How sensitive are they? For results in figure 3, what is the algha value/diffusion steps used? No explicit reports on computation complexity (for example, wall-clock time, GPU, memory).\n>\n\nWe thank the reviewer for their comments and time. The sensitivity of the crucial hyperparameters $\\alpha$ and the number of diffusion steps are analyzed in Fig.1a-b and Fig.2a.\nFig.1 shows the sensitivity to $\\alpha$. As expected, too high and too low $\\alpha$ values lead to worse performance, while $\\alpha=0.001$ performs best. This trend aligns with previous maximum entropy RL papers [1], where tuning $\\alpha$ per environment is suggested. However, inspired by previous works [2] we propose optimizing $\\alpha$ for a target entropy (Section 4.4), eliminating the need for tuning $\\alpha$ for each environment. Please note that the target entropy depends on the action dimension of the considered environment and does not need to be tuned (please see App. E).\n\nSimilarly, the number of diffusion steps is task-depended, but generally, we do not observe a decrease in performance with more diffusion steps (Fig. 3a). However, a more complex task might need more diffusion steps. Our experiments found that using 16 diffusion steps performs well for all tasks. \n\nWhile we agree that evaluating a set of hyperparameters on all tasks leads to the optimal performance, this would also require a large amount of compute. We have therefore opted to optimize for the hyperparameters during learning, as done for $\\alpha$, or fix the number of diffusion steps to a reasonable value, which works well for most of the tasks.\n\nFig. 2 b plots the wall-clock time for the whole training process (including testing the policy) for 1M environment interactions. The plot shows that the higher the number of diffusion steps is, the longer the training takes, which makes sense given the increased computation. However, with an average training time of around 4.5 h with 16 diffusion steps DIME runs reasonably fast compared to high-performing baselines like BRO that need an average training time of 8.5h [3] on the same hardware. \n\n[1] Haarnoja, T. et al. Soft Actor Critic. ICML 2018\n\n[2] Haarnoja, T et al. Soft actor-critic algorithms and applications. 2018\n\n[3] Naumann et al. Bigger, regularized, optimistic. Neurips 2024.", "claims_and_evidence": "Main claims made in the submission are generally supported by clear evidence.", "methods_and_evaluation_criteria": "Yes, the proposed algorithm DIME is tested on various standard RL environments with 10 seeds and compared against different baselines. Ablation study on important hyper-parameters is also provided for analysis.", "theoretical_claims": "Theoretical part looks ok to me. Though theoretical part is a very important part of the paper. Math is not my strength so chairs may need to rely more on others' reviews.", "experimental_designs_or_analyses": "Some ablations is reported on just one environment. It would be better to include all. For example, for hyper-parameter alpha and diffusion steps, are the optimal values the same across different environments? How sensitive are they? \nNo explicit reports on computation complexity (for example, wall-clock time, GPU, memory). As you claimed that DIME is computational", "relation_to_broader_scientific_literature": "The paper extends the MaxEnt-RL literature by introducing diffusion policies as a more expressive policy class over traditional Gaussian policies.  Unlike previous diffusion-based methods, this paper uniquely addresses entropy maximization explicitly through a theoretically justified lower-bound entropy objective.", "essential_references_not_discussed": "The paper provides a strong background on diffusion-based RL and MaxEnt RL literature, citing the key relevant works.", "other_strengths_and_weaknesses": "See sections above.", "other_comments_or_suggestions": "See sections above.", "questions_for_authors": "As mentioned above, for hyper-parameter alpha and diffusion steps, are the optimal values the same across different environments? How sensitive are they? For results in figure 3, what is the algha value/diffusion steps used?\nAlso, any computational complexity analysis (wall-clock time, GPU)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2EV1Vvc0gR", "reviewer": "Reviewer_Xhj7", "summary": "This paper combines maximum entropy reinforcement learning with the diffusion model, leading to a coherent algorithm that addresses several critical problems that exist in previous methods. The main contribution of this paper is that it derives a tractable lower bound of the entropy of the action distribution for diffusion-based policies. Based on this foundation, this paper transforms the traditional MaxEnt RL framework to accommodate diffusion-based policies, which enhances the exploration ability as compared to other diffusion-based algorithms. The empirical evaluations are conducted on several complex continuous control tasks, and the proposed method, DIME, does demonstrate improved performance.\n\n## update after rebuttal\nMost of my concerns are resolved by the rebuttal, except for the cost of back-propagation through the diffusion chain. I would encourage the authors to consider addressing this limitation since this would be necessary to apply DIME to large-scale applications such as generative model fine-tuning. I will maintain my positive evaluation of this paper.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised.\n \n---\n\n> The comparison with baseline methods may not be strictly fair. For non-diffusion methods, hyper-parameters such as the target entropy, the width of the critic networks are not aligned with DIME. It would be better to include additional experiments to ablate the effect of these parameters to demonstrate the benefit of diffusion compared to Gaussian policies.\n> \n\nWe thank the reviewer for the comments and for their time. We agree that additional experiments that analyze various architecture designs are interesting. However, we would like to note that  Fig. 1c-d show the learning curves for DIME and a Gaussian policy under the same setup with the same critic and actor networks for 3M environment interactions until convergence. These experiments showcase the benefit of using a diffusion policy over a Gaussian policy.  \n\nAdditionally, we have tuned CrossQ in the benchmarks to obtain the best performance and have chosen the same or comparable networks as DIME. Please note, while DIME uses simple MLPs for the critic, BRO uses more complex network structures (BRONET) with residual connections, boosting the performance [1] and potentially providing a benefit against DIME. We believe that more sophisticated network structures might even help boosting DIME’s performance and consider this research as a promising field for future work.  Please note that the target entropy values for DIME need to vary compared to a Gaussian-based policy, as we consider the sum of log ratios of distributions (Eq.15) instead of the log probability of a single Gaussian as entropy. These instances can therefore not be directly compared and need other values as a target.\n\n[1] Naumann et al. Bigger, regularized, optimistic. Neurips 2024.\n \n---\n\n> Most of the relevant literature are properly cited. However, in the field of offline RL, there are a few articles that also employs diffusion-based policies, such as Diffusion-DICE [2], DAC [3], and QGPO [4]. The authors are encouraged to include a discussion about these works.\n> \n\nWe thank the reviewer for pointing us to recent works on diffusion-based offline RL. We will extend the related work accordingly in the final version. Could the reviewer please provide a reference to QGPO, as this is not listed in the review?\n \n---\n\n> One of the main weaknesses of this paper is that the policy improvement step (Eq. 27) seems to require gradient back-propagation through the whole diffusion generation path, similar to Diffusion-QL. In practice this incurs significant memory cost and computational overhead.\n> \n\nWe thank the reviewer for sharing their concerns and agree that gradient back-propagation through the chain can be limiting. However, in practice, this only becomes problematic in very high dimensions, such as when fine-tuning large image-models [1]. In fact, most recent works on diffusion-based sampling/inference employ equal update schemes while scaling to problems with 1000+ dimensions using hundreds of diffusion steps [2,3,4]. \n\nHowever, if one wants to avoid having to compute gradients through the simulation, one could use the adjoint ODE for computing gradients, which describes the evolution of the gradient of the objective with respect to the parameters (see [5]). Another alternative is to use different loss functions such as the log-variance loss [6] or adjoint matching [1].\n\n[1] Domingo-Enrich, Carles, et al. \"Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control.\" ICLR 2025\n\n[2] Vargas, Francisco, Will Grathwohl, and Arnaud Doucet. \"Denoising diffusion samplers.\" ICLR 2023\n\n[3] Vargas, Francisco, et al. \"Transport meets variational inference: Controlled monte carlo diffusions.\" ICLR 2024\n\n[4] Blessing, Denis, et al. \"Underdamped diffusion bridges with applications to sampling.\" ICLR 2025\n\n[5] Li, Xuechen, et al. \"Scalable gradients for stochastic differential equations.\" AISTATS 2020\n\n[6] Richter, Lorenz, and Julius Berner. \"Improved sampling via learned diffusions.\" ICLR 2024\n \n---\n\n> Eq. 33, the RHS of the equation should be …\n> \n\nWe thank the reviewer for pointing out this typo, which we have fixed.\n \n---\n\n> In Section 4.4, the authors mentioned that the diffusion coefficient $\\beta$ is optimized alongwith the RL process. Could the authors provide illustrations about the learned $\\beta$  at the end of RL training?\n>\nAs requested, we provide the adaptation of the $\\beta$ value during the training under this link: https://imgur.com/a/ASqSHCe . We have extended the updated version of the paper with the corresponding figure.", "claims_and_evidence": "Yes, main claims in this paper are well supported by theory or empirical evidence.", "methods_and_evaluation_criteria": "Yes. The authors test DIME on several tasks with high-dimentional observations and actions selected from three benchmarks. These are widely acknowledged benchmarks for online reinforcemenet learning.", "theoretical_claims": "Yes, I have checked the theory. The main point of the theory is the lower bound of the entropy of the action distribution, based on which the rest largely follows the theory established by SAC.", "experimental_designs_or_analyses": "The authors designed two types of experiments to demonstrate the effectiveness of their design:\n- Three **ablation studies** examined the role of the reward scaling parameter $\\alpha$, the impact of diffusion steps, and the improvement brought by diffusion policies compared to Gaussian policies.  \n- **Performance comparisons** with other reinforcement learning algorithms. \n\nIssues:\n- The comparison with baseline methods may not be strictly fair. For non-diffusion methods, hyper-parameters such as the target entropy, the width of the critic networks are not aligned with DIME. It would be better to include additional experiments to ablate the effect of these parameters to demonstrate the benefit of diffusion compared to Gaussian policies.", "relation_to_broader_scientific_literature": "This article is related to two major fields: MaxEnt RL and diffusion-based policies.\n\nMaxEnt RL augments the traditional RL objective with the policy entropy to promote exploration. The most famous algorithm in MaxEnt RL is SAC [1]. On the other hand, Diffusion-based policies have stronger expressive power compared to those based on simple distributions such as Gaussian or Dirac distribution. Although the first attempt towards diffusion-based policies merged in offline RL [2], diffusion policies have so far been widely applied in online RL (e.g. QSM [3], DACER [4], and QVPO[5]). However, combing diffusion policies with MaxEnt RL is challenging because it is generally difficult to complute the log-probability of the diffusion policy analytically, and DIME offers a nice approach to this through a tractable lower bound of the entropy term. \n\n[1] Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft  actor-critic: Off-policy maximum entropy deep reinforcement learning  with a stochastic actor. ICML. 2018.\n\n[2] Wang, Z., Hunt, J. J., and Zhou, M. Diffusion policies as an  expressive policy class for offline reinforcement learning. ICLR. 2023\n\n[3] Psenka, M., Escontrela, A., Abbeel, P., and Ma, Y. Learning a  diffusion model policy from rewards via q-score matching. ICML. 2024. \n\n[4] Wang, Y., Wang, L., Jiang, Y., Zou, W., Liu, T., Song, X., Wang, W., Xiao,  L., WU, J., Duan, J., and Li, S. E. Diffusion actor-critic with entropy  regulator. NeurIPS. 2024.\n\n[5] Ding, S., Hu, K., Zhang, Z., Ren, K., Zhang, W., Yu, J., Wang, J., and Shi, Y. Diffusion-based reinforcement learning via q-weighted  variational policy optimization. NeurIPS. 2024.", "essential_references_not_discussed": "Most of the relevant literature are properly cited. However, in the field of offline RL, there are a few articles that also employs diffusion-based policies, such as Diffusion-DICE [2], DAC [3], and QGPO [4]. The authors are encouraged to include a discussion about these works.\n\n[1] Mao, L., Xu, H., Zhan, X., Zhang, W.,  & Zhang, A. Diffusion-dice: In-sample diffusion guidance for offline reinforcement learning. \n\n[2] Fang, L., Liu, R., Zhang, J., Wang,  W., & Jing, B. Y. Diffusion Actor-Critic: Formulating  Constrained Policy Iteration as Diffusion Noise Regression for Offline  Reinforcement Learning. \n\n[3] Lu, C., Chen, H., Chen, J., Su, H.,  Li, C., & Zhu, J. Contrastive energy prediction for  exact energy-guided diffusion sampling in offline reinforcement learning.", "other_strengths_and_weaknesses": "One of the main weaknesses of this paper is that the policy improvement step (Eq. 27) seems to require gradient back-propagation through the whole diffusion generation path, similar to Diffusion-QL. In practice this incurs significant memory cost and computational overhead.", "other_comments_or_suggestions": "Eq. 33, the RHS of the equation should be $\\frac{\\pi_{0:N}(a^{0:N}|s)}{\\pi_{1:N|0}(a^{1:N}|s, a^0)}$.", "questions_for_authors": "In Section 4.4, the authors mentioned that the diffusion coefficient $\\beta$ is optimized alongwith the RL process. Could the authors provide illustrations about the learned $\\beta$ at the end of RL training?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Lgm3EvGOkH", "reviewer": "Reviewer_TJZx", "summary": "This paper introduces DIME, a novel online RL algorithm using diffusion policies. The key innovation lies in proposing a new method for maximizing the entropy of a diffusion policy with more rigorous theoretical justification. The authors establish a lower bound for policy entropy and derive a practical diffusion RL algorithm based on this foundation. The algorithm can be interpreted as minimizing the KL divergence between the marginal distributions of the reverse process defined by the diffusion model and the forward diffusion process by adding noise to the maximum entropy optimal action distribution at each diffusion step. Extensive experiments conducted across multiple benchmark environments demonstrate superior sample efficiency and converged performance compared to baseline methods in most cases.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised.\n \n---\n\n> The experimental section lacks comparisons with some of the latest diffusion-based online RL methods, such as QVPO and DACER.\n> \nWe thank the reviewer for the comments and their time. We have run additional experiments comparing DIME against the most recent diffusion-based RL baselines QVPO and DACER. We report the results over 10 seeds here: https://imgur.com/a/3h8hMkB . DIME converges faster to a similar performance as DACER on the Ant-v3 task and significantly outperforms both methods on the Humanoid-v3 task. Please note that the results reported in the DACER paper differ from our results. In their work, the X-axis corresponds to the number of updates, which is not the number of environment interactions. To obtain the number of environment interactions, the X-Axis would need to be multiplied by 20.\n \n---\n\n> The authors employ distributional RL in DIME but do not specify whether the baseline methods also use it. Additionally, they do not provide corresponding ablation studies, making it unclear how much improvement is attributed to the use of distributional RL.\n> \nWe thank the reviewer for this comment. We agree that the descriptions of the baselines are a bit short. We have extended the updated version of the paper. Except for BRO and BRO (Fast), all methods do not employ distributional Q by default. A comparison to a distributional variant of CrossQ has been done in Fig. 1c)- d) already. \n\nWe agree that additional ablations strengthen the paper. Therefore, we have conducted additional experiments that provide additional evidence: https://imgur.com/a/N6sUqxQ. DIME with distributional Q improves performance slightly, but DIME w/o distributional Q performs on par or better than BRO, and outperforms other baselines, especially the diffusion-based methods. Please note that BRO employs quantile distributional RL.\n\nAdditionally, we analyzed the effect of distributional Q on the Ant-v3 and Humanoid-v3 tasks. More concretely, we compare DIME against variants of diffusion-based baselines with distributional Q, i.e., we compare against DACER, Diffusion-QL with distributional Q, and Consistency-AC with distributional Q. Please note that DACER employs distributional Q by default. Similarly, we compare DIME without distributional Q against the diffusion-based variants without distributional Q. The results can be seen here: https://imgur.com/a/xBGkPl7\nDIME w/o distributional performs on par with baselines on the Ant-v3 task and performs significantly better on the humanoid-v3 task. DIME with distributional Q converges slightly faster to the same end performance as DACER on the Ant-v3 but significantly outperforms on the Humanoid-v3 task.\n \n---\n\n> In Equation $Q_\\theta$ should be $Q_\\phi$\n> \n\nWe thank the reviewer for pointing out this typo. \n\nWe thank the reviewer for raising these questions, which led us to conduct additional experiments. We will include those results in the future version.", "claims_and_evidence": "Yes, the experimental results clearly demonstrate the effectiveness of the proposed algorithm.", "methods_and_evaluation_criteria": "Yes, the proposed algorithm effectively maximizes the entropy of the diffusion policy.", "theoretical_claims": "No, most of the proofs align with standard derivations in maximum entropy RL, but I have not thoroughly verified their correctness.", "experimental_designs_or_analyses": "Yes, the authors’ experimental design is reasonable, and they conducted extensive experiments with multiple random seeds to mitigate statistical uncertainty.", "relation_to_broader_scientific_literature": "This paper is closely related to diffusion-based offline and online RL, as well as standard maximum entropy RL.", "essential_references_not_discussed": "I did not find any.", "other_strengths_and_weaknesses": "### Strengths\n\n1. The writing of the paper is very clear, with well-structured theoretical derivations that progress logically.\n\n2. This paper addresses a crucial issue in applying diffusion models to online RL: policy entropy estimation. It presents a novel solution distinct from previous works and enables non-Gaussian exploration patterns.\n\n3. The experiments in this paper are comprehensive, demonstrating superior performance compared to baseline methods like DIPO.\n\n### Weaknesses\n\n1. The experimental section lacks comparisons with some of the latest diffusion-based online RL methods, such as QVPO and DACER.\n\n2. The authors employ distributional RL in DIME but do not specify whether the baseline methods also use it. Additionally, they do not provide corresponding ablation studies, making it unclear how much improvement is attributed to the use of distributional RL.", "other_comments_or_suggestions": "In Equation 26, the $Q_\\theta$ should be $Q_\\phi$.", "questions_for_authors": "See weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0t3hbGxeeQ", "reviewer": "Reviewer_gx83", "summary": "This paper proposed a diffusion-based maximum entropy RL. The authors first derived the lower bound of entropy by data processing inequality. Then the soft policy iteration is recased to diffusion-based version with new regularization term with the lower bound. The policy optmization is reformulated to the lower bound of soft policy iteration, which gives a tractble loss function to optimize diffusion policy. Experimental results have shown the performance is improved compared to previous diffusion-based RL.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our work and for the many helpful comments and suggestions. We hope the following replies address the questions and concerns raised.\n \n---\n\n> The claim 'significantly outperforming other diffusion-based methods on challenging high-dimensional control benchmarks' is not that well supported, given that the authors only compare DIME (proposed) with QSM and DIPO, leaving DACER, QVPO […]\n\nWe have conducted additional experiments and compared DIME to QVPO and DACER. We kindly refer the reviewer to our answers to Reviewer TJZx who shared a similar concern.\n \n---\n\n> […] the authors made two modifications to the current maximum entropy RL , one is the lower-bound modified soft policy interation, another is the Q-learning via crossQ. The authors does not provide ablation studies that ablates either of these two modifications. How they contributes to the performance improvement is not clear.\n\nPlease note that a direct comparison to the common max. entr. objective is not possible, as the marginal entropy of the diffusion process is not tractable.  Therefore, DIME proposes a lower-bound objective that is tractable and can be used to optimize the diffusion policy by adapting the policy iteration scheme. Please note that we only use CrossQ’s Q-network architecture which does not entail a modification to the maximum entropy objective. Our work proposes a framework for training diffusion policies that is orthogonal to improvements for critics/Q-functions.\n \n---\n\n> I have concerns on the equation (25) which is the policy loss. [….] Then why all the previous lower bounds are necessary? Can we just do normal Q-learning and optimize (25)?\n\nPlease note that optimizing the maximum entropy objective for diffusion models is problematic as the marginal entropy of the diffusion process is not tractable and requires further treatment. As such, our work proposes (25), which is justified through the lower bound in (16) or the data processing inequality in (20). Both perspectives show that optimizing (25) is equivalent to maximizing the maximum entropy objective in (1).\n\nWe kindly ask the reviewer to elaborate on what is meant by 'normal Q-learning'. If the reviewer means defining the Q-function as the sum of the discounted rewards without taking the expected entropies into account, we would like to mention that this objective is equivalent to entropy-regularized RL [1, 2] and does not optimize the maximum entropy objective as we consider in this work.  We hope that this addresses the reviewers’ concerns. If we misunderstand the question, we are more than happy to provide further clarification.   \n\n[1] Neu, G. A unified view of entropy-regularized markov decision processes. \n\n[2] Levine, S. Reinforcement Learning and Control as Probabilistic Inference.\n \n---\n\n\n> […] However, in (27), you write  $\\overrightarrow{\\pi}$  with $n-1|n$, which is not consistent since $n-1|n$ should be the reverse process Gaussian.\n\nWe thank the reviewer for spotting the mistake in (27) and apologize for any confusion it may have caused.  The correct transition densities read $\\pi^{\\theta}_{n-1|n}$ for the denoising process and $\\overrightarrow{\\pi} _{n|n-1}$ for the noising process.\n \n---\n\n> The claim near line 295 […] In (27), no terms need reparameterization since it only involves the log probability. How exactly the reparameterization wokrs here? […] The implementation is actually mismatched with the actual policy loss in (27).\n\nPlease note that the reparameterization trick is necessary as we are computing gradients with respect to the parameters of the distribution that we are taking expectations with. For $L(\\theta$) (Eq. 27) he expectation is taken w.r.t. the parameterized denoising process $\\pi^{\\theta}\\_{0:N}$. \nFor diffusion models, we need to iteratively apply the reparameterization trick for every transition density $\\pi^{\\theta}\\_{n-1|n}$  in order to reparameterize the joint distribution $\\pi\\_{0:N}^{\\theta}$. For a single transition $\\pi^{\\theta}\\_{n-1|n}$ the reparameterization trick reads\n$$\na^{n-1}  = a^{n} + (\\beta_{n} a^{n} + 2\\eta^2\\beta_{n} f_n^{\\theta}(a^{n}|s))\\delta + \\xi_{n},\n$$\n\nwith $\\xi\\_n \\sim \\mathcal{N}(0,2\\eta^2\\beta\\_{n}\\delta I)$, which agrees with the Euler-Maruyama discretization.  As such, the denoised action $a^0$ is a function of all previous transitions and therefore depends on $\\theta$.  This is inline with (27) where the expectation is taken under the diffusion policy $\\pi^{\\theta}\\_{0:N}$ which agrees with the implementation in our code base, where $a^0$ is coming from $\\pi^{\\theta}_{0:N}$ .\n\nWe thank the reviewer for the comment and agree that the reparameterization trick needs further clarification. In response, we added additional details to the updated version of the paper..", "claims_and_evidence": "1. The lower bound claim is correct.\n2. The policy iteration scheme is correct.\n3. The claim \n> significantly outperforming other diffusion-based methods on challenging high-dimensional control benchmarks\nis not that well supported, given that the authors only compare DIME (proposed) with QSM and DIPO, leaving DACER, QVPO, and other related works mentioned in the paper not involved.", "methods_and_evaluation_criteria": "1. The evaluation metrics make sense.\n2. There is one thing not clear: the authors made two modifications to the current maximum entropy RL, one is the lower-bound modified soft policy interation, another is the Q-learning via crossQ. The authors does not provide ablation studies that ablates either of these two modifications. How they contributes to the performance improvement is not clear.", "theoretical_claims": "1. The lower bound and policy iteration convergence is correct.\n---\n2. I have concerns on the equation (25), which is the policy loss. It can stands on it self since it just fits the reverse diffusion trajectories with the forward diffusion trajectories. Then why all the previous lower bounds are necessary? Can we just do normal $Q$-learning and optimize (25)?\n---\n3. The claim near line 295\n> .... stochastic opitmization of L(θ) does not need access to samples a0 ∼ exp Qϕ /Zϕ , instead relying\non stochastic gradients obtained via reparameterization trick\n(Kingma, 2013) using samples from the diffusion model π θ .\n\nseems problematic. In (27), no terms need reparametrization since it only involves the log probability.", "experimental_designs_or_analyses": "There is one thing not clear: the authors made two modifications to the current maximum entropy RL, one is the lower-bound modified soft policy interation, another is the Q-learning part via crossQ. The authors does not provide ablation studies that ablates either of these two modifications. How they contributes to the performance improvement is not clear.\n\n**Update after rebuttal**: What I mean is the effect of crossQ, and the authors have resolved it in the rebuttal.", "relation_to_broader_scientific_literature": "There might be some but none is worth mentioning here.", "essential_references_not_discussed": "No, the authors provides sufficient references.", "other_strengths_and_weaknesses": "Strength:\n1. The baselines keep up with some most recent RL algorithms, which is good.\n\nWeakness:\n1. The notations are extremely comfusing and have some errors. For example, I believe $\\rightarrow$ means forward diffusion process and $\\leftarrow$ means reverse diffusion process. However, in (27), you write $\\pi^\\rightarrow$ with $n-1|n$, which is not consistent since $n-1|n$ should be the reverse process Gaussian. I didn't fully understand.\n2. As I said the policy loss stands on itself, the value of all the previous discussion remain unclear.", "other_comments_or_suggestions": "In line 351, There is a typo on the env names An-tv3, it should be Ant-v3.", "questions_for_authors": "I will repeat some questions I listed before,\n\n---\n1. I have concerns on the equation (25), which is the policy loss. It can stands on it self since it just fits the reverse diffusion trajectories with the forward diffusion trajectories. Then why all the previous lower bounds are necessary? Can we just do normal $Q$-learning and optimize (25)?\n---\n2. The claim near line 295\n> .... stochastic opitmization of L(θ) does not need access to samples a0 ∼ exp Qϕ /Zϕ , instead relying\non stochastic gradients obtained via reparameterization trick\n(Kingma, 2013) using samples from the diffusion model π θ .\n\nseems problematic. In (27), no terms need reparametrization since it only involves the log probability. How exactly the reparametrization works here?\n\n---\n3. Continuing on 3, I check the code the authors provided, it shows the policy loss as\n```python\nactor_loss = (- min_qf_pi + ent_coef_value * (run_costs.squeeze() + sto_costs.squeeze() + terminal_costs.squeeze())).mean()\n```\nwhere the first term is $Q_\\phi(s, \\tilde a)$ where $\\tilde a$ is sampled from $\\pi^\\theta$. However, in (27) the $a^0$ in $Q_\\phi(s,a^0)$ term should not be a function of $\\theta$ since it is derived from a part of $\\pi^{\\rightarrow}(a^{0:N}|s)$. Therefore, the implementation is actually mismatched with the actual policy loss in (27). Please kindly explain the gap between theoritical results and your implentation and how it affects the contribution of this paper.\n\n**Update after rebuttal**: it is a misunderstanding from my side. The derivation is clear.\n\nOverall, I do think there are some insights on the loss function in (25), which somehow flipped the KL divergence to make the diffusion policy training tractable. However, as the importance of major theoritical results remain unclear and the implementation are not matching the theoritical derivation and statements, I lean to reject the paper unless the authors have covincing answers to all these questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents DIME, a novel approach leveraging diffusion-based policies within maximum entropy reinforcement learning. It proposes a tractable objective, i.e. the entropy lower bound. All reviewers recommended acceptance, a decision I endorse.
Decision: Accept (poster)