Allow nameless bind in MJX.

PiperOrigin-RevId: 742592221
Change-Id: I9490aebd64e3a31219c552d5b6eb55236680f10d

doc/changelog.rst

@@ -23,10 +23,14 @@ Bug fixes
 - Fixed a bug that caused the parent frame of elements in the child worldbody to be incorrectly set when attaching an
   mjSpec to a frame or a site.
 
+Python bindings
+^^^^^^^^^^^^^^^
+- Added support for nameless :ref:`mjSpec` objects in the ``bind`` method, see the corresponding :ref:`section<PyMJCF>`
+  in the documentation.
+
 Version 3.3.0 (Feb 26, 2025)
 ----------------------------
 
-
 Feature promotion
 ^^^^^^^^^^^^^^^^^
 .. youtube:: qJFbx-FR7Bc

doc/python.rst

@@ -628,8 +628,10 @@ Parent:
   The parent body of a given element -- including bodies and frames -- can be accessed via the ``parent`` property.
   For example, the parent of a site can be accessed via ``site.parent``.
 
-Relationship to ``PyMJCF``
+.. _PyMJCF:
---------------------------
+
+Relationship to ``PyMJCF`` and ``bind``
+---------------------------------------
 
 `dm_control <https://github.com/google-deepmind/dm_control/tree/main>`__'s
 `PyMJCF <https://github.com/google-deepmind/dm_control/blob/main/dm_control/mjcf/README.md>`__ module provides similar
@@ -645,9 +647,8 @@ includes a reimplementation of the ``PyMJCF`` example in the ``dm_control``
 
 ``PyMJCF`` provides a notion of "binding", giving access to :ref:`mjModel` and :ref:`mjData` values via a helper class.
 In the native API, the helper class is not needed, so it is possible to directly bind an ``mjs`` object to
-:ref:`mjModel` and :ref:`mjData`. This requires the objects to have a non-empty name. For example, say we have multiple
+:ref:`mjModel` and :ref:`mjData`. For example, say we have multiple geoms containing the string "torso" in their name.
-geoms containing the string "torso" in their name. We want to get their Cartesian positions in the XY plane from
+We want to get their Cartesian positions in the XY plane from ``mjData``. This can be done as follows:
-``mjData``. This can be done as follows:
 
 .. code-block:: python
 
@@ -655,6 +656,9 @@ geoms containing the string "torso" in their name. We want to get their Cartesia
    pos_x = [torso.xpos[0] for torso in torsos]
    pos_y = [torso.xpos[1] for torso in torsos]
 
