diff --git a/python/mujoco/minimize.py b/python/mujoco/minimize.py index b06f3eb1..d3cda293 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: Optional[Union[float, np.ndarray, str]] = None, ) -> Tuple[np.ndarray, List[IterLog]]: """Nonlinear Least Squares minimization with box bounds. @@ -178,6 +179,12 @@ 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 (matches scipy's x_scale). Setting + x_scale=D solves the problem in the change of variables z = x / D + and un-scales the result. None (default) or 1.0 is no scaling. + 'jac' sets D_i = 1 / ||J(:, i)|| at each iteration. 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 +210,24 @@ 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; + # None means no scaling (D = 1). + adaptive_scale = isinstance(x_scale, str) + if adaptive_scale and x_scale != 'jac': + raise ValueError( + f"x_scale must be None, 'jac', a positive scalar, or array, got " + f'{x_scale!r}.' + ) + if x_scale is None or 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 +309,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 +334,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 +359,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 +392,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 +408,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: @@ -388,7 +419,7 @@ def least_squares( status = Status.DX_TOL break - # Modify regularizer like in (Bazaraa, Sherali, and Shetty) + # Modify regularizer like in (Fletcher, 1971) if reduction_ratio > 0.75: mu, n_reduc = decrease_mu(mu, n_reduc) elif reduction_ratio < 0.25: diff --git a/python/mujoco/minimize_test.py b/python/mujoco/minimize_test.py index 1792643d..a3081c0f 100644 --- a/python/mujoco/minimize_test.py +++ b/python/mujoco/minimize_test.py @@ -319,6 +319,110 @@ 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) + + def test_x_scale_with_bounds(self) -> None: + """x_scale combined with bounds reaches the constrained optimum.""" + def residual(x): + return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)]) + + x0 = np.array((2.0, 0.0)) + bounds = [np.array((1.5, -10.0)), np.array((10.0, 10.0))] + expected = np.array([1.5, 2.25]) + + # Without scaling. + x_noscale, _ = minimize.least_squares( + x0, residual, bounds=bounds, verbose=minimize.Verbosity.SILENT) + np.testing.assert_allclose(x_noscale, expected, atol=1e-4) + + # With fixed x_scale array. + x_scaled, _ = minimize.least_squares( + x0, residual, bounds=bounds, x_scale=np.array([1.0, 1e-5]), + verbose=minimize.Verbosity.SILENT) + np.testing.assert_allclose(x_scaled, expected, atol=1e-4) + + # With adaptive x_scale='jac'. + x_jac, _ = minimize.least_squares( + x0, residual, bounds=bounds, x_scale='jac', + verbose=minimize.Verbosity.SILENT) + np.testing.assert_allclose(x_jac, expected, atol=1e-4) + + def test_x_scale_jac_convergence(self) -> None: + """'jac' scaling converges faster on Powell's badly scaled function.""" + # Powell's badly scaled function (Moré, Garbow, Hillstrom #3): + # r1 = 1e4 * x1 * x2 - 1 + # r2 = exp(-x1) + exp(-x2) - 1.0001 + # The 1e4 multiplier creates Jacobian columns with wildly different + # norms, making the unscaled LM regularizer ineffective. + def residual(x): + return np.stack([ + 1e4 * x[0, :] * x[1, :] - 1, + np.exp(-x[0, :]) + np.exp(-x[1, :]) - 1.0001, + ]) + + x0 = np.array([0.0, 1.0]) + # Analytical solution: x1*x2 = 1e-4, exp(-x1)+exp(-x2) = 1.0001. + x_star = np.array([1.098159e-5, 9.106146]) + + x_default, trace_default = minimize.least_squares( + x0, residual, verbose=minimize.Verbosity.SILENT, max_iter=400) + x_jac, trace_jac = minimize.least_squares( + x0, residual, x_scale='jac', verbose=minimize.Verbosity.SILENT) + + # Both reach the correct minimum. + np.testing.assert_allclose(x_default, x_star, rtol=1e-4) + np.testing.assert_allclose(x_jac, x_star, rtol=1e-4) + + # Without scaling the solver needs 331 iterations to converge. + # With 'jac' scaling it converges in 55: a 6x improvement. + self.assertGreater(len(trace_default) - 1, 300) + self.assertLess(len(trace_jac) - 1, 70) + if __name__ == '__main__': absltest.main()