PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Tonghe Zhang, Chao Yu, Sichang Su, Yu Wang

arXiv:2505.22094v7cs.ROcs.LG

TL;DR

Imperfect demonstrations and limited exploration make online RL fine-tuning of flow policies difficult, particularly when few denoising steps complicate likelihood computation. ReinFlow injects learnable noise to form a discrete-time Markov process with exact likelihoods and policy-gradient training. Across locomotion and manipulation tasks, it improves rewards and success rates while reducing wall-clock time, but sample efficiency, real-world adaptation, noise sensitivity and scaling remain open challenges.

  • Problem

    Flow policies trained on mixed-quality demonstrations lack built-in exploration, while their likelihoods are difficult to compute stably with few denoising steps.

  • Method

    ReinFlow injects learnable noise into flow trajectories, converting them into discrete-time Markov processes with exact likelihoods for online policy-gradient optimization.

  • Results

    ReinFlow improves manipulation success by 40.34% and locomotion reward by 135.36% on average, with a 62.82% wall-time reduction versus DPPO.

  • Takeaways & Limitations

    The framework fine-tunes diverse flow policies for continuous robotic control, including very few or one denoising step, while reducing wall-clock time by over 50% versus diffusion RL.

  • Takeaways & Limitations

    The current implementation prioritizes wall-time efficiency over sample efficiency, and future work must address real-world RL, noise sensitivity and scaling to large VLA models.

Abstract

from arXiv · show

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, including Rectified Flow [35] and Shortcut Models [19], particularly at very few or even one denoising step. We benchmark ReinFlow in representative locomotion and manipulation tasks, including long-horizon 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 [43]. 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%. Project webpage: https://reinflow.github.io/

1 Introduction

ReinFlow addresses the difficulty of online RL fine-tuning for flow policies by injecting learnable noise and enabling tractable likelihood computation. It improves locomotion rewards and manipulation success while reducing wall time.

  • Mixed-quality demonstrations and limited exploration can leave flow policies unable to surpass expert performance on challenging tasks.
  • Online RL can address imperfect demonstrations, but flow policies remain technically challenging to fine-tune because their likelihoods are difficult to compute stably.
  • ReinFlow injects learnable noise to convert flow policies into discrete-time Markov processes with Gaussian transitions and exact likelihood computation.
  • 135.36% average reward increase was achieved for locomotion policies, while manipulation success rate improved by 40.34%.
  • 62.82% wall-time reduction was achieved across tasks compared to DPPO.
  • ReinFlow’s design choices were examined through sensitivity analysis covering pretrained-data scale, denoising steps, noise conditioning and regularization.

2 Related Work

Prior work has adapted RL to diffusion policies and begun studying flow models, but existing approaches differ in exploration, task domain, or optimization framework. ReinFlow targets online RL for continuous robotic control with learnable noise and exact likelihoods.

  • Diffusion-policy methods use Q-learning, action gradients, or policy gradients for online or adaptable RL fine-tuning.
  • Directly transferring diffusion-policy RL methods to flow policies is difficult because the two model families have different mathematical structures.
  • Flow Q-Learning lacks exploration during online adaptation, whereas ReinFlow injects bounded, learnable noise into flow trajectories.
  • Flow-GRPO and ORW-CFM-W2 study online RL for flow models in vision tasks, while ReinFlow addresses continuous robotic control with policy gradients and exact likelihood computation.

3 Problem Formulation

The paper formulates robot learning as a continuous-state, continuous-action POMDP and represents action generation with flow matching dynamics. Fewer denoising steps accelerate inference but increase discretization error.

  • Robot learning is modeled as an infinite-horizon POMDP with continuous states, actions and observations, discounted rewards, unknown transitions and emissions.
  • Flow matching transports samples from p0 to p1 through a time-dependent mapping governed by a neural ODE.
  • Rectified Flow uses a straight path Xt = tX1 +(1−t)X0 with constant velocity X1 −X0.
  • Shortcut Models align two-step and one-step velocities to improve generation quality at very few denoising steps.
  • In robot action generation, the velocity field conditions on observations and produces denoised actions from the action space.
  • Reducing denoising steps speeds flow-policy inference, but larger integration steps increase discretization error and reduce quality.

4 Algorithm Design