+Using the ``bind`` method requires the :ref:`mjModel` and :ref:`mjData` to be compiled from the :`ref:`mjSpec`. If
+objects are added or removed from the :ref:`mjSpec` since the last compilation, an error is raised.
+
 Notes
 -----
 

mjx/mujoco/mjx/_src/support.py

@@ -295,73 +295,56 @@ class BindModel(object):
     self.prefix = ''
     ids = []
     for spec in specs:
-      if not spec.name:
+      if model.signature != spec.signature:
-        raise KeyError(f'cannot bind spec with empty name')
+        raise ValueError(
+            'mjSpec signature does not match mjx.Model signature:'
+            f' {spec.signature} != {model.signature}'
+        )
+      elif spec.id < 0:
+        raise KeyError(f'invalid id: {spec.id}')
       elif isinstance(spec, mujoco.MjsBody):
         self.prefix = 'body_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_BODY, spec.name)
       elif isinstance(spec, mujoco.MjsJoint):
         self.prefix = 'jnt_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_JOINT, spec.name)
       elif isinstance(spec, mujoco.MjsGeom):
         self.prefix = 'geom_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_GEOM, spec.name)
       elif isinstance(spec, mujoco.MjsSite):
         self.prefix = 'site_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_SITE, spec.name)
       elif isinstance(spec, mujoco.MjsLight):
         self.prefix = 'light_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, spec.name)
       elif isinstance(spec, mujoco.MjsCamera):
         self.prefix = 'cam_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, spec.name)
       elif isinstance(spec, mujoco.MjsMesh):
         self.prefix = 'mesh_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_MESH, spec.name)
       elif isinstance(spec, mujoco.MjsHField):
         self.prefix = 'hfield_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_HFIELD, spec.name)
       elif isinstance(spec, mujoco.MjsPair):
         self.prefix = 'pair_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_PAIR, spec.name)
       elif isinstance(spec, mujoco.MjsTendon):
         self.prefix = 'tendon_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_TENDON, spec.name)
       elif isinstance(spec, mujoco.MjsActuator):
         self.prefix = 'actuator_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, spec.name)
       elif isinstance(spec, mujoco.MjsSensor):
         self.prefix = 'sensor_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, spec.name)
       elif isinstance(spec, mujoco.MjsNumeric):
         self.prefix = 'numeric_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_NUMERIC, spec.name)
       elif isinstance(spec, mujoco.MjsText):
         self.prefix = 'text_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_TEXT, spec.name)
       elif isinstance(spec, mujoco.MjsTuple):
         self.prefix = 'tuple_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_TUPLE, spec.name)
       elif isinstance(spec, mujoco.MjsKey):
         self.prefix = 'key_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_KEY, spec.name)
       elif isinstance(spec, mujoco.MjsEquality):
         self.prefix = 'eq_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, spec.name)
       elif isinstance(spec, mujoco.MjsExclude):
         self.prefix = 'exclude_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_EXCLUDE, spec.name)
       elif isinstance(spec, mujoco.MjsSkin):
         self.prefix = 'skin_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_SKIN, spec.name)
       elif isinstance(spec, mujoco.MjsMaterial):
         self.prefix = 'material_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_MATERIAL, spec.name)
       else:
         raise ValueError('invalid spec type')
-      if idx < 0:
+      ids.append(spec.id)
-        raise KeyError(f'invalid name: {spec.name}')  # pytype: disable=attribute-error
-      ids.append(idx)
     if len(ids) == 1:
       self.id = ids[0]
     else:
@@ -402,42 +385,36 @@ class BindData(object):
     self.prefix = ''
     ids = []
     for spec in specs:
-      if not spec.name:
+      if model.signature != spec.signature:
-        raise KeyError(f'cannot bind spec with empty name')
+        raise ValueError(
+            'mjSpec signature does not match mjx.Model signature:'
+            f' {spec.signature} != {model.signature}'
+        )
+      if spec.id < 0:
+        raise KeyError(f'invalid id: {spec.id}')
       elif isinstance(spec, mujoco.MjsBody):
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_BODY, spec.name)
+        pass
       elif isinstance(spec, mujoco.MjsJoint):
         self.prefix = 'jnt_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_JOINT, spec.name)
       elif isinstance(spec, mujoco.MjsGeom):
         self.prefix = 'geom_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_GEOM, spec.name)
       elif isinstance(spec, mujoco.MjsSite):
         self.prefix = 'site_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_SITE, spec.name)
       elif isinstance(spec, mujoco.MjsLight):
         self.prefix = 'light_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, spec.name)
       elif isinstance(spec, mujoco.MjsCamera):
         self.prefix = 'cam_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, spec.name)
       elif isinstance(spec, mujoco.MjsTendon):
         self.prefix = 'ten_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_TENDON, spec.name)
       elif isinstance(spec, mujoco.MjsActuator):
         self.prefix = 'actuator_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, spec.name)
       elif isinstance(spec, mujoco.MjsSensor):
         self.prefix = 'sensor_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, spec.name)
       elif isinstance(spec, mujoco.MjsEquality):
         self.prefix = 'eq_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, spec.name)
       else:
         raise ValueError('invalid spec type')
-      if idx < 0:
+      ids.append(spec.id)
-        raise KeyError(f'invalid name: {spec.name}')  # pytype: disable=attribute-error
-      ids.append(idx)
     if len(ids) == 1:
       self.id = ids[0]
     else:

mjx/mujoco/mjx/_src/support_test.py

@@ -176,15 +176,15 @@ class SupportTest(parameterized.TestCase):
         </worldbody>
 
         <actuator>
-          <motor name="actuator1" joint="joint1"/>
+          <motor joint="joint1"/>
-          <motor name="actuator2" joint="joint2"/>
+          <motor joint="joint2"/>
-          <motor name="actuator3" joint="joint3"/>
+          <motor joint="joint3"/>
         </actuator>
 
         <sensor>
-          <framepos name="sensor1" objtype="body" objname="body1"/>
+          <framepos objtype="body" objname="body1"/>
-          <framepos name="sensor2" objtype="body" objname="body2"/>
+          <framepos objtype="body" objname="body2"/>
-          <framepos name="sensor3" objtype="body" objname="body3"/>
+          <framepos objtype="body" objname="body3"/>
         </sensor>
     </mujoco>
     """
@@ -309,29 +309,28 @@ class SupportTest(parameterized.TestCase):
           dx7.bind(mx, body).xfrc_applied, [0, 0, 0, 0, 0, 0]
       )
 
-    # test invalid name
+    # test attribute and type mismatches
-    with self.assertRaises(
+    with self.assertRaisesRegex(
-        AttributeError, msg='ctrl is not available for this type'
+        AttributeError, 'ctrl is not available for this type'
     ):
       print(dx.bind(mx, s.geoms).ctrl)
-    with self.assertRaises(
+    with self.assertRaises(KeyError):
-        KeyError, msg='actuator_actuator_ctrl'
-    ):
       print(dx.bind(mx, s.actuators).actuator_ctrl)
-    with self.assertRaises(
+    with self.assertRaisesRegex(
-        AttributeError, msg='actuator_actuator_ctrl'
+        AttributeError,
+        "'Data' object has no attribute 'actuator_actuator_ctrl'",
     ):
       print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3]))
-    with self.assertRaises(
+    with self.assertRaisesRegex(
-        AttributeError, msg='qpos, qvel, qacc are not available for this type'
+        AttributeError, 'qpos, qvel, qacc are not available for this type'
     ):
       print(dx.bind(mx, s.geoms).qpos)
-    with self.assertRaises(KeyError, msg='invalid name: invalid_actuator_name'):
+
-      s.actuators[0].name = 'invalid_actuator_name'
+    # test that modified names do not raise an error
-      print(dx.bind(mx, s.actuators).set('ctrl', [1, 2, 3]))
+    s.actuators[0].name = 'modified_actuator_name'
-    with self.assertRaises(KeyError, msg='invalid name: invalid_geom_name'):
+    np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl)
-      s.geoms[0].name = 'invalid_geom_name'
+    s.geoms[0].name = 'modified_geom_name'
-      print(mx.bind(s.geoms).pos)
+    np.testing.assert_array_equal(mx.bind(s.geoms[0]).pos, m.geom_pos[0, :])
 
     # test batched data
     batch_size = 16
@@ -343,12 +342,15 @@ class SupportTest(parameterized.TestCase):
           vdx.bind(mx, s.bodies[i]).xpos, [d.xpos[i, :]] * batch_size
       )
 
-    # test emtpy name
+    # test that adding a body requires recompilation
     s.worldbody.add_body()
-    m = s.compile()
+    with self.assertRaises(ValueError) as e:
-    mx = mjx.put_model(m)
-    with self.assertRaises(KeyError, msg='cannot bind spec with empty name'):
       mx.bind(s.bodies)
+    self.assertEqual(
+        str(e.exception),
+        'mjSpec signature does not match mjx.Model signature:'
+        ' 5495345807332648606 != 270010677651259353',
+    )
 
   _CONTACTS = """
     <mujoco>

mjx/mujoco/mjx/_src/types.py

@@ -844,6 +844,7 @@ class Model(PyTreeNode):
     name_tupleadr: tuple name pointers                        (ntuple,)
     name_keyadr: keyframe name pointers                       (nkey,)
     names: names of all objects, 0-terminated                 (nnames,)
+    signature: compilation signature
   """
 
   nq: int
@@ -1187,6 +1188,7 @@ class Model(PyTreeNode):
   name_tupleadr: np.ndarray
   name_keyadr: np.ndarray
   names: bytes
+  signature: np.uint64
   _sizes: jax.Array
 
 

python/mujoco/codegen/generate_spec_bindings.py

@@ -627,6 +627,22 @@ def generate_signature() -> None:
     print(code)
 
 
+def generate_id() -> None:
+  """Generate id functions."""
+  for key, _, _, _, _ in SPECS:
+    if key == 'mjsPlugin':
+      continue
+    elem = key.removeprefix('mjs')
+    titlecase = 'Mjs' + elem
+    code = f"""\n
+      {key}.def_property_readonly("id",
+      [](raw::{titlecase}& self) -> int {{
+        return mjs_getId(self.element);
+      }});
+    """
+    print(code)
+
+
 def main(argv: Sequence[str]) -> None:
   if len(argv) > 1:
     raise app.UsageError('Too many command-line arguments.')
@@ -634,6 +650,7 @@ def main(argv: Sequence[str]) -> None:
   generate_add()
   generate_find()
   generate_signature()
+  generate_id()
 
 
 if __name__ == '__main__':