Energy-Weighted Flow Matching for Offline Reinforcement Learning
Authors: Shiyuan Zhang
Weitong Zhang
Quanquan Gu
Venue: ICLR 2025 Poster (2025)
Keywords: Flow Matching Models
Diffusion Models
Energy-guidance Generative Models
Offline Reinforcement Learning
Primary area: reinforcement learning
TL;DR: We study the energy-guidance in flow matching models and diffusion models and applied the proposed algorithm in offline RL

Abstract: This paper investigates energy guidance in generative modeling, where the target distribution is defined as $q(\mathbf x) \propto p(\mathbf x)\exp(-\beta \mathcal E(\mathbf x))$, with $p(\mathbf x)$ being the data distribution and $\mathcal E(\mathbf x)$ as the energy function. To comply with energy guidance, existing methods often require auxiliary procedures to learn intermediate guidance during the diffusion process. To overcome this limitation, we explore energy-guided flow matching, a generalized form of the diffusion process. We introduce energy-weighted flow matching (EFM), a method that directly learns the energy-guided flow without the need for auxiliary models. Theoretical analysis shows that energy-weighted flow matching accurately captures the guided flow. Additionally, we extend this methodology to energy-weighted diffusion models and apply it to offline reinforcement learning (RL) by proposing the Q-weighted Iterative Policy Optimization (QIPO). Empirically, we demonstrate that the proposed QIPO algorithm improves performance in offline RL tasks. Notably, our algorithm is the first energy-guided diffusion model that operates independently of auxiliary models and the first exact energy-guided flow matching model in the literature.

Published as a conference paper at ICLR 2025
ENERGY-WEIGHTE FLOW MATCHING FOR OFFLINE REINFORCEMENT LEARNING
Shiyuan Zhang' Weitong Zhang" Quanquan Gu
'Tsinghua University 2hNC-Chapel Hill Unniversity of California, Los Angeles sh i yuan-22-21 @ma ls. w...siinguaaa...mmm edwa. edu.cn, weitongz@unc. qguêca. ucla. edu
ABSTRACT
This paper investigates energy guidance in generative modeling, where the get distribution is defined as ((x) x p(x) (xp(-3E(x)), with p(x) being the data distribution and E(x) as the energy function. To comply with energy guid- ance, existing methods often require auxiliary procedures to learn intermediate guidance during the diffusion process. To overcome this limitation, we explore energy energy-guided flow matching, a generalized form of the diffusion process. We in- troduce energy-weighted flow matching (EFM), a method that directly learns the energy-guidke flow without the need for auxiliary models. Theoretical analysis shows that energy-weighte Additionally, we extend this methodology to energy-weighter diffusion models and apply it to offline reinforcement learning (RL) by proposing the Q-weighted Iterative Policy Optimization (QIPO). Empirically, we demonstrate that the pro- posed QIPO algorithm improves performance in offline RL tasks. Notably, our algorithm is the first energy-guie guided diffusion model that operates independently of auxiliary models and the first exact energy- guidded flow matching model in the literature.
flow matching accurately captures the guided flow.
INTRODUCTION
Recent years have witnessed the success of applying diffusion models (Ho ct al. 2020 Song et al,
2020) and flow matching models (Chen al. 2018 Lipman et al.) 2022) generative models.
Given this success, another important aspect is to guide generative models to achieve specific, con-
trolled outputs, such generating images for certain class (Ho & Salimans 2021 Dhariwal & Nichol, 2021), designing molecular structures with desired properties (Wang et al. [024]] Hooge-
boom et al. 2022), improving policies for reinforcement learning (Wang al.] 2022] Lu al.
2023). Guidance can come from various sources, such as classifiers, including both classifier guid-
ance (Dhariwal & Nichol 202T) and classifier-free guidance (Ho & Salimans (02I]. In addition, Lu
eet al. (2023) proposed using guidance from an energy function, where the distribution is generated from gxx) X p(x) exp(- (xxp--3E(x)) (x)), where the model is guided generate data x with lower energy Exx) from the original data distribution.
Several recent efforts have been made to learn and sample from the guided distribution g(x) using diffusion models. Chen et al. (2022) performed rejection sampling from the learned data distribution p(x). Lu et al. (2023) introduced intermediate energy function E-)), allowing the score function of q to be decomposed as Vx log ge (x) Vx log px(x) VxE((x) within the diffusion process. Lu et al. (223) also proposed contrastive energy prediction for training the intermediate energy E, relying on back propagation to calculate its gradient with respect to x. Wang et al. (2024) proposed to directly approximate the gradient of this intermediate energy function &, as the firreefield' guid- ance. However, all these methods require either additional neural networks, back propaaaopng,tion, or post-processing to compose the guided distribution (x), which introduces unnecessary errors and complexity. Therefore, the following question arises:
Q1. Can directly obtain an energy-guided diffusion model without auxiliary models?
Another challenge for energy-guided generative models lies in providing guidance in flow matching models (Chen et al., 2018) Lipman et al. (022),, which is a more general, simulation- -free counter- part diffusion models. Zheng al. (2023) explored the use of classifier-fres guidance for flow

---

Published as a conference paper at ICLR 2025
matching in offline RL. However, since flow matching models approximate the velocity field u, (x) for the dynamics of the probability density path pa (x()), it is highly non-trivial to obtain the guided velocity field ü, (x) for the distribution qe (x) under energy guidance. This presents the second key question:
Q2. Can we inject exact energy guidance into general flow matching models?
In this paper, answer the aforementioned two questions affirmatively by proposing an energy- guided velocity field and an energy-weighted flow matching objective, with extensions to energy- weighted diffusion models and applications in offline reinforcement learning. Our contributions summarized as follows:
In response to Q2., for general flow matching, we propose the energy-guided velocity field (x), based on the conditional velocity field ueo (x(xx))).. The proposed ū (x) theoretically guaranteed to generate the energy-guided distribution g(x) x p(x) exx))xxp(388(x
We introduce the tnergy-weighter flow matching loss to train neural network vg that approx- imates ü, (x). The energy weighted flow matching only requires the conditional vector field u, (x)x)) and the energy E(xo) for Xo from the dataset. As the answer to Q1., we extend this approach to diffusion models, proposing the snergy-weighher diffusion model. Energy-weighter diffusion model learns an eeergyyguided diffusion model directly without any auxiliary model.
We apply these methods to offline reinforcement learning tasks to evaluate the performance of the energy-weighte flow matching and diffusion models. Under this framework, we introduce an iterative policy refinement technique for offline reinforcement learning. Empirically, we demon- strate that the proposed method achieves superior performance across various offline RL tasks.
Vectors are denoted by lowercase boldface letters, such x, and matrices by uppercase boldface letters, such as A. For any positive integer k; the set 1,2,..... l denoted by [k)) and we define [k] = [k]] U {0}. The natural logarithm of x is denoted by log z. We use Pt. to represent the marginal distribution of time i, and Poe represent the conditional distribution of x given Xx. Similarly, Po denotes the original data distribution for the diffusion model at = 0, while pro represents the conditional distribution of x, given Xy the forward process of the diffusion model.
Notations.
2 RELATED WORK
Diffusion Models and Flow Matching Models. Diffusion models (Ho et al. 2020) and score matching (Song al., 2020) have emerged powerful generative modeling techniques in tasks such as image synthesis (Dhariwal& & Nichol (021). text-to-image generation (Podell et al. 2023), and video generation (Ho al., 2022). In addition to these frontier applications, the success of diffusion models has been further enhanced by accelerated sampling processes (Lu et al. 2022a b) and the extension of diffusion models discrete value spaces (Austin al.) 2021). Alongside the success of diffusion models, flow matching models (Lipman et al., 2022] Chen et al. 2018) provide an alternative for simulation-free generation. Unlike score-based approaches, flow matching models aim to learn the velocity field that transports data points from the initial noise distribution to the target data distribution. This velocity field can be viewed as a generalized form of the reverse process in diffusion models and can be extended to optimal transport (Villani et al. 2009), rectified flow (Liu) 2022), and more complex flows
Guidance in Diffusion and Flow Matching Models. Beyond learning and generating the orig- inal data distribution with diffusion flow matching models, significant efforts have been made to control the generation process to produce data with specific desired properties. Dhariwal & Nichol (202T) introduced classifier guidance, which decomposes the conditional score function V log p x into the sum of the data distribution gradient V log pxx) and the gradient from a classi- fier V log p(g)() To simplify this, Ho & Salimans (2021) proposed classifier-fre. guidance, which directly integrates V log p(y]x) into the score function. Sendera et al. (2024) studied diffusion- structured samplers by introducing the inductive bias in Langevin process. [u et al. (2023); Chen et al. (2022) further explored energy-based guidance, where the target distribution is defined as g(x) X p(x) cx)(-BE(x)). Unlike classifier guidance, energy-based guidance extends to real- valued energy functions E, making it particularly relevant for tasks such molecular structure generation. Specifically, [Cen et (202)); Cremer et al. (2024) employed rejection sampling to implement energy guidance, while Lu et al." 2023 : et al. (2024) used auxiliary models to esti-
2

---

