Aligning Flow Map Policies with Optimal $Q$-Guidance

Christos Ziakas, Alessandra Russo, Joey Bose

Keywords: Generative Policies, Flow Matching, Offline-to-Online Reinforcement Learning

TL;DR: Aligning 1-step Flow-matching Policies with Optimal $Q$-Guidance

Abstract: Generative policies based on expressive model classes, such as diffusion and flow matching, are well-suited to complex control problems with highly multimodal action distributions. Their expressivity, however, comes at a significant inference cost: generating each action typically requires simulating many steps of the generative process, compounding latency across sequential decision-making rollouts. We introduce flow map policies, a novel class of generative policies designed for fast action generation by learning to take arbitrary-size jumps---including one-step jumps---across the generative dynamics of existing flow-based policies. We instantiate flow map policies for offline-to-online reinforcement learning (RL) and formulate online adaptation as a trust-region optimization problem that improves the critic's $Q$-value while remaining close to the offline policy. We theoretically derive Flow Map $Q$-Guidance Training (FMQ), a principled closed-form learning target that is optimal for adapting offline flow map policies under a critic-guided trust-region constraint. We further introduce $Q$-Guided Beam Search (QGBS), a stochastic flow-map sampler that combines renoising with beam search to enable iterative inference-time refinement. Across $12$ challenging robotic manipulation and locomotion tasks from OGBench and RoboMimic, FMQ achieves state-of-the-art performance in offline-to-online RL, outperforming the previous one-step policy MVP by a relative improvement of $21.3\%$ on the average success rate.

Submission Number: 105