PaperlayerView exact revision on arXiv

Source-linked AI summary

Q-learning with Adjoint Matching

Qiyang Li, Sergey Levine

arXiv:2601.14234v4cs.LGcs.AIcs.ROstat.ML

TL;DR

Expressive flow and diffusion policies are difficult to optimize against critics because denoising-process backpropagation is unstable, while existing work either discards gradients or compromises expressivity. QAM uses adjoint matching to form a step-wise objective from critic action gradients without unstable backpropagation, and author-reported results show strong performance across offline and offline-to-online RL. Its main scope boundary is behavior-policy support, since actions with extremely low behavior probability may be difficult to represent.

  • Problem

    Continuous-action RL lacks an optimization approach that preserves expressive flow or diffusion policies while effectively using critic action gradients.

  • Method

    QAM is a TD-based RL algorithm that applies adjoint matching to transform critic action gradients into a step-wise flow-policy objective without backpropagation through denoising.

  • Results

    QAM consistently achieves strong performance across offline RL and offline-to-online RL benchmarks, outperforming prior methods in the reported empirical study.

  • Takeaways & Limitations

    QAM provides an expressive policy-extraction method that uses critic gradients while avoiding common instability, approximation, and distillation limitations.

  • Takeaways & Limitations

    QAM can struggle with support mismatch when optimal actions have extremely low probability under the behavior distribution.

Abstract

from arXiv · show

We propose Q-learning with Adjoint Matching (QAM), a novel TD-based reinforcement learning (RL) algorithm that tackles a long-standing challenge in continuous-action RL: efficient optimization of an expressive diffusion or flow-matching policy with respect to a parameterized Q-function. Effective optimization requires exploiting the first-order information of the critic, but it is challenging to do so for flow or diffusion policies because direct gradient-based optimization via backpropagation through their multi-step denoising process is numerically unstable. Existing methods work around this either by only using the value and discarding the gradient information, or by relying on approximations that sacrifice policy expressivity or bias the learned policy. QAM sidesteps both of these challenges by leveraging adjoint matching, a recently proposed technique in generative modeling, which transforms the critic's action gradient to form a step-wise objective function that is free from unstable backpropagation, while providing an unbiased, expressive policy at the optimum. Combined with temporal-difference backup for critic learning, QAM consistently outperforms prior approaches on hard, sparse reward tasks in both offline and offline-to-online RL.

1 INTRODUCTION

Continuous-action RL must balance expressive flow or diffusion policies against tractable critic-based optimization. QAM uses adjoint matching to retain critic action gradients without unstable denoising-process backpropagation, while targeting the optimal behavior-constrained policy.

  • Flow and diffusion policies represent complex, multi-modal action distributions but require optimization through multi-step denoising processes.
  • Directly using critic action gradients through denoising is often unstable, forcing prior methods to discard gradients or sacrifice policy expressivity.
  • Noisy-action guidance can avoid backpropagation instability but depends on critic gradients being reliable away from the dataset’s narrow noiseless-action support.
  • QAM applies adjoint matching to use critic action gradients for flow-policy optimization under a behavior constraint.
  • QAM’s objective directly uses critic gradients and recovers the optimal behavior-regularized policy at convergence.

2 RELATED WORK

Prior methods for optimizing diffusion and flow policies against critics use post-processing, backpropagation, or intermediate fine-tuning. These approaches trade computation, stability, approximation, or expressivity in different ways, while offline-to-online methods seek sample-efficient online improvement from offline initialization.

  • Post-processing methods refine base policies using critic-value rejection sampling, gradient steps, or residual policies, often adding computation during training or evaluation.
  • Backpropagation-based methods directly optimize critic and policy but may suffer denoising-process instability or distill policies at the expense of expressivity.
  • Intermediate fine-tuning constructs step-wise objectives for denoising states, removing the need to backpropagate through the denoising process.
  • Offline-to-online RL pretrains on offline data, then uses the resulting policy and value functions to accelerate online RL.

3 PRELIMINARIES

The paper formulates offline and offline-to-online RL in an MDP and uses flow-matching models to represent denoising trajectories. Adjoint matching converts the constrained optimization problem into an objective whose optimum preserves the desired tilted action distribution without changing the optimal solution.

  • The RL setup is an MDP with state and action spaces, transitions, discounting, rewards, and an initial-state distribution, together with a transition dataset.
  • Offline RL learns a policy from the transition dataset, while offline-to-online RL fine-tunes that policy through environment interaction with bounded episodes.
  • Flow matching models generate data from noise through a time-varying velocity field and an ODE, with training based on matching the flow velocity.
  • Adjoint matching modifies a base flow model so its induced distribution is tilted according to a value function.
  • The stochastic optimal-control formulation and its adjoint formulation produce the correct marginal tilt distribution at the denoised endpoint.
  • The lean adjoint state is computed using the base flow model, and omitted terms are zero at the optimum, preserving the correct tilt distribution.

