Updates to minimize.least_squares. Fixes #1585

- Residual callable is vectorized for easy multithreading by the user. Internally all vectors are now explicitly column vectors.
- Removed central findiff option, it wasn't applicable at the bounds anyway and just complicated the code.
- Added optional user-provided norm function for non-quadratic (robust) norms.
- Added an iter_callback callable for user convenience.
- Added option to internally check user-provided Jacobian and norm against finite differences.
- Updated the notebook accordingly.

PiperOrigin-RevId: 631205889
Change-Id: I3b9f8893756e329640de464e6f2e26a39e7dbd9e

python/mujoco/minimize.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Nonlinear Least Squares minimization with box bounds."""
 
+import abc
 import dataclasses
 import enum
 import time
@@ -32,14 +33,14 @@ class Verbosity(enum.Enum):
 
 class Status(enum.Enum):
   FACTORIZATION_FAILED = enum.auto()
-  NO_IMPORVEMENT = enum.auto()
+  NO_IMPROVEMENT = enum.auto()
   MAX_ITER = enum.auto()
   DX_TOL = enum.auto()
 
 
 _STATUS_MESSAGE = {
     Status.FACTORIZATION_FAILED: 'factorization failed.',
-    Status.NO_IMPORVEMENT: 'insufficient reduction.',
+    Status.NO_IMPROVEMENT: 'insufficient reduction.',
     Status.MAX_ITER: 'maximum iterations reached.',
     Status.DX_TOL: 'norm(dx) < tol.',
 }
@@ -68,90 +69,70 @@ class IterLog:
   step: Optional[np.ndarray] = None
 
 
-def jacobian_fd(
+class Norm(abc.ABC):
-    residual: Callable[[np.ndarray], np.ndarray],
+  """Abstract interface for norm functions, measuring the magnitude of vectors.
-    x: np.ndarray,
-    r: np.ndarray,
-    eps: np.float64,
-    central: bool,
-    n_res: int,
-    bounds: Optional[List[np.ndarray]] = None,
-):
-  """Finite-difference Jacobian of a residual function.
 
-  Args:
+  Key Concepts:
-    residual: function that returns the residual for a given point.
-    x: point at which to evaluate the Jacobian.
-    r: residual at x.
-    eps: finite-difference step size.
-    central: whether to use central differences.
-    n_res: number or residual evaluations so far.
-    bounds: optional pair of lower and upper bounds.
 
-  Returns:
+  * Norm Value: The value of the norm for a given input vector.
-    jac: Jacobian of the residual at x.
+  * Gradient and Hessian: The gradient (first derivative) and Hessian (second
-    n_res: updated number of residual evaluations.
+    derivative) of the norm function with respect to the input vector.
 
+  Subclasses Must Implement:
+
+  * `value(self, r: np.ndarray)`: Computes and returns the norm value for the
+     input vector `r`.
+  * `grad_hess(self, r: np.ndarray, proj: np.ndarray)`: Computes and returns
+     both  the gradient and Hessian of the norm at `r`, projected onto `proj`.
+     The reason we ask the user to perform the projection themselves is that
+     norm Hessians are often large and sparse, and the "sandwich" projection
+     operator `proj.T @ hess @ proj` can be computed efficiently by taking the
+     specific norm structure into account.
   """
