PaperlayerView exact revision on arXiv

Source-linked AI summary

Q-Learning with Scalar Adjoint Matching

Yonghoon Dong, Minsung Yoon, Jaehyuk Kim, Jungwoo Park, Changyeon Kim, Jinwoo Shin

arXiv:2610.10437v1cs.LGcs.AIcs.RO

TL;DR

Flow-policy fine-tuning against learned value functions is challenging because actions arise over many flow steps and adjoint matching requires costly per-step vector–Jacobian products. SQAM uses a closed-form scalar adjoint and a value penalty at policy-generated actions. It improves results on the hardest OGBench domains and on three real-robot tasks.

  • Problem

    Adjoint matching can directly update a flow policy using learned value functions, but its per-step vector–Jacobian products make computation grow with the number of flow steps and policy size.

    Source: §1 · p. 1

  • Method

    SQAM replaces those products with a closed-form scalar adjoint and adds a value penalty at policy-generated actions.

    Source: §1 · pp. 1–2

  • Results

    Across 50 OGBench tasks, SQAM outperforms prior methods in offline and offline-to-online RL; on the four hardest domains, it adds 18 to 35 percentage points of offline success rate over the best baseline in each.

    Source: §1 · p. 2

  • Takeaways & Limitations

    SQAM also extends to offline-to-online fine-tuning of a vision-language-action policy on a real bimanual robot, improving over supervised fine-tuning on all three manipulation tasks.

    Source: §1 · p. 2, §6 · p. 9

    SQAM’s effectiveness depends on the value-learning scheme, and a flow-policy-specific value-learning method may make that dependence more reliable than adding a penalty to an existing critic.

    Source: H. Limitations and future work · p. 38

Abstract

from arXiv · show

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.

1. Introduction

Flow-policy fine-tuning with learned value functions is costly when updates propagate through every flow step. SQAM replaces those per-step products with a scalar adjoint and pairs it with a value penalty at policy-generated actions, improving results on hard OGBench domains and real-robot tasks.

1. Introduction

  • Flow policies model rich, multimodal action distributions and are used in offline and offline-to-online RL as well as in large pretrained robot policies.Fine-tuning them with off-policy RL aims to improve beyond their training data using a learned action-value function.Source: §1 · p. 1
  • Fine-tuning flow policies is difficult because actions are generated over multiple sampling steps; adjoint matching updates the flow model directly but requires a policy vector–Jacobian product at every step.The resulting computational cost grows with both the number of flow steps and policy size; alternatives that freeze the policy rely on a separate component and can limit expressiveness.Source: §1 · p. 1
  • The scalar adjoint approximates the batch-averaged velocity Jacobian as a scalar multiple of the identity, yielding a closed-form update that scales the final-action value gradient by flow time without policy vector–Jacobian products.It needs a single critic gradient at the sampled action and preserves the qualitative trends of exact propagation, as reported in Section 4.2.Source: §1 · p. 1
  • SQAM pairs the scalar adjoint with a value penalty at policy-generated actions because scalar-adjoint effectiveness depends particularly on the critic’s value at those actions.Regularization at those actions improves performance, whereas dataset-only and action-distance penalties provide little improvement.Source: §1 · pp. 1–2
  • On the four hardest OGBench domains, SQAM adds 18 to 35 percentage points of offline success rate over the best baseline in each, and improves over supervised fine-tuning on all three robot tasks.Across 50 OGBench tasks, it outperforms prior methods in offline and offline-to-online RL; overall, it reaches 75% against 68% for the strongest baseline.Source: §1 · p. 2

2. Background

The background introduces flow matching as a way to transport samples with a learned velocity field, then describes how flow policies use that process to generate actions. It also sets out the off-policy actor-critic setting and adjoint matching’s propagation of value gradients through the flow steps.

