Efficient Online Reinforcement Learning for Diffusion Policy
Authors: Haitong Ma
Tianyi Chen
Kai Wang
Na Li
Bo Dai
Venue: ICML 2025 poster (2025)
Keywords: reinforcement learning
diffusion models
diffusion policy
Primary area: reinforcement_learning->deep_rl
TL;DR: We generalize denoising score matcing by reweighting the loss function, enabling efficient online RL with diffusion policies.

Abstract: Diffusion policies have achieved superior performance in imitation learning and offline reinforcement learning (RL) due to their rich expressiveness. However, the conventional diffusion training procedure requires samples from target distribution, which is impossible in online RL since we cannot sample from the optimal policy. Backpropagating policy gradient through the diffusion process incurs huge computational costs and instability, thus being expensive and not scalable. To enable efficient training of diffusion policies in online RL, we generalize the conventional denoising score matching by reweighting the loss function. The resulting Reweighted Score Matching (RSM) preserves the optimal solution and low computational cost of denoising score matching, while eliminating the need to sample from the target distribution and allowing learning to optimize value functions. We introduce two tractable reweighted loss functions to solve two commonly used policy optimization problems, policy mirror descent and max-entropy policy, resulting in two practical algorithms named Diffusion Policy Mirror Descent (DPMD) and Soft Diffusion Actor-Critic (SDAC). We conducted comprehensive comparisons on MuJoCo benchmarks. The empirical results show that the proposed algorithms outperform recent diffusion-policy online RLs on most tasks, and the DPMD improves more than 120% over soft actor-critic on Humanoid and Ant.

