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

sysid: named ic construction, bug fixes, docstrings, README, more tests

Browse Source 
Co-authored-by: Kevin Zakka <kevinarmandzakka@gmail.com>
This commit is contained in:
Levi Burner
2026-02-10 15:15:17 -05:00
parent 210cf86486
commit e89dae359e
15 changed files with 768 additions and 1630 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/sysid/_src/io.py
+10 -11
View File
@@ -35,7 +35,16 @@ def save_results(
opt_result: scipy_optimize.OptimizeResult,
residual_fn,
):
"""Save optimization results and confidence intervals to disk."""
"""Save optimization results and confidence intervals to disk.
Args:
experiment_results_folder: Directory where results are written.
models_sequences: Model/sequence groups; identified XMLs are saved here.
initial_params: Parameters before optimization.
opt_params: Parameters after optimization.
opt_result: Scipy OptimizeResult from the optimizer.
residual_fn: Residual function used to compute confidence intervals.
"""
experiment_results_folder = pathlib.Path(experiment_results_folder)
if not experiment_results_folder.exists():
experiment_results_folder.mkdir(parents=True, exist_ok=True)
@@ -70,13 +79,3 @@ def save_results(
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
),
)
python/mujoco/sysid/_src/model_modifier.py
+63 -8
View File
@@ -111,7 +111,12 @@ def is_position_actuator(actuator) -> bool:
def get_actuator_pd_gains(
model: mujoco.MjModel, actuator_name: str
) -> tuple[float, float]:
"""Return the (P, D) gains of a position actuator."""
"""Return the (P, D) gains of a position actuator.
Args:
model: MuJoCo model.
actuator_name: Name of the actuator.
"""
actuator_id = mujoco.mj_name2id(
model, mujoco.mjtObj.mjOBJ_ACTUATOR.value, actuator_name
)
@@ -128,7 +133,13 @@ def apply_pgain(
actuator_name: str,
value: float | np.ndarray,
) -> mujoco.MjSpec:
"""Set the proportional gain for a position actuator."""
"""Set the proportional gain for a position actuator.
Args:
spec: MuJoCo model specification.
actuator_name: Name of the actuator.
value: Proportional gain value.
"""
# TODO(b/0): assert scalar
actuator = _get_obj_or_raise(spec, "actuator", actuator_name)
assert isinstance(actuator, mujoco.MjsActuator)
@@ -144,7 +155,13 @@ def apply_dgain(
actuator_name: str,
value: float | np.ndarray,
) -> mujoco.MjSpec:
"""Set the derivative gain for a position actuator."""
"""Set the derivative gain for a position actuator.
Args:
spec: MuJoCo model specification.
actuator_name: Name of the actuator.
value: Derivative gain value.
"""
# TODO(b/0): assert scalar
actuator = _get_obj_or_raise(spec, "actuator", actuator_name)
assert isinstance(actuator, mujoco.MjsActuator)
@@ -159,7 +176,13 @@ def apply_pdgain(
actuator_name: str,
value: np.ndarray,
) -> mujoco.MjSpec:
"""Set both proportional and derivative gains for a position actuator."""
"""Set both proportional and derivative gains for a position actuator.
Args:
spec: MuJoCo model specification.
actuator_name: Name of the actuator.
value: 2-element array ``[P_gain, D_gain]``.
"""
if value.size != 2:
raise ValueError(f"pdgain must be a 2-element array, got {value.size}.")
apply_pgain(spec, actuator_name, value[0])
@@ -174,7 +197,16 @@ def apply_body_mass_ipos(
ipos: np.ndarray | None = None,
rot_inertia_scale: bool = False,
) -> mujoco.MjSpec:
"""Apply mass and center-of-mass position to a body."""
"""Apply mass and center-of-mass position to a body.
Args:
spec: MuJoCo model specification.
body_name: Name of the body.
mass: Optional new mass value.
ipos: Optional new center-of-mass position.
rot_inertia_scale: If True, scale rotational inertia proportionally to
mass change.
"""
# TODO(b/0): assert mass and ipos shapes
body = _infer_inertial(spec, body_name)
mass_original = body.mass
@@ -200,7 +232,18 @@ def scale_body_inertia(
def pi_from_theta(theta: np.ndarray) -> np.ndarray:
"""Convert base parameters θ to inertial parameters π."""
"""Convert base parameters θ to inertial parameters π.
Args:
theta: 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
Returns:
10-D array π = [m, hx, hy, hz, Ixx, Iyy, Izz, Ixy, Iyz, Ixz].
"""
alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3 = theta
exp_alpha = np.exp(alpha)
exp_d1 = np.exp(d1)
@@ -361,7 +404,13 @@ def apply_body_theta_inertia(
body_name: str,
theta: np.ndarray,
) -> mujoco.MjSpec:
"""Apply base-parameter inertia θ to a body in the spec."""
"""Apply base-parameter inertia θ to a body in the spec.
Args:
spec: MuJoCo model specification.
body_name: Name of the body.
theta: 10-element array [alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3].
"""
if theta.size != 10:
raise ValueError(f"theta must be a 10-element array, got {theta.size}.")
pi = pi_from_theta(theta)
@@ -392,7 +441,13 @@ def apply_body_theta_inertia(
def apply_body_inertia(spec: mujoco.MjSpec, name: str, param: Parameter):
"""Apply inertia parameters to a body based on the parameter type."""
"""Apply inertia parameters to a body based on the parameter type.
Args:
spec: MuJoCo model specification.
name: Name of the body.
param: Parameter with an ``inertia_type`` attribute.
"""
if not hasattr(param, "inertia_type"):
raise ValueError(
f"Parameter {param.name} does not have inertia_type attribute."
python/mujoco/sysid/_src/optimize.py
+28 -3
View File
@@ -44,7 +44,7 @@ def _scipy_least_squares(
jac_arg: str | Callable[..., Any]
if use_mujoco_jac:
# This is the default step sized for finite difference used in
# This is the default step size for finite difference used in
# scipy's least_squares and mujoco's minimize finite difference
# https://github.com/scipy/scipy/blob/91e18f3bd355477b
# 8b7747ec82d70ac98ffd2422/scipy/optimize/_numdiff.py#L404
@@ -143,6 +143,7 @@ def optimize(
initial_params: parameter.ParameterDict,
residual_fn: Callable[..., Any],
optimizer: Literal["scipy", "mujoco", "scipy_parallel_fd"] = "mujoco",
verbose: bool = True,
**optimizer_kwargs,
) -> tuple[parameter.ParameterDict, scipy_optimize.OptimizeResult]:
"""Run nonlinear least-squares optimization on the residual.
@@ -153,11 +154,12 @@ def optimize(
returned by :func:`build_residual_fn`.
optimizer: Backend — ``"mujoco"`` (default), ``"scipy"``, or
``"scipy_parallel_fd"`` (scipy with MuJoCo finite-difference Jacobian).
verbose: If True, log parameter comparison table after optimization.
**optimizer_kwargs: Forwarded to the backend (e.g. ``max_iters``,
``verbose``, ``loss``).
Returns:
``(opt_params, opt_result)`` — the optimised ParameterDict and a
``(opt_params, opt_result)`` — the optimized ParameterDict and a
``scipy.optimize.OptimizeResult`` with at least ``x``, ``jac``, ``grad``.
"""
x0 = initial_params.as_vector()
@@ -187,6 +189,16 @@ def optimize(
opt_params.update_from_vector(opt_result.x)
if verbose:
logging.info(
"\n%s",
opt_params.compare_parameters(
initial_params.as_vector(),
opt_params.as_vector(),
measured_params=initial_params.as_nominal_vector(),
),
)
return opt_params, opt_result
@@ -197,7 +209,20 @@ def calculate_intervals(
lambda_zero_thresh=1e-15,
v_zero_thresh=1e-8,
):
"""Calculate confidence intervals from the Jacobian at the optimum."""
"""Calculate confidence intervals from the Jacobian at the optimum.
Args:
residuals_star: List of residual arrays at the optimum.
J: Jacobian matrix at the optimum, shape ``(n_residuals, n_params)``.
alpha: Significance level for the confidence intervals.
lambda_zero_thresh: Threshold below which eigenvalues are treated as zero.
v_zero_thresh: Threshold below which eigenvector elements are treated as
zero.
Returns:
``(Sigma_X, intervals)`` — the parameter covariance matrix and the
half-width confidence intervals for each parameter.
"""
if J is None or J.size == 0:
return np.empty((0, 0)), np.empty((0,))
python/mujoco/sysid/_src/parameter.py
+27 -5
View File
@@ -105,6 +105,11 @@ class Parameter:
return self.nominal.flatten()
def update_from_vector(self, vector: np.ndarray) -> None:
"""Update the current value from a flat vector.
Args:
vector: Flat array of length ``self.size``.
"""
vector_array = np.atleast_1d(vector)
if len(vector_array) != self.size:
raise ValueError(
@@ -125,7 +130,11 @@ class Parameter:
self.value = self.nominal.copy()
def sample(self, rng: np.random.Generator | None = None) -> np.ndarray:
"""Sample a random value uniformly within bounds."""
"""Sample a random value uniformly within bounds.
Args:
rng: Optional numpy random generator. Uses default if None.
"""
if rng is None:
rng = np.random.default_rng()
return rng.uniform(self.min_value.flatten(), self.max_value.flatten())
@@ -197,7 +206,7 @@ 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,
vectorized access (``as_vector`` / ``update_from_vector``), serialization,
and tabular comparison of parameter estimates.
Frozen parameters are silently skipped by vector/bounds methods so that the
@@ -271,7 +280,12 @@ class ParameterDict:
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."""
"""Update all non-frozen parameters from a flat vector.
Args:
vector: Flat array whose length equals the total size of non-frozen
parameters.
"""
start = 0
for param in self.parameters.values():
if not param.frozen:
@@ -333,14 +347,22 @@ class ParameterDict:
param.reset()
def sample(self, rng: np.random.Generator | None = None) -> np.ndarray:
"""Sample parameter values within bounds for non-frozen parameters."""
"""Sample parameter values within bounds for non-frozen parameters.
Args:
rng: Optional numpy random generator. Uses default if None.
"""
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."""
"""Randomize parameter values for non-frozen parameters.
Args:
rng: Optional numpy random generator. Uses default if None.
"""
for param in self.parameters.values():
if not param.frozen:
param.value = param.sample(rng)
python/mujoco/sysid/_src/plotting.py
-617
View File
@@ -19,627 +19,10 @@ from __future__ import annotations
from collections.abc import Sequence
from matplotlib.lines import Line2D
import matplotlib.pyplot as plt
import mujoco
from mujoco.sysid._src import parameter
import numpy as np
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.
title_prefix: Optional prefix for subplot titles.
sensor_ids: Optional list of sensor indices to plot.
"""
# 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 sensor_id in 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),
):
"""Plot the objective value over optimization iterations."""
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,
):
"""Plot candidate parameter values and their diffs over iterations."""
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(b/0) support pages, they are currently broken because
# saving to disk overwrites 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,
):
"""Plot a heatmap of candidate parameter values over iterations."""
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,
):
"""Plot parameter estimates with confidence intervals."""
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,
python/mujoco/sysid/_src/signal_modifier.py
+43 -6
View File
@@ -68,7 +68,13 @@ def apply_bias(
sensor_name: str,
bias: parameter.Parameter,
) -> timeseries.TimeSeries:
"""Apply a bias to a sensor in a timeseries."""
"""Apply a bias to a sensor in a timeseries.
Args:
ts: Input timeseries.
sensor_name: Name of the sensor to modify.
bias: Parameter whose ``.value`` is added to the sensor columns.
"""
indices = ts.get_indices(sensor_name)[1]
data_out = ts.data.copy()
data_out[..., indices] += bias.value
@@ -80,7 +86,13 @@ def apply_gain(
sensor_name: str,
gain: parameter.Parameter,
) -> timeseries.TimeSeries:
"""Apply a gain to a sensor in a timeseries."""
"""Apply a gain to a sensor in a timeseries.
Args:
ts: Input timeseries.
sensor_name: Name of the sensor to modify.
gain: Parameter whose ``.value`` multiplies the sensor columns.
"""
indices = ts.get_indices(sensor_name)[1]
data_out = ts.data.copy()
data_out[..., indices] *= gain.value
@@ -92,7 +104,13 @@ def apply_delay(
sensor_name: str,
delay: parameter.Parameter,
) -> timeseries.TimeSeries:
"""Apply a delay to a sensor in a timeseries."""
"""Apply a delay to a sensor in a timeseries.
Args:
ts: Input timeseries.
sensor_name: Name of the sensor to delay.
delay: Parameter whose ``.value`` is the delay in seconds.
"""
indices = ts.get_indices(sensor_name)[1]
ts_sensor = timeseries.TimeSeries(
@@ -205,7 +223,15 @@ def apply_resample_and_delay(
sensor_delays: dict[str, float] | None = None,
predicted_data: bool = True,
) -> timeseries.TimeSeries:
"""Resample a timeseries and apply per-sensor delays."""
"""Resample a timeseries and apply per-sensor delays.
Args:
ts: Input timeseries to resample.
times: Target timestamps.
default_delay: Default delay applied to all columns.
sensor_delays: Optional per-sensor delay overrides.
predicted_data: If True, negate delays (shift predicted to match measured).
"""
delays = _build_per_column_delays(
ts, default_delay, sensor_delays, predicted_data
)
@@ -234,7 +260,13 @@ def prepare_sensor_weights(
n_sensors: int,
model: mujoco.MjModel,
) -> np.ndarray:
"""Prepare sensor weights array from a dict or numpy array."""
"""Prepare sensor weights array from a dict or numpy array.
Args:
sensor_weights: Mapping from sensor name to weight, or a flat array.
n_sensors: Total number of sensor columns.
model: MuJoCo model for resolving sensor names to indices.
"""
if isinstance(sensor_weights, np.ndarray):
if sensor_weights.ndim != 1 or sensor_weights.shape[0] != n_sensors:
raise ValueError(
@@ -284,5 +316,10 @@ def normalize_residual(
residual: np.ndarray,
measured_data: np.ndarray,
) -> np.ndarray:
"""Normalize the residual by the standard deviation of the measured data."""
"""Normalize the residual by the standard deviation of the measured data.
Args:
residual: Residual array, shape ``(n_timesteps, n_sensors)``.
measured_data: Measured data array, same shape as *residual*.
"""
return residual / (np.linalg.norm(measured_data, axis=0) / np.sqrt(2))
python/mujoco/sysid/_src/signal_transform.py
+23 -4
View File
@@ -51,7 +51,12 @@ class SignalTransform:
self.normalize = normalize
def delay(self, pattern: str, param: parameter.Parameter) -> None:
"""Register a delay for sensors matching *pattern* (fnmatch)."""
"""Register a delay for sensors matching *pattern* (fnmatch).
Args:
pattern: fnmatch pattern matched against sensor names.
param: Parameter whose ``.value`` is the delay in seconds.
"""
self._delays.append((pattern, param.name, param))
def gain(
@@ -79,7 +84,13 @@ class SignalTransform:
param: parameter.Parameter,
target: str = "both",
) -> None:
"""Register an additive bias for sensors matching *pattern*."""
"""Register an additive bias for sensors matching *pattern*.
Args:
pattern: fnmatch pattern matched against sensor names.
param: Parameter whose ``.value`` is the additive bias.
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}"
@@ -87,11 +98,19 @@ class SignalTransform:
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."""
"""Only include these sensors in the returned residual/timeseries.
Args:
sensor_names: Sensor names to keep in the output.
"""
self._enabled_sensors = list(sensor_names)
def set_sensor_weights(self, weights: Mapping[str, float]) -> None:
"""Set per-sensor weights for the weighted diff."""
"""Set per-sensor weights for the weighted diff.
Args:
weights: Mapping from sensor name to weight.
"""
self._sensor_weights = weights
# Private methods.
python/mujoco/sysid/_src/timeseries.py
+33 -12
View File
@@ -220,22 +220,26 @@ class TimeSeries:
nq = model.nq
nv = model.nv
# Bodies
# Bodies with free joints, named by body rather than joint.
for body_id in range(model.nbody):
b = model.body(body_id)
body_name = b.name
start_index = model.body_dofadr[body_id]
dof_adr = model.body_dofadr[body_id]
if start_index >= 0 and b.dofnum[0] == 6:
qpos_indices = np.arange(start_index, start_index + 7)
if dof_adr >= 0 and b.dofnum[0] == 6:
# Use the body's first joint qposadr for qpos. Free joints have
# 7 qpos elements but only 6 dofs, so dofadr and qposadr diverge
# for subsequent entries.
first_jnt = model.body_jntadr[body_id]
qpos_adr = model.jnt_qposadr[first_jnt]
qpos_indices = np.arange(qpos_adr, qpos_adr + 7)
qpos_map[f"{body_name}_qpos"] = (SignalType.MjStateQPos, qpos_indices)
qvel_indices = np.arange(start_index + nq, start_index + nq + 6)
qvel_indices = np.arange(dof_adr + nq, dof_adr + nq + 6)
qvel_map[f"{body_name}_qvel"] = (SignalType.MjStateQVel, qvel_indices)
# Joints
# Joints, excluding free joints which are handled above.
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
@@ -246,9 +250,11 @@ class TimeSeries:
elif jnt_type == mujoco.mjtJoint.mjJNT_FREE:
continue
qpos_indices = np.arange(start_index, start_index + qpos_width)
qpos_adr = model.jnt_qposadr[jnt_id]
qpos_indices = np.arange(qpos_adr, qpos_adr + qpos_width)
qpos_map[f"{jnt_name}_qpos"] = (SignalType.MjStateQPos, qpos_indices)
qvel_indices = np.arange(start_index + nq, start_index + nq + qvel_width)
dof_adr = model.jnt_dofadr[jnt_id]
qvel_indices = np.arange(dof_adr + nq, dof_adr + nq + qvel_width)
qvel_map[f"{jnt_name}_qvel"] = (SignalType.MjStateQVel, qvel_indices)
# Actuators
@@ -445,7 +451,11 @@ class TimeSeries:
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."""
"""Look up the signal type and column indices for a named observation.
Args:
obs_name: Name of the observation signal.
"""
assert self.signal_mapping is not None
if obs_name not in self.signal_mapping:
raise ValueError(
@@ -462,7 +472,14 @@ class TimeSeries:
str, tuple[SignalType, np.ndarray | list[int] | int]
],
) -> TimeSeries:
"""Construct a TimeSeries, normalising index entries to ``np.ndarray``."""
"""Construct a TimeSeries, normalizing index entries to ``np.ndarray``.
Args:
times: 1-D timestamp array.
data: Data array with first axis corresponding to time.
signal_mapping: Dict mapping signal names to ``(type, indices)`` tuples.
Index entries are coerced to ``np.ndarray``.
"""
normalized: SignalMappingType = {}
for key in signal_mapping:
signal_type, indices = signal_mapping[key]
@@ -560,7 +577,11 @@ class TimeSeries:
)
def save_to_csv(self, path: str | pathlib.Path) -> None:
"""Save the time series data to a CSV file."""
"""Save the time series data to a CSV file.
Args:
path: Path where the CSV file will be written.
"""
np.savetxt(
path,
np.concatenate([self.times[:, None], self.data], axis=1),
python/mujoco/sysid/_src/trajectory.py
+81 -7
View File
@@ -302,12 +302,23 @@ class SystemTrajectory:
height=height,
)
def _map_states(from_array, to_array, from_names, to_mapping, map_offset):
i = 0
for name in from_names:
_, indices = to_mapping[name]
width = indices.shape[0]
to_array[indices - map_offset] = from_array[i:i+width]
i += width
return to_array
def create_initial_state(
model: mujoco.MjModel,
qpos: np.ndarray,
qvel: np.ndarray | None = None,
act: np.ndarray | None = None,
qpos_names: Sequence[str] | None = None,
qvel_names: Sequence[str] | None = None,
act_names: Sequence[str] | None = None,
) -> np.ndarray:
"""Build a ``mjSTATE_FULLPHYSICS`` initial-state vector from components.
@@ -316,14 +327,46 @@ def create_initial_state(
qpos: Joint positions, shape ``(nq,)``.
qvel: Joint velocities, shape ``(nv,)``. Defaults to zero.
act: Actuator activations, shape ``(na,)``. Defaults to zero.
qpos_names: Names to map elements of qpos to specific MuJoCo states.
If None, assumes qpos is in MuJoCo's order.
qvel_names: Names to map elements of qvel to specific MuJoCo states.
If None, assumes qvel is in MuJoCo's order.
act_names: Actuator names to map elements of act to specific MuJoCo
actuators. If None, assumes act is in MuJoCo's order.
Returns:
Flat state vector suitable for ``mujoco.rollout``.
"""
data = mujoco.MjData(model)
initial_state = np.empty((
mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS.value),
))
if qpos_names is not None and len(qpos_names) != qpos.shape[0]:
raise ValueError(
f"Expected qpos to have shape {len(qpos_names)}, got {qpos.shape[0]}"
)
if qvel is None and qvel_names is not None:
raise ValueError("Expected qvel to not be None when qvel_names is not None")
if qvel_names is not None and qvel is not None and len(qvel_names) != qvel.shape[0]:
raise ValueError(
f"Expected qvel to have shape {len(qvel_names)}, got {qvel.shape[0]}"
)
if act_names is not None and len(act_names) != act.shape[0]:
raise ValueError(
f"Expected act to have shape {len(act_names)}, got {act.shape[0]}"
)
if (qpos_names is not None
or qvel_names is not None
or act_names is not None):
qpos_map, qvel_map, act_map, _ = timeseries.TimeSeries.compute_all_state_mappings(model)
if qpos_names is not None:
qpos = _map_states(qpos, np.copy(data.qpos), qpos_names, qpos_map, 0)
if qvel_names is not None:
indices_offset = data.qpos.shape[0]
qvel = _map_states(qvel, np.copy(data.qvel), qvel_names, qvel_map, indices_offset)
if act_names is not None:
indices_offset = data.qpos.shape[0] + data.qvel.shape[0]
act = _map_states(act, np.copy(data.act), act_names, act_map, indices_offset)
if qpos.shape[0] != model.nq:
raise ValueError(
f"Expected qpos to have shape {model.nq}, got {qpos.shape[0]}."
@@ -341,6 +384,10 @@ def create_initial_state(
f"Expected act to have shape {model.na}, got {act.shape[0]}."
)
data.act[:] = act
initial_state = np.empty((
mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS.value),
))
mujoco.mj_getState(
model, data, initial_state, mujoco.mjtState.mjSTATE_FULLPHYSICS.value
)
@@ -433,7 +480,14 @@ class ModelSequences:
def timeseries2array(
control_signal: timeseries.TimeSeries | Sequence[timeseries.TimeSeries],
) -> tuple[np.ndarray, np.ndarray]:
"""Convert control TimeSeries to stacked arrays, dropping the last step."""
"""Convert control TimeSeries to stacked arrays, dropping the last step.
Args:
control_signal: Control TimeSeries or sequence of TimeSeries.
Returns:
``(control_array, control_times)`` with the last time step removed.
"""
if isinstance(control_signal, timeseries.TimeSeries):
control = control_signal.data
control_times = control_signal.times
@@ -457,7 +511,11 @@ def timeseries2array(
def sequence2array(
initial_states: np.ndarray | Sequence[np.ndarray],
) -> np.ndarray:
"""Stack a sequence of initial-state vectors into a single array."""
"""Stack a sequence of initial-state vectors into a single array.
Args:
initial_states: Single state array or sequence of state arrays.
"""
if isinstance(initial_states, np.ndarray):
return initial_states
return np.stack(initial_states, axis=0)
@@ -474,12 +532,28 @@ def arrays2traj(
state_mapping: timeseries.SignalMappingType,
ctrl_mapping: timeseries.SignalMappingType,
) -> Sequence[SystemTrajectory]:
"""Convert raw rollout arrays into a list of SystemTrajectory objects."""
"""Convert raw rollout arrays into a list of SystemTrajectory objects.
Args:
models: Single model or sequence of models (one per batch element).
initial_states: Initial state array(s).
control: Control array, shape ``(nbatch, nsteps, nu)``.
control_times: Control timestamps, shape ``(nbatch, nsteps)``.
state: State array, shape ``(nbatch, nsteps, nstate)``.
sensordata: Sensor data array, shape ``(nbatch, nsteps, nsensordata)``.
signal_mapping: Signal mapping for sensor data.
state_mapping: Signal mapping for state data.
ctrl_mapping: Signal mapping for control data.
"""
nbatch = state.shape[0]
# TODO(kevin): When is np.tile necessary?
# TODO(kevin): When is np.tile/atleast_2d/etc necessary?
# initial_states = np.tile(initial_states, (nbatch, 1))
# control = np.tile(control, (nbatch, 1, 1))
# control_times = np.tile(control_times, (nbatch, 1))
initial_states = np.atleast_2d(initial_states)
if control.ndim == 2:
control = control[np.newaxis, :, :]
control_times = np.atleast_2d(control_times)
if isinstance(models, mujoco.MjModel):
models_list = [models] * nbatch