Add graph mode to MJX, in preparation of supporting staged graph capture.

PiperOrigin-RevId: 865380826
Change-Id: I840b48027e1142f87eb745d0d19b0cd3e87904b2

doc/changelog.rst

@@ -84,7 +84,7 @@ General
 MJX
 ^^^
 - Added ``actuator_length``, ``cdof`` and ``cdof_dof`` fields to ``mjx.Data``.
-
+- Add ``graph_mode`` argument to ``put_model`` to support multiple Warp graph capture modes.
 
 Documentation
 ^^^^^^^^^^^^^

mjx/mujoco/mjx/_src/io.py

@@ -50,6 +50,11 @@ def _is_cuda_gpu_device(device: jax.Device) -> bool:
   return device in jax.devices('cuda')
 
 
+def _check_warp_installed():
+  if not mjxw.WARP_INSTALLED:
+    raise RuntimeError('warp-lang is not installed. Cannot use WARP implementation of MJX.')
+
+
 def _resolve_impl(
     device: jax.Device,
 ) -> types.Impl:
@@ -119,10 +124,7 @@ def _check_impl_device_compatibility(
           'Warp implementation requires a CUDA GPU device, got '
           f'{device}.'
       )
-    if not mjxw.WARP_INSTALLED:
-      raise RuntimeError(
-          'Warp is not installed. Cannot use Warp implementation of MJX.'
-      )
+    _check_warp_installed()
 
   is_cpu_device = device.platform == 'cpu'
   if impl == types.Impl.C or impl == types.Impl.CPP:
@@ -448,12 +450,10 @@ def _put_model_c(
 
 def _put_model_warp(
     m: mujoco.MjModel,
+    graph_mode: mjxw.types.GraphMode,
     device: Optional[jax.Device] = None,
 ) -> types.Model:
   """Puts mujoco.MjModel onto a device, resulting in mjx.Model."""
-  if not mjxw.WARP_INSTALLED:
-    raise RuntimeError('Warp not installed.')
-
   with wp.ScopedDevice('cpu'):  # pylint: disable=undefined-variable
     mw = mjwp.put_model(m)  # pylint: disable=undefined-variable
 
@@ -464,7 +464,10 @@ def _put_model_warp(
   option_keys = {f.name for f in mjxw.types.OptionWarp.fields()} - {
       f.name for f in types.Option.fields()
   }
+  # graph_mode is MJX-specific, not from mujoco.mjx.third_party.mujoco_warp.
+  option_keys = option_keys - {'graph_mode'}
   private_options = {k: getattr(mw.opt, k) for k in option_keys}
+  private_options['graph_mode'] = graph_mode
   fields['opt'] = _put_option(m.opt, types.Impl.WARP, private_options)
   fields['stat'] = _put_statistic(m.stat, types.Impl.WARP)
 
@@ -527,6 +530,7 @@ def put_model(
     m: mujoco.MjModel,
     device: Optional[jax.Device] = None,
     impl: Optional[Union[str, types.Impl]] = None,
+    graph_mode: Optional[mjxw.types.GraphMode] = None,
 ) -> types.Model:
   """Puts mujoco.MjModel onto a device, resulting in mjx.Model.
 
@@ -534,12 +538,15 @@ def put_model(
     m: the model to put onto device
     device: which device to use - if unspecified picks the default device
     impl: implementation to use
+    graph_mode: CUDA graph capture mode (for Warp only). Use GraphMode enum from
+      warp._src.jax_experimental.ffi. GraphMode.WARP is the default mode.
 
   Returns:
     an mjx.Model placed on device
 
   Raises:
     ValueError: if impl is not supported
+    RuntimeError: if impl is WARP and warp-lang is not installed
   """
 
   impl, device = _resolve_impl_and_device(impl, device)
@@ -548,7 +555,9 @@ def put_model(
   elif impl == types.Impl.C:
     return _put_model_c(m, device)
   elif impl == types.Impl.WARP:
-    return _put_model_warp(m, device)
+    _check_warp_installed()
+    graph_mode = graph_mode or getattr(mjxw.types.GraphMode, 'WARP')
+    return _put_model_warp(m, graph_mode, device)
   elif impl == types.Impl.CPP:
     return _put_model_cpp(m, device)
   else:
@@ -861,9 +870,6 @@ def _make_data_warp(
         f' {type(m)}.'
     )
 
-  if not mjxw.WARP_INSTALLED:
-    raise RuntimeError('Warp is not installed.')
-
   with wp.ScopedDevice('cpu'):  # pylint: disable=undefined-variable
     dw = mjwp.make_data(m, nworld=1, naconmax=naconmax, njmax=njmax)  # pylint: disable=undefined-variable
 
@@ -1003,6 +1009,7 @@ def make_data(
   elif impl == types.Impl.CPP:
     return _make_data_cpp(m, device)
   elif impl == types.Impl.WARP:
+    _check_warp_installed()
     naconmax = nconmax if naconmax is None else naconmax
     return _make_data_warp(m, device, naconmax, njmax)
 

mjx/mujoco/mjx/_src/io_test.py

@@ -334,6 +334,25 @@ class ModelIOTest(parameterized.TestCase):
 
     _ = jax.tree.map_with_path(check_ndim, mx)
 
+  @parameterized.parameters('JAX', 'WARP', None)
+  def test_put_model_warp_graph_mode(self, mode: str | None):
+    """Tests that put_model accepts graph_mode parameter."""
+    if not mjxw.WARP_INSTALLED:
+      self.skipTest('Warp not installed.')
+    if not mjx_io.has_cuda_gpu_device():
+      self.skipTest('No CUDA GPU device available.')
+
+    if mode is None:
+      graph_mode = None
+    else:
+      graph_mode = getattr(mjxw_types.GraphMode, mode)
+
+    m = mujoco.MjModel.from_xml_string(_SIMPLE_BODY)
+    mx = mjx.put_model(m, impl='warp', graph_mode=graph_mode)
+
+    expected = graph_mode or mjxw_types.GraphMode.WARP
+    self.assertEqual(mx.opt._impl.graph_mode, expected)
+
   @parameterized.parameters('c', 'jax')
   def test_unsupported_contact_types(self, impl):
     """Tests that unsupported contact types raise an error."""

mjx/mujoco/mjx/warp/collision_driver.py

@@ -42,7 +42,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _collision_shim(
     # Model
@@ -279,6 +278,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
           'contact__type',
           'contact__worldid',
       },
+      graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
       d.qpos.shape[0],

mjx/mujoco/mjx/warp/forward.py

@@ -1220,6 +1220,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'efc__type',
           'efc__vel',
       },
+      graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
       d.qpos.shape[0],
@@ -3018,6 +3019,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'efc__type',
           'efc__vel',
       },
+      graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
       d.qpos.shape[0],

mjx/mujoco/mjx/warp/forward_test.py

@@ -216,8 +216,11 @@ class StepTest(parameterized.TestCase):
           'pendula.xml',
       ),
       batch_size=(1, 7),
+      # NOTE: GraphMode.JAX is incompatible with MuJoCo Warp at the moment,
+      # even when setting graph_conditional=False.
+      graph_mode=('WARP',),
   )
