Q-Learning with Scalar Adjoint Matching

AI summary

Overview

Research area: Reinforcement learning, specifically off-policy fine-tuning of flow-matching (generative) policies for robot control, with an emphasis on value-based actor-critic learning.

Technical level: Advanced. The paper builds on flow matching, the lean adjoint ODE, vector–Jacobian products, stochastic optimal control, and path-space KL trust regions. Readers need comfort with diffusion/flow generative models and off-policy actor-critic RL.

Scope (one sentence): The paper derives a closed-form "scalar adjoint" that removes the per-flow-step vector–Jacobian product from adjoint matching, pairs it with a value penalty on the critic at policy-generated actions, and evaluates the resulting algorithm (SQAM) on 50 OGBench tasks and on a real bimanual robot.

What This Paper Is About

Flow policies can represent rich, multimodal action distributions and serve as the backbone of recent large pretrained robot policies, but improving them beyond their demonstration data requires fine-tuning against a learned value function. The standard adjoint-matching approach for doing this propagates the critic's value gradient back to every intermediate flow step, which requires a vector–Jacobian product through the policy at each step and therefore scales poorly with the number of flow steps and policy size. SQAM's goal is to make that propagation cheap and closed-form while keeping the fine-tuned policy competitive.

Key Contributions

1. A closed-form scalar adjoint. Motivated by the empirical observation that the batch-averaged velocity Jacobian of pretrained flow policies concentrates on its diagonal, the authors approximate the Jacobian as a scalar multiple of the identity and derive a closed-form solution to the lean adjoint ODE: the adjoint at flow time τ is simply the critic's final-action gradient scaled by τ. This eliminates vector–Jacobian products entirely.

2. Identification of the value-learning dependency. The authors show that the scalar adjoint's performance depends on how the critic is learned, and isolate the critic's value at policy-generated actions (rather than at dataset actions or in terms of action distance) as the decisive factor.

3. The SQAM algorithm. SQAM combines the scalar adjoint with a value penalty that penalizes the critic's value at the policy-generated action relative to the value at the dataset action. It inherits the optional path-space KL trust region from TRQAM and uses bridge states re-noised along the linear interpolation path instead of a long stochastic rollout.

4. Empirical validation in simulation and on hardware. Gains concentrate on the four hardest OGBench domains (18 to 35 percentage points over the strongest baseline per domain), overall offline success is 75% versus 68% for the strongest baseline, and SQAM improves over supervised fine-tuning on all three real-robot manipulation tasks with a vision-language-action policy.

Main Findings

• The pretrained velocity Jacobian is diagonally dominated. Averaged over a batch of states, the velocity Jacobian of the pretrained flow policy is dominated by its diagonal. The paper shows this on four domains at two flow times (Figure 2) and states that all ten domains are covered in Section F.4.

• The scalar adjoint tracks the exact adjoint qualitatively. Figure 4 (three seeds per domain at 0.3M steps, without the value penalty) shows that both adjoints contract similarly in norm as flow time approaches zero, and that the scalar adjoint stays positively aligned with the exact adjoint at every flow time. The angle between them is larger on the two manipulation domains than on the two locomotion domains — precisely where the scalar adjoint alone underperforms.

• The scalar adjoint alone is insufficient. Using the scalar adjoint without regularization does not work well, especially on manipulation domains (Section 4.3).

• The value penalty is the regularization that matters. Among the compared value-learning schemes — an IQL-style in-sample critic, the ReBRAC bootstrap penalty, a critic action-gradient-norm penalty, and the proposed value penalty — the value penalty performs best on cube-triple and cube-quadruple. The in-sample critic and ReBRAC stay close to the unregularized scalar adjoint. The gradient-norm penalty also improves over the unregularized critic, most clearly on cube-quadruple.

• The value penalty helps the exact adjoint much less. Applying the value penalty to TRQAM (which uses the exact adjoint) improves it as well, but "far less than it improves SQAM" (Section F.2).

• Main OGBench results. Offline fine-tuning for 1M steps with 8 seeds over 50 tasks in 10 domains: SQAM adds 21 points over the best baseline on antmaze-giant, 18 on humanoidmaze-large, 22 on cube-triple, and 35 on cube-quadruple. On cube-quadruple the strongest baseline reaches 19 and SQAM reaches 54. Over all fifty tasks SQAM reaches 75% versus 68% for TRQAM. Online runs continue to 1.5M steps and the trend holds.

• Value penalty coefficient behavior. On the two cube manipulation tasks success rises sharply as the coefficient c grows, while on two locomotion tasks gains are small. The authors use c = 0.3 on every domain except cube-double, where the penalty is off because a large c lowers offline performance (Section F.5).

• KL budget behavior. On most domains performance stays flat or decreases steadily as the budget widens, making a small budget a reliable default. The exception is puzzle-4x4, which prefers a larger budget, consistent with its larger state space.

• Real-robot results. Fine-tuning RLDX-1 (a pretrained vision-language-action policy) on a real bimanual robot with two 20-DoF dexterous hands and an action chunk of 16 steps: after supervised fine-tuning, then 15K offline steps and a further 10K online steps, evaluated over 30 held-out episodes per task, SQAM reaches 25/30, 15/30, and 13/30 (offline-to-online) versus SFT's 15/30, 7/30, and 4/30. EXPO-FT, a residual fine-tuning method, reaches 15/30, 0/30, and 3/30 offline-to-online. SQAM improves over SFT on all three tasks; EXPO-FT does not appear to work well in this setting.

• Compute comparison. The paper states that Section F.1 compares the computational cost of the scalar and exact adjoint updates; specific cost figures are not reported in the content provided.