2. Background

  • A flow policy generates an action by integrating its state-conditioned velocity field from Gaussian noise and taking the endpoint of the flow.The described flow-matching setup uses a linear interpolation path between source and target distributions.Source: §2 · pp. 2–3
  • The actor-critic setup learns an action-value function, or critic, and improves the policy by maximizing that function.The background frames reinforcement learning as maximizing expected discounted return in a Markov decision process.Source: §2 · p. 3
  • In off-policy fine-tuning, replay transitions come from a buffer initialized with behavior-policy data, and online rollouts later grow that buffer.The pretrained flow policy is first trained by behavior cloning on the data, then fine-tuned offline and online.Source: §2 · p. 3
  • Adjoint matching propagates the critic’s final-action gradient backward through the flow chain so each intermediate step receives a direction for improving the final action’s value.The adjoint is integrated backward along the sampling chain; Dong et al. (2026c) stabilize optimization with a path-space KL budget.Source: §2 · p. 3

3. SQAM: Q-Learning with Scalar Adjoint Matching

The section derives a scalar approximation to adjoint matching from diagonal concentration in pretrained-flow velocity Jacobians. It then introduces a critic value penalty at policy-generated actions, where the scalar update uses the critic.

3.1. A closed-form scalar adjoint under an isotropic velocity Jacobian

  • The batch-averaged velocity Jacobian concentrates on its diagonal, motivating an isotropic scalar approximation that yields a closed-form adjoint.Under the assumption that the Jacobian is a scalar multiple of the identity, the lean adjoint ODE has a closed-form solution.Source: §3.1 · p. 3, Figure 2 · p. 4
  • The scalar adjoint uses the final-action critic gradient scaled by flow time, requiring one critic gradient and no per-step policy vector–Jacobian products.The exact adjoint instead requires a vector–Jacobian product at every flow step, with cost growing with flow-step count and policy size.Source: §3.1 · pp. 3–4

3.2. Regularizing the critic at the policy’s own actions

  • The authors hypothesize scalar-adjoint failures, especially on manipulation domains, arise from critic errors at policy-generated actions.The scalar adjoint directly uses the critic’s gradient at those actions.Source: §3.2 · p. 4
  • SQAM penalizes critic value at the policy-generated action relative to the dataset-action value, controlling the critic where the actor update uses it.The penalty acts on the action already drawn for the actor update and adds no sampling cost.Source: §3.2 · p. 4
  • Unlike in-sample maximization (Kostrikov et al., 2022; Garg et al., 2023) and behavioral regularization (Tarasov et al., 2023; Fujimoto and Gu, 2021), the value penalty controls the critic at the actor-update action.The other schemes respectively fit dataset actions or constrain action distance; experiments support the value penalty’s advantage.Source: §3.2 · p. 4

4. Experiments

The experiments compare SQAM with prior methods on OGBench, examine the scalar adjoint and value-learning choices, and test extensions to large pretrained policies. SQAM performs best on the hardest benchmark domains, while analyses identify conditions that shape its gains.

4.1. Main Experiments

  • SQAM adds 21 points on antmaze-giant, 18 on humanoidmaze-large, 22 on cube-triple, and 35 on cube-quadruple; overall it reaches 75% versus TRQAM’s 68%.On cube-quadruple, the strongest baseline reaches 19% and SQAM reaches 54%.Source: §4.1 · p. 5, Table 1 · p. 6

4.2. Comparison between Adjoints

  • The exact and scalar adjoints contract similarly and remain positively aligned, but their angle is larger on manipulation domains where the scalar adjoint alone underperforms.The comparison uses the same rollout states and is reported without the value penalty.Source: §4.2 · p. 5, Figure 4 · p. 7

4.3. Analyses over Value Learning Schemes

  • On the two manipulation domains where the scalar adjoint alone underperforms, the value penalty works best; in-sample critic and ReBRAC stay close to the unregularized scalar adjoint.The value penalty acts on the critic’s value at the policy-generated action, unlike the in-sample critic’s dataset-action training and ReBRAC’s action-distance penalty.Source: §4.3 · p. 6
  • The gradient-norm penalty also improves over the unregularized critic, most clearly on cube-quadruple, and is linked to the value penalty through the critic’s gradient at the policy-generated action.The link follows from a first-order expansion of the value penalty.Source: §4.3 · p. 6

4.4. Ablation Studies

  • Increasing the value-penalty coefficient sharply raises success on the two cube manipulation tasks but yields small gains on the two locomotion tasks; the penalty is off on cube-double.A large coefficient lowers cube-double’s offline performance; the chosen coefficient is 0.3 on every other domain.Source: §4.4 · p. 7, Figure 6 · p. 8
  • On most domains, performance stays flat or declines as the KL budget widens, while puzzle-4x4 prefers a larger budget.The budget sets how far the fine-tuned policy may deviate from the pretrained policy.Source: §4.4 · p. 7, Figure 7 · p. 8

