2025年6月3日星期二

RL训练太慢?UC伯克利重磅开源FastTD3:并行+大批次+分布Critic,人形机器人3小时搞定

点击下方卡片,关注"AI生成未来"

如您有工作需要分享,欢迎联系:aigc_to_future

作者:Haoran Geng等

解读:AI生成未来

图片

论文链接:https://arxiv.org/pdf/2505.22642
项目链接:https://younggyo.me/fast_td3/

亮点直击

  • 提出了FastTD3,一种简单、快速且高效的强化学习算法,能够解决多种运动和操作任务,而之前的RL算法需要数十小时才能完成或无法解决。
  • 可以通过一个极其简单的方案实现:使用大批次更新、并行仿真、分布式强化学习以及精心调优的超参数来训练TD3智能体。展示了各种设计选择的有效性。
  • 发布了一个易于使用的开源FastTD3实现,以加速机器人领域的强化学习研究。该实现支持流行的测试平台,如HumanoidBench、IsaacLab和MuJoCo Playground。

总结速览

解决的问题

  • 训练速度慢:现有RL算法(如PPO、PQL等)在复杂机器人任务(如HumanoidBench、IsaacLab)中训练时间过长(如48小时未收敛)。
  • 实现复杂性高:部分高效算法(如PQL)依赖异步并行等复杂设计,难以广泛应用。
  • 样本效率与稳定性的权衡:PPO等on-policy算法样本效率低,而off-policy算法(如TD3)虽高效但训练不稳定或速度不足。
效果展示
FastTD3解决了MuJoCo游乐场的一系列任务,实现了与PPO相似或更快的挂钟时间:

FastTD3在单个GPU上在不到3小时的时间内解决了HumanoidBench中的各种任务,而之前的算法即使经过数十小时的训练也很困难。然而可以看到,由于奖励设计的局限性,行为非常不自然:

证明了FastTD3可用于模拟到真实的传输。使用非常快速和方便的MuJoCo Playground在模拟中训练Booster T1策略,并使用Booster Gym将其转移到真实的Booster T1机器人上:

不同的 RL 算法可能需要不同的奖励函数

提出的方案

  • 并行仿真(Parallel Simulation):加速数据收集。
  • 大批次更新(Large-Batch Updates):提升训练稳定性。
  • 分布化Critic(Distributional Critic):改进价值估计。
  • 精细调参:针对机器人任务优化的超参数。

应用的技术

  • 算法基础:TD3(双延迟确定性策略梯度)的off-policy框架。
  • 加速技术:GPU并行化仿真(单A100实现)、大批次梯度更新。
  • 性能优化:分布化Critic、简化异步流程。

达到的效果

  • 训练速度:在HumanoidBench等任务中,3小时内完成训练(对比PPO更快,尤其在复杂地形任务中)。
  • 稳定性:训练过程稳定,无需复杂调参。
  • 易用性:提供轻量级PyTorch实现,支持快速部署和扩展(兼容HumanoidBench/IsaacLab/MuJoCo Playground)。
  • 兼容性:与前沿RL研究(如SR-SAC、TDMPC2等)正交,可灵活集成新技术。

FastTD3:面向人形机器人控制的简单、快速、高效强化学习算法

FastTD3是基于双延迟深度确定性策略梯度算法的高性能变体,专为复杂机器人任务优化。这些优化借鉴了Li等人的研究发现:并行仿真、大批次训练和分布化评论家对离线策略强化学习算法的性能提升至关重要。

设计选择

并行环境
与Li等人的观察一致,发现大规模并行环境能显著加速TD3训练。假设:确定性策略梯度算法与并行仿真的结合具有特殊优势,因为并行环境产生的随机性增加了数据分布的多样性。这使得TD3既能充分发挥其价值函数高效利用的优势,又能缓解探索能力不足的缺陷。

大批次训练
实验表明,使用32,768的超大批次规模训练FastTD3智能体效果显著。推测:在并行环境下,大批次更新通过确保每次梯度更新的数据多样性,为评论家提供了更稳定的学习信号。否则,除非保持高更新数据比(UTD),否则大量数据可能永远不被智能体利用。虽然增大批次规模会增加单次更新的耗时,但由于训练效率提升,总体训练时间往往反而缩短。

分布化强化学习
分布化评论家在多数情况下能提升性能。但需注意这会引入额外超参数——。尽管实验表明它们并不难调节,未来可考虑集成SimbaV2提出的评论家奖励归一化方案。

