强化学习(RL) 是一种教导机器人导航和操纵物理世界的流行方法，其本身可以简化并表示为刚体1（即施加力时不会变形的固体物理对象）之间的相互作用。为了便于在实际时间内收集训练数据，RL 通常利用模拟，其中任意数量的复杂对象的近似值由许多刚体组成，这些刚体通过关节连接并由执行器驱动。但这带来了一个挑战：RL 代理通常需要数百万到数十亿个模拟帧才能熟练完成即使是简单的任务，例如行走、使用工具或组装玩具积木。
虽然通过回收模拟帧 来提高训练效率已经取得了进展，但一些强化学习工具却通过将模拟帧的生成分布在许多模拟器上来规避这个问题。这些分布式模拟平台产生了令人印象深刻的结果，训练速度非常快，但它们必须在具有数千个 CPU 或 GPU 的计算集群上运行，而大多数研究人员无法访问这些集群。
在“ Brax - 用于大规模刚体模拟的可微分物理引擎”中，我们介绍了一种新的物理模拟引擎，它仅使用单个 TPU 或 GPU 即可达到大型计算集群的性能。该引擎旨在高效地在单个加速器上运行数千个并行物理模拟以及机器学习 (ML) 算法，并在互连加速器 pod 之间无缝扩展数百万个模拟。我们已经开源了引擎以及可通过Colab访问的参考 RL 算法和模拟环境。使用这个新平台，我们展示了与传统工作站设置相比 100-1000 倍更快的训练速度。
三种典型的 RL 工作流程。左侧显示了典型的工作站流程：在单台机器上，在 CPU 环境中，训练需要数小时或数天。中间显示了典型的分布式模拟流程：通过将模拟扩展到数千台机器，训练只需几分钟。右侧显示了 Brax 流程：学习和大批量模拟在单个 CPU/GPU 芯片上并行进行。
物理模拟引擎设计机会
刚体物理学用于视频游戏、机器人技术、分子动力学、生物力学、图形和动画等领域。为了准确地模拟此类系统，模拟器会整合重力、电机驱动、关节约束、物体碰撞等力，以模拟物理系统随时间的运动。
模拟三个球体、一堵墙、两个关节和一个执行器。对于每个模拟时间步长，力和扭矩被整合在一起以更新每个物理体的位置、旋转和速度。
仔细观察当今大多数物理模拟引擎的设计，可以发现有很多提高效率的机会。如上所述，典型的机器人学习流程将单个学习者置于紧密的反馈中，并同时进行许多模拟，但分析这种架构后，我们发现：
这种布局会带来巨大的延迟瓶颈。由于数据必须通过数据中心内的网络传输，因此学习者必须等待 10,000 多纳秒才能从模拟器中获取经验。如果这种经验已经与学习者的神经网络位于同一台设备上，延迟将降至 <1 纳秒。
训练代理所需的计算（一个模拟步骤，然后是代理神经网络的一次更新）被打包数据所花费的计算所掩盖（即，在引擎内编组数据，然后将其转换为protobuf等有线格式，然后转换为TCP缓冲区，然后在学习者端撤消所有这些步骤）。
每个模拟器中发生的计算非常相似，但并不完全相同。
Brax 设计
针对这些观察，Brax 的设计确保了模拟没有分支（即模拟“ if”逻辑会因环境状态而出现分歧），从而使其物理计算在其数千个并行环境中完全相同。物理引擎中分支的一个例子是球和墙之间接触力的应用：根据球是否接触墙，将执行不同的代码路径。也就是说，如果球接触到墙，将执行单独的代码来模拟球从墙上弹起。Brax 采用了以下三种策略的组合来避免分支：
用连续函数代替离散分支逻辑，例如使用有符号距离函数近似球壁接触力。这种方法可最大程度地提高效率。
在 JAX 的即时编译期间评估分支。许多基于环境静态属性的分支（例如两个对象是否可能发生碰撞）可以在模拟时间之前进行评估。
在模拟过程中运行分支的两侧，然后仅选择所需的结果。因为这会执行一些最终不会使用的代码，因此与上述方法相比，它浪费了操作。
一旦保证计算完全一致，整个训练架构的复杂性就可以降低，以便在单个 TPU 或 GPU 上执行。这样做可以消除跨机器通信的计算开销和延迟。实际上，对于同等工作负载，这些变化将训练成本降低了 100 倍至 1000 倍。
Brax 环境
环境是封装好的微型世界，用于定义强化学习代理要学习的任务。环境不仅包含模拟世界的手段，还包含各种功能，例如如何观察世界以及该世界中目标的定义。
近年来，出现了一些标准基准环境，用于测试新的 RL 算法并使用研究科学家普遍理解的指标来评估这些算法的影响。Brax 包括四个来自流行的OpenAI gym 的即用型环境： Ant、HalfCheetah、Humanoid和Reacher。
从左到右：Ant、HalfCheetah、Humanoid 和 Reacher 是 RL 研究的流行基线环境。
Brax 还包括三个新颖的环境：灵巧地操纵物体（机器人技术中的热门挑战）、广义运动（前往放置在其周围任何地方的目标的代理）和工业机器人手臂的模拟。
左图： Grasp，一只学习灵巧操作的爪手。中图： Fetch，一只像盒子一样的玩具狗，学习基于目标的一般运动策略。右图：UR5e的模拟，一种工业机器人手臂。
性能基准
分析 Brax 性能的第一步是测量它模拟大批量环境的速度，因为这是学习者获取足够的经验以快速学习所需要克服的关键瓶颈。
下面这两个图表显示了 Brax 在并行模拟越来越多环境时可以产生多少物理步骤（环境状态更新）。左侧的图表显示，Brax 每秒的步数与并行环境的数量线性相关，仅在 10,000 个环境时达到内存带宽瓶颈，这不仅足以训练单个代理，也适合训练整个代理群体。右侧的图表显示了两件事：首先，Brax 不仅在 TPU 上表现良好，而且在高端 GPU 上也表现良好（参见V100和P100曲线）；其次，通过利用 JAX 的设备并行原语，Brax 可以在多个设备之间无缝扩展，达到每秒数亿个物理步骤（参见TPUv3 8x8曲线，即 64 个TPUv3芯片通过高速互连直接相互连接）。
左图： 4x2 TPU v3 上每个 Brax 环境的每秒模拟步数的缩放。右图： Ant 环境上多个加速器的每秒模拟步数的缩放。
分析 Brax 性能的另一种方法是测量其对在单个工作站上运行强化学习实验所需时间的影响。在这里，我们将 Brax 训练流行的Ant基准环境与由MuJoCo 物理引擎提供支持的OpenAI 对应环境进行比较。
在下图中，蓝线代表标准工作站设置，其中学习器在 GPU 上运行，模拟器在 CPU 上运行。我们看到，训练一只蚂蚁以合理的熟练程度（y 轴上的分数为 4000）运行所需的时间从蓝线的约 3 小时下降到加速器硬件上使用 Brax 的约 10 秒。值得注意的是，即使仅在 CPU 上（灰线），Brax 的执行速度也要快一个数量级以上，这得益于学习器和模拟器都处于同一进程中。
Brax 的优化PPO与标准 GPU 支持的 PPO 学习 MuJoCo-Ant-v2 环境，评估了 1000 万步。请注意，x 轴是以秒为单位的对数挂钟时间。阴影区域表示 5 个副本中表现最低和最高的种子，实线表示平均值。
物理保真度
设计一个与现实世界行为相匹配的模拟器是一个已知的难题，这项工作并没有解决。尽管如此，将 Brax 与参考模拟器进行比较还是很有用的，以确保它产生的输出至少同样有效。在这种情况下，我们再次将 Brax 与MuJoCo进行比较，后者因其模拟质量而备受好评。我们期望看到，在其他所有条件相同的情况下，无论是在 MuJoCo 还是 Brax 中训练，策略都有类似的奖励轨迹。
MuJoCo-Ant-v2 与 Brax Ant 对比，显示环境步骤数与环境平均情节得分的关系。两种环境均使用相同的SAC标准实现进行训练。阴影区域表示五次运行中表现最低和最高的种子，实线表示平均值。
这些曲线表明，当两个模拟器的奖励以大致相同的速度增加时，两个引擎计算的物理复杂度或难度也相当。由于两条曲线的最高奖励大致相同，因此我们相信，相同的一般物理限制适用于在任一模拟中尽其所能运作的代理。
我们还可以测量布拉克斯保存线性动量、角动量和能量的能力。
Brax 以及其他几个物理引擎的线性动量（左）、角动量（中）和能量（右）非守恒缩放。y 轴表示与预期计算的漂移（漂移越大，越小，越好），x 轴表示模拟的时间量。
这种物理模拟质量衡量标准最初是由 MuJoCo 的作者提出的，旨在了解模拟在计算越来越大的时间步长时如何偏离轨道。在这里，Brax 的表现与其邻居相似。
结论
我们邀请研究人员通过在Brax Training Colab 中训练自己的策略来对 Brax 的物理保真度进行更定性的测量。学习到的轨迹与在 OpenAI Gym 中看到的轨迹明显相似。
我们的工作使快速、可扩展的强化学习和机器人研究变得更加容易实现——以前只能通过大型计算集群才能实现的研究现在可以在工作站上运行，或者通过托管的 Google Colaboratory免费运行。我们的 Github 存储库不仅包含 Brax 模拟引擎，还包括一系列用于快速训练的参考强化学习算法。我们迫不及待地想看看 Brax 能带来什么样的新研究。
致谢
我们要感谢我们的论文合著者：Anton Raichuk、Sertan Girgin、Igor Mordatch 和 Olivier Bachem。我们还要感谢 Erwin Coumans 就构建物理引擎提供的建议、Blake Hechtman 和 James Bradbury 就 JAX 和 XLA 提供的优化帮助以及 Luke Metz 和 Shane Gu 提供的建议。我们还要感谢 Vijay Sundaram、Wright Bagwell、Matt Leffler、Gavin Dodd、Brad Mckee 和 Logan Olson 帮助孵化这个项目。
评论