Skip to content

Commit ec32264

Browse files
yuvaltassacopybara-github
yuvaltassa
authored and
copybara-github
committed
Fix bug in mj_mulM2.
PiperOrigin-RevId: 710756899 Change-Id: I30c0725859ba0f1a59750dc89eebb9eba84fa71b
1 parent e42370c commit ec32264

File tree

5 files changed

+176
-65
lines changed

5 files changed

+176
-65
lines changed

‎doc/changelog.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ Python bindings
1616
Bug fixes
1717
^^^^^^^^^
1818
- Fixed a bug in the box-sphere collider, depth was incorrect for deep penetrations (:github:issue:`2206`).
19+
- Fixed a bug in :ref:`mj_mulM2` and added a test.
1920

2021
Version 3.2.6 (Dec 2, 2024)
2122
---------------------------

‎src/engine/engine_support.c‎

Lines changed: 11 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1085,51 +1085,28 @@ void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum
10851085

10861086
// multiply vector by M^(1/2)
10871087
void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
1088-
int adr, nv = m->nv;
1088+
int nv = m->nv;
10891089
const mjtNum* qLD = d->qLD;
10901090
const mjtNum* qLDiagSqrtInv = d->qLDiagSqrtInv;
10911091
const int* dofMadr = m->dof_Madr;
10921092

10931093
mju_zero(res, nv);
10941094

1095+
// res = L * vec
10951096
for (int i=0; i < nv; i++) {
1096-
#ifdef mjUSEAVX
1097-
// simple: diagonal division, AVX
1098-
if (m->dof_simplenum[i] >= 4) {
1099-
// init
1100-
__m256d result, val1, val2;
1101-
1102-
// parallel computation
1103-
val1 = _mm256_loadu_pd(vec+i);
1104-
val2 = _mm256_set_pd(qLDiagSqrtInv[dofMadr[i+3]],
1105-
qLDiagSqrtInv[dofMadr[i+2]],
1106-
qLDiagSqrtInv[dofMadr[i+1]],
1107-
qLDiagSqrtInv[dofMadr[i+0]]);
1108-
result = _mm256_div_pd(val1, val2);
1109-
1110-
// store result
1111-
_mm256_storeu_pd(res+i, result);
1112-
1113-
// skip rest of block
1114-
i += 3;
1115-
continue;
1116-
}
1117-
#endif
1118-
1119-
// simple: diagonal division
1097+
// simple: diagonal
11201098
if (m->dof_simplenum[i]) {
1121-
res[i] = vec[i]/qLDiagSqrtInv[i];
1099+
res[i] = vec[i];
11221100
}
11231101

11241102
// regular: full multiplication
11251103
else {
11261104
// diagonal
1127-
adr = dofMadr[i];
1128-
res[i] += vec[i]/qLDiagSqrtInv[i];
1105+
res[i] += vec[i];
11291106

11301107
// off-diagonal
11311108
int j = m->dof_parentid[i];
1132-
adr++;
1109+
int adr = dofMadr[i] + 1;
11331110
while (j >= 0) {
11341111
res[i] += qLD[adr]*vec[j];
11351112

@@ -1139,6 +1116,11 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
11391116
}
11401117
}
11411118
}
1119+
1120+
// res = sqrt(D) * res
1121+
for (int i=0; i < nv; i++) {
1122+
res[i] /= qLDiagSqrtInv[i];
1123+
}
11421124
}
11431125

11441126

‎test/engine/engine_core_smooth_test.cc‎

Lines changed: 53 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -407,39 +407,61 @@ TEST_F(CoreSmoothTest, SolveMIsland) {
407407
mj_deleteModel(model);
408408
}
409409

