警告 自 0.13.0 版本起，只有 brax/training 模块仍在积极维护。用户应改用 MuJoCo Playground，而不是 brax/envs；所有这些环境都可以很好地与 brax/training 配合训练。如果您希望使用 Brax 进行物理仿真，请使用位于 github.com/google-deepmind/mujoco 的 MJX（通过 pip install mujoco_mjx 安装）或 MuJoCo Warp，而非将 Brax 作为 MuJoCo 物理引擎的封装层。未来我们可能会将 brax 单纯重新定位为一个强化学习库。

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 中进行训练。
• 使用 PyTorch 在 GPU 上进行 Brax 训练 演示了如何将 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——用于大规模刚体仿真的可微分物理引擎》（arXiv:2106.13281），该论文发表于 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.14.2},
  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，以及场景编排功能。
• 我们出色的 开源协作者与贡献者。感谢你们！

Brax 快速上手指南

⚠️ 重要提示 (v0.13.0+) 目前仅 brax/training（强化学习训练部分）处于活跃维护状态。

• 环境仿真：不再推荐使用 brax/envs。请使用 MuJoCo Playground，它可与 brax/training 完美配合。
• 物理引擎：如需纯物理仿真，请直接使用 MJX (pip install mujoco_mjx) 或 MuJoCo Warp。
• Brax 未来可能专注于作为纯强化学习库。

Brax 是一个基于 JAX 构建的高速、可微分物理引擎，专为机器人学、强化学习及大规模并行仿真设计。它能在 TPU/GPU 上实现每秒百万级的物理步长仿真，并内置了 PPO、SAC 等主流强化学习算法。

环境准备

• 操作系统：Linux (推荐), macOS, Windows (WSL2)
• 硬件要求：
  • CPU：支持 AVX2 指令集。
  • GPU：NVIDIA GPU (需安装 CUDA 和 cuDNN) 可获得最佳加速效果；也支持 TPU。
• 软件依赖：
  • Python 3.9+
  • pip 或 conda/mamba 包管理器
  • GPU 用户前置步骤：在安装 Brax 前，请确保已正确安装对应 CUDA 版本的 JAX。参考 JAX GPU 安装指南。

安装步骤

方法一：使用 pip 安装（推荐）

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

国内加速建议：如果下载速度慢，可使用清华或阿里镜像源： pip install brax -i https://pypi.tuna.tsinghua.edu.cn/simple

方法二：使用 Conda/Mamba 安装

conda install -c conda-forge brax
# 如果使用 mamba:
# mamba install -c conda-forge brax

方法三：从源码安装（开发者模式）

git clone https://github.com/google/brax.git
cd brax
python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install -e .

基本使用

Brax 的核心功能主要分为物理仿真和强化学习训练。由于新版架构调整，推荐结合 MuJoCo Playground 进行环境交互，或使用 Brax 内置的训练管道。

1. 运行预置训练示例

安装完成后，Brax 提供了一个简单的命令行工具来启动训练。以下命令将使用默认配置训练一个智能体（通常在几秒到几分钟内完成）：

learn

注：若在 GPU 环境下运行，请确保 JAX 已正确识别 GPU。

2. Python API 快速示例

以下是最基础的代码结构，展示如何导入库并检查可用环境（具体环境加载建议参考 MuJoCo Playground 文档）：

import brax
from brax import envs

# 注意：brax/envs 在新版本中已不再作为主要维护对象
# 以下为概念性演示，实际训练建议结合 mujoco_playground 使用
print("Available environments in legacy brax:", envs.registered_environment_names())

# 引入训练组件
from brax.training import ppo

# 典型的训练流程通常涉及创建环境、初始化网络、然后运行训练循环
# 详细用法强烈建议参考官方 Colab 教程

3. 进阶学习：Colab 交互式教程

对于初学者，最直观的上手方式是通过 Google Colab 运行官方提供的 Notebook，无需本地配置即可体验全部功能：

• 基础入门：Brax Basics - 了解 API 及基础物理原语仿真。
• 强化学习训练：Brax Training - 学习如何使用 PPO/SAC 等算法训练策略，以及模型的保存与加载。
• 结合 MJX 仿真：Brax Training with MJX - 使用最新的 MJX 物理后端进行训练。
• PyTorch 集成：Brax Training with PyTorch - 在 PyTorch 框架下利用 Brax 进行快速训练。