PaperlayerView exact revision on arXiv

Source-linked AI summary

Trust Region Q Adjoint Matching

Yonghoon Dong, Kyungmin Lee, Changyeon Kim, Jaehyuk Kim, Jinwoo Shin

arXiv:2605.27079v1cs.LGcs.AIcs.RO

TL;DR

Off-policy fine-tuning of pretrained flow policies is unstable because multi-step sampling and critic errors can cause destructive policy drift. TRQAM enforces a path-space KL trust region through adaptive SOC dynamics, outperforming prior methods across 50 OGBench tasks with a 68% offline-RL success rate versus 46% for the strongest baseline.

  • Problem

    Off-policy fine-tuning must improve pretrained flow policies despite unstable multi-step optimization and critic-guided updates that can cause destructive drift.

  • Method

    TRQAM adaptively updates a trust-region parameter λ inside SOC sampling dynamics, where Girsanov makes path-space KL an exact function of λ.

  • Results

    Across 50 OGBench tasks, TRQAM consistently outperforms prior methods in offline and offline-to-online RL, achieving 68% offline-RL success versus 46% for the strongest baseline.

  • Takeaways & Limitations

    TRQAM provides sampling-level control over deviation from pretrained policies and remains stable on Robomimic, where fixed-temperature adjoint matching collapses.

  • Takeaways & Limitations

    Computing the adjoint matching loss requires a vector-Jacobian product through the velocity field at each backward-ODE step, with cost scaling with model size.

Abstract

from arXiv · show

Off-policy reinforcement learning of pretrained flow policies remains challenging due to the instability of optimization arising from the multi-step sampling process. Recently, Q-learning with Adjoint Matching (QAM) addressed this issue by reformulating into a memoryless stochastic optimal control (SOC) problem with a learned critic. However, QAM inherits a fundamental fragility of critic-guided improvement: small critic errors are amplified when critics are ill-conditioned, often leading to model collapse. This paper introduces Trust Region Q-Adjoint Matching (TRQAM), a stable off-policy fine-tuning algorithm that adaptively controls the path-space KL with pretrained flow policies through projected dual descent. Specifically, we optimize the trust-region parameter $λ$ in SOC dynamics, and theoretically show that the path-space KL can be represented by a closed-form function of $λ$. As a result, our method can precisely control the exact deviation from pretrained flow policies, achieving stable off-policy RL. Through experiments on 50 OGBench tasks, TRQAM consistently outperforms prior arts in both offline RL and offline-to-online RL. In particular, TRQAM achieves an overall success rate of 68% in offline RL, substantially improves the strongest baseline at 46%.

1 Introduction

TRQAM addresses critic-induced instability in off-policy fine-tuning of pretrained flow policies by controlling path-space deviation through a trust-region parameter in the SOC dynamics. Its closed-form KL control and projected dual descent enable stable offline and offline-to-online reinforcement learning.

  • Motivation: Pretrained flow policies provide rich action distributions, useful skills, behavioral constraints, and broad coverage for downstream off-policy RL fine-tuning.Their implicit multi-step denoising process makes direct gradient-based improvement expensive and unstable.
  • Motivation: QAM reformulates fine-tuning as memoryless SOC, but imperfect critics can compound errors through TD bootstrapping and amplify deviations from the pretrained prior.Gradient clipping does not prevent adjoint-loss instability or task collapse on Robomimic.
  • Method: TRQAM introduces λ into SOC sampling dynamics and adapts it via projected dual descent to enforce a prescribed KL bound between fine-tuned and pretrained policies.The trust-region constraint is imposed directly through sampling dynamics rather than only through a loss-level regularizer.
  • Results: TRQAM tightly tracks the target KL bound during offline and online training, unlike loss-level KL regularization that can allow excessive realized deviation and collapse.The prescribed bound controls how much the fine-tuned policy can deviate from its pretrained prior.
  • Theory: Girsanov’s theorem shows that scaling the diffusion coefficient by √λ makes path-space KL an exact closed-form function of λ.This turns λ into a principled trust-region parameter whose dual update directly enforces the target KL bound.
  • Results: TRQAM consistently outperforms prior methods across offline RL and offline-to-online RL on 50 OGBench tasks.The reported evaluation covers 50 OGBench tasks.

2 Background

