Flow policy optimization improves control tasks using q function guidance

Refinement-based Flow Policy Optimization

Machine LearningArtificial Intelligence

Summary

Some computer programs learn to make decisions by figuring out the best actions to take in different situations. Usually, these programs need examples of good actions to learn from, but sometimes it's hard to get them directly. The authors developed a new method where the program starts with random guesses and then improves those guesses step-by-step using a guide that estimates how good each action is. By repeating this process, the program learns to pick better and more diverse actions without needing direct examples. This method works well in tests involving controlling things smoothly and handling complicated choices.

What this means in practice

  • For robotics engineers: Use refined flow policies to improve robotic control actions when direct sampling of optimal moves is difficult.
  • For autonomous vehicle developers: Apply stochastic refinement guided by value functions to generate diverse driving actions for complex scenarios.

Authors

Bumgeun Park, Hyukjun Yang, Donghwan Lee

Abstract

Flow-based policies offer an expressive representation for online reinforcement learning, but conventional flow matching requires samples drawn from the distribution to be modeled. This poses a challenge when the desired action distribution is defined only implicitly by a Q-function, since directly sampling actions from the resulting distribution is generally intractable. We propose Refinement-Based Flow Policy Optimization (RFPO), a novel framework for training a flow policy in online reinforcement learning by alternating between Q-guided sample refinement and self-target flow matching. RFPO first generates actions from Gaussian noise using the current flow policy and then uses a finite-step stochastic refinement procedure to move them toward an energy-based distribution induced by the Q-function. Each refined action is then paired with its corresponding initial noise sample and used as a fixed target for flow-matching training. By repeatedly refining its own outputs and learning from the resulting targets, RFPO incorporates Q-guidance into the policy without requiring direct samples from the target distribution, while retaining the capacity to represent multiple action modes. We further provide a theoretical analysis of the distributional dynamics induced by RFPO. Across six continuous-control tasks, RFPO matches or outperforms a standard Gaussian-policy baseline on almost every task. Experiments on six synthetic two-dimensional target distributions with diverse geometries demonstrate that RFPO captures complex multimodal structure without mode collapse.