PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2609.34497v1cs.LG

TL;DR

Flow policies can represent complex action distributions, but iterative sampling limits action-generation speed, and prior adjoint matching supervises instantaneous rather than average velocities. QAMM converts critic-derived adjoints into MeanFlow average-velocity supervision for finite-interval transport. It produces competitive two-call policies on HumanoidMaze tasks, with a higher descriptive five-task mean than TRQAM on Large.

  • Problem

    Flow policies model complex action distributions, but iterative sampling limits action-generation speed, while adjoint matching’s instantaneous-velocity supervision still requires flow integration for accurate action generation.

    Source: §1 · p. 1

  • Method

    QAMM propagates the critic’s terminal gradient backward into an adjoint-corrected velocity and uses it in a stop-gradient MeanFlow average-velocity target.

    Source: §1 · p. 1, §6 · p. 6

  • Results

    QAMM learns competitive two-call policies on Medium and reaches a higher descriptive five-task mean than TRQAM on Large.

    Source: §5.2 · p. 5

  • Takeaways & Limitations

    The resulting average-velocity representation supports two-call refinement without differentiating through the full sampling trajectory.

    Source: §6 · p. 6

    The theory establishes conditional compatibility, not convergence of the jointly trained algorithm, and the implemented controller does not certify the KL of the few-step policy.

    Source: §4.4 · p. 5, §6 · p. 6

Abstract

from arXiv · show

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.

1 Introduction

Flow policies model complex offline-action distributions, but iterative sampling limits action-generation speed, while prior critic-guided approaches either differentiate through sampling, discard the critic’s action gradient, or guide noisy actions. QAMM adapts adjoint supervision to MeanFlow’s average velocity, aiming to learn finite-interval transport for two-call action generation.

1 Introduction

  • Flow policies model complex action distributions, but their iterative samplers constrain decision speed; prior methods differentiate through sampling, discard the critic’s action gradient, or guide noisy actions.The cited approaches are sampler optimization or auxiliary policies [2] [7] [19], value-weighted regression [3] [24], and intermediate score guidance [8] [20].Source: §1 · p. 1
  • QAMM propagates the critic’s terminal gradient backward to form an adjoint-corrected velocity, then uses it in a stop-gradient average-velocity target.The training procedure initializes from a behavior flow, jointly trains actor and critic, and constructs trajectories by noising policy-generated endpoints.Source: §1 · p. 1
  • QAMM derives an adjoint MeanFlow target that converts the critic’s terminal action gradient into finite-interval average-velocity supervision with explicit gradient boundaries.The target uses an adjoint-corrected velocity to supervise both the instantaneous term of the MeanFlow identity and the direction of its derivative.Source: §1 · p. 1
  • The offline procedure couples policy-generated endpoints to Markov-noised trajectories and supervises both diagonal and off-diagonal time pairs while retaining two-call deployment.The adapted path-space controller stabilizes Q guidance; deployment uses two-call refinement with independent re-noising.Source: §1 · pp. 1–2
  • The evaluation covers five HumanoidMaze-medium and five HumanoidMaze-large tasks against flow-policy baselines, with initial backup and sampling diagnostics.Figure 1 depicts adjoint-corrected velocities supervising MeanFlow predictions and reports offline success on five HumanoidMaze-large tasks; QAMM uses two seeds and baselines use eight.Source: §1 · p. 2, Figure 1 · p. 2

2 Related Work

The related work distinguishes QAMM from prior ways of incorporating value information into generative policies and situates it among adjoint, average-velocity, and few-step methods. Its training target combines critic-derived adjoint supervision with MeanFlow, while independent re-noising is a sampling choice.

Generative policies for offline RL

  • Prior diffusion and flow policies incorporate Q through action- or noise-space optimization, value- or energy-weighted regression, or intermediate score guidance.These approaches are distinguished from propagating a terminal action gradient into average-velocity supervision.Source: §2 · p. 2

Adjoint guidance and average velocity

  • QAMM combines the lean adjoint of Adjoint Matching [5] and QAM [14] with MeanFlow’s average-velocity identity [9], adapting TRQAM’s path-space control [6] for stability.Flow matching and rectified flow [15] [16] provide the instantaneous-velocity background.Source: §2 · p. 2

