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

• URL: http://arxiv.org/abs/2106.13281v1
• Date: Thu, 24 Jun 2021 19:09:12 GMT
• Title: Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation
• Abstract summary: We present Brax, an open source library for rigid body simulation written in JAX. We present results on a suite of tasks inspired by the existing reinforcement learning literature, but remade in our engine.
• Score: 33.36244621210259
• License: http://creativecommons.org/licenses/by-nc-sa/4.0/
• 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.

Related papers

• UniLab: A Heterogeneous Architecture for Robot RL Beyond GPU-Dominant Paradigms [68.40808444348735]
  In simulation-dominated robot control, the essential question is not which processor runs physics, but whether simulation throughput, policy learning, and runtime synchronization form an efficient end-to-end loop.<n>We present UniLab, a heterogeneous CPU-resident / GPU-learning architecture that decouples CPU-parallel simulation from GPU policy updates through a unified runtime for data movement, buffering, and synchronization.
  arXiv  Detail & Related papers  (2026-05-28T17:53:50Z)
• TABX: A High-Throughput Sandbox Battle Simulator for Multi-Agent Reinforcement Learning [4.254850120280717]
  We introduce the Totally Accelerated Battle Simulator in JAX (TABX) for reconfigurable multi-agent tasks.<n>TABX enables massive parallelization and significantly reduces computational overhead.<n>By providing a fast, scalable, and easily customized framework, TABX serves as a scalable foundation for future research.
  arXiv  Detail & Related papers  (2026-02-02T05:34:38Z)
• Isaac Lab: A GPU-Accelerated Simulation Framework for Multi-Modal Robot Learning [72.43357471969564]
  Isaac Lab combines high-fidelity GPU parallel physics, rendering, and a modular, composable architecture for designing environments and training robot policies.<n>We highlight its application to a diverse set of challenges, including whole-body control, cross-embodiment mobility, contact-rich and dexterous manipulation, and the integration of human demonstrations for skill acquisition.<n>We believe Isaac Lab's combination of advanced simulation capabilities, rich sensing, and data-center scale execution will help unlock the next generation of breakthroughs in robotics research.
  arXiv  Detail & Related papers  (2025-11-06T21:43:02Z)
• GEM: A Gym for Agentic LLMs [88.36970707762424]
  General Experience Maker ( GEM) is an open-source environment simulator designed for the age of large language models (LLMs)<n> GEM provides a standardized framework for the environment-agent interface, including asynchronous vectorized execution for high throughput.<n>We conduct apple-to-apple benchmarking of PPO, GRPO and REINFORCE in both single- and multi-turn settings using GEM to shed light on the algorithmic designs.
  arXiv  Detail & Related papers  (2025-10-01T15:55:57Z)
• AXLearn: Modular Large Model Training on Heterogeneous Infrastructure [64.33868455931301]
  AXLearn is a production deep learning system that facilitates scalable and high-performance training of large deep learning models.<n>Compared to other state-of-the-art deep learning systems, AXLearn has a unique focus on modularity and support for heterogeneous hardware infrastructure.
  arXiv  Detail & Related papers  (2025-07-07T18:50:58Z)
• Optimisation of the Accelerator Control by Reinforcement Learning: A Simulation-Based Approach [0.615163395430594]
  This study aims to create a simulation-based framework integrated with Reinforcement Learning (RL) Using textttElegant as the simulation backend, we developed a Python wrapper that simplifies the interaction between RL algorithms and accelerator simulations. The proposed RL framework acts as a co-pilot for physicists, offering intelligent suggestions to enhance beamline performance, reduce tuning time, and improve operational efficiency.
  arXiv  Detail & Related papers  (2025-03-12T16:57:52Z)
• AutoHete: An Automatic and Efficient Heterogeneous Training System for LLMs [68.99086112477565]
  Transformer-based large language models (LLMs) have demonstrated exceptional capabilities in sequence modeling and text generation. Existing heterogeneous training methods significantly expand the scale of trainable models but introduce substantial communication overheads and CPU workloads. We propose AutoHete, an automatic and efficient heterogeneous training system compatible with both single- GPU and multi- GPU environments.
  arXiv  Detail & Related papers  (2025-02-27T14:46:22Z)
• minimax: Efficient Baselines for Autocurricula in JAX [42.270349368227436]
  This work introduces the minimax library for UED training on accelerated hardware. Using JAX to implement fully-tensorized environments and autocurriculum algorithms, minimax allows the entire training loop to be compiled for hardware acceleration.
  arXiv  Detail & Related papers  (2023-11-21T16:43:13Z)
