From 71333938ecf9ebb5fd6c05033a828ea373b93e46 Mon Sep 17 00:00:00 2001 From: Erik Frey Date: Mon, 22 Apr 2024 13:19:57 -0700 Subject: [PATCH] Add support for condim to MJX. PiperOrigin-RevId: 627136131 Change-Id: Ia62eb20cb80395b569c7fc7f4d4ef87e12eaaf40 --- doc/changelog.rst | 7 +- doc/mjx.rst | 8 +-- mjx/mujoco/mjx/_src/collision_driver.py | 8 ++- mjx/mujoco/mjx/_src/collision_driver_test.py | 3 +- mjx/mujoco/mjx/_src/constraint.py | 71 +++++++++++++------- mjx/mujoco/mjx/_src/constraint_test.py | 34 +++++++--- mjx/mujoco/mjx/_src/io.py | 35 ++++++---- mjx/mujoco/mjx/_src/io_test.py | 35 +++------- mjx/mujoco/mjx/_src/support.py | 13 ---- mjx/mujoco/mjx/_src/types.py | 4 +- mjx/mujoco/mjx/test_data/constraints.xml | 22 +++++- 11 files changed, 143 insertions(+), 97 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 052cae76..f931baed 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -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 `: 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. diff --git a/doc/mjx.rst b/doc/mjx.rst index b53708f6..1628c8ea 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -191,7 +191,7 @@ The following features are **fully supported** in MJX: * - :ref:`Geom ` - ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH`` * - :ref:`Constraint ` - - ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_PYRAMIDAL`` + - ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_FRICTIONLESS``, ``CONTACT_PYRAMIDAL`` * - :ref:`Equality ` - ``CONNECT``, ``WELD``, ``JOINT`` * - :ref:`Integrator ` @@ -199,7 +199,7 @@ The following features are **fully supported** in MJX: * - :ref:`Cone ` - ``PYRAMIDAL`` * - :ref:`Condim ` - - 3 + - 1, 3, 4, 6 * - :ref:`Solver ` - ``CG``, ``NEWTON`` * - Fluid Model @@ -217,10 +217,8 @@ The following features are **in development** and coming soon: - Feature * - :ref:`Geom ` - ``SDF``, ``HFIELD``, ``ELLIPSOID``, ``CYLINDER`` - * - :ref:`Condim ` - - 1, 4, 6 * - :ref:`Constraint ` - - :ref:`Frictionloss `, ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF`` + - :ref:`Frictionloss `, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF`` * - :ref:`Integrator ` - ``IMPLICIT``, ``IMPLICITFAST`` * - Dynamics diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index b59fe1a8..7dca95e2 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -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) diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index afc076f7..294c6a67 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -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') diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index e4264aee..c174dd9f 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -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 diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index 99346828..66d1721f 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -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') diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index f0872fd4..1870e2dd 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -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()}) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 323714eb..72f8b5c2 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -63,7 +63,7 @@ _MULTIPLE_CONVEX_OBJECTS = """ _MULTIPLE_CONSTRAINTS = """ - + @@ -173,22 +173,6 @@ class ModelIOTest(parameterized.TestCase): """)) - def test_condim_not_implemented(self): - with self.assertRaises(NotImplementedError): - mjx.put_model(mujoco.MjModel.from_xml_string(""" - - - - - - - - - - - - """)) - 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) diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 7e1bfb3d..d3ada834 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -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, diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 1df81837..3986d064 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -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 diff --git a/mjx/mujoco/mjx/test_data/constraints.xml b/mjx/mujoco/mjx/test_data/constraints.xml index 42294b50..445e134f 100644 --- a/mjx/mujoco/mjx/test_data/constraints.xml +++ b/mjx/mujoco/mjx/test_data/constraints.xml @@ -14,7 +14,7 @@ - + @@ -37,6 +37,26 @@ + + + + + + + + + + + + + + + + + + + +