Few-step generative control

  • Independent re-noising is QAMM’s sampling choice, whereas its training contribution is the Q-adjoint average-velocity target.Other approaches address sampling cost with different objectives, including consistency policies [4], shortcut models [7], and one-step flow policy mirror descent [2]; broader scalar-Q versus Q-gradient comparisons remain future work.Source: §2 · p. 2

3 Preliminaries

The preliminaries define the offline learning setting, the time-indexed average-velocity representation, and the memoryless control and adjoint framework used by QAMM. They distinguish internal generation time from environment steps and state the assumptions and endpoint treatment relevant to the control interpretation.

3.1 Offline reinforcement learning

  • Offline RL seeks a policy maximizing expected discounted return from a fixed transition dataset without additional environment interaction during training.A critic ensemble is trained by temporal-difference regression, and a behavior model approximates dataset actions conditioned on state.Source: §3.1 · p. 2

3.2 Average velocity

  • MeanFlow represents transport with a state-conditioned average velocity uθ(s, x, r, t), where t is current generation time and r is destination time.Generation runs from Gaussian noise at t = 0 to an action at t = 1; on the diagonal, average velocity equals instantaneous velocity.Source: §3.2 · p. 3

3.3 Memoryless control and lean adjoints

  • The memoryless control framework uses an ideal SDE whose drift is 2vβ − x/t and diffusion squared is 2(1 − t)/t, with memorylessness referring to initial and terminal variables.The path-KL interpretation requires matched diffusion and initial laws plus suitable change-of-measure assumptions; lean adjoints propagate terminal reward derivatives backward through reference-drift input Jacobians.Source: §3.3 · p. 3

4 Adjoint MeanFlow Matching

QAMM turns critic-derived adjoint signals into finite-interval average-velocity supervision for a MeanFlow policy. Its training combines adjoint targets, path construction, and practical actor–critic stabilization while retaining two-call action generation.

4 Adjoint MeanFlow Matching

  • QAMM learns a state-conditioned average-velocity field uθ(s, x, r, t), with t as current generation time and r as destination time.The behavior model, trust-region controller, and mixed critic backup support the construction; their underlying principles are inherited from prior work.Source: §4 · p. 3

4.1 From critic gradients to velocity targets

  • The ideal memoryless AM/QAM control problem yields a terminal tilt proportional to pβ(a | s) exp(Q(s, a)/λ), while the lean adjoint converts the terminal critic gradient into an instantaneous-velocity label.The implementation differentiates the clipped critic at the generated endpoint and discretizes the adjoint; these choices are not exact continuous-time identities.Source: §4.1 · p. 3

4.2 Average-velocity supervision

  • MeanFlow supervision substitutes the adjoint label into both velocity terms, using a JVP of the live student while stopping gradients through the path, adjoint, label, and target.The conditional target compatibility claim assumes a specified path measure and label-noise independence; it does not establish convergence of student updates.Source: §4.2 · pp. 3–4
  • Three equally weighted destination choices—diagonal, terminal, and random intermediate—supervise diagonal velocity and terminal or intermediate displacements.The diagonal term reduces exactly to the adjoint velocity label.Source: §4.2 · p. 4

4.3 Training paths and few-step deployment

  • Training constructs a backward trajectory from the generated endpoint with independent Gaussian noise at each transition, whereas deployment uses fresh-noise endpoint refinement in two calls, not two sequential intervals or Euler steps.The training path is detached before actor regression and still requires multiple transitions; the ideal memoryless reference uses a reverse conditional Markov kernel.Source: §4.3 · p. 4

4.4 Practical stabilization and critic learning

  • Because the weight depends on the residual, the practical optimum need not equal the conditional mean of the unweighted target; experiments use p = 0.3 and c = 10^-3.The experiments use no additional time weighting and no target clipping.Source: §4.4 · p. 4
  • The adapted path-space controller tracks bD as a surrogate, not a certified KL bound for the deployed two-call policy.Surrogate tracking does not establish convergence or exact feasibility.Source: §4.4 · pp. 4–5
  • The critic backup mixes current-policy and EMA-behavior Q values rather than actions; with α = 0.25, α > 0 makes it a mixed-policy backup.Each branch uses one sampled action clipped to the action bounds.Source: §4.4 · p. 5

