Brax

A differentiable rigid-body physics engine from Google written in JAX, bundled with a reinforcement-learning training library.

Brax was open-sourced by a Google team in 2021 (lead paper author C. Daniel Freeman), written in JAX (a numerical computing library that supports automatic differentiation and compiles to run on GPU/TPU). Its selling point is that both the physics simulation and the learning algorithm compile to run on the same accelerator, eliminating the back-and-forth data transfer between CPU and GPU, and it can train a usable policy on MuJoCo-like Gym tasks in a few minutes; because the simulation is differentiable, it can also compute gradients through the simulation directly to optimize a policy. Brax later offered four physics backends: MJX (MuJoCo's JAX reimplementation), generalized coordinates, position-based dynamics, and a spring model. According to its repository, as of version 0.13.0 only brax/training (its PPO, SAC, and similar training code) is still actively maintained; for physics simulation it now recommends switching to MJX or MuJoCo Warp, and MuJoCo Playground for environments.

ExampleMuJoCo Playground's repository notes that to reproduce results exactly as reported in its paper, users can run Brax's training scripts directly, while the environments themselves are implemented with MJX or MuJoCo Warp.

Related

MuJoCo XLA · MuJoCo Playground · JAX · Differentiable Simulation · GPU-Accelerated Parallel Simulation · Massively Parallel Reinforcement Learning

Sources

Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation (arXiv 2106.13281)
google/brax (GitHub)
google-deepmind/mujoco_playground (GitHub)

As of

2026-09