Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Q-learning with Adjoint Matching

Qiyang Li, Sergey Levine

TL;DR

QAM addresses the challenge of leveraging the critic's action gradient to optimize expressive flow/diffusion policies in TD-based RL by applying adjoint matching to derive a stable, step-wise objective that preserves the correct tilt distribution ${p_\theta( X_1 ) \propto p_\beta( X_1 ) e^{Q_\phi(X_1)}}$. It combines this policy extraction with TD updates over an ensemble of critics, providing theoretical guarantees that the learned policy converges to the behavior-constrained optimum ${\pi^*(a|s) \propto \pi_\beta(a|s) \exp(\tau Q_\phi(s,a))}$. Empirically, QAM and its variants outperform strong baselines on offline RL and offline-to-online RL across long-horizon, sparse-reward tasks, including setups with action chunking, while maintaining stability through lean adjoint states. The approach offers a scalable, principled path to deploying highly expressive diffusion/flow policies in practical RL scenarios with limited data.

Abstract

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.

Q-learning with Adjoint Matching

TL;DR

QAM addresses the challenge of leveraging the critic's action gradient to optimize expressive flow/diffusion policies in TD-based RL by applying adjoint matching to derive a stable, step-wise objective that preserves the correct tilt distribution . It combines this policy extraction with TD updates over an ensemble of critics, providing theoretical guarantees that the learned policy converges to the behavior-constrained optimum . Empirically, QAM and its variants outperform strong baselines on offline RL and offline-to-online RL across long-horizon, sparse-reward tasks, including setups with action chunking, while maintaining stability through lean adjoint states. The approach offers a scalable, principled path to deploying highly expressive diffusion/flow policies in practical RL scenarios with limited data.

Abstract

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.
Paper Structure (13 sections, 1 theorem, 91 equations, 8 figures, 7 tables, 1 algorithm)

This paper contains 13 sections, 1 theorem, 91 equations, 8 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Q-learning with Adjoint Matching (QAM)
  5. Experiments
  6. Results
  7. Discussion
  8. Full Results
  9. Additional Discussions for Related Work
  10. Domain and experiment details
  11. Baselines
  12. Hyperparameters
  13. Theoretical guarantees

Key Result

Proposition 1

Take $L_{\mathrm{AM}}(\theta)$ in eq:am, there is a unique $f_\theta$ such that and for all $s \in \mathrm{supp}(D)$,

Figures (8)

  • Figure 1: QAM: Q-learning with Adjoint Matching.Left: QAM uses adjoint matching domingo-enrich2025adjoint that leverages the critic's action gradient directly to fine-tune a flow policy towards the optimal behavior-constrained policy: $\pi_\theta(\cdot \mid s) \propto \pi_\beta(\cdot\mid s) e^{Q(s, \cdot)}$. Right: Aggregated score for offline RL on 50 OGBench ogbench_park2024 tasks.
  • Figure 2: QAM online fine-tunes more effectively than prior methods (50 tasks, 12 seeds). We use the QAM-EDIT (QAM-E) variant for QAM and the CGQL-Linex variant for CGQL (CGQL-L) because they work well for online fine-tuning. For results on individual tasks and full training curves offline, see \ref{['fig:full-main']} in the Appendix.
  • Figure 3: Sensitivity analysis of QAM-EDIT on the representative task from each domain (12 seeds). For all locomotions tasks (e.g., antmaze-{large, giant}, humanoidmaze-{medium, large}), we use task1. For all manipulation tasks (e.g., scene, p33, p44, c2, c3, c4), we use task2. We report the aggregated performance across these 10 representative tasks. Gradient Clipping: whether to use gradient clipping in our optimizer; Flow Steps ($T$): this parameter indicates the number of numerical integration steps that we use for the flow model. Critic Ensemble Size ($K$): number of critic network in the ensemble; Temperature ($\tau$): the parameter that modulates the influence of the prior. We rerun our method with $0.1\times, 0.3\times, 3\times$, and $10\times$ the best $\tau$ from tuning. For results on individual tasks and full training curves offline, see \ref{['fig:full-sensitivity']} in the Appendix.
  • Figure 4: Data quality analysis (12 seeds).Top: performance on the original {navigate/play}-style datasets we use in our main experiments; Bottom: performance on the {stitch/noisy}-style datasets. Each subplot reports the aggregated score over 5 tasks in each domain. For results on individual tasks and full training curves offline, see \ref{['fig:full-quality']} in the Appendix.
  • Figure 5: Full training curves for our main results (12 seeds).
  • ...and 3 more figures

Theorems & Definitions (2)

  • Proposition 1: Extension of Proposition 7 in domingo-enrich2025adjoint to Policy Optimization.
  • proof