作者：C. Daniel Freeman，Erik Frey，Anton Raichuk，Sertan Girgin，Igor Mordatch and Olivier Bachem
来源：arXiv:2106.13281v1 [cs.RO] 24 Jun 2021
原文：https://arxiv.org/pdf/2106.13281.pdf
代码、数据和视频：https://github.com/google/brax
系列文章目录：
上一篇：
论文笔记(十六)：Learning to Walk in Minutes Using Massively Parallel Deep Reinforcement Learning
下一篇：

Brax运动和灵巧的操作策略可以在几秒钟到几分钟内用现代加速器训练。Brax广泛使用JAX[1]库的自动矢量化、设备并行化、及时编译和自动差异化基元。这样，它就可以在数百个相连的加速器上的数千个独立环境中模拟简单的刚体物理系统。对于单个加速器，Brax在OpenAI Gym的MuJoCo Ant[2]在其他环境中达到每秒数百万的模拟步骤。更多细节见第六节或我们Colab[3]互动训练策略。

本文的结构如下：我们首先在第二节提供了发动机。 在第三节，我们描述了Brax结构，从低层次的物理基元开始，它们是如何相互作用的，它们是如何扩展对基于物理模拟感兴趣的从业者的。在第四节，我们回顾了我们ProtoBuf环境规范，并详细介绍了如何使用它来构建丰富的物理模拟任务，包括在初始版本中捆绑的任务套件。在第五节，我们参观了Brax一些捆绑的强化学习算法。在第六节，我们是对的Brax加速器上的扩展行为和Brax和MuJoCo在OpenAI Gym对风格学习中的性能进行了分类。最后，在第七节，我们讨论了引擎的局限性和可能的扩展。

在过去的几年里，强化学习社区在物理模拟环境的研究和控制方面取得了重大进展。这一进展源于强大的算法技术[4-9]与可获得的模拟软件[10-14]的融合。在算法方面，无模型优化技术，如近似策略优化（PPO）[6]以及软演员的批评方法（SAC）[5]爆炸性流行，可以轻松解决前十年的很多问题 "硬 "控制问题。在模拟方面，员工可以选择各种引擎后端来支持他们对模拟环境的研究，包括MuJoCo[10]、pybullet[15]和physX等等，其中许多是微分[16-22， 14]。

这些引擎和算法虽然相当强大，多年来为算法创新提供了坚实的基础，但并非没有缺点。正如它所实践的那样，由于样本的高度复杂性，对于许多用例来说，强化学习仍然是昂贵和缓慢的：RL在探索过程中，只有数百维的状态空间需要数百万到数十亿的模拟步骤。作为环境步骤的一部分，这个问题只会越来越严重，因为环境越来越需要交互式物理计算。

这些引擎和算法虽然相当强大，多年来为算法创新提供了坚实的基础，但并非没有缺点。正如它所实践的那样，由于样本的高度复杂性，对于许多用例来说，强化学习仍然是昂贵和缓慢的：RL在探索过程中，只有数百维的状态空间需要数百万到数十亿的模拟步骤。作为环境步骤的一部分，这个问题只会越来越严重，因为环境越来越需要交互式物理计算。

模拟引擎的设计在三个方面导致了这个不可访问的问题：

首先，目前使用的大多数仿真引擎都在CPU上运行，而RL算法则在GPU或TPU在另一个过程或另一台机器上运行。由于跨机数据传输和网络流量的延迟RL实验所需时间的主导因素。

其次，大多数模拟引擎都是黑盒子：它们不提供采样环境的梯度，这使得它们只适合无模型RL方法。这种缺乏可分性迫使研究人员使用速度慢、效率低的优化方法。

最后，大多数模拟引擎在另一个方面是黑盒子：它们要么是封闭的，要么是基于与强化学习算法完全不同的技术栈。这种缺乏内省的做法不仅限制了快速迭代和调试，损害了生产力，而且使研究人员无法理解环境状态和行动空间之间的关系，从而指导新的RL研究往往至关重要。

Brax在最大坐标中模拟物理交互[33]，场景中每个可以自由移动的独立实体都被单独跟踪。这些数据–位置、旋转方向、速度和角速度–它通常是模拟过程中唯一动态变化的数据。所有其他动态关系，如关节、执行器、碰撞和集成步骤，都是基本状态数据的转换。这是数据基元QP中，以flax[34]实现数据类，并以其跟踪的典型坐标q和p命名。使矢量化容易，QP在一个场景中，平行场景的数量和身体的数量都有领先的批量尺寸。例如，对于4个平行场景，每个场景有10个身体，QP.pos[4， 10, 3]。

