Fix bug in mj_mulM2.

PiperOrigin-RevId: 710756899
Change-Id: I30c0725859ba0f1a59750dc89eebb9eba84fa71b

doc/changelog.rst

@@ -16,6 +16,7 @@ Python bindings
 Bug fixes
 ^^^^^^^^^
 - Fixed a bug in the box-sphere collider, depth was incorrect for deep penetrations (:github:issue:`2206`).
+- Fixed a bug in :ref:`mj_mulM2` and added a test.
 
 Version 3.2.6 (Dec 2, 2024)
 ---------------------------

src/engine/engine_support.c

@@ -1085,51 +1085,28 @@ void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum
 
 // multiply vector by M^(1/2)
 void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
-  int adr, nv = m->nv;
+  int  nv = m->nv;
   const mjtNum* qLD = d->qLD;
   const mjtNum* qLDiagSqrtInv = d->qLDiagSqrtInv;
   const int* dofMadr = m->dof_Madr;
 
   mju_zero(res, nv);
 
+  // res = L * vec
   for (int i=0; i < nv; i++) {
-#ifdef mjUSEAVX
+    // simple: diagonal
-    // simple: diagonal division, AVX
-    if (m->dof_simplenum[i] >= 4) {
-      // init
-      __m256d result, val1, val2;
-
-      // parallel computation
-      val1 = _mm256_loadu_pd(vec+i);
-      val2 = _mm256_set_pd(qLDiagSqrtInv[dofMadr[i+3]],
-                           qLDiagSqrtInv[dofMadr[i+2]],
-                           qLDiagSqrtInv[dofMadr[i+1]],
-                           qLDiagSqrtInv[dofMadr[i+0]]);
-      result = _mm256_div_pd(val1, val2);
-
-      // store result
-      _mm256_storeu_pd(res+i, result);
-
-      // skip rest of block
-      i += 3;
-      continue;
-    }
-#endif
-
-    // simple: diagonal division
     if (m->dof_simplenum[i]) {
-      res[i] = vec[i]/qLDiagSqrtInv[i];
+      res[i] = vec[i];
     }
 
     // regular: full multiplication
     else {
       // diagonal
-      adr = dofMadr[i];
+      res[i] += vec[i];
-      res[i] += vec[i]/qLDiagSqrtInv[i];
 
       // off-diagonal
       int j = m->dof_parentid[i];
-      adr++;
+      int adr = dofMadr[i] + 1;
       while (j >= 0) {
         res[i] += qLD[adr]*vec[j];
 
@@ -1139,6 +1116,11 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
       }
     }
   }
+
+  // res = sqrt(D) * res
+  for (int i=0; i < nv; i++) {
+    res[i] /= qLDiagSqrtInv[i];
+  }
 }
 
 

test/engine/engine_core_smooth_test.cc

@@ -407,39 +407,61 @@ TEST_F(CoreSmoothTest, SolveMIsland) {
   mj_deleteModel(model);
 }
 
