Merge pull request #3079 from aftersomemath:sysid-pr

PiperOrigin-RevId: 868229512
Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea

python/mujoco/sysid/_src/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+

python/mujoco/sysid/_src/io.py

@@ -0,0 +1,82 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""I/O utilities for saving system identification results."""
+
+from collections.abc import Sequence
+import os
+import pathlib
+import pickle
+
+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
+import scipy.optimize as scipy_optimize
+
+
+def save_results(
+    experiment_results_folder: str | os.PathLike[str],
+    models_sequences: Sequence[ModelSequences],
+    initial_params: parameter.ParameterDict,
+    opt_params: parameter.ParameterDict,
+    opt_result: scipy_optimize.OptimizeResult,
+    residual_fn,
+):
+  """Save optimization results and confidence intervals to disk."""
+  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(b/0): 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
+      ),
+  )

python/mujoco/sysid/_src/model_modifier.py

@@ -0,0 +1,543 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Model modifiers."""
+
+from typing import Any
+
+import mujoco
+from mujoco.sysid._src.parameter import InertiaType
+from mujoco.sysid._src.parameter import ModifierFn
+from mujoco.sysid._src.parameter import Parameter
+from mujoco.sysid._src.parameter import ParameterDict
+import numpy as np
+
+
+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(b/0): 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.
+
+  Args:
+    actuator: An actuator object from model or spec.
+
+  Returns:
+    True if the actuator is a position actuator.
+  """
+  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]:
+  """Return the (P, D) gains of a position actuator."""
+  actuator_id = mujoco.mj_name2id(
+      model, mujoco.mjtObj.mjOBJ_ACTUATOR.value, 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:
+  """Set the proportional gain for a position actuator."""
+  # TODO(b/0): 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:
+  """Set the derivative gain for a position actuator."""
+  # TODO(b/0): 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:
+  """Set both proportional and derivative gains for a position actuator."""
+  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:
+  """Apply mass and center-of-mass position to a body."""
+  # TODO(b/0): 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(b/0): assert scalar
+  body = _infer_inertial(spec, body_name)
+  body.inertia *= value
+  return spec
+
+
+def pi_from_theta(theta: np.ndarray) -> np.ndarray:
+  """Convert base parameters θ to inertial parameters π."""
+  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].
+
+  Returns:
+    A 4x4 pseudoinertia matrix J.
+  """
+  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.
+
+  Args:
+    J: A 4x4 positive-definite matrix.
+
+  Returns:
+    An upper-triangular matrix U.
+  """
+  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:
+  """Convert diagonal inertia with quaternion to full inertia matrix."""
+  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:
+  """Extract base parameters θ from a body's inertial properties."""
+  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:
+  """Apply base-parameter inertia θ to a body in the spec."""
+  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):
+  """Apply inertia parameters to a body based on the parameter type."""
+  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)
+
+
+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:
+    spec: MuJoCo model specification object.
+    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.
+
+  Returns:
+    A Parameter configured for the body's inertia.
+  """
+
+  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

python/mujoco/sysid/_src/optimize.py

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

python/mujoco/sysid/_src/parameter.py

@@ -0,0 +1,651 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Parameter utilities."""
+
+from __future__ import annotations
+
+import copy
+from enum import Enum
+import pathlib
+from typing import Callable, TYPE_CHECKING, TypeAlias
+
+import colorama
+import mujoco
+import numpy as np
+import numpy.typing as npt
+from tabulate import tabulate
+import yaml
+
+if TYPE_CHECKING:
+  from typing_extensions import Self
+
+Fore = colorama.Fore
+Style = colorama.Style
+
+ModifierFn: TypeAlias = Callable[[mujoco.MjSpec, "Parameter"], object]
+
+
+class InertiaType(Enum):
+  Mass = 0
+  MassIpos = 1
+  Pseudo = 2
+
+
+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 the number of non-frozen parameters"
+          f" {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 the number of non-frozen parameters"
+            f" {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}"
+        f" {'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}"
+          f" {frozen_str}{Style.RESET_ALL}"
+      )
+
+    return "\n".join(info)

python/mujoco/sysid/_src/plotting.py

