PaperlayerView exact revision on arXiv

Source-linked AI summary

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

C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, Olivier Bachem

arXiv:2106.13281v1cs.ROcs.AI

TL;DR

Reinforcement learning is costly and slow because it requires many simulation steps and often separates CPU-based simulation from accelerator-based learning. Brax combines accelerator-oriented rigid-body simulation and learning in JAX, achieving rapid policy training and scaling to hundreds of millions of simulation steps per second.

  • Problem

    Reinforcement learning remains expensive and slow because of high sample requirements and latency between CPU simulation and accelerator-based learning.

  • Method

    Brax is a JAX-based rigid-body simulator that uses vectorized physical state representations and accelerator-compatible parallel computation for reinforcement learning environments.

  • Results

    Brax trains locomotion and dexterous manipulation policies in seconds to minutes, reaches millions of simulation steps per second on one accelerator, and scales to hundreds of millions across accelerators.

  • Takeaways & Limitations

    On Ant, Brax reaches performant locomotion in about ten seconds, compared with close to half an hour for standard PPO.

  • Takeaways & Limitations

    Large-cluster comparisons are not directly equivalent to typical single-machine use, and Brax’s JIT compilation can sometimes take minutes.

Abstract

from arXiv · show

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.

1 Summary of Contributions

Brax is an open-source, JAX-based differentiable physics engine designed for large-scale rigid-body simulation on accelerators. It uses JAX compilation, parallelism, vectorization, and automatic differentiation to train locomotion and dexterous manipulation policies in seconds to minutes on one modern accelerator while scaling across many environments and connected accelerators.

  • Brax trains locomotion and dexterous manipulation policies in seconds to minutes using one modern accelerator.
  • Brax relies on JAX auto-vectorization, device parallelism, just-in-time compilation, and automatic differentiation for performance.
  • Brax enables simple rigid-body physics simulation across thousands of independent environments and hundreds of connected accelerators.

2 Motivation

Brax is proposed to address RL’s high cost and latency by combining a physics engine and optimizer on one accelerator, while also providing differentiability and open-source accessibility.

  • Motivation: RL remains prohibitively expensive and slow because environments with hundreds of state-space dimensions require millions to billions of simulation steps.This high sample complexity drives substantial exploration costs.
  • Motivation: CPU-based simulation paired with GPU/TPU-based RL creates data-marshalling and network latency that can dominate experiment runtime.The components often run in separate processes, machines, or both.
  • Motivation: Black-box engines provide no environment-state gradients, restricting researchers to model-free RL and slower, less efficient optimization methods.Their lack of differentiability prevents gradient-based approaches to sampled environment states.
  • Motivation: Closed-source or technically incompatible engines limit iteration, debugging, and understanding of relationships between environment state and action spaces.Introspection is described as important for guiding new RL research.
  • Motivation: Brax combines physics simulation and RL optimization on one GPU/TPU, is differentiable, and is open source with Colab packaging.The paper claims 100-1000x improvements in RL training speed/cost and free access to RL research.

3 Using Brax: The core physics loop

Brax represents independently moving entities with dynamically changing position, orientation, velocity, and angular velocity data, while modeling relationships through transformations. Its physics loop parallelly accumulates updates from joints, actuators, and colliders before integrating the state.

  • State representation: Brax tracks each freely moving entity separately in maximal coordinates using position, rotational orientation, velocity, and angular velocity as its dynamic state.Joints, actuators, collisions, and integration steps are built as transformations of this state data.
  • Physics step: Each physics step applies a kinematic integrator, collects joint, actuator, and collider updates in parallel, then applies potential and collision integrators.Actions supply torques or target angles required by actuators.
  • Physical abstractions: Bodies, joints, actuators, and colliders bundle physical metadata and provide transformations that calculate forces, torques, or other differential updates.A revolute joint, for example, encodes a one-degree-of-freedom constraint between parent and child bodies.
  • Physics step: Brax sums differential updates over a short timestep and transforms system state with a second-order symplectic Euler update, parallelizing across components and scenes.The framework notes that extensions to higher-order integrators are straightforward.
  • Extensibility: New control-flow behavior can be added by implementing a Brax_transformation and inserting it into the physics step function.An overarching system class coordinates updates, metadata, and single-step simulation through its step function.

