Fix multiple bugs related to connect and weld constraints with site semantics. Fixes #2179

The introduction of site specification to connects and welds in 3.2.3 conditionally changed the semantics of `mjData.eq_obj1id` and `mjData.eq_obj2id`. These changes were not properly propagated in several places leading to incorrect computations of constraint inertia, readings of affected force/torque sensors and runtime enabling/disabling of such constraints.

PiperOrigin-RevId: 690670420
Change-Id: I55ee8a013cbee8457f8d6c7f33c2981aedafbab6

src/engine/engine_core_constraint.c

@@ -1117,16 +1117,30 @@ void mj_diagApprox(const mjModel* m, mjData* d) {
       // process according to equality-constraint type
       switch (m->eq_type[id]) {
       case mjEQ_CONNECT:
-        // body translation
         b1 = m->eq_obj1id[id];
         b2 = m->eq_obj2id[id];
+
+        // get body ids if using site semantics
+        if (m->eq_objtype[id] == mjOBJ_SITE) {
+          b1 = m->site_bodyid[b1];
+          b2 = m->site_bodyid[b2];
+        }
+
+        // body translation
         dA[i] = m->body_invweight0[2*b1] + m->body_invweight0[2*b2];
         break;
 
       case mjEQ_WELD:  // distinguish translation and rotation inertia
-        // body translation or rotation depending on weldcnt
         b1 = m->eq_obj1id[id];
         b2 = m->eq_obj2id[id];
+
+        // get body ids if using site semantics
+        if (m->eq_objtype[id] == mjOBJ_SITE) {
+          b1 = m->site_bodyid[b1];
+          b2 = m->site_bodyid[b2];
+        }
+
+        // body translation or rotation depending on weldcnt
         dA[i] = m->body_invweight0[2*b1 + (weldcnt > 2)] +
                 m->body_invweight0[2*b2 + (weldcnt > 2)];
         weldcnt = (weldcnt + 1) % 6;
@@ -1650,6 +1664,12 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
           break;
         }
 
+        // get body ids if using site semantics
+        if (m->eq_objtype[i] == mjOBJ_SITE) {
+          id[0] = m->site_bodyid[id[0]];
+          id[1] = m->site_bodyid[id[1]];
+        }
+
         NV = mj_jacDifPairCount(m, chain, id[1], id[0], issparse);
         break;
 
@@ -1659,6 +1679,12 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
           break;
         }
 
+        // get body ids if using site semantics
+        if (m->eq_objtype[i] == mjOBJ_SITE) {
+          id[0] = m->site_bodyid[id[0]];
+          id[1] = m->site_bodyid[id[1]];
+        }
+
         NV = mj_jacDifPairCount(m, chain, id[1], id[0], issparse);
         break;
 

src/engine/engine_core_smooth.c

@@ -1902,7 +1902,7 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
       }
     }
 
-  // cfrc_ext += connect and weld constraints
+  // cfrc_ext += connect, weld, flex constraints
   int i = 0, ne = d->ne;
   while (i < ne) {
     if (d->efc_type[i] != mjCNSTR_EQUALITY)
@@ -1910,8 +1910,8 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
 
     int id = d->efc_id[i];
     mjtNum* eq_data = m->eq_data + mjNEQDATA*id;
-    mjtNum pos[3];
-    int k;
+    mjtNum pos[3], *offset;
+    int k, obj1, obj2, body_semantic;
     switch ((mjtEq) m->eq_type[id]) {
     case mjEQ_CONNECT:
     case mjEQ_WELD:
@@ -1923,10 +1923,17 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
         mju_zero3(cfrc);  // no torque from connect
       }
 
+      body_semantic = m->eq_objtype[id] == mjOBJ_BODY;
+
       // body 1
-      if ((k = m->eq_obj1id[id])) {
+      obj1 = m->eq_obj1id[id];
+      k = body_semantic ? obj1 : m->site_bodyid[obj1];
+      if (k) {
+        offset = body_semantic ? eq_data + 3 * (m->eq_type[id] == mjEQ_WELD) :
+                                 m->site_pos + 3 * obj1;
+
         // transform point on body1: local -> global
-        mj_local2Global(d, pos, 0, eq_data + 3*(m->eq_type[id] == mjEQ_WELD), 0, k, 0);
+        mj_local2Global(d, pos, 0, offset, 0, k, 0);
 
         // tmp = subtree CoM-based torque_force vector
         mju_transformSpatial(cfrc_com, cfrc, 1, d->subtree_com+3*m->body_rootid[k], pos, 0);
@@ -1936,9 +1943,14 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
       }
 
       // body 2
-      if ((k = m->eq_obj2id[id])) {
+      obj2 = m->eq_obj2id[id];
+      k = body_semantic ? obj2 : m->site_bodyid[obj2];
+      if (k) {
+        offset = body_semantic ? eq_data + 3 * (m->eq_type[id] == mjEQ_CONNECT) :
+                                 m->site_pos + 3 * obj2;
+
         // transform point on body2: local -> global
-        mj_local2Global(d, pos, 0, eq_data + 3*(m->eq_type[id] == mjEQ_CONNECT), 0, k, 0);
+        mj_local2Global(d, pos, 0, offset, 0, k, 0);
 
         // tmp = subtree CoM-based torque_force vector
         mju_transformSpatial(cfrc_com, cfrc, 1, d->subtree_com+3*m->body_rootid[k], pos, 0);

src/engine/engine_island.c

@@ -257,8 +257,17 @@ static int treeFirst(const mjModel* m, const mjData* d, int tree[2], int i) {
   if (efc_type == mjCNSTR_EQUALITY) {
     mjtEq eq_type = m->eq_type[efc_id];
     if (eq_type == mjEQ_CONNECT || eq_type == mjEQ_WELD) {
-      tree[0] = m->body_treeid[m->eq_obj1id[efc_id]];
-      tree[1] = m->body_treeid[m->eq_obj2id[efc_id]];
+      int b1 = m->eq_obj1id[efc_id];
+      int b2 = m->eq_obj2id[efc_id];
+
+      // get body ids if using site semantics
+      if (m->eq_objtype[efc_id] == mjOBJ_SITE) {
+        b1 = m->site_bodyid[b1];
+        b2 = m->site_bodyid[b2];
+      }
+
+      tree[0] = m->body_treeid[b1];
+      tree[1] = m->body_treeid[b2];
 
       // handle static bodies
       if (tree[0] < 0) {