Q-learning with Adjoint Matching

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.

• The paper introduces QAM, which leverages adjoint matching to extract critic gradients without backpropagation instability over full denoising processes.
• It reformulates policy optimization into a step-wise stochastic control objective, bypassing the numerical challenges of direct backpropagation in flow models.
• Empirical evaluations on 50 tasks show that QAM outperforms prior methods in both offline and online RL with enhanced stability and sample efficiency.

Q-learning with Adjoint Matching: Stable and Expressive Policy Optimization for RL with Flow Models

Introduction and Motivation

The integration of expressive policy distributions, such as those parameterized by diffusion and flow-matching models, has become standard in deep reinforcement learning (RL) for continuous action spaces. Expressive policies are crucial for solving complex tasks with sparse rewards and multi-modal action structures. However, optimizing such policies using a learned Q-function (critic) presents a longstanding challenge. Specifically, optimizing flow or diffusion policies with respect to a Q-function requires leveraging first-order critic information (i.e., $\nabla_a Q(s,a)$), but direct backpropagation through the full denoising or flow process is numerically unstable. Prior methods usually either ignore the Q-gradient and optimize only the critic value at the output action, or introduce substantial approximation and bias, sacrificing either learning efficiency or policy expressivity. This paper introduces Q-learning with Adjoint Matching (QAM), which leverages adjoint matching—a technique originally developed for generative modeling—to enable high-fidelity, gradient-based policy extraction in RL without backpropagation instability while preserving the full expressivity of flow policies (2601.14234).

Figure 1: QAM: Q-learning with Adjoint Matching. Left: QAM uses adjoint matching to leverage the critic's action gradient directly for expressive policy optimization.

Method: QAM for Gradient-based Policy Optimization

Problem Setup and Limitation of Direct Optimization

In continuous-action RL, flow and diffusion policies are typically optimized as stochastic optimal control (SOC) objectives with a prior constraint, often a KL-regularization w.r.t. a behavior policy. The optimal prior-regularized policy with respect to a critic $Q_\phi(s,a)$ is:

$$\pi^*(a|s) \propto \pi_\beta(a|s) \exp(\tau Q_\phi(s,a))$$

where $\pi_\beta$ is a behavior policy.

The standard approach would parameterize $\pi_\theta$ as a flow-based model and directly solve the resulting SOC, using backpropagation through the multi-step flow (denoising) process. However, this is numerically unstable and can suffer from poor optimization landscapes, particularly when action gradients through deep generative models become ill-conditioned.

Adjoint Matching Formulation

QAM circumvents these issues by utilizing adjoint matching [domingo-enrich2025adjoint]. In adjoint matching, the loss landscape is rearranged such that stochastic optimal control objectives can be optimized step-wise, decomposing the process to avoid direct backpropagation through the full trajectory. The key insight is that the critic's action gradient is only required at the denoised action, and adjoint state propagation can be performed using the (fixed) base flow model, not the model being optimized.

For QAM, the learning objective for the policy becomes:

$$L_{\mathrm{AM}}(\theta) = \mathbb{E}_{s, \{a_t\}_t} \left[\int_0^1 \left\| \frac{2(f_\theta(s, a_t, t) - f_\beta(s, a_t, t))}{\sigma_t} + \sigma_t \tilde{g}_t \right\|_2^2 dt \right],$$

where:

• $f_\theta$ and $f_\beta$ are the trainable ("fine-tuned") and behavior policy’s flow velocity fields respectively,
• $\tilde{g}_t$ is the "lean" adjoint state, initialized with $-\tau \nabla_{a_1} Q_\phi(s, a_1)$ at $t=1$, and backward-propagated via reverse-mode ODEs using only $f_\beta$,
• $\sigma_t$ is the noise schedule for the SDE discretization.

This approach provides an unbiased, stable gradient for training expressive policies without either discarding first-order information or yielding to the instability of deep backpropagation through time.