4 Using Brax: Creating and evaluating environments

Brax lets users define physically simulated scenes through text or programmatic configurations, then build gym-like sequential decision problems with its env abstraction. The release includes MuJoCo-like locomotion tasks, a dexterous Grasp task, and a goal-directed Fetch environment.

  • System definition: ProtoBuf specifications define scene bodies, joints, actuators, and pairwise colliders.Users can also define systems programmatically, with equivalent example configurations provided in both forms.
  • System definition: Brax automatically computes valid body positions and rotations for joint-connected body trees through system.default_qp.The method determines the qp placing each body in a valid joint configuration.
  • Decision environments: The env class tracks initialization, resetting, observations, actions, and rewards, with a wrapper exposing an OpenAI Gym-style interface.This abstraction supplies the metadata needed for sequential decision problems beyond a physics update.
  • Packaged environments: The initial release includes MuJoCo-like Ant, Humanoid, and Halfcheetah tasks, plus Grasp and Fetch environments.These examples demonstrate locomotion, dexterous manipulation, and goal-based locomotion settings.
  • Packaged environments: Grasp is a pick-and-place proof of concept in which a 4-fingered claw hand moves a ball to a target location.The environment is intended to demonstrate that Brax contact physics support nontrivial manipulation tasks.
  • Packaged environments: Fetch supports training varied morphologies for locomotion within 50 million environment frames, using a toy boxy dog-like quadruped by default.The scene can be modified straightforwardly for new body morphologies.

5 Using Brax: Solving locomotion and manipulation problems

Brax provides JAX-implemented PPO, SAC, ES, and APG, with training procedures designed to keep environment processing and learning on accelerators. Its included environments typically solve in seconds to minutes with standard accelerators, while APG remains immature and does not produce locomotion gaits.

  • Algorithms: Brax includes PPO, SAC, ES, and APG implementations that exploit JAX parallelism and just-in-time compilation.PPO is on-policy, SAC is off-policy, ES is black-box optimization, and APG exploits differentiable environment rewards.
  • PPO: PPO generates rollouts and performs synchronized SGD updates entirely on accelerators without context switches.The batch is distributed across accelerator cores, with synchronized normalization statistics and gradient updates.
  • PPO: 75% of PPO training time goes to running the Ant environment, making environment processing the primary bottleneck under the best hyperparameters.The implementation is efficient enough that the environment, despite being fast, dominates throughput.
  • SAC: 78% of SAC training time goes to SGD updates, so a single accelerator core provides the most cost-efficient setup.The remaining breakdown is 12% for running the environment and 10% for working with the replay buffer; SAC uses an accelerator-resident replay buffer and a single jitted training function.
  • ES: > 99% of ES running time is spent evaluating environment steps.ES generates perturbations on a lead accelerator, distributes them across accelerator cores, and updates the policy from evaluation scores.
  • APG: APG compiles gradients through short trajectories but does not currently produce locomotive gaits and is prone to local minima.The paper describes APG as less mature than the other three algorithms and defers more advanced differentiable algorithm work.
  • Performance: The released hyperparameters typically solve each environment with a standard accelerator in seconds to minutes.The release also provides exhaustive hyperparameter experiments and performance plots for SAC and PPO.

6 Performance Benchmarking

