Maximum Entropy Reinforcement Learning with Diffusion Policy
Authors: Xiaoyi Dong
Jian Cheng
Xi Sheryl Zhang
Venue: ICML 2025 poster (2025)
Keywords: Diffusion models
online reinforcement learning
maximum entropy reinforcement learning
soft actor-critic
Primary area: reinforcement_learning->deep_rl
TL;DR: We employ the diffusion model as the policy representation to achieve the Maximum Entropy RL objective, which enables efficient exploration and brings the policy closer to the optimal MaxEnt policy.

Abstract: The Soft Actor-Critic (SAC) algorithm with a Gaussian policy has become a mainstream implementation for realizing the Maximum Entropy Reinforcement Learning (MaxEnt RL) objective, which incorporates entropy maximization to encourage exploration and enhance policy robustness. While the Gaussian policy performs well on simpler tasks, its exploration capacity and potential performance in complex multi-goal RL environments are limited by its inherent unimodality.  In this paper, we employ the diffusion model, a powerful generative model capable of capturing complex multimodal distributions, as the policy representation to fulfill the MaxEnt RL objective, developing a method named MaxEnt RL with Diffusion Policy (MaxEntDP). Our method enables efficient exploration and brings the policy closer to the optimal MaxEnt policy. Experimental results on Mujoco benchmarks show that MaxEntDP outperforms the Gaussian policy and other generative models within the MaxEnt RL framework, and performs comparably to other state-of-the-art diffusion-based online RL algorithms. Our code is available at https://github.com/diffusionyes/MaxEntDP.

