Aligning Flow Map Policies with Optimal Q-Guidance

TL;DR

This paper introduces flow map policies for fast, flexible action generation in complex control tasks, combining theoretical insights with practical algorithms to improve offline-to-online reinforcement learning performance.

Contribution

It proposes flow map policies for rapid action sampling, derives a new Q-guidance learning target, and develops a stochastic sampler for iterative inference, advancing offline-to-online RL.

Findings

FMQ outperforms previous methods with a 21.3% success rate improvement.

Flow map policies enable arbitrary jump actions, reducing inference latency.

The approach achieves state-of-the-art results across 12 robotic tasks.

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 (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.