4.5. Extension to Large Pretrained Policies

  • On a real-world bimanual robot, SQAM improves over supervised fine-tuning on all three tasks after 15K offline and 10K online steps.The evaluation fine-tunes the pretrained RLDX-1 vision-language-action policy and reports held-out success over 30 episodes per task.Source: §4.5 · pp. 7–8

5. Related work

Prior work steers flow policies with critic information through intermediate actions, adjoint propagation, or auxiliary policies. Critic regularizers differ in where they act, while offline-to-online RL faces value instability and policy forgetting at the transition.

5. Related work

  • Adjoint matching propagates value gradients through a lean adjoint ODE, while other methods guide intermediate actions or learn values at each noise level.Other approaches use a Tweedie posterior estimate, match the policy score to the critic gradient, or train a separate value at each noise level.Source: §5 · p. 8
  • Critic regularizers act on policy-preferred actions, dataset-action value targets, action distance, critic gradients, or calibration to a reference policy.CQL, IQL, ReBRAC, Fisher-BRC, and Cal-QL exemplify these differing regularization choices.Source: §5 · p. 8
  • Offline-to-online RL pretrains a policy and critic on offline data before online fine-tuning, but the transition can destabilize value estimates and lead to policy forgetting.Source: §5 · p. 9

6. Conclusion

SQAM replaces exact adjoint propagation with a closed-form scalar adjoint and pairs it with a value penalty at policy-generated actions. The conclusion reports its strongest OGBench gains on the four hardest domains and transfer to a vision-language-action policy on a real robot.

6. Conclusion

  • SQAM improves most on the four hardest OGBench domains and extends to a vision-language-action policy on a real robot.It replaces the exact adjoint with a closed-form scalar adjoint and pairs this with a value penalty at policy-generated actions.Source: §6 · p. 9
  • The authors suggest that value learning, rather than the policy update, may be the main bottleneck for off-policy RL with flow policies.They state that a rough adjoint approximation suffices while the value-learning scheme largely determines performance.Source: §6 · p. 9

Reproducibility statement

  • The reproducibility statement identifies pseudocode, proofs, benchmark and dataset details, baseline descriptions, and robot protocols in the appendices, and says SQAM’s implementation will be released.Source: Reproducibility statement · p. 10

Ethics statement

  • Robot experiments took place in a controlled laboratory with an operator supervising throughout; they involved no human subjects or personal data.The authors say deployment beyond supervised settings calls for usual safeguards when improving policies against learned value functions.Source: Ethics statement · p. 10

AI use statement

  • The authors report using generative AI for writing, editing, and routine coding and plotting, but not for research ideas, results, or claims.Source: AI use statement · p. 10

A. Full algorithm

  • The full SQAM algorithm includes a trust region inherited from TRQAM, and its scalar adjoint permits loss states to be built from the final action alone.The method samples the final action with an ODE sampler, re-noises it along a linear interpolation path to form bridge states, and uses the trust region by default for stable training.Source: A. Full algorithm · p. 15

B. Baselines

The baselines span frozen-policy edits, noise-space or guidance methods, flow distillation, and adjoint-matching updates. The value-learning comparisons vary where critic regularization acts, including directly at policy-generated actions.

B. Baselines

  • The OGBench comparison includes seven off-policy flow-policy fine-tuning methods, all sampling transitions uniformly from a shared replay buffer.The buffer contains offline data during offline training and grows with collected rollouts online.Source: B. Baselines · p. 16
  • FQL avoids differentiating through the sampling chain by distilling the multi-step flow into a one-step policy.The one-step policy trains jointly with the velocity field, maximizing critic value while an L2 term holds it near the rollout.Source: B. Baselines · p. 16
  • DSRL keeps the flow frozen and learns a noise-space policy, while IFQL selects the highest-valued candidate from a flow policy using an IQL-style value scheme.DSRL also fine-tunes the behavior-cloning velocity online in the offline-to-online setting.Source: B. Baselines · p. 16
  • QAM uses adjoint matching with a fixed temperature, whereas TRQAM adapts the sampling-dynamics trust region using a path-space KL estimate.QAM-E adds an entropy-regularized residual edit policy to QAM-generated actions.Source: B. Baselines · p. 17
  • EXPO leaves its expressive base policy intact and trains a residual edit policy; EXPO-FT applies this recipe to a pretrained vision-language-action policy.CGQL-L combines the behavior-cloning velocity with Q-function guidance based on an auxiliary critic trained on intermediate noisy actions.Source: B. Baselines · pp. 16–17