@@ -0,0 +1,731 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Plotting utilities."""
+
+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,
+    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.value
+
+  # 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.value,
+          )
+          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

python/mujoco/sysid/_src/residual.py

@@ -0,0 +1,449 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Residual computation for system identification."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping, Sequence
+import copy
+import os
+from typing import TypeAlias
+
+import mujoco
+from mujoco.sysid._src import model_modifier
+from mujoco.sysid._src import parameter
+from mujoco.sysid._src import signal_modifier
+from mujoco.sysid._src import timeseries
+from mujoco.sysid._src.trajectory import ModelSequences
+from mujoco.sysid._src.trajectory import SystemTrajectory
+from mujoco.sysid._src.trajectory import sysid_rollout
+import numpy as np
+
+_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}"
+          f" but declared as {enabled_obs_type}."
+      )
+
+    if obs_type == timeseries.SignalType.CustomObs:
+      raise ValueError(
+          "You are attempting to use the default SysID's modify_residual with"
+          f" 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])]
+    for i in range(x_reshaped.shape[1]):
+      param_dicts[i].update_from_vector(x_reshaped[:, i])
+    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.
+
+  Args:
+    x: Current parameter vector.
+    params: Parameter dictionary with bounds and metadata.
+    models_sequences: List of model-sequence groups.
+    build_model: Callable to compile a model from parameters.
+    modify_residual: Optional post-processing for the residual.
+    custom_rollout: Optional custom rollout function.
+    n_threads: Number of threads for parallel rollout.
+    return_pred_all: Whether to return all predicted signals.
+    resample_true: Whether to resample ground-truth to predicted times.
+    sensor_weights: Per-sensor weighting for the residual.
+    enabled_observations: Additional observation signals to include.
+
+  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

python/mujoco/sysid/_src/signal_modifier.py

@@ -0,0 +1,288 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Common signal modifiers."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import mujoco
+from mujoco.sysid._src import parameter
+from mujoco.sysid._src import timeseries
+import numpy as np
+
+
+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.
+    sort: If True, sort the indices before returning.
+
+  Returns:
+    A list of sensor indices.
+  """
+  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:
+  """Apply a bias to a sensor in a 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:
+  """Apply a gain to a sensor in a 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:
+  """Apply a delay to a sensor in a 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 within a time window.
+
+  The selected timestamps, plus the max delay, can be sampled from ts_sample.
+
+  Args:
+    ts: The timeseries to window.
+    min_t: The minimum time.
+    max_t: The maximum time.
+
+  Returns:
+    A new timeseries with the timestamps windowed.
+  """
+  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 to fit within the bounds of a delayed series.
+
+  Included timestamps must 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:
+  """Resample a timeseries and apply per-sensor delays."""
+  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:
+  """Prepare sensor weights array from a dict or numpy array."""
+  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 weighted difference between measured and predicted data.
+
+  Calculates ``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).
+    model: Optional mujoco model. This argument is required if sensor_weights
+      is not None.
+    sensor_weights: An optional dict mapping sensor name to weight. Unspecified
+      sensors are assumed to have a weight of 1.
+
+  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))

python/mujoco/sysid/_src/signal_transform.py

@@ -0,0 +1,282 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Declarative signal transformation for system identification residuals."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from fnmatch import fnmatch
+
+import mujoco
+from mujoco.sysid._src import parameter
+from mujoco.sysid._src import signal_modifier
+from mujoco.sysid._src import timeseries
+import numpy as np
+
+
+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:
+      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.
+
+    Args:
+      params: Current parameter dictionary.
+      sensordata_predicted: Predicted sensor TimeSeries from rollout.
+      sensordata_measured: Measured sensor TimeSeries (ground truth).
+      model: MuJoCo model for sensor metadata lookup.
+      return_pred_all: Whether to return all predicted signals.
+      state: Unused; part of the ModifyResidualFn signature.
+      sensor_weights: Per-sensor weighting for the residual.
+
+    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

python/mujoco/sysid/_src/timeseries.py

@@ -0,0 +1,773 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Time series utilities."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import Enum
+import pathlib
+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.
+
+  Args:
+    model: MuJoCo model used to look up sensor metadata.
+    names: Signal names or (name, type) tuples.
+    allowed_types: Set of acceptable signal types.
+
+  Returns:
+    A dictionary mapping signal names to (type, indices) tuples.
+  """
+  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"
+          f" {[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.value, 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.value, 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.value, 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.value, 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.value, 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.value, 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``.
+
+    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 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"
+            f" ({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 indices.size:
+          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"
+            f" ({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"
+            f" ({model.nu})."
+        )
+    else:
+      signal_mapping = _resolve_signals(
+          model, names, allowed_types={SignalType.MjCtrl}
+      )
+      max_idx = 0
+      for _, indices in signal_mapping.values():
+        if indices.size:
+          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"
+            f" ({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] | 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.
+
+    Args:
+      ts: The source TimeSeries.
+      enabled_sensors: Names of signals to keep.
+    """
+    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"
+            f" {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
+        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 not 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)

python/mujoco/sysid/_src/trajectory.py

@@ -0,0 +1,571 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Trajectory data containers for system identification."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import dataclasses
+import pathlib
+
+from absl import logging
+import mujoco
+import mujoco.rollout as mj_rollout
+from mujoco.sysid._src import timeseries
+import numpy as np
+
+
+@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 this trajectory to disk as a compressed NumPy archive."""
+    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:
+    """Load a trajectory from a compressed NumPy archive."""
+    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(
+            "Sensor data dimension %d does not match model sensor dimension %d",
+            self.sensordata.data.shape[1],
+            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.value
+    state_size = mujoco.mj_stateSize(self.model, state_spec)
+    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).
+
+    Args:
+      chunk_size: Number of time steps per chunk.
+
+    Returns:
+      A list of SystemTrajectory chunks.
+    """
+    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`.
+
+    Args:
+      height: Height of the rendered frames.
+      width: Width of the rendered frames.
+      camera: Camera index or name.
+      fps: Frames per second.
+
+    Returns:
+      A list of RGB frames as numpy arrays.
+    """
+    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.value),
+  ))
+  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.value
+  )
+  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]:
+  """Convert control TimeSeries to stacked arrays, dropping the last step."""
+  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:
+  """Stack a sequence of initial-state vectors into a single array."""
+  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]:
+  """Convert raw rollout arrays into a list of SystemTrajectory objects."""
+  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).
+    rollout_signal_mapping: Optional signal mapping for sensordata.
+    rollout_state_mapping: Optional signal mapping for state.
+    ctrl_mapping: Optional signal mapping for controls.
+
+  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,
+  )