截断双Q学习(CDQ)
Nauman等人指出,当结合层归一化时,使用Q值均值比CDQ采用的Q值最小值效果更好。实验显示,在没有层归一化时趋势相反——CDQ仍是关键设计选择,采用最小值通常表现更优。这表明CDQ仍是需要按任务调节的重要超参数。

网络架构
采用1024-512-256单元数的多层感知机(MLP),行动者采用512-256-128单元数。实验发现,更小的模型会降低时间效率和样本效率。虽然尝试过类似BRO或Simba的残差连接与层归一化,但这些设计反而会拖慢训练且无显著收益。推测:并行仿真和大批次训练提供的数据多样性降低了更新的"离线策略性",从而缓解了自助法、函数逼近和离线策略学习"致命三角"带来的不稳定性。因此即使没有残差连接或层归一化等稳定结构,训练过程仍能保持稳定。

探索噪声调度
与PQL采用混合噪声(为每个环境从采样不同高斯噪声尺度)的方案不同,发现混合噪声未带来显著增益。但因其代码实现简便,仍采用该方案。实验证实的大噪声对FastTD3有益。

更新数据比(UTD)
已有研究表明:提高UTD(即每环境步长的梯度更新次数)通常需要额外技术或架构调整。但FastTD3在标准3层MLP(无归一化)下,样本效率随UTD提升而改善,仅需付出更长的训练耗时。这归因于FastTD3工作在极低UTD(通常每128-4096个并行环境步长仅2/4/8次更新),降低了高UTD导致的早期过拟合风险。

经验回放池大小
不同于全局固定回放池大小,设定为。实际使用中发现:增大能提升性能,但会增加GPU内存占用(所有数据存储在GPU上)。

实现细节

并行环境
对于IsaacLab和MuJoCo Playground,直接使用它们原生的并行仿真支持。然而,HumanoidBench不支持基于 GPU 的并行化,因此使用Stable Baselines3的 SubprocVecEnv。HumanoidBench 的默认配置会为每个仿真启动一个基于 GPU 的渲染器,这使得运行超过 100 个环境变得困难。已向 HumanoidBench 提交了一个Pull Request,添加了禁用默认 GPU 渲染器的支持,并已合并到主分支。

图片

环境封装器
为了构建一个支持不同测试套件(各自采用不同配置)的易用代码库,为每个套件开发或使用了封装器:

  • MuJoCo Playground:使用原生 RSLRLBraxWrapper,将 JAX 张量转换为 PyTorch 张量,并遵循RSL-RL的 API。由于该封装器不支持在重置环境前保存最终观测值,在一个独立的分支中实现了这一功能,并计划将其合并到主仓库。
  • IsaacLab:由于原生支持 PyTorch,实现了一个符合 RSL-RL API 的简单封装器。目前,本文的实现不支持训练期间的渲染,因为 IsaacLab 不允许并发运行多个仿真。
  • HumanoidBench:开发了一个遵循 RSL-RL API 的封装器,并将 NumPy 数组转换为 PyTorch 张量。

非对称 Actor-Critic对于IsaacLab和MuJoCo Playground(它们通常为评论家网络提供特权状态),实现了非对称 Actor-Critic的支持。

AMP 与 torch.compile尽管基于JAX的 RL 实现在近年来因速度优势而流行,但仍选择PyTorch作为基础,因其简单性和灵活性。在单块A100 GPU上,使用AMP(自动混合精度)和bfloat16可将训练速度提升高达40%,且未观察到不稳定性。还基于LeanRL  支持 torch.compile,实验显示其可带来35%的加速。当同时使用AMP和 torch.compile 时,训练速度最高可提升70%。

经验回放池
没有固定全局回放池大小,而是设定为。这种方式能更好地解耦回放池大小与并行环境数量的影响。例如:

  • 如果全局回放池固定为100万,任务片段长度为1000,而用户将并行环境数从1000增加到2000,则回放池只能保存每个环境轨迹的一半(在 500 步 后就会开始丢弃早期样本),这可能影响性能。
  • 但如果为每个环境设定 ****,则回放池能完整保存轨迹,不受并行环境数量的影响。

由于专注于非视觉领域,所有数据均存储在GPU上,以避免CPU-GPU 数据传输的开销。

实验

实验设置
对于DreamerV3、SAC和TDMPC2基线算法,使用 HumanoidBench 代码库中提供的三次运行学习曲线。由于每次训练耗时48 小时,通过插值计算实际时间戳来绘制曲线。对于SimbaV2,使用官方代码库在涉及灵巧手的任务上进行实验。SimbaV2 采用单次运行结果,未来计划增加更多实验。所有 FastTD3 结果均为三次运行的平均值。实验在配备单块 NVIDIA A100 GPU和16 CPU 核心的云实例上完成。