This section defines the off-policy reinforcement-learning setting and pretrained flow policies, then introduces SOC and QAM as critic-guided fine-tuning methods. It also highlights QAM’s empirical instability under fixed-temperature adjoint matching.

  • Reinforcement Learning: Reinforcement learning seeks a policy maximizing expected discounted return in an MDP, using an action-value function Qπ(s, a) as a critic for policy improvement.The MDP includes state and action spaces, transitions, rewards, an initial-state distribution, and discount factor γ.
  • Reinforcement Learning: Off-policy fine-tuning updates a pretrained policy πbase using replay-buffer transitions collected by behavior policies different from the current policy.This setting focuses on adapting pretrained policies while learning from previously collected data.
  • Flow matching and flow policy: Flow matching learns a velocity field that transports Gaussian source samples to a target distribution through an ODE, while a flow policy generates actions by integrating state-conditioned dynamics.The flow policy starts from X0 ∼ N(0, Id) and integrates dXτ = vθ(Xτ, τ; s)dτ.
  • Stochastic optimal control and Q adjoint matching: SOC fine-tunes a pretrained flow policy by adding a drift perturbation that steers sampling trajectories toward higher-reward regions without backpropagating through the multi-step sampling chain.The deterministic flow ODE is replaced by an equivalent SDE, and the control steers the process away from the pretrained-policy dynamics.
  • Stochastic optimal control and Q adjoint matching: QAM uses the critic as the terminal cost, steering the pretrained policy toward critic-preferred actions while penalizing deviation from πbase through a quadratic term.QAM solves the SOC objective via adjoint matching and updates the fine-tuned velocity field without backpropagating through the sampling chain.
  • Stochastic optimal control and Q adjoint matching: Over 80% to near zero: on Robomimic-can, QAM and QAM-E losses grew above 10^20 despite gradient clipping, collapsing task success while TRQAM remained stable.The collapse persisted across most tested hyperparameter settings on Robomimic-lift and Robomimic-can.

3 Method

TRQAM stabilizes off-policy fine-tuning by internalizing a trust-region parameter λ in the stochastic optimal-control sampling dynamics. It makes path-space KL an exact function of λ and adapts λ through projected dual descent to enforce a prescribed KL bound.

  • Method overview: TRQAM scales the diffusion coefficient by λ, making path-space KL an exact function of λ through Girsanov.This converts a prescribed KL bound εKL into a structural constraint enforced during sampling.
  • Why fixed λ is fragile: Fixed-temperature QAM and QAM-E show diverging adjoint loss and collapsing task success, reflecting exponential amplification of critic errors under aggressive policy updates.The corresponding trust-region interpretation is that small λ permits aggressive deviation from the base policy, whereas large λ keeps the sampler closer to it.
  • What λ controls: Theorem 1 ties λ to trajectory-level path-space KL, while Proposition 1 upper-bounds terminal-policy KL by path-space KL.Together with Lemma 1, these results link λ to critic-error amplification bounded by 2βε.
  • How λ adapts: TRQAM estimates path-space KL from closed-form per-step Gaussian-transition KLs, smooths the estimate with an exponential moving average, and updates λ by projected dual descent.When realized KL exceeds εKL, λ increases; when it falls below the bound, λ decreases.
  • Why λ is internalized: Internalizing λ in the SOC sampling dynamics changes the diffusion scale and reshapes the controlled SDE, pulling it toward base dynamics as λ increases.By contrast, the external form treats λ as a conventional regularization weight; TRQAM enforces the trust region directly through sampling dynamics.

4 Experiments