C. Value learning schemes

  • The value-learning schemes differ in whether they penalize critic values at policy actions, constrain critic gradients, or change the bootstrap target or action distribution.The value penalty compares critic values at policy and dataset actions; the gradient-norm penalty constrains the gradient used by the scalar adjoint, while ReBRAC penalizes action distance in the bootstrap target and the in-sample critic uses expectile-fitted state values.Source: C. Value learning schemes · pp. 17–18

D.1. Domains and tasks

The appendix specifies the OGBench domains and data, then documents the real-robot training and evaluation setup. The robot study uses a dual-arm platform, vision-language-action policy inputs, and matched task protocols for SQAM and EXPO-FT.

D.1. Domains and tasks

  • OGBench evaluation covers 10 reward-based single-task domains, with five tasks per domain across navigation, multi-object manipulation, and combinatorial planning.The benchmark experiments run eight random seeds per method and task; reported tables give mean success rate ± one standard deviation unless noted otherwise.Source: §D.1 · p. 18

D.2. Hyperparameters

  • The appendix reports dataset sizes, episode lengths, and action dimensions by domain, and describes domain-specific hyperparameter selection and evaluation over eight seeds.For antmaze-giant, the KL budget opens from 0.5 offline to 3.0 at the online transition; the value-penalty coefficient is tuned on two tasks per domain.Source: §D.1 · p. 18, §D.2 · p. 19, Table 3 · p. 19

E. Real-robot experiments: RB-Y1 with Wuji Hand 2

E.1. Platform, observations and actions

  • The platform is an RB-Y1 humanoid with two 20-DoF hands; bimanual control uses its two 7-DoF arms and hands, while torso, neck, and mobile base remain unused.A head-mounted ZED Mini stereo camera supplies two 30 Hz RGB views, alongside a 62-dimensional joint-position vector and task instruction.Source: §E.1 · p. 20, Figure 9 · p. 20

E.2. Tasks and evaluation protocol

  • SQAM updates the pretrained policy’s action head on 16 executed steps of a 24-step action chunk, while EXPO-FT keeps the base policy frozen and adds a residual.Each task uses a fixed initial configuration, and checkpoints are scored by operator-judged success over 30 held-out episodes.Source: §E.1 · p. 20, §E.2 · p. 21

E.3. Training details

  • The real-robot comparison modifies EXPO-FT’s released recipe to match the shared protocol, including omitting corrective human actions and using the same mean-pooled VLM feature for both critics.SQAM uses a small actor learning rate; EXPO-FT continues online training from its own 15K-step offline checkpoint.Source: §E.2 · p. 21, §E.3 · p. 21

E.4. Additional results

The additional results examine robot-task outcomes, update cost, the value penalty, trust-region behavior, and the batch-averaged velocity Jacobian. They report persistent computational savings from the scalar adjoint and domain-dependent effects of removing the trust region or adding the penalty.

E.4. Additional results

  • On all three robot tasks, SQAM rollout success tends to rise during online training; successful episodes are faster than SFT on flip-plastic-bag and put-and-close, while place-straw is comparable.The evaluation frames also show fewer incomplete flips under SQAM than SFT on flip-plastic-bag.Source: §E.4 · p. 21, Figure 10 · p. 22

F.1. The cost of one update

  • The exact adjoint remains about half again as slow at every actor width, and its update-time gap from the scalar adjoint grows with the number of flow steps.The exact adjoint uses a vector–Jacobian product through the policy at each flow step; the scalar adjoint uses none.Source: §F.1 · p. 25, Figure 15 · p. 25

F.2. The value penalty inside TRQAM

  • At the offline endpoint, the value penalty raises TRQAM by 18 points on cube-triple and 23 on cube-quadruple, versus 48.5 and 78 points for SQAM.The comparison uses four seeds and keeps the trust region active in all configurations.Source: §F.2 · p. 25, Figure 16 · p. 26

