Improve warmstarting if island structure exists.

PiperOrigin-RevId: 752901506
Change-Id: I062d14df7da7d9dc1a2c1c30a16d0db7215af336
This commit is contained in:
Yuval Tassa
2025-04-29 15:17:13 -07:00
committed by Copybara-Service
parent 4c92c60988
commit 1766a388cc
3 changed files with 127 additions and 46 deletions
+9
View File
@@ -629,6 +629,15 @@ static void warmstart(const mjModel* m, mjData* d) {
}
}
// have island structure: unconstrained qacc = qacc_smooth
if (d->nisland > 0) {
for (int i=0; i < nv; i++) {
if (d->dof_island[i] < 0) {
d->qacc[i] = d->qacc_smooth[i];
}
}
}
mj_freeStack(d);
}
+100 -39
View File
@@ -29,88 +29,149 @@ namespace {
using ::testing::DoubleNear;
using ::testing::NotNull;
using ::testing::Pointwise;
using ::std::vector;
using ::std::abs;
using ::std::max;
// compare two vectors, relative error (reduces size of large vector elements)
// 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());
// make scale vector
int n = v1.size();
vector<mjtNum> scale(n);
for (int i = 0; i < n; i++) {
scale[i] = max(1.0, abs(v1[i]) + abs(v2[i]));
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));
}
// scale and compare
for (int i = 0; i < n; i++) {
v1[i] /= scale[i];
v2[i] /= scale[i];
}
EXPECT_THAT(v1, Pointwise(DoubleNear(rtol), v2));
}
using SolverTest = MujocoTest;
static const char* const kIlslandEfcPath =
"engine/testdata/island/island_efc.xml";
static const char* const kModelPath =
"testdata/model.xml";
// compare accelerations produced by CG solver with and without islands
TEST_F(SolverTest, IslandsEquivalent) {
const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
const std::string xml_path = GetTestDataFilePath(kModelPath);
char error[1024];
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
ASSERT_THAT(model, NotNull()) << error;
model->opt.solver = mjSOL_CG; // use CG solver
model->opt.jacobian = mjJAC_SPARSE; // use sparse
model->opt.tolerance = 0; // set tolerance to 0
model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands
model->opt.ls_tolerance = 0; // set ls_tolerance to 0
int nv = model->nv;
int state_size = mj_stateSize(model, mjSTATE_INTEGRATION);
mjtNum* state = (mjtNum*) mju_malloc(sizeof(mjtNum)*state_size);
mjtNum* qacc_diff = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv);
mjData* data_island = mj_makeData(model);
mjData* data_noisland = mj_makeData(model);
mjtNum rtol = 1e-5;
// Below are 3 tolerances associated with 3 different iteration counts,
// they are only moderately tight, 2x higher than x86-64 failure on Linux,
// i.e. in that case the test fails with rtol smaller than {5e-3, 5e-4, 5e-5}.
// The point of this test is to show that CG convergence is actually not very
// precise, simply changing whether islands are used changes the solution by
// quite a lot, even at high iteration count and zero {ls_}tolerance.
// Increasing the iteration count higher than 60 does not improve convergence.
constexpr int kNumTol = 3;
mjtNum maxiter[kNumTol] = {30, 40, 60};
mjtNum rtol[kNumTol] = {1e-2, 1e-3, 1e-4};
for (bool warmstart : {true, false}) {
if (warmstart) {
model->opt.disableflags |= mjDSBL_WARMSTART;
} else {
model->opt.disableflags &= ~mjDSBL_WARMSTART;
}
mj_resetData(model, data_noisland);
for (int i = 0; i < kNumTol; ++i) {
model->opt.iterations = maxiter[i];
model->opt.ls_iterations = maxiter[i];
while (data_noisland->time < .3) {
mj_step(model, data_noisland);
for (bool coldstart : {true, false}) {
mj_resetDataKeyframe(model, data_noisland, 0);
mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION);
mj_setState(model, data_island, state, mjSTATE_INTEGRATION);
if (coldstart) {
model->opt.disableflags |= mjDSBL_WARMSTART;
} else {
model->opt.disableflags &= ~mjDSBL_WARMSTART;
}
mj_forward(model, data_noisland);
while (data_noisland->time < .1) {
mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION);
mj_setState(model, data_island, state, mjSTATE_INTEGRATION);
model->opt.enableflags |= mjENBL_ISLAND; // enable islands
mj_forward(model, data_island);
model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands
model->opt.enableflags |= mjENBL_ISLAND; // enable islands
mj_forward(model, data_island);
ExpectEqRel(AsVector(data_noisland->qacc, nv),
AsVector(data_island->qacc, nv), rtol);
model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands
mj_forward(model, data_noisland);
auto time = std::to_string(data_noisland->time);
for (int j = 0; j < nv; j++) {
// increase tolerance for large elements
mjtNum scale = 0.5 * max(2.0, abs(data_noisland->qacc[j]) +
abs(data_island->qacc[j]));
EXPECT_THAT(data_noisland->qacc[j],
DoubleNear(data_island->qacc[j], scale * rtol[i]))
<< "time: " << time << '\n'
<< "dof: " << j << '\n'
<< "maxiter: " << maxiter[i] << '\n'
<< "rtol: " << scale * rtol[i];
}
mj_step(model, data_noisland);
}
}
}
mj_deleteData(data_noisland);
mj_deleteData(data_island);
mju_free(qacc_diff);
mju_free(state);
mj_deleteModel(model);
}
// compare accelerations produced by CG solver with and without islands
TEST_F(SolverTest, IslandsEquivalentForward) {
const std::string xml_path = GetTestDataFilePath(kModelPath);
char error[1024];
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
ASSERT_THAT(model, NotNull()) << error;
model->opt.solver = mjSOL_CG; // use CG solver
model->opt.tolerance = 0; // set tolerance to 0
model->opt.ls_tolerance = 0; // set ls_tolerance to 0
mjtNum rtol = 2e-6;
mjData* data_island = mj_makeData(model);
mjData* data_noisland = mj_makeData(model);
for (bool coldstart : {true, false}) {
mj_resetDataKeyframe(model, data_island, 0);
mj_resetDataKeyframe(model, data_noisland, 0);
if (coldstart) {
model->opt.disableflags |= mjDSBL_WARMSTART;
} else {
model->opt.disableflags &= ~mjDSBL_WARMSTART;
}
model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands
mj_forward(model, data_noisland);
model->opt.enableflags |= mjENBL_ISLAND; // enable islands
mj_forward(model, data_island);
for (int j = 0; j < model->nv; j++) {
mjtNum scale = 0.5 * max(2.0, abs(data_noisland->qacc[j]) +
abs(data_island->qacc[j]));
EXPECT_THAT(data_noisland->qacc[j],
DoubleNear(data_island->qacc[j], scale * rtol))
<< "dof: " << j << '\n'
<< "rtol: " << scale * rtol;
}
}
mj_deleteData(data_noisland);
mj_deleteData(data_island);
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);
+18 -7
View File
@@ -45,24 +45,35 @@ TEST_F(PipelineTest, SparseDenseEquivalent) {
mjtNum tol = 1e-11;
for (mjtSolver solver : {mjSOL_NEWTON, mjSOL_PGS, mjSOL_CG}) {
model->opt.solver = solver;
const char* sname[4] = {"NEWTON", "PGS", "CG", "NOSLIP"};
mjtSolver solver[4] = {mjSOL_NEWTON, mjSOL_PGS, mjSOL_CG, mjSOL_NEWTON};
// set dense jacobian, call mj_forward, save accelerations
for (int i : {0, 1, 2, 3}) {
model->opt.solver = solver[i];
if (i == 3) {
model->opt.noslip_iterations = 2;
}
// set dense jacobian, call mj_step, save qacc and new qpos
model->opt.jacobian = mjJAC_DENSE;
mj_resetDataKeyframe(model, data, 0);
mj_forward(model, data);
mj_step(model, data);
std::vector<mjtNum> qacc_dense = AsVector(data->qacc, model->nv);
std::vector<mjtNum> qpos_dense = AsVector(data->qpos, model->nq);
// set sparse jacobian, call mj_forward, save accelerations
// set sparse jacobian, call mj_step, save qacc and new qpos
model->opt.jacobian = mjJAC_SPARSE;
mj_resetDataKeyframe(model, data, 0);
mj_forward(model, data);
mj_step(model, data);
std::vector<mjtNum> qacc_sparse = AsVector(data->qacc, model->nv);
std::vector<mjtNum> qpos_sparse = AsVector(data->qpos, model->nq);
// expect accelerations to be insignificantly different
EXPECT_THAT(qacc_dense, Pointwise(DoubleNear(tol), qacc_sparse))
<< "failed equivalence for solver=" << solver;
<< "failed qacc equivalence for solver=" << sname[i];
// expect positions to be insignificantly different
EXPECT_THAT(qpos_dense, Pointwise(DoubleNear(tol), qpos_sparse))
<< "failed qpos equivalence for solver=" << sname[i];
}
mj_deleteData(data);