chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add per-island PGS solver dispatch.

Browse Source 
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
This commit is contained in:
Yuval Tassa
2026-04-28 05:13:22 -07:00
committed by Copybara-Service
parent 00fff8c780
commit 647af382c1
7 changed files with 146 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+5
View File
@@ -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
-6
View File
@@ -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
+1 -1
View File
@@ -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
+57 -36
View File
@@ -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) &&
nisland > 0 &&
m->opt.noslip_iterations == 0 &&
(m->opt.solver == mjSOL_CG || m->opt.solver == mjSOL_NEWTON);
int islands_supported = !mjDISABLED(mjDSBL_ISLAND) && nisland > 0;
// run solver over constraint islands
if (islands_supported) {
int nidof = d->nidof;
// 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
switch ((mjtSolver) m->opt.solver) {
case mjSOL_PGS:
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);
}
mj_solPGS_island(m, d, island, m->opt.iterations);
}
} else {
// have threadpool, solve using threads
solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON);
break;
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)
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);
// run noslip solver per island if enabled
if (m->opt.noslip_iterations > 0) {
for (int island=0; island < nisland; island++) {
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:
mjERROR("unknown solver type %d", m->opt.solver);
// run noslip solver if enabled
if (m->opt.noslip_iterations > 0) {
mj_solNoSlip(m, d, m->opt.noslip_iterations);
}
}
// run noslip solver if enabled
if (m->opt.noslip_iterations > 0) {
mj_solNoSlip(m, d, m->opt.noslip_iterations);
// dual solvers: map efc_force to joint space (always monolithic)
if (m->opt.solver == mjSOL_PGS || m->opt.noslip_iterations > 0) {
mj_dualFinish(m, d);
}
TM_END(mjTIMER_CONSTRAINT);
src/engine/engine_solver.c
+37 -14
View File
@@ -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];
saveStats(m, d, island_stat, stats_iter, improvement, 0, 0, nactive, nchange, 0, 0);
if (island_stat < mjNISLAND) {
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
+9
View File
@@ -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
+37
View File
@@ -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