From 55e3ca3acfbb787e8c697f88a2813791ba1ba178 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Tue, 8 Apr 2025 11:21:15 -0700 Subject: [PATCH] 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 | 51 ++++++++++++++++++++++++++-------- python/mujoco/minimize_test.py | 9 +++--- 2 files changed, 43 insertions(+), 17 deletions(-) diff --git a/python/mujoco/minimize.py b/python/mujoco/minimize.py index 449d4bb8..665c4fd9 100644 --- a/python/mujoco/minimize.py +++ b/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, - tol: float = 1e-6, + mu_factor: float = 10.0 ** 0.1, + 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. - if dx_norm < tol: + # Check termination condition on step norm. + 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() - xh = x + np.diag(eps_vec) + eps_vec = np.where(x > mid, -eps, eps) + 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 diff --git a/python/mujoco/minimize_test.py b/python/mujoco/minimize_test.py index e642a8e1..5ba90b2f 100644 --- a/python/mujoco/minimize_test.py +++ b/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():