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

Use the word "primal" for things that relate to both Newton and CG.

Browse Source 
PiperOrigin-RevId: 852634373
Change-Id: I2c7641bb92d9c01252f861b9c36e72ef6ddc6e98
This commit is contained in:
Yuval Tassa
2026-01-05 23:55:45 -08:00
committed by Copybara-Service
parent fa70f73691
commit fa044fe0e5
1 changed files with 49 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_solver.c
+49 -52
View File
@@ -747,10 +747,10 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
} }
//------------------------- CG and Newton solver -------------------------------------------------- //------------------------- Primal solvers ---------------------------------------------------------
// CG context // Primal context
struct _mjCGContext { typedef struct {
int is_sparse; // 1: sparse, 0: dense int is_sparse; // 1: sparse, 0: dense
int is_elliptic; // 1: elliptic, 0: pyramidal int is_elliptic; // 1: elliptic, 0: pyramidal
int island; // current island index, -1 if monolithic int island; // current island index, -1 if monolithic
@@ -801,7 +801,7 @@ struct _mjCGContext {
int* JT_colind; int* JT_colind;
mjtNum* JT; mjtNum* JT;
// common arrays (CGallocate) // common arrays (PrimalAllocate)
mjtNum* Jaref; // Jac*qacc - aref (nefc x 1) mjtNum* Jaref; // Jac*qacc - aref (nefc x 1)
mjtNum* Jv; // Jac*search (nefc x 1) mjtNum* Jv; // Jac*search (nefc x 1)
mjtNum* Ma; // M*qacc (nv x 1) mjtNum* Ma; // M*qacc (nv x 1)
@@ -811,7 +811,7 @@ struct _mjCGContext {
mjtNum* search; // linesearch vector (nv x 1) mjtNum* search; // linesearch vector (nv x 1)
mjtNum* quad; // quadratic polynomials for constraint costs (nefc x 3) mjtNum* quad; // quadratic polynomials for constraint costs (nefc x 3)
// Newton arrays, known-size (CGallocate) // Newton arrays, known-size (PrimalAllocate)
mjtNum* D; // constraint inertia (nefc x 1) mjtNum* D; // constraint inertia (nefc x 1)
int* H_rowadr; // Hessian row addresses (nv x 1) int* H_rowadr; // Hessian row addresses (nv x 1)
int* H_rownnz; // Hessian row nonzeros (nv x 1) int* H_rownnz; // Hessian row nonzeros (nv x 1)
@@ -848,15 +848,13 @@ struct _mjCGContext {
int LSiter; // number of linesearch iterations int LSiter; // number of linesearch iterations
int LSresult; // linesearch result int LSresult; // linesearch result
mjtNum LSslope; // linesearch slope at solution mjtNum LSslope; // linesearch slope at solution
}; } mjPrimalContext;
typedef struct _mjCGContext mjCGContext;
// set sizes and pointers to mjData arrays in mjPrimalContext
// set sizes and pointers to mjData arrays in mjCGContext static void PrimalPointers(const mjModel* m, const mjData* d, mjPrimalContext* ctx, int island) {
static void CGpointers(const mjModel* m, const mjData* d, mjCGContext* ctx, int island) {
// clear everything // clear everything
memset(ctx, 0, sizeof(mjCGContext)); memset(ctx, 0, sizeof(mjPrimalContext));
// globals // globals
ctx->is_sparse = mj_isSparse(m); ctx->is_sparse = mj_isSparse(m);
@@ -957,9 +955,9 @@ static void CGpointers(const mjModel* m, const mjData* d, mjCGContext* ctx, int
} }
// allocate fixed-size arrays in mjCGContext // allocate fixed-size arrays in mjPrimalContext
// mj_{mark/free}Stack in calling function! // mj_{mark/free}Stack in calling function!
static void CGallocate(mjData* d, mjCGContext* ctx, int flg_Newton) { static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
// local sizes // local sizes
int nv = ctx->nv; int nv = ctx->nv;
int nefc = ctx->nefc; int nefc = ctx->nefc;
@@ -1009,7 +1007,7 @@ static void CGallocate(mjData* d, mjCGContext* ctx, int flg_Newton) {
// update efc_force, qfrc_constraint, cost-related // update efc_force, qfrc_constraint, cost-related
static void CGupdateConstraint(mjCGContext* ctx, int flg_HessianCone) { static void PrimalUpdateConstraint(mjPrimalContext* ctx, int flg_HessianCone) {
int nefc = ctx->nefc, nv = ctx->nv; int nefc = ctx->nefc, nv = ctx->nv;
// update constraints // update constraints
@@ -1046,7 +1044,7 @@ static void CGupdateConstraint(mjCGContext* ctx, int flg_HessianCone) {
// update grad, Mgrad // update grad, Mgrad
static void CGupdateGradient(mjCGContext* ctx, int flg_Newton) { static void PrimalUpdateGradient(mjPrimalContext* ctx, int flg_Newton) {
int nv = ctx->nv; int nv = ctx->nv;
// grad = M*qacc - qfrc_smooth - qfrc_constraint // grad = M*qacc - qfrc_smooth - qfrc_constraint
@@ -1074,12 +1072,12 @@ static void CGupdateGradient(mjCGContext* ctx, int flg_Newton) {
// prepare quadratic polynomials and contact cone quantities // prepare quadratic polynomials and contact cone quantities
static void CGprepare(mjCGContext* ctx) { static void PrimalPrepare(mjPrimalContext* ctx) {
int nv = ctx->nv, nefc = ctx->nefc; int nv = ctx->nv, nefc = ctx->nefc;
const mjtNum* v = ctx->search; 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) // 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 // quadGauss[0] already computed in PrimalUpdateConstraint
ctx->quadGauss[1] = mju_dot(v, ctx->Ma, nv) - mju_dot(ctx->qfrc_smooth, v, nv); 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); ctx->quadGauss[2] = 0.5*mju_dot(v, ctx->Mv, nv);
@@ -1150,17 +1148,16 @@ static void CGprepare(mjCGContext* ctx) {
// linesearch evaluation point // linesearch evaluation point
struct _mjCGPnt { struct _mjPrimalPnt {
mjtNum alpha; mjtNum alpha;
mjtNum cost; mjtNum cost;
mjtNum deriv[2]; mjtNum deriv[2];
}; };
typedef struct _mjCGPnt mjCGPnt; typedef struct _mjPrimalPnt mjPrimalPnt;
// evaluate linesearch cost, return first and second derivatives // evaluate linesearch cost, return first and second derivatives
static void CGeval(mjCGContext* ctx, mjCGPnt* p) { static void PrimalEval(mjPrimalContext* ctx, mjPrimalPnt* p) {
int ne = ctx->ne, nf = ctx->nf, nefc = ctx->nefc; int ne = ctx->ne, nf = ctx->nf, nefc = ctx->nefc;
// clear result // clear result
@@ -1298,8 +1295,8 @@ static void CGeval(mjCGContext* ctx, mjCGPnt* p) {
// update bracket point given 3 candidate points // update bracket point given 3 candidate points
static int updateBracket(mjCGContext* ctx, static int updateBracket(mjPrimalContext* ctx,
mjCGPnt* p, const mjCGPnt candidates[3], mjCGPnt* pnext) { mjPrimalPnt* p, const mjPrimalPnt candidates[3], mjPrimalPnt* pnext) {
int flag = 0; int flag = 0;
for (int i=0; i < 3; i++) { for (int i=0; i < 3; i++) {
// negative deriv // negative deriv
@@ -1320,7 +1317,7 @@ static int updateBracket(mjCGContext* ctx,
// compute next point if updated // compute next point if updated
if (flag) { if (flag) {
pnext->alpha = p->alpha - p->deriv[0]/p->deriv[1]; pnext->alpha = p->alpha - p->deriv[0]/p->deriv[1];
CGeval(ctx, pnext); PrimalEval(ctx, pnext);
} }
return flag; return flag;
@@ -1328,9 +1325,9 @@ static int updateBracket(mjCGContext* ctx,
// line search // line search
static mjtNum CGsearch(mjCGContext* ctx, mjtNum tolerance, mjtNum ls_iterations) { static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_iterations) {
int nv = ctx->nv, nefc = ctx->nefc; int nv = ctx->nv, nefc = ctx->nefc;
mjCGPnt p0, p1, p2, pmid, p1next, p2next; mjPrimalPnt p0, p1, p2, pmid, p1next, p2next;
// clear results // clear results
ctx->LSiter = 0; ctx->LSiter = 0;
@@ -1361,15 +1358,15 @@ static mjtNum CGsearch(mjCGContext* ctx, mjtNum tolerance, mjtNum ls_iterations)
} }
// prepare quadratics and cones // prepare quadratics and cones
CGprepare(ctx); PrimalPrepare(ctx);
// init at alpha = 0, save // init at alpha = 0, save
p0.alpha = 0; p0.alpha = 0;
CGeval(ctx, &p0); PrimalEval(ctx, &p0);
// always attempt one Newton step // always attempt one Newton step
p1.alpha = p0.alpha - p0.deriv[0]/p0.deriv[1]; p1.alpha = p0.alpha - p0.deriv[0]/p0.deriv[1];
CGeval(ctx, &p1); PrimalEval(ctx, &p1);
// check for initial convergence // check for initial convergence
if (mju_abs(p1.deriv[0]) < gtol) { if (mju_abs(p1.deriv[0]) < gtol) {
@@ -1429,7 +1426,7 @@ static mjtNum CGsearch(mjCGContext* ctx, mjtNum tolerance, mjtNum ls_iterations)
// move to Newton point w.r.t current // move to Newton point w.r.t current
p1.alpha -= p1.deriv[0]/p1.deriv[1]; p1.alpha -= p1.deriv[0]/p1.deriv[1];
CGeval(ctx, &p1); PrimalEval(ctx, &p1);
// check for convergence // check for convergence
if (mju_abs(p1.deriv[0]) < gtol) { if (mju_abs(p1.deriv[0]) < gtol) {
@@ -1455,16 +1452,16 @@ static mjtNum CGsearch(mjCGContext* ctx, mjtNum tolerance, mjtNum ls_iterations)
// compute next-points for bracket // compute next-points for bracket
p2next = p1; p2next = p1;
p1next.alpha = p1.alpha - p1.deriv[0]/p1.deriv[1]; p1next.alpha = p1.alpha - p1.deriv[0]/p1.deriv[1];
CGeval(ctx, &p1next); PrimalEval(ctx, &p1next);
// bracketed search // bracketed search
while (ctx->LSiter < ls_iterations) { while (ctx->LSiter < ls_iterations) {
// evaluate at midpoint // evaluate at midpoint
pmid.alpha = 0.5*(p1.alpha + p2.alpha); pmid.alpha = 0.5*(p1.alpha + p2.alpha);
CGeval(ctx, &pmid); PrimalEval(ctx, &pmid);
// make list of candidates // make list of candidates
mjCGPnt candidates[3] = {p1next, p2next, pmid}; mjPrimalPnt candidates[3] = {p1next, p2next, pmid};
// check candidates for convergence // check candidates for convergence
mjtNum bestcost = 0; mjtNum bestcost = 0;
@@ -1516,7 +1513,7 @@ static mjtNum CGsearch(mjCGContext* ctx, mjtNum tolerance, mjtNum ls_iterations)
// allocate and compute Hessian given efc_state // allocate and compute Hessian given efc_state
// mj_{mark/free}Stack in caller function! // mj_{mark/free}Stack in caller function!
static void MakeHessian(mjData* d, mjCGContext* ctx) { static void MakeHessian(mjData* d, mjPrimalContext* ctx) {
int nv = ctx->nv, nefc = ctx->nefc; int nv = ctx->nv, nefc = ctx->nefc;
// compute constraint inertia // compute constraint inertia
@@ -1605,11 +1602,11 @@ static void MakeHessian(mjData* d, mjCGContext* ctx) {
} }
// forward declaration of HessianCone (readability) // forward declaration of HessianCone (for readability)
static void HessianCone(mjData* d, mjCGContext* ctx); static void HessianCone(mjData* d, mjPrimalContext* ctx);
// factorize Hessian: L = chol(H), maybe (re)compute H given efc_state // factorize Hessian: L = chol(H), maybe (re)compute H given efc_state
static void FactorizeHessian(mjData* d, mjCGContext* ctx, int flg_recompute) { static void FactorizeHessian(mjData* d, mjPrimalContext* ctx, int flg_recompute) {
int nv = ctx->nv, nefc = ctx->nefc; int nv = ctx->nv, nefc = ctx->nefc;
// maybe compute constraint inertia // maybe compute constraint inertia
@@ -1674,7 +1671,7 @@ static void FactorizeHessian(mjData* d, mjCGContext* ctx, int flg_recompute) {
// elliptic case: Hcone = H + cone_contributions // elliptic case: Hcone = H + cone_contributions
static void HessianCone(mjData* d, mjCGContext* ctx) { static void HessianCone(mjData* d, mjPrimalContext* ctx) {
int nv = ctx->nv, nefc = ctx->nefc; int nv = ctx->nv, nefc = ctx->nefc;
mjtNum local[36]; mjtNum local[36];
@@ -1746,7 +1743,7 @@ static void HessianCone(mjData* d, mjCGContext* ctx) {
// incremental update to Hessian factor due to changes in efc_state // incremental update to Hessian factor due to changes in efc_state
static void HessianIncremental(mjData* d, mjCGContext* ctx, const int* oldstate) { static void HessianIncremental(mjData* d, mjPrimalContext* ctx, const int* oldstate) {
int rank, nv = ctx->nv, nefc = ctx->nefc; int rank, nv = ctx->nv, nefc = ctx->nefc;
mj_markStack(d); mj_markStack(d);
@@ -1811,16 +1808,16 @@ static void HessianIncremental(mjData* d, mjCGContext* ctx, const int* oldstate)
// driver // driver
static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter, int flg_Newton) { static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, int flg_Newton) {
int iter = 0; int iter = 0;
mjtNum alpha, beta; mjtNum alpha, beta;
mjtNum *gradold = NULL, *Mgradold = NULL, *Mgraddif = NULL; mjtNum *gradold = NULL, *Mgradold = NULL, *Mgraddif = NULL;
mjCGContext ctx; mjPrimalContext ctx;
mj_markStack(d); mj_markStack(d);
// make context // make context
CGpointers(m, d, &ctx, island); PrimalPointers(m, d, &ctx, island);
CGallocate(d, &ctx, flg_Newton); PrimalAllocate(d, &ctx, flg_Newton);
// local copies // local copies
int nv = ctx.nv; int nv = ctx.nv;
@@ -1849,13 +1846,13 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
mju_subFrom(ctx.Jaref, ctx.efc_aref, nefc); mju_subFrom(ctx.Jaref, ctx.efc_aref, nefc);
// first update // first update
CGupdateConstraint(&ctx, flg_Newton & (m->opt.cone == mjCONE_ELLIPTIC)); PrimalUpdateConstraint(&ctx, flg_Newton & (m->opt.cone == mjCONE_ELLIPTIC));
if (flg_Newton) { if (flg_Newton) {
// compute and factorize Hessian // compute and factorize Hessian
MakeHessian(d, &ctx); MakeHessian(d, &ctx);
FactorizeHessian(d, &ctx, /*flg_recompute=*/0); FactorizeHessian(d, &ctx, /*flg_recompute=*/0);
} }
CGupdateGradient(&ctx, flg_Newton); PrimalUpdateGradient(&ctx, flg_Newton);
// start both with preconditioned gradient // start both with preconditioned gradient
mju_scl(ctx.search, ctx.Mgrad, -1, nv); mju_scl(ctx.search, ctx.Mgrad, -1, nv);
@@ -1877,7 +1874,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
// main loop // main loop
while (iter < maxiter) { while (iter < maxiter) {
// perform linesearch // perform linesearch
alpha = CGsearch(&ctx, m->opt.tolerance * m->opt.ls_tolerance, m->opt.ls_iterations); alpha = PrimalSearch(&ctx, m->opt.tolerance * m->opt.ls_tolerance, m->opt.ls_iterations);
// no improvement: done // no improvement: done
if (alpha == 0) { if (alpha == 0) {
@@ -1898,11 +1895,11 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
mjtNum oldcost = ctx.cost; mjtNum oldcost = ctx.cost;
// update // update
CGupdateConstraint(&ctx, flg_Newton & (m->opt.cone == mjCONE_ELLIPTIC)); PrimalUpdateConstraint(&ctx, flg_Newton & (m->opt.cone == mjCONE_ELLIPTIC));
if (flg_Newton) { if (flg_Newton) {
HessianIncremental(d, &ctx, oldstate); HessianIncremental(d, &ctx, oldstate);
} }
CGupdateGradient(&ctx, flg_Newton); PrimalUpdateGradient(&ctx, flg_Newton);
// count state changes // count state changes
int nchange = 0; int nchange = 0;
@@ -1973,23 +1970,23 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
// CG entry point // CG entry point
void mj_solCG(const mjModel* m, mjData* d, int maxiter) { void mj_solCG(const mjModel* m, mjData* d, int maxiter) {
mj_solCGNewton(m, d, /*island=*/-1, maxiter, /*flg_Newton=*/0); mj_solPrimal(m, d, /*island=*/-1, maxiter, /*flg_Newton=*/0);
} }
// CG entry point (one island) // CG entry point (one island)
void mj_solCG_island(const mjModel* m, mjData* d, int island, int maxiter) { void mj_solCG_island(const mjModel* m, mjData* d, int island, int maxiter) {
mj_solCGNewton(m, d, island, maxiter, /*flg_Newton=*/0); mj_solPrimal(m, d, island, maxiter, /*flg_Newton=*/0);
} }
// Newton entry point // Newton entry point
void mj_solNewton(const mjModel* m, mjData* d, int maxiter) { void mj_solNewton(const mjModel* m, mjData* d, int maxiter) {
mj_solCGNewton(m, d, /*island=*/-1, maxiter, /*flg_Newton=*/1); mj_solPrimal(m, d, /*island=*/-1, maxiter, /*flg_Newton=*/1);
} }
// Newton entry point (one island) // Newton entry point (one island)
void mj_solNewton_island(const mjModel* m, mjData* d, int island, int maxiter) { void mj_solNewton_island(const mjModel* m, mjData* d, int island, int maxiter) {
mj_solCGNewton(m, d, island, maxiter, /*flg_Newton=*/1); mj_solPrimal(m, d, island, maxiter, /*flg_Newton=*/1);
} }