Iterative Refinement of Flow Policies in Probability Space for Online Reinforcement Learning

TL;DR

This work addresses the challenge of online fine-tuning for pre-trained flow-based policies in continuous-control RL, where iterative inference and distributional shift hinder stability. It introduces Stepwise Flow Policy (SWFP), which reinterprets flow matching as a cascade of fixed-step JKO updates in probability space, enforcing Wasserstein trust regions and entropic regularization to stabilize policy improvement. The method trains short, composable flow blocks with a soft critic in a parallel-block framework, guaranteeing monotonic improvement and providing theoretical convergence toward the MaxEnt policy. Empirically, SWFP delivers superior stability, computational efficiency, and adaptation performance across toy tasks and diverse robotic benchmarks, including significant gains in offline-to-online transfer.

Abstract

While behavior cloning with flow/diffusion policies excels at learning complex skills from demonstrations, it remains vulnerable to distributional shift, and standard RL methods struggle to fine-tune these models due to their iterative inference process and the limitations of existing workarounds. In this work, we introduce the Stepwise Flow Policy (SWFP) framework, founded on the key insight that discretizing the flow matching inference process via a fixed-step Euler scheme inherently aligns it with the variational Jordan-Kinderlehrer-Otto (JKO) principle from optimal transport. SWFP decomposes the global flow into a sequence of small, incremental transformations between proximate distributions. Each step corresponds to a JKO update, regularizing policy changes to stay near the previous iterate and ensuring stable online adaptation with entropic regularization. This decomposition yields an efficient algorithm that fine-tunes pre-trained flows via a cascade of small flow blocks, offering significant advantages: simpler/faster training of sub-models, reduced computational/memory costs, and provable stability grounded in Wasserstein trust regions. Comprehensive experiments demonstrate SWFP's enhanced stability, efficiency, and superior adaptation performance across diverse robotic control benchmarks.

Figures (5)

• Figure 1: The illustration of continuous-time flow (left) and discrete-time block flow (right).
• Figure 2: Illustration of SWFP's parallel block training. Each horizontal chain of black arrows represents the forward composition of flow blocks within a training epoch. Double-headed arrows indicate the Wasserstein distance minimization objectives: $\mathcal{W}^2_2(\pi^{k+1}_{n-1}, \pi^{k}_{n-1})$ and $\mathcal{W}^2_2(\pi^{k}_{n-1}, \pi^{k+1}_{n})$ during optimization of block $n$.
• Figure 3: Illustration of SWFP performance on a bandit toy example: (left) source distribution and behavior data in the replay buffer; (right) implicit SWFP steers the source distribution toward the high-performance behavior policy, heatmap shows target policy.
• Figure 4: Learning curves of online fine-tuning with various methods. Observe that SWFP largely always dominates or attains similar performance to the next best method.
• Figure 5: Sensitivity of Hyperparameters.