Skip to content

Replace the flex metric factorization with a block preconditioner - #3439

Open
smallquail wants to merge 1 commit into
google-deepmind:mainfrom
smallquail:flex-metric-blocks
Open

Replace the flex metric factorization with a block preconditioner#3439
smallquail wants to merge 1 commit into
google-deepmind:mainfrom
smallquail:flex-metric-blocks

Conversation

@smallquail

@smallquail smallquail commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

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:

ms/step share
build + factorize the metric 9.64 50%
constraint solve, 39 CG iterations 9.16 48%
everything else 0.43 2%
total 19.23

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_smooth becomes 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 the qacc_smooth solve, and both apply the same operator.

metric state upstream this PR
inactive: 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
per-step stiffness, exactly factorizable M^-1, plus a per-step sparse Cholesky of the flex block M^-1, plus the per-vertex 3x3 blocks
per-step stiffness, M couples across the flex block an inner PCG of up to 50 iterations over the row above, run once per outer CG iteration M^-1, plus the per-vertex 3x3 blocks

The constant bending factor is factored at model compile time and costs nothing per step, so it stays.

Models affected

model upstream this PR
bag.xml, per-step stretch 19.23 ms/step, 39 CG 8.29 ms/step, 86 CG
poncho.xml, bending 2.96 ms/step, 16.8 CG 2.96 ms/step, 14.6 CG
poncho_edgeequality.xml, bending 2.11 ms/step, 9.8 CG 2.11 ms/step, 9.1 CG
flag.xml, trampoline.xml, metric inactive unchanged path unchanged path

Only 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^-1 and 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_effSolve is renamed mjd_effPrec, since it is a preconditioner rather than a solve, and the name mjd_effSolve now belongs to the function that solves (M + K)x = b. The qacc_smooth call site in mj_fwdAcceleration is textually unchanged but now reaches the iterative solve.
  • That solve converges on opt.tolerance rather than a hardcoded 1e-10, and raises mjWARN_INERTIA on reaching its iteration cap rather than returning an under-converged qacc_smooth.
  • Covered dofs are located by walking the covered rows of the stiffness matrix. They are not necessarily 3-aligned from dof 0: any joint declared before a flexcomp shifts them.
  • 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.

In mjUSESINGLE builds the default opt.tolerance of 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.

Comment thread model/flex/bag.xml Outdated
Comment thread src/engine/engine_derivative.c Outdated
Comment thread src/engine/engine_derivative.c Outdated
Comment thread src/engine/engine_derivative.c Outdated
Comment thread src/engine/engine_derivative.c Outdated
Comment thread src/engine/engine_derivative.c
Comment thread src/engine/engine_derivative.c Outdated
Comment thread src/engine/engine_derivative.c Outdated
Comment thread test/xml/xml_write_read_test.cc Outdated
Comment thread test/engine/engine_derivative_test.cc
@smallquail
smallquail requested a review from yuvaltassa July 29, 2026 13:06
@quagla
quagla self-requested a review July 29, 2026 13:09
@smallquail
smallquail force-pushed the flex-metric-blocks branch 2 times, most recently from 5f90404 to 55d13ae Compare July 29, 2026 13:36
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
smallquail force-pushed the flex-metric-blocks branch from 55d13ae to cfcd826 Compare July 29, 2026 14:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants