diff --git a/python/mujoco/sysid/_src/optimize.py b/python/mujoco/sysid/_src/optimize.py index 65227355..b0e22982 100644 --- a/python/mujoco/sysid/_src/optimize.py +++ b/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) diff --git a/python/mujoco/sysid/_src/parameter.py b/python/mujoco/sysid/_src/parameter.py index b5e7d46a..ab3877d7 100644 --- a/python/mujoco/sysid/_src/parameter.py +++ b/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). diff --git a/python/mujoco/sysid/tests/test_integration.py b/python/mujoco/sysid/tests/test_integration.py index d6b7eeb7..0ecfbc6c 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,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 diff --git a/python/mujoco/sysid/tests/test_parameter.py b/python/mujoco/sysid/tests/test_parameter.py index 069a771f..9274c1b2 100644 --- a/python/mujoco/sysid/tests/test_parameter.py +++ b/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])