一个使用JAX编写的可微分的物理引擎。
项目描述
Brax是一个快速且完全可微分的物理引擎，用于机器人、人类感知、材料科学、强化学习和其他仿真密集型应用的研究和开发。
Brax是用JAX编写的，旨在用于加速硬件。它既适用于单设备仿真，又可扩展到多设备上的大规模并行仿真，无需昂贵的数据中心。
Brax在TPU上每秒模拟数百万个物理步骤，并包含一系列学习算法，这些算法可以在几秒到几分钟内训练智能体。
一个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如何在其他ML框架中用于快速训练，在此例中是PyTorch。
本地使用Brax
要从pypi安装Brax，请使用以下命令安装：
python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install brax
conda install -c conda-forge brax # s/conda/mamba for mamba
或者，要从源代码安装Brax，请克隆此存储库，然后cd到它，然后
python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install -e .
要训练一个模型
learn
NVIDIA GPU支持在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，以及场景合成器。
- 我们出色的开源协作者和贡献者。感谢你们！
下载文件
下载适合您平台的文件。如果您不确定选择哪个，请了解更多关于安装包的信息。
源分发
brax-0.10.5.tar.gz (767.5 kB 查看散列)
构建分发
brax-0.10.5-py3-none-any.whl (998.9 kB 查看散列)
关闭
brax-0.10.5.tar.gz的散列
| 算法 | 散列摘要 | |
|---|---|---|
| SHA256 | e7563130c2b08bf0c9453d87602126732f20afc4624cb8574b3577fa62fdbcec |
|
| MD5 | 3e76ede8e34c93a341e11a46b7c58130 |
|
| BLAKE2b-256 | 20f483784481156932925a5b1d6b33b44aaeba46aeca5789f3fa7bf64c62f87a |
关闭
brax-0.10.5-py3-none-any.whl的散列
| 算法 | 散列摘要 | |
|---|---|---|
| SHA256 | 304fe6e5e266e42a18f197f2b7b6a9bb03a87bd97928e385c51c874b56f95866 |
|
| MD5 | e121fb3de65d8c50921ed344901eaf2b |
|
| BLAKE2b-256 | 39c5c61ff2ce0688e3349e2fbeff763f8a3ba396f714a7115b4c5163b6cb2118 |
