diff --git a/python/mujoco/minimize.py b/python/mujoco/minimize.py index b06f3eb1..fb829398 100644 --- a/python/mujoco/minimize.py +++ b/python/mujoco/minimize.py @@ -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: diff --git a/python/mujoco/minimize_test.py b/python/mujoco/minimize_test.py index 1792643d..77694a3f 100644 --- a/python/mujoco/minimize_test.py +++ b/python/mujoco/minimize_test.py @@ -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()