// the find
google/brax
Massively parallel rigidbody physics simulation on accelerator hardware.
Brax is a JAX-based rigid body physics simulator built for RL and robotics research, offering differentiable simulation that runs at millions of steps per second on TPU/GPU without needing a datacenter. It's aimed at researchers doing sim-to-real robotics, policy gradient methods, or anyone who wants to train PPO/SAC agents in minutes instead of hours.
The four pipelines (MJX, Generalized, Positional, Spring) share one API and can run side by side, which is genuinely useful for sim-to-real transfer experiments. Differentiability through the whole simulator enables analytic policy gradients (APG), something you can't do with most physics engines. Baseline RL algorithms (PPO, SAC, ARS, ES) are included and work out of the box rather than requiring glue code to a separate training library. Scaling from one device to many is handled by JAX's pmap/vmap rather than custom distributed infrastructure.
The README opens with a warning that most of the repo is no longer actively maintained — brax/envs is deprecated in favor of MuJoCo Playground, and brax as a physics engine is being pointed toward MJX instead. Anyone adopting this today for simulation (not just the training code) is building on a component the maintainers are walking away from. Documentation beyond the colab notebooks is thin; there's no real API reference, so understanding pipeline internals means reading source. The four-pipeline abstraction implies interchangeability, but Generalized/Positional/Spring differ enough in contact handling and accuracy that swapping pipelines can change training outcomes, and the README doesn't flag this. GPU setup requires manually matching CUDA/CuDNN versions to JAX, which is a common source of broken installs.