MuJoCo Warp (MJWarp)#

MuJoCo Warp (MJWarp) is an implementation of MuJoCo written in Warp and optimized for NVIDIA hardware and parallel simulation. MJWarp lives in the google-deepmind/mujoco_warp GitHub repository and is currently in beta.

MJWarp is developed and maintained as a joint effort by NVIDIA and Google DeepMind.

Beta software

  • MJWarp is beta software and is under active development.

  • MJWarp developers will triage and respond to bug reports and feature requests.

  • MJWarp is mostly feature complete but requires performance optimization, documentation, and testing.

  • The intended audience during Beta are physics engine enthusiasts and learning framework integrators.

Installation#

The beta version of MuJoCo Warp is installed from GitHub. Please note that the beta version of MuJoCo Warp does not support all versions of MuJoCo, Warp, CUDA, NVIDIA drivers, etc.

git clone https://github.com/google-deepmind/mujoco_warp.git
cd mujoco_warp
python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install uv
uv pip install -e .[dev,cuda]

Test the Installation

pytest

Basic Usage#

Once installed, the package can be imported via import mujoco_warp as mjw. Structs, functions, and enums are available directly from the top-level mjw module.

Structs#

Before running MJWarp functions on an NVIDIA GPU, structs must be copied onto the device via mjw.put_model and mjw.make_data or mjw.put_data functions. Placing an mjModel on device yields an mjw.Model. Placing an mjData on device yields an mjw.Data.

mjm = mujoco.MjModel.from_xml_string("...")
mjd = mujoco.MjData(mjm)
m = mjw.put_model(mjm)
d = mjw.put_data(mjm, mjd)

These MJWarp variants mirror their MuJoCo counterparts but have a few key differences:

  1. mjw.Model and mjw.Data contain Warp arrays that are copied onto device.

  2. Some fields are missing from mjw.Model and mjw.Data for features that are unsupported.

Batch sizes#

MJWarp is optimized for parallel simulation. A batch of simulations can be specified with three parameters:

  • nworld: Number of worlds to simulate.

  • nconmax: Expected number of contacts per world. The maximum number of contacts for all worlds is nconmax * nworld.

  • naconmax: Alternative to nconmax, maximum number of contacts over all worlds. If nconmax and naconmax are both set then nconmax is ignored.

  • njmax: Maximum number of constraints per world.

Semantic difference for nconmax and njmax.

It is possible for the number of contacts per world to exceed nconmax if the total number of contacts for all worlds does not exceed nworld x nconmax. However, the number of constraints per world is strictly limited by njmax.

XML parsing

Values for nconmax and njmax are not parsed from size/nconmax and size/njmax (these parameters are deprecated). Values for these parameters must be provided to mjw.make_data or mjw.put_data.

Functions#

MuJoCo functions are exposed as MJWarp functions of the same name, but following PEP 8-compliant names. Most of the main simulation and some of the sub-components for forward simulation are available from the top-level mjw module.

Minimal example#

# Throw a ball at 100 different velocities.

import mujoco
import mujoco_warp as mjw
import warp as wp

_MJCF=r"""
<mujoco>
  <worldbody>
    <body>
      <freejoint/>
      <geom size=".15" mass="1" type="sphere"/>
    </body>
  </worldbody>
</mujoco>
"""

mjm = mujoco.MjModel.from_xml_string(_MJCF)
m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=100)

# initialize velocities
wp.copy(d.qvel, wp.array([[float(i) / 100, 0, 0, 0, 0, 0] for i in range(100)], dtype=float))

# simulate physics
mjw.step(m, d)

print(f'qpos:\n{d.qpos.numpy()}')

Command line scripts#

Benchmark an environment with testspeed

mjwarp-testspeed benchmark/humanoid/humanoid.xml

Interactive environment simulation with MJWarp

mjwarp-viewer benchmark/humanoid/humanoid.xml

Feature Parity#

MJWarp supports most of the main simulation features of MuJoCo, with a few exceptions. MJWarp will raise an exception if asked to copy to device an mjModel with field values referencing unsupported features.

The following features are not supported in MJWarp:

Category

Feature

Equality

FLEX

Integrator

IMPLICIT, IMPLICITFAST not supported with fluid drag

Solver

PGS, noslip, islands

Fluid Model

Ellipsoid model

Sensors

GEOMDIST, GEOMNORMAL, GEOMFROMTO

Flex

VERTCOLLIDE=false, INTERNAL=true, nflex > 1

Jacobian format

SPARSE

Option

contact override

Plugins

All except SDF

User parameters

All

Performance Tuning#

The following are considerations for optimizing the performance of MJWarp.

Graph capture#

MJWarp functions, for example mjw.step, often comprise a collection of kernel launches. Warp will launch these kernels individually if the function is called directly. To improve performance, especially if the function will be called multiple times, it is recommended to capture the operations that comprise the function as a CUDA graph

with wp.ScopedCapture() as capture:
  mjw.step(m, d)

The graph can then be launched or re-launched

wp.capture_launch(capture.graph)

and will typically be significantly faster compared to calling the function directly. Please see the Warp Graph API reference for details.

Batch sizes#

The maximum numbers of contacts and constraints, nconmax / naconmax and njmax respectively, are specified when creating mjw.Data with mjw.make_data or mjw.put_data. Memory and computation scales with the values of these parameters. For best performance, the values of these parameters should be set as small as possible while ensuring the simulation does not exceed these limits.

