Add eq_objtype to mjx.Model and support for connect constraint using sites in MJX.

PiperOrigin-RevId: 671225109
Change-Id: I8e2e6eb4c539d318658d3d7dfe0a7f03b411aabc

doc/changelog.rst

@@ -61,6 +61,7 @@ MJX
   :doc:`fluid drag <computation/fluid>`.
 - Fixed a bug where ``qLDiagInv`` had the wrong size for sparse mass matrices.
 - Added support for joint and tendon :ref:`frictionloss <coFriction>`.
+- Added support for :ref:`connect<equality-connect>` equality constraints using two sites.
 
 Bug fixes
 ^^^^^^^^^

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -887,7 +887,7 @@ class DimTest(parameterized.TestCase):
   def test_ncon(self):
     m = test_util.load_test_file('constraints.xml')
     dim = collision_driver.make_condim(m)
-    expected = [1] * 4 + [3] * 20 + [4] * 4 + [6] * 4
+    expected = [1] * 4 + [3] * 24 + [4] * 4 + [6] * 4
     np.testing.assert_array_equal(dim, np.array(expected))
 
   def test_disable_contact(self):

mjx/mujoco/mjx/_src/constraint.py

@@ -32,6 +32,7 @@ from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import EqType
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ObjType
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -106,25 +107,47 @@ def _efc_equality_connect(m: Model, d: Data) -> Optional[_Efc]:
     return None
 
   @jax.vmap
-  def rows(obj1id, obj2id, data, solref, solimp):
+  def rows(is_site, obj1id, obj2id, body1id, body2id, data, solref, solimp):
     anchor1, anchor2 = data[0:3], data[3:6]
 
-    # error is difference in global positions
     pos1 = d.xmat[obj1id] @ anchor1 + d.xpos[obj1id]
     pos2 = d.xmat[obj2id] @ anchor2 + d.xpos[obj2id]
+
+    if m.nsite:
+      pos1 = jp.where(is_site, d.site_xpos[obj1id], pos1)
+      pos2 = jp.where(is_site, d.site_xpos[obj2id], pos2)
+
+    # error is difference in global positions
     pos = pos1 - pos2
 
     # compute Jacobian difference (opposite of contact: 0 - 1)
-    jacp1, _ = support.jac(m, d, pos1, obj1id)
+    jacp1, _ = support.jac(m, d, pos1, body1id)
-    jacp2, _ = support.jac(m, d, pos2, obj2id)
+    jacp2, _ = support.jac(m, d, pos2, body2id)
     j = (jacp1 - jacp2).T
     pos_imp = math.norm(pos)
-    invweight = m.body_invweight0[obj1id, 0] + m.body_invweight0[obj2id, 0]
+    invweight = m.body_invweight0[body1id, 0] + m.body_invweight0[body2id, 0]
     zero = jp.zeros_like(pos)
 
     return _row(j, pos, pos_imp, invweight, solref, solimp, zero, zero)
 
-  args = (m.eq_obj1id, m.eq_obj2id, m.eq_data, m.eq_solref, m.eq_solimp)
+  is_site = m.eq_objtype == ObjType.SITE
+
+  body1id = np.copy(m.eq_obj1id)
+  body2id = np.copy(m.eq_obj2id)
+  if m.nsite:
+    body1id[is_site] = m.site_bodyid[body1id[is_site]]
+    body2id[is_site] = m.site_bodyid[body2id[is_site]]
+
+  args = (
+      is_site,
+      m.eq_obj1id,
+      m.eq_obj2id,
+      body1id,
+      body2id,
+      m.eq_data,
+      m.eq_solref,
+      m.eq_solimp,
+  )
   args = jax.tree_util.tree_map(lambda x: x[eq_id], args)
   # concatenate to drop row grouping
   return jax.tree_util.tree_map(jp.concatenate, rows(*args))

mjx/mujoco/mjx/_src/constraint_test.py

@@ -103,31 +103,31 @@ class ConstraintTest(parameterized.TestCase):
     self.assertEqual(ne, 0)
     self.assertEqual(nf, 2)
     self.assertEqual(nl, 5)
