chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Updates to minimize.least_squares. Fixes #1585

Browse Source 
- 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
This commit is contained in:
Yuval Tassa
2024-05-06 15:34:56 -07:00
committed by Copybara-Service
parent 93b57c1421
commit 34e537e557
3 changed files with 1223 additions and 803 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/minimize_test.py
+158 -52
View File
@@ -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:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)], dtype=np.float64)
def residual(x):
return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
for central in [False, True]:
out = io.StringIO()
x0 = np.array((0.0, 0.0))
x, _ = minimize.least_squares(x0, residual, output=out, central=central)
expected_x = np.array((1.0, 1.0))
np.testing.assert_array_almost_equal(x, expected_x)
self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol')
out = io.StringIO()
x0 = np.array((0.0, 0.0))
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.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol')
def test_start_at_minimum(self) -> None:
def residual(x: np.ndarray) -> np.ndarray:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
def residual(x):
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:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
def residual(x):
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.
def jac_bad1(x: np.ndarray, r: np.ndarray) -> Tuple[float, np.ndarray]:
return -jacobian(x, r)
out1 = io.StringIO()
minimize.least_squares(x0, residual, jacobian=jac_bad1, output=out1)
self.assertContainsSubsequence(out1.getvalue(), 'insufficient reduction')
# Try with bad Jacobian, ask least_squares to check it.
def bad_jacobian(x, r):
del r # Unused.
return np.array([[-1, 0], [-20 * x[0, 0], 15]])
with self.assertRaisesRegex(ValueError, r'\bJacobian does not match\b'):
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:
res0 = [1 - x[i] for i in range(dim - 1)]
res1 = [10 * (x[i] - x[i + 1] ** 2) for i in range(dim - 1)]
def residual(x):
res0 = [1 - x[i, :] 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:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
def residual(x):
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 central in [False, True]:
for bounds in bounds_types.values():
out = io.StringIO()
x, trace = minimize.least_squares(
x0,
residual,
bounds=bounds,
output=out,
central=central,
verbose=verbose,
)
self.assertContainsSubsequence(out.getvalue(), ' < tol')
grad = trace[-2].jacobian.T @ trace[-2].residual
# If x_i is on the boundary, gradient points out, otherwise it is 0.
for i, xi in enumerate(x):
if xi == bounds[0][i]:
self.assertGreater(grad[i], 0)
elif xi == bounds[1][i]:
self.assertLess(grad[i], 0)
else:
self.assertAlmostEqual(grad[i], 0, places=4)
for bounds in bounds_types.values():
out = io.StringIO()
x, trace = minimize.least_squares(
x0,
residual,
bounds=bounds,
output=out,
verbose=minimize.Verbosity.FULLITER,
)
self.assertContainsSubsequence(out.getvalue(), ' < tol')
grad = trace[-2].jacobian.T @ trace[-2].residual
# If x_i is on the boundary, gradient points out, otherwise it is 0.
for i, xi in enumerate(x):
if xi == bounds[0][i]:
self.assertGreater(grad[i], 0)
elif xi == bounds[1][i]:
self.assertLess(grad[i], 0)
else:
self.assertAlmostEqual(grad[i].item(), 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()