ReFPO: Reflow Regularization for Flow Matching Policy Gradients

Abstract: We present Reflow-regularized Flow Matching Policy Gradients (ReFPO), a simple online RL method that adds explicit Reflow regularization to FPO for efficient flow-based control. We uncover a key structural property: the gradient updates in Flow Matching Policy Gradients (FPO) can be interpreted as an implicit advantage-weighted Reflow process, providing a new geometric perspective on flow-based policy gradients. Building on this insight, ReFPO introduces an explicit geometric regularizer that can be implemented with a single line of code change without incurring additional computational overhead or auxiliary distillation stages. By synergizing advantage-guided updates with path rectification, our method reduces CFM proxy-ratio spikes, stabilizes PPO-style training, and enables high-fidelity one-step inference that often matches or exceeds multi-step performance. We experimentally demonstrate that ReFPO improves average performance and discretization robustness across GridWorld, MuJoCo Playground, and high-dimensional Humanoid Control tasks, providing a scalable and stable approach for generative policies in complex physical simulations.

• The paper introduces an explicit geometric regularizer to stabilize flow matching policy gradients and improve one-step inference efficiency.
• It leverages conditional flow matching losses to compute a proxy advantage, ensuring smoother PPO updates and reducing policy ratio spikes.
• Empirical results across tasks like MuJoCo and humanoid control demonstrate enhanced multimodality, training stability, and reduced inference latency.

ReFPO: Explicit Geometric Regularization for Stable and Efficient Flow-Based Policy Gradients

Introduction and Motivation

Flow matching and diffusion models have emerged as powerful policy parameterizations within reinforcement learning (RL), offering significant expressivity for modeling complex, multimodal action distributions. Unlike conventional Gaussian policies, flow-based policies can capture non-convexity and multi-modality vital for high-dimensional tasks. However, their inherent iterative sampling—requiring multiple ODE integration steps per action—impedes real-time deployment due to substantial inference latency. Attempts to accelerate inference using rectified flows, consistency models, or distillation have had limited impact in online RL settings; these methods typically introduce multi-stage pipelines and additional complexity, failing to robustly address the non-stationarity inherent in online optimization.

This paper introduces Reflow-regularized Flow Matching Policy Gradients (ReFPO), a principled approach that leverages an explicit geometric regularizer grounded in flow rectification. The core insight is the geometric structure in Flow Matching Policy Gradients (FPO): FPO’s updates can be interpreted as an implicit advantage-weighted Reflow process operating on self-generated sample trajectories. ReFPO formalizes this geometric connection, making it explicit via lightweight regularization that enhances both training stability and inference efficiency, all with negligible computational cost.

Flow Matching Policy Gradients: Geometric Interpretation

FPO adapts the Proximal Policy Optimization (PPO) methodology to likelihood-free flow policies by substituting the clipped likelihood ratio with a proxy derived from Conditional Flow Matching (CFM) losses. This proxy ratio is computed via the exponential difference in CFM losses under current and previous parameters, acting as a tractable surrogate for the ELBO-based likelihood ratio. Crucially, FPO’s optimization is performed on policy-induced distributions—the model generates its own action-noise pairs, resulting in a closed feedback loop where the policy rectifies its own trajectories under reward-based guidance.

Mathematically, gradient updates in FPO for positive-advantage actions align the velocity field toward the straight direction connecting sample pairs, while negative-advantage actions invert this direction. This update mechanism closely resembles an advantage-weighted Reflow, which recursively straightens the probability flow according to environmental rewards during online RL.

Explicit Reflow Regularization: Algorithmic Formulation

While FPO implicitly straightens flows for rewarded actions, it can induce distortions for negatively weighted trajectories, locally amplifying curvature in the velocity field and increasing proxy-ratio volatility. ReFPO mitigates these issues by introducing an explicit, unweighted Reflow regularizer: it directly penalizes curvature in the learned velocity field across all sampled trajectories, regardless of advantage sign. The Reflow regularization term is constructed from the same CFM loss already used in FPO and added to the objective through a coefficient $\lambda$, requiring only a one-line code modification.

Two complementary effects arise: the regularizer reduces local truncation error for one-step ODE solvers, improving one-step inference performance; and it stabilizes the inner-loop PPO optimization, yielding smoother parameter updates and reducing the frequency of extreme policy-ratio spikes. These effects are evidenced empirically by improved one-step action quality, increased training stability, and reduced flow collapse risk.

Empirical Evaluation: Multimodality, Robustness, and Scaling

Multimodal Grid World

Visualizations in a bifurcated GridWorld demonstrate that ReFPO improves trajectory straightness and spatial consistency without sacrificing policy multimodality. Whereas FPO’s one-step action distribution retains fluctuations and directional drift, ReFPO yields coherent one-step trajectories closely matching multi-step rollouts.

Figure 1: Learned velocity fields and rollout trajectories in multimodal GridWorld for FPO (top) and ReFPO (bottom), showing improved straightness and one-step consistency under explicit regularization.

MuJoCo Playground: Continuous Control and Diagnostics

In the DeepMind Control Suite/MuJoCo Playground, ReFPO consistently achieves higher mean reward than FPO and baseline Gaussian PPO across ten tasks. Notably, ReFPO’s one-step ($N=1$) inference achieves rewards comparable to or exceeding multi-step rollouts, with reduced straightness error and explosion rate for the proxy ratio.

Figure 2: FPO and ReFPO rewards on ten MuJoCo Playground tasks; shaded regions show the gap between multi-step and one-step evaluation, which narrows significantly under ReFPO regularization.

Policy-ratio diagnostics reveal that reward drops in FPO coincide with sharp proxy-ratio spikes, evidencing instability during PPO updates. ReFPO regularization directly stabilizes the flow geometry, resulting in smaller proxy-ratio fluctuations and smoother reward curves.

Figure 3: Reward curves and policy ratio spikes on PointMass, FingerSpin, and BallInCup; ReFPO maintains stable rewards by suppressing proxy-ratio spikes and explosion events.

Humanoid Control: High-Dimensional Robustness

In SMPL-based humanoid MoCap tracking tasks, ReFPO maintains strong one-step performance, achieving high success rates and low MPJPE (Mean Per Joint Position Error) even under sparse goal conditioning (root+hands, root-only). One-step inference with ReFPO reduces latency substantially without trade-off in accuracy, and ReFPO outperforms both FPO and Gaussian PPO under incomplete conditioning.

Figure 4: Performance comparison for FPO and ReFPO on humanoid control, highlighting superior stability, tracking accuracy, and efficient one-step inference.

Practical and Theoretical Implications

ReFPO’s explicit geometric regularizer offers a scalable, general-purpose stabilization mechanism for online reward-weighted flow-matching policy optimization. The approach decouples path straightness from reward guidance, preventing mode collapse and enhancing robustness under coarse discretization. ReFPO’s empirical gains—particularly in sparse, high-dimensional tasks—illustrate the practical value of multimodal flow policies coupled with efficient, stable one-step inference.

Theoretically, ReFPO’s geometric perspective enriches the understanding of online flow-based RL, linking advantage-weighted updates to recursive path rectification and highlighting the importance of explicit geometric priors. The regularizer’s simplicity suggests potential for broader transferability to other CFM-based RL and fine-tuning methods.

Future work may extend ReFPO beyond PPO-style surrogates to a wider class of online flow-policy optimization algorithms, adaptively tune geometric regularization coefficients, and further explore integration with path-efficient flow models in RL and generative modeling.

Conclusion

ReFPO formalizes and exploits the implicit Reflow geometry within flow-matching policy gradients, introducing an explicit geometric regularizer with minimal overhead. The resulting method achieves improved one-step inference fidelity, training stability, and multimodal expressivity in online RL across diverse domains and task complexities. These advances underline the importance of geometric structure in flow-based policy optimization and pave the way for future developments in robust, low-latency generative control.