-    self.assertEqual(nc, 148)
+    self.assertEqual(nc, 164)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 155)  # only joint/tendon limit, contact
+    self.assertEqual(dx.efc_J.shape[0], 171)  # only joint/tendon limit, contact
 
   def test_disable_contact(self):
     m = test_util.load_test_file('constraints.xml')
     m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT
     ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m))
-    self.assertEqual(ne, 11)
+    self.assertEqual(ne, 14)
     self.assertEqual(nf, 2)
     self.assertEqual(nl, 5)
     self.assertEqual(nc, 0)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 18)  # only equality, joint/tendon limit
+    self.assertEqual(dx.efc_J.shape[0], 21)  # only equality, joint/tendon limit
 
   def test_disable_frictionloss(self):
     m = test_util.load_test_file('constraints.xml')
     m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.FRICTIONLOSS
     ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m))
-    self.assertEqual(ne, 11)
+    self.assertEqual(ne, 14)
     self.assertEqual(nf, 0)
     self.assertEqual(nl, 5)
-    self.assertEqual(nc, 148)
+    self.assertEqual(nc, 164)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 164)
+    self.assertEqual(dx.efc_J.shape[0], 183)
 
   def test_margin(self):
     """Test margin."""

mjx/mujoco/mjx/_src/types.py

@@ -681,6 +681,7 @@ class Model(PyTreeNode):
     eq_type: constraint type (mjtEq)                          (neq,)
     eq_obj1id: id of object 1                                 (neq,)
     eq_obj2id: id of object 2                                 (neq,)
+    eq_objtype: type of both objects (mjtObj)                 (neq,)
     eq_active0: initial enable/disable constraint state       (neq,)
     eq_solref: constraint solver reference                    (neq, mjNREF)
     eq_solimp: constraint solver impedance                    (neq, mjNIMP)
@@ -991,6 +992,7 @@ class Model(PyTreeNode):
   eq_type: np.ndarray
   eq_obj1id: np.ndarray
   eq_obj2id: np.ndarray
+  eq_objtype: np.ndarray
   eq_active0: np.ndarray
   eq_solref: jax.Array
   eq_solimp: jax.Array

mjx/mujoco/mjx/test_data/constraints.xml

@@ -16,6 +16,13 @@
   <worldbody>
     <geom pos="0 0 -1" type="plane" size="20 20 .01" condim="1"/>
 
+    <site name="0" pos="0 0 1"/>
+    <body pos="0 0 1">
+      <freejoint/>
+      <geom class="box"/>
+      <site name="1"/>
+    </body>
+
     <body name="anchor1" pos="-3 0 0"/>
     <body name="beam1" pos="-3 0 0">
       <joint name="joint1" type="ball" range="0 45" solreflimit="0.03 0.9" solimplimit="0.89 0.9 0.01 2.1"/>
@@ -76,6 +83,7 @@
   </tendon>
 
   <equality>
+    <connect name="site"     site1="0"          site2="1"/>
     <connect name="connect"  body1="anchor1"    body2="beam1"   anchor="1 0 -1" />
     <weld    name="weld"     body1="anchor2"    body2="beam2"   relpose="0 0 0 1 -.3 0 0" torquescale="0.002" anchor="0 -2 0"/>
     <joint   name="joint"    joint1="joint3"    joint2="joint4" polycoef="0.5 -1 0.1 0.15 0.2" />
@@ -90,10 +98,10 @@
 
   <keyframe>
     <!-- keyframe 0: default position with some motion, zero contacts -->
-    <key qpos='1 0 0 0 -1 0 0 1 0 0 0 0 0 1 0 0 0 4 0 0 1 0 0 0 5 0 0 1 0 0 0 6 0 0 1 0 0 0 7 0 0 1 0 0 0' qvel='1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1'/>
+    <key qpos=' -1 -1 1 1 0 0 0 1 0 0 0 -1 0 0 1 0 0 0 0 0 1 0 0 0 4 0 0 1 0 0 0 5 0 0 1 0 0 0 6 0 0 1 0 0 0 7 0 0 1 0 0 0' qvel='1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1'/>
     <!-- keyframe 1: some contacts but constraints are quadratic -->
