Refactor islands to be memory contiguous.

PiperOrigin-RevId: 755803476
Change-Id: I41972b07e0d5ef5d0117c94f565b93367b87458b
This commit is contained in:
Yuval Tassa
2025-05-07 05:05:34 -07:00
committed by Copybara-Service
parent 449de73430
commit ecb769fc3a
30 changed files with 1742 additions and 1116 deletions
-85
View File
@@ -17,7 +17,6 @@
#include <algorithm>
#include <cstdlib>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
@@ -29,19 +28,9 @@ namespace {
using ::testing::DoubleNear;
using ::testing::NotNull;
using ::std::vector;
using ::std::abs;
using ::std::max;
// compare two vectors, relative error (increase tolerance for large elements)
inline void ExpectEqRel(vector<mjtNum> v1, vector<mjtNum> v2, mjtNum rtol) {
ASSERT_TRUE(v1.size() == v2.size());
for (int i = 0; i < v1.size(); i++) {
mjtNum scale = 0.5 * max(2.0, abs(v1[i]) + abs(v2[i]));
EXPECT_THAT(v1[i], DoubleNear(v2[i], scale*rtol));
}
}
using SolverTest = MujocoTest;
static const char* const kModelPath =
@@ -169,79 +158,5 @@ TEST_F(SolverTest, IslandsEquivalentForward) {
mj_deleteModel(model);
}
static const char* const kIlslandEfcPath =
"engine/testdata/island/island_efc.xml";
// compare qacc from 1 iteration of monolithic CG solver and one big island
TEST_F(SolverTest, OneBigIsland) {
const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
ASSERT_THAT(model, NotNull());
model->opt.solver = mjSOL_CG; // use CG solver
model->opt.disableflags |= mjDSBL_WARMSTART; // disable warmstart
model->opt.tolerance = 0; // set tolerance to 0
model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands
int state_size = mj_stateSize(model, mjSTATE_INTEGRATION);
mjtNum* state = (mjtNum*) mju_malloc(sizeof(mjtNum)*state_size);
mjData* data_island = mj_makeData(model);
mjData* data_noisland = mj_makeData(model);
int nv = model->nv;
mjtNum rtol = 1e-7;
// save current (default) iterations
int iterations_default = model->opt.iterations;
while (data_noisland->time < .2) {
// step and copy the state to data_island
mj_step(model, data_noisland);
mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION);
mj_setState(model, data_island, state, mjSTATE_INTEGRATION);
// set small number of iterations
model->opt.iterations = 1;
// call forward on data_noisland
mj_forward(model, data_noisland);
// enable islands
model->opt.enableflags |= mjENBL_ISLAND;
// call forward (just for smooth dynamics and to allocate islands)
mj_forward(model, data_island);
// overwrite island structure with one big island
data_island->nisland = 1;
data_island->island_dofnum[0] = nv;
data_island->island_dofadr[0] = 0;
for (int i = 0; i < nv; i++) {
data_island->island_dofind[i] = data_island->dof_islandind[i] = i;
}
int nefc = data_island->nefc;
data_island->island_efcnum[0] = nefc;
data_island->island_efcadr[0] = 0;
for (int i = 0; i < nefc; i++) data_island->island_efcind[i] = i;
// solve using using one big island
mj_fwdConstraint(model, data_island);
// re-disable islands and reset iterations
model->opt.enableflags &= ~mjENBL_ISLAND;
model->opt.iterations = iterations_default;
// compare accelerations (relative error)
ExpectEqRel(AsVector(data_noisland->qacc, nv),
AsVector(data_island->qacc, nv), rtol);
}
mj_deleteData(data_noisland);
mj_deleteData(data_island);
mju_free(state);
mj_deleteModel(model);
}
} // namespace
} // namespace mujoco