Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation
This work provides a tool for researchers and practitioners in reinforcement learning and robotics to accelerate simulation and policy training, though it is incremental as it builds on existing simulation and RL methods.
The authors introduced Brax, a differentiable physics engine for large-scale rigid body simulation built in JAX, focusing on performance and parallelism on accelerators, and demonstrated its application in training reinforcement learning policies on tasks inspired by existing literature, with results showing training in minutes on common benchmarks.
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.