Replace the flex metric factorization with a block preconditioner

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.
This commit is contained in:
Alessio
2026-07-29 14:19:58 +01:00
parent d400914e97
commit 55d13aec5f
16 changed files with 4914 additions and 485 deletions
+70 -8
View File
@@ -2104,9 +2104,11 @@ TEST_F(DerivativeTest, FlexStiffAssembleInterp) {
}
}
// mjd_effSolve: exact-preconditioner fast path solves (M+K)x = b directly; the general
// refinement path stays within its tolerance when exactness does not hold
TEST_F(DerivativeTest, EffSolveExact) {
// mjd_effSolve drives (M+K)x = b to opt.tolerance on the relative residual,
// for every metric coverage case. It is a PCG whose preconditioner
// (mjd_effPrec) is only approximate, so the accuracy comes from the iteration
// and not from the preconditioner being exact.
TEST_F(DerivativeTest, EffSolve) {
// relative residual of (M+K)x - b after mjd_effSolve
auto solve_residual = [](const mjModel* m, mjData* d) {
int nv = m->nv;
@@ -2142,8 +2144,7 @@ TEST_F(DerivativeTest, EffSolveExact) {
ASSERT_GE(data->efm_active, 1);
EXPECT_GT(data->nefmK, 0);
EXPECT_GT(data->nefmdof, 0);
EXPECT_EQ(data->efm_active, 2);
EXPECT_LT(solve_residual(model.get(), data.get()), MjTol(1e-10, 1e-6));
EXPECT_LT(solve_residual(model.get(), data.get()), MjTol(1e-8, 1e-4));
// bending-only cloth: no CSR or per-step factor, constant factor covers, exact
static const char* const kXmlBend = R"(
@@ -2166,8 +2167,7 @@ TEST_F(DerivativeTest, EffSolveExact) {
EXPECT_EQ(data->nefmK, 0);
EXPECT_EQ(data->nefmdof, 0);
EXPECT_GT(model->nefm0dof, 0);
EXPECT_EQ(data->efm_active, 2);
EXPECT_LT(solve_residual(model.get(), data.get()), MjTol(1e-10, 1e-6));
EXPECT_LT(solve_residual(model.get(), data.get()), MjTol(1e-8, 1e-4));
// cloth under a jointed parent: M couples across the covered block, not exact,
// the refinement path must still meet its tolerance
@@ -2192,9 +2192,71 @@ TEST_F(DerivativeTest, EffSolveExact) {
data = MakeData(model);
mj_forward(model.get(), data.get());
ASSERT_GE(data->efm_active, 1);
EXPECT_EQ(data->efm_active, 1);
EXPECT_LT(solve_residual(model.get(), data.get()), 1e-4);
}
// A cloth with per-step stretch stiffness, used by the two tests below.
static const char* const kStretchCloth = R"(
<mujoco>
<option solver="CG" integrator="implicitfast"/>
<worldbody>
<body name="base" pos="0 0 1">
<joint type="slide" axis="0 0 1"/>
<geom type="sphere" size=".01" mass="1" contype="0" conaffinity="0"/>
<flexcomp name="cloth" type="grid" count="6 6 1" spacing="0.05 0.05 0.05"
radius=".005" dim="2" mass="0.5" pos="0 0 0" dof="full">
<contact selfcollide="none" contype="0" conaffinity="0"/>
<elasticity young="1e3" poisson="0.2" damping="0.1" elastic2d="both"
thickness="0.01"/>
</flexcomp>
</body>
</worldbody>
</mujoco>
)";
// An unreachable tolerance drives mjd_effSolve to its iteration cap; it must
// report that rather than return an under-converged qacc_smooth silently.
TEST_F(DerivativeTest, EffSolveCapWarns) {
char error[1024];
MjModelPtr model = LoadModelFromString(kStretchCloth, error, sizeof(error));
ASSERT_THAT(model.get(), NotNull()) << error;
model->opt.tolerance = 0; // unreachable: the PCG can never meet it
MjDataPtr data = MakeData(model);
MockWarningHandler warning_handler;
// both: the specific cause, then the mjWARN_INERTIA it is reported through
warning_handler.ExpectWarnings("Flex stiffness is too ill-conditioned");
warning_handler.ExpectWarnings("Inertia matrix is too close to singular");
mj_forward(model.get(), data.get());
testing::Mock::VerifyAndClearExpectations(&warning_handler);
}
// PCG requires a symmetric preconditioner. mjd_effPrec must satisfy
// u.P(v) == v.P(u); it did not when the covered and uncovered dofs shared a
// kinematic tree, which is what the flex-under-a-slider model here exercises.
TEST_F(DerivativeTest, EffPrecIsSymmetric) {
char error[1024];
MjModelPtr model = LoadModelFromString(kStretchCloth, error, sizeof(error));
ASSERT_THAT(model.get(), NotNull()) << error;
MjDataPtr data = MakeData(model);
mj_forward(model.get(), data.get());
ASSERT_GE(data->efm_active, 1);
int nv = model->nv;
std::vector<mjtNum> u(nv), v(nv), Pu(nv), Pv(nv);
for (int trial = 0; trial < 5; trial++) {
for (int i = 0; i < nv; i++) {
u[i] = mju_Halton(i + trial*nv, 2) - 0.5;
v[i] = mju_Halton(i + trial*nv, 5) - 0.5;
}
mjd_effPrec(model.get(), data.get(), Pu.data(), u.data());
mjd_effPrec(model.get(), data.get(), Pv.data(), v.data());
mjtNum a = mju_dot(v.data(), Pu.data(), nv);
mjtNum b = mju_dot(u.data(), Pv.data(), nv);
EXPECT_THAT(a, MjNear(b, 1e-10, 1e-4))
<< "preconditioner is not symmetric, trial " << trial;
}
}
} // namespace
} // namespace mujoco