chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Address reviewer feedback.

Browse Source
This commit is contained in:
Kevin Zakka
2026-05-27 13:07:35 -07:00
parent ab4102ea42
commit 2b9960a617
3 changed files with 42 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/sysid/_src/optimize.py
+25 -28
View File
@@ -26,18 +26,17 @@ import scipy.optimize as scipy_optimize
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,
eps: float | None = None,
) -> None:
"""Warn if cond(JᵀJ) at the starting point exceeds ``threshold``.
"""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).
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()
@@ -46,7 +45,8 @@ def _warn_if_ill_conditioned(
residuals, _, _ = residual_fn(x, initial_params)
return np.concatenate(residuals)
eps = np.finfo(np.float64).eps ** 0.5
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(
@@ -59,13 +59,10 @@ def _warn_if_ill_conditioned(
)[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")
cond_jtj = float(np.linalg.cond(jac)) ** 2
if cond_jtj > threshold:
logging.warning(
"cond(JᵀJ) ≈ %.1e at the starting point; the problem may be "
"cond(J^T J) ~ %.1e at the starting point; the problem may be "
"ill-conditioned. Consider x_scale='jac' or regularizing.",
cond_jtj,
)
@@ -84,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]
@@ -117,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,
@@ -128,10 +123,13 @@ def _mujoco_least_squares(
x0: np.ndarray,
residual_fn: Callable[..., Any],
bounds: tuple[np.ndarray, np.ndarray],
x_scale: XScale = 1.0,
**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:
@@ -144,7 +142,6 @@ def _mujoco_least_squares(
residual=residual_fn,
verbose=verbose,
max_iter=max_iter,
x_scale=x_scale,
**kwargs,
)
@@ -203,7 +200,7 @@ 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
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:
@@ -212,12 +209,9 @@ def optimize(
* ``verbose``: per-backend verbosity flag (separate from this
function's ``verbose``).
* ``loss``: scipy loss function name (scipy backends only).
* ``x_scale``: per-parameter scaling. ``"jac"`` is adaptive
``D_i = 1/||J(:,i)||`` per iteration; an explicit array or
positive scalar is used as ``D`` directly. Defaults: ``"jac"``
for scipy backends, ``1.0`` (no scaling) for the mujoco backend.
See :func:`scipy.optimize.least_squares` and
:func:`mujoco.minimize.least_squares` for details.
* ``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
@@ -241,16 +235,19 @@ def optimize(
)
if check_conditioning:
_warn_if_ill_conditioned(initial_params, residual_fn)
_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. Optimization can stall in that corner on ill-conditioned or
# rank-deficient problems; both the mujoco and scipy backends are affected.
# 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)
at_bound = ((x0 - lo) <= 1e-3 * safe_rng) | ((hi - x0) <= 1e-3 * safe_rng)
at_bound &= rng > 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 "
python/mujoco/sysid/_src/parameter.py
+12 -7
View File
@@ -139,15 +139,18 @@ 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) -> None:
"""Shift values within 0.1% of a bound to ``lo + fraction*(hi-lo)`` (or
symmetric for the upper bound). Interior components are unchanged."""
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) <= 1e-3 * safe_rng
at_hi = (hi - v) <= 1e-3 * safe_rng
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)
@@ -306,13 +309,15 @@ class ParameterDict:
param.update_from_vector(vector[start : start + size])
start += size
def move_off_bounds(self, fraction: float = 0.05) -> Self:
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)
param.move_off_bound(fraction=fraction, at_bound_tol=at_bound_tol)
return self
def save_to_disk(self, path: str | pathlib.Path) -> None:
python/mujoco/sysid/tests/test_integration.py
+5 -4
View File
@@ -238,11 +238,12 @@ def test_arm_recover_armature():
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] + x[1] - 2.0])
r = np.array([x[0] - 1.0, 2.0 * x[0] - 2.0])
else:
r = (x[0] + x[1] - 2.0).reshape(1, -1)
r = np.stack([x[0] - 1.0, 2.0 * x[0] - 2.0])
return [r], None, None
@@ -260,7 +261,7 @@ def _full_rank_residual_fn(x, p):
(_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."""
"""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))
@@ -271,5 +272,5 @@ def test_check_conditioning(residual_fn, expect_warning, caplog):
optimizer="scipy", verbose=False, check_conditioning=True,
max_iters=1,
)
fired = any("cond(JᵀJ)" in r.message for r in caplog.records)
fired = any("cond(J^T J)" in r.message for r in caplog.records)
assert fired is expect_warning