It is expected that good values for these limits will be environment specific. In practice, selecting good values typically involves trial-and-error. mjwarp-testspeed with the flag --measure_alloc for printing the number of contacts and constraints at each simulation step and interacting with the simulation via mjwarp-viewer and checking for overflow errors can both be useful techniques for iteratively testing values for these parameters.

Solver iterations#

MuJoCo’s default solver settings for the maximum numbers of solver iterations and linesearch iterations are expected to provide reasonable performance. Reducing MJWarp’s settings Option.iterations and/or Option.ls_iterations limits may improve performance and should be secondary considerations after tuning nconmax / naconmax and njmax.

Reducing these limits too much may prevent the constraint solver from converging and can lead to inaccurate or unstable simulation.

Impact on Performance: MJX (JAX) and MJWarp

In MJX these solver parameters are key for controlling simulation performance. With MJWarp, in contrast, once all worlds have converged the solver can early exit and avoid unnecessary computation. As a result, the values of these settings have comparatively less impact on performance.

Contact sensor matching#

Scenes that include contact sensors have a parameter that specifies the maximum number of matched contacts per sensor Option.contact_sensor_max_match. For best performance, the value of this parameter should be as small as possible while ensuring the simulation does not exceed the limit. Matched contacts that exceed this limit will be ignored.

The value of this parameter can be set directly, for example model.opt.contact_sensor_maxmatch = 16, or via an XML custom numeric field

<custom>
  <numeric name="contact_sensor_maxmatch" data="16"/>
</custom>

Similar to the maximum numbers of contacts and constraints, a good value for this setting is expected to be environment specific. mjwarp-testspeed and mjwarp-viewer may be useful for tuning the value of this parameter.

Parallel linesearch#

In addition to the constraint solver’s iterative linesearch, MJWarp provides a parallel linesearch routine that evaluates a set of step sizes in parallel and selects the best one. The step sizes are spaced logarithmically from Model.opt.ls_parallel_min_step to 1 and the number of step sizes to evaluate is set via Model.opt.ls_iterations.

In some cases the parallel routine may provide improved performance compared to the constraint solver’s default iterative linesearch.

To enable this routine set Model.opt.ls_parallel=True or add a custom numeric field to the XML

<custom>
  <numeric name="ls_parallel" data="1"/>
</custom>

Experimental feature

The parallel linesearch is currently an experimental feature.

Batched Model Fields#

To enable batched simulation with different model parameter values, many mjw.Model fields have a leading batch dimension. By default, the leading dimension is 1 (i.e., field.shape[0] == 1) and the same value(s) will be applied to all worlds. It is possible to override one of these fields with a wp.array that has a leading dimension greater than one. This field will be indexed with a modulo operation of the world id and batch dimension: field[worldid % field.shape[0]]. Importantly, the field shape should be overridden prior to graph capture (i.e., wp.ScopedCapture)

# override shape and values
m.dof_damping = wp.array([[0.1], [0.2]], dtype=float)

with wp.ScopedCapture() as capture:
  mjw.step(m, d)

It is possible to override the field shape and set the field values after graph capture

# override shape
m.dof_damping = wp.empty((2, 1), dtype=float)

with wp.ScopedCapture() as capture:
  mjw.step(m, d)

# set batched values
dof_damping_batch = wp.array([[0.1], [0.2]], dtype=float)
wp.copy(m.dof_damping, dof_damping_batch)  # m.dof = dof_damping_batch will not update the captured graph

Heterogeneous worlds

Heterogeneous worlds, for example: per-world meshes or number of degrees of freedom, are not currently available.

Frequently Asked Questions#

Learning frameworks#

Does MJWarp work with JAX?

Yes. MJWarp is interoperable with JAX. Please see the Warp Interoperability documentation for details.

Additionally, MJX provides a JAX API for a subset of MJWarp’s API. The backend is specified with impl='warp'.

Does MJWarp work with PyTorch?

Yes. MJWarp is interoperable with PyTorch. Please see the Warp Interoperability documentation for details.

How to train policies with MJWarp physics?

For examples that train policies with MJWarp physics, please see:

Features#

Is MJWarp differentiable?

No. MJWarp is not currently differentiable via Warp’s automatic differentiation functionality. Updates from the team related to enabling automatic differentiation for MJWarp are tracked in this GitHub issue.

Does MJWarp work with multiple GPUs?

Yes. Warp’s wp.ScopedDevice enables multi-GPU computation

# create a graph for each device
graph = {}
for device in wp.get_cuda_devices():
  with wp.ScopedDevice(device):
    m = mjw.put_model(mjm)
    d = mjw.make_data(mjm)
    with wp.ScopedCapture(device) as capture:
      mjw.step(m, d)
    graph[device] = capture.graph

# launch a graph on each device
for device in wp.get_cuda_devices():
  wp.capture_launch(graph[device])

Please see the Warp documentation for details and mjlab distributed training for a reinforcement learning example.

Is MJWarp on GPU deterministic?

No. There may be ordering or small numerical differences between results computed by different executions of the same code. This is characteristic of non-deterministic atomic operations on GPU. Set device to CPU with wp.set_device("cpu") for deterministic results.

Developments for deterministic results on GPU are tracked in this GitHub issue.

How are orientations represented?

Orientations are represented as unit quaternions and follow MuJoCo’s conventions: w, x, y, z or scalar, vector.

wp.quaternion

MJWarp utilizes Warp’s built-in type wp.quaternion. Importantly however, MJWarp does not utilize Warp’s x, y, z, w quaternion convention or operations and instead implements quaternion routines that follow MuJoCo’s conventions. Please see math.py for the implementations.