-TEST_F(CoreSmoothTest, SolveLD2) {
+static const char* const kInertiaPath = "engine/testdata/inertia.xml";
-  static constexpr char xml[] = R"(
+
-  <mujoco>
+TEST_F(CoreSmoothTest, FactorI) {
-    <default>
+  const std::string xml_path = GetTestDataFilePath(kInertiaPath);
-      <geom type="capsule" size="0.1"/>
+  char error[1024];
-      <joint axis="0 1 0"/>
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
-    </default>
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error;
+
+  mjData* data = mj_makeData(model);
+  mj_forward(model, data);
+
+  // dense L matrix
+  int nv = model->nv;
+  vector<mjtNum> Ldense(nv*nv);
+  mj_fullM(model, Ldense.data(), data->qLD);
+  // clear upper triangle, set diagonal to 1
+  for (int i=0; i < nv; i++) {
+    for (int j=i; j < nv; j++) {
+      Ldense[i*nv+j] = i == j ? 1 : 0;
+    }
+  }
+
+  // dense D matrix
+  vector<mjtNum> Ddense(nv*nv);
+  mj_fullM(model, Ddense.data(), data->qLD);
+  // clear everything but the diagonal
+  for (int i=0; i < nv; i++) {
+    for (int j=0; j < nv; j++) {
+      if (i != j) Ddense[i*nv+j] = 0;
+    }
+  }
+
+  // perform multiplication: M = L^T * D * L
+  vector<mjtNum> tmp(nv*nv);
+  vector<mjtNum> M(nv*nv);
+  mju_mulMatMat(tmp.data(), Ddense.data(), Ldense.data(), nv, nv, nv);
+  mju_mulMatTMat(M.data(), Ldense.data(), tmp.data(), nv, nv, nv);
+
+  // dense M matrix
+  vector<mjtNum> Mexpected(nv*nv);
+  mj_fullM(model, Mexpected.data(), data->qM);
+
+  // expect matrices to match to floating point precision
+  EXPECT_THAT(M, Pointwise(DoubleNear(1e-12), Mexpected));
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+TEST_F(CoreSmoothTest, SolveLD2) {
+  const std::string xml_path = GetTestDataFilePath(kInertiaPath);
+  char error[1024];
+  mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error;
 
-    <worldbody>
-      <body>
-        <geom fromto="0 0 0 0 0 1"/>
-        <joint/>
-        <body pos="0 0 1">
-          <geom fromto="0 0 0 1 0 1"/>
-          <joint/>
-        </body>
-        <body pos="0 0 1">
-          <geom fromto="0 0 0 -1 0 1"/>
-          <joint/>
-          <body pos="-1 0 1">
-            <geom fromto="0 0 0 1 0 1"/>
-            <joint/>
-          </body>
-          <body pos="-1 0 1">
-            <geom fromto="0 0 0 -1 0 1"/>
-            <joint/>
-          </body>
-        </body>
-      </body>
-    </worldbody>
-  </mujoco>
-  )";
-  mjModel* m = LoadModelFromString(xml);
   mjData* d = mj_makeData(m);
   mj_forward(m, d);
 

test/engine/engine_support_test.cc

@@ -29,12 +29,13 @@
 namespace mujoco {
 namespace {
 
-using ::testing::DoubleNear;
+using ::std::vector;
-using ::testing::Eq;
 using ::testing::ContainsRegex;  // NOLINT
-using ::testing::MatchesRegex;
+using ::testing::DoubleNear;
-using ::testing::Pointwise;
 using ::testing::ElementsAreArray;
+using ::testing::Eq;
+using ::testing::MatchesRegex;
+using ::testing::NotNull;
 using ::testing::Pointwise;
 
 using AngMomMatTest = MujocoTest;
@@ -685,9 +686,9 @@ TEST_F(SupportTest, GetSetStateStepEqual) {
   mj_deleteModel(model);
 }
 
-using AddMTest = MujocoTest;
+using InertiaTest = MujocoTest;
 
-TEST_F(AddMTest, DenseSameAsSparse) {
+TEST_F(InertiaTest, DenseSameAsSparse) {
   mjModel* m = LoadModelFromPath("humanoid/humanoid100.xml");
   mjData* d = mj_makeData(m);
   int nv = m->nv;
@@ -732,6 +733,72 @@ TEST_F(AddMTest, DenseSameAsSparse) {
   mj_deleteModel(m);
 }
 
+static const char* const kInertiaPath = "engine/testdata/inertia.xml";
+
+TEST_F(InertiaTest, mulM) {
+  const std::string xml_path = GetTestDataFilePath(kInertiaPath);
+  char error[1024];
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error;
+  int nv = model->nv;
+
+  mjData* data = mj_makeData(model);
+  mj_forward(model, data);
+
+  // dense M matrix
+  vector<mjtNum> Mdense(nv*nv);
+  mj_fullM(model, Mdense.data(), data->qM);
+
+  // arbitrary RHS vector
+  vector<mjtNum> vec(nv);
+  for (int i=0; i < nv; i++) vec[i] = vec[i] = 20 + 30*i;
+
+  // multiply directly
+  vector<mjtNum> res1(nv, 0);
+  mju_mulMatVec(res1.data(), Mdense.data(), vec.data(), nv, nv);
+
+  // multiply with mj_mulM
+  vector<mjtNum> res2(nv, 0);
+  mj_mulM(model, data, res2.data(), vec.data());
+
+  // expect vectors to match to floating point precision
+  EXPECT_THAT(res1, Pointwise(DoubleNear(1e-10), res2));
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+TEST_F(InertiaTest, mulM2) {
+  const std::string xml_path = GetTestDataFilePath(kInertiaPath);
+  char error[1024];
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error;
+  int nv = model->nv;
+
+  mjData* data = mj_makeData(model);
+  mj_forward(model, data);
+
+  // arbitrary RHS vector
+  vector<mjtNum> vec(nv);
+  for (int i=0; i < nv; i++) vec[i] = .2 + .3*i;
+
+  // multiply sqrtMvec = M^1/2 * vec
+  vector<mjtNum> sqrtMvec(nv);
+  mj_mulM2(model, data, sqrtMvec.data(), vec.data());
+
+  // multiply Mvec = M * vec
+  vector<mjtNum> Mvec(nv);
+  mj_mulM(model, data, Mvec.data(), vec.data());
+
+  // compute vec' * M * vec in two different ways, expect them to match
+  mjtNum sqrtMvec2 = mju_dot(sqrtMvec.data(), sqrtMvec.data(), nv);
+  mjtNum vecMvec = mju_dot(vec.data(), Mvec.data(), nv);
+  EXPECT_FLOAT_EQ(sqrtMvec2, vecMvec);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 static const char* const kIlslandEfcPath =
     "engine/testdata/island/island_efc.xml";
 

test/engine/testdata/inertia.xml

@@ -0,0 +1,39 @@
+<mujoco model="model for inertia-related tests">
+  <default>
+    <geom type="capsule" size="0.1"/>
+    <joint axis="0 1 0"/>
+  </default>
+
+  <worldbody>
+    <body pos="1 -.5 0" euler="10 20 30">
+      <geom type="box" size="0.15 .2 .25" pos=".1 .2 .3" euler="10 20 30"/>
+      <freejoint name="6 non-simple dofs" align="false"/>
+    </body>
+
+    <body pos="0 -.5 0">
+      <geom type="sphere" size="0.2"/>
+      <joint type="ball" name="3 simple dofs"/>
+    </body>
+
+    <body>
+      <geom fromto="0 0 0 0 0 1"/>
+      <joint/>
+      <body pos="0 0 1">
+        <geom fromto="0 0 0 1 0 1"/>
+        <joint/>
+      </body>
+      <body pos="0 0 1">
+        <geom fromto="0 0 0 -1 0 1"/>
+        <joint/>
+        <body pos="-1 0 1">
+          <geom fromto="0 0 0 1 0 1"/>
+          <joint/>
+        </body>
+        <body pos="-1 0 1">
+          <geom fromto="0 0 0 -1 0 1"/>
+          <joint/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>