Published as a conference paper at ICLR 2025
Table 1: Comparison between guidance methods. Exact Guidance? ? means if the model can generate p(x)p" (c/x) when 3 o 1. w/o Auxiliary Model? means if the method can direct learn the guidance without auxiliary model (v) or not (x).
Guidance
Exact Guidance?
wo Auxiliary Model?
Classifier gureeereggudddeeeeee (Dhariwal Nichol 2021
Classifier-fre free guidance (Ho & Salimans 2021
V
Contrastive energy prediction (Lu et al. 2023
X
Energy-weighte diffusion (ours)
mate the guidance from the energy function. We defer a more formal, technical comparison between the energy-based guidance and classifier-based guidance Table [ in Section [.3 the context of flow matching, [Zheng et al. (2023) introduced classifier-free guidance for flow matching in the domain of offline reinforcement learning.
Diffusion and Flow Matching Models in Reinforcement Learning. Recent advances in diffusion models and flow matching models have enabled a range of applications in reinforcement learning (RL). Janner et al.] (2022); Wang et al. (2022) explore modeling behavior policies using diffusion models. Building on these results, Chen al. [2022]
al. (2023) model the offline RL objective as energy- mnergy-guided diffusion process, while Ajay et al. (2022); Zheng et al. (2023) apply the same policy optimization using classifier-ff free diffusion and flow matching models. Chen al.] (2023); Hansen-Fstrund et al.] (2023) use diffusion models to regularize the distance between the optimal policy and the behavioral policy, and Fang Lan (2024): He et al. (2023) leverage diffusion models for constrained policy optimization. Another line of research (Jackson et al. [022]] Lee et al. 2024
Lu et al. 2024) focuses on using generative models to augment synthetic datasets.
3
PRELIMINARIES
3.1
CONDITIONAL FLOW MATCHING FOR GENERATIVE MODELING
Continuous Normalizing Flows (CNFs) (Chen et al. 2018 considers the dynamic of the probability
density function by probability density path p [0, I x gd + R20 which transmits between the data
distribution Po and the initial distribution (e.g.. Gaussian distribution) p1. The flo Φ [0, 1] x Rd ->
Rd is constructed by a vector field v: [0, 1] x Rd + Rd describing the velocity of the particle at position x, i.e., d
ae (x) = (x)) where Φ (x) = x. |
In order to ensure that the vector field generates the probability density path Pt, the following continuity equation (Villani et al.; 2009) is required:
dt (x) + div (x)v = O, Vx e Rd.
(3.1)
The objective of flow matching is to learn neural network vi to learn the ground truth vector field u, by minimizing their differences, i.e., LrM(O) = E, (x) (x) (x()]] with respect to the network parameter O. However, it is infeasible calculate the ground truth vector field u To address this issue, Lipman et al. (202)) suggests to match the conditional vector field (x) instead of the vector field u, (x). presented by the following theorem:
Theorem 3.1 (Theorem 1, 2; Lipman et al. 2022). Given the conditional vector field u (x that generates the conditional distribution Peo (x Xy). then the "marginal" vector field u, fx,, Pus (x) |dxo generates the marginal distribution (x).. In addition, up to a constant factor independent of 0, Flow Matching loss (P) and Conditional Flow Matching loss (C) are equal, where
LPx (0) Ex, 1 (x) (x) (CF = Ex,x0, Il, (xx) uxo(x((x))1
(3.2)
where 1 A(t), Xo follows the data distribution po(-) and x pro where Pro is generated by conditional vector field uan- Hence VeLFM (0) = (O)C
We adapt the notation of diffusion to unify the diffusion and flow matching. The notation here is dif- ferent from flow matching notations in Chen et al. (2018); [iiman et al. (2022) where pı represent the data distribution. The flow matching result remains unchanged except switching with 0.
3

---

Published as a conference paper at ICLR 2025
In practice, the conditional distribution (x xo) is usually modeled Caussian path with (eo(x(Xo) = N (1 I. Zheng et al. (2023) suggests that in this case, conditional flow matching equivalent to the score matching (Song al. 2020):
Lemma 3.2 (Lemma 1, Zheng et al. 2023). Let pxo(x(Xo) be Gaussian path with scheduler (14, st); i.e., Pro (x) N (/4x0: of I), then the velocity field uo xo) is related to the score function Vx log Pro (x (x(x0) by
u (x|xi) p x (jioo ( V log (x)
(3.3)
where ji and di are both the derivative of and o, with respect to time t.
In addition, Zheng al. (202)) proved that the reverse process of this diffusion process with Gaus- sian path Ht, can be written by
dx He x (pi0 - marr V log = un =x))
(3.4)
dt
P4
3.2 ENERGY-GUIDER DIFFUSION MODELS
The standard diffusion model aims learn and generate from the data distribution Po- However, instead of generating from Po. there are series of applications consider sampling from an energy- guided distribution (x() oc exp( (x) where & IR + is the energy function and is the strength of the guidance. Lu al.
2023 suggested to construct the score function Vx log q: (x) from the original score function Vx ( by introducing the intermediate energy
3 E R
function (x) through the following theorem:
Theorem 3.3 (Theorem 3. L. et al.
2023
Let qu (x) x (x() exp(- exp(-6E(x)) and define the for- (wixo, of I), and the marginal distribution q: (x(x,, (x)
ward process as qro (x Xo) = Pro (x Xo)
at time / defined by
(x(x) = /
Pc(x)= = =f( X
Let the intermediate energy function be
{x)x log E (x0/))) (xxppp(peexxpxxxxxx_
(3.5)
then the marginal distribution Pt and gr satisfy
qe x Pa (x(x)))xxxx)xxxppp exx()xxx({
Vx (xg = Vx logp:(x) - (x).
(3.6)
Therefore, Lu al. (2023) suggests firstly learn the intermediate energy function & using con- trastive energy prediction (CEP) and to learn the score function Vx (x() using standard diffu- sion models (e.g., DDPM (Ho al. 2020)). Then the score function of the energy-guided distribu-
tion Vx log (x) can therefore be composed according to (3.6).
METHODOLOGY
In this section, propose energy-weighted method for training both CNFs and diffusion models to generate the energy-guided distribution q(x) x p(x) exp(-3E Compared with Lu et al. (023); Wang et al. (2024), the energy-weighte method provide more straightforward way to ob- lain the energy-guided generative models and removes the necessaries of estimating the intermediate energy function & (x) and its gradient (x).
4.1 ENERGY- WEIGHTED FLOW MATCHING
In this subsection, we construct new energy guided flow generate the energy probability distribution. We also proposed two equivalent loss function to train the neural networks for approx- imating the energy-guided flow. We start by the first theorem suggesting a anergy-guides flow to generate the mergy-gguiee probability distribution q (x) & px(x) exp( (x))).
We swap the notation p and g align with our notation systems.
4

---

Published as a conference paper at ICLR 2025
Theorem 4.1. Given energy function E() and conditional flow (xo(x(xx) that generates the probability distribution Pto (x(xx))), the energy guided distribution qe (x(x) o pr(x) exp( is generated by the flow
= /
Por (xx(Xx)) xx exp(-BE(xo)) dxo: exp(- (x((x))
(x)
(4.1)
%0
which will generate distribution go(xx) oc po(x) exp(- The intermediate energy function is defined in (3.5).
Remark 4.2. Theorem suggests method construct the vector field (x(x) from the condi-
tional vector field uto (x Xo) and the intermediate energy function (x(x) in the closed-form solution.
holds universally to any conditional flow including the the optimal transport, Gaussian path or rec- tify flow. We will extend the discussion on the diffusion models in the next subsection.
Despite the closed-fome expression for the energy-gut energy-guided flow, it remains challenging to learn neural network v to match û, since the following two reasons. First, û, in (4.1) requires to sample over data distribution Xo. Secondly, the expression of ú, still requires the estimation of the interme- diate energy function &s. Previous methods are both using auxiliary neural networks to approximate either & (Lu al. 2023) or its gradient Vx& (Wang al. 2024). To overcome these two chal- lenges, the following theorem suggests a weighted flow matching objective which can be directly used learn ü, without the aforementioned procedures.
Theorem 4.3. Define the tnergy-weighter Flow Matching loss LEPM as
cxp(-E (x))
CEFM (0) Ex,x
Ex. |exp( E [v((x) (x)
(4.2)
and the Conditional Energy- Flow Matching loss CCEFM
exp(- (xxp-))(xx))
LCEFM Ex,xxxo
|[][(x) uso (x(xx)))
(4.3)
(x
(xx)))))
where the expectation on is taken over some predefined distribution A(t), xo is sampled from the data distribution Po(-) and x at time t is sampled by (x()) with conditional distribution Pto (x(xx)) generated by the flow uæo (xx xx). (e) and LCErM equal up constant factor. Hence VeLFFM 0) (aCCErM(())
Theorem suggests that minimizing LCEFM is equivalent to minimizing LHM- It is obvious that the global minimum of LCEFM is v (x) = ü, (x), given enough neural network capacity and infinite data. Therefore, one can use LCEFM to directly learn the guided flow (x(x., without calculating the intermediate energy function (x) or gradient.
Besides the aforementioned message, Theorem 4.3 suggests several understandings and intuitions in training the neural network vg which are discussed as follows
Remark 4.4 (Regarding the weighted energy guided loss CεMM). Instead of directly minimizing ExxlI (x) - ú, (x) 2. LErM places higher weight on the imput x with lower intermediate energy (x). Intuitively speaking, exp(-E(x)) can be viewed as prior distribution in generating 4 (x) x (x(x) (x)) Therefore, for all time I, areas with higher exp(-E(x)) will be more likely to be visited. As result, it would be more efficient placing more importance on x in these areas instead of learning ū uniformly for x E ddd
Remark 4.5 (Regarding the conditional weighted energy guided loss LCERM). The weight suggests how the energy "guides" the conditional flow matching. Fixing t and x and changing the form of expectations in 4.3) (C becomes
exp(- (x)(-)B(x)))
(CEFM(0; = (xpx(xxox)
| -xoo))x003
(o ( (Ex((-B((x)))
Intuitively speaking, velocity field (x)() will move the particle x to Xo- Therefore, when the energy guidance does not exist (i.e, E(x) 0), vx() is essentially finding the "center" of all Xu possibly generated from x following p(x, xx). In the presence the energy function E(x)), the learnt vector field vf (x) is biased to the conditional vector field (x(xx) with higher weight (xp((BBExx)) As a result, the learnt velocity field v (x) will generate Xg with lower energy E(xo)-
5

---

Published as a conference paper at ICLR 2025
Remark 4.6 (Connection with the importance sampling). The conditional weighted energy guided loss CCEPM can be also interpreted from the importance sampling techniques. Suppose we can sample directly from the data qo (x) x po(x) exp(- minimizing the following loss Lg will get a velocity field V, for generating distribution ạ0
(0) (xxx
[Ilv (x) Uzo (x xxxx))]l
where (xo(x|xo) (x) Since, where Z is constant, changing the data distribution from Q0 to po yields that
xx(x) (x)
Lg (0) = E xxxxx--x(xxxxxxxxxxxxxxx
(x( l (x) usa (x(
exp( exp(-3E(Xo))
R (oo()o)(exp( BE (xo) 11 llv((x) u (x) = LCFFM(();
where the second equation is given by (x(x) = po(x) -erp-1 according Lemma
4.2 WEIGHTED DIFFUSION MODELS
Theorem 4.3 suggests a general method learn an energy-guided flow v given any condition flow (x ). including diffusion flow (Song et al., (020). optimal transport (Lipman et al. 20022). rectified flow (Liu) 2022) or even more complicated u.o (x)x)). In this subsection, we restrict the analysis to the diffusion flow and present several useful analysis for the diffusion and score matching
models. The first corollary provides the closed-form score function Vx log ge (x) for the energy- guided distribution (x( x (x( exp(-E
Corollary 4.7. Under the assumptions claimed in Lemma 3.2 X
when Pro Xx) is Gaussian path (x(x)))) (xE((x) and
with scheduler (pe,ot), ot), we have log (xx)) V log
exp( (xx))) dxos exp(-Ex(x))
Vx log (x) = (x x) V (
(4.4)
Xx
where Vx log (x(xx) =
- = (0,,
Corollary 4.7 suggests a method to estimate the guided score function without calculating the gra- dient of the intermediate energy function (xE((x) as conducted in Lu et al. (2023). Then the following corollary suggests a similar mnergy-weighte diffusion model to train this score function vx log 9, (x) in practice.
Corollary 4.8. Define the Energy-weighted Diffusion loss LED and the Conditional Energy- weighted Diffusion loss LCED separately as
exp(-Ex (x))
(0) Re,x
|s (x))
Vx log lx)
E Ep.(
exp( (xp()))())))
CCED (0) Ex
12 (x) Vx log (x] (x
(xo (exp( - 38 (xx))
where the expectation is taken from t X(t), Xg (X())) and Pro (x Thus the marginal
distribution x is (x).. (f) is equal with LCED ( up to a constant and thus VeLED = (OLCED(O).
Remark 4.9. A similar approach is proposed in Wang et al. 2024 for estimating VxE((x) using neural network by
VxE((x)= = Epos [exp( (P0(-BB((x)) (Vx log pr (x) - Vx (xo
(4.5)
exp(-E((x))
Ex) for generating the score function Vx log qr (x). However, in order to obtain VxE, (x), Wang al. (2024) requires to estimate (x)( by sampling and approximate VxEx(x) via another neural network. In contrast, using LCED discussed in Corollary 4.8 removes the necessity of estimating both VxE((x) and (x). Therefore, Energy-went weighted diffusion can directly obtain the score function for guided distribution without additional sampling back-propagation.
and then plugging (45 back to Vx log (x) Vx log (x)

---

Published as a conference paper at ICLR 2025
Algorithm Training Energy- Weighted Diffusion Model
Input: Score function s! schedule (/4, ot), guidance scale 3, batch size B, time weight X(1
1: for batch {xi, {(x))
]] J. do
2: for index E [B] do
3: Calculate guidance Di softmax(-$E((x) = exp(-8E(x6)) Σj exp(-3E(x))
4: Sample 4 (0, 1), calculate Me, t sample E; ~N(0 1a) and X Mc. x ot.
5: end for
6: Calculate and take gradient step using (CD Σ, X(t)) (xx.) <c/or, le/o..13
7:
end for
In the implementation of the diffusion models, since x ( I) the conditional score func- tion Vx log Pro (x|x)) -c/o; where < = (x, - -oxo)//x/--- (N, In addition, the denominator Bpo(xo) (xxxp(-8((x))) can be approximated by the empirical average Σien (xp(-3E(xx) (LN in a batch. A practical algorithm is presented in Algorithm The training process is similar with the standard DDPM (Ho al. (020) or matching (Song al., 20200). The only difference is that we incorporate the weight by calculating a energy guidance using the softmax value of using
(xt) from the current batch in Line 3
4.3 COMPARISON BETWEEN CEP AND CLASSIFIER (FREE) GUIDANCE
In this subsection, we compare method with Contrastive Energy Prediction (CEP, Lu al. 2023), Gluuiaee (CG, Dhariwal & Nichol 2021) and Classifier-Free Guidance (CFG, Ho & Salimans 2021). We consider the guided distribution qo (x( x po (x (cs)p((ch)) (clx) where po(clx) the classifier, 3 is the guidance scale, is the desired class which we fix during the analysis. Comparing
with the formulation of the energy guidance qı (x) P1 (x) exp(-3E(x) (x)),, the "energy function"
can be interpreted {(x) = log (c)xx. To begin with, the following lemma provides the closed- form solution for the energy energy-guided diffusion and classifier guided diffusion
Lemma 4.10. Given the same guidance scale B and the same diffusion process, let the energy function defined by {xx) log p(c[x). the score function for CG and CFG are both:
Vx log (x) v log (x) + Vx log [Epo Pp) [c))) while the score function for energy-weighted diffusion and CEP both
(4.6)
Vx (
Vx log + Vx log Epo. (No (x)p (c/xo).
(4.7)
3
B=1
B=3
10
Negative
Positive
-2o 0. 2.5 (a) Ground truth of pxx)p" 1|x]; the distribution of p(x/c (Negative), p(x|c 1) (Positive).
0
-25 0.0 a/s
(b) Data sampled from classifier-free diffusion.
Data sampled from energy-weighted diffusion.
Figure 1: Visualization of the ground-truth distribution p(x)p²(( 2 1 1|x) with different values of 3, the posterior distribution p(x[c) with E (0, 1}. and the data sampled from classifier- diffusion and energy-weighted diffusion. The energy-weighted diffusion process demonstrates better perfor- mance when > 1. More examples and details of this experiments are provided in Appendix C
The score function (47) is exactly the score function that generates qn (x) oc (x exp(- (38E(x)) = po(x))(c(x)) according to Theorem 3.3 In sharp contrast, (4.6) is not guaranteed to generate

---

Published as a conference paper at ICLR 2025
the desired distribution du when B + 1 because [xpp F p (xo (x)))) (xx/ As demonstrated in Figure [ when B = 1, the distributions generated by CFG and energy-guided guided diffusion are both similar to the ground-truth distribution. However, when >> the distribution generated by CFG differs from the ground-truth distribution, whereas the energy-guide still generate the ground-truth distribution. Finally, the following lemma also verifies that when 3 = I, the snergy-weighte diffusion process is the same with the classifier-free guidance to learn the score function of the posterior distribution Vx log Pt (x(c):
diffusion
Lemma 4.11. Let 3 = 1 and assume that E(x) log p(c/x) for some fixed c, then (CED(O) is
(0) E.,xxxolelle (x) V Po xo)
(4.8)
where the expectation is taken over t ~ X(t), po(-le) and x ~ p=o(- The global optimal for (48) is s (x) Vx
In addition to the exact guidance provided by the CEP and energy-weighted diffusion, is important to highlight that the energy-weighter diffusion model eliminates the necessity for the intermediate energy model. This advantage is similar to the simplicity provided by the CFG, compared with CG. We summarized the difference and connection between energy energy-weighted diffusion, CEP, CFG and CG in [
Q- WEIGHTED ITERATIVE POLICY OPTIMIZATION FOR OFFLINE RL
We consider the episodic Markov Decision Processes denoted by M(S, A,,, P, 7) with A denot- ing the state and action space respectively. r is the reward function, P(-|s, a) is the transition kernel, and y is the discount factor. In offline RL, the data is collected by a behavioral policy p. The policy optimization with KL regularization (Peters et al. 2010 Peng et al. 2019) is formulated as
argmax Fa (- ( a)
[T. I),
(5.1)
where x denotes the state and denotes the action. Q(x. a) is the estimation of the state-action value function Q (x. a) = E = mx,(x, |x)0 = zb, ag = a, TT.. The closed-fome foroddffommmm-oo solution to (5.1) is (alx) pu(a(x) exp(
(BQQssQ(x) (x, (xx,)) a))- Following the procedure of the energy energy-weighted diffusion model
discussed in Section 4 we propose learning the score function s/ -a, x) or the velocity field v (a. x)
that generates the optimal policy (-+(x) using the Q-weighted diffusion loss LoD or Q-weighted flow matching loss CQF. respectively, defined as:
exp(3Q(x, a))
(0) = (x,(x,,)
5/ (a;::)) - Va, log Pin (a (ayaa)))
(5.2)
Baxp(-1))) (exp(((((xx a)]
exp(3Q(x, a))
(Qr(O) = (x,a)
v. (a,: u
(5.3)
E 1x) (3 ]
where the expectation is taken over t X(t), (x, a) is sampled from the offline RL dataset, and a, ~ (aa)- Two components are essential for training either (5.2) or (5.3): First, the behavioral policy (--(x) can be trained via standard diffusion or flow matching models. Second, any Q function derived from offline RL algorithms can be used as Q(x, a) in (5. or (.3)
5.1 PROPOSED ALGORITHM
We present the algorithm sketch for the Q-weighted diffusion process in Algorithm 2 and defer the detailed implementation to Appendix [ From high-level overview, the algorithm first trains the score model to match the behavioral policy (a x) for K1 episodes, then trains the Q func- tion for K₂ episodes. The algorithm then performs the Q-weighted iterative policy optimization (QIPO) as follows: First, in Line 10 using the current score function S the algorithm samples several support actions a, to estimate the expectation Eū-p(((x) (5.2). In Line 13 the algorithm optimizes O with respect to the empirical estimation of (0). Therefore, assuming the score function (-Ix) generates the behavior distribution ju(1xx after the warm-up for K2 episodes, using the sampled support action a.j. the optimal score function for Line [5] corresponds to the policy Ti = ((-(x) (-((((((((l)))) When the number of episodes & = Keenew + 1, the algorithm revis- its Line 10 with the new score function and regenerates the support action set with the new policy
8

---

Published as a conference paper at ICLR 2025
Algorithm 2 2Q Q-weighted iterative policy optimization for offline RL (diffusion)
Input: Score function 8 (-), schedule (je, at), guidance scale 3, batch size B, weight on time X(1
Input: Offline RI. dataset D {(x, a. x)),,, number of epochs K1, K2 and K3. Q function Q
Input: Support action set size M, support action set renew frequency Krenew
1: for diffusion model warm-up epoch & [K]] do
2: Train sy for each batch {xx, a, x,, r: r:) C 'D following standard diffusion model.
3: end for
4: for Q-learning warm-up epoch k: E [K2] do
5: Train Qo for each batch {(x,, a,, rx)) C
6: end for
7: for policy improvement step [K3] do
8: for batch (xx, a, rx) c D do
9: k mod Kennew then
10:
Sample support action set a.j using score function s (.)x;) for all i [B].j j [M]
Denote a:n a,, sample lij U(0, and aijit, (u aij
aij, o. 1) for all j [M]
Calculate guidance gij = softmax (x, a:y)) = exp(AQ &.y))
12:
Σj-o (x.,)
13:
end if
14:
Calculate loss (0) E
1. M(ti5)965 S (aigity) Va, log log (aig,ty la:;)
15:
Update I using the gradient of (0)
16:
end for
17:
end for
T1. Thus the target score function for Line to optimize is x²(+(x) x (+(xx) (- x ((-1x) exp(2,3Q (x)). As result, denoting = (k - 1) mod Krenow the policy generated by the score function S; is:
0x
(a/x) x (aaxx) exp(3Q* (a.x)) x))
u(alx) exp((1 (a.
(5.4)
which will implicitly inerease the guidance scale B.
Similar weighted policy optimization approaches have been applied in Kang al. (2024); Ding et al.] (2024). However, QIPO builds the relationship between the KL-regularized policy opti-
mization so that QIPO can iteratively improve the policy as described (5.4). Compared with directly setting a large guidance scale 3, QIPO makes the support action set a to be more con- centrated the space with higher Q values.
result, QIPO learns a more robust Q-weighted score function se compared to one-step Q-weighted diffusion with larger B. Second, QGPO Lu (ct al.
2023 introduces scaling factor and composes the score function as Va. log Ty (a, [x]
Va. Va log 14 ag x) sVa. Qx a,), where Q, is the intermediate function, similar the E, in Section However, as we discussed in Section 4.3 about the comparison of the CFG, since
sQi(a,, x) = - log (ex.:nlllaeee(((((((e x)] * log (exp:al x)].
using guidance scale > does not guarantee generating policy strictly regularized by the behavioral policy u(a]x). In contrast, as (5.4) suggests, our approach strictly follows the formulation o (a(x) & p(a(x) (xppxx(((xXX)) (a(x)) regularized by the behavioral policy.
5.2 EXPERIMENT RESULTS
We evaluate the performance of QIPO with flow matching and diffusion model on the D4RL. tasks (Fu et al. 2020) in this subsection.
Experiment configurations We implement the flow matching model QIPO-OT using the optimal- transport conditional velocity fields (Lipman et al. 2022) and the diffusion model QIPO-DifT using VP-SDE (Song et al. [220)p We use the same network structure as QGPO for a fair comparison of the efficiency with QGPO. We defer the detailed experiment configurations in Appendix E.I
We compare results with other state-of-the-a benchmarks, including Diffusion- QL (Wang et al.
2022), QGPO (Lu et al.) (2023), IDQL Glansen-Estrucee et al., 2003). SRPO (Chen et 2023) and
Guided Flows (Zheng al. 2023) and present the results Table 2 As the experiment results

---

Published as a conference paper at ICLR 2025
Table 2: Evaluation numbers of D4RL, benchmarks (normalized as suggested by Fu et al. (2220) We report mean + standard deviation of algorithm performance across random seeds. The highest performance is boldfaced highlighted. The performance within 5% of the maximum absolute value in every individual task are highlighted.
Dataset
Enviromment
SIBC
QGPO
IDQL
SRPO
Guided Flows
QIPO-Diff (ours)
QIPO-OT (ours)
Medium-Expee
HalfChectah
92.6
93.3
95.9
92.2
97
91.14 = 0. 48
91.45 = 0.49
Mediam- Expert
Hopper
108.0
108.6
100.1
112.12 0.42
108.02 5.19
Medium- Expert
Walker2d
109.8
110.7
112.7
114.0
94
110.14 + 0.51
110.87 + 1.04
Medium
HallCheetah
45.9
54.1
51.0
60.4
49
48.19 +: 0.20
54.16 1: 1.27
Medium
Hopper
57.1
98.0
65.4
95.5
50.53 + 9.96
94.05 + 13.27
Medium
Walker2d
77.9
86.0
82.5
84.4
77
81.99 = 0.46
87.61 =: 1.46
Medium-Replan
HalfChoctah
37.1
47.6
45.9
51.4
45.27 + 0.42
48.04 = 0.79
Medium-Repe Replay
Hopper
86.2
96.9
92.1
101.2
89
101.23 = 0.47
101.25 = 2.18
Medium- Replay
Walker2d
65.1
84.4
85.1
84.6
78
90.08 # 4.53
78.57 + 26.00
Average (1 Locomotion)
75.6
86.6
82.1
87.1
79.4
50.3
Default
AntMaze-umare
92.0
96..4
94.0
97.1
97.5 0.53
93.62 7.05
Diverse
AntMaze-umaze
85.3
74.4
80.2
82.1
73.58 + 6.42
76. = 9.93
Play Diverse
AntMaze-medium
83.6
84.5
80.7
52. +
80.00 13.66
AntMaze-medium
82.0
83,8
84.8
76.0
86,00 = 8.65
86,42 5.44
Play
AntMazz-large
59.3
66.6
63.5
53.6
73.25 10.90
55.5 29.39
Diverse
AntMaze-large
45.5
6188
67.9
53.6
20.40
32.13 =: 23.16
Average (AntMaze)
71.2
78.3
79.1
73.6
77.3
71.96
suggests, QIPO-Diff and QIPO-OT consistently outperform the baselines in various tasks. We defer
more baseline algorithms for comparison Table I in Appendix D
Among these benchmark algorithms, we would like to highlight the comparison between Guided Flows (Zheng et al., (023) and Q-Guided Policy Optimization (QGPO, Lu et al. 2023). Firstly, compared with Guided Flows (Zheng et al. 2023), QIPO-OT enjoys higher performance across many different tasks. This improved performance is due to the fact that energy-based guidance will provide more accurate guidance compared with classifier free guidance, as discussed in Section 4.3
Secondly, compared with QGPO (Lu et al. 2023). QIPO-Diff directly learns the energy-guided score function without estimating the intermediate energy function. As a result, QIPO does not require the back- propagation to calculate the gradient of the interme- diate energy function (x) and therefore enjoys faster sampling speed compared with QGPO when us- ing the same score network, as presented in Table In addition, since the QIPO guarantees the strict formu- lation π (a]x) x) x p(a[x) exp(s3Q* (a|x)) as shown in (.4)) QIPO enjoys better performance compared with QGPO on various tasks. We defer more detailed discussions on the advantage of QIPO-OT
Table 3: Comparison of the running time for action generation between QGPO (Lu et al. 2023) and QIPO, averaged over 1500
runs. The percentage reduction in time
compared to QGPO is also reported.
Method
Time (ms)
QGPO (Lu et al. QIPO-OT (ours)
2023
75.05
27.26 (43.68%)
QIPO-Diff (ours)
55,86 (-25.599)
to Appendix E22
Ablation Study. We conduct ablation study on changing the support action set M, policy renew period Krenw and the guidance scale B. We defer the detailed ablation study result in Appendix E.3
6 CONCLUSION AND FUTURE WORK
In this paper, we explored the energy guidance in both flow matching and diffusion models. We introduced Energy-weighted Flow Matching (EFM) and Energy-weighted Diffussinn (ED) by in- corporating the energy guidance directly into these generative models without relying on auxiliary models or proseprocessing steps. We applied the proposed methods in offline RL and introduced Q-weighted Iterative Policy Optimization (QIPO), which enjoys competitive empirical performance on the D4RI. benchmark. To the best of our knowledge, this work is the first to present an energy- guided flow matching model and the first algorithm to directly learn an energy-guided diffusion model. While the current QIPO focuses offline RL without interacting with the environment, we leave the extension to online RL, where guidance from the Q-function can be updated through online interactions
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
This paper focuses on the methodology of generative models and strictly adheres to ethical guide- lines and standards. As with other deep generative modeling techniques, energy-weighted diffusion and flow matching model have the potential to generate harmful content, such as "deepfakes", and may reflect or amplify unwanted social biases present in the training data. However, these broader societal impacts are not directly relevant to the specific contributions of our work. Therefore, we do not believe any unique ethical concerns negative social impacts arise from this paper.
REPRODUCIBILITY STATEMENT
We provide detailed descriptions of our experimental setups, dataset construction processes, and code implementation in the supplementary materials to ensure the reproducibility the Energy- weighted Diffusion and QIPO methods. The full experimental configurations presented in Ap- pendix [II] To further support research in the community, we will release the model checkpoints following the de-anonymization process.
ACKNOWLEDGMENTS
We thank the anonymous reviewers for their helpful comments. Part of this work done while WZ was a PhD student at UCLA. WZ and QG are supported in part by the NSF grants DMS-23231 113, CPS-2312094,, IIS-2403400 and the research fund from UCLA-Amazon Science Hub. WZ is also supported by UCLA dissertation year fellowship. Part of this work is supported by the Google Cloud Research Credits program with the award GCP376319164 The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies.
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B Tenenbaum, Tommi S Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? The Eleventh Interna- tional Conference on Learning Representations, 2022.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 14:11791117999 2021.
Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. In The Eleventh International Conference on Learning Representations, 2022.
Huayu Chen, Cheng Lu, Zhengyi Wang, Hang Su, and Jun Zhu. Score regularized policy optimiza- tion through diffusion behavior. In The Twelfth International Conference on Learning Represen- tations, 2023.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances neural information processing systems, 31, 2018.
Julian Cremer, Tuan Le, Frank Noé, Djork-Arné Clevert, and Kristof T Schütt. Pilot: Equivariant diffusion for pocket conditioned de novo ligand generation with multi-objective guidance via importance sampling. arXiv preprint 4xX.::2405.1492. 2024.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in neural information processing systems, 34:8780-877 877944 2021.
Shutong Ding, Ke Hu, Zhenhao Zhang, Kan Ren, Weinan Zhang, Jingyi Yu. Jingya Wang, and Shi. Diffusion-based reinforcement learning via w-weighted variational policy optimization. arXiv preprint arX 2024.
11

---

Published as a conference paper at ICLR 2025
Zeyu Fang and Tian Lan. Learning from random demonstrations: Offline reinforcement learning with importance importance-sampled diffusion models. arXiv preprint arXiv:2405. 2024.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep
data-driven reinforcement learning, 2020.
Philippe Hansen- Esttrcc, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304. 10573, 2023.
Longxiang He, Linrui Zhang, Junbo Tan, and Xueqian Wang. Diffeps: Diffusion model based constrained policy search for offline reinforcement learning. arXiv preprint arXiv:331000533,, 2023.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NewrIPS 2021 Workshop Deep Generative Models and Downstream Applications, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-66 68000,6811100 2020.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances Neural Information Processing Systems, 35:8633- 8646, 2022.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3d. In International conference on machine learning, pp. 8867- 8887. PMLR, 2022.
Matthew Thomas Jackson, Michael Tryfan Matthews, Cong Lu, Benjamin Ellis, Shimon Whiteson, and Jakob Foerster. Policy-guize guided diffusion. arXiv preprint MxXiv:2404.06356 2024.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine Planning with diffusion for flexible behavior synthesis. In International Conference Machine Learning, pp. 9902-9915. PMLR, 2022.
Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Jaewoo Lee, Sujin Yun, Taeyoung Yun, and Jinkyoo Park. Gta: Generative trajectory augmentation with guidance for offline reinforcement learning. arXiv preprint Stti:::240.166000 2024.
TP Lillicrap. Continuous control with deep reinforcement
learning.
arXiv preprint
150.009990027, 2015.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow match- ing for generative modeling. In The Eleventh International Conference on Learning Representa- fions, 2022.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint hDl:22009.14777 2022.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5788 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arX1v:2211.01095, 2022b.
Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive energy prediction for exact energy- guided diffusion sampling in offline reinforcement learning. In Inter- national Conference on Machine Learning, pp. 22825-22855 PMLR, 2023.
Cong Lu, Philip Ball, Yee Whye Teh, and Jack Barker-Holder Synthetic experience replay. Ad- vances in Neural Information Processing Systems, 36, 2024.
12

---

Published as a conference paper at ICLR 2025
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Alvantage-weightede regression: Simple and scalable off oof-po reinforcement learning. arXiv preprint 1rX00:1910.000177, 2019.
Jan Peters and Stefan Schaal, Reinforcement learning by neward-weighted regression for operational space control. In Proceedings of the 24th international conference Machine learning, pp. 745- 750, 2007.
Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, pp. 1607-1612, 2010.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. In The Twelfth International Conference on Learning Representations, 2023.
Victor Gareia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural net- works. In International conference on machine learning, pp. 9323-9332. PMLR, 2021.
Marcin Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Seimeca, Jarrid Rector-Brooks Alexandre Adam, Yoshua Bengio, and Nikolay Malkin. Improved off off-policy training of diffusion samplers. In The Fighth-Fighth Annual Conference on Neural Information Processing Systems, pp. 1-30. ACM, 2024.
Yang Song, Jascha Sohl- Sihk-Dickstten, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- fional Conference on Learning Representations, 2020.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Yan Wang, Lihao Wang, Yuning Shen, Yiqun Wang, Huizhuo Yuan, Yue Wu, and Quanquan Gu. Protein conformation generation via force-guided se (3) diffusion models. arXiv preprint arXiv:2403. 14088, 2024.
Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2022.
Qinqing Zheng, Matt Le, Neta Shaul, Yaron Lipman, Aditya Grover, and Ricky TQ Chen. Guided flows for generative modeling and decision making. arXiv preprint arXiv:2311. 2023.
13

---

Published as a conference paper at ICLR 2025
A COMPARISON WITH OTHER DIFFUSION-BASED RL
Besides the aforementioned works RL, there is another series of work using diffusion mod- els in RL by considering the reward- weighhed regression (Peters & Schaal 2007) where the score- matching loss function formulates (s,(((s, a)) |E (a'; ; (),0)-55 e]] In particular, EDP (Kang al. 2024 considers the f(Q(s, a)) in the format of CRR or IQL formulated as fCRR = expl QQ 8, (s a') a))) TCRR] or fIOL expl(Qa (s, a) (=(s)) TioL and QVPO (Ding et al. 2024 aggressive exploration in online RL with f(Q(s, a)) = Q(, a.. Compared with QIPO, FCrr assumes {aal ss [ā (s) with āg vI/a=a* v1/a* Ieg(a k)s). This oversimplifies the policy as a Gaussian policy and requires the score-network
considers
for calculate the weight. figr. requires an additional value network for (s). Most importantly, both EDP (Kang al., 2024) and QVPO (Ding et al. 2024), are designed to optimize the diffusion policy to maxi- mize the Q function. However, in the energy-guided generation for r{ass) o p(aas) exp(((s.))))))))) a)), where both maximizing the value and aligning with the Q function are equally important, these works cannot align with the original distribution μ and therefore cannot be applied to guide broader generative models applications like image synthesis molecular conformation generation.
B THEORETICAL PROOF OF THE RESULTS
We provide the theoretical proof for all theorems, corollaries and lemmas in this section.
B.1 PROOF OF THEOREM 4.1
We start by the following lemma which provides the closed-form expression of the distribution q. (x) Pe (x) (x)(
Lemma B.1. The closed-form expression of (x() x exx)((((())) is
(x()) exp(
(x(x) Exo /exp( ( (exp(=88(xx)) (x,)))
(B.1)
where the expectation is taken with respect to the data distribution Po The denominator Ex, [xxp(={Exxx)] is constant with respect and
Proof. is easy to verify that
(xp((x))p((EE(x)) = { x
= =Sxo SandxPulf PP (x, xo) (xo))
=x exp( 38 (xx) )dxo
Ex
(B.2)
where the second equation is due to the fact that pt (x (x) (Xo (x) p(x. Xn) pu(xo)Por (xo (x);
and the third equation is due to the fact that J Po (x xo ) = 1. Using B.2 to normalize ge (x) x Pa (xx) exp(-Ex (xx))) yields the closed-fom form solution presented in (B.1).
Then we ready to proof Theorem
We start by considering the condition flow (x(xx) that generates the con- ditional distribution sequence Pro (x)xx(()) According to the continuity equation (3.1). we have
Proof of Theorem 4.1
(xpo(x|Xo) div [pio(x(xo) (x|xo)]
(B.3)
dt
14

---

Published as a conference paper at ICLR 2025
Considering the dynamic of probability distribution (x(x), Lemma B.1 yields
dx)(x)
d (x( exp(-&((x))
dt
dt Ex (exp(=xB(xx)))
(x() x, Poe (x (x) exp( (xxxppppp3xxxdsaaaee
dt
Ex, (xxp(=BE(xx))
c ʃ Xxx P(0,, x) exp(-3E(xx))
Exx (exp( (xn))
S /xxx Peo((((Xx) x (olSooPPloee (X()) (x)) exp(=3E (dx)
dt
Ex, exp( - 13-5 (3E3(() (xn))]
/ X dt
d
Piu (xoxx))))o(xx)) exp(-3E(x)))dax
(B.4)
Ex, [expp-
where the second equation is due to the definition of (x)): the third and the forth equation is due to the Bayesian equality pe(x)pox (xx))ooeXool (Xo|x) p(Xg. x) Po (o Poo (x (o): and the last equation is due
to the fact that the denominator Ex. exp(pppee (3E X )] is constant with respect to time t. Plugging
the definition of continuity equation (3.1 into B.4 suggests that
(x.(x) dt
Pio x (xx) exp(-3E(x))
Ex, [exp(- - (xx))) xo
Jx div [pro xo uo (x|x expl (x)(3
Ex [epp( 38 (xo))
div
[x Peo (x) u (o (x)) (xp(-xX(xxx))))))
Exp [exp( - 3E (E
(x( (x)(=E((x))
por (xxxx)))uanxxx/xx exp(-38(xo)) exp(-Ex(1) dxo
div
(B.5)
Fxx (exp(-epppp (38(())))
Xo
div
(1) (x)],
where the third and the fourth equation is due the linearity of the divergence operator div [-]. The fifth equality is due to the definition of ge (x) and (x). Thus we conclude our proof by showing (x) defined by (B.5) can generate the guided distribution sequence (x).
B.2 PROOF OF THEOREM
Proof of Theorem 4.3
We start by fixing time t and writing the gradient of []v â (x) with
respect to as
Vollv (x)
Vo | (x) 2 ((Vv { 0,(xx) (x)
= Ve
| (x) (x),
(B.6)
where we drop the third term in the last equation due to the fact that Vo (x) 0. Consider the gradient of loss (C),M(O), we have
VoLCEFM(()
//
exp(-3E(xx))
Vo/
(o(X))Pto(x(()))
|// uxo(x((x)) || dxdxxo
(xo [exp(- (xo))]
x.Xo
= //
exp( (xx)))
Pe (x)))xx E (xo) (E
Vo || (x) dxdxx
(xp((o))
x.X0
2 2/ pr(x)poc(xolxx Epp((e))
exp( - 3E (xx)
- 2
VVov (xx), u (x(xx)) dxdxo.
(B.7)
15

---

Published as a conference paper at ICLR 2025
where the second equation leverages the Bayesian rule (x (o(x))Peo(x(xoo) and the fact that Ve (x = O. Noticing the fact that
p(xx xo)
(xxxx)(((((x)) (xx(xx)exp((EE(xx Px (x) (xp( (3) = Bpo(xo) lexp(- (x
exp(- (xp(-&x(x)) = (x.e(xo)xx)) (exp
Poe (xo (x)
(B.8)
(B.9)
By reorganizing the integral, B.
yields
VOLCEFM(()
/ (x( E (x) [xp((-E(())) Ve
(x))
(x) exp( (x))
(ov](xx), / (xo(xx)x) exp(-3E(xx))
exp(-Ex(x)) uso (x/xx) ) dx
Bp)( (exp( (-
(x)))
(x(x)
(Vollv (x) 2 2(Vov ūx(x))) dx
E (x)) eexp( (x,((()()))]]
exp( Exp( (x))
(x(
- (x,(x)
Ep (x) (exp( (x(())
= (OL
(B.10)
where the second equation is the combination of B.9 and the definition of ūx(x) declared in
4.1
in Theorem 4.1 Thus we complete the proof.
B.3 PROOF OF COROLLARY
Proof of Corollary 4.7
We start by denoting fs = /i//u and of = (ji:or - pode)os/ux According
to Lemma 3.2 we have
i fix + of Vx log (x),, [oo(x|xo) fix o6 Vx (x
(B.11)
then for the RHS in 4.4)
Por (Xo (xxx exp(-38(xo)) dxo
(xp(-Er(x))
Xo
exp(=3E(xo))
poc (fex + gt Vx log Peo (x(x0))
dxo
exp(-Ex
xo
exp(=3E(xx)) dxo
fix
Poo
Epm. (expp-xXp(eepeeeeeeeeee
Xo
exp( (xp(-)E(xx)) dxo
+ 92
log
exp(- (Ex(()) (xp(-BE(xo)) dxo. exp(-E((x))
= fix + 9h Px x V log Pxx(xx))
(B.12)
where the first equation due to equation (.III). The second equation is due to the linearity of the integral and the third equation utilizes the fact that [x Por x)) eep(-=E((x (xx))
Epoc (xalx)))) (exp(-3E(xu)] exp(-E, Plugging (B.12) back into Theorem 4.1 and utiliz- ing BB.1 yields
Poc Poc x X log (o exp(- (xx))) dxo, exp(-E,(x))
fix + gi log ū(x) = fix + 9f
which completes by canceling the fix term both side of the equation.
16

---

Published as a conference paper at ICLR 2025
B.4 PROOF OF COROLLARY 4.8
Proof. The proof is similar with the proof of Theorem 4.3 by replacing v? with sf. replacing (x(xo) with Vx log Pro (xxx)). Following the same rule in (B77 and (B.10) yields
VeLCED(O)
/f
exp(-3E(xo))
Vo
Po (xxx))poo(x(xx))
[s](x)) Vx Poo dxdxo
E (No) (exp(-(E(xx))
X.X0
/f
exp( 38 (xo))
pe (x())oo(xxl>> x)
V | (x)]11 dxdxo
(x)) eexp( - (xx)))
x,80
exp(- (38 (x)())
2
Pe(x)poo X x)
(Vos (x). Vx log p(o(x(xo) dxdxo
E8 (xe) (exp(- - 38 (xxx))
X,X0
exp(-Ex(x))
Pt (x)
Vo ||.?(( (x) dx
(x) (xp((=((x)))
Pe (x) exp (x()
Vos{(x), Pod (xo(xolx) exp(-3E(xo))
2
exp(-Ex(x)) Vx log log dx
x (x)) exp [x))
exp( exp(- (x()))
pe (x)
2(Vos (x), Vx log ge (x))) dx
(xx) [ex((-E(x)))
exp(-Ex(x))
(x(
Volls Vx (x)
(x) (exp( (ex((=E((x))
(O)
B.5 PROOF OF LEMMA 4.10
We present the proof for the classifier guidance (CG) and classifier-free guidance (CFG) presented in (.6)) in this subsection. The score function for Contrastive Energy Prediction (CEP) and Energy- weighted Diffusion (ED) presented in (4.7) has been extensively discussed in Corollary 4.7
Proof of (4.6) in Lemma 4.10 In order to sample from the guided distribution q, CG and CFG esti- mate the score function xxq Vx)(()) (x) as
Vx (x) Vx lxg + vV (cgp)(c(x) -(1 logxx + VV log Pr (x
(B.13)
where the second equation is given by the classifier free guidance. Applying the Bayesian rule, Pe(c(x) can be decomposed by
p(x, Pe (c(c)) = p(x (x)p(c|x))xxx Pppxxxxp(((x))))
(B.14)
Plugging BB.4) into 1.).), the score function Vx log (x)) is written by
Vx log (x)
= Vx (x) + PV log (p (c)()))
= Vx (x) + Vx log E [pxx)
where the last equation absorbs the into logarithmic operator, which concludes the proof.
17

---

Published as a conference paper at ICLR 2025
B.6
PROOF OF LEMMA 4.11
Proof. Plugging {(xa)
log p(c(xo) and B into the definition of LCED suggests that
p(c|xo)
(0) Rx,xxx0
Boo(Mo) [p(c|xu) 19.
Vx log Po pxo
A(1)p(c Xo popppeeoooppxxxt (o xoo
[| Vx log Pro (ro (x|
dtdxxdx
p(c)
1,
Vx log ( (t)p((xoc))/x0(/xxx [x(()))
1,x,80
x
|| (x)
Vx px
(B.15)
where the second equation expands the expectation and write Ex [p(c|xo]] p(c) due to Bayesian rule. The third equation is due to (x0)ppx)))p(((c)) -p(xo xxooec). The fourth equation rewrites the integral back as expectation. The optimal solution for [B. is s (x) = x log pr (x c).
given the analysis for the standard matehmat process (Song et al. 2020).
C COMPARISON BETWEEN CLASSIFIER GUIDANCE AND ENERGY GUIDANCE
We visualize the differences between the classifier- -easiiier-guidance and energy guidance in this section.
Experiment setup. We follow the setting in Lu et al. 2023 to sample the data distribution p(x)
and energy function Exx). We shift the energy function by Exx) <- E(x) minx {xx) so that
all energy {(x) 2 0. We note that this translation does not change the distribution of ((x) x
p(x) ex)(=)((x)) According
{xx)
log p(clx) discussed in Section 4.3 generate the
binary class such that p(c- 1|x) = exp(-E(x)). The ground truth distribution of q(x) x
p(x) exp(- ex)(=8E(x)
p(x)p²(c (c
1|x) with different B, the posterior distribution p(x|c 0),
c(x|c = 1) is presented in Figure 2
From Figure 2 we can find that when 3 1, p(x)p(c 1|x) p(x/c 1). However, when
2 1. (()) (xp(=3E(x)) = p(x))p((c 1|x) will be more concentrated on the region with higher likelihood p(c 1|x). We train the score function by optimizing LCED of Algorithm [ and the result is presented in Figure Bal In parallel to the implementation of our methods, we train the score
function s (x, 0) and 8 (x, 1), and implement the classifier-free guidance by
Vo log (x) (1 3)s (x. 0) + Bs (x.c= c 1)
(C.1)
presented in Figure 3b
18

---

Published as a conference paper at ICLR 2025
8
B-1
3-3
B-
Negative
Positive
is Negative
25 Positive
B=1
8=3
p= 10
0
2
2
2
-2.5 0.0 2.5 Negative
-25 0.0 Positive
2.5
8
8-1
8-3
8- 10
2
0
-55 Negative
-25 0.0
3=10
-2.5 2.5 Negative
-25 Positive
8=0
B=1
8=3 3
== 10
0
0
2 6 2
2
2
-25 0.0 2 Negative
-25 0. 2.5 Positive
4=
8=1
$=3
$= 10
-2-
-3
-2s - -(x)) 1/x) with
Figure 2: The ground truth distribution of 7(x) x p(x) different 8, the posterior distribution p(x/c 0) (Negative) and p(x/c 1) (Positive) over 6 data distributions.
19

---

Published as a conference paper at ICLR 2025
P.
P.1
(a) Emergy-weigghee diffuusion
(b) Classifier-Fre free guidance
Figure 3: The data sampled from the score function trained by energy weerg-weeghhtd diffusion and the score function composed by classifier-free guidance with different guidance scale 3. We sample 2000 data points for each state.
20

---

Published as a conference paper at ICLR 2025
D DETAILED ALGORITHMS
We present the detail of the algorithms that are skipped in Section 4 and Section 5 energy-weighted flow matching model in Algorithms
We present the
Algorithm 3 Training
Flow Matching Model
Input: Velocity vector field v.). guidance scale 3, batch size B, time weight X(1)
Input: Conditional vector field (x conditional distribution pro generated by uæo-
1: for batch {xi), ((x)) ) do
2: for index E [ do
3: Sample & ~ (0, 1), sample x; ~ Pr,0(-1x())
4: Calculate guidance 9: voftmax(-8E(xx) = exp(-=3((x)) E3 exp((3E(xx)
5: end for
6: Calculate and take gradient step using CCEFM(() =, Atti)aillv (xx))
7:
end for
Q-weighted Iterative Policy Optimization (QIPO). We use the in-support softmax Q-learning
et al. 2023) to implement the QIPO. We would like to emphasize that any Q function given by the
offline RL algorithm is compatible with our framework. The flow matching version is presented in
Algorithm and the diffusion version is presented in Algorithm We use Q to represent the Q function not contributing gradient to the loss, which is commonly used in TD-learning.
Algorithm Q-weighted iterative policy optimization for offline RL (Flow Matching)
Input: Vector field function vg (-), guidance scale 3, batch size B, weight on time X(()
Input: Conditional vector field Uyo [x xx)) conditional distribution Pro generated by U2o-
Input: Offline RL dataset
{ (xx, a, x', rr), number of epochs K1, K2 and K3, Q function Q
Input: Support action set size M, support action set renew frequency K Krenew
for diffusion model warm-up epoch & [K]] do
2: for batch (x,, a, x(.r:)
D do
3: Sample 4 ~ (0, 1), ai, ~ Ot (a;) for E [B]
4: Calculate and optimize Cor, warmup ( = Em X(ta)ll?(( - ue,o(ai,((a))1
5: end for
6: end for
7: for Q-learning warm-up epoch E [2]]
8: for batch {xx, a, r.)r do
9: Sample support action set al, using vector field vi for j E [M), denote a//0 at end for
11: Calculate (Q(()= - -; |0( a.) r; -7 EMe exp(SQ) (1, (x(..()))
MM (xp(
12: Update using the gradient of (1(b))
13: end for
14: for policy improvement step l: [Kg] do
15: for batch {(x; a,, x(,r;)] C D do
if & mod Krnnew = then
17:
Sample support action set a, using vector field v-Ix)) for all e [B] E [M]
18:
Denote a a, sample tij U(0, and aij, P(u,o(-laay) for all j E [M]
Calculate guidance gij 53ft (xi, anj)) Σx (X.P(((((()))))))
exp(AQ " (x.,,;)))
20: end if
21: Calculate loss LQF(O)= = EBn أ A(ti))9ts)) (ag,, xx) un | 11
22:
Update A using the gradient of (ar(O)
end for
24:
end for
21

---

Published as a conference paper at ICLR 2025
Algorithm 5 --weighted iterative policy optimization for offline RL (Diffusion)
Input: Score function 5 (-), guidance scale (3, batch size B, weight on time X()), schedule (we, of
Input: Offline RI. dataset D = {(x, a. x)),,, number of epochs K1, K2 and K3, Q function Q
Input: Support action set size M, support action set renew frequency Krenew
1: for diffusion model warm-up epoch & [K1] do
2: for batch fx;, a, r.) C "D do
3: Sample & (0, 1), calculate Ot, Ote sample C; -V Ia) and acs, Me, oz, (.), 4: Calculate loss LOD, warmup (0) = - (a +
e.112
5: Update O using the gradient of LQD. warmup (0)
6: end for
7: end for
8: for Q-learning warm-up epoch k E [K2] do
9: for batch {(x,, ai; r:)] D do
10: Sample support action set al, using score function sf for j E [M), denote a/o aq
11:
end for
Q (x;, - rs -y In exp (X,
Calculate - -
M exp(BX) (x(,a(,))
Update &p using the gradient of LQ (1b) end for
14:
for policy improvement step E [Kg] do
16:
for batch {(x,, a. x,,r;) D do
17:
k mod Krenew = then
18:
Sample support action set a,y using score function (-1x,) for all i E [B], [M]
19:
Denote ag = a, sample túj ~ U(0, 1). calculate Hly Ot, Sample Eij (0, Ia) for E [B], [M]. let asj.e,, Me. aij oee Eij exp(AQ (x..a))) Calculate guidance gij softmax (3Q (x,, any) ΣE (xp(a(((xxxa))))
20:
21:
22:
end if
23:
Calculate loss (0) م A(tig)gis s2 (aij,,,) + EEi e
24:
Update O using the gradient of (Q)(0)
25:
end for
26:
end for
E ADDITIONAL EXPERIMENT RESULTS
We first present the comparison between QIPO and other algorithm in Table
E.1 EXPERIMENT CONFIGURATION
QIPO-Diff configurations We adopt the same network structure and training configuration Lu et al. 2023) during the score network warmup K; and Q-function training K₂ for a fair comparison and ensure consistency across experiments. We finetune the score network sf obtained after warm-up up (Line 5 Algorithm 5) with learning rate 10 to perform the Q-weighted iterative policy optimization. The schedule of the diffusion the same with Lu et al. (2023) and we use the DPM- Solver (Lu et al.] 2022a) with di fuu innsste 15 to accelerate the generation step, which also the same as (Lu ct al. 2023).
QIPO-OT configurations We use the optimal transport VF (Lipman et al. 2022 defined by (x Xo (1 miin x
0 minn) m 0.0054,
note we switched the notation between t 0 and = compared with Lipman al. 2022
According to Lemma 3.2 we defined the network of vp as
(x) (1 1)- log ix + (1 - 1)- log - (1 ()=-)--)a vx)),
where sf is the same with QIPO-Diff implementation. Similar with QIPO-Diff, employ the DPM-Solver (Lu et al. 2022a) for the second-order Runge-Kutta-FFne Kutta-Fehlberg method for solving the reverse-time ODE (x).
22

---

Published as a conference paper at ICLR 2025
Table 4: Evaluation numbers of D4RL. benchmarks (normalized as suggested by Fu et al. (020 We report mean + standard deviation algorithm performance across 8 random seeds. The result for ablation study with 3 10 are the averaged performance over random seeds. Together with the benchmark algorithms presented in Table 2 the highest performance is boldfaced highlighted. The performance within 5% of the maximum absolute value in every individual task highlighted, To ensure fairness, ablation results with & = 10 are excluded from the best performance comparison. However they are still highlighted if they meet the 5% performance criterion.
Dataset
Environment
BEAR
TD3+BC
IQ1
Diffuser
DiffusionQQI
QIPO-DII (ours)
QIPO-Dif tours,
Medium Expert
HaliCheetah
53.4
90.7
86.7
79.8
96.5
01:1::0.15
91.96
Medium- Expert
Hupper
96.3
98.0
91.5
107:2
111.1
112.12
0.43
103.56
Eedlimm-Eppert
Walker2d
40.1
110.1
108.6
108.4
110.1
110.14
112.00
Medium
HaliCheetah
41.7
48.3
47.4
44.2
51.1
48.19
0.20
50.11
Medium
Happer
52.1
50.3
06.3
68.5
00.5
89.33
9.96
92.01
Medium
Walker2d
59.1
83.7
78.3
79.7
86.36
Medium Replay
HalfCheetah
38.6
44.6
44.2
42.2
47.8
45.27 +: 0.42
50.10
Medium Replay
33.7
98.67
Medium-Reppe Replay
Walker2d
19.2
81.8
73.9
61.2
95.5
90.08
= 1.53
89.45
Average (Lacomotion)
$1.9
75.3
76.9
75.3
88.0
86.20
86.80
AmMaaze-umaa
73.0
Diverse
AmtMaze-umaee
61.0
71.4
02.2
06.2
73.88 + 6.42
83.00
Play
AntMaze-man medium
0.0
10.6
71.2
78.6
82:75 L: 3.24
92.67
Diverse
LutMaze-mediung
8.0
70.0
Play
AntMaze-large
0.0
0.2
39.6
46.1
73.25
10,90
62.67
Diverse
AntMaze large
0.0
0.0
47.5
56.6
40.5 4. 20.40
39.00
Average (AntMaze)
23.7
27.3
63.0
09.6
77.3
80.78
E.2 INDIVIDUAL TIMING COMPARISON BETWEEN QGPO AND QIPO
We present the timing comparison between QGPO (Lu et al. 2023) and QIPO this section, as shown in Figure 4 It is evident that QIPO-OT consistently outperforms QIPO-Diff in terms of speed, with QGPO being the slowest due to the additional back-pre preck-poopggaion. The advantage of ortimal-transport flow matching over diffusion models aligns with the claims made by Lipman al. 2002) and Zheng et al. (2023), which suggest that OT can accelerate the generation process.
EELL
(a) Locomotion Tasks.
(b) AntMaze Tasks
Figure 4: Comparison of the time complexity of the action selecting function for CEP, QIPO-Diff and QIPO-OT across various tasks. The results averaged over 3 individual runs.
Iterative Policy Optimization We perform the iterative policy optimization (Line 7 Algorithm 2 with K3 100 and evaluate the performance of the agent every 5 epochs. We renew the support set with period Kroww 10. The rewards reported in Table Table the best performance the reported performance over K3 = 100 episodes. Since the QIPO procedure progressively improve the guidance scale beta, as discussed in (5.4), this similar to automatically tune the guidance scale in Lu et al. (2023) as paramm-rrrmtter. We use the soft update (Lillicrap) 2015)with À = 0.005 to stabilize the update of the score network,
E.3 ABLATION STUDIES FOR QIPO
Changing the guidance scale 3. We conduct an ablation study by changing the guidance scale from B 3 B 10 for QIPO-Diff, with the results included in the last column of Table
23

---

Published as a conference paper at ICLR 2025
The algorithm demonstrates robustness to changes in the guidance scale, as the iterative policy optimization process naturally adapts and improves the guidance. Additionally, we observed that the average performance over 3 random seeds sometimes surpasses the original setting with / 3, while in other cases it is slightly below the original performance.
Changing the support action set M. We ablate the size of the support action set, increasing it from M 16 to .M = 32 and M 64, as shown in Figure The results suggest that a larger support set M may lead to slightly improved performance. However, since sampling the support action set requires running the reverse process of the diffusion model or the score matching model, we choose to keep M 16 to balance efficiency and effectiveness.
Locomotion Tasks.
(b) AntMaze Tasks
Figure 5: Comparison of normalized cumulative rewards for different support set sizes (M = 16, M 32, and M 64) across various tasks. The results are averaged over 3 individual runs.
Changing the renew support action period Krenow- We ablate the period for changing the support action set within Krenew E (5. 10, 20, 50]. Intuitively, larger Krenew may result in a better-learned score network but slower updates to the guidance scale. shown in Figure [6] we observe that smaller K values generally lead to improved performance. We hypothesize that this is because the score network for the D4RL tasks is relatively easy for the neural network to learn.
(a) Locomotion Tasks.
(b) AntMaze Tasks
Figure 6: Comparison of normalized cumulative rewards for different support action set renewal
period (Krenw (5, 10, 20, 50}) across various tasks. The results are averaged over 3 individual runs.
F ADDITIONAL EXPERIMENTS
F.I CHANGING THE ENERGY LANDSCAPE
We conduct ablation study to analyze the performance of the energy-weighted diffusion model under varying energy landscapes. Specifically, we use the 3-Gaussian example shown in Figure 3bb and modify the energy landscape
illustrated in Figure 7a The sampling results obtained from running Algorithm | are depicted in Figure [7b) These results suggest that the proposed energy- weighted diffusion model is robust to changes in the energy landscape.
24

---

Published as a conference paper at ICLR 2025
Ground truth of the distribution
Sampled result the new distribution
Figure 7: Ground truth and sampling results on the 8 8-Gaussian dataset with a modified energy landscape.
F.2 TAKE 3 AS A SCORE NETWORK INPUT
Although Algorithms 3 and require the guidance scale as an input for simplicity of presenta- tion, we propose an additional algorithm that unifies different choices of 3 across the range from 3max. As described in Algorithm 6 this algorithm treats B as network input and, for each minibatch, randomly selects 33 value from the uniform distribution (0,,3max)) 3max). The encoding of follows the same method used for encoding time in the network, represented as:
x x + emb(t) @
where emb(-) denotes the time embedding function and @ indicates concatenation. The reverse (gen- eration) process remains unchanged, except that the score function is modified to s (-: ,8), using specific guidance scale 3. We conducted experiments on the bandit case illustrated in Figure g The results indicate that the performance of the proposed algorithm is comparable to that of Algorithms [ suggesting that the score model can effectively learn multiple 3 values simultancously.
Algorithm Training Energy- Weerghed Diffusion Model with 33 as Input
Input: Score function s0 (:3-,) -), maximum guidance scale 3max
Input: Batch size B, time weight A(t), schedule (Hu, ot)
1: for batch {x) [xx) }; do
2: Sample B (0, 3max)
3: for index E [B] do
4: Caleulate guidance 9; softmax((38(xx)) cxp((38(x))) E, exp(-3E(x())
5: Sample t; ~ (0, 1), calculate Pt., T.. sample E, ~ (o( and X1; = -1,xb + oz,
6: end for
7: Calculate and take a gradient step using CEED Σ, A(tr)g/lls (x B) + /
8: end for
P-12
Figure 8: Data sampled from the score function trained using energy weighted diffusion with B as an input. A total of 2,000 data points were sampled for each state.
25

---

Published as a conference paper at ICLR 2025
F.3 IMAGE SYNTHESIS
To evaluate the performance of the proposed methods in guiding the generative pro- cess for more complex tasks, we conduct image synthesis experiments using the -anidinn-bbttefliee-eeee
datase Similar to Lu al. (2023), the energy function is defined the average hue distance between the image and the target hues red, blue, and green, expressed as {(x) [[[(x) Margatl 1 E [0, 1]. We use the meegy--eiight to train for the target distribution described as q(x) x p(x) cxp(-5E(x)) with 3 = 50 in the experiments. The results, presented in Figure demonstrate that methods effectively guide the image generation process.
diffusion
F.4 MOLECULAR STRUCTURE GENERATION
We further evaluate the performance of the proposed weighted diffusion algorithms on molecu- lar structure generation using B(3) diffusion networks (Hoogeboom et al., 2022). The normalized properties-namely, isotropic polarizability (a), highest occupied molecular orbital energy (éhOMо), lowest unoccupied molecular orbital energy (CLUMO), and the energy gap A. between (номо and (uuMo-ar are used to guide the generation process. For instance, when the energy function is set as E(x) a. a, the diffusion model is guided to generate molecular structures with higher isotropic po- larizability. The energy-guided is target for generating the distribution g(x) p(x) exp (x)) in this case. Similar to (Hoogeboom et al. 2022). we train a separate EGNN (Satorras et al., 2021) I) to predict the normalized properties of the generated molecules. The property normalization method follows (Hoogeboom et al., 2022), and we report the performance for 3 E 0, 1, 10 Table 5 The results indicate that the energy-weighted diffusion model successfully generate molecular structures with enhanced desired properties.
Property
8 = 0
B
8 = 3
3 10
error
0.1997
0.6196
2.2965
2.2956
0.1034
A
0.1910
0.4687
1.7025
1.7129
0.0667
CHOMO
0.1142
1.6074
2.7396
2.6798
0.0413
GAMMO
0.2150
0.4089
1.7605
1.7669
0.0362
Table 5: Averaged predicted normalized properties for molecular generation. The values are aver- aged over 1,000 generated molecular structures. The column error represents the prediction error of the classifier. All properties are normalized by their mean and standard deviation from the training dataset.
F.5 TRAINING CURVES OF OT AND DIFFUSION
In this subsection plot the training curves for learning the behavioral policy ((a)s) follow- ing the flow matching (OT) and diffusion loss to further demonstrate the comparison between OT and diffusion when applied in offline RL tasks. In particular, Figure 10 compares the diffu- sion and OT training loss on the ant-maze tasks. Since is impossible to find a va(a;)) (or st (a;; x) matching all ugg (a, (aa) (or Va, log Peo (e.(a,)aa (a.)aa)) for all ag. neither the flow matching loss (a): u u (ag 2l nor the score matching (diffusion) loss ||s (a;: x) V a. log Peo (a,) (aaaaaaaa)))))a)a))())) _ converge to zero, and the remaining training loss is different between OT and diffusion. Therefore we normalize the loss and presented in Figure[T]] 1T by shifting the loss at the last epoch (600-th epoch) as and the loss at the first epoch as 1.
The experiment results as presented in Figure 10 suggest the following findings. First, on the Antmaze tasks, flow matching using OT loss converges similar but slightly faster to the stable point. Second, despite of this difference on the convergence rate, both OT and diffusion can correctly learn the behavioral policy.
subset
26

---

Published as a conference paper at ICLR 2025
(a) Original (without guidance)
(b) Red guidance
(c) Green guidance
(d) Blue guidance
Figure 9: Generated butterflies, the samples are generated from the same 16 random seeds with different guidance or no guidance using reverse process with 1000
27

---

Published as a conference paper at ICLR 2025
Low
50-
(a) Antmaze- Umaze
(b) Antmaze- Antmaze-Medium Medium-Play
(c) Antmaze Large- Play
Los ture
Los CUME
7.-
(d) Antmazc-Umaan Umaze- -ivverse
(c) Antmazo-Medium Diverse
(f) Antmaze- Large- Diverse
Figure 10:
Unnormalized loss w.r.t number of epochs for the first 600 epochs in learning the
behavioral policy μ. The blue curve indicates the OT schedule for flow matching and the red curve indicates the diffusion schedule.
6--
2.5
420
260
(a) Antmaze- Umazze
(b) Antmaze- Antmaze-Medium:
(c) Haimazz-Largee Play
Lse
53
(d) Antmazc- Antmaze-Umaze-
(c) Antmazz--Mdduu--Dieee
(f) Antmaze- Haammaze-aarg--De Large-Diverss Diverse
Figure 11: Normalized loss w.r.t. number of epochs for the first 600 epochs in learning the behavioral policy μ.. The blue curve indicates the OT schedule for flow matching and the red curve indicates the diffusion schedule. The normalization is conducted by rescaling the loss at last iteration to O and the loss at the first iteration 1.
28


Reviews:
[{"review_id": "yG3JjnfVQw", "reviewer": "Reviewer_MVYW", "summary": "This paper introduces the energy-based flow matching to offline reinforcement learning without the need for auxiliary models, and then propose a novel method called Q-weighted Iterative Policy Optimization based on this framework.", "questions": "1.\tWhat is the motivation of applying flow matching to offline RL?\n2.\tCompared with the previous works for diffusion policy with weighted loss, what is the real contribution of this paper?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "1.\tThe manuscript is generally well-written and well-organized.\n2.\tThe paper presents comprehensive theoretical proof.", "weaknesses": "1.\tIn the reviewer’s opinion, the authors missed some existing works including EDP, and QVPO [1, 2], which also train the diffusion policy with the weighted loss. In that case, the paper does not propose anything particularly novel.\n2.\tGiven weakness 1, the authors overclaim their contribution: “Our algorithm is the first energy-guided diffusion model that operates independently of auxiliary models and the first exact energy-guided flow matching model in the literature”.\n3.\tThe reviewer thinks the proposed method should not be viewed as an energy-guided diffusion model but as an extension of RWR (reward-weighted regression) [3] via flow matching.\n4.\tResults in Table 2 does not show a distinct superiority of the proposed QIPO compared with previous methods.\n5.\tThe motivation for applying flow matching to offline RL is not clear. It seems the multimodality of diffusion policy cannot be obviously improved with flow matching loss compared with normal diffusion loss.\n\n[1] Kang B, Ma X, Du C, et al. Efficient diffusion policies for offline reinforcement learning[J]. Advances in Neural Information Processing Systems, 2024, 36.\n\n[2] Ding S, Hu K, Zhang Z, et al. Diffusion-based Reinforcement Learning via Q-weighted Variational Policy Optimization[J]. arXiv preprint arXiv:2405.16173, 2024.\n\n[3] Peters J, Schaal S. Reinforcement learning by reward-weighted regression for operational space control[C]//Proceedings of the 24th international conference on Machine learning. 2007: 745-750.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BUKbg79mZY", "reviewer": "Reviewer_6Qqj", "summary": "This paper proposes a training scheme for energy-guided continuous generative models, i.e. models that learn $q(x)\\propto p(x)\\exp(-E(x))$. To do so it leverages something akin to the $x_0$-conditional trick of Lipman to rewrite a loss that would depend on an intermediate energy $E(x_t)$ as simply depending on $E(x_0)$. Functionally, this is essentially a softmax-of-energy weighted loss for both flow matching and score matching models. This provably converges in the limit to an unbiased model that, unlike prior work, does not require a second model to learn a guidance.\n\nThe authors demonstrate their method on simple demonstrative tasks and a set of benchmark offline-RL continuous control tasks, where the method performs on-par with state of the art methods.", "questions": "One of the assumptions made seems to be that there is a prior (data distribution) $p_0(x_0)$ we can sample from, but results will differ greatly if $p_0$ is an actual sampler (e.g. a pretrained model) or a dataset (in which case you'll get some peaky behaviors). It would be good for the authors to expand on this empirically.\n\nRelatedly, this method depends on a softmax of a potentially peaky value to learn with an off-policy objective. If the prior is flat enough and the energy peaky enough, this may very well mean that most of the data is effectively thrown away and most compute wasted because its weight is 0. There are other methods (although the ones I think of would count as concurrent work, so I'm not holding it against the authors) that learn similar amortized inference policies $q \\propto p(x) \\exp(-\\beta E[x])$ that are able to sample on-$q$-policy (some are also able to do both off-policy and on-policy), therefore potentially wasting less data. Again not holding it against the authors to compare to those methods, but I'd love to see a deeper empirical analysis of when the proposed method works with respect to different energy landscapes; this could inform someone picking a method for their problem (and it is likely that the proposed method would be a top choice for certain problems, but not others).\n\nFinally, I wonder what is the effect of estimating an expectation with a minibatch rather than a more precise method. For example, in the offline RL case, assuming all trajectories in the dataset come from the same policy $\\mu$, then it would be much more precise to estimate $\\mathbb{E}_\\mu[Q^\\psi]$ once over the dataset (especially if $\\psi$ rarely changes), but maybe this doesn't change much.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper proposes a novel solution to a very important problem, and the solution itself is fairly elegant (as most unbiased methods are in this reviewer's opinion). The paper was easy to read and used an appropriate level of detail to explain and illustrate the method. The empirical work also shows good performance on non-trivial problems.", "weaknesses": "The main weakness of the paper is it doesn't really go into the weeds of the method. Unbiased methods are great, but sometimes they come at a cost. For example, importance sampling methods are commonly thought of as high variance. It's not clear if this is the case here, or more generally what are the trade offs that would inform choosing this method over others. \n\nI know it's easy to write this, but there could be more empirical work done. Specifically, two very common class of problems for which these methods are used are missing, image generation and physics problems (e.g. 3d molecular conformer generation). RL is a challenging problem but the method is much more general than that, and so it feels like a missed opportunity that results on these other problems are not presented.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wqXn5wDgMf", "reviewer": "Reviewer_mBzj", "summary": "This paper presents a novel diffusion method combined with offline RL. The major contribution is that it combines flow matching loss with an energy function, specifically weighting the flow matching term with the regularized energy, which avoids calculating the gradient of the intermediate energy function. This idea is integrated into offline RL, leading to the development of Q-weighted iterative policy optimization (QIPO) with two variant velocity fields, OT and VP-SDE. Experiments on toy examples reveal the effectiveness of ED and QIPO on D4RL datasets compared to SOTA methods.", "questions": "Please answer my questions mentioned above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The authors provide substantial theoretical analysis.\n\n2. The experiments reveal that QIPO-Diff and QIPO-OT outperform current SOTA offline RL methods.", "weaknesses": "1.The authors should discuss more details on the benefits of QIPO. The current methods only discuss the advantage of ED and CED over CEP, which is mainly about the diffusion models. It seems that the advantage of QIPO over other offline RL methods (e.g., SRPO, SfBC) is missing.\n\n\n2. Only toy experiments on ED are given. If the authors want to claim the advantages of ED and CED, more experiments should be included, e.g., image synthesis tasks on ImageNet. I understand that verification of the diffusion model in different tasks might not be a major topic in this paper, but in its current form, QIPO appears to be a direct application of ED in the offline RL setting due to ED's advantages. Thus, empirical verification of ED is important.\n\n\n\n3. The theoretical analysis should add more details and fix typos, e.g., in Eq. (a.5), how the div operator is eliminated should be discussed in detail.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q1BeCpTs6H", "reviewer": "Reviewer_JYA4", "summary": "This manuscript discussed how to estimate the score function with an energy guidance without axillary model but can still maintain the exact guidance. The main idea is decomposing the supervision from the data population and energy guidance. The authors also extended the proposed methods to KL-regularized offline reinforcement learning.", "questions": "* Will Monte-Carlo estimation on the denominator lead to a high variance on the loss/gradient estimator? It would be generally not a good idea to do something like this.\n* Do CG/CFG and CEP/proposed methods only differ from the diffusion path and do we have an understanding on why this two diffusion path will lead to different empirical results?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The derivations are sound and the final algorithm is intuitive.", "weaknesses": "* The presentation is not so straightforward.\n* If I understand the proposal correctly, for each $\\beta$ the proposed method needs to at least fine tune the score function and do the sampling. I don’t think it will have significant benefits on we directly estimate the score function without leveraging the proposed methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)