Levi Burner
302 lines
10 KiB
Python
302 lines
10 KiB
Python
# Copyright 2026 DeepMind Technologies Limited
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
|
|
"""Declarative signal transformation for system identification residuals."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from fnmatch import fnmatch
|
|
|
|
import mujoco
|
|
from mujoco.sysid._src import parameter
|
|
from mujoco.sysid._src import signal_modifier
|
|
from mujoco.sysid._src import timeseries
|
|
import numpy as np
|
|
|
|
|
|
class SignalTransform:
|
|
"""Declarative signal transformation replacing boilerplate modify_residual callbacks.
|
|
|
|
Usage::
|
|
|
|
transform = SignalTransform()
|
|
transform.delay("*_pos", params["delay_pos"])
|
|
transform.delay("*_torque", params["delay_torque"])
|
|
transform.gain("*_torque", params["torque_scale"], target="predicted")
|
|
transform.enable_sensors(cfg.sensors_enabled)
|
|
|
|
The ``apply`` method has the same signature as ``ModifyResidualFn`` and can
|
|
be passed directly to ``build_residual_fn(signal_transform=transform)``.
|
|
"""
|
|
|
|
def __init__(self, normalize: bool = True):
|
|
self._delays: list[tuple[str, str, parameter.Parameter]] = []
|
|
self._gains: list[tuple[str, str, str]] = []
|
|
self._biases: list[tuple[str, str, str]] = []
|
|
self._enabled_sensors: list[str] | None = None
|
|
self._sensor_weights: Mapping[str, float] | None = None
|
|
self.normalize = normalize
|
|
|
|
def delay(self, pattern: str, param: parameter.Parameter) -> None:
|
|
"""Register a delay for sensors matching *pattern* (fnmatch).
|
|
|
|
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(
|
|
self,
|
|
pattern: str,
|
|
param: parameter.Parameter,
|
|
target: str = "both",
|
|
) -> None:
|
|
"""Register a multiplicative gain for sensors matching *pattern*.
|
|
|
|
Args:
|
|
pattern: fnmatch pattern matched against sensor names.
|
|
param: Parameter whose ``.value`` is the gain factor.
|
|
target: One of ``"predicted"``, ``"measured"``, or ``"both"``.
|
|
"""
|
|
if target not in ("predicted", "measured", "both"):
|
|
raise ValueError(
|
|
f"target must be 'predicted', 'measured', or 'both', got {target!r}"
|
|
)
|
|
self._gains.append((pattern, param.name, target))
|
|
|
|
def bias(
|
|
self,
|
|
pattern: str,
|
|
param: parameter.Parameter,
|
|
target: str = "both",
|
|
) -> None:
|
|
"""Register an additive bias for sensors matching *pattern*.
|
|
|
|
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}"
|
|
)
|
|
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.
|
|
|
|
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.
|
|
|
|
Args:
|
|
weights: Mapping from sensor name to weight.
|
|
"""
|
|
self._sensor_weights = weights
|
|
|
|
# Private methods.
|
|
|
|
def _resolve_delays(
|
|
self,
|
|
sensor_names: list[str],
|
|
params: parameter.ParameterDict,
|
|
) -> dict[str, float]:
|
|
"""Resolve delay patterns to concrete sensor name -> delay value (last match wins)."""
|
|
resolved: dict[str, float] = {}
|
|
for pattern, param_name, _ in self._delays:
|
|
delay_value = params[param_name].value[0]
|
|
for name in sensor_names:
|
|
if fnmatch(name, pattern):
|
|
resolved[name] = delay_value
|
|
return resolved
|
|
|
|
def _compute_delay_bounds(self) -> tuple[float, float]:
|
|
"""Compute min/max delay across all registered delay params (deduplicated by name)."""
|
|
if not self._delays:
|
|
return 0.0, 0.0
|
|
seen: set[str] = set()
|
|
min_vals: list[float] = []
|
|
max_vals: list[float] = []
|
|
for _, param_name, param in self._delays:
|
|
if param_name in seen:
|
|
continue
|
|
seen.add(param_name)
|
|
min_vals.append(float(param.min_value[0]))
|
|
max_vals.append(float(param.max_value[0]))
|
|
return min(min_vals), max(max_vals)
|
|
|
|
def _get_sensor_names(self, ts: timeseries.TimeSeries) -> list[str]:
|
|
"""Extract sensor names from a TimeSeries signal_mapping."""
|
|
if ts.signal_mapping is None:
|
|
return []
|
|
return list(ts.signal_mapping.keys())
|
|
|
|
def _apply_gains_biases_reference(
|
|
self,
|
|
ts: timeseries.TimeSeries,
|
|
target_label: str,
|
|
params: parameter.ParameterDict,
|
|
) -> timeseries.TimeSeries:
|
|
"""Reference implementation: one full copy per gain/bias application."""
|
|
sensor_names = self._get_sensor_names(ts)
|
|
for pattern, param_name, target in self._gains:
|
|
if target != target_label and target != "both":
|
|
continue
|
|
for name in sensor_names:
|
|
if fnmatch(name, pattern):
|
|
ts = signal_modifier.apply_gain(ts, name, params[param_name])
|
|
for pattern, param_name, target in self._biases:
|
|
if target != target_label and target != "both":
|
|
continue
|
|
for name in sensor_names:
|
|
if fnmatch(name, pattern):
|
|
ts = signal_modifier.apply_bias(ts, name, params[param_name])
|
|
return ts
|
|
|
|
_VERIFY_GAINS_BIASES = False
|
|
|
|
def _apply_gains_biases(
|
|
self,
|
|
ts: timeseries.TimeSeries,
|
|
target_label: str,
|
|
params: parameter.ParameterDict,
|
|
) -> timeseries.TimeSeries:
|
|
"""Apply matching gains and biases to a timeseries for the given target label."""
|
|
sensor_names = self._get_sensor_names(ts)
|
|
data = ts.data.copy()
|
|
|
|
for pattern, param_name, target in self._gains:
|
|
if target != target_label and target != "both":
|
|
continue
|
|
for name in sensor_names:
|
|
if fnmatch(name, pattern):
|
|
indices = ts.get_indices(name)[1]
|
|
data[..., indices] *= params[param_name].value
|
|
|
|
for pattern, param_name, target in self._biases:
|
|
if target != target_label and target != "both":
|
|
continue
|
|
for name in sensor_names:
|
|
if fnmatch(name, pattern):
|
|
indices = ts.get_indices(name)[1]
|
|
data[..., indices] += params[param_name].value
|
|
|
|
result = timeseries.TimeSeries(ts.times, data, ts.signal_mapping)
|
|
|
|
if self._VERIFY_GAINS_BIASES:
|
|
ref = self._apply_gains_biases_reference(ts, target_label, params)
|
|
np.testing.assert_array_equal(result.data, ref.data)
|
|
|
|
return result
|
|
|
|
def apply(
|
|
self,
|
|
params: parameter.ParameterDict,
|
|
sensordata_predicted: timeseries.TimeSeries,
|
|
sensordata_measured: timeseries.TimeSeries,
|
|
model: mujoco.MjModel,
|
|
return_pred_all: bool,
|
|
state: np.ndarray | None = None,
|
|
sensor_weights: Mapping[str, float] | None = None,
|
|
) -> tuple[np.ndarray, timeseries.TimeSeries, timeseries.TimeSeries]:
|
|
"""Apply all registered transforms and compute the residual.
|
|
|
|
Signature matches :data:`ModifyResidualFn` so this method can be passed
|
|
directly as ``modify_residual`` to :func:`model_residual`.
|
|
|
|
Pipeline: window measured data, resample + delay predicted data, apply
|
|
gains/biases, weighted diff, normalise, slice to enabled sensors.
|
|
|
|
Args:
|
|
params: Current parameter dictionary.
|
|
sensordata_predicted: Predicted sensor TimeSeries from rollout.
|
|
sensordata_measured: Measured sensor TimeSeries (ground truth).
|
|
model: MuJoCo model for sensor metadata lookup.
|
|
return_pred_all: Whether to return all predicted signals.
|
|
state: Unused; part of the ModifyResidualFn signature.
|
|
sensor_weights: Per-sensor weighting for the residual.
|
|
|
|
Returns:
|
|
``(residual_array, predicted_ts, measured_ts)`` — the residual matrix
|
|
and the (possibly sliced) predicted/measured TimeSeries.
|
|
"""
|
|
del state # Part of ModifyResidualFn signature but unused here.
|
|
sensor_names = self._get_sensor_names(sensordata_predicted)
|
|
|
|
# 1. Resolve delays and compute bounds.
|
|
sensor_delays = self._resolve_delays(sensor_names, params)
|
|
min_delay, max_delay = self._compute_delay_bounds()
|
|
|
|
# 2. Window measured data.
|
|
sensordata_measured = signal_modifier.apply_delayed_ts_window(
|
|
sensordata_measured, sensordata_predicted, min_delay, max_delay
|
|
)
|
|
|
|
# 3. Resample and delay predicted data.
|
|
if sensor_delays:
|
|
sensordata_predicted = signal_modifier.apply_resample_and_delay(
|
|
sensordata_predicted,
|
|
sensordata_measured.times,
|
|
0.0,
|
|
sensor_delays=sensor_delays,
|
|
)
|
|
else:
|
|
sensordata_predicted = sensordata_predicted.resample(
|
|
sensordata_measured.times
|
|
)
|
|
|
|
# 4. Apply gains and biases.
|
|
sensordata_predicted = self._apply_gains_biases(
|
|
sensordata_predicted, "predicted", params
|
|
)
|
|
sensordata_measured = self._apply_gains_biases(
|
|
sensordata_measured, "measured", params
|
|
)
|
|
|
|
# 5. Weighted diff.
|
|
weights = sensor_weights or self._sensor_weights
|
|
res = signal_modifier.weighted_diff(
|
|
predicted_data=sensordata_predicted.data,
|
|
measured_data=sensordata_measured.data,
|
|
model=model,
|
|
sensor_weights=weights,
|
|
)
|
|
|
|
# 6. Normalize.
|
|
if self.normalize:
|
|
res = signal_modifier.normalize_residual(res, sensordata_measured.data)
|
|
|
|
# 7. Slice to enabled sensors.
|
|
if not return_pred_all and self._enabled_sensors is not None:
|
|
indices = signal_modifier.get_sensor_indices(model, self._enabled_sensors)
|
|
sensordata_predicted = timeseries.TimeSeries(
|
|
sensordata_predicted.times,
|
|
sensordata_predicted.data[:, indices],
|
|
)
|
|
sensordata_measured = timeseries.TimeSeries(
|
|
sensordata_measured.times,
|
|
sensordata_measured.data[:, indices],
|
|
)
|
|
res = res[:, indices]
|
|
|
|
return res, sensordata_predicted, sensordata_measured
|