Q-learning with Adjoint Matching
This addresses a long-standing problem in RL for researchers and practitioners by enabling stable, unbiased policy optimization without sacrificing expressivity, though it builds on existing techniques like adjoint matching.
The paper tackled the challenge of efficiently optimizing expressive diffusion or flow-matching policies in continuous-action reinforcement learning by proposing Q-learning with Adjoint Matching (QAM), which outperformed prior methods on hard, sparse reward tasks in offline and offline-to-online RL.
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.