Merge branch 'google-deepmind:main' into main

mjx/mujoco/mjx/_src/collision_driver.py

@@ -56,7 +56,10 @@ from mujoco.mjx._src.collision_primitive import plane_ellipsoid
 from mujoco.mjx._src.collision_primitive import plane_sphere
 from mujoco.mjx._src.collision_primitive import sphere_capsule
 from mujoco.mjx._src.collision_primitive import sphere_sphere
+from mujoco.mjx._src.collision_sdf import capsule_cylinder
 from mujoco.mjx._src.collision_sdf import capsule_ellipsoid
+from mujoco.mjx._src.collision_sdf import cylinder_cylinder
+from mujoco.mjx._src.collision_sdf import ellipsoid_cylinder
 from mujoco.mjx._src.collision_sdf import ellipsoid_ellipsoid
 from mujoco.mjx._src.collision_types import FunctionKey
 from mujoco.mjx._src.types import Contact
@@ -82,8 +85,11 @@ _COLLISION_FUNC = {
     (GeomType.CAPSULE, GeomType.CAPSULE): capsule_capsule,
     (GeomType.CAPSULE, GeomType.BOX): capsule_convex,
     (GeomType.CAPSULE, GeomType.ELLIPSOID): capsule_ellipsoid,
+    (GeomType.CAPSULE, GeomType.CYLINDER): capsule_cylinder,
     (GeomType.CAPSULE, GeomType.MESH): capsule_convex,
     (GeomType.ELLIPSOID, GeomType.ELLIPSOID): ellipsoid_ellipsoid,
+    (GeomType.ELLIPSOID, GeomType.CYLINDER): ellipsoid_cylinder,
+    (GeomType.CYLINDER, GeomType.CYLINDER): cylinder_cylinder,
     (GeomType.BOX, GeomType.BOX): convex_convex,
     (GeomType.BOX, GeomType.MESH): convex_convex,
     (GeomType.MESH, GeomType.MESH): convex_convex,
@@ -102,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
@@ -345,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

@@ -264,6 +264,29 @@ class EllipsoidCollisionTest(parameterized.TestCase):
       _assert_attr_eq(
           dx.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-4)
 
+  _ELLIPSOID_CYLINDER = """
+    <mujoco>
+      <worldbody>
+        <body>
+          <geom size=".15 .05" type="cylinder"/>
+        </body>
+        <body pos="0 0 0.09">
+          <freejoint/>
+          <geom size=".15 .03 .05" type="ellipsoid"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  def test_ellipsoid_cylinder(self):
+    """Tests ellipsoid cylinder contact."""
+    d, dx = _collide(self._ELLIPSOID_CYLINDER)
+    d.contact.pos[0][2] = 0.04  # MJX finds the deepest point on the surface
+    self.assertLess(dx.contact.dist[0], 0)
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(
+          dx.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4)
+
 
 class CapsuleCollisionTest(parameterized.TestCase):
   _CAP_PLANE = """
@@ -775,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/collision_sdf.py

@@ -85,6 +85,43 @@ def _ellipsoid(pos: jax.Array, size: jax.Array) -> jax.Array:
   return k0 * (k0 - 1.0) / (k1 + (k1 == 0.0) * 1e-12)
 
 
+@jax.custom_jvp
+def _cylinder(pos: jax.Array, size: jax.Array) -> jax.Array:
+  a0 = jp.sqrt(pos[0] * pos[0] + pos[1] * pos[1]) - size[0]
+  a1 = jp.abs(pos[2]) - size[1]
+  b0 = jp.maximum(a0, 0)
+  b1 = jp.maximum(a1, 0)
+  return jp.minimum(jp.maximum(a0, a1), 0) + jp.sqrt(b0 * b0 + b1 * b1)
+
+
+def _cylinder_grad(x: jax.Array, size: jax.Array) -> jax.Array:
+  """Gradient of the cylinder SDF wrt query point and singularities removed."""
+  c = jp.sqrt(x[0]*x[0]+x[1]*x[1])
+  e = jp.abs(x[2])
+  a = jp.array([c - size[0], e - size[1]])
+  b = jp.array([jp.maximum(a[0], 0), jp.maximum(a[1], 0)])
+  j = jp.argmax(a)
+  bnorm = jp.sqrt(b[0]*b[0] + b[1]*b[1])
+  bnorm += jp.allclose(bnorm, 0) * 1e-12
+  grada = jp.array([
+      x[0] / (c + jp.allclose(c, 0) * 1e-12),
+      x[1] / (c + jp.allclose(c, 0) * 1e-12),
+      x[2] / (e + jp.allclose(e, 0) * 1e-12),
+  ])
+  gradm = jp.array([[grada[0], grada[1], 0], [0, 0, grada[2]]])
+  gradb = grada * b[jp.array([0, 0, 1])] / bnorm
+  return jp.where(a[j] < 0, gradm[j], gradb)
+
+
+@_cylinder.defjvp
+def cylinder_jvp(primals, tangents):
+  x, y = primals
+  x_dot, _ = tangents
+  primal_out = _cylinder(x, y)
+  tangent_out = jp.dot(_cylinder_grad(x, y), x_dot)
+  return primal_out, tangent_out
+
+
 def _to_local(f: SDFFn, pos: jax.Array, mat: jax.Array)-> SDFFn:
   return lambda p: f(mat.T @ (p - pos))
 
@@ -134,29 +171,73 @@ def _gradient_descent(
 
 
 def _optim(
-    d1, d2, info1: GeomInfo, info2: GeomInfo
-) -> Tuple[jax.Array, jax.Array, jax.Array]:
+    d1, d2, info1: GeomInfo, info2: GeomInfo, x0: jax.Array,
+) -> Collision:
   """Optimizes the clearance function."""
   d1 = functools.partial(d1, size=info1.size)
   d1 = _to_local(d1, info1.pos, info1.mat)
   d2 = functools.partial(d2, size=info2.size)
   d2 = _to_local(d2, info2.pos, info2.mat)
   fn = _clearance(d1, d2)
-  _, pos = _gradient_descent(fn, 0.5 * (info1.pos + info2.pos), 10)
+  _, pos = _gradient_descent(fn, x0, 10)
   dist = d1(pos) + d2(pos)
-  n = jax.grad(d1)(pos)
-  return pos, dist, n
+  n = jax.grad(d1)(pos) - jax.grad(d2)(pos)
+  return dist, pos, math.make_frame(n)
 
 
 @collider(ncon=1)
 def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Collision:
   """"Calculates contact between a capsule and an ellipsoid."""
-  pos, dist, n = _optim(_capsule, _ellipsoid, c, e)
-  return dist, pos, math.make_frame(n)
+  x0 = 0.5 * (c.pos + e.pos)
+  return _optim(_capsule, _ellipsoid, c, e, x0)
+
+
+@collider(ncon=2)
+def capsule_cylinder(ca: GeomInfo, cy: GeomInfo) -> Collision:
+  """"Calculates contact between a capsule and a cylinder."""
+  # TODO: improve robustness
+  # Near sharp corners, the SDF might give the penetration depth with respect
+  # to a surface that is not in collision. Possible solutions is to find the
+  # contact points analytically or to change the SDF depending on the relative
+  # pose of the bodies.
+  mid = 0.5 * (ca.pos + cy.pos)
+  vec = ca.mat[:, 2] * ca.size[1]
+  x0 = jp.array([mid - vec, mid + vec])
+  optim_ = functools.partial(_optim, _capsule, _cylinder, ca, cy)
+  return jax.vmap(optim_)(x0)
 
 
 @collider(ncon=1)
 def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Collision:
   """"Calculates contact between two ellipsoids."""
-  pos, dist, n = _optim(_ellipsoid, _ellipsoid, e1, e2)
-  return dist, pos, math.make_frame(n)
+  x0 = 0.5 * (e1.pos + e2.pos)
+  return _optim(_ellipsoid, _ellipsoid, e1, e2, x0)
+
+
+@collider(ncon=1)
+def ellipsoid_cylinder(e: GeomInfo, c: GeomInfo) -> Collision:
+  """"Calculates contact between and ellipsoid and a cylinder."""
+  x0 = 0.5 * (e.pos + c.pos)
+  return _optim(_ellipsoid, _cylinder, e, c, x0)
+
+
+@collider(ncon=4)
+def cylinder_cylinder(c1: GeomInfo, c2: GeomInfo) -> Collision:
+  """"Calculates contact between a cylinder and a cylinder."""
+  # TODO: improve robustness
+  # Near sharp corners, the SDF might give the penetration depth with respect
+  # to a surface that is not in collision. Possible solutions is to find the
+  # contact points analytically or to change the SDF depending on the relative
+  # pose of the bodies.
+  basis = math.make_frame(c2.pos - c1.pos)
+  mid = 0.5 * (c1.pos + c2.pos)
+  r = jp.maximum(c1.size[0], c2.size[0])
+  x0 = jp.array([
+      mid + r * basis[1],
+      mid + r * basis[2],
+      mid - r * basis[1],
+      mid - r * basis[2],
+  ])
+  optim_ = functools.partial(_optim, _cylinder, _cylinder, c1, c2)
+  return jax.vmap(optim_)(x0)
+

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
 

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

mjx/mujoco/mjx/_src/dataclasses.py

@@ -18,7 +18,7 @@ import copy
 import dataclasses
 
 import typing
-from typing import Any, Dict, Optional, Sequence, TypeVar, Union
+from typing import Dict, Optional, Sequence, Tuple, TypeVar, Union
 import jax
 import numpy as np
 
@@ -122,7 +122,7 @@ class PyTreeNode:
     raise NotImplementedError
 
   @classmethod
-  def fields(cls) -> tuple[dataclasses.Field[Any], ...]:
+  def fields(cls) -> Tuple[dataclasses.Field, ...]:  # pylint: disable=g-bare-generic
     return dataclasses.fields(cls)
 
   def tree_replace(

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()})

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

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/test_util.py

@@ -17,7 +17,7 @@
 import os
 import sys
 import time
-from typing import Dict, Tuple
+from typing import Dict, Optional, Tuple
 from xml.etree import ElementTree as ET
 
 from etils import epath
@@ -201,9 +201,9 @@ def _make_geom(
 
 def _make_actuator(
     actuator_type: str,
-    joint: str | None = None,
-    site: str | None = None,
-    refsite: str | None = None,
+    joint: Optional[str] = None,
+    site: Optional[str] = None,
+    refsite: Optional[str] = None,
 ) -> Dict[str, str]:
   """Returns attributes for an actuator."""
   if joint:

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>