Reverse Flow Matching: A Unified Framework for Online Reinforcement Learning with Diffusion and Flow Policies
Zeyang Li, Sunbochen Tang, Navid Azizan
TL;DR
This work addresses the challenge of online RL with diffusion and flow policies when the target policy is a Boltzmann distribution defined by the Q-function. It introduces Reverse Flow Matching (RFM), a reverse-inference framework that recasts learning as posterior mean estimation and leverages Langevin Stein-based control variates to reduce variance, unifying noise- and gradient-based approaches. RFM extends Boltzmann targeting to flow policies and provides a principled way to combine Q-values and Q-gradients to achieve minimum-variance estimators, improving training efficiency and stability. The method is instantiated for online RL with flow policies and shows superior, more stable performance across continuous-control benchmarks compared to diffusion baselines and related methods. Overall, RFM offers a versatile, variance-reducing framework that can enhance sample efficiency and stability for expressive policy classes in online reinforcement learning.
Abstract
Diffusion and flow policies are gaining prominence in online reinforcement learning (RL) due to their expressive power, yet training them efficiently remains a critical challenge. A fundamental difficulty in online RL is the lack of direct samples from the target distribution; instead, the target is an unnormalized Boltzmann distribution defined by the Q-function. To address this, two seemingly distinct families of methods have been proposed for diffusion policies: a noise-expectation family, which utilizes a weighted average of noise as the training target, and a gradient-expectation family, which employs a weighted average of Q-function gradients. Yet, it remains unclear how these objectives relate formally or if they can be synthesized into a more general formulation. In this paper, we propose a unified framework, reverse flow matching (RFM), which rigorously addresses the problem of training diffusion and flow models without direct target samples. By adopting a reverse inferential perspective, we formulate the training target as a posterior mean estimation problem given an intermediate noisy sample. Crucially, we introduce Langevin Stein operators to construct zero-mean control variates, deriving a general class of estimators that effectively reduce importance sampling variance. We show that existing noise-expectation and gradient-expectation methods are two specific instances within this broader class. This unified view yields two key advancements: it extends the capability of targeting Boltzmann distributions from diffusion to flow policies, and enables the principled combination of Q-value and Q-gradient information to derive an optimal, minimum-variance estimator, thereby improving training efficiency and stability. We instantiate RFM to train a flow policy in online RL, and demonstrate improved performance on continuous-control benchmarks compared to diffusion policy baselines.