4 Q-LEARNING WITH ADJOINT MATCHING (QAM)

QAM uses adjoint matching to optimize expressive flow policies against a critic without unstable backpropagation through the denoising process, while targeting behavior-constrained optimal policies.

  • Policy optimization: The method parameterizes both behavior and optimized policies as flow models and formulates policy optimization as a prior-regularized stochastic optimal control problem.The behavior policy supplies the prior, while the optimized policy is trained to maximize critic value under a behavior constraint.
  • Adjoint matching: QAM replaces unstable backpropagation through the SDE or denoising process with an adjoint matching objective built from lean adjoint states.The lean adjoint states are computed through vector-Jacobian products and used in a squared loss.
  • TD learning: QAM combines adjoint-based policy optimization with standard temporal-difference learning for critic training.The critic uses an ensemble and pessimistic target backup, while policy and critic losses use replay-buffer transitions.
  • Theoretical guarantees: At convergence, QAM recovers the optimal behavior-constrained policy π⋆(a | s) ∝πβ(a | s) exp(τQ(s, a)).This guarantee assumes the adjoint matching loss is optimized to convergence.
  • Constraint scope: Under the standard behavior constraint, QAM may struggle when optimal actions have extremely low probability under the behavior distribution.The paper considers relaxing the constraint to permit actions close to behavior actions in action-space geometry.
  • Practical variants: QAM provides QAM-FQL and QAM-EDIT variants that enforce Wasserstein-style proximity to the QAM-fine-tuned policy through a distilled policy or bounded action edits.QAM-EDIT restricts edits to ||∆a||∞≤σa, while QAM-FQL constrains a one-step policy through an upper bound on squared Wasserstein distance.

5 EXPERIMENTS

The experiments evaluate QAM on complex, long-horizon sparse-reward tasks designed to stress policy extraction, using broad baseline comparisons and controlled evaluation procedures.

  • Domains and datasets: The study evaluates QAM across 10 OGBench domains with five tasks each, including navigation, manipulation, puzzle, cube, and scene domains.The resulting benchmark contains 50 tasks and includes action-chunked policies for selected domains.
  • Domains and datasets: Action chunking produces high-dimensional actions with complex behavior distributions, making these domains a testbed for comparing policy extraction methods.The selected domains are intended to emphasize the policy extraction problem rather than only environment interaction.
  • Comparisons: QAM is compared with 13 representative baselines spanning Gaussian, backpropagation, advantage-weighted, guidance, and other policy-extraction categories.The comparison includes methods using direct flow-policy backpropagation, one-step distillation, and critic-based alternatives.

6 RESULTS

QAM outperforms prior methods in offline evaluation and shows robust offline-to-online fine-tuning, while its performance depends most strongly on temperature tuning.

  • Offline RL: 44 aggregated score: QAM outperforms all prior methods across 10 domains and 50 offline RL tasks.FAWAC scores 8 and BAM scores 35; the table also reports that QAM-FQL and QAM-EDIT outperform prior baselines.
  • Offline RL: The offline results indicate that loss design matters: BAM retains the SOC guarantee but its backpropagation-through-the-SDE optimization performs below QAM.BAM has aggregated score 35, whereas QAM has aggregated score 44.
  • Offline RL: FQL performs similarly to BAM in offline evaluation, suggesting that implementation details affect the practical disadvantage of backpropagation through time.The paper also reports that behavior-cloning augmentation makes QSM effective in the evaluated setting.
  • Offline-to-online fine-tuning: QAM-EDIT is the most robust method across offline-to-online tasks and outperforms all prior methods on cube-triple.QSM performs better on antmaze-giant, while FQL is slightly better there but slower on puzzle-4x4 and cube-triple.
  • Sensitivity analysis: Temperature has the largest performance impact among the tested QAM-EDIT components and therefore needs tuning.Gradient clipping and a critic ensemble size of K = 10 help, while T = 10 flow steps performs sufficiently and further increases do not improve performance.

7 DISCUSSION

QAM effectively leverages the critic’s action gradient to extract an optimal prior-constrained policy, while practical challenges remain when the critic is ill-conditioned.

  • QAM extracts an optimal prior-constrained policy by leveraging the critic’s action gradient.The method circumvents approximation, instability, and distillation-related expressivity limitations described for prior approaches.
  • Ill-conditioned critics can create optimization stability issues because QAM relies strongly on the critic’s action gradient.Gradient clipping can alleviate this issue, but combining value and gradient information is proposed as a more principled extension.
  • Applying QAM to real-world robotic settings with action chunking policies remains a possible extension.

