PaperlayerView exact revision on arXiv

Source-linked AI summary

Refinement-Based Flow Policy Optimization

Bumgeun Park, Hyukjun Yang, Donghwan Lee

arXiv:2609.15123v2cs.LGcs.AI

TL;DR

Flow-based policy training needs samples from the target distribution, but online RL defines that distribution implicitly through a Q-function, making direct sampling difficult. RFPO alternates Q-guided refinement of its own actions with self-target flow matching; it captures multimodal synthetic targets and matches or outperforms representative baselines in almost every evaluated continuous-control environment.

  • Problem

    Flow matching requires samples from the distribution being modeled, but the desired online-RL action distribution is defined implicitly by an evolving Q-function and is difficult to sample directly.

    Source: §1 · p. 1

  • Method

    RFPO refines current-flow-policy actions toward a Q-induced target with finite-step MALA, then trains the policy by flow matching refined actions paired with their original noise.

    Source: §1 · p. 2

  • Results

    Across six continuous-control MuJoCo environments, RFPO matches or outperforms representative online RL baselines in almost every environment, and captures multimodal structure without mode collapse on six synthetic targets.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    RFPO shifts its policy toward high-value action regions while retaining the capacity to represent multiple action modes without direct samples from the Q-induced target.

    Source: §7 · p. 9

    The theoretical analysis assumes a fixed critic, exact self-target flow matching, and exact flow integration; practical training instead uses a changing critic and finite optimization steps.

    Source: §7 · p. 10

Abstract

from arXiv · show

Flow-based policies offer an expressive representation for online reinforcement learning, but conventional flow matching requires samples drawn from the distribution to be modeled. This poses a challenge when the desired action distribution is defined only implicitly by a Q-function, since directly sampling actions from the resulting distribution is generally intractable. We propose Refinement-Based Flow Policy Optimization (RFPO), a novel framework for training a flow policy in online reinforcement learning by alternating between Q-guided sample refinement and self-target flow matching. RFPO first generates actions from Gaussian noise using the current flow policy and then uses a finite-step stochastic refinement procedure to move them toward an energy-based distribution induced by the Q-function. Each refined action is then paired with its corresponding initial noise sample and used as a fixed target for flow-matching training. By repeatedly refining its own outputs and learning from the resulting targets, RFPO incorporates Q-guidance into the policy without requiring direct samples from the target distribution, while retaining the capacity to represent multiple action modes. We further provide a theoretical analysis of the distributional dynamics induced by RFPO. Across six continuous-control tasks, RFPO matches or outperforms a standard Gaussian-policy baseline on almost every task. Experiments on six synthetic two-dimensional target distributions with diverse geometries demonstrate that RFPO captures complex multimodal structure without mode collapse.

1 INTRODUCTION

Flow matching can represent multimodal action distributions, but its conventional training requires target samples that are difficult to obtain when the desired distribution is defined implicitly by a changing Q-function. RFPO alternates Q-guided refinement with self-target flow matching, and the paper reports theoretical and empirical support for this approach.

1 INTRODUCTION

  • Conventional diffusion and flow-matching training needs target-distribution samples, which are difficult to obtain when the desired action distribution is defined implicitly by an evolving Q-function.Prior work addresses the challenge with Q-weighted objectives, critic-derived targets, or explicit action refinement.Source: §1 · p. 1
  • RFPO alternates Q-guided MCMC refinement of current-policy actions toward a Q-induced energy distribution with flow matching on refined actions paired with their original noise samples.The refined actions serve as fixed targets, so the flow update does not differentiate through the iterative sampler.Source: §1 · p. 2
  • For a fixed critic, the theoretical analysis establishes monotonic improvement of the entropy-regularized policy objective and identifies the energy-based target distribution as a fixed point.The paper motivates MCMC refinement as policy improvement toward a potentially multimodal Q-induced target.Source: §1 · pp. 1–2
  • RFPO matches or outperforms representative online RL baselines in almost every evaluated continuous-control environment and captures multimodal structure across synthetic targets.The synthetic experiments cover diverse geometries and show no collapse onto a single mode.Source: §1 · p. 2

2 RELATED WORK

Prior Q-guided generative policies use weighted objectives, critic-derived gradients, or refined actions to train diffusion and flow models. RFPO applies iterative refinement to actions from the current flow policy and pairs each refined action with its source noise for online policy learning.

