diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 6b23111f..65dad884 100644 --- a/src/engine/engine_forward.c +++ b/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) diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index adb1e23d..be4c5c3b 100644 --- a/src/engine/engine_solver.c +++ b/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); +} diff --git a/src/engine/engine_solver.h b/src/engine/engine_solver.h index 6947489a..7ee007de 100644 --- a/src/engine/engine_solver.h +++ b/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_