block_cholesky / tile_cholesky lack pivot floor, leads to NaNs

[jessicayin]

Summary

With solver=newton, the constraint solver builds the Hessian

H = qM + Σ_k efc_D[k] · J[k]ᵀ J[k]

and factors it with a Cholesky (block_cholesky for nv > 32, tile_cholesky for nv ≤ 32, in mujoco_warp/_src/solver.py). Both the build and the Cholesky run in float32, with no pivot floor. Reference MuJoCo's mju_cholFactor floors each pivot with a mindiag; the Warp port dropped it.

With efc_D > 0 on every active row, H is mathematically SPD. But on stiff-contact states (large efc_D; impratio inflates the tangential rows of elliptic contacts), ‖JᵀD J‖ reaches ~1e8, so the float32 assembly of H carries an absolute error ~ ‖H‖·eps32 that swamps H's true small eigenvalues (the rigid-body dynamics modes, ~3e-4). The float32 H flips numerically indefinite, and the unguarded Cholesky takes sqrt() of a non-positive pivot → NaN, which then propagates qfrc_constraint → qacc → qpos/qvel. solver=cg (matrix-free, no Hessian factorization) is clean on the exact same state.

Tagging @ooctipus who also experienced the same issue independently.

Expected vs. actual

• Expected: a mathematically-SPD constraint Hessian factors without producing NaN, as in reference MuJoCo (whose mju_cholFactor floors pivots with mindiag).
• Actual: float32 block_cholesky / tile_cholesky return NaN; the NaN propagates into qacc/qpos/qvel and corrupts the step.

Root cause

block_cholesky.py / the tile_cholesky path do an unguarded factorization in float32. For a Hessian with a wide eigenvalue range — large ‖JᵀD J‖ (≈ efc_D_max · ‖J‖²) over small dynamics modes — the float32 assembly error (~ ‖H‖·eps32) exceeds the smallest true eigenvalue, so the stored matrix is indefinite and the factor square-roots a non-positive pivot.

Reproduction (pure numpy, self-contained)

The same H = qM + Σ efc_D·JᵀJ is SPD when assembled in float64 but indefinite when accumulated in float32, so the float32 Cholesky fails:

import numpy as np
rng = np.random.default_rng(0); nv = 35          # nv > 32 -> the blocked path

# qM: SPD mass matrix with a fragile small dynamics mode (3e-4) up to O(1)
q, _ = np.linalg.qr(rng.standard_normal((nv, nv)))
qM = (q * np.linspace(3e-4, 1.0, nv)) @ q.T

# stiff contacts: rank-1 efc_D * J Jᵀ terms, ||Jᵀ D J|| ~ 1e8
J = rng.standard_normal((12, nv)); D = np.geomspace(2e5, 2e7, 12)

H64 = qM.copy()                                  # assemble in float64
for k in range(12):
    H64 += D[k] * np.outer(J[k], J[k])

H32 = qM.astype(np.float32)                      # accumulate in float32 (what the solver does)
for k in range(12):
    jk = J[k].astype(np.float32)
    H32 = (H32 + np.float32(D[k]) * np.outer(jk, jk)).astype(np.float32)
H32 = 0.5 * (H32 + H32.T)

print("eigmin, float64 assembly   :", np.linalg.eigvalsh(H64).min())   # +9.7e-02  (SPD)
print("eigmin, float32 accumulation:", np.linalg.eigvalsh(H32).min())  # -1.7e+01  (INDEFINITE)
np.linalg.cholesky(H32)   # -> LinAlgError (numpy);  block_cholesky -> NaN (unguarded sqrt)

eigmin, float64 assembly    : +0.0973
eigmin, float32 accumulation : -16.51
numpy.linalg.LinAlgError: Matrix is not positive definite

A fuller standalone script is attached — it adds (a) a per-pivot mindiag floor that recovers the factor, and (b) a dump of a real ctx.h (48×48, nv=35) from a contact-rich state where the float32 factor NaNs (eigmin = −1.48, efc_D up to 2700, ‖JᵀD J‖ ≈ 8.6e7) alongside a clean state from the same scene (efc_D ≈ 8, eigmin = +3e-4, factors fine). Both halves run with only numpy.

solver_hessian_env0_env221.txt
repro_chol_floor_standalone.py

Proposed fix

Add a scale-relative mindiag pivot floor to block_cholesky and tile_cholesky, matching mju_cholFactor (clamp each pivot to max(pivot, mindiag) with mindiag proportional to the matrix scale, e.g. ~ eps · max_diag). A small absolute floor is insufficient — when ‖H‖ is large the floor is float32-swallowed; the clamp must scale with the matrix.

I've validated a diagonal regularization injected before the factor that eliminates the NaN at negligible cost; a per-pivot relative mindiag (as upstream MuJoCo does) is the more robust form. (A float64 assembly/factor of the constraint Hessian, as reference MuJoCo uses, would also remove the root cause.)

Environment

• mujoco-warp: observed on 3.5.0.2; block_cholesky.py is byte-identical in 3.6.0, so the gap is present there too
• mujoco: 3.5.0 / 3.6.0
• warp-lang: 1.12.0
• newton (as the calling solver): git@2684d75 (the solver=newton MJWarp path)

[thowell]

@jessicayin thank you for reporting this issue!

@adenzler-nvidia would it be possible to add support for a minimum diagonal value, similar to mujoco's factorization routine, to warp's tile_cholesky?

[adenzler-nvidia]

I'm looking into whether this should be a first-level warp feature and what that would buy us compared to just enforcing it during the matrix assembly at the moment. Will report back here.

[adenzler-nvidia]

ok so short update - it's not trivially possible to just add a pivot floor because cuSolverDx which this maps to assumes SPD matrices. Our scalar versions are unfortunately much slower but I have an agents try to improve that, so maybe there is some hope.

I think this only leaves us the option of regularizing the diagonal. That unfortunately can impact performance a bit, a few % on the humanoid benchmark. I'll go through some more options, see if I can minimize the impact.