Add per-island PGS solver dispatch.

Total testspeed runtime for `2humanoids100.xml` reduced by 27.6% (63.4 -> 49.5s) due to early termination on small islands

PiperOrigin-RevId: 906910915
Change-Id: If55ad468c3680ef44eda7000455a77f8003b3122

doc/changelog.rst

@@ -2,6 +2,11 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+- Added island support for the :ref:`PGS solver<soAlgorithms>`.
+
 Version 3.8.0 (April 24, 2026)
 ------------------------------
 

doc/computation/index.rst

@@ -1424,12 +1424,6 @@ While islanding is not free (see implementation in `engine_island.c
 - Unconstrained DOFs are completely untouched by the solver, which otherwise needs to discover that they are unaffected.
 - Solving separate islands can be multi-threaded.
 
-.. admonition:: Known issues
-   :class: note
-
-   Islanding is not yet supported by the PGS solver.
-
-
 .. _soParameters:
 
 Parameters

simulate/simulate.cc

@@ -349,7 +349,7 @@ void UpdateProfiler(mj::Simulate* sim, const mjModel* m, const mjData* d) {
       sim->figcost.linepnt[start + 2] = 0;
     }
 
-    for (int i=0; i<sim->figcost.linepnt[0]; i++) {
+    for (int i=0; i<npoints; i++) {
       // x
       sim->figcost.linedata[start + 0][2*i] = i;
       sim->figcost.linedata[start + 1][2*i] = i;

src/engine/engine_forward.c

@@ -953,7 +953,7 @@ static void solve_threaded(const mjModel* m, mjData* d, int flg_Newton) {
 // compute efc_b, efc_force, qfrc_constraint; update qacc
 void mj_fwdConstraint(const mjModel* m, mjData* d) {
   TM_START;
-  int nv = m->nv, nefc = d->nefc, nisland = d->nisland;
+  int nv = m->nv, nefc = d->nefc, nisland = d->nisland, nidof;
 
   // always clear qfrc_constraint
   mju_zero(d->qfrc_constraint, nv);
@@ -970,50 +970,69 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
   mj_mulJacVec(m, d, d->efc_b, d->qacc_smooth);
   mju_subFrom(d->efc_b, d->efc_aref, nefc);
 
+  // check for invalid solver type
+  if (m->opt.solver != mjSOL_PGS && m->opt.solver != mjSOL_CG && m->opt.solver != mjSOL_NEWTON) {
+    mjERROR("unknown solver type %d", m->opt.solver);
+  }
+
   // warmstart solver
   warmstart(m, d);
   mju_zeroInt(d->solver_niter, mjNISLAND);
 
   // check if islands are supported
-  int islands_supported = !mjDISABLED(mjDSBL_ISLAND)    &&
+  int islands_supported = !mjDISABLED(mjDSBL_ISLAND) && nisland > 0;
-                          nisland > 0                   &&
-                          m->opt.noslip_iterations == 0 &&
-                          (m->opt.solver == mjSOL_CG || m->opt.solver == mjSOL_NEWTON);
 
   // run solver over constraint islands
   if (islands_supported) {
-    int nidof = d->nidof;
+    switch ((mjtSolver) m->opt.solver) {
-
+    case mjSOL_PGS:
-    // copy inputs to islands (vel+acc deps, pos-dependent already copied in mj_island)
-    mju_gather(d->ifrc_smooth,     d->qfrc_smooth,     d->map_idof2dof, nidof);
-    mju_gather(d->ifrc_constraint, d->qfrc_constraint, d->map_idof2dof, nidof);
-    mju_gather(d->iacc_smooth,     d->qacc_smooth,     d->map_idof2dof, nidof);
-    mju_gather(d->iacc,            d->qacc,            d->map_idof2dof, nidof);
-    mju_gather(d->iefc_force,      d->efc_force,       d->map_iefc2efc, nefc);
-    mju_gather(d->iefc_aref,       d->efc_aref,        d->map_iefc2efc, nefc);
-
-    // solve per island, with or without threads
-    if (!d->threadpool) {
-      // no threadpool, loop over islands
       for (int island=0; island < nisland; island++) {
-        if (m->opt.solver == mjSOL_NEWTON) {
+        mj_solPGS_island(m, d, island, m->opt.iterations);
-          mj_solNewton_island(m, d, island, m->opt.iterations);
-        } else {
-          mj_solCG_island(m, d, island, m->opt.iterations);
-        }
       }
-    } else {
+      break;
-      // have threadpool, solve using threads
+
-      solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON);
+    case mjSOL_CG:
+    case mjSOL_NEWTON:
+      // copy inputs to islands (vel+acc deps, pos-dependent already copied in mj_island)
+      nidof = d->nidof;
+      mju_gather(d->ifrc_smooth,     d->qfrc_smooth,     d->map_idof2dof, nidof);
+      mju_gather(d->ifrc_constraint, d->qfrc_constraint, d->map_idof2dof, nidof);
+      mju_gather(d->iacc_smooth,     d->qacc_smooth,     d->map_idof2dof, nidof);
+      mju_gather(d->iacc,            d->qacc,            d->map_idof2dof, nidof);
+      mju_gather(d->iefc_force,      d->efc_force,       d->map_iefc2efc, nefc);
+      mju_gather(d->iefc_aref,       d->efc_aref,        d->map_iefc2efc, nefc);
+
+      // solve per island, with or without threads
+      if (!d->threadpool) {
+        // no threadpool, loop over islands
+        for (int island=0; island < nisland; island++) {
+          if (m->opt.solver == mjSOL_NEWTON) {
+            mj_solNewton_island(m, d, island, m->opt.iterations);
+          } else {
+            mj_solCG_island(m, d, island, m->opt.iterations);
+          }
+        }
+      } else {
+        // have threadpool, solve using threads
+        solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON);
+      }
+
+      // copy back solver outputs (scatter dofs since ni <= nv)
+      mju_scatter(d->qacc,            d->iacc,            d->map_idof2dof, nidof);
+      mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof);
+      mju_gather(d->efc_force, d->iefc_force, d->map_efc2iefc, nefc);
+      break;
     }
 
-    // copy back solver outputs (scatter dofs since ni <= nv)
+    // run noslip solver per island if enabled
-    mju_scatter(d->qacc,            d->iacc,            d->map_idof2dof, nidof);
+    if (m->opt.noslip_iterations > 0) {
-    mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof);
+      for (int island=0; island < nisland; island++) {
-    mju_gather(d->efc_force, d->iefc_force, d->map_efc2iefc, nefc);
+        mj_solNoSlip_island(m, d, island, m->opt.noslip_iterations);
+      }
+    }
   }
 
-  // run solver over all constraints
+  // run solver over all constraints (monolithic)
   else {
     switch ((mjtSolver) m->opt.solver) {
     case mjSOL_PGS:                     // PGS
@@ -1027,15 +1046,17 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
     case mjSOL_NEWTON:                  // Newton
       mj_solNewton(m, d, m->opt.iterations);
       break;
+    }
 
-    default:
+    // run noslip solver if enabled
-      mjERROR("unknown solver type %d", m->opt.solver);
+    if (m->opt.noslip_iterations > 0) {
+      mj_solNoSlip(m, d, m->opt.noslip_iterations);
     }
   }
 
-  // run noslip solver if enabled
+  // dual solvers: map efc_force to joint space (always monolithic)
-  if (m->opt.noslip_iterations > 0) {
+  if (m->opt.solver == mjSOL_PGS || m->opt.noslip_iterations > 0) {
-    mj_solNoSlip(m, d, m->opt.noslip_iterations);
+    mj_dualFinish(m, d);
   }
 
   TM_END(mjTIMER_CONSTRAINT);

src/engine/engine_solver.c

@@ -66,7 +66,6 @@ static void saveStats(const mjModel* m, mjData* d, int island, int iter,
 
 
 // finalize dual solver: map to joint space
-// TODO: b/295296178 - add island support to Dual solvers
 static void dualFinish(const mjModel* m, mjData* d) {
   // map constraint force to joint space
   mj_mulJacTVec(m, d, d->qfrc_constraint, d->efc_force);
@@ -77,10 +76,15 @@ static void dualFinish(const mjModel* m, mjData* d) {
 }
 
 
+// PGS: map efc_force to joint space
+void mj_dualFinish(const mjModel* m, mjData* d) {
+  dualFinish(m, d);
+}
+
+
 // compute 1/diag(AR)
 //   res[c] = 1 / AR[efclist[c], efclist[c]] for c = 0..nefc-1
 //   efclist is NULL for monolithic (sequential) iteration
-// TODO: b/295296178 - add island support to Dual solvers
 static void ARdiaginv(const mjModel* m, const mjData* d, mjtNum* res,
                       int nefc, const int* efclist, int flg_subR) {
   const mjtNum *AR = d->efc_AR;
@@ -118,7 +122,6 @@ static void ARdiaginv(const mjModel* m, const mjData* d, mjtNum* res,
 
 
 // extract diagonal block from AR, clamp diag to 1e-10 if flg_subR
-// TODO: b/295296178 - add island support to Dual solvers
 static void extractBlock(const mjModel* m, const mjData* d, mjtNum* Ac,
                          int start, int n, int flg_subR) {
   int nefc = d->nefc;
@@ -178,7 +181,6 @@ static void extractBlock(const mjModel* m, const mjData* d, mjtNum* Ac,
 
 
 // compute residual for one block
-// TODO: b/295296178 - add island support to Dual solvers
 static void residual(const mjModel* m, const mjData* d, mjtNum* res, int i, int dim, int flg_subR) {
   int nefc = d->nefc;
 
@@ -208,7 +210,6 @@ static void residual(const mjModel* m, const mjData* d, mjtNum* res, int i, int
 
 
 // compute cost change
-// TODO: b/295296178 - add island support to Dual solvers
 static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
                          const mjtNum* res, int dim) {
   mjtNum change;
@@ -235,7 +236,6 @@ static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
 
 // set efc_state to dual constraint state; return nactive
 //   iterates over efclist (or sequentially if NULL), classifies by ne/nf ranges
-// TODO: b/295296178 - add island support to Dual solvers
 static int dualState(const mjData* d, int* state,
                      int ne, int nf, int nefc, const int* efclist) {
   const mjtNum* force = d->efc_force;
@@ -384,7 +384,6 @@ static void solveQCQP(mjtNum* force, int i, int dim,
 //   island: island index for stats (use -1 for monolithic, mapped to 0)
 //   ne, nf, nefc: constraint type counts
 //   efclist: maps list position c to monolithic efc index (NULL for sequential)
-// TODO: b/295296178 - add island support to Dual solvers
 static void solPGS(const mjModel* m, mjData* d, int island,
                    int ne, int nf, int nefc,
                    const int* efclist, int maxiter) {
@@ -572,10 +571,20 @@ static void solPGS(const mjModel* m, mjData* d, int island,
 }
 
 
-// PGS entry point (monolithic)
+// PGS entry point (monolithic, no dualFinish — caller handles it)
 void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
   solPGS(m, d, /*island=*/-1, d->ne, d->nf, d->nefc, /*efclist=*/NULL, maxiter);
-  dualFinish(m, d);
+}
+
+
+// PGS entry point (one island)
+void mj_solPGS_island(const mjModel* m, mjData* d, int island, int maxiter) {
+  int ne = d->island_ne[island];
+  int nf = d->island_nf[island];
+  int nefc = d->island_nefc[island];
+  int iefcadr = d->island_iefcadr[island];
+
+  solPGS(m, d, island, ne, nf, nefc, d->map_iefc2efc + iefcadr, maxiter);
 }
 
 
@@ -754,8 +763,10 @@ static void solNoSlip(const mjModel* m, mjData* d, int island,
     improvement *= scale;
 
     // save noslip stats after all the entries from regular solver
-    int stats_iter = iter + d->solver_niter[island_stat];
+    if (island_stat < mjNISLAND) {
-    saveStats(m, d, island_stat, stats_iter, improvement, 0, 0, nactive, nchange, 0, 0);
+      int stats_iter = iter + d->solver_niter[island_stat];
+      saveStats(m, d, island_stat, stats_iter, improvement, 0, 0, nactive, nchange, 0, 0);
+    }
 
     // increment iteration count
     iter++;
@@ -767,16 +778,28 @@ static void solNoSlip(const mjModel* m, mjData* d, int island,
   }
 
   // update solver iterations
-  d->solver_niter[island_stat] += iter;
+  if (island_stat < mjNISLAND) {
+    d->solver_niter[island_stat] += iter;
+  }
 
   mj_freeStack(d);
 }
 
 
-// NoSlip entry point (monolithic)
+// NoSlip entry point (monolithic, no dualFinish — caller handles it)
 void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
   solNoSlip(m, d, /*island=*/-1, d->ne, d->nf, d->nefc, /*efclist=*/NULL, maxiter);
-  dualFinish(m, d);
+}
+
+
+// NoSlip entry point (one island)
+void mj_solNoSlip_island(const mjModel* m, mjData* d, int island, int maxiter) {
+  int ne = d->island_ne[island];
+  int nf = d->island_nf[island];
+  int nefc = d->island_nefc[island];
+  int iefcadr = d->island_iefcadr[island];
+
+  solNoSlip(m, d, island, ne, nf, nefc, d->map_iefc2efc + iefcadr, maxiter);
 }
 
 

src/engine/engine_solver.h

@@ -35,10 +35,19 @@ void mj_solNewton(const mjModel* m, mjData* d, int maxiter);
 
 //------------------------------ per-island solvers ------------------------------------------------
 
+// PGS solver (one island, no dualFinish — caller handles it)
+void mj_solPGS_island(const mjModel* m, mjData* d, int island, int maxiter);
+
+// NoSlip solver (one island, no dualFinish — caller handles it)
+void mj_solNoSlip_island(const mjModel* m, mjData* d, int island, int maxiter);
+
 // CG solver
 void mj_solCG_island(const mjModel* m, mjData* d, int island, int maxiter);
 
 // Newton entry point
 void mj_solNewton_island(const mjModel* m, mjData* d, int island, int maxiter);
 
+// map efc_force to joint space (used after dual island dispatch)
+void mj_dualFinish(const mjModel* m, mjData* d);
+
 #endif  // MUJOCO_SRC_ENGINE_ENGINE_SOLVER_H_

test/engine/engine_island_test.cc

@@ -15,6 +15,7 @@
 // Tests for engine/engine_island.c.
 
 #include <string>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -602,5 +603,41 @@ TEST_F(IslandTest, EqualityConstraintOfTendons) {
   mj_deleteModel(model);
 }
 
+TEST_F(IslandTest, PGSIslandExact) {
+  const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
+  char error[1024];
+  mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(m, NotNull()) << error;
+  mjData* d = mj_makeData(m);
+
+  // simulate to get a non-trivial state
+  while (d->time < 0.5) {
+    mj_step(m, d);
+  }
+
+  // switch to PGS, disable early termination
+  m->opt.solver = mjSOL_PGS;
+  m->opt.tolerance = 0;
+
+  // solve with islands
+  m->opt.disableflags &= ~mjDSBL_ISLAND;
+  mj_forward(m, d);
+  ASSERT_GT(d->nisland, 1);
+  std::vector<mjtNum> qfrc_island(d->qfrc_constraint,
+                                  d->qfrc_constraint + m->nv);
+
+  // solve without islands
+  m->opt.disableflags |= mjDSBL_ISLAND;
+  mj_forward(m, d);
+  std::vector<mjtNum> qfrc_mono(d->qfrc_constraint,
+                                d->qfrc_constraint + m->nv);
+
+  // expect exact match
+  EXPECT_EQ(qfrc_island, qfrc_mono);
+
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
 }  // namespace
 }  // namespace mujoco