2.1 Q-GUIDED GENERATIVE POLICY LEARNING

  • Q-guided generative policy methods include Q-weighted objectives, score matching, critic-gradient updates, and Q-derived velocity targets for diffusion or flow policies.Examples include QVPO, DPMD and SDAC, QSM, DACER, and RFM.Source: §2.1 · p. 2
  • Unlike DIPO and DACER-F, which refine replay-buffer actions, RFPO initializes MALA from current-policy actions and pairs refined actions with their source noise for self-target flow matching.DIPO uses Q-gradient ascent and fits a diffusion policy; DACER-F uses unadjusted Langevin updates and combines flow matching with direct Q-value maximization (Yang et al., 2023; Zhu et al., 2026).Source: §2.1 · p. 2

2.2 AMORTIZING ITERATIVE SAMPLE REFINEMENT

  • RFPO adapts the refinement–learning structure used in amortized MCMC, amortized SVGD, and MCMC teaching to online policy learning with a state-dependent Q-function target.Those preceding methods apply iterative refinement to model-generated samples and train the generating model to reproduce the refined samples.Source: §2.2 · p. 3

3 PRELIMINARIES

The preliminaries define a Q-induced Boltzmann action target, conditional flow matching as transport from base noise to a target distribution, and MALA as a corrected sampling procedure. Together these concepts establish the target and tools used by RFPO.

3.1 PROBLEM FORMULATION

  • The problem formulation uses a state-dependent Boltzmann distribution induced by the differentiable Q-function as the target action distribution.A smaller temperature α concentrates probability more strongly around high-value actions, while a larger α spreads it more broadly.Source: §3.1 · p. 3

3.2 CONDITIONAL FLOW MATCHING

  • Conditional flow matching learns a time-dependent vector field that transports base samples toward a conditional target distribution using source–target pairs and interpolated paths.Under the standard formulation, minimizing the objective recovers a marginal vector field whose flow transports the base distribution toward the target.Source: §3.2 · p. 4

3.3 METROPOLIS-ADJUSTED LANGEVIN ALGORITHM

  • MALA proposes gradient-informed moves and applies a Metropolis–Hastings correction that leaves the target distribution invariant.The method uses log-density gradients and density ratios, so it does not require evaluating the target's normalizing constant.Source: §3.3 · p. 4

4 REFINEMENT-BASED FLOW POLICY OPTIMIZATION

RFPO generates actions from Gaussian noise, refines them toward the Q-induced target, and trains the flow policy on refined-action/noise pairs. Its critic update uses flow-generated next-state actions in the temporal-difference target.

4 REFINEMENT-BASED FLOW POLICY OPTIMIZATION

  • RFPO alternates Q-guided refinement of current-policy actions with self-target flow matching on refined actions paired with their original noise.The refined targets are held fixed during the flow update, avoiding backpropagation of critic gradients through the iterative flow sampler.Source: §4 · p. 5

4.1 Q-GUIDED SAMPLE REFINEMENT

  • MALA refines each current-policy action toward the Q-induced target using a finite number of steps, bringing the policy distribution closer without assuming the result is an exact target sample.The refined action initializes from the flow-generated action and is produced using the current Q-function.Source: §4.1 · pp. 5–6

4.2 SELF-TARGET FLOW MATCHING

  • Self-target flow matching pairs each refined action with the base noise that generated its original action and treats the refined action as a fixed training target.The objective uses conditional flow matching on interpolated source–target pairs, with gradients applied to the velocity field rather than the refinement process.Source: §4.2 · p. 6

4.3 CRITIC LEARNING

  • The clipped double-Q critic update uses replay-buffer transitions and a flow-generated next-state action in its temporal-difference target, without additional smoothing noise.The target uses the minimum of two target critics evaluated at that action.Source: §4.3 · p. 6

5 THEORETICAL ANALYSIS

Under exact self-target flow matching and flow integration, RFPO reproduces the refined-action distribution, and each cycle corresponds to K MALA transitions. For a fixed critic, these dynamics monotonically improve the entropy-regularized objective and make the energy-based target a fixed point.

5 THEORETICAL ANALYSIS

  • The results assume a realizable conditional-mean velocity field, exact recovery through flow matching, and error-free flow integration.The analysis also holds the critic, temperature, and refinement step size fixed across cycles.Source: §5 · p. 6
  • Under the idealized analysis, self-target flow matching reproduces the refined-action distribution even when each action is paired with its original source noise.The pairing changes the transport path, not the terminal distribution.Source: §5 · p. 7
  • Each RFPO cycle has the same distributional effect as applying K MALA transitions to the current policy distribution.The MALA kernel leaves the target distribution invariant, so finite-step refinement can accumulate across policy updates.Source: §5 · p. 7
  • For a fixed critic, each cycle monotonically improves the entropy-regularized policy objective and leaves the energy-based target as a fixed point, but convergence requires additional MALA conditions.The monotonicity holds for every step size η > 0 and K ≥ 1; the guarantee alone does not establish convergence to the target.Source: §5 · p. 7

