信任区域Q伴随匹配 Trust Region Q Adjoint Matching
通过自适应信任区域控制稳定离策略流策略微调
前置知识
流策略
流策略使用流匹配来建模动作分布,通过从简单源分布(如高斯分布)逐步转换到目标分布来生成动作。具体来说,给定状态s,流策略通过积分dX_tau = v_theta(X_tau, tau; s)d_tau来采样动作,其中X_0服从标准正态分布是初始噪声,v_theta是学习的速度场,tau在[0,1]之间是时间参数。这种方法可以建模多峰、复杂的动作分布,比传统的单峰高斯策略更具表达能力。
TRQAM的核心就是对预训练的流策略进行微调,理解流策略的生成机制是理解本文采样动力学和控制的基础
随机最优控制
SOC是一个通过添加漂移扰动u来微调预训练流策略的框架,它引导轨迹向更高奖励的区域发展,而不需要反向传播通过多步采样链。具体来说,我们有受控随机微分方程 dX_tau^u = (b(X_tau^u, tau) + sigma(tau)u(X_tau^u, tau))d_tau + sqrt(lambda)sigma(tau)dB_tau,其中u是控制信号,b是基础漂移,sigma(tau)是扩散系数,B_tau是布朗运动。优化目标是最小化控制成本的同时最大化终端奖励。
TRQAM的核心贡献就是在这个SOC动力学中引入信任区域参数lambda,理解SOC是理解TRQAM方法原理的关键
伴随匹配
伴随匹配是解决SOC问题的高效算法,它通过一个精简的伴随常微分方程来计算梯度。具体来说,伴随变量e_a_tau通过向后积分计算,其中e_a_1等于负的Q函数对动作的梯度。然后通过最小化伴随匹配损失来更新速度场,而不需要反向传播通过采样链。这种方法避免了通过多步采样链进行反向传播的计算成本和不稳定性。
TRQAM基于QAM方法,QAM就是Q-learning with Adjoint Matching,理解伴随匹配是理解TRQAM更新规则的基础
Girsanov定理
Girsanov定理是测度论中的一个重要结果,它描述了如何通过添加一个漂移项来改变一个随机过程的概率测度。在本文的上下文中,Girsanov定理建立了一个精确的恒等式,将控制成本和路径空间KL散度联系起来。具体来说,定理表明路径空间KL散度等于控制成本的期望除以2lambda,这意味着将扩散系数缩放sqrt(lambda)使得路径空间KL成为lambda的精确函数。
Girsanov定理是TRQAM的理论基础,它使得信任区域参数lambda可以通过对偶下降来精确控制路径空间KL散度
路径空间KL散度
路径空间KL散度衡量的是两个随机过程生成的轨迹分布之间的差异。对于受控SDE和基础SDE生成的轨迹分布P_u和P_base,路径空间KL散度量化了控制扰动如何使受控轨迹偏离基础轨迹。TRQAM的一个关键洞察是,通过控制路径空间KL,我们可以控制终端动作分布之间的差异,因为终端策略KL散度小于等于路径空间KL散度。
TRQAM通过控制路径空间KL散度来实现信任区域约束,理解这个概念是理解TRQAM如何稳定训练的关键
研究动机
现有的流策略微调方法存在严重的稳定性问题。流策略虽然能建模丰富的动作分布,但它们的多步去噪过程导致基于梯度的策略改进需要通过多步采样链进行反向传播,这既昂贵又不稳定。最近提出的Q-learning with Adjoint Matching通过将微调重新表述为无记忆随机最优控制问题来解决这个问题,它使用学习的critic通过伴随匹配控制采样过程向更高价值的动作发展。然而,QAM继承了一个基本脆弱性:当critic病态时,小的critic错误会被放大,经常导致模型崩溃。具体来说,在Robomimic-can任务上,即使使用梯度裁剪,QAM和QAM-E的伴随损失也会增长到超过10的20次方,将任务成功率从80%以上推至接近零。这种崩溃在Robomimic-lift和Robomimic-can的大多数超参数设置中持续存在。
本文的目标是本文的目标是提出一种稳定的离策略微调算法,能够自适应地控制与预训练流策略的路径空间KL散度,从而在不破坏预训练策略有用行为的前提下改进下游性能。具体来说,作者希望解决如何在没有从预训练先验的破坏性漂移的情况下改进下游性能这一核心挑战,类似于on-policy RL中的信任区域原理。
与已有工作不同的是,本文的独特切入角度是将信任区域参数lambda直接内化到随机最优控制采样动力学中,并通过投影对偶下降自适应地调整它。这与现有方法有本质区别:传统方法如KL正则化只在损失水平上施加约束,而TRQAM在采样动力学层面强制执行信任区域约束。具体来说,TRQAM通过将扩散系数缩放sqrt(lambda),利用Girsanov定理证明路径空间KL散度是lambda的精确闭式函数,从而将对偶更新直接通过采样动力学强制执行目标KL边界,而不是通过传统的损失级KL正则化软约束来实现。这种区别在实践中很重要:传统的损失级KL正则化只在损失水平上与critic指导竞争,允许强的critic信号将实现的KL推到远超目标边界,使策略容易崩溃;而TRQAM在整个离线和在线训练过程中紧密跟踪目标边界。
核心方法
TRQAM的整体思路是在随机最优控制采样动力学中引入一个信任区域参数lambda,并通过投影对偶下降自适应地调整它,以强制执行微调策略和预训练策略之间的预定KL边界。直觉上,小的lambda允许对预训练策略的激进偏离(对应于大的逆温度beta),而大的lambda保持受控采样器接近预训练策略。TRQAM通过将扩散系数缩放sqrt(lambda),利用Girsanov定理证明路径空间KL散度是lambda的精确闭式函数,从而将对偶更新直接通过采样动力学强制执行目标KL边界。技术路线上,TRQAM首先建立为什么固定的lambda在离策略RL中是脆弱的,然后证明当lambda内化在SOC采样动力学中时它控制什么,接着提出如何通过投影对偶下降在整个训练过程中适应lambda,最后解释为什么lambda必须内化在SOC采样动力学中而不是作为传统的KL正则化添加。
TRQAM的核心创新点是将信任区域参数lambda直接内化到随机最优控制采样动力学中,而不是作为外部的损失级正则化。具体来说,TRQAM使用受控SDE,其中扩散系数被缩放sqrt(lambda)。通过Girsanov定理,作者证明了控制成本和路径空间KL散度之间的精确恒等式,即路径空间KL散度等于控制成本的期望除以2lambda。这意味着lambda成为了一个有原则的信任区域参数,使得对偶更新可以直接通过采样动力学强制执行目标KL边界。与已有方法的本质区别在于:传统方法如KL正则化只在损失水平上施加约束,允许强的critic信号将实现的KL推到远超目标边界;而TRQAM在采样动力学层面强制执行信任区域约束,紧密跟踪目标边界。
方法步骤详情
TRQAM的完整算法步骤如下。首先,初始化微调速度场等于基础速度场,信任区域参数lambda_0,KL估计器D_0和对偶步长eta_lambda。然后,对于每个训练步骤n从0到N-1,执行以下操作:第一,使用当前微调速度场和lambda_n通过受控SDE采样轨迹X;第二,通过向后积分伴随常微分方程求解伴随变量;第三,通过最小化伴随匹配损失更新参数theta;第四,使用公式估计路径空间KL,并通过指数移动平均平滑;第五,通过投影对偶下降更新信任区域参数lambda_{n+1}等于max{0, lambda_n + eta_lambda(D_n - epsilon_KL)}。当实现的KL超过边界时,lambda上升,受控动力学变得更保守;当它低于边界时,lambda下降,允许更激进的改进。
技术新颖性
TRQAM的技术新颖性体现在三个方面。首先,它识别了固定温度伴随匹配中critic错误的指数放大作为一个基本脆弱性,并通过引理形式化,在Robomimic上得到实证确认。其次,它证明了将扩散系数缩放sqrt(lambda)使得路径空间KL通过Girsanov成为lambda的精确函数,将lambda转变为一个有原则的信任区域参数。第三,它提出了信任区域Q伴随匹配,将lambda内化在SOC采样动力学中,并通过投影对偶下降自适应调整它,在采样级别而不是作为损失级惩罚强制执行目标KL边界。这种设计在50个OGBench任务上持续优于先前的艺术,在离策略RL和离策略到在线RL中都取得了优异的性能,特别是在离策略RL中达到68%的整体成功率,明显优于最强基线的46%。
实验结果
TRQAM在50个OGBench任务上的实验结果证明了其优越性能。在离策略RL设置下,TRQAM达到68%的聚合成功率,比QAM提高33个百分点,比其最强变体QAM-E提高23个百分点,比最强非伴随基线DSRL提高22个百分点,在长视界和组合套件上获得最大增益。具体来说,在antmaze-large任务上,TRQAM达到89加减4%,而QAM为62加减9%,QAM-E为86加减3%,DSRL为53加减2%;在antmaze-giant任务上,TRQAM达到41加减4%,而QAM为29加减4%,QAM-E为6加减8%,DSRL为1加减1%;在cube-triple任务上,TRQAM达到50加减5%,而QAM为19加减6%,QAM-E为11加减4%,DSRL为34加减6%。在Robomimic上,TRQAM在固定温度伴随匹配崩溃的地方保持稳定。在机制研究方面,作者发现自适应KL优于常数KL:TRQAM和QAM带外部KL正则化都明显优于QAM(常数lambda),这与固定温度下critic错误指数放大的理论一致。作者还发现内部KL优于外部KL正则化:在Robomimic-lift和Robomimic-can上,当epsilon_KL等于0.1时,TRQAM在整个训练过程中紧密跟踪边界,而QAM带外部KL让实现的KL漂移到epsilon_KL以上,并伴随成功率下降。敏感性分析表明,epsilon_KL是最重要的超参数,改变它会产生可预测的成功率变化:在所有四个代表性任务上,成功率随着epsilon_KL平滑变化,并且在所有四个任务上,紧预算最好。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Offline RL on 50 OGBench tasks (aggregate) | Success rate (%) | 68 | 46 (DSRL, strongest baseline) | +22 percentage points (47.8% relative improvement) |
| Offline RL on antmaze-large | Success rate (%) | 89加减4 | 86加减3 (QAM-E) | +3 percentage points (3.5% relative improvement) |
| Offline RL on antmaze-giant | Success rate (%) | 41加减4 | 6加减8 (QAM-E) | +35 percentage points (583% relative improvement) |
| Offline RL on cube-triple | Success rate (%) | 50加减5 | 11加减4 (QAM-E) | +39 percentage points (355% relative improvement) |
| Offline RL on humanoidmaze-large | Success rate (%) | 36加减4 | 4加减5 (QAM-E) | +32 percentage points (800% relative improvement) |
局限与改进
作者承认的主要局限性是计算伴随匹配损失需要在向后常微分方程的每一步通过速度场计算向量-雅可比积,这个VJP成本随模型大小缩放。除此之外,TRQAM的另一个潜在局限性是它需要预先训练的流策略,这意味着它不适用于从零开始学习策略的场景。此外,TRQAM的性能对KL预算epsilon_KL的选择敏感,虽然作者展示了这个超参数的行为是可预测的,但它仍然需要根据任务结构进行调整,这可能需要一些调参经验。最后,TRQAM的理论结果依赖于Girsanov定理,该定理要求一定的技术条件,虽然在实践中这些条件通常满足,但在理论上仍需要考虑这些限制。
独立分析的弱点
TRQAM的一个潜在弱点是它依赖于预训练的流策略,这意味着在预训练策略编码无意义行为的情况下,TRQAM可能无法有效改进性能。具体来说,如果预训练策略在任务相关区域覆盖率差或包含有害的行为模式,TRQAM的信任区域约束可能会限制策略改进的空间。另一个弱点是TRQAM需要调优KL预算epsilon_KL,虽然作者展示了这个超参数的行为是可预测的,但在新环境中选择最佳epsilon_KL仍然需要一些先验知识或试错。此外,TRQAM的VJP计算成本可能成为扩展到更大模型的瓶颈,特别是在计算资源有限的情况下。改进方向包括:第一,研究自适应选择epsilon_KL的方法,基于任务结构或在线反馈自动调整KL预算;第二,优化VJP计算,例如使用近似方法或并行化;第三,探索TRQAM与其他预训练方法的结合,以改进预训练策略的质量。
未来方向
作者提出的未来工作方向包括将TRQAM扩展到更广泛的策略类,不仅限于流策略,以及探索TRQAM在其他RL设置中的应用,如多智能体RL或元RL。基于TRQAM的成果,可以延伸的研究方向包括:第一,研究TRQAM在更大规模和更复杂的环境中的性能,例如真实机器人任务或具有高维状态和动作空间的场景;第二,探索TRQAM与其他信任区域方法的结合,例如将TRQAM的KL约束与TRPO或PPO的约束结合,以获得更稳定和高效的训练;第三,研究TRQAM在离策略RL设置中的理论性质,例如收敛性保证和样本复杂度分析;第四,探索TRQAM在连续控制之外的应用,例如离散决策问题或自然语言处理中的策略微调;第五,研究TRQAM的安全性,例如在安全关键的RL任务中如何信任区域约束来保证策略不违反安全限制。
复现评估
作者已经开源了代码(https://github.com/yonghdong/trqam),并在博客(https://yonghdong.github.io/blog/trqam)上提供了额外的实验细节和可视化。实验使用标准的开源基准OGBench(50个任务)和Robomimic,这些基准是广泛可用的。作者报告了8个种子的平均结果,并提供了标准差,这有助于评估结果的统计显著性。超参数的详细设置在附录中提供,这有助于复现实验。然而,作者没有明确报告所需的计算资源,例如GPU类型、训练时间或总计算成本。虽然代码开源和基准可用有助于复现,但缺少计算资源细节可能使精确复现具有挑战性。总体而言,TRQAM的复现性评估为中等:代码开源和基准可用是有利的,但缺少计算资源细节和VJP成本的具体量化是不利的。
论文图表
Figure 1包含两个子图:(a)展示了不同方法的优化形式,说明TRQAM将信任区域控制内化到SOC采样动力学中,而QAM和QAM-E的最优解采用指数倾斜形式,受到破坏性漂移的影响,其中小的critic错误可以被指数放大成从预训练先验的大偏离;(b)展示了在50个OGBench任务上的离策略RL成功率,TRQAM达到68%,明显优于伴随匹配基线(QAM为35%,QAM-E为45%)和其他流策略微调范式(FQL为28%,IFQL为35%,DSRL为46%,CGQL-L为23%),这些方法缺乏这样的收敛保证。
这张图对理解论文很重要,因为它直观地展示了TRQAM与现有方法的区别:部分说明了固定温度伴随匹配的破坏性漂移问题,部分展示了TRQAM在50个OGBench任务上的优越性能,为本文的核心贡献提供了实证支持。
Figure 2展示了在Robomimic-can任务上TRQAM、QAM和QAM-E的训练曲线,包括伴随损失(左y轴,对数尺度)和成功率(右y轴,线性尺度)。QAM和QAM-E的伴随损失即使在梯度裁剪下也增长到超过10的20次方,将任务成功率从80%以上推至接近零(加减1个标准差),而TRQAM保持稳定。这种崩溃在Robomimic-lift和Robomimic-can的大多数超参数设置中持续存在。
这张图对理解论文很重要,因为它提供了固定温度伴随匹配脆弱性的实证证据,支持了指数放大理论,并展示了TRQAM如何解决这个问题。这是支持本文核心动机的关键实验证据。