Q-Learning with Scalar Adjoint Matching

Q-Learning with Scalar Adjoint Matching. By Yonghoon Dong and five others, from KAIST and RLWRLD. Published on the seventh of October.

Flow policies capture rich and diverse action distributions, and there is growing interest in fine-tuning them with off-policy reinforcement learning so they can improve beyond the demonstrations they were trained on. Doing that against a learned value function is not trivial. A flow policy does not emit its action in one go: it generates it over many flow steps, and the value function only speaks to the final action at the end of that chain.

Adjoint matching is a principled way to bridge the distance. It updates the flow model itself by propagating value information from the final action back to each flow step. The cost is a vector-Jacobian product through the policy at every single step, and that cost grows with both the number of flow steps and the size of the policy. For large policies with many steps, the principled route becomes an expensive one.

The paper starts from an observation about pretrained flow policies. Their velocity Jacobian, averaged over a batch, concentrates on its diagonal. Motivated by that finding, the authors derive a closed-form scalar adjoint. Instead of carrying the value gradient back step by step through the network, it scales the value gradient at the final action by the flow time. The per-step vector-Jacobian products disappear.

A second finding shapes the method. Under the scalar adjoint, controlling the critic's value at actions the policy itself generates turns out to matter a great deal. So the proposed algorithm, Q-learning with Scalar Adjoint Matching, or SQAM, pairs two pieces: the scalar adjoint for updating the flow policy, and a value penalty applied at those policy-generated actions. The first piece removes the expensive computation. The second controls the critic's value at exactly the actions the policy produces.

The gains are concentrated rather than uniform. They sit on the four hardest OGBench domains, where SQAM's success rate exceeds that of the strongest baseline in each domain by between eighteen and thirty-five percentage points. The comparison is made domain by domain, against whichever baseline was strongest there.

To check whether the approach carries over to large pretrained policies, the authors also fine-tune a vision-language-action policy on a real bimanual robot. There, SQAM improves over supervised fine-tuning on all three tasks.

The arc runs from measurement to simplification. Look at what the Jacobian of a pretrained flow policy actually is, notice that its mass sits on the diagonal, and an expensive per-step backward pass collapses into a single scaling by flow time. The value penalty is the companion that shortcut was found to need, and the robot experiment is the test that it survives at the scale of a vision-language-action model.