diff --git a/doc/changelog.rst b/doc/changelog.rst index fd565e76..ad3f10cc 100644 --- a/doc/changelog.rst +++ b/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` + in the documentation. + Version 3.3.0 (Feb 26, 2025) ---------------------------- - Feature promotion ^^^^^^^^^^^^^^^^^ .. youtube:: qJFbx-FR7Bc diff --git a/doc/python.rst b/doc/python.rst index 62541701..85a54b73 100644 --- a/doc/python.rst +++ b/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 `__'s `PyMJCF `__ 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 -geoms containing the string "torso" in their name. We want to get their Cartesian positions in the XY plane from -``mjData``. This can be done as follows: +:ref:`mjModel` and :ref:`mjData`. For example, say we have multiple geoms containing the string "torso" in their name. +We want to get their Cartesian positions in the XY plane from ``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 ----- diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 591ee0ee..90c517ef 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -295,73 +295,56 @@ class BindModel(object): self.prefix = '' ids = [] for spec in specs: - if not spec.name: - raise KeyError(f'cannot bind spec with empty name') + if model.signature != spec.signature: + 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: - raise KeyError(f'invalid name: {spec.name}') # pytype: disable=attribute-error - ids.append(idx) + ids.append(spec.id) 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: - raise KeyError(f'cannot bind spec with empty name') + if model.signature != spec.signature: + 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: - raise KeyError(f'invalid name: {spec.name}') # pytype: disable=attribute-error - ids.append(idx) + ids.append(spec.id) if len(ids) == 1: self.id = ids[0] else: diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index 5948f9e1..1fd1971b 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -176,15 +176,15 @@ class SupportTest(parameterized.TestCase): - - - + + + - - - + + + """ @@ -309,29 +309,28 @@ class SupportTest(parameterized.TestCase): dx7.bind(mx, body).xfrc_applied, [0, 0, 0, 0, 0, 0] ) - # test invalid name - with self.assertRaises( - AttributeError, msg='ctrl is not available for this type' + # test attribute and type mismatches + with self.assertRaisesRegex( + AttributeError, 'ctrl is not available for this type' ): print(dx.bind(mx, s.geoms).ctrl) - with self.assertRaises( - KeyError, msg='actuator_actuator_ctrl' - ): + with self.assertRaises(KeyError): print(dx.bind(mx, s.actuators).actuator_ctrl) - with self.assertRaises( - AttributeError, msg='actuator_actuator_ctrl' + with self.assertRaisesRegex( + AttributeError, + "'Data' object has no attribute 'actuator_actuator_ctrl'", ): print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3])) - with self.assertRaises( - AttributeError, msg='qpos, qvel, qacc are not available for this type' + with self.assertRaisesRegex( + 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' - print(dx.bind(mx, s.actuators).set('ctrl', [1, 2, 3])) - with self.assertRaises(KeyError, msg='invalid name: invalid_geom_name'): - s.geoms[0].name = 'invalid_geom_name' - print(mx.bind(s.geoms).pos) + + # test that modified names do not raise an error + s.actuators[0].name = 'modified_actuator_name' + np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl) + s.geoms[0].name = 'modified_geom_name' + 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() - mx = mjx.put_model(m) - with self.assertRaises(KeyError, msg='cannot bind spec with empty name'): + with self.assertRaises(ValueError) as e: mx.bind(s.bodies) + self.assertEqual( + str(e.exception), + 'mjSpec signature does not match mjx.Model signature:' + ' 5495345807332648606 != 270010677651259353', + ) _CONTACTS = """ diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 228f3f3d..515d15f4 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/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 diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index e4617eab..bb52643d 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/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__':