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

Improvements to minimize.least_squares:

Browse Source 
- Added central differencing option.
- Changed `mu` semantics from log10 to actual values.
- `mu` control is now more aggressive, saves a few iterations.
- Use Armijo sufficient reduction criterion.
- Added some checks for float64 and NaNs.
- Added log of final result to trace.
- Removed unhelpful QP timing.

PiperOrigin-RevId: 610096950
Change-Id: Idf1a7c3155ce8e82fc2d5fe0f2da2f16636ff63d
This commit is contained in:
Yuval Tassa
2024-02-24 20:41:33 -08:00
committed by Copybara-Service
parent c6a41fbfe6
commit 03a8fa9ca9
2 changed files with 188 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/minimize.py
+148 -81
View File
@@ -39,9 +39,9 @@ class Status(enum.Enum):
_STATUS_MESSAGE = {
Status.FACTORIZATION_FAILED: 'factorization failed.',
Status.NO_IMPORVEMENT: 'no improvement found.',
Status.NO_IMPORVEMENT: 'insufficient reduction.',
Status.MAX_ITER: 'maximum iterations reached.',
Status.DX_TOL: 'norm(step) < tolerance.',
Status.DX_TOL: 'norm(dx) < tol.',
}
@@ -72,7 +72,9 @@ def jacobian_fd(
residual: Callable[[np.ndarray], np.ndarray],
x: np.ndarray,
r: np.ndarray,
eps: float,
eps: np.float64,
central: bool,
n_res: int,
bounds: Optional[List[np.ndarray]] = None,
):
"""Finite-difference Jacobian of a residual function.
@@ -82,37 +84,74 @@ def jacobian_fd(
x: point at which to evaluate the Jacobian.
r: residual at x.
eps: finite-difference step size.
bounds: optional pair of lower and upper bounds of the solution.
central: whether to use central differences.
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.
"""
nx = x.size
nr = r.size
jac = np.zeros((nr, nx))
xh = x.copy()
for i in range(nx):
if bounds is not None:
# Have bounds: scale eps, don't cross bounds.
lower, upper = bounds
eps_i = eps * (upper[i] - lower[i])
if xh[i] < upper[i] - eps_i:
# Not near upper bound, use forward.
xh[i] += eps_i
rh = residual(xh)
jac[:, i] = (rh - r) / eps_i
if bounds is None:
# No bounds, simple forward or central differencing.
for i in range(nx):
xh[i] = x[i] + eps
rp = residual(xh)
if central:
xh[i] = x[i] - eps
rm = residual(xh)
jac[:, i] = (rp - rm) / (2*eps)
else:
# Near upper bound, use backward.
xh[i] -= eps_i
rh = residual(xh)
jac[:, i] = (r - rh) / eps_i
else:
# No bounds, just use forward fin-diff.
xh[i] += eps
rh = residual(xh)
jac[:, i] = (rh - r) / eps
xh[i] = x[i]
return jac
jac[:, i] = (rp - r) / eps
xh[i] = x[i]
n_res += 2*nx if central else nx
else:
lower, upper = bounds
midpoint = 0.5 * (upper - lower)
for i in range(nx):
# Scale eps, don't cross bounds.
eps_i = eps * (upper[i] - lower[i])
if central:
# Use central differencing if away from bounds.
if x[i] - eps_i < lower[i]:
# Near lower bound, use forward.
xh[i] = x[i] + eps_i
rp = residual(xh)
jac[:, i] = (rp - r) / eps_i
n_res += 1
elif x[i] + eps_i > upper[i]:
# Near upper bound, use backward.
xh[i] = x[i] - eps_i
rm = residual(xh)
jac[:, i] = (r - rm) / eps_i
n_res += 1
else:
# Use central.
xh[i] = x[i] + eps_i
rp = residual(xh)
xh[i] = x[i] - eps_i
rm = residual(xh)
jac[:, i] = (rp - rm) / (2*eps_i)
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(
@@ -120,13 +159,14 @@ 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,
eps: Optional[float] = -6,
mu_min: Optional[float] = -6,
mu_max: Optional[float] = 8,
mu_delta: Optional[float] = 0.5,
tol: Optional[float] = 1e-7,
max_iter: Optional[int] = 100,
verbose: Optional[Union[Verbosity, int]] = Verbosity.ITER,
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,
max_iter: int = 100,
verbose: Union[Verbosity, int] = Verbosity.ITER,
output: Optional[TextIO] = None,
) -> Tuple[np.ndarray, List[IterLog]]:
"""Nonlinear Least Squares minimization with box bounds.
@@ -137,10 +177,11 @@ def least_squares(
bounds: optional pair of lower and upper bounds on the solution.
jacobian: optional function that returns Jacobian of the residual at a given
point and residual. If not given, `residual` will be finite-differenced.
eps: log10 of the perurbation used for automatic finite-differencing.
mu_min: log10 of the minimum value of the regularizer.
mu_max: log10 of the maximum value of the regularizer.
mu_delta: log10 of the factor increasing or decreasing the regularizer.
eps: perurbation used for automatic finite-differencing.
central: whether to use central differences.
mu_min: minimum value of the regularizer.
mu_max: maximum value of the regularizer.
mu_factor: factor increasing or decreasing the regularizer.
tol: termination tolerance on the step size.
max_iter: maximum number of iterations.
verbose: verbosity level.
@@ -155,29 +196,43 @@ def least_squares(
# Convert verbosity to int.
verbose = Verbosity(verbose).value
# Constant for Armijo's sufficient-reduction rule
armijo_c1 = 1e-2
# Initialize locals.
x = x0.copy()
mu = -np.inf # Optimistically start with no regularization.
n = x.size
status = Status.MAX_ITER
i = 0
trace = []
x = x0.astype(np.float64)
n = x.size
xnew = np.zeros((n,))
dx = np.zeros((n,))
scratch = np.zeros((n, n + 7))
dx_norm = 0.0
xnew = np.zeros((n,))
status = Status.MAX_ITER
eps = np.float64(eps)
mu = np.float64(0.0) # Optimistically start with no regularization.
n_reduc = 0 # Number of sequential mu reductions.
# Initialize logging.
trace = []
n_res = 0
n_jac = 0
t_res = 0.0
t_jac = 0.0
t_qp = 0.0
# Regularization control functions.
if mu_factor <= 1:
raise ValueError('mu_factor must be > 1.')
# Decrease mu agressively: sequential decreases grow exponentially.
def decrease_mu(mu, n_reduc):
dmu = (1/mu_factor) ** (2**n_reduc)
mu = 0.0 if mu * dmu < mu_min else mu * dmu
n_reduc += 1
return mu, n_reduc
# Increase mu carefully: always increase by mu_factor.
def increase_mu(mu):
return min(mu_max, max(mu_min, mu_delta + mu))
def decrease_mu(mu):
return -np.inf if mu - mu_delta < mu_min else mu - mu_delta
mu = max(mu_min, mu_factor * mu)
n_reduc = 0 # Reset n_reduc.
return mu, n_reduc
if bounds is not None:
# Checks bounds.
@@ -192,6 +247,10 @@ def least_squares(
# Clip.
np.clip(x, bounds[0], bounds[1], out=x)
# Check for NaNs.
if not np.all(np.isfinite(x)):
raise ValueError('x0 must be finite.')
# Get initial residual.
t_start = time.time()
r = residual(x)
@@ -199,6 +258,9 @@ def least_squares(
t_res += time.time() - t_start
n_res += 1
if r.dtype != np.float64:
raise ValueError('residual function must return float64 arrays.')
# Minimize.
for i in range(max_iter):
if status != Status.MAX_ITER:
@@ -210,9 +272,8 @@ def least_squares(
# Get Jacobian jac.
t_start = time.time()
if jacobian is None:
jac = jacobian_fd(residual, x, r, 10**eps, bounds)
jac, n_res = jacobian_fd(residual, x, r, eps, central, n_res, bounds)
t_res += time.time() - t_start
n_res += n
else:
jac = jacobian(x, r)
t_jac += time.time() - t_start
@@ -221,31 +282,28 @@ def least_squares(
# Get gradient, Gauss-Newton Hessian.
grad = jac.T @ r
hess = jac.T @ jac
gnorm = np.linalg.norm(grad)
# Bounds relative to x
dbounds = [None, None] if bounds is None else [bounds[0] - x, bounds[1] - x]
dlower = None if bounds is None else bounds[0] - x
dupper = None if bounds is None else bounds[1] - x
# Find some reduction.
reduction = -1
while reduction < 0:
# Increase mu until factorizabl.
# Find reduction satisfying Armijo's rule.
armijo = -1
reduction = 0.0
while armijo < 0:
# Increase mu until factorizable.
factorizable = False
while not factorizable:
# Formula from https://arxiv.org/abs/2112.02089
reg = np.sqrt(gnorm * 10**mu) * np.eye(n)
t_start = time.time()
nfree = mujoco.mju_boxQP(
dx, scratch, None, hess + reg, grad, dbounds[0], dbounds[1]
n_free = mujoco.mju_boxQP(
dx, scratch, None, hess + mu * np.eye(n), grad, dlower, dupper
)
t_qp += time.time() - t_start
if nfree > -1:
if n_free >= 0:
factorizable = True
elif mu >= mu_max:
status = Status.FACTORIZATION_FAILED
break
else:
mu += mu_delta
mu, n_reduc = increase_mu(mu)
if status != Status.MAX_ITER:
break
@@ -260,12 +318,13 @@ def least_squares(
# New objective, evaluate reduction.
ynew = 0.5 * rnew.dot(rnew)
reduction = y - ynew
armijo = reduction + armijo_c1*grad.dot(dx)
if reduction < 0:
if armijo < 0:
if mu >= mu_max:
status = Status.NO_IMPORVEMENT
break
mu = increase_mu(mu)
mu, n_reduc = increase_mu(mu)
if status != Status.MAX_ITER:
break
@@ -273,19 +332,23 @@ def least_squares(
# Compute reduction ratio.
expected_reduction = -(grad.dot(dx) + 0.5 * dx.T @ hess @ dx)
reduction_ratio = 0.0
if expected_reduction == 0:
print('Zero expected reduction: exact minimum found?', file=output)
elif expected_reduction < 0:
print('Negative expected reduction: should not occur.', file=output)
if expected_reduction <= 0:
if verbose > Verbosity.SILENT.value:
if expected_reduction == 0:
print('Zero expected reduction: exact minimum found?', file=output)
elif expected_reduction < 0:
print('Negative expected reduction: should not occur.', file=output)
else:
reduction_ratio = reduction / expected_reduction
# Iteration message.
dx_norm = np.linalg.norm(dx)
if verbose >= Verbosity.ITER.value:
logmu = np.log10(mu) if mu > 0 else -np.inf
message = (
f'iter: {i:<3d} y: {y:<8.3g} mu: {mu:>4.1f} '
f'ratio: {reduction_ratio:<5.2g} '
f'dx: {dx_norm:<8.3g} reduction: {reduction:<8.3g}'
f'iter: {i:<3d} y: {y:<9.4g} log10mu: {logmu:>4.1f} '
f'ratio: {reduction_ratio:<7.2g} '
f'dx: {dx_norm:<7.2g} reduction: {reduction:<7.2g}'
)
print(message, file=output)
@@ -296,39 +359,43 @@ def least_squares(
trace.append(log)
# Check for success.
dx_norm = np.linalg.norm(dx)
if dx_norm < tol:
status = Status.DX_TOL
break
# Modify regularizer like in (Bazaraa, Sherali, and Shetty)
if reduction_ratio > 0.75:
mu = decrease_mu(mu)
mu, n_reduc = decrease_mu(mu, n_reduc)
elif reduction_ratio < 0.25:
mu = increase_mu(mu)
mu, n_reduc = increase_mu(mu)
# Accept proposal.
x = xnew
r = rnew
# Append final log to trace.
# Note: unlike other iter logs, this is at the end point.
yfinal = 0.5 * r.dot(r)
red = np.float64(0.0)
log = IterLog(candidate=x, objective=yfinal, reduction=red, regularizer=mu)
trace.append(log)
# Print final diagnostics.
if verbose > Verbosity.SILENT.value:
message = f'Terminated after {i} iterations: '
message += _STATUS_MESSAGE[status]
message += f' Residual evals: {n_res:d}'
message += f' y: {yfinal:<.4g}, Residual evals: {n_res:d}'
if n_jac > 0:
message += f', Jacobian evals: {n_jac:d}'
print(message, file=output)
time_total = time.time() - t_start_total
if time_total > 0:
qp_percent = 100 * t_qp / time_total
r_percent = 100 * t_res / time_total
time_scale = 1 if time_total > 1 else 1000
time_units = 's' if time_total > 1 else 'ms'
message = f'total time {time_scale * time_total:<.1f}{time_units}'
message += f' of which QP {qp_percent:<.1f}%, residual {r_percent:<.1f}%'
message += f' of which residual {r_percent:<.1f}%'
if t_jac > 0:
jac_percent = 100 * t_jac / time_total
message += f' Jacobian {jac_percent:<.1f}%'
python/mujoco/minimize_test.py
+40 -31
View File
@@ -25,18 +25,19 @@ import numpy as np
class MinimizeTest(absltest.TestCase):
def test_basic(self) -> None:
def residual(x: np.ndarray) -> float:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
def residual(x: np.ndarray) -> np.ndarray:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)], dtype=np.float64)
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(step) < tol')
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')
def test_start_at_minimum(self) -> None:
def residual(x: np.ndarray) -> float:
def residual(x: np.ndarray) -> np.ndarray:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
out = io.StringIO()
@@ -44,11 +45,11 @@ 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.assertContainsSubsequence(out.getvalue(), 'norm(step) < tol')
self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol')
self.assertContainsSubsequence(out.getvalue(), 'exact minimum found')
def test_jac_callback(self) -> None:
def residual(x: np.ndarray) -> float:
def residual(x: np.ndarray) -> np.ndarray:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
def jacobian(x: np.ndarray, r: np.ndarray) -> Tuple[float, np.ndarray]:
@@ -60,19 +61,19 @@ class MinimizeTest(absltest.TestCase):
x, _ = minimize.least_squares(x0, residual, jacobian=jacobian, output=out)
expected_x = np.array((1.0, 1.0))
np.testing.assert_array_almost_equal(x, expected_x)
self.assertContainsSubsequence(out.getvalue(), 'norm(step) < tol')
self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol')
# 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(), 'no improvement found.')
self.assertContainsSubsequence(out1.getvalue(), 'insufficient reduction')
def test_max_iter(self) -> None:
dim = 20 # High-D Rosenbrock
def residual(x: np.ndarray) -> float:
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)]
return np.asarray(res0 + res1)
@@ -89,7 +90,7 @@ class MinimizeTest(absltest.TestCase):
np.testing.assert_array_almost_equal(x, expected_x)
def test_bounds(self) -> None:
def residual(x: np.ndarray) -> float:
def residual(x: np.ndarray) -> np.ndarray:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
out = io.StringIO()
@@ -104,27 +105,35 @@ class MinimizeTest(absltest.TestCase):
x, _ = minimize.least_squares(x0, residual, bounds=bounds_types['inbounds'],
output=out)
np.testing.assert_array_almost_equal(x, expected_x)
self.assertContainsSubsequence(out.getvalue(), 'norm(step) < tol')
self.assertContainsSubsequence(out.getvalue(), 'norm(dx) < tol')
# Test different bounds conditions.
verbose = minimize.Verbosity.FULLITER
for bounds in bounds_types.values():
out = io.StringIO()
x, trace = minimize.least_squares(x0, residual, bounds=bounds, output=out,
verbose=verbose)
self.assertContainsSubsequence(out.getvalue(), 'norm(step) < tol')
grad = trace[-1].jacobian.T @ trace[-1].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 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)
def test_bad_bounds(self) -> None:
def residual(x: np.ndarray) -> float:
def residual(x: np.ndarray) -> np.ndarray:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
out = io.StringIO()