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

doc/changelog.rst

@@ -24,6 +24,11 @@ MJX
 
 Bug fixes
 ^^^^^^^^^
+- Fixed several bugs related to connect and weld constraints with site semantics (fixes :github:issue:`2179`, reported
+  by :github:user:`yinfanyi`). 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`, but 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.
 - Fixed a bug in slider-crank :ref:`transmission<geTransmission>`. The bug was introduced in 3.0.0.
 
 Version 3.2.4 (Oct 15, 2024)

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];
+    mjtNum pos[3], *offset;
-    int k;
+    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]];
+      int b1 = m->eq_obj1id[efc_id];
-      tree[1] = m->body_treeid[m->eq_obj2id[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) {

test/engine/engine_core_constraint_test.cc

@@ -17,6 +17,7 @@
 #include <cstddef>
 #include <cstring>
 #include <string>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -249,6 +250,41 @@ TEST_F(CoreConstraintTest, JacobianPreAllocate) {
   }
 }
 
+TEST_F(CoreConstraintTest, EqualityBodySite) {
+  const std::string xml_path =
+      GetTestDataFilePath("engine/testdata/equality_site_body_compare.xml");
+
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+  mjData* data = mj_makeData(model);
+
+  // simulate, get diag(A)
+  while (data->time < 0.1) {
+    mj_step(model, data);
+  }
+  int nefc_site = data->nefc;
+  std::vector<mjtNum> dA = AsVector(data->efc_diagApprox, nefc_site);
+
+  // reset
+  mj_resetData(model, data);
+
+  // turn site-defined equalities off, equivalent body-defined equalities on
+  for (int e=0; e < 4; e++) data->eq_active[e] = 1 - data->eq_active[e];
+
+  // simulate again, get diag(A)
+  while (data->time < 0.1) {
+    mj_step(model, data);
+  }
+
+  // compare
+  EXPECT_EQ(nefc_site, data->nefc);
+  EXPECT_THAT(AsVector(data->efc_diagApprox, data->nefc),
+              Pointwise(DoubleNear(1e-12), dA));
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+
 static const char* const kIlslandEfcPath =
     "engine/testdata/island/island_efc.xml";
 

test/engine/engine_core_smooth_test.cc

@@ -39,9 +39,6 @@ using ::testing::DoubleNear;
 using ::testing::NotNull;
 using CoreSmoothTest = MujocoTest;
 
-std::vector<mjtNum> GetVector(const mjtNum* array, int length) {
-  return std::vector<mjtNum>(array, array + length);
-}
 
 constexpr bool EndsWith(std::string_view str, std::string_view suffix) {
   return str.size() >= suffix.size() &&
@@ -104,7 +101,7 @@ TEST_F(CoreSmoothTest, MjKinematicsWorldXipos) {
 
   mj_resetDataDebug(model, data, 'd');
   mj_kinematics(model, data);
-  EXPECT_THAT(GetVector(&data->xipos[0], 3), ElementsAre(0, 0, 0));
+  EXPECT_THAT(AsVector(&data->xipos[0], 3), ElementsAre(0, 0, 0));
 
   mj_deleteData(data);
   mj_deleteModel(model);
@@ -241,16 +238,50 @@ TEST_F(CoreSmoothTest, WeldRatioTorqueFree) {
 
 TEST_F(CoreSmoothTest, WeldRatioForceSlideRotated) {
   constexpr char kModelFilePath[] =
-      "engine/testdata/core_smooth/rne_post/weld/tfratio0_force_slide_rotated.xml";
+      "engine/testdata/core_smooth/rne_post/weld/"
+      "tfratio0_force_slide_rotated.xml";
   TestConnect(kModelFilePath);
 }
 
 TEST_F(CoreSmoothTest, WeldRatioMultipleConstraints) {
   constexpr char kModelFilePath[] =
-      "engine/testdata/core_smooth/rne_post/weld/tfratio0_multiple_constraints.xml";
+      "engine/testdata/core_smooth/rne_post/weld/"
+      "tfratio0_multiple_constraints.xml";
   TestConnect(kModelFilePath);
 }
 
+TEST_F(CoreSmoothTest, EqualityBodySite) {
+  const std::string xml_path =
+      GetTestDataFilePath("engine/testdata/equality_site_body_compare.xml");
+
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+  mjData* data = mj_makeData(model);
+
+  // simulate, get sensordata
+  while (data->time < 0.1) {
+    mj_step(model, data);
+  }
+  std::vector<mjtNum> sdata = AsVector(data->sensordata, model->nsensordata);
+
+  // reset
+  mj_resetData(model, data);
+
+  // turn site-defined equalities off, equivalent body-defined equalities on
+  for (int e=0; e < 4; e++) data->eq_active[e] = 1 - data->eq_active[e];
+
+  // simulate again, get sensordata
+  while (data->time < 0.1) {
+    mj_step(model, data);
+  }
+
+  // compare
+  EXPECT_THAT(AsVector(data->sensordata, model->nsensordata),
+              Pointwise(DoubleNear(1e-8), sdata));
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 // --------------------------- site actuators ----------------------------------
 
 // Test Cartesian position control using site transmission with refsite

test/engine/engine_island_test.cc

@@ -14,20 +14,15 @@
 
 // Tests for engine/engine_island.c.
 
-#include <array>
-#include <cstring>
 #include <string>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjthread.h>
-#include <mujoco/mjxmacro.h>
 #include <mujoco/mujoco.h>
 #include "src/engine/engine_island.h"
 #include "src/engine/engine_util_sparse.h"
-#include "src/thread/thread_pool.h"
 #include "test/fixture.h"
 
 namespace mujoco {
@@ -357,10 +352,10 @@ TEST_F(IslandTest, IslandEfc) {
 
   // expect island structure to correspond to comment at top of xml
   EXPECT_EQ(data->nisland, 4);
-  EXPECT_EQ(data->ne, 4);
+  EXPECT_EQ(data->ne, 7);
   EXPECT_EQ(data->nf, 2);
   EXPECT_EQ(data->nl, 1);
-  EXPECT_EQ(data->nefc, 27);
+  EXPECT_EQ(data->nefc, 30);
 
   mj_deleteData(data);
   mj_deleteModel(model);
@@ -378,10 +373,10 @@ TEST_F(IslandTest, IslandEfcElliptic) {
   mj_forward(model, data);
 
   EXPECT_EQ(data->nisland, 4);
-  EXPECT_EQ(data->ne, 4);
+  EXPECT_EQ(data->ne, 7);
   EXPECT_EQ(data->nf, 2);
   EXPECT_EQ(data->nl, 1);
-  EXPECT_EQ(data->nefc, 22);
+  EXPECT_EQ(data->nefc, 25);
 
   mj_deleteData(data);
   mj_deleteModel(model);

test/engine/engine_solver_test.cc

@@ -14,7 +14,10 @@
 
 // Tests for engine/engine_solver.c
 
+#include <algorithm>
+#include <cstdlib>
 #include <string>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -27,6 +30,28 @@ namespace {
 using ::testing::DoubleNear;
 using ::testing::NotNull;
 using ::testing::Pointwise;
+using ::std::vector;
+using ::std::abs;
+using ::std::max;
+
+// compare two vectors, relative error (reduces size of large vector elements)
+inline void ExpectEqRel(vector<mjtNum> v1, vector<mjtNum> v2, mjtNum rtol) {
+  ASSERT_TRUE(v1.size() == v2.size());
+
+  // make scale vector
+  int n = v1.size();
+  vector<mjtNum> scale(n);
+  for (int i = 0; i < n; i++) {
+    scale[i] = max(1.0, abs(v1[i]) + abs(v2[i]));
+  }
+
+  // scale and compare
+  for (int i = 0; i < n; i++) {
+    v1[i] /= scale[i];
+    v2[i] /= scale[i];
+  }
+  EXPECT_THAT(v1, Pointwise(DoubleNear(rtol), v2));
+}
 
 using SolverTest = MujocoTest;
 
@@ -51,7 +76,7 @@ TEST_F(SolverTest, IslandsEquivalent) {
   mjData* data_island = mj_makeData(model);
   mjData* data_noisland = mj_makeData(model);
 
-  mjtNum tol = 2e-4;
+  mjtNum rtol = 1e-5;
 
   for (bool warmstart : {true, false}) {
     if (warmstart) {
@@ -73,8 +98,8 @@ TEST_F(SolverTest, IslandsEquivalent) {
       mj_forward(model, data_island);
       model->opt.enableflags &= ~mjENBL_ISLAND;  // disable islands
 
-      EXPECT_THAT(AsVector(data_noisland->qacc, nv),
+      ExpectEqRel(AsVector(data_noisland->qacc, nv),
-                  Pointwise(DoubleNear(tol), AsVector(data_island->qacc, nv)));
+                  AsVector(data_island->qacc, nv), rtol);
     }
   }
 
@@ -102,7 +127,7 @@ TEST_F(SolverTest, OneBigIsland) {
   mjData* data_noisland = mj_makeData(model);
 
   int nv = model->nv;
-  mjtNum tol = 1e-8;
+  mjtNum rtol = 1e-7;
 
   // save current (default) iterations
   int iterations_default = model->opt.iterations;
@@ -144,9 +169,9 @@ TEST_F(SolverTest, OneBigIsland) {
     model->opt.enableflags &= ~mjENBL_ISLAND;
     model->opt.iterations = iterations_default;
 
-    // compare accelerations
+    // compare accelerations (relative error)
-    EXPECT_THAT(AsVector(data_noisland->qacc, nv),
+    ExpectEqRel(AsVector(data_noisland->qacc, nv),
-                Pointwise(DoubleNear(tol), AsVector(data_island->qacc, nv)));
+                AsVector(data_island->qacc, nv), rtol);
   }
 
   mj_deleteData(data_noisland);

test/engine/testdata/equality_site_body_compare.xml

@@ -0,0 +1,58 @@
+<mujoco model="Equivalent Connects and Welds defined with sites and bodies">
+  <option viscosity="1">
+    <flag contact="disable"/>
+  </option>
+
+  <visual>
+    <scale framewidth="0.05" framelength="0.4"/>
+  </visual>
+
+  <default>
+    <default class="box">
+      <geom type="box" size=".2" fromto=".2 .2 0 .2 .2 -1" rgba=".4 .7 .6 1"/>
+    </default>
+    <site rgba=".0 .7 0 1" size="0.05"/>
+  </default>
+
+  <worldbody>
+    <geom pos="0 0 -2" type="plane" size="10 10 .01"/>
+    <light pos="0 0 20"/>
+
+    <site name="a1" pos="-2 0 0" euler="-90 0 0"/>
+    <body name="a" pos="-2 1 0">
+      <freejoint/>
+      <geom class="box"/>
+      <site name="a2" pos="0 -1 0"/>
+      <body pos="0 0 -1">
+        <geom class="box" size=".15"/>
+        <site name="a_sensor"/>
+      </body>
+    </body>
+
+    <site name="b1" euler="-90 0 0"/>
+    <body name="b" pos="0 1 0">
+      <freejoint/>
+      <geom class="box"/>
+      <site name="b2" pos="0 -1 0"/>
+      <body pos="0 0 -1">
+        <geom class="box" size=".15"/>
+        <site name="b_sensor"/>
+      </body>
+    </body>
+  </worldbody>
+
+  <equality>
+    <!-- equivalent connect and weld constraints defined with sites (enabled) and bodies (disabled) -->
+    <connect name="connect site" site1="a1"    site2="a2"    active="true"/>
+    <connect name="connect body" body1="a"     body2="world" active="false" anchor="0 -1 0"/>
+    <weld    name="weld site"    site1="b1"    site2="b2"    active="true"/>
+    <weld    name="weld body"    body1="world" body2="b"     active="false" anchor="0 -1 0" relpose="0 0 0 1 -1 0 0"/>
+  </equality>
+
+  <sensor>
+    <force  name="a_force"  site="a_sensor"/>
+    <torque name="a_torque" site="a_sensor"/>
+    <force  name="b_force"  site="b_sensor"/>
+    <torque name="b_torque" site="b_sensor"/>
+  </sensor>
+</mujoco>

test/engine/testdata/island/island_efc.xml

@@ -1,6 +1,6 @@
 <!--
   test model for constraint islanding:
-- 4 equalities (1:joint, 3:connect)
+- 7 equalities (1:joint, 6:connect)
 - 1 limit
 - 2 friction constraints
 - 1 unconstrained dof
@@ -65,14 +65,17 @@
       <geom type="box" size=".03 .03 .03" pos="0.01 0.01 0.01"/>
     </body>
 
+    <site name="0" pos="-.45 -.05 .35"/>
     <body pos="-.5 0 .3" name="connect">
       <freejoint/>
       <geom type="box" size=".05 .05 .05"/>
+      <site name="1" pos=".05 -.05 .05"/>
     </body>
   </worldbody>
 
   <equality>
     <joint joint1="hinge1" joint2="hinge2"/>
     <connect body1="connect" body2="world" anchor="-.05 -.05 .05"/>
+    <connect site1="0" site2="1"/>
   </equality>
 </mujoco>