ReinFlow turns deterministic flow trajectories into noise-injected discrete-time Markov processes, yielding exact likelihoods for policy-gradient optimization even with few denoising steps. Joint training of velocity and noise networks supports exploration and regularization.

  • Likelihood Computation over a Short Denoising Trajectory: Numerical likelihood estimation suffers Monte Carlo and discretization errors, especially with few denoising steps, while deterministic transitions make discrete-process probabilities impossible to compute.
  • Likelihood Computation over a Short Denoising Trajectory: ReinFlow injects learnable, observation- and time-conditioned Gaussian noise into each flow transition, creating closed-form Markov-process transitions.
  • Likelihood Computation over a Short Denoising Trajectory: The resulting joint log probability is exact for arbitrarily large step sizes, avoiding trace estimation and supporting stable one-step or few-step fine-tuning.
  • Policy Optimization: A policy-gradient theorem for discrete-time Markov-process policies connects the joint transition likelihood to optimization of the final executed action.
  • Policy Optimization: ReinFlow combines the exact likelihood with a clipped-surrogate policy-gradient objective and jointly optimizes the velocity and noise networks.
  • Exploration and Regularization: The noise network increases trajectory diversity during training and is discarded afterward, restoring the deterministic flow-matching policy.
  • Exploration and Regularization: Minimizing the negative per-symbol entropy-rate regularizer promotes more diverse actions and enhances exploration.

5 Experiments

ReinFlow is evaluated on locomotion and manipulation tasks with state, pixel, and sparse-reward settings, including very few-step fine-tuning. Across benchmarks, it improves performance and efficiency while using PPO-based online RL.

  • Experimental setup: ReinFlow fine-tunes Rectified Flow and Shortcut Model policies with very few or even one denoising step.The two policy variants share the same fine-tuning hyperparameters.
  • Implementation scope: PPO is used as the ReinFlow implementation because of its stability, prioritizing wall-time efficiency over sample cost in simulated environments.More sample-efficient RL algorithms, especially for expensive real-world data collection, remain future work.
  • Experimental setup: ReinFlow is evaluated in locomotion and manipulation tasks with state or pixel observations and, for manipulation, sparse rewards.The benchmarks include OpenAI Gym, Franka Kitchen long-horizon planning, and Robomimic visual manipulation.
  • Overall results: ReinFlow demonstrates strong training stability and significant success-rate increases across tasks, with the best overall efficiency and performance improvement in Gym and Franka Kitchen.The comparison includes diffusion- and flow-policy fine-tuning methods, particularly DPPO and FQL.
  • Overall results: 45.77% average improvement in success rate is reported across three Robomimic visual manipulation tasks.ReinFlow achieves success rates comparable to DPPO while using one denoising step for Can and Square and four for Transport, versus five-step DDIM.

6 The Design Choice and Key Factors Affecting ReinFlow

The study examines how data scale, denoising steps, noise conditioning, noise magnitude, and regularization affect ReinFlow. Fine-tuning gains persist across scales and denoising choices, while exploration benefits depend on noise and regularization choices.

  • Scaling: Increasing pretraining data or inference steps does not consistently improve reward, whereas ReinFlow improves reward or success rate across pretraining scales.The results support RL as an additional scaling path beyond data quantity or inference computation.
  • Scaling: The gain from RL remains invariant to denoising steps at four steps in Hopper and one step in Square.The figure presents RL as an orthogonal scaling path beyond data or inference.
  • Noise-network inputs: Conditioning the noise network on both observations and time often yields higher success rates than conditioning only on observations.Time conditioning lets the noise network alter noise intensity across denoising steps.
  • Noise level and exploration: Moderate noise enables rapid improvement, reaching up to three times higher rewards than small-noise settings, while performance becomes less sensitive beyond that threshold.Reducing noise benefits visuomotor policies, precision-critical tasks, longer denoising chains, and weakly pretrained models.
  • Regularization and exploration: Entropy regularization is generally more effective than W2 regularization in locomotion, and lowering the W2 coefficient can help surpass the behavior-cloning baseline.The comparison links stronger exploration behavior to the regularization choice.

7 Conclusion, Limitations, and Future Work

ReinFlow fine-tunes flow matching policies for continuous robotic control, surpassing existing RL fine-tuning methods while reducing wall-clock time. The authors identify implementation efficiency, sample efficiency, real-world adaptation, hyperparameter sensitivity, and scaling to large VLA models as future challenges.

  • Conclusion: ReinFlow stably fine-tunes a family of flow matching policies for continuous robotic control and surpasses existing diffusion- or flow-based RL fine-tuning methods.It reduces wall-clock time by over 50% compared to the state-of-the-art diffusion RL algorithm.
  • Conclusion: The authors conducted sensitivity analysis to identify key factors affecting ReinFlow’s performance.
  • Limitations and Future Work: The current on-policy implementation leaves sample-efficient training and adaptation to real-world RL for future work.The authors note that parallelism saves wall-time, but the implementation is not yet sample-efficient or adapted to real-world RL.
  • Limitations and Future Work: ReinFlow remains sensitive to noise magnitude, motivating automatic tuning or removal of these hyperparameters.
  • Limitations and Future Work: Experiments use relatively small networks, leaving scaling to large flow-based vision-language-action models as an open challenge.

