Reverse Flow Matching: A Unified Framework for Online Reinforcement Learning with Diffusion and Flow Policies

• Online RL lacks direct samples from the target action distribution, while the Boltzmann policy is unnormalized and generally intractable to sample directly.Source: §1 · p. 1
• Noise-expectation methods weight noise by Q-values, while gradient-expectation methods weight Q-function gradients; their formal relationship and general synthesis remain unclear.Source: §1 · pp. 1–2
• RFM addresses unavailable target samples with a tractable loss that casts training as posterior mean estimation.Source: §1 · p. 2
• Langevin Stein operators yield zero-mean control variates, unifying noise- and gradient-expectation methods and enabling Q-value/Q-gradient combination for improved efficiency and stability.Source: §1 · p. 2
• RFM extends Boltzmann-distribution targeting from diffusion to flow policies, which can use source distributions beyond the standard Gaussian to incorporate domain knowledge.Source: §1 · p. 2
• RFM-trained flow policies demonstrate superior performance on continuous-control benchmarks compared with diffusion-policy baselines.Source: §1 · p. 2

• Optimization via differentiable sampling is one category of prior methods (Wang et al., 2024; Lv et al., 2025; Celik et al., 2025).Source: §2 · p. 2
• Other categories use Langevin dynamics for sampling (Psenka et al., 2024) or iterative weighted regression (Ding et al., 2024; Fan et al., 2025; Ma et al., 2025).Source: §2 · p. 2
• A fourth category targets Boltzmann distributions (Akhound-Sadegh et al., 2024; Jain et al., 2025; Ma et al., 2025; Dong et al., 2025).Source: §2 · p. 2

3.1. Reinforcement Learning

• An MDP specifies states, actions, transitions, rewards, initial-state distribution, and discount; maximum-entropy RL adds entropy regularization, with policy improvement yielding a Boltzmann action distribution.Source: §3.1 · p. 2

3.2. Diffusion and Flow

• Diffusion and flow learn an evolution from source to target through a probability path; conditional flow matching regresses to conditional velocities rather than an often-intractable marginal velocity.Under the stated equivalence, conditional and marginal flow-matching objectives share the same minimizers and gradients.Source: §3.2 · pp. 2–3

3.3. Control Variates

• A control variate adds an auxiliary function with known mean to a Monte Carlo estimator, with its coefficient chosen to reduce variance.Source: §3.3 · p. 3

4. Reverse Flow Matching

• RFM trains diffusion and flow models for targets lacking efficient samplers by replacing direct target samples with posterior-mean estimation.Direct policy gradient would require backpropagating through the full sampling procedure, which is computationally expensive and numerically unstable.Source: §4 · pp. 3–4

4.1. From Forward Construction to Reverse Inference

• Treating intermediate Xt as evidence, RFM infers endpoint posteriors and shares conditional flow matching’s global minimizers under absolute-continuity and realizability assumptions.The optimization dynamics can still differ because reverse-flow-matching gradients depend on the proposal distribution.Source: §4.1 · pp. 4–5
• RFM supports alternative prediction parameterizations and standard diffusion schedules, including reverse score matching for non-Gaussian source distributions.Source: §4.1 · p. 5

4.2. Langevin Stein Operators

• Langevin Stein operators yield vector-valued zero-mean control variates that preserve posterior means and can reduce estimation variance.Matrix-valued test functions extend the scalar Stein construction to vector-valued posterior means.Source: §4.2 · pp. 5–6
• Posterior means also enable training-free sampling through ODE or SDE integration, while Stein control variates can reduce variance and extend to MCMC and SMC.Source: §4.2 · p. 7

4.3. Application to Boltzmann Distributions

• For Boltzmann targets, η=0 and η=1 recover noise-expectation and gradient-expectation methods, respectively, unifying prior online-RL approaches (Akhound-Sadegh et al., 2024; Jain et al., 2025; Ma et al., 2025; Dong et al., 2025).The estimator’s variance can be controlled through the coefficients, for which variance-minimizing choices are available.Source: §4.3 · p. 7, §4.2 · p. 7

4.4. Online Reinforcement Learning with Flow Policies

• RFM is instantiated for flow policies in online RL, with extension to diffusion policies described as straightforward.Source: §4.4 · p. 7
• Bounded actions are handled by learning the flow in unconstrained latent space and mapping latents to actions with tanh.The latent variable u is in R^d and the action is a = tanh(u).Source: §4.4 · p. 8
• The latent-space Boltzmann target includes a tanh change-of-variables Jacobian, which is essential for enforcing the correct action-space distribution.Source: §4.4 · p. 8
• The policy estimates the posterior mean of base noise given the state, time, and noisy latent, then uses it to form the RFM velocity target and actor loss.The posterior is defined using the base noise distribution and the latent-space soft policy-improvement target.Source: §4.4 · p. 8
• A policy-induced proposal samples noisy latents, and the method selects the highest-Q action among M candidates.The proposal forms u_t by interpolating an ODE-generated u_1 with Gaussian noise.Source: §4.4 · p. 8

