Update MJX rne_postconstraint with contributions from equality connect and equality weld constraints.
PiperOrigin-RevId: 777491801 Change-Id: Ib32dda4178def7a96d31e91787f3d0866a56cd76
This commit is contained in:
Taylor Howell
committed by
Copybara-Service
Copybara-Service
parent
57b3d6d0f4
commit
2ce5caa239
@@ -248,23 +248,6 @@ def _put_model_jax(
|
||||
if t == mujoco.mjtGeom.mjGEOM_MESH:
|
||||
mesh_geomid.add(g)
|
||||
|
||||
# check for unsupported sensor and equality constraint combinations
|
||||
sensor_rne_postconstraint = (
|
||||
np.any(m.sensor_type == types.SensorType.ACCELEROMETER)
|
||||
| np.any(m.sensor_type == types.SensorType.FORCE)
|
||||
| np.any(m.sensor_type == types.SensorType.TORQUE)
|
||||
| np.any(m.sensor_type == types.SensorType.FRAMELINACC)
|
||||
| np.any(m.sensor_type == types.SensorType.FRAMEANGACC)
|
||||
)
|
||||
eq_connect_weld = np.any(m.eq_type == types.EqType.CONNECT) | np.any(
|
||||
m.eq_type == types.EqType.WELD
|
||||
)
|
||||
if sensor_rne_postconstraint and eq_connect_weld:
|
||||
raise NotImplementedError(
|
||||
'rne_postconstraint not implemented with equality constraints:'
|
||||
' connect, weld.'
|
||||
)
|
||||
|
||||
for enum_field, enum_type, mj_type in (
|
||||
(m.actuator_biastype, types.BiasType, mujoco.mjtBias),
|
||||
(m.actuator_dyntype, types.DynType, mujoco.mjtDyn),
|
||||
|
||||
@@ -625,38 +625,6 @@ class DataIOTest(parameterized.TestCase):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
mjx.make_data(m)
|
||||
|
||||
@parameterized.product(
|
||||
sensor=['accelerometer', 'force', 'torque'], equality=['connect', 'weld']
|
||||
)
|
||||
def test_sensor_constraint_compatibility(self, sensor, equality):
|
||||
"""Test unsupported sensor and equality constraint combinations."""
|
||||
equality_constraint = f'{equality} body1="body1" body2="body2"'
|
||||
if equality == 'connect':
|
||||
equality_constraint += ' anchor="0 0 0"'
|
||||
m = mujoco.MjModel.from_xml_string(f"""
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="body1">
|
||||
<freejoint/>
|
||||
<geom size="0.1"/>
|
||||
<site name="site1"/>
|
||||
</body>
|
||||
<body name="body2">
|
||||
<freejoint/>
|
||||
<geom size="0.1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<equality>
|
||||
<{equality_constraint}/>
|
||||
</equality>
|
||||
<sensor>
|
||||
<{sensor} site="site1"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
""")
|
||||
with self.assertRaises(NotImplementedError):
|
||||
mjx.put_model(m, impl='jax')
|
||||
|
||||
@parameterized.parameters(JacobianType.DENSE, JacobianType.SPARSE)
|
||||
def test_qm_mapm2m(self, jacobian):
|
||||
"""Test that qM is mapped to M."""
|
||||
|
||||
@@ -29,6 +29,7 @@ from mujoco.mjx._src.types import EqType
|
||||
from mujoco.mjx._src.types import JointType
|
||||
from mujoco.mjx._src.types import Model
|
||||
from mujoco.mjx._src.types import ModelJAX
|
||||
from mujoco.mjx._src.types import ObjType
|
||||
from mujoco.mjx._src.types import TrnType
|
||||
from mujoco.mjx._src.types import WrapType
|
||||
# pylint: enable=g-importing-member
|
||||
@@ -659,11 +660,126 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
|
||||
cfrc_contact.reshape((-1, 6))
|
||||
)
|
||||
|
||||
# TODO(taylorhowell): connect and weld constraints
|
||||
if np.any(m.eq_type == EqType.CONNECT):
|
||||
raise NotImplementedError('Connect constraints are not implemented.')
|
||||
if np.any(m.eq_type == EqType.WELD):
|
||||
raise NotImplementedError('Weld constraints are not implemented.')
|
||||
# cfrc_ext += connect, weld
|
||||
cfrc_ext_equality = []
|
||||
cfrc_ext_equality_adr = []
|
||||
|
||||
connect_id = m.eq_type == EqType.CONNECT
|
||||
nconnect = connect_id.sum()
|
||||
|
||||
if nconnect:
|
||||
cfrc_connect_force = d._impl.efc_force[: 3 * nconnect].reshape(
|
||||
(nconnect, 3)
|
||||
)
|
||||
|
||||
is_site = m.eq_objtype == ObjType.SITE
|
||||
body1id = np.copy(m.eq_obj1id)
|
||||
body2id = np.copy(m.eq_obj2id)
|
||||
pos1 = m.eq_data[:, :3]
|
||||
pos2 = m.eq_data[:, 3:6]
|
||||
|
||||
if m.nsite:
|
||||
body1id[is_site] = m.site_bodyid[m.eq_obj1id[is_site]]
|
||||
body2id[is_site] = m.site_bodyid[m.eq_obj2id[is_site]]
|
||||
pos1 = jp.where(is_site[:, None], m.site_pos[m.eq_obj1id], pos1)
|
||||
pos2 = jp.where(is_site[:, None], m.site_pos[m.eq_obj2id], pos2)
|
||||
|
||||
# body 1
|
||||
k1_connect = body1id[connect_id]
|
||||
k1_connect_mask = k1_connect != 0
|
||||
offset1_connect = pos1[connect_id]
|
||||
|
||||
pos1_connect = jax.vmap(lambda pnt, mat, vec: mat @ pnt + vec)(
|
||||
offset1_connect, d.xmat[k1_connect], d.xpos[k1_connect]
|
||||
)
|
||||
subtree_com1_connect = d.subtree_com[jp.array(m.body_rootid)[k1_connect]]
|
||||
cfrc_com1_connect = jax.vmap(
|
||||
lambda dif, frc, mask: mask * jp.concatenate([-jp.cross(dif, frc), frc])
|
||||
)(subtree_com1_connect - pos1_connect, cfrc_connect_force, k1_connect_mask)
|
||||
|
||||
# body 2
|
||||
k2_connect = body2id[connect_id]
|
||||
k2_connect_mask = -1 * (k2_connect != 0)
|
||||
offset2_connect = pos2[connect_id]
|
||||
|
||||
pos2_connect = jax.vmap(lambda pnt, mat, vec: mat @ pnt + vec)(
|
||||
offset2_connect, d.xmat[k2_connect], d.xpos[k2_connect]
|
||||
)
|
||||
subtree_com2_connect = d.subtree_com[jp.array(m.body_rootid)[k2_connect]]
|
||||
cfrc_com2_connect = jax.vmap(
|
||||
lambda dif, frc, mask: mask * jp.concatenate([-jp.cross(dif, frc), frc])
|
||||
)(subtree_com2_connect - pos2_connect, cfrc_connect_force, k2_connect_mask)
|
||||
|
||||
cfrc_ext_equality.append(jp.vstack([cfrc_com1_connect, cfrc_com2_connect]))
|
||||
cfrc_ext_equality_adr.append(jp.concatenate([k1_connect, k2_connect]))
|
||||
|
||||
weld_id = m.eq_type == EqType.WELD
|
||||
nweld = weld_id.sum()
|
||||
|
||||
if nweld:
|
||||
cfrc_weld = d._impl.efc_force[
|
||||
3 * nconnect : 3 * nconnect + 6 * nweld
|
||||
].reshape((nweld, 6))
|
||||
cfrc_weld_force = cfrc_weld[:, :3]
|
||||
cfrc_weld_torque = cfrc_weld[:, 3:]
|
||||
|
||||
is_site = m.eq_objtype == ObjType.SITE
|
||||
body1id = np.copy(m.eq_obj1id)
|
||||
body2id = np.copy(m.eq_obj2id)
|
||||
pos1 = m.eq_data[:, 3:6]
|
||||
pos2 = m.eq_data[:, :3]
|
||||
|
||||
if m.nsite:
|
||||
body1id[is_site] = m.site_bodyid[m.eq_obj1id[is_site]]
|
||||
body2id[is_site] = m.site_bodyid[m.eq_obj2id[is_site]]
|
||||
pos1 = jp.where(is_site[:, None], m.site_pos[m.eq_obj1id], pos1)
|
||||
pos2 = jp.where(is_site[:, None], m.site_pos[m.eq_obj2id], pos2)
|
||||
|
||||
# body 1
|
||||
k1_weld = body1id[weld_id]
|
||||
k1_weld_mask = k1_weld != 0
|
||||
offset1_weld = pos1[weld_id]
|
||||
|
||||
pos1_weld = jax.vmap(lambda pnt, mat, vec: mat @ pnt + vec)(
|
||||
offset1_weld, d.xmat[k1_weld], d.xpos[k1_weld]
|
||||
)
|
||||
subtree_com1_weld = d.subtree_com[jp.array(m.body_rootid)[k1_weld]]
|
||||
cfrc_com1_weld = jax.vmap(
|
||||
lambda dif, frc, trq, mask: mask
|
||||
* jp.concatenate([trq - jp.cross(dif, frc), frc])
|
||||
)(
|
||||
subtree_com1_weld - pos1_weld,
|
||||
cfrc_weld_force,
|
||||
cfrc_weld_torque,
|
||||
k1_weld_mask,
|
||||
)
|
||||
|
||||
# body 2
|
||||
k2_weld = body2id[weld_id]
|
||||
k2_weld_mask = -1 * (k2_weld != 0)
|
||||
offset2_weld = pos2[weld_id]
|
||||
|
||||
pos2_weld = jax.vmap(lambda pnt, mat, vec: mat @ pnt + vec)(
|
||||
offset2_weld, d.xmat[k2_weld], d.xpos[k2_weld]
|
||||
)
|
||||
subtree_com2_weld = d.subtree_com[jp.array(m.body_rootid)[k2_weld]]
|
||||
cfrc_com2_weld = jax.vmap(
|
||||
lambda dif, frc, trq, mask: mask
|
||||
* jp.concatenate([trq - jp.cross(dif, frc), frc])
|
||||
)(
|
||||
subtree_com2_weld - pos2_weld,
|
||||
cfrc_weld_force,
|
||||
cfrc_weld_torque,
|
||||
k2_weld_mask,
|
||||
)
|
||||
|
||||
cfrc_ext_equality.append(jp.vstack([cfrc_com1_weld, cfrc_com2_weld]))
|
||||
cfrc_ext_equality_adr.append(jp.concatenate([k1_weld, k2_weld]))
|
||||
|
||||
if nconnect or nweld:
|
||||
cfrc_ext = cfrc_ext.at[jp.concatenate(cfrc_ext_equality_adr)].add(
|
||||
jp.vstack(cfrc_ext_equality)
|
||||
)
|
||||
|
||||
# forward pass over bodies: compute cacc, cfrc_int
|
||||
def _forward(carry, cfrc_ext, cinert, cvel, body_dofadr, body_dofnum):
|
||||
|
||||
@@ -230,21 +230,98 @@ class SmoothTest(absltest.TestCase):
|
||||
|
||||
|
||||
class RnePostConstraintTest(parameterized.TestCase):
|
||||
_CONNECT_SITE = """
|
||||
<equality>
|
||||
<connect site1="site1" site2="site2"/>
|
||||
</equality>
|
||||
"""
|
||||
_CONNECT_BODY = """
|
||||
<equality>
|
||||
<connect body1="body1" body2="body2" anchor="1 2 3"/>
|
||||
</equality>
|
||||
"""
|
||||
_WELD_SITE = """
|
||||
<equality>
|
||||
<weld site1="site1" site2="site2"/>
|
||||
</equality>
|
||||
"""
|
||||
_WELD_BODY = """
|
||||
<equality>
|
||||
<weld body1="body1" body2="body2"/>
|
||||
</equality>
|
||||
"""
|
||||
_CONNECT_SITE_WELD_SITE = """
|
||||
<equality>
|
||||
<connect site1="site1" site2="site2"/>
|
||||
<weld site1="site1" site2="site2"/>
|
||||
</equality>
|
||||
"""
|
||||
_WELD_SITE_CONNECT_SITE = """
|
||||
<equality>
|
||||
<weld site1="site1" site2="site2"/>
|
||||
<connect site1="site1" site2="site2"/>
|
||||
</equality>
|
||||
"""
|
||||
_WELD_SITE_CONNECT_SITE_WELD_BODY = """
|
||||
<equality>
|
||||
<weld site1="site1" site2="site2"/>
|
||||
<connect site1="site1" site2="site2"/>
|
||||
<weld body1="body1" body2="body2"/>
|
||||
</equality>
|
||||
"""
|
||||
_CONNECT_SITE_WELD_SITE_WELD_BODY = """
|
||||
<equality>
|
||||
<connect site1="site1" site2="site2"/>
|
||||
<weld site1="site1" site2="site2"/>
|
||||
<weld body1="body1" body2="body2"/>
|
||||
</equality>
|
||||
"""
|
||||
_CONNECT_SITE_CONNECT_BODY_CONNECT_WELD = """
|
||||
<equality>
|
||||
<connect site1="site1" site2="site2"/>
|
||||
<connect body1="body1" body2="body2" anchor="1 2 3"/>
|
||||
<weld body1="body1" body2="body2"/>
|
||||
</equality>
|
||||
"""
|
||||
|
||||
@parameterized.parameters(ConeType)
|
||||
def test_rnepostconstraint(self, cone_type):
|
||||
@parameterized.parameters(
|
||||
('', ConeType.PYRAMIDAL, None),
|
||||
('', ConeType.ELLIPTIC, None),
|
||||
(_CONNECT_SITE, ConeType.PYRAMIDAL, None),
|
||||
(_CONNECT_BODY, ConeType.PYRAMIDAL, None),
|
||||
(_WELD_SITE, ConeType.PYRAMIDAL, None),
|
||||
(_WELD_BODY, ConeType.PYRAMIDAL, None),
|
||||
(_CONNECT_SITE_WELD_SITE, ConeType.PYRAMIDAL, None),
|
||||
(
|
||||
_WELD_SITE_CONNECT_SITE,
|
||||
ConeType.PYRAMIDAL,
|
||||
np.array([6, 7, 8, 0, 1, 2, 3, 4, 5]),
|
||||
),
|
||||
(
|
||||
_WELD_SITE_CONNECT_SITE_WELD_BODY,
|
||||
ConeType.PYRAMIDAL,
|
||||
np.array([6, 7, 8, 0, 1, 2, 3, 4, 5]),
|
||||
),
|
||||
(_CONNECT_SITE_WELD_SITE_WELD_BODY, ConeType.PYRAMIDAL, None),
|
||||
(_CONNECT_SITE_CONNECT_BODY_CONNECT_WELD, ConeType.PYRAMIDAL, None),
|
||||
)
|
||||
def test_rnepostconstraint(self, equality, cone_type, efc_map):
|
||||
"""Tests MJX rne_postconstraint function to match MuJoCo mj_rnePostConstraint."""
|
||||
|
||||
m = mujoco.MjModel.from_xml_string("""
|
||||
m = mujoco.MjModel.from_xml_string(f"""
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<geom name="floor" size="10 10 .05" type="plane"/>
|
||||
<body pos="0 0 1">
|
||||
<site name="site1"/>
|
||||
<body name="body1">
|
||||
</body>
|
||||
<body pos="0 0 1" name="body2">
|
||||
<joint type="ball" damping="1"/>
|
||||
<geom type="capsule" size="0.1 0.5" fromto="0 0 0 0.5 0 0" condim="1"/>
|
||||
<body pos="0.5 0 0">
|
||||
<joint type="ball" damping="1"/>
|
||||
<geom type="capsule" size="0.1 0.5" fromto="0 0 0 0.5 0 0" condim="3"/>
|
||||
<site name="site2"/>
|
||||
</body>
|
||||
</body>
|
||||
<body pos="0 1 1">
|
||||
@@ -256,6 +333,7 @@ class RnePostConstraintTest(parameterized.TestCase):
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
{equality}
|
||||
<keyframe>
|
||||
<key qpos='0.424577 0.450592 0.451703 -0.642391 0.729379 0.545151 0.407756 0.0674697 0.424577 1.450592 0.451703 -0.642391 0.729379 0.545151 0.407756 0.0674697'/>
|
||||
</keyframe>
|
||||
@@ -273,6 +351,11 @@ class RnePostConstraintTest(parameterized.TestCase):
|
||||
mx = mjx.put_model(m)
|
||||
dx = mjx.put_data(m, d)
|
||||
|
||||
if efc_map is not None:
|
||||
efc_force = d.efc_force.copy()
|
||||
efc_force[: len(efc_map)] = d.efc_force[efc_map]
|
||||
dx = dx.tree_replace({'_impl.efc_force': jp.array(efc_force)})
|
||||
|
||||
# rne postconstraint
|
||||
mujoco.mj_rnePostConstraint(m, d)
|
||||
dx = jax.jit(mjx.rne_postconstraint)(mx, dx)
|
||||
|
||||
Reference in New Issue
Block a user