Add island support to Newton solver

PiperOrigin-RevId: 758618486
Change-Id: I257d78c7dc9aa4dbf7cd6d8849d8ec48201edf24

src/engine/engine_forward.c

@@ -723,16 +723,16 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
   mju_zeroInt(d->solver_niter, mjNISLAND);
 
   // check if islands are supported
-  int islands_supported = mjENABLED(mjENBL_ISLAND)  &&
-                          nisland > 0               &&
-                          m->opt.solver == mjSOL_CG &&
-                          m->opt.noslip_iterations == 0;
+  int islands_supported = mjENABLED(mjENBL_ISLAND)      &&
+                          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;
 
-    // copy CG inputs to islands (vel+acc deps, pos-dependent already copied in mj_island)
+    // 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);
@@ -741,14 +741,20 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
     mju_gather(d->iefc_aref,       d->efc_aref,        d->map_iefc2efc, nefc);
 
     // solve per island
-    if (!d->threadpool) {
-      // no threadpool, loop over islands
-      for (int island=0; island < nisland; island++) {
-        mj_solCG_island(m, d, island, m->opt.iterations);
+    if (m->opt.solver == mjSOL_CG) {
+      if (!d->threadpool) {
+        // no threadpool, loop over islands
+        for (int island=0; island < nisland; island++) {
+          mj_solCG_island(m, d, island, m->opt.iterations);
+        }
+      } else {
+        // have threadpool, solve using threads
+        mj_solCG_island_multithreaded(m, d);
       }
     } else {
-      // have threadpool, solve using threads
-      mj_solCG_island_multithreaded(m, d);
+      for (int island=0; island < nisland; island++) {
+        mj_solNewton_island(m, d, island, m->opt.iterations);
+      }
     }
 
     // copy back solver outputs (scatter dofs since ni <= nv)

src/engine/engine_solver.c

@@ -1058,7 +1058,6 @@ static void CGupdateGradient(mjCGContext* ctx, int flg_Newton) {
   }
 
   // Newton: Mgrad = H \ grad
-  // TODO: b/295296178 - add island support to Newton solver
   if (flg_Newton) {
     if (ctx->is_sparse) {
       mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? ctx->Lcone : ctx->L),
@@ -2042,3 +2041,10 @@ void mj_solCG_island(const mjModel* m, mjData* d, int island, int maxiter) {
 void mj_solNewton(const mjModel* m, mjData* d, int maxiter) {
   mj_solCGNewton(m, d, /*island=*/-1, maxiter, /*flg_Newton=*/1);
 }
+
+
+
+// Newton entry point (one island)
+void mj_solNewton_island(const mjModel* m, mjData* d, int island, int maxiter) {
+  mj_solCGNewton(m, d, island, maxiter, /*flg_Newton=*/1);
+}

src/engine/engine_solver.h

@@ -32,9 +32,13 @@ void mj_solCG(const mjModel* m, mjData* d, int maxiter);
 // Newton solver
 void mj_solNewton(const mjModel* m, mjData* d, int maxiter);
 
+
 //------------------------------ per-island solvers ------------------------------------------------
 
 // 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);
+
 #endif  // MUJOCO_SRC_ENGINE_ENGINE_SOLVER_H_