QAMM: Adjoint MeanFlow Matching for Few-Step Offline Reinforcement Learning

Yuehu Gong; Shutong Ding; Mokai Pan; Yimiao Zhou; Jiashu Hou; Ye Shi; Yanwei Fu

cs.LGSep 28, 2026

QAMM: Adjoint MeanFlow Matching for Few-Step Offline Reinforcement Learning

Authors: Yuehu Gong, Shutong Ding, Mokai Pan, Yimiao Zhou, Jiashu Hou, Ye Shi, Yanwei Fu

Organizations: School of Data Science, Fudan University · ShanghaiTech University · Shanghai Innovation Institute

Abstract

Flow policies can model rich action distributions, but their iterative sampling limits decision speed. Adjoint matching uses the critic's action gradient to improve a flow policy without backpropagating through its sampling trajectory, yet its supervision is defined for instantaneous velocities. We propose QAMM, a method that turns the critic-derived adjoint signal into supervision for MeanFlow's average velocity. The resulting policy learns finite-interval transport directly and generates actions with few network evaluations. We derive the adjoint MeanFlow target, specify its gradient boundaries, and train it with an offline actor-critic. On ten HumanoidMaze tasks, QAMM produces effective two-call policies and achieves competitive performance against strong flow-policy baselines. These results show that adjoint-based Q optimization can be combined with average-velocity learning to obtain expressive offline policies with few-step action generation.

Figures & tables

Figure 1: QAMM: learning average velocities from critic gradients. Left: adjoint-corrected velocities (orange) originate at fixed states on the gray reference path. At the selected

(x_{t},t)

, the corrected velocity supervises MeanFlow through the target below. The purple curve schematically represents predictions

\widehat{x}_{r}=x_{t}+(r-t)u_{\theta}(s,x_{t},r,t)

as the destination

r

varies; the dotted arrow denotes regression supervision. Contours illustrate the ideal behavior-regularized Q tilt. Right: offline success on the five HumanoidMaze-large tasks, using the means and seed standard deviations in Table 1 . QAMM uses two seeds and the tuned setting; the supplied baselines use eight seeds.

Dataset	Task	FQL	CGQL-L	DSRL	IFQL	QAM	QAM-E	TRQAM	QAMM
Medium	1	$82.5\pm 13.4$	$81.2\pm 4.3$	$15.6\pm 14.0$	$\boldsymbol{91.0\pm 1.4}$	$31.9\pm 24.8$	$20.7\pm 19.5$	$\underline{86.8\pm 2.4}$	$67.8\pm 3.0$
	2	$98.8\pm 0.8$	$\boldsymbol{99.2\pm 0.4}$	$84.6\pm 5.5$	$94.0\pm 1.6$	$\underline{99.2\pm 0.8}$	$98.2\pm 0.8$	$87.7\pm 2.5$	$81.6\pm 6.8$
	3	$80.2\pm 14.5$	$0.1\pm 0.1$	$57.6\pm 22.6$	$\boldsymbol{94.6\pm 1.3}$	$89.3\pm 5.1$	$70.9\pm 21.3$	$88.4\pm 1.6$	$\underline{90.0\pm 1.2}$
	4	$0.0\pm 0.0$	$0.0\pm 0.0$	$0.0\pm 0.0$	$\boldsymbol{81.4\pm 2.5}$	$0.0\pm 0.0$	$0.0\pm 0.0$	$58.9\pm 9.7$	$\underline{60.2\pm 7.4}$
	5	$97.6\pm 3.2$	$\underline{98.8\pm 0.7}$	$85.5\pm 2.3$	$98.2\pm 1.0$	$\boldsymbol{98.8\pm 0.6}$	$98.5\pm 1.0$	$96.2\pm 2.3$	$93.6\pm 2.0$
	Mean	$71.8\pm 4.8$	$55.9\pm 0.9$	$48.6\pm 6.2$	$\boldsymbol{91.8\pm 0.6}$	$63.9\pm 5.4$	$57.7\pm 5.2$	$\underline{83.6\pm 2.5}$	$78.6\pm 2.8$

Table 1: Offline success rate (%, mean

\pm

seed standard deviation), averaged over 800k–1M updates. QAMM uses two seeds, other methods eight. Large QAMM uses the tuned setting, selected on one reported seed. Domain means weight five tasks equally. Bold marks the highest and underline the second-highest unrounded mean in each row.

Figure 2: Per-task offline-to-online success curves. Columns are tasks 1–5; rows are HumanoidMaze-medium and HumanoidMaze-large. The vertical line at 1M updates separates offline and online training. Lines show seed means and bands show $\pm 1$ seed standard deviation. Large QAMM uses the tuned setting.

Figure 3: Fixed-backup-mixture ablation on HumanoidMaze-medium task 1 and HumanoidMaze-large task 4. Red uses $\alpha=0.25$ and purple uses $\alpha=0$ ; all other center settings and training budgets are matched. Lines and bands show two-seed mean and $\pm 1$ standard deviation. TRQAM’s eight-seed curve is context, not an ablation arm.

Figure 4: Sampling diagnostics on HumanoidMaze-medium task 1 (50 episodes per setting, one historical stage-1 checkpoint). Left: sequential sampling and endpoint refinement; the star marks $t=0.5$ , QAMM’s default two-call sampler. Right: critic-selected Best-of- $N$ with that sampler. The checkpoint uses a different backup mixture from the final configuration, so these comparisons are exploratory.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Eight-mode QAMM toy diagnostic for two seeds (rows). Columns show the target and generated densities with one call, two-call endpoint re-noising at noise scale $\sigma=0.125$ , 128-step ODE, and 128-step SDE sampling. For the two-call sampler, $x_{t}=(1-\sigma)\widehat{x}_{1}+\sigma z$ with $z\sim\mathcal{N}(0,I)$ , so the internal forward time is $t=0.875$ . Panel annotations give grid Jensen–Shannon divergence and covered modes. This training protocol differs from the HumanoidMaze main experiments.

Setting	Value
Behavior pretraining / offline updates	300k / 1M
Batch size / critic ensemble size	256 / 10
Action horizon / discount	1 / 0.999
Policy hidden layers	$512\times 4$
Actor learning rate	$10^{-4}$
Behavior and critic learning rate	$3\times 10^{-4}$

Appendix

Table 2: Recorded Humanoid center configuration.

Explore similar work

Aligning Flow Map Policies with Optimal Q-Guidance
May 12, 2026Christos Ziakas, Alessandra Russo, Avishek Joey BoseFlow Policies Geometric Guidance
Imperial College London · Mila
Trust Region Q Adjoint Matching
May 26, 2026Yonghoon Dong, Kyungmin Lee, Changyeon Kim +2Flow Policies Q-Learning
1KAIST AI · 2Seoul National University · 3RLWRLD
Entropy-Regularized Adjoint Matching for Offline Reinforcement Learning
May 7, 2026Abdelghani Ghanem, Mounir GhoghoEntropy Regularized Reinforcement Learning Model-Based Reinforcement Learning
College of Computing Mohammed VI Polytechnic University Rabat 11103, Morocco

May 12, 2026, cs.LGJ/K move · Enter open · S save

Aligning Flow Map Policies with Optimal Q-Guidance

Christos Ziakas, Alessandra Russo, Avishek Joey Bose

Imperial College London · Mila

Flow Policies Geometric Guidance Trust Region Mean Average Precision

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Aligning Flow Map Policies with Optimal Q-Guidance

Trust Region Q Adjoint Matching

Entropy-Regularized Adjoint Matching for Offline Reinforcement Learning