Reverse Flow Matching: A Unified Framework for Online Reinforcement Learning with Diffusion and Flow Policies
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 that distinguishes online RL from standard generative modeling is the lack of direct samples from the target 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 uses a weighted average of noise as the training target, and a gradient-expectation family, which employs a weighted average of Q-function gradients. However, it remains unclear how these objectives are formally related, or whether 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 share the same expectation. We show that existing noise-expectation and gradient-expectation methods are simply 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 it enables the principled combination of Q-value and Q-gradient information to form an effective 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.
Lay Summary
How can an AI system learn to make good decisions when no expert can show it the right answer? We study this challenge for generative decision-making models: systems that generate possible actions, much like image generators create possible images, but must improve from feedback on their own attempts rather than from demonstrations to imitate. We introduce Reverse Flow Matching, a trial-and-error learning method for these models. Instead of requiring a finished expert solution, our method learns from rough, noisy guesses about what a good solution might look like, together with scores that indicate how useful those guesses are. Our method also reveals that two existing approaches to this problem are actually two versions of the same underlying idea. This theoretical unification lets us combine their strengths, producing steadier learning signals that lead to better performance. Together, these findings contribute to the broader goal of embodied intelligence, by enabling generative AI systems to learn through action, feedback, and self-improvement when expert demonstrations are unavailable.