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

Refactor islands to be memory contiguous.

Browse Source 
PiperOrigin-RevId: 755803476
Change-Id: I41972b07e0d5ef5d0117c94f565b93367b87458b
This commit is contained in:
Yuval Tassa
2025-05-07 05:05:34 -07:00
committed by Copybara-Service
parent 449de73430
commit ecb769fc3a
30 changed files with 1742 additions and 1116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_solver.c
+279 -136
View File
@@ -766,13 +766,55 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
// CG context
struct _mjCGContext {
int flg_Newton; // 1: Newton, 0: CG
// island-related
int island; // current island index, -1 if monolithic
// sizes
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
int ne; // number of equalities
int nf; // number of friction constraints
int nefc; // number of all constraints
// contact array
mjContact* contact;
// dof arrays
const mjtNum* qfrc_smooth;
const mjtNum* qacc_smooth;
mjtNum* qfrc_constraint;
mjtNum* qacc;
// inertia
const int* M_rownnz;
const int* M_rowadr;
const int* M_diagnum;
const int* M_colind;
const int* dof_Madr;
const int* dof_parentid;
const mjtNum* qM;
const mjtNum* qLD;
const mjtNum* qLDiagInv;
// efc arrays
const mjtNum* efc_D;
const mjtNum* efc_R;
const mjtNum* efc_frictionloss;
const mjtNum* efc_aref;
const int* efc_id;
const int* efc_type;
mjtNum* efc_force;
int* efc_state;
// Jacobians
const int* J_rownnz;
const int* J_rowadr;
const int* J_rowsuper;
const int* J_colind;
const int* JT_rownnz;
const int* JT_rowadr;
const int* JT_rowsuper;
const int* JT_colind;
const mjtNum* J;
const mjtNum* JT;
// common arrays (CGallocate)
mjtNum* Jaref; // Jac*qacc - aref (nefc x 1)
@@ -793,7 +835,7 @@ struct _mjCGContext {
int* L_rownnz; // Hessian factor row nonzeros (nv x 1)
int* L_rowadr; // Hessian factor row addresses (nv x 1)
// Newton arrays, computed-size (HessianMake)
// Newton arrays, computed-size (MakeHessian)
int nH; // number of nonzeros in Hessian H
int* H_colind; // Hessian column indices (nH x 1)
mjtNum* H; // Hessian (nH x 1)
@@ -818,23 +860,130 @@ struct _mjCGContext {
typedef struct _mjCGContext mjCGContext;
// set sizes and pointers to mjData arrays in mjCGContext
static void CGpointers(const mjModel* m, const mjData* d, mjCGContext* ctx, int island) {
int is_sparse = mj_isSparse(m);
ctx->contact = d->contact;
ctx->island = island;
// set sizes and pointers (monolithic)
if (island < 0) {
// sizes
ctx->nv = m->nv;
ctx->ne = d->ne;
ctx->nf = d->nf;
ctx->nefc = d->nefc;
// dof arrays
ctx->qfrc_smooth = d->qfrc_smooth;
ctx->qfrc_constraint = d->qfrc_constraint;
ctx->qacc_smooth = d->qacc_smooth;
ctx->qacc = d->qacc;
// inertia
ctx->M_rownnz = d->M_rownnz;
ctx->M_rowadr = d->M_rowadr;
ctx->M_diagnum = m->dof_simplenum;
ctx->M_colind = d->M_colind;
ctx->dof_Madr = m->dof_Madr;
ctx->dof_parentid = m->dof_parentid;
ctx->qM = d->qM;
ctx->qLD = d->qLD;
ctx->qLDiagInv = d->qLDiagInv;
// efc arrays
ctx->efc_D = d->efc_D;
ctx->efc_R = d->efc_R;
ctx->efc_frictionloss = d->efc_frictionloss;
ctx->efc_aref = d->efc_aref;
ctx->efc_id = d->efc_id;
ctx->efc_type = d->efc_type;
ctx->efc_force = d->efc_force;
ctx->efc_state = d->efc_state;
// Jacobians
ctx->J = d->efc_J;
if (is_sparse) {
ctx->J_rownnz = d->efc_J_rownnz;
ctx->J_rowadr = d->efc_J_rowadr;
ctx->J_rowsuper = d->efc_J_rowsuper;
ctx->J_colind = d->efc_J_colind;
ctx->JT_rownnz = d->efc_JT_rownnz;
ctx->JT_rowadr = d->efc_JT_rowadr;
ctx->JT_rowsuper = d->efc_JT_rowsuper;
ctx->JT_colind = d->efc_JT_colind;
ctx->JT = d->efc_JT;
}
}
// set sizes and pointers (per-island)
else {
// sizes
ctx->nv = d->island_nv[island];
ctx->ne = d->island_ne[island];
ctx->nf = d->island_nf[island];
ctx->nefc = d->island_nefc[island];
// dof arrays
int idofadr = d->island_idofadr[island];
ctx->qfrc_smooth = d->ifrc_smooth + idofadr;
ctx->qfrc_constraint = d->ifrc_constraint + idofadr;
ctx->qacc_smooth = d->iacc_smooth + idofadr;
ctx->qacc = d->iacc + idofadr;
// inertia
ctx->M_rownnz = d->iM_rownnz + idofadr;
ctx->M_rowadr = d->iM_rowadr + idofadr;
ctx->M_diagnum = d->iM_diagnum + idofadr;
ctx->M_colind = d->iM_colind;
ctx->qM = d->iM;
ctx->qLD = d->iLD;
ctx->qLDiagInv = d->iLDiagInv + idofadr;
// efc arrays
int iefcadr = d->island_iefcadr[island];
ctx->efc_D = d->iefc_D + iefcadr;
ctx->efc_R = d->iefc_R + iefcadr;
ctx->efc_frictionloss = d->iefc_frictionloss + iefcadr;
ctx->efc_aref = d->iefc_aref + iefcadr;
ctx->efc_id = d->iefc_id + iefcadr;
ctx->efc_type = d->iefc_type + iefcadr;
ctx->efc_force = d->iefc_force + iefcadr;
ctx->efc_state = d->iefc_state + iefcadr;
// Jacobians
if (!is_sparse) {
ctx->J = d->iefc_J + d->nidof * iefcadr;
} else {
ctx->J_rownnz = d->iefc_J_rownnz + iefcadr;
ctx->J_rowadr = d->iefc_J_rowadr + iefcadr;
ctx->J_rowsuper = d->iefc_J_rowsuper + iefcadr;
ctx->J_colind = d->iefc_J_colind;
ctx->JT_rownnz = d->iefc_JT_rownnz + idofadr;
ctx->JT_rowadr = d->iefc_JT_rowadr + idofadr;
ctx->JT_rowsuper = d->iefc_JT_rowsuper + idofadr;
ctx->JT_colind = d->iefc_JT_colind;
ctx->J = d->iefc_J;
ctx->JT = d->iefc_JT;
}
}
}
// allocate fixed-size arrays in mjCGContext
// mj_{mark/free}Stack in calling function!
static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
int island, int flg_Newton) {
static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, int island, int flg_Newton) {
// clear everything
memset(ctx, 0, sizeof(mjCGContext));
// get sizes
int nv = island < 0 ? m->nv : d->island_dofnum[island];
int nefc = island < 0 ? d->nefc : d->island_efcnum[island];
// set sizes and pointers
CGpointers(m, d, ctx, island);
// 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];
// local sizes
int nv = ctx->nv;
int nefc = ctx->nefc;
// common arrays
ctx->Jaref = mjSTACKALLOC(d, nefc, mjtNum);
@@ -849,7 +998,7 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
// Newton only, known-size arrays
ctx->flg_Newton = flg_Newton;
if (flg_Newton) {
ctx->D = mjSTACKALLOC(d, nefc, mjtNum);
ctx->D = mjSTACKALLOC(d, nefc, mjtNum);
// sparse Newton only
if (mj_isSparse(m)) {
@@ -866,28 +1015,35 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
// update efc_force, qfrc_constraint, cost-related
static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) {
static void CGupdateConstraint(mjCGContext* ctx) {
int nefc = ctx->nefc, nv = ctx->nv;
const int* dofind = ctx->dofind;
const int* efcind = ctx->efcind;
// update constraints
mj_constraintUpdate_island(m, d, ctx->Jaref, &(ctx->cost), ctx->flg_Newton, ctx->island);
mj_constraintUpdate_impl(ctx->ne, ctx->nf, ctx->nefc, ctx->efc_D, ctx->efc_R,
ctx->efc_frictionloss, ctx->Jaref, ctx->efc_type, ctx->efc_id,
ctx->contact, ctx->efc_state, ctx->efc_force,
&(ctx->cost), ctx->flg_Newton);
// compute qfrc_constraint (dense or sparse)
if (!ctx->JT) {
mju_mulMatTVec(ctx->qfrc_constraint, ctx->J, ctx->efc_force, nefc, nv);
} else {
mju_mulMatVecSparse(ctx->qfrc_constraint, ctx->JT, ctx->efc_force, nv,
ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper);
}
// count active and cone
ctx->nactive = 0;
ctx->ncone = 0;
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);
for (int i=0; i < nefc; i++) {
ctx->nactive += (ctx->efc_state[i] != mjCNSTRSTATE_SATISFIED);
ctx->ncone += (ctx->efc_state[i] == mjCNSTRSTATE_CONE);
}
// add Gauss cost, set in quadratic[0]
mjtNum Gauss = 0;
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]);
for (int i=0; i < nv; i++) {
Gauss += 0.5 * (ctx->Ma[i] - ctx->qfrc_smooth[i]) * (ctx->qacc[i] - ctx->qacc_smooth[i]);
}
ctx->quadGauss[0] = Gauss;
@@ -895,22 +1051,20 @@ static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) {
}
// TODO(tassa): Restore mjData const-ness.
// update grad, Mgrad
static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) {
static void CGupdateGradient(mjCGContext* ctx) {
int nv = ctx->nv;
const int* dofind = ctx->dofind;
// grad = M*qacc - qfrc_smooth - qfrc_constraint
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];
for (int i=0; i < nv; i++) {
ctx->grad[i] = ctx->Ma[i] - ctx->qfrc_smooth[i] - ctx->qfrc_constraint[i];
}
// Newton: Mgrad = H \ grad
// TODO: b/295296178 - add island support to Newton solver
if (ctx->flg_Newton) {
if (mj_isSparse(m)) {
if (ctx->L_rowadr) {
mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? ctx->Lcone : ctx->L),
ctx->grad, nv, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind);
} else {
@@ -921,44 +1075,32 @@ static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) {
// CG: Mgrad = M \ grad
else {
mju_copy(ctx->Mgrad, ctx->grad, nv);
mj_solveM_island(m, d, ctx->Mgrad, ctx->island);
mj_solveLD(ctx->Mgrad, ctx->qLD, ctx->qLDiagInv, nv, 1,
ctx->M_rownnz, ctx->M_rowadr, ctx->M_diagnum, ctx->M_colind);
}
}
// prepare quadratic polynomials and contact cone quantities
static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) {
int nv = ctx->nv, nefc = ctx->nefc, island = ctx->island;
const int* dofind = ctx->dofind;
const int* efcind = ctx->efcind;
static void CGprepare(mjCGContext* ctx) {
int nv = ctx->nv, nefc = ctx->nefc;
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
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[1] = mju_dot(v, ctx->Ma, nv) - mju_dot(ctx->qfrc_smooth, v, nv);
ctx->quadGauss[2] = 0.5*mju_dot(v, ctx->Mv, nv);
// process constraints
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
for (int i=0; i < nefc; i++) {
// pointers to numeric data
const mjtNum* Jv = ctx->Jv + c;
const mjtNum* Jaref = ctx->Jaref + c;
const mjtNum* D = d->efc_D + i;
const mjtNum* Jv = ctx->Jv + i;
const mjtNum* Jaref = ctx->Jaref + i;
const mjtNum* D = ctx->efc_D + i;
// pointer to this quadratic
mjtNum* quad = ctx->quad + 3*c;
mjtNum* quad = ctx->quad + 3*i;
// init with scalar quadratic
mjtNum DJ0 = D[0]*Jaref[0];
@@ -967,12 +1109,12 @@ static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) {
quad[2] = Jv[0]*D[0]*Jv[0];
// elliptic cone: extra processing
if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) {
if (ctx->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) {
// extract contact info
mjContact* con = d->contact + d->efc_id[i];
const mjContact* con = ctx->contact + ctx->efc_id[i];
int dim = con->dim;
mjtNum U[6], V[6], UU = 0, UV = 0, VV = 0, mu = con->mu;
mjtNum* friction = con->friction;
const mjtNum* friction = con->friction;
// complete vector quadratic (for bottom zone)
for (int j=1; j < dim; j++) {
@@ -1006,7 +1148,7 @@ static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) {
quad[8] = D[0] / ((mu*mu) * (1 + (mu*mu)));
// advance to next constraint
c += (dim-1);
i += (dim-1);
}
// apply scaling
@@ -1028,9 +1170,8 @@ typedef struct _mjCGPnt mjCGPnt;
// evaluate linesearch cost, return first and second derivatives
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;
static void CGeval(mjCGContext* ctx, mjCGPnt* p) {
int ne = ctx->ne, nf = ctx->nf, nefc = ctx->nefc;
// clear result
mjtNum cost = 0, alpha = p->alpha;
@@ -1041,26 +1182,24 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt*
mju_copy3(quadTotal, ctx->quadGauss);
// process constraints
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
for (int i=0; i < nefc; i++) {
// equality
if (i < ne) {
mju_addTo3(quadTotal, ctx->quad+3*c);
mju_addTo3(quadTotal, ctx->quad+3*i);
continue;
}
// friction
if (i < ne + nf) {
// search point, friction loss, bound (Rf)
mjtNum start = ctx->Jaref[c], dir = ctx->Jv[c];
mjtNum start = ctx->Jaref[i], dir = ctx->Jv[i];
mjtNum x = start + alpha*dir;
mjtNum f = d->efc_frictionloss[i];
mjtNum Rf = d->efc_R[i]*f;
mjtNum f = ctx->efc_frictionloss[i];
mjtNum Rf = ctx->efc_R[i]*f;
// -bound < x < bound : quadratic
if (-Rf < x && x < Rf) {
mju_addTo3(quadTotal, ctx->quad+3*c);
mju_addTo3(quadTotal, ctx->quad+3*i);
}
// x < -bound : linear negative
@@ -1078,10 +1217,10 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt*
}
// limit and contact
if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // elliptic cone
if (ctx->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // elliptic cone
// extract contact info
mjContact* con = d->contact + d->efc_id[i];
mjtNum* quad = ctx->quad + 3*c;
const mjContact* con = ctx->contact + ctx->efc_id[i];
mjtNum* quad = ctx->quad + 3*i;
int dim = con->dim;
mjtNum mu = con->mu;
@@ -1137,14 +1276,14 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt*
}
// advance to next constraint
c += (dim-1);
i += (dim-1);
} else { // inequality
// search point
mjtNum x = ctx->Jaref[c] + alpha*ctx->Jv[c];
mjtNum x = ctx->Jaref[i] + alpha*ctx->Jv[i];
// active
if (x < 0) {
mju_addTo3(quadTotal, ctx->quad+3*c);
mju_addTo3(quadTotal, ctx->quad+3*i);
}
}
}
@@ -1170,7 +1309,7 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt*
// update bracket point given 3 candidate points
static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx,
static int updateBracket(mjCGContext* ctx,
mjCGPnt* p, const mjCGPnt candidates[3], mjCGPnt* pnext) {
int flag = 0;
for (int i=0; i < 3; i++) {
@@ -1192,7 +1331,7 @@ static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx,
// compute next point if updated
if (flag) {
pnext->alpha = p->alpha - p->deriv[0]/p->deriv[1];
CGeval(m, d, ctx, pnext);
CGeval(ctx, pnext);
}
return flag;
@@ -1201,8 +1340,8 @@ static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx,
// line search
static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
int nv = ctx->nv;
static mjtNum CGsearch(mjCGContext* ctx, mjtNum tolerance, mjtNum ls_iterations) {
int nv = ctx->nv, nefc = ctx->nefc;
mjCGPnt p0, p1, p2, pmid, p1next, p2next;
// clear results
@@ -1218,23 +1357,36 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
}
// compute scaled gradtol and slope scaling
mjtNum gtol = m->opt.tolerance * m->opt.ls_tolerance * snorm / ctx->scale;
mjtNum gtol = tolerance * snorm / ctx->scale;
mjtNum slopescl = ctx->scale / snorm;
// compute Mv, Jv
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);
// compute Mv = M * v (island or monolithic)
if (ctx->island >= 0) {
mju_mulSymVecSparse(ctx->Mv, ctx->qM, ctx->search, nv,
ctx->M_rownnz, ctx->M_rowadr, ctx->M_diagnum, ctx->M_colind);
} else {
mj_mulM_impl(ctx->Mv, ctx->search, nv, ctx->qM,
ctx->dof_Madr, ctx->dof_parentid, ctx->M_diagnum);
}
// compute Jv = J * search (dense or sparse)
if (!ctx->J_rowadr) {
mju_mulMatVec(ctx->Jv, ctx->J, ctx->search, nefc, nv);
} else {
mju_mulMatVecSparse(ctx->Jv, ctx->J, ctx->search, nefc,
ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, ctx->J_rowsuper);
}
// prepare quadratics and cones
CGprepare(m, d, ctx);
CGprepare(ctx);
// init at alpha = 0, save
p0.alpha = 0;
CGeval(m, d, ctx, &p0);
CGeval(ctx, &p0);
// always attempt one Newton step
p1.alpha = p0.alpha - p0.deriv[0]/p0.deriv[1];
CGeval(m, d, ctx, &p1);
CGeval(ctx, &p1);
if (p0.cost < p1.cost) {
p1 = p0;
}
@@ -1289,14 +1441,14 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
// one-sided search
int p2update = 0;
while (p1.deriv[0]*dir <= -gtol && ctx->LSiter < m->opt.ls_iterations) {
while (p1.deriv[0]*dir <= -gtol && ctx->LSiter < ls_iterations) {
// save current
p2 = p1;
p2update = 1;
// move to Newton point w.r.t current
p1.alpha -= p1.deriv[0]/p1.deriv[1];
CGeval(m, d, ctx, &p1);
CGeval(ctx, &p1);
// check for convergence
if (mju_abs(p1.deriv[0]) < gtol) {
@@ -1306,7 +1458,7 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
}
// check for failure to bracket
if (ctx->LSiter >= m->opt.ls_iterations) {
if (ctx->LSiter >= ls_iterations) {
ctx->LSresult = 3; // could not bracket
ctx->LSslope = mju_abs(p1.deriv[0])*slopescl;
return p1.alpha;
@@ -1322,13 +1474,13 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
// compute next-points for bracket
p2next = p1;
p1next.alpha = p1.alpha - p1.deriv[0]/p1.deriv[1];
CGeval(m, d, ctx, &p1next);
CGeval(ctx, &p1next);
// bracketed search
while (ctx->LSiter < m->opt.ls_iterations) {
while (ctx->LSiter < ls_iterations) {
// evaluate at midpoint
pmid.alpha = 0.5*(p1.alpha + p2.alpha);
CGeval(m, d, ctx, &pmid);
CGeval(ctx, &pmid);
// make list of candidates
mjCGPnt candidates[3] = {p1next, p2next, pmid};
@@ -1349,8 +1501,8 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
}
// update brackets
int b1 = updateBracket(m, d, ctx, &p1, candidates, &p1next);
int b2 = updateBracket(m, d, ctx, &p2, candidates, &p2next);
int b1 = updateBracket(ctx, &p1, candidates, &p1next);
int b2 = updateBracket(ctx, &p2, candidates, &p2next);
// no update possible: numerical accuracy reached, use midpoint
if (!b1 && !b2) {
@@ -1730,8 +1882,6 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
// 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) {
@@ -1741,27 +1891,32 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
}
int* oldstate = mjSTACKALLOC(d, nefc, int);
// initialize matrix-vector products
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);
// compute Ma = M * qacc (island or monolithic)
if (island >= 0) {
mju_mulSymVecSparse(ctx.Ma, ctx.qM, ctx.qacc, nv,
ctx.M_rownnz, ctx.M_rowadr, ctx.M_diagnum, ctx.M_colind);
} else {
for (int c=0; c < nefc; c++) {
ctx.Jaref[c] -= d->efc_aref[efcind[c]];
}
mj_mulM_impl(ctx.Ma, ctx.qacc, nv, ctx.qM,
ctx.dof_Madr, ctx.dof_parentid, ctx.M_diagnum);
}
// compute Jaref = J * qacc - aref (dense or sparse)
if (!ctx.J_rownnz) {
mju_mulMatVec(ctx.Jaref, ctx.J, ctx.qacc, nefc, nv);
} else {
mju_mulMatVecSparse(ctx.Jaref, ctx.J, ctx.qacc, nefc,
ctx.J_rownnz, ctx.J_rowadr, ctx.J_colind, ctx.J_rowsuper);
}
mju_subFrom(ctx.Jaref, ctx.efc_aref, nefc);
// first update
CGupdateConstraint(m, d, &ctx);
CGupdateConstraint(&ctx);
if (flg_Newton) {
// compute and factorize Hessian
MakeHessian(m, d, &ctx);
FactorizeHessian(m, d, &ctx, /*flg_recompute=*/0);
}
CGupdateGradient(m, d, &ctx);
CGupdateGradient(&ctx);
// start both with preconditioned gradient
mju_scl(ctx.search, ctx.Mgrad, -1, nv);
@@ -1772,8 +1927,9 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
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]]];
for (int i=0; i < nv; i++) {
int* map2dof = d->map_idof2dof + d->island_idofadr[island];
island_inertia += d->qM[m->dof_Madr[map2dof[i]]];
}
scale = 1 / island_inertia;
}
@@ -1782,7 +1938,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
// main loop
while (iter < maxiter) {
// perform linesearch
alpha = CGsearch(m, d, &ctx);
alpha = CGsearch(&ctx, m->opt.tolerance * m->opt.ls_tolerance, m->opt.ls_iterations);
// no improvement: done
if (alpha == 0) {
@@ -1790,13 +1946,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
}
// move to new solution
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.qacc, ctx.search, alpha, nv);
mju_addToScl(ctx.Ma, ctx.Mv, alpha, nv);
mju_addToScl(ctx.Jaref, ctx.Jv, alpha, nefc);
@@ -1805,27 +1955,20 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
mju_copy(gradold, ctx.grad, nv);
mju_copy(Mgradold, ctx.Mgrad, nv);
}
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]];
}
}
mju_copyInt(oldstate, ctx.efc_state, nefc);
mjtNum oldcost = ctx.cost;
// update
CGupdateConstraint(m, d, &ctx);
CGupdateConstraint(&ctx);
if (flg_Newton) {
HessianIncremental(m, d, &ctx, oldstate);
}
CGupdateGradient(m, d, &ctx);
CGupdateGradient(&ctx);
// count state changes
int nchange = 0;
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
nchange += (d->efc_state[i] != oldstate[c]);
for (int i=0; i < nefc; i++) {
nchange += (ctx.efc_state[i] != oldstate[i]);
}
// scale improvement, gradient, save stats
@@ -1857,8 +2000,8 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
}
// update
for (int c=0; c < nv; c++) {
ctx.search[c] = -ctx.Mgrad[c] + beta*ctx.search[c];
for (int i=0; i < nv; i++) {
ctx.search[i] = -ctx.Mgrad[i] + beta*ctx.search[i];
}
}
}