实验结果
在下图3中提供了每个套件的所有任务的汇总结果,在下图4中提供了选定任务集的单独结果。在下图5和图6中提供了广泛的实验结果,研究了各种设计选择的影响。

图片
图片
图片
图片

不同RL算法需要不同的奖励函数
在 MuJoCo Playground 训练人形机器人运动策略时,发现PPO和FastTD3即使使用相同奖励函数也会产生显著不同的步态(见下图7a-b)。推测:现有奖励函数通常针对 PPO 调优,不同算法可能需要不同的奖励结构来产生理想行为。为此专门为 FastTD3 调整了奖励函数——增加惩罚项强度。得益于 FastTD3 的快速训练,该调参过程非常高效。如图7c 所示,调整后的奖励使 FastTD3 学习到比图7a 更稳定美观的步态。而将 FastTD3 调优的奖励用于 PPO 训练时,却会导致步态缓慢不可用(图7d)。这表明标准指标——片段回报——可能无法有效衡量策略的实际可用性。

图片

FastSAC 实验
为验证方法通用性,本文将 FastTD3 方案迁移至 SAC开发出FastSAC。实验显示 FastSAC 显著快于原始 SAC(下图8),但训练稳定性较差,推测源于高维动作空间中动作熵最大化的固有难度。鉴于 SimbaV2在主实验中明显快于原始 SAC,未来可将此类最新进展融入 FastTD3/FastSAC。

图片

FastTD3 的仿真到现实迁移
下图1的仿真到现实实验使用Booster Gym ,该平台支持12自由度 Booster T1 人形机器人控制(固定手臂/腰部/朝向)。为简化流程,将 Booster Gym 的机器人配置和奖励函数移植到原本支持23自由度控制的 MuJoCo Playground。相比基于 IsaacGym 的 Booster Gym,在 MuJoCo Playground 中训练 FastTD3 能大幅简化和加速迭代周期。

图片

讨论

本文提出了FastTD3——一种简洁、高效且性能优异的强化学习算法,该算法成功解决了HumanoidBench、IsaacLab和MuJoCo Playground中的各类运动控制与机械操作任务。研究表明,在未引入新架构或训练技术的前提下,仅通过精心调优的超参数结合简洁算法,即可为复杂机器人任务构建出人意料的强基线。随本报告同步开源的FastTD3实现具有轻量化、易用化特点,提供用户友好功能与预设超参数配置。

需要特别说明的是,本研究目标并非宣称算法创新性或优越性。本文的方法直接建立在学界已有洞见之上:并行Q学习(PQL)揭示了如何通过大规模并行仿真扩展离策略强化学习,后续的并行Q网络(PQN)在离散控制中得出了相似结论。Raffin与Shukla同样证明SAC算法通过并行仿真与精细超参数调优可实现高效扩展。本工作旨在将这些洞见提炼为简洁算法,提供全面的设计选择实验分析,并发布开箱即用的实现方案。

本文对多个未来方向充满期待:首先,FastTD3与近期强化学习进展具有正交性,这些改进可轻松集成以推动技术前沿,预计整合过程将高效顺畅;其次,FastTD3作为离策略算法,非常适用于人形机器人控制的示范驱动强化学习场景,以及通过真实环境交互微调仿真训练策略;最后,FastTD3的快速迭代特性可赋能基于语言模型生成奖励的逆向强化学习框架,为长期困扰人形控制领域的奖励设计难题提供新思路。

参考文献

[1] FastTD3: Simple, Fast, and Capable Reinforcement Learning for Humanoid Control

致谢

如果您觉得这篇文章对你有帮助或启发,请不吝点赞、在看、转发,让更多人受益。同时,欢迎给个星标⭐,以便第一时间收到我的最新推送。每一个互动都是对我最大的鼓励。让我们携手并进,共同探索未知,见证一个充满希望和伟大的未来!



技术交流

加入「AI生成未来社区」群聊,一起交流讨论,涉及 图像生成、视频生成、3D生成、具身智能等多个不同方向,备注不同方向邀请入群!可添加小助手备注方向加群!

图片

没有评论:

发表评论

Apple提出UniGen!多模态理解生成统一xii新架构!CoT - V提升图像生成质量!

最新论文解读系列论文名:UniGen: Enhanced Training\x26amp;Test-Time Strate 最新论文解读系列 论文名: UniGen: Enhanced Training&Test-Time Strategies for Unified...