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.

翻译：我们介绍Brax,这是一个开放源库,用于僵硬体模拟,重点是加速器的性能和平行性能,以JAX书写。我们介绍了一套受现有强化学习文献启发、但又在引擎中重新制作的任务的结果。此外,我们还在JAX中重新实施PPO、SAC、ES和直接政策优化,与我们的环境一起汇编,使学习算法和环境处理能够在同一设备上进行,并且无缝地在加速器上进行。最后,我们包括一些笔记本,便利在几分钟内对类似于OpenAI Gym MuJoCo(OpenAI Gym MuJoCo-sol)的共同任务进行表演者政策培训。