-  def test_step(self, xml: str, batch_size: int):
+  def test_step(self, xml: str, batch_size: int, graph_mode: str):
     if not _FORCE_TEST:
       if not mjxw.WARP_INSTALLED:
         self.skipTest('Warp not installed.')
@@ -227,7 +230,9 @@ class StepTest(parameterized.TestCase):
     m = test_util.load_test_file(xml)
     m.opt.iterations = 10
     m.opt.ls_iterations = 10
-    mx = mjx.put_model(m, impl='warp')
+    mx = mjx.put_model(
+        m, impl='warp', graph_mode=getattr(mjxw.types.GraphMode, graph_mode)
+    )
 
     d = mujoco.MjData(m)
     worldids = jp.arange(batch_size)

mjx/mujoco/mjx/warp/smooth.py

@@ -170,6 +170,7 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
           'xpos',
           'xquat',
       },
+      graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
       d.qpos.shape[0],
@@ -379,6 +380,7 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
           'wrap_obj',
           'wrap_xpos',
       },
+      graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
       d.qpos.shape[0],

mjx/mujoco/mjx/warp/types.py

@@ -16,12 +16,21 @@
 DO NOT EDIT. This file is auto-generated.
 """
 import dataclasses
+import typing
 from typing import Tuple
 import jax
 from jax import tree_util
 from jax.interpreters import batching
 from mujoco.mjx._src import dataclasses as mjx_dataclasses
 import numpy as np
+
+if typing.TYPE_CHECKING:
+  GraphMode = int
+else:
+  try:
+    from warp._src.jax_experimental.ffi import GraphMode
+  except ImportError:
+    GraphMode = int
 PyTreeNode = mjx_dataclasses.PyTreeNode
 
 @dataclasses.dataclass(frozen=True)
@@ -35,7 +44,6 @@ class TileSet:
     adr: address of each tile in the set
     size: size of all the tiles in this set
   """
-
   adr: np.ndarray
   size: int
 
@@ -95,6 +103,7 @@ class OptionWarp(PyTreeNode):
   ccd_tolerance: jax.Array
   contact_sensor_maxmatch: int
   graph_conditional: bool
+  graph_mode: GraphMode
   has_fluid: bool
   impratio_invsqrt: jax.Array
   is_sparse: bool