Brax scales environment simulation to hundreds of millions of steps per second across accelerators, while its compiled PPO reaches performant Ant locomotion much faster than standard PPO. Its environments also show qualitatively similar SAC reward progression to MuJoCo counterparts and competitive linear momentum conservation scaling.

  • Simulation throughput: Brax scales to hundreds of millions of environment steps per second by distributing computation within and across accelerators.The reported scaling uses JAX vectorization and device-parallelism primitives across accelerator clusters.
  • Comparison limitations: Direct engine comparisons are difficult because widely used engines commonly rely on custom CPU multithreading or bespoke distributed accelerator setups.The authors caution that comparing Brax on a TPUv3 8x8 with MuJoCo-Ant on a single-threaded machine is not apples-to-apples.
  • Training speed: Brax’s Ant reaches performant locomotion in about ten seconds, whereas standard uncompiled PPO takes close to half an hour.The comparison uses Brax’s compiled, parallelism-optimized PPO against a traditional standard PPO implementation.
  • MuJoCo comparison: For fixed SAC hyperparameters, Brax environments achieve similar reward in a similar number of environment steps as their MuJoCo counterparts.The comparison is qualitative and concerns progression of reward gain rather than higher reward.
  • Simulation quality: Brax achieves competitive linear momentum conservation scaling in the astronaut diagnostic, attributed to maximal Cartesian position coordinates and symplectic integration.The diagnostic evaluates momentum and energy nonconservation as a function of simulation fidelity.

7 Limitations and Future Work

Brax’s spring-constraint design introduces stability, jitter, collision-detection, and compilation limitations that require tuning or constrain scalability. The release also leaves some algorithms insufficiently tested and raises broader benchmark, misuse, and compute-consumption concerns.

  • Spring constraints can make new scenes unstable, requiring tuning of damping forces, mass and inertia scale, and integration step size.Instabilities arise from a small integrator radius of convergence and worsen with differences in mass scale.
  • Brax simulations exhibit more jitter than hypothetical Featherstone simulations, while stronger joint springs reduce the maximum stable integration step size.The release chooses spring constants to maximize simulation speed while retaining qualitative behavior.
  • Naive collision detection scales quadratically and may become a bottleneck as tasks grow more complex, motivating more advanced collision methods.Current scenes can still parallelize collision primitives without straining modern accelerator memory buffers.
  • JIT compilation for complicated environments can take minutes, sometimes approaching or exceeding training time despite design efforts to reduce development friction.Brax compiles Pythonic physics environments and learning algorithms side-by-side to XLA.
  • APG and ES are less thoroughly tested than PPO and SAC, and future work should more fully leverage the engine’s differentiability.The authors also note benchmark proliferation, potential misuse, and the possibility that faster reinforcement learning increases compute expenditure.

A Appendix - Brax System Specification

The appendix demonstrates constructing the same Brax scene through a ProtoBuf specification and Python code. Both representations define a system with timestep .01 and gravity z = -9.8, then add parent and child bodies connected by a joint.

  • Scene construction: The appendix presents ProtoBuf and Pythonic constructions of the same Brax scene.The ProtoBuf example and a short Python snippet are both provided.
  • Bodies: The parent body is named Parent, frozen at position and rotation (1, 1, 1), with mass 1 and inertia (1, 1, 1).The Python construction begins adding the corresponding Parent body after creating the system.
  • Bodies: The child body is named Child and has mass 1 with inertia (1, 1, 1).The child is added to the system as a separate body in the Python example.

B Appendix - Grasp Trajectory · C Appendix - Hyperparameters for Figures

The appendices show a performant Brax grasping trajectory and document the hardware and hyperparameters used for training curves. The grasping policy carries a ball between randomly respawned red targets, while separate configurations are given for PPO, braxppo, and SAC experiments.

  • B Appendix - Grasp Trajectory: The grasping appendix presents a performant policy trained and simulated within Brax.The appendix includes a figure depicting the policy’s grasp trajectory.
  • B Appendix - Grasp Trajectory: During the first 300 steps, the hand picks up a ball and carries it to successive red targets.Each target respawns at a different random location after the ball approaches it.
  • C Appendix - Hyperparameters for Figures: The appendices detail hardware and hyperparameters used for all training curves and figures.Configurations cover standard PPO, braxppo, and SAC experiments.
  • C Appendix - Hyperparameters for Figures: Figure 3 compares a 128 Core Intel Xeon Processor at 2.2 Ghz for MuJoCo with a 32 Core Intel Xeon Processor at 2.0 GhZ for the 32x-CPU curve.Standard PPO uses [28], while braxppo uses the authors’ repository.
  • C Appendix - Hyperparameters for Figures: Standard PPO uses 10000000 training steps, evaluation every 10000 steps, and a 3e-4 learning rate.Its configuration also specifies 16-step unrolls, batch size 2048 // 16, 32 minibatches, and 10 epochs.
  • C Appendix - Hyperparameters for Figures: Braxppo uses 10000000 total environment steps, evaluation every 20 steps, reward scaling 10, and 2048 environments.Its learning rate is 3e-4, with unroll length 5, batch size 1024, 16 minibatches, and 4 update epochs.
  • C Appendix - Hyperparameters for Figures: Figure 4 uses a 32 Core Intel Xeon Processor at 2.2 Ghz for environments and a 2x2 TPUv2 for the learning algorithm.SAC settings for humanoid and ant include learning rate 3e-4, reward scale 0.1, replay size 10000, 5000000 steps, and 64 gradient updates per batch.
  • C Appendix - Hyperparameters for Figures: SAC for halfcheetah uses a 6e-4 learning rate, reward scale 10.0, replay size 10000, 5000000 steps, and 32 gradient updates per batch.The discount is .97, and evaluation occurs every 10000 steps.

