借助新的物理模拟引擎加速强化学习

2021 年 8 月 16 日 TensorFlow

发布人：Google Research 高级软件工程师 C. Daniel Freeman 和资深软件工程师 Erik Frey

强化学习 (Reinforcement Learning) 是一种指导机器人在现实世界完成导航和执行操作的热门方法，其本身可以简化并表示为刚性物体 [1]（即受外力作用时不会变形的固体物理对象）之间的相互作用。RL 通常会利用模拟，在有效的时间段内收集训练数据。在模拟过程中，我们会对任意数量的复杂对象做近似处理，将其视作由关节组成、且由致动器提供动力的许多刚体。但挑战也应运而生：即便是行走、使用工具或组装积木等简单的任务，RL 智能体也通常需要数百万甚至数十亿的模拟帧才能熟练地完成。

虽然通过回收模拟帧来提高训练效率的方法已经取得进展，但一些 RL 工具需仍通过模拟帧的生成分布到许多模拟器来回避这个问题。这些分布式模拟平台已经取得了亮眼表现，训练速度也非常快。但它们必须在拥有数千个 CPU 或 GPU 的计算集群上运行，往往大多数研究人员无法访问这些 CPU 或 GPU。

通过回收模拟帧来提高训练效率
https://ai.googleblog.com/2020/08/tackling-open-challenges-in-offline.html

在“Brax：用于大规模刚体模拟的微分物理引擎 (Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation)”中，我们提出了一种新的物理模拟引擎，该引擎可以在单个 TPU 或 GPU 下可提供与大型计算集群相媲美的性能。该引擎旨在保障在单个加速器上高效并行运行数千个物理模拟和机器学习 (ML) 算法，以及在相互连接的加速器 pod 中无缝扩展数百万个模拟。我们开源了引擎、参考 RL 算法和模拟环境，这些都可以通过 Colab 进行访问。通过使用这个新的平台，我们展现了比传统工作站安装快 100-1000 倍的训练速度。

三种典型 RL 工作流。左边显示了一个标准的工作站流程：一台配备 CPU 的机器，在这样的环境中，训练需要数小时或数天。中间显示了一个标准的分布式模拟流程：通过将模拟移交给数千台机器，训练需要几分钟。右图显示了 Brax 流程：在单个 CPU/GPU 芯片上并行推进学习和大批量模拟

Brax：用于大规模刚体模拟的微分物理引擎
https://arxiv.org/abs/2106.13281
引擎、参考 RL 算法和模拟环境
https://github.com/google/brax
Colab
https://colab.sandbox.google.com/github/google/brax/blob/main/notebooks/training.ipynb

物理模拟引擎设计机会

刚体物理在电子游戏、机器人、分子动力学、生物力学、图形和动画等领域均有应用。为了对这类系统进行准确建模，模拟器集成了来自重力、电机驱动、关节约束、物体碰撞等方面的力用于模拟物理系统根据时间所做出的动作。

模拟情形：三个球体、一个壁、两个关节和一个致动器。系统会整合每个模拟时步中的力和力矩，以更新每个物体的位置、旋转和速度

在仔细研究当今大多数物理模拟引擎的设计方式过后，不难看出，目前存在一些提高效率的潜在机会。正如我们上面所提到的，典型的机器人学习流水线会将单个学习器置于已有的、许多且并行模拟的紧密反馈中，但在分析这个架构后，人们发现：

1. 这种架构带来了巨大的延迟瓶颈。由于数据必须通过数据中心内的网络进行传输，因此学习器必须等待 10,000 多纳秒才能从模拟器中提取经验。如果这种经验在与学习器神经网络相同的设备上运行，延迟时间则会下降到不足 1 纳秒。

2. 训练智能体（一模拟步后随之更新智能体的神经网络）所需的计算量远不及打包数据（即在引擎内编组数据，转换为诸如 protobuf 之类的有线格式，然后进入 TCP 缓冲区，在学习器端撤消所有这些步骤）所需的计算量。

3. 每个模拟器中发生的计算非常相似，但并不完全相同。

Brax 设计

针对这些观察结果，Brax 旨在使物理计算在其数千个并行环境中的每一个环境中保持完全相同，方法是确保模拟没有分支（即模拟会因环境状态而发生相应变化的 “if” 逻辑）。物理引擎的分支示例之一是在球和墙壁之间施加接触力：系统将根据球是否接触墙，执行不同的代码路径。也就是说，如果 (if) 球接触墙壁，则将执行用于模拟球从墙壁反弹的单独代码。Brax 混合使用以下三种策略来避免分支：

