Source-linked AI summary
Flow Matching Policy Gradients
David McAllister, Songwei Ge, Brent Yi, Chung Min Kim, Ethan Weber, Hongsuk Choi, Haiwen Feng, Angjoo Kanazawa
TL;DR
Flow-based policies offer expressive distributions, but policy-gradient training traditionally relies on likelihood-based objectives and Gaussian policies. FPO uses an advantage-weighted ratio from conditional flow matching as a PPO-compatible surrogate, avoiding exact likelihood computation while preserving flow-based sampling flexibility. It trains flow policies from scratch across control tasks, including under-conditioned humanoid control where Gaussian policies fail to learn.
Problem
Policy-gradient reinforcement learning commonly uses Gaussian policies, while direct likelihood optimization for flow models requires computationally prohibitive divergence estimation.
Method
FPO optimizes flow-based policies with an advantage-weighted ratio computed from the conditional flow matching objective in a PPO-inspired surrogate.
Results
FPO trains flow-based policies from scratch across diverse control tasks, learns multimodal actions, and enables under-conditioned humanoid training where Gaussian policies fail to learn.
Takeaways & Limitations
FPO provides a simple, flexible bridge between flow-based generative models and online reinforcement learning while preserving standard flow-based sampling, distillation, and fine-tuning.
Takeaways & Limitations
Flow-based policy training and deployment are generally more computationally intensive than for Gaussian policies, and FPO lacks established KL-estimation and entropy-regularization machinery.
Abstract
from arXiv · showhide
Flow-based generative models, including diffusion models, excel at modeling continuous distributions in high-dimensional spaces. In this work, we introduce Flow Policy Optimization (FPO), a simple on-policy reinforcement learning algorithm that brings flow matching into the policy gradient framework. FPO casts policy optimization as maximizing an advantage-weighted ratio computed from the conditional flow matching loss, in a manner compatible with the popular PPO-clip framework. It sidesteps the need for exact likelihood computation while preserving the generative capabilities of flow-based models. Unlike prior approaches for diffusion-based reinforcement learning that bind training to a specific sampling method, FPO is agnostic to the choice of diffusion or flow integration at both training and inference time. We show that FPO can train diffusion-style policies from scratch in a variety of continuous control tasks. We find that flow-based models can capture multimodal action distributions and achieve higher performance than Gaussian policies, particularly in under-conditioned settings.
1 Introduction
FPO brings flow-based generative policies into on-policy policy gradients using a flow-matching-based surrogate rather than explicit likelihoods. Across control settings, it supports flexible sampling, multimodal actions, and training in challenging under-conditioned tasks.
- 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.
- FPO: 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.
- 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.
- Empirical findings: 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.
2 Related Work
Flow matching provides a tractable route to modeling complex continuous distributions and can be incorporated into reinforcement learning without expensive likelihood computation. FPO extends policy-gradient training to richer flow-based policies while retaining compatibility with standard actor-critic methods.
- 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.
- 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.
- 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.
- Diffusion Models: 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.
- 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.
- 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.
3 Flow Matching Policy Gradients
FPO replaces PPO’s exact likelihood ratio with an advantage-weighted ratio derived from conditional flow-matching losses, enabling policy-gradient training of flow-based policies. Its ELBO-based formulation preserves compatibility with PPO-style updates while supporting flexible flow or diffusion implementations and practical Monte Carlo estimation.
- 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.
- 3.2 Flow Policy Optimization: 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.
- 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.
- 3.3 FPO Surrogate Objective: 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.
- 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.
- 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.
4 Experiments
Across GridWorld, continuous control, and humanoid tasks, FPO trains flow-based policies from scratch and models diverse actions. It outperforms Gaussian PPO and DPPO on most MuJoCo tasks and is especially robust under sparse humanoid conditioning.
- 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.
- 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.
- GridWorld: 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.
- 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.
- MuJoCo Playground: 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.
- 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.
5 Discussion and Limitations
FPO provides a simple bridge between flow-based generative models and online reinforcement learning, but flow-based policies are generally more computationally intensive and some applications remain unstable.
- 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.
- Flow-based policy training and deployment are generally more computationally intensive than corresponding Gaussian policies.
- FPO lacks established machinery for KL divergence estimation used in adaptive learning rates and entropy regularization.
- 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.
- Despite these limitations, FPO offers a simple and flexible bridge to online reinforcement learning and may benefit fine-tuning of pretrained flow-based policies.
A.1 FPO Derivation
The derivation connects advantage-weighted conditional flow matching losses to expected ELBO optimization, showing how FPO increases the likelihood of high-reward actions and their noisy counterparts.
- The derivation develops the connection between the conditional flow matching objective and the ELBO, including the FPO ratio.
- FPO’s general result is that it optimizes the advantage-weighted expected ELBO of noise-perturbed data for monotonic weighting functions.
- Advantage weighting guides probability flow toward action neighborhoods associated with higher reward.
- For uniform weighting w(λτ) = 1, the weighted loss directly corresponds to maximizing the ELBO of clean actions.
A.1.1 Loss Weighting Choices
Flow-matching losses can be viewed through a unified weighted denoising framework, where monotonic noise-level weightings preserve the expected-ELBO relationship.
- Popular flow-based and diffusion models can be reparameterized as instances of a weighted denoising loss.
- The function w(λτ) determines the relative importance assigned to different noise levels.
- For loss weights varying monotonically with noise timestep τ, the weighted loss has the stated relationship with expected ELBO.
- 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.
A.1.2 Flow Matching as Expected ELBO Optimization
FPO replaces direct likelihood maximization with ELBO-based optimization for noise-perturbed data and forms a current-to-old policy ratio from expected ELBOs estimated by Monte Carlo sampling.
- FPO derives its general flow-matching formulation by replacing direct likelihood maximization with maximization of the ELBO for noise-perturbed data.
- The policy ratio compares the relative likelihood of actions and their noisy versions under the current and old policies.
- The ELBO-equivalence constant is independent of model parameters and cancels when forming the FPO ratio within one trust-region episode.
- FPO estimates the ratio through Monte Carlo sampling of diffusion or flow timesteps and noise.
- The ratio estimate overestimates scale but remains unbiased in gradient direction, while additional timestep and noise samples can reduce or eliminate scale bias.
A.1.3 FPO with Diffusion Schedules
Under uniform diffusion weighting, FPO’s unweighted flow-matching loss is equivalent to maximizing the ELBO of clean actions, connecting the method to maximum likelihood while retaining likelihood-free computation.
- 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.
- ELBO connection: 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.
- 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.
- 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.
- Practical implication: 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.
A.1.4 Advantage-Weighed Flow Matching Discussion
The discussion addresses positive advantage weighting, Gaussian-policy diversity, PPO tuning, and image fine-tuning. It reports richer diffusion-policy behavior in GridWorld but instability when training image diffusion models on their own outputs.
- 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.
- 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.
- 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.
- 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.