Improvements to minimize.least_squares.

All changes make functionality more similar to SciPy least squares:
- Use adaptive findiff epsilon.
- Make termination on step size relative to norm(x).
- Add termination on gradient norm.
- Make default tolerances like SciPy's.

PiperOrigin-RevId: 745221614
Change-Id: Iee93256651fca8154c97fa3bdaa9c67ede28e573

python/mujoco/minimize.py

@@ -36,6 +36,7 @@ class Status(enum.Enum):
   NO_IMPROVEMENT = enum.auto()
   MAX_ITER = enum.auto()
   DX_TOL = enum.auto()
+  G_TOL = enum.auto()
 
 
 _STATUS_MESSAGE = {
@@ -43,6 +44,7 @@ _STATUS_MESSAGE = {
     Status.NO_IMPROVEMENT: 'insufficient reduction.',
     Status.MAX_ITER: 'maximum iterations reached.',
     Status.DX_TOL: 'norm(dx) < tol.',
+    Status.G_TOL: 'norm(gradient) < tol.',
 }
 
 
@@ -57,6 +59,7 @@ class IterLog:
     regularizer: Value of the regularizer used for this iteration.
     residual: Optional value of the residual at the candidate.
     jacobian: Optional value of the Jacobian at the candidate.
+    grad: Optional value of the gradient at the candidate.
     step: Optional change in decision variable during this iteration.
   """
 
@@ -66,6 +69,7 @@ class IterLog:
   regularizer: np.float64
   residual: Optional[np.ndarray] = None
   jacobian: Optional[np.ndarray] = None
+  grad: Optional[np.ndarray] = None
   step: Optional[np.ndarray] = None
 
 
@@ -141,11 +145,12 @@ def least_squares(
     bounds: Optional[Sequence[np.ndarray]] = None,
     jacobian: Optional[Callable[[np.ndarray, np.ndarray], np.ndarray]] = None,
     norm: Norm = Quadratic(),
-    eps: float = 1e-6,
+    eps: float = np.finfo(np.float64).eps ** 0.5,
     mu_min: float = 1e-6,
     mu_max: float = 1e8,
-    mu_factor: float = 10.0**0.1,
+    mu_factor: float = 10.0 ** 0.1,
-    tol: float = 1e-6,
+    xtol: float = 1e-8,
+    gtol: float = 1e-8,
     max_iter: int = 100,
     verbose: Union[Verbosity, int] = Verbosity.ITER,
     output: Optional[TextIO] = None,
@@ -166,7 +171,8 @@ def least_squares(
     mu_min: Minimum value of the regularizer.
     mu_max: Maximum value of the regularizer.
     mu_factor: Factor for increasing or decreasing the regularizer.
-    tol: Termination tolerance on the step size.
+    xtol: Termination tolerance on relative step size.
+    gtol: Termination tolerance on gradient norm.
     max_iter: Maximum number of iterations.
     verbose: Verbosity level.
     output: Optional file or StringIO to which to print messages.
@@ -281,6 +287,23 @@ def least_squares(
     # Get gradient, Gauss-Newton Hessian.
     grad, hess = norm.grad_hess(r, jac)
 
+    # Get free (unclamped) gradient.
+    if bounds is None:
+      grad_free = grad
+    else:
+      clamped_lower = (x == bounds[0]) & (grad > 0)
+      clamped_upper = (x == bounds[1]) & (grad < 0)
+      clamped = clamped_lower | clamped_upper
+      grad_free = grad[~clamped]
+
+    # Check termination condition on gradient norm.
+    g_norm = np.linalg.norm(grad_free)
+    if g_norm <= gtol:
+      status = Status.G_TOL
+      if g_norm == 0:
+        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
@@ -353,13 +376,15 @@ def least_squares(
     # Append log to trace, call iter_callback.
     log = IterLog(candidate=x, objective=y, reduction=reduction, regularizer=mu)
     if verbose >= Verbosity.FULLITER.value:
-      log = dataclasses.replace(log, residual=r, jacobian=jac, step=dx)
+      log = dataclasses.replace(
+          log, residual=r, jacobian=jac, grad=grad, step=dx
+      )
     trace.append(log)
     if iter_callback is not None:
       iter_callback(trace)
 
-    # Check for success.
+    # Check termination condition on step norm.
-    if dx_norm < tol:
+    if dx_norm < xtol * (xtol + np.linalg.norm(x)):
       status = Status.DX_TOL
       break
 
@@ -376,7 +401,7 @@ def least_squares(
   # Append final log to trace, call iter_callback.
   # Note: unlike other iter logs, values are computed at the end point.
   yfinal = norm.value(r)
-  red = np.float64(0.0)  # No reduction sice we didn't take a step.
+  red = np.float64(0.0)  # No reduction since we didn't take a step.
   log = IterLog(candidate=x, objective=yfinal, reduction=red, regularizer=mu)
   trace.append(log)
   if iter_callback is not None:
@@ -430,13 +455,15 @@ def jacobian_fd(
   """
   n = x.size
   if bounds is None:
-    eps_vec = eps * np.ones(n)
+    eps_vec = eps * np.ones((n, 1))
   else:
     mid = 0.5 * (bounds[1] - bounds[0])
-    eps_vec = np.where(x > mid, -eps, eps).flatten()
+    eps_vec = np.where(x > mid, -eps, eps)
-  xh = x + np.diag(eps_vec)
+  eps_vec *= np.maximum(1.0, np.abs(x))
+  eps_vec = (eps_vec + x) - x
+  xh = x + np.diag(eps_vec.flatten())
   rh = residual(xh)
-  jac = (rh - r) / eps_vec
+  jac = (rh - r) / eps_vec.T
   return jac, n_res + n
 
 

python/mujoco/minimize_test.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for minimize.py."""
 
 import io
 
@@ -32,7 +31,7 @@ class MinimizeTest(absltest.TestCase):
     x, _ = minimize.least_squares(x0, residual, output=out)
     expected_x = np.array((1.0, 1.0))
     np.testing.assert_array_almost_equal(x, expected_x)
-    self.assertIn('norm(dx) < tol', out.getvalue())
+    self.assertIn('norm(gradient) < tol', out.getvalue())
 
   def test_start_at_minimum(self) -> None:
     def residual(x):
@@ -43,7 +42,7 @@ class MinimizeTest(absltest.TestCase):
     x, _ = minimize.least_squares(x0, residual, output=out)
     expected_x = np.array((1.0, 1.0))
     np.testing.assert_array_almost_equal(x, expected_x)
-    self.assertIn('norm(dx) < tol', out.getvalue())
+    self.assertIn('norm(gradient) < tol', out.getvalue())
     self.assertIn('exact minimum found', out.getvalue())
 
   def test_jac_callback(self) -> None:
@@ -61,7 +60,7 @@ class MinimizeTest(absltest.TestCase):
     )
     expected_x = np.array((1.0, 1.0))
     np.testing.assert_array_almost_equal(x, expected_x)
-    self.assertIn('norm(dx) < tol', out.getvalue())
+    self.assertIn('norm(gradient) < tol', out.getvalue())
     self.assertIn('Jacobian matches', out.getvalue())
 
     # Try with bad Jacobian, ask least_squares to check it.
@@ -116,7 +115,7 @@ class MinimizeTest(absltest.TestCase):
         x0, residual, bounds=bounds_types['inbounds'], output=out
     )
     np.testing.assert_array_almost_equal(x, expected_x)
-    self.assertIn('norm(dx) < tol', out.getvalue())
+    self.assertIn('norm(gradient) < tol', out.getvalue())
 
     # Test different bounds conditions.
     for bounds in bounds_types.values():