6 EXPERIMENTS

Across six continuous-control environments, RFPO matches or outperforms representative online RL baselines in almost every environment. On six synthetic two-dimensional targets, it captures diverse multimodal geometries without collapsing onto a single high-value region.

6.1 LEARNING MULTIMODAL DISTRIBUTIONS WITH DIVERSE GEOMETRIES

  • Across six two-dimensional targets, RFPO captures multiple high-density regions and varied geometries without collapsing onto a single high-value region.The targets include separated and anisotropic components and nonlinear ridge structures; Figure 2 compares target densities, learned samples, and component-region masses.Source: Figure 2 · p. 8, §6.1 · p. 8

6.2 PERFORMANCE ON CONTINUOUS-CONTROL BENCHMARKS

  • RFPO matches or outperforms representative online RL baselines in almost every environment; Table 2 reports final smoothed returns across six MuJoCo tasks.Methods were trained for 1M environment steps.Source: §6.2 · p. 8, Table 2 · p. 9

7 CONCLUSION

RFPO alternates Q-guided refinement of its own policy samples with self-target flow matching, avoiding direct samples from the Q-induced target. Its analysis and experiments examine the resulting distributional dynamics and performance.

7 CONCLUSION

  • RFPO iteratively refines actions sampled from its current policy toward the Q-induced energy-based target, then uses the refined actions as self-generated flow-matching targets.This avoids requiring direct samples from the target while retaining multiple action modes.Source: §7 · p. 9
  • Across six MuJoCo environments, RFPO matches or outperforms representative online RL baselines in almost every environment.On six synthetic two-dimensional targets, it captures complex multimodal structure without collapsing onto a single mode.Source: §7 · p. 9
  • Under exact self-target flow matching and exact flow integration, each RFPO cycle corresponds to K MALA transitions, with the energy-based target as a fixed point.The analysis assumes a fixed critic; practical training instead has an evolving critic and finite optimization steps.Source: §7 · pp. 9–10
  • RFPO’s refinement kernel, acceptance correction, and current-policy initialization are design choices whose individual contributions are not fully isolated in the experiments.The authors identify systematic ablation of these choices and their combinations as future work.Source: §7 · p. 10

AI USE STATEMENT

The authors disclose generative AI assistance across research and manuscript work, while distinguishing tasks for which it was not used. They also describe their review of AI-assisted material and provide implementation, experimental, and theoretical details for reproducibility.

AI USE STATEMENT

  • Generative AI tools assisted with brainstorming, methodology and design feedback, theoretical discussion and proof checking, literature work, implementation, and manuscript preparation.The authors state that AI did not generate or clean datasets, formulate the mathematical claims, refine hypotheses, or interpret results.Source: AI USE STATEMENT · p. 10
  • The authors reviewed and revised AI-assisted material, used proof checking only as a supplement to their own verification, and take responsibility for the final work..Source: AI USE STATEMENT · p. 10
  • Implementation, hyperparameters, evaluation procedures, synthetic targets, baseline configurations, and theoretical assumptions and proofs are specified in the appendices.The authors state that these details are sufficient to reproduce the method and evaluate its claims.Source: REPRODUCIBILITY STATEMENT · p. 10

A COMPARISON WITH DIRECT Q MAXIMIZATION

Direct Q maximization sends critic gradients through the flow sampler, whose trajectory sensitivities can attenuate or amplify them. RFPO instead uses Q-guided fixed regression targets; its bounded-action implementation performs refinement and flow matching in transformed action coordinates.

A COMPARISON WITH DIRECT Q MAXIMIZATION

  • Direct Q maximization backpropagates critic gradients through the flow sampler, where products of solver-step state Jacobians can attenuate or amplify the gradient reaching flow parameters.The sensitivity also varies along trajectories and across temperatures.Source: A COMPARISON WITH DIRECT Q MAXIMIZATION · p. 13
  • Across five temperatures, Figure 4 reports trajectory sensitivity factors Ai ranging from 0.08 to 1.31; early perturbations can be attenuated, while later ones remain near scale or are mildly amplified.These measurements illustrate sampler-dependent Jacobian factors in the direct Q-maximization gradient, not gradient norms or optimization stability.Source: A COMPARISON WITH DIRECT Q MAXIMIZATION · p. 13, Figure 4 · p. 13
  • RFPO places Q-gradient information in a fixed velocity-regression target, with no derivative through flow sampling or critic gradient backpropagated through the sampler.The comparison identifies a gradient path RFPO avoids; it does not claim direct Q maximization necessarily diverges.Source: A COMPARISON WITH DIRECT Q MAXIMIZATION · p. 14

