Empty out Model fields restricted to MuJoCo in the same way we do for Data fields.

PiperOrigin-RevId: 719705508
Change-Id: I4b07433481eafaea1d95a5bdd8fd83786939188f

mjx/mujoco/mjx/_src/io.py

@@ -39,7 +39,7 @@ def _strip_weak_type(tree):
 
 
 def _make_option(
-    o: mujoco.MjOption, _full_compat: bool = False
+    o: mujoco.MjOption, _full_compat: bool = False  # pylint: disable=invalid-name
 ) -> types.Option:
   """Returns mjx.Option given mujoco.MjOption."""
   if not _full_compat:
@@ -183,6 +183,15 @@ def put_model(
       if f.metadata.get('restricted_to') != 'mjx'
   }
   fields = {f: getattr(m, f) for f in mj_field_names}
+
+  # zero out fields restricted to MuJoCo
+  if not _full_compat:
+    for f in types.Model.fields():
+      if f.metadata.get('restricted_to') == 'mujoco' and isinstance(
+          fields[f.name], np.ndarray
+      ):
+        fields[f.name] = np.zeros((0,), dtype=fields[f.name].dtype)
+
   fields['dof_hasfrictionloss'] = fields['dof_frictionloss'] > 0
   fields['tendon_hasfrictionloss'] = fields['tendon_frictionloss'] > 0
   fields['geom_rbound_hfield'] = fields['geom_rbound']
@@ -522,7 +531,7 @@ def _make_contact(
   # if we have fewer Contacts for a condim range, pad the range with zeros
 
   # build a map for where to find a dim-matching contact, or -1 if none
-  contact_map = np.zeros_like(dim) - 1
+  contact_map = -np.ones_like(dim)
   for i, di in enumerate(fields['dim']):
     space = [j for j, dj in enumerate(dim) if di == dj and contact_map[j] == -1]
     if not space:
@@ -672,7 +681,9 @@ def put_data(
     fields['_qLDiagInv_sparse'] = jp.zeros(0, dtype=float)
     # otherwise clear out unused arrays
     for f in types.Data.fields():
-      if f.metadata.get('restricted_to') == 'mujoco':
+      if f.metadata.get('restricted_to') == 'mujoco' and isinstance(
+          fields[f.name], np.ndarray
+      ):
         fields[f.name] = np.zeros(0, dtype=fields[f.name].dtype)
 
   fields['contact'] = contact

mjx/mujoco/mjx/_src/io_test.py

@@ -20,7 +20,9 @@ import jax
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
+# pylint: disable=g-importing-member
 from mujoco.mjx._src.types import ConeType
+# pylint: enable=g-importing-member
 import numpy as np
 
 
@@ -117,6 +119,9 @@ class ModelIOTest(parameterized.TestCase):
     self.assertEqual(mx.nM, m.nM)
     self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep)
 
+    # fields restricted to MuJoCo should not be populated
+    self.assertEqual(mx.bvh_aabb.shape, (0,))
+
     np.testing.assert_allclose(mx.body_parentid, m.body_parentid)
     np.testing.assert_allclose(mx.geom_type, m.geom_type)
     np.testing.assert_allclose(mx.geom_bodyid, m.geom_bodyid)