From e10853e52e8af8279c6ec0c5663ee887148ed405 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Mon, 6 Oct 2025 04:02:35 -0700 Subject: [PATCH] MuJoCo Warp basic usage and minimal example documentation. PiperOrigin-RevId: 815654625 Change-Id: Ifa90a7469ef827777317bd3c4cda216130a5c7e9 --- doc/mjwarp/index.rst | 107 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 88 insertions(+), 19 deletions(-) diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst index 85aad6bc..e5436ca7 100644 --- a/doc/mjwarp/index.rst +++ b/doc/mjwarp/index.rst @@ -14,23 +14,6 @@ for `Nvidia `__ GPUs. MJWarp lives in the `google-deepmind/mujoco_warp `__ GitHub repository and is currently in beta. -.. _MJW_getstarted: - -Getting Started -=============== - -.. code-block:: python - - import mujoco - import mujoco_warp as mjw - - mjm = mujoco.MjModel.from_file("example.xml") - - m = mjw.put_model(mjm) - d = mjw.make_data(mjm) - - mjw.step(m, d) - .. _MJW_install: Installation @@ -55,10 +38,96 @@ Test the Installation pytest +.. _MJW_Usage: + +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 :ref:`mjModel` on device yields an ``mjw.Model``. Placing +an :ref:`mjData` on device yields an ``mjw.Data``: + +.. code-block:: python + + 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 :ref:`main simulation ` and +some of the :ref:`sub-components ` for forward simulation are available from the top-level ``mjw`` +module. + +Minimal example +--------------- + +.. code-block:: python + + # Throw a ball at 100 different velocities. + + import mujoco + import mujoco_warp as mjw + import warp as wp + + _MJCF=r""" + + + + + + + + + """ + + 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 + +.. code-block:: python + + with wp.ScopedCapture() as capture: + mjw.step(m, d) + +The graph can then be launched or re-launched + +.. code-block:: python + + 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. + .. _MJW_Cli: -Utilities -========= +Helpful Command Line Scripts +---------------------------- Benchmark an environment with testspeed