ReinFlow: Fine-tuning Flow Matching Policy with Online Reinforcement Learning

Tonghe Zhang, Chao Yu, Sichang Su, Yu Wang

Abstract:

We propose ReinFlow, a simple yet effective online reinforcement learning (RL)
framework that fine-tunes a family of flow matching policies for continuous robotic
control. Derived from rigorous RL theory, ReinFlow injects learnable noise into a
flow policy’s deterministic path, converting the flow into a discrete-time Markov
Process for exact and straightforward likelihood computation. This conversion
facilitates exploration and ensures training stability, enabling ReinFlow to fine-tune
diverse flow model variants stably, including Rectified Flow [34] and Shortcut
Models [18], particularly at very few or even one denoising step. We benchmark
ReinFlow in representative locomotion and manipulation tasks, including longhorizon planning with visual input and sparse reward. The episode reward of
Rectified Flow policies obtained an average net growth of 135.36% after fine-tuning
in challenging legged locomotion tasks while saving denoising steps and 82.63% of
wall time compared to state-of-the-art diffusion RL fine-tuning method DPPO [42].
The success rate of the Shortcut Model policies in state and visual manipulation
tasks achieved an average net increase of 40.34% after fine-tuning with ReinFlow
at four or even one denoising step, whose performance is comparable to fine-tuned
DDIM policies while saving computation time for an average of 23.20% . Code,
model, and checkpoints available on the project website: https://reinflow.github.io/