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.
+    # Refresh D for adaptive ('jac') scaling: column-norm preconditioner.
-    grad, hess = norm.grad_hess(r, jac)
+    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
+    # Bounds relative to x, expressed in scaled coords (dz = dx / D).
-    dlower = None if bounds is None else bounds[0] - x
+    dlower = None if bounds is None else (bounds[0] - x) / D
-    dupper = None if bounds is None else bounds[1] - x
+    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.
+      # New candidate (D * dx is the x-space step).
-      xnew = x + dx
+      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.
+    # Iteration message. Step printed in x-space (D * dx).
-    dx_norm = np.linalg.norm(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:

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()