Add mjSpec binding to MJX.
PiperOrigin-RevId: 714026500 Change-Id: I5cb3fe6cc975b3aa43d44185e3bc186b45845cf3
This commit is contained in:
Alessio Quaglino
committed by
Copybara-Service
Copybara-Service
parent
6f7fba794b
commit
3f32cc2479
@@ -13,7 +13,8 @@
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
"""Engine support functions."""
|
||||
from typing import Optional, Tuple, Union
|
||||
from collections.abc import Sequence
|
||||
from typing import Optional, Tuple, Union, Any
|
||||
|
||||
import jax
|
||||
from jax import numpy as jp
|
||||
@@ -236,7 +237,7 @@ def _getadr(
|
||||
def id2name(
|
||||
m: Union[Model, mujoco.MjModel], typ: mujoco._enums.mjtObj, i: int
|
||||
) -> Optional[str]:
|
||||
"""Gets the name of an object with the specified mjtObj type and id.
|
||||
"""Gets the name of an object with the specified mjtObj type and ids.
|
||||
|
||||
See mujoco.id2name for more info.
|
||||
|
||||
@@ -284,6 +285,174 @@ def name2id(
|
||||
return names_map.get(name, -1)
|
||||
|
||||
|
||||
class BindModel(object):
|
||||
"""Class holding the requested MJX Model and spec id for binding a spec to Model."""
|
||||
|
||||
def __init__(self, model: Model, specs: Sequence[Any]):
|
||||
self.model = model
|
||||
try:
|
||||
iter(specs)
|
||||
except TypeError:
|
||||
specs = [specs]
|
||||
ids = []
|
||||
for spec in specs:
|
||||
match spec:
|
||||
case mujoco.MjsBody():
|
||||
self.prefix = 'body_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_BODY, spec.name))
|
||||
case mujoco.MjsJoint():
|
||||
self.prefix = 'jnt_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_JOINT, spec.name))
|
||||
case mujoco.MjsGeom():
|
||||
self.prefix = 'geom_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_GEOM, spec.name))
|
||||
case mujoco.MjsSite():
|
||||
self.prefix = 'site_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_SITE, spec.name))
|
||||
case mujoco.MjsLight():
|
||||
self.prefix = 'light_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, spec.name))
|
||||
case mujoco.MjsCamera():
|
||||
self.prefix = 'cam_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, spec.name))
|
||||
case mujoco.MjsMesh():
|
||||
self.prefix = 'mesh_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_MESH, spec.name))
|
||||
case mujoco.MjsHfield():
|
||||
self.prefix = 'hfield_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_HFIELD, spec.name))
|
||||
case mujoco.MjsPair():
|
||||
self.prefix = 'pair_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_PAIR, spec.name))
|
||||
case mujoco.MjsTendon():
|
||||
self.prefix = 'tendon_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_TENDON, spec.name))
|
||||
case mujoco.MjsActuator():
|
||||
self.prefix = 'actuator_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, spec.name))
|
||||
case mujoco.MjsSensor():
|
||||
self.prefix = 'sensor_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, spec.name))
|
||||
case mujoco.MjsNumeric():
|
||||
self.prefix = 'numeric_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_NUMERIC, spec.name))
|
||||
case mujoco.MjsText():
|
||||
self.prefix = 'text_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_TEXT, spec.name))
|
||||
case mujoco.MjsTuple():
|
||||
self.prefix = 'tuple_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_TUPLE, spec.name))
|
||||
case mujoco.MjsKey():
|
||||
self.prefix = 'key_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_KEY, spec.name))
|
||||
case mujoco.MjsEquality():
|
||||
self.prefix = 'eq_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, spec.name))
|
||||
case mujoco.MjsExclude():
|
||||
self.prefix = 'exclude_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_EXCLUDE, spec.name))
|
||||
case mujoco.MjsSkin():
|
||||
self.prefix = 'skin_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_SKIN, spec.name))
|
||||
case mujoco.MjsMaterial():
|
||||
self.prefix = 'material_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_MATERIAL, spec.name))
|
||||
case _:
|
||||
raise ValueError('invalid spec type')
|
||||
if len(ids) == 1:
|
||||
self.id = ids[0]
|
||||
else:
|
||||
self.id = ids
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
return getattr(self.model, self.prefix + name)[self.id, :]
|
||||
|
||||
|
||||
def _bind_model(self: Model, obj: Sequence[Any]) -> BindModel:
|
||||
"""Bind a Mujoco spec to an MJX Model."""
|
||||
return BindModel(self, obj)
|
||||
|
||||
|
||||
class BindData(object):
|
||||
"""Class holding the requested MJX Data and spec id for binding a spec to Data."""
|
||||
|
||||
def __init__(self, data: Data, model: Model, specs: Sequence[Any]):
|
||||
self.data = data
|
||||
try:
|
||||
iter(specs)
|
||||
except TypeError:
|
||||
specs = [specs]
|
||||
ids = []
|
||||
for spec in specs:
|
||||
match spec:
|
||||
case mujoco.MjsBody():
|
||||
self.prefix = ''
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_BODY, spec.name))
|
||||
case mujoco.MjsJoint():
|
||||
self.prefix = 'jnt_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_JOINT, spec.name))
|
||||
case mujoco.MjsGeom():
|
||||
self.prefix = 'geom_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_GEOM, spec.name))
|
||||
case mujoco.MjsSite():
|
||||
self.prefix = 'site_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_SITE, spec.name))
|
||||
case mujoco.MjsLight():
|
||||
self.prefix = 'light_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, spec.name))
|
||||
case mujoco.MjsCamera():
|
||||
self.prefix = 'cam_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, spec.name))
|
||||
case mujoco.MjsTendon():
|
||||
self.prefix = 'tendon_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_TENDON, spec.name))
|
||||
case mujoco.MjsActuator():
|
||||
self.prefix = 'actuator_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, spec.name))
|
||||
case mujoco.MjsSensor():
|
||||
self.prefix = 'sensor_'
|
||||
ids.append(name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, spec.name))
|
||||
case _:
|
||||
raise ValueError('invalid spec type')
|
||||
if len(ids) == 1:
|
||||
self.id = ids[0]
|
||||
else:
|
||||
self.id = ids
|
||||
|
||||
def __getname(self, name: str):
|
||||
try:
|
||||
getattr(self.data, self.prefix + name)
|
||||
return self.prefix + name
|
||||
except AttributeError:
|
||||
try:
|
||||
getattr(self.data, name)
|
||||
return name
|
||||
except AttributeError as e:
|
||||
raise ValueError(f'invalid name: {name}') from e
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
return getattr(self.data, self.__getname(name))[self.id, ...]
|
||||
|
||||
def set(self, name: str, value: jax.Array) -> Data:
|
||||
"""Set the value of an array in an MJX Data."""
|
||||
array = getattr(self.data, self.__getname(name))
|
||||
if len(value) == 1:
|
||||
array = array.at[self.id].set(value[0])
|
||||
else:
|
||||
for i, v in enumerate(value):
|
||||
array = array.at[self.id[i]].set(v)
|
||||
return self.data.replace(**{self.__getname(name): array})
|
||||
|
||||
|
||||
def _bind_data(self: Data, model: Model, obj: Sequence[Any]) -> BindData:
|
||||
"""Bind a Mujoco spec to an MJX Data."""
|
||||
return BindData(self, model, obj)
|
||||
|
||||
|
||||
Model.bind = _bind_model
|
||||
Data.bind = _bind_data
|
||||
|
||||
|
||||
def _decode_pyramid(
|
||||
pyramid: jax.Array, mu: jax.Array, condim: int
|
||||
) -> jax.Array:
|
||||
|
||||
@@ -157,6 +157,92 @@ class SupportTest(parameterized.TestCase):
|
||||
i = i if n is not None else -1
|
||||
self.assertEqual(support.name2id(mx, obj, n), i)
|
||||
|
||||
def test_bind(self):
|
||||
xml = """
|
||||
<mujoco model="test_bind_model">
|
||||
<worldbody>
|
||||
<body pos="1 2 3" name="body1">
|
||||
<joint axis="1 0 0" type="slide" name="joint1"/>
|
||||
<geom size="1 2 3" type="box" name="geom1"/>
|
||||
</body>
|
||||
<body pos="4 5 6" name="body2">
|
||||
<joint axis="0 1 0" type="slide" name="joint2"/>
|
||||
<geom size="4 5 6" type="box" name="geom2"/>
|
||||
</body>
|
||||
<body pos="7 8 9" name="body3">
|
||||
<joint axis="0 0 1" type="slide" name="joint3"/>
|
||||
<geom size="7 8 9" type="box" name="geom3"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
|
||||
<actuator>
|
||||
<motor name="actuator1" joint="joint1"/>
|
||||
<motor name="actuator2" joint="joint2"/>
|
||||
<motor name="actuator3" joint="joint3"/>
|
||||
</actuator>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
s = mujoco.MjSpec.from_string(xml)
|
||||
m = s.compile()
|
||||
d = mujoco.MjData(m)
|
||||
mx = mjx.put_model(m)
|
||||
dx = mjx.put_data(m, d)
|
||||
mujoco.mj_step(m, d)
|
||||
dx = mjx.step(mx, dx)
|
||||
|
||||
# test getting
|
||||
np.testing.assert_array_equal(mx.bind(s.bodies).pos, m.body_pos)
|
||||
np.testing.assert_array_equal(dx.bind(mx, s.bodies).xpos, d.xpos)
|
||||
for i in range(m.nbody):
|
||||
np.testing.assert_array_equal(m.bind(s.bodies[i]).pos, m.body_pos[i, :])
|
||||
np.testing.assert_array_equal(mx.bind(s.bodies[i]).pos, m.body_pos[i, :])
|
||||
np.testing.assert_array_equal(d.bind(s.bodies[i]).xpos, d.xpos[i, :])
|
||||
np.testing.assert_array_equal(
|
||||
dx.bind(mx, s.bodies[i]).xpos, d.xpos[i, :]
|
||||
)
|
||||
|
||||
np.testing.assert_array_equal(mx.bind(s.geoms).size, m.geom_size)
|
||||
np.testing.assert_array_equal(dx.bind(mx, s.geoms).xpos, d.geom_xpos)
|
||||
for i in range(m.ngeom):
|
||||
np.testing.assert_array_equal(m.bind(s.geoms[i]).size, m.geom_size[i, :])
|
||||
np.testing.assert_array_equal(mx.bind(s.geoms[i]).size, m.geom_size[i, :])
|
||||
np.testing.assert_array_equal(d.bind(s.geoms[i]).xpos, d.geom_xpos[i, :])
|
||||
np.testing.assert_array_equal(
|
||||
dx.bind(mx, s.geoms[i]).xpos, d.geom_xpos[i, :]
|
||||
)
|
||||
|
||||
np.testing.assert_array_equal(mx.bind(s.joints).axis, m.jnt_axis)
|
||||
for i in range(m.njnt):
|
||||
np.testing.assert_array_equal(m.bind(s.joints[i]).axis, m.jnt_axis[i, :])
|
||||
np.testing.assert_array_equal(mx.bind(s.joints[i]).axis, m.jnt_axis[i, :])
|
||||
|
||||
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl)
|
||||
for i in range(m.nu):
|
||||
np.testing.assert_array_equal(d.bind(s.actuators[i]).ctrl, d.ctrl[i])
|
||||
np.testing.assert_array_equal(
|
||||
dx.bind(mx, s.actuators[i]).ctrl, d.ctrl[i]
|
||||
)
|
||||
|
||||
# test setting
|
||||
np.testing.assert_array_equal(d.ctrl, [0, 0, 0])
|
||||
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl)
|
||||
dx2 = dx.bind(mx, s.actuators).set('ctrl', [1, 2, 3])
|
||||
np.testing.assert_array_equal(dx2.bind(mx, s.actuators).ctrl, [1, 2, 3])
|
||||
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, [0, 0, 0])
|
||||
dx3 = dx.bind(mx, s.actuators[1:]).set('ctrl', [4, 5])
|
||||
np.testing.assert_array_equal(dx3.bind(mx, s.actuators).ctrl, [0, 4, 5])
|
||||
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, [0, 0, 0])
|
||||
dx4 = dx.bind(mx, s.actuators[1]).set('ctrl', [6])
|
||||
np.testing.assert_array_equal(dx4.bind(mx, s.actuators).ctrl, [0, 6, 0])
|
||||
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, [0, 0, 0])
|
||||
|
||||
# test invalid name
|
||||
with self.assertRaises(ValueError):
|
||||
print(dx.bind(mx, s.actuators).actuator_ctrl)
|
||||
with self.assertRaises(ValueError):
|
||||
print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3]))
|
||||
|
||||
_CONTACTS = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
|
||||
Reference in New Issue
Block a user