Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation

文章概括
摘要
1 捐款摘要
2 动机
3 使用Brax：核心物理学循环
4 使用Brax：创建和评估环境
- 4.1 系统规格
- 4.2 类似健身房的环境
- - 4.2.1 MuJoCo Gym-Likes
  - 4.2.2 抓取
  - 4.2.3 获取
5 使用Brax：解决运动和操作问题
- 5.1 与Brax捆绑的学习算法
- - 5.1.1 近似的政策优化（PPO）
  - 5.1.2 Soft Actor Critic (SAC)
  - 5.1.3 进化策略(ES)
  - 5.1.4 分析性政策梯度（APG）
- 5.2 培训表现
6 性能标杆
- 6.1 在加速器上实现并行化
- 6.2 引擎比较
7 局限性和未来工作
- 7.1 弹簧接头
- 7.2 碰撞
- 7.3 Jitting, JAX 和 XLA
- 7.4 算法
- 7.5 社会影响
鸣谢和资金披露
A Appendix - Brax System Specification
B Appendix - Grasp Trajectory
C Appendix - 数字的超参数
D Appendix - 超参数扫频
E Appendix - 与Mujoco的主要区别
- E.1 Halfcheetah
- E.2 Ant
- E.3 Humanoid
F Appendix - License

文章概括

作者：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
代码、数据和视频：
系列文章目录：
上一篇：
论文笔记（十六）：Learning to Walk in Minutes Using Massively Parallel Deep Reinforcement Learning
下一篇：

摘要

我们介绍了Brax，一个用JAX编写的用于刚体仿真的开源库，重点是加速器上的性能和并行性。我们介绍了一套任务的结果，这些任务受到现有强化学习文献的启发，但在我们的引擎中重新制作。此外，我们还提供了PPO、SAC、ES和直接策略优化的重新实现，在JAX中与我们的环境一起编译，允许学习算法和环境处理发生在同一设备上，并在加速器上无缝扩展。最后，我们还包括一些笔记本，这些笔记本有助于在几分钟内对常见的OpenAI健身房的类似MuJoCo的任务进行高性能的策略训练。

图1：Brax初始版本中包含的环境实例套件。从左到右：蚂蚁、取物、抓握、半猎豹和人形。

1 捐款摘要

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

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

2 动机

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

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

仿真引擎的设计在三个方面促成了这个不可访问的问题：

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

第二，大多数仿真引擎是黑盒子：它们不提供采样环境状态的梯度，这使得它们只适合于无模型的RL方法。这种可分性的缺乏迫使研究者使用速度较慢、效率较低的优化方法。

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

我们提交Brax作为同时解决这三个问题的建议。Brax将物理引擎和RL优化器放在同一个GPU/TPU芯片上，将RL训练的速度/成本提高了100-1000倍。它是可微分的，为新的优化技术打开了大门。而且它是一个开源库，被打包在Colabs中运行，因此任何人都可以免费进行RL研究。

3 使用Brax：核心物理学循环

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

算法1：Brax中物理学步骤结构的伪代码。脉冲更新(dp_i)被平行地收集到每种类型的关节、致动器和碰撞器。然后，积分器转换将这些更新应用到qp上。

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

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

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

qpt+δt=Brax_system.step(qpt,actions)qp_{t+δt} = Brax\_system.step(qp_t, actions) qpt+δt=Brax_system.step(qpt,actions)

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

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

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

4 使用Brax：创建和评估环境

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

4.1 系统规格

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

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

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

4.2 类似健身房的环境

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

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

表1: Brax中包含的环境的观察和行动空间数据。

4.2.1 MuJoCo Gym-Likes

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

4.2.2 抓取

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

4.2.3 获取

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

5 使用Brax：解决运动和操作问题

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

5.1 与Brax捆绑的学习算法

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

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

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

5.1.1 近似的政策优化（PPO）

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

训练的过程如下：

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

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

5.1.2 Soft Actor Critic (SAC)

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

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

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

5.1.3 进化策略(ES)

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

训练的过程如下：

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

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

5.1.4 分析性政策梯度（APG）

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

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

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

5.2 培训表现

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

6 性能标杆

6.1 在加速器上实现并行化

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

图2：（左）本版本中每个环境在4x2 TPU v3上的每秒有效环境步数的缩放。(右图)Ant环境中几个加速器每秒有效环境步数的比例。误差条在这个比例下是不可见的。

6.2 引擎比较

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

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

图3: Brax的编译和优化的PPO实现与标准PPO实现的训练曲线的定性比较[28]。注意X轴是对数墙时间，以秒为单位。所有带有 "brax "标签的曲线是Brax的Ant版本，而MuJoCo的曲线是MuJoCo-Ant-v2。两种ppo的实现都对1000万个环境步骤进行了评估。阴影区域表示在5个副本中表现最低和最高的种子，实线表示平均值。所用的超参数见附录C。所用的超参数见附录C。

接下来，为了验证Brax的MuJoCo环境版本在质量上与MuJoCo的环境相似，我们将SAC在我们环境中的标准实现的训练曲线与MuJoCo版本的训练曲线并排描述。从质量上看，对于一组固定的SAC超参数，Brax环境与MuJoCo的对应环境相比，在类似的环境步骤数量中实现了类似的奖励。请注意，这并不是说我们促进了 “更高的奖励”，因为比较不同的奖励函数在理论上是有问题的（尽管Brax的奖励函数与MuJoCo的健身房定义非常接近，更多细节见附录E）。我们只想证明奖励收益的进展是相似的，而且Brax环境在相似的学习步骤中取得了质量上相似的性能。

