Merge pull request #3294 from kevinzakka:mujoco-minimize-x-scale

PiperOrigin-RevId: 923519000
Change-Id: I4a8fa1f031428cf41d194d17ea4c41975f93f745

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: