Flow Matching Policy Gradients

FPO

• FPO optimizes flow-based generative models with policy gradients using an advantage-weighted ratio derived from conditional flow matching.The method uses a PPO-inspired surrogate objective as a drop-in replacement for Gaussian policies.Source: §1 · p. 1
• FPO avoids explicit likelihood calculations and remains compatible with deterministic or stochastic sampling procedures with varying integration orders and step counts.Unlike denoising-as-MDP approaches, it treats the rollout sampling procedure as a black box.Source: §1 · p. 1

Empirical findings

• Flow policies learn multimodal action distributions in states with multiple optimal actions, unlike Gaussian policies.This behavior is observed in the GridWorld analysis.Source: §1 · p. 2
• FPO enables training from scratch for under-conditioned humanoid control, where only root-level commands are provided.The humanoid setting is among the high-dimensional tasks used to validate the method.Source: §1 · p. 2

Policy Gradients

• Policy gradients optimize policies without differentiable environment rewards by weighting action log-probabilities with observed rewards or advantages.They are widely used in continuous control, locomotion, and manipulation.Source: §2 · p. 2

Diffusion Policies

• FPO trains flow-based generative policies that represent richer distributions than diagonal Gaussians while remaining compatible with standard actor-critic training.This positions flow-based policies as generative alternatives for continuous-control reinforcement learning.Source: §2 · p. 2

Diffusion Models

• Diffusion models model complex continuous distributions and are widely used for images, videos, audio, and robot actions.Their applications include data collected through human teleoperation.Source: §2 · p. 2
• Flow matching learns a vector field that transports samples from a tractable prior distribution to a target data distribution.Its conditional objective trains denoising from Gaussian-perturbed data.Source: §2 · p. 2

Reinforcement Learning with Diffusion Policies

• Direct likelihood optimization for flow models requires divergence estimation and is computationally prohibitive, whereas flow matching optimizes a variational likelihood lower bound with a denoising loss.The paper leverages this loss within policy gradients to train diffusion models from rewards.Source: §2 · p. 2

Diffusion Policies

• Prior on-policy diffusion methods frame denoising as a Markov decision process, whereas FPO avoids binding training to a specific sampling method.This distinction supports flexible integration choices during training and inference.Source: §2 · p. 3

3.2 Flow Policy Optimization

• FPO replaces exact action likelihoods in PPO with a proxy ratio computed from transformed conditional flow-matching losses.The proxy is designed as a drop-in replacement for the PPO likelihood ratio.Source: §3.2 · pp. 3–4
• FPO remains compatible with GAE, GRPO, noise-prediction and clean-action parameterizations, and flow or diffusion implementations.Different loss parameterizations can be reweighted for mathematical equivalence to the conditional flow-matching objective.Source: §3.2 · p. 4, §3.4 · pp. 4–5

3.3 FPO Surrogate Objective

• The FPO ratio is computed from current and old policies’ ELBOs, which approximate log-likelihood through flow-matching denoising objectives.This avoids computationally expensive exact likelihood estimation in flow-based models.Source: §3.1 · p. 3, §3.3 · p. 4
• Maximizing the ratio increases modeled likelihood while reducing the KL gap between the ELBO and true log-likelihood.The likelihood component favors positive-advantage actions, while the correction term tightens the approximation.Source: §3.3 · p. 4

3.4 Estimating the FPO Ratio with Flow Matching

• The ratio estimator averages Monte Carlo draws of flow timestep and noise, recovering the exact FPO ratio in the limit while using only a few draws in practice.The sample count Nmc controls learning efficiency; with small Nmc, the ratio can be biased upward, although gradient estimates remain directionally unbiased.Source: §3.4 · p. 5

3.5 Denoising MDP Comparison

• Unlike denoising-MDP approaches, FPO treats rollout sampling as a black box, avoiding denoising-step horizon expansion, initial-noise observations, and dependence on stochastic sampling procedures.Prior denoising-MDP formulations increase credit-assignment difficulty and learning dimensionality, while restricting sampling procedures by construction.Source: §3.5 · p. 5

Experimental scope

• FPO is evaluated in GridWorld, 10 MuJoCo Playground environments, and physics-based humanoid control spanning varied dimensionality, sparsity, horizons, and simulators.The experiments use Gymnasium, MuJoCo Playground, and Isaac Gym.Source: §4 · p. 6

GridWorld

• At GridWorld saddle points, FPO transforms an initial Gaussian into a bimodal action distribution and produces varied trajectories reaching different goals.Trajectories from the same starting state can reach different goals and follow different paths.Source: §4.1 · p. 7
• Gaussian PPO also reaches the goals but behaves more deterministically, favoring the nearest goal with less trajectory variation than FPO.This comparison is reported in the supplemental material.Source: §4.1 · p. 7

MuJoCo Playground

• 8 of 10 MuJoCo Playground tasks are outperformed by FPO relative to both Gaussian PPO and DPPO.The comparison uses five seeds per task and 60M environment steps.Source: §4.2 · pp. 7–8, Figure 2 · p. 7, Figure 3 · p. 7
• Sampling choices and clipping materially affect FPO performance: more sampled (τ, ϵ) pairs generally help, ϵ-MSE outperforms u-MSE, and εclip selection is important.The reported FPO configuration uses 8 (τ, ϵ) pairs, ϵ-MSE, and εclip = 0.05.Source: §4.2 · p. 8, Table 1 · p. 8

