Merge branch 'google-deepmind:main' into tendons

This commit is contained in:
Abhishek Joshi
2024-06-03 11:14:21 -05:00
committed by GitHub
173 changed files with 9083 additions and 40007 deletions
+16 -22
View File
@@ -10,27 +10,22 @@
"\n",
"# <h1><center>LQR tutorial <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a></center></h1>\n",
"\n",
"This notebook provides an example of an LQR controller using [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LBAvTJ0xHKy7"
},
"source": [
"### Copyright notice"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_UbO9uhtBSX5"
},
"source": [
"> <p><small><small>Copyright 2022 DeepMind Technologies Limited</small></p>\n",
"> <p><small><small>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 <a href=\"http://www.apache.org/licenses/LICENSE-2.0\">http://www.apache.org/licenses/LICENSE-2.0</a>.</small></small></p>\n",
"> <p><small><small>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.</small></small></p>"
"This notebook provides an example of an LQR controller using [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme).\n",
"\n",
"<!-- Copyright 2021 DeepMind Technologies Limited\n",
"\n",
" Licensed under the Apache License, Version 2.0 (the \"License\");\n",
" you may not use this file except in compliance with the License.\n",
" You may obtain a copy of the License at\n",
"\n",
" http://www.apache.org/licenses/LICENSE-2.0\n",
"\n",
" Unless required by applicable law or agreed to in writing, software\n",
" distributed under the License is distributed on an \"AS IS\" BASIS,\n",
" WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
" See the License for the specific language governing permissions and\n",
" limitations under the License.\n",
"-->"
]
},
{
@@ -935,7 +930,6 @@
"accelerator": "GPU",
"colab": {
"collapsed_sections": [
"LBAvTJ0xHKy7",
"QPdJNe3k62mx"
],
"private_outputs": true
+3 -1
View File
@@ -40,7 +40,9 @@ numpy==1.26.0; python_version >= '3.9' \
--hash=sha256:914b28d3215e0c721dc75db3ad6d62f51f630cb0c277e6b3bcb39519bed10bd8 \
--hash=sha256:c78a22e95182fb2e7874712433eaa610478a3caf86f28c621708d35fa4fd6e7f \
--hash=sha256:86f737708b366c36b76e953c46ba5827d8c27b7a8c9d0f471810728e5a2fe57c \
--hash=sha256:020cdbee66ed46b671429c7265cf00d8ac91c046901c55684954c3958525dab2
--hash=sha256:020cdbee66ed46b671429c7265cf00d8ac91c046901c55684954c3958525dab2 \
--hash=sha256:d6fa6d17727169ff1385ad3cb8f290bbcc3f2097322d90507c1956a4f9f870fc \
--hash=sha256:79d9c3c363cbd919d879dd38aeb19d96be8699eda2af5f3f3c97bf774e1e6438
pip==23.3.1 \
--hash=sha256:55eb67bb6171d37447e82213be585b75fe2b12b359e993773aca4de9247a052b
PyOpenGL==3.1.7 \
+827 -635
View File
File diff suppressed because it is too large Load Diff
+3 -3
View File
@@ -84,7 +84,7 @@ if(NOT TARGET mujoco)
if(MUJOCO_FRAMEWORK)
message("MuJoCo framework is at ${MUJOCO_FRAMEWORK}/mujoco.framework")
set(MUJOCO_LIBRARY
${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.1.5.dylib
${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.1.6.dylib
)
target_compile_options(mujoco INTERFACE -F${MUJOCO_FRAMEWORK})
endif()
@@ -92,7 +92,7 @@ if(NOT TARGET mujoco)
if(NOT MUJOCO_FRAMEWORK)
find_library(
MUJOCO_LIBRARY mujoco mujoco.3.1.5 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
MUJOCO_LIBRARY mujoco mujoco.3.1.6 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
)
find_path(MUJOCO_INCLUDE mujoco/mujoco.h HINTS ${MUJOCO_INCLUDE_DIR} REQUIRED)
message("MuJoCo is at ${MUJOCO_LIBRARY}")
@@ -140,7 +140,7 @@ findorfetch(
GIT_REPO
https://github.com/abseil/abseil-cpp
GIT_TAG
2f9e432cce407ce0ae50676696666f33a77d42ac # LTS 20240116.1
d7aaad83b488fd62bd51c81ecf16cd938532cc0a # LTS 20240116.2
TARGETS
${MUJOCO_PYTHON_ABSL_TARGETS}
EXCLUDE_FROM_ALL
+8
View File
@@ -1194,6 +1194,14 @@ Euler integrator, semi-implicit in velocity.
mujoco.mjd_inverseFD(self.model, self.data, eps, flg_centered,
None, None, None, None, None, None, None)
def test_geom_distance(self):
mujoco.mj_forward(self.model, self.data)
fromto = np.empty(6, np.float64)
dist = mujoco.mj_geomDistance(self.model, self.data, 0, 2, 200, fromto)
self.assertEqual(dist, 41.9)
np.testing.assert_array_equal(fromto,
np.array((42., 0., 0., 42., 0., 41.9)))
def test_inverse_fd(self):
eps = 1e-6
flg_centered = 0
+12
View File
@@ -531,6 +531,18 @@ PYBIND11_MODULE(_functions, pymodule) {
Def<traits::mj_objectVelocity>(pymodule);
Def<traits::mj_objectAcceleration>(pymodule);
Def<traits::mj_contactForce>(pymodule);
Def<traits::mj_geomDistance>(
pymodule,
[](const raw::MjModel* m, const raw::MjData* d,
int geom1, int geom2, mjtNum distmax,
std::optional<Eigen::Ref<EigenArrayXX>> fromto) {
if (fromto.has_value() && fromto->size() != 6) {
throw py::type_error("fromto should be of size 6");
}
return InterceptMjErrors(::mj_geomDistance)(
m, d, geom1, geom2, distmax,
fromto.has_value() ? fromto->data() : nullptr);
});
Def<traits::mj_differentiatePos>(
pymodule,
[](const raw::MjModel* m, Eigen::Ref<EigenVectorX> qvel,
+240 -118
View File
@@ -14,10 +14,11 @@
# ==============================================================================
"""Nonlinear Least Squares minimization with box bounds."""
import abc
import dataclasses
import enum
import time
from typing import Callable, List, Optional, TextIO, Tuple, Union
from typing import Callable, List, Optional, Sequence, TextIO, Tuple, Union
import mujoco
import numpy as np
@@ -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,124 +69,109 @@ class IterLog:
step: Optional[np.ndarray] = None
def jacobian_fd(
residual: Callable[[np.ndarray], np.ndarray],
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.
class Norm(abc.ABC):
"""Abstract interface for norm functions, measuring the magnitude of vectors.
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.
central: whether to use central differences.
n_res: number or residual evaluations so far.
bounds: optional pair of lower and upper bounds.
Key Concepts:
Returns:
jac: Jacobian of the residual at x.
n_res: updated number of residual evaluations.
* Norm Value: The value of the norm for a given input vector.
* Gradient and Hessian: The gradient (first derivative) and Hessian (second
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
jac = np.zeros((nr, nx))
xh = x.copy()
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:
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
@abc.abstractmethod
def value(self, r: np.ndarray) -> np.float64:
"""Returns the value of the norm at the input vector `y = norm(r)`."""
pass
@abc.abstractmethod
def grad_hess(self, r: np.ndarray, proj: np.ndarray):
"""Computes the projected gradient and Hessian of the norm at `r`.
Args:
r: A NumPy column vector (nr x 1).
proj: A pre-computed projection matrix (nr x nx).
Returns:
A tuple containing:
* Projected gradient: proj.T @ (d_norm/d_r).
* Projected Hessian: proj.T @ (d^2_norm/d_r^2) @ proj.
"""
pass
class Quadratic(Norm):
"""Implementation of the quadratic norm."""
def value(self, r: np.ndarray):
"""Returns the quadratic norm of `r`."""
return 0.5 * (r.T @ r).item()
def grad_hess(self, r: np.ndarray, proj: np.ndarray):
"""Computes the projected gradient and Hessian of the quadratic norm at `r`.
Args:
r: A NumPy column vector (nr x 1).
proj: A pre-computed projection matrix (nr x nx).
Returns:
A tuple containing:
* Projected gradient: `proj.T @ r`.
* Projected Hessian: `proj.T @ proj`.
"""
grad = proj.T @ r
hess = proj.T @ proj # Notionally proj.T @ np.eye(r.size) @ proj
return grad, hess
def least_squares(
x0: np.ndarray,
residual: Callable[[np.ndarray], np.ndarray],
bounds: Optional[List[np.ndarray]] = None,
bounds: Optional[Sequence[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
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
x0: Initial guess
residual: Vectorized function returning the residual for 1 or more points.
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: 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.
output: optional file or StringIO to which to print messages.
norm: Norm object returning norm scalar or its projected gradient and
Hessian. See Norm class for detailed documentation.
eps: Perurbation used for automatic finite-differencing.
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.
max_iter: Maximum number of iterations.
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 = x.size
xnew = np.zeros((n,))
dx = np.zeros((n,))
n = x0.size
x = x0.astype(np.float64).reshape((n, 1))
xnew = np.zeros((n, 1))
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 list(bounds)
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
hess = jac.T @ jac
grad, hess = norm.grad_hess(r, 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.
# Note: unlike other iter logs, this is at the end point.
yfinal = 0.5 * r.dot(r)
red = np.float64(0.0)
# 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.
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')
+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()
+4 -4
View File
@@ -7,13 +7,13 @@
<key>CFBundleIdentifier</key>
<string>org.mujoco.mjpython</string>
<key>CFBundleVersion</key>
<string>3.1.5</string>
<string>3.1.6</string>
<key>CFBundleGetInfoString</key>
<string>3.1.5</string>
<string>3.1.6</string>
<key>CFBundleLongVersionString</key>
<string>3.1.5</string>
<string>3.1.6</string>
<key>CFBundleShortVersionString</key>
<string>3.1.5</string>
<string>3.1.6</string>
<key>CFBundleExecutable</key>
<string>mjpython</string>
<key>CFBundleIconFile</key>
+3 -3
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "mujoco"
version = "3.1.5"
version = "3.1.6"
authors = [
{name = "Google DeepMind", email = "mujoco@deepmind.com"},
]
@@ -36,9 +36,9 @@ dynamic = ["readme", "scripts"]
[project.urls]
Homepage = "https://github.com/google-deepmind/mujoco"
Documentation = "https://mujoco.readthedocs.io/en/3.1.5"
Documentation = "https://mujoco.readthedocs.io/en/3.1.6"
Repository = "https://github.com/google-deepmind/mujoco"
Changelog = "https://mujoco.readthedocs.io/en/3.1.5/changelog.html"
Changelog = "https://mujoco.readthedocs.io/en/3.1.6/changelog.html"
[tool.setuptools]
include-package-data = false
+16 -22
View File
@@ -10,27 +10,22 @@
"\n",
"# <h1><center>Tutorial <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a></center></h1>\n",
"\n",
"This notebook provides an introductory tutorial for [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme), using the native Python bindings."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xBSdkbmGN2K-"
},
"source": [
"### Copyright notice"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_UbO9uhtBSX5"
},
"source": [
"> <p><small><small>Copyright 2022 DeepMind Technologies Limited.</small></p>\n",
"> <p><small><small>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 <a href=\"http://www.apache.org/licenses/LICENSE-2.0\">http://www.apache.org/licenses/LICENSE-2.0</a>.</small></small></p>\n",
"> <p><small><small>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.</small></small></p>"
"This notebook provides an introductory tutorial for [**MuJoCo** physics](https://github.com/google-deepmind/mujoco#readme), using the native Python bindings.\n",
"\n",
"<!-- Copyright 2021 DeepMind Technologies Limited\n",
"\n",
" Licensed under the Apache License, Version 2.0 (the \"License\");\n",
" you may not use this file except in compliance with the License.\n",
" You may obtain a copy of the License at\n",
"\n",
" http://www.apache.org/licenses/LICENSE-2.0\n",
"\n",
" Unless required by applicable law or agreed to in writing, software\n",
" distributed under the License is distributed on an \"AS IS\" BASIS,\n",
" WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
" See the License for the specific language governing permissions and\n",
" limitations under the License.\n",
"-->"
]
},
{
@@ -2122,7 +2117,6 @@
"accelerator": "GPU",
"colab": {
"collapsed_sections": [
"xBSdkbmGN2K-",
"YvyGCsgSCxHQ"
],
"gpuClass": "premium",