Add island support to CG solver.

BEGIN_PUBLIC

Add island support to CG solver.

END_PUBLIC

PiperOrigin-RevId: 565653886
Change-Id: Ib971ae37dd011b8f2cbe1f88d8875e0ed657c7af

test/engine/CMakeLists.txt

@@ -72,6 +72,9 @@ target_link_libraries(engine_resource_test fixture gmock)
 mujoco_test(engine_sensor_test)
 target_link_libraries(engine_sensor_test fixture gmock)
 
+mujoco_test(engine_solver_test)
+target_link_libraries(engine_solver_test fixture gmock)
+
 mujoco_test(engine_support_test)
 target_link_libraries(engine_support_test fixture gmock)
 

test/engine/engine_solver_test.cc

@@ -0,0 +1,165 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Tests for engine/engine_solver.c
+
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <mujoco/mujoco.h>
+#include "test/fixture.h"
+
+namespace mujoco {
+namespace {
+
+std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
+  return std::vector<mjtNum>(array, array + n);
+}
+
+using ::testing::DoubleNear;
+using ::testing::NotNull;
+using ::testing::Pointwise;
+
+using SolverTest = MujocoTest;
+
+static const char* const kIlslandEfcPath =
+    "engine/testdata/island/island_efc.xml";
+
+// compare accelerations produced by CG solver with and without islands
+TEST_F(SolverTest, IslandsEquivalent) {
+  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.tolerance = 0;                     // set tolerance to 0
+  model->opt.enableflags &= ~mjENBL_ISLAND;     // disable islands
+
+  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 tol = 2e-4;
+
+  for (bool warmstart : {true, false}) {
+    if (warmstart) {
+      model->opt.disableflags |= mjDSBL_WARMSTART;
+    } else {
+      model->opt.disableflags &= ~mjDSBL_WARMSTART;
+    }
+    mj_resetData(model, data_noisland);
+
+    while (data_noisland->time < .3) {
+      mj_step(model, data_noisland);
+
+      mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION);
+      mj_setState(model, data_island, state, mjSTATE_INTEGRATION);
+
+      mj_forward(model, data_noisland);
+
+      model->opt.enableflags |= mjENBL_ISLAND;  // enable islands
+      mj_forward(model, data_island);
+      model->opt.enableflags &= ~mjENBL_ISLAND;  // disable islands
+
+      EXPECT_THAT(AsVector(data_noisland->qacc, nv),
+                  Pointwise(DoubleNear(tol), AsVector(data_island->qacc, nv)));
+    }
+  }
+
+  mj_deleteData(data_noisland);
+  mj_deleteData(data_island);
+  mju_free(qacc_diff);
+  mju_free(state);
+  mj_deleteModel(model);
+}
+
+// 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 tol = 1e-9;
+
+  // 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
+    EXPECT_THAT(AsVector(data_noisland->qacc, nv),
+                Pointwise(DoubleNear(tol), AsVector(data_island->qacc, nv)));
+  }
+
+  mj_deleteData(data_noisland);
+  mj_deleteData(data_island);
+  mju_free(state);
+  mj_deleteModel(model);
+}
+
+
+}  // namespace
+}  // namespace mujoco