Yuval Tassa
f8462a156d
The line search now requires a negative cost (improvement) in addition to a small derivative to declare convergence, preventing premature termination when no actual improvement has been made. Follows the proposal in github.com/google-deepmind/mujoco_warp/pull/1471 PiperOrigin-RevId: 941579503 Change-Id: I8fcd20f7b959e50d77cd5d6de0a3c6d95f86b9e1
2507 lines
74 KiB
C
2507 lines
74 KiB
C
// 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 <stddef.h>
|
|
#include <stdint.h>
|
|
#include <string.h>
|
|
|
|
#include <mujoco/mjdata.h>
|
|
#include <mujoco/mjmacro.h>
|
|
#include <mujoco/mjmodel.h>
|
|
#include <mujoco/mjsan.h> // 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_macro.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<n; j++ )
|
|
for( k=0; k<rownnz[start+j]; k++ )
|
|
{
|
|
int col = colind[rowadr[start+j]+k];
|
|
if( col>=start && col<start+n )
|
|
Ac[j*n+col-start] = AR[rowadr[start+j]+k];
|
|
}
|
|
*/
|
|
// assume full sub-matrix, find starting k: same for all rows
|
|
int k;
|
|
for (k=0; k < rownnz[start]; k++) {
|
|
if (colind[rowadr[start]+k] == start) {
|
|
break;
|
|
}
|
|
}
|
|
|
|
// SHOULD NOT OCCUR
|
|
if (k >= 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;
|
|
}
|
|
|
|
|
|
// project onto friction ellipsoid, write to force[i+1..i+dim-1]
|
|
// project tangential force onto friction ellipsoid: sum(f_t[j]^2/mu[j]^2) <= f_n^2
|
|
// if feasible is true, only scale down if outside the ellipsoid
|
|
// if feasible is false, always scale to the boundary
|
|
static void projectEllipsoid(mjtNum* friction, mjtNum normal, const mjtNum* mu,
|
|
int dim, int feasible) {
|
|
mjtNum s = 0;
|
|
for (int j=0; j < dim-1; j++) {
|
|
s += friction[j]*friction[j] / (mu[j]*mu[j]);
|
|
}
|
|
|
|
mjtNum normal2 = normal*normal;
|
|
if (!feasible || s > normal2) {
|
|
mjtNum scl = mju_sqrt(normal2 / mju_max(mjMINVAL, s));
|
|
for (int j=0; j < dim-1; j++) {
|
|
friction[j] *= scl;
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
// 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) {
|
|
projectEllipsoid(v, force[i], mu, dim, /*feasible=*/0);
|
|
}
|
|
|
|
// assign
|
|
mju_copy(force+i+1, v, dim-1);
|
|
}
|
|
|
|
|
|
// project contact force block onto friction cone (pyramidal or elliptic)
|
|
static void projectCone(mjtNum* force, const mjtNum* mu, int dim, int type) {
|
|
// elliptic cone: project onto friction ellipsoid
|
|
if (type == mjCNSTR_CONTACT_ELLIPTIC) {
|
|
// clamp normal force
|
|
if (force[0] < 0) {
|
|
mju_zero(force, dim);
|
|
} else {
|
|
projectEllipsoid(force+1, force[0], mu, dim, /*feasible=*/1);
|
|
}
|
|
}
|
|
|
|
// pyramidal or scalar: clamp to non-negative
|
|
else {
|
|
if (force[0] < 0) {
|
|
force[0] = 0;
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
// global variable to toggle Nesterov momentum (for benchmarks/tests)
|
|
mjTHREADLOCAL int mj_nesterov_momentum = 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;
|
|
|
|
// Nesterov momentum
|
|
mjtBool nesterov = (mj_nesterov_momentum != 0);
|
|
mjtNum* force_prev = NULL;
|
|
mjtNum* force_momentum = NULL;
|
|
if (nesterov) {
|
|
force_prev = mjSTACKALLOC(d, nefc, mjtNum);
|
|
force_momentum = mjSTACKALLOC(d, nefc, mjtNum);
|
|
mju_gather(force_prev, force, efclist, 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;
|
|
int nesterov_k = 0; // Nesterov counter (resets on adaptive restart)
|
|
while (iter < maxiter) {
|
|
// Nesterov momentum extrapolation
|
|
if (nesterov) {
|
|
mjtNum beta = 0;
|
|
if (iter > 0) {
|
|
beta = (mjtNum)(nesterov_k - 1) / (mjtNum)(nesterov_k + 2);
|
|
}
|
|
|
|
// update with momentum, save pre-extrapolation value
|
|
if (beta > 0) {
|
|
for (int c=0; c < nefc; c++) {
|
|
int i = efclist ? efclist[c] : c;
|
|
mjtNum f_save = force[i];
|
|
force[i] += beta*(force[i] - force_prev[c]);
|
|
force_prev[c] = f_save;
|
|
}
|
|
|
|
// friction loss: project onto bounds
|
|
for (int c=ne; c < ne+nf; c++) {
|
|
int i = efclist ? efclist[c] : c;
|
|
force[i] = mju_clip(force[i], -floss[i], floss[i]);
|
|
}
|
|
|
|
// contact force: project onto friction cone
|
|
for (int c=ne+nf; c < nefc; ) {
|
|
int i = efclist ? efclist[c] : c;
|
|
int dim = 1;
|
|
int type = d->efc_type[i];
|
|
const mjtNum* mu = NULL;
|
|
|
|
if (type == mjCNSTR_CONTACT_ELLIPTIC) {
|
|
dim = d->contact[d->efc_id[i]].dim;
|
|
mu = d->contact[d->efc_id[i]].friction;
|
|
}
|
|
|
|
projectCone(force+i, mu, dim, type);
|
|
c += dim;
|
|
}
|
|
}
|
|
|
|
// iter == 0 or beta <= 0 (nesterov_k <= 1): just save current force
|
|
else {
|
|
mju_gather(force_prev, force, efclist, nefc);
|
|
}
|
|
|
|
// save extrapolated point for gradient restart check
|
|
mju_gather(force_momentum, force, efclist, nefc);
|
|
}
|
|
|
|
// 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);
|
|
|
|
// Nesterov gradient restart (O'Donoghue-Candès): reset when correction opposes extrapolation
|
|
if (nesterov) {
|
|
mjtBool restart = false;
|
|
if (iter > 0) {
|
|
mjtNum dot_corr_extr = 0;
|
|
for (int c=0; c < nefc; c++) {
|
|
int i = efclist ? efclist[c] : c;
|
|
mjtNum correction = force[i] - force_momentum[c];
|
|
mjtNum extrapolation = force_momentum[c] - force_prev[c];
|
|
dot_corr_extr += correction * extrapolation;
|
|
}
|
|
restart = (dot_corr_extr < 0);
|
|
}
|
|
|
|
if (restart) {
|
|
nesterov_k = 0;
|
|
} else {
|
|
nesterov_k++;
|
|
}
|
|
}
|
|
|
|
// 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 && (p1.alpha == 0 || p1.cost < 0)) {
|
|
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 && p1.cost < 0) {
|
|
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] = ctx.nJ;
|
|
}
|
|
}
|
|
|
|
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);
|
|
}
|