diff --git a/python/mujoco/sysid/README.md b/python/mujoco/sysid/README.md
index b96b9f0f..db96769a 100644
--- a/python/mujoco/sysid/README.md
+++ b/python/mujoco/sysid/README.md
@@ -1,637 +1,95 @@
-# Practical System Identification
+# System Identification Toolbox
-A toolbox for system identification built on top of MuJoCo.
+Given a MuJoCo model and recorded sensor data, find parameters
+that make simulation match reality. By default, the library uses
+nonlinear least-squares with box constraints to minimize the difference
+between measured and simulated (predicted) outputs. Residuals
+can be modified by static or optimized parameters, such as
+weights and time-delays.
-## API Overview
+The optimizer uses Gauss-Newton with finite-difference Jacobians. Each
+parameter perturbation requires an independent simulation rollout. All
+of them execute in a single batched call to `mujoco.rollout`, parallelized
+across threads.
-The library solves a **box-constrained nonlinear least-squares** problem. Given
-a parameter vector `θ`, simulated sensor readings `ȳ(θ)`, and recorded sensor
-data `y`, the objective is:
+## Pipeline
-```
-min ½ ‖W (ȳ(θ) − y)‖²
- θ
+**You provide:**
+- One or more `ModelSequences` each bundling a single `MjSpec` with one or
+ more sequences of measured data. All will be optimized jointly.
+- A `ParameterDict` defining differentiable `Parameter`'s with bounds.
+- Callbacks to apply `Parameter`'s to an `MjSpec` (individually or jointly)
+- (optional) Functions (`build_model`, `custom_rollout`, `modify_residual`)
+ that override the default residual function behaviour.
-subject to θ_min ≤ θ ≤ θ_max
-```
+**The framework:**
+- `build_residual_fn` which composes user code in `ModelSequences` and overrides
+- Optimizes (`optimize`) the residual function returned by `build_residual_fn`
+- Optimizes parameters via batched parallel rollouts (`optimize`).
+- Saves results and generates an HTML report (`save_results`, `default_report`).
-where `W` is a diagonal weighting matrix and the box constraints enforce
-physical plausibility (e.g., positive masses).
+## What Can You Identify?
-The optimizer uses the **Gauss-Newton** method. The residual Jacobian
-`J = ∂r/∂θ` is computed by **finite differences**: each column of `J` requires
-one perturbed simulation rollout, and these evaluations are independent across
-parameters and parallelize naturally across threads. The Gauss-Newton
-approximate Hessian is `H ≈ JᵀJ` and the gradient is `g = Jᵀr`, yielding the
-update `Δθ = −H⁻¹g`. Box constraints are handled by projected steps.
+You can optimize any parameter that differentiably modifies the final
+residuals. Common use cases include:
-The pipeline has five stages:
+**Physics parameters** settable on `MjSpec`. Most parameters in MjSpec
+can be easily set directly by the user provided callbacks:
-```
-Define Parameters ──> Package Data ──> Build Residual ──> Optimize ──> Save / Report
- ParameterDict ModelSequences build_residual_fn optimize save_results
-```
+| Target | Approach |
+|---|---|
+| Contact sliding friction | `spec.pair("cp").friction[0] = p.value[0]` |
+| Joint damping | `spec.joint("j1").damping = p.value[0]` |
----
-
-### What Can You Identify?
-
-**Anything settable on `MjSpec` can be identified** via modifier callbacks. The
-convenience functions handle common cases with correct bounds; for everything
-else, write a `modifier` lambda that sets the quantity on the spec.
-
-**Physics parameters** — these change the model before simulation:
+Convenience functions are provided for common system identification
+parameterizations that cannot be trivially appled to an MjSpec:
| Target | Approach |
|---|---|
| Body mass | `body_inertia_param(..., InertiaType.Mass)` |
| Body mass + center of mass | `body_inertia_param(..., InertiaType.MassIpos)` |
-| Full body inertia (10-D) | `body_inertia_param(..., InertiaType.Pseudo)` |
-| Actuator P/D gains | `Parameter(..., modifier=lambda s, p: apply_pgain(s, "act1", p.value[0]))` |
-| Contact friction / solref | `Parameter(..., modifier=lambda s, p: s.pair("cp").friction.__setitem__(0, p.value[0]))` |
-| Joint damping / stiffness | `Parameter(..., modifier=lambda s, p: setattr(s.joint("j1"), "damping", p.value[0]))` |
+| Full inertia (10-D) | `body_inertia_param(..., InertiaType.Pseudo)` |
+| Actuator P/D gains | `apply_pdgain(spec, "act1", p.value)` |
-**Measurement parameters** — real sensors aren't perfect. They may lag behind
-the simulation clock, have an unknown scale factor, or sit at a nonzero offset.
-These can't be set on `MjSpec` because they aren't physics — they're artifacts
-of the measurement system. `SignalTransform` (Section 4) adjusts the simulated
-or recorded signals *after* rollout to account for these:
+Full inertia uses the pseudo-inertia Cholesky parameterization
+([Rucker & Wensing 2022](https://ieeexplore.ieee.org/document/9690029)),
+guaranteeing physical consistency without singularities.
-| Target | Approach |
-|---|---|
-| Sensor delay | `transform.delay("*_pos", params["delay"])` |
-| Sensor gain/scale | `transform.gain("*_torque", params["scale"])` |
-| Sensor bias/offset | `transform.bias("*_vel", params["bias"])` |
+**Measurement parameters** such as sensor delays, gains, and biases are
+properties of the measurement system, not the physics model. The library
+provides utilities for applying these corrections to the residual after
+rollout, they are functionally complete but their API is not yet final.
-#### Common recipes
-
-**Identify link masses of a robot arm:**
-
-```python
-from mujoco.sysid import body_inertia_param, InertiaType, ParameterDict
-
-params = ParameterDict()
-for link in ["link1", "link2", "link3"]:
- params.add(body_inertia_param(spec, model, link, inertia_type=InertiaType.Mass))
-```
-
-**Identify contact friction:**
-
-```python
-from mujoco.sysid import Parameter
-
-params.add(Parameter(
- "floor_friction",
- nominal=1.0, min_value=0.1, max_value=3.0,
- modifier=lambda s, p: s.pair("foot_floor").friction.__setitem__(0, p.value[0]),
-))
-```
-
----
-
-### 1. Define Parameters
-
-A **`Parameter`** is a named value (scalar or array) with bounds and an optional
-**modifier callback** that knows how to apply itself to a MuJoCo spec. A
-**`ParameterDict`** collects parameters into the single vector that the
-optimizer sees — it handles flattening them into one array, writing optimizer
-updates back, and enforcing bounds:
-
-```python
-from mujoco.sysid import Parameter, ParameterDict
-
-params = ParameterDict()
-
-params.add(Parameter(
- "box_mass",
- nominal=5.0, # starting value
- min_value=1.0,
- max_value=10.0,
- modifier=lambda spec, p: setattr(spec.body("box"), "mass", p.value[0]),
-))
-
-params.add(Parameter(
- "friction",
- nominal=[1.6, 0.005],
- min_value=[0.0, 0.0],
- max_value=[3.0, 0.01],
- frozen=True, # excluded from optimization
- modifier=lambda spec, p: spec.pair("contact").friction.__setitem__(slice(0, 2), p.value),
-))
-```
-
-**`frozen`**: A frozen parameter is completely invisible to the optimizer — it
-is excluded from `as_vector()`, `update_from_vector()`, `get_bounds()`, and
-`randomize()`. Its modifier callback is also **not called** during
-`apply_param_modifiers`, so the model uses whatever value is already in the XML
-spec for that quantity. The intended workflow: define all parameters you might
-ever want to identify up front, then toggle `frozen` on and off as you
-iteratively narrow which parameters matter.
-
-Key `ParameterDict` methods:
-
-| Method | Description |
-|---|---|
-| `as_vector()` | Flatten all non-frozen parameters into a 1-D array |
-| `update_from_vector(x)` | Write a flat array back into the parameters |
-| `get_bounds()` | Returns `(lower, upper)` bound arrays |
-| `randomize(rng)` | Sample each non-frozen parameter uniformly within bounds |
-| `reset()` | Restore every parameter to its nominal value |
-| `copy()` | Deep copy (preserves modifier lambdas) |
-| `save_to_disk(path)` / `load_from_disk(path)` | YAML serialization (schema + values) |
-
-#### Body inertia parameterization
-
-Rigid-body inertia is tricky to identify: mass, center-of-mass, and the rotational
-inertia tensor are coupled, and naively optimizing the 6 independent entries of
-the inertia tensor can produce physically impossible results (e.g. negative
-eigenvalues). The library implements three parameterizations of increasing
-fidelity:
-
-| `InertiaType` | Params | What it identifies |
-|---|---|---|
-| `Mass` | 1 | Mass only. Optionally scales the existing rotational inertia proportionally (`scale_rot_inertia=True`). |
-| `MassIpos` | 4 | Mass + center-of-mass position (3-D). Optionally scales rotational inertia. |
-| `Pseudo` | 10 | Full inertia via the pseudo-inertia Cholesky factor from [Rucker & Wensing 2022](https://ieeexplore.ieee.org/document/9690029). The 10 parameters `θ = [α, d₁, d₂, d₃, s₁₂, s₂₃, s₁₃, t₁, t₂, t₃]` are the entries of a lower-triangular matrix whose product `LLᵀ` is the 4×4 pseudo-inertia matrix. Physical consistency (positive mass, positive-definite inertia tensor) is guaranteed by construction for any `θ`. |
-
-Use `body_inertia_param` to create a `Parameter` with the right nominal values,
-bounds, and modifier already wired up:
-
-```python
-from mujoco.sysid import body_inertia_param, InertiaType
-
-param = body_inertia_param(
- spec, model, "link1",
- inertia_type=InertiaType.Pseudo,
-)
-params.add(param)
-```
-
----
-
-### 2. Package Data
-
-Measured data is stored as **`TimeSeries`** objects (frozen dataclass: `times`,
-`data`, optional `signal_mapping`):
-
-```python
-from mujoco.sysid import TimeSeries, SignalType
-
-# For sensor/state observations:
-sensordata = TimeSeries.from_names(times, sensor_data, model) # all sensors
-sensordata = TimeSeries.from_names(times, data, model, names=["joint1_pos", "joint2_pos"])
-
-# Explicit type disambiguation (useful when sensor/state names overlap):
-sensordata = TimeSeries.from_names(times, data, model, names=[
- ("joint1_pos", SignalType.MjSensor), # sensor named "joint1_pos"
- ("joint1_qpos", SignalType.MjStateQPos), # joint state
-])
-
-# For control signals:
-control = TimeSeries.from_control_names(times, control_data, model) # all actuators
-control = TimeSeries.from_control_names(times, data, model, names=["motor1_ctrl"])
-
-# For custom/raw data:
-ts = TimeSeries.from_custom_map(times, data, ["signal1", "signal2"])
-```
-
-`signal_mapping` is a dict `{name: (SignalType, indices)}` that labels which
-columns of `data` correspond to which sensor/actuator.
-
-**`TimeSeries` factory methods:**
-
-| Constructor | Use case |
-|---|---|
-| `TimeSeries.from_names(times, data, model, names=None)` | Sensor/state data. If `names=None`, maps all model sensors. |
-| `TimeSeries.from_control_names(times, data, model, names=None)` | Control signals. If `names=None`, maps all actuators. |
-| `TimeSeries.from_custom_map(times, data, signals)` | Custom data with explicit signal definitions. |
-| `TimeSeries(times, data)` | Raw arrays, no signal mapping. |
-| `TimeSeries(times, data, signal_mapping)` | Named signals with explicit mapping. |
-
-**`TimeSeries` methods:** `resample(new_times=, target_dt=)`, `interpolate(t)`,
-`get(t)`, `save_to_disk(path)`, `load_from_disk(path)`, `dt_statistics()`,
-`remove_from_beginning(t)`, `slice_by_name(ts, names)`.
-
-Bundle a spec with one or more data sequences into a **`ModelSequences`**:
-
-```python
-from mujoco.sysid import ModelSequences, create_initial_state
-
-initial_state = create_initial_state(model, qpos, qvel, act)
-
-ms = ModelSequences(
- name="robot",
- spec=spec,
- sequence_name=["traj_1", "traj_2"], # or a single string
- initial_state=[initial_state_1, initial_state_2],
- control=[control_1, control_2],
- sensordata=[sensordata_1, sensordata_2],
-)
-```
-
-You can pass a single sequence (not wrapped in a list) and it will be
-auto-promoted.
-
-**Multiple `ModelSequences`:** Each `ModelSequences` carries its own `spec`, but
-the optimizer applies the **same parameter vector `θ`** to all of them. This
-enables joint optimization across different physical configurations. For example,
-you might have the same robot arm recorded with and without a known payload
-attached — two different specs (one has the payload body), two sets of recorded
-data, but the inertial parameters of the arm links are shared. The residuals
-from all `ModelSequences` are stacked and minimized jointly, giving a better-
-conditioned problem than fitting each dataset independently.
-
----
-
-### 3. Build the Residual Function
-
-The **residual** is the vector of differences between simulated sensor readings
-and recorded sensor data: `r(θ) = W(ȳ(θ) − y)`. Each element measures how
-much the simulation with parameters `θ` disagrees with reality for one sensor
-at one timestep. The optimizer's job is to find the `θ` that makes this vector
-as small as possible (in the least-squares sense).
-
-**`build_residual_fn`** captures data and configuration, returning a closure
-that the optimizer will call repeatedly:
-
-```python
-from mujoco.sysid import build_residual_fn
-
-residual_fn = build_residual_fn(
- models_sequences=[ms],
- # Optional overrides:
- modify_residual=..., # custom residual logic
- custom_rollout=..., # custom simulation
- sensor_weights=..., # per-sensor weighting
- enabled_observations=..., # subset of sensors to use
-)
-```
-
-#### How `residual_fn` works internally
-
-The returned `residual_fn(x, params)` accepts `x` as either a **1-D vector**
-(plain function evaluation) or a **2-D matrix** of shape `(n_params, n_fd)`
-(batched finite-difference Jacobian evaluation, where each column is a
-perturbed parameter vector). This is the key to parallelism.
-
-For each column `i` of `x`:
-
-1. `params.update_from_vector(x[:, i])` — writes the optimizer's current
- candidate values back into the `Parameter` objects so that each
- parameter's `.value` attribute reflects column `i` of `x`
-2. `model_i = apply_param_modifiers(params, spec)` — iterates over every
- non-frozen parameter and calls its `modifier(spec, param)` callback,
- then compiles the mutated spec into an `MjModel`
-3. Replicate `model_i` once per trajectory chunk (if you have `C` data
- sequences, you get `C` copies)
-
-This produces a flat list of `n_fd * C` models. All of them are rolled out in
-a **single call** to `mujoco.rollout.rollout`:
-
-```python
-datas = [mujoco.MjData(models[0]) for _ in range(n_threads)] # one per thread
-
-state, sensordata = mujoco.rollout.rollout(
- models, # n_fd * C models
- datas, # K thread-local scratch MjData objects
- initial_states, # n_fd * C initial states
- control, # n_fd * C control sequences
-)
-```
-
-MuJoCo's rollout engine distributes the `n_fd * C` independent rollouts across
-`K` threads using the `MjData` objects as thread-local scratch space (each
-thread gets its own `MjData` to avoid data races). **This is why the Jacobian
-computation is fast**: all `n_params + 1` perturbed rollouts (times `C`
-trajectory chunks) execute in one batched, multithreaded call.
-
-After rollout, residuals are computed per-trajectory (predicted vs. measured
-sensor data), then stacked and returned.
-
-#### Concrete example
-
-Suppose you have `p = 10` parameters and `C = 3` trajectory chunks:
-
-- **Function eval** (`x` is 1-D): `1 * 3 = 3` rollouts, distributed across
- threads.
-- **Jacobian eval** (`x` is `(10, 11)` — nominal + 10 perturbations): `11 * 3
- = 33` rollouts in one batched call. On a 16-core machine this is ~2x wall
- time of a single rollout.
-
-#### Three tiers of customization
-
-| Tier | What you provide | When to use |
-|---|---|---|
-| **Default** | Nothing extra (or `SignalTransform`) | Standard MuJoCo sensors, optional delays/gains |
-| **Custom rollout** | `custom_rollout=fn` | Non-standard simulation (e.g. task-space control) |
-| **Custom residual** | `modify_residual=fn` | State-based observations, exotic loss functions |
-
----
-
-### 4. SignalTransform (Declarative Residual Configuration)
-
-After simulation, the residual pipeline compares predicted sensor readings to
-recorded data. But real sensors aren't ideal — position encoders may lag by a
-few milliseconds, torque sensors may have an unknown scale factor, and velocity
-estimates may sit at a nonzero offset. These aren't physics parameters (you
-can't set "delay" on an `MjSpec`), so they need to be corrected *after* the
-rollout, before the residual is computed.
-
-**`SignalTransform`** lets you declare these corrections and which sensors to
-use, without writing a custom residual callback. Internally it:
-
-1. **Time-shifts** the predicted (or measured) signals by per-sensor delay
- parameters, resampling onto a common time grid.
-2. **Scales** sensor columns by gain parameters (`target="predicted"` scales
- the simulation output, `target="measured"` scales the recording — useful
- when the sensor's scale factor is unknown on either side).
-3. **Offsets** sensor columns by bias parameters.
-4. Computes the weighted difference and normalizes by RMS.
-
-Patterns use `fnmatch` syntax, so `"*_pos"` matches all sensors whose name
-ends in `_pos`:
-
-```python
-from mujoco.sysid import SignalTransform
-
-transform = SignalTransform()
-transform.delay("*_pos", params["delay_pos"]) # fnmatch pattern
-transform.delay("*_torque", params["delay_torque"])
-transform.gain("*_torque", params["torque_scale"], target="predicted")
-transform.bias("*_vel", params["vel_bias"])
-transform.enable_sensors(["joint1_pos", "joint2_pos", "joint1_torque"])
-transform.set_sensor_weights({"joint1_torque": 0.5})
-
-residual_fn = build_residual_fn(
- models_sequences=[ms],
- modify_residual=transform.apply, # drop-in replacement
-)
-```
-
-`SignalTransform.apply` has the same signature as `ModifyResidualFn`, so it
-plugs directly into `build_residual_fn`.
-
-#### What this replaces
-
-Without `SignalTransform`, you'd write the same logic by hand as a
-`modify_residual` callback using the low-level `signal_modifier` functions
-(Section 8):
-
-```python
-from mujoco.sysid._src import signal_modifier
-
-def modify_residual(params, predicted, measured, model, return_pred_all, **kw):
- # 1. Apply delays and resample onto a common time grid.
- min_d, max_d = -0.02, 0.05 # must track delay bounds yourself
- measured = signal_modifier.apply_delayed_ts_window(measured, predicted, min_d, max_d)
- sensor_delays = {"joint1_pos": params["delay_pos"].value[0], ...}
- predicted = signal_modifier.apply_resample_and_delay(
- predicted, measured.times, default_delay=0.0, sensor_delays=sensor_delays,
- )
- # 2. Apply gains and biases.
- predicted = signal_modifier.apply_gain(predicted, "joint1_torque", params["torque_scale"])
- predicted = signal_modifier.apply_bias(predicted, "joint1_vel", params["vel_bias"])
- # 3. Compute residual.
- diff = signal_modifier.weighted_diff(predicted.data, measured.data, model, weights)
- diff = signal_modifier.normalize_residual(diff, measured.data)
- return diff, predicted, measured
-```
-
-`SignalTransform` does all of this — including tracking delay bounds, expanding
-fnmatch patterns to sensor names, and handling the windowing/resampling
-bookkeeping — from a few declarative lines.
-
----
-
-### 5. Optimize
-
-**`optimize`** runs box-constrained Gauss-Newton least-squares on the residual
-function:
-
-```python
-from mujoco.sysid import optimize
-
-opt_params, opt_result = optimize(
- initial_params=params,
- residual_fn=residual_fn,
- optimizer="mujoco", # "mujoco", "scipy", or "scipy_parallel_fd"
- max_iters=200,
-)
-```
-
-#### Optimizer backends
-
-| Backend | Jacobian | Description |
-|---|---|---|
-| `"mujoco"` (recommended) | Parallel FD, batched | `mujoco.minimize.least_squares`. Calls `residual_fn(x)` with `x` as a 2-D matrix `(n_params, n_params+1)` — the nominal point plus one perturbation per parameter — so the entire Jacobian is computed in a single batched, multithreaded rollout call. |
-| `"scipy"` | Sequential 2-point FD | `scipy.optimize.least_squares`. Computes the Jacobian column-by-column (sequential). Slower for problems with many parameters. |
-| `"scipy_parallel_fd"` | Parallel FD via MuJoCo, scipy outer loop | Scipy's trust-region solver but with `mujoco.minimize.jacobian_fd` for the Jacobian. Gives scipy's convergence control with MuJoCo's batched FD speed. |
-
-All three backends solve box-constrained nonlinear least-squares using the
-Gauss-Newton Hessian approximation `H ≈ JᵀJ`. They differ in how the
-finite-difference Jacobian is computed and how the trust-region step is
-handled: `"mujoco"` uses projected Gauss-Newton steps, while `"scipy"` and
-`"scipy_parallel_fd"` use scipy's trust-region reflective algorithm (`trf`).
-
-#### Return value
-
-Returns `(opt_params, OptimizeResult)` where `opt_params` is a deep copy of
-the input `ParameterDict` with `.value` set to the solution, and
-`OptimizeResult` contains:
-- `.x` — the solution vector
-- `.jac` — the Jacobian at the solution (used for confidence intervals)
-- `.grad` — the gradient at the solution
-- `.extras` — (**mujoco backend only**) dict with `"objective"` (cost per
- iteration) and `"candidate"` (parameter vector per iteration), when verbose
-
----
-
-### 6. Save Results and Report
-
-**`save_results`** writes everything to disk:
-
-```python
-from mujoco.sysid import save_results
-
-save_results(
- experiment_results_folder="results/exp01",
- models_sequences=[ms],
- initial_params=params,
- opt_params=opt_params,
- opt_result=opt_result,
- residual_fn=residual_fn,
-)
-```
-
-This creates:
-- `params_x_0.yaml` — initial parameter values
-- `params_x_hat.yaml` — optimized parameter values
-- `results.pkl` — full `OptimizeResult`
-- `confidence.pkl` — parameter covariance matrix `Σ_θ = σ²_r H⁻¹` and
- per-parameter confidence intervals, computed from the eigendecomposition of
- `H = JᵀJ` at the solution. Parameters in near-null-space directions of `H`
- receive infinite confidence intervals, making identifiability issues
- immediately visible.
-- `{model_name}.xml` — identified MuJoCo XML for each model
-
-**`default_report`** generates an HTML report with sensor comparisons, parameter
-tables, and videos:
-
-```python
-from mujoco.sysid import default_report
-
-default_report(
- models_sequences=[ms],
- initial_params=params,
- opt_params=opt_params,
- opt_result=opt_result,
- residual_fn=residual_fn,
- save_dir="results/exp01",
-)
-```
-
----
-
-### 7. Model Modification
-
-`apply_param_modifiers` is the default `build_model` implementation — it
-iterates over all non-frozen parameters, calls each one's `modifier` callback
-on the spec, and compiles. Most users never need to call it directly; it runs
-automatically inside the residual pipeline.
-
-The remaining functions are useful when writing a **custom `build_model`**
-(e.g. the box case study manually mutates the spec instead of using modifier
-callbacks):
-
-| Function | Description |
-|---|---|
-| `apply_param_modifiers(params, spec)` | Run all modifier callbacks, return compiled `MjModel` |
-| `apply_param_modifiers_spec(params, spec)` | Run all modifier callbacks, return the `MjSpec` |
-| `apply_pgain(spec, name, value)` | Set proportional gain on a position actuator |
-| `apply_dgain(spec, name, value)` | Set derivative gain on a position actuator |
-| `apply_pdgain(spec, name, value)` | Set both P and D gains |
-| `apply_body_inertia(spec, name, param)` | Apply Mass / MassIpos / Pseudo inertia |
-| `body_inertia_param(spec, model, name, ...)` | Create a `Parameter` for body inertia |
-| `remove_visuals(spec)` | Strip textures, materials, and visual-only geoms |
-
----
-
-### 8. Signal Modification (Power-User API)
-
-Low-level functions used internally by `SignalTransform` and the default
-residual pipeline. Useful when writing a custom `modify_residual`:
-
-| Function | Description |
-|---|---|
-| `get_sensor_indices(model, name)` | Column indices for a named sensor |
-| `apply_gain(ts, name, param)` | Multiply sensor columns by `param.value` |
-| `apply_bias(ts, name, param)` | Add `param.value` to sensor columns |
-| `apply_delay(ts, name, param)` | Time-shift sensor columns |
-| `apply_delayed_ts_window(ts, ts_ref, min_d, max_d)` | Crop `ts` to the valid time window |
-| `apply_resample_and_delay(ts, times, default_delay, ...)` | Resample with per-sensor delays |
-| `weighted_diff(pred, meas, model, weights)` | `measured - predicted`, optionally weighted |
-| `normalize_residual(residual, measured)` | Divide by column-wise RMS of measured data |
-
----
-
-### 9. Additional Utilities
-
-| Function / Class | Module | Description |
-|---|---|---|
-| `create_initial_state(model, qpos, qvel, act)` | trajectory | Pack qpos/qvel/act into a flat state vector |
-| `SystemTrajectory` | trajectory | Frozen dataclass holding a single rollout (model, control, sensordata, state) |
-| `sysid_rollout(models, datas, control, initial_states)` | trajectory | Parallel MuJoCo rollout returning `SystemTrajectory` list |
-| `render_rollout(model, data, state, framerate)` | plotting | Render state trajectories to pixel frames |
-| `calculate_intervals(residuals, J, alpha)` | optimize | Confidence intervals from Jacobian at the solution |
-| `sweep_parameter(params, name, values, residual_fn)` | optimize | 1-D parameter sweep returning cost curve |
-| `plot_sensor_comparison(model, ...)` | plotting | Matplotlib overlay of predicted vs. measured sensors |
-| `SignalType` | timeseries | Enum: `MjSensor`, `CustomObs`, `MjStateQPos`, `MjStateQVel`, `MjStateAct`, `MjCtrl` |
-
----
-
-## Type Aliases
-
-```python
-ModifyResidualFn = Callable[
- ..., tuple[np.ndarray, TimeSeries, TimeSeries]
-]
-# (params, sensordata_predicted, sensordata_measured, model, return_pred_all, state=..., sensor_weights=...)
-# Returns (residual_array, pred_timeseries, measured_timeseries)
-
-CustomRolloutFn = Callable[..., Sequence[SystemTrajectory]]
-# (models, datas, control_signal, initial_states, param_dicts, ...)
-# Returns list of SystemTrajectory
-
-BuildModelFn = Callable[[ParameterDict, MjSpec], MjModel]
-# Default: apply_param_modifiers
-```
-
----
-
-## Skeleton Case Study
-
-Pseudocode showing the five-stage pipeline. Replace the data-loading step with
-your own hardware logs or simulation data. For a complete runnable example, see
-`case_studies/box/`.
+## Example
```python
import mujoco
-import numpy as np
+from mujoco import sysid
-from mujoco.sysid import (
- Parameter,
- ParameterDict,
- TimeSeries,
- ModelSequences,
- build_residual_fn,
- create_initial_state,
- optimize,
- save_results,
-)
-
-# 1. Load model.
+# 1. Load model and define parameters.
spec = mujoco.MjSpec.from_file("robot.xml")
model = spec.compile()
-# 2. Define parameters with modifier callbacks.
-params = ParameterDict()
-params.add(Parameter(
- "link1_mass",
- nominal=2.0,
- min_value=0.5,
- max_value=5.0,
- modifier=lambda spec, p: setattr(spec.body("link1"), "mass", p.value[0]),
-))
+def set_link1_mass(spec, p):
+ spec.body("link1").mass = p.value[0]
-# 3. Package recorded data.
-# times: (N,) timestamps
-# ctrl_array: (N, model.nu) control inputs
-# sensor_array: (N, model.nsensordata) recorded sensor readings
-# qpos_0, qvel_0: initial joint positions and velocities
-control = TimeSeries.from_control_names(times, ctrl_array, model)
-sensordata = TimeSeries.from_names(times, sensor_array, model)
-initial_state = create_initial_state(model, qpos_0, qvel_0)
+params = sysid.ParameterDict()
+params.add(sysid.Parameter(
+ "link1_mass", nominal=2.0, min_value=0.5, max_value=5.0,
+ modifier=set_link1_mass))
-ms = ModelSequences(
- name="robot",
- spec=spec,
- sequence_name="traj_1",
- initial_state=initial_state,
- control=control,
- sensordata=sensordata,
-)
+# 2. Load and package measured data.
+# arrays assumed to be in MuJoCo order, otherwise pass names argument
+control = sysid.TimeSeries.from_control_names(times, ctrl_array, model)
+measureddata = sysid.TimeSeries.from_names(times, measurement_array, model)
+initial_state = sysid.create_initial_state(model, qpos_0, qvel_0)
+ms = sysid.ModelSequences("robot", spec, "traj_1", initial_state, control, measureddata)
-# 4. Build the residual function and optimize.
-# models_sequences is a list because you can jointly optimize across
-# multiple ModelSequences with different specs (see Section 2).
-residual_fn = build_residual_fn(models_sequences=[ms])
-opt_params, opt_result = optimize(
- initial_params=params,
- residual_fn=residual_fn,
- optimizer="mujoco",
-)
-
-# 5. Inspect results.
-print(opt_params)
-save_results("results/", [ms], params, opt_params, opt_result, residual_fn)
+# 3. Build residual, optimize, save.
+residual_fn = sysid.build_residual_fn(models_sequences=[ms])
+opt_params, opt_result = sysid.optimize(initial_params=params, residual_fn=residual_fn)
+sysid.save_results("results/", [ms], params, opt_params, opt_result, residual_fn)
```
+
+`default_report` generates an interactive HTML report with videos, measurement comparisons,
+parameter tables, and confidence intervals.
diff --git a/python/mujoco/sysid/__init__.py b/python/mujoco/sysid/__init__.py
index a5e865df..93b8f8ea 100644
--- a/python/mujoco/sysid/__init__.py
+++ b/python/mujoco/sysid/__init__.py
@@ -12,49 +12,46 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""Practical system identification for MuJoCo."""
+"""System identification toolbox."""
-from mujoco.sysid._src import model_modifier
-from mujoco.sysid._src import parameter
-from mujoco.sysid._src import plotting
-from mujoco.sysid._src import signal_modifier
-from mujoco.sysid._src.io import save_results
-from mujoco.sysid._src.model_modifier import apply_body_inertia
-from mujoco.sysid._src.model_modifier import apply_dgain
-from mujoco.sysid._src.model_modifier import apply_param_modifiers
-from mujoco.sysid._src.model_modifier import apply_param_modifiers_spec
-from mujoco.sysid._src.model_modifier import apply_pdgain
-from mujoco.sysid._src.model_modifier import apply_pgain
-from mujoco.sysid._src.model_modifier import body_inertia_param
-from mujoco.sysid._src.model_modifier import remove_visuals
-from mujoco.sysid._src.optimize import calculate_intervals
-from mujoco.sysid._src.optimize import optimize
-from mujoco.sysid._src.parameter import InertiaType
-from mujoco.sysid._src.parameter import Parameter
-from mujoco.sysid._src.parameter import ParameterDict
-from mujoco.sysid._src.plotting import plot_sensor_comparison
-from mujoco.sysid._src.plotting import render_rollout
-from mujoco.sysid._src.residual import build_residual_fn
-from mujoco.sysid._src.residual import BuildModelFn
-from mujoco.sysid._src.residual import construct_ts_from_defaults
-from mujoco.sysid._src.residual import CustomRolloutFn
-from mujoco.sysid._src.residual import model_residual
-from mujoco.sysid._src.residual import ModifyResidualFn
-from mujoco.sysid._src.residual import residual
-from mujoco.sysid._src.signal_modifier import apply_bias
-from mujoco.sysid._src.signal_modifier import apply_delay
-from mujoco.sysid._src.signal_modifier import apply_delayed_ts_window
-from mujoco.sysid._src.signal_modifier import apply_gain
-from mujoco.sysid._src.signal_modifier import apply_resample_and_delay
-from mujoco.sysid._src.signal_modifier import get_sensor_indices
-from mujoco.sysid._src.signal_modifier import normalize_residual
-from mujoco.sysid._src.signal_modifier import weighted_diff
-from mujoco.sysid._src.signal_transform import SignalTransform
-from mujoco.sysid._src.timeseries import SignalType
-from mujoco.sysid._src.timeseries import TimeSeries
-from mujoco.sysid._src.trajectory import create_initial_state
-from mujoco.sysid._src.trajectory import ModelSequences
-from mujoco.sysid._src.trajectory import sysid_rollout
-from mujoco.sysid._src.trajectory import SystemTrajectory
-from mujoco.sysid.report.defaults import default_report
-from mujoco.sysid.report.defaults import default_report_matplotlib
+from mujoco.sysid._src import model_modifier as model_modifier
+from mujoco.sysid._src import parameter as parameter
+from mujoco.sysid._src import signal_modifier as signal_modifier
+from mujoco.sysid._src.io import save_results as save_results
+from mujoco.sysid._src.model_modifier import apply_body_inertia as apply_body_inertia
+from mujoco.sysid._src.model_modifier import apply_dgain as apply_dgain
+from mujoco.sysid._src.model_modifier import apply_param_modifiers as apply_param_modifiers
+from mujoco.sysid._src.model_modifier import apply_param_modifiers_spec as apply_param_modifiers_spec
+from mujoco.sysid._src.model_modifier import apply_pdgain as apply_pdgain
+from mujoco.sysid._src.model_modifier import apply_pgain as apply_pgain
+from mujoco.sysid._src.model_modifier import body_inertia_param as body_inertia_param
+from mujoco.sysid._src.model_modifier import remove_visuals as remove_visuals
+from mujoco.sysid._src.optimize import calculate_intervals as calculate_intervals
+from mujoco.sysid._src.optimize import optimize as optimize
+from mujoco.sysid._src.parameter import InertiaType as InertiaType
+from mujoco.sysid._src.parameter import Parameter as Parameter
+from mujoco.sysid._src.parameter import ParameterDict as ParameterDict
+from mujoco.sysid._src.plotting import render_rollout as render_rollout
+from mujoco.sysid._src.residual import build_residual_fn as build_residual_fn
+from mujoco.sysid._src.residual import BuildModelFn as BuildModelFn
+from mujoco.sysid._src.residual import construct_ts_from_defaults as construct_ts_from_defaults
+from mujoco.sysid._src.residual import CustomRolloutFn as CustomRolloutFn
+from mujoco.sysid._src.residual import model_residual as model_residual
+from mujoco.sysid._src.residual import ModifyResidualFn as ModifyResidualFn
+from mujoco.sysid._src.residual import residual as residual
+from mujoco.sysid._src.signal_modifier import apply_bias as apply_bias
+from mujoco.sysid._src.signal_modifier import apply_delay as apply_delay
+from mujoco.sysid._src.signal_modifier import apply_delayed_ts_window as apply_delayed_ts_window
+from mujoco.sysid._src.signal_modifier import apply_gain as apply_gain
+from mujoco.sysid._src.signal_modifier import apply_resample_and_delay as apply_resample_and_delay
+from mujoco.sysid._src.signal_modifier import get_sensor_indices as get_sensor_indices
+from mujoco.sysid._src.signal_modifier import normalize_residual as normalize_residual
+from mujoco.sysid._src.signal_modifier import weighted_diff as weighted_diff
+from mujoco.sysid._src.signal_transform import SignalTransform as SignalTransform
+from mujoco.sysid._src.timeseries import SignalType as SignalType
+from mujoco.sysid._src.timeseries import TimeSeries as TimeSeries
+from mujoco.sysid._src.trajectory import create_initial_state as create_initial_state
+from mujoco.sysid._src.trajectory import ModelSequences as ModelSequences
+from mujoco.sysid._src.trajectory import sysid_rollout as sysid_rollout
+from mujoco.sysid._src.trajectory import SystemTrajectory as SystemTrajectory
+from mujoco.sysid.report.defaults import default_report as default_report
diff --git a/python/mujoco/sysid/_src/io.py b/python/mujoco/sysid/_src/io.py
index aa1fb721..cc431499 100644
--- a/python/mujoco/sysid/_src/io.py
+++ b/python/mujoco/sysid/_src/io.py
@@ -35,7 +35,16 @@ def save_results(
opt_result: scipy_optimize.OptimizeResult,
residual_fn,
):
- """Save optimization results and confidence intervals to disk."""
+ """Save optimization results and confidence intervals to disk.
+
+ Args:
+ experiment_results_folder: Directory where results are written.
+ models_sequences: Model/sequence groups; identified XMLs are saved here.
+ initial_params: Parameters before optimization.
+ opt_params: Parameters after optimization.
+ opt_result: Scipy OptimizeResult from the optimizer.
+ residual_fn: Residual function used to compute confidence intervals.
+ """
experiment_results_folder = pathlib.Path(experiment_results_folder)
if not experiment_results_folder.exists():
experiment_results_folder.mkdir(parents=True, exist_ok=True)
@@ -70,13 +79,3 @@ def save_results(
model_sequences.spec.to_file(
(experiment_results_folder / f"{model_sequences.name}.xml").as_posix()
)
-
- # Log nominal compared to initial.
- x0 = initial_params.as_vector()
- x_nominal = initial_params.as_nominal_vector()
- logging.info(
- "Initial Parameters\n%s",
- initial_params.compare_parameters(
- x0, opt_result.x, measured_params=x_nominal
- ),
- )
diff --git a/python/mujoco/sysid/_src/model_modifier.py b/python/mujoco/sysid/_src/model_modifier.py
index 3dc6287b..1f266c5e 100644
--- a/python/mujoco/sysid/_src/model_modifier.py
+++ b/python/mujoco/sysid/_src/model_modifier.py
@@ -111,7 +111,12 @@ def is_position_actuator(actuator) -> bool:
def get_actuator_pd_gains(
model: mujoco.MjModel, actuator_name: str
) -> tuple[float, float]:
- """Return the (P, D) gains of a position actuator."""
+ """Return the (P, D) gains of a position actuator.
+
+ Args:
+ model: MuJoCo model.
+ actuator_name: Name of the actuator.
+ """
actuator_id = mujoco.mj_name2id(
model, mujoco.mjtObj.mjOBJ_ACTUATOR.value, actuator_name
)
@@ -128,7 +133,13 @@ def apply_pgain(
actuator_name: str,
value: float | np.ndarray,
) -> mujoco.MjSpec:
- """Set the proportional gain for a position actuator."""
+ """Set the proportional gain for a position actuator.
+
+ Args:
+ spec: MuJoCo model specification.
+ actuator_name: Name of the actuator.
+ value: Proportional gain value.
+ """
# TODO(b/0): assert scalar
actuator = _get_obj_or_raise(spec, "actuator", actuator_name)
assert isinstance(actuator, mujoco.MjsActuator)
@@ -144,7 +155,13 @@ def apply_dgain(
actuator_name: str,
value: float | np.ndarray,
) -> mujoco.MjSpec:
- """Set the derivative gain for a position actuator."""
+ """Set the derivative gain for a position actuator.
+
+ Args:
+ spec: MuJoCo model specification.
+ actuator_name: Name of the actuator.
+ value: Derivative gain value.
+ """
# TODO(b/0): assert scalar
actuator = _get_obj_or_raise(spec, "actuator", actuator_name)
assert isinstance(actuator, mujoco.MjsActuator)
@@ -159,7 +176,13 @@ def apply_pdgain(
actuator_name: str,
value: np.ndarray,
) -> mujoco.MjSpec:
- """Set both proportional and derivative gains for a position actuator."""
+ """Set both proportional and derivative gains for a position actuator.
+
+ Args:
+ spec: MuJoCo model specification.
+ actuator_name: Name of the actuator.
+ value: 2-element array ``[P_gain, D_gain]``.
+ """
if value.size != 2:
raise ValueError(f"pdgain must be a 2-element array, got {value.size}.")
apply_pgain(spec, actuator_name, value[0])
@@ -174,7 +197,16 @@ def apply_body_mass_ipos(
ipos: np.ndarray | None = None,
rot_inertia_scale: bool = False,
) -> mujoco.MjSpec:
- """Apply mass and center-of-mass position to a body."""
+ """Apply mass and center-of-mass position to a body.
+
+ Args:
+ spec: MuJoCo model specification.
+ body_name: Name of the body.
+ mass: Optional new mass value.
+ ipos: Optional new center-of-mass position.
+ rot_inertia_scale: If True, scale rotational inertia proportionally to
+ mass change.
+ """
# TODO(b/0): assert mass and ipos shapes
body = _infer_inertial(spec, body_name)
mass_original = body.mass
@@ -200,7 +232,18 @@ def scale_body_inertia(
def pi_from_theta(theta: np.ndarray) -> np.ndarray:
- """Convert base parameters θ to inertial parameters π."""
+ """Convert base parameters θ to inertial parameters π.
+
+ Args:
+ theta: 10-D array [alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3] where:
+ alpha: Scale parameter (log of U[3,3])
+ [d1, d2, d3]: Log of diagonal elements
+ [s12, s23, s13]: Shear parameters from upper triangle
+ [t1, t2, t3]: Translation parameters from last column
+
+ Returns:
+ 10-D array π = [m, hx, hy, hz, Ixx, Iyy, Izz, Ixy, Iyz, Ixz].
+ """
alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3 = theta
exp_alpha = np.exp(alpha)
exp_d1 = np.exp(d1)
@@ -361,7 +404,13 @@ def apply_body_theta_inertia(
body_name: str,
theta: np.ndarray,
) -> mujoco.MjSpec:
- """Apply base-parameter inertia θ to a body in the spec."""
+ """Apply base-parameter inertia θ to a body in the spec.
+
+ Args:
+ spec: MuJoCo model specification.
+ body_name: Name of the body.
+ theta: 10-element array [alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3].
+ """
if theta.size != 10:
raise ValueError(f"theta must be a 10-element array, got {theta.size}.")
pi = pi_from_theta(theta)
@@ -392,7 +441,13 @@ def apply_body_theta_inertia(
def apply_body_inertia(spec: mujoco.MjSpec, name: str, param: Parameter):
- """Apply inertia parameters to a body based on the parameter type."""
+ """Apply inertia parameters to a body based on the parameter type.
+
+ Args:
+ spec: MuJoCo model specification.
+ name: Name of the body.
+ param: Parameter with an ``inertia_type`` attribute.
+ """
if not hasattr(param, "inertia_type"):
raise ValueError(
f"Parameter {param.name} does not have inertia_type attribute."
diff --git a/python/mujoco/sysid/_src/optimize.py b/python/mujoco/sysid/_src/optimize.py
index c8c7b091..65227355 100644
--- a/python/mujoco/sysid/_src/optimize.py
+++ b/python/mujoco/sysid/_src/optimize.py
@@ -44,7 +44,7 @@ def _scipy_least_squares(
jac_arg: str | Callable[..., Any]
if use_mujoco_jac:
- # This is the default step sized for finite difference used in
+ # This is the default step size for finite difference used in
# scipy's least_squares and mujoco's minimize finite difference
# https://github.com/scipy/scipy/blob/91e18f3bd355477b
# 8b7747ec82d70ac98ffd2422/scipy/optimize/_numdiff.py#L404
@@ -143,6 +143,7 @@ def optimize(
initial_params: parameter.ParameterDict,
residual_fn: Callable[..., Any],
optimizer: Literal["scipy", "mujoco", "scipy_parallel_fd"] = "mujoco",
+ verbose: bool = True,
**optimizer_kwargs,
) -> tuple[parameter.ParameterDict, scipy_optimize.OptimizeResult]:
"""Run nonlinear least-squares optimization on the residual.
@@ -153,11 +154,12 @@ def optimize(
returned by :func:`build_residual_fn`.
optimizer: Backend — ``"mujoco"`` (default), ``"scipy"``, or
``"scipy_parallel_fd"`` (scipy with MuJoCo finite-difference Jacobian).
+ verbose: If True, log parameter comparison table after optimization.
**optimizer_kwargs: Forwarded to the backend (e.g. ``max_iters``,
``verbose``, ``loss``).
Returns:
- ``(opt_params, opt_result)`` — the optimised ParameterDict and a
+ ``(opt_params, opt_result)`` — the optimized ParameterDict and a
``scipy.optimize.OptimizeResult`` with at least ``x``, ``jac``, ``grad``.
"""
x0 = initial_params.as_vector()
@@ -187,6 +189,16 @@ def optimize(
opt_params.update_from_vector(opt_result.x)
+ if verbose:
+ logging.info(
+ "\n%s",
+ opt_params.compare_parameters(
+ initial_params.as_vector(),
+ opt_params.as_vector(),
+ measured_params=initial_params.as_nominal_vector(),
+ ),
+ )
+
return opt_params, opt_result
@@ -197,7 +209,20 @@ def calculate_intervals(
lambda_zero_thresh=1e-15,
v_zero_thresh=1e-8,
):
- """Calculate confidence intervals from the Jacobian at the optimum."""
+ """Calculate confidence intervals from the Jacobian at the optimum.
+
+ Args:
+ residuals_star: List of residual arrays at the optimum.
+ J: Jacobian matrix at the optimum, shape ``(n_residuals, n_params)``.
+ alpha: Significance level for the confidence intervals.
+ lambda_zero_thresh: Threshold below which eigenvalues are treated as zero.
+ v_zero_thresh: Threshold below which eigenvector elements are treated as
+ zero.
+
+ Returns:
+ ``(Sigma_X, intervals)`` — the parameter covariance matrix and the
+ half-width confidence intervals for each parameter.
+ """
if J is None or J.size == 0:
return np.empty((0, 0)), np.empty((0,))
diff --git a/python/mujoco/sysid/_src/parameter.py b/python/mujoco/sysid/_src/parameter.py
index fcb47d66..b5e7d46a 100644
--- a/python/mujoco/sysid/_src/parameter.py
+++ b/python/mujoco/sysid/_src/parameter.py
@@ -105,6 +105,11 @@ class Parameter:
return self.nominal.flatten()
def update_from_vector(self, vector: np.ndarray) -> None:
+ """Update the current value from a flat vector.
+
+ Args:
+ vector: Flat array of length ``self.size``.
+ """
vector_array = np.atleast_1d(vector)
if len(vector_array) != self.size:
raise ValueError(
@@ -125,7 +130,11 @@ class Parameter:
self.value = self.nominal.copy()
def sample(self, rng: np.random.Generator | None = None) -> np.ndarray:
- """Sample a random value uniformly within bounds."""
+ """Sample a random value uniformly within bounds.
+
+ Args:
+ rng: Optional numpy random generator. Uses default if None.
+ """
if rng is None:
rng = np.random.default_rng()
return rng.uniform(self.min_value.flatten(), self.max_value.flatten())
@@ -197,7 +206,7 @@ class ParameterDict:
"""An ordered collection of :class:`Parameter` objects.
Behaves like a ``dict[str, Parameter]`` with convenience methods for
- vectorised access (``as_vector`` / ``update_from_vector``), serialisation,
+ vectorized access (``as_vector`` / ``update_from_vector``), serialization,
and tabular comparison of parameter estimates.
Frozen parameters are silently skipped by vector/bounds methods so that the
@@ -271,7 +280,12 @@ class ParameterDict:
return np.concatenate(vectors) if vectors else np.array([])
def update_from_vector(self, vector: np.ndarray) -> None:
- """Update all non-frozen parameters from a flat vector."""
+ """Update all non-frozen parameters from a flat vector.
+
+ Args:
+ vector: Flat array whose length equals the total size of non-frozen
+ parameters.
+ """
start = 0
for param in self.parameters.values():
if not param.frozen:
@@ -333,14 +347,22 @@ class ParameterDict:
param.reset()
def sample(self, rng: np.random.Generator | None = None) -> np.ndarray:
- """Sample parameter values within bounds for non-frozen parameters."""
+ """Sample parameter values within bounds for non-frozen parameters.
+
+ Args:
+ rng: Optional numpy random generator. Uses default if None.
+ """
if rng is None:
rng = np.random.default_rng()
lower_bounds, upper_bounds = self.get_bounds()
return rng.uniform(lower_bounds, upper_bounds)
def randomize(self, rng: np.random.Generator | None = None) -> None:
- """Randomize parameter values for non-frozen parameters."""
+ """Randomize parameter values for non-frozen parameters.
+
+ Args:
+ rng: Optional numpy random generator. Uses default if None.
+ """
for param in self.parameters.values():
if not param.frozen:
param.value = param.sample(rng)
diff --git a/python/mujoco/sysid/_src/plotting.py b/python/mujoco/sysid/_src/plotting.py
index a21a8722..484a3a3b 100644
--- a/python/mujoco/sysid/_src/plotting.py
+++ b/python/mujoco/sysid/_src/plotting.py
@@ -19,627 +19,10 @@ from __future__ import annotations
from collections.abc import Sequence
-from matplotlib.lines import Line2D
-import matplotlib.pyplot as plt
import mujoco
-from mujoco.sysid._src import parameter
import numpy as np
-def plot_sensor_comparison(
- model: mujoco.MjModel,
- predicted_times: np.ndarray | None = None,
- predicted_data: np.ndarray | None = None,
- real_data: np.ndarray | None = None,
- real_times: np.ndarray | None = None,
- preid_data: np.ndarray | None = None,
- preid_times: np.ndarray | None = None,
- commanded_data: np.ndarray | None = None,
- commanded_times: np.ndarray | None = None,
- size_factor: float = 1.0,
- title_prefix: str = "",
- sensor_ids: list[int] | None = None,
-):
- """Plots sensor trajectories from simulation and real data.
-
- Args:
- model: The model object providing sensor information.
- predicted_times: Optional 1D array of timestamps corresponding to
- simulation data.
- predicted_data: Optional 2D array of simulation sensor data with shape
- (num_timesteps, sensor_data_dimension).
- real_data: Optional 2D array of real sensor data with the same shape as
- predicted_data.
- real_times: A 1D array of timestamps corresponding to real data. If None
- and real_data is provided, the first available timestamp array is used.
- preid_data: Optional 2D array of pre-identification sensor data.
- preid_times: A 1D array of timestamps for pre-identification data.
- commanded_data: Optional 2D array of commanded sensor data.
- commanded_times: A 1D array of timestamps for commanded data.
- size_factor: A scaling factor for the figure size.
- title_prefix: Optional prefix for subplot titles.
- sensor_ids: Optional list of sensor indices to plot.
- """
- # Define a more appealing color palette
- predicted_color = "#1f77b4" # Steel blue
- real_color = "#ff7f0e" # Safety orange
- preid_color = "#2ca02c" # Forest green
- commanded_color = "#9467bd" # Purple
-
- # Determine the reference time array to use
- reference_times = None
- if predicted_times is not None:
- reference_times = predicted_times
- elif real_times is not None:
- reference_times = real_times
- elif preid_times is not None:
- reference_times = preid_times
- elif commanded_times is not None:
- reference_times = commanded_times
- else:
- raise ValueError("At least one time array must be provided")
-
- # Set times for data sources that don't have their own time arrays
- if real_data is not None and real_times is None:
- real_times = reference_times
- if preid_data is not None and preid_times is None:
- preid_times = reference_times
- if commanded_data is not None and commanded_times is None:
- commanded_times = reference_times
- if predicted_data is not None and predicted_times is None:
- predicted_times = reference_times
-
- if sensor_ids is None:
- sensor_ids = list(range(model.nsensor))
- assert predicted_data is not None
- n_plots = predicted_data.shape[1]
-
- fig, axes = plt.subplots(
- n_plots,
- 1,
- figsize=(10 * size_factor, 2.5 * n_plots * size_factor),
- sharex=True,
- )
- if n_plots == 1:
- axes = [axes]
- axes = list(axes) # pyright: ignore[reportArgumentType]
-
- # Set an overall title for the figure.
- fig.suptitle(title_prefix + " Sensors", fontsize=14) # , y=1.02)
-
- # Loop over each sensor.
- plot_i = 0
- sensor_dim = 1
- j = 0
- dim_str = ""
- for sensor_id in sensor_ids:
- sensor = model.sensor(sensor_id)
- sensor_name = sensor.name
- sensor_dim = int(sensor.dim[0])
- sensor_addr = int(sensor.adr[0])
-
- for j in range(sensor_dim):
- ax = axes[plot_i]
- plot_i += 1
- dim_str = "" if sensor_dim == 1 else f" {j}"
- if predicted_data is not None:
- assert predicted_times is not None
- predicted_signal = predicted_data[
- :, sensor_addr : sensor_addr + sensor_dim
- ]
- ax.plot(
- predicted_times,
- predicted_signal[:, j],
- lw=2,
- color=predicted_color,
- alpha=0.8,
- label="Sim" + dim_str,
- )
- if real_data is not None:
- assert real_times is not None
- real_signal = real_data[:, sensor_addr : sensor_addr + sensor_dim]
- ax.plot(
- real_times,
- real_signal[:, j],
- lw=2,
- color=real_color,
- linestyle="--",
- alpha=0.7,
- label="Real" + dim_str,
- )
- if preid_data is not None:
- assert preid_times is not None
- preid_signal = preid_data[:, sensor_addr : sensor_addr + sensor_dim]
- ax.plot(
- preid_times,
- preid_signal[:, j],
- lw=2,
- color=preid_color,
- linestyle=":",
- alpha=0.6,
- label="Pre-ID" + dim_str,
- )
- if commanded_data is not None:
- assert commanded_times is not None
- commanded_signal = commanded_data[
- :, sensor_addr : sensor_addr + sensor_dim
- ]
- ax.plot(
- commanded_times,
- commanded_signal[:, j],
- lw=2,
- color=commanded_color,
- linestyle="-.",
- alpha=0.6,
- label="Commanded" + dim_str,
- )
- # Place the sensor name in a white box in the top-left corner.
- ax.text(
- 0.02,
- 0.9,
- sensor_name + dim_str,
- transform=ax.transAxes,
- fontsize=10,
- weight="bold",
- verticalalignment="top",
- horizontalalignment="left",
- bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
- )
-
- # Enable a dashed grid.
- ax.grid(True, linestyle="--", alpha=0.7)
-
- # Loop over "extra" sensors from the user
- for _ in range(plot_i, n_plots):
- sensor_name = "user_sensor"
- dim_str = "" if sensor_dim == 1 else f" {j}"
- ax = axes[plot_i]
- plot_i += 1
- if predicted_data is not None:
- assert predicted_times is not None
- predicted_signal = predicted_data[:, plot_i - 1]
- ax.plot(
- predicted_times,
- predicted_signal,
- lw=2,
- color=predicted_color,
- alpha=0.8,
- label="Sim",
- )
- if real_data is not None:
- assert real_times is not None
- real_signal = real_data[:, plot_i - 1]
- ax.plot(
- real_times,
- real_signal,
- lw=2,
- color=real_color,
- linestyle="--",
- alpha=0.7,
- label="Real",
- )
- if preid_data is not None:
- assert preid_times is not None
- preid_signal = preid_data[:, plot_i - 1]
- ax.plot(
- preid_times,
- preid_signal,
- lw=2,
- color=preid_color,
- linestyle=":",
- alpha=0.6,
- label="Pre-ID",
- )
- if commanded_data is not None:
- assert commanded_times is not None
- commanded_signal = commanded_data[:, plot_i - 1]
- ax.plot(
- commanded_times,
- commanded_signal,
- lw=2,
- color=commanded_color,
- linestyle="-.",
- alpha=0.6,
- label="Commanded",
- )
- # Place the sensor name in a white box in the top-left corner.
- ax.text(
- 0.02,
- 0.9,
- sensor_name + dim_str,
- transform=ax.transAxes,
- fontsize=10,
- weight="bold",
- verticalalignment="top",
- horizontalalignment="left",
- bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
- )
-
- # Enable a dashed grid.
- ax.grid(True, linestyle="--", alpha=0.7)
-
- # Add a unified, figure-level legend if any data is provided.
- legend_handles = []
- if predicted_data is not None:
- legend_handles.append(
- Line2D([0], [0], color=predicted_color, lw=2, label="Simulation")
- )
- if real_data is not None:
- legend_handles.append(
- Line2D([0], [0], color=real_color, lw=2, linestyle="--", label="Real")
- )
- if preid_data is not None:
- legend_handles.append(
- Line2D([0], [0], color=preid_color, lw=2, linestyle=":", label="Pre-ID")
- )
- if commanded_data is not None:
- legend_handles.append(
- Line2D(
- [0],
- [0],
- color=commanded_color,
- lw=2,
- linestyle="-.",
- label="Commanded",
- )
- )
-
- if legend_handles:
- fig.legend(
- handles=legend_handles,
- loc="upper center",
- bbox_to_anchor=(0.5, 0.935),
- ncol=len(legend_handles),
- fancybox=True,
- shadow=True,
- fontsize=10,
- title="Data Source",
- )
-
- fig.supxlabel("Time (s)", fontsize=8)
- plt.tight_layout(rect=(0, 0.03, 1, 0.9))
-
-
-def plot_objective(
- objective: Sequence[float],
- figsize: tuple[float, float] = (8, 5),
-):
- """Plot the objective value over optimization iterations."""
- plt.figure(figsize=figsize)
- plt.plot(objective, linewidth=2, marker="o", markersize=4)
- final_value = objective[-1]
- if abs(final_value) < 1e-3 or abs(final_value) > 1e3:
- final_str = f"{final_value:.2e}"
- else:
- final_str = f"{final_value:.4f}"
- plt.title(f"Objective Over Time (Final: {final_str})", fontsize=14, pad=10)
- plt.grid(True, linestyle="--", alpha=0.6)
- plt.xlabel("Iteration", fontsize=12)
- plt.ylabel("Objective", fontsize=12)
- plt.xticks(fontsize=10)
- plt.yticks(fontsize=10)
- plt.tight_layout()
-
-
-def plot_candidate(
- candidate: Sequence[np.ndarray],
- bounds: (
- tuple[Sequence[float] | np.ndarray, Sequence[float] | np.ndarray] | None
- ) = None,
- param_names: Sequence[str] | None = None,
- figsize: tuple[float, float] = (12, 2.5),
- dims_per_page: int = 6,
- log_diff: bool = True,
- bound_eps: float = 1e-3,
-):
- """Plot candidate parameter values and their diffs over iterations."""
- values = np.array(candidate) # shape: (n_iter, n_dim)
- n_iter, n_dim = values.shape
- diffs = np.diff(values, axis=0)
-
- mins = np.full(n_dim, -np.inf)
- maxs = np.full(n_dim, np.inf)
- if bounds is not None:
- mins = np.array(bounds[0])
- maxs = np.array(bounds[1])
- assert mins.shape == (n_dim,) and maxs.shape == (n_dim,)
-
- if param_names is not None:
- assert len(param_names) == n_dim
-
- # TODO(b/0) support pages, they are currently broken because
- # saving to disk overwrites the pages
- # n_pages = math.ceil(n_dim / dims_per_page)
- n_pages = 1
- for _page in range(n_pages):
- # start = page * dims_per_page
- # end = min((page + 1) * dims_per_page, n_dim)
- start = 0
- end = n_dim
- dims_in_page = end - start
-
- fig, axes = plt.subplots(
- dims_in_page,
- 2,
- figsize=(figsize[0], figsize[1] * dims_in_page),
- sharex="col",
- )
- if dims_in_page == 1:
- axes = np.expand_dims(axes, 0)
-
- for i, dim in enumerate(range(start, end)):
- label = param_names[dim] if param_names is not None else f"Dim {dim}"
- ax_val, ax_diff = axes[i]
-
- vals = values[:, dim]
- ax_val.set_ylabel(label, fontsize=10)
- ax_val.grid(True, linestyle="--", alpha=0.6)
- ax_val.tick_params(labelsize=9)
-
- if bounds is not None:
- lower, upper = mins[dim], maxs[dim]
- ax_val.axhspan(lower, upper, color="gray", alpha=0.08)
- ax_val.plot(
- [0, n_iter - 1],
- [lower, lower],
- color="gray",
- linestyle="--",
- alpha=0.3,
- linewidth=1,
- )
- ax_val.plot(
- [0, n_iter - 1],
- [upper, upper],
- color="gray",
- linestyle="--",
- alpha=0.3,
- linewidth=1,
- )
- near_lower = np.abs(vals - lower) < bound_eps
- near_upper = np.abs(vals - upper) < bound_eps
- near_bound = near_lower | near_upper
- for t in range(1, n_iter):
- is_near_prev = near_bound[t - 1]
- is_near_curr = near_bound[t]
- color = "#d62728" if is_near_prev and is_near_curr else "#1f77b4"
- ax_val.plot(
- [t - 1, t], [vals[t - 1], vals[t]], color=color, linewidth=2
- )
- ax_val.plot(t, vals[t], marker="o", markersize=3, color=color)
- # Overlay triangle markers for near-bound points
- for t in range(n_iter):
- if near_lower[t]:
- ax_val.plot(t, vals[t], marker="v", markersize=6, color="#d62728")
- elif near_upper[t]:
- ax_val.plot(t, vals[t], marker="^", markersize=6, color="#d62728")
- else:
- ax_val.plot(vals, linewidth=2, marker="o", markersize=3)
-
- # Annotate final value
- final_val = vals[-1]
- final_str = (
- f"{final_val:.2e}"
- if abs(final_val) < 1e-3 or abs(final_val) > 1e3
- else f"{final_val:.4f}"
- )
- ax_val.text(
- n_iter - 1,
- final_val,
- final_str,
- ha="right",
- va="bottom",
- fontsize=9,
- color="blue",
- )
-
- # Annotate final value.
- final_val = values[-1, dim]
- final_str = (
- f"{final_val:.2e}"
- if abs(final_val) < 1e-3 or abs(final_val) > 1e3
- else f"{final_val:.4f}"
- )
- ax_val.text(
- n_iter - 1,
- final_val,
- final_str,
- ha="right",
- va="bottom",
- fontsize=9,
- color="blue",
- )
-
- # Plot diffs
- if log_diff:
- eps = 1e-12
- ax_diff.plot(
- np.log10(np.abs(diffs[:, dim]) + eps),
- linewidth=2,
- marker="x",
- markersize=4,
- color="tab:orange",
- )
- ax_diff.set_ylabel("log Δ", fontsize=9)
- else:
- ax_diff.plot(
- diffs[:, dim],
- linewidth=2,
- marker="x",
- markersize=4,
- color="tab:orange",
- )
-
- ax_diff.grid(True, linestyle="--", alpha=0.6)
- ax_diff.tick_params(labelsize=9)
-
- # Set common labels/titles
- axes[-1, 0].set_xlabel("Iteration", fontsize=12)
- axes[-1, 1].set_xlabel("Iteration", fontsize=12)
- axes[0, 0].set_title("Candidate Value", fontsize=12)
- axes[0, 1].set_title("Δ Candidate (Diff)", fontsize=12)
-
- fig.suptitle(
- f"Candidate Values and Changes (Dims {start}-{end - 1})", fontsize=14
- )
- fig.tight_layout(rect=(0, 0, 1, 0.96))
-
-
-def plot_candidate_heatmap(
- candidate: Sequence[np.ndarray],
- param_names: Sequence[str] | None = None,
- bounds: (
- tuple[Sequence[float] | np.ndarray, Sequence[float] | np.ndarray] | None
- ) = None,
- normalize: bool = True,
- figsize: tuple[float, float] = (10, 6),
- cmap: str = "RdBu",
- show_colorbar: bool = True,
- bound_eps: float = 1e-3,
-):
- """Plot a heatmap of candidate parameter values over iterations."""
- data = np.array(candidate).T # shape: (n_dim, n_iter)
- n_dim = data.shape[0]
-
- if normalize and bounds is not None:
- min_bounds, max_bounds = bounds
- assert len(min_bounds) == len(max_bounds) == n_dim
- norm_data = np.empty_like(data)
- for i in range(n_dim):
- min_val = min_bounds[i]
- max_val = max_bounds[i]
- denom = max_val - min_val if max_val > min_val else 1.0
- norm_data[i] = (data[i] - min_val) / denom
- else:
- norm_data = data
-
- fig, ax = plt.subplots(figsize=figsize)
- im = ax.imshow(norm_data, aspect="auto", cmap=cmap)
-
- ax.set_xlabel("Iteration", fontsize=12)
- ax.set_ylabel("Parameter", fontsize=12)
-
- # Y-axis labels.
- if param_names is not None:
- assert len(param_names) == n_dim
- ax.set_yticks(np.arange(n_dim))
- ax.set_yticklabels(param_names, fontsize=10)
- else:
- ax.set_yticks(np.arange(n_dim))
- ax.set_yticklabels([f"Dim {i}" for i in range(n_dim)], fontsize=10)
-
- # Plot Xs where values are at bounds.
- if bounds is not None:
- min_bounds, max_bounds = bounds
- for dim in range(n_dim):
- min_val = min_bounds[dim]
- max_val = max_bounds[dim]
- for iter_idx, val in enumerate(data[dim]):
- if abs(val - min_val) < bound_eps or abs(val - max_val) < bound_eps:
- ax.plot(iter_idx, dim, "kx", markersize=6, markeredgewidth=1.5)
-
- if show_colorbar:
- cbar = fig.colorbar(im, ax=ax)
- label = "Normalized Value" if normalize else "Value"
- cbar.set_label(label, fontsize=12)
-
- ax.set_title("Candidate Heatmap", fontsize=14)
- fig.tight_layout()
-
-
-def parameter_confidence(
- all_exp_names: Sequence[str],
- all_params: Sequence[parameter.ParameterDict],
- all_intervals: Sequence[np.ndarray],
- cols: int = 5,
- gt_params: parameter.ParameterDict | None = None,
-):
- """Plot parameter estimates with confidence intervals."""
- named_estimates = {}
- # Create an entry for every non-frozen parameter
- for params in all_params:
- param_names = params.get_non_frozen_parameter_names()
- for name in param_names:
- if name not in named_estimates:
- named_estimates[name] = {
- "x": [],
- "intervals": [],
- "min_bounds": [],
- "max_bounds": [],
- "plot_labels": [],
- }
-
- for exp_name, params, intervals in zip(
- all_exp_names, all_params, all_intervals, strict=True
- ):
- param_names = params.get_non_frozen_parameter_names()
- xs = params.as_vector()
- bounds = params.get_bounds()
- assert xs.shape[0] == len(param_names)
- if gt_params is not None:
- for name in param_names:
- if name in gt_params:
- named_estimates[name]["xgt"] = gt_params[name].value[0]
- else:
- assert name[-1] == "]"
- left_bracket_i = name[::-1].find("[")
- index = int(name[-left_bracket_i:-1])
- named_estimates[name]["xgt"] = gt_params[
- name[: -left_bracket_i - 1]
- ].value[index]
-
- for i, (name, x, interval) in enumerate(
- zip(param_names, xs, intervals, strict=True)
- ):
- named_estimates[name]["x"].append(x)
- named_estimates[name]["intervals"].append(interval)
- named_estimates[name]["min_bounds"].append(bounds[0][i])
- named_estimates[name]["max_bounds"].append(bounds[1][i])
- named_estimates[name]["plot_labels"].append(exp_name)
-
- rows = len(named_estimates) // cols + 1
- fig, axs = plt.subplots(
- rows, cols, figsize=(20, 2 * (len(named_estimates) // cols + 1))
- )
- if rows == 1:
- axs = [axs]
-
- for i, name in enumerate(named_estimates):
- x_list = named_estimates[name]["x"]
- intervals = named_estimates[name]["intervals"]
- plot_labels = named_estimates[name]["plot_labels"]
-
- row = i % rows
- col = i // rows
-
- min_bound = np.min(named_estimates[name]["min_bounds"])
- max_bound = np.min(named_estimates[name]["max_bounds"])
-
- for j, (x, interval, plot_label) in enumerate(
- zip(x_list, intervals, plot_labels, strict=True)
- ):
- if not np.isfinite(interval) or 2.0 * interval > 2.0 * (
- max_bound - min_bound
- ):
- interval = 2.0 * (max_bound - min_bound)
- eb = axs[row][col].errorbar(x, -j, xerr=interval)
- eb[-1][0].set_linestyle("--")
- else:
- axs[row][col].errorbar(x, -j, xerr=interval)
- axs[row][col].scatter(x, -j, marker="x", label=plot_label)
-
- axs[row][col].set_xlim([min_bound, max_bound])
- axs[row][col].yaxis.set_ticklabels([])
- axs[row][col].set_title(name)
- axs[row][col].grid(True)
- axs[row][col].legend(
- fontsize=5, loc="upper right", bbox_to_anchor=(1.4, 1.0)
- )
- if gt_params is not None:
- axs[row][col].axvline(named_estimates[name]["xgt"], color="b", ls="--")
-
- fig.tight_layout()
-
-
def render_rollout(
model: mujoco.MjModel | Sequence[mujoco.MjModel],
data: mujoco.MjData,
diff --git a/python/mujoco/sysid/_src/signal_modifier.py b/python/mujoco/sysid/_src/signal_modifier.py
index 09c4b974..faf5938a 100644
--- a/python/mujoco/sysid/_src/signal_modifier.py
+++ b/python/mujoco/sysid/_src/signal_modifier.py
@@ -68,7 +68,13 @@ def apply_bias(
sensor_name: str,
bias: parameter.Parameter,
) -> timeseries.TimeSeries:
- """Apply a bias to a sensor in a timeseries."""
+ """Apply a bias to a sensor in a timeseries.
+
+ Args:
+ ts: Input timeseries.
+ sensor_name: Name of the sensor to modify.
+ bias: Parameter whose ``.value`` is added to the sensor columns.
+ """
indices = ts.get_indices(sensor_name)[1]
data_out = ts.data.copy()
data_out[..., indices] += bias.value
@@ -80,7 +86,13 @@ def apply_gain(
sensor_name: str,
gain: parameter.Parameter,
) -> timeseries.TimeSeries:
- """Apply a gain to a sensor in a timeseries."""
+ """Apply a gain to a sensor in a timeseries.
+
+ Args:
+ ts: Input timeseries.
+ sensor_name: Name of the sensor to modify.
+ gain: Parameter whose ``.value`` multiplies the sensor columns.
+ """
indices = ts.get_indices(sensor_name)[1]
data_out = ts.data.copy()
data_out[..., indices] *= gain.value
@@ -92,7 +104,13 @@ def apply_delay(
sensor_name: str,
delay: parameter.Parameter,
) -> timeseries.TimeSeries:
- """Apply a delay to a sensor in a timeseries."""
+ """Apply a delay to a sensor in a timeseries.
+
+ Args:
+ ts: Input timeseries.
+ sensor_name: Name of the sensor to delay.
+ delay: Parameter whose ``.value`` is the delay in seconds.
+ """
indices = ts.get_indices(sensor_name)[1]
ts_sensor = timeseries.TimeSeries(
@@ -205,7 +223,15 @@ def apply_resample_and_delay(
sensor_delays: dict[str, float] | None = None,
predicted_data: bool = True,
) -> timeseries.TimeSeries:
- """Resample a timeseries and apply per-sensor delays."""
+ """Resample a timeseries and apply per-sensor delays.
+
+ Args:
+ ts: Input timeseries to resample.
+ times: Target timestamps.
+ default_delay: Default delay applied to all columns.
+ sensor_delays: Optional per-sensor delay overrides.
+ predicted_data: If True, negate delays (shift predicted to match measured).
+ """
delays = _build_per_column_delays(
ts, default_delay, sensor_delays, predicted_data
)
@@ -234,7 +260,13 @@ def prepare_sensor_weights(
n_sensors: int,
model: mujoco.MjModel,
) -> np.ndarray:
- """Prepare sensor weights array from a dict or numpy array."""
+ """Prepare sensor weights array from a dict or numpy array.
+
+ Args:
+ sensor_weights: Mapping from sensor name to weight, or a flat array.
+ n_sensors: Total number of sensor columns.
+ model: MuJoCo model for resolving sensor names to indices.
+ """
if isinstance(sensor_weights, np.ndarray):
if sensor_weights.ndim != 1 or sensor_weights.shape[0] != n_sensors:
raise ValueError(
@@ -284,5 +316,10 @@ def normalize_residual(
residual: np.ndarray,
measured_data: np.ndarray,
) -> np.ndarray:
- """Normalize the residual by the standard deviation of the measured data."""
+ """Normalize the residual by the standard deviation of the measured data.
+
+ Args:
+ residual: Residual array, shape ``(n_timesteps, n_sensors)``.
+ measured_data: Measured data array, same shape as *residual*.
+ """
return residual / (np.linalg.norm(measured_data, axis=0) / np.sqrt(2))
diff --git a/python/mujoco/sysid/_src/signal_transform.py b/python/mujoco/sysid/_src/signal_transform.py
index fe84e307..38ef83d9 100644
--- a/python/mujoco/sysid/_src/signal_transform.py
+++ b/python/mujoco/sysid/_src/signal_transform.py
@@ -51,7 +51,12 @@ class SignalTransform:
self.normalize = normalize
def delay(self, pattern: str, param: parameter.Parameter) -> None:
- """Register a delay for sensors matching *pattern* (fnmatch)."""
+ """Register a delay for sensors matching *pattern* (fnmatch).
+
+ Args:
+ pattern: fnmatch pattern matched against sensor names.
+ param: Parameter whose ``.value`` is the delay in seconds.
+ """
self._delays.append((pattern, param.name, param))
def gain(
@@ -79,7 +84,13 @@ class SignalTransform:
param: parameter.Parameter,
target: str = "both",
) -> None:
- """Register an additive bias for sensors matching *pattern*."""
+ """Register an additive bias for sensors matching *pattern*.
+
+ Args:
+ pattern: fnmatch pattern matched against sensor names.
+ param: Parameter whose ``.value`` is the additive bias.
+ target: One of ``"predicted"``, ``"measured"``, or ``"both"``.
+ """
if target not in ("predicted", "measured", "both"):
raise ValueError(
f"target must be 'predicted', 'measured', or 'both', got {target!r}"
@@ -87,11 +98,19 @@ class SignalTransform:
self._biases.append((pattern, param.name, target))
def enable_sensors(self, sensor_names: list[str]) -> None:
- """Only include these sensors in the returned residual/timeseries."""
+ """Only include these sensors in the returned residual/timeseries.
+
+ Args:
+ sensor_names: Sensor names to keep in the output.
+ """
self._enabled_sensors = list(sensor_names)
def set_sensor_weights(self, weights: Mapping[str, float]) -> None:
- """Set per-sensor weights for the weighted diff."""
+ """Set per-sensor weights for the weighted diff.
+
+ Args:
+ weights: Mapping from sensor name to weight.
+ """
self._sensor_weights = weights
# Private methods.
diff --git a/python/mujoco/sysid/_src/timeseries.py b/python/mujoco/sysid/_src/timeseries.py
index f9d2bfa9..8fef0228 100644
--- a/python/mujoco/sysid/_src/timeseries.py
+++ b/python/mujoco/sysid/_src/timeseries.py
@@ -220,22 +220,26 @@ class TimeSeries:
nq = model.nq
nv = model.nv
- # Bodies
+ # Bodies with free joints, named by body rather than joint.
for body_id in range(model.nbody):
b = model.body(body_id)
body_name = b.name
- start_index = model.body_dofadr[body_id]
+ dof_adr = model.body_dofadr[body_id]
- if start_index >= 0 and b.dofnum[0] == 6:
- qpos_indices = np.arange(start_index, start_index + 7)
+ if dof_adr >= 0 and b.dofnum[0] == 6:
+ # Use the body's first joint qposadr for qpos. Free joints have
+ # 7 qpos elements but only 6 dofs, so dofadr and qposadr diverge
+ # for subsequent entries.
+ first_jnt = model.body_jntadr[body_id]
+ qpos_adr = model.jnt_qposadr[first_jnt]
+ qpos_indices = np.arange(qpos_adr, qpos_adr + 7)
qpos_map[f"{body_name}_qpos"] = (SignalType.MjStateQPos, qpos_indices)
- qvel_indices = np.arange(start_index + nq, start_index + nq + 6)
+ qvel_indices = np.arange(dof_adr + nq, dof_adr + nq + 6)
qvel_map[f"{body_name}_qvel"] = (SignalType.MjStateQVel, qvel_indices)
- # Joints
+ # Joints, excluding free joints which are handled above.
for jnt_id in range(model.njnt):
jnt_name = model.joint(jnt_id).name
- start_index = model.jnt_qposadr[jnt_id]
jnt_type = model.jnt_type[jnt_id]
qpos_width = 1
@@ -246,9 +250,11 @@ class TimeSeries:
elif jnt_type == mujoco.mjtJoint.mjJNT_FREE:
continue
- qpos_indices = np.arange(start_index, start_index + qpos_width)
+ qpos_adr = model.jnt_qposadr[jnt_id]
+ qpos_indices = np.arange(qpos_adr, qpos_adr + qpos_width)
qpos_map[f"{jnt_name}_qpos"] = (SignalType.MjStateQPos, qpos_indices)
- qvel_indices = np.arange(start_index + nq, start_index + nq + qvel_width)
+ dof_adr = model.jnt_dofadr[jnt_id]
+ qvel_indices = np.arange(dof_adr + nq, dof_adr + nq + qvel_width)
qvel_map[f"{jnt_name}_qvel"] = (SignalType.MjStateQVel, qvel_indices)
# Actuators
@@ -445,7 +451,11 @@ class TimeSeries:
return cls(times=times, data=data, signal_mapping=signal_mapping)
def get_indices(self, obs_name: str) -> tuple[SignalType, np.ndarray]:
- """Look up the signal type and column indices for a named observation."""
+ """Look up the signal type and column indices for a named observation.
+
+ Args:
+ obs_name: Name of the observation signal.
+ """
assert self.signal_mapping is not None
if obs_name not in self.signal_mapping:
raise ValueError(
@@ -462,7 +472,14 @@ class TimeSeries:
str, tuple[SignalType, np.ndarray | list[int] | int]
],
) -> TimeSeries:
- """Construct a TimeSeries, normalising index entries to ``np.ndarray``."""
+ """Construct a TimeSeries, normalizing index entries to ``np.ndarray``.
+
+ Args:
+ times: 1-D timestamp array.
+ data: Data array with first axis corresponding to time.
+ signal_mapping: Dict mapping signal names to ``(type, indices)`` tuples.
+ Index entries are coerced to ``np.ndarray``.
+ """
normalized: SignalMappingType = {}
for key in signal_mapping:
signal_type, indices = signal_mapping[key]
@@ -560,7 +577,11 @@ class TimeSeries:
)
def save_to_csv(self, path: str | pathlib.Path) -> None:
- """Save the time series data to a CSV file."""
+ """Save the time series data to a CSV file.
+
+ Args:
+ path: Path where the CSV file will be written.
+ """
np.savetxt(
path,
np.concatenate([self.times[:, None], self.data], axis=1),
diff --git a/python/mujoco/sysid/_src/trajectory.py b/python/mujoco/sysid/_src/trajectory.py
index 7154bf3b..afd3679c 100644
--- a/python/mujoco/sysid/_src/trajectory.py
+++ b/python/mujoco/sysid/_src/trajectory.py
@@ -302,12 +302,23 @@ class SystemTrajectory:
height=height,
)
+def _map_states(from_array, to_array, from_names, to_mapping, map_offset):
+ i = 0
+ for name in from_names:
+ _, indices = to_mapping[name]
+ width = indices.shape[0]
+ to_array[indices - map_offset] = from_array[i:i+width]
+ i += width
+ return to_array
def create_initial_state(
model: mujoco.MjModel,
qpos: np.ndarray,
qvel: np.ndarray | None = None,
act: np.ndarray | None = None,
+ qpos_names: Sequence[str] | None = None,
+ qvel_names: Sequence[str] | None = None,
+ act_names: Sequence[str] | None = None,
) -> np.ndarray:
"""Build a ``mjSTATE_FULLPHYSICS`` initial-state vector from components.
@@ -316,14 +327,46 @@ def create_initial_state(
qpos: Joint positions, shape ``(nq,)``.
qvel: Joint velocities, shape ``(nv,)``. Defaults to zero.
act: Actuator activations, shape ``(na,)``. Defaults to zero.
+ qpos_names: Names to map elements of qpos to specific MuJoCo states.
+ If None, assumes qpos is in MuJoCo's order.
+ qvel_names: Names to map elements of qvel to specific MuJoCo states.
+ If None, assumes qvel is in MuJoCo's order.
+ act_names: Actuator names to map elements of act to specific MuJoCo
+ actuators. If None, assumes act is in MuJoCo's order.
Returns:
Flat state vector suitable for ``mujoco.rollout``.
"""
data = mujoco.MjData(model)
- initial_state = np.empty((
- mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS.value),
- ))
+
+ if qpos_names is not None and len(qpos_names) != qpos.shape[0]:
+ raise ValueError(
+ f"Expected qpos to have shape {len(qpos_names)}, got {qpos.shape[0]}"
+ )
+ if qvel is None and qvel_names is not None:
+ raise ValueError("Expected qvel to not be None when qvel_names is not None")
+ if qvel_names is not None and qvel is not None and len(qvel_names) != qvel.shape[0]:
+ raise ValueError(
+ f"Expected qvel to have shape {len(qvel_names)}, got {qvel.shape[0]}"
+ )
+ if act_names is not None and len(act_names) != act.shape[0]:
+ raise ValueError(
+ f"Expected act to have shape {len(act_names)}, got {act.shape[0]}"
+ )
+
+ if (qpos_names is not None
+ or qvel_names is not None
+ or act_names is not None):
+ qpos_map, qvel_map, act_map, _ = timeseries.TimeSeries.compute_all_state_mappings(model)
+ if qpos_names is not None:
+ qpos = _map_states(qpos, np.copy(data.qpos), qpos_names, qpos_map, 0)
+ if qvel_names is not None:
+ indices_offset = data.qpos.shape[0]
+ qvel = _map_states(qvel, np.copy(data.qvel), qvel_names, qvel_map, indices_offset)
+ if act_names is not None:
+ indices_offset = data.qpos.shape[0] + data.qvel.shape[0]
+ act = _map_states(act, np.copy(data.act), act_names, act_map, indices_offset)
+
if qpos.shape[0] != model.nq:
raise ValueError(
f"Expected qpos to have shape {model.nq}, got {qpos.shape[0]}."
@@ -341,6 +384,10 @@ def create_initial_state(
f"Expected act to have shape {model.na}, got {act.shape[0]}."
)
data.act[:] = act
+
+ initial_state = np.empty((
+ mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS.value),
+ ))
mujoco.mj_getState(
model, data, initial_state, mujoco.mjtState.mjSTATE_FULLPHYSICS.value
)
@@ -433,7 +480,14 @@ class ModelSequences:
def timeseries2array(
control_signal: timeseries.TimeSeries | Sequence[timeseries.TimeSeries],
) -> tuple[np.ndarray, np.ndarray]:
- """Convert control TimeSeries to stacked arrays, dropping the last step."""
+ """Convert control TimeSeries to stacked arrays, dropping the last step.
+
+ Args:
+ control_signal: Control TimeSeries or sequence of TimeSeries.
+
+ Returns:
+ ``(control_array, control_times)`` with the last time step removed.
+ """
if isinstance(control_signal, timeseries.TimeSeries):
control = control_signal.data
control_times = control_signal.times
@@ -457,7 +511,11 @@ def timeseries2array(
def sequence2array(
initial_states: np.ndarray | Sequence[np.ndarray],
) -> np.ndarray:
- """Stack a sequence of initial-state vectors into a single array."""
+ """Stack a sequence of initial-state vectors into a single array.
+
+ Args:
+ initial_states: Single state array or sequence of state arrays.
+ """
if isinstance(initial_states, np.ndarray):
return initial_states
return np.stack(initial_states, axis=0)
@@ -474,12 +532,28 @@ def arrays2traj(
state_mapping: timeseries.SignalMappingType,
ctrl_mapping: timeseries.SignalMappingType,
) -> Sequence[SystemTrajectory]:
- """Convert raw rollout arrays into a list of SystemTrajectory objects."""
+ """Convert raw rollout arrays into a list of SystemTrajectory objects.
+
+ Args:
+ models: Single model or sequence of models (one per batch element).
+ initial_states: Initial state array(s).
+ control: Control array, shape ``(nbatch, nsteps, nu)``.
+ control_times: Control timestamps, shape ``(nbatch, nsteps)``.
+ state: State array, shape ``(nbatch, nsteps, nstate)``.
+ sensordata: Sensor data array, shape ``(nbatch, nsteps, nsensordata)``.
+ signal_mapping: Signal mapping for sensor data.
+ state_mapping: Signal mapping for state data.
+ ctrl_mapping: Signal mapping for control data.
+ """
nbatch = state.shape[0]
- # TODO(kevin): When is np.tile necessary?
+ # TODO(kevin): When is np.tile/atleast_2d/etc necessary?
# initial_states = np.tile(initial_states, (nbatch, 1))
# control = np.tile(control, (nbatch, 1, 1))
# control_times = np.tile(control_times, (nbatch, 1))
+ initial_states = np.atleast_2d(initial_states)
+ if control.ndim == 2:
+ control = control[np.newaxis, :, :]
+ control_times = np.atleast_2d(control_times)
if isinstance(models, mujoco.MjModel):
models_list = [models] * nbatch
diff --git a/python/mujoco/sysid/report/defaults.py b/python/mujoco/sysid/report/defaults.py
index 60a20bde..06a2b7e7 100644
--- a/python/mujoco/sysid/report/defaults.py
+++ b/python/mujoco/sysid/report/defaults.py
@@ -15,13 +15,10 @@
"""Default report generation for system identification results."""
from collections.abc import Sequence
-import os
import pathlib
-import matplotlib.pyplot as plt
from mujoco.sysid._src import model_modifier
from mujoco.sysid._src import parameter
-from mujoco.sysid._src import plotting
from mujoco.sysid._src.optimize import calculate_intervals
from mujoco.sysid._src.residual import BuildModelFn
from mujoco.sysid._src.trajectory import ModelSequences
@@ -91,7 +88,7 @@ def default_report(
generate_video_from_trajectories(
initial_params=initial_params,
opt_params=opt_params,
- build_model=build_model,
+ _build_model=build_model,
trajectories=all_trajectories,
model_spec=model_spec_to_render,
output_filepath=video_all_path,
@@ -103,7 +100,7 @@ def default_report(
generate_video_from_trajectories(
initial_params=initial_params,
opt_params=opt_params,
- build_model=build_model,
+ _build_model=build_model,
trajectories=all_trajectories,
model_spec=model_spec_to_render,
output_filepath=video_init_path,
@@ -116,7 +113,7 @@ def default_report(
generate_video_from_trajectories(
initial_params=initial_params,
opt_params=opt_params,
- build_model=build_model,
+ _build_model=build_model,
trajectories=all_trajectories,
model_spec=model_spec_to_render,
output_filepath=video_opt_path,
@@ -292,105 +289,3 @@ def default_report(
if save_path:
rb.save(save_path / "report.html")
return rb
-
-
-# TODO(nimrod): Consider deleting this function, given we can export plots from
-# plotly either on the web or with fig.write_image.
-def default_report_matplotlib(
- experiment_results_folder: os.PathLike[str],
- models_sequences: Sequence[ModelSequences],
- params: parameter.ParameterDict,
- sysid_residual,
- x0: np.ndarray,
- opt_result: scipy_optimize.OptimizeResult,
- build_model: BuildModelFn | None = model_modifier.apply_param_modifiers,
-):
- """Outputs PNG plots to the experiment results folder."""
- experiment_results_folder = pathlib.Path(experiment_results_folder)
- if not experiment_results_folder.exists():
- experiment_results_folder.mkdir(parents=True, exist_ok=True)
-
- x_hat = opt_result.x
- params.update_from_vector(x_hat)
-
- # Save the ID'd models out
- assert build_model is not None
- model_hat = None
- for model_sequences in models_sequences:
- model_hat = build_model(params, model_sequences.spec)
- assert model_hat is not None
-
- # Get predictions for initial solution.
- params.update_from_vector(x0)
- names = [
- f"{model_sequences.name}\n{sequence}"
- for model_sequences in models_sequences
- for sequence in model_sequences.sequence_name
- ]
- _, pred0s, record0s = sysid_residual(x0, return_pred_all=True)
-
- for name, pred0, record0 in zip(names, pred0s, record0s, strict=True):
- plotting.plot_sensor_comparison(
- model_hat,
- predicted_times=pred0[0].times,
- predicted_data=pred0[0].data,
- real_times=record0[0].times,
- real_data=record0[0].data,
- title_prefix=f"x0 {name}",
- size_factor=0.5,
- )
- name_fig = name.replace("/", " ")
- name_fig = name_fig.replace("\n", " ")
- plt.savefig(os.path.join(experiment_results_folder, f"x0 {name_fig}.png"))
-
- residuals_star, preds_star, records_star = sysid_residual(
- x_hat, return_pred_all=True
- )
- for name, pred, record, _ in zip(
- names, preds_star, records_star, pred0s, strict=True
- ):
- plotting.plot_sensor_comparison(
- model_hat,
- predicted_times=pred[0].times,
- predicted_data=pred[0].data,
- real_times=record[0].times,
- real_data=record[0].data,
- title_prefix=f"x* {name}",
- size_factor=0.5,
- )
- name_fig = name.replace("/", " ")
- name_fig = name_fig.replace("\n", " ")
- plt.savefig(experiment_results_folder / f"xstar {name_fig}.png")
-
- # Add diagnostic optimization trace plots.
- if "extras" in opt_result:
- # Objective value over iterations.
- objective = opt_result.extras["objective"]
- plotting.plot_objective(objective)
- plt.savefig(experiment_results_folder / "loss.png", dpi=300)
-
- # Candidate parameter values over iterations.
- candidate = opt_result.extras["candidate"]
-
- # Candidate parameter values over iterations.
- # Candidate heatmap over iterations.
- plotting.plot_candidate_heatmap(
- candidate,
- param_names=params.get_non_frozen_parameter_names(),
- bounds=params.get_bounds(),
- )
- plt.savefig(experiment_results_folder / "candidate_heatmap.png", dpi=300)
-
- plotting.plot_candidate(
- candidate,
- bounds=params.get_bounds(),
- param_names=params.get_non_frozen_parameter_names(),
- )
- plt.savefig(experiment_results_folder / "candidate.png", dpi=300)
-
- _, intervals = calculate_intervals(residuals_star, opt_result.jac)
- plotting.parameter_confidence(
- all_exp_names=["trial"], all_params=[params], all_intervals=[intervals]
- )
- # plotting.parameter_confidence(["trial"], [params], [x_hat], [intervals])
- plt.savefig(experiment_results_folder / "params.png")
diff --git a/python/mujoco/sysid/tests/test_integration.py b/python/mujoco/sysid/tests/test_integration.py
index a2c5d8f1..d6b7eeb7 100644
--- a/python/mujoco/sysid/tests/test_integration.py
+++ b/python/mujoco/sysid/tests/test_integration.py
@@ -12,230 +12,216 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
-"""End-to-end integration test using the box model."""
+"""End-to-end integration tests for mujoco.sysid."""
import pathlib
import tempfile
import mujoco
-import mujoco.rollout as mj_rollout
-from mujoco.sysid._src import signal_modifier
-from mujoco.sysid._src import timeseries
-from mujoco.sysid._src.io import save_results
-from mujoco.sysid._src.model_modifier import _infer_inertial
-from mujoco.sysid._src.optimize import optimize
-from mujoco.sysid._src.parameter import Parameter
-from mujoco.sysid._src.parameter import ParameterDict
-from mujoco.sysid._src.residual import build_residual_fn
-from mujoco.sysid._src.trajectory import ModelSequences
-from mujoco.sysid._src.trajectory import create_initial_state
-from mujoco.sysid.tests.conftest import BOX_XML
+import mujoco.rollout as rollout
+from mujoco import sysid
import numpy as np
-def _generate_box_data(
- spec: mujoco.MjSpec,
- duration: float = 1.0,
-) -> tuple[timeseries.TimeSeries, timeseries.TimeSeries, np.ndarray]:
- """Generate synthetic box-pushing data via rollout."""
+# ---------------------------------------------------------------------------
+# Models
+# ---------------------------------------------------------------------------
+
+SPRING_MASS_XML = """\
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+"""
+
+ARM_XML = """\
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+"""
+
+JOINT_NAMES = ["joint1", "joint2", "joint3", "joint4", "joint5"]
+TRUE_ARMATURE = {"joint1": 0.5, "joint2": 0.4, "joint3": 0.3,
+ "joint4": 0.2, "joint5": 0.1}
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _generate_data(xml, ctrl_fn, duration):
+ """Rollout a model and return (spec, initial_state, control_ts, sensor_ts)."""
+ spec = mujoco.MjSpec.from_string(xml)
model = spec.compile()
data = mujoco.MjData(model)
n_steps = int(duration / model.opt.timestep)
t = np.arange(n_steps) * model.opt.timestep
- force = (np.sin(t) * 3.0).reshape(-1, 1)
- control_ts = timeseries.TimeSeries(t, force)
+ ctrl = ctrl_fn(t)
- initial_state = create_initial_state(model, data.qpos, data.qvel, data.act)
-
- control_applied = force[:-1]
- state, _ = mj_rollout.rollout(model, data, initial_state, control_applied)
+ initial_state = sysid.create_initial_state(
+ model, data.qpos, data.qvel, data.act
+ )
+ state, sensor = rollout.rollout(model, data, initial_state, ctrl[:-1])
state = np.squeeze(state, axis=0)
+ sensor = np.squeeze(sensor, axis=0)
+ times = state[:, 0]
- sensor_ids = [1, 8]
- signal_mapping = {
- "pos_x": (timeseries.SignalType.MjStateQPos, np.array([0])),
- "vel_x": (timeseries.SignalType.MjStateQVel, np.array([1])),
- }
- state_times = state[:, 0]
- sensordata = timeseries.TimeSeries(
- state_times,
- state[:, sensor_ids],
- signal_mapping,
+ control_ts = sysid.TimeSeries(t, ctrl)
+ sensor_ts = sysid.TimeSeries.from_names(times, sensor, model)
+ return spec, initial_state, control_ts, sensor_ts
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+def test_spring_mass_recover_mass():
+ """Recover true mass=1.0 starting from initial guess of 2.0."""
+ spec, initial_state, control_ts, sensor_ts = _generate_data(
+ SPRING_MASS_XML,
+ ctrl_fn=lambda t: (
+ 5.0 * np.sin(2 * np.pi * 1.5 * t)
+ + 3.0 * np.sin(2 * np.pi * 3.7 * t)
+ ).reshape(-1, 1),
+ duration=3.0,
)
- return control_ts, sensordata, initial_state
+ params = sysid.ParameterDict()
+ params.add(sysid.Parameter(
+ "mass", nominal=1.0, min_value=0.3, max_value=3.0,
+ modifier=lambda s, p: setattr(s.body("ball"), "mass", p.value[0]),
+ ))
+ params["mass"].value[:] = 2.0
-
-def _build_box_params() -> ParameterDict:
- """Build parameter dict with modifier callbacks matching box config."""
- pdict = ParameterDict()
-
- pdict.add(
- Parameter(
- "box_mass",
- [5],
- min_value=[4.5],
- max_value=[5.5],
- modifier=lambda s, p: setattr(
- _infer_inertial(s, "box"), "mass", p.value[0]
- ),
- )
+ ms = sysid.ModelSequences(
+ "spring_mass", spec, "measured", initial_state, control_ts, sensor_ts,
)
- pdict.add(
- Parameter(
- "solref1",
- [0.01],
- min_value=[0.002],
- max_value=[0.02],
- modifier=lambda s, p: s.pair("box_floor").solref.__setitem__(
- 0, p.value[0]
- ),
- )
+ residual_fn = sysid.build_residual_fn(models_sequences=[ms])
+ opt_params, _ = sysid.optimize(
+ initial_params=params, residual_fn=residual_fn, optimizer="mujoco",
)
- pdict.add(
- Parameter(
- "solref2",
- [1.0],
- min_value=[0.3],
- max_value=[1.7],
- frozen=True,
- modifier=lambda s, p: s.pair("box_floor").solref.__setitem__(
- 1, p.value[0]
- ),
- )
+
+ np.testing.assert_allclose(opt_params["mass"].value[0], 1.0, atol=1e-4)
+
+
+def test_arm_recover_armature():
+ """Recover 5 joint armature values and verify save_results."""
+ spec, initial_state, control_ts, sensor_ts = _generate_data(
+ ARM_XML,
+ ctrl_fn=lambda t: np.column_stack([
+ 5.0 * np.sin(2 * np.pi * 0.5 * t),
+ 4.0 * np.sin(2 * np.pi * 0.7 * t + 0.5),
+ 3.0 * np.sin(2 * np.pi * 0.4 * t + 1.0),
+ 2.0 * np.sin(2 * np.pi * 0.9 * t + 1.5),
+ 1.0 * np.sin(2 * np.pi * 0.6 * t + 2.0),
+ ]),
+ duration=2.0,
)
- pdict.add(
- Parameter(
- "friction1",
- [1.6],
- min_value=[0],
- max_value=[3.0],
- frozen=True,
- modifier=lambda s, p: s.pair("box_floor").friction.__setitem__(
- 0, p.value[0]
- ),
- )
+
+ params = sysid.ParameterDict()
+ for name in JOINT_NAMES:
+ params.add(sysid.Parameter(
+ f"{name}_armature",
+ nominal=TRUE_ARMATURE[name],
+ min_value=0.001,
+ max_value=1.0,
+ modifier=lambda s, p, n=name: setattr(
+ s.joint(n), "armature", p.value[0]
+ ),
+ ))
+ params[f"{name}_armature"].value[:] = 0.01
+
+ ms = sysid.ModelSequences(
+ "arm", spec, "sinusoidal", initial_state, control_ts, sensor_ts,
)
- pdict.add(
- Parameter(
- "friction2",
- [0.005],
- min_value=[0],
- max_value=[0.01],
- modifier=lambda s, p: s.pair("box_floor").friction.__setitem__(
- 1, p.value[0]
- ),
- )
+ residual_fn = sysid.build_residual_fn(models_sequences=[ms])
+ opt_params, opt_result = sysid.optimize(
+ initial_params=params, residual_fn=residual_fn, optimizer="mujoco",
)
- pdict.add(
- Parameter(
- "friction3",
- [0.0001],
- min_value=[0],
- max_value=[0.001],
- frozen=True,
- modifier=lambda s, p: s.pair("box_floor").friction.__setitem__(
- 2, p.value[0]
- ),
- )
- )
- return pdict
-
-def test_box_end_to_end():
- """Full 5-stage pipeline: generate data, build residual, optimize 3 iters, save."""
- spec = mujoco.MjSpec.from_string(BOX_XML)
-
- # 1. Generate synthetic ground-truth data.
- control, sensordata, initial_state = _generate_box_data(spec, duration=1.0)
-
- # 2. Build config with known parameters.
- params = _build_box_params()
-
- # 3. Create ModelSequences.
- models_sequences = [
- ModelSequences(
- "box",
- spec,
- "push",
- initial_state,
- control,
- sensordata,
- allow_missing_sensors=True,
- )
- ]
-
- # 4. Define modify_residual (box uses state-based residual).
- def modify_residual(
- _params,
- _sensordata_predicted,
- sensordata_measured,
- model,
- _return_pred_all,
- state=None,
- **_kwargs,
- ):
- assert state is not None
- sensor_ids = [1, 8]
- statedata_predicted = timeseries.TimeSeries(
- state[:, 0],
- state[..., sensor_ids],
- {
- "pos_x": (timeseries.SignalType.MjStateQPos, np.array([0])),
- "vel_x": (timeseries.SignalType.MjStateQVel, np.array([1])),
- },
+ for name in JOINT_NAMES:
+ np.testing.assert_allclose(
+ opt_params[f"{name}_armature"].value[0],
+ TRUE_ARMATURE[name],
+ atol=1e-4,
)
- sensordata_measured = signal_modifier.apply_delayed_ts_window(
- sensordata_measured, statedata_predicted, 0.0, 0.0
- )
- statedata_predicted = statedata_predicted.resample(
- sensordata_measured.times
- )
- res = signal_modifier.weighted_diff(
- predicted_data=statedata_predicted.data,
- measured_data=sensordata_measured.data,
- model=model,
- )
- res = signal_modifier.normalize_residual(res, sensordata_measured.data)
- return res, statedata_predicted, sensordata_measured
- residual_fn = build_residual_fn(
- models_sequences=models_sequences,
- modify_residual=modify_residual,
- )
-
- # 5. Perturb params and optimize (3 iters to verify).
- rng = np.random.default_rng(42)
- params.randomize(rng=rng)
-
- # Compute initial cost.
- initial_residuals, _, _ = residual_fn(params.as_vector(), params)
- initial_cost = sum(np.sum(r**2) for r in initial_residuals)
-
- opt_params, opt_result = optimize(
- initial_params=params,
- residual_fn=residual_fn,
- optimizer="mujoco",
- max_iters=3,
- verbose=False,
- )
-
- # 6. Assert basic properties.
- assert opt_result.x.shape == params.as_vector().shape
-
- # Compute final cost.
- final_residuals, _, _ = residual_fn(opt_result.x, opt_params)
- final_cost = sum(np.sum(r**2) for r in final_residuals)
- assert (
- final_cost <= initial_cost
- ), f"Cost should decrease: {final_cost} > {initial_cost}"
-
- # 7. Save results to a temp dir.
+ # Verify save_results produces expected files.
with tempfile.TemporaryDirectory() as tmpdir:
- save_results(
+ sysid.save_results(
experiment_results_folder=tmpdir,
- models_sequences=models_sequences,
+ models_sequences=[ms],
initial_params=params,
opt_params=opt_params,
opt_result=opt_result,
@@ -246,4 +232,4 @@ def test_box_end_to_end():
assert (result_dir / "params_x_hat.yaml").exists()
assert (result_dir / "results.pkl").exists()
assert (result_dir / "confidence.pkl").exists()
- assert (result_dir / "box.xml").exists()
+ assert (result_dir / "arm.xml").exists()
diff --git a/python/mujoco/sysid/tests/test_timeseries.py b/python/mujoco/sysid/tests/test_timeseries.py
index 6ae8cdbd..4809e290 100644
--- a/python/mujoco/sysid/tests/test_timeseries.py
+++ b/python/mujoco/sysid/tests/test_timeseries.py
@@ -329,3 +329,125 @@ def test_from_custom():
assert "b" in ts.signal_mapping
assert ts.signal_mapping["a"][1].size == 1
assert ts.signal_mapping["b"][1].size == 2
+
+
+# ---------------------------------------------------------------------------
+# compute_all_state_mappings correctness tests
+# ---------------------------------------------------------------------------
+
+
+def _verify_state_mapping(model):
+ """Set known qpos/qvel values and verify mappings recover them correctly."""
+ qpos_map, qvel_map, act_map, _ = TimeSeries.compute_all_state_mappings(model)
+
+ data = mujoco.MjData(model)
+ # Fill with distinct values so misalignment is detectable.
+ data.qpos[:] = 100 + np.arange(model.nq)
+ data.qvel[:] = 200 + np.arange(model.nv)
+
+ state = np.empty(
+ mujoco.mj_stateSize(
+ model, mujoco.mjtState.mjSTATE_FULLPHYSICS.value
+ )
+ )
+ mujoco.mj_getState(
+ model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS.value
+ )
+ # Strip the leading time element (mjSTATE_FULLPHYSICS includes time).
+ state_no_time = state[1:]
+
+ # Verify every qpos mapping entry.
+ for name, (sig_type, indices) in qpos_map.items():
+ assert sig_type == SignalType.MjStateQPos
+ values = state_no_time[indices]
+ # All qpos values should be in [100, 100+nq).
+ assert np.all(values >= 100) and np.all(values < 100 + model.nq), (
+ f"{name}: got {values}"
+ )
+
+ # Verify every qvel mapping entry.
+ for name, (sig_type, indices) in qvel_map.items():
+ assert sig_type == SignalType.MjStateQVel
+ values = state_no_time[indices]
+ # All qvel values should be in [200, 200+nv).
+ assert np.all(values >= 200) and np.all(values < 200 + model.nv), (
+ f"{name}: got {values}"
+ )
+
+ # Verify total coverage.
+ all_qpos_indices = np.concatenate([v[1] for v in qpos_map.values()])
+ all_qvel_indices = np.concatenate([v[1] for v in qvel_map.values()])
+ assert len(all_qpos_indices) == model.nq
+ assert len(all_qvel_indices) == model.nv
+ assert len(np.unique(all_qpos_indices)) == model.nq, "Duplicate qpos indices"
+ assert len(np.unique(all_qvel_indices)) == model.nv, "Duplicate qvel indices"
+
+
+def test_state_mapping_hinge_only():
+ """All-hinge model: qposadr == dofadr, so mapping is straightforward."""
+ xml = """
+
+
+
+
+
+
+
+
+ """
+ _verify_state_mapping(mujoco.MjModel.from_xml_string(xml))
+
+
+def test_state_mapping_free_plus_hinge():
+ """Free body + hinge joints: jnt_qposadr != jnt_dofadr for the hinges."""
+ xml = """
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ """
+ _verify_state_mapping(mujoco.MjModel.from_xml_string(xml))
+
+
+def test_state_mapping_two_free_bodies():
+ """Two free bodies: body_dofadr != jnt_qposadr for the second body."""
+ xml = """
+
+
+
+
+
+
+
+
+
+
+ """
+ _verify_state_mapping(mujoco.MjModel.from_xml_string(xml))
+
+
+def test_state_mapping_ball_plus_hinge():
+ """Ball joint + hinge: ball takes 4 qpos / 3 qvel, offsetting the hinge."""
+ xml = """
+
+
+
+
+
+
+
+
+
+
+ """
+ _verify_state_mapping(mujoco.MjModel.from_xml_string(xml))
diff --git a/python/mujoco/sysid/tests/test_trajectory.py b/python/mujoco/sysid/tests/test_trajectory.py
index 06380cb5..f411c2bd 100644
--- a/python/mujoco/sysid/tests/test_trajectory.py
+++ b/python/mujoco/sysid/tests/test_trajectory.py
@@ -132,6 +132,51 @@ def test_create_initial_state_wrong_qpos(box_model):
create_initial_state(box_model, np.zeros(999))
+def test_create_initial_state_with_names():
+ """Named mapping places qpos/qvel into correct slots for a subset of joints."""
+ xml = """
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ """
+ model = mujoco.MjModel.from_xml_string(xml)
+ data = mujoco.MjData(model)
+
+ # Map only the two hinge joints by name, leaving the free body at defaults.
+ qpos_subset = np.array([1.1, 2.2])
+ qvel_subset = np.array([3.3, 4.4])
+ state = create_initial_state(
+ model, qpos_subset, qvel_subset,
+ qpos_names=["h1_qpos", "h2_qpos"],
+ qvel_names=["h1_qvel", "h2_qvel"],
+ )
+
+ # Unpack the state to verify values ended up in the right slots.
+ mujoco.mj_setState(
+ model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS.value
+ )
+ # Free body qpos (indices 0-6) should be at model defaults.
+ # Hinge joints at qposadr 7 and 8.
+ np.testing.assert_allclose(data.qpos[7], 1.1)
+ np.testing.assert_allclose(data.qpos[8], 2.2)
+ # Hinge joints at dofadr 6 and 7.
+ np.testing.assert_allclose(data.qvel[6], 3.3)
+ np.testing.assert_allclose(data.qvel[7], 4.4)
+
+
def test_split(sample_trajectory):
"""A long trajectory can be split into smaller chunks for batched optimization."""
traj, *_ = sample_trajectory