🤖 AI Summary
This work addresses the challenge of end-to-end optimization when combining flow matching with value-gradient-based reinforcement learning, a setting where sampling instability often undermines performance. Existing approaches typically compromise either representational capacity or the iterative generative nature of the process. To overcome this, we propose VINE, a stable sampling method tailored for reinforcement learning that constructs differentiable trajectories by reconstructing interpolated states at each denoising step. Our analysis reveals that the observed instability stems from the original sampling strategy rather than the iterative generation mechanism itself. VINE is the first method to enable end-to-end value-gradient optimization while preserving high expressivity, achieving substantial improvements over state-of-the-art baselines on both the OGBench offline reinforcement learning benchmark and real-world robotic manipulation tasks.
📝 Abstract
Flow-matching policies have emerged as an effective policy parameterization for robot learning. They iteratively generate actions from noise, enabling highly expressive modeling of complex and multimodal action distributions. However, prior works observed that scaling these policies with value-gradient reinforcement learning (RL) often leads to training instability. Existing methods attribute this instability to iterative generation and therefore avoid end-to-end value-gradient optimization by sacrificing iterative generation, high expressiveness, or value-gradient optimization. Contrary to prior belief, we show the instability does not stem from iterative generation itself, but from the vanilla sampling strategy originally designed for behavior cloning, which becomes brittle under value-gradient RL. Motivated by this insight, we propose VINE, an RL-oriented sampling method that enables stable end-to-end value-gradient optimization for flow-matching policies. Instead of following a single flow trajectory, VINE reconstructs a new interpolation state at every denoising step, creating a stable differentiable path for value-gradient propagation while remaining compatible with the original flow-matching denoising process. As a result, VINE preserves the expressiveness and iterative generation of flow-matching without sacrificing end-to-end value-gradient optimization. Despite performing end-to-end backpropagation through all ten denoising steps, VINE achieves stable policy improvement and consistently outperforms state-of-the-art RL methods on the OGBench offline RL benchmark and real-world robotic manipulation task. Videos are available on our website: https://agibottech.github.io/vine.