PaperlayerView exact revision on arXiv

Source-linked AI summary

Energy-Weighted Flow Matching for Offline Reinforcement Learning

Shiyuan Zhang, Weitong Zhang, Quanquan Gu

arXiv:2503.04975v1cs.LG

TL;DR

Existing energy-guided diffusion methods require auxiliary networks, back-propagation, or intermediate guidance procedures. This paper introduces energy-weighted flow matching and diffusion models that learn guided dynamics directly, then applies them to offline RL through QIPO, which consistently outperforms baselines across various tasks.

  • Problem

    Existing methods for sampling energy-guided distributions require additional neural networks, back-propagation, or intermediate energy guidance, motivating a direct approach.

  • Method

    The paper proposes an energy-guided velocity field, energy-weighted flow-matching and diffusion objectives, and applies them to offline RL with QIPO.

  • Results

    QIPO-Diff and QIPO-OT consistently outperform baselines across various offline RL tasks, while QIPO-Diff samples faster than QGPO with the same score network.

  • Takeaways & Limitations

    Energy guidance can be incorporated directly into flow matching and diffusion models without auxiliary models or post-processing, achieving competitive D4RL performance.

  • Takeaways & Limitations

    The generative models may generate harmful content or reflect and amplify unwanted social biases in training data.

Abstract

from arXiv · show

