Add support for condim to MJX.

PiperOrigin-RevId: 627136131 Change-Id: Ia62eb20cb80395b569c7fc7f4d4ef87e12eaaf40
2024-04-22 13:19:57 -07:00
parent 3ddf1b492a
commit 71333938ec
11 changed files with 143 additions and 97 deletions
@@ -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)
@@ -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')
@@ -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

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

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

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