B BOUNDED ACTION IMPLEMENTATION

  • In the bounded-action implementation, the flow produces a latent endpoint that is mapped to bounded actions by an elementwise tanh transformation.The latent flow state is clipped to [−10, 10] after each Euler sampling step; MALA refinement does not apply that clipping.Source: B BOUNDED ACTION IMPLEMENTATION · p. 14
  • Refinement in latent u-space targets the action-space distribution transformed by the action map’s Jacobian determinant.The log-target gradient includes the Q-function differentiated through the map and a log-Jacobian term.Source: B BOUNDED ACTION IMPLEMENTATION · p. 14
  • After K MALA steps, each refined latent endpoint is paired with its original noise and treated as a fixed target for flow matching in u-space.The training interpolation runs from the source noise to the refined endpoint.Source: B BOUNDED ACTION IMPLEMENTATION · p. 15

C CRITIC LEARNING AND TRAINING PROCEDURE

The training procedure combines critic updates with repeated flow-matching updates on a reused batch of policy-generated and refined samples. The appendix also specifies evaluation asymmetries and practical hyperparameter choices, including temperature and MALA step-size adaptation.

C CRITIC LEARNING AND TRAINING PROCEDURE

  • The critic target action comes from the online flow policy; RFPO maintains no target flow policy and adds neither an entropy term nor target-policy smoothing noise.Time-limit truncations reset the environment but retain the critic bootstrap term.Source: C CRITIC LEARNING AND TRAINING PROCEDURE · p. 15
  • MALA refinement uses the online clipped-double critic for both proposal scoring and the Metropolis–Hastings acceptance probability.The refinement target is the minimum of the two online critics.Source: C CRITIC LEARNING AND TRAINING PROCEDURE · p. 15
  • After 5,000 uniformly sampled warm-up actions, each environment step performs one critic update and five flow-matching updates using a batch constructed once and reused.Only interpolation time is resampled across the five flow-matching updates.Source: C CRITIC LEARNING AND TRAINING PROCEDURE · p. 15
  • RFPO uses stochastic flow-policy actions for evaluation, whereas each baseline uses deterministic actions; the authors say this asymmetry works against RFPO.RFPO and SAC use 5,000 warm-up actions, compared with 10,000 for diffusion- and flow-policy baselines.Source: C CRITIC LEARNING AND TRAINING PROCEDURE · p. 16

D HYPERPARAMETERS

  • In the two-dimensional experiments, Q is specified in closed form rather than learned; grid evaluation supplies ground-truth mode weights, while RFPO estimates weights from generated samples.The method accesses the target only through Q.Source: Table 3 · p. 17
  • For HumanoidStandup, temperature starts at 2.0, reaches 0.2 after the first 20% of training, and stays there; other environments use constant temperatures.The final temperature is 0.2.Source: D HYPERPARAMETERS · p. 17
  • MALA step size adapts toward a 0.6 acceptance rate, increasing when the estimate exceeds that target and decreasing otherwise.The authors describe this as a practical choice; their theoretical analysis uses a fixed refinement kernel.Source: D HYPERPARAMETERS · p. 17
  • Clipping the MALA proposal drift modifies the proposal while preserving the target distribution as invariant for a fixed target and step size.Forward and reverse proposal densities use the clipped drift, while the target-density ratio remains unmodified.Source: D HYPERPARAMETERS · pp. 17–18

E TARGET GEOMETRIES

The target landscapes span separated Gaussian bumps and curve-shaped ridges, with shared reference heights that do not imply equal component masses. The appendix states the kernel, assumptions, and integrability result used in the proofs.

E TARGET GEOMETRIES

  • The six landscapes share nominal reference heights hmin = 0.55 and hmax = 1.00, but component masses also depend on widths, lengths, and overlap.The ring’s baseline height h0 = 0.40 maintains a continuous ridge between higher-value regions.Source: E TARGET GEOMETRIES · pp. 18, 20
  • Blob targets use Gaussian-shaped bumps, while curve targets form ridges around sampled curves with normalized distance-based weights.Curve-ridge widths preserve visible geometric features, including the ring’s central low-density region.Source: E TARGET GEOMETRIES · pp. 18–20
  • Curve-landscape rescaling sets the maximum over sampled curve points to hmax = 1.00, not the global maximum over the action domain.Component evaluation assigns actions to segments using their nearest sampled curve points.Source: E TARGET GEOMETRIES · p. 19