-    <key qpos='1 0.0011 -8.2e-07 -0.0011 -0.82 1.9 -0.58 0.97 -0.14 0.17 -0.069 0.22 0.29 0.93 0.2 0.15 0.28 4.3 0.23 -0.26 0.97 0.12 0.11 0.18 5.4 0.25 -0.26 0.97 0.11 0.16 0.17 6.4 0.25 -0.26 0.97 0.11 0.16 0.17 7.4 0.25 -0.26 0.97 0.11 0.16 0.17' qvel='-0.14 3.3e-05 0.14 -0.14 -0.93 -3.6 -1.8 1 -0.57 -0.16 0.21 1.1 1 2.2 0.65 1.3 -4.2 -2.6 -5.6 0.54 1.6 2.8 -3.7 -2.3 0.7 -1.8 1.6 2.8 -3.7 -2.3 0.62 -1.8 1.6 2.8 -3.7 -2.3 0.61 -1.8'/>
+    <key qpos=' -1 -1 1 1 0 0 0 1 0.0011 -8.2e-07 -0.0011 -0.82 1.9 -0.58 0.97 -0.14 0.17 -0.069 0.22 0.29 0.93 0.2 0.15 0.28 4.3 0.23 -0.26 0.97 0.12 0.11 0.18 5.4 0.25 -0.26 0.97 0.11 0.16 0.17 6.4 0.25 -0.26 0.97 0.11 0.16 0.17 7.4 0.25 -0.26 0.97 0.11 0.16 0.17' qvel='1 1 1 1 1 1 -0.14 3.3e-05 0.14 -0.14 -0.93 -3.6 -1.8 1 -0.57 -0.16 0.21 1.1 1 2.2 0.65 1.3 -4.2 -2.6 -5.6 0.54 1.6 2.8 -3.7 -2.3 0.7 -1.8 1.6 2.8 -3.7 -2.3 0.62 -1.8 1.6 2.8 -3.7 -2.3 0.61 -1.8'/>
     <!-- keyframe 2: some contacts and some constraints are in cone state (for elliptic) -->
-    <key qpos='1 0.0087 2.4e-07 -0.0086 -0.89 1.8 -0.77 0.98 -0.2 -0.0022 -0.026 0.19 0.33 0.86 0.32 0.064 0.38 4.4 0.36 -0.81 0.97 -0.0013 -0.0011 0.25 5.6 0.52 -0.75 0.98 -0.018 0.17 0.094 6.6 0.52 -0.75 0.98 -0.017 0.16 0.094 7.6 0.52 -0.76 0.98 -0.017 0.16 0.094' qvel='0.2 -1.8e-05 -0.2 -0.72 0.072 0.025 0.015 -4.9 0.35 -0.22 0.26 0.99 -4.6 1.7 1.1 0.52 -0.73 0.16 7.9 1 0.043 -0.042 1.3 1.6 -1.6 0.98 0.027 -0.024 1.3 1.6 -1.8 0.96 0.025 -0.022 1.3 1.6 -1.8 0.96'/>
+    <key qpos=' -1 -1 1 1 0 0 0 1 0.0087 2.4e-07 -0.0086 -0.89 1.8 -0.77 0.98 -0.2 -0.0022 -0.026 0.19 0.33 0.86 0.32 0.064 0.38 4.4 0.36 -0.81 0.97 -0.0013 -0.0011 0.25 5.6 0.52 -0.75 0.98 -0.018 0.17 0.094 6.6 0.52 -0.75 0.98 -0.017 0.16 0.094 7.6 0.52 -0.76 0.98 -0.017 0.16 0.094' qvel='1 1 1 1 1 1 0.2 -1.8e-05 -0.2 -0.72 0.072 0.025 0.015 -4.9 0.35 -0.22 0.26 0.99 -4.6 1.7 1.1 0.52 -0.73 0.16 7.9 1 0.043 -0.042 1.3 1.6 -1.6 0.98 0.027 -0.024 1.3 1.6 -1.8 0.96 0.025 -0.022 1.3 1.6 -1.8 0.96'/>
   </keyframe>
 </mujoco>