# Online Reward-Weighted Fine-Tuning of Flow Matching with Wasserstein Regularization

arXiv: https://arxiv.org/abs/2502.06061 · HTML: https://mlsanity.com/paper/2502.06061/

## Key facts

- Cited by 43 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 11 of its 46 references that are in the map.
- First posted on arXiv 9 Feb 2025.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Builds on (most cited first, 11 of 11)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis](https://mlsanity.com/paper/2306.09341/) (2023)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [Reinforced Self-Training (ReST) for Language Modeling](https://mlsanity.com/paper/2308.08998/) (2023)
- [DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models](https://mlsanity.com/paper/2305.16381/) (2023)
- [Aligning Text-to-Image Models using Human Feedback](https://mlsanity.com/paper/2302.12192/) (2023)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [Understanding the performance gap between online and offline alignment algorithms](https://mlsanity.com/paper/2405.08448/) (2024)
- [Sequence-Augmented SE(3)-Flow Matching For Conditional Protein Backbone Generation](https://mlsanity.com/paper/2405.20313/) (2024)

## Built on by (most cited first, 30 of 43)

- [ReinFlow: Fine-tuning Flow Matching Policy with Online Reinforcement Learning](https://mlsanity.com/paper/2505.22094/) (2025)
- [Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models](https://mlsanity.com/paper/2509.25050/) (2025)
- [Flow-Based Policy for Online Reinforcement Learning](https://mlsanity.com/paper/2506.12811/) (2025)
- [One-Step Flow Policy Mirror Descent](https://mlsanity.com/paper/2507.23675/) (2025)
- [Taming Preference Mode Collapse via Directional Decoupling Alignment in Diffusion Reinforcement Learning](https://mlsanity.com/paper/2512.24146/) (2025)
- [Reinforcement Learning for Flow-Matching Policies](https://mlsanity.com/paper/2507.15073/) (2025)
- [Value Gradient Guidance for Flow Matching Alignment](https://mlsanity.com/paper/2512.05116/) (2025)
- [DistDF: Time-Series Forecasting Needs Joint-Distribution Wasserstein Alignment](https://mlsanity.com/paper/2510.24574/) (2025)
- [UniGRPO: Unified Policy Optimization for Reasoning-Driven Visual Generation](https://mlsanity.com/paper/2603.23500/) (2026)
- [Reverse Flow Matching: A Unified Framework for Online Reinforcement Learning with Diffusion and Flow Policies](https://mlsanity.com/paper/2601.08136/) (2026)
- [Inference-Time Alignment Control for Diffusion Models with Reinforcement Learning Guidance](https://mlsanity.com/paper/2508.21016/) (2025)
- [HapticVLA: Contact-Rich Manipulation via Vision-Language-Action Model without Inference-Time Tactile Sensing](https://mlsanity.com/paper/2603.15257/) (2026)
- [Reinforcing Few-step Generators via Reward-Tilted Distribution Matching](https://mlsanity.com/paper/2605.26108/) (2026)
- [PaCo-RL: Advancing Reinforcement Learning for Consistent Image Generation with Pairwise Reward Modeling](https://mlsanity.com/paper/2512.04784/) (2025)
- [ProteinZero: Self-Improving Protein Generation via Online Reinforcement Learning](https://mlsanity.com/paper/2506.07459/) (2025)
- [Adaptive Divergence Regularized Policy Optimization for Fine-tuning Generative Models](https://mlsanity.com/paper/2510.18053/) (2025)
- [π-StepNFT: Wider Space Needs Finer Steps in Online RL for Flow-based VLAs](https://mlsanity.com/paper/2603.02083/) (2026)
- [Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning](https://mlsanity.com/paper/2606.11087/) (2026)
- [Variational Supervised Contrastive Learning](https://mlsanity.com/paper/2506.07413/) (2025)
- [Finite Difference Flow Optimization for RL Post-Training of Text-to-Image Models](https://mlsanity.com/paper/2603.12893/) (2026)
- [Fine-tuning Flow Matching Generative Models with Intermediate Feedback](https://mlsanity.com/paper/2510.18072/) (2025)
- [FM-IRL: Flow-Matching for Reward Modeling and Policy Regularization in Reinforcement Learning](https://mlsanity.com/paper/2510.09222/) (2025)
- [GoldenStart: Q-Guided Priors and Entropy Control for Distilling Flow Policies](https://mlsanity.com/paper/2603.14245/) (2026)
- [FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling](https://mlsanity.com/paper/2604.06916/) (2026)
- [Boosting Maximum Entropy Reinforcement Learning via One-Step Flow Matching](https://mlsanity.com/paper/2602.01606/) (2026)
- [Improving Text-to-Image Generation with Intrinsic Self-Confidence Rewards](https://mlsanity.com/paper/2603.00918/) (2026)
- [Flow Matching with Injected Noise for Offline-to-Online Reinforcement Learning](https://mlsanity.com/paper/2602.18117/) (2026)
- [AdvantageFlow: Advantage-Weighted Least Squares for RL in Flow Models](https://mlsanity.com/paper/2605.26013/) (2026)
- [GROW: Group-Relative Advantage-Weighted On-Policy Reinforcement Learning of Autoregressive-Diffusion Text-to-Speech model](https://mlsanity.com/paper/2608.03215/) (2026)
- [SpeechGym: An Audio-Native Gym for Training Voice Agents via Reinforcement Learning](https://mlsanity.com/paper/2608.26432/) (2026)

## Abstract (arXiv)

Recent advancements in reinforcement learning (RL) have achieved great success in fine-tuning diffusion-based generative models. However, fine-tuning continuous flow-based generative models to align with arbitrary user-defined reward functions remains challenging, particularly due to issues such as policy collapse from overoptimization and the prohibitively high computational cost of likelihoods in continuous-time flows. In this paper, we propose an easy-to-use and theoretically sound RL fine-tuning method, which we term Online Reward-Weighted Conditional Flow Matching with Wasserstein-2 Regularization (ORW-CFM-W2). Our method integrates RL into the flow matching framework to fine-tune generative models with arbitrary reward functions, without relying on gradients of rewards or filtered datasets. By introducing an online reward-weighting mechanism, our approach guides the model to prioritize high-reward regions in the data manifold. To prevent policy collapse and maintain diversity, we incorporate Wasserstein-2 (W2) distance regularization into our method and derive a tractable upper bound for it in flow matching, effectively balancing exploration and exploitation of policy optimization. We provide theoretical analyses to demonstrate the convergence properties and induced data distributions of our method, establishing connections with traditional RL algorithms featuring Kullback-Leibler (KL) regularization and offering a more comprehensive understanding of the underlying mechanisms and learning behavior of our approach. Extensive experiments on tasks including target image generation, image compression, and text-image alignment demonstrate the effectiveness of our method, where our method achieves optimal policy convergence while allowing controllable trade-offs between reward maximization and diversity preservation.