一个物理模拟对象通常包括额外的数据，因此我们将其他信息–质量、惯性、物体的尺寸等–捆绑在与特定QPs相关的抽象中。这些抽象概念是体、关节、执行器和碰撞器。举个例子，$joint.revolute$类将所有相关的元数据捆绑在一起，管理一对父体和子体的1度自由度约束。然后，这个类的应用函数计算力和扭矩–即速度、角速度、位置和旋转的变化–这对于将两个体约束到1自由度的关节配置是必要的。这两个体与$QP$对象中的特定指数相关。因此，调用$joint.revolute.apply(qp)$从整个$QP$对象–即被约束的两个$QP$实体–收集相关的物理数据，并返回对整个$QP$状态的矢量化、差分化更新。所有的Brax变换都遵循这种模式，即应用函数以这种方式变换一个$QP$。

为了完成物理学步骤，Brax然后在一个短的时间步骤中对QP数据的所有差分更新进行汇总，并通过二阶共轭欧拉更新对系统状态进行转换（对高阶积分器的扩展是直接的，但更多细节见7）。在整个过程中，我们尽可能地在执行器、关节、碰撞器，甚至整个模拟场景中进行并行化。参见Alg. 1中关于这个循环结构的伪代码，或者[35]中关于这个循环的代码。

一个总体的系统类处理所有这些更新和物理元数据的协调和记账。这个类还提供了一种通过步骤函数来执行单一模拟步骤的方法。

$$q{p}_{t+\delta t}=Brax\_system.step(q{p}_t,actions)$$

其中行动是系统中任何执行器需要的任何扭矩或目标角度。

修改或扩展这个控制流程就像实现一个符合这个结构的新的$Brax\_transformation$一样简单，然后在物理学步骤函数中适当地插入这个转化。

为了更好地可视化和探索Brax的核心物理循环，请看我们的基础知识Colab [36] 。

Brax提供了一个额外的抽象层来与物理模拟场景进行交互。在第4.1节中，我们描述了用于定义Brax系统的ProtoBuf规范–即描述系统中任何物理约束的最低级别数据。接下来，在第4.2节中，我们激励了env类，它允许实践者在Brax系统之上构建类似健身房的决策问题。最后，我们讨论了已经与Brax预先打包的环境。

我们的ProtoBuf文本规范允许用户定义场景中的所有体，它们是如何通过关节相互连接的，以及对象之间的任何执行器或碰撞器，成对的。

对于任何由关节连接的身体树，Brax的系统类将自动确定qp的位置和旋转，通过$system.default\_qp$方法将每个身体置于有效的关节配置中。

让人想起MuJoCo的基于xml的系统定义，用户可以用文本定义系统，也可以用程序定义系统。我们在附录A中提供了一个定义父体和子体之间联合的配置实例，包括纯文本形式和程序化形式。类似的配置文件在各自的环境文件中定义了Brax repo中的每个系统，例如[37]。请参阅我们的介绍性Colab笔记本，了解这两个apis的交互式游览。

对于顺序决策问题，我们必须跟踪超出物理学更新所需的额外元数据。我们提供了一个env类，用于处理任何初始化、重置、观察、行动或奖励函数定义的簿记，这些都是完全指定一个顺序决策问题所需要的。我们还提供了一个围绕这个类的包装器，以作为OpenAI体育馆式的接口直接与之连接。

为了说明Brax作为一个引擎的多功能性，我们在最初的版本中包括并解决了几个例子环境。MuJoCo-likes（蚂蚁、人形、半猎豹）、Grasp（一个灵巧的操纵环境）和Fetch（一个基于目标的运动环境）。这些环境的尺寸数据见表1。

在过去的几年里，强化学习和控制社区一直使用OpenAI体育馆MuJoCo任务作为开发算法的基准。虽然这些任务已被充分理解，并且基本上得到了解决，但我们提供了我们自己对其中三个环境的相当忠实的重构，作为比较的基准点，以帮助从业者的预期。由于引擎的微妙差异，这些环境与它们所基于的MuJoCo场景并不完全相同，我们在附录E中指出了主要的差异。

由于算法和硬件的进步使机器人能够解决更复杂的问题，灵巧的操纵任务已经爆炸性地流行起来。抓取 "是一个简单的取放环境，一只四指爪手必须拿起一个球并将其移动到目标位置。我们把这个环境主要作为一个概念证明，以证明我们引擎的接触物理学足以支持不复杂的操纵任务。关于一个成功策略的代表性轨迹样本，见附录B的图6。

