Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 941234486
Change-Id: Ie5d8cc0c9975e66f9a81bb6a86e33f640a5d19e3

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 @@ _NESTED_DATACLASS_MANUAL_METHODS = {
 
 
 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 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
         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)