Algorithmic Schematic

The overall learning process alternates between optimizing the critic via temporal-difference backups and the policy via adjoint matching. The computational pipeline involves samples from the replay buffer (offline or a mixture with online data), rollout through the SDE using $f_\theta$, calculation of adjoint states using only $f_\beta$, and loss minimization on the policy. For sample efficiency and stability, an ensemble of critics and gradient clipping are used.

Empirical Evaluation

Offline RL Performance and Compared Baselines

QAM is evaluated on 50 challenging tasks from the OGBench benchmark, which includes diverse long-horizon sparse-reward domains (scene, puzzle, Rubik’s Cube, humanoid and ant mazes). Empirical results demonstrate that QAM and its variants (QAM-EDIT, QAM-FQL) consistently outperform all prior diffusion/flow-based, Gaussian-policy, advantage-weighted, and classifier-guidance baselines in offline RL (aggregate top-1 task scores: QAM-E 46/50; best prior baseline: 42/50 for QSM).

Online Fine-tuning and Robustness

QAM is further shown to enable robust and sample-efficient offline-to-online RL: after offline pretraining, continued fine-tuning achieves higher and more consistent scores than state-of-the-art alternatives across manipulation and navigation domains.

Figure 2: QAM online fine-tunes more effectively across domains compared to baselines; sample efficiency is highlighted.

Sensitivity and Ablation

Analysis includes hyperparameter sensitivity for gradient clipping, number of flow steps, critic ensemble size, and temperature coefficient $\tau$. The method is shown to be robust across a broad range of configurations, with the temperature parameter requiring domain-dependent tuning for best performance.

Figure 3: QAM-EDIT exhibits robust performance across sensitivity variations for critical hyperparameters.

Training Stability and Data Quality

QAM demonstrates stability to data corruption/augmentation, maintaining superior aggregate scores on action-noisy and "stitched" datasets where prior methods degrade substantially.

Figure 4: QAM is robust to data quality degradation compared with baselines, with little performance loss on corrrupted datasets.

Theoretical Guarantees

QAM inherits theoretical optimality properties provided by adjoint matching. When the policy extraction objective is optimized to convergence, the learned policy provably coincides with the optimal behavior-constrained solution. Unlike prior methods that bias policy updates or use value-only proxies, QAM’s adjoint matching objective ensures that the policy remains expressive and unbiased, conditional on the critic’s quality.

Extensions of QAM to relax the strict KL constraint use Wasserstein regularization or editing policies, allowing exploration outside the narrow support of the behavior policy when offline data is suboptimal or narrow.

Implications and Future Directions

The QAM framework fundamentally addresses the optimization-expressivity trade-off in learning policies for continuous-action RL with flow/diffusion models. Practically, it enables high-performing, stable, and sample-efficient policy extraction in domains that demand multimodal, expressive action distributions, such as robotic manipulation, planning, and long-horizon navigation.

On the theoretical side, QAM establishes a template for integrating recent advances in stochastic optimal control and generative modeling into RL, offering provable policy improvement without backpropagation pathologies.

Key directions for future research include:

• Principles for automatic tuning of behavior/prior regularization when the offline dataset's coverage is poor,
• Extensions to hierarchical and compositional RL using flow-based policy structures,
• Deployment and validation at scale in real-world robotic and simulated domains with partial observability and action chunking.

Conclusion

Q-learning with Adjoint Matching defines a principled, practical pathway for stable, expressive, and effective policy optimization in RL with flow and diffusion models. By leveraging adjoint matching, it uniquely balances the exploitation of critic gradients and expressive policy parameterizations, outperforming prior methods both in offline and offline-to-online RL settings. The demonstrated empirical sample efficiency and robustness, combined with strong theoretical guarantees, position QAM as a state-of-the-art approach in continuous-action RL, with significant implications for scaling up deep RL in real-world tasks and for further synthesis between generative modeling and control.