Methodology in Plain English

A flow policy produces an action by starting from random noise and integrating a learned velocity field over many steps until an action emerges. The critic only scores the final action, so to improve the policy you need to tell each intermediate step which way to move. Adjoint matching does this by running an "adjoint" ODE backward from the final action's value gradient, but each step requires a Jacobian of the velocity field times a vector — expensive and growing with the number of flow steps and the network size.

The authors first inspect that velocity Jacobian averaged across a batch of states and find it is dominated by its diagonal — the off-diagonal entries are small. If you assume the Jacobian is just a scalar times the identity, the adjoint ODE collapses to a single scalar ODE per coordinate, which can be solved in closed form: the adjoint at time τ equals τ times the value gradient at the final action. Since estimating that scalar would reintroduce the Jacobian, they drop the exponential factor and keep only the τ scaling, calling it the scalar adjoint. This requires exactly one critic gradient at the sampled final action and no Jacobian products.

Because the scalar adjoint now leans directly on the critic's gradient at the action the policy itself produced, errors in the critic at that specific action feed straight into the policy update. The authors therefore add a penalty that pushes the critic's value at the policy-generated action to be no larger than its value at the dataset action — a penalty applied at the very action the actor already sampled, so no extra sampling is required. They keep TRQAM's path-space KL trust region, estimated via a path-space KL surrogate and tuned by dual descent, and replace the long, numerically unstable stochastic rollout for adjoint states with "bridge states": sample the endpoint with the ODE sampler, then re-noise along the linear interpolation path with fresh Gaussian noise.

On OGBench, all methods start from the same flow policy pretrained by behavior cloning for 300K steps, are fine-tuned offline for 1M steps, and continue online to 1.5M steps; manipulation domains use action-chunked policies of chunk size 5. Baselines are FQL, CGQL-L, DSRL, IFQL, QAM, QAM-E, and TRQAM.

Why This Matters

• Impact on research. The paper argues that a rough approximation of the adjoint suffices and that the value-learning scheme largely determines performance, suggesting the bottleneck of off-policy RL with flow policies may lie in the critic rather than the policy update. It also makes adjoint matching scale more gracefully, since the cost no longer grows with the number of flow steps and policy size — relevant as pretrained flow policies grow larger.

• Real-world applications:

  • Fine-tuning large pretrained robot policies (such as vision-language-action models) on new manipulation tasks after supervised fine-tuning plateaus.
  • Dexterous manipulation with multi-fingered hands, as demonstrated with two 20-DoF hands on a bimanual platform.
  • Long-horizon goal-conditioned control, where the benchmark's hardest domains (antmaze-giant, humanoidmaze-large, cube-triple, cube-quadruple) show the largest gains.
  • General offline-to-online RL pipelines where a policy is pretrained on demonstrations and then improved from its own rollouts.

• Industry relevance. Robot foundation-model developers who want to keep improving a deployed policy from its own experience without abandoning the pretrained generative model stand to benefit, along with anyone paying for per-update compute in flow-based RL. The paper's result that a residual edit policy struggled on the real robot while direct flow-policy updates succeeded is directly relevant to practitioners choosing between the two paradigms.

Future Directions

1. Better critic learning for flow policies. The authors explicitly conclude by hoping their findings encourage further work on value learning for flow policies, given that the value penalty was the decisive ingredient.

2. Scaling to larger and more capable pretrained policies. The real-robot study covers one vision-language-action policy (RLDX-1) on three tasks; whether the pattern holds for larger models and broader task suites is open.

3. Characterizing when the diagonal-Jacobian approximation holds. The approximation is justified empirically on the studied domains, and the angle between the scalar and exact adjoints is larger on manipulation domains — understanding when the approximation degrades would sharpen the method.

4. Quantifying the compute savings in practice. The paper reports a comparison of computational cost in Section F.1, but the concrete numbers are not given in the content provided; wall-clock and memory savings at large policy scale remain an empirical question.

Target Audience

Researchers and engineers working on off-policy reinforcement learning with generative (flow and diffusion) policies, offline-to-online RL, and robot learning from demonstrations. It will be most valuable to readers already familiar with adjoint matching, flow matching, or diffusion policy fine-tuning who want a cheaper, more scalable policy-update rule and a clear demonstration of why critic regularization at policy-generated actions matters. Readers seeking an introductory treatment of flow-based RL would need substantial background first.

Authors’ abstract

Flow policies capture rich and diverse action distributions, and fine-tuning them with off-policy RL to improve beyond the demonstrations has drawn growing interest. However, fine-tuning a flow policy against a learned value function is not trivial, because the policy generates its action over many flow steps. Adjoint matching offers a principled way to update the flow model itself by propagating value information from the final action back to each flow step, but it requires a vector--Jacobian product through the policy at every step, a cost that grows with the number of flow steps and the policy size. We observe that the batch-averaged velocity Jacobian of pretrained flow policies concentrates on its diagonal. Motivated by this finding, we derive a closed-form scalar adjoint that scales the value gradient at the final action by the flow time, eliminating the per-step vector--Jacobian products. We further find that controlling the critic's value at policy-generated actions is particularly important under the scalar adjoint. Based on these findings, we propose Q-learning with Scalar Adjoint Matching (SQAM), which combines the scalar adjoint with a value penalty at those actions. SQAM's gains concentrate on the four hardest OGBench domains, where its success rate exceeds that of the strongest baseline in each domain by 18 to 35 percentage points. To test whether SQAM extends to large pretrained policies, we also fine-tune a vision-language-action policy on a real bimanual robot. SQAM improves over supervised fine-tuning on all three tasks.