Add support for condim to MJX.

PiperOrigin-RevId: 627136131
Change-Id: Ia62eb20cb80395b569c7fc7f4d4ef87e12eaaf40

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 <coContact>`: 1, 3, 4, 6.
 
 Bug fixes
 ^^^^^^^^^
-13. Defaults of lights were not being saved, now fixed.
+14. 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. 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
+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.
 

doc/mjx.rst

@@ -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

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)

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')

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 contact_efc(c: Contact, condim: int):
-  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]
 
-    # rotate Jacobian differences to contact frame
+    @jax.vmap
-    diff_con = c.frame @ diff.T
+    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
+      if condim == 1:
-    # 4 pyramidal friction directions
+        return diff[0] * active, tran, dist * active, c.solref, c.solimp
-    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)
 
-    active = dist < 0
+      # a pair of opposing pyramid edges per friction dimension
-    j, invweight = jp.stack(js) * active, jp.stack(invweights)
+      # repeat friction directions with positive and negative sign
-    pos = jp.repeat(dist, 4) * active
+      fri = jp.repeat(c.friction[: condim - 1], 2, axis=0).at[1::2].mul(-1)
-    solref, solimp = jp.tile(c.solref, (4, 1)), jp.tile(c.solimp, (4, 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)
+  # group efc calculations by condim
-  # remove contact grouping dimension:
+  dims, begs = np.unique(d.contact.dim, return_index=True)
-  j, invweight, pos, solref, solimp = jax.tree_util.tree_map(jp.concatenate, res)
+  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
 

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)
+    d_efc_j = d.efc_J.reshape((-1, m.nv))
-    _assert_eq(d.efc_J, dx.efc_J[nnz].reshape(-1), 'efc_J')
+
-    _assert_eq(d.efc_D, dx.efc_D[nnz], 'efc_D')
+    # ne, nf, nl order matches
-    _assert_eq(d.efc_aref, dx.efc_aref[nnz], 'efc_aref')
+    efl = d.ne + d.nf + d.nl
-    _assert_eq(d.efc_frictionloss, dx.efc_frictionloss[nnz], 'efc_frictionloss')
+    _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')

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]
+  # reorder contacts so that their condims match those specified in dim.
-  pad_fn = lambda x: np.concatenate(
+  # if we have fewer Contacts for a condim range, pad the range with zeros
-      (x, np.zeros((pad_size,) + x.shape[1:], dtype=x.dtype))
+
-  )
+  # build a map for where to find a dim-matching contact, or -1 if none
-  fields = jax.tree_util.tree_map(pad_fn, fields)
+  contact_map = np.zeros_like(dim) - 1
-  fields['dist'][-pad_size:] = np.inf
+  for i, di in enumerate(fields['dim']):
-  # TODO(erikfrey): move contacts to appropriate dim index
+    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(
+  fields.update(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type)
-      dict(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()})

mjx/mujoco/mjx/_src/io_test.py

@@ -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[5:15], d_efc_j[4:14])  # contact
-    np.testing.assert_allclose(dx.efc_J[9:], 0)  # no 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[5:15], d.efc_aref[4:14])
-    np.testing.assert_allclose(dx.efc_aref[9:], 0)
+    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)
 

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,

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
 

mjx/mujoco/mjx/test_data/constraints.xml

@@ -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>