Technical Summary: Aligning Flow Map Policies with Optimal Q-Guidance

Problem Statement

Generative policies based on expressive model classes, such as diffusion models and flow matching, are well-suited for complex control problems characterized by highly multimodal action distributions. However, their expressivity incurs a significant inference cost: generating a single action typically requires simulating many steps of the generative process (e.g., solving an Ordinary Differential Equation numerically). This compounding latency hinders deployment in online and real-world settings.

While recent efforts have attempted to learn one-step offline policies to mitigate this latency, these approaches often remain insufficient for offline-to-online reinforcement learning (RL). In the online phase, the policy must improve beyond the static behavioral prior. Current one-step generative policies frequently rely on heuristic "generate-and-select" procedures (e.g., best-of-$N$ sampling) to bias actions toward high critic $Q$-values. This approach offsets the computational advantage of one-step generation by requiring numerous policy and critic evaluations per decision and fails to guarantee optimal local improvement for any finite $N$.

Methodology

The paper introduces 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.

1. Flow Map Policies

Unlike standard flow-matching policies that learn an instantaneous velocity field $v_t$ to be integrated over time, Flow Map Policies learn a two-time jump operator $X_{r,t}$ that evolves action dynamics between any two time points $r, t \in [0, 1]$.

• Definition: The policy is defined by a flow map $X_{r,t}$ such that $X_{r,t}(a_r | s) = a_t$.
• Parametrization: The map is parameterized via an average action velocity $u_{r,t}$, where $X_{r,t}(a_r | s) = a_r + (t-r)u_{r,t}(a_r | s)$.
• Training Objectives: To train the flow map on off-diagonal time pairs ($r < t$), the authors enforce consistency rules derived from the flow-map jump condition and the semi-group property of the ODE. This yields three policy-distillation objectives:
  1. Lagrangian: Distills the time derivative of the map.
  2. Eulerian: Distills the spatial gradient and time derivative.
  3. Progressive: Ensures consistency across intermediate jumps.
     Note: Mean-flow policies (e.g., MVP) are shown to be a specific instantiation of the Eulerian objective.

2. Flow Map Q-Guidance (FMQ)

For online adaptation, the authors formulate the problem as a trust-region optimization that improves the critic's $Q$-value while remaining close to the offline policy.

• Optimization Problem: Maximize $Q_\phi(s, a_1 + \Delta)$ subject to $\|\Delta\|_2 \le \eta$, where $a_1$ is the action generated by the offline flow map.
• Closed-Form Solution (Theorem 3.2): By approximating the $Q$-function with a first-order Taylor expansion, the authors derive an analytically optimal target displacement. The optimal average velocity $u^*_{r,1}$ is:
  $$u^*_{r,1} = u^{off}_{r,1} + \eta \frac{\nabla_a Q_\phi(s, a_1)}{\|\nabla_a Q_\phi(s, a_1)\|_2}$$
• Learning Target: The policy is trained to regress against this optimal target using a stop-gradient operator, creating a self-bootstrapped learning target that aligns the action distribution with $Q$-guidance without distillation networks.
• Adaptive Trust Region: To handle critic uncertainty, the radius $\eta$ is made adaptive per sample based on the disagreement between a twin-critic ensemble. Higher disagreement leads to a smaller trust region, prioritizing the offline prior.

3. Q-Guided Beam Search (QGBS)

To further refine actions at inference time without additional training or costly ODE simulations, the authors introduce QGBS.

• Mechanism: It combines stochastic renoising (corrupting the current action to an intermediate state $t' < 1$ based on Signal-to-Noise Ratio) with beam search.
• Process:
  1. Exploration: Renoise the current action to generate a batch of intermediate states.
  2. Denoising & Selection: Use the flow map to denoise these states to $t=1$, generating multiple candidate actions.
  3. Exploitation: Select the top candidates based on the critic $Q$-value, apply the optimal trust-region displacement (Eq. 11), and iterate.
• This allows for iterative inference-time refinement, balancing exploration and exploitation.

Key Contributions

1. Flow Map Policies: A framework for learning one-step policies as two-time jump operators for flow-based generative actors, generalizing existing one-step policies.
2. Algorithms (FMQ & QGBS):
   • FMQ: An efficient adaptation method using optimal $Q$-guidance derived from a trust-region constraint, eliminating the need for best-of-$N$ heuristics during training.
   • QGBS: A stochastic inference-time refinement algorithm combining renoising, beam selection, and trust-region $Q$-guidance.
3. State-of-the-Art Performance: The method achieves superior performance across 12 robotic manipulation and locomotion tasks.

Experimental Results

The authors evaluated FMQ and QGBS on 12 tasks from OGBench and RoboMimic, comparing against baselines including QC (multi-step flow matching) and MVP (state-of-the-art one-step mean flow policy).

• Performance: FMQ achieved a state-of-the-art Interquartile Mean (IQM) success rate of 0.91, outperforming MVP (0.75) by a relative improvement of 21.3% and QC (0.86) by 5.8%.
• Efficiency: FMQ is approximately 2.77$\times$ more efficient during online adaptation than MVP, converging to the same success targets significantly faster.
• Inference Scaling: Applying QGBS to FMQ further improved the IQM to 0.93. Notably, QGBS achieved this with fewer network forward evaluations (NFE) than standard best-of-32 sampling in certain configurations, demonstrating that diversifying candidates via renoising is more efficient than simply sampling more candidates.
• Robustness: The method showed particular gains in challenging environments (e.g., cube-trl-t4 and amaze-gnt-t4) where baselines struggled, and maintained stability in sparse-reward locomotion tasks.

Significance

The paper claims to bridge the gap between the expressivity of generative policies and the low-latency requirements of offline-to-online RL. By formulating online adaptation as a trust-region optimization problem, the authors derive a principled, closed-form learning target (FMQ) that aligns actions with $Q$-value guidance without the computational overhead of iterative sampling or the sub-optimality of heuristic selection. Furthermore, the introduction of QGBS provides a mechanism to refine policies at inference time, leveraging stochastic sampling and beam search to improve performance without additional training costs. The work demonstrates that one-step generative policies, when properly guided by optimal $Q$-gradients, can outperform both multi-step diffusion models and existing one-step baselines in complex continuous control tasks.