Experiments evaluate TRQAM on 50 OGBench tasks and Robomimic stability studies against off-policy fine-tuning baselines. TRQAM achieves the strongest aggregate offline performance, benefits from pretrained policies, and gains stability through adaptive, internal KL control with tight budgets.

  • Experimental setup: TRQAM is evaluated on 50 OGBench tasks for main comparisons and Robomimic for ablation and mechanism studies, with methods sharing pretrained policies and training schedules.OGBench spans 10 suites; manipulation tasks use action-chunked policies with chunk size h = 5.
  • Offline and offline-to-online results: 68% aggregate success across 50 tasks: TRQAM improves on QAM (35%) by 33 points, QAM-E (45%) by 23 points, and DSRL (46%) by 22 points.The largest gains occur on long-horizon and combinatorial suites, and TRQAM remains strongest through offline-to-online training.
  • Pretrained-policy benefit: TRQAM benefits substantially from pretrained flow policies, whereas QAM and QAM-E show little to no benefit from the same pretrained initialization.Under an identical offline-to-online protocol, TRQAM reaches high success much earlier from a pretrained policy than from scratch.
  • Ablation studies: Both adaptive KL variants outperform QAM on cube-triple-task1 and humanoidmaze-medium-task1, while QAM exhibits diverging adjoint loss and collapsing task success on Robomimic.The comparison isolates adaptation from constant KL and finds adaptive regulation more stable.
  • Ablation studies: With εKL = 0.1, TRQAM tightly tracks the prescribed KL bound on Robomimic-lift and Robomimic-can, while QAM + External KL drifts above the bound with degraded success.Only TRQAM’s internal parameterization ties λ to realized KL through an exact identity.
  • KL-budget sensitivity: Across four representative OGBench tasks, success rate changes smoothly with εKL, and tight budgets are best across all four tasks.The sweep ranges from 0.5 to 4 in steps of 0.5, making εKL a predictable control knob.

5 Related Work

Prior work studies offline-to-online RL, RL fine-tuning of expressive flow and diffusion policies, and KL-based trust-region constraints. These approaches address distribution shift, policy-update stability, and the tradeoff among expressivity, computational cost, and training stability.

  • Offline-to-online RL: Offline-to-online RL commonly pretrains policies and value functions offline before online fine-tuning for sample-efficient reinforcement learning.The transition introduces distribution shift that can destabilize value functions and induce catastrophic forgetting, motivating trust-region-style constraints.
  • Fine-tuning flow and diffusion policies: Flow matching and diffusion policies parameterize multi-modal action distributions and are increasingly pretrained at scale, with growing efforts to train them using RL.Existing methods navigate tradeoffs among policy expressivity, computational cost, and training stability.
  • KL trust regions in RL: KL regularization stabilizes policy updates through hard constraints or soft penalties, with some methods adapting strength through dual updates.These approaches enforce the constraint via auxiliary losses external to the policy.

6 Conclusion · A Algorithm

TRQAM is presented as a stable off-policy fine-tuning method that controls deviation from pretrained flow policies through a trust-region parameter in the SOC dynamics. The algorithm initializes from a pretrained velocity field, performs critic and adjoint-matching updates, estimates and smooths a KL surrogate, and applies dual descent to update the trust-region parameter.

  • 6 Conclusion: TRQAM scales diffusion through λ, making path-space KL an exact function of λ and enforcing the target bound at the sampling level via dual descent.This mechanism is established through Girsanov in Theorem 1.
  • 6 Conclusion: 68% overall success across 50 OGBench tasks improves the strongest baseline at 46%, with the largest gains on long-horizon and combinatorial domains.On Robomimic, TRQAM remains stable where fixed-temperature adjoint matching collapses.
  • 6 Conclusion: The method’s limitation is that adjoint matching requires a vector-Jacobian product through the velocity field at each backward-ODE step, with cost scaling with model size.The VJP requirement is identified as a computational limitation.
  • A Algorithm: Algorithm 3 fine-tunes Flow Matching policies with TRQAM additions over QAM [25].The algorithm is explicitly framed as a fine-tuning procedure for Flow Matching policies.
  • A Algorithm: The algorithm takes a replay buffer, pretrained behavior velocity field, fine-tuned velocity field, critic, step size, KL budget, dual stepsize, EMA coefficient, and iteration count as inputs.These inputs define the data, policy, critic, optimization, trust-region, smoothing, and iteration settings.
  • A Algorithm: Initialization copies the pretrained velocity field into the fine-tuned field, sets λ0 > 0 and D0 ← 0, and specifies memoryless SDE diffusion scaling √λn σn(τ) = g(τ).The displayed initialization and SDE relation provide the algorithm’s starting state and diffusion parameterization.
  • A Algorithm: Each policy update samples trajectories with a memoryless Euler scheme, computes the critic action gradient, and solves the lean adjoint ODE backward.The displayed procedure then applies stop-gradient operations and optimizes the adjoint-matching objective.
  • A Algorithm: The trust-region stage estimates a path-space KL surrogate, smooths it with an EMA, updates λ by dual descent, and outputs the fine-tuned velocity field.The displayed steps include EMA smoothing with coefficient ρ, dual descent, and final output vft.

