From ea230a950c4fcaefb22abcac3ff8b4c6685abec9 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Wed, 15 Jul 2026 14:56:57 -0700 Subject: [PATCH] Implicit flex elasticity in the CG constraint solver via an effective metric MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- doc/includes/references.h | 24 + include/mujoco/mjdata.h | 17 + include/mujoco/mjmodel.h | 7 + include/mujoco/mjxmacro.h | 28 +- model/flex/bunny.xml | 2 +- model/flex/bunny_multicell.xml | 2 +- model/flex/bunny_quadratic.xml | 2 +- model/flex/bunny_with_uv.xml | 96 +-- model/flex/jelly.xml | 2 +- model/flex/pancake.xml | 2 +- model/flex/plate.xml | 2 +- model/flex/quadratic.xml | 2 +- model/flex/sphere_trilinear.xml | 2 +- model/flex/trampoline.xml | 2 +- model/flex/trilinear.xml | 2 +- python/mujoco/functions.cc | 1 + python/mujoco/introspect/structs.py | 158 ++++ python/mujoco/structs.cc | 1 + src/engine/engine_derivative.c | 1138 +++++++++++++++++++++++++ src/engine/engine_derivative.h | 33 + src/engine/engine_forward.c | 186 ++-- src/engine/engine_forward.h | 4 + src/engine/engine_inverse.c | 13 + src/engine/engine_io.c | 14 +- src/engine/engine_io.h | 3 +- src/engine/engine_island.c | 41 + src/engine/engine_setconst.c | 164 ++++ src/engine/engine_solver.c | 50 +- src/engine/engine_util_solve.c | 22 +- src/engine/engine_util_solve.h | 1 + src/user/user_model.cc | 84 +- src/user/user_model.h | 2 + test/engine/engine_derivative_test.cc | 372 ++++++++ test/engine/engine_forward_test.cc | 57 +- test/engine/engine_island_test.cc | 33 + unity/Runtime/Bindings/MjBindings.cs | 22 + wasm/codegen/generated/bindings.cc | 22 + wasm/codegen/generated/bindings.h | 84 ++ wasm/codegen/generators/constants.py | 7 + 39 files changed, 2501 insertions(+), 203 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index 52323950..324ad35a 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -111,6 +111,10 @@ typedef struct mjData_ { int nl; // number of limit constraints int nefc; // number of constraints int nJ; // number of non-zeros in constraint Jacobian + int efm_active; // implicit effective metric M+K: 0 inactive, 1 active, 2 active + preconditioner exact + int nefmK; // number of non-zeros in effective-stiffness CSR + int nefmdof; // number of rows in effective-metric factor + int nefmL; // number of non-zeros in the effective-metric factor int nY; // number of non-zeros in constraint inverse inertia square root int nA; // number of non-zeros in constraint inverse inertia matrix int nisland; // number of detected constraint islands @@ -200,6 +204,7 @@ typedef struct mjData_ { // computed by mj_fwdPosition/mj_flex mjtNum* flexvert_xpos; // Cartesian flex vertex positions (nflexvert x 3) mjtNum* flexelem_aabb; // flex element bounding boxes (center, size) (nflexelem x 6) + mjtNum* flexelem_krot; // corotated element stiffness (implicit only) (nflexstiffness x 1) mjtNum* flexedge_J; // flex edge Jacobian (nJfe x 1) mjtNum* flexedge_length; // flex edge lengths (nflexedge x 1) mjtNum* flexvert_J; // flex vertex Jacobian (nJfv x 2) @@ -369,6 +374,18 @@ typedef struct mjData_ { mjtNum* efc_vel; // velocity in constraint space: J*qvel (nefc x 1) mjtNum* efc_aref; // reference pseudo-acceleration (nefc x 1) + // computed by mj_fwdPosition/mj_invPosition when the implicit effective metric M+K is active + mjtNum* efm_c; // smooth-force shift h*K*qvel (nv x 1) + int* efm_K_rownnz; // effective-stiffness CSR row nonzeros (nv x 1) + int* efm_K_rowadr; // effective-stiffness CSR row addresses (nv x 1) + int* efm_K_colind; // effective-stiffness CSR column indices (nefmK x 1) + mjtNum* efm_K_val; // effective-stiffness CSR values (nefmK x 1) + int* efm_dofid; // factor row -> dof address (nefmdof x 1) + int* efm_L_rownnz; // factor row nonzeros (nefmdof x 1) + int* efm_L_rowadr; // factor row addresses (nefmdof x 1) + int* efm_L_colind; // factor column indices (nefmL x 1) + mjtNum* efm_L; // Cholesky factor of diag(M)+K, covered dofs (nefmL x 1) + //-------------------- arena-allocated: POSITION, VELOCITY, CONTROL/ACCELERATION dependent // computed by mj_fwdConstraint/mj_inverse @@ -584,6 +601,8 @@ typedef struct mjModel_ { mjtSize nflexelemdata; // number of element vertex ids in all flexes mjtSize nflexstiffness; // number of stiffness parameters in all flexes mjtSize nflexbending; // number of bending parameters in all flexes + mjtSize nefm0dof; // number of dofs covered by the constant metric factor + mjtSize nefm0L; // number of non-zeros in the constant metric factor mjtSize nflexelemedge; // number of element edge ids in all flexes mjtSize nflexshelldata; // number of shell fragment vertex ids in all flexes mjtSize nflexevpair; // number of element-vertex pairs in all flexes @@ -902,6 +921,11 @@ typedef struct mjModel_ { mjtNum* flex_size; // vertex bounding box half sizes in qpos0 (nflex x 3) mjtNum* flex_stiffness; // finite element stiffness matrix (nflexstiffness x 1) mjtNum* flex_bending; // bending stiffness (nflexbending x 1) + int* efm0_dofid; // constant metric factor row->dof address (nefm0dof x 1) + int* efm0_L_rownnz; // constant metric factor row nonzeros (nefm0dof x 1) + int* efm0_L_rowadr; // constant metric factor row addresses (nefm0dof x 1) + int* efm0_L_colind; // constant metric factor column indices (nefm0L x 1) + mjtNum* efm0_L; // factor of M + (dt^2+dt*d)*K_bend (nefm0L x 1) mjtNum* flex_damping; // Rayleigh's damping coefficient (nflex x 1) mjtNum* flex_edgestiffness; // edge stiffness (nflex x 1) mjtNum* flex_edgedamping; // edge damping (nflex x 1) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 15afe8b2..621b1129 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -135,6 +135,10 @@ typedef struct mjData_ { int nl; // number of limit constraints int nefc; // number of constraints int nJ; // number of non-zeros in constraint Jacobian + int efm_active; // implicit effective metric M+K: 0 inactive, 1 active, 2 active + preconditioner exact + int nefmK; // number of non-zeros in effective-stiffness CSR + int nefmdof; // number of rows in effective-metric factor + int nefmL; // number of non-zeros in the effective-metric factor int nY; // number of non-zeros in constraint inverse inertia square root int nA; // number of non-zeros in constraint inverse inertia matrix int nisland; // number of detected constraint islands @@ -224,6 +228,7 @@ typedef struct mjData_ { // computed by mj_fwdPosition/mj_flex mjtNum* flexvert_xpos; // Cartesian flex vertex positions (nflexvert x 3) mjtNum* flexelem_aabb; // flex element bounding boxes (center, size) (nflexelem x 6) + mjtNum* flexelem_krot; // corotated element stiffness (implicit only) (nflexstiffness x 1) mjtNum* flexedge_J; // flex edge Jacobian (nJfe x 1) mjtNum* flexedge_length; // flex edge lengths (nflexedge x 1) mjtNum* flexvert_J; // flex vertex Jacobian (nJfv x 2) @@ -393,6 +398,18 @@ typedef struct mjData_ { mjtNum* efc_vel; // velocity in constraint space: J*qvel (nefc x 1) mjtNum* efc_aref; // reference pseudo-acceleration (nefc x 1) + // computed by mj_fwdPosition/mj_invPosition when the implicit effective metric M+K is active + mjtNum* efm_c; // smooth-force shift h*K*qvel (nv x 1) + int* efm_K_rownnz; // effective-stiffness CSR row nonzeros (nv x 1) + int* efm_K_rowadr; // effective-stiffness CSR row addresses (nv x 1) + int* efm_K_colind; // effective-stiffness CSR column indices (nefmK x 1) + mjtNum* efm_K_val; // effective-stiffness CSR values (nefmK x 1) + int* efm_dofid; // factor row -> dof address (nefmdof x 1) + int* efm_L_rownnz; // factor row nonzeros (nefmdof x 1) + int* efm_L_rowadr; // factor row addresses (nefmdof x 1) + int* efm_L_colind; // factor column indices (nefmL x 1) + mjtNum* efm_L; // Cholesky factor of diag(M)+K, covered dofs (nefmL x 1) + //-------------------- arena-allocated: POSITION, VELOCITY, CONTROL/ACCELERATION dependent // computed by mj_fwdConstraint/mj_inverse diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index a27b6bf0..48eadb4a 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -272,6 +272,8 @@ typedef struct mjModel_ { mjtSize nflexelemdata; // number of element vertex ids in all flexes mjtSize nflexstiffness; // number of stiffness parameters in all flexes mjtSize nflexbending; // number of bending parameters in all flexes + mjtSize nefm0dof; // number of dofs covered by the constant metric factor + mjtSize nefm0L; // number of non-zeros in the constant metric factor mjtSize nflexelemedge; // number of element edge ids in all flexes mjtSize nflexshelldata; // number of shell fragment vertex ids in all flexes mjtSize nflexevpair; // number of element-vertex pairs in all flexes @@ -590,6 +592,11 @@ typedef struct mjModel_ { mjtNum* flex_size; // vertex bounding box half sizes in qpos0 (nflex x 3) mjtNum* flex_stiffness; // finite element stiffness matrix (nflexstiffness x 1) mjtNum* flex_bending; // bending stiffness (nflexbending x 1) + int* efm0_dofid; // constant metric factor row->dof address (nefm0dof x 1) + int* efm0_L_rownnz; // constant metric factor row nonzeros (nefm0dof x 1) + int* efm0_L_rowadr; // constant metric factor row addresses (nefm0dof x 1) + int* efm0_L_colind; // constant metric factor column indices (nefm0L x 1) + mjtNum* efm0_L; // factor of M + (dt^2+dt*d)*K_bend (nefm0L x 1) mjtNum* flex_damping; // Rayleigh's damping coefficient (nflex x 1) mjtNum* flex_edgestiffness; // edge stiffness (nflex x 1) mjtNum* flex_edgedamping; // edge damping (nflex x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 845c9ffd..fc78487a 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -189,6 +189,8 @@ X( nflexelemdata ) \ X( nflexstiffness ) \ X( nflexbending ) \ + X( nefm0dof ) \ + X( nefm0L ) \ X( nflexelemedge ) \ X( nflexshelldata ) \ X( nflexevpair ) \ @@ -501,6 +503,11 @@ X ( mjtNum, flex_size, nflex, 3 ) \ X ( mjtNum, flex_stiffness, nflexstiffness, 1 ) \ X ( mjtNum, flex_bending, nflexbending, 1 ) \ + X ( int, efm0_dofid, nefm0dof, 1 ) \ + X ( int, efm0_L_rownnz, nefm0dof, 1 ) \ + X ( int, efm0_L_rowadr, nefm0dof, 1 ) \ + X ( int, efm0_L_colind, nefm0L, 1 ) \ + X ( mjtNum, efm0_L, nefm0L, 1 ) \ X ( mjtNum, flex_damping, nflex, 1 ) \ X ( mjtNum, flex_edgestiffness, nflex, 1 ) \ X ( mjtNum, flex_edgedamping, nflex, 1 ) \ @@ -866,6 +873,7 @@ X ( mjtNum, cinert, nbody, 10 ) \ X ( mjtNum, flexvert_xpos, nflexvert, 3 ) \ X ( mjtNum, flexelem_aabb, nflexelem, 6 ) \ + X ( mjtNum, flexelem_krot, nflexstiffness, 1 ) \ X ( mjtNum, flexedge_J, nJfe, 1 ) \ X ( mjtNum, flexedge_length, nflexedge, 1 ) \ X ( mjtNum, flexvert_J, nJfv, 2 ) \ @@ -995,11 +1003,25 @@ X ( mjtNum, ifrc_constraint, MJ_D(nidof), 1 ) // array fields of mjData that live in d->arena +#define MJDATA_ARENA_POINTERS_EFM \ + X ( mjtNum, efm_c, MJ_M(nv), 1 ) \ + X ( int, efm_K_rownnz, MJ_M(nv), 1 ) \ + X ( int, efm_K_rowadr, MJ_M(nv), 1 ) \ + X ( int, efm_K_colind, MJ_D(nefmK), 1 ) \ + X ( mjtNum, efm_K_val, MJ_D(nefmK), 1 ) \ + X ( int, efm_dofid, MJ_D(nefmdof), 1 ) \ + X ( int, efm_L_rownnz, MJ_D(nefmdof), 1 ) \ + X ( int, efm_L_rowadr, MJ_D(nefmdof), 1 ) \ + X ( int, efm_L_colind, MJ_D(nefmL), 1 ) \ + X ( mjtNum, efm_L, MJ_D(nefmL), 1 ) + + #define MJDATA_ARENA_POINTERS \ MJDATA_ARENA_POINTERS_CONTACT \ MJDATA_ARENA_POINTERS_SOLVER \ MJDATA_ARENA_POINTERS_DUAL \ - MJDATA_ARENA_POINTERS_ISLAND + MJDATA_ARENA_POINTERS_ISLAND \ + MJDATA_ARENA_POINTERS_EFM // scalar fields of mjData @@ -1021,6 +1043,10 @@ X( int, nl ) \ X( int, nefc ) \ X( int, nJ ) \ + X( int, efm_active ) \ + X( int, nefmK ) \ + X( int, nefmdof ) \ + X( int, nefmL ) \ X( int, nY ) \ X( int, nA ) \ X( int, nisland ) \ diff --git a/model/flex/bunny.xml b/model/flex/bunny.xml index 0d10cad9..3160530f 100644 --- a/model/flex/bunny.xml +++ b/model/flex/bunny.xml @@ -16,7 +16,7 @@ -