Efficient Online Reinforcement Learning for Diffusion Policy
Haitong Ma Tianyi Chen
Kai Wang? Bo Dai
Na Li
Abstract
Ermon, 2019; Ho al., 2020). With the iterative denoising design, diffusion models achieved superior expressiveness and multimodality in representing complex probability dis- tributions, demonstrating remarkable performance in image and video generation (Ramesh et al., 2021; Saharia et al., 2022). The superior expressiveness and multimodality nat- urally benefit the policies in sequential decision-makee problems. In fact, diffusion policy has been introduced in imitation learning and offline reinforcement learning (RL), where expert datasets are presented. Diffusion policies im- proved significantly over previous deterministic or unimodal policies on manipulation (Chi et al., 2023; Ke et al., 2024; Scheikl et al., 2024) and locomotion tasks (Huang et al., 2024).
Diffusion policies have achieved superior perfor- mance in imitation learning and offline reinforce- ment learning (RL) due to their rich expressive- ness. However, the conventional diffusion train- ing procedure requires samples from target dis- tribution, which is impossible in online RL since we cannot sample from the optimal policy. Back- propagating policy gradient through the diffusion process incurs huge computational costs and in- stability, thus being expensive and not scalable. To enable efficient training of diffusion policies in online RL., we generalize the conventional de- noising score matching by reweighting the loss function. The resulting Reweighted Score Match- ing (RSM) preserves the optimal solution and low computational cost of denoising score matching, while eliminating the need to sample from the target distribution and allowing learning to opti- mize value functions. We introduce two tractable reweighted loss functions to solve two commonly used policy optimization problems, policy mir- ror descent and max-entropy policy, resulting in two practical algorithms named Diffusion Pol- icy Mirror Descent (DPMD) and Soft Diffusion Actor-Critic (SDAC). We conducted comprehen- sive comparisons on MuJoCo benchmarks. The empirical results show that the proposed algo- rithms outperform recent diffusion-policy online RLs on most tasks, and the DPMD improves more than 120% over soft actor-critic on Humanoid and Ant.
Meanwhile, online RL has been long seeking expressive policy families. One promising direction lies in energy- based models (EBMs)-a class of probabilistic models that represent distributions via unnormalized densities. When applied to RL, energy-based policies ppolicies modeled as EBMs have been shown to arise as optimal solutions in proximity-based policy optimization (Nachum et al., 2017: Mei al., 2020) and max-entropy RL (Neu et al., 2017; Haarnoja al., 2017). Despite their theoretical appeal, training and sampling from such unnormalized models in continuous action spaces are notoriously difficult due to their intractable likelihood (Song & Kingma, 2021). To mitigate this, a variety of probabilistic models have been introduced for efficient sampling and learning, but with the cost of approximation error. In practice, many algo- rithms (Schulman, 2015; Schulman al., 2017; Haarnoja et al., 2018; Mansen-Estruce al., 2023) project the energy- based policies onto the Gaussian policies. However, this projection severely limits the expressiveness of the origi- nal energy-based formulations, often resulting in degraded performance.
1. Introduction
Many successes of diffusion-basse generative models have been witnessed recently (Sohl-Dickstein et al., 2015; Song &
Diffusion models are closely related to EBMs. as they can be regarded as EBMs perturbed by series of noise (Song & Ermon, 2019; Shribak et al., 2024), thus being the perfect candidate to represent energy-based policies in RL. Unfor- tunately, it is highly non-trivial to train diffusion policies in online RL. The commonly used diffusion model train- ing procedure, denoising score matching (Ho et al., 2020), requires data samples from the target data distribution (usu- ally large image dataset in image generation). However,
'Equal supervision 'Harvard University. Georgia Institute of Technology. <haitongma www.haaaarddddee. edu>, Na Li <nali @ sessecsshhe het.aeaaaaaaaa...dd...... had Bo Dai <bodai@ cc gatech.edu>>.
Emails: Haitong
Ma
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Efficient Online Reinforcement Learning for Diffusion Policy
cannot sample from the optimal policy in online RL, where the policy is learned by optimizing the returns value functions. There exist several preliminary studies try- ing to bypass the sampling issue (Psenka et al., 2023; Jain et al., 2024; Yang et al., 2023; Wang et al., 2024; Ding et al., 2024a; Ren et al., 2024), but all these methods suffer from biased estimations and/or huge memory and computation costs, resulting in suboptimal policies and limiting the true potential of diffusion policies in online RL..
A is the action space, r: S x A R is reward function, P : 8 x A - A(S) is the transition operator with (SS) as the family of distributions over S, pa E A(S) is the ini- tial distribution and E (0, is the discount factor. We consider two types of commonly seen policy optimization problems in RL, (a) Policy mirror descent and (b) Max- entropy policy.
Policy Mirror Descent closely related to practical proximity- algorithms such as TRPO (Schulman, 2015) and PPO (Schulman et al., 2017), but with dif- ferent approach to enforce the proximity constraints. We consider policy mirror descent with Kullback-Leible (KI.) divergence proximal term (Tomar et al., 2021; Lan, 2023; Peters et al., 2010) updates the policy with TMD (a|s) (a]s)==/)) = argmax Eant IQ Hold (s, a XD (Tol s) +(A))((A)
To handle these challenges, we propose to generalize diffu- sion model training by reweighting the conventional denois- ing score matching loss, resulting in two efficient algorithms to train diffusion policies in online RL without sampling from optimal policies. Specifically,
Building upon the viewpoint of difffusion models noise- perturbed EBMs, we propose Reweighted Score Match- ing (RSM), a family of loss functions to train diffusion models, which generalizes the denoising score matching by reweighting the loss function while preserving the opti- mal solution as noise-perturber EBMs.
(1)
where Q (s, a)
Exgia 2 az)lso
s, au = aj is the state-action value function and Told is the current policy. The additional KI. divergence objective constrains the updated policy to be approximately within the trust region. The closed-form solution of policy mirror descent (1) satisfies
RSM leads to computationally tractable and efficient algo- rithms to train diffusion policies in online RL. We show that, by choosing different reweighting functions, we can train diffusion policies to solve two policy optimization problems, policy mirror descent and max-entropy pol- icy, resulting in two practical algorithms named Diffu- sion Policy Mirror Descent (DPMD) and Soft Diffusion Actor-Critic (SDAC). Both problems commonly seen in theoretical studies but empirically challenging in the con- tinuous action space, and the proposed algorithms bridge this gap between the theory and practice of online RI..
TMD (a) = Told (als) exp (QTaM (8, a) /A)
(2)
ZMD
and ZMD(s) = J Told (als) exp (Q (8, a) (A) da is the par- tition function.
Max-entropy RL. Maximum entropy RI. considers the entropy-regularized expected return as the policy learning objective to justify the optimal stochastic policy
arg max J(x) == E,
Σ r (r(ST, ar) + (+))))1)))))
We conduct extensive empirical evaluation on MuJoCo, showing that the proposed algorithms outperform recent diffusion-basee based online RL baselines in most tasks. More- over, both algorithms improve more than 100% over SAC on Humanoid and DPMD improves more than 100% over SAC on Ant, demonstrating the potential of diffusion pol- icy in online RL.
where H (x(-ss)) = Iall log (ra is the entropy, A is a regularization coefficient for the entropy. The soft pol- icy iteration algorithm (Haarnoja et al., 2017) is proposed to solve the optimal max-entropy policy. Soft policy iteration algorithm iteratively conducts soft policy evaluation and soft policy improvement, where soft policy evaluation updates the soft Q-function by repeatedly applying soft Bellman up- date operator 7 to current value function Q i.e.
2. Preliminaries
S A + R.
We introduce the necessary preliminaries in this section. First, we introduce reinforcement learning and two com- monly seen policy optimization problems in online RL. Then we briefly recap the diffusion models and energy- based models.
TEQ(s,, = r(s,, ax) + + (V(++)) (4) where ((s)) Ea.~s [Q(sr, a,) - X Alog (a, st)]. Then in the soft policy improvement stage, the policy is updated to fit the target max-entropy policy
cxp Rold (s, a) /A)
2.1. Reinforcement Learning
Max (als (als)
(5)
7(x)
Markoy Decision Processes (MDPs) We consider Markov decision process (Puterman, 2014) specified by a tuple M = Ar,, P. H4: 7), where is the state space,
where Quilid (s, a) is the converged result of (4) with Thaid Z(s) = f exp(( Told (s, a)/A)da. Max-entropy RL shows foundational concept in the
trade-
2

---

Efficient Online Reinforcement Learning for Diffusion Policy
off with stochastic policies, leading to practical algorithms with strong performance even with the restrictive Gaussian policies such as soft actor-critic (Haarnoja et al., 2018).
where 1xje-1 == v Bix,-1, 3,D), x is the perturbed data at step, and p, are probability distri- butions'. As the perturbations at every step are independent and additive Gaussian, we can directly sample the t t-step perturbed data x, by
2.2. Energy-Based Models
X ~ pc - / 10000//xxo/0//0000 po((x)))10(x//x0))
The closed-form solutions of both policy mirror descent (2) and max-entropy policy (5) have unknown normaliza- tion constants. Such probabilistic models with unknown normalization constants are known as energy-based mod- els (EBMs), whose density functions can be abstracted as
where frjo(xx [x)] == N(xx; Vorxo (1 - (1) T))
(1)
1=1
(6)
Po(x) = exp (x)())
The backward process recovers the data distribution from a noise distribution pr with a series of reverse kernels Pc-1)(((-1(xx)) The kernel of reverse pro- cess is also Gaussian and can be parametrized as NV (1
Z
where 2 = f exp (- Ex dx is the unknown normaliza- tion constant or partition function. We only know the energy functions Exx) = log po (x), i.e., the negative log density. The gradient of log density Vx log pox) E()) is called the score functions.
Vas (xt + Prso(x;;t), ofi) with score networks
se (xe:() and fixed covariance o, 1- network sy tx); is trained by the denoising score match- ing to match the forward and reverse processes (Ho et al., 2020), whose loss function is
3, The score
I-a.
The unknown normalization constants Z raise difficulties in training and sampling of EBMs (Song & Kingma, 2021). One of the commonly used approaches is the score-based methods, which first learns the score function via score matching (Hyvärinen & Dayan, 2005; Song et al., 2020) and then draws samples via Markov chain Monte Carlo (MCMC) such as Langevin dynamics with the learned score functions (Neal et al., 2011). 1). However, the MCMC sam- pling is inefficient due to the lack of finite-time guarantees, preventing based EBMs from being widely used in practice.
T
M 1=
E
[1|50 (xx:t) Vx, log 2t (xe xx(xx))
[|]]
Norupo
(7)
The score function Xx, log 9:j0 (x0) can be computed from the sampled Gaussian noise perturbing ag to be a., thus the loss in (7) is tractable and easy to implement. After learning the Se via (7), we can draw samples via the reverse diffusion process by the iterative formulation
(xt + Bis0 (xe, f)) + 1 ---1 3,2t
In the practice of online RL, projection onto Gaussian policies is commonly used in policy optimization with EBMs. For example, the well-known soft actor-critic (SAC, Haarnoja et al., 2018) parameterize the policy as Guassian To (als) = V (Jo, o8 (s)) and updates the parameters 0= 10,, 0g] by optimizing the K L-iivergeecee to the target max-entropy policy ming DKL. [T Mox However, the projection loses expressiveness, and the resulting policies might be sub-optimal, leaving huge gap between theory and practice of energy-based policies.
X = Vat
(8)
1- ar
for t = T,T ,,..., and 2t N (0,I).
3. Reweighted Score Matching: A General Loss Family for Diffusion Models
In this section, we first present the connection between energy-based models and diffusion models, justifying the expressiveness of diffusion policy to represent energy-based policies such as (5). Then we identify the difficulties in training of diffusion policy in the context of online RL, where the conventional denoising score matching is in- tractable. To mitigate this, we propose our core contribution, Reweighted Score Matching (RSM), by reweighting the de- noising score matching loss.
2.3. Denoising Diffusion Probabilistic Models
Denoising diffusion probabilistic models (DDPMs, Sohl- Dickstein et al., 2015; Ho et al., 2020) are composed of forward diffusion process that gradually perturbs the data distribution Xg ~ Po to a noise distribution XT ~> Pr, and a reverse diffusion process that reconstructs the data dis- tribution Po from the noise distribution pr. The forward corruption kernels are Gaussian with a variance schedule $1,..., Br, resulting in the forward trajectories with joint distributions
Notation. To fit the diffusion policy context, we consider the diffusion policy notations, where the diffusion is on actions a conditioned on state 8. The score function has the additional input of states so(a;; 8, t). The data distribu-
We use and q interchangeably as density function in this paper. Generally, represents intractable distributions (like the t-step marginal pt (xx)). and q represents tractable distributions such as the Gaussian corruption qe 1 (x-1).
(x:T(Xo:T) = (xx) II
g 1 (xx xx-1)
3

---

Efficient Online Reinforcement Learning for Diffusion Policy
Noise perturbation
Score matching
Reweighting
with 80 X 1)
Baw Tusbin Pan
Sample (DDPM)
Sample Uniform
Sample Gaussian
I Bx rpp ||
V log p(|
9r (x, || V log prll" (Reweighted, ours)
Data distribution Po
Perturbed data Pt.
(DDPM)
Figure 1. Diffuusin model aims match score network (xx, f) with noise-perturbed score function Vx, log (x by minimizing the expectation of error 1.2-norm [xs,(x.: f) - Xx, lx,))(x,)] over distribution ps. RSM generalize to other weight function ge to enable diffusion policy training in online RL.
tion (-1s)) refers the policy in (2) (5), and we will explicitly mention which one refer as po if needed.
where refer vanilla score matching (VSM) loss. It is obvious that the minimizer of LVSM is sæ- (a:; 8, Va, log Pa (r (s) for any az Then we show the following loss equivalence,
3.1. Diffusion Models Noise-Perturbed Energy- -ased Models
LDSM(0; s,t) LVsM(0; s,t) constant
(11)
We first revisit the energy-based view of diffusion models, i.e., diffusion models are noise-perturbed EBMs (Song& & Er- mon, 2019; Shribak et al., 2024), to justify that the diffusion policy can efficiently represent the energy-based policies.
where constant is a constant irrelevant with 0. The de- tailed derivations are deferred to Appendix B.I.
Therefore, minimizing LDSM in DDPMs is equiva- lent to minimizing LVSM, whose optimal solution is sa- s,t) Va, log anywhere on the a space.
Proposition 3.1 (Diffusion models as noise-perturbed EBMs). Consider a single term in loss function of DDPM (7) at given state s and time t,
CDSM(0; s,t) ==
Proposition 3.1 indicates that the underlying learning tar- get of the score function so (ag; s, t) is the noise perturperturbed score functions Va, lo (s). At the sampling stage, as the noise gets close to zero when t goes from T to the reverse process (8), the noise-perturbed EBMs gradually resemble the original noiseless target data distribution Po- Therefore, diffusion models can be regarded as series of noise-perturbed EBMs and can use the diffusion policy to express the energy energy-based policies such as (2) and (5).
E
||ss0 (ag 8, Xx, log (ogggpaaaaa)) (9)
(-1))
We refer to it as the denoising score matching (DSM) loss. The optimal solution 6* achieved when the following holds for all a
Sg- (a;; = Va. log (a.ls)
where Pr(a:|s) = q: oaa ao (ooo)o(ac (aols)) is the noise- perturbed policy with perturbation kernel q (n)) = N(a; Vājag, (1 11--x) āt) I), for noise schedule index 1.2.....
Revisiting the challenges in online RL. Proposition 3.1 shows that the DSM loss (9) is tractable and efficient way to train the score network to match Va. log Pt when we have access samples from po. However, training diffusion policy is highly non-trivial in online RL because of two major challenges:
The connection first revealed in Vincent (2011), and we revisit it in this paper.
Sampling challenge: In online RL, we do not have data samples from the policies such as (2) or (5). the DSM
Proof. Consider the following loss function minimizing the squared error between so (a;; s, f) and Va, log (rrasss).
"The reason to add noise perturbations in score functions is to encourage exploration on the energy landscape, which significantly improves the sampling quality and makes diffusion-like models the key breakthrough EBMs (Song & Ermon, 2019).
LVSM(( 8, ))=
E a-)p)(-1e)
[Isa(ar: 8, Va, Pogpr(a {]]
(10)

---

Efficient Online Reinforcement Learning for Diffusion Policy
loss (9) is longer tractable.
4.1. Tractable Reweighted Loss Functions
Another possible solution is to treat the reverse process as policy parameterizations and backpropagate policy gradient through the whole reverse diffusion process (8) like Wang et al. (2024); Celik et al. (2025). However, this recursive gradient propagation not only incurs huge computational and memory costs, making diffusion policy learning expensive and unstable. Moreover, this policy parametrization viewpoint is limited to the max- entropy policy formulation (5).
Computational challenge:
4.1.1. DIFFUSION POLICY MIRROR DESCENT
Consider the mirror descent policy TMD(-Is) in (2) and set po(-Is) (MD((1s)) we define the reweighting function as
JMD = ZMD (()))))(a))))))
where (M) J
Tota (a ()) exp (Q (s, a) /A) da. Then show the reweighted loss LOMD (0: s, t) is tractable via the following derivation, [ (8: s,
f gMD (1; (am 8) [|50(a:;) 8,1) - Va, logpr(a:1s)]] da, =
These challenges hinder the feasibility and performance of diffusion-baan bassdd policies online RL. We need a principal way to train diffusion policies when we have the energy function as partial or full knowledge of the data distribution, which we reveal in the following.
Q (s. ag) exp
|1:0 (as; 8, V ao loganj(ar(aa))
A
Corum(0....)
+ constant
3.2. Reweighted Score Matching
(13)
We develop our core contribution, Reweighted Score Match- ing (RSM), general loss family leading to efficient diffusion policy learning algorithms that eliminate the aforementioned difficulties.
where LDPMD s,t) is tractable through unbiased sampling- based approximation. The derivation is similar to Proposition 3.1, and we defer it to Appendix B.2.
1.2. SOFT DIFFUSION
A key observation from VSM loss (10) is that, integrat- ing the square error (as; (e; s, t) - Va. log logp:(a;) (r(a;|s)] over distribution pe (-(ss) is not the only option to perform score matching that matches sa(a, 8, f) with Vo, log (r(a:ss) Our core idea is to generalize this by reweighting the VSM loss (10), allowing integrations with respect to any strictly positive function g(ar; s) AA x S + (0, oo) as long as the following loss function is well-defined,
The max-entropy policy TaaxEnt in (5) is more challenging as we only know the energy function. is also closely re- lated to the Boltzmann sampling problem AkhounddSaddee et al., 2024; Midgley et al., 2022). We need special sam- pling protocol to handle it. First, we define the reweighting function as
IMaxEnt ha(a (s) (s)))
(² t) g( ax; (a)ss s, V aa, (c(s)]]
=
da;
where he (a(a)(s) is sampling distribution we choose. We require he (arls) to have full support on a, space. Then we can show the following equivalence,
(12)
where the superscript, y indicates the reweighting function. The optimal solution of minimizing (0) remains the same as matching Va, log Pe (arls) everywhere on at space in Proposition 3.1, 8o (a; 8, t). (12) indicates a general loss family, where the VSM loss (10) lies in it as [P (0; 8, l).
La (0; 8,()
f MMaxEnt (ag s) [sso(a;;st) Va. logp((alss)) da,
=constantx
( ao)
The reweighting technique gives more flexibility in loss function design. We will show tractable equivalent formula- tions in the next section.
E a r 1o
exp
|| (ar;st)) V log
(aDAc((.aaa)
+ constant
4. Diffusion Policy Optimization using Reweighted Score Matching
(14)
where doje is a conditional Gaussian distribution defined as
1
1-at
1-
In this section, we show two different reweighting functions, with which the loss _ can be converted to tractable loss functions, to train diffusion policies to represent both the mirror descent policy (2) and softmax policy (5). We also discuss practical issues, such as the exploration-exploitatio tradeoff, sampling distributions.
éoje la:) := N āo: Var
aj
(15)
ar
The reason we introduce doje @ reverse sampling trick.
(a:) to use the following
Remark 4.1. (Reverse sampling trick.) The density func-
5

---

Efficient Online Reinforcement Learning for Diffusion Policy
tions of qrjo and Poje are
to balance exploration and exploitation, similar with Wang et al. (2024).
(oje(aolat) 2T 1 -
ai
4/2
ag -0 )
Ot
exp
-2 (1 (tt)
Log likelihood computation. The soft policy evaluation step in SDAC requires explicit log- ligeliielhhod that is non- trivial for diffusion policies. However, we observe that the action after batch action sampling is of low stochasticity, thus, we can use the log probability of the additive Gaussian to approximate the log probability of the policy.
(a: o2
qrjo(ar an) = (2 (1 ae)) 6/2
exp
-2 (1 - an)
where these two density functions only differ by a constant. We show an abstract example of the reverse sampling trick here. Consider the following integral that is well-defined
Numerical stability. In practice, the exponential of large Q functions in (13) and (14) might cause the loss to explode. We handle the numerical stabilities by (a) Normalization. In DPMD, we normalize the Q(s, aa) with the exponen- tial moving average (EMA) of mean and standard devia- tion over the sampled minibatch. (b) The logsumexp trick. In SDAC, we sample multiple a( ~ doje, i E (1.2..... K} for every 8, a and use the Logsumes xp trick to avoid explosion of the weights, which means replacing the exp(Q(s, a2) in (14) with
(s) := h4(arls) ()po s ge (a (a)(a)))(( (a)))a)) aoi s s)daada;
E Eagr-po qojo [a [s) (a) au; s)]
where 1 A x A x S + R is an integrable function. No- tice that we can equivalently compute the integral J(s) by another expectation,
(s)
x 1i (1r(ar)))pe 8) (oo(s) Ooje (aaolaa))) Maar On; 8) s)danda;
= -ar-hharrdoa
(olaas))((a; ao: 8)]
This trick helps bypass sampling from po and sample from the distribution he instead. The detailed derivations are in Appendix B.2.
exp [(s,, (((5,a0)) )/> - log Σ exp (Q(s, (((s.ab)/A))
(4)
The trick does not conflict with our theoretical derivation, which is another reweighting on the space.
Summary. We can see both loss functions (13) and (14) handles the aforementioned sampling and computational challenges. First, we avoid sampling from the target policy TMD or Tsoftmax and sampling from either the current pol- icy Told a distribution hú we can choose. Second, we have similar computation with denoising score matching (7), avoiding extra computational cost induced by diffusion pol- icy learning. These benefits perfectly echo the difficulties of sampling and computations in applying vanilla diffusion model training to online RL, enabling efficient diffusion policy learning.
Reverse sampling distribution selection. In SDAC, we can choose the sampling distribution hy. Empirically we tried uniform distribution, last policy Told (-1s), and the perturhed data distributions f oota (aols)gej((- ) All these dis- tributions show similar performance.
Combining all the discussions above, we present the prac- tical algorithm of DPMD in Algorithm 1, while SDAC are detailed in Appendix C.2.
Remark 4.2 (Broader applications). We emphasize that al- though we develop RSM with the reweighting techniques for online RL problems, the RSM has its own merit and can
4.3. Comparison with Recent Diffusion-based Online RLs
be applied to enable diffusion models on any probabilistic modeling problem with known energy functions, such as Saahound-Sadege et al., 2024; Midg- ley et al., 2022). We also show a toy example of Boltzmann sampling in Section 5.1 where use RSM to train toy dif- fusion model to generate samples from a Gaussian mixture distribution with only access to the energy functions.
We say both proposed algorithms are efficient because of similar computation and memory cost with denoising score matching (7) while bypassing the sampling issues, while recent diffusion- -bsed online RL either incur huge com- putational or memory cost or induce approximation errors. Recent works on diffusion policy online RLs can be cat- egorized into these families: i) Score-based Bolzmann sampling. With the known energy functions in (5), Psenka et al. (2023); Jain al. (2024) differentiated it to get the non-noisy score function and use Langevin dynamics or diffusion to sample from (5). The empirical performance is not good due to the inaccurate score function obtained by differentiating a learned energy function. ii) Reverse diffu- sion as policy parametrizations. The reverse process (8) can also be directly regarded as a complex parametrization of A. Wang et al. (2024) backpropagate policy gradients through the reverse diffusion process, resulting in huge com-
Boltzmann samplers
4.2. Practical Issues of Diffusion Policy Training
Batch action sampling. Ding et al. (2024a) revealed that diffusion models are too random for efficient exploitation, and proposed to sample a batch of cations and choose the one with the highest Q-value as the behavior policy,
(16) We leverage this trick in both of our algorithms and add a Gaussian noise whose noise level is automatically tuned
arg max Q(s, a ()).
6

---

Efficient Online Reinforcement Learning for Diffusion Policy
Algorithm 1 Diffusion Policy Mirror Descent (DPMD)
known as Boltzmann sampling) to verify the effectiveness of the proposed reweighted score matching. The Gaussian mixture model is composed of two modes whose mean values are [3, 3] and [-3, -3] and mixing coefficients are 0.8 and 0.2 shown in Figure 2(a). The detailed training setup can be found in Appendix C.3, while another Boltzmann sampling task named Two Moon is shown in Appendix C.4.
Require: Diffusion noise schedule Gr Or for E ..)....T. MDP M, initial policy parameters 8g, ini- tial Q-function parameters Co. replay buffer D 0, learning rate 3, KL-divergence coefficient Ao and target Atargot, HQQ (0) = 0.0, (0) := 1.0, EMA parameter &
1: for epoch e = 1, ..... do
2: Sample M actions with policy - and choose one
As we only know the energy function, we select the SDAC- like loss function in (14) as our training objective. We compare three diffusion models trained with two types of loss functions: a. proposed SDAC-like loss function in (14) with sampling distribution h being Gaussian and uniform distributions in Figure 2(b) and Figure 2(c), which have access to the true energy function but cannot sample directly from the Gaussian mixture. c. Denoising score matching loss (9) in Figure 2(d), which has access to sample from the Gaussian mixture. Empirical results showed that all diffu- sion models can approximately recover both the two modes and the mixing coefficients, which verifies the effectiveness of the proposed RSM approach to train diffusion models.
according to (16).
3: Interact with M, store the data in update replay buffer D.
4: Sample minibatch of (s, a.r, from D.
5: # Policy Pllicy evaluation
6: Sample a' via reverse diffusion process (8) with S0.-1'
7:
Update Qe by minimizing the Bellman residual (28).
8: # Di ffusion on Policy Mirror Descent.
9:
Sample uniformly from (1,2....,T} Sample do using 56.-, and a, (ojo--lao)
10: Compute Qc(s. and normalize Qe(s,a))
Qa(s.a)) -mole-1)
(2(a-1)
Moreover, we also show the naive Langevin dynamics (Neal et al., 201 I) samples access to the true score function without noise perturbations. It shows that even with the true score function, Langevin dynamics can not correctly recover the mixing coefficient in finite steps (20 steps in this case), demonstrating the neces- sity of diffusion models even with given energy functions.
update 0g with score matching
Ees [CDPMD s, t)] in (13) with Qe- Update KL-divergence coefficient de f- Ae-1 +
reference in Figure 2(e), which has
11:
8(Ae-1- - -ttrget).
12: Update EMA uo(e) = (1 - E)uelE - +)+ +
{ mmean(Qe). oo(e)= = (1 ()))((e11+((
13: end for
putation costs. Ding et al. (2024a) approximate the policy learning as maximum likelihood estimation for the reverse process, which incurs approximation errors and can not handle negative Q-values. iii) Others. Yang et al. (2023) maintained a separate diffusion buffer to approximate the policy distribution and fit it with the diffusion model. Ren et al. (2024) combined the reverse process MDP with MDP in RL and conducted policy optimizations. They all induce huge memory and computation costs, thus being impractical and unnecessary. More general related works can be found in Appendix A.
5.2. OpenAl Gym MuJoCo Tasks
5.2.1. EXPERIMENTAL SETUP
We implemented the proposed DPMD and SDAC algorithms with the JAX package" and evaluated the performance on 10 OpenAl Gym MuJoCo v4 tasks. All environments except Humanoid-ve v4 are trained over 200K iterations with total of I million environment interactions, while Humanoid-v4 has five times more.
Baselines. The baselines include two families of model- free RL algorihhms. The first family is diffusion policy RL, which includes a collection of recent diffusion-policy online RLs, including QSM (Psenka et al., 2023), QVPO (Ding et al., 2024a). DACER (Wang et al., 2024). DIPO (Yang et al., 2023) and DPPO (Ren et al., 2024). The second family is classic model- -ree online RL baselines including PPO (Schulman et al., 2017), TD3 (Fujimoto et al., 2018) and SAC (Haarnoja et al., 2018). A more detailed explana- tion to the baselines can be found in Appendix C.5.
5. Experimental Results
This section presents the experimental results. We first use a toy example, generating a 2D Gaussian mixture, to verify the effectiveness of the proposed reweighted score matching as diffusion model training. Then we show the empirical re- sults of the proposed DPMD and SDAC algorithms evaluated with OpenAl Gym MuJoCo tasks.
5.1. Toy Example
'The implementation can be found at http://jjil...ccmmmmmmmooogggg/t ott.a/oioaa.aliiiiiooooooooooaa. .onleee.
We first show a toy example of generating a 2D Gaussian mixture distribution from the known energy function (also
7

---

Efficient Online Reinforcement Learning for Diffusion Policy
True data distribution
REM wutth Ea Dausdnn sample
RSM with Lniiferns cample
Denesiing score mansting
Lar gevin Dyna
30%
92%
80%
20%
20%
20%G
90%
10%
20
0
6
-3 20% 85%
-3 9.20% 50%
-6
0 5.20% 90%
-6
-3
5.20% 50%
-3
0.0% 80%
(a)
(b)
(c)
(d)
(e)
Figure 2. The seatter plots of generating 2D Gaussian mixture, the histograms show the partition on each axis. Figure 2(a) shows the true data samples with mixing coefficients [0.8, 0.2]. Figures 2(b) to 2(d) show that the proposed reweighted score matching and denoising
score matching can approximately recover the true data distribution. Figure 2(e) shows the slow mixing of Langevin dynamics that the mixing coefficients can not be correctly recovered.
HalfCheetah-v4
Reacherv44
Humanold
Pusher-u4
1es InverteddPnndullum-u
1.00
2.0
-25
6000
-100
0.75
Return
Return
0.5
-200
0.50
-75
2000
0.25
0.0
0.00
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
0.00
0.25
0.50
0.75
1.00
0.0
0.5
3.0
3.5
2.0 Lus
0.0
0.5
1.0
1.5
2.0
Iterations
Los
Iter reoration
Itteations
145
iterations
Iteraiions
1w.5
Ant v4
1+3
Hopper v4
Swimmerrvv
1w]
WbikerZddv v4
InvertedDDoualePenouuu
150
1c3
Retum
Retur
300
Retur
50
1
2.0
0.0
0.5
1.5
0.0
0.5
2.0
0.0
0.5
Iterationns
Morations
Itorations
165
Iterations
Iterations
DPMD (ours)
SDAC (ours)
QSM
DIPO
DACER
QVPO
SAC
Figure 3. Average return over 20 evaluation episodes every 25k iterations (125k for Humanoid) during training. We seleet the top 5 baselines ranked by average performance over all tasks for clarity. The error bars are standard deviations over 5 random seeds.
5.2.2. EXPERIMENTAL RESULTS
mations since QSM matches the score function by differen- tiating the Q-functions. QVPO failed Reacher and Pusher since it cannot handle negative Q-functions. DACER failed InvertedPendulum despite its good performance in some complex tasks, probably due to the gradient instability when backpropagated recursively.
The performance and training curves are shown in Table and Figure 3, which shows that our proposed algorithm outperforms all the basclines in all OpenAI Gym MuJoCo environments. Especially, for those complex locomotion tasks including the Ha fCheethh, Walker2d, Ant, and Humanoida d, our top-performing algorithm variant obtained 36.0%, 41.7%, 127.3%, 143.5% performance improve- ment compared to SAC. Specifically, DPMD achieved at least 4,4 43.4%, %, 32.1 % 23.1 % performance improve- ment compared to other diffusion-policy online RL. base- lines (not the same for all environments), respectively, while SDAC shows comparable performance with DPMD except Ant. The empirical results demonstrate the superior and consistent performance of our proposed algorithm and the true potential of diffusion policies in online RL.
Computation and memory cost. We count the GPU mem- ory allocations and total computation time listed in Table 2. The computation is conducted on a desktop workstation with AMD Ryzen 9 7950X CPU, 96 GB memory, and NVIDIA RTX 4090 GPU. We achieve low memory consumption and faster computations compared to other diffusion--olliy baselines. Note that the QSM essentially does not involve the denoising process, thus has the lowest computation re- quirements. We can still achieve comparable computation time and memory cost with QSM, indicating the proposed RSM does not add much extra computational cost due to the diffusion policies. SDAC does not need to sample from the current policy to perform reweighted score matching, so it runs faster than DPMD.
Moreover, the performance of DPMD is very stable and con- sistently good for all the tasks. Other algorithms, including SDAC, performed badly on one or some tasks. For example, QSM failed the InvertedDoubieeendulum possibly because its true value function is known be highly non-smooth. The non smooth nature results in bad score function esti-
Sensitivity analysis. In Figure 4, we perform sensitivity analyses of different diffusion steps and diffusion noise
8

---

Efficient Online Reinforcement Learning for Diffusion Policy
Table 1. Performance on OpenAT Gym MuloCo environments. The numbers show the best mean returns and standard deviations over 200k iterations and 5 random seeds.
HALFCHEETAH
REACHER
HUMANOID
PUSHER
INVERTEDPENDULUM
PPO
4852 732
-8.69
11.50
952 + 259
-25.52 + 2.60
1000 +
Classic
TD3
8149 688
-3.10 0.07
5816 + 358
-25.07 = 1.01
1000 = 0
Model-Free RI
SAC
8981 370
-65.35 + 56.42
2858 2637
-31.22 + 0.26
1000 =
QSM
10740 + 444
-4.16
5652 + 435
80.78 + 2.20
1000 +
DIPO
9063
-3.29
4880 1072
32.89 0.34
1000 =
Diffusion Policy RL
DACER
11203 +: 246
0.07
2755 + 3599
30.82 0.13
801 - 446
QVPO
7321 1087
+30.59 16.57
121 + 75
129.00 0.96
1000 -
DPPO
1173 = 392
-6.22
484 + 64
89.31 + 17.32
1000 -
DPMD
11924 + 609
3.14
6959 + 460
30.43 + 0.37
1000 =
SDAC
12210 = 964
3.47
6437 177
32.53 + 5.27
1000 =
ANT
HOPPER
SWIMMER
WALKER2D
INVERTED2
PPO
3442 851
3227 164
84.5 12.4
4114 + 806
9358 +
Classic
TD3
3733 1336
1934 1079
71.9 15.3
2476 + 1357
9360 +
Model-Free RL
SAC
2500 767
3197 294
63.5 10.2
3233 + 871
1
QSM
938 + 164
2804 466
5..0
2523 + 872
2186 + 234
DIPO
965 +
1191 770
4.9
1961 + 1509
9352 +
Diffusion Policy RL
DACER
4301 524
3212 + 86
103.0
3194 + 1822
6289 + 3977
QVPO
718 + 336
2873
5..4
2337 1215
7603 + 3910
DPPO
60 + 15
2175 + 556
106.1
1130 + 686
9346 +
DPMD
5683 138
3275 + 55
4365 = 266
9360 0
SDAC
1391 202
2955 370
119.1 + 41.9
3995 + 498
9360 +
6. Conclusion
Table 2. GPU memory allocation and total compute time of 200K iterations and million environment interactions. *QSM did not learn diffusion policies essentially thus the computation is lightweight.
In this paper, proposed Reweighted Score Match- ing (RSM), an efficient diffusion policy training algorithm tailored for online RL. Regarding diffusion models as noise- perturbed EBMs, we develop the reweighted score matching to train diffusion models with access only to the energy func- tions and bypass sampling from the data distribution. In this way, we can train diffusion policy with only access the Q-function the energy functions in online RL. Empirical results have shown superior performance compared to SAC and other recent diffusion policy online RLs. Possible fu- ture directions include improving the stability of diffusion policies and efficient exploration scheme design.
Algorithm
GPU Memory (MB)
Training time (min)
QSM
997
14.23
QVPO
5219
30,90
DACER
1371
27.61
DIPO
5096
19.31
DPPO
1557
95.16
DPMD
21.10
SDAC
1113
16.10
schedules on the DPMD variant. Results show that 10 and 20 diffusion steps obtain comparable results, both outperform- ing the 30-step setting. The linear and cosine noise sched- ules perform similarly, and both outperform the variance- preserving schedule. Therefore, we choose 20 steps and cosine schedules for all tasks. The results also show the robustness to the diffusion process hyperparameters.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Sanaltiiit Plasion stace
Sendtbilty nobr schecute
Acknowledgment
This paper is supported by NSF AI institute: NSF ECCS-2401390, NSF ECCS-2401391, ONR N000142512173, NSF ASCENT 2328241, NSF IIS- 2403240, Schmidt Sciences A12050 Fellowship.
2112085,
1.5
Figure Sensitivity analysis on diffusion steps and diffusion noise schedule on Ant-v4.
9

---

Efficient Online Reinforcement Learning for Diffusion Policy
References
Ho, J., Jain, A., and Abbeel, P. Denoising Diffusion Prob- abilistic Models. In Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851. Curran Associates, Inc., 2020.
Akhound-Sadegh. T., Rector-Brooks, J., Bose, A. J., Mittal, S., Lemos, P. Liu, C.-H., Sendera, M., Ravanbakhsh, S., Gidel, G., Bengio. Y., et al. Iterated denoising energy matching for sampling from boltzmann densities. arXiv preprint 1rXiv:2402.061211 2024.
Huang, X., Chi, Y., Wang, R., Li, Z., Peng. X. B., Shao, S., Nikolic, B., and Sreenath, K. Diffuseloco: Real-time legged locomotion control with diffusion from offline datasets. arXiv preprint arXiv:2404. 19264, 2024.
Celik, Oi, Z.,, Blessing, D., Li, G., Palanicek, D., Peters, J., Chalvatzaki, G., and Neumann, G. Dime: Diffusion- based maximum entropy reinforcement learning. arXiv preprint ar (rXi:25502002366. 2025.
Hyvärinen, A. and Dayan, P. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Chen. H., Lu, C., Ying, C., Su, H., and Zhu, J. Offline rein- forcement learning via high-fidelitty generative behavior modeling. arXiv preprint arXiv:2209. 14548, 2022.
Jain, V., Akhound-Sadegh, T., and Ravanbakhsh, S. Sam- pling from energy-based policies using diffusion. arXiv preprint urXiv:2410.01312. 2024.
Chi, C., Xu, Z., Feng, S., Cousineau, E., Du, Y., Burchfiel, B., Tedrake, R., and Song. S. Diffusion policy: Visuomo- tor policy learning via action diffusioon. The International Journal of Robotics Research, pp. 02783649241223366. 2023.
Janner, M., Du, Y., Tenenbaum, J. B., and Levine, S. Plan- ning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09900 2022.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the Design Space of Diffusion Diffusion-BBsed Generative Models, October 2022.
Ding, S., Hu, K., Zhang, Z. Ren, K., Zhang. W., Yu, J., Wang, J., and Shi, Y. Diffusion-baie based reinforcement learning via g-weighted variational policy optimization. arXiv preprint arXiv:2405. 16173, 2024a.
Ke, T.-W., Gkanatsios, N., and Fragkiadaki, K. 3d diffuser actor: Policy diffusion with 3d scene representations. arXiv preprint arXiv:2402.10885 2024.
Ding. Z., Zhang. A., Tian, Y., and Zheng, Q. Diffusion world model: Future modeling beyond step-by-step roll- out for offline reinforcement learning. arXiv preprint orX5v:2402.03370, 2024b.
Lan, G. Policy mirror descent for reinforcement learning: Linear convergence, new sampling complexity, and gen- eralized problem classes. Mathematical programming, 198(1):1059-111) 1106, 2023.
Du, Y., Yang, S., Dai, B., Dai, H., Nachum, O., Tenenbaum, J., Schuurmans, D., and Abbeel, P. Learning universal policies via text-guided video generation. Advances in Neural Information Processing Systems, 36, 2024.
Mei, J., Xiao, C., Szepesvari, C., and Schuurmans, D. On the global convergence rates of softmax policy gradient methods. In International conference on machine learn- ing. pp. 6820-6829. PMLR, 2020.
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.
Midgley, L. I., Stimper, V., Simm, G. N., Schölkopf, B., and Hernández- Lobato, J. M. Flow annealed importance sampling bootstrap. arXiv preprint arXiv:2208.01 2022.
Haarnoja, T.. Tang, H., Abbeel, P., and Levine, S. Rein- forcement Learning with Deep Energy-Band -ased Policies, July 2017.
Nachum, O., Norouzi, M., Xu, K., and Schuurmans, D. Bridging the gap between value and policy based rein- forcement learning. Advances in neural information pro- cessing systems, 30, 2017.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforce- ment learning with stochastic actor. In International conference machine learning, pp. 1861-1870. PMLR, 2018.
Neal, R. M. et al. Meme using hamiltonian dynamics. Hand- book of markov chain monte carlo, 2(11):2, 2011.
Hansen-Estruck P., Kostrikov, I., Janner, M., Kuba, J. G., and Levine, S. Idql: Implicit q-learning as an actor- critic method with diffusion policies. arXiv preprint arXiv:2304. 10573, 2023.
Neu, G., Jonsson, A., and Gómez, A unified view of regraae-rg markov decision processes. arXiv preprint arXiv: / 1705.0779, 2017.
10

---

Efficient Online Reinforcement Learning for Diffusion Policy
Peters, J., Mulling, K., and Altun, Y. Relative entropy policy search. In Proceedings the AAAI Conference Artificial Intelligence, volume 24, pp. 1607-1612, 2010.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Psenka, M., Escontrela, A., Abbeel, P., and Ma, Y. Learn- ing diffusion model policy from rewards via q-score matching, arXiv preprint arXiv:2312.11752. 2023.
Song, Y. and Kingma, D. How train your energy-based models. arXiv preprint arXiv:2 2021.
Song, Y., Garg, S., Shi, J., and Ermon, S. Sliced score matching: scalable approach to density and score es- timation. Uncertainty in Artificial Intelligence, pp.
Puterman, M. L. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
574-584. PMLR, 2020.
Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Rad- ford, A., Chen, M., and Sutskever, I. Zero-shot text-to- image generation. In International conference on ma- chine learning, pp. 8821-8831. Pmlr, 2021.
Song, Y. Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-Based Generative Model- ing through Stochastic Differential Equations, February 2021.
Ren, Z., Lidard, J., Ankile, L. L., Simeonov, A., Agrawal, P., Majumdar, A., Burchfiel, B., Dai, H., and Simehowitz, M. Diffusion policy policy optimization. arXiv preprinl arXiv:2409.00588, 2024.
Tomar, M.. Shani, L., Efroni, Y. and Ghavamzadeh, M. Mirror descent policy optimization, 2021. URL http::/aarrv.orr/aa///055....11
Vincent, P. A Connection Between Score Matching and Denoising Autoencoders. Neural Computation, 23(7): 1661-1674, July 2011. ISSN 0899-7667, 1530-888X. doi: 10.1162/NECO a 00142.
Rigter, M., Yamada, J., and Posner, I. World models via policy-guided trajectory diffusion. arXiv preprint arXiv:2312.08533, 2023.
Saharia, C., Chan, W., Saxena, 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. Ad- vances in neural information processing systems, 35: 36479-36494, 2022.
Wang. Y., Wang. L.., Jiang, Y., Zou, W., Liu, T., Song, X., Wang, W., Xiao, L., Wu, J., Duan, J., and Li, S. E. Diffusion Actor-Critin with Entropy Regulator, December 2024.
Wang. Z., Hunt, J. J., and Zhou, M. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.066
Scheikl, P. M., Schreiber, N., Haas, C., Freymuth, N., Neu- mann, G., Lioutikov, R., and Mathis-Ullrich, F. Move- ment primitive diffusion: Learning gentle robotic ma- nipulation of deformable objects. IEEE Robotics and Automation Letters, 2024.
2022.
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. 13122, 2023.
Schulman, J. Trust region policy optimization. arXiv preprint drX1v:1502.05477,2000 2015.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms.
arXiv preprint arXiv: (K17.0707.6644 2017.
Shribak, D., Gao, C.-X., Li, Y., Xiao, C., and Dai, B. Dif- fusion spectral representation for reinforcement learning. arXiv preprint arXiv:2406. 00..00666662/1 16121, 2024.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, Deep Unsupervised Learning using Nonequi- librium Thermodynamics. In Proceedings of the 32nd International Conference Machine Learning, pp. 2256- 2265. PMLR, June 2015.
Song. J., Meng, C., and Ermon, S. Denoising Diffusson Implicit Models, October 2022.

---

Efficient Online Reinforcement Learning for Diffusion Policy
A. Related works
Diffusion models for decision making. Due to their rich expressiveness in modeling complex and multimodal distributions, diffusion models have been leveraged to represent stochastic policies (Wang et al., 2022: Chen al., 2022; Gaansn--Estrcee al., 2023), plan trajectories (Janner et al., 2022; Chi et al., 2023; Du et al., 2024) and capture transition dynamics (Rigter et al., 2023; Ding et al., 2024b; Shribak al., 2024). Specifically, we focus on the diffusion policies. Diffusion policies have been primarily used on offline RL with expert datasets, where the denoising score matching (7) is still available and the learned Q-function only provides extra guidance such as regularization (Wang et al., 2022) or multiplication in the energy function. However, in online RL we do not have the dataset, thus denoising score matching is impossible.
Diffusion Models. Diffusion models have dual interpretation of EBMs and latent variable models. The latent variable interpretation is motivated by the solving reverse-time diffusion thermodynamics via multiple layers of decoder networks (Sohl-Dickstein et al., 2015). It was later refined by Ho al. (2020) via simplified training loss. The EBM interpretation aims to solve pitfalls in Langevin dynamics sampling by adding progressively decreasing noise (Song & Ermon, 2019). Then the two viewpoints are merged together with viewpoints from stochastic differential equations (Song et al., 2021), followed by numerous improvements on the training and sampling design (Song et al., 2022; Karras et al., 2022).
Noise-conditioned score networks. A equivalent approaches developed by Song & Ermon (2019) simultancously with diffusion models is to fit the score function of a series of noise-perturber data distribution A/ (x;;x T , 2 K}
with a noise schedule d1 > 02 > > σk. The resulting models, named the codde-condiiioeed score networks (NCSN) So (x:; oj), take the noise level into the inputs and are learned by denoising score matching (Vincent, 2011)
E
(xx.xxx((xxo.11
[Ilfa (x;; a;) - Vx. log (xx[
(17)
Then in the sampling stage, Song & Ermon (2019) uses the Langevin dynamics x,+1 = = V log q(x; xx) + /2nzz
to sample from energy function. Song& & Ermon (2019) additionally replace the original score function Vx, log q(x x) in the Langevin dynamics with the learned noisy score function fo( o.):
X;+1 f- x; + f σ.) + 22yz, 0,
(18)
named as annealed Langevin dynamics. The scheduled noise perturbation design significantly improved the image generation performance to match the state-of-the-an (SOTA) at that time (Song & Ermon, 2019), which further refined by DDPM.
We can that the annealed Langevin dynamics (18) resembles the DDPM sampling (8) with different scale factors, and the denoising score matching loss (9) is equivalent to (7). Therefore, DDPM can be interpreted as EBMs with multi-level noise perturbations. A more thorough discussion on their equivalency can also be found in (Ho et al., 2020; Song et al., 2021).
B. Derivations
B.1. Derivations of Proposition 3.1
We repeat Proposition 3.1 here,
Proposition B.1 (Diffusion models as noise-perturbed EBMs). Consider a single term in loss function of DDPM (7) at given state s and time
CDSM s, ==
E
|0 (ax; 8,t) Vx, log grjo (ar [0)
(9)
apr-po(-1s)
a.--ujol-lan
We refer to it as the denoising score matching (DSM) loss. The optimal solution H* is achieved when the following holds for all at
sg- (ar; (, Va, log Pogge (a.ls)
where parls) [ ())(ar) (o)po is the noise-perturbed policy with perturbation kernel quo (ao) N(a, Va,an, (1 ( I), for noise schedule index t 1.2....,
Proof. We first check the square error (1s((ar,,.)) - Va, logpr(an (a:ls)]l²
12

---

Efficient Online Reinforcement Learning for Diffusion Policy
[so(ar: ae. 8, - Va, log ())]s)])
s( Va / 9110(ala0)Pp(0)) s)da =1s00 8, 1)12
1
s 8.t)
so(al, 8,t), / Va, log (o constant
(19)
ps(ar(s) (s)
1
(as 8, 8,1)
(ejo(arlao)pol 8)da0
pt (()ls)))
(ppl(alaa)pp(anbl 8) (so(ar, s,t), aa log (a)a((aaaa))) das + constant
pr(a, 8)
(e [Is s, t)]l (so 8, t), Va, log (= das + constant
s:(a,
(eju(arlao)po (sols) (a 8, Va, log udo (ar dag + constant
(20)
where the constant denotes constants that are irrelevant with 8. (20) builds the foundation of transferring the intractable vanilla score matching loss (10) to tractable denoising score matching that can be easily computed by the sampled Gaussian noise.
Therefore, we can integrate both sides on pe (r(aass),
[|so(ar;stt) Va, log ( [o)o|² |1]] = Earrp: [[so(arsstt Va, logr((a))s))) + constant
E Go P0 a.
where the LHS is (O; 8,1) and RHS is LVsM(0;; 8, constant defined in (10). This concludes the proof of Proposition 3.1.
B.2. Derivations of Section 4.1
Section 4.1 shows that we can match the score network s, with noise-perturbed policy score function Va, log (e(a(()) without sampling from Po like denoising score matching (7).
B.2.1. DERIVATIONS DIFFUSION POLICY MIRROR DESCENT
First, restate the results. Consider the mirror descent policy (MD(+)s) in (2) and po(-ls) TMD(-1s), we define the reweighting function
JMD =
(8) (xD((s)p(au
where Z(s) = f Tota (als) exp (Q (s, a) (A) da. Then we can show the reweighted loss LoAD (0; t) is tractable the following derivation,
LAMD (f: - gMD (a;; 8) | [|s; 8, - Va. ()) da;
E
exp (Q (s, au) /A) so (ar;s,) 8, Va. log arj (a, [0) +constant
(13)
ao-d a.
[oemm(0.e.))
where (DPMD(O, s,t, is tractable through sampling based approximation.
13

---

Efficient Online Reinforcement Learning for Diffusion Policy
Proof. (13)
L (0; s,t)
= =f 9MD((a:: s) (,;; - Va. logp)(arls)] dac
Sbustute in (20),
1
/ ( s) [ss( s,t) Va, log log dagda; + constant
(an((aL:)) 8) pi(ails)
/ ZMD((S) / (:)((arla)) (xp((
exp(Q(ss a) ||so(ar,stt) Va. log a o (a ao)] dayda; + constant ZMD(S)
Fod // O:jo () (ttp:/aalaaaccp(((...cco
(s) ||50 8, Va, log q:o a |o)la dagda; + constant.
eexp (Q (8, an) (A) 80 (a:; 8, t) Va, log auu (az |n)) constant
E Gpontlld gejo
B.2.2. DERIVATIONS OF SOFT DIFFUSION ACTOR-CRITIC
We first restate the results. First, we define the reweighting function as Jaoftmax he(a, (s)p())
where (a, s) is a sampling distribution we choose. We require hy (aals)) to have full support on a space. Then we can show the following equivalence,
CosaeE (0; 8, 1) / gun((rr;) [100(ar; Va, la: da,
= =constant
E
exp (Q (s, ān) /A) so s, Va. log Φnj (ānl (ānla:)ll
+constant
(sDAC((,a..)
(14)
where doje is conditional Gaussian distribution defined as
1 Üi I at. āz
Φoje (āolaa) = (āo: vāc
(15)
Proof. First, we substitute (20) into (14) to get that
LMMwEoon (0; 1) = 1/ ha((a (10))/)))) - Va, logp:(ar|s)] daz
Sbustute in (20).
=lʃ hr (() ()ppp((rlaa) ||so 8t) Va, log 4:j0 (a; ao)l danda; constant
(21)
Then leverage the reverse sampling trick,
There exists reverse sampling distribution Φoje satisfying
Φoj aa)) = NV aa;; Vou at āi =n1T) x anj(a an) -N (at; Varao, (1)
(22)
and their score functions match
Va, log (a an) = Va. log ( 1 aa:)= a, Voiao
(23)
- āq
This is achieved by examining the density function,
14

---

Efficient Online Reinforcement Learning for Diffusion Policy
1 āt
dd/2
ag
vār
a2
Pole (āo aa)= = 2T
exp
( x(e
(24)
while
llVo a=1l"
Gejo (a:
an) (2 4/2 exp
(a=)-
Poje (ao aa)
(25)
2 - (it))
Leveraging the reverse sampling trick, change the weighting function in (21) to
h,(a,ls) Z (gg ()))))))))))))))
= 8) (s) exp ((,, ao)
9:jo (alo(ar)a))
(26)
Z(s)
=(a)) 4/2 ( an) (s
Considering the score equivalence, (21) can be further derived to
L Geotsmaax (0)
/f
(222)//2
he (āo exp (Q(s, aa) ||s0 (ari s,t) Va. log Ooje (āo az)ll dāgda; + constant
=(a,)a/2 4/2
(Q ao) (A) sq (a, s, Va. log Pope (āo a.) +constant
ac-
(SDAc(9.8.1)
(27)
C. Experiments
C.1. Policy Evaluation for Intropy-Reeularizee MDPs
C.I.I. POLICY EVALUATION FOR DPMD
The policy evaluation of DPMD minimizes the Bellman residual to learn the Q-function parameters,
- (T; Ex (Qc(s. a) (r(s, a) + YEg Ea Ba'aa(a [Qr (s', a')1))
(28)
C. POLICY EVALUATION FOR SDAC
Following the soft policy evaluation (4), we learn the Q-function parameters with the Bellman residual
- (C: m) E.,a~ (Q)(s, (r(s, a) À log (a|s) Tig (aaa(("aa")) [Qq (s', a')))))
(29)
In practice, as we sample batch of actions and select the one with the highest Q-value like (16) and add Gaussian noise with tunable standard deviations, directly select the log probability of additive noise as log (als).
C.2. Algorithm Details
We show the detailed pseudocode of SDAC here in Algorithm
C.3. Training Setups for the Toy Example
Consider Gaussian mixture model with density function
|xx
[3:3]]
1
||x +
[3; [1:
(x) = 0.8 2T
+ 0.2
exp
(30)
2m
2
15

---

Efficient Online Reinforcement Learning for Diffusion Policy
Algorithm 2 Soft Diffusion Actor-Critic (SDAC)
Require: Diffusion noise schedule Bt, 0g for E (1. T). MDP M, initial policy parameters 8g, initial Q-function parameters Co, replay buffer D 0, learning rate 3, KL-divergence coefficient A and target Atarget for epoch e = 1,2,... do
2: Sampling and experience replay.
3: Interact with M using policy $9,., thorough algorithm update replay buffer D.
4: Sample minibatch of (s, a, x') from D.
5: # Policy evaluation.
6: Sample a' via reverse diffusion process (8) with $6.
7: Update Qe with soft policy evaluation (4).
8: Policy cy Improvement for ffusion policies.
9: Sample uniformly from {1,2,...,?) Sample a, from by
10: Sample K a ~ Φoje for i 1.2,...,
11:
Compute Qe @u) and update 0g with empirical loss of
Qc s, (0 12c) -log] Σ exp (ee (s, n22)
Eg,t
exp
||s0 (ax; Va, log Φe (āp) |2))
K
12:
Update KL-divergence coefficient Ae + de-1 + 3(1e-1 Ararget).
end for
where the RSSM optimizes
E Xanup
(u(o)) |lso (xG;t) - Vx, log Φoje (xo]>;))]
(31)
Xor-
for the Gaussian sampling, f(x,) (0, 4I) for all and for uniform sampling, Pt a uniform distribution from [-6, 6)
on both dimensions. The score network is trained via the hyperparameters listed in Table The Langevin dynamics has direct access to the true score function Vx log (x))..
Table 3. Hyperparameters for the toy example.
Name
Value
Name
Value
Learning rate
3e-4
Diffusion noise schedule
linear
Diffusion steps
20
Diffusion noise schedule start
0.001
Hidden layers
2
Diffusion noise schedule end
0.999
Hidden layer neurons
128
Training batch size
1024
Activation Function
LeakyReLU
Training epoches
300
C.4. Additional Toy Examples
Additionally, we show a more complex two-moon distribution with known energy function as
logp(=)
|| -
+ log
exp
21
+ exp
+
(32)
0.2
2
0.3
0.3
We compare the proposed RSM with DDPM (Ho al., 2020) that has access to the data samples, and two other Boltzmann samplers that have access to the energy functions (32), iDEM (Akhound-Sadegh et al., 2024) based on diffusion with approximated noise-perturbed score functions, and FAB (Midgley et al., 2022) based on normalizing flows. The results are shown in Figure We can see that three diffusion- methods, RSM, DDPM, and iDEM, show similar performance. However, the two set of data samples are not separate enough for FAB, showing the advantage of the diffusion model and the limitations of normalizing flow from the restrictive invertible mappings.
16

---

Efficient Online Reinforcement Learning for Diffusion Policy
DOPM Generated
IDM Generated
FAB (wi 20hparane aigarmm.. than RSSK)
(a) PDF of two moon dis- (b) RSM Generation. tribution.
(c) DDPM Generation.
(d) iDEM Generation.
(e) FAB Generation.
Figure 5. Results to Fit Two Moon distribution, commonly used Boltzmann sampler benchmark. We compare RSM, DDPM, iDEM, and FAB. RSM, DDPM, and iDEM all recover two separate modes, while FAB shows connections between the two modes.
C.5. Baselines
We include two families of methods as our baselines. For the first family methods, we select 5 online diffusion-politc RL algorithms: QSM (Psenka et al., 2023), QVPO (Ding et al., 2024a), DACER (Wang et al., 2024), DIPO (Yang et al., 2023) and DPPO (Ren al., 2024). We include both off-policy (QSM, QVPO, DACER, DIPO) and on-policy (DPPO) diffusion RL methods among this group of algorithms. QSM uses the Langevin dynamics, one of the MCMC methods, with derivatives of learned Q-function as the score function. QVPO derives Q-weighted variational objective for diffusion policy training, yet this objective cannot handle negative rewards properly. DACER directly backward the gradient through the reverse diffusion process and proposes a GMM entropy regulator to balance exploration and exploitation. DIPO utilizes a two-stage strategy, which maintains a large number of state-action particles updated by the gradient of the Q- Q-function, and then fit the particles with diffusion model. DPPO constructs a two-layer MDP with diffusion steps and environment steps, respectively, and then performs Proximal Policy Optimization on the overall MDP. In our experiments, we use the trainin--from-seeee scratch setting of DPPO to ensure consistency with other methods.
The second family of baselines includes 3 classic model-free RL methods: PPO (Schulman et al., 2017), TD3 (Fujimoto et al., 2018) and SAC (Haarnoja et al., 2018). For PPO, we set the replay buffer size as 4096 and use every collected sample 10 times for gradient update. Across all baselines, we collect samples from 5 parallel environments in a total of million environment interactions and 200k epoches/iteration. The results are evaluated with the average return of 20 episodes across 5 random seeds.
C.6. Hyperparameters
Table 4. Hyperparameters
Name
Value
Critic learning rate
3e-4 3e-4, linear annealing to 3e-5
Policy learning rate
Diffusion steps
20
Diffusion noise schedules
Cosine
Policy network hidden layers
Policy network hidden neurons
256
Policy network activation
Mish
Value network hidden layers
3
Value network hidden neurons
256
Value network activation
Mish
Replay buffer size (off-policy only)
million
where the Cosine noise schedule means 3r = -1 ai ar-1
with ā = f(0) f(t) = cos
f()
(t/T+s Its
22)
17


Reviews:
[{"review_id": "VHxrEnadzR", "reviewer": "Reviewer_VLk9", "summary": "The paper proposes a novel method that leverages diffusion models to enhance SAC, but in a nontrivial way. To address the challenges of using diffusion policies—such as the need to track gradients through the entire reverse chain—the paper introduces RSSM, a new approach for estimating the score function. \n\nThe key idea is to reverse the standard diffusion sampling process. Instead of training a conventional diffusion model that samples from $p(a_0)$ and then generates $p(a_t | a_0)$, the proposed method samples from $p(a_t)$ and then estimates $p(a_0 | a_t)$. The paper also provides a theoretical justification that, under the RSSM formulation, it is still valid to learn the score function.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful comments. Here are the responses to the questions, \n\n## Theoretical claims\n\n**Q1:** \n\n> *How should $\\tilde{p_t}$ be chosen? Specifically, what is the appropriate distribution over $t$, and what form should $\\tilde{p_t}$ take? it seems that any choice might work for the algorithm—is that correct?*\n\n **A1:** Yes. Theoretically, as long as the support of $\\tilde{p}_t$ covers the support of $p_t(a_t)$, i.e., the distribution of the forward process at time $t$, we can recover the noisy score function by optimizing the RSSM loss (9). \n \n We discussed our empirical choice of $\\tilde{p}_t$ in Section 4.1. In online RL, RSSM loss depends on the Q-function, which might not be accurate for any action $a$. Therefore, we design $\\tilde{p}_t$ by sampling from the reverse process induced by the current policy, ensuring alignment with regions where the Q-function is more reliable.\n\n  ----\n\n**Q2:** \n> *Is the assumption in Eq. (10) too strong? Does it limit the model’s expressiveness? Furthermore, does this assumption still yield a valid ELBO? ...*\n\n**A2:** \nWe would like to clarify that the reverse sampling distribution $\\tilde q_{0|t}(a_0|a_t)$ in eq.(10) is not the unknown posterior distribution $q_{0|t}(a_0|a_t)$ in the reverse process, or its approximations.\nWe did not change the reverse process, so its expressiveness and the ELBO to learn it are not affected.  $\\tilde{q}_{0|t}(a_0|a_t)$ is a known distribution $\\mathcal{N}(\\frac{1}{\\sqrt{\\bar{\\alpha_t}}}a_t, \\frac{1-\\bar\\alpha_t}{\\alpha_t}I)$, only used in computing RSSM loss function.\n\nThe proposed RSSM loss and original DDPM loss to train diffusion models [1] are both derived from the same ELBO. However, DDPM loss is not tractable when we cannot sample from dataset $p_0$ and only know its energy function, which is the case in online RL. The novelty of RSSM is a tractable sampling-based method to train a diffusion model/policy when only the energy function is available. \n\n ----\n\n**Q3:** \n> *How is entropy estimated using this method?*\n\n**A3**:  When we implement the diffusion policy, we leverage a trick called efficient behavior policy, where we sample multiple actions and choose the one with the highest $Q$ function, which is also used in many diffusion policy papers like [1,2]. We notice that after the efficient behavior policy trick, the policy stochasticity is low, so we can estimate the entropy by only the entropy of additive Gaussian noise. \n\nOur strong empirical performance shows that this simple entropy estimation works well enough.\n\n[1] \nDing, Shutong, et al. Diffusion-based reinforcement learning via q-weighted variational policy optimization.\n\n[2] Janner, Michael, et al. Planning with diffusion for flexible behavior synthesis.\n\n ----\n\n## Experimental Designs\n\n**Q4:** \n> *Clarification of reverse sampling distribution $\\tilde{q}(a_0 | a_t)$ since action generation still requires reverse sampling, is it possible to directly use $\\tilde{q}(a_0 | a_t)$ instead?*\n\n**A4:** As the **A2** pointed out, $\\tilde{q}_{0|t}$ is not the reverse process posterior, and we did not change the reverse denoising process. Therefore, we still need to sample from the reverse process, which is the key to the rich expressiveness of diffusion models. \n\n----\n\n**Q5:** \n> *When normalizing Q-values, do you sample multiple actions for a given state, or is the normalization performed within a batch?*\n\n**A5:** We just normalize within the batch, no additional sampling is needed.\n\n----\n\n**Q6:** \n> *I noticed that DACER exhibits extremely high variance in your experiments. However, this phenomenon was not present in the original DACER paper, which used the v3 dataset. Do you have any insights into why this discrepancy occurs? Could you also provide some experimental results on the v3 dataset for a more straightforward comparison?*\n\n**A6:** Thank you for the keen observations. The differences stem from hyperparameter choices: the original DACER uses a smaller learning rate (1e-4) and 20 parallel environments, resulting in significantly more fresh data and lower variance. In contrast, we use 3e-4 with 5 environments across all algorithms for a fair comparison. \n\nWe provide DACER (using both our and original hyperparameters) and SDAC in both V3 and V4 in [Figure 2 and Table 2](https://anonymous.4open.science/r/rebuttal_pdf-D49E/rebuttal_figures.pdf). The original DACER parameters show great stability but much slower learning. The variances of V3 and V4 envs are similar.\n\nThe proposed SDAC can achieve stable learning with a larger learning rate and fewer envs (as 1 is commonly used), demonstrating better robustness compared to DACER.\n\n----\n\nWe appreciate your thoughtful review and hope our responses have clarified the concerns raised. We would be grateful if you would consider updating your score in light of these clarifications.", "claims_and_evidence": "All claims are well supported by rigorous proofs and illustrative examples.", "methods_and_evaluation_criteria": "The benchmarks used in the paper are standard and well-accepted.", "theoretical_claims": "I reviewed the theoretical proofs and did not notice any apparent issues. The theoretical results also make sense to me. However, I do have a few minor questions:\n\n1. How should $\\tilde{p}_t$ be chosen? Specifically, what is the appropriate distribution over $t$, and what form should $\\tilde{p}_t$ take? From the current description, it seems that any choice might work for the algorithm—is that correct?\n\n2. Is the assumption in Eq. (10) too strong? Does it limit the model’s expressiveness? Furthermore, does this assumption still yield a valid ELBO? In classical diffusion models, the assumption is typically on $q(a_{t-1} | a_t)$ being Gaussian, not $q(a_0 | a_t)$. I’m concerned that this assumption may be too restrictive.\n\n3. How is entropy estimated using this method? After learning the score function with RSSM, do we still need to integrate along the ODE to compute entropy? Since entropy is required to learn the Q-function, clarification on this step would be helpful.", "experimental_designs_or_analyses": "I have a few minor questions regarding the algorithm and experimental setup:\n\n1. Since action generation still requires reverse sampling, is it possible to directly use $\\tilde{q}(a_0 | a_t)$ instead?\n\n2. When normalizing Q-values, do you sample multiple actions for a given state, or is the normalization performed within a batch?\n\n3. I noticed that DACER exhibits extremely high variance in your experiments. However, this phenomenon was not present in the original DACER paper, which used the v3 dataset. Do you have any insights into why this discrepancy occurs? Could you also provide some experimental results on the v3 dataset for a more straightforward comparison?", "relation_to_broader_scientific_literature": "None.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "The paper is easy to follow, and the proposed approach is novel. For potential weaknesses and questions, please refer to my previous section.", "other_comments_or_suggestions": "None.", "questions_for_authors": "None.", "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": "H9ecm4tYAU", "reviewer": "Reviewer_pogC", "summary": "This paper highlights the challenges in mimicking an energy-based policy, primarily due to two key reasons: the intractability of the energy function caused by the partition term and the inherent limitation of online RL, where optimal policies are not directly accessible. While existing online diffusion policy algorithms attempt to address these issues, some other limitations remain. In this line, the authors propose RSSM, which samples a one-step estimate of the denoised action to precisely follow the score of the energy-based policy, along with SDAC, a novel diffusion policy algorithm built upon RSSM. The effectiveness of SDAC is demonstrated through experiments conducted on both toy examples and Gym benchmarks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful and constructive feedback. Below, we first provide detailed derivations of the reverse sampling trick, followed by detailed answers to the specific questions.\n\n## Reverse Sampling Trick Derivations\n\nCheck [rebuttal_proof](https://anonymous.4open.science/r/rebuttal_pdf-D49E/) if the equations do not show up.\n\nWe first explain (14) from Gaussian probability density function (PDF), \n\n$\\tilde{q}\\_{0|t}(\\\\tilde a\\_0|a\\_t)$ **is defined as** $\\mathcal{N}(\\tilde a_0;\\frac{1}{\\sqrt{\\bar\\alpha_t}}a_t, \\frac{1-\\bar\\alpha_t}{\\bar\\alpha_t}I)$ with PDF\n\n$$\n\\tilde{q}_{0|t}(\\tilde a_0|a_t)  =  (2\\pi\\frac{1-\\bar\\alpha_t}{\\bar\\alpha_t})^{-d/2}\\exp(-\\frac{||\\tilde a_0 - \\frac{1}{\\sqrt{\\bar\\alpha_t}}a_t ||^2}{2\\frac{(1-\\bar\\alpha_t)}{\\bar\\alpha_t}}) \n$$\n\nThe PDF of forward process perturbation kernel $q_{t|0}(a_t|a_0)$ is\n\n$$\nq_{t|0}(a_t|a_0)=\\mathcal{N}(a_t;\\sqrt{\\bar\\alpha_t}{a_0}, (1-\\bar\\alpha_t)I)  = (2\\pi(1-\\bar\\alpha_t))^{-d/2} \\exp(-\\frac{||\\sqrt{\\bar\\alpha}_ta_0 - a_t ||^2}{2(1-\\bar\\alpha_t)}) \n$$\n\n**Note that $\\tilde q_{0|t}$ is NOT the posterior distribution of $q_{t|0}$, which is usually intractable.**\nNotice that\n\n$$\n\\\\nabla_{a_t}\\log \\\\tilde{q}\\_{0|t}(a_0|a_t) = \\\\nabla_{a_t}\\\\log q_{t|0}(a_t|a_0) = - \\\\frac{a_t - \\\\sqrt{\\bar\\\\alpha_t}a_0}{1 - \\\\bar\\\\alpha_t} \\\\quad (*)\n$$\n\nWith these properties, we start from (13)\n\n$$\n\\iint \\tilde{p}\\_t(a_t | s) q\\_{t | 0}(a_t | a_0) \\exp (Q(s, a_0) / \\lambda) ||s_\\theta(a_t;s,t) - \\nabla_{a_t}\\log q_{t|0}(a_t|a_0)||^2 d a_0 d a_t\\quad (13)\n$$\n\nWe first substitute in the PDF of $q_{t|0}$,\n\n$$\n= \\iint \\tilde{p}\\_t(a_t \\mid s) (2\\pi(1-\\bar\\alpha_t))^{-d/2} \\exp(-\\frac{||a_t - \\sqrt{\\bar\\alpha_t}a_0||^2}{2(1-\\bar\\alpha_t)})\\exp (Q(s, a_0) / \\lambda) ||s_\\theta(a_t;s,t) - \\nabla_{a_t}\\log q_{t|0}(a_t|a_0)||^2 d a_0 d a_t\n$$\n\n$$\n= (\\bar\\alpha_t)^{d/2}\\iint \\tilde{p}\\_t(a_t \\mid s) \\underbrace{ (2\\pi\\frac{1-\\bar\\alpha_t}{\\bar\\alpha_t})^{-d/2}\\exp(-\\frac{||a_0 - \\frac{1}{\\sqrt{\\bar\\alpha_t}}a_t ||^2}{2\\frac{(1-\\bar\\alpha_t)}{\\bar\\alpha_t}}) }_{ \\tilde{q}\\_{0|t}(a_0|a_t)}\\exp (Q(s, a_0) / \\lambda) ||s\\_\\theta(a_t;s,t) - \\nabla\\_{a_t}\\log q\\_{t|0}(a_t|a_0)||^2 d a_0 d a_t\n$$\n\n$$\n= (\\bar\\alpha_t)^{d/2}\\iint \\tilde{p}\\_t(a_t | s) \\tilde{q}\\_{0|t}(a_0|a_t) \\exp (Q(s, a_0) / \\lambda)  ||s_\\theta(a_t;s,t) - \\nabla_{a_t}\\log \\tilde{q}_{0|t}(a_0|a_t)||^2 d a_0 d a_t \\quad (**)\n$$\n\nwhere the final equality we leverage $(*)$.\n\nThe $(**)$ equals the equation (15) times constant $ \\bar\\alpha_t^{-d/2}$, not affecting the optimal solution at given $t$. \n\n[1] revealed that ignoring the $t$-scale constant weights does not affect the performance of diffusion models when learned jointly like (16), thus we ignore the constant and sample uniformly from timestep $t$.\n\n[1] Ho, J., Jain, A., & Abbeel, P. Denoising diffusion probabilistic models. \n\n---\n\n### Experimental results or analysis\n\nPlease refer to [Figure 6](https://anonymous.4open.science/r/rebuttal_pdf-D49E/rebuttal_figures.pdf) for ablation study of $\\lambda$ learning rate.\n\n---\n\n### Essential References Not Discussed\n\n**Q2:** Relation to QIPO.\n\n**A2:** Our loss function looks similar to QIPO. However, the partition function estimation in QIPO is expensive and biased. The proposed SDAC does not need partition function estimation, and our estimation is unbiased. \n\nMoreover, the QIPO paper was published **after** ICML submission, we will add QIPO to the revision. \n\n---\n\n### Other Strengths And Weaknesses\n\n**Q3:** \"efficient behavior policy\" trick:\n\n**A3:** Yes, we use it. We have the same observation as QVPO, the diffusion policy is too random, thus we use the efficient behavior policy trick, sampling multiple actions and selecting the one with the highest Q-function. It is a standard trick in diffusion policy and diffusion-based planners. We skip it due to space limit and will clarify it in the revision.\n\n---\n\n**Q4:** \n> Log probability is intractable in diffusion policy, ..., standard PD update without incorporating log probability terms.\n\n**A4:** As Q3 mentioned, we use the efficient behavior policy trick and additive Gaussian noise. We notice the policy stochasticity is low after the efficient behavior policy trick, thus we can approximate the entropy using only the log density of the additive Gaussian. We provide results fixing the log density term in [Table 3](https://anonymous.4open.science/r/rebuttal_pdf-D49E/rebuttal_figures.pdf).\n\n---\n\n**Q5:** Entropy update.\n\n**A5:** Our $\\lambda$ updates are adaptive in line 185 of `/relax/algorithm/sdac.py`. We will fix the typo.\n\n**Q6:** Performance with 1M updates\n\n**A6:** We provide 4 envs that are not saturated with 200K updates in [Figure 1](https://anonymous.4open.science/r/rebuttal_pdf-D49E/rebuttal_figures.pdf). SDAC outperforms all the baselines.\n\n----\n\nWe appreciate your thoughtful review and hope our responses have clarified the concerns raised. We would be grateful if you would consider updating your score in light of these clarifications.", "claims_and_evidence": "While most claims appear to be well-founded, the derivation of RSSM is somewhat unclear, which could impact the overall clarity and validity of the statement. Providing a more detailed and structured explanation of this derivation would help strengthen the argument.", "methods_and_evaluation_criteria": "The ten Gym environments used for evaluation are appropriate and commonly used.", "theoretical_claims": "I checked all of them and they seem mostly sound. However, I remain uncertain whether the reverse sampling trick can be applied in this manner, specifically from eq. (13) to eq. (15). It would be helpful if the proof for this part were presented in a more comprehensive and detailed manner.", "experimental_designs_or_analyses": "I reviewed most aspects of the experimental design and analysis, including the experimental setups and ablations. Most parts of the main experiments seem appropriate, but the use of five parallel online environments for sampling is relatively uncommon. Nonetheless, since the authors have provided baseline performance under the same settings, fairness does not appear to be a significant concern. The toy experiments and ablation studies are generally well-designed, though an additional ablation study for the entropy coefficient would be beneficial. Moreover, given that QSM inevitably differs in its derivation from the original paper, both in implementation details (ex) transitioning from DDPM to Langevin) and hyperparameters, it would be better to carefully verify these aspects to ensure consistency.", "relation_to_broader_scientific_literature": "The key contribution of SDAC, particularly through its derivation from RSSM, is its ability to leverage the gradient of the Q function using approximately denoised actions when the score function is well-aligned. In contrast, most diffusion policy algorithms designed to follow energy-based policies inherently rely on the gradient of the Q function computed with noisy actions. However, in most cases, only the standard Q function is accessible, which may not effectively handle such noisy actions.", "essential_references_not_discussed": "One of the key contributions is exponential Q-weighted score matching, which closely resembles QIPO (https://openreview.net/pdf?id=HA0oLUvuGI), published in ICLR 2025. While there are differences in application areas (online RL vs. offline RL) and specific methodologies (reverse sampling vs. forward sampling), the fundamental idea appears to be quite similar.", "other_strengths_and_weaknesses": "Based on the provided code, there are some implementation components that do not fully align with or are not explicitly disclosed in the explanations presented in the paper. These include:\n\n1. The implementation involves sampling multiple actions and selecting the one with the highest Q value for both online interaction and evaluation. This approach is referred to as \"efficient behavior policy\" in the QVPO paper and is also widely adopted in diffusion policy algorithms for offline RL, as initially introduced by the IDQL paper. Disclosing this aspect is important, not only to ensure fairness in evaluation but also because the impact of exploratory sampling, i.e. adding Gaussian noise appears to be significantly influenced by this choice.\n2. Critic update. The paper states that SDAC employs soft policy evaluation. However, this does not seem feasible for diffusion policies due to the intractability of the log probability. Additionally, the provided code uses a standard TD update without incorporating log probability terms. Clarifying this discrepancy would help us understand the method more accurately.\n3. Entropy update. The provided code indicates that the entropy coefficient $\\lambda$ is updated via gradient descent using the difference between the target entropy and the entropy of $\\mathcal{N}(\\mu, (0.1\\lambda)^2 I)$, which corresponds to the distribution of online samples induced by exploratory sampling for given action from the diffusion policy. Notably, $ \\lambda $ at each timestep follows a scheduled approach rather than an adaptive one. Additionally, in line 13 of Algorithm 1, even when using the simplified notation, it seems more appropriate to use $\\lambda_{e-1}$ instead of $\\lambda_e$ immediately after $\\beta$. Clarifying these aspects would improve the transparency and accuracy of the explanation.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Can you clarify or provide additional details on the points mentioned in \"Other Strengths and Weaknesses\" that you consider essential?\n2. Can you provide a more comprehensive and detailed proof for the transition from eq. (13) to eq. (15) to improve clarity and understanding?\n3. Since the benchmark settings, using five parallel online environments for sampling over 200K steps, is relatively uncommon, can you provide partial results using a single online environment with 1M steps for comparison?", "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": "pDLxesyGF6", "reviewer": "Reviewer_e72f", "summary": "This work presents a diffusion-based online RL method called Soft Diffusion Actor-Critic (SDAC). The authors highlight the difficulty of training online RL methods due to inability to sample from the target distribution (optimal policy) and the computationally intensive nature of training some diffusion-based approaches. The authors utilize the connection between diffusion models and EBMs to motivate their approach of sampling from energy-based policies. Their approach, reverse sampling score matching, is theoretically motivated by deriving the loss function and showing it learns the correct score function. Experiments are performed on a simple 2-D Gaussian mixture to validate their diffusion-based approach, then online RL experiments are performed on the Mujoco benchmark which demonstrates the improved performance of SDAC compared to existing classical and diffusion-based model-free methods.\n\n## Update after rebuttal\nThe authors' response addressed most of my major concerns, hence I raised my score. The explanation for RSSM in the original paper was incorrect/misleading, hence the authors should incorporate their latest response below in the update manuscript.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback. Below, we provide our response and have merged some questions and reindexed the references for clarity.\n\n\n## Claims And Evidence\n\n**Q1:** \n> *reasons why other baseline methods have higher memory/time than SDAC.*\n\n> *benefit compared to other Boltzmann sampling methods such as iDEM [3] and FAB [4]?*\n\n  1. DQS/iDEM [1,2] use $K$ Monte-Carlo samples and $K$ energy function evaluations for one $\\nabla\\log(x_t)$. RSSM only needs one sample and energy function evaluation, saving memory and time. \n  2. The sampling in RSSM is unbiased, while iDEM sampling is biased according to Proposition 1 in [2].\n  3. DQS/iDEM/QSM use the gradient of energy function, which might be inaccurate when the energy function (Q-function) is learned in RL.\n  4. FAB [4] leverages flow models whose expressiveness is limited compared to diffusion-based models. FAB also suffers from high variance due to importance sampling.\n   \n> *Add Boltzmann samplers as toy example baseline.*\n\n  We add iDEM and FAB to the toy example, as well as the two-moon distribution in [Figure 5, 6](https://anonymous.4open.science/r/rebuttal_pdf-D49E/rebuttal_figures.pdf). We also count the training time and memory footprint in Table 5. \n\n> *The claim [3] “induce huge memory and computation costs” is not elaborated upon sufficiently. [3] updates the actions stored in the replay buffer (NOT as a set of particles as claimed by the authors)...*\n\nWe respectfully disagree with the reviewer’s claim. According to the official DIPO implementation [3], `line 116` in `main.py` explicitly calls `diffusion_memory = DiffusionMemory(state_size, action_size, memory_size, device)`, clearly indicating the use of an additional buffer, inducing memory and time cost.\n\n- QVPO,DACER,DIPO: discussed in Section 4.2.\n\n---\n\n**Q2:** Performance not comparable to spinning up SAC at 1 million updates.\n\n**A2:**  Spinning Up’s results run for **3 million** updates, not 1M as claimed. At 1M, it often underperforms SDAC, which uses only 200K updates.\n\n---\n\n## Methods And Evaluation Criteria\n**Q3:** Experiments on other tasks.\n\n- Please refer to [Figure 4 & Table 4](https://anonymous.4open.science/r/rebuttal_pdf-D49E/rebuttal_figures.pdf). We select three tasks from DeepMind Control Suite, and a PushT task that is common in diffusion policy.\n\n---\n\n## Theoritical claims\n\n**Q4:** RSSM does not satisfy Bayes' rule.\n\n**A4:** We clarify that the reverse sampling distribution **$\\tilde q_{0|t}$ is not, and does not intend to approximate the posterior distribution $q_{0|t}$ in the reverse process mentioned by the reviewer.** $ q_{0|t}$ is usually intractable, while $\\tilde q_{0|t}(\\cdot|a_t)$ is a Gaussian.\n\nWe respectfully argue that we did not ignore any terms. In fact, our derivation does not use the Bayes' rule, it is based on the probability density function of Gaussians. **We provide the detailed derivations in our response to Reviewer pogC**, including the score function equivalence. \n\nThe novelty of RSSM is integrating score matching loss (12) over measure $g$ in (31), so that it is tractable through unbiased sampling-based approximation, when we cannot sample from $p_0$. We realize that the current notations easily cause causality confusion. We will carefully revise the notations in the revision.\n\n---\n\n## Experimental Designs\n\nPlease see [Figure 1 and Table 1](https://anonymous.4open.science/r/rebuttal_pdf-D49E/rebuttal_figures.pdf) for performance of 1M updates.\n\n## Essential References Not Discussed \n\n**Q5:** \n> *The description of [1] is completely incorrect, in that it does not use Langevin sampling, rather it uses a diffusion process...*\n\n> *the authors fail to acknowledge that several existing approaches [1,2] that tackle these issues.*\n\nWe acknowledge a partial inaccuracy in our claim about [1]. Nonetheless, one of the core issues we highlight remains: the gradient of learned Q-functions might be inaccurate, affecting RL performance.\n\nWe did miss the literature on Boltzmann samplers. We will address this in the revision.\n\n---\n\n## Other questions:\n\n**Q6** Computing the log probability, the necessity of additive Gaussian noise, and using temperature parameters.\n\n**A6** Please refer to **A3,A4** in rebuttal to reviewer pogC. We use temperature parameters to control the additive noise scale.\n\nThanks for pointing out the typos and writing comments. We will revise accordingly. \n\n[1] Jain, Vineet, et al. Sampling from Energy-based Policies using Diffusion.\n\n[2] Akhound-Sadegh, Tara, et al. Iterated denoising energy matching for sampling from boltzmann densities.\n\n[3] Midgley, Laurence Illing, et al. Flow Annealed Importance Sampling Bootstrap.\n\n[4] Yang, Long, et, al. Policy representation via diffusion probability model for reinforcement learning.\n\n---\n\nWe appreciate your valuable comments and hope our responses have clarified the concerns raised. We would be grateful if you would consider updating your score in light of these clarifications.", "claims_and_evidence": "- The claim that SDAC is “more efficient compared to recent diffusion policies for online RL” (line 315-317) is not supported by enough theoretical or empirical evidence. The authors claim that their method avoids sampling from $\\pi_\\text{target}$ and the training loss has similar cost as the denoising score-matching loss, but both of these statements apply for [1] (which uses iDEM [2] and has been mischaracterized by the authors in line 321-325) as well as [3]. The conclusions drawn from the experiment in Table 2 which compares memory and wall clock time of different methods also have several issues, see ‘Experimental Designs Or Analyses’ below.\n- The claim that “performance is increased by more than 120% over SAC” is slightly misleading and there are certain caveats that must be clarified. First, based on the results in Section 5.2, this statement seems to be true for only 2 out of the 10 environments considered in the experiments. Second, the standard for off-policy methods is to perform one network update per environment step, however, in the experiments all methods perform 200k iterations for 1 million environment steps. Based on public benchmarks of SAC on Mujoco (https://spinningup.openai.com/en/latest/spinningup/bench.html), it seems after 1 million updates, SAC does reach performance levels on par or higher than those claimed for SDAC.\n\n[1] Jain, Vineet, Tara Akhound-Sadegh, and Siamak Ravanbakhsh. \"Sampling from Energy-based Policies using Diffusion.\" *arXiv preprint arXiv:2410.01312* (2024).\n\n[2] Akhound-Sadegh, Tara, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera et al. \"Iterated denoising energy matching for sampling from boltzmann densities.\" In *Proceedings of the 41st International Conference on Machine Learning*, pp. 760-786. 2024.\n\n[3] Yang, Long, Zhixiong Huang, Fenghao Lei, Yucun Zhong, Yiming Yang, Cong Fang, Shiting Wen, Binbin Zhou, and Zhouchen Lin. \"Policy representation via diffusion probability model for reinforcement learning.\" *arXiv preprint arXiv:2305.13122* (2023).", "methods_and_evaluation_criteria": "The authors consider the Mujoco benchmark as their main experimental setting, which is a popular suite of environments used for comparing RL algorithms. The baselines used seem to be exhaustive. There are two minor improvements I can suggest:\n\n- Mujoco benchmark has been considered saturated for some time, so it would be nice if the authors could demonstrate performance improvements on more settings to strengthen their claims.\n- The 2-D Gaussian experiments were meant as a proof-of-concept, but it would be beneficial to include other Boltzmann samplers [1,2,3] as baselines.\n\n[1] Akhound-Sadegh, Tara, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera et al. \"Iterated denoising energy matching for sampling from boltzmann densities.\" In *Proceedings of the 41st International Conference on Machine Learning*, pp. 760-786. 2024.\n\n[2] Midgley, Laurence Illing, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato. \"Flow Annealed Importance Sampling Bootstrap.\" In *The Eleventh International Conference on Learning Representations*.\n\n[3] He, Jiajun, Wenlin Chen, Mingtian Zhang, David Barber, and José Miguel Hernández-Lobato. \"Training Neural Samplers with Reverse Diffusive KL Divergence.\" In *The 28th International Conference on Artificial Intelligence and Statistics*.", "theoretical_claims": "I checked the proofs provided in the main paper as well as the appendix. One of the key ideas presented in the paper, the reverse sampling trick which is also used in Theorem 3.2, incorrectly replaces the distribution $q_{t | 0} (\\cdot | a_0)$ with $\\tilde{q}_{0|t}(\\cdot | a_t)$ in the loss function based on algebraic manipulation of the Gaussian pdf. However, from Bayes’ rule, \n\n$$\nq_{t|0}(a_t | a_0) = \\frac{q_{0|t}(a_0 | a_t) q(a_t)}{q(a_0)},\n$$\n\nwhere $q(a_t)$ and $q(a_0)$ are marginal distributions for $a_t$ and $a_0$ respectively. Since equation (13) is integrated wrt both $a_t$ and $a_0$, we cannot simply ignore these terms. In addition, when comparing their score functions, the term $q(a_t)$ cannot simply be ignored and hence the score function $\\nabla_{a_t} \\log q_{t|0} (a_t | a_0)$ cannot be replaced with  $\\nabla_{a_t} \\log \\tilde{q}_{0|t} (a_0 | a_t)$. \n\nOverall, it seems the authors fail to realize that the pdf $q_{t | 0} (\\cdot | a_0)$ is not just an algebraic expression but has a semantic meaning associated with it, where $a_0$ is the mean and $a_t$ is the random variable, which is NOT the same as $\\tilde{q}_{0|t}(\\cdot | a_t)$ where $a_t$ is the fixed mean and $a_0$ is the random variable.", "experimental_designs_or_analyses": "- The overall experimental design seems to be sound. There is the question of how would SDAC perform in relation to baselines if the standard practice of 1 million updates for 1 million environment steps were used.\n- The conclusions drawn from the benchmarking of memory and wall clock time seem misleading. In any implementation, there are several factors that can influence the memory and time, including the deep learning framework used, third party libraries, how optimized the implementations are etc. I believe it is incorrect to directly draw the conclusion that some method is more efficient without accounting for these factors. The authors should provide specific reasons why other baseline methods have higher memory/time that SDAC.", "relation_to_broader_scientific_literature": "The paper is one of many which aims to tackle challenges in online RL, specifically related to learning more expressive policy functions, by using a diffusion model. This paper attempts to solve this problem by proposing a novel loss function that bypasses the issue of sampling from a target distribution, and aims to reduce computation during the backward pass. \n\nWhile the specific ideas introduced in this work are novel, there are other existing works which tackle the same problems in diffusion-based policies for online RL [1,2].\n\n[1] Yang, Long, Zhixiong Huang, Fenghao Lei, Yucun Zhong, Yiming Yang, Cong Fang, Shiting Wen, Binbin Zhou, and Zhouchen Lin. \"Policy representation via diffusion probability model for reinforcement learning.\" *arXiv preprint arXiv:2305.13122* (2023).\n\n[2] Jain, Vineet, Tara Akhound-Sadegh, and Siamak Ravanbakhsh. \"Sampling from Energy-based Policies using Diffusion.\" *arXiv preprint arXiv:2410.01312* (2024).", "essential_references_not_discussed": "The paper cites most relevant works in the area of diffusion models applied for online RL. However, the descriptions of some of these works are completely incorrect, and the authors fail to acknowledge that some of the issues they tackle in this work have also been addressed by existing works. \n\n- The desription of [1] is completely incorrect, in that it does not use Langevin sampling, rather it uses a diffusion process (specifically iDEM) to learn a policy that samples from the Boltzmann distribution of the Q-function.\n- The claim of the authors that [2] “induce huge memory and computation costs” is not elaborated upon sufficiently. The approach updates the actions stored in the replay buffer (NOT as a set of particles as claimed by the authors) based on the gradient of the Q-function, and fits a diffusion model using score matching loss to these updated actions. The memory footprint should not be significantly different from the proposed method.\n- The authors mention that the necessity of sampling from target distribution and backpropagating through the diffusion chain are two main issues with employing diffusion-based methods. However, the authors fail to acknowledge that several existing approaches [1,2] have been proposed that tackle these issues.\n\n[1] Jain, Vineet, Tara Akhound-Sadegh, and Siamak Ravanbakhsh. \"Sampling from Energy-based Policies using Diffusion.\" *arXiv preprint arXiv:2410.01312* (2024).\n\n[2] Yang, Long, Zhixiong Huang, Fenghao Lei, Yucun Zhong, Yiming Yang, Cong Fang, Shiting Wen, Binbin Zhou, and Zhouchen Lin. \"Policy representation via diffusion probability model for reinforcement learning.\" *arXiv preprint arXiv:2305.13122* (2023).", "other_strengths_and_weaknesses": "- The proposed method claims to use the maximum entropy RL formulation via soft policy evaluation, but it is not clear how $\\log \\pi(a_t | s_t)$ is calculated efficiently in equation (2), since calculating  exact log likelihoods for diffusion models is a non-trivial (and computationally demanding) problem.\n- The overall structuring and presentation of the paper could be improved. Section 1 and 2 do a good job of introducing the main problem and describing the setting. From Section 3 onwards, it is easy to get lost in the mathematical details of the proposed method. It would be easier for the reader if the authors first present their method clearly describing the sampling and training procedures, then proceed to describe the theoretical results.", "other_comments_or_suggestions": "- The writing at the sentence level, particularly in the introduction could be greatly improved. Some examples of poor writing,\n    - Line 42-43: Huge successes of diffusion-based generative models have been witnessed recently\n    - Line 46-47: diffusion models achieved superior expressiveness and multimodality\n    - Line 13: naturally benefit ~~the~~ policies\n    - Line 16: offline RL, where expert datasets are presented\n    - There are many more examples, I suggest the authors revise the writing of the paper carefully\n- Missing one citation reference in Appendix A.", "questions_for_authors": "- In Section 4.1, why is there a need to add Gaussian noise for exploration? In principle, using a temperature parameter to scale the energy function should be enough since high temperatures would correspond to more random behavior. This is also the approach used in other energy-based methods like [1,2].\n- What is the benefit of the proposed RSSM method compared to other Boltzmann sampling methods such as iDEM [3] and FAB [4]?\n- How does SDAC compare to other methods when the number of iterations are increased to 1M steps (which is commonly used in online RL experiments)?\n\n[1] Haarnoja, Tuomas, Aurick Zhou, Pieter Abbeel, and Sergey Levine. \"Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.\" In *International conference on machine learning*, pp. 1861-1870. Pmlr, 2018.\n\n[2] Jain, Vineet, Tara Akhound-Sadegh, and Siamak Ravanbakhsh. \"Sampling from Energy-based Policies using Diffusion.\" *arXiv preprint arXiv:2410.01312* (2024).\n\n[3] Akhound-Sadegh, Tara, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera et al. \"Iterated denoising energy matching for sampling from boltzmann densities.\" In *Proceedings of the 41st International Conference on Machine Learning*, pp. 760-786. 2024.\n\n[4] Midgley, Laurence Illing, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato. \"Flow Annealed Importance Sampling Bootstrap.\" In *The Eleventh International Conference on Learning Representations*.", "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": "fG4VpxQwfR", "reviewer": "Reviewer_bxjo", "summary": "This paper studies training diffusion policy in the online RL setting and proposes an algorithm based on the energy-based view of diffusion models. It then evaluates the proposed algorithm on Mujoco tasks provided in Gym.\n\n# Main Ideas\n* To learn a diffusion policy, we need to learn the score function $s_\\theta(a_t;s,t)$.\n* This paper leverage the energy-based view of diffusion models: the score function $s_\\theta(a_t;s,t)$ matches the noise-perturbed score function $\\nabla_{a_t}\\log\\tilde{\\pi}_t(a_t|s)$  in expectation when considering an energy-based model. So the score function can be learned with samples $a_0$ and $a_t$.\n\n* The noise-perturbed score function can be expressed as the expectation of $ q_{t|0} (a_t | a_0) $  with regard to the distribution of the energy based model.\n\n* Given samples $a_t$, we can approximate such expectation with the reverse sampling distribution $\\tilde{q}_{0|t}$ and the Q-function.\n\n# Main Results\n\n* The authors show that their proposed reverse sampling score matching can be used to fit a diffusion model on a 2D Gaussian dataset.\n* Their proposed policy learning algorithm SDAC outperforms non-diffusion RL methods and recent methods for training diffusion policy in online RL.\n* The SDAC is more robust than other diffusion-based methods in the sense that it achieves consistent performance on 10 tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s thoughtful comments. Here are the responses, where we grouped relevant questions together,\n\n## Q1: Explanation about the reverse sampling trick.\n\n> *I found the reverse sampling trick (14) introduced in line 213 hard to understand and seems lacking any evidence or explanations.*\n\nWe provide the detailed derivations of the reverse sampling trick in our response to Reviewer pogC.\n\nThe proposed RSSM loss and original DDPM loss [1] both aim to train a diffusion model to generate data from the target distribution $p_0$. However, DDPM loss requires sampling from $p_0$, which is not possible in online RL since we only know the energy function of $p_0$. \n\nThe core novelty of RSSM is a tractable approach to train diffusion model/policy purely from the energy function rather than samples. The reverse sampling trick is an algebric operation to use sampling-based approximation to calculate the loss. \n\n\n\n> *The sampling distribution $\\tilde{p}_t$ in Theorem 3.2 is introduced without any explanation, which is very confusion. It is better to, at least explain that, it can be approximated with diffusion samples.*\n\nThe sampling distribution $\\tilde{p}_t$ is a distribution we can choose. Any distributions covering the support $p_t(x_t)$ can be used, such as Gaussian, Cauchy, etc.\n\nBeing approximated by diffusion models is not required for $\\tilde{p}_t$.\n\n**The choice of $\\tilde{p}_t$**. Our choice of $\\tilde{p}_t$ is based on the fact that the RSSM loss relies on the energy function ($Q$-function in RL) that is more accurate as samples of $a$ are close to the current policy. Therefore, we directly use the current reverse sampling at the $t$ step, which is exactly the current policy. \n\n[1] Ho, J., Jain, A., & Abbeel, P. Denoising diffusion probabilistic models. NeurIPS 2020.\n\n\n## Q2: \n\n> *In section 5.2, the authors do not provide results for DPPO, TD3 or PPO. I think it is better to provide results for all the baselines...*\n\nWe report the results for all the baselines, the reviewer might mean we did not provide the training curves in Figure 2 for DPPO, TD3, and PPO. We reported in [Figure 8 here](https://anonymous.4open.science/r/rebuttal_pdf-D49E/rebuttal_figures.pdf).\n\n\n## Q3 Performance comparison between TD3 and SAC\n\n> *because TD3 is reported to be consistently outperforming SAC on Mujoco Tasks.*\n\n> *Could you please analyze why TD3 cannot outperform SAC?*\n\n\nIn Table 1,\n- TD3 is better in Reacher, Humanoid, Pusher, Ant, and Swimmer.\n- SAC is better in HalfCheetah, Hopper, and Walker2d,\n- They are very close in InvertedDoublePendulum and InvertedPendulum. \n\nTherefore, it is difficult to say one algorithm outperforms another one.\n\nGenerally speaking, the performance gap between TD3 and SAC is sensitive to various hyperparameters. We reviewed some of our references using both as baselines and summarized the comparison in the figure below.\n\n\n| Reference | Tasks                  | TD3 is better  | SAC is better | Total |\n|-----------|------------------------|----------------|---------------|-------|\n| QVPO   | Gym MuJoCo             | 2              | 3             | 5     |\n| QSM    | Deepmind Control Suite | 4              | 4             | 8     |\n| DACER  | Gym MuJoCo             | 1              | 7             | 8     |\n| DIPO   | Gym MuJoCo             | 2              | 3             | 5     |\n\nTherefore, it is generally not straightforward to conclude that SAC consistently outperforms TD3, or vice versa.\n\n\n\n## Q4: about the toy example\n\n> *The author argue that RSSM can apply to any probablistic model with known enery function. The Gaussian mixture task is way to simple for support this... it seems that the proposed methods are quite restricted to RL problems...*\n\n**A4:** Our toy example on Gaussian mixture is just a proof-of-concept, not for showing the capability limit of RSSM.\nMoreover, modeling a Gaussian mixture is non-trivial since [1] highlights its slow-mixing issue, and our Figure 2(e) further shows that naive Langevin dynamics fails to recover the correct mixture within finite steps.\n\nTo further demonstrate RSSM’s effectiveness, we include results on the Two Moon distribution, a standard benchmark for Boltzmann samplers, in [Figure 6](https://anonymous.4open.science/r/rebuttal_pdf-D49E/rebuttal_figures.pdf). We also add iDEM [2] and FAB [3] as baselines. RSSM performs well and achieves the lowest KL divergence among all methods.\n\n[1] Song, Yang, and Stefano Ermon. Generative modeling by estimating gradients of the data distribution.\n\n[2] Akhound-Sadegh, Tara, et al. Iterated denoising energy matching for sampling from boltzmann densities.\n\n[3] Midgley, Laurence Illing, et al. Flow Annealed Importance Sampling Bootstrap.\n\n\n---\n\nWe appreciate your valuable comments and hope our responses have clarified the concerns raised. We would be grateful if you would consider updating your score in light of these clarifications.", "claims_and_evidence": "The main claims are reverse sampling score matching (RSSM) and a practical algorithm SDAC. \nAmong the explanation of RSSM, I found the reverse sampling trick (14) introduced in line 213 hard to understand and seems lacking any evidence or explanations.", "methods_and_evaluation_criteria": "The authors use ten continuous control tasks provided by the GYM repository, which is standard for online RL problems.", "theoretical_claims": "No, I have not check the correctness of any proofs. The proof scratch in page 4 is hard to follow due to the missing explanation for the reverse sampling trick.", "experimental_designs_or_analyses": "Yes, I have checked the results in section 5.1 and 5.2. In section 5.2, the authors do not provide results for DPPO, TD3 or PPO. I think it is better to provide results for all the baselines, at least for TD3, because TD3 is reported to be consistently outperforming SAC on Mujoco Tasks.", "relation_to_broader_scientific_literature": "While the authors argue that the proposed RSSM can applied to *any probabilistic model problems with known energy function*, the Gaussian mixture task is way to simple for support this. In this sense, it seems that the proposed methods are quite restricted to RL problems that learn a soft Q function.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "My remaining comments are mostly about the clarity of the paper. \n\n1. The paper lacks an intuitive explanation for why the proposed RSSM can work. Figure 1 is actually not very informative, as its messages are merely (i) the Q-function is approximated with policy evaluation (ii) the loss of the actor involves the approximated Q-function. I would suggest the authors to provide a figure that explains the main ideas summarized above and emphasize that the loss of the actor can be computed with only diffusion samples.\n\n2. The sampling distribution $\\tilde{p}_t$ in Theorem 3.2 is introduced without any explanation, which is very confusion. It is better to, at least explain that, it can be approximated with diffusion samples.\n\n3. The related work of this paper is scattered throughout the paper, which breaks the flow of this paper.\n\n4. It is better to put the \"Difficulties to train diffusion model in online RL setup.\" before 3.1, as they are not related to the energy-based view of diffusion models but with the main challenge to be resolved in this paper.", "other_comments_or_suggestions": "No.", "questions_for_authors": "1. Could you please elaborate on the reverse sampling trick?\n2. Could you please provide results for all the baselines mentioned in 5.2.1?\n3. Could you please analyze why TD3 cannot outperform SAC?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents SDAC, a novel diffusion-policy method for online reinforcement learning. The approach avoids unstable Q score matching and costly multi-step backpropagation by introducing reverse sampling, a technique that nonparametrically maps noisy actions to the denoised domain. The paper establishes a theoretical connection between reverse sampling score matching (RSSM) and regression to an energy-based policy under a custom measure, supported by a toy example using a Gaussian mixture target. SDAC also demonstrates strong empirical performance on the MuJoCo benchmark.

The paper's key strengths include consistent empirical improvements across diverse tasks, a simple yet novel core methodology, and solid theoretical grounding. However, some limitations remain, such as limited clarity in parts of the theoretical derivation, missing implementation details, and minor inaccuracies in the discussion of related work.

The rebuttal addresses most concerns through detailed explanations and additional experiments. These clarifications should be sufficiently reflected in the final version. In particular, the authors are encouraged to clearly articulate the role of reverse sampling and the custom measure, include missing implementation details, such as QVPO-style efficient behavior, present key results in the main text, and carefully verify baseline information by consulting both papers and official codebases, such as DQS and QSM. 

Overall, this paper is recommended for acceptance, provided that the raised issues and limitations are incorporated into the final version.
Decision: Accept (poster)