Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 941234486
Change-Id: Ie5d8cc0c9975e66f9a81bb6a86e33f640a5d19e3

Author: thowell

mjx/mujoco/mjx/_src/io.py

@@ -174,9 +174,6 @@ def f(leaf):
 
 def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
   """Converts a warp type to an MJX compatible numpy type."""
-  if hasattr(wp_field, '_is_batched'):
-    wp_field.strides = wp_field.strides[1:]
-    wp_field.shape = wp_field.shape[1:]
   # warp scalars
   wp_dtype = type(wp_field)
   if wp_dtype in wp._src.types.warp_type_to_np_dtype:
@@ -202,7 +199,13 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
       wp_field[0], mjwp_types.TileSet
   ):
     return tuple(
-        mjxw.types.TileSet(wp_field[i].adr.numpy(), wp_field[i].size)
+        mjxw.types.TileSet(
+            wp_field[i].adr.numpy(),
+            wp_field[i].size,
+            wp_field[i].elemid.numpy()
+            if wp_field[i].elemid is not None
+            else np.zeros(0, dtype=np.int32),
+        )
         for i in range(len(wp_field))
     )
   if isinstance(wp_field, mjwp_types.BlockDim):
@@ -272,11 +275,15 @@ def _put_option(
 
 
   if impl == types.Impl.WARP:
-    if not mjxw.mjwp_io.ENABLE_ISLANDS:
-      fields['disableflags'] = types.DisableBit(
-          fields['disableflags'] | mjwp_types.DisableBit.ISLAND
-      )
-    impl_fields = {k: _wp_to_np_type(v) for k, v in impl_fields.items()}
+    impl_fields = {
+        k: (
+            v_np.reshape(v_np.shape[1:])
+            if isinstance(v_np := _wp_to_np_type(v, k), np.ndarray)
+            and mjxw.types._BATCH_DIM['Option'].get(k, False)  # pylint: disable=protected-access
+            else v_np
+        )
+        for k, v in impl_fields.items()
+    }
     return types.Option(**fields, _impl=mjxw.types.OptionWarp(**impl_fields))
 
   raise NotImplementedError(f'Unsupported implementation: {impl}')
@@ -460,6 +467,8 @@ def _put_model_warp(
     if not hasattr(mw, k) or k in ('stat', 'opt'):
       continue
     field = _wp_to_np_type(getattr(mw, k), k)
+    if mjxw.types._BATCH_DIM['Model'].get(k, False):  # pylint: disable=protected-access
+      field = field.reshape(field.shape[1:])
     if k == 'geom_dataid' and field.ndim > 1:
       # Batched geom_dataid is not supported in MJX.
       field = field[0]
@@ -468,6 +477,8 @@ def _put_model_warp(
   impl_fields = {}
   for k in mjxw.types.ModelWarp.__annotations__.keys():
     field = _wp_to_np_type(getattr(mw, k), k)
+    if mjxw.types._BATCH_DIM['Model'].get(k, False):  # pylint: disable=protected-access
+      field = field.reshape(field.shape[1:])
     impl_fields[k] = field
 
   model = types.Model(
@@ -1298,18 +1309,19 @@ def _get_data_into_warp(
           'ten_J',
           'flexedge_J',
           'M',
+          'map_efc2iefc', 'map_iefc2efc'
       ):
         continue
       if field.name.startswith('efc_'):
         continue
 
+      # Skip island fields; host MjData arena memory cannot be reallocated from Python if not stepped.
+      if field.name.startswith('island_'):
+        continue
+
       if isinstance(value, np.ndarray) and value.shape:
         result_field = getattr(result_i, field.name)
         if result_field.shape != value.shape:
-          # When ENABLE_ISLANDS is False, mujoco_warp allocates island fields with width 0,
-          # while the host MjData sizes to nv/ntree. Skip to prevent mismatch.
-          if value.size == 0 and not mjxw.mjwp_io.ENABLE_ISLANDS:
-            continue
           raise ValueError(
               f'Input field {field.name} has shape {value.shape}, but output'
               f' has shape {result_field.shape}'

mjx/mujoco/mjx/_src/io_test.py

@@ -550,7 +550,15 @@ def test_put_data(self, impl: str):
     elif impl == 'warp':
       qm = np.zeros((m.nv, m.nv), dtype=np.float64)
       mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
-      np.testing.assert_allclose(dx._impl.M[:m.nv, :m.nv], qm)
+      warp_M = np.zeros((m.nv, m.nv))
+      mujoco.mju_sym2dense(
+          warp_M,
+          np.array(dx._impl.M),
+          m.M_rownnz,
+          m.M_rowadr,
+          m.M_colind,
+      )
+      np.testing.assert_allclose(warp_M, qm)
       # TODO(taylorhowell): test efc__J
       np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3])
 
@@ -811,12 +819,6 @@ def test_get_data_into_warp(self):
     dx = mjx.make_data(m, impl='warp')
     mjx.get_data_into(d, mx, dx)
 
-    # Island data is not populated when ENABLE_ISLANDS is False, so the host
-    # MjData island fields should be left at their default (nv,)/(ntree,)
-    # shapes rather than triggering a shape mismatch.
-    self.assertEqual(d.dof_island.shape, (m.nv,))
-    self.assertEqual(d.tree_island.shape, (m.ntree,))
-
   @parameterized.parameters(('jax',))
   def test_get_data_into_wrong_shape(self, impl):
     """Tests that get_data_into throwsif input and output shapes don't match."""

mjx/mujoco/mjx/codegen/generate_warp_shim.py

@@ -327,7 +327,12 @@ def _jax_shim_fn(
           f'Unknown param source: {mjwarp_field_info[arg].param_source}'
       )
 
-    if arg in field_usage.data_out_fields:
+    # Only treat as output if it is a Warp array (excludes scalars like
+    # naconmax).
+    if (
+        arg in field_usage.data_out_fields
+        and 'array' in mjwarp_field_info[arg].expected_type
+    ):
       # all out fields are in_out, since JAX already allocated them
       in_out_argnames.append(f"'{arg}'")
       num_outputs += 1
@@ -344,7 +349,11 @@ def _jax_shim_fn(
   if field_usage.render_context_in_caller:
     jax_args.append('ctx.key')
 
-  needs_dummy_output = not field_usage.data_out_fields
+  # If there are no Warp array outputs, we need a dummy output for JAX FFI.
+  needs_dummy_output = not any(
+      'array' in mjwarp_field_info[f].expected_type
+      for f in field_usage.data_out_fields
+  )
   if needs_dummy_output and fn_name != 'render':
     num_outputs = 1
     if field_usage.render_context_in_caller:

mjx/mujoco/mjx/codegen/generate_warp_types.py

@@ -181,11 +181,13 @@ def _get_annotations_recursive(
 
 
 def _build_new_class_body_ast(
-    keys: Set[str],
+    keys: typing.Iterable[str],
     cls_name: str,
     target_annotations: Dict[str, Any],
+    defaults: Optional[Dict[str, Any]] = None,
     shape_property: Optional[str] = None,
     add_docstring: bool = True,
+    sort_keys: bool = True,
 ) -> List[ast.AST]:
   """Builds the list of AST nodes for the new class body."""
   new_body_nodes: List[ast.AST] = []
@@ -194,16 +196,29 @@ def _build_new_class_body_ast(
     docstring = f'Derived fields from {cls_name}.'
     new_body_nodes.append(ast.Expr(value=ast.Constant(value=docstring)))
 
-  # Sort keys alphabetically before creating AST nodes
-  sorted_keys = sorted(list(keys))
-  for key in sorted_keys:
+  if sort_keys:
+    maybe_sorted_keys = sorted(list(keys))
+  else:
+    maybe_sorted_keys = list(keys)
+
+  for key in maybe_sorted_keys:
     annotation_node = _get_target_annotation_node(key, target_annotations)
 
+    value_node = None
+    if defaults and key in defaults:
+      value_node = ast.Constant(value=defaults[key])
+
+    if value_node and value_node.value is None:
+      annotation_node = _ast_parse_type(
+          f'typing.Optional[{ast.unparse(annotation_node)}]'
+      )
+
     new_body_nodes.append(
         ast.AnnAssign(
             target=ast.Name(id=key, ctx=ast.Store()),
             annotation=annotation_node,
-            simple=1,  # No value assignment
+            value=value_node,
+            simple=1,
         )
     )
 
@@ -316,11 +331,21 @@ def __hash__(self) -> int:
 
 
 def write_nested_dataclass(target_fpath: epath.Path, cls: Any):
+  fields = dataclasses.fields(cls)
+  keys = [f.name for f in fields]
+  defaults = {
+      f.name: f.default
+      for f in fields
+      if f.default is not dataclasses.MISSING
+  }
+
   new_class_body = _build_new_class_body_ast(
-      set(cls.__annotations__.keys()),
+      keys,
       cls.__name__,
       dict(cls.__annotations__),
+      defaults=defaults,
       add_docstring=False,
+      sort_keys=False,
   )
   cls_str = '\n'.join(['  ' + ast.unparse(node) for node in new_class_body])
   cls_str = cls_str.replace('jax.Array', 'np.ndarray')

mjx/mujoco/mjx/codegen/trace.py

@@ -63,6 +63,10 @@ def _get_imported_module_fpaths(fpath: epath.Path) -> Dict[str, str]:
 
   all_resolved_fpaths = {}
   for fully_qualified_name, alias in all_imported_names:
+    # only trace mujoco and mujoco warp
+    if not (fully_qualified_name.startswith('mujoco_warp') or
+            fully_qualified_name.startswith('mujoco')):
+      continue
     fpath = _resolve_module_name_to_fpath(fully_qualified_name)
     if fpath:
       all_resolved_fpaths[alias] = fpath
@@ -181,9 +185,17 @@ def visit_Call(self, node: ast.Call):
         self.add_field_usage(in_node, False)
         return
 
-      for arg in node.args:
-        if isinstance(arg, ast.Attribute):
-          self.add_field_usage(arg, is_output=True)
+      # do not trace array creation
+      is_input_only_wp_fn = (
+          len(parts) == 2
+          and parts[0] == 'wp'
+          and parts[1] in ('empty', 'zeros', 'array', 'clone')
+      )
+
+      if not is_input_only_wp_fn:
+        for arg in node.args:
+          if isinstance(arg, ast.Attribute):
+            self.add_field_usage(arg, is_output=True)
 
       called_fn_name = '.'.join(parts[1:])
       key = (hash(self._current_fpath), called_fn_name)

mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py

@@ -61,6 +61,7 @@
 from mujoco.mjx.third_party.mujoco_warp._src.io import set_const as set_const
 from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_0 as set_const_0
 from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_fixed as set_const_fixed
+from mujoco.mjx.third_party.mujoco_warp._src.io import set_const_spring as set_const_spring
 from mujoco.mjx.third_party.mujoco_warp._src.io import set_length_range as set_length_range
 from mujoco.mjx.third_party.mujoco_warp._src.island import island as island
 from mujoco.mjx.third_party.mujoco_warp._src.passive import passive as passive

mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py

@@ -25,7 +25,6 @@
 from etils import epath
 
 import mujoco.mjx.third_party.mujoco_warp as mjw
-from mujoco.mjx.third_party.mujoco_warp._src import io
 from mujoco.mjx.third_party.mujoco_warp._src import warp_util
 from mujoco.mjx.third_party.mujoco_warp._src.io import load_trajectory
 from mujoco.mjx.third_party.mujoco_warp._src.io import override_model
@@ -43,12 +42,12 @@
 EVENT_TRACE = flags.DEFINE_bool("event_trace", False, "print an event trace report")
 NOISE_STD = flags.DEFINE_float("noise_std", 0.01, "add noise to ctrl signal (standard deviation)")
 NOISE_RATE = flags.DEFINE_float("noise_rate", 0.1, "add noise to ctrl signal (noise rate)")
-ENABLE_ISLANDS = flags.DEFINE_bool(
-  "enable_islands",
-  False,
-  "Enable constraint islands solver",
-)
+NVMAX = flags.DEFINE_integer("nvmax", None, "maximum active DOFs per world")
+
 
+INIT_ASLEEP = flags.DEFINE_bool(
+  "init_asleep", False, "initialize all trees as asleep before simulation (requires sleep enabled)"
+)
 
 DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
 REPLAY = flags.DEFINE_string("replay", None, "NPZ file with ctrl sequence to replay")
@@ -141,8 +140,6 @@ def init_structs(
   fn: Callable[..., None], mjm: mujoco.MjModel
 ) -> Tuple[mjw.Model, mjw.Data, mjw.RenderContext | None, list[np.ndarray] | None]:
   """Initialize device structs."""
-  io.ENABLE_ISLANDS = ENABLE_ISLANDS.value
-
   mjd = mujoco.MjData(mjm)
   ctrls = None
   if REPLAY.value:
@@ -158,8 +155,18 @@ def init_structs(
     m = mjw.put_model(mjm)
     if OVERRIDE.value:
       override_model(m, OVERRIDE.value)
+    if INIT_ASLEEP.value:
+      mjd.tree_asleep[:] = np.arange(mjm.ntree, dtype=np.int32)
+
     d = mjw.put_data(
-      mjm, mjd, nworld=NWORLD.value, nconmax=NCONMAX.value, njmax=NJMAX.value, njmax_nnz=NJMAX_NNZ.value, nccdmax=NCCDMAX.value
+      mjm,
+      mjd,
+      nworld=NWORLD.value,
+      nconmax=NCONMAX.value,
+      njmax=NJMAX.value,
+      njmax_nnz=NJMAX_NNZ.value,
+      nccdmax=NCCDMAX.value,
+      nvmax=NVMAX.value,
     )
 
     if mjw.RenderContext not in get_type_hints(fn).values():