[论文解读] Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation
简要结论：Brax 是一个开源、可微分的刚体物理引擎，采用 JAX 构建，在加速器上运行，支持大规模并行环境仿真，并可在设备端学习，配备打包的强化学习算法。
We present Brax, an open source library for rigid body simulation with a focus on performance and parallelism on accelerators, written in JAX. We present results on a suite of tasks inspired by the existing reinforcement learning literature, but remade in our engine. Additionally, we provide reimplementations of PPO, SAC, ES, and direct policy optimization in JAX that compile alongside our environments, allowing the learning algorithm and the environment processing to occur on the same device, and to scale seamlessly on accelerators. Finally, we include notebooks that facilitate training of performant policies on common OpenAI Gym MuJoCo-like tasks in minutes.
研究动机与目标
- 通过将物理与学习在加速器上同址，推动更快、更易获取的 RL 研究。
- 提供一个可微分、开源的引擎，以实现基于梯度的 RL 方法。
- 提供基于 ProtoBuf 的环境规范和 Gym 风格接口，适用于可扩展任务。
- 展示 Brax 在加速器上的可扩展性和性能，并与 MuJoCo 进行比较。
- 对可微分物理引擎的局限性和未来工作进行综述，以指导改进。
提出的方法
- 使用带有 QP 数据原语和向量化变换的极坐标自由度刚体动力学实现。
- 使用 JAX 实现自动向量化、设备级并行、JIT 和自动微分，以在加速器上运行成千上万的环境。
- 将物理步作为并行变换（关节、执行器、碰撞体）应用于 Brax_system，并使用二阶辛欧拉积分器。
- 提供基于 ProtoBuf 的系统规范和 gym 风格 Env，用于构建和求解物理任务。
- 将 RL 算法（PPO、SAC、ES、APG）打包在 JAX 中，以在加速器上完全运行，进行设备端回滚和梯度更新。
- 展示环境套件，包括 MuJoCo 风格的 Ant、Humanoid、Halfcheetah、Grasp 和 Fetch。
实验结果
研究问题
- RQ1一个可在加速器上运行的可微分物理引擎，能否通过将环境仿真与学习同址显著加速 RL 训练？
- RQ2在 TPU/GPU 上，Brax 的性能和可扩展性如何与传统的 CPU 引擎如 MuJoCo 在常见 RL 基准上进行比较？
- RQ3在 Brax 实现时，微分性对学习算法（PPO、SAC、ES、APG）的影响是什么？
- RQ4大规模、可微分的刚体仿真所需的实际限制与调优有哪些？
- RQ5ProtoBuf 环境规范与 gym 风格接口用于构建多样的物理任务的有效性如何？
主要发现
- Brax 在单个加速器上对某些环境实现了每秒数百万次仿真步。
- 在 4x2 TPUv3 上，Brax 通过在加速器之间分布计算，扩展到每秒数亿次环境步。
- 为 Brax 编译的 PPO 和 SAC 实现实现了在设备端回滚和梯度更新，数据传输开销极小，训练速度显著快于非加速基线。
- 在使用 SAC 评估时，Brax 环境在学习轨迹和奖励上与 MuJoCo 相当，超参数可比。
- 与 MuJoCo A nt 相比，使用优化的 PPO 的 Brax 约 10 秒即可实现 locomotion，而标准 PPO 实现约 30 分钟。
- Brax 显示出良好的动量守恒以及与其他引擎相对竞争的能量/角动量行为。
更好的研究，从现在开始
从阅读论文到最终审阅，大幅缩短您的研究时间。
无需绑定信用卡
本解读由 AI 生成，并经人工编辑审核。