Merge pull request #3079 from aftersomemath:sysid-pr
PiperOrigin-RevId: 868229512 Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea
This commit is contained in:
Copybara-Service
@@ -0,0 +1,259 @@
|
||||
# Copyright 2026 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.
|
||||
# ==============================================================================
|
||||
|
||||
"""Optimization routines for system identification."""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Literal
|
||||
|
||||
from absl import logging
|
||||
from mujoco import minimize as mujoco_minimize
|
||||
from mujoco.sysid._src import parameter
|
||||
import numpy as np
|
||||
import scipy.optimize as scipy_optimize
|
||||
import scipy.special
|
||||
|
||||
|
||||
def _scipy_least_squares(
|
||||
x0: np.ndarray,
|
||||
residual_fn: Callable[..., Any],
|
||||
bounds: tuple[np.ndarray, np.ndarray],
|
||||
use_mujoco_jac: bool = False,
|
||||
**kwargs,
|
||||
) -> scipy_optimize.OptimizeResult:
|
||||
"""Run scipy least_squares with optional MuJoCo finite-difference Jacobian."""
|
||||
max_nfev = kwargs.pop("max_iters", 200)
|
||||
if kwargs.pop("verbose", True):
|
||||
verbose = 2
|
||||
else:
|
||||
verbose = 0
|
||||
x_scale = kwargs.pop("x_scale", "jac")
|
||||
loss = kwargs.pop("loss", "linear")
|
||||
|
||||
jac_arg: str | Callable[..., Any]
|
||||
if use_mujoco_jac:
|
||||
# This is the default step sized for finite difference used in
|
||||
# scipy's least_squares and mujoco's minimize finite difference
|
||||
# https://github.com/scipy/scipy/blob/91e18f3bd355477b
|
||||
# 8b7747ec82d70ac98ffd2422/scipy/optimize/_numdiff.py#L404
|
||||
eps = np.finfo(np.float64).eps ** 0.5
|
||||
if "diff_step" in kwargs:
|
||||
eps = kwargs.pop("diff_step")
|
||||
|
||||
def _jac_fn(x):
|
||||
return mujoco_minimize.jacobian_fd(
|
||||
residual=residual_fn,
|
||||
x=x.reshape((-1, 1)),
|
||||
r=residual_fn(x).reshape((-1, 1)),
|
||||
eps=eps,
|
||||
n_res=0,
|
||||
bounds=[bounds[0].reshape((-1, 1)), bounds[1].reshape((-1, 1))],
|
||||
)[0]
|
||||
|
||||
jac_arg = _jac_fn
|
||||
else:
|
||||
jac_arg = "2-point"
|
||||
|
||||
return scipy_optimize.least_squares(
|
||||
residual_fn,
|
||||
x0,
|
||||
bounds=bounds,
|
||||
max_nfev=max_nfev,
|
||||
verbose=verbose,
|
||||
x_scale=x_scale,
|
||||
loss=loss,
|
||||
jac=jac_arg, # pyright: ignore[reportArgumentType]
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _mujoco_least_squares(
|
||||
x0: np.ndarray,
|
||||
residual_fn: Callable[..., Any],
|
||||
bounds: tuple[np.ndarray, np.ndarray],
|
||||
**kwargs,
|
||||
) -> scipy_optimize.OptimizeResult:
|
||||
"""Run MuJoCo's native least_squares optimizer."""
|
||||
if kwargs.pop("verbose", True):
|
||||
verbose = mujoco_minimize.Verbosity.FULLITER
|
||||
else:
|
||||
verbose = mujoco_minimize.Verbosity.SILENT
|
||||
max_iter = kwargs.pop("max_iters", 200)
|
||||
x, log = mujoco_minimize.least_squares(
|
||||
x0=x0,
|
||||
bounds=bounds,
|
||||
residual=residual_fn,
|
||||
verbose=verbose,
|
||||
max_iter=max_iter,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# If verbose, return the full optimization log.
|
||||
extras = {}
|
||||
if verbose == mujoco_minimize.Verbosity.FULLITER:
|
||||
extras["objective"] = [entry.objective for entry in log]
|
||||
extras["candidate"] = [entry.candidate[:, 0] for entry in log]
|
||||
|
||||
return scipy_optimize.OptimizeResult(
|
||||
x=x,
|
||||
jac=log[-1].jacobian,
|
||||
grad=log[-1].grad,
|
||||
extras=extras,
|
||||
)
|
||||
|
||||
|
||||
def _dispatch_optimizer(
|
||||
x0: np.ndarray,
|
||||
residual_fn: Callable[..., Any],
|
||||
bounds: tuple[np.ndarray, np.ndarray],
|
||||
optimizer: Literal["scipy", "mujoco", "scipy_parallel_fd"],
|
||||
**kwargs,
|
||||
) -> scipy_optimize.OptimizeResult:
|
||||
"""Dispatch to the appropriate least-squares backend."""
|
||||
if optimizer in ["scipy", "scipy_parallel_fd"]:
|
||||
return _scipy_least_squares(
|
||||
x0,
|
||||
residual_fn,
|
||||
bounds,
|
||||
use_mujoco_jac=optimizer == "scipy_parallel_fd",
|
||||
**kwargs,
|
||||
)
|
||||
elif optimizer == "mujoco":
|
||||
return _mujoco_least_squares(x0, residual_fn, bounds, **kwargs)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported optimizer: '{optimizer}'. Expected one of: 'scipy',"
|
||||
" 'scipy_parallel_fd', or 'mujoco'."
|
||||
)
|
||||
|
||||
|
||||
def optimize(
|
||||
initial_params: parameter.ParameterDict,
|
||||
residual_fn: Callable[..., Any],
|
||||
optimizer: Literal["scipy", "mujoco", "scipy_parallel_fd"] = "mujoco",
|
||||
**optimizer_kwargs,
|
||||
) -> tuple[parameter.ParameterDict, scipy_optimize.OptimizeResult]:
|
||||
"""Run nonlinear least-squares optimization on the residual.
|
||||
|
||||
Args:
|
||||
initial_params: Starting parameter values and bounds.
|
||||
residual_fn: Callable with signature ``(x, params) -> (residuals, ...)`` as
|
||||
returned by :func:`build_residual_fn`.
|
||||
optimizer: Backend — ``"mujoco"`` (default), ``"scipy"``, or
|
||||
``"scipy_parallel_fd"`` (scipy with MuJoCo finite-difference Jacobian).
|
||||
**optimizer_kwargs: Forwarded to the backend (e.g. ``max_iters``,
|
||||
``verbose``, ``loss``).
|
||||
|
||||
Returns:
|
||||
``(opt_params, opt_result)`` — the optimised ParameterDict and a
|
||||
``scipy.optimize.OptimizeResult`` with at least ``x``, ``jac``, ``grad``.
|
||||
"""
|
||||
x0 = initial_params.as_vector()
|
||||
bounds = initial_params.get_bounds()
|
||||
opt_params = initial_params.copy()
|
||||
|
||||
# Check if there are any parameters to optimize.
|
||||
if not opt_params or opt_params.size == 0:
|
||||
logging.warning(
|
||||
"The ParameterDict is empty or contains only frozen Parameters. "
|
||||
"Please declare all Parameters that need to be optimized."
|
||||
)
|
||||
return opt_params, scipy_optimize.OptimizeResult(
|
||||
x=x0,
|
||||
jac=np.zeros((0, x0.shape[0])),
|
||||
grad=np.zeros_like(x0),
|
||||
extras={},
|
||||
)
|
||||
|
||||
def optimized_residual_fn(x):
|
||||
residuals, _, _ = residual_fn(x, opt_params)
|
||||
return np.concatenate(residuals)
|
||||
|
||||
opt_result = _dispatch_optimizer(
|
||||
x0, optimized_residual_fn, bounds, optimizer, **optimizer_kwargs
|
||||
)
|
||||
|
||||
opt_params.update_from_vector(opt_result.x)
|
||||
|
||||
return opt_params, opt_result
|
||||
|
||||
|
||||
def calculate_intervals(
|
||||
residuals_star,
|
||||
J,
|
||||
alpha=0.05,
|
||||
lambda_zero_thresh=1e-15,
|
||||
v_zero_thresh=1e-8,
|
||||
):
|
||||
"""Calculate confidence intervals from the Jacobian at the optimum."""
|
||||
if J is None or J.size == 0:
|
||||
return np.empty((0, 0)), np.empty((0,))
|
||||
|
||||
# TODO(levi): account for per sensor variance
|
||||
# Estimate sensor variance by assuming a good model fit, so
|
||||
# remaining variance in the residual is due to sensor noise.
|
||||
# Dividing by n - p is an unbiased estimate of the noise.
|
||||
final_r = np.concatenate(residuals_star)
|
||||
s2 = np.dot(final_r, final_r) / (final_r.size - J.shape[1])
|
||||
H = J.T @ J
|
||||
|
||||
# Calculate the diagonals of the inverse of H
|
||||
# using the observation that division by zero
|
||||
# of eig(H) close to zero is canceled by numerically
|
||||
# zero elements of the eigenvectors
|
||||
# That is numerically zero eigenvalues only
|
||||
# cause a confidence bound to be infinite if that eigenvalue
|
||||
# has a numerically non-zero effect on the considered parameter
|
||||
lamb, V = np.linalg.eigh(H)
|
||||
lamb_max = np.max(lamb)
|
||||
diag_inv_H = []
|
||||
for j in range(H.shape[0]):
|
||||
inv_H_jj = 0.0
|
||||
v_j_max = np.max(np.abs(V[:, j]))
|
||||
for i in range(H.shape[0]):
|
||||
lambda_i = lamb[i]
|
||||
if lambda_i / lamb_max < lambda_zero_thresh:
|
||||
lambda_i = 0.0
|
||||
|
||||
v_j_i = V[j, i]
|
||||
if np.abs(v_j_i / v_j_max) < v_zero_thresh:
|
||||
v_j_i = 0.0
|
||||
|
||||
if lambda_i == 0.0 and v_j_i != 0.0:
|
||||
inv_H_jj += np.inf
|
||||
elif lambda_i == 0.0 and v_j_i == 0.0:
|
||||
pass
|
||||
else:
|
||||
inv_H_jj += v_j_i**2 / lambda_i
|
||||
diag_inv_H.append(inv_H_jj)
|
||||
diag_inv_H = np.array(diag_inv_H)
|
||||
|
||||
# In general eigenvalue decomposition should be more accurate
|
||||
# than calculating the inverse of H using a general method
|
||||
# TODO(levi): expand the eigenvalue/eigenvector element
|
||||
# cancelation above to the full inverse matrix
|
||||
# inv_H = V @ np.diag(np.divide(
|
||||
# 1, lamb, out=np.inf*np.zeros_like(lamb),
|
||||
# where=lamb != 0.0)) @ V.T
|
||||
lamb[lamb == 0] = lambda_zero_thresh
|
||||
inv_H = V @ np.diag(1 / lamb) @ V.T
|
||||
# print('inv test')
|
||||
# print(np.diag(inv_H @ H))
|
||||
# print(np.diag(np.linalg.inv(H) @ H)))
|
||||
Sigma_X = s2 * inv_H
|
||||
intervals = np.sqrt(diag_inv_H * s2) * scipy.special.stdtrit(
|
||||
final_r.size - J.shape[1], 1 - alpha / 2
|
||||
)
|
||||
return Sigma_X, intervals
|
||||
Reference in New Issue
Block a user