Fix bug in mj_mulM2.

PiperOrigin-RevId: 710756899
Change-Id: I30c0725859ba0f1a59750dc89eebb9eba84fa71b
This commit is contained in:
Yuval Tassa
2024-12-30 11:14:00 -08:00
committed by Copybara-Service
parent e42370c982
commit ec322641b7
5 changed files with 178 additions and 67 deletions
+1
View File
@@ -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)
---------------------------
+11 -29
View File
@@ -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];
}
}
+54 -32
View File
@@ -407,39 +407,61 @@ TEST_F(CoreSmoothTest, SolveMIsland) {
mj_deleteModel(model);
}
TEST_F(CoreSmoothTest, SolveLD2) {
static constexpr char xml[] = R"(
<mujoco>
<default>
<geom type="capsule" size="0.1"/>
<joint axis="0 1 0"/>
</default>
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<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);
+73 -6
View File
@@ -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<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";
+39
View File
@@ -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>