NeurIPS Paper Checklist

The NeurIPS Paper Checklist documents responsible research practices, including assumptions, reproducibility, statistical reporting, and societal impact. It is a required submission component reviewed alongside the paper and published with the final version.

  • The checklist addresses reproducibility, transparency, research ethics, and societal impact, and papers omitting it will be desk rejected.It follows the references and optional supplemental material and does not count toward the page limit.
  • Checklist answers must be Yes, No, or NA, with a short justification provided immediately after each answer.NA means the question is not applicable or the relevant information is unavailable.
  • The checklist is visible to reviewers and ethics reviewers, is used as one evaluation factor, and is included with the published final paper.A No or NA answer is acceptable when properly justified and is generally not grounds for rejection.
  • Authors should discuss limitations, strong assumptions, robustness to assumption violations, and potential societal impacts such as malicious use, unfairness, or privacy risks.The checklist specifically encourages a separate Limitations section and requires explanations when authors answer NA or No for societal impact.
  • Theory assumptions and proofs: Theoretical results should state their assumptions and provide complete, correct proofs, with theorems, formulas, and proofs numbered and cross-referenced.Proofs may appear in the main paper or supplemental material, with a proof sketch encouraged when they appear only in the supplement.
  • Experimental result reproducibility: Experimental papers should disclose information needed for reproducibility and report error bars, confidence intervals, or statistical significance tests for experiments supporting their main claims.The sources of variability represented by error bars should also be identified.

A.1 Proof of Theorem 4.1

This section develops the policy-gradient result for discrete-time Markov-process policies in POMDPs and connects it to ReinFlow’s flow-policy formulation and action chunking.

  • Policy-gradient theorem: The proof targets a policy-gradient theorem for reactive policies represented as discrete-time Markov processes.The derivation extends policy-gradient results from MDPs to POMDPs.
  • Policy-gradient theorem: The RL objective uses discounted visitation and value, Q, and advantage functions to express policy gradients through action log-probability gradients.The proof introduces the observation visitation measure and substitutes the advantage function into the gradient expression.
  • ReinFlow specialization: For ReinFlow, general policy parameters combine the velocity-net parameters θ and noise-net parameters θ′.The combined parameter vector is written as ¯θ = [θ, θ′].
  • Assumptions: The derivation assumes stationary transition kernels and a reactive, stationary policy, making Aπ, Qπ, and Vπ time independent.The theorem is explicitly limited to stationary and reactive policies.
  • Action chunking: With action chunking, actions fixed after the initial observation are conditionally independent within a chunk, so chunk log probability sums internal action log probabilities.The implementation flattens each chunk into a single actor output while preserving conditional independence given the initial observation and network inputs.
  • Comparison with prior methods: ReinFlow avoids DPPO’s denoised-step advantage calculations, while FQL instead distills a one-step policy from a multi-step flow using offline Q-learning.The comparison motivates ReinFlow’s direct online fine-tuning of multi-step flow policies.

C Environment and Dataset Configuration

The experiments cover four locomotion tasks and four manipulation tasks across state and pixel inputs, using specified environment configurations and dataset choices aligned with prior work.

  • Tasks: The evaluated locomotion tasks are Hopper, Walker2d, Ant, and Humanoid, listed in ascending difficulty.The tasks are OpenAI Gym environments used for flow-policy fine-tuning.
  • Tasks: The manipulation evaluation includes state-input Franka Kitchen and pixel-input Robomimic environments.Figures 8 and 9 summarize the locomotion and manipulation task sets.
  • Environment configuration: The environment configuration records proprioception dimension, action and chunk dimensions, maximum rollout length, and reward sparsity.Sparse rewards assign +1 only upon task completion and otherwise 0, directly matching manipulation success.
  • Dataset configuration: Robomimic experiments follow DPPO’s configuration, whose supplied datasets are described as lower quality and/or quantity than the official datasets.The authors note that the simplified configuration reduces state-vector size and dataset size.
  • Dataset configuration: Gym experiments use the behavior-cloning dataset provided by Ren et al. to replicate DPPO, while FQL comparisons use D4RL offline RL datasets.The DPPO Gym dataset lacks offline rewards and its construction was not specified.
  • Dataset configuration: Minor differences in the OpenAI Gym task setup produced minimal discrepancies in reward curves.This is the reported consequence of the dataset and environment differences described for the comparisons.

D.1 Model Architecture

