From ec322641b78188fcd15c98dc60ea0f9dc3f6223e Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 30 Dec 2024 11:14:00 -0800 Subject: [PATCH] Fix bug in `mj_mulM2`. PiperOrigin-RevId: 710756899 Change-Id: I30c0725859ba0f1a59750dc89eebb9eba84fa71b --- doc/changelog.rst | 1 + src/engine/engine_support.c | 40 ++++-------- test/engine/engine_core_smooth_test.cc | 86 ++++++++++++++++---------- test/engine/engine_support_test.cc | 79 +++++++++++++++++++++-- test/engine/testdata/inertia.xml | 39 ++++++++++++ 5 files changed, 178 insertions(+), 67 deletions(-) create mode 100644 test/engine/testdata/inertia.xml diff --git a/doc/changelog.rst b/doc/changelog.rst index a4988dfb..c5ec24be 100644 --- a/doc/changelog.rst +++ b/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) --------------------------- diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 12538817..94a03de4 100644 --- a/src/engine/engine_support.c +++ b/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 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 + // simple: diagonal 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]/qLDiagSqrtInv[i]; + res[i] += vec[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]; + } } diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index f27d1670..a6d2d0b3 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -407,39 +407,61 @@ TEST_F(CoreSmoothTest, SolveMIsland) { mj_deleteModel(model); } -TEST_F(CoreSmoothTest, SolveLD2) { - static constexpr char xml[] = R"( - - - - - +static const char* const kInertiaPath = "engine/testdata/inertia.xml"; + +TEST_F(CoreSmoothTest, FactorI) { + 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; + + mjData* data = mj_makeData(model); + mj_forward(model, data); + + // dense L matrix + int nv = model->nv; + vector 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 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 tmp(nv*nv); + vector 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 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; - - - - - - - - - - - - - - - - - - - - - - - - )"; - mjModel* m = LoadModelFromString(xml); mjData* d = mj_makeData(m); mj_forward(m, d); diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc index 36dd6e06..456143cd 100644 --- a/test/engine/engine_support_test.cc +++ b/test/engine/engine_support_test.cc @@ -29,12 +29,13 @@ namespace mujoco { namespace { -using ::testing::DoubleNear; -using ::testing::Eq; +using ::std::vector; using ::testing::ContainsRegex; // NOLINT -using ::testing::MatchesRegex; -using ::testing::Pointwise; +using ::testing::DoubleNear; 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 Mdense(nv*nv); + mj_fullM(model, Mdense.data(), data->qM); + + // arbitrary RHS vector + vector vec(nv); + for (int i=0; i < nv; i++) vec[i] = vec[i] = 20 + 30*i; + + // multiply directly + vector res1(nv, 0); + mju_mulMatVec(res1.data(), Mdense.data(), vec.data(), nv, nv); + + // multiply with mj_mulM + vector 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 vec(nv); + for (int i=0; i < nv; i++) vec[i] = .2 + .3*i; + + // multiply sqrtMvec = M^1/2 * vec + vector sqrtMvec(nv); + mj_mulM2(model, data, sqrtMvec.data(), vec.data()); + + // multiply Mvec = M * vec + vector 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"; diff --git a/test/engine/testdata/inertia.xml b/test/engine/testdata/inertia.xml new file mode 100644 index 00000000..dc6a4a15 --- /dev/null +++ b/test/engine/testdata/inertia.xml @@ -0,0 +1,39 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +