Merge pull request #3294 from kevinzakka:mujoco-minimize-x-scale
PiperOrigin-RevId: 923519000 Change-Id: I4a8fa1f031428cf41d194d17ea4c41975f93f745
This commit is contained in:
Copybara-Service
+42
-11
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user