Add support for condim to MJX.

PiperOrigin-RevId: 627136131
Change-Id: Ia62eb20cb80395b569c7fc7f4d4ef87e12eaaf40
This commit is contained in:
Erik Frey
2024-04-22 13:19:57 -07:00
committed by Copybara-Service
parent 3ddf1b492a
commit 71333938ec
11 changed files with 143 additions and 97 deletions
+4 -3
View File
@@ -38,12 +38,13 @@ MJX
geoms.
11. Fixed a bug where capsules might be ignored in broadphase colliision checking.
12. Added cylinder collisions using SDFs.
13. Added support for all :ref:`condim <coContact>`: 1, 3, 4, 6.
Bug fixes
^^^^^^^^^
13. Defaults of lights were not being saved, now fixed.
14. Prevent overwriting of frame names by body names when saving an XML. Bug introduced in 3.1.4.
15. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually
14. Defaults of lights were not being saved, now fixed.
15. Prevent overwriting of frame names by body names when saving an XML. Bug introduced in 3.1.4.
16. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually
not optional.
16. Fixed bug that prevented memory allocations larger than 2.15 GB.
+3 -5
View File
@@ -191,7 +191,7 @@ The following features are **fully supported** in MJX:
* - :ref:`Geom <mjtGeom>`
- ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH``
* - :ref:`Constraint <mjtConstraint>`
- ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_PYRAMIDAL``
- ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_FRICTIONLESS``, ``CONTACT_PYRAMIDAL``
* - :ref:`Equality <mjtEq>`
- ``CONNECT``, ``WELD``, ``JOINT``
* - :ref:`Integrator <mjtIntegrator>`
@@ -199,7 +199,7 @@ The following features are **fully supported** in MJX:
* - :ref:`Cone <mjtCone>`
- ``PYRAMIDAL``
* - :ref:`Condim <coContact>`
- 3
- 1, 3, 4, 6
* - :ref:`Solver <mjtSolver>`
- ``CG``, ``NEWTON``
* - Fluid Model
@@ -217,10 +217,8 @@ The following features are **in development** and coming soon:
- Feature
* - :ref:`Geom <mjtGeom>`
- ``SDF``, ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``
* - :ref:`Condim <coContact>`
- 1, 4, 6
* - :ref:`Constraint <mjtConstraint>`
- :ref:`Frictionloss <coFriction>`, ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``
- :ref:`Frictionloss <coFriction>`, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``
* - :ref:`Integrator <mjtIntegrator>`
- ``IMPLICIT``, ``IMPLICITFAST``
* - Dynamics
+6 -2
View File
@@ -108,7 +108,7 @@ def has_collision_fn(t1: GeomType, t2: GeomType) -> bool:
def geom_pairs(
m: Union[Model, mujoco.MjModel],
) -> Iterator[Tuple[int, int, int]]:
"""Returns geom pairs to check for collisions.
"""Yields geom pairs to check for collisions.
Args:
m: a MuJoCo or MJX model
@@ -351,7 +351,11 @@ def collision(m: Model, d: Data) -> Data:
for key, contact in groups.items():
# determine which contacts we'll use for collision testing by running a
# broad phase cull if requested
if max_geom_pairs > -1 and contact.geom.shape[0] > max_geom_pairs:
if (
max_geom_pairs > -1
and contact.geom.shape[0] > max_geom_pairs
and not set(key.types) & _GEOM_NO_BROADPHASE
):
pos1, pos2 = d.geom_xpos[contact.geom.T]
size1, size2 = m.geom_rbound[contact.geom.T]
dist = jax.vmap(jp.linalg.norm)(pos2 - pos1) - (size1 + size2)
+2 -1
View File
@@ -798,7 +798,8 @@ class DimTest(parameterized.TestCase):
def test_ncon(self):
m = test_util.load_test_file('constraints.xml')
dim = collision_driver.make_condim(m)
np.testing.assert_array_equal(dim, np.array([3] * 16))
expected = [1] * 4 + [3] * 20 + [4] * 4 + [6] * 4
np.testing.assert_array_equal(dim, np.array(expected))
def test_disable_contact(self):
m = test_util.load_test_file('constraints.xml')
+46 -25
View File
@@ -284,35 +284,54 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
if d.ncon == 0:
return None
@jax.vmap
def fn(c: Contact):
dist = c.dist - c.includemargin
geom_bodyid = jp.array(m.geom_bodyid)
body1, body2 = geom_bodyid[c.geom1], geom_bodyid[c.geom2]
diff = support.jac_dif_pair(m, d, c.pos, body1, body2)
t = m.body_invweight0[body1, 0] + m.body_invweight0[body2, 0]
def contact_efc(c: Contact, condim: int):
# rotate Jacobian differences to contact frame
diff_con = c.frame @ diff.T
@jax.vmap
def fn(c: Contact):
dist = c.dist - c.includemargin
active = dist < 0
body1, body2 = jp.array(m.geom_bodyid)[c.geom]
jac1p, jac1r = support.jac(m, d, c.pos, body1)
jac2p, jac2r = support.jac(m, d, c.pos, body2)
diff = c.frame @ (jac2p - jac1p).T
if condim > 3: # only calculate rotational diff if needed
diff = jp.concatenate((diff, c.frame @ (jac2r - jac1r).T), axis=0)
tran = m.body_invweight0[body1, 0] + m.body_invweight0[body2, 0]
# TODO(robotics-simulation): add support for other friction dimensions
# 4 pyramidal friction directions
js, invweights = [], []
for diff_tan, friction in zip(diff_con[1:], c.friction[:2]):
for f in (friction, -friction):
js.append(diff_con[0] + diff_tan * f)
invweights.append((t + f * f * t) * 2 * f * f / m.opt.impratio)
if condim == 1:
return diff[0] * active, tran, dist * active, c.solref, c.solimp
active = dist < 0
j, invweight = jp.stack(js) * active, jp.stack(invweights)
pos = jp.repeat(dist, 4) * active
solref, solimp = jp.tile(c.solref, (4, 1)), jp.tile(c.solimp, (4, 1))
# a pair of opposing pyramid edges per friction dimension
# repeat friction directions with positive and negative sign
fri = jp.repeat(c.friction[: condim - 1], 2, axis=0).at[1::2].mul(-1)
# repeat condims of jacdiff to match +/- friction directions
j = diff[0] + jp.repeat(diff[1:condim], 2, axis=0) * fri[:, None]
# pyramidal has common invweight across all edges
diag_approx = tran + fri[0] * fri[0] * tran
inv_w = diag_approx * 2 * fri[0] * fri[0] / m.opt.impratio
repeat_fn = lambda x: jp.repeat(x[None], (condim - 1) * 2, axis=0)
inv_w, pos, solref, solimp = jax.tree_util.tree_map(
repeat_fn, (inv_w, dist, c.solref, c.solimp)
)
return j * active, inv_w, pos * active, solref, solimp
return j, invweight, pos, solref, solimp
return fn(c)
res = fn(d.contact)
# remove contact grouping dimension:
j, invweight, pos, solref, solimp = jax.tree_util.tree_map(jp.concatenate, res)
# group efc calculations by condim
dims, begs = np.unique(d.contact.dim, return_index=True)
efcs = []
for i in range(len(dims)):
dim, beg = dims[i], begs[i]
end = begs[i + 1] if i < len(dims) - 1 else None
c = jax.tree_util.tree_map(lambda x, b=beg, e=end: x[b:e], d.contact)
efc = contact_efc(c, dim)
if dim > 1:
# remove efc grouping dimension
efc = jax.tree_util.tree_map(jp.concatenate, efc)
efcs.append(efc)
efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs)
j, invweight, pos, solref, solimp = efc
frictionloss = jp.zeros_like(pos)
return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
@@ -323,7 +342,9 @@ def counts(efc_type: np.ndarray) -> Tuple[int, int, int, int]:
ne = (efc_type == ConstraintType.EQUALITY).sum()
nf = 0 # no support for friction loss yet
nl = (efc_type == ConstraintType.LIMIT_JOINT).sum()
nc = (efc_type == ConstraintType.CONTACT_PYRAMIDAL).sum()
nc_f = (efc_type == ConstraintType.CONTACT_FRICTIONLESS).sum()
nc_p = (efc_type == ConstraintType.CONTACT_PYRAMIDAL).sum()
nc = nc_f + nc_p
return ne, nf, nl, nc
+26 -8
View File
@@ -48,13 +48,31 @@ class ConstraintTest(absltest.TestCase):
mujoco.mj_forward(m, d)
mx = mjx.put_model(m)
dx = mjx.put_data(m, d)
dx = mjx.make_constraint(mx, dx)
nnz = dx.efc_J.any(axis=1)
_assert_eq(d.efc_J, dx.efc_J[nnz].reshape(-1), 'efc_J')
_assert_eq(d.efc_D, dx.efc_D[nnz], 'efc_D')
_assert_eq(d.efc_aref, dx.efc_aref[nnz], 'efc_aref')
_assert_eq(d.efc_frictionloss, dx.efc_frictionloss[nnz], 'efc_frictionloss')
d_efc_j = d.efc_J.reshape((-1, m.nv))
# ne, nf, nl order matches
efl = d.ne + d.nf + d.nl
_assert_eq(d_efc_j[:efl], dx.efc_J[:efl], 'efc_J')
_assert_eq(d.efc_D[:efl], dx.efc_D[:efl], 'efc_D')
_assert_eq(d.efc_aref[:efl], dx.efc_aref[:efl], 'efc_aref')
_assert_eq(dx.efc_frictionloss, 0, 'efc_frictionloss')
# contact order might not match, so check efcs contact by contact
for i in range(d.ncon):
geom_match = (dx.contact.geom == d.contact.geom[i]).all(axis=-1)
geom_match &= (dx.contact.pos == d.contact.pos[i]).all(axis=-1)
self.assertTrue(geom_match.any(), f'contact {i} not found in MJX contact')
j = np.nonzero(geom_match)[0][0]
self.assertEqual(d.contact.dim[i], dx.contact.dim[j])
nc = max(1, (d.contact.dim[i] - 1) * 2)
d_beg, dx_beg = d.contact.efc_address[i], dx.contact.efc_address[j]
d_end, dx_end = d_beg + nc, dx_beg + nc
_assert_eq(d_efc_j[d_beg:d_end], dx.efc_J[dx_beg:dx_end], 'efc_J')
_assert_eq(d.efc_D[d_beg:d_end], dx.efc_D[dx_beg:dx_end], 'efc_D')
d_efc_aref = d.efc_aref[d_beg:d_end]
dx_efc_aref = dx.efc_aref[dx_beg:dx_end]
_assert_eq(d_efc_aref, dx_efc_aref, 'efc_aref')
def test_disable_refsafe(self):
m = test_util.load_test_file('constraints.xml')
@@ -87,9 +105,9 @@ class ConstraintTest(absltest.TestCase):
self.assertEqual(ne, 0)
self.assertEqual(nf, 0)
self.assertEqual(nl, 2)
self.assertEqual(nc, 64)
self.assertEqual(nc, 148)
dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
self.assertEqual(dx.efc_J.shape[0], 66) # only joint range, contact
self.assertEqual(dx.efc_J.shape[0], 150) # only joint range, contact
def test_disable_contact(self):
m = test_util.load_test_file('constraints.xml')
+22 -13
View File
@@ -68,9 +68,6 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
if m.ntendon:
raise NotImplementedError('tendons are not supported')
if (m.geom_condim != 3).any() or (m.pair_dim != 3).any():
raise NotImplementedError('only condim=3 is supported')
if m.body_gravcomp.any():
raise NotImplementedError('gravcomp is not supported')
@@ -321,13 +318,27 @@ def _make_contact(
"""Converts mujoco.structs._MjContactList into mjx.Contact."""
fields = {f.name: getattr(c, f.name) for f in types.Contact.fields()}
fields['frame'] = fields['frame'].reshape((-1, 3, 3))
pad_size = dim.size - c.dist.shape[0]
pad_fn = lambda x: np.concatenate(
(x, np.zeros((pad_size,) + x.shape[1:], dtype=x.dtype))
)
fields = jax.tree_util.tree_map(pad_fn, fields)
fields['dist'][-pad_size:] = np.inf
# TODO(erikfrey): move contacts to appropriate dim index
# reorder contacts so that their condims match those specified in dim.
# if we have fewer Contacts for a condim range, pad the range with zeros
# build a map for where to find a dim-matching contact, or -1 if none
contact_map = np.zeros_like(dim) - 1
for i, di in enumerate(fields['dim']):
space = [j for j, dj in enumerate(dim) if di == dj and contact_map[j] == -1]
if not space:
# this can happen if max_geom_pairs or max_contact_points is too low
raise ValueError(f'unable to place Contact[{i}], no space in condim {di}')
contact_map[space[0]] = i
if contact_map.size > 0:
# reorganize contact according, with a zero contact at the end for -1
zero = jax.tree_util.tree_map(
lambda x: np.zeros((1,) + x.shape[1:], dtype=x.dtype), fields
)
zero['dist'][:] = np.finfo(float).max
fields = jax.tree_util.tree_map(lambda *x: np.concatenate(x), fields, zero)
fields = jax.tree_util.tree_map(lambda x: x[contact_map], fields)
fields['dim'] = dim
fields['efc_address'] = efc_address
@@ -398,9 +409,7 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
fields['qLDiagInv'] = np.zeros(0)
fields['contact'] = _make_contact(d.contact, dim, efc_address)
fields.update(
dict(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type)
)
fields.update(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type)
# copy because device_put is async:
data = types.Data(**{k: copy.copy(v) for k, v in fields.items()})
+10 -25
View File
@@ -63,7 +63,7 @@ _MULTIPLE_CONVEX_OBJECTS = """
_MULTIPLE_CONSTRAINTS = """
<mujoco>
<worldbody>
<geom type="plane" size="3 3 .01"/>
<geom type="plane" size="3 3 .01" condim="6"/>
<body name="cap1" pos="-.3 -.3 .2">
<freejoint/>
<geom type="capsule" size=".2 .05"/>
@@ -173,22 +173,6 @@ class ModelIOTest(parameterized.TestCase):
</tendon>
</mujoco>"""))
def test_condim_not_implemented(self):
with self.assertRaises(NotImplementedError):
mjx.put_model(mujoco.MjModel.from_xml_string("""
<mujoco>
<worldbody>
<body>
<freejoint/>
<geom size="0.05" condim="1"/>
</body>
<body>
<freejoint/>
<geom size="0.05" condim="1"/>
</body>
</worldbody>
</mujoco>"""))
def test_gravcomp_not_implemented(self):
with self.assertRaises(NotImplementedError):
mjx.put_model(mujoco.MjModel.from_xml_string("""
@@ -345,21 +329,21 @@ class DataIOTest(parameterized.TestCase):
np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat)
# efc_ are also shape transformed and padded
self.assertEqual(dx.efc_J.shape, (21, 8)) # nefc, nv
self.assertEqual(dx.efc_J.shape, (45, 8)) # nefc, nv
d_efc_j = d.efc_J.reshape((-1, 8))
np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3]) # connect eq
np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3]) # one active limit
np.testing.assert_allclose(dx.efc_J[4], 0) # one inactive limit
np.testing.assert_allclose(dx.efc_J[5:9], d_efc_j[4:8]) # contact
np.testing.assert_allclose(dx.efc_J[9:], 0) # no contact
np.testing.assert_allclose(dx.efc_J[5:15], d_efc_j[4:14]) # contact
np.testing.assert_allclose(dx.efc_J[15:], 0) # no contact
# check another efc_ too
self.assertEqual(dx.efc_aref.shape, (21,)) # nefc
self.assertEqual(dx.efc_aref.shape, (45,)) # nefc
np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3])
np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3])
np.testing.assert_allclose(dx.efc_aref[4], 0)
np.testing.assert_allclose(dx.efc_aref[5:9], d.efc_aref[4:8])
np.testing.assert_allclose(dx.efc_aref[9:], 0)
np.testing.assert_allclose(dx.efc_aref[5:15], d.efc_aref[4:14])
np.testing.assert_allclose(dx.efc_aref[15:], 0)
# check sparse transform is correct
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
@@ -416,9 +400,10 @@ class DataIOTest(parameterized.TestCase):
np.testing.assert_allclose(d_2.site_xmat, d.site_xmat)
# efc_* are also shape transformed and filtered
self.assertEqual(d_2.efc_J.shape, (64,)) # nefc * nv
self.assertEqual(d_2.nefc, 14)
self.assertEqual(d_2.efc_J.shape, (112,)) # nefc * nv
np.testing.assert_allclose(d_2.efc_J, d.efc_J)
self.assertEqual(d_2.efc_aref.shape, (8,)) # nefc
self.assertEqual(d_2.efc_aref.shape, (14,)) # nefc
np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
-13
View File
@@ -149,19 +149,6 @@ def jac(
return jacp, jacr
def jac_dif_pair(
m: Model,
d: Data,
pos: jax.Array,
body_1: jax.Array,
body_2: jax.Array,
) -> jax.Array:
"""Compute Jacobian difference for two body points."""
jacp2, _ = jac(m, d, pos, body_2)
jacp1, _ = jac(m, d, pos, body_1)
return jacp2 - jacp1
def apply_ft(
m: Model,
d: Data,
+3 -1
View File
@@ -217,12 +217,14 @@ class ConstraintType(enum.IntEnum):
Attributes:
EQUALITY: equality constraint
LIMIT_JOINT: joint limit
CONTACT_FRICTIONLESS: frictionless contact
CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone
"""
EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY
# unsupported: FRICTION_DOF, FRICTION_TENDON
LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT
# unsupported: LIMIT_TENDON, CONTACT_FRICTIONLESS
# unsupported: LIMIT_TENDON
CONTACT_FRICTIONLESS = mujoco.mjtConstraint.mjCNSTR_CONTACT_FRICTIONLESS
CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
# unsupported: CONTACT_ELLIPTIC
+21 -1
View File
@@ -14,7 +14,7 @@
</default>
<worldbody>
<geom pos="0 0 -1" type="plane" size="10 10 .01"/>
<geom pos="0 0 -1" type="plane" size="20 20 .01" condim="1"/>
<body name="anchor1" pos="-3 0 0"/>
<body name="beam1" pos="-3 0 0">
@@ -37,6 +37,26 @@
<joint name="joint4" axis="1 0 0" type="hinge" damping="1.0"/> <!-- tests no joint range -->
<geom class="box"/>
</body>
<body name="box_condim1" pos="4 0 0">
<freejoint/>
<geom class="box" condim="1"/>
</body>
<body name="box_condim3" pos="5 0 0">
<freejoint/>
<geom class="box" condim="3"/>
</body>
<body name="box_condim4" pos="6 0 0">
<freejoint/>
<geom class="box" condim="4"/>
</body>
<body name="box_condim6" pos="6 0 0">
<freejoint/>
<geom class="box" condim="6"/>
</body>
</worldbody>
<equality>