DOCSEARCHbyAlgolia
Submit Docs

Brax Documentation

repository·main·Indexed 25 days ago

https://github.com/google/brax

A fast, fully differentiable physics engine written in JAX, designed for large-scale reinforcement learning and robotics research. Brax is optimized for TPUs and GPUs and provides four physics pipelines: MJX (MuJoCo XLA), Generalized, Positional, and Spring. It includes tools for training policies, such as the brax.training.learner module and the bin/learn CLI, and supports environments like the Barkour benchmark.

Tokens
7.2K
Snippets
12
Records
64
Agent score
85%

What's inside brax

  1. Overview of Brax physics pipelines

    main

    Brax provides four distinct physics pipelines that share the same API and can run side-by-side for experiments like transfer learning:

    • MuJoCo XLA - MJX: A JAX reimplementation of the MuJoCo physics engine.
    • Generalized: Calculates motion in generalized coordinates using dynamics algorithms similar to MuJoCo and TDS.
    • Positional: Uses Position Based Dynamics (PBD) for fast and stable resolution of joint and collision constraints.
    • Spring: Provides fast, impulse-based simulation for rapid experimentation, similar to methods used in video games.

  2. Access hyperparameter sweep results for SAC and PPO

    main
    The datasets/ directory contains results from hyperparameter sweeps for Soft Actor-Critic (SAC) and Proximal Policy Optimization (PPO) algorithms. The results are provided as zipped JSON files. Inside these archives, hyperparameters are organized first by environment and then sorted by performance.

  3. Quickstart via Google Colab

    main

    You can explore Brax without local installation using several Colab notebooks:

    • Brax Basics: Introduces the Brax API and basic physics primitives.
    • Brax Training: Introduces training algorithms, policy training, and loading/saving policies.
    • Brax Training with MuJoCo XLA - MJX: Demonstrates training using the MJX physics simulator.
    • Brax Training with PyTorch on GPU: Demonstrates using Brax with PyTorch for fast training.

  4. Install Brax from source

    main

    To install Brax from the source repository, clone the repository, navigate to the directory, and install in editable mode.

    python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install -e .

  9. Install Brax via pip

    main

    To install Brax from PyPI, use a virtual environment and pip. It is recommended to upgrade pip before installation.

    python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install brax

  11. Configure GPU rendering and XLA for Brax/MuJoCo

    main

    When using a GPU (e.g., in Google Colab), follow these steps to ensure proper rendering and performance:

    1. Enable Nvidia EGL driver: Create an ICD config so glvnd can pick up the Nvidia EGL driver.
    2. Optimize XLA: Tell XLA to use Triton GEMM to improve steps/sec by ~30% on some GPUs.
    3. Set MuJoCo backend: Configure MuJoCo to use the EGL rendering backend.
    # Add ICD config for Nvidia EGL
NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'
if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):
  with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:
    f.write('''{
    "file_format_version" : "1.0.0",
    "ICD" : {
        "library_path" : "libEGL_nvidia.so.0"
    }
}''')

# Enable Triton GEMM for XLA
xla_flags = os.environ.get('XLA_FLAGS', '')
xla_flags += ' --xla_gpu_triton_gemm_any=True'
os.environ['XLA_FLAGS'] = xla_flags

# Set MuJoCo to use EGL
%env MUJOCO_GL=egl
    #@title Check if MuJoCo installation was successful

from google.colab import files

import distutils.util
import os
import subprocess
if subprocess.run('nvidia-smi').returncode:
  raise RuntimeError(
      'Cannot communicate with GPU. '
      'Make sure you are using a GPU Colab runtime. '
      'Go to the Runtime menu and select Choose runtime type.')

# Add an ICD config so that glvnd can pick up the Nvidia EGL driver.
# This is usually installed as part of an Nvidia driver package, but the Colab
# kernel doesn't install its driver via APT, and as a result the ICD is missing.
# (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)
NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'
if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):
  with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:
    f.write("""
{
    "file_format_version" : "1.0.0",
    "ICD" : {
        "library_path" : "libEGL_nvidia.so.0"
    }
}
""")

# Tell XLA to use Triton GEMM, this improves steps/sec by ~30% on some GPUs
xla_flags = os.environ.get('XLA_FLAGS', '')
xla_flags += ' --xla_gpu_triton_gemm_any=True'
os.environ['XLA_FLAGS'] = xla_flags

# Configure MuJoCo to use the EGL rendering backend (requires GPU)
print('Setting environment variable to use GPU rendering:')
%env MUJOCO_GL=egl

try:
  print('Checking that the installation succeeded:')
  import mujoco
  mujoco.MjModel.from_xml_string('<mujoco/>')
except Exception as e:
  raise e from RuntimeError(
      'Something went wrong during installation. Check the shell output above '
      'for more information.\n' 
      'If using a hosted Colab runtime, make sure you enable GPU acceleration ' 
      'by going to the Runtime menu and selecting "Choose runtime type".')

print('Installation successful.')