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

Add island support to CG solver.

Browse Source 
BEGIN_PUBLIC

Add island support to CG solver.

END_PUBLIC

PiperOrigin-RevId: 565653886
Change-Id: Ib971ae37dd011b8f2cbe1f88d8875e0ed657c7af
This commit is contained in:
Yuval Tassa
2023-09-15 05:46:35 -07:00
committed by Copybara-Service
parent 492b6a72ec
commit 2cc1205498
7 changed files with 394 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+10 -8
View File
@@ -12,11 +12,11 @@ New features
:align: right
:width: 240px
1. Added constraint island discovery in :ref:`mj_island`. Constraint islands are disjoint sets of constraints
and degrees-of-freedom that do not interact. In a future release the constraint solver will be refactored to
exploit this disjoint structure. Island discovery can be activated using a new :ref:`enable flag<option-flag-island>`
which will be removed after the refactor. If island discovery is enabled, geoms, contacts and
tendons will be colored according to the corresponding island, see video.
1. Added constraint island discovery with :ref:`mj_island`. Constraint islands are disjoint sets of constraints
and degrees-of-freedom that do not interact. The only solver which currently supports islands is
:ref:`CG<option-solver>`. Island discovery can be activated using a new :ref:`enable flag<option-flag-island>`.
If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island,
see video.
.. youtube:: QewlEqIZi1o
:align: right
@@ -74,8 +74,9 @@ General
12. Added the flag :ref:`invdiscrete<option-flag-invdiscrete>`, which enables discrete-time inverse dynamics for all
:ref:`integrators<option-integrator>` other than ``RK4``. See the flag documentation for more details.
13. Added :ref:`ls_iterations<option-ls_iterations>` and :ref:`ls_tolerance<option-ls_tolerance>` options for adjusting
linesearch stopping criteria in CG and Newton solvers. This can be useful for performance tuning.
14. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store normalizing transformation.
linesearch stopping criteria in CG and Newton solvers. These can be useful for performance tuning.
14. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store the normalizing transformation applied to
mesh assets. Fixes `#409 <https://github.com/google-deepmind/mujoco/issues/409>`__ .
15. Added camera :ref:`resolution<body-camera-resolution>` attribute and :ref:`camprojection<sensor-camprojection>`
sensor. If camera resolution is set to positive values, the camera projection sensor will report the location of a
target site, projected onto the camera image, in pixel coordinates.
@@ -89,7 +90,8 @@ Python bindings
Bug fixes
^^^^^^^^^
17. Fixed a bug that was causing the geom margins to be ignored during the midphase.
17. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
midphase collision trees.
Version 2.3.7 (July 20, 2023)
src/engine/engine_forward.c
+37 -19
View File
@@ -494,13 +494,15 @@ static void warmstart(const mjModel* m, mjData* d) {
// 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;
int nv = m->nv, nefc = d->nefc, nisland = d->nisland;
// always clear qfrc_constraint
mju_zero(d->qfrc_constraint, nv);
// no constraints: copy unconstrained acc, clear forces, return
if (!nefc) {
mju_copy(d->qacc, d->qacc_smooth, nv);
mju_copy(d->qacc_warmstart, d->qacc_smooth, nv);
mju_zero(d->qfrc_constraint, nv);
mju_zeroInt(d->solver_niter, mjNISLAND);
TM_END(mjTIMER_CONSTRAINT);
return;
@@ -514,26 +516,42 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
warmstart(m, d);
mju_zeroInt(d->solver_niter, mjNISLAND);
// run main solver
switch ((mjtSolver) m->opt.solver) {
case mjSOL_PGS: // PGS
mj_solPGS(m, d, m->opt.iterations);
break;
// check if islands are supported
int islands_supported = mjENABLED(mjENBL_ISLAND) &&
m->opt.solver == mjSOL_CG &&
m->opt.noslip_iterations == 0;
case mjSOL_CG: // CG
mj_solCG(m, d, m->opt.iterations);
break;
case mjSOL_NEWTON: // Newton
mj_solNewton(m, d, m->opt.iterations);
break;
default:
mjERROR("unknown solver type %d", m->opt.solver);
// run solver over constraint islands
if (islands_supported) {
// loop over islands
for (int island=0; island < nisland; island++) {
mj_solCG_island(m, d, island, m->opt.iterations);
}
d->solver_nisland = nisland;
}
// one (monolithic) island
d->solver_nisland = 1;
// run solver over all constraints
else {
switch ((mjtSolver) m->opt.solver) {
case mjSOL_PGS: // PGS
mj_solPGS(m, d, m->opt.iterations);
break;
case mjSOL_CG: // CG
mj_solCG(m, d, m->opt.iterations);
break;
case mjSOL_NEWTON: // Newton
mj_solNewton(m, d, m->opt.iterations);
break;
default:
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_solver.c
+172 -78
View File
@@ -30,15 +30,9 @@
#include "engine/engine_util_solve.h"
#include "engine/engine_util_sparse.h"
//---------------------------------- utility functions ---------------------------------------------
// rescale cost and gradient
static mjtNum rescale(const mjModel* m, mjtNum x) {
return x / (m->stat.meaninertia * mjMAX(1, m->nv));
}
// save solver statistics
static void saveStats(const mjModel* m, mjData* d, int island, int iter,
mjtNum improvement, mjtNum gradient, mjtNum lineslope,
@@ -68,6 +62,7 @@ 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);
@@ -80,6 +75,7 @@ static void dualFinish(const mjModel* m, mjData* d) {
// compute 1/diag(AR)
// TODO: b/295296178 - add island support to Dual solvers
static void ARdiaginv(const mjModel* m, mjData* d, mjtNum* res, int flg_subR) {
int nefc = d->nefc;
const int *rowadr = d->efc_AR_rowadr;
@@ -109,6 +105,7 @@ static void ARdiaginv(const mjModel* m, mjData* d, mjtNum* res, int flg_subR) {
// 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, mjData* d, mjtNum* Ac,
int start, int n, int flg_subR) {
int nefc = d->nefc;
@@ -166,6 +163,7 @@ static void extractBlock(const mjModel* m, mjData* d, mjtNum* Ac,
// compute residual for one block
// TODO: b/295296178 - add island support to Dual solvers
static void residual(const mjModel* m, mjData* d, mjtNum* res, int i, int dim, int flg_subR) {
int nefc = d->nefc;
@@ -196,6 +194,7 @@ static void residual(const mjModel* m, mjData* d, mjtNum* res, int i, int dim, i
// 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 delta[6], change;
@@ -221,6 +220,7 @@ static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
// set efc_state to dual constraint state; return nactive
// TODO: b/295296178 - add island support to Dual solvers
static int dualState(const mjModel* m, mjData* d) {
int nactive, ne = d->ne, nf = d->nf, nefc = d->nefc;
const mjtNum *force = d->efc_force, *floss = d->efc_frictionloss;
@@ -308,6 +308,7 @@ static int dualState(const mjModel* m, mjData* d) {
//---------------------------- PGS solver ----------------------------------------------------------
// TODO: b/295296178 - add island support to Dual solvers
void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
int dim, iter = 0, ne = d->ne, nf = d->nf, nefc = d->nefc;
const mjtNum *floss = d->efc_frictionloss;
@@ -321,6 +322,7 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
// TODO: b/295296178 - Use island index (currently hardcoded to 0)
int island = 0;
mjtNum scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv));
// precompute inverse diagonal of AR
ARdiaginv(m, d, ARinv, 0);
@@ -481,7 +483,7 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
}
// scale improvement, save stats
improvement = rescale(m, improvement);
improvement *= scale;
saveStats(m, d, island, iter, improvement, 0, 0, nactive, nchange, 0, 0);
// increment iteration count
@@ -520,6 +522,7 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
//---------------------------- NoSlip solver -------------------------------------------------------
// TODO: b/295296178 - add island support to Dual solvers
void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
int dim, iter = 0, ne = d->ne, nf = d->nf, nefc = d->nefc;
const mjtNum *floss = d->efc_frictionloss;
@@ -533,6 +536,7 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
// TODO: b/295296178 - Use island index (currently hardcoded to 0)
int island = 0;
mjtNum scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv));
// precompute inverse diagonal of A
ARdiaginv(m, d, ARinv, 1);
@@ -707,7 +711,7 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
}
// scale improvement, save stats
improvement = rescale(m, improvement);
improvement *= scale;
saveStats(m, d, island, iter, improvement, 0, 0, nactive, nchange, 0, 0);
// increment iteration count
@@ -734,6 +738,13 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
// CG context
struct _mjCGContext {
// island-related
int island; // current island index, -1 if monolithic
int nv; // number of dofs
int nefc; // number of constraints
int* dofind; // dof indices of this island, NULL if monolithic
int* efcind; // constraint indices of this island, NULL if monolithic
// arrays
mjtNum* Jaref; // Jac*qacc - aref (nefc x 1)
mjtNum* Jv; // Jac*search (nefc x 1)
@@ -756,6 +767,7 @@ struct _mjCGContext {
// globals
mjtNum cost; // constraint + Gauss cost
mjtNum quadGauss[3]; // quadratic polynomial for Gauss cost
mjtNum scale; // scaling factor for improvement and gradient
int nactive; // number of active constraints
int ncone; // number of contacts in cone state
int nupdate; // number of Cholesky updates
@@ -770,13 +782,22 @@ typedef struct _mjCGContext mjCGContext;
// allocate mjCGContext: mjMARK/FREE in caller function!
static void CGallocate(const mjModel* m, mjData* d,
mjCGContext* ctx, int flg_Newton) {
int nv = m->nv, nefc = d->nefc;
static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
int island, int flg_Newton) {
// get sizes
int nv = island < 0 ? m->nv : d->island_dofnum[island];
int nefc = island < 0 ? d->nefc : d->island_efcnum[island];
// clear everything
memset(ctx, 0, sizeof(mjCGContext));
// island-related
ctx->island = island;
ctx->nv = nv;
ctx->nefc = nefc;
ctx->dofind = island < 0 ? NULL : d->island_dofind + d->island_dofadr[island];
ctx->efcind = island < 0 ? NULL : d->island_efcind + d->island_efcadr[island];
// common arrays
ctx->Jaref = mj_stackAllocNum(d, nefc);
ctx->Jv = mj_stackAllocNum(d, nefc);
@@ -802,24 +823,29 @@ static void CGallocate(const mjModel* m, mjData* d,
// update efc_force, qfrc_constraint, cost-related
static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) {
int nefc = d->nefc, nv = m->nv;
int nefc = ctx->nefc, nv = ctx->nv;
const int* dofind = ctx->dofind;
const int* efcind = ctx->efcind;
// update constraints
mj_constraintUpdate(m, d, ctx->Jaref, &(ctx->cost), ctx->flg_Newton);
mj_constraintUpdate_island(m, d, ctx->Jaref, &(ctx->cost), ctx->flg_Newton, ctx->island);
// count active and cone
ctx->nactive = 0;
ctx->ncone = 0;
for (int i=0; i < nefc; i++) {
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
ctx->nactive += (d->efc_state[i] != mjCNSTRSTATE_SATISFIED);
ctx->ncone += (d->efc_state[i] == mjCNSTRSTATE_CONE);
}
// add Gauss cost, set in quadratic[0]
mjtNum Gauss = 0;
for (int i=0; i < nv; i++) {
Gauss += 0.5*(ctx->Ma[i]-d->qfrc_smooth[i])*(d->qacc[i]-d->qacc_smooth[i]);
for (int c=0; c < nv; c++) {
int i = dofind ? dofind[c] : c;
Gauss += 0.5 * (ctx->Ma[c] - d->qfrc_smooth[i]) * (d->qacc[i] - d->qacc_smooth[i]);
}
ctx->quadGauss[0] = Gauss;
ctx->cost += Gauss;
}
@@ -827,15 +853,18 @@ static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) {
// update grad, Mgrad
static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) {
int nv = m->nv;
static void CGupdateGradient(const mjModel* m, const mjData* d, mjCGContext* ctx) {
int nv = ctx->nv;
const int* dofind = ctx->dofind;
// grad = M*qacc - qfrc_smooth - qfrc_constraint
for (int i=0; i < nv; i++) {
ctx->grad[i] = ctx->Ma[i] - d->qfrc_smooth[i] - d->qfrc_constraint[i];
for (int c=0; c < nv; c++) {
int i = dofind ? dofind[c] : c;
ctx->grad[c] = ctx->Ma[c] - d->qfrc_smooth[i] - d->qfrc_constraint[i];
}
// Newton: Mgrad = H \ grad
// TODO: b/295296178 - add island support to Newton solver
if (ctx->flg_Newton) {
if (mj_isSparse(m)) {
mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? ctx->Hcone : ctx->H),
@@ -847,7 +876,8 @@ static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) {
// CG: Mgrad = M \ grad
else {
mj_solveM(m, d, ctx->Mgrad, ctx->grad, 1);
mju_copy(ctx->Mgrad, ctx->grad, nv);
mj_solveM_island(m, d, ctx->Mgrad, ctx->island);
}
}
@@ -855,23 +885,36 @@ static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) {
// prepare quadratic polynomials and contact cone quantities
static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) {
int nv = m->nv, nefc = d->nefc;
int nv = ctx->nv, nefc = ctx->nefc, island = ctx->island;
const int* dofind = ctx->dofind;
const int* efcind = ctx->efcind;
const mjtNum* v = ctx->search;
// Gauss: alpha^2*0.5*v'*M*v + alpha*v'*(Ma-qfrc_smooth) + 0.5*(a-qacc_smooth)'*(Ma-qfrc_smooth)
// quadGauss[0] already computed in CGupdateConstraint
ctx->quadGauss[1] = mju_dot(v, ctx->Ma, nv) - mju_dot(v, d->qfrc_smooth, nv);
mjtNum v_dot_smooth;
if (island < 0) {
v_dot_smooth = mju_dot(d->qfrc_smooth, v, nv);
} else {
v_dot_smooth = 0;
for (int c=0; c < nv; c++) {
v_dot_smooth += d->qfrc_smooth[dofind[c]] * v[c];
}
}
ctx->quadGauss[1] = mju_dot(v, ctx->Ma, nv) - v_dot_smooth;
ctx->quadGauss[2] = 0.5*mju_dot(v, ctx->Mv, nv);
// process constraints
for (int i=0; i < nefc; i++) {
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
// pointers to numeric data
mjtNum* Jv = ctx->Jv + i;
mjtNum* Jaref = ctx->Jaref + i;
mjtNum* D = d->efc_D + i;
const mjtNum* Jv = ctx->Jv + c;
const mjtNum* Jaref = ctx->Jaref + c;
const mjtNum* D = d->efc_D + i;
// pointer to this quadratic
mjtNum* quad = ctx->quad + 3*i;
mjtNum* quad = ctx->quad + 3*c;
// init with scalar quadratic
mjtNum DJ0 = D[0]*Jaref[0];
@@ -916,10 +959,10 @@ static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) {
quad[5] = UU;
quad[6] = UV;
quad[7] = VV;
quad[8] = D[0]/(mu*mu*(1+mu*mu));
quad[8] = D[0] / ((mu*mu) * (1 + (mu*mu)));
// advance to next constraint
i += (dim-1);
c += (dim-1);
}
// apply scaling
@@ -941,8 +984,9 @@ typedef struct _mjCGPnt mjCGPnt;
// evaluate linesearch cost, return first and second derivatives
static void CGeval(const mjModel* m, mjData* d, mjCGContext* ctx, mjCGPnt* p) {
int ne = d->ne, nf = d->nf, nefc = d->nefc;
static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* p) {
int ne = d->ne, nf = d->nf, nefc = ctx->nefc;
const int* efcind = ctx->efcind;
// clear result
mjtNum cost = 0, alpha = p->alpha;
@@ -953,25 +997,26 @@ static void CGeval(const mjModel* m, mjData* d, mjCGContext* ctx, mjCGPnt* p) {
mju_copy3(quadTotal, ctx->quadGauss);
// process constraints
for (int i=0; i < nefc; i++) {
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
// equality
if (i < ne) {
mju_addTo3(quadTotal, ctx->quad+3*i);
mju_addTo3(quadTotal, ctx->quad+3*c);
continue;
}
// friction
if (i < ne + nf) {
// search point, friction loss, bound (Rf)
mjtNum start = ctx->Jaref[i], dir = ctx->Jv[i];
mjtNum start = ctx->Jaref[c], dir = ctx->Jv[c];
mjtNum x = start + alpha*dir;
mjtNum f = d->efc_frictionloss[i];
mjtNum Rf = d->efc_R[i]*f;
// -bound < x < bound : quadratic
if (-Rf < x && x < Rf) {
mju_addTo3(quadTotal, ctx->quad+3*i);
mju_addTo3(quadTotal, ctx->quad+3*c);
}
// x < -bound : linear negative
@@ -992,7 +1037,7 @@ static void CGeval(const mjModel* m, mjData* d, mjCGContext* ctx, mjCGPnt* p) {
if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // elliptic cone
// extract contact info
mjContact* con = d->contact + d->efc_id[i];
mjtNum* quad = ctx->quad + 3*i;
mjtNum* quad = ctx->quad + 3*c;
int dim = con->dim;
mjtNum mu = con->mu;
@@ -1048,14 +1093,14 @@ static void CGeval(const mjModel* m, mjData* d, mjCGContext* ctx, mjCGPnt* p) {
}
// advance to next constraint
i += (dim-1);
c += (dim-1);
} else { // inequality
// search point
mjtNum x = ctx->Jaref[i] + alpha*ctx->Jv[i];
mjtNum x = ctx->Jaref[c] + alpha*ctx->Jv[c];
// active
if (x < 0) {
mju_addTo3(quadTotal, ctx->quad+3*i);
mju_addTo3(quadTotal, ctx->quad+3*c);
}
}
}
@@ -1081,8 +1126,8 @@ static void CGeval(const mjModel* m, mjData* d, mjCGContext* ctx, mjCGPnt* p) {
// update bracket point given 3 candidate points
static int updateBracket(const mjModel* m, mjData* d, mjCGContext* ctx,
mjCGPnt* p, mjCGPnt candidates[3], mjCGPnt* pnext) {
static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx,
mjCGPnt* p, const mjCGPnt candidates[3], mjCGPnt* pnext) {
int flag = 0;
for (int i=0; i < 3; i++) {
// negative deriv
@@ -1112,7 +1157,8 @@ static int updateBracket(const mjModel* m, mjData* d, mjCGContext* ctx,
// line search
static mjtNum CGsearch(const mjModel* m, mjData* d, mjCGContext* ctx) {
static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
int nv = ctx->nv;
mjCGPnt p0, p1, p2, pmid, p1next, p2next;
// clear results
@@ -1121,19 +1167,19 @@ static mjtNum CGsearch(const mjModel* m, mjData* d, mjCGContext* ctx) {
ctx->LSslope = 1; // means not computed
// save search vector length, check
mjtNum snorm = mju_norm(ctx->search, m->nv);
mjtNum snorm = mju_norm(ctx->search, nv);
if (snorm < mjMINVAL) {
ctx->LSresult = 1; // search vector too small
return 0;
}
// compute scaled gradtol and slope scaling
mjtNum gtol = m->opt.tolerance * m->opt.ls_tolerance * snorm * m->stat.meaninertia * mjMAX(1, m->nv);
mjtNum slopescl = 1 / (snorm * m->stat.meaninertia * mjMAX(1, m->nv));
mjtNum gtol = m->opt.tolerance * m->opt.ls_tolerance * snorm / ctx->scale;
mjtNum slopescl = ctx->scale / snorm;
// compute Mv, Jv
mj_mulM(m, d, ctx->Mv, ctx->search);
mj_mulJacVec(m, d, ctx->Jv, ctx->search);
mj_mulM_island(m, d, ctx->Mv, ctx->search, ctx->island, /*flg_vecunc=*/0);
mj_mulJacVec_island(m, d, ctx->Jv, ctx->search, ctx->island, /*flg_resunc=*/0, /*flg_vecunc=*/0);
// prepare quadratics and cones
CGprepare(m, d, ctx);
@@ -1293,6 +1339,7 @@ static mjtNum CGsearch(const mjModel* m, mjData* d, mjCGContext* ctx) {
// elliptic case: Hcone = H + cone_contributions
// TODO: b/295296178 - add island support to Newton solver
static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
int nv = m->nv, nefc = d->nefc;
mjtNum local[36];
@@ -1337,8 +1384,7 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
// update
mju_cholUpdateSparse(ctx->Hcone, LTJ_row, nv, 1,
ctx->rownnz, ctx->rowadr, ctx->colind, nnz, LTJ_ind,
d);
ctx->rownnz, ctx->rowadr, ctx->colind, nnz, LTJ_ind, d);
}
}
@@ -1372,6 +1418,7 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
// compute and factorize Hessian: direct method
// TODO: b/295296178 - add island support to Newton solver
static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
int nv = m->nv, nefc = d->nefc;
mj_markStack(d);
@@ -1412,8 +1459,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
// factorize H, uncompressed layout
int rank = mju_cholFactorSparse(ctx->H, nv, mjMINVAL,
ctx->rownnz, ctx->rowadr, ctx->colind,
d);
ctx->rownnz, ctx->rowadr, ctx->colind, d);
// rank-defficient, SHOULD NOT OCCUR
if (rank != nv) {
@@ -1460,6 +1506,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
// incremental update to Hessian
// TODO: b/295296178 - add island support to Newton solver
static void HessianIncremental(const mjModel* m, mjData* d,
mjCGContext* ctx, const int* oldstate) {
int rank, nv = m->nv, nefc = d->nefc;
@@ -1529,18 +1576,21 @@ static void HessianIncremental(const mjModel* m, mjData* d,
// driver
static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_Newton) {
int iter = 0, nv = m->nv, nefc = d->nefc;
static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, int flg_Newton) {
int iter = 0;
mjtNum alpha, beta;
mjtNum *gradold = NULL, *Mgradold = NULL, *Mgraddif = NULL;
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);
CGallocate(m, d, &ctx, island, flg_Newton);
// local copies
int nv = ctx.nv;
int nefc = ctx.nefc;
const int* dofind = ctx.dofind;
const int* efcind = ctx.efcind;
// allocate local storage
if (!flg_Newton) {
@@ -1551,9 +1601,17 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
int* oldstate = mj_stackAllocInt(d, nefc);
// initialize matrix-vector products
mj_mulM(m, d, ctx.Ma, d->qacc);
mj_mulJacVec(m, d, ctx.Jaref, d->qacc);
mju_subFrom(ctx.Jaref, d->efc_aref, nefc);
int flg_vecunc = 1; // d->qacc is uncompressed
mj_mulM_island(m, d, ctx.Ma, d->qacc, island, flg_vecunc);
int flg_resunc = 0; // ctx.Jaref is compressed
mj_mulJacVec_island(m, d, ctx.Jaref, d->qacc, island, flg_resunc, flg_vecunc);
if (island < 0) {
mju_subFrom(ctx.Jaref, d->efc_aref, nefc);
} else {
for (int c=0; c < nefc; c++) {
ctx.Jaref[c] -= d->efc_aref[efcind[c]];
}
}
// first update
CGupdateConstraint(m, d, &ctx);
@@ -1565,6 +1623,19 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
// start both with preconditioned gradient
mju_scl(ctx.search, ctx.Mgrad, -1, nv);
// compute and save scaling factor
mjtNum scale;
if (island < 0) {
scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv));
} else {
mjtNum island_inertia = 0;
for (int c=0; c < nv; c++) {
island_inertia += d->qM[m->dof_Madr[dofind[c]]];
}
scale = 1 / island_inertia;
}
ctx.scale = scale;
// main loop
while (iter < maxiter) {
// perform linesearch
@@ -1576,7 +1647,13 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
}
// move to new solution
mju_addToScl(d->qacc, ctx.search, alpha, nv);
if (island < 0) {
mju_addToScl(d->qacc, ctx.search, alpha, nv);
} else {
for (int c=0; c < nv; c++) {
d->qacc[dofind[c]] += alpha * ctx.search[c];
}
}
mju_addToScl(ctx.Ma, ctx.Mv, alpha, nv);
mju_addToScl(ctx.Jaref, ctx.Jv, alpha, nefc);
@@ -1585,7 +1662,13 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
mju_copy(gradold, ctx.grad, nv);
mju_copy(Mgradold, ctx.Mgrad, nv);
}
mju_copyInt(oldstate, d->efc_state, nefc);
if (island < 0) {
mju_copyInt(oldstate, d->efc_state, nefc);
} else {
for (int c=0; c < nefc; c++) {
oldstate[c] = d->efc_state[efcind[c]];
}
}
mjtNum oldcost = ctx.cost;
// update
@@ -1597,13 +1680,14 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
// count state changes
int nchange = 0;
for (int i=0; i < nefc; i++) {
nchange += (d->efc_state[i] != oldstate[i]);
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
nchange += (d->efc_state[i] != oldstate[c]);
}
// scale improvement, save stats
mjtNum improvement = rescale(m, oldcost-ctx.cost);
mjtNum gradient = rescale(m, mju_norm(ctx.grad, nv));
// scale improvement, gradient, save stats
mjtNum improvement = scale * (oldcost - ctx.cost);
mjtNum gradient = scale * mju_norm(ctx.grad, nv);
saveStats(m, d, island, iter, improvement, gradient, ctx.LSslope,
ctx.nactive, nchange, ctx.LSiter, ctx.nupdate);
@@ -1630,26 +1714,29 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
}
// update
for (int i=0; i < nv; i++) {
ctx.search[i] = -ctx.Mgrad[i] + beta*ctx.search[i];
for (int c=0; c < nv; c++) {
ctx.search[c] = -ctx.Mgrad[c] + beta*ctx.search[c];
}
}
}
// finalize statistics
if (island < mjNISLAND) {
// if island is -1 (monolithic), clamp to 0
int island_stat = island < 0 ? 0 : island;
// update solver iterations
d->solver_niter[island] += iter;
d->solver_niter[island_stat] += iter;
// set solver_nnz
if (flg_Newton) {
if (mj_isSparse(m)) {
d->solver_nnz[island] = 2*ctx.nnz - nv;
d->solver_nnz[island_stat] = 2*ctx.nnz - nv;
} else {
d->solver_nnz[island] = nv*nv;
d->solver_nnz[island_stat] = nv*nv;
}
} else {
d->solver_nnz[island] = 0;
d->solver_nnz[island_stat] = 0;
}
}
@@ -1660,12 +1747,19 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
// CG entry point
void mj_solCG(const mjModel* m, mjData* d, int maxiter) {
mj_solCGNewton(m, d, maxiter, 0);
mj_solCGNewton(m, d, /*island=*/-1, maxiter, /*flg_Newton=*/0);
}
// CG entry point (one island)
void mj_solCG_island(const mjModel* m, mjData* d, int island, int maxiter) {
mj_solCGNewton(m, d, island, maxiter, /*flg_Newton=*/0);
}
// Newton entry point
void mj_solNewton(const mjModel* m, mjData* d, int maxiter) {
mj_solCGNewton(m, d, maxiter, 1);
mj_solCGNewton(m, d, /*island=*/-1, maxiter, /*flg_Newton=*/1);
}
src/engine/engine_solver.h
+6
View File
@@ -18,6 +18,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
//------------------------------ monolithic solvers ------------------------------------------------
// PGS solver
void mj_solPGS(const mjModel* m, mjData* d, int maxiter);
@@ -31,4 +32,9 @@ 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);
#endif // MUJOCO_SRC_ENGINE_ENGINE_SOLVER_H_
test/engine/CMakeLists.txt
+3
View File
@@ -72,6 +72,9 @@ target_link_libraries(engine_resource_test fixture gmock)
mujoco_test(engine_sensor_test)
target_link_libraries(engine_sensor_test fixture gmock)
mujoco_test(engine_solver_test)
target_link_libraries(engine_solver_test fixture gmock)
mujoco_test(engine_support_test)
target_link_libraries(engine_support_test fixture gmock)
test/engine/engine_solver_test.cc
+165
View File
@@ -0,0 +1,165 @@
// Copyright 2023 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Tests for engine/engine_solver.c
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <mujoco/mujoco.h>
#include "test/fixture.h"
namespace mujoco {
namespace {
std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
return std::vector<mjtNum>(array, array + n);
}
using ::testing::DoubleNear;
using ::testing::NotNull;
using ::testing::Pointwise;
using SolverTest = MujocoTest;
static const char* const kIlslandEfcPath =
"engine/testdata/island/island_efc.xml";
// compare accelerations produced by CG solver with and without islands
TEST_F(SolverTest, IslandsEquivalent) {
const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
ASSERT_THAT(model, NotNull());
model->opt.solver = mjSOL_CG; // use CG solver
model->opt.tolerance = 0; // set tolerance to 0
model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands
int nv = model->nv;
int state_size = mj_stateSize(model, mjSTATE_INTEGRATION);
mjtNum* state = (mjtNum*) mju_malloc(sizeof(mjtNum)*state_size);
mjtNum* qacc_diff = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv);
mjData* data_island = mj_makeData(model);
mjData* data_noisland = mj_makeData(model);
mjtNum tol = 2e-4;
for (bool warmstart : {true, false}) {
if (warmstart) {
model->opt.disableflags |= mjDSBL_WARMSTART;
} else {
model->opt.disableflags &= ~mjDSBL_WARMSTART;
}
mj_resetData(model, data_noisland);
while (data_noisland->time < .3) {
mj_step(model, data_noisland);
mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION);
mj_setState(model, data_island, state, mjSTATE_INTEGRATION);
mj_forward(model, data_noisland);
model->opt.enableflags |= mjENBL_ISLAND; // enable islands
mj_forward(model, data_island);
model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands
EXPECT_THAT(AsVector(data_noisland->qacc, nv),
Pointwise(DoubleNear(tol), AsVector(data_island->qacc, nv)));
}
}
mj_deleteData(data_noisland);
mj_deleteData(data_island);
mju_free(qacc_diff);
mju_free(state);
mj_deleteModel(model);
}
// compare qacc from 1 iteration of monolithic CG solver and one big island
TEST_F(SolverTest, OneBigIsland) {
const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
ASSERT_THAT(model, NotNull());
model->opt.solver = mjSOL_CG; // use CG solver
model->opt.disableflags |= mjDSBL_WARMSTART; // disable warmstart
model->opt.tolerance = 0; // set tolerance to 0
model->opt.enableflags &= ~mjENBL_ISLAND; // disable islands
int state_size = mj_stateSize(model, mjSTATE_INTEGRATION);
mjtNum* state = (mjtNum*) mju_malloc(sizeof(mjtNum)*state_size);
mjData* data_island = mj_makeData(model);
mjData* data_noisland = mj_makeData(model);
int nv = model->nv;
mjtNum tol = 1e-9;
// save current (default) iterations
int iterations_default = model->opt.iterations;
while (data_noisland->time < .2) {
// step and copy the state to data_island
mj_step(model, data_noisland);
mj_getState(model, data_noisland, state, mjSTATE_INTEGRATION);
mj_setState(model, data_island, state, mjSTATE_INTEGRATION);
// set small number of iterations
model->opt.iterations = 1;
// call forward on data_noisland
mj_forward(model, data_noisland);
// enable islands
model->opt.enableflags |= mjENBL_ISLAND;
// call forward (just for smooth dynamics and to allocate islands)
mj_forward(model, data_island);
// overwrite island structure with one big island
data_island->nisland = 1;
data_island->island_dofnum[0] = nv;
data_island->island_dofadr[0] = 0;
for (int i = 0; i < nv; i++) {
data_island->island_dofind[i] = data_island->dof_islandind[i] = i;
}
int nefc = data_island->nefc;
data_island->island_efcnum[0] = nefc;
data_island->island_efcadr[0] = 0;
for (int i = 0; i < nefc; i++) data_island->island_efcind[i] = i;
// solve using using one big island
mj_fwdConstraint(model, data_island);
// re-disable islands and reset iterations
model->opt.enableflags &= ~mjENBL_ISLAND;
model->opt.iterations = iterations_default;
// compare accelerations
EXPECT_THAT(AsVector(data_noisland->qacc, nv),
Pointwise(DoubleNear(tol), AsVector(data_island->qacc, nv)));
}
mj_deleteData(data_noisland);
mj_deleteData(data_island);
mju_free(state);
mj_deleteModel(model);
}
} // namespace
} // namespace mujoco
test/pipeline_test.cc
+1 -1
View File
@@ -64,7 +64,7 @@ TEST_F(PipelineTest, SparseDenseEquivalent) {
mj_deleteModel(model);
}
// mj_forward should be deterministic when warm starts are disabled
// mj_forward should be idempotent when warm starts are disabled
TEST_F(PipelineTest, DeterministicNoWarmstart) {
const std::string xml_path = GetTestDataFilePath(kDefaultModel);
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);