Replace the flex metric factorization with a block preconditioner - #3439
Open
smallquail wants to merge 1 commit into
Open
Replace the flex metric factorization with a block preconditioner#3439smallquail wants to merge 1 commit into
smallquail wants to merge 1 commit into
Conversation
smallquail
force-pushed
the
flex-metric-blocks
branch
from
July 29, 2026 08:28
7f8a0a5 to
8878cd5
Compare
yuvaltassa
reviewed
Jul 29, 2026
quagla
self-requested a review
July 29, 2026 13:09
smallquail
force-pushed
the
flex-metric-blocks
branch
2 times, most recently
from
July 29, 2026 13:36
5f90404 to
55d13ae
Compare
quagla
approved these changes
Jul 29, 2026
Every step, the flex block of the implicit effective metric M + K was factorized by sparse Cholesky, because K depends on the configuration. On model/flex/bag.xml, added here, that is roughly half the step, against a comparable share for the constraint solve it exists to accelerate. Keep only the metric's per-vertex 3x3 diagonal blocks, prefactored. Neither consumer needs the exact inverse: the CG constraint solver only wants a preconditioner, and qacc_smooth can come from an iterative solve using those blocks. They are O(n) to build and to apply, but weaker, so CG runs about twice the iterations and qacc_smooth becomes an iteration rather than a direct solve. Net, the bag model steps roughly twice as fast. The preconditioner, by metric state. Inactive, meaning no flex elasticity or an explicit integrator: M^-1, unchanged. Bending only (nefmK == 0): M^-1 plus the exact constant bending factor from mj_setConst on the dofs it covers, unchanged; that factor is built at model compile time and costs nothing per step. Per-step stiffness: M^-1 plus the 3x3 blocks, where before it was a per-step sparse Cholesky, or, when M couples across the flex block, an inner PCG of up to 50 iterations run once per outer CG iteration. Only models carrying per-step stretch stiffness change in wall-clock. Both ponchos hold their timing and take slightly fewer CG iterations than before, because the preconditioner is now symmetric: it applies M^-1 and the covered blocks to disjoint sets of dofs, where previously the two overlapped and the operator was not symmetric, which PCG requires. mjd_effSolve is the accurate solve of (M + K)x = b; what used to carry that name only preconditions and is now mjd_effPrec. Its CG guarded the division by pAp with mjMINVAL, an absolute floor on a quantity that scales with the square of the right-hand side, so a small b aborted the solve while the curvature was healthy: four flex models were quietly left short of tolerance. For an SPD metric the guard is positivity, and with that the same solves converge. The qacc_smooth call site in mj_fwdAcceleration is textually unchanged but now reaches the iterative solve, which converges on opt.tolerance rather than a hardcoded threshold, floored in mjUSESINGLE builds where the squared target is unreachable in float. Reaching the iteration cap names the ill-conditioned flex stiffness and then reports it through mjWARN_INERTIA, rather than returning an under-converged result. Covered dofs are located by walking the covered rows of the stiffness matrix, as they need not be 3-aligned from dof 0: any joint declared before a flexcomp shifts them. mjData.efm_L_rownnz, efm_L_rowadr and efm_L_colind described the sparsity of the deleted factorization and are removed: left NULL with nonzero mjxmacro extents they made the Python bindings hand back uninitialized arrays. efm_active loses the value 2 for the same reason, nothing selects a solve path on preconditioner exactness any more. Both are recorded under breaking changes. model/flex/bag.xml is added because no shipped model carried per-step stretch stiffness. The ponchos are bending-only and trampoline.xml uses an explicit integrator, so the metric never activates there. It is excluded from WriteReadCompareTest: stretch stiffness amplifies rest geometry that XML rounds on save.
smallquail
force-pushed
the
flex-metric-blocks
branch
from
July 29, 2026 14:30
55d13ae to
cfcd826
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Every step, the flex block of the effective metric M + K was factorized by sparse Cholesky, because K depends on the configuration. This replaces that factorization with the metric's per-vertex 3x3 diagonal blocks, prefactored.
Where the time goes
model/flex/bag.xml(added here, 2868 dof), on upstream:The factorization goes away entirely. The blocks are O(n) to build and to apply, but weaker: CG runs about twice the iterations, and
qacc_smoothbecomes an iteration rather than a direct solve. Net 19.23 -> 8.29 ms/step.Preconditioner, by case
The metric has two consumers, the CG constraint solver (
Mgrad = Mtilde \ grad) and theqacc_smoothsolve, and both apply the same operator.M^-1nefmK == 0M^-1, plus the exact constant bending factor frommj_setConston the dofs it coversM^-1, plus a per-step sparse Cholesky of the flex blockM^-1, plus the per-vertex 3x3 blocksM^-1, plus the per-vertex 3x3 blocksThe constant bending factor is factored at model compile time and costs nothing per step, so it stays.
Models affected
bag.xml, per-step stretchponcho.xml, bendingponcho_edgeequality.xml, bendingflag.xml,trampoline.xml, metric inactiveOnly models carrying per-step stretch stiffness change in wall-clock. Bending-only models hold their timing and take slightly fewer CG iterations, because the preconditioner is now symmetric: it applies
M^-1and the covered blocks to disjoint sets of dofs, where before the two overlapped and the operator was not symmetric, which PCG requires.Also in this change
mjd_effSolveis renamedmjd_effPrec, since it is a preconditioner rather than a solve, and the namemjd_effSolvenow belongs to the function that solves (M + K)x = b. Theqacc_smoothcall site inmj_fwdAccelerationis textually unchanged but now reaches the iterative solve.opt.tolerancerather than a hardcoded 1e-10, and raisesmjWARN_INERTIAon reaching its iteration cap rather than returning an under-convergedqacc_smooth.model/flex/bag.xmlis added because no shipped model carried per-step stretch stiffness. The ponchos are bending-only, andtrampoline.xmluses an explicit integrator, so the metric never activates there.In
mjUSESINGLEbuilds the defaultopt.toleranceof 1e-8 is below float epsilon, so the solve cannot reach it and would warn on every flex model. A floor on the tolerance would address that; it is not included here.