# Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation

arXiv: https://arxiv.org/abs/2106.13281 · HTML: https://mlsanity.com/paper/2106.13281/

## Key facts

- Cited by 407 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 57 references that are in the map.
- First posted on arXiv 24 Jun 2021.

## Builds on (most cited first, 23 of 23)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Soft Actor-Critic Algorithms and Applications](https://mlsanity.com/paper/1812.05905/) (2018)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Solving Rubik's Cube with a Robot Hand](https://mlsanity.com/paper/1910.07113/) (2019)
- [Evolution Strategies as a Scalable Alternative to Reinforcement Learning](https://mlsanity.com/paper/1703.03864/) (2017)
- [Human-level performance in first-person multiplayer games with population-based deep reinforcement learning](https://mlsanity.com/paper/1807.01281/) (2018)
- [DiffTaichi: Differentiable Programming for Physical Simulation](https://mlsanity.com/paper/1910.00935/) (2019)
- [Unity: A General Platform for Intelligent Agents](https://mlsanity.com/paper/1809.02627/) (2018)
- [Acme: A Research Framework for Distributed Reinforcement Learning](https://mlsanity.com/paper/2006.00979/) (2020)
- [A Differentiable Physics Engine for Deep Learning in Robotics](https://mlsanity.com/paper/1611.01652/) (2016)
- [Understanding and correcting pathologies in the training of learned optimizers](https://mlsanity.com/paper/1810.10180/) (2018)
- [NeuralSim: Augmenting Differentiable Simulators with Neural Networks](https://mlsanity.com/paper/2011.04217/) (2020)
- [On the Importance of Hyperparameter Optimization for Model-based Reinforcement Learning](https://mlsanity.com/paper/2102.13651/) (2021)
- [Fast and Feature-Complete Differentiable Physics for Articulated Rigid Bodies with Contact](https://mlsanity.com/paper/2103.16021/) (2021)
- [Unbiasing Truncated Backpropagation Through Time](https://mlsanity.com/paper/1705.08209/) (2017)
- [Deluca -- A Differentiable Control Library: Environments, Methods, and Benchmarking](https://mlsanity.com/paper/2102.09968/) (2021)
- [Lyceum: An efficient and scalable ecosystem for robot learning](https://mlsanity.com/paper/2001.07343/) (2020)

## Built on by (most cited first, 30 of 407)

- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [AlpacaFarm: A Simulation Framework for Methods that Learn from Human Feedback](https://mlsanity.com/paper/2305.14387/) (2023)
- [Learning to Walk in Minutes Using Massively Parallel Deep Reinforcement Learning](https://mlsanity.com/paper/2109.11978/) (2021)
- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [Rewarded soups: towards Pareto-optimal alignment by interpolating weights fine-tuned on diverse rewards](https://mlsanity.com/paper/2306.04488/) (2023)
- [ManiSkill3: GPU Parallelized Robotics Simulation and Rendering for Generalizable Embodied AI](https://mlsanity.com/paper/2410.00425/) (2024)
- [Waymax: An Accelerated, Data-Driven Simulator for Large-Scale Autonomous Driving Research](https://mlsanity.com/paper/2310.08710/) (2023)
- [Accelerated Policy Learning with Parallel Differentiable Simulation](https://mlsanity.com/paper/2204.07137/) (2022)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Flow Matching Policy Gradients](https://mlsanity.com/paper/2507.21053/) (2025)
- [Do Differentiable Simulators Give Better Policy Gradients?](https://mlsanity.com/paper/2202.00817/) (2022)
- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [Gradients are Not All You Need](https://mlsanity.com/paper/2111.05803/) (2021)
- [AirfRANS: High Fidelity Computational Fluid Dynamics Dataset for Approximating Reynolds-Averaged Navier-Stokes Solutions](https://mlsanity.com/paper/2212.07564/) (2022)
- [Automated Reinforcement Learning (AutoRL): A Survey and Open Problems](https://mlsanity.com/paper/2201.03916/) (2022)
- [VMAS: A Vectorized Multi-Agent Simulator for Collective Robot Learning](https://mlsanity.com/paper/2207.03530/) (2022)
- [Robust Autonomy Emerges from Self-Play](https://mlsanity.com/paper/2502.03349/) (2025)
- [JaxMARL: Multi-Agent RL Environments and Algorithms in JAX](https://mlsanity.com/paper/2311.10090/) (2023)
- [EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine](https://mlsanity.com/paper/2206.10558/) (2022)
- [TorchRL: A data-driven decision-making library for PyTorch](https://mlsanity.com/paper/2306.00577/) (2023)
- [Hyperparameters in Reinforcement Learning and How To Tune Them](https://mlsanity.com/paper/2306.01324/) (2023)
- [Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2306.09884/) (2023)
- [VeLO: Training Versatile Learned Optimizers by Scaling Up](https://mlsanity.com/paper/2211.09760/) (2022)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [RoboVerse: Towards a Unified Platform, Dataset and Benchmark for Scalable and Generalizable Robot Learning](https://mlsanity.com/paper/2504.18904/) (2025)
- [Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning](https://mlsanity.com/paper/2303.17503/) (2023)
- [safe-control-gym: a Unified Benchmark Suite for Safe Learning-based Control and Reinforcement Learning in Robotics](https://mlsanity.com/paper/2109.06325/) (2021)
- [NeuPhysics: Editable Neural Geometry and Physics from Monocular Videos](https://mlsanity.com/paper/2210.12352/) (2022)
- [Contextualize Me -- The Case for Context in Reinforcement Learning](https://mlsanity.com/paper/2202.04500/) (2022)
- [Gymnasium: A Standard Interface for Reinforcement Learning Environments](https://mlsanity.com/paper/2407.17032/) (2024)

## Abstract (arXiv)

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.