5.1. Toy Example

• RFM maintains the minimum discrepancy on SWD, MMD2, and Sinkhorn distance while using 20 inference steps, compared with 100 for iDEM and QNE.Methods use the same posterior-estimation budget; the target is a two-moon distribution whose density is known up to a normalizing constant.Source: Figure 1 · p. 8, §5.1 · pp. 8–9
• For the toy experiment, RFM is compared with iDEM and QNE using the same posterior-mean Monte Carlo sampling number, hence the same training computation.RFM uses 20-step ODE integration at inference, while the diffusion baselines use 100 denoising steps.Source: §5.1 · p. 9

5.2. RL Tasks

• The RL evaluation compares RFM with SAC, QSM, MaxEntDP, DQS, and QVPO on eight DeepMind Control Suite environments.The baselines include Gaussian-policy maximum-entropy RL, score matching, noise- and gradient-expectation diffusion policies, and weighted regression.Source: §5.2 · p. 9
• RFM is the only method reported to perform consistently well across all eight environments, with substantially better stability and fewer inference steps than the diffusion policy baselines.RFM uses 10 flow steps; the other diffusion policy baselines use 20 diffusion steps.Source: §5.2 · p. 9

6. Conclusion

• Reverse flow matching provides a unified framework for training diffusion and flow policies to sample from Boltzmann distributions in online RL.It turns unavailable target samples into a posterior mean estimation problem using a tractable objective.Source: §6 · p. 9
• Langevin Stein operators yield control variates for a general class of posterior mean estimators that reduce estimation variance and improve training stability.Existing methods for training diffusion policies to sample Boltzmann distributions arise as special cases of the formulation.Source: §6 · p. 9
• The paper reports stronger performance and greater stability across continuous-control tasks than state-of-the-art baselines.Source: §6 · p. 9

Impact Statement

• RFM targets unnormalized distributions without direct samples and is presented as supporting stable, efficient learning of expressive policies for continuous-control tasks.The impact statement frames this as support for deploying RL methods in increasingly complex real-world applications.Source: Impact Statement · p. 10

A. Related Works

• The related-work review focuses on off-policy RL algorithms with diffusion and flow policies, citing sample efficiency and strong empirical performance as its rationale.It notes that on-policy algorithms are also an active research area.Source: A. Related Works · p. 12
• The review groups prior methods into differentiable sampling, Langevin dynamics, iterative weighted regression, and targeting Boltzmann distributions.The Boltzmann-targeting group includes gradient-expectation methods that average energy or Q-function gradients and noise-expectation methods that average noise [Akhound-Sadegh et al., 2024; Jain et al., 2025; Ma et al., 2025; Dong et al., 2025].Source: A. Related Works · p. 12

B. Prior Methods as Special Cases

• RFM recovers prior noise- and gradient-expectation methods as estimators from the same posterior mean problem, distinguished by their control-variate test functions.The noise-expectation family uses Q evaluations and averages over noise; the gradient-expectation family averages over ∇Q [Ma et al., 2025; Dong et al., 2025; Akhound-Sadegh et al., 2024; Jain et al., 2025].Source: B. Prior Methods as Special Cases · p. 12
• For a standard Gaussian source and isotropic test function Φt = ηId, the section gives a posterior mean estimator for X0 and derives the corresponding estimator for X1.The X1 estimator is obtained from the data-posterior form by deriving its control variate.Source: B. Prior Methods as Special Cases · p. 12
• The section identifies schedule choices for iDEM and QNE: variance-exploding for iDEM and variance-preserving for QNE.QNE's proposal distribution for x0 is p0 = N(0, Id); RSM uses that proposal as well and a discrete-time DDPM.Source: B. Prior Methods as Special Cases · p. 13
• The noise-expectation methods QNE and RSM correspond to the noise-posterior estimator with η = 0.QNE uses the noise-posterior estimator to train a noise-prediction network, while RSM uses the same estimator to train a score-prediction network.Source: B. Prior Methods as Special Cases · p. 13
• The gradient-expectation methods iDEM and DQS correspond to the data-posterior estimator with η = 1.iDEM uses this estimator to predict the score; DQS follows the iDEM approach.Source: B. Prior Methods as Special Cases · p. 13

