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.
Description
Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation
%0 Generic
%1 freeman2021differentiable
%A Freeman, C. Daniel
%A Frey, Erik
%A Raichuk, Anton
%A Girgin, Sertan
%A Mordatch, Igor
%A Bachem, Olivier
%D 2021
%K engine physics
%T Brax -- A Differentiable Physics Engine for Large Scale Rigid Body
Simulation
%U http://arxiv.org/abs/2106.13281
%X 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.
@misc{freeman2021differentiable,
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.},
added-at = {2023-08-13T11:37:25.000+0200},
author = {Freeman, C. Daniel and Frey, Erik and Raichuk, Anton and Girgin, Sertan and Mordatch, Igor and Bachem, Olivier},
biburl = {https://www.bibsonomy.org/bibtex/20a4797feeda4475705829b26d5c90e9a/parismic},
description = {Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation},
interhash = {2f0adb61d31c74f49dddf480312121b0},
intrahash = {0a4797feeda4475705829b26d5c90e9a},
keywords = {engine physics},
note = {cite arxiv:2106.13281Comment: 9 pages + 12 pages of appendices and references. In submission at NeurIPS 2021 Datasets and Benchmarks Track},
timestamp = {2023-08-13T11:37:25.000+0200},
title = {Brax -- A Differentiable Physics Engine for Large Scale Rigid Body
Simulation},
url = {http://arxiv.org/abs/2106.13281},
year = 2021
}