MuJoCo Warp (MJWarp)#

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

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:

mjw.Model and mjw.Data contain Warp arrays that are copied onto device.
Some fields are missing from mjw.Model and mjw.Data for features that are unsupported.

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.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()}')

A call to mjw.step is comprised of a collection of kernel launches. Warp will launch these kernels individually if this 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 mjw.step directly. Please see the Warp Graph API reference for details.

Helpful 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`