Brax: A Differentiable Physics Engine for Large Scale Rigid Body Simulation

https://arxiv.org/abs/2106.13281

JAX rigid-body engine; millions of steps per second on one accelerator, hundreds of millions across accelerators; Ant locomotion in about ten seconds versus about half an hour for standard PPO.

Differentiable simulation, Massively parallel GPU simulation