410+
static const char* const kInertiaPath = "engine/testdata/inertia.xml";
411+
412+
TEST_F(CoreSmoothTest, FactorI) {
413+
const std::string xml_path = GetTestDataFilePath(kInertiaPath);
414+
char error[1024];
415+
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
416+
ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error;
417+
418+
mjData* data = mj_makeData(model);
419+
mj_forward(model, data);
420+
421+
// dense L matrix
422+
int nv = model->nv;
423+
vector<mjtNum> Ldense(nv*nv);
424+
mj_fullM(model, Ldense.data(), data->qLD);
425+
// clear upper triangle, set diagonal to 1
426+
for (int i=0; i < nv; i++) {
427+
for (int j=i; j < nv; j++) {
428+
Ldense[i*nv+j] = i == j ? 1 : 0;
429+
}
430+
}
431+
432+
// dense D matrix
433+
vector<mjtNum> Ddense(nv*nv);
434+
mj_fullM(model, Ddense.data(), data->qLD);
435+
// clear everything but the diagonal
436+
for (int i=0; i < nv; i++) {
437+
for (int j=0; j < nv; j++) {
438+
if (i != j) Ddense[i*nv+j] = 0;
439+
}
440+
}
441+
442+
// perform multiplication: M = L^T * D * L
443+
vector<mjtNum> tmp(nv*nv);
444+
vector<mjtNum> M(nv*nv);
445+
mju_mulMatMat(tmp.data(), Ddense.data(), Ldense.data(), nv, nv, nv);
446+
mju_mulMatTMat(M.data(), Ldense.data(), tmp.data(), nv, nv, nv);
447+
448+
// dense M matrix
449+
vector<mjtNum> Mexpected(nv*nv);
450+
mj_fullM(model, Mexpected.data(), data->qM);
451+
452+
// expect matrices to match to floating point precision
453+
EXPECT_THAT(M, Pointwise(DoubleNear(1e-12), Mexpected));
454+
455+
mj_deleteData(data);
456+
mj_deleteModel(model);
457+
}
458+
410459
TEST_F(CoreSmoothTest, SolveLD2) {
411-
static constexpr char xml[] = R"(
412-
<mujoco>
413-
<default>
414-
<geom type="capsule" size="0.1"/>
415-
<joint axis="0 1 0"/>
416-
</default>
460+
const std::string xml_path = GetTestDataFilePath(kInertiaPath);
461+
char error[1024];
462+
mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
463+
ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error;
417464

418-
<worldbody>
419-
<body>
420-
<geom fromto="0 0 0 0 0 1"/>
421-
<joint/>
422-
<body pos="0 0 1">
423-
<geom fromto="0 0 0 1 0 1"/>
424-
<joint/>
425-
</body>
426-
<body pos="0 0 1">
427-
<geom fromto="0 0 0 -1 0 1"/>
428-
<joint/>
429-
<body pos="-1 0 1">
430-
<geom fromto="0 0 0 1 0 1"/>
431-
<joint/>
432-
</body>
433-
<body pos="-1 0 1">
434-
<geom fromto="0 0 0 -1 0 1"/>
435-
<joint/>
436-
</body>
437-
</body>
438-
</body>
439-
</worldbody>
440-
</mujoco>
441-
)";
442-
mjModel* m = LoadModelFromString(xml);
443465
mjData* d = mj_makeData(m);
444466
mj_forward(m, d);
445467

‎test/engine/engine_support_test.cc‎

