chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allow nameless bind in MJX.

Browse Source 
PiperOrigin-RevId: 742592221
Change-Id: I9490aebd64e3a31219c552d5b6eb55236680f10d
This commit is contained in:
Alessio Quaglino
2025-04-01 01:44:42 -07:00
committed by Copybara-Service
parent 05f3e914b0
commit ebd30493c8
6 changed files with 78 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+5 -1
View File
@@ -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
+9 -5
View File
@@ -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
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
-----
mjx/mujoco/mjx/_src/support.py
+17 -40
View File
@@ -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:
mjx/mujoco/mjx/_src/support_test.py
+28 -26
View File
@@ -176,15 +176,15 @@ class SupportTest(parameterized.TestCase):
</worldbody>
<actuator>
<motor name="actuator1" joint="joint1"/>
<motor name="actuator2" joint="joint2"/>
<motor name="actuator3" joint="joint3"/>
<motor joint="joint1"/>
<motor joint="joint2"/>
<motor joint="joint3"/>
</actuator>
<sensor>
<framepos name="sensor1" objtype="body" objname="body1"/>
<framepos name="sensor2" objtype="body" objname="body2"/>
<framepos name="sensor3" objtype="body" objname="body3"/>
<framepos objtype="body" objname="body1"/>
<framepos objtype="body" objname="body2"/>
<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
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 = """
<mujoco>
mjx/mujoco/mjx/_src/types.py
+2
View File
@@ -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
+17
View File
@@ -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__':