From c12dc23852996ba666b58b136ac61548202ed84a Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Tue, 26 May 2026 17:06:09 -0700 Subject: [PATCH 1/2] Add x_scale to mujoco.minimize.least_squares. 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. --- python/mujoco/minimize.py | 50 +++++++++++++++++++++++++++------- python/mujoco/minimize_test.py | 47 ++++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 10 deletions(-) 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() From aa4a43da1732d73c6fb89ae60893fac5ece72fef Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Wed, 27 May 2026 12:55:14 -0700 Subject: [PATCH 2/2] Address reviewer feedback: default x_scale=None, drop TRF qualifier. --- python/mujoco/minimize.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/python/mujoco/minimize.py b/python/mujoco/minimize.py index fb829398..520fef11 100644 --- a/python/mujoco/minimize.py +++ b/python/mujoco/minimize.py @@ -156,7 +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, + x_scale: Optional[Union[float, np.ndarray, str]] = None, ) -> Tuple[np.ndarray, List[IterLog]]: """Nonlinear Least Squares minimization with box bounds. @@ -179,13 +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. 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. + 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 @@ -211,13 +210,15 @@ 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. + # 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 'jac', a positive scalar, or array, got {x_scale!r}." + f"x_scale must be None, 'jac', a positive scalar, or array, got " + f'{x_scale!r}.' ) - if adaptive_scale: + if x_scale is None or adaptive_scale: D = np.ones((n, 1)) else: D = np.asarray(x_scale, dtype=np.float64)