我们对各种目标导向的运动任务进行了广泛的实验。Fetch代表了一个普遍稳定的环境定义，能够在5000万个环境帧内训练各种形态的人进行运动。在这个版本中，我们包括了一个玩具，盒状狗的四足形态作为基础身体，但对于新的身体形态来说，修改这个场景是很简单的。

要在这个版本所包含的环境中训练高性能策略并对其进行交互式评估，请参见我们的训练Colab[3]。

Brax包括几种常见的强化学习算法，这些算法的实现是为了利用JAX的并行性和即时编译能力。这些算法是：

• 近似政策优化（PPO） [6]
• Soft Actor Critic（SAC）[4]
• 进化策略（ES） [32］
• 分析性政策梯度（APG）

每种算法在某些方面都是独特的。PPO是一种政策性RL算法，SAC是非政策性的，ES是一种黑箱优化算法，而APG利用了环境提供的奖励的可差异性。这种算法覆盖的广度显示了Brax的灵活性，以及它加速研究和降低成本的潜力。在这项工作中，我们把实验分析的重点放在PPO和SAC上（见第6节），而把ES和APG的分析推迟到未来的工作中。

为了获得基于JAX的分批环境的所有好处，可以在加速器上运行，我们建立了一个定制的PPO实现。特别是环境数据（rollouts）是在加速器上生成的，随后由SGD优化器处理。这些数据不需要离开加速器，也不需要在不同进程之间进行上下文切换。整个训练循环（环境滚动+SGD更新）发生在一个不间断的jitted函数中。

训练的过程如下：

• 该批数据在每个可用的加速器核心之间平均分配，并收集环境扩展。
• 基于这批数据计算归一化统计，在所有核心之间同步统计，然后对观测数据进行归一化处理
• 每个加速器核心将该批数据分成适当数量的小批，计算梯度更新，在所有核心之间同步进行，然后同步应用。

算法的性能/吞吐量在很大程度上取决于超参数（如批处理量、小批处理量、优化历时的数量）。我们注意到，对于最佳的超参数，我们对PPO的实现足够高效，主要的瓶颈来自于环境（例如，75%的时间用于运行蚂蚁的环境），尽管环境本身相当快。

与PPO不同的是，SAC使用一个重放缓冲器来进行批次采样。为了利用Brax的全部潜力，我们实现了一个带有重放缓冲器的自定义SAC，它完全生活在一个加速器上。这使得整个训练过程可以被编译成一个单一的jitted函数，并在没有任何干扰的情况下运行。训练的过程大致如下：

• 每个可用的加速器核心在环境中运行几步，并将这些数据添加到每个核心的重放缓冲区中
• 根据新生成的数据计算出正常化的统计数据，统计数据在所有核心之间同步
• 进行若干SGD更新，每个加速器核心从自己的重放缓冲器中取样，计算梯度更新，并与其他核心同步进行最终更新

SAC比PPO的样本效率高得多，因此我们观察到训练吞吐量现在变成了SGD更新的瓶颈（12%用于运行环境，10%用于重放缓冲区，78%用于SGD更新）。由于SGD更新对多个核心的扩展性较差，使用1个以上的加速器核心带来的好处微乎其微，因此最经济的设置是用一个加速器核心实现的。

为了实现ES，我们遵循与PPO/SAC相同的范式：我们在一个加速器上不间断地运行一切，使所有的处理都包含在加速器中。

训练的过程如下：

• 主导加速器产生政策参数扰动
• 政策参数扰动被平均分配到所有可用的加速器核心上进行评估
• 主导加速器根据评估分数计算梯度并更新策略

该算法花了>99%的运行时间来评估环境步骤。

作为如何利用我们引擎的可区分性的概念证明，我们提供了一个APG实现。训练比以前的算法要简单得多：

• 编译一个函数，通过一个短轨迹获取损失的梯度
• 用这个函数进行梯度下降

编译完梯度更新后，该算法将剩余的大部分时间用于评估梯度函数。这种算法没有前三种算法成熟，目前没有产生运动步态，相反，似乎很容易在我们提供的环境中陷入局部最小值。通过长轨迹进行微分是一个活跃的研究领域[38, 21, 18]，并且已知很难优化[39, 40]，因此我们将更先进的可微分算法推迟到未来的版本。

作为我们发布的一部分，我们为所有的环境包括了高性能的超参数。这些超参数通常在几秒钟到几分钟内用标准加速器解决他们的环境。关于我们的超参数实验的详尽列表，见我们的 repo[41]。关于每个环境的最佳20个超参数设置在SAC和PPO上详尽的超参数扫描的性能图，见附录D。

通过利用JAX的矢量化和设备并行化基元，我们可以通过在加速器内和跨加速器之间分配环境计算，轻松地将Brax扩展到每秒数亿步的性能。图2描述了本版本中的环境套件在一个特定的快速、现代加速器集群（TPUv3的4x2拓扑结构）上的扩展曲线，以及Ant环境在各种加速器和TPU拓扑结构下的性能扩展。作为参考，Colab TPU实例目前提供2x2 TPUv2加速器的有限免费使用。

在引擎之间进行完全的苹果对苹果的比较是困难的，主要是因为最广泛使用的引擎并行化的主要方式是通过在CPU上定制多线程线束，或者通过分布式聚集无头工人与附加加速器–通常是定制的设置，大多数从业者无法使用。因此，将Brax的Ant环境编译到TPUv3 8x8加速器上运行（-每秒数亿步）与从业者在单线程机器上运行OpenAI健身房MuJoCo-ant的典型用例（-每秒数千步）相比，可能并不公平。虽然我们包括了在大型TPU集群上部署Brax的结果，但我们强调，Brax在单个1x1 TPUv2上的性能明显优于绝大多数从业者到现在为止能够以大幅降低的成本实现的性能。

为了明确这种性能上的差距，我们首先考虑在各种架构上用Brax的PPO实现对Ant环境的训练速度进行定性比较。我们将其与传统的设置进行比较，使用PPO[28]的标准实现–即没有编译也没有针对并行性进行优化，在图3中可以看出。请注意，Brax在10秒左右就达到了高性能的运动，而标准的PPO实现则需要将近半个小时。

在这一节中，我们详细介绍了我们的引擎的几个重要限制和缺陷。

众所周知，依靠弹簧约束而不是更复杂的费瑟斯顿式方法的物理引擎可能很脆，而且可能需要仔细调整阻尼力。实际上，这些不稳定性是以积分器的小收敛半径出现的，需要小的积分步长。更糟糕的是，这些不稳定性随着问题中存在的质量尺度的不同而增长。虽然对弹簧约束的依赖大大简化了我们引擎的核心基元，但它确实意味着在一个新的物理场景中确保稳定性可能需要对阻尼力、质量和惯性比例以及积分步长进行相当数量的调整。

此外，由于我们的系统基本上是大型耦合弹簧-质量配置，我们的模拟痕迹比假设的相应费瑟斯通模拟有更多的 “抖动”。这可以通过增加关节弹簧约束的强度来缓解，但这是以减少最大稳定积分步长为代价的。对于这个版本的环境，我们选择了这些弹簧常数，以便最大限度地提高仿真速度，同时仍然保留定性的平滑仿真，我们将在未来的工作中研究Featherstone方法。

受Tiny Differentiable Simulator[16]的启发，我们对所有的碰撞基元使用速度级的碰撞更新和Baumgarte稳定。我们确实试验过完全弹性的、冲动的碰撞，但发现运动质量和稳定性受到影响。由于这种选择，我们继承了这些方法的已知调整要求和内在的非物理性[43]。我们试验了基于撞击时间的碰撞检测，但是，与DiffTaichi[17]的作者类似，我们发现它对代码库增加的复杂度没有提供什么精度优势。

此外，我们目前只对场景中的任何碰撞器使用二次缩放的、天真的碰撞检测。典型的基于物理的顺序决策问题不涉及足够多的碰撞器，这不会成为一个显著的瓶颈，因为我们仍然可以很容易地对场景中的所有碰撞基元进行并行化，而不会对现代加速器的内存缓冲区造成压力，但我们设想随着任务的复杂性增加，这将变得更加紧张。我们把更高级的碰撞物理学，例如基于LCP的求解器，以及更有效的碰撞修剪留给未来的版本。

虽然我们吹嘘我们能够将pythonic物理环境和学习算法与XLA并排编译，作为我们的库从JAX继承的一个强大的比较优势，但这并不是没有任何开发摩擦。对于Brax的终端用户来说，最突出的是JIT编译时间有时会接近或超过复杂环境的训练时间（即，编译可能需要几分钟）。我们对Brax的核心设计模式进行了广泛的迭代，以改善这种情况，在某些情况下，直接与JAX开发团队合作，在TPU上调整XLA编译启发式方法，以提高编译速度和性能。最终，编译时间仍然是一个小瓶颈，特别是对于利用可分性的学习算法。

这项工作介绍了我们的PPO和SAC实现的结果。虽然我们在这个版本中包括了APG和ES，但它们还没有被彻底地测试过，我们也没有用它们进行过那么多的超参数探索。我们把它留给未来的工作，以充分利用我们引擎的可分性。

从定义上来说，制作从业人员常用的另一个版本使现有的基准更加复杂，但我们希望我们的库所释放的开发速度能够弥补这一额外的摩擦。同时，发布一个可以快速解决控制问题的引擎的民主化效应可能是双刃剑：一项民主化的技术和一件武器之间的区别完全取决于谁在挥舞它。掌握对机器人的控制权代表了一个改变社会的机会，因此我们希望我们的引擎只有助于改善和加速我们未来的公平自动化。

然而，仍然有一种可能性，即通过发布一个明显更快的引擎，我们无意中大大增加了用于强化学习问题的计算量，就像在城市中建造一条新的高速公路会反直觉地增加交通一样[44]。至少对于我们自己的能源支出来说，我们所做的实验是在那些有望在2030年实现完全可再生资源的数据中心中进行的[45]。

作者感谢Erwin Coumans在物理引擎的微妙实现细节方面提供的宝贵建议，感谢Blake Hechtman和James Bradbury回答了作者的许多问题，并提供了JAX和XLA的优化帮助，感谢Luke Metz和Shane Gu在本项目开发过程中提供的激励性反馈和有益的讨论，感谢Yuval Tassa对本稿的早期草案提供的特殊反馈。作者还感谢Vijay Sundaram、Wright Bagwell、Matthew Leffler、Gavin Dodd、Brad Mckee和Logan Olson帮助孵化了这个项目。

在这一节中，我们演示了如何使用ProtoBuf规范构建一个Brax场景，以及用python构建同一场景的一个简短片段。

在这个附录中，我们提供了一个图，描述了一个用于把握的执行策略。

在本附录中，我们详细介绍了所有数字中所有训练曲线所使用的硬件和超参数。

图3：硬件：128核英特尔至强处理器，频率为2.2Ghz，用于运行MuJoCo环境。32核英特尔至强处理器，2.0 GhZ，用于运行32x-CPU曲线。标准PPO使用[28]，braxppo使用我们的 repo。

图4：硬件：32核英特尔至强处理器，频率为2.2Ghz，用于运行环境。2x2 TPUv2用于学习算法，使用[28]

在这个附录中，我们绘制了在我们详尽的超参数扫描中发现的表现最好的20条训练曲线。超参数的精确值可以在这里找到压缩的、分类的json文件。

所有的图都是在1x1的TPUv2上生成的，即Colab的免费TPU层上的硬件。

在这个附录中，我们指出了我们对halfcheetah、ant和humanoid的实现与最初的MuJoCo-*-v2设想之间的主要差异。随着时间的推移，我们将努力使我们的实现与MuJoCo的实现更加接近，但我们将详尽的分析（例如，尝试在Brax和MuJoCo之间转移策略，sim2real风格）推迟到未来的工作。

MuJoCo使用与世界的关节来实现2d平面运动。相比之下，Brax直接屏蔽了旋转和平移运动的集成更新，因此半身人只能在一个平面内移动，并且只能围绕一个轴旋转。质量、惯性和致动器的比例被选择为尽可能地接近MuJoCo的半猎豹。我们推测，图4中剩余的性能差异归结于我们如何实现接触物理学的引擎特定差异，以及我们特定的执行模型。这种环境也被认为是有点病态的[46]，在mujoco halfcheetah中发现的最高性能的策略是物理学上的破坏，因此在两个实现之间比较极高性能的策略（例如，>10,000分）在理论上是有问题的，因为它相当于比较这两个引擎的破坏方式。

因为我们在ACME实现中使用的标准SAC超参数已经针对现有的MuJoCo环境进行了积极的优化，所以也有可能我们的Brax搜索只是在超参数空间中的一个糟糕的超容积，尽管我们确实相当积极地搜索了高性能的超参数。不管两个实现之间的奖励差异的绝对大小，braxppo和braxsac都能找到性能相当好的火车头步态。我们将在未来的版本中继续减少性能差异。

Brax的蚂蚁有经过调整的质量、惯性和执行器的强度。Brax的奖励函数忽略了接触成本（见[47]）。除此以外，这种环境在两个引擎之间实现了质量上最相似的步态。

除了对质量、惯性和执行器强度的常规调整外，类人猿的奖励函数也有轻微的修改。在我们的环境中，扭动关节的正则化惩罚较低（0.01，而MuJoCo的0.1），而且在Brax中，躯干触发完成的重置条件从0.6到2.1，而MuJoCos 1.0和2.0。

此外，我们实现三自由度执行器的方式与MuJoCo略有不同。更多细节，请看我们的关节实现[48]。

Brax是在Apache许可证2.0下发布的，据我们所知，并没有侵犯任何版权。它将被托管在github上，网址是https://github.com/google/brax。