B Baselines · C Experimental details

The paper evaluates TRQAM against six flow-policy fine-tuning baselines spanning distillation, critic guidance, noise-space optimization, implicit value learning, and adjoint matching. These baselines differ in how they approximate, guide, or perturb pretrained flow policies during offline and online fine-tuning.

  • B Baselines: Six baselines span distinct fine-tuning paradigms for flow-matching policies, using uniformly sampled replay-buffer transitions without re-weighting.The replay buffer contains the offline dataset during offline training and adds online rollouts during online fine-tuning.
  • B Baselines: FQL distills a multi-step flow policy into a one-step policy, jointly maximizing the critic while remaining close to the behavior-cloning rollout.The environment policy is the learned one-step policy, with α controlling closeness to the rollout.
  • B Baselines: CGQL-Linex combines a behavior-cloning velocity field with classifier-free Q-guidance from an auxiliary intermediate critic on noisy actions.Its guidance strength is modulated by w, and Huber-style stabilization prevents exponential blow-up in the Linex loss.
  • B Baselines: DSRL performs SAC directly in the noise space of a frozen flow policy, then pushes sampled noise through the flow policy to obtain actions.The modified version also fine-tunes the behavior-cloning velocity online using a target network for stability.
  • B Baselines: IFQL uses IQL-style expectile regression for value learning and extracts policies by selecting the highest-Q action among candidates sampled from a behavior-cloning flow policy.The value network fits an upper expectile of the critic, which bootstraps the critic.
  • B Baselines: QAM solves a memoryless SOC problem with λ = 1, trains velocity against the behavior-cloning velocity using lean adjoint matching, and applies element-wise gradient clipping.QAM-E adds an entropy-regularized SAC residual edit policy that perturbs QAM actions by at most σa in L∞ distance.

C.1 Domains and tasks … D.1 Proof of Lemma 1: exponential amplification of critic errors

The paper evaluates TRQAM across OGBench and Robomimic with specified tuning, compute, and reporting protocols, then proves that critic errors can cause exponentially amplified policy divergence. The theoretical bound quantifies this sensitivity through KL and total variation guarantees.

  • C.1 Domains and tasks: OGBench covers 10 domains spanning long-horizon navigation, multi-object manipulation, and combinatorial planning, while Robomimic tests lift, can, and square manipulation tasks.OGBench uses reward-based single-task variants; Robomimic is demonstration-based and serves as a stability stress-test.
  • C.1 Domains and tasks: Experiments use 8 random seeds per method and task, reporting mean success rate ± standard deviation unless otherwise stated.Plots show means with shaded standard-deviation regions.
  • C.2 Compute resources: The reported experiments consumed approximately 47,000 wallclock-run-hours across OGBench and Robomimic, including main comparisons and supporting ablations.The OGBench category alone accounts for approximately 38,400 wallclock-run-hours.
  • C.3 Hyperparameters: Baseline hyperparameters on OGBench use QAM’s per-domain optima, while TRQAM tunes εKL on two tuning tasks per domain.Robomimic instead receives independent full-range sweeps for DSRL, QAM, QAM-E, and TRQAM, using 8 seeds per configuration.
  • D Proofs: The proofs establish Theorem 1’s path-space KL identity, Proposition 1’s terminal-KL upper bound, and Lemma 1’s exponential critic-error amplification.They operate under standard controlled-diffusion regularity assumptions that justify Girsanov’s theorem.
  • D.1 Proof of Lemma 1: exponential amplification of critic errors: For policies exponentially tilted by Q and eQ relative to πbase, a uniform critic error ||Q − eQ||∞ ≤ ε implies DKL(πQ || πeQ) ≤ 2βε.The policy construction uses πbase(a | s)e^(βeQ(a)) for the perturbed critic.
  • D.1 Proof of Lemma 1: exponential amplification of critic errors: The KL guarantee follows pointwise from log(πQ/πeQ) ≤ 2βε and integration under πQ.This links bounded critic error directly to a quantitative divergence bound between the induced policies.