Maximum Entropy Reinforcement Learning with Diffusion Policy
Xiaoyi Dong 12 Jian Cheng 134
Xi Sheryl Zhang 14
Abstract
ever, this approach often overlooks the inherent uncertainty and variabilityy of real-world environments, which can lead to suboptimal or overly deterministic policies. To address these limitations, Maximum Entropy Reinforcement Learn- ing (MaxEnt RI) incorporates entropy maximization into the standard RL objective, encouraging exploration and im- proving robustness during policy learning (Toussaint, 2009; Ziebart, 2010; Haarnoja et al., 2017).
The Soft Actor-Critic (SAC) algorithm with a Gaussian policy has become a mainstream im- plementation for realizing the Maximum Entropy Reinforcement Learning (MaxEnt RL) objective. which incorporates entropy maximization to en- courage exploration and enhance policy robust- ness. While the Gaussian policy performs well on simpler tasks, its exploration capacity and po- tential performance complex multi-goal RL environments are limited by its inherent uni- modality. In this paper, we employ the diffusion model, powerful generative model capable of capturing complex multimodal distributions, the policy representation to fulfill the MaxEnt RL objective, developing method named Max- Ent RL with Diffusion Policy (MaxEntDP). Our method enables efficient exploration and brings the policy closer to the optimal MaxEnt pol- icy. Experimental results on Mujoco benchmarks show that MaxEntDP outperforms the Gaussian policy and other generative models within the MaxEnt RL framework, and performs compara- bly to other state-of-the-art diffusion-based on- line RL algorithms. Our code is available at http:://tttb..oooooo//ooooo////a
The Soft Actor-Critic (SAC) algorithm (Haarnoja al., 2018) 1)) is effective method for achieving the MaxEnt RL objective, which alternates between policy evaluation and policy improvement to progressively refine the policy. With high-capacity neural network approximators and suitable optimization techniques, SAC can provably converge to the optimal MaxEnt policy within the chosen policy set. The choice of policy representation in SAC is crucial, as it influences the exploration behavior during training and determines the proximity of the candidate policies to the optimal MaxEnt policy. In complex multi-goal RL tasks, where multiple feasible behavioral modes exist, the com- monly used Gaussian policy typically explores only single mode, which can cause the agent get trapped local optimum and fail to approach the optimal MaxEnt policy that captures all possible behavioral modes.
In this paper, we propose using diffusion models (Sohl- Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020; Song et al., 2021b), a powerful generative model, as the pol- icy representation within the SAC framework. This allows for the exploration of all promising behavioral modes and facilitates convergence to the optimal MaxEnt policy. Diffu- sion models transform the original data distribution into tractable Gaussian by progressively adding Gaussian noise, which is known as the forward diffusion process. After train- ing neural network to predict the noise added to the noisy samples, the original data can be recovered by solving the reverse diffusion process with the noise prediction network. While several generative models, e.g., variational autoen- coders (Kingma, 2013), generative adversarial networks (Goodfellow et al., 2020), and normalizing flows (Rezende & Mohamed, 2015) could serve as the policy representation, we choose diffusion models due to their balance between expressiveness and inference speed, achieving remarkable performance with affordable training and inference costs.
1. Introduction
Reinforcement Learning (RL) has emerged as a powerful paradigm for training intelligent agents make decisions in complex control tasks (Silver et al., 2016; Mnih et al., 2015; Kaufmann et al., 2023; Kiran et al., 2021; Ibarz et al., 2021). Traditionally, RL focuses on maximizing the expected cu- mulative reward, where the agent selects actions that yield the highest return in each state (Sutton & Barto, 1999). How-
DL, Institute of Automation, Chinese Academy of Sciences "School of Artificial Intelligence, University of Chinese Academy of Sciences "School of Future Technology, University of Chi- nese Academy of Sciences 'AiRiA. Correspondence Xi Sheryl Zhang <sheryl zhangxxi @ gmai..com>>
Proceedings of the nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
However, integrating diffusion models into the SAC frame- work presents two key challenges: 1) How to train diffu- sion model to approximate the exponential of the Q-function in the policy improvement step? 2) How to compute the log probability of the diffusion policy when evaluating the soft Q-function? To address the first challenge, we analyze the training target of the noise prediction network in diffu- sion models and propose a Q-weighted Noise Estimation method. For the second challenge, we introduce numeri- cal integration technique to approximate the log probability of the diffusion model. We evaluate the effectiveness of our approach on Mujoco benchmarks. The experimental re- sults demonstrate that our method outperforms the Gaussian policy and other generative models within the MaxEnt RL framework, and performs comparably to other state-of-the- art diffusion-based online RL algorithms.
and the policy, where O and & represent the parameters of the respective networks. These networks are optimized by alternating between policy evaluation and policy improve- ment.
In the policy evaluation step, the soft Q-function of the current policy T is learned by minimizing the soft Bellman error:
[ {/2 2 (olts.a-Q(s..)))
l(0) = B(o,a)>D
(2)
where 'D is the replay buffer, and the target value Q a) r(s, a) + +.='s'ppa''' [Qo(s', a') - 3 log T [s)].
In the policy improvement step, the old policy Te. up- dated towards the exponential of the new Q-function, whose soft value is guaranteed higher than the old policy. How- ever, the target policy may be too complex to be exactly represented by any policy within the parameterized policy set IT = { E $), where t is the parameter space of the policy. Therefore, the new policy is obtained by pro- jecting the target policy onto the policy set II based on the divergence:
2. Preliminary
2.1. Maximum Entropy Reinforcement Learning
In this paper, we focus on policy learning in continuous action spaces. We consider a Markov Decision Process (MDP) defined by the tuple (S, A. p, Po: 7), where S rep- resents the state space, A is the continuous action space, p: P:SSX XSSSSSxxx x A +> [0, + is the probability density function of the next state 8:+1 E S given the current state s; E S and the action a; E A., r S x A -> [Fmin; T'max] is the bounded reward function, Po : S + [0, +o0] is the distri- bution of the initial state So and 7 E [o, 1] is the discount factor. The marginals of the trajectory distribution induced by a policy x (a,(s;) are denoted as pr(ss, a:).
exp( exp(220(s.)))
[(Ø) = DKL (T((-Is)
(3)
(s)
Theorem 2.1. (Soft Policy Iteration) In the tabular setting, ler ((O)) and L(ōk) be minimized for each k. Repeated application of policy evaluation and policy improvement, i.e., k oc. Tée will converge to a policy x' such that Q" (s.a) 2 (*((s,)) a) for all E TI and (s, E x with |A| <0.
ESX A
The standard RL aims to learn policy that maximizes the expected cumulative reward. To encourage stochastic policies, Maximum Entropy RL augments this objective by incorporating the expected entropy of the policy:
Theorem 2.1 suggests that if the Bellman error can be re- duced to zero and the policy loss is minimized at each opti- mization step, the soft actor-critic algorithm will converge to the optimal maximum entropy policy within the policy set II. This indicates that the choice of the policy set II significantly affects the performance of the soft actor-critic algorithm. Specifically, more expressive policy class will yield a policy closer to the optimal MaxEnt policy. Inspired by this intuition, we employ the diffusion model to represent the policy, as is highly expressive and well-suited cap- ture the complex multimodal distribution (Chi al., 2023; Wang et al., 2023; Chen et al., 2023; Ajay et al., 2023).
J(TT) = =20 7 (er. [r(*t, ax) + (3H(T(((+)))]] (1)
where H((T(-1&:) = (()))())ss)) [- log (az) ())]) and B is
the temperature parameter that controls the trade-off be- tween the entropy and reward terms. A higher value of 3 drives the optimal policy to be more stochastic, which is advantageous for RL tasks requiring extensive exploration. In contrast, the standard RL objective can be seen the limiting case where B - 0.
2.3. Diffusion Models
2.2. Soft Actor Critic
Diffusion models are powerful generative models. Given an unknown data distribution p(æ.)), which is typically mix- ture of Dirac delta measures over the training dataset, diffu- sion models transform this data distribution into a tractable Gaussian distribution by progressively adding Gaussian noise (Ho et al., 2020). In the context of a Variance- Preserving (VP) diffusion process (Ho et al., 2020; Song
The optimal maximum entropy policy can be derived by applying the Soft Actor-Critic (SAC) algorithm (Haarnoja et al., 2018). In this subsection, we will briefly introduce the framework of SAC, and the relevant proofs are provided in Appendix A.I. The SAC algorithm utilizes two parame- terized networks, Qo and ne. to model the soft Q-function
2

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
al., 2021b), the transition from the original sample wo at time 0 to the noisy sample a at time t [0, 1] follows the distribution:
policies perform well in relatively simple single-goal RL environments, they often struggle with more complex multi- goal tasks.
pla, mo) - (ol
(4)
Consider a typical RL task that involves multiple behavior modes. The most efficient solution is to explore all be- havior modes until one obviously outperforms the others. However, this exploration strategy is difficult to achieve with Gaussian policies. In the training process of soft actor-critic algorithm with Gaussian policies, minimizing the KL divergence between the Gaussian policy and the exponential of the O-function--wiie in multi-goal tasks tends to push the Gaussian policy to allocate most of the probability mass to the action region with the highest Q value (Chen et al., 2024a). Consequently, other promising action regions with slightly lower Q values will be neglected, which may cause the agent to become stuck at local optimal policy.
where or represents the log of the Signal-to-Noise Ratio (SNR) at time t, and (-) is the sigmoid function. a, deter- mines the amount of noise added at each time and is referred to as the noise schedule of a diffusion model. Denote the marginal distribution of ar as p(arz). The noise schedule should be designed ensure that p20) p(xx1) P2 (2x [0, I). and that a₁ is strictly decreasing w.r.t. t. Then, starting from 1 ~ N(x (0, I), the original data samples can be recovered by reversing the diffusion process from t= 1 10 / = 0. For sample generation, we can also employ the following probability flow ordinary differential equation (ODE) that shares the same marginal distribution with the diffusion process (Song et al., 2021b):
is often multimodal
However, an efficient exploration strategy can be achieved by replacing the Gaussian policy with more expressive pol- icy representation class. If accurately fitting the multimodal target policy (i.e., the exponential of the Q-function), the agent will explore all high-return action regions at a high probability, thus reducing the risk of converging to local optimum. Moreover, recall that when the assumptions on loss optimization are met, the soft actor-critic algorithm is guaranteed to converge to the optimal maximum entropy pol- icy within the chosen policy class. Therefore, with sufficient network capacity and appropriate optimization techniques, we can obtain the true optimal maximum entropy policy, as long the selected policy representation class is expressive enough to capture it.
dæ: = f(t)x; 2º Vz. p(z:),
(5)
dt
where [(1) = 1 dlogg(mo)) Va, log p(x)). known as the score function, is the only un- known term. Consequently, diffusion models train a neural network E c0 a1) to approximate the scaled score func- tion- - (-)0=)) log p(z)). The training loss L(Ø) is defined as:
2 df g²(1)= = alogg(oo) and de
[(Ø)= E, We (e(()))))) + (o)- p(g:(
(6)
Ex [ure (x)) - E112] + C
The above analysis emphasizes the importance of applying an expressive policy class to achieve efficient exploration as well as a higher performance upper bound. Since diffusion models have demonstrated remarkable performance in cap- turing complex multimodal distributions, we adopt them to represent the policy within the soft actor-critic framework. However, integrating diffusion-based policy into the soft actor-critic algorithm presents several challenges: (I) In the policy improvement step, the new diffusion policy is updated to approximate the exponential of the Q-function. However, existing methods for training diffusion models rely on samples from the target distribution, which are un- available in this case. (2) In the policy evaluation step, computing the soft Q-function requires access to the proba- bility of the diffusion policy. Nevertheless, diffusion models implicitly model data distributions by estimating their score functions, making it intractable to compute the exact proba- bility.
where To ~ p(zo), c I)..)). t u((0, 1),, xz
Vo(a,)ag + (a(-or) we is a weighting function and usually set to W/ = 1, and C' is constant independent of ø. In this setup, the network cg ar) target at predicting the expectation of noise added to the noisy sample 2,, and is therefore called the noise prediction network. Minimizing the loss function (Ø)) results in the following relationship:
Cg(22,0r) Vo(
V logp(z.)
(8)
Then we can solve the probability flow ODE in Equation with the assistance of existing ODE solvers (Ho et al., 2020; Song et al., 10; Lu et al., 2022; Karras et al., 2022: Zheng et al., 2023) to generate data samples.
3. Methodology
The remainder of this section addresses these challenges and describes how to incorporate diffusion models into the soft actor-critic algorithm for efficient policy learning. We first
In the soft actor-critic algorithms, Gaussian policies have become the most widely used class of policy representation due to their simplicity and efficiency. Although Gaussian
3

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
propose the Q-weighted Noise Estimation approach to fit the exponential of the Q-function in Section 3.1, then introduce method for probability approximation in diffusion policies in Section 3.2, and finally present the complete algorithm in Section 3.3. We name this method MaxEntDP because it can fulfill the MaxEnt RL objective with diffusion policies.
The proof is provided in Appendix A.2. Lemma 3.1 demon- strates that the conditional distribution p(an (a)) can be seen as a Gaussian distribution of do weighted by the exponential of the Q-function. Sampling from the Gaussian distribu- tion is straightforward, we can apply importance sampling (Bishop, 2006) to estimate the expectation in Equation 11. Theorem 3.2. (Importance Sampling Estimate for the Score Function) The score function can be estimated by
3.1. Q-weighted Noise Estimation
Given a Q-function ((s, a). below we will analyze how to train a noise prediction network Eg in the diffusion model to approximate the target distribution:
k
Va. logp(ar) ~ Vol-ai)
w(ab)e'
(13)
K
1=1
ग [a]ss) = exp(
((((ss))) (xp()((s,aa) a))
where et e',..., N(O,1), ab 1 a:
+ vol-oe)
Volari
Volar)
((x)
and the importance ratio w(as)- = exp(+Q(an))
with ((a))
2(a.)
Omitting the state in the condition for simplicity and follow- ing the symbol convention of diffusion models, we rewrite r (als) as p(an). The transition from the original action samples ag at time = 0 to the noisy actions a, at time t E [0, 1] is defined as:
being the normalizing constant of p(ag) a.).
The derivation is detailed in Appendix A.3. Although this importance sampling estimate is unbiased, it exhibits high variance when the variance of the Q-function is large. To address this issue, we employ the weighted importance sam- pling approach (Bishop, 2006) to reduce variance and stabi- lize the training process.
p(a,la)) = N(a.l V (an,aan, oaa))an r-
(10)
Note that the symbol stands for the time of diffusion mod- els if not specified.
Theorem 3.3. (Weighted Importance Sampling Estimate for the Score Function) The score function can be estimated by
The marginal distribution of noisy actions a at time t is de- noted by p(a:). To sample from p(ag), we need to estimate the score function aa, log p(a;) each intermediate time t during the diffusion process. The score function can be reformulated as:
w (aa)
Va, log logp(a.) R
(14)
vo(-au)
K w(a))
j=1
K
softmax { [ (((e
(15)
Va, log p(a.) = lo(a Va, log (11)
vo(-ou an)
which is an expectation with respect to the conditional distri- bution (aa)aa)) a.k.a. the reverse transition distribution of
where softmax( (+Q(() ES= exp( (((a)))
expl 3 Q(a)))
the diffusion process. If samples from p(ao) are available, as is often the case in the application scenarios of diffusion models (Saharia al., 2022; Ho et al., 2022; Chi et al., 2023; Xu et al., 2023; Huang et al., 2023), we can first sample original actions ag ~ p(ao), and then sample noisy actions a, ~ p(a, (a)) to obtain several sample pairs following the joint distribution p(a., a:). Then for a fixed noisy action at, the corresponding agg will conform the conditional distribu- tion p(aglas (a.), which can serve as Monte Carlo samples to estimate the expectation in Equation 11. Conversely, in the context of the soft actor-critic algorithm, we lack samples from the target distribution p(a.) but instead have access to Q-function. Therefore, we must establish the relation- ship between the conditional distribution p(aglay (a;) and the f-function.
The normalizing constant (as)) is canceled out in Equation 15, eliminating the need for its explicit computation. Since the bias of the weighted importance sampling method de- creases as the number of Monte Carlo samples increases, larger value of K is preferred in practice given adequate computation budgets.
Then the training target of the noise prediction network is
E" ca(a,, an) = o az
Va, log p(a)(
(16)
K
Q(a Die,,
softmax
(17)
This target can be interpreted as a weighted sum of noise, with the weights being the exponential of the Q-value. Con- sequently, we refer to this method as Q-weighted Noise Estimation for training the noise prediction network. The overall training loss is
Lemma 3.1. (Decomposition of the Reverse Transition Distribution) The conditional distribution p(as) (a) can be decomposed as
d(-a.) I) (12) (:) x ((aa))() volas) a. ala))
[(Ø) = Ep((a) [Il colar, 04) - 6" ca(a,, az) [
(18)

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
While the true distribution of noisy actions p(a;) may be inaccessible, we can substitute it with other distributions with full support, as the loss will still be minimized for each at given sufficient network capacity.
this expression is intractable because both the integral in e and the integral of the noise prediction error diverge, with only their difference converging (Kong et al., 2023). We attempt to approximate the integral using numerical inte- gration techniques. However, we observe that using the log SNR as the integration variable results in a high variance, as it spans from -00 to +oo. Therefore, we instead utilize o (as) with a narrower integration domain of (0. 1).
We briefly compare our method with two previous ap- proaches that approximate the exponential of a given func- tion Q(a). The QSM method (Psenka et al., 2024) esti- mates the score function as Va. (ogp(a))
Va. {Q(a)). This approximation requires p( p()z) x exp( 3 (a:)), which is true only when the time t is close to 0. There- fore, the score function estimation in QSM is impre- cise for most values of L. Another method iDEM (Akhound-Sadegh et al., 2024) proposes Va, log p(a;) = 1-k )), ab 8 Q((a)) and the
Theorem 3.5. (The Probability Approximation of Diffu- sion Policy) The log probability of diffusion policy can be approximated by
logPelan) + + Wr. (an.) - Es(a,, (21)
vo(aa) == softmax ( 3/3 Q(ap
where d' = - -()og(27T)) to:r are uniformly spaced
derivation is included in Appendix A.4 for completion. Al- though the expressions of iDEM and our method appear similar and both can approach the true score function as K - 0o, our method does not require computing the gra- dient of the Q-function, which more computationally efficient, especially when the Q-function is evaluated on neural network. Furthermore, the experiments in Section 5.2 demonstrate that the variance of the score estimation in our method is significantly lower than the other two methods that rely on gradient computation, leading to a more stable training process.
timesteps in [tmins tmarl Wt. (ton=-))-(oo)) is the weight
o((a,)
at Li (ar,,, - to4 + ΣE || - 6o(a,, 1/2 is the noise prediction error estimation at ty.
The detailed derivation is provided in Appendix A.5.
3.3. MaxEnt RL with Diffussin Policy
After addressing the critical challenges in training and proba- bility estimation for the diffusion policy, we present the com- plete algorithm for achieving the MaxEnt RL objective with a diffusion policy. Our approach is based on the soft actor- critic framework. We utilize two neural networks: Q9 a) to model the Q-function, and Ca at, 8) to model the noise prediction network for the diffusion policy TH (a(aal))) (s).
3.2. Probability Approximation of Diffusion Policy
Diffusion models approximate the desired distributions by estimating their score function. Although this implicit mod- cling enhances the expressiveness of the model, enabling it to approximate any distribution with differentiable prob- ability density function, it also introduces challenges in computing the exact likelihood of the distribution.
The training process alternates between policy evaluation and policy improvement. In the policy evaluation step, the Q-network is trained by minimizing the soft Bellman error, as defined in Equation 2. Here, the actions a' ~ TTo (ss') are sampled by solving the probability flow ODE in Equation with the noise prediction network co (a,, Qr s), and the log probality log (-Is) is approximated using Equation 21. In the policy improvement step, the noise prediction network is optimized using the loss function in Equation 18 with the training target computed in Equation 17. The pseudocode for our method is presented in Algorithm 1.
Previous study (Kong et al., 2023; Wu et al., 2024) proved that the log- ligeliiolhood of p(as) can be written exactly as an expression that depends only on the true noise prediction target, i.e.
log p(an) = =c 2 Be Ill e- ca an) [2]] da, (19)
where e e= - log(2TT) + & fto (ax)da, with d being the dimension of ag. E ~ NO0, 1). a = vola:Jao + Va(-ai)e, and e² c'(ar, a.) = - -Va(-ar)Va a. log pl at) is the training target of the noise prediction network.
In addition, we adopt several techniques to improve the training and inference of our method:
Truncated Gaussian Noise Distribution for Bounded Ac- tion Space. In RL tasks with bounded action spaces, the Q-function is undefined outside the action space. To avoid evaluating Q-values for illegal actions, the noise distribution in Equation 17 is modified from a standard Gaussian to a
Corollary 3.4. (The Exact Probability of Diffusion Policy) Let eo be a well-trained noise prediction network, iee., it can induce a probability density function po(as) satisfying Eg(a,, at) = --vooaaa V
a. log (at), then
(gao)(an)= = 1 E.e E Eay, ox) [12] do.. (20)
'The minimizers of Equation 18 and 3 will be equal when the exponential of the Q-function can be exactly expressed by the chosen policy set, the capacity of the noise prediction network is preferred to be large if allowed.
This corollary can be inferred from Equation 19, However,
5

---

21: end for Maximum Entropy Reinforcement Learning with Diffusion Policy
Algorithm 1 MaxEnt RL with Diffusion Policy
Gaussian policies. To improve the sample efficiency of SAC, CrossQ (Bhatt et al., 2024) and BRO (Nauman al., 2024) construct larger critic networks and apply suite of regular- ization techniques to stabilize training. MEow (Chao et al., 2024) employs energy-based normalizing flows as unified policies to represent both the actor and the critic, simplifying the training process for MaxEnt RL. This paper highlights the importance of policy representation within the MaxEnt RL framework: a more expressive policy representation enhances exploration and facilitates closer convergence to the optimal MaxEnt policy. Diffusion models, which are more expressive than Gaussian distributions and energy- based normalizing flows and easier to train and sample than energy-based models, present an ideal policy representation that effectively balances expressiveness and the complexity of training and inference.
1: Initialize critic networks Qe Q0 and the noise pre- diction network Eg with random parameters 81, B2, Φ.
2: Initialize target networks 0' <- 01, 0/2 +- 02
3: Initialize replay buffer D>
4: for each iteration do
5: for each sampling step do
6: Sample a ~ (-1s) according to Equation 5
7: Step environment: , I + env(a)
8: Store as...r.s) in
9: end for
10: for each update step do
11:
Sample B transitions (s, a, s') from D
12:
Sample a' ~ (-(s-)s) according to Equation 5
13:
Compute log To(aa (a'ls') using Equation 21
14:
Compute the target Q-value: Q(s, = a( +
(min;=1,2 Qo. (s', a') - Boog TT (a)))'s))
Diffusion Policies for Offiie RL. Offline RL attempts to learn a well-performing policy from pre-collected dataset. Collected by multiple policies, the offline datasets may ex- hibit high skewness and multi-modality. Diffusion Policy (Chi et al., 2023) trains a diffusion model to approximate the multi-modal expert behavior by behavior cloning. To optimize the policy for higher performance, Diffusion-QI. (Wang et al., 2023) combines the diffusion loss with Q- value loss evaluated on the generated actions, CEP (Lu et al., 2023) trains a separate guidance network using Q-function to guide the actions to regions with high Q values, and EDA (Chen et al., 2024b) employs direct preference optimization to align the diffusion policy with Q-function. To improve the training and inference speed of diffusion policy, EDP (Kang et al., 2024) adopts action approximation and efficient ODE sampler DPM-solver for action generation, and CPQL (Chen et al., 2024c) utilizes the consistency policy (Song et al., 2023), a one-step diffusion policy. Due to the lack of online samples, the above approaches require staying close to the behavior policy to prevent -o---iiiittiittion actions whose performances are unpredictable. However, in this paper, we focus on online RL, where online interactions are accessible to correct the errors in value evaluation. There- fore, different techniques should be developed to employ diffusion models in online RL.
15:
Update critics: 0; = arg ming, B (Q0, ((Q,(188) (8,
((s, a))²
16:
Sample t ~ U([Imin: /max) and the noisy action a2 ((ael/a((a)) ((aa)),,o(-a))) ((-a))
17:
Estimate E" 0r, 8) with Equation 17
18:
Update the noise prediction network: Φ
arg minş B Σ ll c0(a,, Okt, 8) - E" (ar, U1t, 8) 11
19:
Updtae target networks: 0, + T0, + (1 7)
20:
end for
truncated standard Gaussian. This modification still gen- erates samples according to the Gaussian function, but all samples are bounded in the specified range.
Action Selection for Inference. Previous studies (Chao et al., 2024; Wang et al., 2023; Mao et al., 2024; Chen et al., 2024b) have found that deterministic policy typi- cally outperforms its stochastic counterpart during testing. Consequently, we employ an action selection technique to further refine the policy after training. Specifically, M ac- tion candidates are sampled from the diffusion policy, and the action with the highest Q-value is selected to interact with the RL environment.
Diffusion Policies for Online RL. In online RL., a key chal- lenge lies in balancing exploration and exploitation. Pre- vious studies (Psenka et al., 2024; Yang et al., 2023; Ding et al., 2024; Wang et al., 2024) apply expressive diffusion models as policy representations to promote the exploration of the state-action space. QSM (Psenka et al., 2024) fits the exponential of the Q-function by training a score network to approximate the action gradient of the Q-function. DIPO (Yang et al., 2023) improves the actions by applying the action gradient of the Q-function and clones the improved actions. QVPO (Ding al., 2024) weights the diffusion loss with the value, assigning probabilities to actions that
4. Related Work
MaxEnt RL A variety of approaches have been proposed to achieve the MaxEnt RL objective. SQL (Haarnoja et al., 2017) introduces soft Q-learning to learn the optimal soft Q-function and trains an energy- based model using the amor- tized Stein variational gradient method to generate actions according to the exponential of the optimal soft Q-function. SAC (Haarnoja et al., 2018) presents the soft actor-critic algorithm, which iteratively improves the policy towards a higher soft value, and provides an implementation using
6

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
(a) Humanoldiv +3
(b) Anth3
Walker
Netr
5000
Metri1"
EAG
CAU
1DD
TD3
Natur
4000
+200
3000
0000
Spauds
2300
Thining Step
Traning Step
Thaining Step
(d) HailCheesah-id
Hopper-u3
Swimmer--33
19900
CAC
400
19300
$300
Eplcdd Retur
8300
0000
TaKning Skip
-10P
Taaking Shop
Takinng Ekop
X10P
Figure 1. I Learning curves on Mujoco benchmarks. The solid lines are the means and the shaded regions represent the standard deviations over five runs.
are linearly proportional to the Q-value. DACER (Wang et al., 2024) optimizes the Q-value loss to generate actions with high Q values and adds extra noise to the generated actions to keep constant policy entropy. Unlike previous approaches, we employ the MaxEnt RL objective to encour- age exploration and enhance policy robustness. Similar to QSM, we train the diffusion model to fit the exponential of the Q-function. However, our Q-weighted noise estimation method more accurate and stable. Furthermore, we in- clude policy entropy when computing the Q-function, which can further promote exploration.
Figure 2. The generated trajectories during the training process. From left to right are trajectories generated after 2k, 4k, and 6k training steps. The goals are denoted by the red points.
from the closest goal. Under the MaxEnt RL objective, the optimal policy to choose one goal randomly and then move toward it. Figure 2 shows the trajectories generated by the diffusion policy during the training process. We can see that MaxEntDP can effectively explore the state-action space and learn a multi-modal policy that approaches the optimal MaxEnt policy.
5. Experiments
In this section, we conduct experiments to address the fol- lowing questions: (1) Can MaxEntDP effectively learn a multi-modal policy in a multi-goal task? (2) Does the dif- fusion policy outperform the Gaussian policy and other generative models within the MaxEnt RL framework? (3) How does performance vary when replacing the Q-weighted Noise Estimation method with competing approaches, such as QSM and iDEM? (4) How does MaxEntDP compare to other diffusion-based online RI. algorithms? (5) Does the MaxEnt RL objective benefit policy training?
5.2. Comparative Evaluation
Policy Representations. To reveal the superiority of apply- ing diffusion models as policy representations to achieve the MaxEnt RL objective, we compare the performance of Max- EntDP on Mujoco benchmarks (Todorov et al., 2012) with other algorithms. Our chosen baseline algorithms include SAC (Haarnoja al., 2018), MEow (Chao et al., 2024). and TD3 (Fujimoto et al., 2018). SAC and MEow two methods to pursue the same MaxRnt RI. objective using Gaussian policy and based normalizing flow policy, and TD3 provides a contrast to the deterministic policy. Fig- ure shows that MaxEntDP surpasses (a-d) or matches (e-f) baseline algorithms on all tasks, and its evaluation variance
5.1. A Toy Multi-goal Environment
In this subsection, we use a 2-D multi-goal environment (Haarnoja et al., 2017) to demonstrate the effectiveness of MaxEntDP. In this environment, the agent is 2-D point mass trying to reach one of four symmetrically placed goals. The state and action are position and velocity, respectively. And the reward is a penalty for the velocity and distance

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
(6) Learning Curve
(8) Standard Deviation of Target Nolso Esiimation
ONE DCM
1980
150
(DEM
ficm
USe
00
07
Baining Stop
Tasking Shep
151 Repparod
Figure 3. Comparison between QNE and two competing methods, QSM and iDEM on HalfCheetah--3 benchmark. (a) Learning curves. (b) Standard deviations target noise (a.k.a, the sealed score function) estimation computed on batch of noisy actions.
Cifal
much smaller than other algorithms. This result indicates that the combination of MaxEntRL and diffusion policy effectively balances exploration and exploitation, enabling rapid convergence to a robust and high-performing policy.
Taring
Figure 4. learning curves of diffusion- based online RL algorithms.
entropy in the Q-function. This implies that the MaxEnt RL objective can benefit policy learning: it not only encour- ages the action distribution at the current step be more stochastic (by fitting the exponential of Q-function), but also encourages transferring to the states with higher entropy (by computing the soft Q-function). Therefore, the MaxFnt RI. objective shows a stronger exploration ability of the whole sate-action space, leading to an efficient and stable training process.
Diffusion Models Training Methods. In this subsection, we demonstrate the advantages of the proposed Q-weighted Noise Estimation method (QNE) on training diffusion mod- els. compared two competing methods, QSM and iDEM. We replace the QNE module with QSM and iDEM to ob- serve performance differences. As shown in Figure 3(a), the performance of QSM and iDEM improves initially but then fluctuates after reaching a high level. This may be due to both methods relying on the gradient computation of the Q-function to estimate the score function. When the Q- D-value is large, its gradient typically varies much across different actions, leading high variance in score function estimation for QSM and iDEM, as illustrated in Figure 3(b). This increased variance causes instability in the training of the noise prediction network. In contrast, QNE exhibits significantly lower variance, and its performance improves steadily throughout the training process.
Ari-d
Ledt 3P
800
Figure 5. The learning curve of MaxEntDP with and without en- tropy in Q-function computation.
Disfusion-based Online RL Algorithms. We also compare MaxEntDP with state-of-the-ar diffusion-based online RL algorithms: QSM, DIPO, QVPO, and DACER. These algo- rithms adopt different techniques to seek balance between exploration and exploitation. Since the performances of different exploration strategies depend quite a lot on the characteristics of the RL tasks, none of the competing meth- ods performs consistently well on all tasks, as shown by Figure 4. In contrast, our MaXEntDP outperforms or per- forms comparably to the top method each task, showing consistent sample efficiency and stability.
6. Conclusion
This paper proposes MaxEntDP, method to achieve the MaxEnt RL objective with diffusion policies. Compared to the Gaussian policy, the diffusion policy shows stronger ex- ploration ability and expressivenes: to approach the optimal MaxEnt policy. To address challenges in applying diffusion policies, we propose Q-weighted noise estimation to train the diffusion model and introduce the numerical integration technique to approximate the probability of diffusion policy, Experiments on Mujoco benchmarks demonstrate that Max- EnDP outperforms Gaussian policy and other generative models within the MaxEnt RL framework, and performs comparably to other diffusion-based online RI. algorithms.
5.3. Ablation Analysis
In addition, we analyze the function of the MaxEnt RL ob- jective by removing the probability approximation module in MaxEntDP. After doing this, we compute the original Q-function rather than the soft Q-function in the policy evaluation step. As shown in Figure 5, the performance de- creases and exhibits greater variance after excluding policy
Limitations and Future Work. Since different RL tasks require varying levels of exploration, we adjust the tem- perature coefficient for each task and keep it fixed during
8

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
training. Future work will explore how to automatically adapt this parameter, making MaxEntDP easier apply in real-world applications.
Chen, H., Lu, C., Wang. Z., Su, H., and Zhu, J. Score regu- larized policy optimization through diffusion behavior. In International Conference on Learning Representations, 2024a.
Acknowledgements
Chen, H., Zheng, K., Su, H.,, and Zhu, J. Aligning diffu- sion behaviors with q-functions for efficient continuous control. In Advances in Neural Information Processing Systems, 2024b.
We extend our gratitude to the ICML reviewers who eval- uated MaxEntDP for their valuable insights and feedback. This work supported by the National Key Research and Development Program of China (No. 20212D0201504).
Chen, Y., Li, H., and Zhao, D. Boosting continuous control with consistency policy. In Proceedings of the Interna- tional Conference on Autonomous Agents and Multiagent Systems, pp. 335-344, 2024c.
Impact Statement
This paper focuses on achieving the MaxEnt RL objective, which particularly effective for reinforcement learning tasks that require extensive exploration policy robust- Beyond advancing RL, our proposed Q Q-weighted noise estimation and numerical integration techniques ad- dress two fundamental issues in diffusion models: fitting the exponential of given energy function and computing exact probabilities. These two modules can be seamlessly integrated into diffusion-basee studies that involve these issues.
Chi, C., Xu, Z.., Feng, S., Cousineau, E.,, Du, Y., Burchfiel, B., Tedrake, R., and Song, Diffusion policy: Visuomo- tor policy learning via action diffusion. The International Journal of Robotics Research, pp. 02783649241273666 2023.
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. In Advances in Neural Information Processing Systems, 2024.
References
Frostig, R., Johnson, M. J., and Leary, C. Compiling chine learning programs via high-level tracing. Systems for Machine Learning, 4(9), 2018.
Ajay, A., Du, Y.. Gupta, A., Tenenbaum, J. B., Jaakkola, T. S., and Agrawal, P. Is conditional generative model- ing all you need for decision making? In International Conference on Learning Representations, 2023.
Fujimoto, S., Hoof, H., and Meger, D. Addressing function approximation error in actor-critic methods. In Interna- tional Conference Machine Learning, pp. 1587-1596. PMLR, 2018
Shhooudd-Sddeg.. T., Rector-Brooks, J., Bose, 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. In International Conference Machine Learning, 2024.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial networks. Communications of the ACM, 63(1 1339-14131144, 2020.
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 sample efficiency and simplicity. In The Twelfth Interna- tional Conference on Learning Representations, 2024.
Haarnoja, T., Tang, H., Abbeel, P., and Levine, S. Re- inforcement learning with deep energy-basee based policies. In International Conference on Machine Learning, pp. 1352-1361. PMLR, 2017.
Bishop, C. M. Puttern Recognition and Machine Learning (Information Science and Statistics). Springer- Veriag,r-Verlag, Berlin, Heidelberg, 2006. ISBN 0387310738.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, Soft actor- critic: Off Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International Con- ference Machine Learning, pp. 1861-1870. PMLR, 2018.
Chao, C.-H., Feng, C., Sun, W.-F., Lee, C.-K., See, S., and Lee, C.-Y. Maximum entropy reinforcement learning via energy-based normalizing flow. In Advances in Neural Information Processing Systems, 2024.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion prob- abilistic models. In Advances in Neural Information Processing Systems, 2020.
Chen, H., Lu, C., Ying, C., Su, H., and Zhu, J. Offline reinforcement learning via high-fidelity generative behav- ior modeling. In International Conference on Learning Representations, 2023.
Ho, J.. Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., and Fleet, D. Video diffusion models, Advances

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
Neural Information Processing Systems, volume 35, pp. 8633-8646, 2022.
Mao, L., Xu, H., Zhan, X., Zhang, W., and Zhang, A. Diffusion-dice: In-sample diffusion guidance for offine reinforcement learning, In Advances in Neural Informa- tion Processing Systems, 2024.
Huang, R., Huang, J., Yang, D., Ren, Y., Liu, L.,, Li, M., Ye, Z., Liu, J., Yin, X., and Zhao, Z. Make-an-audio: Text-to- audio generation with prompt-enhanced diffusion models. In International Conference on Machine Learning, pp. 13916-13932. PMLR, 2023.
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidje- land, A. K., Ostrovski, G., et al. Human-level control through deep reinforcement learning. nature, 518(7540): 529-533, 2015.
Ibarz, J., Tan, J., Finn, C., Kalakrishnan, M., Pastor, P., and Levine, S. How to train your robot with deep reinforce- ment learning: lessons we have learned. The International Journal of Roboties Research, 40(4-5):698-- 721, 2021.
Nauman, M., Ostaszewski, M., Jankowski, K., Mitoš, P., and Cygan, M. Bigger, regularized, optimistic: scal- ing for compute and sample efficient continuous control. Advances in Neural Information Processing Systems, 2024.
Kang, B., Ma, X., Du, C., Pang, T., and Yan, Effi- cient diffusion policies for offline reinforcement learning. In Advances in Neural Information Processing Systems, 2024.
Pan, C., Yi, Z., Shi, G., and Qu, G. Model-based diffusion for trajectory optimization. In Advances in Neural Infor- mation Processing Systems, volume 37, pp. 57914-57943, 2024.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-basee generative models. In Advances in Neural Information Processing Systems, 2022.
Psenka, M., Escontrela, A., Abbeel, P., and Ma, Y. Learning a diffusion model policy from rewards via q-score match- ing. In International Conference on Machine Learning. 2024.
Kaufmann, E., Bauersfeld, L., Loquercio, A., Müller, M., Koltun, V., and Scaramuzza, D. Champion-level drone racing using deep reinforcement learning. Nature, 620 (7976)9982 987, 2023.
Rezende, D. and Mohamed, S. Variational inference with normalizing flows. In International conference on ma- chine learning, pp. 1530-1538, PMLR, 2015.
Kingma, D. P. Auto-encoding variational bayes. arXiv preprint 13Xiv:1312.6111 2013.
Kiran, B. R., Sobh, I., Talpaert, V., Mannion, P., Al Sal- lab, A. A., Yogamani, S., and Pérez, P. Deep reinforce- ment learning for autonomous driving: survey. IEEE Transactions on Intelligent Transportation Systems, 23 (6):4909 4926, 2021.
Saharia, C., Chan, W., Saxena, S., Li, L., Whang. J., Denton, E. L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., et al. Photorealistic text-to-image dif- fusion models with deep language understanding. In Advances in Neural Information Processing Systems, vol- ume 35, pp. 36479-36494 2022.
Kong, X., Brekelmans, R., and Ver Steeg, G. Information- theoretic diffusion. In International Conference on Learn- ing Representations, 2023.
Sendera, M., Kim, M., Mittal, S., Lemos, P., Scimeca, L., Rector-Brooks, J., Adam, A., Bengio, Y., and Malkin, N. Improved off off-policy training of diffusion samplers. In Advances in Neural Information Processing Systems, volume 37, pp. 81016-81045, 2024.
Li, S., Krohn, R., Chen, T., Ajay. A., Agrawal, P., and Chalvatzaki, G. Learning multimodal behaviors from scratch with diffusion policy gradient. In Advances in Neural Information Processing Systems, volume 37, pp. 38456-38479, 2024.
Silver, D., Huang. A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., and Zhu, J. Dpm- solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. In Advances in Neural Information Processing Systems, volume 35, pp. 5775- 2022.
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-2266 PMLR, 2015.
Lu, C., Chen, H., Chen, J., Su, H., Li, C., and Zhu, J. Con- trastive energy prediction for exact energy-guided diffu- sion sampling in offline reinforcement learning. In Inter- national Conference on Machine Learning, pp. 22825- 22855. PMLR, 2023.
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021a.
10

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, volume 32, 2019.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score- basre- generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. Consis- tency models. In International Conference on Machine Learning, pp. 3221 1-32252. PMLR, 2023.
Sutton, R. S. and Barto, A. G. Reinforcement learning: An introduction. Robotica, (7(2):229-233 1999.
Todorov, E., Erez, T., and Tassa, Y. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ inter- national conference on intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Toussaint, M. Robot trajectory optimization using approxi- mate inference. In International Conference on Machine Learning, pp. 1056. PMLR, 2009.
Wang, Y., Wang, L., Jiang, Y., Zou, W., Liu, T., Song, X., Wang, W., Xiao, L., Wu, J., Duan, J., et al. Diffusion actor-critic with entropy regulator. In Advances in Neural Information Processing Systems, 2024.
Wang, Z., Hunt, J. J., and Zhou, M. Diffusion policies as an expressive policy class for offline reinforcement learning. In International Conference on Learning Representations, 2023.
Wu, Y., Luo, Y., Kong, X., Papalexakis, E. E., and Steeg, G. V. Your diffusion model secretly noise classifier and benefits from contrastive training. In Advances in Neural Information Processing Systems, 2024.
Xu, J., Wang, X., Cheng, W., Cao, Y.-P.. Shan, Y., Qie, X., and Gao, S. Dream3d: Zero-shot text-to-3d synthesis us- ing 3d shape prior and text-to-imase to-iea diffusion models. In Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pp. 20908-20918, 2023.
Yang, L., Huang, Z., Lei, F., 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.1 /3/22, 2023.
Zheng, K., Lu, C., Chen, J., and Zhu, J. Dpm-solver-v3: Im- proved diffusion ode solver with empirical model statis- tics. Advances in Neural Information Processing Sys- tems, volume 36, pp- 55502-55542, 2023.
Ziebart, B. D. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. Carnegie Mel- lon University, 2010.

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
A. Theoretic Proofs.
A.1. Proofs for Soft Actor-Critic Algorithm
Our proof is based on the tabular setting, i.c., ||| < 0o, |Al < oo and the replay buffer DD covers all (s, E [S] |A].
The soft Q-function of policy x is defined as:
20
(s,,,, = r((ss a+) + Bpp
log क ((att1/s)1)))
(22)
which satisfies:
(T = r((8, az) + E.p.
(r((8)1,, 3 (
= =r(s az) + Ep.
log
+
Σ
(r(stt) a+t))
Blgg π ((a(41
1=2
(23)
=r(st, (1)) pBppm
BBlog (ar41)/441) r(841, E
(1(044141))4444)) 3 log π. (4
(24)
= r((s,, az) + 1 1-3 log ग (ac+118(+1) + Q" (8:11, a441)]
(25)
Equation 25 is called the soft Bellman equation.
Lemma A.1. (Soft Policy Evaluation) Qo converges to the soft Q-function of FS as L(0) 0.
Proof. Define the soft Bellman operator T as:
r r(s, r((,
'pp 1Q(6',
Blog r(aals')]
(26)
For two Q-function Q and Q', we have
IT"((s,)) - -T'Q'sss = 1/ ap.''' 1Q(s' a') - Q'(s', a')l
(27)
S "yEs' apan' 11((' - Q'(s', a')ll
(28)
S y max 1Q(5 (''s', aa)l
(29)
(d',a')
= = Q Q' 11:00
(30)
Then
|| ITQ-T"Q' 1b0 < - ||
(31)
which proves that the soft Bellman operator 7 is contraction. It has a unique fixed point Q=, When Q-function loss L(0) 0, the Q-function Qo satisfies the soft Bellman equation Qa a) = r(s, a) =.p [Qo(s', a') - 3 log T [s)] for all (s, a) E (S] x [A], indicating that Qa converges to the true soft function of Td.
Lemma A.2. (Soft Policy Improvement) Let The E II and assume Qo = =ex after soft policy evaluation. If Toues is the minimizer of the loss defined Equation 3, then To (8, a) 2 TT (8, for all (s, a) E A with |A| < co.
12

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
Proof. Since the new policy #8et1 is the minimizer of the loss defined in Equation 3, it holds that
exp()Qo(s,))
Toks1 (-18) = arg min DKL
T(-ss)
(32)
Tell
Zo(8)
exp( 10TH (s,
arg min DKL Tell
T(|ss)
(33)
Zd (s)
arg min well
(1
log (a)8)
n (s, a) + log nm0 (8)
(34)
arg min tell
log - (a)8) (-
n (8, a)
+ +ogg n0 (8)
(35)
arg min aell
(-
(a)8( 3 n (8, a) }
(36)
Since Tée E II, we have
1 0 (s, a)) ( s Ba (1x)
1
Ba (1 (+la) logTor_(als))
log To)
(37)
B
According the soft Bellman equation, the Q-function of Toz satisfies
The (8t, at) = r((t at) + B.
1Q (8t+1, a1+1) - log To, (0++1 8841)]
(38)
< r((s ax) + y
o pop. ppp. movts T (81+1, ac+1) 3 log moe+s (at-1(8041)
(39)
= r((, az) + yI.., pmpp ~
[ a.41) - 3 log ग. To (ar+118841))
(40)
+ 7 R R -p,a m0 12 TM (81+2, log To, (01+2/8142]
(41)
(42)
< 5 as),
(43)
which is proved by repeatedly expanding TT using the soft Bellman equation and applying Equation 37. Then the proof for Lemma A.2 is completed.
Theorem A.3. (Soft Policy Iteration) In the tabular setting, let L(Ok) 0 and L(ox) be minimized for each k. Repeated application of policy evaluation and policy improvement, i.e., k oo, The will converge to policy x² such that QF' (s,a) a) 2 (s a) for all Il and (s,a) E S x A with [A] < oc.
Proof. According to Lemma A.I and A.2, when L(Ox) 0 and L(dx) be minimized for each kt we have Vk, * > QF This indicates that the sequence TT is monotonically increasing. Furthermore, the Q-function is bounded since both the reward and entropy are bound. Therefore, when A: + oo, the policy sequence converges to some π². Below we will prove that xt is the optimal MaxEnt policy within ग.
Since n² has already converged, it satisfies
1 log logr(alss) Ea-r(1s))
Q (s.a)
म (-Is) = arg min well
(44)
3
following Equation 36. Then for all ग E II. it holds that
Ea-r (11) log "agm"((ajp (als)
1 Q (s, (0) < Eawn(())
[ (agg())))
3 Q (s, (s,a)
(45)
3
Using the same iterative argument as in the proof of Equation 43, we can derive Q (s, a) 2 Q= a) for all (s. E SxA.. Consequently, x* is the optimal MaxEnt policy within IT. The proof is completed.
13

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
A.2. The Decomposition of the Condition Distribution p(apla:) According to the Bayesian rule, the conditional distribution p(aola;) satisfies:
p(aylar) p(an)p(arlao
(46)
p(x))
x p(aa))p(a
(47)
Maa ((-a))
(48)
where Equation 48 derived by substituting Equation 9 and 10. For the same ao and a,, the probability density of as following the Gaussian distribution N(a:/ (o(o/ f)) equal to the probability density of ag following the Gaussian distribution /N (aol Va(a.) dr, d(-ae) I) up constant compensate for the scale difference between the two a(a))
random variables. Then we have
o(-ar) I).
pla ar) & expl Q(an))
vola)) at o(ar)
(49)
A.3. Estimating the Score Function with Importance Sampling
The score function satisfies
V Va, log p(ae) Bpiaao a.) Va, logr(arlan)
(50)
p(aola:) az))
E agn:Waal
Volue)
ae aa Tj olayg
N(av va(a)) ae o(a))
a(-ai) I)
Va, log logp(ay
(51)
E Value)
(w(a)) V a. log logp(arlan))
(52)
ag M(a
ae [lae]
-w(a)) a Volaaa
E aonNiaal
x(-n 1) [(a))
o(-a()
(53)
Volae)
where the importance ratio w(aa)) = exp( Q(ao)) Z(a.)
with 7(a;) exp( ((xp()Q(a)))))))) Q(a)))
Voia.
a(-a)) T)das being the at a(o.)
a, + Va(-ai) E, then Equation 53 can be rewritten Vola.)
normalizing constant of p(agla;) Let ag
1 -
Vo(a.)
1
Va, logp(ar)
=
BenA((o.)) [u(ao)e]
(54)
va(-or)
k
1
w w(ab)e'
(55)
Vo(-ai) a,)
K
where e!,...,, cK ~ ~N(0,,) ab
Volae) a + Va(-ai)
c.
Vola.)
A.4. The Derivation of the iDEM Method
We provide the derivation of the iDEM method to demonstrate the difference and relationship between Q-weighted noise estimation and iDEM. Our derivation is equivalent the official proof of iDEM, although in a different way.
Since Va, log
ar d I) = Vo(ai)- of -())))
/aaa holds that Jao and Va, logp(arlan) = ((aa))
Vo(a.)
1
a(-ar) I).
Va. logp(ar(an)
Van Nagl volar)
as o(ar)
(56)
Va(as)
14

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
Substitute Equation 56 into Equation 52, have
an g(-ai))
Va. log logpp())) Voluu) NanM Joror Volo ai storee
w(ao)Va, log log/Vaal
o(ar) I]
(57)
1)
va(ar)
Volas) /
volaoVaa Va
Valai)
1
at.
a(-ai) IJda-- a(ar)
(58)
After applying the integration by parts formula, Equation 58 can be expanded to
Va. log p(a)( = vo(as) / (Vao)(a))) NN vo(ai)
1
o(-ai) Ndda
ar
(59)
o(a,)
Since w(ao)= Q((o)) it satisfies that V Va.w(an) w(ao) V au 3 (aao). Then we have %(a,)
1 vo(as) (a(aa) 00 1 Q(av)) Naall va(ar)
/
au a(-or) I)das
Va, logg(as)
(60)
a(ar)
wran Vao B Q((o)
E Vo(as) aagMinal
(61)
Voccg)
Equation 61 appears similar to Equation 53, except that the random variable in the expectation transfers from --weighted noise Q Q-weighted gradient. Utilizing the same weighted importance sampling method as Q-weighted noise, the score function can be estimated by
w(ab) j=1 w(a))
I Q(ab)) 3
Va logp(a:) 2 Va(a.)
Vas
(62)
K
1 3 Q(aa NiV ab Q(ai).
softmax
(63)
Va(an
A.5. Probabity Approximation of Difftusion Policy Using Numerical Integration Techniques
We use numerical integration techniques to estimate the following integral:
1 +oc 2
logpo(ao)
/ Ex
[I]
E ar)) [2 da
(64)
where - log(( log(2 a(ar)day with d being the dimension of ag, (O, a Vol
a(a.) First, using the equation 0g log I-o(as]. we change the integral variable from a to o(as) as o(a)) has narrow integration domain of (0, 1):
d logPa(an) = log(2 (2g( +
2 /. (d.) - Ee [l] ca an) (2
da(a))
(65)
2
(a.))(--a))
In practice, calculate the integral between (a ( for numerical stability, where in our experiments, tmin 1e - 3 and mmax 0.9946. Obtain T + discrete timesteps by setting 4 = 'min + + (4max /min), i 1,.., T. Then the integration domain of (o ) o (o is split into T intervals, where the range of the i-th segment is (o ), o (ot; )} Using the left-hand endpoints to represent the function value of each interval, the integral can be approximated by
d log (ao) x log (Tee
1 Σ (d (an.) EE [l] Es(ar, a. [13]
aaa, ,) - o (cxs.)
(66)
2
(av:)o(-ann
Estimating the noise predicting error Ee [II c - cs(as, ac) [[]]
using Monte Carlo samples, we have
log Ps (ao) log((mre) (2Tre) +
d
1
Σ
d d.
1
Σ
N
é' es(af, ax)) 113
a(an, 1) -a(()) a(ar:)o(-aa,
(67)
N
1=1
j=1
15

---

olu,, ) Maximum Entropy Reinforcement Learning with Diffusion Policy
where e,, ,....,"" E ~ (). af. Vo + Vo(-ou; Je. The equation 67 can be short for
log = d
Σ wi, o(as) Ep(at, at.)
(68)
where d' - 1og(2 wt, ((an_))-0(00) is the weight tús Co(ar, as.) ES el co(al, and) 112 the
noise prediction error estimation at f;-
B. Supplementary Related Work on Diffusion-based Energy Models
line of work focuses on applying the expressive diffusion models to approximate the exponential of a given energy function. QSM (Psenka et al., 2024) trains the score function by aligning with the gradient of the energy function. iDEM et al., 2024) proposes weighted sum of the gradient of the energy function to estimate the true score function. These two approaches, which are based gradient computation, suffer from a large estimation variance and demonstrate training instability when used for diffusion policy optimization, evidenced in Section 5.2. Recently, work (Sendera et al., 2024) considers the Mulee-Maruyamn samplers diffusion models as continuous generative flow networks (GFlowNets), and exploits the trajectory balance objective to train diffusion models. In this method, a replay buffer is used store the sample generation trajectories of diffusion models, which may cause a heavy memory burden. The model-based diffusion (Pan al., 2024) proposes the Monte Carlo estimation for computing the score function and uses the Monte Carlo score ascent generate samples following the Boltzmann distribution of given function. The model-basee diffusion is similar to our QNE method, however, QNE has several properties that matter in RL training:
We use parameterized network to approximate the scaled score function, while the model-based diffusion needs to compute the score function using Monte Carlo estimation when generating samples. Therefore, sample generation of model-based diffusion conssunnn,, which will slow the training speed of the RL algorithms.
We adopt ancestral sampling in DDPM to generate samples, that are more diverse than that of Monte Carlo score ascent used in model-based diffusion.
We propose to modify the standard Gaussian to the truncated Gaussian in QNE to model the action distribution with a bounded action space. However, model-based diffusion can not address such a bounded distribution.
These properties make QNE well-suited for the optimization of diffusion policy.
C. Experimental Details
C.1. Hyperparameters Settings
All experiments in this paper are conducted on GPU of Nvidia GeForce RTX 3090 and a CPU of AMD EPYC 7742. Our implementation of SAC, MEow, TD3, QSM, DACER, QVPO, and DIPO follows their official codes: http:///ithbb..ommoossiwww/sssss
matching rl, https//gtthub.com//pyyyyyyDDDDDt Diffusionn-with-nnee Onwittt---nnnneeee RL., http://githbb.co//waaass 9/qvpo, and http://iiiibbcoo//llllaaaaintttt The shared hyperparameters of all algorithms are listed Table 12.
https://ittub..oo//tujjjm/////
http://ggthab..co//Aessssttttttt
C.2. Training Time
The training time for all algorithms presented in Table 2. Leveraging the computation efficiency of JAX (Frostig et al., 2018) and the parallel processing capabilities of GPU, MaxEntDP demonstrates high training efficiency compared to competing methods, only behind TD3 and QSM. This highlights its advantage for real-world applications that require high computation efficiency.
Whee comparing with other diffusion-basee algorithms, MaxEntDP 3-layer ML.Ps the actor and critic networks following the default settings of these algorithms. In other experiments, 2-layer MLPss are used as they can already attain good performance.
16

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
Table
The shared hyperparameters of all algorithms.
Hyperparameler
MaxEntDP
SAC
MEow
TD3
QSM
DACER
QVPO
DIPO
Batch size
256
256
256
256
256
256
256
256
Discount factor 7
0.99
0.99
0.99
0.99
0.99
0.99
0.99
0.99
Target smoothing coefficient r
0.005
0.005
0.005
0.005
0.005
0.005
0.005
0.005
No. of hidden layers
2
2
2
2
3
3
3
No. of hidden nodes
256
256
256
256
256
256
256
256
Actor learning rate
3--4
3e-4
3e-4
3e-4
3e.4
3e-4
3e.4
3e-4
Critic learning rate
3e-4
3e-4
3e-4
3e-4
3e-4
3e-4
3e.4
3e-4
Activation
mish
relu
relu
relu
mish
mish
mish
mish
Replay buffer size
lc6
le6
1e6
1e6
lc6
lc6
1c6
lc6
Diffusion steps
20
N/A
N/A
N/A
20
20
20
20
Action candidate number
10
N/A
N/A
N/A
N/A
N/A
32
N/A
Table 2. The comparison of training time HalfChectah-v3 benehmark.
Algorithm (2-layer MLP)
MaxEmDD (jax)
SAC
MEow
TD3
Training time (h)
3.9
4.6
11
1.7
Algorithm (3-layer MLP)
MaxEmDD (jax)
QSM (jax)
DACER (jax)
QVPO
DIPO
Training time (h)
1.9
5.9
22.6
55.6
D. Supplementary Experiments
D.1. Hyperparameter Analysis
In this subsection, we analyze the effect of different hyperparameter settings on the performance:
Sample Number for Q-weighted Noise Estimation. The Q-weighted noise estimation can be seen as a weighted importance sampling method to estimate the training target of the noise prediction network. With more samples, the estimation will be more accurate and less varied, which benefits the training of diffusion policy. This consistent with the observation in Figure 6(a) that the performance will be better with a larger K. We select K = 500 since it can obtain good performance and cause relatively small computation costs.
Sample Number for Probability Approximation. For probability approximation of diffusion policy, several Monte Carlo samples are utilized to estimate the noise prediction error at each diffusion timestep. This sample number also preferred be large for higher accuracy and less variance. The performance of different sample numbers N is shown in Figure 6(b). We set N 50 after trading off performance and computation efficiency.
Diffusion Steps. Due the discretization error of ODE solvers, the actual distribution of generated actions may be different from the diffusion policy induced by the noise prediction network. Therefore, when the diffusion steps T is small, the neg-nggli discretization error will disrupt the training process and lead performance drop. As shown in Figure 6(c), the performance is higher with larger T. We choose T 20 as the default setting for a balance between performance and computation efficiency.
Candidate Number for Action Selection. By selecting the action with the highest Q-value among several action candidates, the action selection technique can further improve the performance of the diffusion policy when testing. Figure 6(d) demonstrates that larger number of action candidates will result in better performance, Consequently, we set = 10 by default.
Temperature Coefficient. The temperature coefficient 3, which determines the exploration strength, is an important parameter in the MaxEnt RL framework. Since the difficulty and reward scales vary across different tasks, different need to be set for different tasks. We sweep over [0.01, 0.02, 0.05, 0.1, 0.2] to find the optimal setting for each task, displaying the results in Figure?? 7. The temperature coefficient selected for each task is listed in Table 3.
17

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
(a)
17000
K-4000
12000
N-100
E-000
Nad0
12000
10000
Ke100
K-500
Edecste Reeum
Edtoddo Rosum
0.6
0.2
Training Step
x12
Training Shep
X10
(c)
(d)
12000
12000
15000
1-2
Epttods ftetum
Eptoddo Retum
688
Training Step
Traiinng Stop
X1PP
Figure Learning curves of different parameter settings on HalfCheetah-v3 benchmark. (a) Testing different numbers of samples K for Q-weighted noise estimation. (b) Testing different numbers of samples N for probability approximation. (c) Testing different diffusion steps T. (d) Testing different candidate numbers M for action selection.
Table The temperature coefficients adapted for each benchmark.
Benchmark
Temperature cocfficient
Ant-v3
0.05
HalfChcetah-v3
0.2 Hopper-v3
0.05
Humanoid-v3
0.02
Swimmer-v3
0.01
Walker2d-v3
0.01
D.2. Multimodal Policy Learning on the Challenging AntMaze Benchmarks
We adopt the AntMaze benchmarks proposed in DDiffPG (Li et al., 2024) to test the multi-modal policy learning ability of MaxEntDP on the challenging high-dimensional RL tasks. In this environment, the agent is quadruped robot trying reach the specified goals. Instead of the sparse reward employed in DDifiPG, we use dense reward, a penalty for the distance from the closest goal, guide policy learning. We demonstrate the trajectories generated by MaxEntDP and SAC after 1M environment interactions in Figure 8. MaxEntDP can learn diverse behavior modes even in the challenging high-dimensional tasks, while SAC fails to learn different solutions. In addition, we visualize state coverage through the training process for MaxEnDD and SAC, showing the results in Figure 9, We can see that MaxEnIDP can explore multiple behavior modes at the same time, while SAC focuses only on simple mode. This reveals the importance of using the expressive diffusion policy for efficient exploration and multimodal policy learning.
D.3. Comparative Evaluation on the Meep Control Suite
We test MaxEntDP on high-dimensione tasks on the DeepMind Control Suite benchmarks. The performance comparison with SAC, MEow, and TD3 is displayed Figure 10. MaxEntDP outperforms other baseline algorithms on these challenging high-dimensional RI. tasks.
18

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
Art-v3
HalfChoetah-v3
Hopper-v3
1.82
4000
5030
A.01
3.0.05
10000
3+0.05
A-2.1
2503
5.0.00
Episode Reeum
Retum
2500
6000
2003
2000
4000
2020
260-1
503
240.00
2-0.00
250.0
1.0
0.0
Step
X100
Training
x10
Training
Humanotd-vv
(e) Swimmer-va
Walker2d-v3
3-0.2
3-0.2
5.0.2
0000
3+0.1
2-0.1
5-0,1
A-1115
$.2.05
340.05
4003
W
4030
80
Eoloode Retur
8000
80
Eptoads
40
2000
20
1030
1003
Step
Training
Training Step
Figure 7. Learning curves of different temperature coefficients on Mujoco benchmarks.
D.4. Testing the Accuracy of the Proposed Numerical Integration Technique on Probability Approximation
In theory, the numerical integration will be accurate when the diffusion step T and the number of samples N for probability approximation become large enough, according to the Law of Large Numbers. To exhibit the accuracy of different T and N, we conduct experiments on a simple 2-D toy example where the target distribution p(x) is a mixture of four Gaussian distributions with equal weights. Therefore, we set ((x) log p(x)) and utilize the QNE method proposed in our paper to train diffusion model. We display the approximation results of different T and N Figure 11. The setting T 20, N 50 used in the paper can provide effective probability approximation for the diffusion policy. When the samples are less (T 20, N 20), although there is nen-negligible error the ground truth, the numerical integration method can still assign higher values for the region with higher probability. In this case, the estimated log probability can be considered as a kind of intrinsic curiosity reward to promote the exploration of the action region with low policy probability.
19

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
SAC.
MaxEmtDP
Antmaae-vv
Antmaze-v2
Antmaze-v3
Antmaze-vd
Figure 8. Trajecories generated by MaxEntDP and SAC after environment interactions Antmaze benchmarks.
SAC
MaxEntDP
Antmaze-v1
Antmaze-v2
Antmaze-v3
Antmaze-v4
Figure State coverage MaxEntDP and SAC after environment interactions Antmaze benchmarks. MaxEntDP can explore different behavior modes the same time and show broader state coverage than SAC, exhibiting efficient exploration of the high- dimensional state-action space.
1s) Dog
Dog
(c)
Ma
McxDaCPP
MadrriiP
64G
SMC
SHC
Mizw
530
200
300
Eoboce
100
32
05 0.2
Stop
Figure Learning curves on DeepMind Control suite. solid lines are the means, and shaded regions represent the standard deviations over five runs.
20

---

Maximum Entropy Reinforcement Learning with Diffusion Policy
N=10
T-20
T-50
(4) True log probubility
(b) Estimated log probubility
Figure 11. The prohability approximation using numerical integration method 2-D toy example with different diffusion steps and sample numbers N. The target distribution is mixture of four Gaussian distributions, whose means are (-0.5, -0.5), (-0.5, 0.5), (0.5, 0.5) and (0.5, -0.5). The standard deviations and weights of the four components the same, which The setting in the paper (I' = 20, N 50) can provide an effective approximation for the true log probability. When fewer samples (T 20, N 10) are used, despite some estimation errors, our method can still assign higher values to high high-probaaee probability regions.
0.1 and 0.25, respectively.
21


Reviews:
[{"review_id": "XzQePuaOS9", "reviewer": "Reviewer_fZmT", "summary": "This paper introduces MaxEntDP, a new online diffusion-based RL algorithm that integrates diffusion models into the maximum entropy framework. The method proposes a Q-weighted noise estimation for policy improvement and use numerical integration to estimate action probability for policy evaluation. Experiments are conducted on MuJoCo environment, comparing with classic MaxEnt algorithms and other online diffusion-based algorithms, verifying the effectiveness of the proposed method.", "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": "Thank you for your valuable feedback and suggestions. Below we will address each concern raised in the review.\n>**Q1: Can MaxEntDP learn different solutions in high-dimensional multi-goal tasks from DDiffPG? And provide evidence for improved exploration.**\n    \nWe tested MaxEntDP and SAC on four versions of the AntMaze environment in DDiffPG (using dense rewards) and visualized the generated trajectories in Figure 1 of https://anonymous.4open.science/api/repo/pics-E459/file/antmaze.pdf?v=eebbbe1f. The results confirm that MaxEntDP can learn diverse behavior modes even in challenging high-dimensional RL tasks, while SAC fails to learn different solutions.\nIn addition, we visualized state coverage for MaxEntDP and SAC (see Figure 2 in the above link). The results show that MaxEntDP explores multiple behavior modes and exhibits broader state coverage. This highlights the advantage of using a diffusion policy for efficient exploration.\n    \n>**Q2: No evidence supporting numerical integration as an effective approximation.**\n    \nAccording to the Law of Large Numbers, numerical integration accuracy improves with larger diffusion steps $T$ and sample numbers $N$. To exhibit the accuracy of different $T$ and $N$, we conducted experiments on a 2D toy example (a mixture of four Gaussians) and presented the results in https://anonymous.4open.science/api/repo/pics-E459/file/probability.pdf?v=86f81805. As shown in the figure, our setting in the paper ($T=20, N=50$) provides an effective probability approximation. Moreover, when fewer samples ($T=20, N=10$) are used, despite some estimation errors, our method still assigns higher values to high-probability regions, which can be considered as an intrinsic curiosity reward to promote exploration on the action region with low policy probability.\n    \n>**Q3: The claim that MaxEntDP achieves the optimal MaxEnt policy with sufficient model capacity.**\n    \nWe do not make such a claim in the paper. Instead, we argue that incorporating a diffusion policy into MaxEnt RL improves exploration and moves the policy closer to the optimal solution, as proposed in the abstract. The 2D toy example in the paper reveals that MaxEntDP can explore the state-action space efficiently and finally learn a multimodal policy. The Mujoco experiments show performance improvements over other generative models and diffusion-based approaches, supporting our main claim.\n    \n>**Q4: Concern about Monte Carlo sampling.**\n \nYour concern is reasonable since a higher estimation accuracy requires more Monte Carlo samples. However, in our experiments, we find that a small sample number of 1000 (for both diffusion policy optimization and action probability estimation) is enough to obtain a good performance. Additionally, as the samples can be processed parallelly and GPU throughput continues to improve, Monte Carlo sampling is not a bottleneck.\n\n>**Q5: Necessity of action selection and its application to all baselines.**\n    \nAction selection is crucial because our method approximates the exponential of the Q-function, generating both high- and low-return actions. While beneficial for exploration, this can reduce test-time performance. To mitigate this, we apply action selection (only in testing) to pick the action with the highest Q-value. Similar techniques are used in SAC (take the Gaussian mean) and EBFlow. To ensure fairness, we applied action selection to all diffusion-based baselines (with a candidate number of 10) and reported results in https://anonymous.4open.science/api/repo/pics-E459/file/comparison_diffusion_selection.pdf?v=a952c896. Our MaxEntDP continued to demonstrate high sample efficiency and stability across all tasks.\n    \n>**Q6: Add experiments in more challenging environments.**\n    \nWe compared MaxEntDP with SAC on the DMControl and MyoSuite benchmarks. Results are shown in https://anonymous.4open.science/api/repo/pics-E459/file/DMC_myo.pdf?v=538c572b.\n    \n>**Q7: Include a comparison with BRO.**\n    \nBRO improves SAC with advanced network architectures, regularization, and optimistic Q-value estimation. Since BRO's enhancements can be applied to any baseline (non-diffusion and diffusion-based) and MaxEntDP, a direct comparison would be unfair. However, integrating BRO's improvements into MaxEntDP is an interesting future direction.\n    \n>**Q8: Performance gains of MaxEntDP over existing diffusion-based baselines.**\n    \nNone of the competing diffusion-based methods consistently outperforms across all tasks—e.g., DACER struggles on Walker2d, and others underperform on HalfCheetah. In contrast, MaxEntDP exhibits consistent sample efficiency and stability across all tasks.\n    \n>**Q9: Missing reference to DDiffPG.**\n    \nThanks for your valuable suggestion. We will cite it in future revisions. While DDiffPG explicitly distinguishes different behavior modes and learns a Q-function for each mode, MaxEntDP employs a single Q-function, making it simpler and more computationally efficient.", "claims_and_evidence": "Overall, the authors make several claims but only partially supported. First, in Section 3, the authors claim that MaxEnt RL with expressive diffusion model can capture multimodal behaviors. This is evidenced by the multi-goal 2D toy example. However, I am suspicious how this will generalize to high-dimensional tasks because the goals in the toy example is easy to explore. For example, in DDiffPG [1] it provides several high-dim multi-goal tasks, and it is interesting to see if MaxEnt RL can learn different solutions too. Second, the authors claim that by maximizing the entropy MaxEnt RL improves the exploration. However, there is no experiment or analysis to support it. It would be great to include some visualizations or state coverage in the experiment. Moreover, the authors say numerical integration provides an effective approximation and MaxEntDP achieves the optimal MaxEnt policy given sufficient model capacity. Both are not supported without any evidence or comparison. \n\n[1] Li, Steven, et al. \"Learning multimodal behaviors from scratch with diffusion policy gradient.\" Advances in Neural Information Processing Systems 37 (2024): 38456-38479.", "methods_and_evaluation_criteria": "The paper proposes a novel Q-weighted noise estimation for diffusion policy optimization and an approximation for the diffusion model action probability. One problem is that both parts require monte-carlo sampling, which possess a challenge on estimation accuracy. In appendix the authors provide ablations on the key hyper-paramaters, in which the high variance verifies the instability. Moreover, the paper provides two tricks to stabilize and improve the performance. The action selection for inference is a common approach in offline RL but unusual in online setup. However, most of baselines do not use this action selection for a fair comparison as shown in appendix. \n\nThe method is evaluated only on MuJoCo tasks, which has almost no exploration challenge. I would suggest to experiment on more challenging environments, e.g., myosuite to showcase the exploration benefit from the MaxEnt framework.", "theoretical_claims": "I have checked the proofs in the paper.", "experimental_designs_or_analyses": "The experiments confirm the effectiveness of the proposed method. However, it does not include a comparison with the current state-of-the-art BRO [2]. Additionally, its performance shows only a marginal improvement over existing diffusion-based baselines.\n\n[2] Nauman, Michal, et al. \"Bigger, regularized, optimistic: scaling for compute and sample-efficient continuous control.\" arXiv preprint arXiv:2405.16158 (2024).", "relation_to_broader_scientific_literature": "The paper is highly relevant to RL and diffusion model.", "essential_references_not_discussed": "One key claim in Section 3 and Section 5.1 is to learn a multi-modal policy, however one related work DDiffPG is missing [1].\nWhen comparing to learning curve, the current SOTA BRO is missing [2].", "other_strengths_and_weaknesses": "The paper is well-written and easy-to-follow.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "See above.", "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": "ODqrghZBmk", "reviewer": "Reviewer_7ggH", "summary": "This paper introduces Maximum entropy Reinforcement Learning with Diffusion Policy (MaxEntDP). More specifically, this method proposes solutions to the well-known problems on how to approximate the target distribution composed of the exponential of the Q-function and how to calculate the log likelihoods of the marginal distribution of diffusion models. Both problems are essential to train in the maximum entropy RL framework and need careful treatment when using diffusion models.\n\n## Updated after rebuttal \nI adjusted my score after the authors have clarified my questions and concerns", "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": "Thank you for your thoughtful review and constructive feedback. Below, we will address your comments and hope that this clarifies the context of our work.\n\n>**Q1: Consider more sophisticated RL tasks.**\n\nWe compare MaxEntDP with SAC on DMContorl and MyoSuite benchmarks. The results are shown in https://anonymous.4open.science/api/repo/pics-E459/file/DMC_myo.pdf?v=538c572b.\n\n>**Q2: Justification of Eq 59.**\n\nWe are sorry for omitting the parentheses on $\\bigtriangledown_{a_0} w(a_0)$ in Eq 59, which causes trouble in understanding. In fact, we apply the integration by parts formula (line 777) to derive Eq 59, i.e., $\\int_{\\Omega} v \\nabla u \\, d\\Omega = \\int_{\\Gamma} u v n \\, d\\Gamma - \\int_{\\Omega} u \\nabla v \\, d\\Omega$ (consider $w(a_0)$ as $v$ and $N(a_0)$ as $u$; the term $\\int_{\\Gamma} u v n \\, d\\Gamma$ is zero since $w(a_0)$ and $N(a_0)$ decay rapidly at infinity).\n\n>**Q3: Justification of Eq 65.**\n\nIn Eq 65, we change the integration variable from $\\alpha_t$ to $\\sigma(\\alpha_t)$ using the equation $\\alpha_t=\\log \\frac{\\sigma(\\alpha_t)}{1-\\sigma(\\alpha_t)}$, as illustrated in line 807. Then the integration domain of Eq 65 has been changed to the varying range of the new integration variable $\\sigma(\\alpha_t)$, which is $(0,1)$.\n\n>**Q4: The accuracy of the log-likelihood approximation.**\n\nDue to the length limit of rebuttal, please refer to our response to the Q2 of reviewer fZmT.\n    \n>**Q5: The paper lacks recent work such as BRO and CrossQ.**\n\nWe will cite them in future versions. Briefly, the two methods propose some improvements to the SAC algorithm. BRO develops the advanced BroNet architecture, regularization, and optimistic upper-bound Q-value approximation. CrossQ removes the target networks and adopts Batch Normalization for high sample efficiency. Since the improvements of the two methods are also compatible with our MaxEntDP algorithm, it is interesting to see how the performance of MaxEntDP can be enhanced by combining these improvements.\n\n>**Q6: Why can the noisy actions** $p(a_t)$ **be substituted with other distributions?**\n\nAs shown in Eq17, the training target $\\epsilon^{*}$ of the noise prediction network depends only on its input $(a_t,\\alpha_t)$, which means that we can minimize the L2 loss for each input point. Thus, the noisy action samples from other distributions can also be used for training. This can be seen as a kind of \"off-policy\" training. Moreover, in the paper we diffuse the action samples in the replay buffer to obtain the surrogate distribution of $p(a_t)$ (line 295). As the training progresses, the actions generated by the diffusion policy will be closer to the target distribution $p(a_0)$, then the diffused distribution of these actions will also be closer to the true distribution of $p(a_t)$, consequently, the training will be more and more \"on-policy\".\n\n>**Q7: The reason for using a truncated standard Gaussian in RL tasks with bounded action spaces.**\n\nWe would like to point that MaxEntDP and SAC model the bouded policy distribution in different ways. SAC outputs an unbounded distribution and transforms it to a bounded one by applying a tanh function, while MaxEntDP directly models a bounded distribution following the common practices in the image generation domain. To learn this bounded distribution, we need to change the target distribution to $p(a_0)\\propto\\exp(\\frac{1}{\\beta}Q(a_0))I_{a_0\\in[lb,ub]}$, where $I_{a_0\\in [lb,ub]}$ is an indicator function to check if $a_0$ lies within the bound. Similar to lemma 3.1, we can obtain that the reverse transition distribution has be changed to $p(a_0|a_t)\\propto\\exp(\\frac{1}{\\beta}Q(a_0)) I_{a_0\\in[lb,ub]} \n     \\mathcal{N}(a_0|\\frac{1}{\\sqrt{\\sigma(\\alpha_t)}} a_t,\\frac{\\sigma(-\\alpha_t)}{\\sigma(\\alpha_t)}I)$. Taking the indicator function and the Gaussian distribution together, we can derive a truncated Gaussian distribution with a bound of $[lb,ub]$. Then, the conditional distribution $p(a_0|a_t)$ can be seen as a truncated Gaussian distribution of $a_0$ weighted by the exponential of the Q-function. Therefore, for diffusion policy training with a bounded action space, we only need to modify the standard Gaussian distribution in the QNE method to a truncated Gaussian distribution.\n\n>**Q8: Performance comparison with other Diffusion-Based methods.**\n\nWe would like to emphasize that none of the competing diffusion-based methods performs consistently well on all tasks (DACER struggles on Walker2d, and others underperform on HalfCheetah). However, our MaxEntDP displays consistent sample efficiency and stability across all tasks.\n\n>**Q9: The meaning of \"training step\" in the learning curves; K-L in line 119.**\n\nWe use \"training step\" because both DACER and QVPO use it as the x-axis of the learning curve, which means the training steps of the actor/critic network. Since we use a UTD of 1, it is the same as the number of environment interactions. And we will modify the 'K-L' in line 119 in revised versions.", "claims_and_evidence": "This work proposes a solution on approximating the likelihood of the marginal distribution of diffusion models and proposes an objective to train the policy's score function without evaluating the gradient of the Q-function. \n\nThe claims are backed up with mathematical derivations and proofs and ablation studies.", "methods_and_evaluation_criteria": "The considered benchmark makes sense, though there are more sophisticated tasks in RL that should also be considered, e.g., high-dimensional control tasks from the DeepMind control suite are harder learning tasks that usually have a longer horizon than the mujoco environments from gym.", "theoretical_claims": "I did not check the proofs for Theorem 2.1\n\nI checked the proofs to Theorem 3.1, Theorem 3.2, Theorem 3.3, which seem to be fine to me. \n\n\nI also checked the proof for the iDEM derivation (A4) and have a concern:  \nEq 59 pulls the gradient w.r.t. a_0 in front of w(a_0), but this is mathematically not correct as w(a_0) is depending on a_0 as well. How is this justified? \n\n\nI also checked the proof for Theorem 3.5. but have difficulties understanding it: How is the replacement of the integration domain in Eq. 65 justified? I don't understand this step. Intuitively, it is a very strong change integrating from 0 to 1 instead of from -infinity to infinity. It is difficult to judge whether the approximation of calculating the log likelihood is a good approximation. My concern is also strengthened by the fact that the learning curves with the entropy bonus are actually not significantly better than the learning curves without the entropy (Fig. 5 in Section 5).", "experimental_designs_or_analyses": "From my understanding there are no experimental environments designed specifically for this work.", "relation_to_broader_scientific_literature": "The work relates well to prior findings and distinguishes itself from prior works.", "essential_references_not_discussed": "To my knowledge, all essential references are discussed w.r.t. diffusion models. However, the paper lacks recent works in the field of maximum entropy reinforcement learning such as Bigger Regularized optimistic (BRO) [1] and CrossQ [2].\n\n[1] Nauman, Michal, et al. \"Bigger, Regularized, Optimistic: scaling for compute and sample efficient continuous control.\" The Thirty-eighth Annual Conference on Neural Information Processing Systems.\n\n[2] Bhatt, A., et al. \"CrossQ: Batch Normalization in Deep Reinforcement Learning for Greater Sample Efficiency and Simplicity.\" International Conference on Learning Representations (ICLR). 2024.", "other_strengths_and_weaknesses": "In general, the paper is well written and the reader can follow it. It is well-motivated and considers an important aspect of research for using diffusion models in reinforcement learning. \nHowever, there are several points that need clarification: \n\n- Lines 220 following mention that the distribution for the noisy actions p(a_t) is unknown but can be substituted with other distributions with full support. What does this mean? How is this justified given that the expectation is not correctly evaluated? Covering full support can be problematic for high-dimensional action spaces. This seems to be a major issue but was not discussed. Additionally, the considered environments do not consider high-dimensional action spaces, so it is hard to judge whether this is a bottleneck. \n\n- The paper states that a truncated standard Gaussian is used to keep the sampled actions within the bounds. This, however, changes the likelihood to my understanding, but no correction is mentioned. For example, SAC applies the change of variables and corrects the likelihood in this way. This is however not considered here. \n\n- While there are some tasks where the proposed method performs better compared to Gaussian-based policies, it does not perform better than other Diffusion-Based methods.   \n\n- More experiments on more sophisticated tasks with longer horizons and higher dimensional action spaces, such as those from the Deepmind control suite, myo suite, or the humanoid bench, would strengthen the paper. \n\n- The learning curves plot the \"training step\" but it is never mentioned what this means. I am confused as RL papers usually plot the number of environment interactions instead of some other metric. This needs clarification. \n\n- Minor note: K-L instead of KL in line 119 right column", "other_comments_or_suggestions": "please see my comments from before.", "questions_for_authors": "please see my comments from before.\n\n\n------ I adjusted my score after the authors have clarified my questions and concerns ------", "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": "dcnMbMrBjU", "reviewer": "Reviewer_TSPJ", "summary": "This paper focuses on adapting diffusion-based policies to maximum entropy reinforcement learning (MaxEnt RL) for better exploration. The primary obstacles are: 1) policy evaluation involves computing the log-probability over the clean actions, which for diffusion policies is non-trivial; and 2) policy improvement requires aligning the diffusion with the Boltzmann distribution of the Q-value functions, which we cannot readily access. To tackle these challenges, this paper proposes Q-weighted Noise Estimation and borrowed the probability calculation method from previous literature, which leads to a coherent algorithm that accommodates diffusion-based policies while also strictly corresponding to MaxEnt RL in theory. The evaluation is conducted on classical MuJoCo locomotion tasks, and the proposed method, MaxEntDP demonstrates improvements over both diffusion-based and non-diffusion algorithms.\n\n## update after rebuttal\nThe authors demonstrated in their last response that MaxEntDP achieves significantly better performance over SAC on dog tasks from DMControl, which validates the benefit of combining diffusion policies with the MaxEnt RL framework. In light of this, I decided to raise my score, and I encourage the authors to conduct experiments on the DMControl benchmark with other baseline algorithms as well. \n\nHowever, some theoretical concerns still remain. As I mentioned in my initial review, the assumption that $\\epsilon_\\phi$ is a well-trained noise prediction network is too strong since the policy improvement and the policy evaluation steps are interleaved. This compromises the accuracy of the estimated log-probability and may cause potential instabilities during training. \n\nTherefore, given these observations and discussions, I will keep my evaluation as Weak Accept.", "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": "Thank you for your valuable feedback and suggestions. Here, we aim to address the questions raised in the review.\n\n>**Q1: Include results on tasks like MetaWorld or DMControl.**\n\nWe compare MaxEntDP with SAC on DMControl and MyoSuite benchmarks. The results are shown in https://anonymous.4open.science/api/repo/pics-E459/file/DMC_myo.pdf?v=538c572b. \n\n>**Q2: The importance-weighted sampling technique in the paper is similar to iDEM.**\n\nWe would like to emphasize that although the expression of our QNE method appears similar to iDEM [1], our method exhibits significantly lower estimation variance and steady performance improvement throughout the training process over the same method replaced with iDEM (Figure 3). In addition, our method does not require gradient computation\nof the Q-function as in iDEM and thus is more computationally\nefficient. These advantages demonstrate the superiority of our QNE method on diffusion policy optimization.\n\n>**Q3: Concern on the assumption of Corollary 3.4.**\n\nWe use the target Q network to train the diffusion policy. Since we adopt the Exponential Moving Average (EMA) update to smooth the change of the target Q network, it is not difficult to learn a well-trained policy network.\n\n>**Q4: Keep the hyper-parameter tuning effort the same across baseline algorithms.**\n\nFor fair comparison, we unify shared hyperparameters (batch size, discount factor, the depth and width of hidden layers, learning rate and replay buffer size) across baseline algorithms. Other hyperparameters strictly follow the settings specified in the original paper or codebase, which are already tuned by their authors. \n\n>**Q5: One related work model-based diffusion [2] is missing.**\n\nWe will cite it in revised versions of the paper. The model-based diffusion proposes the Monte Carlo estimation for computing the score function and uses the Monte Carlo score ascent to generate samples following the Boltzmann distribution of a given function. This method is similar to our QNE method, however, QNE has several properties which matter in RL training: \n1. We use a parameterized network to approximate the scaled score function, while the model-based diffusion needs to compute the score function using Monte Carlo estimation when generating samples. Therefore, sample generation of diffusion-based diffusion is time-consuming, which will slow the training speed of the RL algorithms.\n2. We adopt the ancestral sampling in DDPM to generate samples, that are more diverse than that of Monte Carlo score ascent used in model-based diffusion.\n3. We propose to modify the standard Gaussian to the truncated Gaussian in QNE to model the action distribution with a bounded action space. However, model-based diffusion can not address such a bounded distribution.", "claims_and_evidence": "Yes, most of the claims in the paper are supported by theory or empirical evidence.", "methods_and_evaluation_criteria": "MaxEntDP inherits the popularized MaxEnt RL framework, which is known to be effective for online RL due to enhanced exploration brought by entropy regularization. As for the benchmarks, although Gym-MuJoCo is the most commonly used benchmark in online RL, it is comparatively simple due to dense rewards and lower DoF.  I encourage the authors to also include results on tasks like MetaWorld or DMControl, where the task is more complex and the necessity of exploration is more pronounced.", "theoretical_claims": "First of all, the theory largely builds upon the analysis from MaxEnt RL. For the policy improvement step, this paper adopts an importance-weighted sampling technique similar to iDEM to reversely sample from the posterior $p(a_0|a_t)$. To reduce the variance of importance sampling, importance re-sampling is used to derive a biased objective but with reduced variance. The variance reduction property has been demonstrated by previous literature. \n\nFor the policy evaluation step, the analytical calculation for log probability is also borrowed from previous literature. However, Corollary 3.4 holds under the assumption that $\\epsilon_\\phi$ is a well-trained noise prediction network. During RL training, **such an assumption may not hold** since the policy improvement and evaluation are interleaved and we cannot guarantee that the policy network is well-trained. Under this circumstance, Eq. 20 does not exactly correspond to the log-probability of the actual action distribution.", "experimental_designs_or_analyses": "Yes, the experiments validate the main claims in this paper, including the superior performance of MaxEntDP, the variance reduction property of QNE, and the necessity of entropy in policy evaluation. \n\nHowever, from current results it seems that the benefits of MaxEntDP is marginal, and in fact some baseline methods like SAC may achieve better performances if allowed to tune its hyper-parameters. I suggest to include experiments on harder online RL benchmarks, such as DMControl, and keeping the hyper-parameter tuning effort the same across baseline algorithms, to further validate the benefits such as improved exploration.", "relation_to_broader_scientific_literature": "1. Using importance sampling to sample from the posterior distribution is similar, although not the same, to the technique used in iDEM [1].\n2. The analytical computation of the log-probability is inspired by ITDiffusion [2]. The authors changed the integration variable for numerical stability.\n\n[1] Akhound-Sadegh, Tara, et al. \"Iterated denoising energy matching for sampling from boltzmann densities.\" arXiv preprint arXiv:2402.06121 (2024).\n\n[2] Kong, Xianghao, Rob Brekelmans, and Greg Ver Steeg. \"Information-theoretic diffusion.\" arXiv preprint arXiv:2302.03792 (2023).", "essential_references_not_discussed": "The idea of sampling from an energy-based distribution, of which we only have the potential function but no samples, is also related to model-based diffusion [1]. \n\n[1] Pan, Chaoyi, et al. \"Model-based diffusion for trajectory optimization.\" Advances in Neural Information Processing Systems 37 (2024): 57914-57943.", "other_strengths_and_weaknesses": "The development of the problem, the motivation, and the method are clear and easy to follow.", "other_comments_or_suggestions": "My main concerns are discussed in \"Theoretical Claims\" and \"Experimental Designs Or Analyses\", please refer to these sections.", "questions_for_authors": "N/A", "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: **Summary:**  
This paper proposes MaxEntDP, a diffusion-based reinforcement learning framework that integrates the maximum entropy principle into the training of expressive diffusion policies. The method addresses key challenges in policy evaluation and improvement via Q-weighted noise estimation and numerical approximation of marginal likelihoods, showing promising empirical results and solid theoretical foundations, although some concerns about scalability, practical assumptions, and comparisons remain.

**Pros:**  
1. Presents a novel integration of diffusion models into the MaxEnt RL framework, enabling more expressive and exploratory policy learning.  
2. Provides sound theoretical analysis and well-motivated algorithmic components, supported by ablations and baseline comparisons.  
3. Demonstrates performance gains on standard MuJoCo benchmarks, with indications of promise in more complex settings.

**Cons:**  
1. Key assumptions (e.g., well-trained noise networks and full support approximations) may not fully hold during interleaved training phases.  
2. Experimental evaluations are limited to relatively simple environments, with marginal gains over existing methods and missing recent SOTA baselines.  
3. Some theoretical derivations and implementation choices (e.g., integration bounds, likelihood estimation) require further justification and clarity.
Decision: Accept (poster)