• Waymax: An Accelerated, Data-Driven Simulator for Large-Scale Autonomous Driving Research [76.93956925360638]
  Waymax is a new data-driven simulator for autonomous driving in multi-agent scenes. It runs entirely on hardware accelerators such as TPUs/GPUs and supports in-graph simulation for training. We benchmark a suite of popular imitation and reinforcement learning algorithms with ablation studies on different design decisions.
  arXiv  Detail & Related papers  (2023-10-12T20:49:15Z)
• JAX-LOB: A GPU-Accelerated limit order book simulator to unlock large scale reinforcement learning for trading [8.884142720013081]
  Financial exchanges use limit order books (LOBs) to process orders and match trades. For many applications, there is a requirement for processing multiple books, either for the calibration of ABMs or for the training of RL agents. We showcase the first GPU-enabled LOB simulator designed to process thousands of books in parallel, with a notably reduced per-message processing time.
  arXiv  Detail & Related papers  (2023-08-25T10:26:43Z)
• Gym-preCICE: Reinforcement Learning Environments for Active Flow Control [0.0]
  Gym-preCICE is a Python adapter fully compliant with Gymnasium (formerly known as OpenAI Gym) API. Gym-preCICE takes advantage of preCICE, an open-source coupling library for partitioned multi-physics simulations. The framework results in a seamless integration of realistic physics-based simulation toolboxes with RL algorithms.
  arXiv  Detail & Related papers  (2023-05-03T10:54:56Z)
• Orbit: A Unified Simulation Framework for Interactive Robot Learning Environments [38.23943905182543]
  We present Orbit, a unified and modular framework for robot learning powered by NVIDIA Isaac Sim. It offers a modular design to create robotic environments with photo-realistic scenes and high-fidelity rigid and deformable body simulation. We aim to support various research areas, including representation learning, reinforcement learning, imitation learning, and task and motion planning.
  arXiv  Detail & Related papers  (2023-01-10T20:19:17Z)
• Braxlines: Fast and Interactive Toolkit for RL-driven Behavior Engineering beyond Reward Maximization [15.215372246434413]
  In reinforcement learning (RL)-driven approaches, the goal of continuous control is to synthesize desired behaviors. In this paper, we introduce braxlines, a toolkit for fast and interactive-driven behavior generation beyond simple reward RL. Our implementations build on a hardware-accelerated Brax simulator in Jax with minimal modifications, enabling behavior within minutes of training.
  arXiv  Detail & Related papers  (2021-10-10T02:41:01Z)
• SurRoL: An Open-source Reinforcement Learning Centered and dVRK Compatible Platform for Surgical Robot Learning [78.76052604441519]
  SurRoL is an RL-centered simulation platform for surgical robot learning compatible with the da Vinci Research Kit (dVRK) Ten learning-based surgical tasks are built in the platform, which are common in the real autonomous surgical execution. We evaluate SurRoL using RL algorithms in simulation, provide in-depth analysis, deploy the trained policies on the real dVRK, and show that our SurRoL achieves better transferability in the real world.
  arXiv  Detail & Related papers  (2021-08-30T07:43:47Z)
• PlasticineLab: A Soft-Body Manipulation Benchmark with Differentiable Physics [89.81550748680245]
  We introduce a new differentiable physics benchmark called PasticineLab. In each task, the agent uses manipulators to deform the plasticine into the desired configuration. We evaluate several existing reinforcement learning (RL) methods and gradient-based methods on this benchmark.
  arXiv  Detail & Related papers  (2021-04-07T17:59:23Z)
• Composable Learning with Sparse Kernel Representations [110.19179439773578]
  We present a reinforcement learning algorithm for learning sparse non-parametric controllers in a Reproducing Kernel Hilbert Space. We improve the sample complexity of this approach by imposing a structure of the state-action function through a normalized advantage function. We demonstrate the performance of this algorithm on learning obstacle-avoidance policies in multiple simulations of a robot equipped with a laser scanner while navigating in a 2D environment.
  arXiv  Detail & Related papers  (2021-03-26T13:58:23Z)
• A User's Guide to Calibrating Robotics Simulators [54.85241102329546]
  This paper proposes a set of benchmarks and a framework for the study of various algorithms aimed to transfer models and policies learnt in simulation to the real world. We conduct experiments on a wide range of well known simulated environments to characterize and offer insights into the performance of different algorithms. Our analysis can be useful for practitioners working in this area and can help make informed choices about the behavior and main properties of sim-to-real algorithms.
  arXiv  Detail & Related papers  (2020-11-17T22:24:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.