The architecture combines velocity, critic, and noise-injection networks across state and pixel inputs, with training and inference choices designed for stable flow-policy fine-tuning.

  • State-input tasks: State-input velocity networks use MLPs conditioned on the action chunk, state, and denoising time.Time is encoded with sinusoidal positional embeddings and linear projections with Mish activation.
  • Pixel-input tasks: Pixel-input actors and critics use a single-layer Visual Transformer with random-shift augmentation and an MLP for proprioception.The actor receives visual features, time embeddings, and raw action chunks; the critic receives time and condition features.
  • Noise injection: The noise-injection network shares actor inputs and outputs bounded per-coordinate action-noise standard deviations through Tanh and affine transformation.The bounds σmax and σmin are important ReinFlow hyperparameters.
  • Noise injection: After fine-tuning, ReinFlow discards the noise network and returns a policy with an ODE inference path; evaluation omits injected noise.The authors observed that the no-noise evaluation version often achieved higher reward.
  • Parameter scale: 135.36% net reward increase accompanied a 31.29% net success-rate increase with less than 6% additional parameters for state-input tasks.The added parameters come from the noise-injection network.
  • Parameter scale: The authors leave efficient noise-injection architectures that balance reward improvement and parameter count for future work.Shared representations are expected to keep scaling overhead minimal for larger visual backbones, but architecture design remains open.
  • Training stability: Critic warm-up is crucial for stable training and rapid convergence, especially for larger visual-input models.Excessive warm-up can overfit the critic and cause violent policy-gradient oscillations.

E Additional Experimental Results

Additional experiments compare ReinFlow with diffusion and flow-based baselines across locomotion and manipulation settings, while examining denoising steps, exploration, regularization, and efficiency. ReinFlow improves performance across tasks, but longer denoising trajectories can plateau and increase simulation time.

  • Performance Increase: 135.36% average episode reward net increase was achieved in OpenAI Gym locomotion tasks using D4RL datasets, while manipulation success rate increased 40.34% overall.The manipulation increases were 31.29% in Franka Kitchen and 45.77% in Robomimic.
  • Sample Complexity and Wall-clock Time: FQL is more sample-efficient in simpler tasks but struggles in complex locomotion, where it is asymptotically inferior to DPPO and ReinFlow.Its shorter per-iteration time does not imply greater overall efficiency because it uses smaller batches, more iterations, and lacks comparable parallel computing.
  • Comparison with Other Diffusion RL Methods: ReinFlow generally outperforms other methods in asymptotic performance and stability across random seeds in representative locomotion tasks.The comparison includes diffusion RL baselines, with figures covering Ant-v0, Hopper-v2, and Walker2d-v2.
  • Changing the Number of Fine-tuned Denoising Steps: Increasing denoising steps improves initial reward, but success rate or episode reward quickly plateaus while longer trajectories consume more simulation time.This trade-off is shown for Shortcut Policies in Franka Kitchen at K = 1, 2, and 4 denoising steps.
  • Noise Level and Exploration: Moderate noise enables rapid improvement up to three times higher rewards, whereas small noise limits exploration and excessive noise offers diminishing sensitivity.Reducing noise is beneficial for visuomotor policies, precision-critical tasks, longer denoising chains, and weakly pretrained models.
  • Regularization and Exploration: Entropy regularization improves sparse-reward Franka Kitchen success from 96.17 ± 3.65% to 99.00 ± 0.75% with α = 0.1.Noise injection makes likelihood tractable and aids exploration, while entropy regularization provides explicit control through its coefficient.

F.3 Additional Ablation Studies

Additional ablations test noise conditioning and noise scale across environments, alongside implementation and hyperparameter details. Conditioning noise on both state and time yields higher rewards, while larger noise levels promote exploration across varied settings.

  • Noise conditioning across environments: Conditioning the noise network on both state and time yields higher rewards than conditioning only on state in Humanoid-v3.The comparison uses three random seeds and evaluates σθ′(s) against σθ′(s, t).
  • Noise scale in sparse-reward manipulation: Larger noise levels promote exploration in sparse-reward Franka Kitchen-complete experiments.The additional experiments report that the identified design trends generalize across environments, reward structures, and task complexities.
  • Regularization and implementation: ReinFlow uses entropy regularization by default for state-input tasks but not for visual manipulation tasks.The implementation also documents shared and task-specific hyperparameters across Gym, Kitchen, and Robomimic benchmarks.
  • Experimental protocol: Experiments use multiple random seeds, with five seeds for variable Franka Kitchen tasks and sequential wall-time measurements under specified GPU rendering settings.The reported setup distinguishes task-specific seed counts and hardware conditions for reproducibility.
Loading 2505.22094v7…