google-deepmind / google-deepmind/mujoco_warp

block_cholesky / tile_cholesky lack pivot floor, leads to NaNs

Open
#1,415 3 comments 0 reactions 1 assignee Claimed by @thowell View on GitHub
Dominant language
Python
Stars
1.5k
Forks
217
Avg merge
2d 11h
Merged PRs (30d)
48

Description

## 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:

```python
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](https://github.com/user-attachments/files/28656097/solver_hessian_env0_env221.txt)
[repro_chol_floor_standalone.py](https://github.com/user-attachments/files/28656101/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)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.