本文档介绍了由盛超胡、子卿范、超钦黄、李森、张亚、王彦峰以及陶大程共同发表于国际机器学习大会(International Conference on Machine Learning, ICML)2024年的研究论文《Q-Value Regularized Transformer for Offline Reinforcement Learning》。该研究旨在解决离线强化学习(Offline Reinforcement Learning)领域中的关键挑战,提出了一种名为Q值正则化Transformer(QT)的新算法,旨在融合条件序列建模(Conditional Sequence Modeling, CSM)与动态规划(Dynamic Programming, DP)两类方法的优势,以提升策略在轨迹缝合(stitching)和稀疏奖励场景下的性能。
第一,研究背景与动机 该研究的学术背景在于离线强化学习,即仅从预先收集的静态数据集中学习策略,而无需与环境进行在线交互。近年来,条件序列建模(CSM)范式,例如决策Transformer(Decision Transformer, DT),通过将离线RL视为序列建模问题,展现了强大的轨迹建模能力和训练稳定性。然而,CSM方法存在一个根本性缺陷:它们难以有效地从次优轨迹中“缝合”出最优轨迹。这是因为CSM依赖于轨迹级别的回报条件(如Return-to-Go, RTG),而单个轨迹内的采样回报可能与跨多个轨迹的最优未来回报不一致。另一方面,基于动态规划(DP)或Q学习的方法通过价值函数来估计每个状态的最优未来回报,从而天然具备缝合能力。但这些方法,特别是面对长时程和稀疏奖励任务时,往往存在学习不稳定、价值函数估计不准确的问题。
现有的一些尝试,如QDT,虽然试图结合两者,但其本质是使用保守的价值函数对数据集中的RTG进行重新标注,这仅是一种数据增强手段。在推理阶段,由于仍然是轨迹级别的建模,RTG值不匹配的问题依然存在,导致性能提升有限。基于这些观察,本研究的目标是提出一个能够有机结合Transformer的轨迹建模能力与DP方法的最优回报预测能力的框架,从而在保持稳定性的同时,显著提升模型的轨迹缝合能力和在稀疏奖励环境下的表现。
第二,研究方法与工作流程 QT算法的核心是一个基于Transformer架构的策略模型,其训练损失由两部分组成:1) 条件行为克隆损失,确保策略输出与离线数据集中行为策略的分布对齐;2) 策略改进损失,根据学到的Q值函数选择高回报的动作。具体工作流程如下:
1. 策略模型架构与基础训练(条件Transformer策略) 研究采用类似于Decision Transformer的Transformer模型作为策略网络。模型的输入是一个包含最近k步历史信息的轨迹序列,序列中的每个token由状态(s_t)、动作(a_t)和回报到目标(r̂_t,即从当前步到轨迹结束的累积奖励)组成。模型以自回归(auto-regressive)的方式,根据历史序列和未来的目标回报(RTG),预测当前状态对应的动作。对于连续动作空间,基础训练目标(L_dt)是最小化预测动作与数据集中真实动作之间的均方误差(MSE)。理论分析(定理3.1)表明,仅使用此损失进行训练,学到的策略将收敛于行为策略,但其性能上限被数据集中的轨迹质量所限制,且缺乏缝合能力。
2. Q值函数的学习与整合(Q值正则化训练) 为了引入缝合能力,QT并行地学习一个Q值函数(动作-价值函数)。研究者构建了两个Q网络(Q_φ1, Q_φ2)及其对应的目标网络,并采用n步贝尔曼方程(n-step Bellman equation)和双Q学习(Double Q-learning)技术来更新Q网络参数。n步估计相比1步估计能提供更准确的价值目标。 关键创新在于将学到的Q值函数整合到策略的训练过程中。策略的最终学习目标函数(L)是基础行为克隆损失(L_dt)和一个Q值最大化损失(L_q)的线性组合。具体公式为:L = L_dt - α * E[ Q_φ(s, π(τ)) ]。其中,α是一个经过归一化的超参数,用于平衡两个损失项。L_q项的作用是引导策略模型在生成动作时,倾向于选择那些具有较高Q值的动作。这种设计使得策略在遵循行为策略分布(通过L_dt)的同时,能够主动寻求更优的动作(通过L_q),从而在分布内实现策略提升。理论证明(定理3.2)表明,最终学习到的策略在价值函数上不会差于行为策略。
3. 基于Q值的推理过程 在推理阶段,QT无需像传统DT那样精心调试或猜测一个合适的初始RTG值。取而代之的是,QT预设一组候选的初始RTG值(例如{r̂0^0, r̂1^0, …, r̂m^0})。对于给定的初始状态,模型会根据这m个不同的RTG值,并行地生成m个候选动作序列(利用GPU并行计算能力,开销很小)。然后,利用训练好的Q值函数评估这m个候选动作在当前状态下的Q值。最终执行的动作是从这m个候选动作中选择Q值最高的那个。这个过程在每个时间步重复进行,并根据实际获得的奖励动态更新每个候选RTG值(r̂_t+1 = r̂_t - r_t)。这种方法将动作选择权交给了更可靠、状态相关的Q值函数,而非固定的、可能不匹配的RTG条件,从而在推理阶段进一步增强了决策的优异性。
4. 实验设计与评估流程 研究在广泛使用的D4RL基准测试集上进行了全面的实证评估,涵盖了五个不同特性的领域:Gym(标准连续控制)、Adroit(复杂机械臂操作)、Kitchen(多任务厨房环境)、Maze2D(评估缝合能力)以及AntMaze(评估稀疏奖励和长时程规划)。研究对比了大量基线方法,包括主流的DP方法(如CQL, IQL)、CSM方法(如DT, QDT, GDT, CGDT)、基于模型的方法(如COMBO)以及扩散模型方法(如Diffuser, DD)。对于每个任务,报告了超过30次随机测试的平均归一化得分(归一化为专家性能的百分比)。此外,研究还进行了深入的消融实验,以剖析QT各个组成部分(条件Transformer策略、Q值正则化训练、n步更新、基于Q值的推理)的贡献,并专门设计了实验来验证其在缝合能力、稀疏奖励处理以及长时程任务上的优势。
第三,主要研究结果 实验结果显示,QT在几乎所有测试的D4RL任务上都达到了最先进的性能(State-of-the-Art),显著超越了传统的DP方法和CSM方法。
在Gym领域的“medium”和“medium-replay”数据集上(这些数据集包含大量次优数据),QT相比其他Transformer-based方法提升显著。这表明QT中的策略改进项有效地引导模型在已探索的动作子空间内选择更优的动作。
在Adroit和Kitchen领域,这些任务对策略正则化的鲁棒性要求很高。QT凭借其强大的Transformer策略表达能力和有效的Q值正则化,表现远超扩散模型等基线,尤其是在需要泛化到未见状态和进行长程价值优化的任务中。
在Maze2D领域(缝合能力测试),QT在所有迷宫复杂度(open, umaze, medium, large)和奖励设置(稀疏、稠密)下都取得了最佳成绩,大幅领先于CQL、DT和QDT。这确凿地证明了QT成功地将DP的缝合能力赋予了CSM模型。
在AntMaze领域(稀疏奖励、长时程挑战),QT同样表现出色,在“medium-diverse”和“large-diverse”等高难度任务上显著优于其他方法。这表明QT的混合架构在需要稳定价值估计和长期规划的稀疏奖励场景中具有独特优势。
消融研究结果进一步验证了设计有效性: * 组件贡献:在Walker2d-medium-replay任务上,逐项添加QT的组件(Q值训练、Q值推理)都带来了性能提升。其中,在训练阶段加入Q值正则化(L_q)是提升最大的因素。使用n步贝尔曼更新优于1步更新。 * 缝合能力:在Maze2D任务上,QT的性能全面优于CQL、DT和QDT,表明其不仅具备了缝合能力,而且综合效果超过了任一类单独的方法。 * 稀疏奖励能力:在延迟奖励(稀疏)设置的Mujoco任务中,CQL表现很差,DT表现尚可,而QDT因使用CQL重新标注RTG而性能下降。QT则凭借其强大的策略正则化(L_dt)缓冲了不准确Q值估计的负面影响,在稀疏和稠密奖励下均表现优异。 * 长时程能力:随着输入序列长度(k)的增加,DT的性能出现波动,而QT的性能则持续提升,显示出其对长时程信息更好的利用能力和稳定性。
第四,研究结论与价值 本研究提出的Q值正则化Transformer(QT)算法,成功地将Transformer的轨迹建模能力与动态规划的最优回报预测能力相结合,为离线强化学习提供了一个新颖且强大的框架。其核心价值在于: 1. 方法学创新:QT创造性地将Q值函数作为正则化项引入条件序列模型的训练目标中,实现了策略正则化与策略改进的内在统一。这不同于简单的数据增强(如QDT),而是在学习机制层面进行了深度融合。 2. 性能突破:在广泛的D4RL基准测试中,QT取得了领先的性能,证明了其在处理轨迹缝合、稀疏奖励、长时程规划以及需要强正则化的复杂任务方面的综合优势。 3. 实用性强:QT在推理阶段采用基于Q值的多候选动作选择机制,免除了手动调整RTG的繁琐,提高了算法的实用性和鲁棒性。
该研究的科学意义在于为离线RL算法设计提供了一个新的思路,即如何有效地融合基于模型/序列的预测能力和基于价值的动态规划能力,以同时获得分布约束、策略提升和稳定学习等理想特性。其应用价值体现在能够从现有的、可能包含大量次优轨迹的静态数据集中,学习出更接近最优性能的策略,这对于机器人学、自动驾驶等数据收集成本高或风险大的领域具有重要潜力。
第五,研究亮点与特色 本研究的亮点包括: 1. 问题定位精准:明确指出并深入分析了CSM方法在“轨迹缝合”上的根本缺陷(RTG不一致性),以及DP方法在长时程和稀疏奖励下的不稳定性,从而提出了有针对性的融合方案。 2. 算法设计优雅而有效:通过一个简单的Q值正则化项,巧妙地将价值估计的“全局最优”信息注入到局部轨迹建模中,理论上有保障(策略提升定理),实践中效果显著。 3. 实证全面且具说服力:在覆盖不同挑战维度的五个标准领域进行测试,并与众多强基线进行对比,结果全面领先。详尽的消融实验清晰揭示了各模块的作用。 4. 推理机制创新:提出的基于Q值的多候选动作选择推理策略,既解决了RTG条件不匹配的难题,又充分利用了现代计算硬件的并行能力,是一个具有实用价值的工程创新。
第六,其他有价值内容 论文还讨论了QT在Atari游戏(离散动作域)上的扩展实验,结果显示其同样具有竞争力。同时,作者也指出了QT的局限性:其有效性依赖于数据集中存在显式的奖励信号。对于仅有状态-动作对而无奖励的人类演示数据集,QT的性能可能会受到限制。此外,附录中提供了详细的理论证明(关于DT的性能边界和QT的策略改进保证)和完整的超参数设置,增加了研究的严谨性和可复现性。