brax - 基于JAX的高性能物理引擎适用于机器人和强化学习仿真

Brax是一个快速且完全可微分的物理引擎，用于机器人、人类感知、材料科学、强化学习和其他需要大量模拟的应用研究与开发。

Brax使用JAX编写，专为加速硬件设计。它不仅能高效进行单设备模拟，还可扩展到多设备上进行大规模并行模拟，无需烦人的数据中心。

Brax在TPU上每秒可模拟数百万步物理过程，并包含一套学习算法，能在几秒到几分钟内训练智能体：

基准学习算法，如PPO、SAC、ARS和进化策略。
利用模拟器可微分性的学习算法，如解析策略梯度。

一个API，四种管线

Brax提供四种易于切换的物理管线：

MuJoCo XLA - MJX - MuJoCo物理引擎的JAX重新实现。
广义使用类似于MuJoCo和TDS的动力学算法，在广义坐标中计算运动。
位置使用基于位置的动力学，一种快速且稳定的解决关节和碰撞约束的方法。
弹簧提供快速且低成本的模拟，适用于快速实验，使用视频游戏中常见的简单基于脉冲的方法。

这些管线共享相同的API，可以在同一模拟中并行运行。这使Brax非常适合进行迁移学习实验和缩小模拟与现实世界之间的差距。

快速入门：云端Colab

通过一系列Colab笔记本轻松快速地探索Brax：

Brax基础介绍了Brax API，并展示了如何模拟基本物理原语。
Brax训练介绍了Brax的训练算法，让你直接在Colab中训练自己的策略。它还演示了如何加载和保存策略。
使用MuJoCo XLA - MJX的Brax训练演示了使用MJX物理模拟器进行Brax训练。
使用GPU上的PyTorch进行Brax训练演示了如何在其他ML框架中使用Brax进行快速训练，这里以PyTorch为例。

本地使用Brax

要从pypi安装Brax，请使用以下命令：

python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install brax

你也可以从Conda或Mamba安装：

conda install -c conda-forge brax  # 使用mamba时将conda替换为mamba

或者，要从源代码安装Brax，克隆此仓库，进入目录，然后：

python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install -e .

要训练模型：

learn

支持在NVIDIA GPU上训练，但你必须首先安装CUDA、CuDNN和支持GPU的JAX。

了解更多

要深入了解Brax的设计和性能特征，请参阅我们的论文Brax -- 用于大规模刚体模拟的可微分物理引擎，该论文发表在NeurIPS 2021的数据集和基准测试专题上。

引用Brax

如果你想在出版物中引用Brax，请使用：

@software{brax2021github,
  author = {C. Daniel Freeman and Erik Frey and Anton Raichuk and Sertan Girgin and Igor Mordatch and Olivier Bachem},
  title = {Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation},
  url = {http://github.com/google/brax},
  version = {0.10.5},
  year = {2021},
}

致谢

自最初发布以来，Brax已经取得了长足的进步。我们向以下人员表示感谢和由衷的赞美：

Manu Orsini和Nikola Momchev对Brax的训练算法进行了重大重构，使其更易于访问和重用。
Erwin Coumans慷慨地提供建议和指导，以及来自Tiny Differentiable Simulator的许多有用参考。
Baruch Tabanpour，一位帮助启动brax v2并彻底改造接触库的同事。
Shixiang Shane Gu和Hiroki Furuta，为Brax贡献了BIG-Gym、Braxlines和场景编辑器。
我们出色的开源合作者和贡献者。非常感谢你们！