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/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