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

@@ -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: