From ab4102ea4252fda03d68e7d91a988aa0e512be23 Mon Sep 17 00:00:00 2001
From: Kevin Zakka <kevinarmandzakka@gmail.com>
Date: Mon, 25 May 2026 20:50:06 -0700
Subject: [PATCH] Add opt-in cond(J^T J) check to sysid optimize().

Pass check_conditioning=True to FD the Jacobian at the starting point
and warn if cond(J^T J) suggests numerical ill-conditioning.
---
 python/mujoco/sysid/_src/optimize.py          | 49 +++++++++++++++++++
 python/mujoco/sysid/tests/test_integration.py | 40 +++++++++++++++
 2 files changed, 89 insertions(+)

diff --git a/python/mujoco/sysid/_src/optimize.py b/python/mujoco/sysid/_src/optimize.py
index ed91a944..860e710b 100644
--- a/python/mujoco/sysid/_src/optimize.py
+++ b/python/mujoco/sysid/_src/optimize.py
@@ -29,6 +29,48 @@ import scipy.special
 XScale = Literal["jac"] | np.ndarray | float
 
 
+def _warn_if_ill_conditioned(
+    initial_params: parameter.ParameterDict,
+    residual_fn: Callable[..., Any],
+    threshold: float = 1e12,
+) -> None:
+  """Warn if cond(JᵀJ) at the starting point exceeds ``threshold``.
+
+  Costs one extra finite-difference Jacobian. ``threshold=1e12`` corresponds
+  to cond(J) ~ 1e6, well below the float64 limit (~1e16).
+  """
+  x0 = initial_params.as_vector()
+  bounds = initial_params.get_bounds()
+
+  def f(x):
+    residuals, _, _ = residual_fn(x, initial_params)
+    return np.concatenate(residuals)
+
+  eps = np.finfo(np.float64).eps ** 0.5
+  r0 = f(x0).reshape(-1, 1)
+  jac = np.asarray(
+      mujoco_minimize.jacobian_fd(
+          residual=f,
+          x=x0.reshape(-1, 1),
+          r=r0,
+          eps=eps,
+          n_res=0,
+          bounds=[bounds[0].reshape(-1, 1), bounds[1].reshape(-1, 1)],
+      )[0],
+      dtype=np.float64,
+  )
+  # eigvalsh on the gram matrix so rank-deficient directions are visible
+  # when n_params > n_residual components (SVD would drop to min(m, n)).
+  ev = np.maximum(np.linalg.eigvalsh(jac.T @ jac), 0.0)
+  cond_jtj = float(ev[-1] / ev[0]) if ev[0] > 0 else float("inf")
+  if cond_jtj > threshold:
+    logging.warning(
+        "cond(JᵀJ) ≈ %.1e at the starting point; the problem may be "
+        "ill-conditioned. Consider x_scale='jac' or regularizing.",
+        cond_jtj,
+    )
+
+
 def _scipy_least_squares(
     x0: np.ndarray,
     residual_fn: Callable[..., Any],
@@ -149,6 +191,7 @@ def optimize(
     residual_fn: Callable[..., Any],
     optimizer: Literal["scipy", "mujoco", "scipy_parallel_fd"] = "mujoco",
     verbose: bool = True,
+    check_conditioning: bool = False,
     **optimizer_kwargs,
 ) -> tuple[parameter.ParameterDict, scipy_optimize.OptimizeResult]:
   """Run nonlinear least-squares optimization on the residual.
@@ -160,6 +203,9 @@ def optimize(
     optimizer: Backend — ``"mujoco"`` (default), ``"scipy"``, or
       ``"scipy_parallel_fd"`` (scipy with MuJoCo finite-difference Jacobian).
     verbose: If True, log parameter comparison table after optimization.
+    check_conditioning: If True, estimate ``cond(JᵀJ)`` at the starting
+      point and emit a warning if it suggests numerical ill-conditioning.
+      Costs one extra finite-difference Jacobian.
     **optimizer_kwargs: Forwarded to the backend. Common ones:
 
       * ``max_iters``: maximum number of optimizer iterations.
@@ -194,6 +240,9 @@ def optimize(
         extras={},
     )
 
+  if check_conditioning:
+    _warn_if_ill_conditioned(initial_params, residual_fn)
+
   # Warn if any non-frozen parameter component starts at (or essentially at)
   # a box bound. Optimization can stall in that corner on ill-conditioned or
   # rank-deficient problems; both the mujoco and scipy backends are affected.
diff --git a/python/mujoco/sysid/tests/test_integration.py b/python/mujoco/sysid/tests/test_integration.py
index d6b7eeb7..ff2688f0 100644
--- a/python/mujoco/sysid/tests/test_integration.py
+++ b/python/mujoco/sysid/tests/test_integration.py
@@ -14,6 +14,7 @@
 # ==============================================================================
 """End-to-end integration tests for mujoco.sysid."""
 
+import logging
 import pathlib
 import tempfile
 
@@ -21,6 +22,7 @@ import mujoco
 import mujoco.rollout as rollout
 from mujoco import sysid
 import numpy as np
+import pytest
 
 
 # ---------------------------------------------------------------------------
@@ -233,3 +235,41 @@ def test_arm_recover_armature():
     assert (result_dir / "results.pkl").exists()
     assert (result_dir / "confidence.pkl").exists()
     assert (result_dir / "arm.xml").exists()
+
+
+def _rank_1_residual_fn(x, p):
+  del p
+  if x.ndim == 1:
+    r = np.array([x[0] + x[1] - 2.0])
+  else:
+    r = (x[0] + x[1] - 2.0).reshape(1, -1)
+  return [r], None, None
+
+
+def _full_rank_residual_fn(x, p):
+  del p
+  if x.ndim == 1:
+    r = np.array([x[0] - 1.0, x[1] - 2.0])
+  else:
+    r = np.stack([x[0] - 1.0, x[1] - 2.0])
+  return [r], None, None
+
+
+@pytest.mark.parametrize("residual_fn,expect_warning", [
+    (_rank_1_residual_fn, True),
+    (_full_rank_residual_fn, False),
+])
+def test_check_conditioning(residual_fn, expect_warning, caplog):
+  """check_conditioning=True warns iff cond(JᵀJ) is large at the starting point."""
+  params = sysid.ParameterDict()
+  params.add(sysid.Parameter("a", 1.0, -10.0, 10.0))
+  params.add(sysid.Parameter("b", 1.0, -10.0, 10.0))
+
+  with caplog.at_level(logging.WARNING, logger="absl"):
+    sysid.optimize(
+        initial_params=params, residual_fn=residual_fn,
+        optimizer="scipy", verbose=False, check_conditioning=True,
+        max_iters=1,
+    )
+  fired = any("cond(JᵀJ)" in r.message for r in caplog.records)
+  assert fired is expect_warning