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

Add mujoco.minimize module.

Browse Source 
PiperOrigin-RevId: 608443678
Change-Id: Ib5dc78fe4523488b6c4e824f226591cf6a8a6216
This commit is contained in:
Yuval Tassa
2024-02-19 18:57:39 -08:00
committed by Copybara-Service
parent 24eb4c9f09
commit b56d4bf8a5
2 changed files with 482 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/minimize.py
+337
View File
@@ -0,0 +1,337 @@
# Copyright 2024 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Nonlinear Least Squares minimization with box bounds."""
import dataclasses
import enum
import time
from typing import Callable, List, Optional, TextIO, Tuple, Union
import mujoco
import numpy as np
class Verbosity(enum.Enum):
SILENT = 0
FINAL = 1
ITER = 2
FULLITER = 3
class Status(enum.Enum):
FACTORIZATION_FAILED = enum.auto()
NO_IMPORVEMENT = enum.auto()
MAX_ITER = enum.auto()
DX_TOL = enum.auto()
_STATUS_MESSAGE = {
Status.FACTORIZATION_FAILED: 'factorization failed.',
Status.NO_IMPORVEMENT: 'no improvement found.',
Status.MAX_ITER: 'maximum iterations reached.',
Status.DX_TOL: 'norm(step) < tolerance.',
}
@dataclasses.dataclass(frozen=True)
class IterLog:
"""Log of a single iteration of the non-linear least-squares solver.
Attributes:
candidate: Value of the decision variable at the beginning this iteration.
objective: Value of the objective at the candidate.
reduction: Reduction of the objective during this iteration.
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.
step: Optional change in decision variable during this iteration.
"""
candidate: np.ndarray
objective: np.float64
reduction: np.float64
regularizer: np.float64
residual: Optional[np.ndarray] = None
jacobian: Optional[np.ndarray] = None
step: Optional[np.ndarray] = None
def jacobian_fd(
residual: Callable[[np.ndarray], np.ndarray],
x: np.ndarray,
r: np.ndarray,
eps: float,
bounds: Optional[List[np.ndarray]] = None,
):
"""Finite-difference Jacobian of a residual function.
Args:
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.
bounds: optional pair of lower and upper bounds of the solution.
Returns:
jac: Jacobian of the residual at x.
"""
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
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
def least_squares(
x0: np.ndarray,
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,
output: Optional[TextIO] = None,
) -> Tuple[np.ndarray, List[IterLog]]:
"""Nonlinear Least Squares minimization with box bounds.
Args:
x0: initial guess
residual: function that returns the residual for a given point x.
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.
tol: termination tolerance on the step size.
max_iter: maximum number of iterations.
verbose: verbosity level.
output: optional file or StringIO to which to print messages.
Returns:
x: best solution found
trace: sequence of solution iterates.
"""
t_start_total = time.time()
# Convert verbosity to int.
verbose = Verbosity(verbose).value
# Initialize locals.
x = x0.copy()
mu = -np.inf # Optimistically start with no regularization.
n = x.size
i = 0
trace = []
dx = np.zeros((n,))
scratch = np.zeros((n, n + 7))
dx_norm = 0.0
xnew = np.zeros((n,))
status = Status.MAX_ITER
n_res = 0
n_jac = 0
t_res = 0.0
t_jac = 0.0
t_qp = 0.0
# Regularization control functions.
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
if bounds is not None:
# Checks bounds.
if len(bounds) != 2:
raise ValueError('bounds must have 2 elements.')
if bounds[0].size != n or bounds[1].size != n:
raise ValueError('bounds must have the same size as x0.')
if not np.all(np.isfinite(bounds[0])) or not np.all(np.isfinite(bounds[1])):
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.
np.clip(x, bounds[0], bounds[1], out=x)
# Get initial residual.
t_start = time.time()
r = residual(x)
rnew = r
t_res += time.time() - t_start
n_res += 1
# Minimize.
for i in range(max_iter):
if status != Status.MAX_ITER:
break
# Get objective y.
y = 0.5 * r.dot(r)
# Get Jacobian jac.
t_start = time.time()
if jacobian is None:
jac = jacobian_fd(residual, x, r, 10**eps, bounds)
t_res += time.time() - t_start
n_res += n
else:
jac = jacobian(x, r)
t_jac += time.time() - t_start
n_jac += 1
# 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]
# Find some reduction.
reduction = -1
while reduction < 0:
# Increase mu until factorizabl.
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]
)
t_qp += time.time() - t_start
if nfree > -1:
factorizable = True
elif mu >= mu_max:
status = Status.FACTORIZATION_FAILED
break
else:
mu += mu_delta
if status != Status.MAX_ITER:
break
# New candidate, residual.
xnew = x + dx
t_start = time.time()
rnew = residual(xnew)
t_res += time.time() - t_start
n_res += 1
# New objective, evaluate reduction.
ynew = 0.5 * rnew.dot(rnew)
reduction = y - ynew
if reduction < 0:
if mu >= mu_max:
status = Status.NO_IMPORVEMENT
break
mu = increase_mu(mu)
if status != Status.MAX_ITER:
break
# 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)
else:
reduction_ratio = reduction / expected_reduction
# Iteration message.
if verbose >= Verbosity.ITER.value:
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}'
)
print(message, file=output)
# Append log to trace.
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)
# 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)
elif reduction_ratio < 0.25:
mu = increase_mu(mu)
# Accept proposal.
x = xnew
r = rnew
# 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}'
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}%'
if t_jac > 0:
jac_percent = 100 * t_jac / time_total
message += f' Jacobian {jac_percent:<.1f}%'
print(message, file=output)
return x, trace
python/mujoco/minimize_test.py
+145
View File
@@ -0,0 +1,145 @@
# Copyright 2024 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for minimize.py."""
import io
from typing import Tuple
from absl.testing import absltest
from mujoco import minimize
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)])
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')
def test_start_at_minimum(self) -> None:
def residual(x: np.ndarray) -> float:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
out = io.StringIO()
x0 = np.array((1.0, 1.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')
self.assertContainsSubsequence(out.getvalue(), 'exact minimum found')
def test_jac_callback(self) -> None:
def residual(x: np.ndarray) -> float:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
def jacobian(x: np.ndarray, r: np.ndarray) -> Tuple[float, np.ndarray]:
del r # Unused.
return np.array([[-1, 0], [-20 * x[0], 10]])
x0 = np.array((0.0, 0.0))
out = io.StringIO()
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')
# 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.')
def test_max_iter(self) -> None:
dim = 20 # High-D Rosenbrock
def residual(x: np.ndarray) -> float:
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.
x0 = np.zeros(dim)
out = io.StringIO()
minimize.least_squares(x0, residual, max_iter=20, output=out)
self.assertContainsSubsequence(out.getvalue(), 'maximum iterations')
# Succeed after 100 iterations (default).
x, _ = minimize.least_squares(x0, residual)
expected_x = np.ones(20)
np.testing.assert_array_almost_equal(x, expected_x)
def test_bounds(self) -> None:
def residual(x: np.ndarray) -> float:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
out = io.StringIO()
x0 = np.array((0.0, 0.0))
expected_x = np.array((1.0, 1.0))
bounds_types = {'inbounds': [np.array((-2.0, -2.0)), np.array((2.0, 2.0))],
'onlower': [np.array((-2.0, 2.0)), np.array((0.5, 3.0))],
'onupper': [np.array((-2.0, -2.0)), np.array((0.5, 2.0))]}
# In bounds finds true minimum.
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')
# 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)
def test_bad_bounds(self) -> None:
def residual(x: np.ndarray) -> float:
return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
out = io.StringIO()
x0 = np.array((0.0, 0.0))
bad_bounds = [
[0, 1, 2],
[np.array((-2, 2, 0)), np.array((0.5, 3, 4))],
[np.array((-2, 2, 0)), np.array((0.5, 3, np.inf))],
[np.array((-2, 2, 0)), np.array((-5, 3, 6))],
]
for bounds in bad_bounds:
with self.assertRaises(ValueError):
minimize.least_squares(x0, residual, bounds=bounds, output=out)
if __name__ == '__main__':
absltest.main()