Implicit flex elasticity in the CG constraint solver via an effective metric
This CL replaces the post-hoc implicit flex correction (`flexInterp_cgsolve`) with a **linearly-implicit effective metric** `M̃ = M + (h² + h·damping)·K` carried by the CG constraint solver itself. Contact/friction forces and implicit flex elasticity are now computed against one consistent metric, instead of the solver seeing `M` and a post-solve correction changing `qacc` behind its back. Gate (unchanged semantics): `solver="CG"` + implicit/implicitfast integrator + pyramidal cones + flex stiffness present. Newton and PGS are untouched. `solver="CG"` remains the user-facing contract — the factorization is an implementation detail of the preconditioner. ### What's in the metric - **mjData `efm_*`** (arena, efc-like lifetime/skip semantics; built in `mj_fwdPosition`, value-refreshed in `mj_fwdVelocity`): the per-step stiffness CSR `efm_B_*`, its reverse-Cholesky factor `efm_dofid` + `efm_L_*` (nested-dissection ordered, separators-first for the reverse factorization), and the smooth-force shift `efm_c = h·K·qvel`. - **`mjd_flexStiff_assemble`** now assembles stretch (Gauss–Newton), standard dim-2 bending, and — via the cached corotated stiffness `d->flexelem_krot` — interp stiffness (all node bodies on simple sliders: point Jacobian is I₃, `flex_centered` not required; fixed nodes drop like pins) into one dof-level CSR. `mjd_effMulAdd`/`mjd_effSolve` apply the metric, with matrix-free operator fallbacks where assembly does not apply. - **mjModel `efm0_*`** (`nefm0dof`/`nefm0L`): the constant part of the metric factor — currently the dim-2 bending factor, computed once in `mj_setConst` — so bending-only models pay zero per-step factorization cost. Naming mirrors mjData's `efm_*` with the standard `0`-suffix (reference/constant) idiom, and is deliberately not bending-specific: future constant contributors extend it without renames. - The solver consumes the metric through pre-shifted `qfrc_smooth` and the metric products `Ma`/`Mv`/`Mgrad`; `qacc_smooth` becomes the unconstrained minimizer of the implicit dynamics, which makes the no-constraint shortcut and the warmstart choice consistent by construction. - **`mj_inverse` adds `B·qacc − c`**, making inverse dynamics discrete-consistent with the gated forward dynamics — exact, since the gated path has no qDeriv term (new test `ForwardTest.GatedFlexInverseConsistency`). ### Performance All numbers: ms/step over the same 2000-step window, models as shipped on each side (old code with the old model settings vs this CL with the new ones). The new solver path activates on exactly two shipped models — the ponchos, the only flex models that need an implicit integrator (poncho on Euler degenerates to >200 ms/step). For them, this CL trades speed for consistency: the implicit bending solve now runs inside every solver iteration, where the contact solve can see the stiffness, instead of once after the solve. Solver iterations drop because the curvature is visible, but each iteration pays for the implicit solve: | model | before | after | solver iters/step | |---|---|---|---| | poncho | 2.47 | 3.30 (1.33×) | 16.8 → 11.8 | | poncho_edgeequality | 1.96 | 2.72 (1.39×) | 13.2 → 10.0 | What that price buys: contact forces consistent with the implicit elasticity (previously the post-hoc correction changed `qacc` after the constraint solve), discrete-consistent inverse dynamics, and the removal of the post-hoc special case from the integration path. Raising poncho's timestep from 2 to 5 ms leaves its per-step cost nearly flat, so the consistency price can be recovered by taking fewer steps where accuracy allows. Every other flex model was measured stable on Euler at its shipped timestep and switches to it (these models predate the post-hoc integrator; implicit was never load-bearing for them). They end up equal or faster than before: bunny_multicell 0.47 → 0.40, trampoline 0.28 → 0.25, plate 1.02 → 0.99, pancake 0.34 → 0.33. Finally, the per-step factorization makes configurations practical that the old code could only integrate explicitly: implicit stretch elasticity (`elastic2d="stretch"`/`"both"`, dim-3 solids) and factorized interp stiffness. No before/after exists for these — stock has no implicit treatment of stretch at all. ### Behavior changes - With the post-hoc correction deleted, interp/bending models running `solver="Newton"` (or elliptic cones, or islands) now integrate flex elasticity **explicitly** (previously: post-hoc implicit). Affects e.g. `gripper_trilinear` (stable, and faster, but different semantics). Follow-up options: Newton-side metric support, or a documented fallback. - With the gate on, `mj_forward` outputs are timestep-dependent for gated models (they answer the linearly-implicit discrete problem); `qacc_smooth` and `mj_inverse` change accordingly. Non-gated models are bit-identical (full suite green throughout). ### Validation - 1737/1737 tests, including new: `FlexStretchDerivatives` (FD-validated GN operator), `FlexStiffAssemble`/`FlexStiffAssembleInterp` (CSR ≡ operators), `GatedFlexInverseConsistency` (fails pre-change), equivalence tests vs the old post-hoc treatment (bending matches to 2e-11). - Fingerprint discipline throughout: bending-only models bit-exact across every refactor; permutation/kernel changes verified iteration-identical. ### Known follow-ups (not in this CL) 3×3-block sparse Cholesky kernel (the numeric factorization is index-bound; projected ~3× on the factor); mjModel persistence of the factor's symbolic pattern (rest-pose ND makes sizes compile-time); the general effective-metric mode (all solvers, all PSD-safe force classes, behind an enable flag). PiperOrigin-RevId: 948561856 Change-Id: I8b8e32ebd0428042af71647d0470d10773bf6daf
This commit is contained in:
Alessio Quaglino
committed by
Copybara-Service
Copybara-Service
parent
f0fa3d8260
commit
ea230a950c
+65
-121
@@ -168,6 +168,11 @@ void mj_fwdPosition(const mjModel* m, mjData* d) {
|
||||
mj_transmission(m, d);
|
||||
TM_ADD(mjTIMER_POS_KINEMATICS);
|
||||
|
||||
// implicit effective metric Mtilde = M + K: build (or deactivate) for this step. Arena
|
||||
// lifetime and skip semantics mirror the constraint data: built once per position stage,
|
||||
// value-refreshed in the velocity stage, consumed downstream.
|
||||
mjd_effBuild(m, d, mj_flexCG(m), /*flg_factor=*/1);
|
||||
|
||||
TM_END1(mjTIMER_POSITION);
|
||||
}
|
||||
|
||||
@@ -214,6 +219,8 @@ void mj_fwdVelocity(const mjModel* m, mjData* d) {
|
||||
// add bias force due to tendon armature
|
||||
mj_tendonBias(m, d, d->qfrc_bias);
|
||||
|
||||
mjd_effShift(m, d);
|
||||
|
||||
TM_END(mjTIMER_VELOCITY);
|
||||
}
|
||||
|
||||
@@ -786,6 +793,18 @@ void mj_fwdAcceleration(const mjModel* m, mjData* d) {
|
||||
// qfrc_smooth += project(xfrc_applied)
|
||||
mj_xfrcAccumulate(m, d, d->qfrc_smooth);
|
||||
|
||||
// implicit effective metric (built in mj_fwdPosition): the smooth acceleration is that of
|
||||
// the linearly-implicit dynamics, (M + K)*qacc_smooth = qfrc_smooth + c, so the constraint
|
||||
// solver, the no-constraint shortcut and the warmstart all see one consistent metric.
|
||||
if (d->efm_active) {
|
||||
mj_markStack(d);
|
||||
mjtNum* qfrc_eff = mjSTACKALLOC(d, nv, mjtNum);
|
||||
mju_add(qfrc_eff, d->qfrc_smooth, d->efm_c, nv);
|
||||
mjd_effSolve(m, d, d->qacc_smooth, qfrc_eff);
|
||||
mj_freeStack(d);
|
||||
return;
|
||||
}
|
||||
|
||||
// copy for in-place solve: qacc_smooth = qfrc_smooth
|
||||
if (!sleep_filter) {
|
||||
mju_copy(d->qacc_smooth, d->qfrc_smooth, nv);
|
||||
@@ -900,6 +919,7 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
|
||||
mju_zero(d->qfrc_constraint, nv);
|
||||
|
||||
// no constraints: copy unconstrained acc, clear forces, return
|
||||
// (with the effective metric active, qacc_smooth is already the implicit answer)
|
||||
if (!nefc) {
|
||||
mju_copy(d->qacc, d->qacc_smooth, nv);
|
||||
mju_zeroInt(d->solver_niter, mjNISLAND);
|
||||
@@ -921,7 +941,16 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
|
||||
mju_zeroInt(d->solver_niter, mjNISLAND);
|
||||
|
||||
// check if islands are supported
|
||||
int islands_supported = !mjDISABLED(mjDSBL_ISLAND) && nisland > 0;
|
||||
// TODO: support islands with the implicit effective metric and remove the mj_flexCG
|
||||
// condition. It is here because the metric machinery is monolithic: the efm_c shift and
|
||||
// the Ma/Mv/Mgrad operators (mjd_effMulAdd, mjd_effSolve) act on global dof vectors with
|
||||
// no island-local form. Discovery is already handled: findEdges unions the trees of every
|
||||
// stiffness-active flex, so a flex always lands in one island together with everything it
|
||||
// touches. Removal therefore needs only the solver side: apply the efm_c shift to that
|
||||
// island's dofs, gather/scatter its island-local vectors around the covered-compact
|
||||
// factor solves (the factors themselves need no change), and enable the metric path
|
||||
// (flg_flex) for the flex-containing island alone.
|
||||
int islands_supported = !mjDISABLED(mjDSBL_ISLAND) && nisland > 0 && !mj_flexCG(m);
|
||||
|
||||
// run solver over constraint islands
|
||||
if (islands_supported) {
|
||||
@@ -1298,7 +1327,7 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
|
||||
|
||||
|
||||
// return 1 if any flex needs implicit stiffness treatment (interp or bending)
|
||||
static int flex_has_implicit_stiffness(const mjModel* m) {
|
||||
static mjtBool flex_has_implicit_stiffness(const mjModel* m) {
|
||||
for (int f=0; f < m->nflex; f++) {
|
||||
if (m->flex_rigid[f]) {
|
||||
continue;
|
||||
@@ -1316,120 +1345,32 @@ static int flex_has_implicit_stiffness(const mjModel* m) {
|
||||
m->flex_bendingadr[f] >= 0) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
// standard flex with stretch
|
||||
if (!m->flex_interp[f] && m->flex_dim[f] >= 2 &&
|
||||
m->flex_stiffnessadr[f] >= 0 &&
|
||||
m->flex_stiffness[m->flex_stiffnessadr[f]] != 0) {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
// preconditioned CG solve for implicit flex interp
|
||||
// solves (M - h*qDeriv - (h^2+h*d)*K) * qacc = qfrc - h*K*qvel
|
||||
// where K is the flex stiffness, using the already-factored standard system
|
||||
// (M - h*qDeriv) as a preconditioner
|
||||
static void flexInterp_cgsolve(const mjModel* m, mjData* d,
|
||||
mjtNum* qacc, const mjtNum* qfrc, int nv) {
|
||||
mjtNum h = m->opt.timestep;
|
||||
int implicit = (m->opt.integrator == mjINT_IMPLICIT);
|
||||
|
||||
mj_markStack(d);
|
||||
|
||||
// allocate CG work vectors
|
||||
mjtNum* rhs = mjSTACKALLOC(d, nv, mjtNum);
|
||||
mjtNum* r = mjSTACKALLOC(d, nv, mjtNum);
|
||||
mjtNum* z = mjSTACKALLOC(d, nv, mjtNum);
|
||||
mjtNum* p = mjSTACKALLOC(d, nv, mjtNum);
|
||||
mjtNum* Ap = mjSTACKALLOC(d, nv, mjtNum);
|
||||
mjtNum* temp = mjSTACKALLOC(d, nv, mjtNum);
|
||||
|
||||
// precompute K_rot cache: same layout as m->flex_stiffness
|
||||
int krot_size = m->nflexstiffness;
|
||||
mjtNum* K_rot_cache = mjSTACKALLOC(d, krot_size, mjtNum);
|
||||
mju_zero(K_rot_cache, krot_size);
|
||||
mjd_flexInterp_cacheKrot(m, d, K_rot_cache);
|
||||
|
||||
// build RHS: rhs = qfrc
|
||||
mju_copy(rhs, qfrc, nv);
|
||||
|
||||
// flex_interp velocity correction: rhs += h*K_interp*qvel (K_interp is NSD)
|
||||
mjd_flexInterp_mul(m, d, rhs, d->qvel, h, 0, K_rot_cache);
|
||||
|
||||
// standard flex bending velocity correction: rhs -= h*K_bend*qvel
|
||||
mjd_flexBend_mul(m, d, rhs, d->qvel, -h, 0); // rhs -= h*K_bend*v
|
||||
|
||||
// --- helper: compute Ap = A*x ---
|
||||
// A*x = (M - h*qDeriv)*x - (h^2+h*d)*K_interp*x + (h^2+h*d)*K_bend*x
|
||||
#define FLEX_CG_MATVEC(Ap_out, x_in) \
|
||||
mju_mulMatVecSparse(Ap_out, d->qDeriv, x_in, nv, m->D_rownnz, m->D_rowadr, \
|
||||
m->D_colind, NULL); \
|
||||
mju_mulSymVecSparse(temp, d->M, x_in, nv, m->M_rownnz, m->M_rowadr, \
|
||||
m->M_colind); \
|
||||
mju_addScl(Ap_out, temp, Ap_out, -h, nv); \
|
||||
mjd_flexInterp_mul(m, d, Ap_out, x_in, -(h*h), -h, K_rot_cache); \
|
||||
mjd_flexBend_mul(m, d, Ap_out, x_in, h*h, h)
|
||||
|
||||
// --- helper: preconditioner solve z = (M - h*qDeriv)^{-1} * r ---
|
||||
#define FLEX_CG_PRECOND(z_out, r_in) \
|
||||
if (implicit) { \
|
||||
mju_solveLUSparse(z_out, d->qLU, r_in, nv, m->D_rownnz, m->D_rowadr, \
|
||||
m->D_diag, m->D_colind, NULL); \
|
||||
} else { \
|
||||
mju_copy(z_out, r_in, nv); \
|
||||
mj_solveLD(z_out, d->qH, d->qHDiagInv, nv, 1, m->M_rownnz, m->M_rowadr, \
|
||||
m->M_colind, NULL); \
|
||||
}
|
||||
|
||||
// initial residual: r = rhs - A*qacc
|
||||
FLEX_CG_MATVEC(Ap, qacc);
|
||||
mju_sub(r, rhs, Ap, nv);
|
||||
|
||||
// check if already converged
|
||||
mjtNum rnorm = mju_dot(r, r, nv);
|
||||
mjtNum tol = 1e-10 * mju_dot(rhs, rhs, nv);
|
||||
if (rnorm < tol || rnorm < mjMINVAL) {
|
||||
mj_freeStack(d);
|
||||
return;
|
||||
}
|
||||
|
||||
// z = precond(r), p = z
|
||||
FLEX_CG_PRECOND(z, r);
|
||||
mju_copy(p, z, nv);
|
||||
mjtNum rz = mju_dot(r, z, nv);
|
||||
|
||||
// CG iterations
|
||||
int maxiter = 50;
|
||||
for (int iter=0; iter < maxiter; iter++) {
|
||||
FLEX_CG_MATVEC(Ap, p);
|
||||
|
||||
// alpha = rz / dot(p, Ap)
|
||||
mjtNum pAp = mju_dot(p, Ap, nv);
|
||||
if (mju_abs(pAp) < mjMINVAL) break;
|
||||
mjtNum alpha = rz / pAp;
|
||||
|
||||
// qacc += alpha * p
|
||||
mju_addToScl(qacc, p, alpha, nv);
|
||||
|
||||
// r -= alpha * Ap
|
||||
mju_addToScl(r, Ap, -alpha, nv);
|
||||
|
||||
// check convergence
|
||||
rnorm = mju_dot(r, r, nv);
|
||||
if (rnorm < tol || rnorm < mjMINVAL) break;
|
||||
|
||||
// z = precond(r)
|
||||
FLEX_CG_PRECOND(z, r);
|
||||
|
||||
// beta = rz_new / rz
|
||||
mjtNum rz_new = mju_dot(r, z, nv);
|
||||
mjtNum beta = rz_new / mju_max(mjMINVAL, rz);
|
||||
|
||||
// p = z + beta * p
|
||||
mju_addScl(p, z, p, beta, nv);
|
||||
rz = rz_new;
|
||||
}
|
||||
|
||||
#undef FLEX_CG_MATVEC
|
||||
#undef FLEX_CG_PRECOND
|
||||
|
||||
mj_freeStack(d);
|
||||
// implicit-flex solve gate: with the CG solver, an implicit integrator and flex stiffness
|
||||
// present, the CG solve carries the implicit flex stiffness itself -- K = (h^2+h*d) times the flex stiffness enters
|
||||
// the objective/gradient/linesearch, and the preconditioned gradient becomes (M+K)\grad by
|
||||
// linear matrix-free CG against the existing M factor (the in-solver form of the old post-hoc
|
||||
// flexInterp_cgsolve treatment, no factorization anywhere); mj_implicitSkip then folds the
|
||||
// implicit flex force of the solver's qacc into qfrc. When active with islands
|
||||
// enabled, mj_fwdConstraint forces a monolithic solve (flex mesh coupling is invisible to
|
||||
// constraint islanding). solver="Newton" keeps its exact-factorization semantics untouched.
|
||||
// Models outside the gate integrate flex elasticity explicitly.
|
||||
int mj_flexCG(const mjModel* m) {
|
||||
return m->opt.solver == mjSOL_CG &&
|
||||
(m->opt.integrator == mjINT_IMPLICIT || m->opt.integrator == mjINT_IMPLICITFAST) &&
|
||||
m->opt.cone != mjCONE_ELLIPTIC && !mjENABLED(mjENBL_SLEEP) &&
|
||||
flex_has_implicit_stiffness(m);
|
||||
}
|
||||
|
||||
|
||||
@@ -1454,11 +1395,16 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
|
||||
mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
|
||||
}
|
||||
|
||||
// check for flex_interp that needs implicit treatment
|
||||
int has_flex_stiffness = !sleep_filter && flex_has_implicit_stiffness(m);
|
||||
// implicit flex stiffness is carried by the constraint solver (see mj_flexCG): use the
|
||||
// solver's qacc directly. The qDeriv treatment is skipped for these models -- flex damping
|
||||
// is already implicit inside the solve (the s2 terms of B), joint damping and other velocity
|
||||
// derivatives integrate explicitly. This avoids both the qDeriv machinery and the
|
||||
// sequential flex-vs-qDeriv splitting. Models outside the gate (non-Newton solver, elliptic
|
||||
// cones, islands, sleep) integrate flex elasticity explicitly.
|
||||
int flexcg = !sleep_filter && mj_flexCG(m);
|
||||
|
||||
// factorization
|
||||
if (!skipfactor) {
|
||||
if (!skipfactor && !flexcg) {
|
||||
// implicit
|
||||
if (m->opt.integrator == mjINT_IMPLICIT) {
|
||||
// compute analytical derivative qDeriv
|
||||
@@ -1495,7 +1441,10 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
|
||||
}
|
||||
|
||||
// standard sparse solve
|
||||
if (m->opt.integrator == mjINT_IMPLICIT) {
|
||||
if (flexcg) {
|
||||
// constraint solver's qacc already carries the implicit flex force
|
||||
mju_copy(qacc, d->qacc, m->nv);
|
||||
} else if (m->opt.integrator == mjINT_IMPLICIT) {
|
||||
mju_solveLUSparse(qacc, d->qLU, qfrc, nv, m->D_rownnz, m->D_rowadr, m->D_diag, m->D_colind,
|
||||
dof_awake_ind);
|
||||
} else {
|
||||
@@ -1508,16 +1457,11 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
|
||||
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1, m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind);
|
||||
}
|
||||
|
||||
// flex: CG correction for implicit flex stiffness
|
||||
if (has_flex_stiffness) {
|
||||
flexInterp_cgsolve(m, d, qacc, qfrc, m->nv);
|
||||
}
|
||||
|
||||
// implicitfast: local unsymmetric solve for standalone free bodies
|
||||
// adds the bias (gyroscopic) derivative, dropped from the global symmetric solve; the
|
||||
// 6x6 block of M - h*D is decoupled from the rest of the system (D sparsity is tree-local),
|
||||
// so overwriting these rows of qacc leaves all other DOFs unaffected
|
||||
if (m->opt.integrator == mjINT_IMPLICITFAST) {
|
||||
if (m->opt.integrator == mjINT_IMPLICITFAST && !flexcg) {
|
||||
for (int j=0; j < m->njnt; j++) {
|
||||
mjtNum A[36];
|
||||
if (!mjd_freeMhat(m, d, j, m->opt.timestep, A)) {
|
||||
|
||||
Reference in New Issue
Block a user