A ADDITIONAL RESULTS

The data-quality analysis evaluates QAM-EDIT on original and corrupted OGBench datasets, aggregating performance across five tasks per domain.

  • QAM-EDIT is evaluated on noisy manipulation datasets and stitch navigation datasets from OGBench.Noisy datasets use larger uncorrelated Gaussian noise, while stitch datasets contain shorter trajectory segments than the original navigation datasets.
  • Figure 4 reports aggregated scores over 5 tasks in each domain for original navigate/play datasets and corrupted stitch/noisy datasets.The top row covers the original datasets used in the main experiments; the bottom row covers stitch/noisy datasets.

B FULL RESULTS

The full-results section provides offline and online training curves, end-of-offline performance, and sensitivity-analysis curves for all methods and tasks.

  • Figure 5 shows full training curves through the offline phase until 1M steps and the online phase after 1M steps.
  • Table 2 reports performance at the end of 1M-step offline training.
  • Figure 6 presents full training curves for the sensitivity analysis.Hyperparameters are provided in Section F.

C ADDITIONAL DISCUSSIONS FOR RELATED WORK

Related work uses guidance and intermediate supervision to optimize diffusion or flow policies, but additive denoising guidance can produce incorrect distributions and some approximations lack convergence guarantees.

  • CEP combines a base diffusion policy with a critic-based Boltzmann guidance term, while CFGRL combines a base flow policy with an optimality-conditioned velocity field.CEP approximates π ∝ (πβ)^α(e^Q)^(1−α), whereas CFGRL approximates π ∝ (πβ)^α(πo=1)^(1−α).
  • Summing denoising processes does not generally produce the correct composite distribution, even when each denoising step is exact.The related-work equations express this issue as a mismatch between the target score and the sum of behavior-policy and critic gradients.
  • DAC seeks a policy proportional to πβ(·|s)e^Q(s,·) but derives its objective under an approximate score-gradient assumption.The supplied passage states that such assumptions provide a convenient approximation without guaranteeing the policy’s convergence point.

D DOMAIN AND EXPERIMENT DETAILS

Experiments span 10 domains with repeated-seed evaluation, confidence intervals, and substantial computational requirements.

  • Experiments cover 10 domains, with dataset size, episode length, and action dimension reported for each domain.
  • Each method and task uses 12 seeds, with means and 95% confidence intervals computed by bootstrapping 5000 samples.
  • Each complete offline-to-online run takes around 3 hours, with total computation reported as 51 000 GPU hours.

E BASELINES

The baselines cover flow- and diffusion-policy extraction methods that use value regression, critic gradients, guidance, approximations, or direct optimization with behavior regularization.

  • FBRAC directly optimizes a multi-step flow policy against the Q-function through BPTT while adding a behavior-cloning flow-matching term.
  • FQL distills a multi-step flow policy into a one-step noise-conditioned policy to avoid BPTT, alongside behavior-cloning flow matching and a standard TD critic backup.
  • QSM trains a diffusion model using the critic’s action gradient and standard diffusion denoising, with behavior regularization added for offline-to-online learning.
  • DAC approximates intermediate-action scores by combining behavior-cloning noise with Q-maximization gradients and uses an approximate Gaussian reconstruction procedure.
  • CGQL combines a behavior-cloning flow velocity field with a Q-derived guidance field, using a coefficient ϑ that often works better below 1.
  • CGQL-MSE and CGQL-Linex use alternative regression objectives for approximating Qt; the Linex objective has the same optimal solution as the desired Q, while MSE yields a lower bound.

G THEORETICAL GUARANTEES

The theoretical result establishes a unique critical point for the adjoint-matching objective, whose velocity field generates the desired behavior-regularized policy.

  • Theoretical guarantees: Proposition 1 extends prior adjoint-matching theory to state-conditioned policy optimization and establishes a unique velocity field.The proof adds a state-conditioned step to the earlier result.
  • Theoretical guarantees: The state-conditioned objective differs from the prior formulation only by a factor of 2, so both have identical critical points.
  • Theoretical guarantees: For each fixed state, the loss has a single critical point whose velocity field generates the target action distribution.
  • Theoretical guarantees: Across states, the global critical point is the Cartesian product of the unique per-state critical points.
  • Theoretical guarantees: The resulting converged policy coincides with the optimal behavior-constrained policy π⋆(a | s) ∝ πβ(a | s) exp(τQ(s, a)).
Loading 2601.14234v4…