F.3. Without the trust region

  • Without the trust region, the best constant λ varies by two orders of magnitude across domains; with it, the KL budget affects performance more predictably and training remains stable.For example, λ=0.1 is best on antmaze-giant but fails on puzzle-3x3, while λ=10 is best on puzzle-3x3 but fails on antmaze-large.Source: §F.3 · p. 25

F.4. Batch-averaged velocity Jacobian on all domains

  • Across every measured domain, the batch-averaged velocity Jacobian of the pretrained flow is dominated by its diagonal at four flow times.Figures 18 and 19 extend the Jacobian analysis to all ten domains.Source: §F.4 · p. 26

F.5. Ablation studies on all domains

Across ten domains, the ablations show that the value penalty matters most where the scalar adjoint alone underperforms, while alternative value-learning schemes have mixed or negative effects. A small KL budget generally performs best or comparably, with a stated domain exception.

F.5. Ablation studies on all domains

  • A small KL budget performs best or equally well on most domains, but puzzle-4x4 prefers a larger budget; the value penalty yields large gains on cube-triple and cube-quadruple and little gain on most other domains.On cube-double, c=0.3 lowers offline success, so the main experiments omit the penalty there.Source: §F.5 · p. 26

F.6. Value learning schemes on all domains

  • The value-penalty sweep completes the all-domain comparison of value-learning schemes, with its largest gains concentrated on cube-triple and cube-quadruple.The value-penalty trends are reported alongside the gradient-norm, ReBRAC, and in-sample-critic sweeps.Source: §F.5 · p. 26, §F.6 · p. 26
  • The gradient-norm penalty helps most visibly on puzzle-3x3 and cube-quadruple but degrades several domains at its largest coefficient; ReBRAC stays near the unregularized critic except at its largest β, while the in-sample critic falls below it on most domains.The sweeps cover all ten domains and compare each regularization scheme with an unregularized-critic reference.Source: §F.6 · p. 26

G. Full OGBench results

The full-results appendix provides offline success rates for all 50 tasks and plots offline-to-online training curves across the domains. It also presents parameter sweeps and Jacobian visualizations, whose captions specify their scope and layout.

G. Full OGBench results

  • The appendix also sweeps the KL budget, value-penalty coefficient, and alternative critic regularizers across ten domains, and compares constant-λ SQAM with trust-region SQAM on representative tasks.The captions specify the parameter ranges and seed counts but do not establish a best setting.Source: Figure 17 · p. 27, Figure 20 · p. 30, Figure 21 · p. 31
  • Table 8 reports offline success rates at 1M training steps for all 50 tasks, with means and standard deviations over eight seeds.It is the per-task breakdown behind Table 1.Source: G. Full OGBench results · p. 26, Table 8 · p. 35
  • Figures 25 and 26 plot success-rate curves for the locomotion, scene, cube, and puzzle domains, with five tasks per domain and a dashed marker at the transition from 1M offline to 0.5M online steps.The curves show eight seeds per method, with bands indicating one standard deviation.Source: G. Full OGBench results · p. 26, Figure 25 · p. 36, Figure 26 · p. 37

H. Limitations and future work

The authors identify open questions about value learning for flow policies, multi-task fine-tuning, and why pretrained flow-policy Jacobians concentrate on their diagonals. Their real-robot experiments fine-tune each task separately, and they leave the structural explanation for future work.

H. Limitations and future work

  • SQAM adds a value penalty to a standard critic, while a value-learning method designed for flow policies from the ground up remains an open direction.The authors suggest such a method may make the scalar adjoint’s dependence on value learning more reliable.Source: H. Limitations and future work · p. 38
  • The real-robot experiments fine-tune the policy separately on each task, leaving single-policy multi-task fine-tuning untested.The authors identify multi-task fine-tuning as a more direct test of the potential benefit of updating the generative policy itself.Source: H. Limitations and future work · p. 38
  • The authors observe diagonal concentration of the batch-averaged Jacobian on every measured domain but do not explain why or establish when it arises.They propose theoretical analysis to clarify the conditions under which the scalar adjoint applies.Source: H. Limitations and future work · p. 38

Related papers

Loading 2610.10437v1…