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

Rejigged solver statistics in preparation for islanding.

Browse Source 
- See public description below.
- Stopped incrementing the iteration count in saveStats().
- `sizeof(mjSolverStat) == 40`, so this ends up costing 160KB, up from 40KB.

BEGIN_PUBLIC

Changed the size of `mjData.solver`, the structure used to collect solver diagnostic information. The array is now of length `mjNISLAND * mjNSOLVER`, where each row of length `mjNSOLVER` contains separate solver statistics for each constraint island. Until solver islanding is implemented, only row 0 is used.

- The new constant `mjNISLAND` is set to 20.
- `mjNSOLVER` is reduced from 1000 to 200.
- Added `mjData.solver_nisland`, the number of islands for which the solver ran.
- `mjData.solver_niter` (renamed from mjData.solver_iter) and `mjData.solver_nnz` are now integer vectors of length `mjNISLAND`.

END_PUBLIC

PiperOrigin-RevId: 565030093
Change-Id: I773e918805c6ced79f0dab5f19ea23956760c8c5
This commit is contained in:
Yuval Tassa
2023-09-13 06:17:45 -07:00
committed by Copybara-Service
parent 8faf47dd16
commit 86d8b912f4
15 changed files with 3303 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.c
+8 -2
View File
@@ -1212,11 +1212,17 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
// in-place sparse backsubstitution for one island: x = inv(L'*D*L)*x
// L is in lower triangle of qLD; D is on diagonal of qLD
void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int island) {
// if no islands, call mj_solveLD
const mjtNum* qLD = d->qLD;
const mjtNum* qLDiagInv = d->qLDiagInv;
if (island < 0) {
mj_solveLD(m, x, 1, qLD, qLDiagInv);
return;
}
// local constants: general
const int* Madr = m->dof_Madr;
const int* parentid = m->dof_parentid;
const mjtNum* qLD = d->qLD;
const mjtNum* qLDiagInv = d->qLDiagInv;
const int* simplenum = m->dof_simplenum;
// local constants: island specific
src/engine/engine_forward.c
+5 -2
View File
@@ -501,7 +501,7 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
mju_copy(d->qacc, d->qacc_smooth, nv);
mju_copy(d->qacc_warmstart, d->qacc_smooth, nv);
mju_zero(d->qfrc_constraint, nv);
d->solver_iter = 0;
mju_zeroInt(d->solver_niter, mjNISLAND);
TM_END(mjTIMER_CONSTRAINT);
return;
}
@@ -512,7 +512,7 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
// warmstart solver
warmstart(m, d);
d->solver_iter = 0;
mju_zeroInt(d->solver_niter, mjNISLAND);
// run main solver
switch ((mjtSolver) m->opt.solver) {
@@ -532,6 +532,9 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
mjERROR("unknown solver type %d", m->opt.solver);
}
// one (monolithic) island
d->solver_nisland = 1;
// save result for next step warmstart
mju_copy(d->qacc_warmstart, d->qacc, nv);
src/engine/engine_io.c
+4 -3
View File
@@ -1428,9 +1428,10 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
// clear solver diagnostics
memset(d->warning, 0, mjNWARNING*sizeof(mjWarningStat));
memset(d->timer, 0, mjNTIMER*sizeof(mjTimerStat));
memset(d->solver, 0, mjNSOLVER*sizeof(mjSolverStat));
d->solver_iter = 0;
d->solver_nnz = 0;
memset(d->solver, 0, mjNSOLVER*mjNISLAND*sizeof(mjSolverStat));
d->solver_nisland = 0;
mju_zeroInt(d->solver_niter, mjNISLAND);
mju_zeroInt(d->solver_nnz, mjNISLAND);
mju_zero(d->solver_fwdinv, 2);
// clear collision diagnostics
src/engine/engine_print.c
+25 -16
View File
@@ -837,24 +837,33 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
}
// SOLVER STAT
if (d->solver_iter) {
if (d->nefc) {
fprintf(fp, "SOLVER STAT\n");
fprintf(fp, " solver_iter = %d\n", d->solver_iter);
fprintf(fp, " solver_nnz = %d\n", d->solver_nnz);
for (int i=0; i < mjMIN(mjNSOLVER, d->solver_iter); i++) {
fprintf(fp, " %d: improvement = ", i);
fprintf(fp, float_format, d->solver[i].improvement);
fprintf(fp, " gradient = ");
fprintf(fp, float_format, d->solver[i].gradient);
fprintf(fp, " lineslope = ");
fprintf(fp, float_format, d->solver[i].lineslope);
fprintf(fp, "\n");
fprintf(fp, " nactive = %d nchange = %d neval = %d nupdate = %d\n",
d->solver[i].nactive, d->solver[i].nchange,
d->solver[i].neval, d->solver[i].nupdate);
fprintf(fp, " solver_nisland = %d\n", d->solver_nisland);
printVector(" solver_fwdinv = ", d->solver_fwdinv, 2, fp, float_format);
int nisland_stat = mjMIN(d->solver_nisland, mjNISLAND);
for (int island=0; island < nisland_stat; island++) {
int niter_stat = mjMIN(mjNSOLVER, d->solver_niter[island]);
if (niter_stat) {
fprintf(fp, " ISLAND %d\n", island);
fprintf(fp, " solver_niter = %d\n", d->solver_niter[island]);
fprintf(fp, " solver_nnz = %d\n", d->solver_nnz[island]);
for (int i=0; i < niter_stat; i++) {
mjSolverStat* stat = d->solver + island*mjNSOLVER + i;
fprintf(fp, " %d: improvement = ", i);
fprintf(fp, float_format, stat->improvement);
fprintf(fp, " gradient = ");
fprintf(fp, float_format, stat->gradient);
fprintf(fp, " lineslope = ");
fprintf(fp, float_format, stat->lineslope);
fprintf(fp, "\n");
fprintf(fp, " nactive = %d nchange = %d neval = %d nupdate = %d\n",
stat->nactive, stat->nchange,
stat->neval, stat->nupdate);
}
fprintf(fp, "\n");
}
}
printVector("solver_fwdinv = ", d->solver_fwdinv, 2, fp, float_format);
fprintf(fp, "\n");
}
printVector("ENERGY = ", d->energy, 2, fp, float_format);
src/engine/engine_solver.c
+71 -40
View File
@@ -39,24 +39,30 @@ static mjtNum rescale(const mjModel* m, mjtNum x) {
// save solver statistics, count
static void saveStats(const mjModel* m, mjData* d, int* piter,
// save solver statistics
static void saveStats(const mjModel* m, mjData* d, int island, int iter,
mjtNum improvement, mjtNum gradient, mjtNum lineslope,
int nactive, int nchange, int neval, int nupdate) {
// compute position, increase iter
int i = d->solver_iter + (*piter);
(*piter)++;
// save if within range
if (i < mjNSOLVER) {
d->solver[i].improvement = improvement;
d->solver[i].gradient = gradient;
d->solver[i].lineslope = lineslope;
d->solver[i].nactive = nactive;
d->solver[i].nchange = nchange;
d->solver[i].neval = neval;
d->solver[i].nupdate = nupdate;
// if out of range, return
if (island >= mjNISLAND) {
return;
}
// if no islands, use first island
island = mjMAX(0, island);
// get mjSolverStat pointer
iter += d->solver_niter[island]; // add current niter (in case of noslip)
mjSolverStat* stat = d->solver + island*mjNSOLVER + iter;
// save stats
stat->improvement = improvement;
stat->gradient = gradient;
stat->lineslope = lineslope;
stat->nactive = nactive;
stat->nchange = nchange;
stat->neval = neval;
stat->nupdate = nupdate;
}
@@ -313,6 +319,9 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
mjtNum* ARinv = mj_stackAllocNum(d, nefc);
int* oldstate = mj_stackAllocInt(d, nefc);
// TODO: b/295296178 - Use island index (currently hardcoded to 0)
int island = 0;
// precompute inverse diagonal of AR
ARdiaginv(m, d, ARinv, 0);
@@ -471,9 +480,13 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
nchange += (oldstate[i] != d->efc_state[i]);
}
// scale improvement, save stats, count
// scale improvement, save stats
improvement = rescale(m, improvement);
saveStats(m, d, &iter, improvement, 0, 0, nactive, nchange, 0, 0);
saveStats(m, d, island, iter, improvement, 0, 0, nactive, nchange, 0, 0);
// increment iteration count
iter++;
// terminate
if (improvement < m->opt.tolerance) {
@@ -481,17 +494,20 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
}
}
// update solver iterations
d->solver_iter += iter;
// finalize statistics
if (island < mjNISLAND) {
// update solver iterations
d->solver_niter[island] += iter;
// set nnz
if (mj_isSparse(m)) {
d->solver_nnz = 0;
for (int i=0; i < nefc; i++) {
d->solver_nnz += d->efc_AR_rownnz[i];
// set nnz
if (mj_isSparse(m)) {
d->solver_nnz[island] = 0;
for (int i=0; i < nefc; i++) {
d->solver_nnz[island] += d->efc_AR_rownnz[i];
}
} else {
d->solver_nnz[island] = nefc*nefc;
}
} else {
d->solver_nnz = nefc*nefc;
}
// map to joint space
@@ -515,6 +531,9 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
mjtNum* ARinv = mj_stackAllocNum(d, nefc);
int* oldstate = mj_stackAllocInt(d, nefc);
// TODO: b/295296178 - Use island index (currently hardcoded to 0)
int island = 0;
// precompute inverse diagonal of A
ARdiaginv(m, d, ARinv, 1);
@@ -687,9 +706,12 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
nchange += (oldstate[i] != d->efc_state[i]);
}
// scale improvement, save stats, count
// scale improvement, save stats
improvement = rescale(m, improvement);
saveStats(m, d, &iter, improvement, 0, 0, nactive, nchange, 0, 0);
saveStats(m, d, island, iter, improvement, 0, 0, nactive, nchange, 0, 0);
// increment iteration count
iter++;
// terminate
if (improvement < m->opt.noslip_tolerance) {
@@ -698,7 +720,7 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
}
// update solver iterations
d->solver_iter += iter;
d->solver_niter[island] += iter;
// map to joint space
dualFinish(m, d);
@@ -1514,6 +1536,9 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
mjCGContext ctx;
mj_markStack(d);
// TODO: b/295296178 - Use island index (currently hardcoded to 0)
int island = 0;
// allocate context
CGallocate(m, d, &ctx, flg_Newton);
@@ -1576,12 +1601,15 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
nchange += (d->efc_state[i] != oldstate[i]);
}
// scale improvement, save stats, count
// scale improvement, save stats
mjtNum improvement = rescale(m, oldcost-ctx.cost);
mjtNum gradient = rescale(m, mju_norm(ctx.grad, nv));
saveStats(m, d, &iter, improvement, gradient, ctx.LSslope,
saveStats(m, d, island, iter, improvement, gradient, ctx.LSslope,
ctx.nactive, nchange, ctx.LSiter, ctx.nupdate);
// increment iteration count
iter++;
// termination
if (improvement < m->opt.tolerance || gradient < m->opt.tolerance) {
break;
@@ -1608,18 +1636,21 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
}
}
// update solver iterations
d->solver_iter += iter;
// finalize statistics
if (island < mjNISLAND) {
// update solver iterations
d->solver_niter[island] += iter;
// set solver_nnz
if (flg_Newton) {
if (mj_isSparse(m)) {
d->solver_nnz = 2*ctx.nnz - nv;
// set solver_nnz
if (flg_Newton) {
if (mj_isSparse(m)) {
d->solver_nnz[island] = 2*ctx.nnz - nv;
} else {
d->solver_nnz[island] = nv*nv;
}
} else {
d->solver_nnz = nv*nv;
d->solver_nnz[island] = 0;
}
} else {
d->solver_nnz = 0;
}
mj_freeStack(d);