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

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),
-                  Pointwise(DoubleNear(tol), AsVector(data_island->qacc, nv)));
+      ExpectEqRel(AsVector(data_noisland->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
-    EXPECT_THAT(AsVector(data_noisland->qacc, nv),
-                Pointwise(DoubleNear(tol), AsVector(data_island->qacc, nv)));
+    // compare accelerations (relative error)
+    ExpectEqRel(AsVector(data_noisland->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>