F.1 DEFINITIONS

  • The appendix defines the MALA transition kernel Ms, and the Metropolis–Hastings correction makes the energy-based target pα(· | s) invariant for that kernel.Applying the kernel K times to π gives πM_s^K.Source: §F.1 · p. 21

F.2 ASSUMPTIONS

  • The theoretical analysis assumes regularity and moment conditions across cycles, including finite Qϕ and action gradients on the refinement domain.The assumptions apply for ρ-almost every state used in the self-target flow-matching objective.Source: §F.2 · p. 21

F.3 INTEGRABILITY OF THE VELOCITY FIELD

  • Under those assumptions, the conditional-mean velocity field is defined almost everywhere and has finite integrated squared norm.The proof uses Gaussian initial noise and a finite second moment for refined actions.Source: §F.3 · pp. 21–22

F.4 PROOF OF LEMMA 5.1

The appendix establishes that exact self-target flow matching reproduces the refined-action distribution, so alternating updates correspond to repeated refinement-kernel transitions. It then derives monotonicity and fixed-point results, discusses bounded-action coordinates, and describes comparison baselines.

F.4 PROOF OF LEMMA 5.1

  • Self-target flow matching recovers the conditional-mean velocity field, whose flow maps Gaussian noise to the refined-action distribution.The proof uses the coupling of Gaussian source noise with refined actions and establishes the endpoint law by taking the interpolation limit.Source: §F.4 · p. 22

F.5 PROOF OF LEMMA 5.2

  • With exact inner optimization and flow integration, each policy update applies K MALA transitions to the current policy distribution.Iterating the identity gives πn = π0M_s^(nK).Source: §F.5 · p. 23

F.6 PROOF OF THEOREM 5.3

  • For a fixed critic, the entropy-regularized objective is nondecreasing across updates, and the energy-based target is a fixed point.The argument does not constrain η or K.Source: §F.6 · pp. 23–24

F.7 BOUNDED ACTIONS

  • The bounded-action analysis does not cover Euler discretization or latent clipping during practical flow sampling.It assumes a fixed refinement kernel and exact flow-ODE integration.Source: §F.7 · p. 24
  • Under the stated integrability conditions, the latent-coordinate and action-coordinate objectives agree, extending the idealized conclusions to the induced bounded-action distribution.Practical bounded actions are recovered through T(u).Source: §F.7 · p. 24

G BASELINES

  • SAC is the representative Gaussian-policy baseline, while DPMD and SDAC use Q-reweighted score matching for diffusion policies.DPMD targets a policy-mirror-descent distribution; SDAC targets a Q-induced Boltzmann distribution.Source: G BASELINES · p. 25
  • DIPO refines replay-buffer actions with Q-gradient ascent, whereas RFM estimates velocity targets; RFPO instead refines current-policy actions with MALA.DIPO trains a diffusion policy on refined actions, while RFM uses posterior-mean estimation for its flow targets.Source: G BASELINES · p. 25

H FULL RESULTS ON MULTIMODAL TARGET DISTRIBUTIONS

The appendix presents sample-distribution and mode-weight evolution for the evaluated synthetic targets over 3,000 cycles. Its cost analysis breaks down RFPO’s network computation and contrasts tensor sizes with particle-based baselines.

H FULL RESULTS ON MULTIMODAL TARGET DISTRIBUTIONS

  • The full-results section presents generated samples at selected cycles and component-region masses over 3,000 cycles, with four-seed standard-deviation bands.Dashed lines indicate ground-truth component-region masses.Source: H FULL RESULTS ON MULTIMODAL TARGET… · p. 25
  • The figures track generated samples and mode weights for isotropic, grid, anisotropic, ring-shaped, spiral-shaped, and arc-shaped targets.The captions identify the target geometry and show its evolution.Source: Figure 5 · p. 26, Figure 6 · p. 26, Figure 7 · p. 26

I COMPUTATIONAL COST

  • RFPO’s two 20-step ODE integrations account for 10,240 of 20,500 total rows, while five-step MALA refinement accounts for 6,144 rows, or 30%.Each refinement uses six critic-ensemble evaluations and six input-gradient computations at batch size 256.Source: I COMPUTATIONAL COST · p. 28
  • The widest RFPO network tensor is 256 rows, compared with 8,192 for DPMD and SDAC and up to 25,600 for RFM.RFPO uses six short critic input-gradient computations to refine each policy-generated action.Source: I COMPUTATIONAL COST · p. 28

Related papers

Loading 2609.15123v2…