• LRFM, LRFM−N, and LRFM−D differ only by parameter-independent constants, so they share global minimizers and parameter gradients.Source: C. Proofs of Theoretical Results · p. 14
• Under mutual absolute continuity and a sufficiently rich function class, LRFM-N, LRFM-D, and conditional flow matching have the same global minimizers.The objectives use the same regression target and differ only in the weighting measure over intermediate samples.Source: C. Proofs of Theoretical Results · pp. 14–15
• The zero-mean control-variate construction supports posterior-mean estimation, and its estimator has zero variance exactly when it is constant on the posterior support.Source: C. Proofs of Theoretical Results · p. 15
• Propositions 4.9 and 4.10 specify diagonal and isotropic coefficient choices that minimize the asymptotic variance of the corresponding SNIS estimators.Source: C. Proofs of Theoretical Results · p. 16
• Theorem 4.14 gives a family of Boltzmann-target posterior-mean estimators controlled by diagonal coefficients, with Gaussian-source and isotropic cases yielding additional identities and a linear-combination form.The isotropic restriction sets all coefficients to η.Source: C. Proofs of Theoretical Results · p. 17

D. Reverse Score Matching

• For arbitrary source distributions, reverse score matching separately learns the score because reverse-flow-matched velocity need not yield a tractable score expression.With a Gaussian source, the score has a closed-form expression linearly related to velocity, so learning velocity is sufficient for SDE-based sampling.Source: D. Reverse Score Matching · pp. 17–18
• Reverse score matching applies reverse inference to train without direct target samples, using conditional score matching under linear interpolation and independent coupling.Source: D. Reverse Score Matching · pp. 17–18
• The reverse score-matching and noise- or data-posterior objectives share global minimizers and identical gradients up to parameter-independent constants; under Theorem 4.2 assumptions, they also share minimizers with conditional score matching.Source: D. Reverse Score Matching · pp. 18–19

E. Algorithm Details

• The implementation estimates the posterior mean with SNIS using K samples and estimates the diagonal variance-control coefficients from the same samples, with a small ridge term.Source: E. Algorithm Details · p. 19

• The toy target is a two-moon Boltzmann distribution with temperature λ=1 and energy E(x).Source: §F.1 · p. 19
• RFM uses the linear schedule αt=t and βt=1−t, making the conditional velocity x1−x0 the target for regression via a posterior-mean estimate.Source: §F.1 · p. 20
• A diagonal-covariance Gaussian-mixture proposal is fitted to the unnormalized target for efficient SNIS posterior estimation.The proposal is fitted by maximizing a variational objective.Source: §F.1 · p. 20
• Sample quality is evaluated with sliced Wasserstein distance, RBF-kernel MMD2, and entropically regularized Sinkhorn distance.The metrics capture complementary notions of distributional discrepancy.Source: §F.1 · p. 20

F.2. RL Tasks

• RFM uses temperature λ=0.02 and 100 Monte Carlo samples for posterior-mean estimation across environments.Source: §F.2 · p. 20
• Average walker-run training time is reported in Table 2, with experiments conducted on an RTX 5090 GPU and Intel Core Ultra 9 285K CPU.Source: §F.2 · p. 21, Table 2 · p. 21

G.1. Sensitivity Analyses

• RFM is robust to K and flow-step choices, and insensitive to action-candidate count M provided M is not too small.The main experiments use K=100, M=32, and N=10.Source: §G.1 · p. 21, Figure 3 · p. 21

G.2. Ablation Studies

• RFM outperforms fixed isotropic noise- and gradient-expectation variants on finger-turn hard and yields better toy samples at equal computation.On the toy task, all three methods can perform well with sufficient training, so the comparison emphasizes limited computation.Source: §G.2 · pp. 21–22, Figure 4 · p. 22, Figure 5 · p. 22
• Standard-Q RFM outperforms its soft-Q variant, and adding the log-likelihood term slightly degrades performance.The paper cites the extra cost of flow-policy likelihood evaluation and suggests ODE likelihood computation may add numerical difficulty and instability.Source: Figure 4 · p. 22, §G.2 · p. 22
• The tested proposal distributions show no major performance differences in either the RL task or the toy example.The tested RL choices are policy-derived, replay-buffer-derived, and Gaussian proposals; the toy comparison uses policy-derived and Gaussian proposals.Source: §G.2 · p. 22, Figure 6 · p. 23

G.3. Extended Comparisons

• RFM outperforms SDAC and DIME on finger-turn hard and walker-run; DIME training is unstable.DIME's instability is attributed tentatively to its long backpropagation chain through the sampling process.Source: Figure 7 · p. 23, §G.3 · p. 23