This paper investigates energy guidance in generative modeling, where the target distribution is defined as $q(\mathbf x) \propto p(\mathbf x)\exp(-β\mathcal E(\mathbf x))$, with $p(\mathbf x)$ being the data distribution and $\mathcal E(\mathcal 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.

1 INTRODUCTION

The introduction frames energy guidance as a way to control generative outputs, while identifying unresolved challenges for diffusion and flow matching models. The paper proposes direct energy-guided velocity and weighted objectives, extending them to diffusion and offline RL.

  • Energy guidance targets controlled outputs, including class-specific images, property-driven molecules, and improved reinforcement-learning policies.
  • Existing diffusion guidance methods require rejection sampling, auxiliary energy networks, back-propagation, or force-field procedures.
  • The paper asks whether energy-guided diffusion can be obtained without auxiliary models and whether exact guidance can be injected into general flow matching.
  • The proposed energy-guided velocity field is theoretically guaranteed to generate q(x) ∝ p(x) exp(−βE(x)).
  • Energy-weighted flow matching trains the guided velocity using conditional vector fields and dataset energies, while its diffusion extension learns guidance without auxiliary models.
  • The methods are applied to offline RL through iterative policy refinement and reportedly achieve superior performance across various offline RL tasks.

2 RELATED WORK

Related work covers diffusion and flow matching models, guidance methods for controlled generation, and their applications to reinforcement learning. It distinguishes energy-based guidance from classifier-based approaches through its real-valued energy formulation and auxiliary-model requirements.

  • Diffusion and flow matching models have been applied to image, text-to-image, video, and reinforcement-learning generation tasks.
  • Guidance in Diffusion and Flow Matching Models: Energy-based guidance defines q(x) ∝ p(x) exp(−βE(x)) and supports real-valued energy functions, unlike classifier guidance.
  • Guidance in Diffusion and Flow Matching Models: Prior energy-guidance methods include rejection sampling and procedures that estimate intermediate energy or its gradient.
  • Table 1 compares guidance methods by exactness and whether they learn guidance without an auxiliary model.
  • Guidance in Diffusion and Flow Matching Models: Classifier-free guidance has also been introduced for flow matching in offline reinforcement learning.

3 PRELIMINARIES

The preliminaries establish continuous flows, conditional flow matching, diffusion paths, and energy-guided distributions. They explain how conditional vector fields can train marginal flows and how prior diffusion guidance composes scores using an intermediate energy.

  • Continuous normalizing flows describe probability-density paths between data and an initial distribution using a time-dependent velocity field.
  • Conditional Flow Matching for Generative Modeling: Flow matching learns a neural velocity field by minimizing differences from the ground-truth vector field.
  • Conditional Flow Matching for Generative Modeling: Because the ground-truth marginal field is infeasible to calculate, conditional flow matching instead matches conditional vector fields generated from data samples.
  • Conditional Flow Matching for Generative Modeling: The conditional path is typically modeled as a Gaussian distribution with time-dependent mean and variance.
  • Energy-Guided Diffusion Models: Energy-guided diffusion targets q0(x) ∝ p0(x) exp(−βE(x)) rather than the original data distribution p0.
  • Energy-Guided Diffusion Models: Prior guidance introduces an intermediate energy Et(x), yielding qt(x) ∝ pt(x) exp(−Et(x)) and a score decomposition involving ∇xEt(x).
  • Energy-Guided Diffusion Models: Contrastive energy prediction is used to learn the intermediate energy before composing the guided score with the standard diffusion score.

4 METHODOLOGY

The methodology constructs energy-guided flows and trains them with energy-weighted objectives, extending the approach from general flow matching to diffusion models without intermediate-energy estimation. The analysis also establishes exact guidance behavior relative to classifier-free guidance.

  • Energy-weighted flow matching: Energy-weighted flow matching trains a velocity field for q(x) ∝ p(x) exp(−βE(x)) using conditional flow fields and dataset energies.The conditional loss uses the conditional vector field and E(x0), while remaining equivalent to the energy-guided flow loss up to a constant factor.
  • Energy-guided flow construction: Theorem 4.1 constructs an energy-guided flow from a conditional flow, and it applies to optimal transport, Gaussian paths, rectified flow, and other conditional flows.The resulting flow generates qt(x) ∝ pt(x) exp(−Et(x)) and q0(x) ∝ p0(x) exp(−βE(x)).
  • Direct training objective: The conditional energy-weighted loss directly learns the guided flow without calculating the intermediate energy function or its gradient, given sufficient capacity and data.This avoids the auxiliary neural networks used by prior methods to estimate Et or ∇xEt.
  • Energy-weighted diffusion: Energy-weighted diffusion trains the guided score function through an equivalent conditional loss and removes additional sampling and back-propagation for intermediate-energy estimation.The energy-weighted and conditional diffusion losses differ only by a constant factor.
  • Comparison with classifier guidance: For classifier-defined energy E(x) = −log p(c|x), energy-weighted diffusion exactly targets p0(x)p(c|x)^β, whereas classifier-free guidance is not guaranteed to do so when β ≠ 1.The two methods are similar at β = 1, but only energy-weighted diffusion remains aligned with the ground-truth distribution for β > 1.

5 Q-WEIGHTED ITERATIVE POLICY OPTIMIZATION FOR OFFLINE RL

QIPO applies Q-weighted diffusion and flow matching to offline RL by iteratively incorporating Q guidance while preserving behavioral-policy regularization. On D4RL tasks, both variants consistently outperform the evaluated baselines, with QIPO also offering speed and guidance advantages over related methods.

  • Proposed method: QIPO learns a score function for the optimal policy using either Q-weighted diffusion loss or Q-weighted flow matching loss.The behavioral policy can be learned with standard diffusion or flow matching, while any Q function from offline RL algorithms can supply the weighting.
  • Algorithm: Each policy-improvement iteration reuses the current score function to sample support actions, applies Q-based guidance, and periodically renews the support set.The procedure warms up the behavior score model and Q function before iterative policy optimization.
  • Iterative policy optimization: π_l+1(a|x) ∝ π_l(a|x) exp(βQψ(a,x)), yielding π_l+1(a|x) ∝ μ(a|x) exp((l+1)βQψ(a,x)).The iterative update implicitly increases the guidance scale β while retaining behavioral-policy weighting.
  • Comparison: QIPO strictly follows behavioral-policy regularization, unlike guidance-scale approaches that do not guarantee the same formulation when the scale exceeds one.The resulting policy has the form π(a|x) ∝ μ(a|x) exp(sβQψ(a|x)).
  • Experiment results: QIPO-Diff and QIPO-OT consistently outperform the baselines across various D4RL tasks.The comparison includes Diffusion-QL, QGPO, IDQL, SRPO, and Guided Flows; Table 2 reports means and standard deviations across 8 random seeds.
  • Experiment results: QIPO-OT has higher performance across many tasks than Guided Flows, while QIPO-Diff directly learns energy guidance without estimating an intermediate energy function.Compared with QGPO, QIPO-Diff avoids back-propagation through the intermediate energy function and is reported to sample faster with the same score network.

6 CONCLUSION AND FUTURE WORK

The paper introduces direct energy guidance for flow matching and diffusion models and applies both to offline RL through QIPO. It reports competitive D4RL performance while identifying online RL as future work.

  • Conclusion: Energy-weighted Flow Matching and Energy-weighted Diffusion incorporate energy guidance directly without auxiliary models or post-processing.The paper presents these methods as its direct energy-guidance approach for both generative-model families.
  • Conclusion: QIPO achieves competitive empirical performance on the D4RL benchmark in offline reinforcement learning.The current method focuses on offline RL without environment interaction.
  • Future work: Extending QIPO to online RL is left for future work so Q-function guidance can be updated through online interactions.The stated future direction concerns updating guidance using environment interaction.

ETHICS STATEMENT

The ethics statement frames the work as generative-model methodology that follows ethical guidelines while acknowledging risks shared with deep generative modeling. It identifies harmful-content generation and unwanted bias as potential concerns but reports no unique impacts.

  • Potential risks: Energy-weighted diffusion and flow matching may generate harmful content such as deepfakes and reflect or amplify unwanted social biases.The statement presents these as potential risks associated with deep generative modeling techniques.
  • Scope: The paper states that these broader societal impacts are not directly relevant to its specific contributions.This is the authors’ stated scope judgment in the ethics statement.
  • Ethics assessment: The authors report adherence to ethical guidelines and standards and identify no unique ethical concerns or negative social impacts.The conclusion is specific to the paper’s methodology-focused scope.

B.1 PROOF OF THEOREM 4.1

The appendix proof establishes that the proposed energy-guided velocity field generates the guided distribution sequence and supports the associated weighted-training result. The derivation uses continuity equations, Bayesian identities, and normalization properties.

  • Guided distribution: The guided distribution is defined as q_t(x) ∝ p_t(x) exp(−E_t(x)), with normalization obtained from the energy-weighted data distribution.The appendix states a closed-form expression for q_t and notes that the denominator is constant with respect to x and t.
  • Proof derivation: Bayesian identities connect p_t(x), the conditional transition, and the joint distribution used in the derivation.These identities allow the proof to rewrite expectations under the guided distribution.
  • Proof of Theorem 4.1: The proof begins from the conditional flow field that generates the conditional distribution sequence and applies the continuity equation.This sets up the velocity-field derivation for the guided distribution.
  • Proof conclusion: The resulting velocity field is shown to generate the guided distribution sequence q_t(x).The proof concludes this by substituting the proposed field into the continuity-equation dynamics.
  • Proof of Theorem 4.3: Theorem 4.3’s proof reorganizes the weighted objective using the same conditional-flow and Bayesian relationships before concluding the training result.The appendix invokes gradient manipulation and the definition of the proposed guided velocity field.
  • Diffusion guidance: The appendix separately derives diffusion guidance comparisons for classifier guidance, classifier-free guidance, contrastive energy prediction, and energy-weighted diffusion.The score-function derivation uses the classifier-based decomposition and standard score-matching analysis.

C COMPARISON BETWEEN CLASSIFIER GUIDANCE AND ENERGY GUIDANCE

The experiments compare energy guidance with classifier-based guidance and describe QIPO implementations for diffusion and flow matching in offline RL. Energy-weighted methods are evaluated through controlled distributions, sampling procedures, and Q-weighted policy optimization.

  • Energy-guided distributions: The experiment constructs q(x) ∝ p(x)exp(−βE(x)) and compares it with posterior distributions for positive and negative classes.The energy is shifted to remain nonnegative, and p(c = 1|x) = exp(−E(x)) is used to define the binary class.
  • Energy-guided distributions: For β = 1, p(x)p(c = 1|x) is proportional to p(x|c = 1), while β ≥ 1 concentrates q(x) on regions with higher p(c = 1|x).
  • QIPO implementation: QIPO uses in-support softmax Q-learning and supports both flow matching and diffusion implementations with an arbitrary compatible offline-RL Q function.The flow-matching and diffusion versions are presented as separate algorithms, while Qψ is treated as gradient-detached in the loss where required.
  • QIPO implementation: The QIPO procedures sample support actions, compute softmax Q-weighting with guidance scale β, and update the policy model through weighted training losses.The flow-matching version uses a conditional vector field, whereas the diffusion version uses a score function and diffusion schedule.
  • Evaluation: The paper compares QIPO with other algorithms in Table 4, using D4RL benchmark evaluations as the main performance comparison.

E.1 EXPERIMENT CONFIGURATIONS

The experiments configure QIPO-Diff and QIPO-OT with matched training choices and compare their action-selection timing and benchmark evaluation procedures. The reported timing comparison favors the optimal-transport flow-matching variant over diffusion and QGPO.

  • QIPO-Diff configurations: QIPO-Diff matches Lu et al. (2023) in network structure and warm-up configurations for fair comparison.
  • QIPO-Diff configurations: QIPO-Diff fine-tunes the warm-started score network with learning rate 10^-4 and uses DPM-Solver with diffusion-step = 15.
  • QIPO-OT configurations: QIPO-OT uses the optimal-transport vector field and applies DPM-Solver for solving the reverse-time ODE.
  • Benchmark evaluation: Table 4 reports normalized D4RL performance as mean ± standard deviation across 8 random seeds, while β = 10 ablations average 3 seeds.
  • Timing comparison: QIPO-OT consistently outperforms QIPO-Diff in speed, while QGPO is slowest because of additional back-propagation.
  • Evaluation procedure: The QIPO procedure evaluates performance every 5 epochs over K3 = 100 and renews the support set with Krenew = 10.

E.3 ABLATION STUDIES FOR QIPO

The ablations examine guidance scale, support-set size, support-set renewal frequency, energy landscapes, and β-conditioned diffusion. Results indicate robustness to guidance-scale changes, modest benefits from larger support sets, and robustness across modified energy landscapes.

  • Guidance scale: QIPO is robust to changing the diffusion guidance scale from β = 3 to β = 10, with performance sometimes above and sometimes slightly below β = 3.The comparison averages performance over 3 random seeds for the ablation setting.
  • Support action set: Larger support action sets M = 32 and M = 64 may slightly improve performance over M = 16, but M = 16 is retained to balance efficiency and effectiveness.Sampling the support set requires running the reverse process of the diffusion or score-matching model.
  • Support renewal period: Smaller support-set renewal periods generally improve performance across Krenew ∈ {5, 10, 20, 50}.The paper hypothesizes that the D4RL score network is relatively easy for the neural network to learn.
  • Energy landscapes: The energy-weighted diffusion model remains robust when the energy landscape is modified in the 8-Gaussian example.
  • β-conditioned diffusion: A β-conditioned score model samples β uniformly from U(0, βmax) and achieves performance comparable to separately trained models.The score model receives β as an input, and the reverse process uses a specified guidance scale.
  • Additional generation tasks: The molecular-generation evaluation guides normalized molecular properties including α, ϵHOMO, ϵLUMO, and Δϵ.

F.5 TRAINING CURVES OF OT AND DIFFUSION

The training-curve comparison evaluates OT flow matching and diffusion for learning the behavioral policy on six Antmaze tasks. OT converges similarly but slightly faster, while both approaches correctly learn the behavioral policy.

  • Training curves: Figure 10 compares unnormalized diffusion and OT flow-matching losses over the first 600 epochs.
  • Training curves: Figure 11 normalizes each loss by setting the first iteration to 1 and the last iteration to 0.The normalization accounts for different remaining losses between OT and diffusion.
  • Comparison: OT flow matching converges similarly but slightly faster to a stable point on the Antmaze tasks.
  • Comparison: Both OT flow matching and diffusion correctly learn the behavioral policy despite their difference in convergence rate.
Loading 2503.04975v1…