Trust Region Q Adjoint Matching¶
会议: NeurIPS 2026
arXiv: 2605.27079
OpenReview: https://openreview.net/forum?id=o0QcqHlW1x
代码: https://github.com/yonghdong/trqam
领域: 强化学习
关键词: 强化学习, 流匹配策略, 随机最优控制, 伴随匹配, 信任域优化
一句话总结¶
针对流匹配策略在离线强化学习微调中由 Critic 误差放大引发的崩溃问题，本文提出了信赖域 Q 伴随匹配（TRQAM），通过在随机最优控制采样过程中内化自适应参数 \(\lambda\) 并执行对偶投影下降，实现了对路径空间 KL 散度的精确控制与高稳定性策略改进。
研究背景与动机¶
流匹配与扩散策略因其能够建模复杂的多模态动作分布，在机器人连续控制和具身智能中成为了主流的预训练策略骨干。然而，这些生成策略的动作生成依赖于多步去噪常微分或随机微分方程，直接使用基于梯度的强化学习算法对多步采样链进行时域反向传播计算代价极高且数值极不稳定。近期提出的基于伴随匹配的 Q 学习（QAM）将策略微调重新表述为无记忆随机最优控制（SOC）问题，利用伴随状态 ODE 避开了对采样链的反向传播。
然而，QAM 继承了 Critic 引导改进机制的一个致命脆弱性：在离线及离线到在线强化学习中，由于时序差分自举与分布偏移，学到的 Critic 难免存在逼近误差。固定温度系数下的指数倾斜机制会成指数级放大这些局部 Critic 误差（在理论上表现为全变差距离随误差指数增长），导致策略伴随损失发散并发生灾难性性能崩溃。即便是 QAM 原文采用的梯度裁剪启发式方法，也无法阻止这种由价值估计过高引发的策略漂移与失败。
针对这一困境，本文没有采用在外部损失中简单添加传统动作级 KL 惩罚的做法，而是探寻如何在连续时间的随机最优控制采样机制内部内化信任域约束。核心 idea：将信任域乘子 \(\lambda\) 直接嵌入随机最优控制的扩散采样漂移项中，利用 Girsanov 定理证明其精确等价于路径空间 KL 散度的权重，并通过投影对偶下降自适应调整 \(\lambda\)，在采样过程内严格将策略漂移限制在预设的 KL 预算之内。
方法详解¶
整体框架¶
TRQAM 面向预训练流匹配策略 \(\pi_{\text{base}}\) 的离线与离线到在线微调，整个系统包含 Critic 价值学习、随机微分方程采样、伴随状态反向求解、伴随匹配损失优化以及基于对偶投影下降的信任域参数更新五个核心环节。输入为环境状态 \(s\) 与重放缓冲区中的离线转移数据，输出为优化后的流策略速度场 \(v_\theta^{\text{ft}}\) 与状态-动作价值函数 \(Q_\phi\)。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["环境状态 s 与重放缓冲数据"] --> B["前向 SDE 采样轨迹<br/>嵌入当前信任域参数 λ_n"]
B --> C["Critic 动作梯度评估<br/>a~_1 = -∇_a Q_ϕ(s, a)"]
C --> D["反向求解精简伴随 ODE<br/>获取整条轨迹伴随状态 a~_τ"]
D --> E["内化伴随匹配更新<br/>优化流策略速度场 v_θ^ft"]
E --> F["闭式路径空间 KL 估计与平滑<br/>利用速度场差值无额外开销统计"]
F --> G["投影对偶下降更新 λ<br/>动态平衡利用与保守信任域"]关键设计¶
1. 随机控制扩散过程中的内化信任域参数化：消除外部惩罚的失控失效
在无记忆最优传输随机控制框架中，基础 SDE 与受控 SDE 共享相同的边际分布。TRQAM 将参数 \(\lambda\) 嵌入扩散尺度 \(\sigma(\tau) = \sqrt{2(1-\tau)/(\lambda \tau)}\) 中，受控 SDE 的漂移项写作 \(b(X_\tau^u, \tau) + \sigma(\tau) u(X_\tau^u, \tau)\)。与以往在策略损失函数外部附加 \(\lambda \cdot \mathcal{D}_{\text{KL}}\) 惩罚项不同，TRQAM 的 \(\lambda\) 作用于生成采样轨迹的最底层动力学。本文通过 Girsanov 定理证明了受控轨迹测度 \(\mathbb{P}^u\) 与基线轨迹测度 \(\mathbb{P}^{\text{base}}\) 之间的路径空间 KL 散度与控制成本满足严格等式： $\(D_{\mathrm{KL}}\bigl(\mathbb{P}^{u}\,\|\,\mathbb{P}^{\mathrm{base}}\bigr) = \mathbb{E}_{\mathbf{X}\sim\mathbb{P}^{u}}\!\bigg[\,\frac{1}{2\lambda}\int_{0}^{1}\|u(X_{\tau},\tau)\|^{2}\,d\tau\,\bigg]\)$ 由于路径空间 KL 自然成为最终动作空间 KL 的上界（\(D_{\mathrm{KL}}(\pi_\theta(\cdot\mid s) \,\|\, \pi_{\text{base}}(\cdot\mid s)) \le D_{\mathrm{KL}}(\mathbb{P}^u \,\|\, \mathbb{P}^{\text{base}})\)），增大 \(\lambda\) 能够自底向上收缩整条轨迹的偏差，从根本上防止强 Critic 梯度冲垮外部正则项。
2. 零额外计算开销的闭式路径空间 KL 估计器：复用前向速度差
为了将路径空间 KL 散度约束在设定预算 \(\varepsilon_{\text{KL}}\) 之内，系统需要实时度量当前的真实漂移量。本文利用在网格时间步 \(\tau_k = k h\) 上离散化的高斯转移特性，由于相邻步的高斯转移共享完全相同的协方差，每一步的转移 KL 具有精确的解析形式。通过沿采样轨迹对局部 KL 求和，构造了路径空间 KL 的蒙特卡洛经验估计量 \(\widehat{D}_n\)： $\(\widehat{D}_{n}=\mathbb{E}_{\mathbf{X}\sim\mathbb{P}^{u}}\!\bigg[\sum_{k=0}^{T-1}\frac{2h}{\nu(\tau_{k})^{2}}\bigl\|v^{\mathrm{ft}}_{\theta}(X_{\tau_{k}},\tau_{k})-v^{\mathrm{base}}(X_{\tau_{k}},\tau_{k})\bigr\|^{2}\bigg]\)$ 其中 \(\nu(\tau) = \sqrt{2(1-\tau)/\tau}\)。该估计量直接复用了在计算伴随匹配损失时已经前向计算出的微调速度场与基础速度场之差 \(v_\theta^{\text{ft}} - v^{\text{base}}\)，因此在训练过程中完全不需要引入额外的神经网络推理或前向微分计算。
3. 投影对偶下降自适应更新机制：稳定追踪预设 KL 预算
将离线策略优化建模为带约束的优化问题：最大化累计 Critic 预期回报，同时约束路径空间 KL 满足 \(D_{\mathrm{KL}}(\mathbb{P}^u \,\|\, \mathbb{P}^{\text{base}}) \le \varepsilon_{\text{KL}}\)。将 \(\lambda\) 视作该约束对应的对偶变量，利用指数滑动平均（EMA）平滑轨迹估计值 \(\overline{D}_n \leftarrow (1-\rho_D)\overline{D}_{n-1} + \rho_D \widehat{D}_n\) 以抑制高方差噪声，随后以步长 \(\eta_\lambda\) 执行投影对偶下降更新： $\(\lambda_{n+1} \leftarrow \max\bigl\{0,\;\lambda_{n} + \eta_{\lambda}(\overline{D}_{n} - \varepsilon_{\mathrm{KL}})\bigr\}\)$ 当策略受 Critic 引导而漂移超出预设预算 \(\varepsilon_{\text{KL}}\) 时，\(\lambda\) 迅速上升并自动压缩受控 SDE 中的控制扰动幅度，强制策略拉回预训练先验的有效支撑集；当策略偏离在安全预算范围内时，\(\lambda\) 适当减小以允许更积极的探索与强化学习改进。
损失函数 / 训练策略¶
- Critic 价值网络优化：利用离线转移批次数据 \(\mathcal{B} = \{(s, a, r, s')\}\) 计算标准 TD 误差损失： $\(\mathcal{L}(\phi) = \frac{1}{|\mathcal{B}|} \sum_{(s,a,r,s') \in \mathcal{B}} \left[ Q_\phi(s,a) - r - \gamma Q_{\bar{\phi}}(s', a' \sim \pi_\theta(\cdot \mid s')) \right]^2\)$
- 伴随匹配损失优化：由状态 \(s\) 采样轨迹 \(X_\tau\) 并反向求解终端为 \(\tilde{a}_1 = -\nabla_{X_1} Q_\phi(s, X_1)\) 的精简伴随 ODE 得到序列伴随状态 \(\tilde{a}_\tau\) 后，更新流策略参数 \(\theta\)： $\(\mathcal{L}_{\mathrm{Adj\text{-}Match}}(\theta) = \frac{1}{|\mathcal{B}|} \sum_{s \in \mathcal{B}} \sum_{\tau} \left\| \frac{2}{\sigma_n(\tau)} \bigl(v^{\mathrm{ft}}_{\theta}(s,X_{\tau},\tau) - v^{\mathrm{base}}(s,X_{\tau},\tau)\bigr) + \sigma_n(\tau)\tilde{a}_{\tau} \right\|^2\)$
- 两阶段微调调度：先使用行为克隆（BC）预训练流策略 300K 步；随后在离线数据集上执行 100 万步（1M）TRQAM 微调；在离线到在线设置中，再继续追加 50 万步（0.5M）环境在线交互与回放更新。
实验关键数据¶
主实验¶
在涵盖 10 个子领域的 50 项 OGBench 离线强化学习任务（1M 训练步数，8 个独立随机种子）的主对比中，TRQAM 展现了显著超越现有六种代表性方法的表现。
| 领域分类 | 任务数 | FQL (ICML'25) | CGQL-Linex | DSRL (ICLR'25) | IFQL | QAM (ICLR'26) | QAM-E | TRQAM (本文) |
|---|---|---|---|---|---|---|---|---|
| antmaze-large (al) | 5 | 38 ± 9 | 48 ± 7 | 53 ± 2 | 29 ± 8 | 62 ± 9 | 86 ± 3 | 89 ± 4 |
| antmaze-giant (ag) | 5 | 2 ± 6 | 7 ± 5 | 1 ± 1 | 12 ± 3 | 29 ± 4 | 6 ± 8 | 41 ± 4 |
| humanoidmaze-med (hm) | 5 | 74 ± 5 | 57 ± 2 | 53 ± 10 | 93 ± 2 | 64 ± 7 | 60 ± 6 | 84 ± 3 |
| humanoidmaze-large (hl) | 5 | 2 ± 1 | 6 ± 3 | 1 ± 1 | 30 ± 7 | 4 ± 3 | 4 ± 5 | 36 ± 4 |
| scene | 5 | 70 ± 5 | 58 ± 1 | 80 ± 0 | 36 ± 1 | 64 ± 4 | 63 ± 6 | 79 ± 1 |
| puzzle-3x3 (p33) | 5 | 25 ± 10 | 0 ± 0 | 100 ± 0 | 64 ± 4 | 15 ± 3 | 89 ± 4 | 100 ± 0 |
| puzzle-4x4 (p44) | 5 | 9 ± 7 | 0 ± 0 | 61 ± 8 | 42 ± 4 | 1 ± 1 | 54 ± 8 | 99 ± 1 |
| cube-double (c2) | 5 | 44 ± 4 | 55 ± 2 | 72 ± 4 | 9 ± 2 | 71 ± 2 | 71 ± 3 | 81 ± 3 |
| cube-triple (c3) | 5 | 7 ± 5 | 0 ± 1 | 34 ± 6 | 24 ± 7 | 19 ± 6 | 11 ± 4 | 50 ± 5 |
| cube-quadruple (c4) | 5 | 9 ± 5 | 1 ± 1 | 9 ± 3 | 6 ± 3 | 18 ± 3 | 9 ± 3 | 19 ± 5 |
| Overall (全任务平均) | 50 | 28% | 23% | 46% | 35% | 35% | 45% | 68% |
消融实验¶
为了严格解耦自适应调节机制与 SDE 内部参数化机制各自带来的增益，本文在代表性任务与 Robomimic 机械臂操作基准上对比了不同 KL 调控范式：
| 机制配置 | Robomimic-can 成功率 | 伴随匹配损失发散情况 | 预设 KL 预算追踪能力 | 说明 |
|---|---|---|---|---|
| TRQAM（完整方法：内化自适应 \(\lambda\)） | 83.1% ± 3.2% | 完全稳定（损失保持正常量级） | 严密贴合（实际 KL 紧密跟随 \(\varepsilon_{\text{KL}}\)） | 理论保证参数精确映射路径空间 KL |
| QAM + External KL（外部损失惩罚） | 61.4% ± 5.6% | 保持稳定 | 严重失控（实际 KL 大幅漂移超出预设预算） | 强 Critic 信号在优化时冲垮外部正则项 |
| QAM（基线：固定 \(\lambda=1\)） | 0.0% ± 0.0% | 极端发散（损失爆炸超过 \(10^{20}\)） | 无约束机制 | 固定温度导致误差指数放大，策略崩溃 |
关键发现¶
- 内化参数化是约束有效的必要条件：实验证明，外部 KL 惩罚虽能缓解崩溃，但实际 KL 会严重漂移超标（达到预设值的数倍），导致成功率显著落后；只有嵌入受控 SDE 的内化 \(\lambda\) 才能忠实执行给定的 KL 预算。
- 预训练先验的杠杆作用被真正激活：当从头训练（Scratch）与微调预训练策略对比时，QAM 与 QAM-E 的微调曲线与从头训练完全重合（无法有效继承先验技能）；而 TRQAM 能够充分利用预训练先验，在训练极早期就建立起巨大优势。
- 长时序与组合操作收益最显著：在 AntMaze-Giant（41% vs 6%）、HumanoidMaze-Large（36% vs 4%）及 Cube-Triple（50% vs 11%）等极度容易出现误差累积的长程任务中，TRQAM 的相对提升幅度最大。
亮点与洞察¶
- 理论推导的严密闭环：通过 Girsanov 变换建立了受控流扩散过程控制代价与路径空间 KL 的严格数学等式，让信任域系数 \(\lambda\) 具有了毫无歧义的物理与信息论意义。
- 零计算代价的工程优雅性：构造的路径空间 KL 经验估计式精确复用了伴随匹配反向传播前必须计算的差值张量，未引入任何额外网络前向或采样步开销。
- 打破固定温度下的脆弱性怪圈：给出了策略优化中 Critic 逼近误差随反温度参数 \(\beta\) 指数级放大的数学上界分析，从机理上阐明了为什么固定参数微调必然在复杂环境中面临崩溃。
局限与展望¶
- 单步静态超参数选择：虽然统一的全局 \(\varepsilon_{\text{KL}} = 0.5\) 能取得 63% 的平均成功率，但在特大状态空间（如 Puzzle-4x4）或特定多任务场景下，静态预算并非完全最优，未来可探索动态感知任务进展的自适应调度机制。
- 动作块边界对齐问题：实验中对于机械臂操作任务采用了固定长度为 5 的动作分块（Action Chunking），路径空间 KL 的约束在跨块衔接处的动态特性仍有进一步理论刻画空间。
相关工作与启发¶
- vs Q-learning with Adjoint Matching (QAM)：QAM 提出了流匹配微调的伴随状态方程，但固定温度导致数值极端脆弱并依赖容易失效的梯度裁剪；TRQAM 在内部将伴随匹配泛化为带自适应参数 \(\lambda\) 的信任域控制，解决了训练崩溃难题。
- vs DSRL (Diffusion Space RL)：DSRL 仅在冻结生成模型的输入噪声空间进行探索和微调，牺牲了流模型自身高阶表达能力的演化；TRQAM 直接端到端微调流速度场本身，突破了预训练策略原有的生成瓶颈。
- vs PPO / TRPO：TRPO/PPO 奠定了在线策略强化学习通过 KL 约束防止更新步过大的稳定性基石；TRQAM 首次将这一思想完整移植到了连续时间生成流策略的离线强化学习中。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次从随机最优控制与 Girsanov 理论出发将路径空间信任域无缝内化至流匹配采样过程。
- 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 50 个 OGBench 任务与 Robomimic 操控任务，跨越 8 个随机种子，进行了深度机制解耦与消融。
- 写作质量: ⭐⭐⭐⭐⭐ 理论证明完整，从现象发现、理论建模、对偶算法到实验验证逻辑脉络高度连贯自洽。
- 价值: ⭐⭐⭐⭐⭐ 为大模型时代基于流匹配与扩散模型的具身智能/VLA 策略强化学习微调提供了坚实的稳定性基础。