Lines changed: 72 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,13 @@
2929
namespace mujoco {
3030
namespace {
3131

32+
using ::std::vector;
33+
using ::testing::ContainsRegex; // NOLINT
3234
using ::testing::DoubleNear;
35+
using ::testing::ElementsAreArray;
3336
using ::testing::Eq;
34-
using ::testing::ContainsRegex; // NOLINT
3537
using ::testing::MatchesRegex;
36-
using ::testing::Pointwise;
37-
using ::testing::ElementsAreArray;
38+
using ::testing::NotNull;
3839
using ::testing::Pointwise;
3940

4041
using AngMomMatTest = MujocoTest;
@@ -685,9 +686,9 @@ TEST_F(SupportTest, GetSetStateStepEqual) {
685686
mj_deleteModel(model);
686687
}
687688

688-
using AddMTest = MujocoTest;
689+
using InertiaTest = MujocoTest;
689690

690-
TEST_F(AddMTest, DenseSameAsSparse) {
691+
TEST_F(InertiaTest, DenseSameAsSparse) {
691692
mjModel* m = LoadModelFromPath("humanoid/humanoid100.xml");
692693
mjData* d = mj_makeData(m);
693694
int nv = m->nv;
@@ -732,6 +733,72 @@ TEST_F(AddMTest, DenseSameAsSparse) {
732733
mj_deleteModel(m);
733734
}
734735

736+
static const char* const kInertiaPath = "engine/testdata/inertia.xml";
737+
738+
TEST_F(InertiaTest, mulM) {
739+
const std::string xml_path = GetTestDataFilePath(kInertiaPath);
740+
char error[1024];
741+
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
742+
ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error;
743+
int nv = model->nv;
744+
745+
mjData* data = mj_makeData(model);
746+
mj_forward(model, data);
747+
748+
// dense M matrix
749+
vector<mjtNum> Mdense(nv*nv);
750+
mj_fullM(model, Mdense.data(), data->qM);
751+
752+
// arbitrary RHS vector
753+
vector<mjtNum> vec(nv);
754+
for (int i=0; i < nv; i++) vec[i] = vec[i] = 20 + 30*i;
755+
756+
// multiply directly
757+
vector<mjtNum> res1(nv, 0);
758+
mju_mulMatVec(res1.data(), Mdense.data(), vec.data(), nv, nv);
759+
760+
// multiply with mj_mulM
761+
vector<mjtNum> res2(nv, 0);
762+
mj_mulM(model, data, res2.data(), vec.data());
763+
764+
// expect vectors to match to floating point precision
765+
EXPECT_THAT(res1, Pointwise(DoubleNear(1e-10), res2));
766+
767+
mj_deleteData(data);
768+
mj_deleteModel(model);
769+
}
770+
771+
TEST_F(InertiaTest, mulM2) {
772+
const std::string xml_path = GetTestDataFilePath(kInertiaPath);
773+
char error[1024];
774+
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
775+
ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error;
776+
int nv = model->nv;
777+
778+
mjData* data = mj_makeData(model);
779+
mj_forward(model, data);
780+
781+
// arbitrary RHS vector
782+
vector<mjtNum> vec(nv);
783+
for (int i=0; i < nv; i++) vec[i] = .2 + .3*i;
784+
785+
// multiply sqrtMvec = M^1/2 * vec
786+
vector<mjtNum> sqrtMvec(nv);
787+
mj_mulM2(model, data, sqrtMvec.data(), vec.data());
788+
789+
// multiply Mvec = M * vec
790+
vector<mjtNum> Mvec(nv);
791+
mj_mulM(model, data, Mvec.data(), vec.data());
792+
793+
// compute vec' * M * vec in two different ways, expect them to match
794+
mjtNum sqrtMvec2 = mju_dot(sqrtMvec.data(), sqrtMvec.data(), nv);
795+
mjtNum vecMvec = mju_dot(vec.data(), Mvec.data(), nv);
796+
EXPECT_FLOAT_EQ(sqrtMvec2, vecMvec);
797+
798+
mj_deleteData(data);
799+
mj_deleteModel(model);
800+
}
801+
735802
static const char* const kIlslandEfcPath =
736803
"engine/testdata/island/island_efc.xml";
737804

‎test/engine/testdata/inertia.xml‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
<mujoco model="model for inertia-related tests">
2+
<default>
3+
<geom type="capsule" size="0.1"/>
4+
<joint axis="0 1 0"/>
5+
</default>
6+
7+
<worldbody>
8+
<body pos="1 -.5 0" euler="10 20 30">
9+
<geom type="box" size="0.15 .2 .25" pos=".1 .2 .3" euler="10 20 30"/>
10+
<freejoint name="6 non-simple dofs" align="false"/>
11+
</body>
12+
13+
<body pos="0 -.5 0">
14+
<geom type="sphere" size="0.2"/>
15+
<joint type="ball" name="3 simple dofs"/>
16+
</body>
17+
18+
<body>
19+
<geom fromto="0 0 0 0 0 1"/>
20+
<joint/>
21+
<body pos="0 0 1">
22+
<geom fromto="0 0 0 1 0 1"/>
23+
<joint/>
24+
</body>
25+
<body pos="0 0 1">
26+
<geom fromto="0 0 0 -1 0 1"/>
27+
<joint/>
28+
<body pos="-1 0 1">
29+
<geom fromto="0 0 0 1 0 1"/>
30+
<joint/>
31+
</body>
32+
<body pos="-1 0 1">
33+
<geom fromto="0 0 0 -1 0 1"/>
34+
<joint/>
35+
</body>
36+
</body>
37+
</body>
38+
</worldbody>
39+
</mujoco>

0 commit comments

Comments
 (0)