Iterative Refinement of Flow Policies in Probability Space for Online Reinforcement Learning
This work addresses a bottleneck in adapting pre-trained policies for robotics, offering a more stable and efficient method for online fine-tuning.
The paper tackles the problem of fine-tuning flow/diffusion policies in online reinforcement learning, which are vulnerable to distributional shift and difficult to adapt with standard methods, by introducing the Stepwise Flow Policy (SWFP) framework that decomposes flow matching into incremental steps aligned with optimal transport principles, resulting in enhanced stability, efficiency, and superior adaptation performance across robotic control benchmarks.
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.