Allow nameless bind in MJX.
PiperOrigin-RevId: 742592221 Change-Id: I9490aebd64e3a31219c552d5b6eb55236680f10d
This commit is contained in:
Alessio Quaglino
committed by
Copybara-Service
Copybara-Service
parent
05f3e914b0
commit
ebd30493c8
@@ -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:
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user