最后，我们通过它在[42]引入的 "宇航员 "诊断中的表现来考虑我们引擎的仿真质量–这是一个改进版的人形场景，它测量动量和能量不守恒作为仿真保真度的函数，如图5所描述。从质量上看，Brax实现了有竞争力的线性动量守恒比例，这是因为它对位置的最大笛卡尔坐标表示和对称积分方案。能量守恒性能与Havok和MuJoCo的欧拉积分器相一致。相对而言，Brax在角动量守恒方面做得特别好。

图5：Brax以及其他几个引擎的线性动量（左）、角动量（中）和能量（右）的非守恒比例。非Brax的数据是经[42]作者许可后改编的，并在此作了比较。按照Erez等人的做法，在动量守恒场景中，我们禁用了阻尼、碰撞和重力，并以每个致动器每步约0.5牛米的扭矩随机致动四肢1秒。在能量方面，我们另外禁用了致动器，给每个身体部位随机踢了一脚，速度为1米/秒，并在模拟1秒后测量能量漂移。所有的测量都是对128个单精度浮点的随机种子的平均测量。

7 局限性和未来工作

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

7.1 弹簧接头

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

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

7.2 碰撞

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

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

7.3 Jitting, JAX 和 XLA

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

7.4 算法

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

7.5 社会影响

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

然而，仍然有一种可能性，即通过发布一个明显更快的引擎，我们无意中大大增加了用于强化学习问题的计算量，就像在城市中建造一条新的高速公路会反直觉地增加交通一样[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帮助孵化了这个项目。

A Appendix - Brax System Specification

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

B Appendix - Grasp Trajectory

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

图6：在Brax中模拟和训练的高性能抓取策略的前300步的快照。手能够拿起球并把它带到一系列的红色目标。一旦球接近红色目标，红色目标就会在不同的随机位置重新出现。

C Appendix - 数字的超参数

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

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

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

图4：MuJoCo和Brax的训练曲线轨迹的定性比较。(左）MuJoCo-Humanoid-v2和Brax-humanoid的训练曲线，（中）MuJoCo-Ant-v2和Brax-ant，以及（右）MuJoCo-HalfCheetah-v2和Brax-halfcheetah。所有的环境都是用相同的SAC[28]的标准实现来评估的，环境在CPU上评估，在2x2 TPUv2上学习，即不是Brax的加速器优化实现。实线表示平均性能，包络线是随机种子的方差。见附录。使用的超参数见附录C。关于halfcheetah的性能差距的简短讨论见附录E。

D Appendix - 超参数扫频

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

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

图7：本版本中5个环境的奖励曲线，经过1000万步的braxppo训练。左栏表示奖励与步数的关系，右栏是相同的数据，以秒为单位的挂钟时间而不是步数来重复。请注意，在1000万步内，抓取和仿人都没有找到成功的策略。

图8：本版本中5个环境的奖励曲线，经过5亿步的braxppo训练。左栏表示奖励与步数的关系，右栏是相同的数据，以秒为单位的挂钟时间而不是步数来重复。除类人猿外，所有政策在这个时间尺度上都可以用ppo解决。

图9：本版本中的5个环境在500万步的brax-sac训练中的奖励曲线。左栏表示奖励与步数的关系，右栏是同样的数据，以秒为单位的挂钟时间代替步数。请注意，人形是通过SAC解决的，但这里的抓取却不是。

E Appendix - 与Mujoco的主要区别

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

E.1 Halfcheetah

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

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

E.2 Ant

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

E.3 Humanoid

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

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

F Appendix - License

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

论文笔记（十七）：Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation相关推荐

论文笔记（十七）：Brax — A Differentiable Physics Engine for Large Scale Rigid Body Simulation
Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation 文章概括摘要 1 捐款摘要 2 动机 3 使 ...
论元笔记 ACL 2017|Automatically Labeled Data Generation for Large Scale Event Extraction
文章目录 1 简介 1.1 创新 2 背景知识 3 数据生成 4 事件抽取 5 实验 1 简介论文题目:Automatically Labeled Data Generation for Large ...
论文笔记——Scaling Up Your Kernels to 31x31: Revisiting Large Kernel Design in CNNs
论文PDF 作者一作:丁霄汉谷歌学术主页通讯:张祥雨谷歌学术主页摘要现代卷积神经网络模型普遍倾向于使用多个小卷积核来代替大卷积核,受到ViT的启发,论文重新审视了大卷积核在现代卷积神经网络 ...
论文笔记目录（ver2.0）
1 时间序列 1.1 时间序列预测论文名称来源主要内容论文笔记:DCRNN (Diffusion Convolutional Recurrent Neural Network: Data-Dr ...
论文笔记【A Comprehensive Study of Deep Video Action Recognition】
论文链接:A Comprehensive Study of Deep Video Action Recognition 目录 A Comprehensive Study of Deep Video A ...
论文笔记（十四）：PoseRBPF: A Rao–Blackwellized Particle Filter for 6-D Object Pose Tracking
PoseRBPF: A Rao–Blackwellized Particle Filter for 6-D Object Pose Tracking 文章概括摘要 1. 介绍 2. 相关工作 3. ...
【论文笔记】—低照度图像增强—ZeroShot—RUAS网络—2021-CVPR
论文介绍题目:Retinex-Inspired Unrolling With Cooperative Prior Architecture Search for Low-Light Image E ...
【论文笔记】Reinforcement and Imitation Learning for Diverse Visuomotor Skills
目录 Abstract Introduction Related Work Model A. Background: GAIL and PPO 1. 行为克隆(Behavior Cloning) 2. ...
论文笔记（六）：A VR Teleoperation Suite with Manipulation Assist
A VR Teleoperation Suite with Manipulation Assist 文章概括摘要 1. 介绍和相关工作 2. 系统概述 2.1 硬件 2.2 软件 2.3 物体位姿跟 ...

论文笔记（十七）：Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation