Improve warmstarting if island structure exists.

PiperOrigin-RevId: 752901506
Change-Id: I062d14df7da7d9dc1a2c1c30a16d0db7215af336

src/engine/engine_forward.c

@@ -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);
   }
 

test/engine/engine_solver_test.cc

@@ -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());
-
+  for (int i = 0; i < v1.size(); i++) {
-  // make scale vector
+    mjtNum scale = 0.5 * max(2.0, abs(v1[i]) + abs(v2[i]));
-  int n = v1.size();
+    EXPECT_THAT(v1[i], DoubleNear(v2[i], scale*rtol));
-  vector<mjtNum> scale(n);
-  for (int i = 0; i < n; i++) {
-    scale[i] = max(1.0, abs(v1[i]) + abs(v2[i]));
   }
-
-  // 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 =
+static const char* const kModelPath =
-    "engine/testdata/island/island_efc.xml";
+    "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}) {
+  for (int i = 0; i < kNumTol; ++i) {
-    if (warmstart) {
+    model->opt.iterations = maxiter[i];
-      model->opt.disableflags |= mjDSBL_WARMSTART;
+    model->opt.ls_iterations = maxiter[i];
-    } else {
-      model->opt.disableflags &= ~mjDSBL_WARMSTART;
-    }
-    mj_resetData(model, data_noisland);
 
-    while (data_noisland->time < .3) {
+    for (bool coldstart : {true, false}) {
-      mj_step(model, data_noisland);
+      mj_resetDataKeyframe(model, data_noisland, 0);
 
-      mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION);
+      if (coldstart) {
-      mj_setState(model, data_island, state, mjSTATE_INTEGRATION);
+        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
+        model->opt.enableflags |= mjENBL_ISLAND;  // enable islands
-      mj_forward(model, data_island);
+        mj_forward(model, data_island);
-      model->opt.enableflags &= ~mjENBL_ISLAND;  // disable islands
 
-      ExpectEqRel(AsVector(data_noisland->qacc, nv),
+        model->opt.enableflags &= ~mjENBL_ISLAND;  // disable islands
-                  AsVector(data_island->qacc, nv), rtol);
+        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);

test/pipeline_test.cc

@@ -45,24 +45,35 @@ TEST_F(PipelineTest, SparseDenseEquivalent) {
 
   mjtNum tol = 1e-11;
 
-  for (mjtSolver solver : {mjSOL_NEWTON, mjSOL_PGS, mjSOL_CG}) {
+  const char* sname[4] = {"NEWTON", "PGS", "CG", "NOSLIP"};
-    model->opt.solver = solver;
+  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);