用连续函数代替离散分支逻辑，例如使用符号距离函数 (Signed distance function) 模拟球壁之间的接触力。这种方法可以最大程度地提高效率。
在 JAX 的即时编译期间评估分支。许多基于环境静态属性的分支，比如两个对象是否可能发生碰撞，系统可能会在模拟时间之前对其进行评估。
在模拟期间运行分支的两侧，但只选择所需的结果。 因为这将执行一些最终没有使用的代码，所以相较于上文，它会浪费操作。

在保证计算一致完全后，整个训练体系结构的复杂度就会降低，且更便于在单个 TPU 或 GPU 上执行。这样即可消除跨机器通信的计算开销和延迟。在实践中，对于相同的工作负载，这些更改将训练成本降低了 100 -1000 倍。

Brax 环境

环境 (Environments) 是微小的封装世界，定义了 RL 智能体要学习的任务。环境不仅包含模拟世界的方法，还包含如何观察这个世界以及如何在该世界中定义目标等功能。

近年来出现了一些用于测试新的 RL 算法，并使用研究员熟悉的指标来评估这些算法影响的基准环境。Brax 内含热门 OpenAI Gym 中的 4 个即用型环境：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 与 OpenAI 基准测试环境中训练，后者由 MuJoCo 物理引擎提供支持。

Ant
https://github.com/google/brax/blob/main/brax/envs/ant.py
OpenAI
https://github.com/openai/gym/blob/master/gym/envs/mujoco/ant.py

在下图中，蓝线代表已安装的标准工作站，其中学习器在 GPU 上运行，模拟器在 CPU 上运行。我们看到通过训练 Ant 按照合理的熟练度（y 轴上的分数为 4000）来运行，所需的时间从蓝线的大约 3 小时下降到在加速器硬件上使用 Brax 时的大约 10 秒。有趣的是，即使仅使用 CPU（灰线），Brax 的执行速度也加快了一个数量级。这得益于处于同一进程中的学习器和模拟器。

Brax 的优化 PPO 与标准 GPU 支持的 PPO 在学习 MuJoCo-Ant-v2 环境时，其执行 1 千万步时的表现对比情况。请注意，x 轴是执行作业所需的实际时间（以秒为单位）。阴影区域表示超过 5 个副本的最低和最高性能种子，实线表示平均值

PPO
https://arxiv.org/abs/1707.06347

物理保真

设计与现实世界行为相匹配的模拟器是一个尚未解决的已知难题。尽管如此，将 Brax 与参考模拟器进行比较，至少确保其产生同样有效的输出。在这种情况下，我们再次比较 Brax 与 MuJoCo，后者因其模拟质量而广受好评。我们希望看到的是，在其他条件相同的情况下，无论是在 MuJoCo 还是 Brax 中训练，策略都具有相似的奖励轨迹。

MuJoCo-Ant-v2 与 Brax Ant 的对比情况，图示为根据环境获得的平均回合分数而针对不同环境绘制的步数。两种环境都采用相同的 SAC 标准进行训练。阴影区域表示五次运行中表现最低和最高的种子，实线表示平均值

SAC
https://arxiv.org/pdf/1812.05905.pdf

这些曲线表明，随着两个模拟器的奖励以大致相同的速度增长，这两个引擎计算物理的复杂性和难度都相当。由于两条曲线都以大致相同的奖励达到顶峰，所以我们相信，相同的常规物理限制适用于在任一模拟中尽其所能操作的智能体。

我们还可以测量 Brax 保持线性动量、角动量和能量守恒的能力。

Brax 和其他几个物理引擎的线性动量（左）、角动量（中）和能量（右）非守恒比例。y 轴表示预期计算的漂移（漂移越小越好），x 轴表示模拟的时间量

这种物理模拟质量的衡量标准由 MuJoCo 的作者率先提出，其目的是了解模拟在计算越来越大的时间步长时的偏航情况。在这里，Brax 的表现与其近邻曲线相似。

结论

我们邀请研究人员通过在 Brax Training Colab 中训练他们自己的策略，对 Brax 的物理保真度进行更定性的测量。学习到的轨迹与在 OpenAI Gym 中看到的轨迹非常相似。

我们的工作使快速、可扩展的 RL 和机器人研究变得更加容易，以前只能通过大型计算集群才能实现的研究现在可以在工作站上运行，或者托管在 Google Colaboratory 免费运行。我们的 Github 代码库不仅包括 Brax 模拟引擎，还包括大量用于快速训练的参考 RL 算法。我们迫不及待地想看看 Brax 能够实现什么样的新研究。

托管在 Google Colaboratory
https://colab.research.google.com/github/google/brax/blob/main/notebooks/training.ipynb
Github 代码库
https://github.com/google/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，感谢他们帮助孵化这个项目。

[1] 由于现实世界的复杂性，我们仍然在开展探索柔性物体的物理研究。