Aligning Flow Map Policies with Optimal Q-Guidance

Summary

Generative policies based on expressive model classes, such as diffusion-models and flow matching, are well-suited to complex control problems with highly multimodal action distributions.

Abstract Summary

Generative policies based on expressive model classes, such as diffusion-models 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.

Key Points

Introduces Aligning Flow Map Policies with Optimal Q-Guidance, advancing rl research.
Detailed evaluation demonstrates effectiveness on benchmarks and real-world scenarios.
Provides insights into methodology and architectural choices.

Flow Map Q Q -Guidance

Generative policies operate over an inner-time axis, i.e., ODE simulation time, that is distinct from the evolution of the MDP. Consequently, for every state along the trajectory $s^{\tau}$, a corresponding action $a^{\tau}_{1}$ must be generated through numerical simulation of the flow-policy ODE in eq. 1 —necessitating large amounts of function evaluations of the policy network.

Experiments

We investigate the application of FMQ across $12$ robotic manipulation and locomotion tasks with varying difficulties across $7$ environments from the OGBench ( Park et al., 2024 ) and Robomimic ( Mandlekar et al., 2021 ) benchmarks. The manipulation tasks include two from Robomimic ( can , square ) and six from OGBench ( cube-dbl-t3/4 , cube-trl-t3/4 , scene-t4/5 ).

Locomotion is evaluated on four OGBench tasks ( hmaze-med-t3/4 , amaze-gnt-t4/5 ).

Experiments — Aligning Flow Map Policies with Optimal Q-Guidance — Figure 3: Convergence speedup of FMQ compared to MVP at success targets ($\xi$), with 95% CIs.

Sources and Demonstrations

Figures are reproduced from Ziakas et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.