// Copyright 2021 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. #include "engine/engine_solver.h" #include #include #include #include #include #include #include // IWYU pragma: keep #include "engine/engine_core_constraint.h" #include "engine/engine_core_smooth.h" #include "engine/engine_core_util.h" #include "engine/engine_memory.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" #include "engine/engine_util_solve.h" #include "engine/engine_util_sparse.h" //---------------------------------- utility functions --------------------------------------------- // 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) { // if island out of range, return if (island >= mjNISLAND) { return; } // if no islands, use first island island = mjMAX(0, island); // if iter out of range, return if (iter >= mjNSOLVER) { return; } // get mjSolverStat pointer 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; } // finalize dual solver: map to joint space static void dualFinish(const mjModel* m, mjData* d) { // map constraint force to joint space mj_mulJacTVec(m, d, d->qfrc_constraint, d->efc_force); // compute constrained acceleration in joint space mj_solveM(m, d, d->qacc, d->qfrc_constraint, 1); mju_addTo(d->qacc, d->qacc_smooth, m->nv); } // 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 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; const mjtNum *R = d->efc_R; // sparse if (mj_isSparse(m)) { const int *rowadr = d->efc_AR_rowadr; const int *rownnz = d->efc_AR_rownnz; const int *colind = d->efc_AR_colind; for (int c=0; c < nefc; c++) { int i = efclist ? efclist[c] : c; int nnz = rownnz[i]; for (int j=0; j < nnz; j++) { int adr = rowadr[i] + j; if (i == colind[adr]) { res[c] = 1 / (flg_subR ? mju_max(mjMINVAL, AR[adr] - R[i]) : AR[adr]); break; } } } } // dense else { int d_nefc = d->nefc; // global nefc for (int c=0; c < nefc; c++) { int i = efclist ? efclist[c] : c; int adr = i * (d_nefc + 1); res[c] = 1 / (flg_subR ? mju_max(mjMINVAL, AR[adr] - R[i]) : AR[adr]); } } } // extract diagonal block from AR, clamp diag to 1e-10 if flg_subR static void extractBlock(const mjModel* m, const mjData* d, mjtNum* Ac, int start, int n, int flg_subR) { int nefc = d->nefc; const mjtNum *AR = d->efc_AR; // sparse if (mj_isSparse(m)) { const int* rownnz = d->efc_AR_rownnz; const int* rowadr = d->efc_AR_rowadr; const int* colind = d->efc_AR_colind; /* // GENERAL CASE mju_zero(Ac, n*n); for( j=0; j=start && col= rownnz[start]) { mjERROR("internal error"); } // copy rows for (int j=0; j < n; j++) { mju_copy(Ac+j*n, AR+rowadr[start+j]+k, n); } } // dense else { for (int j=0; j < n; j++) { mju_copy(Ac+j*n, AR+start+(start+j)*nefc, n); } } // subtract R from diagonal, clamp to 1e-10 from below if (flg_subR) { const mjtNum *R = d->efc_R; for (int j=0; j < n; j++) { Ac[j*(n+1)] -= R[start+j]; Ac[j*(n+1)] = mju_max(1e-10, Ac[j*(n+1)]); } } } // compute residual for one block static void residual(const mjModel* m, const mjData* d, mjtNum* res, int i, int dim, int flg_subR) { int nefc = d->nefc; // sparse if (mj_isSparse(m)) { for (int j=0; j < dim; j++) { res[j] = d->efc_b[i+j] + mju_dotSparse(d->efc_AR + d->efc_AR_rowadr[i+j], d->efc_force, d->efc_AR_rownnz[i+j], d->efc_AR_colind + d->efc_AR_rowadr[i+j]); } } // dense else { for (int j=0; j < dim; j++) { res[j] = d->efc_b[i+j] + mju_dot(d->efc_AR+(i+j)*nefc, d->efc_force, nefc); } } if (flg_subR) { for (int j=0; j < dim; j++) { res[j] -= d->efc_R[i+j]*d->efc_force[i+j]; } } } // compute cost change static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce, const mjtNum* res, int dim) { mjtNum change; // compute change if (dim == 1) { mjtNum delta = force[0] - oldforce[0]; change = 0.5*delta*delta*A[0] + delta*res[0]; } else { mjtNum delta[6]; mju_sub(delta, force, oldforce, dim); change = 0.5*mju_mulVecMatVec(delta, A, delta, dim) + mju_dot(delta, res, dim); } // positive change: restore force if (change > 1e-10) { mju_copy(force, oldforce, dim); change = 0; } return change; } // PCG32 random number generator state typedef struct { uint64_t state; uint64_t inc; } pcg32_state; // generate next 32-bit pseudorandom integer static uint32_t pcg32_next(pcg32_state* rng) { uint64_t oldstate = rng->state; rng->state = oldstate * 6364136223846793005ULL + (rng->inc | 1); uint32_t xorshifted = ((oldstate >> 18u) ^ oldstate) >> 27u; uint32_t rot = oldstate >> 59u; return (xorshifted >> rot) | (xorshifted << ((-rot) & 31)); } // Fisher-Yates shuffle of integer array static void shuffle_int(int* array, int n, pcg32_state* rng) { for (int i = n - 1; i > 0; i--) { uint32_t j = pcg32_next(rng) % (i + 1); int temp = array[i]; array[i] = array[j]; array[j] = temp; } } // set efc_state to dual constraint state; return nactive // iterates over efclist (or sequentially if NULL), classifies by ne/nf ranges static int dualState(const mjData* d, int* state, int ne, int nf, int nefc, const int* efclist) { const mjtNum* force = d->efc_force; const mjtNum* floss = d->efc_frictionloss; // equality and friction always active int nactive = ne + nf; // equality for (int c=0; c < ne; c++) { int i = efclist ? efclist[c] : c; state[i] = mjCNSTRSTATE_QUADRATIC; } // friction for (int c=ne; c < ne+nf; c++) { int i = efclist ? efclist[c] : c; if (force[i] <= -floss[i]) { state[i] = mjCNSTRSTATE_LINEARPOS; // opposite of primal } else if (force[i] >= floss[i]) { state[i] = mjCNSTRSTATE_LINEARNEG; } else { state[i] = mjCNSTRSTATE_QUADRATIC; } } // limit and contact for (int c=ne+nf; c < nefc; c++) { int i = efclist ? efclist[c] : c; // non-negative if (d->efc_type[i] != mjCNSTR_CONTACT_ELLIPTIC) { if (force[i] <= 0) { state[i] = mjCNSTRSTATE_SATISFIED; } else { state[i] = mjCNSTRSTATE_QUADRATIC; nactive++; } } // elliptic else { // get contact dimensionality, friction, mu mjContact* con = d->contact + d->efc_id[i]; int dim = con->dim, result = 0; mjtNum mu = con->mu, f[6]; // f = map force to regular-cone space f[0] = force[i]/mu; for (int j=1; j < dim; j++) { f[j] = force[i+j]/con->friction[j-1]; } // N = normal, T = norm of tangent vector mjtNum N = f[0]; mjtNum T = mju_norm(f+1, dim-1); // top zone if (mu*N >= T) { result = mjCNSTRSTATE_SATISFIED; } // bottom zone else if (N+mu*T <= 0) { result = mjCNSTRSTATE_QUADRATIC; nactive += dim; } // middle zone else { result = mjCNSTRSTATE_CONE; nactive += dim; } // replicate state in all cone dimensions mju_fillInt(state+i, result, dim); // advance c += (dim-1); } } return nactive; } // update constraint state, return nactive and nchange static int dualStateChange(const mjData* d, int* state, int* oldstate, int ne, int nf, int nefc, const int* efclist, int* nchange) { // save old state for (int c=0; c < nefc; c++) { int i = efclist ? efclist[c] : c; oldstate[c] = state[i]; } // update state int nactive = dualState(d, state, ne, nf, nefc, efclist); // count state changes *nchange = 0; for (int c=0; c < nefc; c++) { int i = efclist ? efclist[c] : c; *nchange += (oldstate[c] != state[i]); } return nactive; } // solve QCQP and project onto friction ellipsoid, write to force[i+1..i+dim-1] static void solveQCQP(mjtNum* force, int i, int dim, mjtNum* Ac, mjtNum* bc, const mjtNum* mu) { int flg_active; mjtNum v[6]; // solve if (dim == 3) { flg_active = mju_QCQP2(v, Ac, bc, mu, force[i]); } else if (dim == 4) { flg_active = mju_QCQP3(v, Ac, bc, mu, force[i]); } else { // dim == 5 flg_active = mju_QCQP(v, Ac, bc, mu, force[i], dim-1); } // on constraint: put v on ellipsoid, in case QCQP is approximate if (flg_active) { mjtNum s = 0; for (int j=0; j < dim-1; j++) { s += v[j]*v[j] / (mu[j]*mu[j]); } s = mju_sqrt(force[i]*force[i] / mju_max(mjMINVAL, s)); for (int j=0; j < dim-1; j++) { v[j] *= s; } } // assign mju_copy(force+i+1, v, dim-1); } //---------------------------- PGS solver ---------------------------------------------------------- // core PGS solver: iterates over constraints specified by efclist // 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) static void solPGS(const mjModel* m, mjData* d, int island, int ne, int nf, int nefc, const int* efclist, int maxiter) { const mjtNum *floss = d->efc_frictionloss; mjtNum *force = d->efc_force; mj_markStack(d); mjtNum* ARinv = mjSTACKALLOC(d, nefc, mjtNum); int* oldstate = mjSTACKALLOC(d, 2*nefc, int); int* blockstart = oldstate + nefc; int island_stat = mjMAX(0, island); // island index for diagnostic stats mjtNum scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv)); // precompute inverse diagonal of AR ARdiaginv(m, d, ARinv, nefc, efclist, 0); // initial constraint state dualState(d, d->efc_state, ne, nf, nefc, efclist); // build block-index array: one entry per constraint block int nblocks = 0; for (int c=0; c < nefc; ) { blockstart[nblocks++] = c; int i = efclist ? efclist[c] : c; if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { c += d->contact[d->efc_id[i]].dim; } else { c++; } } // seed PCG32 RNG with a fixed seed pcg32_state rng; rng.state = 0; rng.inc = 1; pcg32_next(&rng); // main iteration int iter = 0; while (iter < maxiter) { // clear improvement mjtNum improvement = 0; // shuffle constraint visitation order shuffle_int(blockstart, nblocks, &rng); // perform one sweep over constraint blocks for (int bi=0; bi < nblocks; bi++) { int c = blockstart[bi]; int i = efclist ? efclist[c] : c; // get constraint dimensionality int dim; if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { dim = d->contact[d->efc_id[i]].dim; } else { dim = 1; } // compute residual for this constraint mjtNum res[6]; residual(m, d, res, i, dim, 0); // save old force mjtNum oldforce[6]; mju_copy(oldforce, force+i, dim); // allocate AR submatrix, required later for costChage mjtNum Athis[36]; // simple constraint if (d->efc_type[i] != mjCNSTR_CONTACT_ELLIPTIC) { // unconstrained minimum force[i] -= res[0]*ARinv[c]; // impose interval and inequality constraints if (c >= ne && c < ne+nf) { if (force[i] < -floss[i]) { force[i] = -floss[i]; } else if (force[i] > floss[i]) { force[i] = floss[i]; } } else if (c >= ne+nf) { if (force[i] < 0) { force[i] = 0; } } } // elliptic cone constraint else { // get friction mjtNum *mu = d->contact[d->efc_id[i]].friction; //-------------------- perform normal or ray update // Athis = AR(this,this) extractBlock(m, d, Athis, i, dim, 0); // normal force too small: normal update if (force[i] < mjMINVAL) { // unconstrained minimum force[i] -= res[0]*ARinv[c]; // clamp if (force[i] < 0) { force[i] = 0; } // clear friction (just in case) mju_zero(force+i+1, dim-1); } // ray update else { // v = ray mjtNum v[6]; mju_copy(v, force+i, dim); // denom = v' * AR(this,this) * v mjtNum v1[6]; mju_mulMatVec(v1, Athis, v, dim, dim); mjtNum denom = mju_dot(v, v1, dim); // avoid division by 0 if (denom >= mjMINVAL) { // x = v' * res / denom mjtNum x = -mju_dot(v, res, dim) / denom; // make sure normal is non-negative if (force[i]+x*v[0] < 0) { x = -v[0]/force[i]; } // add x*v to f for (int j=0; j < dim; j++) { force[i+j] += x*v[j]; } } } //-------------------- perform friction update, keep normal fixed // Ac = AR-submatrix; bc = b-subvector + Ac,rest * f_rest mjtNum bc[5], Ac[25]; mju_copy(bc, res+1, dim-1); for (int j=0; j < dim-1; j++) { mju_copy(Ac+j*(dim-1), Athis+(j+1)*dim+1, dim-1); bc[j] -= mju_dot(Ac+j*(dim-1), oldforce+1, dim-1); bc[j] += Athis[(j+1)*dim]*(force[i]-oldforce[0]); } // guard for f_normal==0 if (force[i] < mjMINVAL) { mju_zero(force+i+1, dim-1); } // QCQP else { solveQCQP(force, i, dim, Ac, bc, mu); } } // accumulate improvement if (dim == 1) { Athis[0] = 1/ARinv[c]; } improvement -= costChange(Athis, force+i, oldforce, res, dim); } // update constraint state int nchange; int nactive = dualStateChange(d, d->efc_state, oldstate, ne, nf, nefc, efclist, &nchange); // scale improvement, save stats improvement *= scale; saveStats(m, d, island_stat, iter, improvement, 0, 0, nactive, nchange, 0, 0); // increment iteration count iter++; // terminate if (improvement < m->opt.tolerance) { break; } } // finalize statistics if (island_stat < mjNISLAND) { // update solver iterations d->solver_niter[island_stat] += iter; // set nnz if (mj_isSparse(m)) { d->solver_nnz[island_stat] = 0; for (int c=0; c < nefc; c++) { d->solver_nnz[island_stat] += d->efc_AR_rownnz[efclist ? efclist[c] : c]; } } else { d->solver_nnz[island_stat] = nefc*nefc; } } mj_freeStack(d); } // 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); } // 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); } //---------------------------- NoSlip solver ------------------------------------------------------- // core NoSlip solver: iterates over constraints specified by efclist // 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) static void solNoSlip(const mjModel* m, mjData* d, int island, int ne, int nf, int nefc, const int* efclist, int maxiter) { int dim, iter = 0; const mjtNum *floss = d->efc_frictionloss; mjtNum *force = d->efc_force; mjtNum *mu, improvement; mjtNum Ac[25], bc[5], res[5], oldforce[5], delta[5], mid, y, K0, K1; mjContact* con; mj_markStack(d); mjtNum* ARinv = mjSTACKALLOC(d, nefc, mjtNum); int* oldstate = mjSTACKALLOC(d, nefc, int); int island_stat = mjMAX(0, island); mjtNum scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv)); // precompute inverse diagonal of A ARdiaginv(m, d, ARinv, nefc, efclist, 1); // initial constraint state dualState(d, d->efc_state, ne, nf, nefc, efclist); // main iteration while (iter < maxiter) { // clear improvement improvement = 0; // correct for cost change at iter 0 if (iter == 0) { for (int c=0; c < nefc; c++) { int i = efclist ? efclist[c] : c; improvement += 0.5*force[i]*force[i]*d->efc_R[i]; } } // perform one sweep: dry friction for (int c=ne; c < ne+nf; c++) { int i = efclist ? efclist[c] : c; // compute residual, save old residual(m, d, res, i, 1, 1); oldforce[0] = force[i]; // unconstrained minimum force[i] -= res[0]*ARinv[c]; // impose interval constraints if (force[i] < -floss[i]) { force[i] = -floss[i]; } else if (force[i] > floss[i]) { force[i] = floss[i]; } // add to improvement delta[0] = force[i] - oldforce[0]; improvement -= 0.5*delta[0]*delta[0]/ARinv[c] + delta[0]*res[0]; } // perform one sweep: contact friction for (int c=ne+nf; c < nefc; c++) { int i = efclist ? efclist[c] : c; // pyramidal contact if (d->efc_type[i] == mjCNSTR_CONTACT_PYRAMIDAL) { // get contact info con = d->contact + d->efc_id[i]; dim = con->dim; mu = con->friction; // loop over pairs of opposing pyramid edges for (int j=i; j < i+2*(dim-1); j+=2) { // compute residual, save old residual(m, d, res, j, 2, 1); mju_copy(oldforce, force+j, 2); // Ac = AR-submatirx extractBlock(m, d, Ac, j, 2, 1); // bc = b-subvector + Ac,rest * f_rest mju_copy(bc, res, 2); for (int k=0; k < 2; k++) { bc[k] -= mju_dot(Ac+k*2, oldforce, 2); } // f0 = mid+y, f1 = mid-y mid = 0.5*(force[j]+force[j+1]); y = 0.5*(force[j]-force[j+1]); // K1 = A00 + A11 - 2*A01, K0 = mid*A00 - mid*A11 + b0 - b1 K1 = Ac[0] + Ac[3] - Ac[1] - Ac[2]; K0 = mid*(Ac[0] - Ac[3]) + bc[0] - bc[1]; // guard against Ac==0 if (K1 < mjMINVAL) { force[j] = force[j+1] = mid; } // otherwise minimize over y \in [-mid, mid] else { // unconstrained minimum y = -K0/K1; // clamp and assign if (y < -mid) { force[j] = 0; force[j+1] = 2*mid; } else if (y > mid) { force[j] = 2*mid; force[j+1] = 0; } else { force[j] = mid+y; force[j+1] = mid-y; } } // accumulate improvement improvement -= costChange(Ac, force+j, oldforce, res, 2); } // skip the rest of this contact c += 2*(dim-1)-1; } // elliptic contact else if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // get contact info con = d->contact + d->efc_id[i]; dim = con->dim; mu = con->friction; // compute residual, save old residual(m, d, res, i+1, dim-1, 1); mju_copy(oldforce, force+i+1, dim-1); // Ac = AR-submatrix extractBlock(m, d, Ac, i+1, dim-1, 1); // bc = b-subvector + Ac,rest * f_rest mju_copy(bc, res, dim-1); for (int j=0; j < dim-1; j++) { bc[j] -= mju_dot(Ac+j*(dim-1), oldforce, dim-1); } // guard for f_normal==0 if (force[i] < mjMINVAL) { mju_zero(force+i+1, dim-1); } // QCQP else { solveQCQP(force, i, dim, Ac, bc, mu); } // accumulate improvement improvement -= costChange(Ac, force+i+1, oldforce, res, dim-1); // skip the rest of this contact c += (dim-1); } } // update constraint state int nchange; int nactive = dualStateChange(d, d->efc_state, oldstate, ne, nf, nefc, efclist, &nchange); // scale improvement, save stats improvement *= scale; // save noslip stats after all the entries from regular solver 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++; // terminate if (improvement < m->opt.noslip_tolerance) { break; } } // update solver iterations if (island_stat < mjNISLAND) { d->solver_niter[island_stat] += iter; } mj_freeStack(d); } // 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); } // 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); } //------------------------- Primal solvers --------------------------------------------------------- // Primal context typedef struct { int is_sparse; // 1: sparse, 0: dense int is_elliptic; // 1: elliptic, 0: pyramidal int island; // current island index, -1 if monolithic // sizes int nv; // number of dofs int ne; // number of equalities int nf; // number of friction constraints int nefc; // number of all constraints int nJ; // number of nonzeros in Jacobian // contact array mjContact* contact; // dof arrays const mjtNum* qfrc_smooth; const mjtNum* qacc_smooth; mjtNum* qfrc_constraint; mjtNum* qacc; // inertia int* M_rownnz; int* M_rowadr; int* M_colind; mjtNum* M; mjtNum* qLD; 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 int* J_rownnz; int* J_rowadr; int* J_rowsuper; int* J_colind; mjtNum* J; int* JT_rownnz; int* JT_rowadr; int* JT_rowsuper; int* JT_colind; mjtNum* JT; // common arrays (PrimalAllocate) mjtNum* Jaref; // Jac*qacc - aref (nefc x 1) mjtNum* Jv; // Jac*search (nefc x 1) mjtNum* Ma; // M*qacc (nv x 1) mjtNum* Mv; // M*search (nv x 1) mjtNum* grad; // gradient of master cost (nv x 1) mjtNum* Mgrad; // M\grad or H\grad (nv x 1) mjtNum* search; // linesearch vector (nv x 1) mjtNum* quad; // quadratic polynomials for constraint costs (nefc x 3) int* oldstate; // previous constraint state (nefc x 1) // CG arrays (PrimalAllocate, CG only) mjtNum* gradold; // previous gradient (nv x 1) mjtNum* Mgradold; // previous preconditioned gradient (nv x 1) mjtNum* graddif; // grad - gradold (nv x 1) mjtNum* Mgraddif; // M\(grad - gradold) (nv x 1) // Newton arrays, known-size (PrimalAllocate) mjtNum* D; // constraint inertia (nefc x 1) mjtNum* cholupd; // scratch for rank-1 Cholesky updates (nv x 1) mjtNum* LTJ; // L'*J for cone Cholesky updates (6 x nv) int* H_rowadr; // Hessian row addresses (nv x 1) int* H_rownnz; // Hessian row nonzeros (nv x 1) int* HT_rownnz; // Hessian transpose row nonzeros (nv x 1) int* HT_rowadr; // Hessian transpose row addresses (nv x 1) int* L_rownnz; // Hessian factor row nonzeros (nv x 1) int* L_rowadr; // Hessian factor row addresses (nv x 1) int* LT_rownnz; // Hessian factor transpose row nonzeros (nv x 1) int* LT_rowadr; // Hessian factor transpose row addresses (nv x 1) // Newton arrays, computed-size (MakeHessian) int nH; // number of nonzeros in Hessian H int* H_colind; // Hessian column indices (nH x 1) int* HT_colind; // Hessian transpose column indices (nH x 1) mjtNum* H; // Hessian (nH x 1) int nL; // number of nonzeros in Cholesky factor L int* L_colind; // Cholesky factor column indices (nL x 1) int* LT_colind; // Cholesky factor transpose column indices (nL x 1) int* LT_map; // CSC-to-CSR index mapping (nL x 1) mjtNum* L; // Cholesky factor (nL x 1) mjtNum* Lcone; // Cholesky factor with cone contributions (nL x 1) // 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 // linesearch diagnostics int LSiter; // number of linesearch iterations int LSresult; // linesearch result mjtNum LSslope; // linesearch slope at solution } mjPrimalContext; // set sizes and pointers to mjData arrays in mjPrimalContext static void PrimalPointers(const mjModel* m, const mjData* d, mjPrimalContext* ctx, int island) { // clear everything memset(ctx, 0, sizeof(mjPrimalContext)); // globals ctx->is_sparse = mj_isSparse(m); ctx->is_elliptic = (m->opt.cone == mjCONE_ELLIPTIC); 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; ctx->nJ = d->nJ; // 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 = m->M_rownnz; ctx->M_rowadr = m->M_rowadr; ctx->M_colind = m->M_colind; ctx->M = d->M; 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 (ctx->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; } } // 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; // 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; } } // allocate fixed-size arrays in mjPrimalContext // mj_{mark/free}Stack in calling function! static void PrimalAllocate(const mjModel* m, mjData* d, mjPrimalContext* ctx, int flg_Newton) { // local sizes and flags int nv = ctx->nv; int nefc = ctx->nefc; int is_sparse = ctx->is_sparse; int is_elliptic = ctx->is_elliptic; int nJ = is_sparse ? d->nJ : 0; // compute island matrix sizes if needed int nC = 0; if (ctx->island >= 0) { // count nC: number of nonzeros in M block of island (always sparse) int island = ctx->island; int idofadr = d->island_idofadr[island]; for (int i = 0; i < nv; i++) { int dof = d->map_idof2dof[idofadr + i]; nC += m->M_rownnz[dof]; } // count nJ: number of nonzeros in J block of island (sparse or dense) if (is_sparse) { nJ = 0; int iefcadr = d->island_iefcadr[island]; for (int i = 0; i < nefc; i++) { int efc = d->map_iefc2efc[iefcadr + i]; nJ += d->efc_J_rownnz[efc]; } } else { nJ = nefc * nv; } ctx->nJ = nJ; } // compute mjtNum block size size_t nNum = 5*nefc + 5*nv; // common arrays if (is_sparse) nNum += nJ; // JT if (flg_Newton) { nNum += nefc + nv; // D, cholupd if (is_elliptic) nNum += 6*nv; // LTJ if (!is_sparse) { nNum += nv*nv; // L (dense) if (is_elliptic) nNum += nv*nv; // Lcone (dense) } } else { nNum += 4*nv; // CG arrays } // add island matrix sizes if (ctx->island >= 0) { nNum += 2 * nC + nv + nJ; // iM, iLD, iLDiagInv, iefc_J } // compute int block size size_t nInt = nefc; // oldstate if (is_sparse) { nInt += 3*nv + nJ; // JT sparse if (flg_Newton) nInt += 8*nv; // Newton sparse } // add island matrix sizes if (ctx->island >= 0) { nInt += 2 * nv + nC; // iM_{rownnz, rowadr, colind} if (is_sparse) { nInt += 3 * nefc + nJ; // iefc_J_{rownnz, rowadr, rowsuper, colind} } } // allocate mjtNum and int blocks mjtNum* numblock = mjSTACKALLOC(d, nNum, mjtNum); int* intblock = mjSTACKALLOC(d, nInt, int); // populate island matrices if needed if (ctx->island >= 0) { int island = ctx->island; int idofadr = d->island_idofadr[island]; int iefcadr = d->island_iefcadr[island]; ctx->M_rownnz = intblock; intblock += nv; ctx->M_rowadr = intblock; intblock += nv; ctx->M_colind = intblock; intblock += nC; ctx->M = numblock; numblock += nC; ctx->qLD = numblock; numblock += nC; ctx->qLDiagInv = numblock; numblock += nv; mju_blockSparse(ctx->qLD, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, d->qLD, m->M_rownnz, m->M_rowadr, m->M_colind, nv, d->map_idof2dof + idofadr, d->map_dof2idof, d->island_idofadr[island], 0, ctx->M, d->M); mju_gather(ctx->qLDiagInv, d->qLDiagInv, d->map_idof2dof + idofadr, nv); ctx->J = numblock; numblock += nJ; if (!is_sparse) { mju_block(ctx->J, d->efc_J, m->nv, nv, nefc, d->map_iefc2efc + iefcadr, d->map_idof2dof + idofadr); } else { ctx->J_rownnz = intblock; intblock += nefc; ctx->J_rowadr = intblock; intblock += nefc; ctx->J_rowsuper = intblock; intblock += nefc; ctx->J_colind = intblock; intblock += nJ; mju_blockSparse(ctx->J, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, nefc, d->map_iefc2efc + iefcadr, d->map_dof2idof, d->island_idofadr[island], 0, NULL, NULL); mju_superSparse(nefc, ctx->J_rowsuper, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind); } } // carve mjtNum block ctx->Jaref = numblock; numblock += nefc; ctx->Jv = numblock; numblock += nefc; ctx->Ma = numblock; numblock += nv; ctx->Mv = numblock; numblock += nv; ctx->grad = numblock; numblock += nv; ctx->Mgrad = numblock; numblock += nv; ctx->search = numblock; numblock += nv; ctx->quad = numblock; numblock += 3*nefc; if (is_sparse) { ctx->JT = numblock; numblock += nJ; } if (flg_Newton) { ctx->D = numblock; numblock += nefc; ctx->cholupd = numblock; numblock += nv; if (is_elliptic) { ctx->LTJ = numblock; numblock += 6*nv; } if (!is_sparse) { ctx->nL = nv*nv; ctx->L = numblock; numblock += ctx->nL; ctx->Lcone = is_elliptic ? numblock : NULL; if (is_elliptic) numblock += ctx->nL; } } else { ctx->gradold = numblock; numblock += nv; ctx->Mgradold = numblock; numblock += nv; ctx->graddif = numblock; numblock += nv; ctx->Mgraddif = numblock; numblock += nv; } // carve int block ctx->oldstate = intblock; intblock += nefc; if (is_sparse) { ctx->JT_rownnz = intblock; intblock += nv; ctx->JT_rowadr = intblock; intblock += nv; ctx->JT_rowsuper = intblock; intblock += nv; ctx->JT_colind = intblock; intblock += nJ; } if (flg_Newton && is_sparse) { ctx->H_rowadr = intblock; intblock += nv; ctx->H_rownnz = intblock; intblock += nv; ctx->HT_rownnz = intblock; intblock += nv; ctx->HT_rowadr = intblock; intblock += nv; ctx->L_rownnz = intblock; intblock += nv; ctx->L_rowadr = intblock; intblock += nv; ctx->LT_rownnz = intblock; intblock += nv; ctx->LT_rowadr = intblock; intblock += nv; } // sparse: compute Jacobian transpose if (is_sparse) { int offset = ctx->J_rowadr[0]; mju_transposeSparse(ctx->JT, ctx->J + offset, nefc, nv, ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind + offset); } } // update efc_force, qfrc_constraint, cost-related static void PrimalUpdateConstraint(mjPrimalContext* ctx, int flg_HessianCone) { int nefc = ctx->nefc, nv = ctx->nv; // update constraints 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), flg_HessianCone); // compute qfrc_constraint (dense or sparse) if (!ctx->is_sparse) { 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 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 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; ctx->cost += Gauss; } // update grad, Mgrad static void PrimalUpdateGradient(mjPrimalContext* ctx, int flg_Newton) { int nv = ctx->nv; // grad = M*qacc - qfrc_smooth - qfrc_constraint 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 if (flg_Newton) { if (ctx->is_sparse) { mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? ctx->Lcone : ctx->L), ctx->grad, nv, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind); } else { mju_cholSolve(ctx->Mgrad, (ctx->ncone ? ctx->Lcone : ctx->L), ctx->grad, nv); } } // CG: Mgrad = M \ grad else { mju_copy(ctx->Mgrad, ctx->grad, nv); mj_solveLD(ctx->Mgrad, ctx->qLD, ctx->qLDiagInv, nv, 1, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, NULL); } } // prepare quadratic polynomials and contact cone quantities static void PrimalPrepare(mjPrimalContext* 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 PrimalUpdateConstraint 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 i=0; i < nefc; i++) { // pointers to numeric data 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*i; // init with scalar quadratic mjtNum DJ0 = D[0]*Jaref[0]; quad[0] = Jaref[0]*DJ0; quad[1] = Jv[0]*DJ0; quad[2] = Jv[0]*D[0]*Jv[0]; // elliptic cone: extra processing if (ctx->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // extract contact info 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; const mjtNum* friction = con->friction; // complete vector quadratic (for bottom zone) for (int j=1; j < dim; j++) { mjtNum DJj = D[j]*Jaref[j]; quad[0] += Jaref[j]*DJj; quad[1] += Jv[j]*DJj; quad[2] += Jv[j]*D[j]*Jv[j]; } // rescale to make primal cone circular U[0] = Jaref[0]*mu; V[0] = Jv[0]*mu; for (int j=1; j < dim; j++) { U[j] = Jaref[j]*friction[j-1]; V[j] = Jv[j]*friction[j-1]; } // accumulate sums of squares for (int j=1; j < dim; j++) { UU += U[j]*U[j]; UV += U[j]*V[j]; VV += V[j]*V[j]; } // store in quad[3-8], using the fact that dim>=3 quad[3] = U[0]; quad[4] = V[0]; quad[5] = UU; quad[6] = UV; quad[7] = VV; quad[8] = D[0] / ((mu*mu) * (1 + (mu*mu))); // advance to next constraint i += (dim-1); } // apply scaling quad[0] *= 0.5; quad[2] *= 0.5; } } // linesearch evaluation point struct _mjPrimalPnt { mjtNum alpha; mjtNum cost; mjtNum deriv[2]; }; typedef struct _mjPrimalPnt mjPrimalPnt; // Huber cost of a single friction constraint at a given point x static mjtNum frictionCost(mjtNum x, mjtNum f, mjtNum Rf, mjtNum D) { // -bound < x < bound : quadratic if (-Rf < x && x < Rf) return 0.5*D*x*x; // x < -bound : linear negative else if (x <= -Rf) return f*(-0.5*Rf - x); // bound < x : linear positive else return f*(-0.5*Rf + x); } // Huber cost difference of a single friction constraint: cost(x) - cost(start) static mjtNum frictionCostDif(mjtNum start, mjtNum x, mjtNum f, mjtNum Rf, mjtNum D) { int state_start = (-Rf < start && start < Rf) ? 0 : (start <= -Rf ? -1 : 1); int state_x = (-Rf < x && x < Rf) ? 0 : (x <= -Rf ? -1 : 1); // both quadratic if (state_start == 0 && state_x == 0) { return 0.5*D*(x - start)*(x + start); } // both linear negative if (state_start == -1 && state_x == -1) { return f*(start - x); } // both linear positive if (state_start == 1 && state_x == 1) { return f*(x - start); } // otherwise different zones: compute absolute costs and subtract return frictionCost(x, f, Rf, D) - frictionCost(start, f, Rf, D); } // compute cost of an elliptic cone at a given alpha static mjtNum ellipticCost(const mjtNum* quad, mjtNum alpha, mjtNum mu, mjtNum Dm) { mjtNum U0 = quad[3], V0 = quad[4], UU = quad[5]; mjtNum UV = quad[6], VV = quad[7]; mjtNum N = U0 + alpha*V0; mjtNum Tsqr = UU + alpha*(2*UV + alpha*VV); // no tangential force : top or bottom zone if (Tsqr <= 0) { // bottom zone: quadratic cost if (N < 0) return alpha*alpha*quad[2] + alpha*quad[1] + quad[0]; // top zone: nothing to do } // otherwise regular processing else { mjtNum T = mju_sqrt(Tsqr); // N>=mu*T : top zone if (N >= mu*T) { // nothing to do } // mu*N+T<=0 : bottom zone else if (mu*N+T <= 0) { return alpha*alpha*quad[2] + alpha*quad[1] + quad[0]; } // otherwise middle zone else { return 0.5*Dm*(N-mu*T)*(N-mu*T); } } return 0; } // compute cost difference of an elliptic cone at a given alpha: cost(alpha) - cost(0) static mjtNum ellipticCostDif(const mjtNum* quad, mjtNum alpha, mjtNum mu, mjtNum Dm) { mjtNum U0 = quad[3], V0 = quad[4], UU = quad[5]; mjtNum UV = quad[6], VV = quad[7]; // determine zone and cost at alpha=0 int zone0 = 0; mjtNum T0 = 0; if (UU <= 0) { zone0 = (U0 < 0) ? 2 : 1; } else { T0 = mju_sqrt(UU); if (U0 >= mu*T0) { zone0 = 1; // top zone } else if (mu*U0 + T0 <= 0) { zone0 = 2; // bottom zone } else { zone0 = 3; // middle zone } } // determine zone and cost at alpha mjtNum N = U0 + alpha*V0; mjtNum Tsqr = UU + alpha*(2*UV + alpha*VV); int zone_alpha = 0; mjtNum T = 0; if (Tsqr <= 0) { zone_alpha = (N < 0) ? 2 : 1; // bottom or top zone } else { T = mju_sqrt(Tsqr); if (N >= mu*T) { zone_alpha = 1; // top zone } else if (mu*N + T <= 0) { zone_alpha = 2; // bottom zone } else { zone_alpha = 3; // middle zone } } // both top zone if (zone0 == 1 && zone_alpha == 1) { return 0; } // both bottom zone if (zone0 == 2 && zone_alpha == 2) { return alpha*alpha*quad[2] + alpha*quad[1]; } // both middle zone if (zone0 == 3 && zone_alpha == 3) { mjtNum diff_alpha = N - mu*T; mjtNum diff0 = U0 - mu*T0; return 0.5*Dm*(diff_alpha - diff0)*(diff_alpha + diff0); } // otherwise different zones: compute absolute costs and subtract return ellipticCost(quad, alpha, mu, Dm) - ellipticCost(quad, 0, mu, Dm); } // evaluate shifted linesearch cost: cost(alpha) - cost(0), and derivatives static void PrimalEval(mjPrimalContext* ctx, mjPrimalPnt* p) { int ne = ctx->ne, nf = ctx->nf, nefc = ctx->nefc; // clear result mjtNum cost = 0, alpha = p->alpha; mjtNum deriv[2] = {0, 0}; // init quad with Gauss, shifted: drop quadGauss[0] mjtNum quadTotal[3] = {0, ctx->quadGauss[1], ctx->quadGauss[2]}; // process constraints for (int i=0; i < nefc; i++) { // equality: shifted quad (skip quad[0]) if (i < ne) { quadTotal[1] += ctx->quad[3*i+1]; quadTotal[2] += ctx->quad[3*i+2]; continue; } // friction: compute cost(alpha) - cost(0) directly if (i < ne + nf) { // search point, friction loss, bound (Rf) mjtNum start = ctx->Jaref[i]; mjtNum dir = ctx->Jv[i]; mjtNum x = start + alpha*dir; mjtNum f = ctx->efc_frictionloss[i]; mjtNum D = ctx->efc_D[i]; mjtNum Rf = ctx->efc_R[i]*f; // cost delta cost += frictionCostDif(start, x, f, Rf, D); // -bound < x < bound : quadratic if (-Rf < x && x < Rf) { deriv[0] += D*x*dir; deriv[1] += D*dir*dir; } // x < -bound : linear negative else if (x <= -Rf) { deriv[0] += -f*dir; } // bound < x : linear positive else { deriv[0] += f*dir; } continue; } // limit and contact if (ctx->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // elliptic cone // extract contact info const mjContact* con = ctx->contact + ctx->efc_id[i]; mjtNum* quad = ctx->quad + 3*i; int dim = con->dim; mjtNum mu = con->mu; // unpack quad mjtNum U0 = quad[3]; mjtNum V0 = quad[4]; mjtNum UU = quad[5]; mjtNum UV = quad[6]; mjtNum VV = quad[7]; mjtNum Dm = quad[8]; // shifted cost cost += ellipticCostDif(quad, alpha, mu, Dm); // compute N, Tsqr for derivatives mjtNum N = U0 + alpha*V0; mjtNum Tsqr = UU + alpha*(2*UV + alpha*VV); // no tangential force : top or bottom zone if (Tsqr <= 0) { // bottom zone: quadratic derivatives if (N < 0) { deriv[0] += 2*alpha*quad[2] + quad[1]; deriv[1] += 2*quad[2]; } // top zone: nothing to do } // otherwise regular processing else { // tangential force mjtNum T = mju_sqrt(Tsqr); // N>=mu*T : top zone if (N >= mu*T) { // nothing to do } // mu*N+T<=0 : bottom zone else if (mu*N+T <= 0) { deriv[0] += 2*alpha*quad[2] + quad[1]; deriv[1] += 2*quad[2]; } // otherwise middle zone else { // derivatives mjtNum N1 = V0; mjtNum T1 = (UV + alpha*VV)/T; mjtNum T2 = VV/T - (UV + alpha*VV)*T1/(T*T); deriv[0] += Dm*(N-mu*T)*(N1-mu*T1); deriv[1] += Dm*((N1-mu*T1)*(N1-mu*T1) + (N-mu*T)*(-mu*T2)); } } // advance to next constraint i += (dim-1); } else { // inequality // search point mjtNum start = ctx->Jaref[i]; mjtNum x = start + alpha*ctx->Jv[i]; mjtNum cost0 = start < 0 ? ctx->quad[3*i] : 0; // active if (x < 0) { // shifted quad: add quad[1], quad[2] and (quad[0] - cost0) quadTotal[0] += ctx->quad[3*i] - cost0; quadTotal[1] += ctx->quad[3*i+1]; quadTotal[2] += ctx->quad[3*i+2]; } else { cost -= cost0; } } } // add total quadratic (quadTotal[0] contains only shifted residuals) cost += alpha*alpha*quadTotal[2] + alpha*quadTotal[1] + quadTotal[0]; deriv[0] += 2*alpha*quadTotal[2] + quadTotal[1]; deriv[1] += 2*quadTotal[2]; // check for convexity; SHOULD NOT OCCUR if (deriv[1] <= 0) { mju_warning("Linesearch objective is not convex"); deriv[1] = mjMINVAL; } // assign and count p->cost = cost; p->deriv[0] = deriv[0]; p->deriv[1] = deriv[1]; ctx->LSiter++; } // update bracket point given 3 candidate points static int updateBracket(mjPrimalContext* ctx, mjPrimalPnt* p, const mjPrimalPnt candidates[3], mjPrimalPnt* pnext) { int flag = 0; for (int i=0; i < 3; i++) { // negative deriv if (p->deriv[0] < 0 && candidates[i].deriv[0] < 0 && p->deriv[0] < candidates[i].deriv[0]) { *p = candidates[i]; flag = 1; } // positive deriv else if (p->deriv[0] > 0 && candidates[i].deriv[0] > 0 && p->deriv[0] > candidates[i].deriv[0]) { *p = candidates[i]; flag = 2; } } // compute next point if updated if (flag) { pnext->alpha = p->alpha - p->deriv[0]/p->deriv[1]; PrimalEval(ctx, pnext); } return flag; } // line search static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_iterations, mjtNum* improvement) { int nv = ctx->nv, nefc = ctx->nefc; mjPrimalPnt p0, p1, p2, pmid, p1next, p2next; // clear results ctx->LSiter = 0; ctx->LSresult = 0; ctx->LSslope = 1; // means not computed *improvement = 0; // save search vector length, check 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 = tolerance * snorm / ctx->scale; mjtNum slopescl = ctx->scale / snorm; // compute Mv = M * v mju_mulSymVecSparse(ctx->Mv, ctx->M, ctx->search, nv, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind); // compute Jv = J * search (dense or sparse) if (!ctx->is_sparse) { 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 PrimalPrepare(ctx); // init at alpha = 0, save p0.alpha = 0; PrimalEval(ctx, &p0); // always attempt one Newton step p1.alpha = p0.alpha - p0.deriv[0]/p0.deriv[1]; PrimalEval(ctx, &p1); // check for initial convergence if (mju_abs(p1.deriv[0]) < gtol) { if (p1.alpha == 0) { ctx->LSresult = 2; // no improvement, initial convergence } else { ctx->LSresult = 0; // SUCCESS } ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; *improvement = -p1.cost; return p1.alpha; } // save direction int dir = (p1.deriv[0] < 0 ? +1 : -1); // SANITY CHECKS /* // descent direction if( mju_dot(ctx->grad, ctx->search, m->nv)>=0 ) printf("NOT A DESCENT: grad %g search %g dot %g\n", mju_norm(ctx->grad, m->nv), mju_norm(ctx->search, m->nv), mju_dot(ctx->grad, ctx->search, m->nv)); // 2nd derivative for Newton cone if( ctx->flg_Newton && ctx->ncone ) { mjtNum dd = -p0.deriv[0]/p0.deriv[1]; if( mju_abs(dd-1)>1e-6 ) printf("2nd DERIVATIVE FAIL: d0 %g d1 %g alpha %g\n", p0.deriv[0], p0.deriv[1], dd); } // cost and gradient at 0: full-space vs. linesearch mjtNum grd = mju_dot(ctx->grad, ctx->search, m->nv); if( mju_abs(p0.cost-ctx->cost)/mjMAX(mjMINVAL,mju_abs(p0.cost+ctx->cost)) > 1e-6 || mju_abs(p0.deriv[0]-grd)/mjMAX(mjMINVAL,mju_abs(p0.deriv[0]+grd)) > 1e-6 ) { printf("LSiter = %d:\n", ctx->LSiter); printf("COST: %g %g %g\n", p0.cost, ctx->cost, mju_abs(p0.cost-ctx->cost)/mjMAX(mjMINVAL,mju_abs(p0.cost+ctx->cost))); printf("GRAD: %g %g %g\n", p0.deriv[0], grd, mju_abs(p0.deriv[0]-grd)/mjMAX(mjMINVAL,mju_abs(p0.deriv[0]+grd))); } */ // one-sided search p2 = p0; int p2update = 1; 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]; PrimalEval(ctx, &p1); // check for convergence if (mju_abs(p1.deriv[0]) < gtol) { ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; *improvement = -p1.cost; return p1.alpha; // SUCCESS } } // check for failure to bracket if (ctx->LSiter >= ls_iterations) { ctx->LSresult = 3; // could not bracket ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; *improvement = -p1.cost; return p1.alpha; } // check for p2 update; SHOULD NOT OCCUR if (!p2update) { ctx->LSresult = 6; // no p2 update ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; *improvement = -p1.cost; return p1.alpha; } // compute next-points for bracket p2next = p1; p1next.alpha = p1.alpha - p1.deriv[0]/p1.deriv[1]; PrimalEval(ctx, &p1next); // bracketed search while (ctx->LSiter < ls_iterations) { // evaluate at midpoint pmid.alpha = 0.5*(p1.alpha + p2.alpha); PrimalEval(ctx, &pmid); // make list of candidates mjPrimalPnt candidates[3] = {p1next, p2next, pmid}; // check candidates for convergence mjtNum bestcost = 0; int bestind = -1; for (int i=0; i < 3; i++) { if (mju_abs(candidates[i].deriv[0]) < gtol && (bestind == -1 || candidates[i].cost < bestcost)) { bestcost = candidates[i].cost; bestind = i; } } if (bestind >= 0) { ctx->LSslope = mju_abs(candidates[bestind].deriv[0])*slopescl; *improvement = -candidates[bestind].cost; return candidates[bestind].alpha; // SUCCESS } // update brackets 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) { if (pmid.cost < 0) { ctx->LSresult = 0; // SUCCESS } else { ctx->LSresult = 7; // no improvement, could not bracket } ctx->LSslope = mju_abs(pmid.deriv[0])*slopescl; *improvement = -pmid.cost; return pmid.alpha; } } // choose bracket with best cost if (p1.cost <= p2.cost && p1.cost < 0) { ctx->LSresult = 4; // improvement but no convergence ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; *improvement = -p1.cost; return p1.alpha; } else if (p2.cost <= p1.cost && p2.cost < 0) { ctx->LSresult = 4; // improvement but no convergence ctx->LSslope = mju_abs(p2.deriv[0])*slopescl; *improvement = -p2.cost; return p2.alpha; } else { ctx->LSresult = 5; // no improvement return 0; } } // allocate and compute Hessian given efc_state // mj_{mark/free}Stack in caller function! static void MakeHessian(mjData* d, mjPrimalContext* ctx) { int nv = ctx->nv, nefc = ctx->nefc; // compute constraint inertia for (int i=0; i < nefc; i++) { ctx->D[i] = ctx->efc_state[i] == mjCNSTRSTATE_QUADRATIC ? ctx->efc_D[i] : 0; } // sparse if (ctx->is_sparse) { // count Hessian nonzeros, initialize rowadr, rownnz ctx->nH = mju_sqrMatTDSparseSymbolic( ctx->H_rownnz, ctx->H_rowadr, NULL, NULL, nefc, nv, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, d); // add M nonzeros to Hessian total (unavoidable overcounting since H_colind is still unknown) ctx->nH += ctx->M_rowadr[nv - 1] + ctx->M_rownnz[nv - 1]; // nH is known: allocate H, H_colind, HT_colind ctx->H = mjSTACKALLOC(d, ctx->nH, mjtNum); int* H_intblock = mjSTACKALLOC(d, 2*ctx->nH, int); ctx->H_colind = H_intblock; ctx->HT_colind = H_intblock + ctx->nH; // shift H row addresses to make room for M int shift = 0; for (int r = 0; r < nv - 1; r++) { shift += ctx->M_rownnz[r]; ctx->H_rowadr[r + 1] += shift; } // compute H = J'*D*J: symbolic phase mju_sqrMatTDSparseSymbolic( ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, NULL, nefc, nv, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, d); // compute H = J'*D*J: numeric phase mju_sqrMatTDSparseNumeric( ctx->H, nv, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, NULL, ctx->J, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, ctx->JT, ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, ctx->D, d); // add mass matrix: H = J'*D*J + M mju_addToMatSparse(ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv, ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind); // compute H' sparse structure (upper triangle, required for symbolic Cholesky) mju_transposeSparse(NULL, NULL, nv, nv, ctx->HT_rownnz, ctx->HT_rowadr, ctx->HT_colind, NULL, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind); // count total and row non-zeros of reverse-Cholesky factors L and LT ctx->nL = mju_cholFactorSymbolic(NULL, ctx->L_rownnz, ctx->L_rowadr, NULL, ctx->LT_rownnz, ctx->LT_rowadr, NULL, ctx->HT_rownnz, ctx->HT_rowadr, ctx->HT_colind, nv, d); // nL is known: allocate blocks and carve L_colind, LT_colind, LT_map, L, Lcone size_t nL_int = 2*ctx->nL + ctx->nL; // L_colind + LT_colind + LT_map size_t nL_num = ctx->is_elliptic ? 2*ctx->nL : ctx->nL; // L + Lcone int* L_intblock = mjSTACKALLOC(d, nL_int, int); mjtNum* L_numblock = mjSTACKALLOC(d, nL_num, mjtNum); ctx->L_colind = L_intblock; ctx->LT_colind = L_intblock + ctx->nL; ctx->LT_map = L_intblock + 2*ctx->nL; ctx->L = L_numblock; ctx->Lcone = ctx->is_elliptic ? L_numblock + ctx->nL : NULL; // symbolic Cholesky: populate L_colind and LT structures mju_cholFactorSymbolic(ctx->L_colind, ctx->L_rownnz, ctx->L_rowadr, ctx->LT_colind, ctx->LT_rownnz, ctx->LT_rowadr, ctx->LT_map, ctx->HT_rownnz, ctx->HT_rowadr, ctx->HT_colind, nv, d); } // dense else { // compute H = M + J'*D*J mju_sqrMatTD_impl(ctx->L, ctx->J, ctx->D, nefc, nv, /*flg_upper=*/ 0); mju_addToSymSparse(ctx->L, ctx->M, ctx->nv, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, /*flg_upper=*/ 0); } } // forward declaration of HessianCone (for readability) static void HessianCone(mjData* d, mjPrimalContext* ctx); // factorize Hessian: L = chol(H), maybe (re)compute H given efc_state static void FactorizeHessian(mjData* d, mjPrimalContext* ctx, int flg_recompute) { int nv = ctx->nv, nefc = ctx->nefc; // maybe compute constraint inertia if (flg_recompute) { for (int i=0; i < nefc; i++) { ctx->D[i] = ctx->efc_state[i] == mjCNSTRSTATE_QUADRATIC ? ctx->efc_D[i] : 0; } } // sparse if (ctx->is_sparse) { // maybe compute H = M + J'*D*J if (flg_recompute) { // compute H = J'*D*J: symbolic phase mju_sqrMatTDSparseSymbolic( ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, NULL, nefc, nv, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, d); // compute H = J'*D*J: numeric phase mju_sqrMatTDSparseNumeric( ctx->H, nv, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, NULL, ctx->J, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, ctx->JT, ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, ctx->D, d); // add mass matrix: H = J'*D*J + C mju_addToMatSparse(ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv, ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind); } // numeric sparse factorization: L = chol(H) using pre-computed sparsity pattern int rank = mju_cholFactorNumeric( ctx->L, nv, mjMINVAL, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, ctx->LT_rownnz, ctx->LT_rowadr, ctx->LT_colind, ctx->LT_map, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, d); // rank-deficient; SHOULD NOT OCCUR if (rank != nv) { mjERROR("rank-deficient sparse Hessian"); } } // dense else { // maybe compute H = M + J'*D*J if (flg_recompute) { mju_sqrMatTD_impl(ctx->L, ctx->J, ctx->D, nefc, nv, /*flg_upper=*/ 0); mju_addToSymSparse(ctx->L, ctx->M, ctx->nv, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, /*flg_upper=*/ 0); } // factorize H mju_cholFactor(ctx->L, nv, mjMINVAL); } // add cones to factor if present if (ctx->ncone) { HessianCone(d, ctx); } // mark full update ctx->nupdate = nefc; } // elliptic case: Hcone = H + cone_contributions static void HessianCone(mjData* d, mjPrimalContext* ctx) { int nv = ctx->nv, nefc = ctx->nefc; mjtNum* LTJ = ctx->LTJ; mjtNum local[36]; // start with Hcone = H mju_copy(ctx->Lcone, ctx->L, ctx->nL); // add contributions for (int i=0; i < nefc; i++) { if (ctx->efc_state[i] == mjCNSTRSTATE_CONE) { mjContact* con = ctx->contact + ctx->efc_id[i]; int dim = con->dim; // Cholesky of local Hessian mju_copy(local, con->H, dim*dim); mju_cholFactor(local, dim, mjMINVAL); // sparse if (ctx->is_sparse) { // get nnz for row i (same for all rows in contact) const int nnz = ctx->J_rownnz[i]; // compute LTJ = L'*J for this contact mju_zero(LTJ, dim*nnz); for (int r=0; r < dim; r++) { for (int c=0; c <= r; c++) { mju_addToScl(LTJ+c*nnz, ctx->J+ctx->J_rowadr[i+r], local[r*dim+c], nnz); } } // update for (int r=0; r < dim; r++) { mju_cholUpdateSparse(ctx->Lcone, LTJ+r*nnz, nv, 1, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, ctx->J_colind+ctx->J_rowadr[i+r], d); } } // dense else { // compute LTJ = L'*J for this contact row mju_zero(LTJ, dim*nv); for (int r=0; r < dim; r++) { for (int c=0; c <= r; c++) { mju_addToScl(LTJ+c*nv, ctx->J+(i+r)*nv, local[r*dim+c], nv); } } // update for (int r=0; r < dim; r++) { mju_cholUpdate(ctx->Lcone, LTJ+r*nv, nv, 1); } } // count updates ctx->nupdate += dim; // advance to next constraint i += (dim-1); } } } // incremental update to Hessian factor due to changes in efc_state static void HessianIncremental(mjData* d, mjPrimalContext* ctx, const int* oldstate) { int rank, nv = ctx->nv, nefc = ctx->nefc; mjtNum* cholupd = ctx->cholupd; // clear update counter ctx->nupdate = 0; // update H factorization for (int i=0; i < nefc; i++) { int flag_update = -1; // add quad if (oldstate[i] != mjCNSTRSTATE_QUADRATIC && ctx->efc_state[i] == mjCNSTRSTATE_QUADRATIC) { flag_update = 1; } // subtract quad else if (oldstate[i] == mjCNSTRSTATE_QUADRATIC && ctx->efc_state[i] != mjCNSTRSTATE_QUADRATIC) { flag_update = 0; } // perform update if flagged if (flag_update != -1) { // update with cholupd = J(i,:)*sqrt(D[i])) if (ctx->is_sparse) { // get nnz and adr of row i const int nnz = ctx->J_rownnz[i], adr = ctx->J_rowadr[i]; // scale cholupd mju_scl(cholupd, ctx->J+adr, mju_sqrt(ctx->efc_D[i]), nnz); // sparse update or downdate rank = mju_cholUpdateSparse(ctx->L, cholupd, nv, flag_update, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, ctx->J_colind+adr, d); } else { mju_scl(cholupd, ctx->J+i*nv, mju_sqrt(ctx->efc_D[i]), nv); rank = mju_cholUpdate(ctx->L, cholupd, nv, flag_update); } ctx->nupdate++; // recompute H directly if accuracy lost if (rank < nv) { FactorizeHessian(d, ctx, /*flg_recompute=*/1); // nothing else to do return; } } } // add cones if present if (ctx->ncone) { HessianCone(d, ctx); } } // driver static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, int flg_Newton) { int iter = 0; mjtNum alpha, beta; mjPrimalContext ctx; mj_markStack(d); // make context PrimalPointers(m, d, &ctx, island); PrimalAllocate(m, d, &ctx, flg_Newton); // local copies int nv = ctx.nv; int nefc = ctx.nefc; int* oldstate = ctx.oldstate; // compute Ma = M * qacc mju_mulSymVecSparse(ctx.Ma, ctx.M, ctx.qacc, nv, ctx.M_rownnz, ctx.M_rowadr, ctx.M_colind); // compute Jaref = J * qacc - aref (dense or sparse) if (!ctx.is_sparse) { 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 PrimalUpdateConstraint(&ctx, flg_Newton & (m->opt.cone == mjCONE_ELLIPTIC)); if (flg_Newton) { // compute and factorize Hessian MakeHessian(d, &ctx); FactorizeHessian(d, &ctx, /*flg_recompute=*/0); } PrimalUpdateGradient(&ctx, flg_Newton); // 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 i=0; i < nv; i++) { int diag_i = ctx.M_rowadr[i] + ctx.M_rownnz[i] - 1; island_inertia += ctx.M[diag_i]; } scale = 1 / island_inertia; } ctx.scale = scale; // main loop while (iter < maxiter) { // perform linesearch mjtNum ls_improvement; alpha = PrimalSearch(&ctx, m->opt.tolerance * m->opt.ls_tolerance, m->opt.ls_iterations, &ls_improvement); // no improvement: done if (alpha == 0) { break; } // move to new solution mju_addToScl(ctx.qacc, ctx.search, alpha, nv); mju_addToScl(ctx.Ma, ctx.Mv, alpha, nv); mju_addToScl(ctx.Jaref, ctx.Jv, alpha, nefc); // save old if (!flg_Newton) { mju_copy(ctx.gradold, ctx.grad, nv); mju_copy(ctx.Mgradold, ctx.Mgrad, nv); } mju_copyInt(oldstate, ctx.efc_state, nefc); // update PrimalUpdateConstraint(&ctx, flg_Newton & (m->opt.cone == mjCONE_ELLIPTIC)); if (flg_Newton) { HessianIncremental(d, &ctx, oldstate); } PrimalUpdateGradient(&ctx, flg_Newton); // count state changes int nchange = 0; for (int i=0; i < nefc; i++) { nchange += (ctx.efc_state[i] != oldstate[i]); } // scale improvement, gradient, save stats mjtNum improvement = scale * ls_improvement; mjtNum gradient = scale * mju_norm(ctx.grad, nv); saveStats(m, d, island, iter, improvement, gradient, ctx.LSslope, ctx.nactive, nchange, ctx.LSiter, ctx.nupdate); // increment iteration count iter++; // termination if ((improvement > 0 && improvement < m->opt.tolerance) || gradient < m->opt.tolerance) { break; } // update direction if (flg_Newton) { mju_scl(ctx.search, ctx.Mgrad, -1, nv); } else { #ifndef mjCG_PRP // Hager-Zhang conjugate direction update mjtNum d_dot_y, y_dot_My, y_dot_Mgrad, d_dot_grad; mjtNum beta_hz; mjtNum d_norm, grad_norm, eta_k; const mjtNum eta = 0.01; // graddif = grad - gradold, Mgraddif = Mgrad - Mgradold mju_sub(ctx.graddif, ctx.grad, ctx.gradold, nv); mju_sub(ctx.Mgraddif, ctx.Mgrad, ctx.Mgradold, nv); // compute d'*y; restart to steepest descent if conjugacy is lost d_dot_y = mju_dot(ctx.search, ctx.graddif, nv); if (d_dot_y < mjMINVAL) { beta = 0; } else { // compute remaining inner products for the HZ formula y_dot_My = mju_dot(ctx.graddif, ctx.Mgraddif, nv); y_dot_Mgrad = mju_dot(ctx.graddif, ctx.Mgrad, nv); d_dot_grad = mju_dot(ctx.search, ctx.grad, nv); // primary Hager-Zhang beta coefficient beta_hz = (y_dot_Mgrad - 2*(y_dot_My/d_dot_y)*d_dot_grad) / d_dot_y; // dynamic truncation threshold to ensure d is not orthogonal to grad d_norm = mju_norm(ctx.search, nv); grad_norm = mju_norm(ctx.grad, nv); eta_k = -1.0 / mju_max(mjMINVAL, d_norm * mju_min(eta, grad_norm)); // apply lower bound beta = mju_max(eta_k, beta_hz); } #else // Polak-Ribiere-Plus conjugate direction update mju_sub(ctx.Mgraddif, ctx.Mgrad, ctx.Mgradold, nv); beta = mju_dot(ctx.grad, ctx.Mgraddif, nv) / mju_max(mjMINVAL, mju_dot(ctx.gradold, ctx.Mgradold, nv)); // reset if negative if (beta < 0) { beta = 0; } #endif // update for (int i=0; i < nv; i++) { ctx.search[i] = -ctx.Mgrad[i] + beta*ctx.search[i]; } } } // 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_stat] += iter; // set solver_nnz if (flg_Newton) { if (mj_isSparse(m)) { // two L factors if Lcone is present int num_factors = 1 + (ctx.Lcone != NULL); d->solver_nnz[island_stat] = num_factors * ctx.nL + ctx.nH; } else { d->solver_nnz[island_stat] = nv*nv; } } else { d->solver_nnz[island_stat] = 0; } } mj_freeStack(d); } // CG entry point void mj_solCG(const mjModel* m, mjData* d, int maxiter) { mj_solPrimal(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_solPrimal(m, d, island, maxiter, /*flg_Newton=*/0); } // Newton entry point void mj_solNewton(const mjModel* m, mjData* d, int maxiter) { mj_solPrimal(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_solPrimal(m, d, island, maxiter, /*flg_Newton=*/1); }