D Appendix - Hyperparameter Sweeps

The appendix presents the top 20 training curves from exhaustive hyperparameter sweeps, with exact hyperparameters provided separately. It also reports Brax PPO and SAC reward curves across five environments over multiple training horizons on a single TPUv2.

  • Hyperparameter sweeps: The appendix plots the top 20 performing training curves from exhaustive hyperparameter sweeps.Precise hyperparameter values are available in zipped, sorted JSON files.
  • Hardware: All plots were generated on a 1x1 TPUv2, matching the hardware available on Colab’s free TPU tier.
  • PPO training: Over 10 million braxppo training steps, grasp and humanoid do not find successful policies.The reward curves show both steps and duplicated wallclock time in seconds.
  • PPO training: Over 500 million braxppo training steps, all policies except humanoid are solvable with ppo.The reward curves are shown against both training steps and wallclock time in seconds.
  • SAC training: Over 5 million brax-sac training steps, humanoid is solved via SAC, but grasp is not.The reward curves are shown against both training steps and wallclock time in seconds.

E Appendix - Major Differences from Mujoco … E.3 Humanoid

The appendix documents major implementation differences between Brax and the original MuJoCo halfcheetah, ant, and humanoid environments, while deferring exhaustive parity and transfer analysis. Differences include constrained integration, reward and reset modifications, and actuator implementation details.

  • E Appendix - Major Differences from Mujoco: Brax identifies major differences from MuJoCo in halfcheetah, ant, and humanoid, deferring exhaustive parity and policy-transfer analysis to future work.The authors plan to improve parity over time but do not analyze sim2real-style transfer here.
  • E.1 Halfcheetah: Brax constrains halfcheetah motion by masking rotational and translational integration updates, allowing planar movement and rotation around one axis.Mass, inertia, and actuator scales were chosen to approximate MuJoCo’s halfcheetah.
  • E.1 Halfcheetah: Both braxppo and braxsac find performant locomotive gaits despite possible hyperparameter mismatch between Brax and aggressively optimized MuJoCo settings.The authors searched fairly aggressively for performant Brax hyperparameters, but acknowledge a potentially poor hyperparameter-space region.
  • E.2 Ant: Brax’s Ant tunes mass, inertia, and actuator strengths, omits contact cost from its reward function, and produces the most qualitatively similar gaits between the engines.The contact-cost omission is cited to prior work.
  • E.3 Humanoid: Humanoid’s joint-torquing regularization penalty is 0.01 in Brax versus 0.1 in MuJoCo.The reward function is slightly modified in addition to tuning mass, inertia, and actuator strengths.
  • E.3 Humanoid: Humanoid torso reset thresholds are 0.6 to 2.1 in Brax, compared with 1.0 to 2.0 in MuJoCo.These thresholds determine where the torso triggers a done condition.
  • E.3 Humanoid: Brax implements three-degree-of-freedom humanoid actuators differently from MuJoCo.The paper directs readers to its joints implementation for details.
Loading 2106.13281v1…