5 Experiments

Across ten HumanoidMaze tasks, QAMM is evaluated as a two-call policy against published flow-policy baselines, with offline results, online continuation, and ablations of backup and sampling choices. Performance varies by task and configuration, and the reported comparisons have important seed and protocol caveats.

5.1 Setup

  • The evaluation covers five HumanoidMaze-medium and five HumanoidMaze-large tasks, comparing QAMM’s two-call policy with published flow-policy baselines using offline success and online learning curves.QAMM uses two training seeds, while the supplied baseline logs have eight; baselines were not re-evaluated under a common harness.Source: §5.1 · p. 5

5.2 HumanoidMaze results

  • Over the final five online evaluations, QAMM averages 98.7% on Medium and 87.2% on Large, versus supplied TRQAM curves at 97.5% and 71.2%, respectively.The comparisons are descriptive because seed counts, selection, and protocols differ.Source: §5.2 · p. 5

5.3 Ablations and sampling analysis

  • Removing the behavior component of the mixed Bellman backup changes Medium task 1’s offline mean from 67.8% to 76.0%, but Large task 4’s from 42.0% to 0.6%.The comparisons use two seeds per setting; the passage characterizes the effects as task-dependent rather than a uniform benefit.Source: §5.3 · p. 5
  • With critic selection among independently generated two-call candidates, Best-of-N success rises from 88% at N = 1 to 100% at N = 16, at a cost of 2N policy calls and extra critic evaluations.The sampling diagnostics use 50 episodes and one historical checkpoint with a different backup mixture from the final configuration, so they are exploratory.Source: §5.3 · pp. 5–6

6 Discussion and Conclusion

QAMM uses adjoint-corrected velocity as both a regression base and the direction of the MeanFlow derivative, with explicit gradient boundaries that avoid differentiating through the full sampling trajectory. The evaluation is limited to two Humanoid datasets, and few network calls alone do not establish an end-to-end speedup.

6 Discussion and Conclusion

  • The adjoint-corrected velocity supervises both the regression base and the direction of the MeanFlow derivative, while explicit gradient boundaries avoid differentiation through the full sampling trajectory.The resulting representation supports two-call refinement.Source: §6 · p. 6
  • The evaluation covers two Humanoid datasets; further tasks and independent seeds are needed to clarify the empirical result’s scope.Few network calls alone do not establish an end-to-end speedup, which requires matched latency measurement.Source: §6 · p. 6
  • The theory establishes conditional compatibility of the target, not convergence of the weighted, jointly trained algorithm; the implemented controller tracks a surrogate rather than certifying the few-step policy’s KL.The learned reference, finite-step adjoint, endpoint-noised path law, and deployed refinement introduce distinct approximations.Source: §6 · p. 6
  • The sampling diagnostics compare sequential sampling and endpoint refinement, and critic-selected Best-of-N, using one historical checkpoint; the authors describe these comparisons as exploratory.The checkpoint uses a different backup mixture from the final configuration.Source: Figure 4 · p. 7

A Derivations and Conditional Results

The appendix derives the forward-time average-velocity identity and states conditions for a conditional target calculation and an endpoint-conditioned path-measure relation. It also distinguishes these results from guarantees about practical training or the learned field.

A.1 Average-velocity identity in forward generation time

  • Holding destination r fixed yields the average-velocity identity for r > t, with the diagonal value following by continuity as r approaches t.Allowing the destination to move with t would add a term to the derivation.Source: §A.1 · p. 9

A.2 Conditional target calculation

  • The detached target has the stated conditional numerical identity when the destination draw is conditionally independent of label noise and the conditional mean equals v∗.Squared-error regression on a frozen square-integrable target is minimized by its conditional mean, but this compatibility argument does not prove repeated training updates converge.Source: §A.2 · p. 9
  • The conditional-mean argument does not establish the optimizer of the practical residual-weighted objective or verify its required conditional-mean premise.Transferring an Adjoint Matching stationary-point result would additionally require its path-measure and reference-field assumptions.Source: §A.2 · p. 9