D.2 Proof of Theorem 1: path-space KL identity · D.3 Proof of Proposition 1: Terminal KL upper-bounded by path-space KL · E Path-space KL surrogate under OT memoryless discretization

The paper establishes that the controlled SDE’s path-space KL has an explicit λ-dependent identity, upper-bounds terminal-policy KL, and admits a tractable estimator under OT memoryless discretization. This supports adaptive KL-constrained improvement using λ as an intrinsic dual variable.

  • D.2 Proof of Theorem 1: path-space KL identity: Theorem 1 derives a closed-form path-space KL identity for the base and controlled SDE path measures using Girsanov’s theorem.The proof uses the Radon–Nikodym derivative on the natural filtration, the full path-space KL definition, and the zero-mean Itô integral property.
  • D.2 Proof of Theorem 1: path-space KL identity: The prior SOC parameterization has no λ inside the SDE to modulate regularization strength, so trust-region adaptation requires an external coefficient on the SOC objective.In that parameterization, path-space KL coincides with the SOC quadratic cost up to a fixed constant.
  • D.2 Proof of Theorem 1: path-space KL identity: λ is an intrinsic parameter of the controlled dynamics, with a 1/λ factor directly modulating the path-space KL.Unlike prior SOC parameterizations, this allows λ to act as an adaptive dual variable for KL-budgeted improvement by reshaping the SDE itself.
  • D.3 Proof of Proposition 1: Terminal KL upper-bounded by path-space KL: Proposition 1 shows that terminal-policy KL is upper-bounded by path-space KL when both path measures share the initial distribution p0.The terminal action distributions are deterministic terminal projections of the path measures, so the result follows from the data-processing inequality for KL divergence.
  • E Path-space KL surrogate under OT memoryless discretization: The OT memoryless discretization section uses the path-space KL identity and terminal-KL upper bound to construct a tractable estimator for Algorithm 3.The estimator targets the path-space KL because it is available in closed form and constrains the terminal-policy KL indirectly.
  • E Path-space KL surrogate under OT memoryless discretization: Under the OT memoryless Euler scheme, fine-tuned and base transitions are Gaussian with shared covariance, so their KL depends on the difference between their velocity-field means.The transition covariance is specified by the step size and schedule, while the means differ only through the velocity fields.
  • E Path-space KL surrogate under OT memoryless discretization: The discrete path-space KL decomposes across Markov transitions by the KL chain rule and is estimated by a Monte Carlo average over trajectories in batch B.The resulting estimator is the one used in Algorithm 3 and may be smoothed with EMA to reduce variance.

F KL-budgeted improvement: primal–dual derivation

This section derives the projected dual update for the KL-budgeted improvement parameter λ. Convexity and Slater’s condition establish strong duality, while the update uses KL estimates to adjust conservativeness around the budget.

  • Primal problem and Slater’s condition: Strong duality applies because the path-space objective is linear, the KL constraint is convex, and u ≡ 0 is strictly feasible when ε_KL > 0.The trivial control recovers P_u = P_base, satisfying Slater’s condition.
  • Lagrangian and dual function: Minimizing the dual function g(λ) over λ ≥ 0 solves the primal KL-budgeted improvement problem.The dual function is defined as the supremum of the Lagrangian over controls.
  • Lagrangian and dual function: The dual subgradient is s(λ) = ε_KL − D_KL(P_uλ ∥ P_base), comparing the KL budget with the inner maximizer’s divergence.Here, u_λ is the inner Lagrangian maximizer at the current dual variable.
  • Projected dual descent: Algorithm 3 evaluates KL at the current control using the Monte Carlo surrogate bD_n, enabling projected dual descent without a closed-form inner-maximizer KL.When realized KL exceeds the budget, λ rises and dynamics become more conservative; below the budget, the trust region relaxes.

G Internal vs. external KL regularization: detailed comparison … H.2 Time-varying KL budget

