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

Add x_scale to mujoco.minimize.least_squares.

Browse Source 
Per-parameter scaling via change of variables z = x / D. Supports
'jac' (adaptive D_i = 1/||J(:,i)|| per iteration, matches scipy's TRF),
explicit array, or a positive scalar. Default 1.0 is a no-op.
This commit is contained in:
Kevin Zakka
2026-05-26 17:06:09 -07:00
parent 583aa5ad4f
commit c12dc23852
2 changed files with 87 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/minimize.py
+40 -10
View File
@@ -156,6 +156,7 @@ def least_squares(
output: Optional[TextIO] = None,
iter_callback: Optional[Callable[[List[IterLog]], None]] = None,
check_derivatives: bool = False,
x_scale: Union[float, np.ndarray, str] = 1.0,
) -> Tuple[np.ndarray, List[IterLog]]:
"""Nonlinear Least Squares minimization with box bounds.
@@ -178,6 +179,13 @@ def least_squares(
output: Optional file or StringIO to which to print messages.
iter_callback: Optional iteration callback, takes trace argument.
check_derivatives: Compare user-defined Jacobian and norm against fin-diff.
x_scale: Per-parameter scaling. Setting ``x_scale=D`` solves the problem in
the change of variables ``z = x / D`` and un-scales the result. ``1.0``
(default) is no scaling. ``'jac'`` sets ``D_i = 1 / ||J(:, i)||`` at each
iteration (matches scipy's ``least_squares(method='trf', x_scale='jac')``).
An array of shape ``(n,)`` (or a positive scalar) is used as ``D``
directly. Note that ``mu``, ``mu_min`` and ``mu_max`` then act on the
scaled subproblem.
Returns:
x: best solution found
@@ -203,6 +211,22 @@ def least_squares(
mu = np.float64(0.0) # Optimistically start with no regularization.
n_reduc = 0 # Number of sequential mu reductions.
# Resolve x_scale -> D of shape (n, 1). For 'jac', D is refreshed each iter.
adaptive_scale = isinstance(x_scale, str)
if adaptive_scale and x_scale != 'jac':
raise ValueError(
f"x_scale must be 'jac', a positive scalar, or array, got {x_scale!r}."
)
if adaptive_scale:
D = np.ones((n, 1))
else:
D = np.asarray(x_scale, dtype=np.float64)
if D.shape not in ((), (n,)):
raise ValueError(f'x_scale array must have shape ({n},), got {D.shape}.')
if not np.all(np.isfinite(D)) or np.any(D <= 0):
raise ValueError('x_scale must be positive and finite.')
D = (np.ones(n) * D).reshape(n, 1)
# Initialize logging.
trace = []
n_res = 0
@@ -284,8 +308,13 @@ def least_squares(
if i == 0 and check_derivatives and not isinstance(norm, Quadratic):
check_norm(r, norm, eps, output)
# Get gradient, Gauss-Newton Hessian.
grad, hess = norm.grad_hess(r, jac)
# Refresh D for adaptive ('jac') scaling: column-norm preconditioner.
if adaptive_scale:
col_norms = np.linalg.norm(jac, axis=0)
D = (1.0 / np.maximum(col_norms, eps)).reshape(n, 1)
# Gradient/Hessian in scaled coords (jac * D.T scales columns).
grad, hess = norm.grad_hess(r, jac * D.T)
# Get free (unclamped) gradient.
if bounds is None:
@@ -304,9 +333,9 @@ def least_squares(
print('Zero gradient norm: exact minimum found?', file=output)
break
# Bounds relative to x
dlower = None if bounds is None else bounds[0] - x
dupper = None if bounds is None else bounds[1] - x
# Bounds relative to x, expressed in scaled coords (dz = dx / D).
dlower = None if bounds is None else (bounds[0] - x) / D
dupper = None if bounds is None else (bounds[1] - x) / D
# Find reduction satisfying Armijo's rule.
armijo = -1
@@ -329,8 +358,8 @@ def least_squares(
if status != Status.MAX_ITER:
break
# New candidate, residual.
xnew = x + dx
# New candidate (D * dx is the x-space step).
xnew = x + D * dx
t_start = time.time()
rnew = residual(xnew)
t_res += time.time() - t_start
@@ -362,8 +391,9 @@ def least_squares(
else:
reduction_ratio = reduction / expected_reduction
# Iteration message.
dx_norm = np.linalg.norm(dx)
# Iteration message. Step printed in x-space (D * dx).
step = D * dx
dx_norm = np.linalg.norm(step)
if verbose >= Verbosity.ITER.value:
logmu = np.log10(mu) if mu > 0 else -np.inf
message = (
@@ -377,7 +407,7 @@ def least_squares(
log = IterLog(candidate=x, objective=y, reduction=reduction, regularizer=mu)
if verbose >= Verbosity.FULLITER.value:
log = dataclasses.replace(
log, residual=r, jacobian=jac, grad=grad, step=dx
log, residual=r, jacobian=jac, grad=grad / D, step=step
)
trace.append(log)
if iter_callback is not None:
python/mujoco/minimize_test.py
+47
View File
@@ -319,6 +319,53 @@ class MinimizeTest(absltest.TestCase):
self.assertIn('User-provided norm gradient matches', out.getvalue())
self.assertIn('User-provided norm Hessian matches', out.getvalue())
def test_x_scale_default_is_no_op(self) -> None:
"""x_scale=1.0 produces an identical trace to the default unscaled run."""
def residual(x):
return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
x0 = np.array((0.0, 0.0))
x_default, trace_default = minimize.least_squares(
x0, residual, verbose=minimize.Verbosity.SILENT)
x_one, trace_one = minimize.least_squares(
x0, residual, x_scale=1.0, verbose=minimize.Verbosity.SILENT)
np.testing.assert_array_equal(x_default, x_one)
self.assertEqual(len(trace_default), len(trace_one))
def test_x_scale_reaches_same_minimum(self) -> None:
"""Both 'jac' and an explicit array reach the unscaled run's minimum."""
def residual(x):
return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
x0 = np.array((0.0, 0.0))
x_unscaled, _ = minimize.least_squares(
x0, residual, verbose=minimize.Verbosity.SILENT)
x_jac, _ = minimize.least_squares(
x0, residual, x_scale='jac', verbose=minimize.Verbosity.SILENT)
x_array, _ = minimize.least_squares(
x0, residual, x_scale=np.array([10.0, 0.1]),
verbose=minimize.Verbosity.SILENT)
np.testing.assert_allclose(x_jac, x_unscaled, atol=1e-6)
np.testing.assert_allclose(x_array, x_unscaled, atol=1e-6)
def test_x_scale_validation(self) -> None:
"""Invalid x_scale arguments raise ValueError."""
def residual(x):
return x
x0 = np.array((1.0, 1.0))
bad_values = [
'bogus', # unknown string
np.array([1.0, -1.0]), # non-positive entry
np.array([np.inf, 1.0]), # non-finite entry
np.array([1.0]), # wrong shape
]
for bad in bad_values:
with self.assertRaises(ValueError):
minimize.least_squares(
x0, residual, x_scale=bad,
verbose=minimize.Verbosity.SILENT)
if __name__ == '__main__':
absltest.main()