A.3 Endpoint-conditioned path measures

  • Endpoint-conditioning the exact reference path measure on q gives the stated path-KL relation when the logarithm is integrable and q is absolutely continuous with respect to pβ.For the ideal terminal tilt proportional to pβeQ/λ, the construction recovers the corresponding terminally tilted reference path measure and preserves the prescribed initial law when reference endpoints are independent.Source: §A.3 · p. 9
  • The path-measure construction does not identify the learned diagonal field as its drift, so measured velocity energy is not thereby an exact KL.The result assumes an exact conditional kernel, not merely a learned reference with approximately correct marginals.Source: §A.3 · p. 9

B.1 Behavior MeanFlow training

Behavior MeanFlow training starts from dataset actions noised along a linear path and trains predictions at three destinations. The offline procedure initializes from the behavior network and constructs detached adjoint targets for the live student.

B.1 Behavior MeanFlow training

  • For dataset action a, the behavior path interpolates independent Gaussian noise z and a, with velocity cβ = a − z, and samples t uniformly from [0, 1].The behavior MeanFlow is trained at three destinations with residual weighting; after 300k updates, its live network and EMA initialize the policy with a fresh optimizer state.Source: §B.1 · p. 9
  • Each offline actor update detaches a two-call policy endpoint and its reverse Markov trajectory from policy-parameter differentiation.The behavior field continues training on dataset actions, while its EMA is fixed during each adjoint-target construction and evolves across updates.Source: §B.1 · p. 9
  • The procedure differentiates the clipped-action target-critic mean with respect to action, then propagates the negative terminal derivative backward through reference-drift VJPs.It forms and freezes the regularized adjoint velocity label before sampling three destination times per point.Source: §B.1 · p. 9
  • The student prediction and input/time JVP build a detached target and residual weight, while optimization differentiates only the prediction branch.Reference parameters are frozen without disabling their input Jacobians; behavior and critic models update with their own losses.Source: §B.1 · p. 10

C Configurations and Reproducibility

Table 2 records the Humanoid center configuration. The accompanying reproducibility details specify the center-run schedule and evaluation aggregation, while noting that supplied baseline logs were not rerun under a common harness.

C Configurations and Reproducibility

  • Each center run uses two seeds, 1M offline updates, 500k online updates, and evaluations of 50 episodes every 50k updates.Each phase averages its final five evaluations per seed, then reports the mean and standard deviation across seeds.Source: C Configurations and Reproducibility · p. 10, Table 2 · p. 10
  • The supplied TRQAM baseline logs use eight seeds, but their implementations and evaluation randomness were not rerun under a common harness.Large tuning changes α to 0.5 and the KL budget to 0.75, selected on seed 10001 before running seed 20002.Source: C Configurations and Reproducibility · p. 10
  • Recorded Table 2 settings differ from historical defaults in the generic agent configuration.The paper identifies source snapshots and launcher manifests, rather than mutable agent defaults alone, as defining a run.Source: C Configurations and Reproducibility · p. 10, Table 2 · p. 10

D Exploratory QAMM Toy Diagnostic

In the two-seed eight-mode diagnostic, all samplers cover the modes, while the two-call sampler has higher grid Jensen–Shannon divergence than the 128-step SDE reference. The diagnostic uses a simplified protocol rather than the full HumanoidMaze procedure.

D Exploratory QAMM Toy Diagnostic

  • The two-call sampler’s grid Jensen–Shannon divergences are 0.093 and 0.097, versus 0.049 and 0.071 for the 128-step SDE reference; all samplers cover eight modes.Some connecting mass remains, and mode-weight total variation is slightly worse than at σ = 0.5; the diagnostic does not show reproduction of the ideal tilted distribution.Source: D Exploratory QAMM Toy Diagnostic · p. 10, Figure 5 · p. 11
  • The toy diagnostic uses a frozen flow-matching reference, randomly initialized actor, and one destination per example, unlike the HumanoidMaze procedure.HumanoidMaze initializes from a behavior MeanFlow, uses three destinations per path point, and constructs an endpoint-noised Markov path.Source: D Exploratory QAMM Toy Diagnostic · p. 10

Related papers

Loading 2609.34497v1…