TRQAM internalizes λ in the SOC dynamics, making path-space KL an exact function of λ and enabling structural trust-region control. Across fixed and time-varying budgets, it tracks prescribed KL targets, while a moderate online relaxation improves exploration without the variance of an aggressive one.

  • G Internal vs. external KL regularization: detailed comparison: Internal TRQAM makes path-space KL an exact function of λ, whereas external regularization leaves realized KL independent of λ.Both formulations can use the same dual update and KL estimator, but λ is a loss weight externally and an SDE coefficient internally.
  • G Internal vs. external KL regularization: detailed comparison: External regularization applies only at the loss level after sampling, so whether realized path-space KL approaches εKL depends on optimizer trade-offs.The KL term competes with the reward gradient instead of directly constraining the generated trajectory distribution.
  • G Internal vs. external KL regularization: detailed comparison: Increasing internal λ shrinks diffusion, weakens the controlled drift contribution, and pulls the controlled SDE toward base dynamics.Theorem 1 then makes the realized path-space KL an exact function of λ, so dual updates reshape trajectories structurally rather than competing with critic guidance.
  • H.1 TRQAM tightly enforces the prescribed KL budget: TRQAM tracks prescribed path-space KL targets εKL ∈{0.5, 1.0, 1.5, 2.0} throughout offline and online training with monotonic ordering preserved across transition.This indicates that the dual update remains effective under distribution shift.
  • H.2 Time-varying KL budget: TRQAM supports resetting εKL during training without retraining or instability because λ is internalized in SOC dynamics.The method accommodates different trust-region needs between offline and online phases, where online fine-tuning may require more exploration.
  • H.2 Time-varying KL budget: On antmaze-giant-task1, Figure 8 compares a static 0.5 schedule with online relaxations 0.5 →3.0 and 0.5 →4.0.The domain has the benchmark’s largest state space and highest exploration demands during online fine-tuning.
  • H.2 Time-varying KL budget: Realized KL adapts to each new target almost immediately after a schedule switch, confirming stable tracking of time-varying bounds.This behavior is demonstrated during the offline-to-online transition.
  • H.2 Time-varying KL budget: ∼98% success is reached rapidly with 0.5 →3.0, while 0.5 →4.0 introduces substantial variance and static 0.5 improves more slowly online.The results identify 3.0 as a sweet spot because tight budgets restrict exploration, whereas aggressive relaxation increases variance.

H.3 Sensitivity Analysis · H.4 Stability stress test on Robomimic · I Broader impacts

Sensitivity sweeps show that εKL trades off performance and online adaptation across OGBench, while Robomimic stress tests reveal fixed-temperature collapse and demonstrate that TRQAM tightly enforces KL budgets. The paper’s broader-impacts discussion frames TRQAM as stable simulated-benchmark fine-tuning with potential data and compute benefits for pretrained robot policies.

  • H.3 Sensitivity Analysis: Tight εKL budgets perform best on five OGBench domains, whereas larger budgets monotonically improve puzzle-4x4.The sweep covered εKL ∈ {0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0} across all ten domains with eight seeds.
  • H.3 Sensitivity Analysis: Smaller εKL slows online adaptation, consistent with tighter trust-region updates and εKL controlling realized deviation from πbase.The same effect appears on antmaze-giant, whose large state space demands broader exploration.
  • H.4 Stability stress test on Robomimic: Robomimic complements OGBench as a standard manipulation benchmark where fixed-temperature adjoint matching is unstable.The evaluation covers simulated manipulation tasks in the Robomimic benchmark.
  • H.4 Stability stress test on Robomimic: Fixed-temperature QAM and QAM-E collapse across most Robomimic lift and can settings as adjoint-matching loss explodes and success rate approaches zero.QAM was swept over six β values, while QAM-E was swept over eight (β, σa) configurations.
  • H.4 Stability stress test on Robomimic: External KL regularization does not reliably enforce the prescribed path-space KL, whereas TRQAM tracks each budget tightly throughout offline training and offline-to-online transition.TRQAM remains stable across εKL ∈ {0.01, 0.03, 0.1, 0.5, 1, 1.5}.
  • I Broader impacts: TRQAM is evaluated entirely on simulated OGBench and Robomimic benchmarks as a method for stable off-policy fine-tuning of pretrained flow-matching policies.The broader-impacts discussion identifies more reliable specialization of pretrained robot policies as a potential benefit.
  • I Broader impacts: More reliable fine-tuning could reduce the data and compute required to specialize useful behaviors and lessen brittleness from pretrained-prior degradation under TD bootstrapping.The passage describes pretrained-prior degradation under TD bootstrapping as a common source of instability in real-world RL deployments.
Loading 2605.27079v1…