sysid: named ic construction, bug fixes, docstrings, README, more tests
Co-authored-by: Kevin Zakka <kevinarmandzakka@gmail.com>
This commit is contained in:
Levi Burner
+66
-608
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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."
|
||||
|
||||
@@ -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,))
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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 = """\
|
||||
<mujoco model="spring_mass">
|
||||
<option timestep="0.002">
|
||||
<flag contact="disable"/>
|
||||
</option>
|
||||
<worldbody>
|
||||
<body name="ball" pos="0 0 0.1">
|
||||
<inertial pos="0 0 0" mass="1.0" diaginertia="0.001 0.001 0.001"/>
|
||||
<joint name="slide" type="slide" axis="1 0 0"
|
||||
stiffness="100" damping="5.0"/>
|
||||
<geom type="sphere" size="0.05"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<actuator>
|
||||
<motor name="push" joint="slide"/>
|
||||
</actuator>
|
||||
<sensor>
|
||||
<jointpos name="position" joint="slide"/>
|
||||
<jointvel name="velocity" joint="slide"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
ARM_XML = """\
|
||||
<mujoco model="arm">
|
||||
<compiler angle="radian" autolimits="true"/>
|
||||
<option integrator="implicitfast" timestep="0.002">
|
||||
<flag contact="disable"/>
|
||||
</option>
|
||||
<worldbody>
|
||||
<body name="link1" pos="0 0 0.1">
|
||||
<inertial pos="0 0 0.05" mass="1.0" diaginertia="0.01 0.01 0.005"/>
|
||||
<joint name="joint1" type="hinge" axis="0 0 1" range="-3.14 3.14"
|
||||
armature="0.5" damping="1.0"/>
|
||||
<geom type="capsule" fromto="0 0 0 0 0 0.1" size="0.04"/>
|
||||
<body name="link2" pos="0 0 0.1">
|
||||
<inertial pos="0 0 0.05" mass="0.8" diaginertia="0.008 0.008 0.004"/>
|
||||
<joint name="joint2" type="hinge" axis="0 1 0" range="-3.14 3.14"
|
||||
armature="0.4" damping="0.8"/>
|
||||
<geom type="capsule" fromto="0 0 0 0 0 0.1" size="0.035"/>
|
||||
<body name="link3" pos="0 0 0.1">
|
||||
<inertial pos="0 0 0.05" mass="0.6" diaginertia="0.006 0.006 0.003"/>
|
||||
<joint name="joint3" type="hinge" axis="0 1 0" range="-3.14 3.14"
|
||||
armature="0.3" damping="0.6"/>
|
||||
<geom type="capsule" fromto="0 0 0 0 0 0.1" size="0.03"/>
|
||||
<body name="link4" pos="0 0 0.1">
|
||||
<inertial pos="0 0 0.04" mass="0.4" diaginertia="0.004 0.004 0.002"/>
|
||||
<joint name="joint4" type="hinge" axis="0 0 1" range="-3.14 3.14"
|
||||
armature="0.2" damping="0.4"/>
|
||||
<geom type="capsule" fromto="0 0 0 0 0 0.08" size="0.025"/>
|
||||
<body name="link5" pos="0 0 0.08">
|
||||
<inertial pos="0 0 0.03" mass="0.2" diaginertia="0.002 0.002 0.001"/>
|
||||
<joint name="joint5" type="hinge" axis="0 1 0" range="-3.14 3.14"
|
||||
armature="0.1" damping="0.2"/>
|
||||
<geom type="capsule" fromto="0 0 0 0 0 0.06" size="0.02"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
<actuator>
|
||||
<motor name="act1" joint="joint1"/>
|
||||
<motor name="act2" joint="joint2"/>
|
||||
<motor name="act3" joint="joint3"/>
|
||||
<motor name="act4" joint="joint4"/>
|
||||
<motor name="act5" joint="joint5"/>
|
||||
</actuator>
|
||||
<sensor>
|
||||
<jointpos name="joint1_pos" joint="joint1"/>
|
||||
<jointpos name="joint2_pos" joint="joint2"/>
|
||||
<jointpos name="joint3_pos" joint="joint3"/>
|
||||
<jointpos name="joint4_pos" joint="joint4"/>
|
||||
<jointpos name="joint5_pos" joint="joint5"/>
|
||||
<jointvel name="joint1_vel" joint="joint1"/>
|
||||
<jointvel name="joint2_vel" joint="joint2"/>
|
||||
<jointvel name="joint3_vel" joint="joint3"/>
|
||||
<jointvel name="joint4_vel" joint="joint4"/>
|
||||
<jointvel name="joint5_vel" joint="joint5"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
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()
|
||||
|
||||
@@ -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 = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body><joint name="j1" type="hinge"/><geom size="0.1" mass="1"/>
|
||||
<body><joint name="j2" type="hinge"/><geom size="0.1" mass="1"/>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
_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 = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="box" pos="0 0 1">
|
||||
<freejoint/>
|
||||
<geom type="box" size=".1 .1 .1"/>
|
||||
</body>
|
||||
<body>
|
||||
<joint name="h1" type="hinge"/><geom size="0.1" mass="1"/>
|
||||
<body>
|
||||
<joint name="h2" type="hinge"/><geom size="0.1" mass="1"/>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
_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 = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="a" pos="0 0 1">
|
||||
<freejoint/><geom type="box" size=".1 .1 .1"/>
|
||||
</body>
|
||||
<body name="b" pos="1 0 1">
|
||||
<freejoint/><geom type="box" size=".1 .1 .1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
_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 = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body>
|
||||
<joint name="ball" type="ball"/><geom size="0.1" mass="1"/>
|
||||
<body>
|
||||
<joint name="h" type="hinge"/><geom size="0.1" mass="1"/>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
_verify_state_mapping(mujoco.MjModel.from_xml_string(xml))
|
||||
|
||||
@@ -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 = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="box" pos="0 0 1">
|
||||
<freejoint/>
|
||||
<geom type="box" size=".1 .1 .1"/>
|
||||
</body>
|
||||
<body>
|
||||
<joint name="h1" type="hinge"/>
|
||||
<geom size="0.1" mass="1"/>
|
||||
<body>
|
||||
<joint name="h2" type="hinge"/>
|
||||
<geom size="0.1" mass="1"/>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user