Merge pull request #3287 from kevinzakka:sysid-issue-3284-fixes

PiperOrigin-RevId: 927249817
Change-Id: I63db037d24238dd0ce82ffe157a2e1bb33c9314a

python/mujoco/sysid/_src/optimize.py

@@ -26,6 +26,48 @@ import scipy.optimize as scipy_optimize
 import scipy.special
 
 
+def _warn_if_ill_conditioned(
+    initial_params: parameter.ParameterDict,
+    residual_fn: Callable[..., Any],
+    threshold: float = 1e12,
+    eps: float | None = None,
+) -> None:
+  """Warn if cond(J^T 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). ``eps`` defaults to
+  the finite-difference step used by the backends.
+  """
+  x0 = initial_params.as_vector()
+  bounds = initial_params.get_bounds()
+
+  def f(x):
+    residuals, _, _ = residual_fn(x, initial_params)
+    return np.concatenate(residuals)
+
+  if eps is None:
+    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,
+  )
+  cond_jtj = float(np.linalg.cond(jac)) ** 2
+  if cond_jtj > threshold:
+    logging.warning(
+        "cond(J^T 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],
@@ -39,7 +81,6 @@ def _scipy_least_squares(
     verbose = 2
   else:
     verbose = 0
-  x_scale = kwargs.pop("x_scale", "jac")
   loss = kwargs.pop("loss", "linear")
 
   jac_arg: str | Callable[..., Any]
@@ -72,7 +113,6 @@ def _scipy_least_squares(
       bounds=bounds,
       max_nfev=max_nfev,
       verbose=verbose,
-      x_scale=x_scale,
       loss=loss,
       jac=jac_arg,  # pyright: ignore[reportArgumentType]
       **kwargs,
@@ -85,12 +125,17 @@ def _mujoco_least_squares(
     bounds: tuple[np.ndarray, np.ndarray],
     **kwargs,
 ) -> scipy_optimize.OptimizeResult:
-  """Run MuJoCo's native least_squares optimizer."""
+  """Run MuJoCo's native least_squares optimizer.
+
+  ``**kwargs`` are forwarded to :func:`mujoco.minimize.least_squares`; see
+  its docstring (notably ``x_scale``).
+  """
   if kwargs.pop("verbose", True):
     verbose = mujoco_minimize.Verbosity.FULLITER
   else:
     verbose = mujoco_minimize.Verbosity.SILENT
   max_iter = kwargs.pop("max_iters", 200)
+
   x, log = mujoco_minimize.least_squares(
       x0=x0,
       bounds=bounds,
@@ -100,7 +145,6 @@ def _mujoco_least_squares(
       **kwargs,
   )
 
-  # If verbose, return the full optimization log.
   extras = {}
   if verbose == mujoco_minimize.Verbosity.FULLITER:
     extras["objective"] = [entry.objective for entry in log]
@@ -144,6 +188,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.
@@ -155,8 +200,18 @@ 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.
-    **optimizer_kwargs: Forwarded to the backend (e.g. ``max_iters``,
-      ``verbose``, ``loss``).
+    check_conditioning: If True, estimate ``cond(J^T 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.
+      * ``verbose``: per-backend verbosity flag (separate from this
+        function's ``verbose``).
+      * ``loss``: scipy loss function name (scipy backends only).
+      * ``x_scale``: per-parameter scaling, forwarded to the backend; see
+        :func:`scipy.optimize.least_squares` and
+        :func:`mujoco.minimize.least_squares`.
 
   Returns:
     ``(opt_params, opt_result)`` — the optimized ParameterDict and a
@@ -179,6 +234,30 @@ def optimize(
         extras={},
     )
 
+  if check_conditioning:
+    _warn_if_ill_conditioned(
+        initial_params, residual_fn, eps=optimizer_kwargs.get("diff_step"),
+    )
+
+  # Warn if any non-frozen parameter component starts at (or essentially at)
+  # a box bound. Only meaningful when x0 is feasible; otherwise the user is
+  # already outside the constraint set and the optimizer will clip first.
+  lo, hi = bounds
+  rng = hi - lo
+  safe_rng = np.where(rng > 0, rng, 1.0)
+  in_bounds = (x0 >= lo) & (x0 <= hi)
+  near_bound = ((x0 - lo) <= 1e-3 * safe_rng) | ((hi - x0) <= 1e-3 * safe_rng)
+  at_bound = in_bounds & near_bound & (rng > 0)
+  if at_bound.any():
+    logging.warning(
+        "%d of %d non-frozen parameter components start at (or essentially "
+        "at) a box bound. Optimization can stall in that corner on "
+        "ill-conditioned problems; consider calling "
+        "initial_params.move_off_bounds() before optimize().",
+        int(at_bound.sum()),
+        x0.size,
+    )
+
   def optimized_residual_fn(x):
     residuals, _, _ = residual_fn(x, opt_params)
     return np.concatenate(residuals)

python/mujoco/sysid/_src/parameter.py

@@ -139,6 +139,22 @@ class Parameter:
       rng = np.random.default_rng()
     return rng.uniform(self.min_value.flatten(), self.max_value.flatten())
 
+  def move_off_bound(
+      self, fraction: float = 0.05, at_bound_tol: float = 1e-3
+  ) -> None:
+    """Shift values within ``at_bound_tol * (hi - lo)`` of a bound to
+    ``lo + fraction*(hi-lo)`` (or symmetric for the upper bound). Interior
+    components are unchanged."""
+    lo, hi = self.get_bounds()
+    rng = hi - lo
+    safe_rng = np.where(rng > 0, rng, 1.0)
+    v = self.as_vector().copy()
+    at_lo = (v - lo) <= at_bound_tol * safe_rng
+    at_hi = (hi - v) <= at_bound_tol * safe_rng
+    v = np.where(at_lo, lo + fraction * rng, v)
+    v = np.where(at_hi, hi - fraction * rng, v)
+    self.update_from_vector(v)
+
   def __str__(self) -> str:
     """Return a string representation of the parameter."""
     if self.size == 1:
@@ -293,6 +309,17 @@ class ParameterDict:
         param.update_from_vector(vector[start : start + size])
         start += size
 
+  def move_off_bounds(
+      self, fraction: float = 0.05, at_bound_tol: float = 1e-3
+  ) -> Self:
+    """Call :meth:`Parameter.move_off_bound` on every non-frozen parameter,
+    returning ``self`` so calls can be chained before
+    :func:`optimize`."""
+    for param in self.parameters.values():
+      if not param.frozen:
+        param.move_off_bound(fraction=fraction, at_bound_tol=at_bound_tol)
+    return self
+
   def save_to_disk(self, path: str | pathlib.Path) -> None:
     """Save the parameter dictionary to disk (schema and data).
 

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,42 @@ 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):
+  # Residual depends only on x[0]: J = [[1, 0], [2, 0]] is 2x2 rank 1.
+  del p
+  if x.ndim == 1:
+    r = np.array([x[0] - 1.0, 2.0 * x[0] - 2.0])
+  else:
+    r = np.stack([x[0] - 1.0, 2.0 * x[0] - 2.0])
+  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^T J) is large at the start."""
+  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^T J)" in r.message for r in caplog.records)
+  assert fired is expect_warning

python/mujoco/sysid/tests/test_parameter.py

@@ -146,3 +146,38 @@ def test_frozen_param_excluded():
   np.testing.assert_array_equal(params["free"].value, [1.5])
   # Frozen param unchanged.
   np.testing.assert_array_equal(params["frozen"].value, [5.0])
+
+
+def test_move_off_bound():
+  """At-bound (or essentially at-bound) values shift inward; interior is untouched."""
+  # At lower bound, at upper bound, essentially-at-lower, interior, vector mixed,
+  # custom fraction, and degenerate (zero-range) bounds.
+  cases = [
+      (0.0, 0.0, 20.0, 0.05, 1.0),    # at lower -> 0.05 * 20
+      (1e-8, 0.0, 20.0, 0.05, 1.0),   # essentially at lower
+      (20.0, 0.0, 20.0, 0.05, 19.0),  # at upper -> 20 - 0.05 * 20
+      (5.0, 0.0, 20.0, 0.05, 5.0),    # interior unchanged
+      (0.0, 0.0, 100.0, 0.1, 10.0),   # custom fraction
+      (3.0, 3.0, 3.0, 0.05, 3.0),     # zero-range bound, pinned
+  ]
+  for nominal, lo, hi, fraction, expected in cases:
+    p = parameter.Parameter("d", nominal, lo, hi)
+    p.move_off_bound(fraction=fraction)
+    np.testing.assert_allclose(p.value, [expected])
+
+  # Vector parameter: only at-bound components are shifted.
+  p = parameter.Parameter(
+      "v", [0.0, 5.0, 10.0], [0.0, 0.0, 0.0], [10.0, 10.0, 10.0]
+  )
+  p.move_off_bound()
+  np.testing.assert_allclose(p.value, [0.5, 5.0, 9.5])
+
+
+def test_move_off_bounds_dict_skips_frozen_and_returns_self():
+  """ParameterDict shifts free params, leaves frozen alone, returns self."""
+  pdict = parameter.ParameterDict()
+  pdict.add(parameter.Parameter("free", 0.0, 0.0, 1.0))
+  pdict.add(parameter.Parameter("frozen", 0.0, 0.0, 1.0, frozen=True))
+  assert pdict.move_off_bounds() is pdict
+  np.testing.assert_allclose(pdict["free"].value, [0.05])
+  np.testing.assert_allclose(pdict["frozen"].value, [0.0])