Reinforcement Learning for Flow-Matching Policies

TL;DR

This paper introduces reinforcement learning methods to enhance flow-matching policies in robotics, significantly outperforming suboptimal demonstrations in simulated control tasks by reducing costs substantially.

Contribution

It presents novel reinforcement learning algorithms, Reward-Weighted Flow Matching and Group Relative Policy Optimization, for training flow-matching policies beyond demonstration quality.

Findings

Both methods outperform suboptimal demonstrations.

GRPO reduces costs by 50-85% compared to ILFM.

Effective in simulated unicycle control tasks.

Abstract

Flow-matching policies have emerged as a powerful paradigm for generalist robotics. These models are trained to imitate an action chunk, conditioned on sensor observations and textual instructions. Often, training demonstrations are generated by a suboptimal policy, such as a human operator. This work explores training flow-matching policies via reinforcement learning to surpass the original demonstration policy performance. We particularly note minimum-time control as a key application and present a simple scheme for variable-horizon flow-matching planning. We then introduce two families of approaches: a simple Reward-Weighted Flow Matching (RWFM) scheme and a Group Relative Policy Optimization (GRPO) approach with a learned reward surrogate. Our policies are trained on an illustrative suite of simulated unicycle dynamics tasks, and we show that both approaches dramatically improve upon the suboptimal demonstrator performance, with the GRPO approach in particular generally incurring between $50\%$ and $85\%$ less cost than a naive Imitation Learning Flow Matching (ILFM) approach.