-  nx = x.size
+
-  nr = r.size
+  @abc.abstractmethod
-  jac = np.zeros((nr, nx))
+  def value(self, r: np.ndarray) -> np.float64:
-  xh = x.copy()
+    """Returns the value of the norm at the input vector `y = norm(r)`."""
-  if bounds is None:
+    pass
-    # No bounds, simple forward or central differencing.
+
-    for i in range(nx):
+  @abc.abstractmethod
-      xh[i] = x[i] + eps
+  def grad_hess(self, r: np.ndarray, proj: np.ndarray):
-      rp = residual(xh)
+    """Computes the projected gradient and Hessian of the norm at `r`.
-      if central:
+
-        xh[i] = x[i] - eps
+    Args:
-        rm = residual(xh)
+        r: A NumPy column vector (nr x 1).
-        jac[:, i] = (rp - rm) / (2*eps)
+        proj: A pre-computed projection matrix (nr x nx).
-      else:
+
-        jac[:, i] = (rp - r) / eps
+    Returns:
-      xh[i] = x[i]
+        A tuple containing:
-    n_res += 2*nx if central else nx
+            * Projected gradient: proj.T @ (d_norm/d_r).
-  else:
+            * Projected Hessian: proj.T @ (d^2_norm/d_r^2) @ proj.
-    lower, upper = bounds
+    """
-    midpoint = 0.5 * (upper - lower)
+    pass
-    for i in range(nx):
+
-      # Scale eps, don't cross bounds.
+
-      eps_i = eps * (upper[i] - lower[i])
+class Quadratic(Norm):
-      if central:
+  """Implementation of the quadratic norm."""
-        # Use central differencing if away from bounds.
+
-        if x[i] - eps_i < lower[i]:
+  def value(self, r: np.ndarray):
-          # Near lower bound, use forward.
+    """Returns the quadratic norm of `r`."""
-          xh[i] = x[i] + eps_i
+    return 0.5 * (r.T @ r).item()
-          rp = residual(xh)
+
-          jac[:, i] = (rp - r) / eps_i
+  def grad_hess(self, r: np.ndarray, proj: np.ndarray):
-          n_res += 1
+    """Computes the projected gradient and Hessian of the quadratic norm at `r`.
-        elif x[i] + eps_i > upper[i]:
+
-          # Near upper bound, use backward.
+    Args:
-          xh[i] = x[i] - eps_i
+        r: A NumPy column vector (nr x 1).
-          rm = residual(xh)
+        proj: A pre-computed projection matrix (nr x nx).
-          jac[:, i] = (r - rm) / eps_i
+
-          n_res += 1
+    Returns:
-        else:
+        A tuple containing:
-          # Use central.
+            * Projected gradient: `proj.T @ r`.
-          xh[i] = x[i] + eps_i
+            * Projected Hessian: `proj.T @ proj`.
-          rp = residual(xh)
+    """
-          xh[i] = x[i] - eps_i
+    grad = proj.T @ r
-          rm = residual(xh)
+    hess = proj.T @ proj  # Notionally proj.T @ np.eye(r.size) @ proj
-          jac[:, i] = (rp - rm) / (2*eps_i)
+    return grad, hess
-          n_res += 2
-      else:
-        # Below midpoint use forward differencing, otherwise backward.
-        if x[i] < midpoint[i]:
-          xh[i] = x[i] + eps_i
-          rp = residual(xh)
-          jac[:, i] = (rp - r) / eps_i
-        else:
-          xh[i] = x[i] - eps_i
-          rm = residual(xh)
-          jac[:, i] = (r - rm) / eps_i
-        n_res += 1
-      # Reset.
-      xh[i] = x[i]
-  return jac, n_res
 
 
 def least_squares(
@@ -159,33 +140,38 @@ def least_squares(
     residual: Callable[[np.ndarray], np.ndarray],
     bounds: Optional[List[np.ndarray]] = None,
     jacobian: Optional[Callable[[np.ndarray, np.ndarray], np.ndarray]] = None,
+    norm: Norm = Quadratic(),
     eps: float = 1e-6,
-    central: bool = False,
     mu_min: float = 1e-6,
     mu_max: float = 1e8,
     mu_factor: float = 10.0**0.1,
-    tol: float = 1e-7,
+    tol: float = 1e-6,
     max_iter: int = 100,
     verbose: Union[Verbosity, int] = Verbosity.ITER,
     output: Optional[TextIO] = None,
+    iter_callback: Optional[Callable[[List[IterLog]], None]] = None,
+    check_derivatives: bool = False,
 ) -> Tuple[np.ndarray, List[IterLog]]:
   """Nonlinear Least Squares minimization with box bounds.
 
   Args:
-    x0: initial guess
+    x0: Initial guess
-    residual: function that returns the residual for a given point x.
+    residual: Vectorized function returning the residual for 1 or more points.
-    bounds: optional pair of lower and upper bounds on the solution.
+    bounds: Optional pair of lower and upper bounds on the solution.
-    jacobian: optional function that returns Jacobian of the residual at a given
+    jacobian: Optional function that returns Jacobian of the residual at a given
       point and residual. If not given, `residual` will be finite-differenced.
-    eps: perurbation used for automatic finite-differencing.
+    norm: Norm object returning norm scalar or its projected gradient and
-    central: whether to use central differences.
+      Hessian. See Norm class for detailed documentation.
-    mu_min: minimum value of the regularizer.
+    eps: Perurbation used for automatic finite-differencing.
-    mu_max: maximum value of the regularizer.
+    mu_min: Minimum value of the regularizer.
-    mu_factor: factor increasing or decreasing the regularizer.
+    mu_max: Maximum value of the regularizer.
-    tol: termination tolerance on the step size.
+    mu_factor: Factor for increasing or decreasing the regularizer.
-    max_iter: maximum number of iterations.
+    tol: Termination tolerance on the step size.
-    verbose: verbosity level.
+    max_iter: Maximum number of iterations.
-    output: optional file or StringIO to which to print messages.
+    verbose: Verbosity level.
+    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.
 
   Returns:
     x: best solution found
@@ -202,10 +188,10 @@ def least_squares(
   # Initialize locals.
   status = Status.MAX_ITER
   i = 0
-  x = x0.astype(np.float64)
+  n = x0.size
-  n = x.size
+  x = x0.astype(np.float64).reshape((n, 1))
-  xnew = np.zeros((n,))
+  xnew = np.zeros((n, 1))
-  dx = np.zeros((n,))
+  dx = np.zeros((n, 1))
   scratch = np.zeros((n, n + 7))
   eps = np.float64(eps)
   mu = np.float64(0.0)  # Optimistically start with no regularization.
@@ -234,6 +220,8 @@ def least_squares(
     n_reduc = 0  # Reset n_reduc.
     return mu, n_reduc
 
+  # Make local copy of bounds to avoid reshaping user input.
+  bounds = None if bounds is None else bounds.copy()
   if bounds is not None:
     # Checks bounds.
     if len(bounds) != 2:
@@ -244,7 +232,10 @@ def least_squares(
       raise ValueError('bounds must be finite.')
     if not np.all(bounds[0] < bounds[1]):
       raise ValueError('bounds[0] must be smaller than bounds[1].')
-    # Clip.
+
+    # Reshape and clip.
+    bounds[0] = bounds[0].reshape(n, 1)
+    bounds[1] = bounds[1].reshape(n, 1)
     np.clip(x, bounds[0], bounds[1], out=x)
 
   # Check for NaNs.
@@ -267,21 +258,28 @@ def least_squares(
       break
 
     # Get objective y.
-    y = 0.5 * r.dot(r)
+    y = norm.value(r)
 
     # Get Jacobian jac.
     t_start = time.time()
     if jacobian is None:
-      jac, n_res = jacobian_fd(residual, x, r, eps, central, n_res, bounds)
+      jac, n_res = jacobian_fd(residual, x, r, eps, n_res, bounds)
       t_res += time.time() - t_start
     else:
       jac = jacobian(x, r)
       t_jac += time.time() - t_start
       n_jac += 1
 
+      # Check user-provided Jacobian
+      if i == 0 and check_derivatives:
+        n_res = check_jacobian(residual, x, r, jac, eps, n_res, bounds, output)
+
+    # Check user-provided norm
+    if i == 0 and check_derivatives and not isinstance(norm, Quadratic):
+      check_norm(r, norm, eps, output)
+
     # Get gradient, Gauss-Newton Hessian.
-    grad = jac.T @ r
+    grad, hess = norm.grad_hess(r, jac)
-    hess = jac.T @ jac
 
     # Bounds relative to x
     dlower = None if bounds is None else bounds[0] - x
@@ -316,13 +314,13 @@ def least_squares(
       n_res += 1
 
       # New objective, evaluate reduction.
-      ynew = 0.5 * rnew.dot(rnew)
+      ynew = norm.value(rnew)
       reduction = y - ynew
-      armijo = reduction + armijo_c1*grad.dot(dx)
+      armijo = reduction + armijo_c1 * (grad.T @ dx).item()
 
       if armijo < 0:
         if mu >= mu_max:
-          status = Status.NO_IMPORVEMENT
+          status = Status.NO_IMPROVEMENT
           break
         mu, n_reduc = increase_mu(mu)
 
@@ -330,7 +328,7 @@ def least_squares(
       break
 
     # Compute reduction ratio.
-    expected_reduction = -(grad.dot(dx) + 0.5 * dx.T @ hess @ dx)
+    expected_reduction = -(grad.T @ dx + 0.5 * dx.T @ hess @ dx).item()
     reduction_ratio = 0.0
     if expected_reduction <= 0:
       if verbose > Verbosity.SILENT.value:
@@ -352,11 +350,13 @@ def least_squares(
       )
       print(message, file=output)
 
-    # Append log to trace.
+    # 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)
     trace.append(log)
+    if iter_callback is not None:
+      iter_callback(trace)
 
     # Check for success.
     if dx_norm < tol:
@@ -373,12 +373,14 @@ def least_squares(
     x = xnew
     r = rnew
 
-  # Append final log to trace.
+  # Append final log to trace, call iter_callback.
-  # Note: unlike other iter logs, this is at the end point.
+  # Note: unlike other iter logs, values are computed at the end point.
-  yfinal = 0.5 * r.dot(r)
+  yfinal = norm.value(r)
-  red = np.float64(0.0)
+  red = np.float64(0.0)  # No reduction sice 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:
+    iter_callback(trace)
 
   # Print final diagnostics.
   if verbose > Verbosity.SILENT.value:
@@ -401,4 +403,124 @@ def least_squares(
         message += f' Jacobian {jac_percent:<.1f}%'
       print(message, file=output)
 
-  return x, trace
+  return x.reshape(x0.shape), trace
+
+
+def jacobian_fd(
+    residual: Callable[[np.ndarray], np.ndarray],
+    x: np.ndarray,
+    r: np.ndarray,
+    eps: np.float64,
+    n_res: int,
+    bounds: Optional[List[np.ndarray]] = None,
+) -> Tuple[np.ndarray, int]:
+  """Finite-difference Jacobian of a residual function.
+
+  Args:
+    residual: vectorized function that returns the residual of a vector array.
+    x: point at which to evaluate the Jacobian.
+    r: residual at x.
+    eps: finite-difference step size.
+    n_res: number or residual evaluations so far.
+    bounds: optional pair of lower and upper bounds.
+
+  Returns:
+    jac: Jacobian of the residual at x.
+    n_res: updated number of residual evaluations (add x.size).
+
+  """
+  n = x.size
+  if bounds is None:
+    eps_vec = eps * np.ones(n)
+  else:
+    mid = 0.5 * (bounds[1] - bounds[0])
+    eps_vec = np.where(x > mid, -eps, eps).flatten()
+  xh = x + np.diag(eps_vec)
+  rh = residual(xh)
+  jac = (rh - r) / eps_vec
+  return jac, n_res+n
+
+
+def check_jacobian(
+    residual: Callable[[np.ndarray], np.ndarray],
+    x: np.ndarray,
+    r: np.ndarray,
+    jac: np.ndarray,
+    eps: np.float64,
+    n_res: int,
+    bounds: Optional[List[np.ndarray]] = None,
+    output: Optional[TextIO] = None,
+    name: Optional[str] = 'Jacobian',
+) -> int:
+  """Check user-provided Jacobian against internal finite-differencing.
+
+  Args:
+    residual: vectorized function that returns the residual of a vector array.
+    x: point at which the r and jac were evaluated.
+    r: residual at x.
+    jac: Jacobian at x.
+    eps: finite-difference step size.
+    n_res: number or residual evaluations so far.
+    bounds: optional pair of lower and upper bounds.
+    output: Optional file or StringIO to which to print messages.
+    name: Optional name of the function being tested.
+
+  Returns:
+    n_res: updated number of residual evaluations.
+
+  """
+  jac_fd, n_res = jacobian_fd(residual, x, r, eps, n_res, bounds)
+  denom = np.abs(jac).sum() + np.abs(jac_fd).sum() + 1e-8
+  rel_diff = np.abs(jac - jac_fd) / denom
+  if np.any(rel_diff > 1e-5):
+    raise ValueError(f'User-provided {name} does not match finite-differences '
+                     'to a relative tolerance of 1e-5.')
+  print(f'User-provided {name} matches finite-differences.', file=output)
+  return n_res
+
+
+def check_norm(
+    r: np.ndarray,
+    norm: Norm,
+    eps: np.float64,
+    output: Optional[TextIO] = None,
+):
+  """Check user-provided norm against internal finite-differencing.
+
+  Args:
+    r: residual vector.
+    norm: Norm function returning either the norm scalar or its gradient
+      and Gauss-Newton Hessian.
+    eps: finite-difference step size.
+    output: Optional file or StringIO to which to print messages.
+  """
+  # Get norm(r) value and 1st, 2nd derivatives.
+  n = np.atleast_2d(norm.value(r))  # norm value as 1x1 array.
+  eye = np.eye(r.size)  # Identity projection.
+  n_g, n_h = norm.grad_hess(r, eye)  # Gradient and Hessian.
+
+  # Check that Hessian is symmetric.
+  if not np.allclose(n_h, n_h.T):
+    raise ValueError('User-provided norm Hessian is not symmetric.')
+
+  # Check that Hessian is positive-definite.
+  if np.any(np.linalg.eigvals(n_h) < 0):
+    h_min = np.min(np.linalg.eigvals(n_h))
+    raise ValueError('User-provided norm Hessian is not positive definite. '
+                     f'Minimum eigenvalue is {h_min:<.4g}')
+
+  # Local function returning norm values (vectorized).
+  def norm_vec(v):
+    norms = [np.atleast_2d(norm.value(v[:, i:i+1])) for i in range(v.shape[1])]
+    return np.hstack(norms)
+
+  # Check the norm gradient.
+  check_jacobian(norm_vec, r, n, n_g.T, eps, 0, None, output, 'norm gradient')
+
+  # Local function returning norm gradients (vectorized).
+  def grad_vec(v):
+    gradients = [norm.grad_hess(v[:, i:i+1], eye)[0] for i in range(v.shape[1])]
+    return np.hstack(gradients)
+
+  # Check the norm Hessian.
+  check_jacobian(grad_vec, r, n_g, n_h, eps, 0, None, output, 'norm Hessian')

python/mujoco/minimize_test.py

@@ -15,7 +15,6 @@
 """Tests for minimize.py."""
 
 import io
-from typing import Tuple
 
 from absl.testing import absltest
 from mujoco import minimize
@@ -25,20 +24,19 @@ import numpy as np
 class MinimizeTest(absltest.TestCase):
 
   def test_basic(self) -> None:
-    def residual(x: np.ndarray) -> np.ndarray:
+    def residual(x):
-      return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)], dtype=np.float64)
+      return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
 
-    for central in [False, True]:
+    out = io.StringIO()
-      out = io.StringIO()
+    x0 = np.array((0.0, 0.0))
-      x0 = np.array((0.0, 0.0))
+    x, _ = minimize.least_squares(x0, residual, output=out)
-      x, _ = minimize.least_squares(x0, residual, output=out, central=central)
+    expected_x = np.array((1.0, 1.0))
-      expected_x = np.array((1.0, 1.0))
+    np.testing.assert_array_almost_equal(x, expected_x)
-      np.testing.assert_array_almost_equal(x, expected_x)
+    self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol')
-      self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol')
 
   def test_start_at_minimum(self) -> None:
-    def residual(x: np.ndarray) -> np.ndarray:
+    def residual(x):
-      return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
+      return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
 
     out = io.StringIO()
     x0 = np.array((1.0, 1.0))
@@ -49,33 +47,36 @@ class MinimizeTest(absltest.TestCase):
     self.assertContainsSubsequence(out.getvalue(), 'exact minimum found')
 
   def test_jac_callback(self) -> None:
-    def residual(x: np.ndarray) -> np.ndarray:
+    def residual(x):
-      return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
+      return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
 
-    def jacobian(x: np.ndarray, r: np.ndarray) -> Tuple[float, np.ndarray]:
+    def jacobian(x, r):
       del r  # Unused.
-      return np.array([[-1, 0], [-20 * x[0], 10]])
+      return np.array([[-1, 0], [-20 * x[0, 0], 10]])
 
     x0 = np.array((0.0, 0.0))
     out = io.StringIO()
-    x, _ = minimize.least_squares(x0, residual, jacobian=jacobian, output=out)
+    x, _ = minimize.least_squares(x0, residual, jacobian=jacobian, output=out,
+                                  check_derivatives=True)
     expected_x = np.array((1.0, 1.0))
     np.testing.assert_array_almost_equal(x, expected_x)
     self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol')
+    self.assertContainsSubsequence(out.getvalue(), 'Jacobian matches')
 
-    # Try with bad Jacobian, expect no improvement.
+    # Try with bad Jacobian, ask least_squares to check it.
-    def jac_bad1(x: np.ndarray, r: np.ndarray) -> Tuple[float, np.ndarray]:
+    def bad_jacobian(x, r):
-      return -jacobian(x, r)
+      del r  # Unused.
-    out1 = io.StringIO()
+      return np.array([[-1, 0], [-20 * x[0, 0], 15]])
-    minimize.least_squares(x0, residual, jacobian=jac_bad1, output=out1)
+    with self.assertRaisesRegex(ValueError, r'\bJacobian does not match\b'):
-    self.assertContainsSubsequence(out1.getvalue(), 'insufficient reduction')
+      minimize.least_squares(x0, residual, jacobian=bad_jacobian, output=out,
+                             check_derivatives=True)
 
   def test_max_iter(self) -> None:
     dim = 20  # High-D Rosenbrock
 
-    def residual(x: np.ndarray) -> np.ndarray:
+    def residual(x):
-      res0 = [1 - x[i] for i in range(dim - 1)]
+      res0 = [1 - x[i, :] for i in range(dim - 1)]
-      res1 = [10 * (x[i] - x[i + 1] ** 2) for i in range(dim - 1)]
+      res1 = [10 * (x[i, :] - x[i + 1, :] ** 2) for i in range(dim - 1)]
       return np.asarray(res0 + res1)
 
     # Fail to reach minimum after 20 iterations.
@@ -90,8 +91,8 @@ class MinimizeTest(absltest.TestCase):
     np.testing.assert_array_almost_equal(x, expected_x)
 
   def test_bounds(self) -> None:
-    def residual(x: np.ndarray) -> np.ndarray:
+    def residual(x):
-      return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
+      return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
 
     out = io.StringIO()
     x0 = np.array((0.0, 0.0))
@@ -108,32 +109,28 @@ class MinimizeTest(absltest.TestCase):
     self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol')
 
     # Test different bounds conditions.
-    verbose = minimize.Verbosity.FULLITER
+    for bounds in bounds_types.values():
-
+      out = io.StringIO()
-    for central in [False, True]:
+      x, trace = minimize.least_squares(
-      for bounds in bounds_types.values():
+          x0,
-        out = io.StringIO()
+          residual,
-        x, trace = minimize.least_squares(
+          bounds=bounds,
-            x0,
+          output=out,
-            residual,
+          verbose=minimize.Verbosity.FULLITER,
-            bounds=bounds,
+      )
-            output=out,
+      self.assertContainsSubsequence(out.getvalue(), ' < tol')
-            central=central,
+      grad = trace[-2].jacobian.T @ trace[-2].residual
-            verbose=verbose,
+      # If x_i is on the boundary, gradient points out, otherwise it is 0.
-        )
+      for i, xi in enumerate(x):
-        self.assertContainsSubsequence(out.getvalue(), ' < tol')
+        if xi == bounds[0][i]:
-        grad = trace[-2].jacobian.T @ trace[-2].residual
+          self.assertGreater(grad[i], 0)
-        # If x_i is on the boundary, gradient points out, otherwise it is 0.
+        elif xi == bounds[1][i]:
-        for i, xi in enumerate(x):
+          self.assertLess(grad[i], 0)
-          if xi == bounds[0][i]:
+        else:
-            self.assertGreater(grad[i], 0)
+          self.assertAlmostEqual(grad[i].item(), 0, places=4)
-          elif xi == bounds[1][i]:
-            self.assertLess(grad[i], 0)
-          else:
-            self.assertAlmostEqual(grad[i], 0, places=4)
 
   def test_bad_bounds(self) -> None:
-    def residual(x: np.ndarray) -> np.ndarray:
+    def residual(x):
       return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
 
     out = io.StringIO()
@@ -150,5 +147,114 @@ class MinimizeTest(absltest.TestCase):
       with self.assertRaises(ValueError):
         minimize.least_squares(x0, residual, bounds=bounds, output=out)
 
+  def test_iter_callback(self) -> None:
+    def residual(x):
+      return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
+
+    out = io.StringIO()
+
+    def iter_callback(trace):
+      print(f'Hello iteration {len(trace)}!', file=out)
+
+    x0 = np.array((0.0, 0.0))
+    x, _ = minimize.least_squares(x0, residual, output=out,
+                                  iter_callback=iter_callback)
+    expected_x = np.array((1.0, 1.0))
+    np.testing.assert_array_almost_equal(x, expected_x)
+    self.assertContainsSubsequence(out.getvalue(), 'Hello iteration 3!')
+
+  def test_norm(self) -> None:
+    def residual(x):
+      return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
+
+    p = 0.01  # Smoothing radius for smooth-L2 norm.
+
+    class SmoothL2(minimize.Norm):
+      def value(self, r):
+        return np.sqrt((r.T @ r).item() + p*p) - p
+
+      def grad_hess(self, r, proj):
+        s = np.sqrt((r.T @ r).item() + p*p)
+        y_r = r / s
+        grad = proj.T @ y_r
+        y_rr = (np.eye(r.size) - y_r @ y_r.T) / s
+        hess = proj.T @ y_rr @ proj
+        return grad, hess
+
+    out = io.StringIO()
+    x0 = np.array((0.0, 0.0))
+    x, _ = minimize.least_squares(x0, residual, norm=SmoothL2(), output=out,
+                                  check_derivatives=True)
+    expected_x = np.array((1.0, 1.0))
+    np.testing.assert_array_almost_equal(x, expected_x)
+    self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol')
+    self.assertContainsSubsequence(out.getvalue(),
+                                   'User-provided norm gradient matches')
+    self.assertContainsSubsequence(out.getvalue(),
+                                   'User-provided norm Hessian matches')
+
+    class SmoothL2BadGrad(minimize.Norm):
+      def value(self, r):
+        return np.sqrt((r.T @ r).item() + p*p) - p
+
+      def grad_hess(self, r, proj):
+        s = np.sqrt((r.T @ r).item() + p*p)
+        y_r = r / s
+        grad = proj.T @ (y_r + 0.001)  # 0.001 is erronous.
+        y_rr = (np.eye(r.size) - y_r @ y_r.T) / s
+        hess = proj.T @ y_rr @ proj
+        return grad, hess
+
+    with self.assertRaisesRegex(ValueError, r'\bgradient does not match\b'):
+      minimize.least_squares(x0, residual, norm=SmoothL2BadGrad(), output=out,
+                             check_derivatives=True)
+
+    class SmoothL2BadHess(minimize.Norm):
+      def value(self, r):
+        return np.sqrt((r.T @ r).item() + p*p) - p
+
+      def grad_hess(self, r, proj):
+        s = np.sqrt((r.T @ r).item() + p*p)
+        y_r = r / s
+        grad = proj.T @ y_r
+        y_rr = (1.001 * np.eye(r.size) - y_r @ y_r.T) / s  # 1.001 is erronous.
+        hess = proj.T @ y_rr @ proj
+        return grad, hess
+
+    with self.assertRaisesRegex(ValueError, r'\bHessian does not match\b'):
+      minimize.least_squares(x0, residual, norm=SmoothL2BadHess(), output=out,
+                             check_derivatives=True)
+
+    class SmoothL2AsymHess(minimize.Norm):
+      def value(self, r):
+        return np.sqrt((r.T @ r).item() + p*p) - p
+
+      def grad_hess(self, r, proj):
+        s = np.sqrt((r.T @ r).item() + p*p)
+        y_r = r / s
+        grad = proj.T @ y_r
+        y_rr = (np.eye(r.size) - (y_r + 0.0001) @ y_r.T) / s
+        hess = proj.T @ y_rr @ proj
+        return grad, hess
+
+    with self.assertRaisesRegex(ValueError, r'\bnot symmetric\b'):
+      minimize.least_squares(x0, residual, norm=SmoothL2AsymHess(), output=out,
+                             check_derivatives=True)
+
+    class SmoothL2NegHess(minimize.Norm):
+      def value(self, r):
+        return np.sqrt((r.T @ r).item() + p*p) - p
+
+      def grad_hess(self, r, proj):
+        s = np.sqrt((r.T @ r).item() + p*p)
+        y_r = r / s
+        grad = proj.T @ y_r
+        y_rr = -(np.eye(r.size) - y_r @ y_r.T) / s  # Negative-definite.
+        hess = proj.T @ y_rr @ proj
+        return grad, hess
+
+    with self.assertRaisesRegex(ValueError, r'\bnot positive definite\b'):
+      minimize.least_squares(x0, residual, norm=SmoothL2NegHess(), output=out,
+                             check_derivatives=True)
 if __name__ == '__main__':
   absltest.main()