Humanoid control

• In humanoid control, FPO matches Gaussian PPO with full-joint conditioning but surpasses it for root or root+hands goals, tracking motions while Gaussian PPO drifts or falls.Terrain-randomized FPO also walks stably across procedurally generated rough ground.Source: §4.3 · p. 9, Figure 4 · p. 9

• FPO trains flow-based generative models with policy gradients while preserving compatibility with standard flow-based mechanisms.The method supports sampling, distillation, and fine-tuning after training.Source: §5 · pp. 9–10
• Flow-based policy training and deployment are generally more computationally intensive than corresponding Gaussian policies.Source: §5 · p. 10
• FPO lacks established machinery for KL divergence estimation used in adaptive learning rates and entropy regularization.Source: §5 · p. 10
• Fine-tuning a pre-trained image diffusion model with reinforcement learning was unstable in practice, with sensitivity to classifier-free guidance compounded by self-generated data.The text attributes this instability to a broader challenge in reinforcement learning for image generation rather than to FPO itself.Source: §5 · p. 10
• Despite these limitations, FPO offers a simple and flexible bridge to online reinforcement learning and may benefit fine-tuning of pretrained flow-based policies.Source: §5 · p. 10

• The derivation develops the connection between the conditional flow matching objective and the ELBO, including the FPO ratio.Source: §A.1 · p. 18
• FPO’s general result is that it optimizes the advantage-weighted expected ELBO of noise-perturbed data for monotonic weighting functions.Source: §A.1 · p. 18
• Advantage weighting guides probability flow toward action neighborhoods associated with higher reward.Source: §A.1 · p. 18
• For uniform weighting w(λτ) = 1, the weighted loss directly corresponds to maximizing the ELBO of clean actions.Source: §A.1 · p. 18

• Popular flow-based and diffusion models can be reparameterized as instances of a weighted denoising loss.Source: §A.1.1 · p. 18
• The function w(λτ) determines the relative importance assigned to different noise levels.Source: §A.1.1 · p. 18
• For loss weights varying monotonically with noise timestep τ, the weighted loss has the stated relationship with expected ELBO.Source: §A.1.1 · p. 18
• The framework includes standard diffusion with uniform weighting, optimal transport linear interpolation, and cosine-schedule velocity prediction.The latter two schedules yield w(λτ) = e−λ/2.Source: §A.1.1 · p. 19

• FPO derives its general flow-matching formulation by replacing direct likelihood maximization with maximization of the ELBO for noise-perturbed data.Source: §A.1.2 · p. 19
• The policy ratio compares the relative likelihood of actions and their noisy versions under the current and old policies.Source: §A.1.2 · p. 19
• The ELBO-equivalence constant is independent of model parameters and cancels when forming the FPO ratio within one trust-region episode.Source: §A.1.2 · p. 19
• FPO estimates the ratio through Monte Carlo sampling of diffusion or flow timesteps and noise.Source: §A.1.2 · p. 19
• The ratio estimate overestimates scale but remains unbiased in gradient direction, while additional timestep and noise samples can reduce or eliminate scale bias.Source: §A.1.2 · p. 19

ELBO connection

• Uniform weighting makes the weighted flow-matching loss correspond directly to the negative ELBO of clean data plus a parameter-independent constant.This stronger result applies to standard diffusion schedules with w(λt) = 1.Source: §A.1.3 · pp. 19–20
• Minimizing the unweighted loss is equivalent to maximizing the ELBO of the clean action at.This provides a direct connection to traditional maximum likelihood estimation.Source: §A.1.3 · p. 20

Advantage weighting

• FPO’s advantage-weighted objective increases a lower bound on the log-likelihood of clean actions weighted by their advantages.The special case therefore weights clean actions directly rather than noise-perturbed actions.Source: §A.1.3 · p. 20

Practical implication

• FPO remains computationally efficient by avoiding explicit likelihood computations despite its close relationship to PPO-style maximum likelihood methods.The FPO ratio in this special case preserves the computational advantage of the general method.Source: §A.1.3 · p. 20
• The Monte Carlo estimator has upward bias in gradient scale, so PPO clipping controls the magnitude of parameter changes.This is a limitation of the estimator rather than a failure of the ELBO connection.Source: §A.1.3 · p. 20

Advantage weighting

• Advantages must remain positive because flow matching learns nonnegative probability flows and uses advantages as loss weights.A constant shift preserves policy-gradient optimization through baseline invariance.Source: §A.1.4 · p. 20

GridWorld behavior

• Gaussian PPO reaches goals in GridWorld but produces less diverse trajectories, with shown initializations converging to a single goal mode.The comparison uses test-time noise perturbations of σ = 0.0, 0.1, and 0.5.Source: §A.2 · pp. 20–21

Training setup

• Flow-matching and diffusion policies use the same MuJoCo training hyperparameters, with adjusted clipping, disabled entropy regularization, and method-specific settings.Humanoid-control training also reports a shared hyperparameter set for Gaussian PPO and diffusion FPO.Source: §A.3 · p. 21, §A.4 · p. 21, Table A.2 · p. 22

Image fine-tuning

• Fine-tuning image diffusion models on self-generated outputs is highly unstable, even with uniform advantages, because CFG artifacts amplify across RL epochs.The authors report quality regression at CFG scales of 2 and 4.Source: §A.5 · pp. 22–23, Figure A.3 · p. 23