Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation
arXiv:2106.13281
Abstract
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.
9 pages + 12 pages of appendices and references. In submission at NeurIPS 2021 Datasets and Benchmarks Track
References in corpus (16)
- Soft Actor-Critic Algorithms and Applications
- High-Dimensional Continuous Control Using Generalized Advantage Estimation
- Dota 2 with Large Scale Deep Reinforcement Learning
- Evolution Strategies as a Scalable Alternative to Reinforcement Learning
- Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems
- Human-level performance in first-person multiplayer games with population-based deep reinforcement learning
- Solving Rubik's Cube with a Robot Hand
- Unity: A General Platform for Intelligent Agents
- DeepMind Control Suite
- D4RL: Datasets for Deep Data-Driven Reinforcement Learning
- DiffTaichi: Differentiable Programming for Physical Simulation
- Acme: A Research Framework for Distributed Reinforcement Learning
- Understanding and correcting pathologies in the training of learned optimizers
- SEED RL: Scalable and Efficient Deep-RL with Accelerated Central Inference
- Fast and Feature-Complete Differentiable Physics for Articulated Rigid Bodies with Contact
- Deluca -- A Differentiable Control Library: Environments, Methods, and Benchmarking
Cited by in corpus (7)
- Automated Reinforcement Learning (AutoRL): A Survey and Open Problems
- Transferring Dexterous Manipulation from GPU Simulation to a Remote Real-World TriFinger
- GRiD: GPU-Accelerated Rigid Body Dynamics with Analytical Gradients
- Neural Networks with Physics-Informed Architectures and Constraints for Dynamical Systems Modeling
- WarpDrive: Extremely Fast End-to-End Deep Multi-Agent Reinforcement Learning on a GPU
- SaLinA: Sequential Learning of Agents
- Exploring the efficacy of neural networks for trajectory compression and the inverse problem