Update MJX rne_postconstraint with contributions from equality connect and equality weld constraints.

thowell · copybara-github · commit 2ce5caa2395a · 2025-06-30T03:21:45.000-07:00
PiperOrigin-RevId: 777491801
Change-Id: Ib32dda4178def7a96d31e91787f3d0866a56cd76
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
@@ -248,23 +248,6 @@ def _put_model_jax(
       if t == mujoco.mjtGeom.mjGEOM_MESH:
         mesh_geomid.add(g)
 
-  # check for unsupported sensor and equality constraint combinations
-  sensor_rne_postconstraint = (
-      np.any(m.sensor_type == types.SensorType.ACCELEROMETER)
-      | np.any(m.sensor_type == types.SensorType.FORCE)
-      | np.any(m.sensor_type == types.SensorType.TORQUE)
-      | np.any(m.sensor_type == types.SensorType.FRAMELINACC)
-      | np.any(m.sensor_type == types.SensorType.FRAMEANGACC)
-  )
-  eq_connect_weld = np.any(m.eq_type == types.EqType.CONNECT) | np.any(
-      m.eq_type == types.EqType.WELD
-  )
-  if sensor_rne_postconstraint and eq_connect_weld:
-    raise NotImplementedError(
-        'rne_postconstraint not implemented with equality constraints:'
-        ' connect, weld.'
-    )
-
   for enum_field, enum_type, mj_type in (
       (m.actuator_biastype, types.BiasType, mujoco.mjtBias),
       (m.actuator_dyntype, types.DynType, mujoco.mjtDyn),
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
@@ -625,38 +625,6 @@ def test_contact_elliptic_condim1(self):
     with self.assertRaises(NotImplementedError):
       mjx.make_data(m)
 
-  @parameterized.product(
-      sensor=['accelerometer', 'force', 'torque'], equality=['connect', 'weld']
-  )
-  def test_sensor_constraint_compatibility(self, sensor, equality):
-    """Test unsupported sensor and equality constraint combinations."""
-    equality_constraint = f'{equality} body1="body1" body2="body2"'
-    if equality == 'connect':
-      equality_constraint += ' anchor="0 0 0"'
-    m = mujoco.MjModel.from_xml_string(f"""
-        <mujoco>
-          <worldbody>
-            <body name="body1">
-              <freejoint/>
-              <geom size="0.1"/>
-              <site name="site1"/>
-            </body>
-            <body name="body2">
-              <freejoint/>
-              <geom size="0.1"/>
-            </body>
-          </worldbody>
-          <equality>
-            <{equality_constraint}/>
-          </equality>
-          <sensor>
-            <{sensor} site="site1"/>
-          </sensor>
-        </mujoco>
-      """)
-    with self.assertRaises(NotImplementedError):
-      mjx.put_model(m, impl='jax')
-
   @parameterized.parameters(JacobianType.DENSE, JacobianType.SPARSE)
   def test_qm_mapm2m(self, jacobian):
     """Test that qM is mapped to M."""
diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py
@@ -29,6 +29,7 @@
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
 from mujoco.mjx._src.types import ModelJAX
+from mujoco.mjx._src.types import ObjType
 from mujoco.mjx._src.types import TrnType
 from mujoco.mjx._src.types import WrapType
 # pylint: enable=g-importing-member
@@ -659,11 +660,126 @@ def _contact_force_to_cfrc_ext(force, pos, frame, id1, id2, com1, com2):
         cfrc_contact.reshape((-1, 6))
     )
 
-  # TODO(taylorhowell): connect and weld constraints
-  if np.any(m.eq_type == EqType.CONNECT):
-    raise NotImplementedError('Connect constraints are not implemented.')
-  if np.any(m.eq_type == EqType.WELD):
-    raise NotImplementedError('Weld constraints are not implemented.')
+  # cfrc_ext += connect, weld
+  cfrc_ext_equality = []
+  cfrc_ext_equality_adr = []
+
+  connect_id = m.eq_type == EqType.CONNECT
+  nconnect = connect_id.sum()
+
+  if nconnect:
+    cfrc_connect_force = d._impl.efc_force[: 3 * nconnect].reshape(
+        (nconnect, 3)
+    )
+
+    is_site = m.eq_objtype == ObjType.SITE
+    body1id = np.copy(m.eq_obj1id)
+    body2id = np.copy(m.eq_obj2id)
+    pos1 = m.eq_data[:, :3]
+    pos2 = m.eq_data[:, 3:6]
+
+    if m.nsite:
+      body1id[is_site] = m.site_bodyid[m.eq_obj1id[is_site]]
+      body2id[is_site] = m.site_bodyid[m.eq_obj2id[is_site]]
+      pos1 = jp.where(is_site[:, None], m.site_pos[m.eq_obj1id], pos1)
+      pos2 = jp.where(is_site[:, None], m.site_pos[m.eq_obj2id], pos2)
+
+    # body 1
+    k1_connect = body1id[connect_id]
+    k1_connect_mask = k1_connect != 0
+    offset1_connect = pos1[connect_id]
+
+    pos1_connect = jax.vmap(lambda pnt, mat, vec: mat @ pnt + vec)(
+        offset1_connect, d.xmat[k1_connect], d.xpos[k1_connect]
+    )
+    subtree_com1_connect = d.subtree_com[jp.array(m.body_rootid)[k1_connect]]
+    cfrc_com1_connect = jax.vmap(
+        lambda dif, frc, mask: mask * jp.concatenate([-jp.cross(dif, frc), frc])
+    )(subtree_com1_connect - pos1_connect, cfrc_connect_force, k1_connect_mask)
+
+    # body 2
+    k2_connect = body2id[connect_id]
+    k2_connect_mask = -1 * (k2_connect != 0)
+    offset2_connect = pos2[connect_id]
+
+    pos2_connect = jax.vmap(lambda pnt, mat, vec: mat @ pnt + vec)(
+        offset2_connect, d.xmat[k2_connect], d.xpos[k2_connect]
+    )
+    subtree_com2_connect = d.subtree_com[jp.array(m.body_rootid)[k2_connect]]
+    cfrc_com2_connect = jax.vmap(
+        lambda dif, frc, mask: mask * jp.concatenate([-jp.cross(dif, frc), frc])
+    )(subtree_com2_connect - pos2_connect, cfrc_connect_force, k2_connect_mask)
+
+    cfrc_ext_equality.append(jp.vstack([cfrc_com1_connect, cfrc_com2_connect]))
+    cfrc_ext_equality_adr.append(jp.concatenate([k1_connect, k2_connect]))
+
+  weld_id = m.eq_type == EqType.WELD
+  nweld = weld_id.sum()
+
+  if nweld:
+    cfrc_weld = d._impl.efc_force[
+        3 * nconnect : 3 * nconnect + 6 * nweld
+    ].reshape((nweld, 6))
+    cfrc_weld_force = cfrc_weld[:, :3]
+    cfrc_weld_torque = cfrc_weld[:, 3:]
+
+    is_site = m.eq_objtype == ObjType.SITE
+    body1id = np.copy(m.eq_obj1id)
+    body2id = np.copy(m.eq_obj2id)
+    pos1 = m.eq_data[:, 3:6]
+    pos2 = m.eq_data[:, :3]
+
+    if m.nsite:
+      body1id[is_site] = m.site_bodyid[m.eq_obj1id[is_site]]
+      body2id[is_site] = m.site_bodyid[m.eq_obj2id[is_site]]
+      pos1 = jp.where(is_site[:, None], m.site_pos[m.eq_obj1id], pos1)
+      pos2 = jp.where(is_site[:, None], m.site_pos[m.eq_obj2id], pos2)
+
+    # body 1
+    k1_weld = body1id[weld_id]
+    k1_weld_mask = k1_weld != 0
+    offset1_weld = pos1[weld_id]
+
+    pos1_weld = jax.vmap(lambda pnt, mat, vec: mat @ pnt + vec)(
+        offset1_weld, d.xmat[k1_weld], d.xpos[k1_weld]
+    )
+    subtree_com1_weld = d.subtree_com[jp.array(m.body_rootid)[k1_weld]]
+    cfrc_com1_weld = jax.vmap(
+        lambda dif, frc, trq, mask: mask
+        * jp.concatenate([trq - jp.cross(dif, frc), frc])
+    )(
+        subtree_com1_weld - pos1_weld,
+        cfrc_weld_force,
+        cfrc_weld_torque,
+        k1_weld_mask,
+    )
+
+    # body 2
+    k2_weld = body2id[weld_id]
+    k2_weld_mask = -1 * (k2_weld != 0)
+    offset2_weld = pos2[weld_id]
+
+    pos2_weld = jax.vmap(lambda pnt, mat, vec: mat @ pnt + vec)(
+        offset2_weld, d.xmat[k2_weld], d.xpos[k2_weld]
+    )
+    subtree_com2_weld = d.subtree_com[jp.array(m.body_rootid)[k2_weld]]
+    cfrc_com2_weld = jax.vmap(
+        lambda dif, frc, trq, mask: mask
+        * jp.concatenate([trq - jp.cross(dif, frc), frc])
+    )(
+        subtree_com2_weld - pos2_weld,
+        cfrc_weld_force,
+        cfrc_weld_torque,
+        k2_weld_mask,
+    )
+
+    cfrc_ext_equality.append(jp.vstack([cfrc_com1_weld, cfrc_com2_weld]))
+    cfrc_ext_equality_adr.append(jp.concatenate([k1_weld, k2_weld]))
+
+  if nconnect or nweld:
+    cfrc_ext = cfrc_ext.at[jp.concatenate(cfrc_ext_equality_adr)].add(
+        jp.vstack(cfrc_ext_equality)
+    )
 
   # forward pass over bodies: compute cacc, cfrc_int
   def _forward(carry, cfrc_ext, cinert, cvel, body_dofadr, body_dofnum):
diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py
@@ -230,21 +230,98 @@ def test_subtree_vel(self):
 
 
 class RnePostConstraintTest(parameterized.TestCase):
+  _CONNECT_SITE = """
+    <equality>
+      <connect site1="site1" site2="site2"/>
+    </equality>
+    """
+  _CONNECT_BODY = """
+    <equality>
+      <connect body1="body1" body2="body2" anchor="1 2 3"/>
+    </equality>
+    """
+  _WELD_SITE = """
+    <equality>
+      <weld site1="site1" site2="site2"/>
+    </equality>
+    """
+  _WELD_BODY = """
+    <equality>
+      <weld body1="body1" body2="body2"/>
+    </equality>
+    """
+  _CONNECT_SITE_WELD_SITE = """
+    <equality>
+      <connect site1="site1" site2="site2"/>
+      <weld site1="site1" site2="site2"/>
+    </equality>
+    """
+  _WELD_SITE_CONNECT_SITE = """
+    <equality>
+      <weld site1="site1" site2="site2"/>
+      <connect site1="site1" site2="site2"/>
+    </equality>
+    """
+  _WELD_SITE_CONNECT_SITE_WELD_BODY = """
+    <equality>
+      <weld site1="site1" site2="site2"/>
+      <connect site1="site1" site2="site2"/>
+      <weld body1="body1" body2="body2"/>
+    </equality>
+    """
+  _CONNECT_SITE_WELD_SITE_WELD_BODY = """
+    <equality>
+      <connect site1="site1" site2="site2"/>
+      <weld site1="site1" site2="site2"/>
+      <weld body1="body1" body2="body2"/>
+    </equality>
+    """
+  _CONNECT_SITE_CONNECT_BODY_CONNECT_WELD = """
+    <equality>
+      <connect site1="site1" site2="site2"/>
+      <connect body1="body1" body2="body2" anchor="1 2 3"/>
+      <weld body1="body1" body2="body2"/>
+    </equality>
+    """
 
-  @parameterized.parameters(ConeType)
-  def test_rnepostconstraint(self, cone_type):
+  @parameterized.parameters(
+      ('', ConeType.PYRAMIDAL, None),
+      ('', ConeType.ELLIPTIC, None),
+      (_CONNECT_SITE, ConeType.PYRAMIDAL, None),
+      (_CONNECT_BODY, ConeType.PYRAMIDAL, None),
+      (_WELD_SITE, ConeType.PYRAMIDAL, None),
+      (_WELD_BODY, ConeType.PYRAMIDAL, None),
+      (_CONNECT_SITE_WELD_SITE, ConeType.PYRAMIDAL, None),
+      (
+          _WELD_SITE_CONNECT_SITE,
+          ConeType.PYRAMIDAL,
+          np.array([6, 7, 8, 0, 1, 2, 3, 4, 5]),
+      ),
+      (
+          _WELD_SITE_CONNECT_SITE_WELD_BODY,
+          ConeType.PYRAMIDAL,
+          np.array([6, 7, 8, 0, 1, 2, 3, 4, 5]),
+      ),
+      (_CONNECT_SITE_WELD_SITE_WELD_BODY, ConeType.PYRAMIDAL, None),
+      (_CONNECT_SITE_CONNECT_BODY_CONNECT_WELD, ConeType.PYRAMIDAL, None),
+  )
+  def test_rnepostconstraint(self, equality, cone_type, efc_map):
     """Tests MJX rne_postconstraint function to match MuJoCo mj_rnePostConstraint."""
 
-    m = mujoco.MjModel.from_xml_string("""
+    m = mujoco.MjModel.from_xml_string(f"""
         <mujoco>
           <worldbody>
             <geom name="floor" size="10 10 .05" type="plane"/>
-            <body pos="0 0 1">
+            <site name="site1"/>
+            <body name="body1">
+            </body>
+            <body pos="0 0 1" name="body2">
               <joint type="ball" damping="1"/>
               <geom type="capsule" size="0.1 0.5" fromto="0 0 0 0.5 0 0" condim="1"/>
               <body pos="0.5 0 0">
                 <joint type="ball" damping="1"/>
                 <geom type="capsule" size="0.1 0.5" fromto="0 0 0 0.5 0 0"  condim="3"/>
+                <site name="site2"/>
               </body>
             </body>
             <body pos="0 1 1">
@@ -256,6 +333,7 @@ def test_rnepostconstraint(self, cone_type):
               </body>
             </body>
           </worldbody>
+          {equality}
           <keyframe>
             <key qpos='0.424577 0.450592 0.451703 -0.642391 0.729379 0.545151 0.407756 0.0674697 0.424577 1.450592 0.451703 -0.642391 0.729379 0.545151 0.407756 0.0674697'/>
           </keyframe>
@@ -273,6 +351,11 @@ def test_rnepostconstraint(self, cone_type):
     mx = mjx.put_model(m)
     dx = mjx.put_data(m, d)
 
+    if efc_map is not None:
+      efc_force = d.efc_force.copy()
+      efc_force[: len(efc_map)] = d.efc_force[efc_map]
+      dx = dx.tree_replace({'_impl.efc_force': jp.array(efc_force)})
+
     # rne postconstraint
     mujoco.mj_rnePostConstraint(m, d)
     dx = jax.jit(mjx.rne_postconstraint)(mx, dx)
