System identification toolbox for MuJoCo.
This resulted from a lengthy collaboration with @kevinzakka, @jonathanembleyriches, @nimrod-gileadi, @gizemozd, @quagla, and @yuval.
This commit is contained in:
Kevin Zakka
@@ -0,0 +1,57 @@
|
||||
"""I/O utilities for saving system identification results."""
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import pickle
|
||||
from collections.abc import Sequence
|
||||
|
||||
import scipy.optimize as scipy_optimize
|
||||
from absl import logging
|
||||
|
||||
from mujoco.sysid._src import parameter
|
||||
from mujoco.sysid._src.optimize import calculate_intervals
|
||||
from mujoco.sysid._src.trajectory import ModelSequences
|
||||
|
||||
|
||||
def save_results(
|
||||
experiment_results_folder: str | os.PathLike,
|
||||
models_sequences: Sequence[ModelSequences],
|
||||
initial_params: parameter.ParameterDict,
|
||||
opt_params: parameter.ParameterDict,
|
||||
opt_result: scipy_optimize.OptimizeResult,
|
||||
residual_fn,
|
||||
):
|
||||
experiment_results_folder = pathlib.Path(experiment_results_folder)
|
||||
if not experiment_results_folder.exists():
|
||||
experiment_results_folder.mkdir(parents=True, exist_ok=True)
|
||||
logging.info("Experiment results will be saved to %s", experiment_results_folder)
|
||||
|
||||
initial_params.save_to_disk(experiment_results_folder / "params_x_0.yaml")
|
||||
opt_params.save_to_disk(experiment_results_folder / "params_x_hat.yaml")
|
||||
|
||||
with open(os.path.join(experiment_results_folder, "results.pkl"), "wb") as handle:
|
||||
pickle.dump(opt_result, handle, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
|
||||
# TODO: these intervals should be part of the params object.
|
||||
residuals_star, _, _ = residual_fn(opt_result.x, opt_params, return_pred_all=True)
|
||||
covariance, intervals = calculate_intervals(residuals_star, opt_result.jac)
|
||||
with open(os.path.join(experiment_results_folder, "confidence.pkl"), "wb") as handle:
|
||||
pickle.dump(
|
||||
{"cov": covariance, "intervals": intervals},
|
||||
handle,
|
||||
protocol=pickle.HIGHEST_PROTOCOL,
|
||||
)
|
||||
|
||||
# Dump identified models to disk.
|
||||
for model_sequences in models_sequences:
|
||||
model_sequences.spec.to_file(
|
||||
(experiment_results_folder / f"{model_sequences.name}.xml").as_posix()
|
||||
)
|
||||
|
||||
# Log nominal compared to initial.
|
||||
x0 = initial_params.as_vector()
|
||||
x_nominal = initial_params.as_nominal_vector()
|
||||
logging.info(
|
||||
"Initial Parameters\n%s",
|
||||
initial_params.compare_parameters(x0, opt_result.x, measured_params=x_nominal),
|
||||
)
|
||||
@@ -0,0 +1,507 @@
|
||||
"""Model modifiers."""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
|
||||
from mujoco.sysid._src.parameter import ModifierFn, Parameter, ParameterDict
|
||||
|
||||
|
||||
def remove_visuals(in_spec: mujoco.MjSpec) -> mujoco.MjSpec:
|
||||
"""Remove visual elements from a Spec."""
|
||||
spec = in_spec.copy()
|
||||
all_geoms = spec.worldbody.find_all("geom")
|
||||
for geom in all_geoms:
|
||||
if geom.contype == 0 and geom.conaffinity == 0:
|
||||
if geom.type == mujoco.mjtGeom.mjGEOM_MESH and geom.meshname != "":
|
||||
meshname = geom.meshname
|
||||
mesh = spec.mesh(meshname)
|
||||
if mesh: # multiple geoms can ref same mesh.
|
||||
spec.delete(mesh)
|
||||
spec.delete(geom)
|
||||
|
||||
for mat in spec.materials:
|
||||
spec.delete(mat)
|
||||
for tex in spec.textures:
|
||||
spec.delete(tex)
|
||||
|
||||
spec.compile() # TODO: is this compile necessary?
|
||||
return spec
|
||||
|
||||
|
||||
def _get_obj_or_raise(spec: mujoco.MjSpec, obj_type: str, obj_name: str) -> Any:
|
||||
getter = getattr(spec, obj_type, None)
|
||||
if not callable(getter):
|
||||
raise AttributeError(f"MjSpec has no method '{obj_type}'")
|
||||
obj = getter(obj_name)
|
||||
if obj is None:
|
||||
raise ValueError(f"{obj_type.capitalize()} '{obj_name}' not found in spec.")
|
||||
return obj
|
||||
|
||||
|
||||
def apply_param_modifiers_spec(
|
||||
params: ParameterDict, spec: mujoco.MjSpec
|
||||
) -> mujoco.MjSpec:
|
||||
for key in params.keys():
|
||||
param = params[key]
|
||||
if not param.frozen:
|
||||
param.apply_modifier(spec)
|
||||
return spec
|
||||
|
||||
|
||||
def apply_param_modifiers(params: ParameterDict, spec: mujoco.MjSpec) -> mujoco.MjModel:
|
||||
return apply_param_modifiers_spec(params, spec).compile()
|
||||
|
||||
|
||||
def _infer_inertial(spec: mujoco.MjSpec, body_name: str) -> mujoco.MjsBody:
|
||||
"""Override spec inertia using inferred inertia from compiled model."""
|
||||
body = _get_obj_or_raise(spec, "body", body_name)
|
||||
assert isinstance(body, mujoco.MjsBody)
|
||||
spec.compiler.inertiafromgeom = 2
|
||||
model = spec.compile()
|
||||
body.explicitinertial = True
|
||||
body.fullinertia = np.full((6, 1), np.nan)
|
||||
body.mass = model.body(body_name).mass[0]
|
||||
body.inertia = model.body(body_name).inertia
|
||||
body.ipos = model.body(body_name).ipos
|
||||
body.iquat = model.body(body_name).iquat
|
||||
return body
|
||||
|
||||
|
||||
def is_position_actuator(actuator) -> bool:
|
||||
"""Check if an actuator is a position actuator.
|
||||
|
||||
This function works on both model.actuator and spec.actuator objects.
|
||||
"""
|
||||
return (
|
||||
actuator.gaintype == mujoco.mjtGain.mjGAIN_FIXED
|
||||
and actuator.biastype == mujoco.mjtBias.mjBIAS_AFFINE
|
||||
and actuator.dyntype in (mujoco.mjtDyn.mjDYN_NONE, mujoco.mjtDyn.mjDYN_FILTEREXACT)
|
||||
and actuator.gainprm[0] == -actuator.biasprm[1]
|
||||
)
|
||||
|
||||
|
||||
def get_actuator_pd_gains(
|
||||
model: mujoco.MjModel, actuator_name: str
|
||||
) -> tuple[float, float]:
|
||||
actuator_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, actuator_name)
|
||||
if actuator_id == -1:
|
||||
raise ValueError(f"Actuator {actuator_name} not found in model.")
|
||||
actuator = model.actuator(actuator_id)
|
||||
if not is_position_actuator(actuator):
|
||||
raise ValueError(f"Actuator {actuator_name} is not a position actuator.")
|
||||
return -actuator.biasprm[1], -actuator.biasprm[2]
|
||||
|
||||
|
||||
def apply_pgain(
|
||||
spec: mujoco.MjSpec,
|
||||
actuator_name: str,
|
||||
value: float | np.ndarray,
|
||||
) -> mujoco.MjSpec:
|
||||
# TODO: assert scalar
|
||||
actuator = _get_obj_or_raise(spec, "actuator", actuator_name)
|
||||
assert isinstance(actuator, mujoco.MjsActuator)
|
||||
if not is_position_actuator(actuator):
|
||||
raise ValueError(f"Actuator {actuator_name} is not a position actuator.")
|
||||
actuator.gainprm[0] = value
|
||||
actuator.biasprm[1] = -value
|
||||
return spec
|
||||
|
||||
|
||||
def apply_dgain(
|
||||
spec: mujoco.MjSpec,
|
||||
actuator_name: str,
|
||||
value: float | np.ndarray,
|
||||
) -> mujoco.MjSpec:
|
||||
# TODO: assert scalar
|
||||
actuator = _get_obj_or_raise(spec, "actuator", actuator_name)
|
||||
assert isinstance(actuator, mujoco.MjsActuator)
|
||||
if not is_position_actuator(actuator):
|
||||
raise ValueError(f"Actuator {actuator_name} is not a position actuator.")
|
||||
actuator.biasprm[2] = -value
|
||||
return spec
|
||||
|
||||
|
||||
def apply_pdgain(
|
||||
spec: mujoco.MjSpec,
|
||||
actuator_name: str,
|
||||
value: np.ndarray,
|
||||
) -> mujoco.MjSpec:
|
||||
if value.size != 2:
|
||||
raise ValueError(f"pdgain must be a 2-element array, got {value.size}.")
|
||||
apply_pgain(spec, actuator_name, value[0])
|
||||
apply_dgain(spec, actuator_name, value[1])
|
||||
return spec
|
||||
|
||||
|
||||
def apply_body_mass_ipos(
|
||||
spec: mujoco.MjSpec,
|
||||
body_name: str,
|
||||
mass: np.ndarray | None = None,
|
||||
ipos: np.ndarray | None = None,
|
||||
rot_inertia_scale: bool = False,
|
||||
) -> mujoco.MjSpec:
|
||||
# TODO: assert mass and ipos shapes
|
||||
body = _infer_inertial(spec, body_name)
|
||||
mass_original = body.mass
|
||||
if mass is not None:
|
||||
body.mass = mass
|
||||
if rot_inertia_scale:
|
||||
scale = mass / mass_original
|
||||
body.inertia *= scale
|
||||
if ipos is not None:
|
||||
body.ipos = ipos
|
||||
return spec
|
||||
|
||||
|
||||
def scale_body_inertia(
|
||||
spec: mujoco.MjSpec,
|
||||
body_name: str,
|
||||
value: np.ndarray,
|
||||
) -> mujoco.MjSpec:
|
||||
# TODO: assert scalar
|
||||
body = _infer_inertial(spec, body_name)
|
||||
body.inertia *= value
|
||||
return spec
|
||||
|
||||
|
||||
def pi_from_theta(theta: np.ndarray) -> np.ndarray:
|
||||
alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3 = theta
|
||||
exp_alpha = np.exp(alpha)
|
||||
exp_d1 = np.exp(d1)
|
||||
exp_d2 = np.exp(d2)
|
||||
exp_d3 = np.exp(d3)
|
||||
U = np.zeros((4, 4))
|
||||
U[0, 0] = exp_d1
|
||||
U[0, 1] = s12
|
||||
U[0, 2] = s13
|
||||
U[0, 3] = t1
|
||||
U[1, 1] = exp_d2
|
||||
U[1, 2] = s23
|
||||
U[1, 3] = t2
|
||||
U[2, 2] = exp_d3
|
||||
U[2, 3] = t3
|
||||
U[3, 3] = 1
|
||||
U *= exp_alpha
|
||||
|
||||
J = U @ U.T
|
||||
|
||||
sigma = J[:3, :3]
|
||||
I_bar = np.trace(sigma) * np.eye(3) - sigma
|
||||
h = J[:3, 3]
|
||||
m = J[3, 3]
|
||||
|
||||
return np.concatenate(([m], h, I_bar.flatten()))
|
||||
|
||||
|
||||
def pseudoinertia_from_pi(pi: np.ndarray) -> np.ndarray:
|
||||
"""Converts inertial parameters π to a 4x4 pseudoinertia matrix J.
|
||||
|
||||
Args:
|
||||
pi: A 10-D array [m, hx, hy, hz, Ixx, Iyy, Izz, Ixy, Iyz, Ixz] where:
|
||||
m: Mass of the body
|
||||
[hx, hy, hz]: First moment of mass
|
||||
[Ixx, Iyy, Izz]: Diagonal elements of inertia tensor
|
||||
[Ixy, Iyz, Ixz]: Off-diagonal elements of inertia tensor
|
||||
|
||||
Returns:
|
||||
A 4x4 pseudoinertia matrix J of the form:
|
||||
[[Σ, h],
|
||||
[hᵀ, m]]
|
||||
where:
|
||||
Σ = (tr(I)/2)I₃ - I:
|
||||
h: The 3x1 first moment of mass vector
|
||||
m: The scalar mass
|
||||
"""
|
||||
m = pi[0]
|
||||
h = pi[1:4]
|
||||
I_bar = pi[4:].reshape((3, 3))
|
||||
|
||||
Sigma = 0.5 * np.trace(I_bar) * np.eye(3) - I_bar
|
||||
|
||||
J = np.zeros((4, 4))
|
||||
J[:3, :3] = Sigma
|
||||
J[:3, 3] = h
|
||||
J[3, :3] = h
|
||||
J[3, 3] = m
|
||||
|
||||
return J
|
||||
|
||||
|
||||
def cholesky_decompose_upper(J: np.ndarray) -> np.ndarray:
|
||||
"""Perform an upper-triangular Cholesky decomposition of J.
|
||||
|
||||
The returned matrix U is such that J = U @ U.T.
|
||||
"""
|
||||
n = J.shape[0]
|
||||
indices = np.arange(n - 1, -1, -1)
|
||||
J_reversed = J[indices][:, indices]
|
||||
L_prime = np.linalg.cholesky(J_reversed)
|
||||
return L_prime[indices][:, indices]
|
||||
|
||||
|
||||
def theta_from_pseudoinertia(J: np.ndarray) -> np.ndarray:
|
||||
"""Extract the 10-D vector of base parameters θ from the pseudoinertia J.
|
||||
|
||||
Args:
|
||||
J: A 4x4 pseudoinertia.
|
||||
|
||||
Returns:
|
||||
A 10-D array θ = [alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3] where:
|
||||
alpha: Scale parameter (log of U[3,3])
|
||||
[d1, d2, d3]: Log of diagonal elements
|
||||
[s12, s23, s13]: Shear parameters from upper triangle
|
||||
[t1, t2, t3]: Translation parameters from last column
|
||||
"""
|
||||
# U: A 4x4 upper-triangular matrix from Cholesky decomposition
|
||||
U = cholesky_decompose_upper(J)
|
||||
|
||||
# Extract exp(α) from the bottom-right element of U.
|
||||
exp_alpha = U[3, 3]
|
||||
alpha = np.log(exp_alpha)
|
||||
|
||||
# Compute the d parameters from the diagonal entries (adjusted by alpha).
|
||||
d1 = np.log(U[0, 0] / exp_alpha)
|
||||
d2 = np.log(U[1, 1] / exp_alpha)
|
||||
d3 = np.log(U[2, 2] / exp_alpha)
|
||||
|
||||
# Extract the shear parameters (off-diagonals in the upper triangle).
|
||||
s12 = U[0, 1] / exp_alpha
|
||||
s13 = U[0, 2] / exp_alpha
|
||||
s23 = U[1, 2] / exp_alpha
|
||||
|
||||
# Extract the translation parameters (last column, except the bottom element).
|
||||
t1 = U[0, 3] / exp_alpha
|
||||
t2 = U[1, 3] / exp_alpha
|
||||
t3 = U[2, 3] / exp_alpha
|
||||
|
||||
return np.array([alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3])
|
||||
|
||||
|
||||
def skew(v: np.ndarray) -> np.ndarray:
|
||||
"""Skew-symmetric matrix from a length-3 vector."""
|
||||
return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
|
||||
|
||||
|
||||
def inertia_to_fullinertia(q: np.ndarray, inertia: np.ndarray) -> np.ndarray:
|
||||
xmat = np.empty(9)
|
||||
mujoco.mju_quat2Mat(xmat, q)
|
||||
R = xmat.reshape(3, 3)
|
||||
return R @ np.diag(inertia) @ R.T
|
||||
|
||||
|
||||
def pi_from_body(spec: mujoco.MjSpec, body_name: str) -> np.ndarray:
|
||||
"""Extracts the 10-D vector of inertial parameters π from a MuJoCo body.
|
||||
|
||||
Args:
|
||||
spec: MuJoCo model specification object.
|
||||
body_name: Name of the body to extract parameters from.
|
||||
|
||||
Returns:
|
||||
A 10-D numpy array π = [m, hx, hy, hz, Ixx, Iyy, Izz, Ixy, Iyz, Ixz] where:
|
||||
m: Mass of the body
|
||||
[hx, hy, hz]: First moment of mass (m * com, where com is center of mass)
|
||||
[Ixx, Iyy, Izz, Ixy, Iyz, Ixz]: Rotational inertia about the origin of the
|
||||
body-fixed reference frame.
|
||||
"""
|
||||
body = _infer_inertial(spec, body_name)
|
||||
mass = body.mass
|
||||
ipos = body.ipos
|
||||
inertia = body.inertia
|
||||
iquat = body.iquat
|
||||
|
||||
fullinertia = inertia_to_fullinertia(iquat, inertia)
|
||||
# Transform inertial from ipos origin to body origin.
|
||||
I_bar = fullinertia - (mass * skew(ipos) @ skew(ipos))
|
||||
|
||||
return np.concatenate([[mass], mass * ipos, I_bar.flatten()])
|
||||
|
||||
|
||||
def theta_inertia_from_body(spec: mujoco.MjSpec, body_name: str) -> np.ndarray:
|
||||
pi = pi_from_body(spec, body_name)
|
||||
J = pseudoinertia_from_pi(pi)
|
||||
return theta_from_pseudoinertia(J)
|
||||
|
||||
|
||||
def apply_body_theta_inertia(
|
||||
spec: mujoco.MjSpec,
|
||||
body_name: str,
|
||||
theta: np.ndarray,
|
||||
) -> mujoco.MjSpec:
|
||||
if theta.size != 10:
|
||||
raise ValueError(f"theta must be a 10-element array, got {theta.size}.")
|
||||
pi = pi_from_theta(theta)
|
||||
|
||||
body = _infer_inertial(spec, body_name)
|
||||
body.mass = pi[0]
|
||||
body.ipos = pi[1:4] / pi[0]
|
||||
|
||||
# This tells the compiler to ignore the diagonal inertia and instead calculate it
|
||||
# from the full inertia.
|
||||
body.inertia[:] = 0.0
|
||||
body.iquat[:] = np.nan
|
||||
|
||||
I_bar = pi[4:].reshape((3, 3))
|
||||
skew_ipos = skew(body.ipos)
|
||||
fullinertia = I_bar + (body.mass * skew_ipos @ skew_ipos)
|
||||
|
||||
# MuJoCo's ordering is: M(1,1), M(2,2), M(3,3), M(1,2), M(1,3), M(2,3) which
|
||||
# corresponds to Ixx, Iyy, Izz, Ixy, Ixz
|
||||
body.fullinertia[0] = fullinertia[0, 0] # Ixx
|
||||
body.fullinertia[1] = fullinertia[1, 1] # Iyy
|
||||
body.fullinertia[2] = fullinertia[2, 2] # Izz
|
||||
body.fullinertia[3] = fullinertia[0, 1] # Ixy
|
||||
body.fullinertia[4] = fullinertia[0, 2] # Ixz
|
||||
body.fullinertia[5] = fullinertia[1, 2] # Iyz
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
def apply_body_inertia(spec: mujoco.MjSpec, name: str, param: Parameter):
|
||||
if not hasattr(param, "inertia_type"):
|
||||
raise ValueError(f"Parameter {param.name} does not have inertia_type attribute.")
|
||||
|
||||
if param.inertia_type == InertiaType.Mass:
|
||||
apply_body_mass_ipos(
|
||||
spec, name, mass=param.value, rot_inertia_scale=param.scale_rot_inertia
|
||||
)
|
||||
|
||||
elif param.inertia_type == InertiaType.MassIpos:
|
||||
apply_body_mass_ipos(
|
||||
spec,
|
||||
name,
|
||||
mass=param.value[0],
|
||||
ipos=param.value[1:4],
|
||||
rot_inertia_scale=param.scale_rot_inertia,
|
||||
)
|
||||
|
||||
elif param.inertia_type == InertiaType.Pseudo:
|
||||
apply_body_theta_inertia(spec, name, param.value)
|
||||
|
||||
|
||||
class InertiaType(Enum):
|
||||
Mass = 0
|
||||
MassIpos = 1
|
||||
Pseudo = 2
|
||||
|
||||
|
||||
def body_inertia_param(
|
||||
spec: mujoco.MjSpec,
|
||||
model: mujoco.MjModel,
|
||||
body_name: str,
|
||||
inertia_type: InertiaType = InertiaType.MassIpos,
|
||||
scale_rot_inertia: bool = False,
|
||||
mass_bound_mult: np.ndarray | None = None,
|
||||
ipos_bound_off: np.ndarray | None = None,
|
||||
stretch_bound_mult: np.ndarray | None = None,
|
||||
shear_bound_off: np.ndarray | None = None,
|
||||
param_name: str | None = None,
|
||||
modifier: ModifierFn | None = None,
|
||||
) -> Parameter:
|
||||
"""Creates Parameter objects for the inertia of a body in a simplified manner.
|
||||
|
||||
Args:
|
||||
model: The MuJoCo model.
|
||||
body_name: Name of the body to create the parameter for.
|
||||
inertia_type: The type of inertia parameterization to use.
|
||||
scale_rot_inertia: Whether to scale the original inertia when mass changes,
|
||||
ignored with pseudo inertia.
|
||||
mass_bound_mult: Multiplicative bounds for the mass parameter.
|
||||
ipos_bound_off: Additive bounds for the ipos parameter.
|
||||
stretch_bound_mult: Multiplicative bounds for the stretch parameters in the
|
||||
pseudo-inertia parameterization.
|
||||
shear_bound_off: Additive bounds for the shear parameters in the pseudo-inertia
|
||||
parameterization.
|
||||
param_name: Optional name for the parameter. Defaults to
|
||||
``"{body_name}_inertia"``.
|
||||
modifier: Optional custom modifier callback. If None, the default
|
||||
:func:`apply_body_inertia` modifier is registered on the Parameter."""
|
||||
|
||||
if mass_bound_mult is None:
|
||||
mass_bound_mult = np.array([0.1, 10.0])
|
||||
if ipos_bound_off is None:
|
||||
ipos_bound_off = np.array([-0.5, 0.5])
|
||||
if stretch_bound_mult is None:
|
||||
stretch_bound_mult = np.array([0.5, 2.0])
|
||||
if shear_bound_off is None:
|
||||
shear_bound_off = np.array([-0.5, 0.5])
|
||||
|
||||
body = model.body(body_name)
|
||||
if param_name is None:
|
||||
param_name = f"{body_name}_inertia"
|
||||
|
||||
if modifier is None:
|
||||
|
||||
def _default_modifier(spec, param):
|
||||
return apply_body_inertia(spec, body_name, param)
|
||||
|
||||
modifier = _default_modifier
|
||||
|
||||
if inertia_type == InertiaType.Mass:
|
||||
param = Parameter(
|
||||
param_name,
|
||||
body.mass,
|
||||
body.mass * mass_bound_mult[0],
|
||||
body.mass * mass_bound_mult[1],
|
||||
modifier=modifier,
|
||||
)
|
||||
param.inertia_type = inertia_type
|
||||
param.scale_rot_inertia = scale_rot_inertia
|
||||
|
||||
elif inertia_type == InertiaType.MassIpos:
|
||||
massipos0 = np.concatenate((body.mass, body.ipos))
|
||||
massipos_low = np.concatenate(
|
||||
(body.mass * mass_bound_mult[0], body.ipos + ipos_bound_off[0])
|
||||
)
|
||||
massipos_high = np.concatenate(
|
||||
(body.mass * mass_bound_mult[1], body.ipos + ipos_bound_off[1])
|
||||
)
|
||||
param = Parameter(
|
||||
param_name, massipos0, massipos_low, massipos_high, modifier=modifier
|
||||
)
|
||||
param.inertia_type = inertia_type
|
||||
param.scale_rot_inertia = scale_rot_inertia
|
||||
|
||||
elif inertia_type == InertiaType.Pseudo:
|
||||
theta_i_0 = theta_inertia_from_body(spec, body_name)
|
||||
|
||||
# mass = exp(2*alpha)
|
||||
alpha = theta_i_0[0]
|
||||
mass = np.exp(2 * alpha)
|
||||
mass_bounds = mass * mass_bound_mult
|
||||
alpha_bounds = 0.5 * np.log(mass_bounds)
|
||||
|
||||
# d1, d2, d3, stretch = exp(2*d)
|
||||
# stretches body along principal axes
|
||||
d = theta_i_0[1 : 1 + 3]
|
||||
stretch = np.exp(2 * d)
|
||||
stretch_bounds = stretch[:, np.newaxis] * np.atleast_2d(stretch_bound_mult)
|
||||
d_bounds = 0.5 * np.log(stretch_bounds)
|
||||
|
||||
# s12, s23, s13
|
||||
# shear the body
|
||||
s_bounds = theta_i_0[4 : 4 + 3, np.newaxis] + np.atleast_2d(shear_bound_off)
|
||||
|
||||
# t1, t2, t3
|
||||
# center of mass
|
||||
t_bounds = theta_i_0[7:10, np.newaxis] + np.atleast_2d(ipos_bound_off)
|
||||
|
||||
theta_bounds = np.vstack(
|
||||
[
|
||||
alpha_bounds,
|
||||
d_bounds,
|
||||
s_bounds,
|
||||
t_bounds,
|
||||
]
|
||||
)
|
||||
param = Parameter(
|
||||
param_name, theta_i_0, theta_bounds[:, 0], theta_bounds[:, 1], modifier=modifier
|
||||
)
|
||||
param.inertia_type = inertia_type
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown inertia_type: {inertia_type}")
|
||||
|
||||
return param
|
||||
@@ -0,0 +1,236 @@
|
||||
"""Optimization routines for system identification."""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Literal
|
||||
|
||||
import numpy as np
|
||||
import scipy.optimize as scipy_optimize
|
||||
from absl import logging
|
||||
from mujoco import minimize as mujoco_minimize
|
||||
from scipy.special import stdtrit
|
||||
|
||||
from mujoco.sysid._src import parameter
|
||||
|
||||
|
||||
def _scipy_least_squares(
|
||||
x0: np.ndarray,
|
||||
residual_fn: Callable,
|
||||
bounds: tuple[np.ndarray, np.ndarray],
|
||||
use_mujoco_jac: bool = False,
|
||||
**kwargs,
|
||||
) -> scipy_optimize.OptimizeResult:
|
||||
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
|
||||
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/91e18f3bd355477b8b7747ec82d70ac98ffd2422/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,
|
||||
bounds: tuple[np.ndarray, np.ndarray],
|
||||
**kwargs,
|
||||
) -> scipy_optimize.OptimizeResult:
|
||||
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,
|
||||
bounds: tuple[np.ndarray, np.ndarray],
|
||||
optimizer: Literal["scipy", "mujoco", "scipy_parallel_fd"],
|
||||
**kwargs,
|
||||
) -> scipy_optimize.OptimizeResult:
|
||||
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,
|
||||
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 len(opt_params) == 0 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,
|
||||
):
|
||||
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) * stdtrit(
|
||||
final_r.size - J.shape[1], 1 - alpha / 2
|
||||
)
|
||||
return Sigma_X, intervals
|
||||
@@ -0,0 +1,606 @@
|
||||
"""Parameter utilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import pathlib
|
||||
from typing import TYPE_CHECKING, Callable, TypeAlias
|
||||
|
||||
import colorama
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import yaml
|
||||
from tabulate import tabulate
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing_extensions import Self
|
||||
|
||||
from mujoco.sysid._src.model_modifier import InertiaType
|
||||
|
||||
Fore = colorama.Fore
|
||||
Style = colorama.Style
|
||||
|
||||
ModifierFn: TypeAlias = Callable[[mujoco.MjSpec, "Parameter"], object]
|
||||
|
||||
|
||||
class Parameter:
|
||||
"""A single (possibly multi-dimensional) parameter for system identification.
|
||||
|
||||
A Parameter holds a current ``value``, a ``nominal`` baseline, and box
|
||||
bounds (``min_value``, ``max_value``). An optional ``modifier`` callback
|
||||
is invoked during model compilation to apply the parameter to an MjSpec.
|
||||
|
||||
Args:
|
||||
name: Human-readable identifier (must be unique within a ParameterDict).
|
||||
nominal: Nominal (initial) value; scalar or array-like.
|
||||
min_value: Lower bound, same shape as *nominal*.
|
||||
max_value: Upper bound, same shape as *nominal*.
|
||||
frozen: If True the parameter is excluded from optimization.
|
||||
modifier: Optional callback ``(MjSpec, Parameter) -> None`` that writes
|
||||
the parameter into a spec during model compilation.
|
||||
"""
|
||||
|
||||
# Type hints for dynamically-added attributes (set by parameter builders).
|
||||
if TYPE_CHECKING:
|
||||
inertia_type: InertiaType | None
|
||||
scale_rot_inertia: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
nominal: float | npt.ArrayLike,
|
||||
min_value: float | npt.ArrayLike,
|
||||
max_value: float | npt.ArrayLike,
|
||||
frozen: bool = False,
|
||||
modifier: ModifierFn | None = None,
|
||||
):
|
||||
self.name = name
|
||||
self.nominal = np.atleast_1d(nominal)
|
||||
self.min_value = np.atleast_1d(min_value)
|
||||
self.max_value = np.atleast_1d(max_value)
|
||||
self.value = self.nominal.copy()
|
||||
self.frozen = frozen
|
||||
self.modifier = modifier
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
return self.nominal.size
|
||||
|
||||
@property
|
||||
def shape(self) -> tuple[int, ...]:
|
||||
return self.nominal.shape
|
||||
|
||||
def apply_modifier(self, spec: mujoco.MjSpec) -> None:
|
||||
"""Apply this parameter's modifier callback to *spec*, if one is set."""
|
||||
if self.modifier:
|
||||
self.modifier(spec, self)
|
||||
|
||||
def as_vector(self) -> np.ndarray:
|
||||
"""Return the current value as a flat 1-D array."""
|
||||
return self.value.flatten()
|
||||
|
||||
def as_nominal_vector(self) -> np.ndarray:
|
||||
"""Return the nominal value as a flat 1-D array."""
|
||||
return self.nominal.flatten()
|
||||
|
||||
def update_from_vector(self, vector: np.ndarray) -> None:
|
||||
vector_array = np.atleast_1d(vector)
|
||||
if len(vector_array) != self.size:
|
||||
raise ValueError(
|
||||
f"Input vector length {vector_array.size} does not match "
|
||||
f"parameter size {self.size}."
|
||||
)
|
||||
self.value = vector_array.reshape(self.shape)
|
||||
|
||||
def get_bounds(self) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Return ``(lower, upper)`` bound arrays, each flat 1-D."""
|
||||
return (
|
||||
self.min_value.flatten(),
|
||||
self.max_value.flatten(),
|
||||
)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset the current value to nominal."""
|
||||
self.value = self.nominal.copy()
|
||||
|
||||
def sample(self, rng: np.random.Generator | None = None) -> np.ndarray:
|
||||
"""Sample a random value uniformly within bounds."""
|
||||
if rng is None:
|
||||
rng = np.random.default_rng()
|
||||
return rng.uniform(self.min_value.flatten(), self.max_value.flatten())
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Return a string representation of the parameter."""
|
||||
if self.size == 1:
|
||||
return (
|
||||
f"{Fore.CYAN}{self.name}{Style.RESET_ALL}: "
|
||||
f"{Fore.GREEN}{float(self.value.item()):.3g}{Style.RESET_ALL} "
|
||||
f"∈ [{Fore.YELLOW}{float(self.min_value.item()):.3g}, "
|
||||
f"{float(self.max_value.item()):.3g}{Style.RESET_ALL}]"
|
||||
)
|
||||
else:
|
||||
return (
|
||||
f"{Fore.CYAN}{self.name}{Style.RESET_ALL}: "
|
||||
f"{Fore.GREEN}array(shape={self.shape}){Style.RESET_ALL} "
|
||||
f"∈ [{Fore.YELLOW}min={np.min(self.min_value):.3g}, "
|
||||
f"max={np.max(self.max_value):.3g}{Style.RESET_ALL}]"
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.__str__()
|
||||
|
||||
def __getstate__(self):
|
||||
return {
|
||||
"name": self.name,
|
||||
"nominal": self.nominal.tolist()
|
||||
if isinstance(self.nominal, np.ndarray)
|
||||
else self.nominal,
|
||||
"min_value": self.min_value.tolist()
|
||||
if isinstance(self.min_value, np.ndarray)
|
||||
else self.min_value,
|
||||
"max_value": self.max_value.tolist()
|
||||
if isinstance(self.max_value, np.ndarray)
|
||||
else self.max_value,
|
||||
"value": self.value.tolist()
|
||||
if isinstance(self.value, np.ndarray)
|
||||
else self.value,
|
||||
"frozen": self.frozen,
|
||||
}
|
||||
|
||||
def __setstate__(self, state):
|
||||
self.name = state["name"]
|
||||
self.nominal = np.array(state["nominal"])
|
||||
self.min_value = np.array(state["min_value"])
|
||||
self.max_value = np.array(state["max_value"])
|
||||
self.value = np.array(state["value"])
|
||||
self.frozen = state["frozen"]
|
||||
|
||||
# Override default deepycopy so lambda references get copied
|
||||
def __deepcopy__(self, memo):
|
||||
cls = self.__class__
|
||||
result = cls.__new__(cls)
|
||||
for k, v in self.__dict__.items():
|
||||
setattr(result, k, copy.deepcopy(v, memo))
|
||||
return result
|
||||
|
||||
|
||||
class ParameterDict:
|
||||
"""An ordered collection of :class:`Parameter` objects.
|
||||
|
||||
Behaves like a ``dict[str, Parameter]`` with convenience methods for
|
||||
vectorised access (``as_vector`` / ``update_from_vector``), serialisation,
|
||||
and tabular comparison of parameter estimates.
|
||||
|
||||
Frozen parameters are silently skipped by vector/bounds methods so that the
|
||||
decision-variable dimension seen by optimizers matches only the free params.
|
||||
"""
|
||||
|
||||
def __init__(self, parameters: dict[str, Parameter] | None = None):
|
||||
if parameters is None:
|
||||
self.parameters = {}
|
||||
else:
|
||||
self.parameters = parameters
|
||||
|
||||
def __getitem__(self, key: str) -> Parameter:
|
||||
return self.parameters[key]
|
||||
|
||||
def __setitem__(self, key: str, value: Parameter) -> None:
|
||||
self.parameters[key] = value
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
return key in self.parameters
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.parameters)
|
||||
|
||||
def copy(self) -> Self:
|
||||
"""Return a deep copy of this ParameterDict."""
|
||||
return copy.deepcopy(self)
|
||||
|
||||
def add(self, param: Parameter) -> None:
|
||||
"""Add a Parameter, keyed by its ``name``."""
|
||||
self.parameters[param.name] = param
|
||||
|
||||
def update(self, pdict: Self) -> None:
|
||||
for keys in pdict.keys():
|
||||
if keys in self.parameters:
|
||||
raise ValueError(f"Parameter '{keys}' already exists in the dictionary.")
|
||||
self.parameters[keys] = pdict[keys]
|
||||
|
||||
def keys(self) -> list[str]:
|
||||
return list(self.parameters.keys())
|
||||
|
||||
def values(self) -> list[Parameter]:
|
||||
return list(self.parameters.values())
|
||||
|
||||
def items(self) -> list[tuple[str, Parameter]]:
|
||||
return list(self.parameters.items())
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
"""Get the total size of all non-frozen parameters."""
|
||||
return sum(p.size for p in self.parameters.values() if not p.frozen)
|
||||
|
||||
def as_vector(self, include_frozen=False) -> np.ndarray:
|
||||
"""Convert all non-frozen parameters to a flat vector."""
|
||||
vectors = [
|
||||
p.as_vector() for p in self.parameters.values() if not p.frozen or include_frozen
|
||||
]
|
||||
return np.concatenate(vectors) if vectors else np.array([])
|
||||
|
||||
def as_nominal_vector(self, include_frozen=False) -> np.ndarray:
|
||||
"""Get the nominal values of parameters as a flat array."""
|
||||
vectors = [
|
||||
p.as_nominal_vector()
|
||||
for p in self.parameters.values()
|
||||
if not p.frozen or include_frozen
|
||||
]
|
||||
return np.concatenate(vectors) if vectors else np.array([])
|
||||
|
||||
def update_from_vector(self, vector: np.ndarray) -> None:
|
||||
"""Update all non-frozen parameters from a flat vector."""
|
||||
start = 0
|
||||
for param in self.parameters.values():
|
||||
if not param.frozen:
|
||||
size = param.size
|
||||
param.update_from_vector(vector[start : start + size])
|
||||
start += size
|
||||
|
||||
def save_to_disk(self, path: str | pathlib.Path) -> None:
|
||||
"""Save the parameter dictionary to disk (schema and data).
|
||||
|
||||
Args:
|
||||
path: Path where the data will be saved.
|
||||
"""
|
||||
parameter_dicts = {
|
||||
name: param.__getstate__() for name, param in self.parameters.items()
|
||||
}
|
||||
with open(path, "w") as handle:
|
||||
yaml.safe_dump(parameter_dicts, handle, default_flow_style=False)
|
||||
|
||||
@classmethod
|
||||
def load_from_disk(cls, path: str | pathlib.Path) -> "ParameterDict":
|
||||
"""Load parameter dictionary from disk (schema and data).
|
||||
|
||||
Args:
|
||||
path: Path to the saved data.
|
||||
|
||||
Returns:
|
||||
A new ParameterDict object.
|
||||
"""
|
||||
with open(path, "r") as handle:
|
||||
parameter_dicts = yaml.safe_load(handle)
|
||||
|
||||
parameters = {}
|
||||
for name, param_dict in parameter_dicts.items():
|
||||
param = Parameter.__new__(Parameter)
|
||||
param.__setstate__(param_dict)
|
||||
parameters[name] = param
|
||||
|
||||
return ParameterDict(parameters)
|
||||
|
||||
def get_bounds(self) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Get the bounds for all non-frozen parameters."""
|
||||
lower_bounds = []
|
||||
upper_bounds = []
|
||||
for param in self.parameters.values():
|
||||
if not param.frozen:
|
||||
lb, ub = param.get_bounds()
|
||||
lower_bounds.append(lb)
|
||||
upper_bounds.append(ub)
|
||||
|
||||
return (
|
||||
np.concatenate(lower_bounds) if lower_bounds else np.array([]),
|
||||
np.concatenate(upper_bounds) if upper_bounds else np.array([]),
|
||||
)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset all parameters to their nominal values."""
|
||||
for param in self.parameters.values():
|
||||
param.reset()
|
||||
|
||||
def sample(self, rng: np.random.Generator | None = None) -> np.ndarray:
|
||||
"""Sample parameter values within bounds for non-frozen parameters."""
|
||||
if rng is None:
|
||||
rng = np.random.default_rng()
|
||||
lower_bounds, upper_bounds = self.get_bounds()
|
||||
return rng.uniform(lower_bounds, upper_bounds)
|
||||
|
||||
def randomize(self, rng: np.random.Generator | None = None) -> None:
|
||||
"""Randomize parameter values for non-frozen parameters."""
|
||||
for param in self.parameters.values():
|
||||
if not param.frozen:
|
||||
param.value = param.sample(rng)
|
||||
|
||||
def compare_parameters(
|
||||
self,
|
||||
init_params: np.ndarray,
|
||||
predicted_params: np.ndarray,
|
||||
measured_params: np.ndarray | None = None,
|
||||
sig_digits: int = 4,
|
||||
) -> str:
|
||||
"""Compare true and predicted parameter values.
|
||||
|
||||
Args:
|
||||
init_params: Initial parameter values as a flat array.
|
||||
predicted_params: Predicted parameter values as a flat array.
|
||||
measured_params: True parameter values as a flat array.
|
||||
sig_digits: Number of significant digits to display.
|
||||
|
||||
Returns:
|
||||
A formatted string with parameter comparison table.
|
||||
"""
|
||||
# Get the vector of non-frozen parameters
|
||||
non_frozen_vector = self.as_vector()
|
||||
|
||||
if non_frozen_vector.size == 0:
|
||||
return "No non-frozen parameters to compare."
|
||||
|
||||
if len(init_params) != non_frozen_vector.size:
|
||||
raise ValueError(
|
||||
f"Initial parameter vector length {len(init_params)} does not match "
|
||||
f"the number of non-frozen parameters {non_frozen_vector.size}."
|
||||
)
|
||||
|
||||
if len(predicted_params) != non_frozen_vector.size:
|
||||
raise ValueError(
|
||||
f"Predicted parameter vector length {len(predicted_params)} does not match "
|
||||
f"the number of non-frozen parameters {non_frozen_vector.size}."
|
||||
)
|
||||
|
||||
if measured_params is not None:
|
||||
if len(measured_params) != non_frozen_vector.size:
|
||||
raise ValueError(
|
||||
f"True parameter vector length {len(measured_params)} does not match "
|
||||
f"the number of non-frozen parameters {non_frozen_vector.size}."
|
||||
)
|
||||
|
||||
# Compute error metrics.
|
||||
rel_deltas = []
|
||||
for i in range(predicted_params.shape[0]):
|
||||
if (
|
||||
init_params[i] == 0
|
||||
or np.abs(predicted_params[i] - init_params[i]) / np.abs(init_params[i]) > 2e1
|
||||
):
|
||||
rel_deltas.append(np.nan)
|
||||
else:
|
||||
rel_deltas.append(
|
||||
np.abs(predicted_params[i] - init_params[i]) / np.abs(init_params[i])
|
||||
)
|
||||
rel_deltas = np.array(rel_deltas)
|
||||
overall_rms_delta = np.sqrt(np.mean((predicted_params - init_params) ** 2))
|
||||
abs_deltas = np.abs(predicted_params - init_params)
|
||||
|
||||
if measured_params is not None:
|
||||
rel_errors = []
|
||||
for i in range(predicted_params.shape[0]):
|
||||
if (
|
||||
measured_params[i] == 0
|
||||
or np.abs(predicted_params[i] - measured_params[i])
|
||||
/ np.abs(measured_params[i])
|
||||
> 2e1
|
||||
):
|
||||
rel_errors.append(np.nan)
|
||||
else:
|
||||
rel_errors.append(
|
||||
np.abs(predicted_params[i] - measured_params[i])
|
||||
/ np.abs(measured_params[i])
|
||||
)
|
||||
rel_errors = np.array(rel_errors)
|
||||
|
||||
overall_rmse = np.sqrt(np.mean((predicted_params - measured_params) ** 2))
|
||||
abs_errors = np.abs(predicted_params - measured_params)
|
||||
else:
|
||||
overall_rmse = np.nan
|
||||
abs_errors = np.full_like(predicted_params, np.nan)
|
||||
rel_errors = np.full_like(predicted_params, np.nan)
|
||||
|
||||
lower_bounds, upper_bounds = self.get_bounds()
|
||||
|
||||
def format_number(x):
|
||||
"""Format number with fixed width for proper table alignment."""
|
||||
if abs(x) < 0.01:
|
||||
return f"{x: .{sig_digits}e}"
|
||||
else:
|
||||
return f"{x: .{sig_digits}f}"
|
||||
|
||||
def get_color_for_error(error):
|
||||
"""Get color code based on relative error magnitude."""
|
||||
if error < 0.02:
|
||||
return Fore.GREEN
|
||||
elif error < 0.1:
|
||||
return Fore.YELLOW
|
||||
else:
|
||||
return Fore.RED
|
||||
|
||||
def create_table_row(param_name, idx):
|
||||
"""Create a formatted table row for a parameter at the given index."""
|
||||
|
||||
true = measured_params[idx] if measured_params is not None else np.nan
|
||||
init = init_params[idx]
|
||||
est = predicted_params[idx]
|
||||
lower_bound = lower_bounds[idx]
|
||||
upper_bound = upper_bounds[idx]
|
||||
delta = abs_deltas[idx]
|
||||
error = abs_errors[idx] if measured_params is not None else np.nan
|
||||
rel_delta = rel_deltas[idx]
|
||||
rel_err = rel_errors[idx] if measured_params is not None else np.nan
|
||||
|
||||
# If a parameter is near the boundary make it magneta
|
||||
if (abs(est - lower_bound) < 1e-8 + 1e-3 * abs(lower_bound)) or (
|
||||
abs(est - upper_bound) < 1e-8 + 1e-3 * abs(upper_bound)
|
||||
):
|
||||
color = Fore.MAGENTA
|
||||
else:
|
||||
if measured_params is None:
|
||||
color = get_color_for_error(rel_delta)
|
||||
else:
|
||||
color = get_color_for_error(error)
|
||||
|
||||
# Format all values with appropriate colors
|
||||
if np.isnan(true):
|
||||
measured_val = ""
|
||||
else:
|
||||
measured_val = f"{Fore.BLUE}{format_number(true)}{Style.RESET_ALL}"
|
||||
init_val = f"{Fore.BLUE}{format_number(init)}{Style.RESET_ALL}"
|
||||
est_val = f"{color}{format_number(est)}{Style.RESET_ALL}"
|
||||
|
||||
lower_bound_val = f"{Fore.BLUE}{format_number(lower_bound)}{Style.RESET_ALL}"
|
||||
upper_bound_val = f"{Fore.BLUE}{format_number(upper_bound)}{Style.RESET_ALL}"
|
||||
|
||||
if np.isnan(error):
|
||||
abs_err_val = ""
|
||||
else:
|
||||
abs_err_val = f"{color}{format_number(error)}{Style.RESET_ALL}"
|
||||
abs_delta_val = f"{color}{format_number(delta)}{Style.RESET_ALL}"
|
||||
|
||||
if np.isnan(rel_err):
|
||||
rel_err_val = ""
|
||||
else:
|
||||
rel_err_val = f"{color}{rel_err * 100:.1f}%{Style.RESET_ALL}"
|
||||
|
||||
if np.isnan(rel_delta):
|
||||
rel_delta_val = ""
|
||||
else:
|
||||
rel_delta_val = f"{color}{rel_delta * 100:.1f}%{Style.RESET_ALL}"
|
||||
|
||||
return [
|
||||
f"{Fore.CYAN}{param_name.ljust(20)}{Style.RESET_ALL}",
|
||||
init_val,
|
||||
measured_val,
|
||||
est_val,
|
||||
lower_bound_val,
|
||||
upper_bound_val,
|
||||
abs_err_val,
|
||||
abs_delta_val,
|
||||
rel_err_val,
|
||||
rel_delta_val,
|
||||
]
|
||||
|
||||
# Build table data.
|
||||
table_data = []
|
||||
non_frozen_idx = 0 # Index for non-frozen parameters in the arrays
|
||||
|
||||
for param_name, param in self.parameters.items():
|
||||
if param.frozen:
|
||||
continue # Skip frozen parameters
|
||||
|
||||
if param.size == 1:
|
||||
table_data.append(create_table_row(param_name, non_frozen_idx))
|
||||
non_frozen_idx += 1
|
||||
else:
|
||||
for i in range(param.size):
|
||||
if param.shape == (param.size,):
|
||||
element_name = f"{param_name}[{i}]"
|
||||
else:
|
||||
multi_idx = np.unravel_index(i, param.shape)
|
||||
idx_str = ",".join(str(x) for x in multi_idx)
|
||||
element_name = f"{param_name}[{idx_str}]"
|
||||
table_data.append(create_table_row(element_name, non_frozen_idx))
|
||||
non_frozen_idx += 1
|
||||
|
||||
# Create and return the formatted table.
|
||||
headers = [
|
||||
"Parameter",
|
||||
"Initial",
|
||||
"Nominal",
|
||||
"Identified",
|
||||
"Lower",
|
||||
"Upper",
|
||||
"Abs Err",
|
||||
"Abs Del",
|
||||
"Rel Err",
|
||||
"Rel Del",
|
||||
]
|
||||
|
||||
table = tabulate(
|
||||
table_data, headers=headers, tablefmt="outline", disable_numparse=True
|
||||
)
|
||||
|
||||
overall_rmse_val = "" if np.isnan(overall_rmse) else f"{overall_rmse:.4g}"
|
||||
overall_rms_delta_val = f"{overall_rms_delta:.4g}"
|
||||
|
||||
return f"{table}\nRMSE: {overall_rmse_val}\nRMS Delta: {overall_rms_delta_val}"
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Return a string representation of all parameters in the dictionary."""
|
||||
if not self.parameters:
|
||||
return f"{Fore.CYAN}ParameterDict{Style.RESET_ALL}(empty)"
|
||||
|
||||
param_strings = []
|
||||
for name, param in self.parameters.items():
|
||||
if param.size == 1:
|
||||
param_strings.append(f" {param}")
|
||||
else:
|
||||
# For multi-dimensional parameters, show each element on its own line
|
||||
param_strings.append(f" {Fore.CYAN}{name}{Style.RESET_ALL}:")
|
||||
if param.shape == (param.size,): # 1D array
|
||||
for i in range(param.size):
|
||||
param_strings.append(
|
||||
f" [{i}]: {Fore.GREEN}{param.value[i]:.3g}{Style.RESET_ALL} "
|
||||
f"∈ [{Fore.YELLOW}{param.min_value[i]:.3g}, "
|
||||
f"{param.max_value[i]:.3g}{Style.RESET_ALL}]"
|
||||
)
|
||||
else: # Multi-dimensional array
|
||||
flat_idx = 0
|
||||
for idx in np.ndindex(param.shape):
|
||||
idx_str = ",".join(str(x) for x in idx)
|
||||
param_strings.append(
|
||||
f" [{idx_str}]:"
|
||||
f" {Fore.GREEN}{param.value[idx]:.3g}{Style.RESET_ALL} ∈"
|
||||
f" [{Fore.YELLOW}{param.min_value.flat[flat_idx]:.3g},"
|
||||
f" {param.max_value.flat[flat_idx]:.3g}{Style.RESET_ALL}]"
|
||||
)
|
||||
flat_idx += 1
|
||||
|
||||
params_str = "\n".join(param_strings)
|
||||
return f"{Fore.CYAN}ParameterDict{Style.RESET_ALL}(\n{params_str}\n)"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.__str__()
|
||||
|
||||
def get_non_frozen_parameter_names(self) -> list[str]:
|
||||
"""Get the names of all non-frozen parameters, expanding multi-dimensional ones."""
|
||||
names = []
|
||||
for name, param in self.parameters.items():
|
||||
if not param.frozen:
|
||||
if param.size == 1:
|
||||
names.append(name)
|
||||
else:
|
||||
if param.shape == (param.size,):
|
||||
for i in range(param.size):
|
||||
names.append(f"{name}[{i}]")
|
||||
else:
|
||||
for idx in np.ndindex(param.shape):
|
||||
idx_str = ",".join(map(str, idx))
|
||||
names.append(f"{name}[{idx_str}]")
|
||||
return names
|
||||
|
||||
def get_parameter_info(self) -> str:
|
||||
"""Get information about all parameters in the dictionary.
|
||||
|
||||
Returns:
|
||||
A formatted string with parameter information.
|
||||
"""
|
||||
if not self.parameters:
|
||||
return "No parameters in dictionary."
|
||||
|
||||
info = []
|
||||
info.append(f"{Fore.CYAN}Parameter Information:{Style.RESET_ALL}")
|
||||
info.append(
|
||||
f"{Fore.CYAN}{'Name':<20} {'Size':<10} {'Shape':<15} {'Frozen':<10}{Style.RESET_ALL}"
|
||||
)
|
||||
info.append("-" * 60)
|
||||
|
||||
for name, param in self.parameters.items():
|
||||
frozen_str = (
|
||||
f"{Fore.RED}Yes{Style.RESET_ALL}"
|
||||
if param.frozen
|
||||
else f"{Fore.GREEN}No{Style.RESET_ALL}"
|
||||
)
|
||||
info.append(
|
||||
f"{Fore.CYAN}{name:<20} {param.size:<10} {str(param.shape):<15} {frozen_str}{Style.RESET_ALL}"
|
||||
)
|
||||
|
||||
return "\n".join(info)
|
||||
@@ -0,0 +1,692 @@
|
||||
"""Plotting utilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import mujoco
|
||||
import numpy as np
|
||||
from matplotlib.lines import Line2D
|
||||
|
||||
from mujoco.sysid._src import parameter
|
||||
|
||||
|
||||
def plot_sensor_comparison(
|
||||
model: mujoco.MjModel,
|
||||
predicted_times: np.ndarray | None = None,
|
||||
predicted_data: np.ndarray | None = None,
|
||||
real_data: np.ndarray | None = None,
|
||||
real_times: np.ndarray | None = None,
|
||||
preid_data: np.ndarray | None = None,
|
||||
preid_times: np.ndarray | None = None,
|
||||
commanded_data: np.ndarray | None = None,
|
||||
commanded_times: np.ndarray | None = None,
|
||||
size_factor: float = 1.0,
|
||||
title_prefix: str = "",
|
||||
sensor_ids: list[int] | None = None,
|
||||
):
|
||||
"""Plots sensor trajectories from simulation and real data.
|
||||
|
||||
Args:
|
||||
model: The model object providing sensor information.
|
||||
predicted_times: Optional 1D array of timestamps corresponding to simulation data.
|
||||
predicted_data: Optional 2D array of simulation sensor data with shape
|
||||
(num_timesteps, sensor_data_dimension).
|
||||
real_data: Optional 2D array of real sensor data with the same shape as
|
||||
predicted_data.
|
||||
real_times: A 1D array of timestamps corresponding to real data.
|
||||
If None and real_data is provided, the first available timestamp array is used.
|
||||
preid_data: Optional 2D array of pre-identification sensor data.
|
||||
preid_times: A 1D array of timestamps for pre-identification data.
|
||||
commanded_data: Optional 2D array of commanded sensor data.
|
||||
commanded_times: A 1D array of timestamps for commanded data.
|
||||
size_factor: A scaling factor for the figure size.
|
||||
"""
|
||||
# Define a more appealing color palette
|
||||
predicted_color = "#1f77b4" # Steel blue
|
||||
real_color = "#ff7f0e" # Safety orange
|
||||
preid_color = "#2ca02c" # Forest green
|
||||
commanded_color = "#9467bd" # Purple
|
||||
|
||||
# Determine the reference time array to use
|
||||
reference_times = None
|
||||
if predicted_times is not None:
|
||||
reference_times = predicted_times
|
||||
elif real_times is not None:
|
||||
reference_times = real_times
|
||||
elif preid_times is not None:
|
||||
reference_times = preid_times
|
||||
elif commanded_times is not None:
|
||||
reference_times = commanded_times
|
||||
else:
|
||||
raise ValueError("At least one time array must be provided")
|
||||
|
||||
# Set times for data sources that don't have their own time arrays
|
||||
if real_data is not None and real_times is None:
|
||||
real_times = reference_times
|
||||
if preid_data is not None and preid_times is None:
|
||||
preid_times = reference_times
|
||||
if commanded_data is not None and commanded_times is None:
|
||||
commanded_times = reference_times
|
||||
if predicted_data is not None and predicted_times is None:
|
||||
predicted_times = reference_times
|
||||
|
||||
if sensor_ids is None:
|
||||
sensor_ids = list(range(model.nsensor))
|
||||
assert predicted_data is not None
|
||||
n_plots = predicted_data.shape[1]
|
||||
|
||||
fig, axes = plt.subplots(
|
||||
n_plots,
|
||||
1,
|
||||
figsize=(10 * size_factor, 2.5 * n_plots * size_factor),
|
||||
sharex=True,
|
||||
)
|
||||
if n_plots == 1:
|
||||
axes = [axes]
|
||||
axes = list(axes) # pyright: ignore[reportArgumentType]
|
||||
|
||||
# Set an overall title for the figure.
|
||||
fig.suptitle(title_prefix + " Sensors", fontsize=14) # , y=1.02)
|
||||
|
||||
# Loop over each sensor.
|
||||
plot_i = 0
|
||||
sensor_dim = 1
|
||||
j = 0
|
||||
dim_str = ""
|
||||
for _i, sensor_id in enumerate(sensor_ids):
|
||||
sensor = model.sensor(sensor_id)
|
||||
sensor_name = sensor.name
|
||||
sensor_dim = int(sensor.dim[0])
|
||||
sensor_addr = int(sensor.adr[0])
|
||||
|
||||
for j in range(sensor_dim):
|
||||
ax = axes[plot_i]
|
||||
plot_i += 1
|
||||
dim_str = "" if sensor_dim == 1 else f" {j}"
|
||||
if predicted_data is not None:
|
||||
assert predicted_times is not None
|
||||
predicted_signal = predicted_data[:, sensor_addr : sensor_addr + sensor_dim]
|
||||
ax.plot(
|
||||
predicted_times,
|
||||
predicted_signal[:, j],
|
||||
lw=2,
|
||||
color=predicted_color,
|
||||
alpha=0.8,
|
||||
label="Sim" + dim_str,
|
||||
)
|
||||
if real_data is not None:
|
||||
assert real_times is not None
|
||||
real_signal = real_data[:, sensor_addr : sensor_addr + sensor_dim]
|
||||
ax.plot(
|
||||
real_times,
|
||||
real_signal[:, j],
|
||||
lw=2,
|
||||
color=real_color,
|
||||
linestyle="--",
|
||||
alpha=0.7,
|
||||
label="Real" + dim_str,
|
||||
)
|
||||
if preid_data is not None:
|
||||
assert preid_times is not None
|
||||
preid_signal = preid_data[:, sensor_addr : sensor_addr + sensor_dim]
|
||||
ax.plot(
|
||||
preid_times,
|
||||
preid_signal[:, j],
|
||||
lw=2,
|
||||
color=preid_color,
|
||||
linestyle=":",
|
||||
alpha=0.6,
|
||||
label="Pre-ID" + dim_str,
|
||||
)
|
||||
if commanded_data is not None:
|
||||
assert commanded_times is not None
|
||||
commanded_signal = commanded_data[:, sensor_addr : sensor_addr + sensor_dim]
|
||||
ax.plot(
|
||||
commanded_times,
|
||||
commanded_signal[:, j],
|
||||
lw=2,
|
||||
color=commanded_color,
|
||||
linestyle="-.",
|
||||
alpha=0.6,
|
||||
label="Commanded" + dim_str,
|
||||
)
|
||||
# Place the sensor name in a white box in the top-left corner.
|
||||
ax.text(
|
||||
0.02,
|
||||
0.9,
|
||||
sensor_name + dim_str,
|
||||
transform=ax.transAxes,
|
||||
fontsize=10,
|
||||
weight="bold",
|
||||
verticalalignment="top",
|
||||
horizontalalignment="left",
|
||||
bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
|
||||
)
|
||||
|
||||
# Enable a dashed grid.
|
||||
ax.grid(True, linestyle="--", alpha=0.7)
|
||||
|
||||
# Loop over "extra" sensors from the user
|
||||
for _ in range(plot_i, n_plots):
|
||||
sensor_name = "user_sensor"
|
||||
dim_str = "" if sensor_dim == 1 else f" {j}"
|
||||
ax = axes[plot_i]
|
||||
plot_i += 1
|
||||
if predicted_data is not None:
|
||||
assert predicted_times is not None
|
||||
predicted_signal = predicted_data[:, plot_i - 1]
|
||||
ax.plot(
|
||||
predicted_times,
|
||||
predicted_signal,
|
||||
lw=2,
|
||||
color=predicted_color,
|
||||
alpha=0.8,
|
||||
label="Sim",
|
||||
)
|
||||
if real_data is not None:
|
||||
assert real_times is not None
|
||||
real_signal = real_data[:, plot_i - 1]
|
||||
ax.plot(
|
||||
real_times,
|
||||
real_signal,
|
||||
lw=2,
|
||||
color=real_color,
|
||||
linestyle="--",
|
||||
alpha=0.7,
|
||||
label="Real",
|
||||
)
|
||||
if preid_data is not None:
|
||||
assert preid_times is not None
|
||||
preid_signal = preid_data[:, plot_i - 1]
|
||||
ax.plot(
|
||||
preid_times,
|
||||
preid_signal,
|
||||
lw=2,
|
||||
color=preid_color,
|
||||
linestyle=":",
|
||||
alpha=0.6,
|
||||
label="Pre-ID",
|
||||
)
|
||||
if commanded_data is not None:
|
||||
assert commanded_times is not None
|
||||
commanded_signal = commanded_data[:, plot_i - 1]
|
||||
ax.plot(
|
||||
commanded_times,
|
||||
commanded_signal,
|
||||
lw=2,
|
||||
color=commanded_color,
|
||||
linestyle="-.",
|
||||
alpha=0.6,
|
||||
label="Commanded",
|
||||
)
|
||||
# Place the sensor name in a white box in the top-left corner.
|
||||
ax.text(
|
||||
0.02,
|
||||
0.9,
|
||||
sensor_name + dim_str,
|
||||
transform=ax.transAxes,
|
||||
fontsize=10,
|
||||
weight="bold",
|
||||
verticalalignment="top",
|
||||
horizontalalignment="left",
|
||||
bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
|
||||
)
|
||||
|
||||
# Enable a dashed grid.
|
||||
ax.grid(True, linestyle="--", alpha=0.7)
|
||||
|
||||
# Add a unified, figure-level legend if any data is provided.
|
||||
legend_handles = []
|
||||
if predicted_data is not None:
|
||||
legend_handles.append(
|
||||
Line2D([0], [0], color=predicted_color, lw=2, label="Simulation")
|
||||
)
|
||||
if real_data is not None:
|
||||
legend_handles.append(
|
||||
Line2D([0], [0], color=real_color, lw=2, linestyle="--", label="Real")
|
||||
)
|
||||
if preid_data is not None:
|
||||
legend_handles.append(
|
||||
Line2D([0], [0], color=preid_color, lw=2, linestyle=":", label="Pre-ID")
|
||||
)
|
||||
if commanded_data is not None:
|
||||
legend_handles.append(
|
||||
Line2D(
|
||||
[0],
|
||||
[0],
|
||||
color=commanded_color,
|
||||
lw=2,
|
||||
linestyle="-.",
|
||||
label="Commanded",
|
||||
)
|
||||
)
|
||||
|
||||
if legend_handles:
|
||||
fig.legend(
|
||||
handles=legend_handles,
|
||||
loc="upper center",
|
||||
bbox_to_anchor=(0.5, 0.935),
|
||||
ncol=len(legend_handles),
|
||||
fancybox=True,
|
||||
shadow=True,
|
||||
fontsize=10,
|
||||
title="Data Source",
|
||||
)
|
||||
|
||||
fig.supxlabel("Time (s)", fontsize=8)
|
||||
plt.tight_layout(rect=(0, 0.03, 1, 0.9))
|
||||
|
||||
|
||||
def plot_objective(
|
||||
objective: Sequence[float],
|
||||
figsize: tuple[float, float] = (8, 5),
|
||||
):
|
||||
plt.figure(figsize=figsize)
|
||||
plt.plot(objective, linewidth=2, marker="o", markersize=4)
|
||||
final_value = objective[-1]
|
||||
if abs(final_value) < 1e-3 or abs(final_value) > 1e3:
|
||||
final_str = f"{final_value:.2e}"
|
||||
else:
|
||||
final_str = f"{final_value:.4f}"
|
||||
plt.title(f"Objective Over Time (Final: {final_str})", fontsize=14, pad=10)
|
||||
plt.grid(True, linestyle="--", alpha=0.6)
|
||||
plt.xlabel("Iteration", fontsize=12)
|
||||
plt.ylabel("Objective", fontsize=12)
|
||||
plt.xticks(fontsize=10)
|
||||
plt.yticks(fontsize=10)
|
||||
plt.tight_layout()
|
||||
|
||||
|
||||
def plot_candidate(
|
||||
candidate: Sequence[np.ndarray],
|
||||
bounds: tuple[Sequence[float] | np.ndarray, Sequence[float] | np.ndarray]
|
||||
| None = None,
|
||||
param_names: Sequence[str] | None = None,
|
||||
figsize: tuple[float, float] = (12, 2.5),
|
||||
dims_per_page: int = 6,
|
||||
log_diff: bool = True,
|
||||
bound_eps: float = 1e-3,
|
||||
):
|
||||
values = np.array(candidate) # shape: (n_iter, n_dim)
|
||||
n_iter, n_dim = values.shape
|
||||
diffs = np.diff(values, axis=0)
|
||||
|
||||
mins = np.full(n_dim, -np.inf)
|
||||
maxs = np.full(n_dim, np.inf)
|
||||
if bounds is not None:
|
||||
mins = np.array(bounds[0])
|
||||
maxs = np.array(bounds[1])
|
||||
assert mins.shape == (n_dim,) and maxs.shape == (n_dim,)
|
||||
|
||||
if param_names is not None:
|
||||
assert len(param_names) == n_dim
|
||||
|
||||
# TODO support pages, they are currently broken because saving to disk overwrites the the pages
|
||||
# n_pages = math.ceil(n_dim / dims_per_page)
|
||||
n_pages = 1
|
||||
for _page in range(n_pages):
|
||||
# start = page * dims_per_page
|
||||
# end = min((page + 1) * dims_per_page, n_dim)
|
||||
start = 0
|
||||
end = n_dim
|
||||
dims_in_page = end - start
|
||||
|
||||
fig, axes = plt.subplots(
|
||||
dims_in_page,
|
||||
2,
|
||||
figsize=(figsize[0], figsize[1] * dims_in_page),
|
||||
sharex="col",
|
||||
)
|
||||
if dims_in_page == 1:
|
||||
axes = np.expand_dims(axes, 0)
|
||||
|
||||
for i, dim in enumerate(range(start, end)):
|
||||
label = param_names[dim] if param_names is not None else f"Dim {dim}"
|
||||
ax_val, ax_diff = axes[i]
|
||||
|
||||
vals = values[:, dim]
|
||||
ax_val.set_ylabel(label, fontsize=10)
|
||||
ax_val.grid(True, linestyle="--", alpha=0.6)
|
||||
ax_val.tick_params(labelsize=9)
|
||||
|
||||
if bounds is not None:
|
||||
lower, upper = mins[dim], maxs[dim]
|
||||
ax_val.axhspan(lower, upper, color="gray", alpha=0.08)
|
||||
ax_val.plot(
|
||||
[0, n_iter - 1],
|
||||
[lower, lower],
|
||||
color="gray",
|
||||
linestyle="--",
|
||||
alpha=0.3,
|
||||
linewidth=1,
|
||||
)
|
||||
ax_val.plot(
|
||||
[0, n_iter - 1],
|
||||
[upper, upper],
|
||||
color="gray",
|
||||
linestyle="--",
|
||||
alpha=0.3,
|
||||
linewidth=1,
|
||||
)
|
||||
near_lower = np.abs(vals - lower) < bound_eps
|
||||
near_upper = np.abs(vals - upper) < bound_eps
|
||||
near_bound = near_lower | near_upper
|
||||
for t in range(1, n_iter):
|
||||
is_near_prev = near_bound[t - 1]
|
||||
is_near_curr = near_bound[t]
|
||||
color = "#d62728" if is_near_prev and is_near_curr else "#1f77b4"
|
||||
ax_val.plot([t - 1, t], [vals[t - 1], vals[t]], color=color, linewidth=2)
|
||||
ax_val.plot(t, vals[t], marker="o", markersize=3, color=color)
|
||||
# Overlay triangle markers for near-bound points
|
||||
for t in range(n_iter):
|
||||
if near_lower[t]:
|
||||
ax_val.plot(t, vals[t], marker="v", markersize=6, color="#d62728")
|
||||
elif near_upper[t]:
|
||||
ax_val.plot(t, vals[t], marker="^", markersize=6, color="#d62728")
|
||||
else:
|
||||
ax_val.plot(vals, linewidth=2, marker="o", markersize=3)
|
||||
|
||||
# Annotate final value
|
||||
final_val = vals[-1]
|
||||
final_str = (
|
||||
f"{final_val:.2e}"
|
||||
if abs(final_val) < 1e-3 or abs(final_val) > 1e3
|
||||
else f"{final_val:.4f}"
|
||||
)
|
||||
ax_val.text(
|
||||
n_iter - 1,
|
||||
final_val,
|
||||
final_str,
|
||||
ha="right",
|
||||
va="bottom",
|
||||
fontsize=9,
|
||||
color="blue",
|
||||
)
|
||||
|
||||
# Annotate final value.
|
||||
final_val = values[-1, dim]
|
||||
final_str = (
|
||||
f"{final_val:.2e}"
|
||||
if abs(final_val) < 1e-3 or abs(final_val) > 1e3
|
||||
else f"{final_val:.4f}"
|
||||
)
|
||||
ax_val.text(
|
||||
n_iter - 1,
|
||||
final_val,
|
||||
final_str,
|
||||
ha="right",
|
||||
va="bottom",
|
||||
fontsize=9,
|
||||
color="blue",
|
||||
)
|
||||
|
||||
# Plot diffs
|
||||
if log_diff:
|
||||
eps = 1e-12
|
||||
ax_diff.plot(
|
||||
np.log10(np.abs(diffs[:, dim]) + eps),
|
||||
linewidth=2,
|
||||
marker="x",
|
||||
markersize=4,
|
||||
color="tab:orange",
|
||||
)
|
||||
ax_diff.set_ylabel("log Δ", fontsize=9)
|
||||
else:
|
||||
ax_diff.plot(
|
||||
diffs[:, dim],
|
||||
linewidth=2,
|
||||
marker="x",
|
||||
markersize=4,
|
||||
color="tab:orange",
|
||||
)
|
||||
|
||||
ax_diff.grid(True, linestyle="--", alpha=0.6)
|
||||
ax_diff.tick_params(labelsize=9)
|
||||
|
||||
# Set common labels/titles
|
||||
axes[-1, 0].set_xlabel("Iteration", fontsize=12)
|
||||
axes[-1, 1].set_xlabel("Iteration", fontsize=12)
|
||||
axes[0, 0].set_title("Candidate Value", fontsize=12)
|
||||
axes[0, 1].set_title("Δ Candidate (Diff)", fontsize=12)
|
||||
|
||||
fig.suptitle(f"Candidate Values and Changes (Dims {start}-{end - 1})", fontsize=14)
|
||||
fig.tight_layout(rect=(0, 0, 1, 0.96))
|
||||
|
||||
|
||||
def plot_candidate_heatmap(
|
||||
candidate: Sequence[np.ndarray],
|
||||
param_names: Sequence[str] | None = None,
|
||||
bounds: tuple[Sequence[float] | np.ndarray, Sequence[float] | np.ndarray]
|
||||
| None = None,
|
||||
normalize: bool = True,
|
||||
figsize: tuple[float, float] = (10, 6),
|
||||
cmap: str = "RdBu",
|
||||
show_colorbar: bool = True,
|
||||
bound_eps: float = 1e-3,
|
||||
):
|
||||
data = np.array(candidate).T # shape: (n_dim, n_iter)
|
||||
n_dim = data.shape[0]
|
||||
|
||||
if normalize and bounds is not None:
|
||||
min_bounds, max_bounds = bounds
|
||||
assert len(min_bounds) == len(max_bounds) == n_dim
|
||||
norm_data = np.empty_like(data)
|
||||
for i in range(n_dim):
|
||||
min_val = min_bounds[i]
|
||||
max_val = max_bounds[i]
|
||||
denom = max_val - min_val if max_val > min_val else 1.0
|
||||
norm_data[i] = (data[i] - min_val) / denom
|
||||
else:
|
||||
norm_data = data
|
||||
|
||||
fig, ax = plt.subplots(figsize=figsize)
|
||||
im = ax.imshow(norm_data, aspect="auto", cmap=cmap)
|
||||
|
||||
ax.set_xlabel("Iteration", fontsize=12)
|
||||
ax.set_ylabel("Parameter", fontsize=12)
|
||||
|
||||
# Y-axis labels.
|
||||
if param_names is not None:
|
||||
assert len(param_names) == n_dim
|
||||
ax.set_yticks(np.arange(n_dim))
|
||||
ax.set_yticklabels(param_names, fontsize=10)
|
||||
else:
|
||||
ax.set_yticks(np.arange(n_dim))
|
||||
ax.set_yticklabels([f"Dim {i}" for i in range(n_dim)], fontsize=10)
|
||||
|
||||
# Plot Xs where values are at bounds.
|
||||
if bounds is not None:
|
||||
min_bounds, max_bounds = bounds
|
||||
for dim in range(n_dim):
|
||||
min_val = min_bounds[dim]
|
||||
max_val = max_bounds[dim]
|
||||
for iter_idx, val in enumerate(data[dim]):
|
||||
if abs(val - min_val) < bound_eps or abs(val - max_val) < bound_eps:
|
||||
ax.plot(iter_idx, dim, "kx", markersize=6, markeredgewidth=1.5)
|
||||
|
||||
if show_colorbar:
|
||||
cbar = fig.colorbar(im, ax=ax)
|
||||
label = "Normalized Value" if normalize else "Value"
|
||||
cbar.set_label(label, fontsize=12)
|
||||
|
||||
ax.set_title("Candidate Heatmap", fontsize=14)
|
||||
fig.tight_layout()
|
||||
|
||||
|
||||
def parameter_confidence(
|
||||
all_exp_names: Sequence[str],
|
||||
all_params: Sequence[parameter.ParameterDict],
|
||||
all_intervals: Sequence[np.ndarray],
|
||||
cols: int = 5,
|
||||
gt_params: parameter.ParameterDict | None = None,
|
||||
):
|
||||
named_estimates = {}
|
||||
# Create an entry for every non-frozen parameter
|
||||
for params in all_params:
|
||||
param_names = params.get_non_frozen_parameter_names()
|
||||
for name in param_names:
|
||||
if name not in named_estimates:
|
||||
named_estimates[name] = {
|
||||
"x": [],
|
||||
"intervals": [],
|
||||
"min_bounds": [],
|
||||
"max_bounds": [],
|
||||
"plot_labels": [],
|
||||
}
|
||||
|
||||
for exp_name, params, intervals in zip(
|
||||
all_exp_names, all_params, all_intervals, strict=True
|
||||
):
|
||||
param_names = params.get_non_frozen_parameter_names()
|
||||
xs = params.as_vector()
|
||||
bounds = params.get_bounds()
|
||||
assert xs.shape[0] == len(param_names)
|
||||
if gt_params is not None:
|
||||
for name in param_names:
|
||||
if name in gt_params:
|
||||
named_estimates[name]["xgt"] = gt_params[name].value[0]
|
||||
else:
|
||||
assert name[-1] == "]"
|
||||
left_bracket_i = name[::-1].find("[")
|
||||
index = int(name[-left_bracket_i:-1])
|
||||
named_estimates[name]["xgt"] = gt_params[name[: -left_bracket_i - 1]].value[
|
||||
index
|
||||
]
|
||||
|
||||
for i, (name, x, interval) in enumerate(
|
||||
zip(param_names, xs, intervals, strict=True)
|
||||
):
|
||||
named_estimates[name]["x"].append(x)
|
||||
named_estimates[name]["intervals"].append(interval)
|
||||
named_estimates[name]["min_bounds"].append(bounds[0][i])
|
||||
named_estimates[name]["max_bounds"].append(bounds[1][i])
|
||||
named_estimates[name]["plot_labels"].append(exp_name)
|
||||
|
||||
rows = len(named_estimates) // cols + 1
|
||||
fig, axs = plt.subplots(
|
||||
rows, cols, figsize=(20, 2 * (len(named_estimates) // cols + 1))
|
||||
)
|
||||
if rows == 1:
|
||||
axs = [axs]
|
||||
|
||||
for i, name in enumerate(named_estimates):
|
||||
x_list = named_estimates[name]["x"]
|
||||
intervals = named_estimates[name]["intervals"]
|
||||
plot_labels = named_estimates[name]["plot_labels"]
|
||||
|
||||
row = i % rows
|
||||
col = i // rows
|
||||
|
||||
min_bound = np.min(named_estimates[name]["min_bounds"])
|
||||
max_bound = np.min(named_estimates[name]["max_bounds"])
|
||||
|
||||
for j, (x, interval, plot_label) in enumerate(
|
||||
zip(x_list, intervals, plot_labels, strict=True)
|
||||
):
|
||||
if not np.isfinite(interval) or 2.0 * interval > 2.0 * (max_bound - min_bound):
|
||||
interval = 2.0 * (max_bound - min_bound)
|
||||
eb = axs[row][col].errorbar(x, -j, xerr=interval)
|
||||
eb[-1][0].set_linestyle("--")
|
||||
else:
|
||||
axs[row][col].errorbar(x, -j, xerr=interval)
|
||||
axs[row][col].scatter(x, -j, marker="x", label=plot_label)
|
||||
|
||||
axs[row][col].set_xlim([min_bound, max_bound])
|
||||
axs[row][col].yaxis.set_ticklabels([])
|
||||
axs[row][col].set_title(name)
|
||||
axs[row][col].grid(True)
|
||||
axs[row][col].legend(fontsize=5, loc="upper right", bbox_to_anchor=(1.4, 1.0))
|
||||
if gt_params is not None:
|
||||
axs[row][col].axvline(named_estimates[name]["xgt"], color="b", ls="--")
|
||||
|
||||
fig.tight_layout()
|
||||
|
||||
|
||||
def render_rollout(
|
||||
model: mujoco.MjModel | Sequence[mujoco.MjModel],
|
||||
data: mujoco.MjData,
|
||||
state: np.ndarray,
|
||||
framerate: int,
|
||||
camera: str | int = -1,
|
||||
width: int = 640,
|
||||
height: int = 480,
|
||||
light_pos: Sequence[float] | None = None,
|
||||
) -> list[np.ndarray]:
|
||||
"""Renders a rollout or batch of rollouts.
|
||||
|
||||
Args:
|
||||
model: Single model or list of models (one per batch).
|
||||
data: MjData scratch object.
|
||||
state: State array of shape (nbatch, nsteps, nstate).
|
||||
framerate: Frames per second to render.
|
||||
camera: Camera name or ID.
|
||||
width: Image width.
|
||||
height: Image height.
|
||||
light_pos: Optional light position [x, y, z] to add a spotlight.
|
||||
|
||||
Returns:
|
||||
List of rendered frames (numpy arrays).
|
||||
"""
|
||||
nbatch = state.shape[0]
|
||||
|
||||
if isinstance(model, mujoco.MjModel):
|
||||
models_list = [model] * nbatch
|
||||
else:
|
||||
models_list = list(model)
|
||||
if len(models_list) == 1:
|
||||
models_list = models_list * nbatch
|
||||
else:
|
||||
assert len(models_list) == nbatch
|
||||
|
||||
# Visual options
|
||||
vopt = mujoco.MjvOption()
|
||||
vopt.geomgroup[3] = 1 # Show visualization geoms
|
||||
|
||||
pert = mujoco.MjvPerturb()
|
||||
catmask = mujoco.mjtCatBit.mjCAT_DYNAMIC
|
||||
|
||||
# Simulate and render.
|
||||
frames = []
|
||||
|
||||
with mujoco.Renderer(models_list[0], height=height, width=width) as renderer:
|
||||
for i in range(state.shape[1]):
|
||||
# Check if we should capture this frame based on framerate
|
||||
if len(frames) < i * models_list[0].opt.timestep * framerate:
|
||||
for j in range(state.shape[0]):
|
||||
# Set state
|
||||
mujoco.mj_setState(
|
||||
models_list[j], data, state[j, i, :], mujoco.mjtState.mjSTATE_FULLPHYSICS
|
||||
)
|
||||
mujoco.mj_forward(models_list[j], data)
|
||||
|
||||
# Use first model to make the scene, add subsequent models
|
||||
if j == 0:
|
||||
renderer.update_scene(data, camera, scene_option=vopt)
|
||||
else:
|
||||
mujoco.mjv_addGeoms(
|
||||
models_list[j], data, vopt, pert, catmask, renderer.scene
|
||||
)
|
||||
|
||||
# Add light, if requested
|
||||
if light_pos is not None:
|
||||
if renderer.scene.nlight < 100: # check limit
|
||||
light = renderer.scene.lights[renderer.scene.nlight]
|
||||
light.ambient = [0, 0, 0]
|
||||
light.attenuation = [1, 0, 0]
|
||||
light.castshadow = 1
|
||||
light.cutoff = 45
|
||||
light.diffuse = [0.8, 0.8, 0.8]
|
||||
light.dir = [0, 0, -1]
|
||||
light.type = mujoco.mjtLightType.mjLIGHT_SPOT
|
||||
light.exponent = 10
|
||||
light.headlight = 0
|
||||
light.specular = [0.3, 0.3, 0.3]
|
||||
light.pos = light_pos
|
||||
renderer.scene.nlight += 1
|
||||
|
||||
# Render and add the frame.
|
||||
pixels = renderer.render()
|
||||
frames.append(pixels)
|
||||
return frames
|
||||
@@ -0,0 +1,408 @@
|
||||
"""Residual computation for system identification."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import os
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import TypeAlias
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
|
||||
from mujoco.sysid._src import (
|
||||
model_modifier,
|
||||
parameter,
|
||||
signal_modifier,
|
||||
timeseries,
|
||||
)
|
||||
from mujoco.sysid._src.trajectory import (
|
||||
ModelSequences,
|
||||
SystemTrajectory,
|
||||
sysid_rollout,
|
||||
)
|
||||
|
||||
_NUM_CPUS: int = os.cpu_count() or 1
|
||||
|
||||
BuildModelFn: TypeAlias = Callable[
|
||||
[parameter.ParameterDict, mujoco.MjSpec], mujoco.MjModel
|
||||
]
|
||||
|
||||
CustomRolloutFn: TypeAlias = Callable[..., Sequence[SystemTrajectory]]
|
||||
"""Replaces the default sysid_rollout. Called with keyword arguments:
|
||||
models, datas, control_signal, initial_states, param_dicts,
|
||||
rollout_signal_mapping, rollout_state_mapping, ctrl_mapping."""
|
||||
|
||||
ModifyResidualFn: TypeAlias = Callable[
|
||||
..., tuple[np.ndarray, timeseries.TimeSeries, timeseries.TimeSeries]
|
||||
]
|
||||
"""Custom residual computation. Called as:
|
||||
modify_residual(params, sensordata_predicted, sensordata_measured,
|
||||
model, return_pred_all, state=..., sensor_weights=...)."""
|
||||
|
||||
|
||||
def construct_ts_from_defaults(
|
||||
state_ts: timeseries.TimeSeries,
|
||||
pred_sensordata: timeseries.TimeSeries,
|
||||
measured_sensordata: timeseries.TimeSeries,
|
||||
enabled_observations: Sequence[tuple[str, timeseries.SignalType]] | None = None,
|
||||
):
|
||||
"""Assemble predicted observations to match the measured signal layout.
|
||||
|
||||
For each enabled observation, copies the predicted values from either
|
||||
``pred_sensordata`` (for MjSensor signals) or ``state_ts`` (for state
|
||||
signals like qpos/qvel/act) into a new array whose columns align with
|
||||
the measured data.
|
||||
|
||||
Args:
|
||||
state_ts: Predicted state TimeSeries (time column already stripped).
|
||||
pred_sensordata: Raw predicted sensor TimeSeries from rollout.
|
||||
measured_sensordata: Measured sensor TimeSeries (defines the target layout).
|
||||
enabled_observations: Subset of observations to include. If None, all
|
||||
observations in ``measured_sensordata`` are used.
|
||||
|
||||
Returns:
|
||||
A ``(measured, predicted)`` tuple of TimeSeries with matching signal
|
||||
mappings, sliced to the enabled observations.
|
||||
"""
|
||||
assert measured_sensordata.signal_mapping is not None
|
||||
|
||||
# Trim measured data enabled observations
|
||||
if enabled_observations:
|
||||
enabled_observations_names = [i[0] for i in enabled_observations]
|
||||
enabled_observations_types = [i[1] for i in enabled_observations]
|
||||
else:
|
||||
enabled_observations_names = list(measured_sensordata.signal_mapping.keys())
|
||||
enabled_observations_types = [
|
||||
v[0] for v in measured_sensordata.signal_mapping.values()
|
||||
]
|
||||
|
||||
selected_measured_sensordata = timeseries.TimeSeries.slice_by_name(
|
||||
measured_sensordata, enabled_observations_names
|
||||
)
|
||||
assert selected_measured_sensordata.signal_mapping is not None
|
||||
selected_measured_signal_mapping = selected_measured_sensordata.signal_mapping
|
||||
|
||||
shape = (pred_sensordata.data.shape[0], selected_measured_sensordata.data.shape[1])
|
||||
predicted_data_out = np.zeros(shape)
|
||||
|
||||
measured_signal_mapping = measured_sensordata.signal_mapping
|
||||
for enabled_obs_name, enabled_obs_type in zip(
|
||||
enabled_observations_names, enabled_observations_types, strict=True
|
||||
):
|
||||
assert state_ts.signal_mapping is not None
|
||||
if (
|
||||
enabled_obs_name not in measured_signal_mapping
|
||||
and enabled_obs_name not in state_ts.signal_mapping
|
||||
):
|
||||
raise ValueError(f"{enabled_obs_name} is missing.")
|
||||
|
||||
obs_type, indices = measured_signal_mapping[enabled_obs_name]
|
||||
|
||||
if obs_type != enabled_obs_type:
|
||||
raise ValueError(
|
||||
f"Observation type error: {enabled_obs_name} is of type {obs_type} but declared as {enabled_obs_type}."
|
||||
)
|
||||
|
||||
if obs_type == timeseries.SignalType.CustomObs:
|
||||
raise ValueError(
|
||||
f"You are attempting to use the default SysID's modify_residual with a custom observation of name {enabled_obs_name}. This is not supported. You must implement your own modify_residual. See documentation at ..."
|
||||
)
|
||||
|
||||
elif obs_type == timeseries.SignalType.MjSensor:
|
||||
target_indices = selected_measured_signal_mapping[enabled_obs_name][1]
|
||||
predicted_data_out[:, ..., target_indices] = pred_sensordata.data[:, ..., indices]
|
||||
|
||||
elif (
|
||||
obs_type == timeseries.SignalType.MjStateQPos
|
||||
or obs_type == timeseries.SignalType.MjStateQVel
|
||||
or obs_type == timeseries.SignalType.MjStateAct
|
||||
):
|
||||
state_indices = state_ts.signal_mapping[enabled_obs_name][1]
|
||||
|
||||
values = state_ts.data[:, ..., state_indices]
|
||||
target_indices = selected_measured_signal_mapping[enabled_obs_name][1]
|
||||
predicted_data_out[:, ..., target_indices] = values
|
||||
|
||||
ts_predicted_data = timeseries.TimeSeries(
|
||||
pred_sensordata.times,
|
||||
predicted_data_out,
|
||||
selected_measured_sensordata.signal_mapping,
|
||||
)
|
||||
|
||||
return selected_measured_sensordata, ts_predicted_data
|
||||
|
||||
|
||||
# Lowest level residual function, works on one model
|
||||
def model_residual(
|
||||
x: np.ndarray,
|
||||
params: parameter.ParameterDict,
|
||||
build_model: Callable[[parameter.ParameterDict], mujoco.MjModel],
|
||||
traj_measured: Sequence[SystemTrajectory] | SystemTrajectory,
|
||||
modify_residual: ModifyResidualFn | None = None,
|
||||
custom_rollout: CustomRolloutFn | None = None,
|
||||
n_threads: int = _NUM_CPUS,
|
||||
return_pred_all: bool = False,
|
||||
resample_true: bool = True,
|
||||
sensor_weights: Mapping[str, float] | None = None,
|
||||
enabled_observations: Sequence[tuple[str, timeseries.SignalType]] = (),
|
||||
):
|
||||
"""Compute residuals for a single model against measured trajectories.
|
||||
|
||||
Builds the model from *x*, rolls out each trajectory, and computes the
|
||||
weighted difference between predicted and measured sensor data.
|
||||
|
||||
Args:
|
||||
x: Decision variable vector (flat, or 2-D for batched finite-difference).
|
||||
params: Parameter dictionary — updated in-place from *x*.
|
||||
build_model: ``(ParameterDict) -> MjModel`` factory.
|
||||
traj_measured: Ground-truth trajectory or sequence of trajectories.
|
||||
modify_residual: Optional custom residual callback (replaces the default
|
||||
resampling / differencing logic).
|
||||
custom_rollout: Optional replacement for :func:`sysid_rollout`.
|
||||
n_threads: Number of ``MjData`` scratch objects for parallel rollout.
|
||||
return_pred_all: If True, return full predicted/measured TimeSeries.
|
||||
resample_true: Whether to resample the measured data at simulation
|
||||
timesteps (ignored when *modify_residual* is provided).
|
||||
sensor_weights: Per-sensor weights for the weighted diff.
|
||||
enabled_observations: Subset of ``(name, SignalType)`` pairs to include.
|
||||
|
||||
Returns:
|
||||
A 3-tuple ``(residuals, pred_sensordatas, measured_sensordatas)``.
|
||||
"""
|
||||
# Convert single trajectory to list for consistent handling.
|
||||
if isinstance(traj_measured, SystemTrajectory):
|
||||
traj_measured = [traj_measured]
|
||||
n_chunks = len(traj_measured)
|
||||
|
||||
# Handle finite difference columns if present.
|
||||
initial_ndim = x.ndim
|
||||
n_fd = 1
|
||||
if x.ndim > 1:
|
||||
n_fd = x.shape[1]
|
||||
x_reshaped = x
|
||||
else:
|
||||
x_reshaped = x.reshape(-1, 1)
|
||||
|
||||
# Process each finite difference column.
|
||||
models = []
|
||||
models_x = []
|
||||
model_0 = None
|
||||
for i in range(n_fd):
|
||||
params.update_from_vector(x_reshaped[:, i])
|
||||
model = build_model(params)
|
||||
if not model_0:
|
||||
model_0 = model
|
||||
models_x.extend([x_reshaped[:, i]] * n_chunks)
|
||||
models.extend([model] * n_chunks)
|
||||
|
||||
assert model_0 is not None
|
||||
qpos_map, qvel_map, act_map, rollout_ctrl_map = (
|
||||
timeseries.TimeSeries.compute_all_state_mappings(model_0)
|
||||
)
|
||||
rollout_state_mapping = qpos_map | qvel_map | act_map
|
||||
rollout_signal_mapping = timeseries.TimeSeries.compute_all_sensor_mapping(model_0)
|
||||
|
||||
# Create data objects for parallel computation.
|
||||
datas = [mujoco.MjData(models[0]) for _ in range(n_threads)]
|
||||
|
||||
# Interpolate control signal.
|
||||
if resample_true:
|
||||
control_chunks = [
|
||||
traj.control.resample(target_dt=models[0].opt.timestep) for traj in traj_measured
|
||||
]
|
||||
else:
|
||||
control_chunks = [traj.control for traj in traj_measured]
|
||||
|
||||
# Rollout trajectories in parallel.
|
||||
if custom_rollout is None:
|
||||
pred_trajectories = sysid_rollout(
|
||||
models=models[: n_fd * n_chunks],
|
||||
datas=datas,
|
||||
control_signal=[control for control in control_chunks] * n_fd,
|
||||
initial_states=[chunk.initial_state for chunk in traj_measured] * n_fd,
|
||||
rollout_signal_mapping=rollout_signal_mapping,
|
||||
rollout_state_mapping=rollout_state_mapping,
|
||||
ctrl_mapping=rollout_ctrl_map,
|
||||
)
|
||||
else:
|
||||
param_dicts = [copy.deepcopy(params) for i in range(x_reshaped.shape[1])]
|
||||
[
|
||||
param_dicts[i].update_from_vector(x_reshaped[:, i])
|
||||
for i in range(x_reshaped.shape[1])
|
||||
]
|
||||
pred_trajectories = custom_rollout(
|
||||
models=models[: n_fd * n_chunks],
|
||||
datas=datas,
|
||||
control_signal=[control for control in control_chunks] * n_fd,
|
||||
initial_states=[chunk.initial_state for chunk in traj_measured] * n_fd,
|
||||
param_dicts=param_dicts,
|
||||
rollout_signal_mapping=rollout_signal_mapping,
|
||||
rollout_state_mapping=rollout_state_mapping,
|
||||
ctrl_mapping=rollout_ctrl_map,
|
||||
)
|
||||
|
||||
# Compute residuals for each trajectory chunk.
|
||||
all_residuals = []
|
||||
pred_sensordatas = []
|
||||
measured_sensordatas = []
|
||||
|
||||
for i in range(len(models)):
|
||||
model = models[i]
|
||||
pred_traj = pred_trajectories[i]
|
||||
assert pred_traj.state is not None
|
||||
pred_state = pred_traj.state.data
|
||||
|
||||
rollout_state_ts = timeseries.TimeSeries(
|
||||
times=pred_state[:, 0],
|
||||
data=pred_state[:, 1:],
|
||||
signal_mapping=rollout_state_mapping,
|
||||
)
|
||||
|
||||
measuredidx = i % n_chunks
|
||||
measuredtraj = traj_measured[measuredidx]
|
||||
|
||||
pred_sensordata = pred_traj.sensordata
|
||||
measured_sensordata = measuredtraj.sensordata
|
||||
|
||||
# If the user passes a residual function allow them to handle all resampling, etc.
|
||||
if modify_residual is not None:
|
||||
params.update_from_vector(models_x[i])
|
||||
res, pred_sensordata, measured_sensordata = modify_residual(
|
||||
params,
|
||||
pred_sensordata,
|
||||
measured_sensordata,
|
||||
model,
|
||||
return_pred_all,
|
||||
state=pred_state,
|
||||
)
|
||||
|
||||
# If the user does not pass a residual function, resample the ground truth data to
|
||||
# match the sime times if requested.
|
||||
else:
|
||||
measured_sensordata, pred_sensordata = construct_ts_from_defaults(
|
||||
rollout_state_ts, pred_sensordata, measured_sensordata, enabled_observations
|
||||
)
|
||||
if resample_true:
|
||||
# Window the true data so that times in it correspond to times spanned by
|
||||
# predicted data.
|
||||
measured_sensordata = signal_modifier.apply_delayed_ts_window(
|
||||
measured_sensordata, pred_sensordata, 0.0, 0.0
|
||||
)
|
||||
# Sample the predicted signal at the true times.
|
||||
pred_sensordata = pred_sensordata.resample(measured_sensordata.times)
|
||||
|
||||
else:
|
||||
# Do not include difference in first sensor outputs in residual vector.
|
||||
# It corresponds to the initial condition and so provides little new
|
||||
# information. Additionally the semantics of rollout make it difficult to
|
||||
# simulate the sensor output corresponding to the initial condition.
|
||||
measured_sensordata = timeseries.TimeSeries(
|
||||
measured_sensordata.times[1:],
|
||||
measured_sensordata.data[1:, :],
|
||||
measured_sensordata.signal_mapping,
|
||||
)
|
||||
|
||||
res = signal_modifier.weighted_diff(
|
||||
predicted_data=pred_sensordata.data,
|
||||
measured_data=measured_sensordata.data,
|
||||
model=model,
|
||||
sensor_weights=sensor_weights,
|
||||
)
|
||||
res = signal_modifier.normalize_residual(res, measured_sensordata.data)
|
||||
|
||||
if pred_sensordata.signal_mapping != measured_sensordata.signal_mapping:
|
||||
raise ValueError(
|
||||
"The observation mapping between the measured data and predicted rollout data"
|
||||
" is not the same. You have not modified the observation data in TimeSeries"
|
||||
" in modify_residual to correctly reflect the measured data."
|
||||
)
|
||||
|
||||
all_residuals.append(res)
|
||||
pred_sensordatas.append(pred_sensordata)
|
||||
measured_sensordatas.append(measured_sensordata)
|
||||
|
||||
res_array = np.stack(all_residuals, axis=0)
|
||||
if initial_ndim == 1:
|
||||
res_array = res_array.ravel()
|
||||
else:
|
||||
res_array = res_array.reshape(res_array.shape[0], -1)
|
||||
|
||||
return res_array.T, pred_sensordatas, measured_sensordatas
|
||||
|
||||
|
||||
def build_residual_fn(**captured_kwargs):
|
||||
"""Create a residual closure with pre-bound keyword arguments.
|
||||
|
||||
Returns a function ``fn(x, params, **overrides)`` that calls
|
||||
:func:`residual` with the captured kwargs merged in. This is the
|
||||
recommended way to construct the callable passed to :func:`optimize`.
|
||||
|
||||
Example::
|
||||
|
||||
residual_fn = build_residual_fn(
|
||||
models_sequences=seqs,
|
||||
signal_transform=transform,
|
||||
)
|
||||
opt_params, result = optimize(params, residual_fn)
|
||||
"""
|
||||
|
||||
def built_residual_fn(x, params, **kwargs):
|
||||
return residual(
|
||||
x,
|
||||
params,
|
||||
**captured_kwargs,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return built_residual_fn
|
||||
|
||||
|
||||
def residual(
|
||||
x: np.ndarray,
|
||||
params: parameter.ParameterDict,
|
||||
models_sequences: list[ModelSequences],
|
||||
build_model: BuildModelFn = model_modifier.apply_param_modifiers,
|
||||
modify_residual: ModifyResidualFn | None = None,
|
||||
custom_rollout: CustomRolloutFn | None = None,
|
||||
n_threads: int = _NUM_CPUS,
|
||||
return_pred_all: bool = False,
|
||||
resample_true: bool = True,
|
||||
sensor_weights: Mapping[str, float] | None = None,
|
||||
enabled_observations: Sequence[tuple[str, timeseries.SignalType]] = (),
|
||||
):
|
||||
"""Top-level residual: iterate over all model-sequence groups.
|
||||
|
||||
Calls :func:`model_residual` for every measured rollout in every
|
||||
:class:`ModelSequences` entry and collects the results.
|
||||
|
||||
Returns:
|
||||
A 3-tuple ``(residuals, preds, records)`` — lists with one entry per
|
||||
measured rollout across all groups.
|
||||
"""
|
||||
residuals = []
|
||||
preds = []
|
||||
records = []
|
||||
for model_sequences in models_sequences:
|
||||
for measured_rollout in model_sequences.measured_rollout:
|
||||
res = model_residual(
|
||||
x,
|
||||
params,
|
||||
lambda p, _spec=model_sequences.spec: build_model(p, _spec),
|
||||
measured_rollout,
|
||||
modify_residual,
|
||||
custom_rollout,
|
||||
n_threads,
|
||||
return_pred_all,
|
||||
resample_true,
|
||||
sensor_weights,
|
||||
enabled_observations,
|
||||
)
|
||||
if isinstance(res, np.ndarray):
|
||||
residuals.append(res)
|
||||
else:
|
||||
residuals.append(res[0])
|
||||
preds.append(res[1])
|
||||
records.append(res[2])
|
||||
|
||||
return residuals, preds, records
|
||||
@@ -0,0 +1,244 @@
|
||||
"""Common signal modifiers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
|
||||
from mujoco.sysid._src import parameter, timeseries
|
||||
|
||||
|
||||
def _get_sensor_indices(model: mujoco.MjModel, sensor_name: str) -> list[int]:
|
||||
sensor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR.value, sensor_name)
|
||||
if sensor_id == -1:
|
||||
raise ValueError(f"Sensor not found in model: {sensor_name}")
|
||||
|
||||
addr = model.sensor_adr[sensor_id]
|
||||
dim = model.sensor_dim[sensor_id]
|
||||
|
||||
return list(range(addr, addr + dim))
|
||||
|
||||
|
||||
def get_sensor_indices(
|
||||
model: mujoco.MjModel,
|
||||
sensor_name: str | list[str],
|
||||
sort: bool = False,
|
||||
) -> list[int]:
|
||||
"""Get sensor indices from a sensor configuration dictionary.
|
||||
|
||||
Args:
|
||||
model: MuJoCo model containing the sensors.
|
||||
sensor_name: sensor name or list of names to return the indices for
|
||||
"""
|
||||
if isinstance(sensor_name, str):
|
||||
return _get_sensor_indices(model, sensor_name)
|
||||
all_indices = []
|
||||
for name in sensor_name:
|
||||
all_indices.extend(_get_sensor_indices(model, name))
|
||||
if sort:
|
||||
return sorted(all_indices)
|
||||
return all_indices
|
||||
|
||||
|
||||
def apply_bias(
|
||||
ts: timeseries.TimeSeries,
|
||||
sensor_name: str,
|
||||
bias: parameter.Parameter,
|
||||
) -> timeseries.TimeSeries:
|
||||
indices = ts.get_indices(sensor_name)[1]
|
||||
data_out = ts.data.copy()
|
||||
data_out[..., indices] += bias.value
|
||||
return timeseries.TimeSeries(ts.times, data_out, ts.signal_mapping)
|
||||
|
||||
|
||||
def apply_gain(
|
||||
ts: timeseries.TimeSeries,
|
||||
sensor_name: str,
|
||||
gain: parameter.Parameter,
|
||||
) -> timeseries.TimeSeries:
|
||||
indices = ts.get_indices(sensor_name)[1]
|
||||
data_out = ts.data.copy()
|
||||
data_out[..., indices] *= gain.value
|
||||
return timeseries.TimeSeries(ts.times, data_out, ts.signal_mapping)
|
||||
|
||||
|
||||
def apply_delay(
|
||||
ts: timeseries.TimeSeries,
|
||||
sensor_name: str,
|
||||
delay: parameter.Parameter,
|
||||
) -> timeseries.TimeSeries:
|
||||
indices = ts.get_indices(sensor_name)[1]
|
||||
|
||||
ts_sensor = timeseries.TimeSeries(ts.times, ts.data[:, indices], ts.signal_mapping)
|
||||
ts_sensor_delayed = ts_sensor.resample(ts.times - delay.value)
|
||||
|
||||
ts_delayed = timeseries.TimeSeries(ts.times, ts.data, ts.signal_mapping)
|
||||
ts_delayed.data[:, indices] = ts_sensor_delayed.data
|
||||
|
||||
return ts_delayed
|
||||
|
||||
|
||||
def apply_time_window(
|
||||
ts: timeseries.TimeSeries,
|
||||
min_t: float,
|
||||
max_t: float,
|
||||
) -> timeseries.TimeSeries:
|
||||
"""Select a subset of a timeseries whose timestamps plus the max delay can be
|
||||
sampled from ts_sample."""
|
||||
min_i = np.searchsorted(ts.times, min_t, side="left")
|
||||
max_i = np.searchsorted(ts.times, max_t, side="right")
|
||||
return timeseries.TimeSeries(
|
||||
ts.times[min_i:max_i], ts.data[min_i:max_i], ts.signal_mapping
|
||||
)
|
||||
|
||||
|
||||
def apply_delayed_ts_window(
|
||||
ts: timeseries.TimeSeries,
|
||||
ts_delayed: timeseries.TimeSeries,
|
||||
min_delay: float,
|
||||
max_delay: float,
|
||||
) -> timeseries.TimeSeries:
|
||||
"""Window a timeseries so that the included timestamps lay within the bounds of a
|
||||
timeseries that may be delayed between min_delay and max_delay.
|
||||
|
||||
Args:
|
||||
ts: The timeseries to window.
|
||||
ts_delayed: The timeseries to use as the bounds.
|
||||
min_delay: The minimum delay. May be negative.
|
||||
max_delay: The maximum delay.
|
||||
|
||||
Returns:
|
||||
A new timeseries with the timestamps windowed.
|
||||
"""
|
||||
if min_delay > max_delay:
|
||||
raise ValueError(
|
||||
"min_delay must be less than or equal to max_delay, "
|
||||
f"received {min_delay} and {max_delay}"
|
||||
)
|
||||
return apply_time_window(
|
||||
ts, ts_delayed.times[0] - min_delay, ts_delayed.times[-1] - max_delay
|
||||
)
|
||||
|
||||
|
||||
def _build_per_column_delays(
|
||||
ts: timeseries.TimeSeries,
|
||||
default_delay: float,
|
||||
sensor_delays: dict[str, float] | None,
|
||||
predicted_data: bool,
|
||||
) -> list[float]:
|
||||
"""Build a per-column delay list, shared by both implementations."""
|
||||
delays = [default_delay] * ts.data.shape[1]
|
||||
if sensor_delays is None:
|
||||
sensor_delays = {}
|
||||
for name, delay in sensor_delays.items():
|
||||
sensor_indices = ts.get_indices(name)[1]
|
||||
for i in sensor_indices:
|
||||
delays[i] = delay
|
||||
if predicted_data:
|
||||
delays = [-d for d in delays]
|
||||
return delays
|
||||
|
||||
|
||||
def _apply_resample_and_delay_columnwise(
|
||||
ts: timeseries.TimeSeries,
|
||||
times: np.ndarray,
|
||||
delays: list[float],
|
||||
) -> np.ndarray:
|
||||
"""Reference implementation: resample each column independently."""
|
||||
resampled_ts = []
|
||||
for i, d in enumerate(delays):
|
||||
ts_sliced = timeseries.TimeSeries(
|
||||
ts.times, ts.data[:, i : i + 1], ts.signal_mapping
|
||||
)
|
||||
ts_sliced_resampled = ts_sliced.resample(times + d)
|
||||
resampled_ts.append(ts_sliced_resampled)
|
||||
return np.concatenate([t.data for t in resampled_ts], axis=1)
|
||||
|
||||
|
||||
_VERIFY_RESAMPLE_GROUPING = False
|
||||
|
||||
|
||||
def apply_resample_and_delay(
|
||||
ts: timeseries.TimeSeries,
|
||||
times: np.ndarray,
|
||||
default_delay: float,
|
||||
sensor_delays: dict[str, float] | None = None,
|
||||
predicted_data: bool = True,
|
||||
) -> timeseries.TimeSeries:
|
||||
delays = _build_per_column_delays(ts, default_delay, sensor_delays, predicted_data)
|
||||
|
||||
# Group columns by delay value to minimize interpolation calls.
|
||||
delay_to_cols: dict[float, list[int]] = {}
|
||||
for i, d in enumerate(delays):
|
||||
delay_to_cols.setdefault(d, []).append(i)
|
||||
|
||||
data_out = np.empty((len(times), ts.data.shape[1]))
|
||||
for d, cols in delay_to_cols.items():
|
||||
group_data = ts.data[:, cols]
|
||||
group_ts = timeseries.TimeSeries(ts.times, group_data, ts.signal_mapping)
|
||||
resampled = group_ts.resample(times + d)
|
||||
data_out[:, cols] = resampled.data
|
||||
|
||||
if _VERIFY_RESAMPLE_GROUPING:
|
||||
reference = _apply_resample_and_delay_columnwise(ts, times, delays)
|
||||
np.testing.assert_array_equal(data_out, reference)
|
||||
|
||||
return timeseries.TimeSeries(times, data_out, ts.signal_mapping)
|
||||
|
||||
|
||||
def prepare_sensor_weights(
|
||||
sensor_weights: Mapping[str, float] | np.ndarray,
|
||||
n_sensors: int,
|
||||
model: mujoco.MjModel,
|
||||
) -> np.ndarray:
|
||||
if isinstance(sensor_weights, np.ndarray):
|
||||
if sensor_weights.ndim != 1 or sensor_weights.shape[0] != n_sensors:
|
||||
raise ValueError(
|
||||
"Expected sensor_weights to be a numpy array of shape (n_sensors,), "
|
||||
f"received {sensor_weights.shape}"
|
||||
)
|
||||
return sensor_weights
|
||||
else:
|
||||
weights = np.ones(n_sensors)
|
||||
ids = get_sensor_indices(model, list(sensor_weights.keys()))
|
||||
for i, w in zip(ids, sensor_weights.values(), strict=True):
|
||||
weights[i] = w
|
||||
return weights
|
||||
|
||||
|
||||
def weighted_diff(
|
||||
predicted_data: np.ndarray,
|
||||
measured_data: np.ndarray,
|
||||
model: mujoco.MjModel | None = None,
|
||||
sensor_weights: Mapping[str, float] | np.ndarray | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Compute the difference `measured_data - predicted_data`, optionally scaled by
|
||||
sensor weights.
|
||||
|
||||
Args:
|
||||
predicted_data: The predicted data, of shape (n_timesteps, n_sensors).
|
||||
measured_data: The measured data, of shape (n_timesteps, n_sensors).
|
||||
sensor_weights: An optional dict mapping sensor name to weight. Unspecified sensors
|
||||
are assumed to have a weight of 1.
|
||||
model: Optional mujoco model. This argument is required if sensor_weights is not
|
||||
None.
|
||||
|
||||
Returns:
|
||||
A numpy array of the weighted difference.
|
||||
"""
|
||||
res = measured_data - predicted_data
|
||||
if sensor_weights is None:
|
||||
return res
|
||||
if model is None:
|
||||
raise ValueError("model is required if sensor_weights is provided")
|
||||
return res * prepare_sensor_weights(sensor_weights, res.shape[-1], model)
|
||||
|
||||
|
||||
def normalize_residual(
|
||||
residual: np.ndarray,
|
||||
measured_data: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
"""Normalize the residual by the standard deviation of the measured data."""
|
||||
return residual / (np.linalg.norm(measured_data, axis=0) / np.sqrt(2))
|
||||
@@ -0,0 +1,257 @@
|
||||
"""Declarative signal transformation for system identification residuals."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from fnmatch import fnmatch
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
|
||||
from mujoco.sysid._src import parameter, signal_modifier, timeseries
|
||||
|
||||
|
||||
class SignalTransform:
|
||||
"""Declarative signal transformation replacing boilerplate modify_residual callbacks.
|
||||
|
||||
Usage::
|
||||
|
||||
transform = SignalTransform()
|
||||
transform.delay("*_pos", params["delay_pos"])
|
||||
transform.delay("*_torque", params["delay_torque"])
|
||||
transform.gain("*_torque", params["torque_scale"], target="predicted")
|
||||
transform.enable_sensors(cfg.sensors_enabled)
|
||||
|
||||
The ``apply`` method has the same signature as ``ModifyResidualFn`` and can
|
||||
be passed directly to ``build_residual_fn(signal_transform=transform)``.
|
||||
"""
|
||||
|
||||
def __init__(self, normalize: bool = True):
|
||||
self._delays: list[tuple[str, str, parameter.Parameter]] = []
|
||||
self._gains: list[tuple[str, str, str]] = []
|
||||
self._biases: list[tuple[str, str, str]] = []
|
||||
self._enabled_sensors: list[str] | None = None
|
||||
self._sensor_weights: Mapping[str, float] | None = None
|
||||
self.normalize = normalize
|
||||
|
||||
def delay(self, pattern: str, param: parameter.Parameter) -> None:
|
||||
"""Register a delay for sensors matching *pattern* (fnmatch)."""
|
||||
self._delays.append((pattern, param.name, param))
|
||||
|
||||
def gain(
|
||||
self,
|
||||
pattern: str,
|
||||
param: parameter.Parameter,
|
||||
target: str = "both",
|
||||
) -> None:
|
||||
"""Register a multiplicative gain for sensors matching *pattern*.
|
||||
|
||||
Args:
|
||||
pattern: fnmatch pattern matched against sensor names.
|
||||
param: Parameter whose ``.value`` is the gain factor.
|
||||
target: One of ``"predicted"``, ``"measured"``, or ``"both"``.
|
||||
"""
|
||||
if target not in ("predicted", "measured", "both"):
|
||||
raise ValueError(
|
||||
f"target must be 'predicted', 'measured', or 'both', got {target!r}"
|
||||
)
|
||||
self._gains.append((pattern, param.name, target))
|
||||
|
||||
def bias(
|
||||
self,
|
||||
pattern: str,
|
||||
param: parameter.Parameter,
|
||||
target: str = "both",
|
||||
) -> None:
|
||||
"""Register an additive bias for sensors matching *pattern*."""
|
||||
if target not in ("predicted", "measured", "both"):
|
||||
raise ValueError(
|
||||
f"target must be 'predicted', 'measured', or 'both', got {target!r}"
|
||||
)
|
||||
self._biases.append((pattern, param.name, target))
|
||||
|
||||
def enable_sensors(self, sensor_names: list[str]) -> None:
|
||||
"""Only include these sensors in the returned residual/timeseries."""
|
||||
self._enabled_sensors = list(sensor_names)
|
||||
|
||||
def set_sensor_weights(self, weights: Mapping[str, float]) -> None:
|
||||
"""Set per-sensor weights for the weighted diff."""
|
||||
self._sensor_weights = weights
|
||||
|
||||
# Private methods.
|
||||
|
||||
def _resolve_delays(
|
||||
self,
|
||||
sensor_names: list[str],
|
||||
params: parameter.ParameterDict,
|
||||
) -> dict[str, float]:
|
||||
"""Resolve delay patterns to concrete sensor name -> delay value (last match wins)."""
|
||||
resolved: dict[str, float] = {}
|
||||
for pattern, param_name, _ in self._delays:
|
||||
delay_value = params[param_name].value[0]
|
||||
for name in sensor_names:
|
||||
if fnmatch(name, pattern):
|
||||
resolved[name] = delay_value
|
||||
return resolved
|
||||
|
||||
def _compute_delay_bounds(self) -> tuple[float, float]:
|
||||
"""Compute min/max delay across all registered delay params (deduplicated by name)."""
|
||||
if not self._delays:
|
||||
return 0.0, 0.0
|
||||
seen: set[str] = set()
|
||||
min_vals: list[float] = []
|
||||
max_vals: list[float] = []
|
||||
for _, param_name, param in self._delays:
|
||||
if param_name in seen:
|
||||
continue
|
||||
seen.add(param_name)
|
||||
min_vals.append(float(param.min_value[0]))
|
||||
max_vals.append(float(param.max_value[0]))
|
||||
return min(min_vals), max(max_vals)
|
||||
|
||||
def _get_sensor_names(self, ts: timeseries.TimeSeries) -> list[str]:
|
||||
"""Extract sensor names from a TimeSeries signal_mapping."""
|
||||
if ts.signal_mapping is None:
|
||||
return []
|
||||
return list(ts.signal_mapping.keys())
|
||||
|
||||
def _apply_gains_biases_reference(
|
||||
self,
|
||||
ts: timeseries.TimeSeries,
|
||||
target_label: str,
|
||||
params: parameter.ParameterDict,
|
||||
) -> timeseries.TimeSeries:
|
||||
"""Reference implementation: one full copy per gain/bias application."""
|
||||
sensor_names = self._get_sensor_names(ts)
|
||||
for pattern, param_name, target in self._gains:
|
||||
if target != target_label and target != "both":
|
||||
continue
|
||||
for name in sensor_names:
|
||||
if fnmatch(name, pattern):
|
||||
ts = signal_modifier.apply_gain(ts, name, params[param_name])
|
||||
for pattern, param_name, target in self._biases:
|
||||
if target != target_label and target != "both":
|
||||
continue
|
||||
for name in sensor_names:
|
||||
if fnmatch(name, pattern):
|
||||
ts = signal_modifier.apply_bias(ts, name, params[param_name])
|
||||
return ts
|
||||
|
||||
_VERIFY_GAINS_BIASES = False
|
||||
|
||||
def _apply_gains_biases(
|
||||
self,
|
||||
ts: timeseries.TimeSeries,
|
||||
target_label: str,
|
||||
params: parameter.ParameterDict,
|
||||
) -> timeseries.TimeSeries:
|
||||
"""Apply matching gains and biases to a timeseries for the given target label."""
|
||||
sensor_names = self._get_sensor_names(ts)
|
||||
data = ts.data.copy()
|
||||
|
||||
for pattern, param_name, target in self._gains:
|
||||
if target != target_label and target != "both":
|
||||
continue
|
||||
for name in sensor_names:
|
||||
if fnmatch(name, pattern):
|
||||
indices = ts.get_indices(name)[1]
|
||||
data[..., indices] *= params[param_name].value
|
||||
|
||||
for pattern, param_name, target in self._biases:
|
||||
if target != target_label and target != "both":
|
||||
continue
|
||||
for name in sensor_names:
|
||||
if fnmatch(name, pattern):
|
||||
indices = ts.get_indices(name)[1]
|
||||
data[..., indices] += params[param_name].value
|
||||
|
||||
result = timeseries.TimeSeries(ts.times, data, ts.signal_mapping)
|
||||
|
||||
if self._VERIFY_GAINS_BIASES:
|
||||
import numpy as _np
|
||||
|
||||
ref = self._apply_gains_biases_reference(ts, target_label, params)
|
||||
_np.testing.assert_array_equal(result.data, ref.data)
|
||||
|
||||
return result
|
||||
|
||||
def apply(
|
||||
self,
|
||||
params: parameter.ParameterDict,
|
||||
sensordata_predicted: timeseries.TimeSeries,
|
||||
sensordata_measured: timeseries.TimeSeries,
|
||||
model: mujoco.MjModel,
|
||||
return_pred_all: bool,
|
||||
state: np.ndarray | None = None,
|
||||
sensor_weights: Mapping[str, float] | None = None,
|
||||
) -> tuple[np.ndarray, timeseries.TimeSeries, timeseries.TimeSeries]:
|
||||
"""Apply all registered transforms and compute the residual.
|
||||
|
||||
Signature matches :data:`ModifyResidualFn` so this method can be passed
|
||||
directly as ``modify_residual`` to :func:`model_residual`.
|
||||
|
||||
Pipeline: window measured data, resample + delay predicted data, apply
|
||||
gains/biases, weighted diff, normalise, slice to enabled sensors.
|
||||
|
||||
Returns:
|
||||
``(residual_array, predicted_ts, measured_ts)`` — the residual matrix
|
||||
and the (possibly sliced) predicted/measured TimeSeries.
|
||||
"""
|
||||
del state # Part of ModifyResidualFn signature but unused here.
|
||||
sensor_names = self._get_sensor_names(sensordata_predicted)
|
||||
|
||||
# 1. Resolve delays and compute bounds.
|
||||
sensor_delays = self._resolve_delays(sensor_names, params)
|
||||
min_delay, max_delay = self._compute_delay_bounds()
|
||||
|
||||
# 2. Window measured data.
|
||||
sensordata_measured = signal_modifier.apply_delayed_ts_window(
|
||||
sensordata_measured, sensordata_predicted, min_delay, max_delay
|
||||
)
|
||||
|
||||
# 3. Resample and delay predicted data.
|
||||
if sensor_delays:
|
||||
sensordata_predicted = signal_modifier.apply_resample_and_delay(
|
||||
sensordata_predicted,
|
||||
sensordata_measured.times,
|
||||
0.0,
|
||||
sensor_delays=sensor_delays,
|
||||
)
|
||||
else:
|
||||
sensordata_predicted = sensordata_predicted.resample(sensordata_measured.times)
|
||||
|
||||
# 4. Apply gains and biases.
|
||||
sensordata_predicted = self._apply_gains_biases(
|
||||
sensordata_predicted, "predicted", params
|
||||
)
|
||||
sensordata_measured = self._apply_gains_biases(
|
||||
sensordata_measured, "measured", params
|
||||
)
|
||||
|
||||
# 5. Weighted diff.
|
||||
weights = sensor_weights or self._sensor_weights
|
||||
res = signal_modifier.weighted_diff(
|
||||
predicted_data=sensordata_predicted.data,
|
||||
measured_data=sensordata_measured.data,
|
||||
model=model,
|
||||
sensor_weights=weights,
|
||||
)
|
||||
|
||||
# 6. Normalize.
|
||||
if self.normalize:
|
||||
res = signal_modifier.normalize_residual(res, sensordata_measured.data)
|
||||
|
||||
# 7. Slice to enabled sensors.
|
||||
if not return_pred_all and self._enabled_sensors is not None:
|
||||
indices = signal_modifier.get_sensor_indices(model, self._enabled_sensors)
|
||||
sensordata_predicted = timeseries.TimeSeries(
|
||||
sensordata_predicted.times,
|
||||
sensordata_predicted.data[:, indices],
|
||||
)
|
||||
sensordata_measured = timeseries.TimeSeries(
|
||||
sensordata_measured.times,
|
||||
sensordata_measured.data[:, indices],
|
||||
)
|
||||
res = res[:, indices]
|
||||
|
||||
return res, sensordata_predicted, sensordata_measured
|
||||
@@ -0,0 +1,700 @@
|
||||
"""Time series utilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pathlib
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import scipy.interpolate
|
||||
|
||||
|
||||
class SignalType(Enum):
|
||||
MjSensor = 0
|
||||
CustomObs = 1
|
||||
MjStateQPos = 2
|
||||
MjStateQVel = 3
|
||||
MjStateAct = 4
|
||||
MjCtrl = 5
|
||||
|
||||
|
||||
SignalMappingType: TypeAlias = dict[str, tuple[SignalType, np.ndarray]]
|
||||
|
||||
InterpolationMethod = Literal[
|
||||
"linear", "cubic", "quadratic", "quintic", "zero_order_hold", "zoh"
|
||||
]
|
||||
|
||||
|
||||
def _resolve_signals(
|
||||
model: mujoco.MjModel,
|
||||
names: Sequence[str | tuple[str, SignalType]],
|
||||
allowed_types: set[SignalType],
|
||||
) -> SignalMappingType:
|
||||
"""Resolves signal names to (canonical_name, type, indices) mappings.
|
||||
|
||||
Each name can be a string or (name, SignalType) tuple for disambiguation.
|
||||
"""
|
||||
result: SignalMappingType = {}
|
||||
idx = 0
|
||||
|
||||
for item in names:
|
||||
name, hint = item if isinstance(item, tuple) else (item, None)
|
||||
resolved = _resolve_one(model, name, hint, allowed_types)
|
||||
|
||||
if resolved is None:
|
||||
if hint is not None and hint not in allowed_types:
|
||||
raise ValueError(f"Signal '{name}' has type {hint.name} which is not allowed.")
|
||||
raise ValueError(
|
||||
f"Could not resolve signal '{item}' with allowed types {[t.name for t in allowed_types]}."
|
||||
)
|
||||
|
||||
canon_name, sig_type, width = resolved
|
||||
result[canon_name] = (sig_type, np.arange(idx, idx + width))
|
||||
idx += width
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# Suffix conventions for state/control signals
|
||||
_SUFFIXES = {
|
||||
SignalType.MjStateQPos: "_qpos",
|
||||
SignalType.MjStateQVel: "_qvel",
|
||||
SignalType.MjStateAct: "_act",
|
||||
SignalType.MjCtrl: "_ctrl",
|
||||
}
|
||||
|
||||
|
||||
def _strip_suffix(name: str, suffix: str) -> str:
|
||||
"""Strip suffix from name if present."""
|
||||
return name[: -len(suffix)] if name.endswith(suffix) else name
|
||||
|
||||
|
||||
def _resolve_one(
|
||||
model: mujoco.MjModel,
|
||||
name: str,
|
||||
hint: SignalType | None,
|
||||
allowed: set[SignalType],
|
||||
) -> tuple[str, SignalType, int] | None:
|
||||
"""Resolve a single signal name to (canonical_name, type, width)."""
|
||||
|
||||
# 1. Sensor
|
||||
if _type_allowed(hint, SignalType.MjSensor, allowed):
|
||||
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name)
|
||||
if sid >= 0:
|
||||
return (name, SignalType.MjSensor, model.sensor_dim[sid])
|
||||
|
||||
# 2. Control
|
||||
if _type_allowed(hint, SignalType.MjCtrl, allowed):
|
||||
base = _strip_suffix(name, "_ctrl")
|
||||
aid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, base)
|
||||
if aid >= 0:
|
||||
return (base + "_ctrl", SignalType.MjCtrl, 1)
|
||||
|
||||
# 3. State (qpos/qvel)
|
||||
for sig_type in (SignalType.MjStateQPos, SignalType.MjStateQVel):
|
||||
if _type_allowed(hint, sig_type, allowed):
|
||||
base = _strip_suffix(name, _SUFFIXES[sig_type])
|
||||
width = _joint_or_body_width(model, base, sig_type)
|
||||
if width > 0:
|
||||
return (base + _SUFFIXES[sig_type], sig_type, width)
|
||||
|
||||
# 4. Actuator state (act)
|
||||
if _type_allowed(hint, SignalType.MjStateAct, allowed):
|
||||
base = _strip_suffix(name, "_act")
|
||||
aid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, base)
|
||||
if aid >= 0 and model.actuator_actnum[aid] > 0:
|
||||
return (base + "_act", SignalType.MjStateAct, model.actuator_actnum[aid])
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _type_allowed(
|
||||
hint: SignalType | None, target: SignalType, allowed: set[SignalType]
|
||||
) -> bool:
|
||||
"""Check if target type is allowed given hint and allowed set."""
|
||||
return (hint is None or hint == target) and target in allowed
|
||||
|
||||
|
||||
def _joint_or_body_width(model: mujoco.MjModel, name: str, sig_type: SignalType) -> int:
|
||||
"""Get state width for a joint or free body."""
|
||||
# Joint widths by type
|
||||
QPOS_WIDTHS = {mujoco.mjtJoint.mjJNT_FREE: 7, mujoco.mjtJoint.mjJNT_BALL: 4}
|
||||
QVEL_WIDTHS = {mujoco.mjtJoint.mjJNT_FREE: 6, mujoco.mjtJoint.mjJNT_BALL: 3}
|
||||
|
||||
jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name)
|
||||
if jid >= 0:
|
||||
jtype = model.jnt_type[jid]
|
||||
widths = QPOS_WIDTHS if sig_type == SignalType.MjStateQPos else QVEL_WIDTHS
|
||||
return widths.get(jtype, 1)
|
||||
|
||||
# Free body
|
||||
bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name)
|
||||
if bid >= 0 and model.body_dofnum[bid] == 6:
|
||||
return 7 if sig_type == SignalType.MjStateQPos else 6
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TimeSeries:
|
||||
"""A utility class for working with time-series data.
|
||||
|
||||
Attributes:
|
||||
times: 1D array of timestamps.
|
||||
data: Array of signal data. The first axis corresponds to time.
|
||||
signal_mapping: Dict of tuples that maps the signal type and its
|
||||
signal fields in data
|
||||
"""
|
||||
|
||||
times: np.ndarray
|
||||
data: np.ndarray
|
||||
signal_mapping: SignalMappingType | None = None
|
||||
|
||||
@staticmethod
|
||||
def compute_all_sensor_mapping(model: mujoco.MjModel) -> SignalMappingType:
|
||||
"""Computes mapping for all sensors in the model."""
|
||||
signal_mapping = {}
|
||||
for sensor_id in range(model.nsensor):
|
||||
addr = model.sensor_adr[sensor_id]
|
||||
dim = model.sensor_dim[sensor_id]
|
||||
name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_SENSOR, sensor_id)
|
||||
indices = np.arange(addr, addr + dim)
|
||||
signal_mapping[name] = (SignalType.MjSensor, indices)
|
||||
return signal_mapping
|
||||
|
||||
@staticmethod
|
||||
def compute_all_control_mapping(model: mujoco.MjModel) -> SignalMappingType:
|
||||
"""Computes mapping for all controls (actuators) in the model."""
|
||||
ctrl_map: SignalMappingType = {}
|
||||
for act_id in range(model.nu):
|
||||
act_name = model.actuator(act_id).name
|
||||
ctrl_indices = np.arange(act_id, act_id + 1)
|
||||
ctrl_map[f"{act_name}_ctrl"] = (SignalType.MjCtrl, ctrl_indices)
|
||||
return ctrl_map
|
||||
|
||||
@staticmethod
|
||||
def compute_all_state_mappings(
|
||||
model: mujoco.MjModel,
|
||||
) -> tuple[
|
||||
SignalMappingType, SignalMappingType, SignalMappingType, SignalMappingType
|
||||
]:
|
||||
"""Computes mappings for all state components (qpos, qvel, act) + ctrl."""
|
||||
qpos_map: SignalMappingType = {}
|
||||
qvel_map: SignalMappingType = {}
|
||||
act_map: SignalMappingType = {}
|
||||
|
||||
nq = model.nq
|
||||
nv = model.nv
|
||||
|
||||
# Bodies
|
||||
for body_id in range(model.nbody):
|
||||
b = model.body(body_id)
|
||||
body_name = b.name
|
||||
start_index = model.body_dofadr[body_id]
|
||||
|
||||
if start_index >= 0 and b.dofnum[0] == 6:
|
||||
qpos_indices = np.arange(start_index, start_index + 7)
|
||||
qpos_map[f"{body_name}_qpos"] = (SignalType.MjStateQPos, qpos_indices)
|
||||
qvel_indices = np.arange(start_index + nq, start_index + nq + 6)
|
||||
qvel_map[f"{body_name}_qvel"] = (SignalType.MjStateQVel, qvel_indices)
|
||||
|
||||
# Joints
|
||||
for jnt_id in range(model.njnt):
|
||||
jnt_name = model.joint(jnt_id).name
|
||||
start_index = model.jnt_qposadr[jnt_id]
|
||||
jnt_type = model.jnt_type[jnt_id]
|
||||
|
||||
qpos_width = 1
|
||||
qvel_width = 1
|
||||
if jnt_type == mujoco.mjtJoint.mjJNT_BALL:
|
||||
qpos_width = 4
|
||||
qvel_width = 3
|
||||
elif jnt_type == mujoco.mjtJoint.mjJNT_FREE:
|
||||
continue
|
||||
|
||||
qpos_indices = np.arange(start_index, start_index + qpos_width)
|
||||
qpos_map[f"{jnt_name}_qpos"] = (SignalType.MjStateQPos, qpos_indices)
|
||||
qvel_indices = np.arange(start_index + nq, start_index + nq + qvel_width)
|
||||
qvel_map[f"{jnt_name}_qvel"] = (SignalType.MjStateQVel, qvel_indices)
|
||||
|
||||
# Actuators
|
||||
for act_id in range(model.nu):
|
||||
act_name = model.actuator(act_id).name
|
||||
start_index = model.actuator_actadr[act_id]
|
||||
num_vals = model.actuator_actnum[act_id]
|
||||
# if index is -1, the actuator is stateless.
|
||||
if start_index != -1:
|
||||
indices = np.arange(start_index + nq + nv, start_index + nq + nv + num_vals)
|
||||
act_map[f"{act_name}_act"] = (SignalType.MjStateAct, indices)
|
||||
|
||||
ctrl_map = TimeSeries.compute_all_control_mapping(model)
|
||||
|
||||
return qpos_map, qvel_map, act_map, ctrl_map
|
||||
|
||||
@classmethod
|
||||
def from_custom_map(
|
||||
cls,
|
||||
times: np.ndarray,
|
||||
data: np.ndarray,
|
||||
signals: Sequence[str | tuple[str, int, SignalType]],
|
||||
) -> TimeSeries:
|
||||
"""Construct a TimeSeries from custom data with explicit signal definitions.
|
||||
|
||||
Use this when you have custom signal types (e.g., from a custom modify_residual
|
||||
function) that are not auto-resolved from a MuJoCo model. You must explicitly
|
||||
specify the signal names, widths, and types.
|
||||
|
||||
Args:
|
||||
times: 1-D timestamp array of length N.
|
||||
data: 2-D array of shape ``(N, D)``.
|
||||
signals: Defines the layout of the columns in `data`.
|
||||
- If a list of strings: Each string is a signal name with width 1 and type ``CustomObs``.
|
||||
- If a list of tuples: Each tuple is ``(name, width, type)``.
|
||||
|
||||
Returns:
|
||||
A TimeSeries object with the constructed signal mapping.
|
||||
"""
|
||||
if data.ndim != 2:
|
||||
raise ValueError("The 'data' array must be 2-dimensional (Time x Features).")
|
||||
|
||||
signal_mapping_dict: SignalMappingType = {}
|
||||
current_index = 0
|
||||
total_width = 0
|
||||
|
||||
for item in signals:
|
||||
if isinstance(item, str):
|
||||
name = item
|
||||
width = 1
|
||||
sig_type = SignalType.CustomObs
|
||||
else:
|
||||
name, width, sig_type = item
|
||||
|
||||
if width <= 0:
|
||||
raise ValueError(f"Signal '{name}' must have positive width, got {width}.")
|
||||
|
||||
indices = np.arange(current_index, current_index + width)
|
||||
signal_mapping_dict[name] = (sig_type, indices)
|
||||
current_index += width
|
||||
total_width += width
|
||||
|
||||
if total_width != data.shape[1]:
|
||||
raise ValueError(
|
||||
f"Total width of signals ({total_width}) does not match "
|
||||
f"data columns ({data.shape[1]})."
|
||||
)
|
||||
|
||||
return cls(times=times, data=data, signal_mapping=signal_mapping_dict)
|
||||
|
||||
@classmethod
|
||||
def from_names(
|
||||
cls,
|
||||
times: np.ndarray,
|
||||
data: np.ndarray,
|
||||
model: mujoco.MjModel,
|
||||
names: Sequence[str | tuple[str, SignalType]] | None = None,
|
||||
) -> TimeSeries:
|
||||
"""Construct a TimeSeries for observations (sensors or state) from the model.
|
||||
|
||||
This method automatically resolves signal names (sensors, qpos, qvel, act) from
|
||||
the MuJoCo model, determining their types and data layout. Use this for standard
|
||||
observation signals that are defined in the model.
|
||||
|
||||
Args:
|
||||
times: 1-D timestamps of length N.
|
||||
data: 2-D array of shape (N, D).
|
||||
model: MuJoCo model used to auto-resolve signal names and types.
|
||||
names: Signal names to map. Can be strings or (name, SignalType) tuples.
|
||||
If None, maps ALL model sensors in sensor address order.
|
||||
|
||||
Warning:
|
||||
When names=None, data columns MUST match the model's sensor layout
|
||||
(i.e., data[:, i] corresponds to model.sensordata[i] during simulation).
|
||||
If your data is in a different order, pass explicit names.
|
||||
|
||||
Raises:
|
||||
ValueError: If MjCtrl signals are passed (use from_control_names).
|
||||
"""
|
||||
if data.ndim != 2:
|
||||
raise ValueError("The 'data' array must be 2-dimensional (Time x Features).")
|
||||
|
||||
if names is None:
|
||||
signal_mapping = cls.compute_all_sensor_mapping(model)
|
||||
# Verify width: assumes data contains ALL sensors in sensor_adr order
|
||||
if model.nsensordata != data.shape[1]:
|
||||
raise ValueError(
|
||||
f"Data columns ({data.shape[1]}) do not match model sensors dim ({model.nsensordata})."
|
||||
)
|
||||
else:
|
||||
signal_mapping = _resolve_signals(
|
||||
model,
|
||||
names,
|
||||
allowed_types={
|
||||
SignalType.MjSensor,
|
||||
SignalType.MjStateQPos,
|
||||
SignalType.MjStateQVel,
|
||||
SignalType.MjStateAct,
|
||||
},
|
||||
)
|
||||
# Verify total resolved width
|
||||
max_idx = 0
|
||||
for _, indices in signal_mapping.values():
|
||||
if len(indices) > 0:
|
||||
max_idx = max(max_idx, indices[-1] + 1)
|
||||
if max_idx != data.shape[1]:
|
||||
raise ValueError(
|
||||
f"Resolved signal width ({max_idx}) does not match data columns ({data.shape[1]})."
|
||||
)
|
||||
|
||||
return cls(times=times, data=data, signal_mapping=signal_mapping)
|
||||
|
||||
@classmethod
|
||||
def from_control_names(
|
||||
cls,
|
||||
times: np.ndarray,
|
||||
data: np.ndarray,
|
||||
model: mujoco.MjModel,
|
||||
names: Sequence[str | tuple[str, SignalType]] | None = None,
|
||||
) -> TimeSeries:
|
||||
"""Construct a TimeSeries for control signals from the model.
|
||||
|
||||
This method automatically resolves control/actuator names from the MuJoCo model,
|
||||
determining their layout. Use this for control signals (MjCtrl type).
|
||||
|
||||
Args:
|
||||
times: 1-D timestamps of length N.
|
||||
data: 2-D array of shape (N, model.nu).
|
||||
model: MuJoCo model used to auto-resolve actuator names.
|
||||
names: Actuator names to map. If None, maps ALL actuators in order.
|
||||
|
||||
Warning:
|
||||
When names=None, data columns MUST match actuator order in the model
|
||||
(i.e., data[:, i] corresponds to actuator i). Pass explicit names
|
||||
if your data is in a different order.
|
||||
"""
|
||||
if data.ndim != 2:
|
||||
raise ValueError("The 'data' array must be 2-dimensional (Time x Features).")
|
||||
|
||||
if names is None:
|
||||
signal_mapping = cls.compute_all_control_mapping(model)
|
||||
if model.nu != data.shape[1]:
|
||||
raise ValueError(
|
||||
f"Data columns ({data.shape[1]}) do not match model controls ({model.nu})."
|
||||
)
|
||||
else:
|
||||
signal_mapping = _resolve_signals(model, names, allowed_types={SignalType.MjCtrl})
|
||||
max_idx = 0
|
||||
for _, indices in signal_mapping.values():
|
||||
if len(indices) > 0:
|
||||
max_idx = max(max_idx, indices[-1] + 1)
|
||||
if max_idx != data.shape[1]:
|
||||
raise ValueError(
|
||||
f"Resolved signal width ({max_idx}) does not match data columns ({data.shape[1]})."
|
||||
)
|
||||
|
||||
return cls(times=times, data=data, signal_mapping=signal_mapping)
|
||||
|
||||
def get_indices(self, obs_name: str) -> tuple[SignalType, np.ndarray]:
|
||||
"""Look up the signal type and column indices for a named observation."""
|
||||
assert self.signal_mapping is not None
|
||||
if obs_name not in self.signal_mapping:
|
||||
raise ValueError(f"{obs_name} observation is not in the observation name map.")
|
||||
return self.signal_mapping[obs_name]
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
times: np.ndarray,
|
||||
data: np.ndarray,
|
||||
signal_mapping: dict[str, tuple[SignalType, np.ndarray | list | int]],
|
||||
) -> TimeSeries:
|
||||
"""Construct a TimeSeries, normalising index entries to ``np.ndarray``."""
|
||||
normalized: SignalMappingType = {}
|
||||
for key in signal_mapping:
|
||||
signal_type, indices = signal_mapping[key]
|
||||
normalized[key] = (signal_type, np.atleast_1d(indices))
|
||||
|
||||
return cls(times, data, normalized)
|
||||
|
||||
@classmethod
|
||||
def slice_by_name(cls, ts: TimeSeries, enabled_sensors: list[str]) -> TimeSeries:
|
||||
"""Return a new TimeSeries containing only the named signals.
|
||||
|
||||
Columns are re-indexed so the resulting ``signal_mapping`` has contiguous
|
||||
indices starting from 0.
|
||||
"""
|
||||
if not ts.signal_mapping:
|
||||
return ts
|
||||
|
||||
original_indices_to_keep = []
|
||||
original_to_new_index_map = {}
|
||||
all_original_indices = []
|
||||
for name in ts.signal_mapping:
|
||||
all_original_indices.extend(ts.signal_mapping[name][1])
|
||||
|
||||
# Build a set of indices to keep for quick lookups
|
||||
kept_indices_set = set()
|
||||
for name in enabled_sensors:
|
||||
if name not in ts.signal_mapping:
|
||||
raise ValueError(
|
||||
f"Attemping to slice TimeSeries failed. {name} is not in {ts.signal_mapping}."
|
||||
)
|
||||
kept_indices_set.update(ts.signal_mapping[name][1])
|
||||
new_index_counter = 0
|
||||
for original_index in all_original_indices:
|
||||
if original_index in kept_indices_set:
|
||||
original_to_new_index_map[original_index] = new_index_counter
|
||||
new_index_counter += 1
|
||||
original_indices_to_keep.append(original_index)
|
||||
|
||||
data = ts.data[..., original_indices_to_keep]
|
||||
|
||||
trimmed_signal_mapping = {}
|
||||
for name in enabled_sensors:
|
||||
metadata, original_indices = ts.signal_mapping[name]
|
||||
|
||||
new_indices = []
|
||||
for original_index in original_indices:
|
||||
new_indices.append(original_to_new_index_map[original_index])
|
||||
|
||||
trimmed_signal_mapping[name] = (metadata, np.asarray(new_indices))
|
||||
|
||||
return cls(ts.times, data, trimmed_signal_mapping)
|
||||
|
||||
def __post_init__(self):
|
||||
"""Validate the time series data after initialization.
|
||||
|
||||
Raises:
|
||||
ValueError: If times is not 1D, if lengths don't match, if times
|
||||
is not strictly increasing, or if arrays are empty.
|
||||
"""
|
||||
if self.times.size == 0:
|
||||
raise ValueError("Empty arrays are not allowed in TimeSeries")
|
||||
if self.times.ndim != 1:
|
||||
raise ValueError(f"times must be a 1D array, got {self.times.ndim}D array")
|
||||
if len(self.times) != len(self.data):
|
||||
raise ValueError(
|
||||
f"Length of times ({len(self.times)}) and data ({len(self.data)}) must match"
|
||||
)
|
||||
if not np.all(np.diff(self.times) > 0):
|
||||
raise ValueError("times must be strictly increasing")
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.data)
|
||||
|
||||
def save_to_disk(self, path: str | pathlib.Path) -> None:
|
||||
"""Save the time series data to disk.
|
||||
|
||||
Args:
|
||||
path: Path where the data will be saved.
|
||||
"""
|
||||
np.savez(
|
||||
path,
|
||||
times=self.times,
|
||||
data=self.data,
|
||||
signal_mapping=np.array(self.signal_mapping, dtype=object),
|
||||
)
|
||||
|
||||
def save_to_csv(self, path: str | pathlib.Path) -> None:
|
||||
"""Save the time series data to a CSV file."""
|
||||
np.savetxt(
|
||||
path,
|
||||
np.concatenate([self.times[:, None], self.data], axis=1),
|
||||
delimiter=",",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def load_from_disk(cls, path: str | pathlib.Path) -> TimeSeries:
|
||||
"""Load time series data from disk.
|
||||
|
||||
Args:
|
||||
path: Path to the saved data.
|
||||
|
||||
Returns:
|
||||
A new TimeSeries object.
|
||||
"""
|
||||
with np.load(path, allow_pickle=True) as npz:
|
||||
times = npz["times"]
|
||||
data = npz["data"]
|
||||
if "signal_mapping" in npz:
|
||||
signal_mapping = npz["signal_mapping"].item()
|
||||
else:
|
||||
signal_mapping = None
|
||||
|
||||
return cls(times=times, data=data, signal_mapping=signal_mapping)
|
||||
|
||||
def interpolate(
|
||||
self, t: float | np.ndarray, method: InterpolationMethod = "linear"
|
||||
) -> np.ndarray:
|
||||
"""Interpolate data at specified time(s).
|
||||
|
||||
This is the core interpolation function used by both get() and resample().
|
||||
|
||||
Args:
|
||||
t: Time point(s) at which to interpolate data.
|
||||
method: Interpolation method to use.
|
||||
|
||||
Returns:
|
||||
Interpolated data values.
|
||||
"""
|
||||
t = np.atleast_1d(np.asarray(t))
|
||||
|
||||
if method in ("zero_order_hold", "zoh"):
|
||||
indices = np.searchsorted(self.times, t, side="right") - 1
|
||||
indices = np.clip(indices, 0, len(self.times) - 1)
|
||||
return self.data[indices]
|
||||
|
||||
return scipy.interpolate.interp1d(
|
||||
self.times,
|
||||
self.data,
|
||||
kind=method,
|
||||
axis=0,
|
||||
bounds_error=False,
|
||||
fill_value=(self.data[0], self.data[-1]), # pyright: ignore[reportArgumentType]
|
||||
assume_sorted=True,
|
||||
)(t)
|
||||
|
||||
def get(
|
||||
self, t: float | np.ndarray, method: InterpolationMethod = "linear"
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Get interpolated data at specified time(s).
|
||||
|
||||
This method is useful for querying data at specific timestamps without
|
||||
creating a new TimeSeries object.
|
||||
|
||||
Args:
|
||||
t: Time point(s) at which to get data.
|
||||
method: Interpolation method to use.
|
||||
|
||||
Returns:
|
||||
Tuple of (times, interpolated_data).
|
||||
"""
|
||||
t_orig = np.asarray(t)
|
||||
t_shape = t_orig.shape
|
||||
result = self.interpolate(t_orig, method=method)
|
||||
if t_shape == ():
|
||||
result = result.squeeze(axis=0)
|
||||
return t_orig, result
|
||||
|
||||
def resample(
|
||||
self,
|
||||
new_times: np.ndarray | None = None,
|
||||
target_dt: float | None = None,
|
||||
method: InterpolationMethod = "linear",
|
||||
) -> TimeSeries:
|
||||
"""Resample the time series to new timestamps or a specific time interval.
|
||||
|
||||
This method creates a new TimeSeries object with data interpolated at the
|
||||
specified timestamps.
|
||||
|
||||
Args:
|
||||
new_times: Optional array of new timestamps. If provided, target_dt is
|
||||
ignored.
|
||||
target_dt: Optional time interval for regular resampling. Only used if
|
||||
new_times is None.
|
||||
method: Interpolation method to use.
|
||||
|
||||
Returns:
|
||||
A new TimeSeries object with resampled data.
|
||||
|
||||
Raises:
|
||||
ValueError: If neither new_times nor target_dt is provided, or if
|
||||
new_times is not strictly increasing.
|
||||
"""
|
||||
# Generate new times if target_dt is provided.
|
||||
if new_times is None:
|
||||
if target_dt is None:
|
||||
raise ValueError("Either new_times or target_dt must be provided")
|
||||
if target_dt <= 0:
|
||||
raise ValueError("target_dt must be a positive float")
|
||||
|
||||
# Create evenly spaced timestamps.
|
||||
new_nsteps = int(np.ceil((self.times[-1] - self.times[0]) / target_dt)) + 1
|
||||
new_times = np.linspace(self.times[0], self.times[-1], new_nsteps, endpoint=True)
|
||||
else:
|
||||
# Make sure new_times is valid.
|
||||
if new_times.ndim != 1:
|
||||
raise ValueError("new_times must be a 1D array")
|
||||
if not np.all(np.diff(new_times) > 0):
|
||||
raise ValueError("new_times must be strictly increasing")
|
||||
|
||||
assert new_times is not None
|
||||
new_data = self.interpolate(new_times, method=method)
|
||||
return TimeSeries(
|
||||
times=new_times, data=new_data, signal_mapping=self.signal_mapping
|
||||
)
|
||||
|
||||
def remove_from_beginning(self, time_to_remove_s: float) -> TimeSeries:
|
||||
"""Remove time from the beginning of the time series.
|
||||
|
||||
Args:
|
||||
time_to_remove_s: Time to remove from the beginning of the time series.
|
||||
|
||||
Returns:
|
||||
A new TimeSeries object with the specified time removed.
|
||||
"""
|
||||
if time_to_remove_s < 0:
|
||||
raise ValueError("time_to_remove_s must be non-negative")
|
||||
if time_to_remove_s > self.times[-1]:
|
||||
raise ValueError(
|
||||
"time_to_remove_s is greater than the duration of the time series"
|
||||
)
|
||||
idx = np.searchsorted(self.times, time_to_remove_s)
|
||||
times_shifted = self.times[idx:] - self.times[idx]
|
||||
return TimeSeries(
|
||||
times=times_shifted, data=self.data[idx:], signal_mapping=self.signal_mapping
|
||||
)
|
||||
|
||||
def dt_statistics(self) -> dict[str, float]:
|
||||
"""Calculate statistics about the time intervals.
|
||||
|
||||
Returns:
|
||||
Dictionary with mean, median, std, min, and max of time intervals.
|
||||
|
||||
Raises:
|
||||
ValueError: If there are fewer than two timestamps.
|
||||
"""
|
||||
if self.times.size < 2:
|
||||
raise ValueError("Must have at least two timestamps to compute dt statistics.")
|
||||
dt_values = np.diff(self.times)
|
||||
stats = {}
|
||||
for fn in ["mean", "median", "std", "min", "max"]:
|
||||
stats[fn] = float(getattr(np, fn)(dt_values))
|
||||
return stats
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""Return a string representation of the TimeSeries object."""
|
||||
t_start, t_end = self.times[0], self.times[-1]
|
||||
duration = t_end - t_start
|
||||
|
||||
data_shape = self.data.shape
|
||||
n_samples = len(self)
|
||||
|
||||
dt_stats = self.dt_statistics()
|
||||
mean_dt = dt_stats["mean"]
|
||||
min_dt = dt_stats["min"]
|
||||
max_dt = dt_stats["max"]
|
||||
|
||||
is_uniform = dt_stats["std"] / mean_dt < 0.01 # Less than 1% variation.
|
||||
|
||||
# Calculate data range (min/max values).
|
||||
data_min = np.min(self.data)
|
||||
data_max = np.max(self.data)
|
||||
data_range = f"[{data_min:.3g}, {data_max:.3g}]"
|
||||
|
||||
parts = [
|
||||
"TimeSeries(",
|
||||
f" samples={n_samples}",
|
||||
f" shape={data_shape}",
|
||||
f" time_range=[{t_start:.3g}, {t_end:.3g}] (duration={duration:.3g})",
|
||||
f" dt={mean_dt:.3g}"
|
||||
+ (" (uniform)" if is_uniform else f" (min={min_dt:.3g}, max={max_dt:.3g})"),
|
||||
f" data_range={data_range}",
|
||||
f" signal_mapping={self.signal_mapping}",
|
||||
")",
|
||||
]
|
||||
|
||||
return "\n".join(parts)
|
||||
@@ -0,0 +1,501 @@
|
||||
"""Trajectory data containers for system identification."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import pathlib
|
||||
from collections.abc import Sequence
|
||||
|
||||
import mujoco
|
||||
import mujoco.rollout as mj_rollout
|
||||
import numpy as np
|
||||
from absl import logging
|
||||
|
||||
from mujoco.sysid._src import timeseries
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SystemTrajectory:
|
||||
"""Encapsulates a trajectory rolled out from a system.
|
||||
|
||||
Attributes:
|
||||
model: MuJoCo model used to simulate the trajectory.
|
||||
control: A TimeSeries instance containing control signals.
|
||||
sensordata: A TimeSeries instance containing sensor data.
|
||||
initial_state: Initial state of the simulation. Shape (n_state,).
|
||||
state: Simulation states over time. Shape (n_steps, n_state). Optional for
|
||||
real robot trajectories.
|
||||
"""
|
||||
|
||||
model: mujoco.MjModel
|
||||
control: timeseries.TimeSeries
|
||||
sensordata: timeseries.TimeSeries
|
||||
initial_state: np.ndarray
|
||||
state: timeseries.TimeSeries | None
|
||||
|
||||
def replace(self, **kwargs) -> SystemTrajectory:
|
||||
"""Return a copy with the specified fields replaced."""
|
||||
return dataclasses.replace(self, **kwargs)
|
||||
|
||||
def get_sensordata_slice(self, sensor: str = "joint_pos") -> np.ndarray:
|
||||
"""Extract contiguous sensor columns by type.
|
||||
|
||||
Args:
|
||||
sensor: One of ``"joint_pos"``, ``"joint_vel"``, or ``"joint_torque"``.
|
||||
|
||||
Returns:
|
||||
2-D array of shape ``(n_steps, total_sensor_dim)``.
|
||||
"""
|
||||
if sensor == "joint_pos":
|
||||
sensor_type = mujoco.mjtSensor.mjSENS_JOINTPOS
|
||||
elif sensor == "joint_vel":
|
||||
sensor_type = mujoco.mjtSensor.mjSENS_JOINTVEL
|
||||
elif sensor == "joint_torque":
|
||||
sensor_type = mujoco.mjtSensor.mjSENS_JOINTACTFRC
|
||||
else:
|
||||
raise ValueError(f"Unsupported sensor type: {sensor}")
|
||||
adr = []
|
||||
dims = []
|
||||
for i in range(self.model.nsensor):
|
||||
if self.model.sensor(i).type == sensor_type:
|
||||
sensor_id = self.model.sensor(i).id
|
||||
adr.append(self.model.sensor_adr[sensor_id])
|
||||
dims.append(self.model.sensor_dim[sensor_id])
|
||||
sensors = sorted(zip(adr, dims, strict=True), key=lambda x: x[0])
|
||||
start = sensors[0][0]
|
||||
total_dim = sum(d for _, d in sensors)
|
||||
end = start + total_dim
|
||||
return self.sensordata.data[:, start:end]
|
||||
|
||||
@property
|
||||
def sensordim(self) -> int:
|
||||
"""Total number of scalar sensor outputs in the model."""
|
||||
return self.model.nsensordata
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Number of time steps in the trajectory."""
|
||||
return len(self.sensordata)
|
||||
|
||||
def save_to_disk(self, path: pathlib.Path) -> None:
|
||||
save_dict = {
|
||||
"control_times": self.control.times,
|
||||
"control_data": self.control.data,
|
||||
"sensordata_times": self.sensordata.times,
|
||||
"sensordata_data": self.sensordata.data,
|
||||
"initial_state": self.initial_state,
|
||||
}
|
||||
if self.state is not None:
|
||||
save_dict["state_times"] = self.state.times
|
||||
save_dict["state_data"] = self.state.data
|
||||
save_dict["state_signal_mapping"] = np.array(
|
||||
self.state.signal_mapping, dtype=object
|
||||
)
|
||||
|
||||
if self.control.signal_mapping:
|
||||
save_dict["control_signal_mapping"] = np.array(
|
||||
self.control.signal_mapping, dtype=object
|
||||
)
|
||||
|
||||
if self.sensordata.signal_mapping:
|
||||
save_dict["sensordata_signal_mapping"] = np.array(
|
||||
self.sensordata.signal_mapping, dtype=object
|
||||
)
|
||||
|
||||
np.savez(path, **save_dict) # type: ignore
|
||||
|
||||
@classmethod
|
||||
def load_from_disk(
|
||||
cls,
|
||||
path: pathlib.Path,
|
||||
model: mujoco.MjModel,
|
||||
allow_missing_sensors: bool = False,
|
||||
) -> SystemTrajectory:
|
||||
with np.load(path, allow_pickle=True) as npz:
|
||||
control_times = npz["control_times"]
|
||||
control_data = npz["control_data"]
|
||||
sensordata_times = npz["sensordata_times"]
|
||||
sensordata_data = npz["sensordata_data"]
|
||||
initial_state = npz["initial_state"]
|
||||
state_times = npz.get("state_times", None)
|
||||
state_data = npz.get("state_data", None)
|
||||
|
||||
control_signal_mapping = None
|
||||
if "control_signal_mapping" in npz:
|
||||
control_signal_mapping = npz["control_signal_mapping"].item()
|
||||
|
||||
sensordata_signal_mapping = None
|
||||
if "sensordata_signal_mapping" in npz:
|
||||
sensordata_signal_mapping = npz["sensordata_signal_mapping"].item()
|
||||
|
||||
state_signal_mapping = None
|
||||
if "state_signal_mapping" in npz:
|
||||
state_signal_mapping = npz["state_signal_mapping"].item()
|
||||
|
||||
predicted_rollout = cls(
|
||||
model=model,
|
||||
control=timeseries.TimeSeries(
|
||||
control_times, control_data, signal_mapping=control_signal_mapping
|
||||
),
|
||||
sensordata=timeseries.TimeSeries(
|
||||
sensordata_times, sensordata_data, signal_mapping=sensordata_signal_mapping
|
||||
),
|
||||
initial_state=initial_state,
|
||||
state=timeseries.TimeSeries(
|
||||
state_times, state_data, signal_mapping=state_signal_mapping
|
||||
)
|
||||
if state_times is not None
|
||||
else None,
|
||||
)
|
||||
predicted_rollout.check_compatible(allow_missing_sensors)
|
||||
return predicted_rollout
|
||||
|
||||
def check_compatible(self, allow_missing_sensors: bool = False) -> None:
|
||||
"""Validate that data dimensions match the model.
|
||||
|
||||
Checks sensor, control, state, and initial-state dimensions.
|
||||
|
||||
Args:
|
||||
allow_missing_sensors: If True, a sensor dimension mismatch is logged
|
||||
as a warning instead of raising.
|
||||
"""
|
||||
if self.sensordata.data.shape[1] != self.model.nsensordata:
|
||||
if not allow_missing_sensors:
|
||||
raise ValueError(
|
||||
f"Sensor data dimension {self.sensordata.data.shape[1]} does not"
|
||||
f" match model sensor dimension {self.model.nsensordata}"
|
||||
)
|
||||
else:
|
||||
logging.warning(
|
||||
f"Sensor data dimension {self.sensordata.data.shape[1]} does not"
|
||||
f" match model sensor dimension {self.model.nsensordata}"
|
||||
)
|
||||
|
||||
if self.control.data.shape[1] != self.model.nu:
|
||||
raise ValueError(
|
||||
f"Control data dimension {self.control.data.shape[1]} does not"
|
||||
f" match model control dimension {self.model.nu}"
|
||||
)
|
||||
|
||||
state_spec = mujoco.mjtState.mjSTATE_FULLPHYSICS
|
||||
state_size = mujoco.mj_stateSize(self.model, state_spec.value)
|
||||
if self.state is not None:
|
||||
if self.state.data.shape[1] != state_size:
|
||||
raise ValueError(
|
||||
f"State dimension {self.state.data.shape[1]} does not match "
|
||||
f"model state dimension {state_size}"
|
||||
)
|
||||
if self.initial_state.shape[0] != state_size:
|
||||
raise ValueError(
|
||||
f"Initial state dimension {self.initial_state.shape[0]} does not"
|
||||
f" match model state dimension {state_size}"
|
||||
)
|
||||
|
||||
def split(self, chunk_size: int) -> list[SystemTrajectory]:
|
||||
"""Split into consecutive non-overlapping chunks of *chunk_size* steps.
|
||||
|
||||
Incomplete trailing steps are discarded. Requires ``state`` to be set
|
||||
(needed to extract the initial state for each chunk).
|
||||
"""
|
||||
if self.state is None:
|
||||
raise ValueError("Cannot split rollout with missing state field.")
|
||||
steps = len(self.sensordata.times)
|
||||
n_complete_chunks = steps // chunk_size
|
||||
control_times = self.control.times
|
||||
control_data = self.control.data
|
||||
sensordata_times = self.sensordata.times
|
||||
sensordata_data = self.sensordata.data
|
||||
trajectories = []
|
||||
for i in range(n_complete_chunks):
|
||||
start_idx = i * chunk_size
|
||||
end_idx = start_idx + chunk_size
|
||||
initial_state = (
|
||||
self.initial_state if start_idx == 0 else self.state.data[start_idx - 1]
|
||||
)
|
||||
control_times_chunk = control_times[start_idx:end_idx]
|
||||
control_data_chunk = control_data[start_idx:end_idx]
|
||||
sensordata_times_chunk = sensordata_times[start_idx:end_idx]
|
||||
sensordata_data_chunk = sensordata_data[start_idx:end_idx]
|
||||
trajectories.append(
|
||||
SystemTrajectory(
|
||||
model=self.model,
|
||||
control=timeseries.TimeSeries(control_times_chunk, control_data_chunk),
|
||||
sensordata=timeseries.TimeSeries(
|
||||
sensordata_times_chunk, sensordata_data_chunk
|
||||
),
|
||||
initial_state=initial_state,
|
||||
state=timeseries.TimeSeries(
|
||||
times=self.state.times[start_idx:end_idx],
|
||||
data=self.state.data[start_idx:end_idx],
|
||||
signal_mapping=self.state.signal_mapping,
|
||||
),
|
||||
)
|
||||
)
|
||||
return trajectories
|
||||
|
||||
def render(
|
||||
self,
|
||||
height: int = 240,
|
||||
width: int = 320,
|
||||
camera: str | int = -1,
|
||||
fps: int = 30,
|
||||
) -> list[np.ndarray]:
|
||||
"""Render this trajectory to a list of RGB frames.
|
||||
|
||||
Requires ``state`` to be set. Delegates to
|
||||
:func:`~mujoco_sysid._src.plotting.render_rollout`.
|
||||
"""
|
||||
if self.state is None:
|
||||
raise ValueError("Cannot render rollout with missing state field.")
|
||||
|
||||
from mujoco.sysid._src.plotting import render_rollout
|
||||
|
||||
# Adapt state to batch format (nbatch=1, nsteps, nstate)
|
||||
state_batch = self.state.data[np.newaxis, :, :]
|
||||
|
||||
data = mujoco.MjData(self.model)
|
||||
|
||||
return render_rollout(
|
||||
model=self.model,
|
||||
data=data,
|
||||
state=state_batch,
|
||||
framerate=fps,
|
||||
camera=camera,
|
||||
width=width,
|
||||
height=height,
|
||||
)
|
||||
|
||||
|
||||
def create_initial_state(
|
||||
model: mujoco.MjModel,
|
||||
qpos: np.ndarray,
|
||||
qvel: np.ndarray | None = None,
|
||||
act: np.ndarray | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Build a ``mjSTATE_FULLPHYSICS`` initial-state vector from components.
|
||||
|
||||
Args:
|
||||
model: MuJoCo model.
|
||||
qpos: Joint positions, shape ``(nq,)``.
|
||||
qvel: Joint velocities, shape ``(nv,)``. Defaults to zero.
|
||||
act: Actuator activations, shape ``(na,)``. Defaults to zero.
|
||||
|
||||
Returns:
|
||||
Flat state vector suitable for ``mujoco.rollout``.
|
||||
"""
|
||||
data = mujoco.MjData(model)
|
||||
initial_state = np.empty(
|
||||
(mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS),)
|
||||
)
|
||||
if qpos.shape[0] != model.nq:
|
||||
raise ValueError(f"Expected qpos to have shape {model.nq}, got {qpos.shape[0]}.")
|
||||
data.qpos[:] = qpos
|
||||
if qvel is not None:
|
||||
if qvel.shape[0] != model.nv:
|
||||
raise ValueError(f"Expected qvel to have shape {model.nv}, got {qvel.shape[0]}.")
|
||||
data.qvel[:] = qvel
|
||||
if act is not None:
|
||||
if act.shape[0] != model.na:
|
||||
raise ValueError(f"Expected act to have shape {model.na}, got {act.shape[0]}.")
|
||||
data.act[:] = act
|
||||
mujoco.mj_getState(model, data, initial_state, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
||||
return initial_state
|
||||
|
||||
|
||||
class ModelSequences:
|
||||
"""A model spec paired with one or more measured trajectory sequences.
|
||||
|
||||
Groups a single ``MjSpec`` (the model to be identified) with the
|
||||
corresponding measured data so that the residual pipeline can iterate
|
||||
over all sequences for that model.
|
||||
|
||||
Args:
|
||||
name: Identifier for this model group (used for file-naming on save).
|
||||
spec: MjSpec that will be recompiled with candidate parameters.
|
||||
sequence_name: Name(s) identifying each measured sequence.
|
||||
initial_state: Initial state(s) for each sequence.
|
||||
control: Measured control TimeSeries for each sequence.
|
||||
sensordata: Measured sensor TimeSeries for each sequence.
|
||||
allow_missing_sensors: Passed through to
|
||||
:meth:`SystemTrajectory.check_compatible`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
spec: mujoco.MjSpec,
|
||||
sequence_name: str | Sequence[str],
|
||||
initial_state: np.ndarray | Sequence[np.ndarray],
|
||||
control: timeseries.TimeSeries | Sequence[timeseries.TimeSeries],
|
||||
sensordata: timeseries.TimeSeries | Sequence[timeseries.TimeSeries],
|
||||
allow_missing_sensors: bool = False,
|
||||
):
|
||||
self.name = name
|
||||
self.spec = spec
|
||||
self.allow_missing_sensors = allow_missing_sensors
|
||||
|
||||
self.gt_model = self.spec.compile()
|
||||
|
||||
self.sequence_name: list[str] = (
|
||||
[sequence_name] if isinstance(sequence_name, str) else list(sequence_name)
|
||||
)
|
||||
self.initial_state: list[np.ndarray] = (
|
||||
[initial_state] if isinstance(initial_state, np.ndarray) else list(initial_state)
|
||||
)
|
||||
self.control: list[timeseries.TimeSeries] = (
|
||||
[control] if isinstance(control, timeseries.TimeSeries) else list(control)
|
||||
)
|
||||
self.sensordata: list[timeseries.TimeSeries] = (
|
||||
[sensordata]
|
||||
if isinstance(sensordata, timeseries.TimeSeries)
|
||||
else list(sensordata)
|
||||
)
|
||||
|
||||
self.measured_rollout: list[SystemTrajectory] = []
|
||||
for initial_state_, control_, sensordata_ in zip(
|
||||
self.initial_state, self.control, self.sensordata, strict=True
|
||||
):
|
||||
measured_rollout_ = SystemTrajectory(
|
||||
model=self.gt_model,
|
||||
control=control_,
|
||||
sensordata=sensordata_,
|
||||
initial_state=initial_state_,
|
||||
state=None,
|
||||
)
|
||||
measured_rollout_.check_compatible(allow_missing_sensors=allow_missing_sensors)
|
||||
self.measured_rollout.append(measured_rollout_)
|
||||
|
||||
def __getitem__(self, key):
|
||||
return ModelSequences(
|
||||
self.name,
|
||||
self.spec,
|
||||
self.sequence_name[key],
|
||||
self.initial_state[key],
|
||||
self.control[key],
|
||||
self.sensordata[key],
|
||||
self.allow_missing_sensors,
|
||||
)
|
||||
|
||||
|
||||
def timeseries2array(
|
||||
control_signal: timeseries.TimeSeries | Sequence[timeseries.TimeSeries],
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
if isinstance(control_signal, timeseries.TimeSeries):
|
||||
control = control_signal.data
|
||||
control_times = control_signal.times
|
||||
else:
|
||||
control = np.stack([ts.data for ts in control_signal], axis=0)
|
||||
control_times = np.stack([ts.times for ts in control_signal], axis=0)
|
||||
# The measured data has N sensor measurements and N controls, where the first sensor
|
||||
# measurement corresponds to the initial condition. Thus we don't have ground truth
|
||||
# for the N+1'th state produced by the N'th control and so there is no point in
|
||||
# simulating it.
|
||||
if control.ndim == 3:
|
||||
control_applied_times = control_times[:, :-1]
|
||||
control_applied = control[:, :-1, :]
|
||||
else:
|
||||
control_applied_times = control_times[:-1]
|
||||
control_applied = control[:-1, :]
|
||||
return control_applied, control_applied_times
|
||||
|
||||
|
||||
def sequence2array(
|
||||
initial_states: np.ndarray | Sequence[np.ndarray],
|
||||
) -> np.ndarray:
|
||||
if isinstance(initial_states, np.ndarray):
|
||||
return initial_states
|
||||
return np.stack(initial_states, axis=0)
|
||||
|
||||
|
||||
def arrays2traj(
|
||||
models: mujoco.MjModel | Sequence[mujoco.MjModel],
|
||||
initial_states: np.ndarray | Sequence[np.ndarray],
|
||||
control: np.ndarray,
|
||||
control_times: np.ndarray,
|
||||
state: np.ndarray,
|
||||
sensordata: np.ndarray,
|
||||
signal_mapping: timeseries.SignalMappingType,
|
||||
state_mapping: timeseries.SignalMappingType,
|
||||
ctrl_mapping: timeseries.SignalMappingType,
|
||||
) -> Sequence[SystemTrajectory]:
|
||||
nbatch = state.shape[0]
|
||||
# TODO(kevin): When is np.tile necessary?
|
||||
# initial_states = np.tile(initial_states, (nbatch, 1))
|
||||
# control = np.tile(control, (nbatch, 1, 1))
|
||||
# control_times = np.tile(control_times, (nbatch, 1))
|
||||
|
||||
if isinstance(models, mujoco.MjModel):
|
||||
models_list = [models] * nbatch
|
||||
else:
|
||||
models_list = list(models)
|
||||
|
||||
return [
|
||||
SystemTrajectory(
|
||||
model=models_list[i],
|
||||
control=timeseries.TimeSeries(
|
||||
control_times[i], control[i], signal_mapping=ctrl_mapping
|
||||
),
|
||||
# NOTE(kevin): When using mjSTATE_FULLPHYSICS, the first element of
|
||||
# the state corresponds to the simulation time. The reason we do not
|
||||
# use control_times[i] is because sensordata times are shifted by
|
||||
# one time step.
|
||||
sensordata=timeseries.TimeSeries(state[i][:, 0], sensordata[i], signal_mapping),
|
||||
initial_state=initial_states[i],
|
||||
state=timeseries.TimeSeries(
|
||||
times=state[i][:, 0], data=state[i], signal_mapping=state_mapping
|
||||
),
|
||||
)
|
||||
for i in range(nbatch)
|
||||
]
|
||||
|
||||
|
||||
def sysid_rollout(
|
||||
models: mujoco.MjModel | Sequence[mujoco.MjModel],
|
||||
datas: mujoco.MjData | Sequence[mujoco.MjData],
|
||||
control_signal: Sequence[timeseries.TimeSeries] | timeseries.TimeSeries,
|
||||
initial_states: np.ndarray | Sequence[np.ndarray],
|
||||
rollout_signal_mapping: timeseries.SignalMappingType | None = None,
|
||||
rollout_state_mapping: timeseries.SignalMappingType | None = None,
|
||||
ctrl_mapping: timeseries.SignalMappingType | None = None,
|
||||
) -> Sequence[SystemTrajectory]:
|
||||
"""Rollout trajectories in parallel for the given models and controls.
|
||||
|
||||
Args:
|
||||
models: MuJoCo model or sequence of models.
|
||||
datas: MuJoCo data or sequence of data.
|
||||
control_signal: Control signals as TimeSeries or sequence of TimeSeries.
|
||||
initial_states: Initial states of the simulation. Shape (n_state,) or
|
||||
(n_batch, n_state).
|
||||
|
||||
Returns:
|
||||
Sequence of SystemTrajectory instances containing the simulation results.
|
||||
"""
|
||||
|
||||
# if the user does not supply it, we create it. Note that this will impact perf.
|
||||
if not rollout_signal_mapping or not rollout_state_mapping or not ctrl_mapping:
|
||||
if isinstance(models, mujoco.MjModel):
|
||||
model0 = models
|
||||
else:
|
||||
model0 = models[0]
|
||||
qpos_map, qvel_map, act_map, ctrl_mapping = (
|
||||
timeseries.TimeSeries.compute_all_state_mappings(model0)
|
||||
)
|
||||
rollout_state_mapping = qpos_map | qvel_map | act_map
|
||||
rollout_signal_mapping = timeseries.TimeSeries.compute_all_sensor_mapping(model0)
|
||||
|
||||
control, control_times = timeseries2array(control_signal)
|
||||
initial_states = sequence2array(initial_states)
|
||||
state, sensordata = mj_rollout.rollout(models, datas, initial_states, control)
|
||||
assert isinstance(state, np.ndarray)
|
||||
assert isinstance(sensordata, np.ndarray)
|
||||
|
||||
return arrays2traj(
|
||||
models,
|
||||
initial_states,
|
||||
control,
|
||||
control_times,
|
||||
state,
|
||||
sensordata,
|
||||
rollout_signal_mapping,
|
||||
rollout_state_mapping,
|
||||
ctrl_mapping,
|
||||
)
|
||||
Reference in New Issue
Block a user