Merge pull request #3079 from aftersomemath:sysid-pr
PiperOrigin-RevId: 868229512 Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea
This commit is contained in:
Copybara-Service
@@ -0,0 +1,15 @@
|
||||
# Copyright 2026 DeepMind Technologies Limited
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
@@ -0,0 +1,224 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""Shared fixtures for mujoco.sysid tests."""
|
||||
|
||||
import mujoco
|
||||
from mujoco.sysid._src import parameter
|
||||
from mujoco.sysid._src import timeseries
|
||||
from mujoco.sysid._src.model_modifier import _infer_inertial
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inline model XML strings — no external file dependencies
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
BOX_XML = """\
|
||||
<mujoco model="box">
|
||||
<compiler angle="radian"/>
|
||||
<option integrator="implicitfast" timestep="0.001" cone="elliptic">
|
||||
<flag contact="enable"/>
|
||||
</option>
|
||||
<asset>
|
||||
<texture type="2d" name="groundplane" builtin="checker" width="300" height="300"/>
|
||||
<material name="groundplane" texture="groundplane"/>
|
||||
</asset>
|
||||
<worldbody>
|
||||
<body name="floor">
|
||||
<geom name="floor" type="box" size="100 100 0.01" pos="0 0 -0.01"/>
|
||||
</body>
|
||||
<body name="box" pos="0 0 0.102">
|
||||
<freejoint name="box_free"/>
|
||||
<inertial pos="0 0 0" mass="5" diaginertia="0.0333 0.0333 0.0333"/>
|
||||
<geom name="box" type="box" size="0.1 0.1 0.1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<actuator>
|
||||
<motor name="push_x" joint="box_free" gear="1 0 0 0 0 0"
|
||||
ctrllimited="true" ctrlrange="-20 20"/>
|
||||
</actuator>
|
||||
<contact>
|
||||
<pair name="box_floor" geom1="box" geom2="floor"
|
||||
friction="1.6 0.005 0.0001" solref="0.01 1.0"/>
|
||||
</contact>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
ARM_XML = """\
|
||||
<mujoco model="test_arm">
|
||||
<compiler angle="radian" autolimits="true"/>
|
||||
<option integrator="implicitfast">
|
||||
<flag contact="disable"/>
|
||||
</option>
|
||||
<asset>
|
||||
<texture name="tex1" type="2d" builtin="checker" width="64" height="64"/>
|
||||
<material name="mat1" texture="tex1"/>
|
||||
</asset>
|
||||
<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"/>
|
||||
<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"/>
|
||||
<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"/>
|
||||
<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"/>
|
||||
<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"/>
|
||||
<geom type="capsule" fromto="0 0 0 0 0 0.06" size="0.02"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
<actuator>
|
||||
<position name="actuator1" joint="joint1" kp="400" kv="40"/>
|
||||
<position name="actuator2" joint="joint2" kp="400" kv="40"/>
|
||||
<position name="actuator3" joint="joint3" kp="400" kv="40"/>
|
||||
<position name="actuator4" joint="joint4" kp="400" kv="40"/>
|
||||
<position name="actuator5" joint="joint5" kp="200" kv="20"/>
|
||||
</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"/>
|
||||
<jointactuatorfrc name="joint1_torque" joint="joint1"/>
|
||||
<jointactuatorfrc name="joint2_torque" joint="joint2"/>
|
||||
<jointactuatorfrc name="joint3_torque" joint="joint3"/>
|
||||
<jointactuatorfrc name="joint4_torque" joint="joint4"/>
|
||||
<jointactuatorfrc name="joint5_torque" joint="joint5"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
OSCILLATOR_XML = """\
|
||||
<mujoco>
|
||||
<compiler autolimits="true"/>
|
||||
<option>
|
||||
<flag gravity="disable" contact="disable" limit="disable"/>
|
||||
</option>
|
||||
<worldbody>
|
||||
<body name="mass">
|
||||
<joint name="j1" axis="1 0 0" type="slide" range="0 0.5"
|
||||
stiffness="10" damping="1"/>
|
||||
<geom type="box" size=".1 .1 .1" mass="0.1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def box_spec() -> mujoco.MjSpec:
|
||||
return mujoco.MjSpec.from_string(BOX_XML)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def box_model(box_spec) -> mujoco.MjModel:
|
||||
return box_spec.compile()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def arm_spec() -> mujoco.MjSpec:
|
||||
"""Minimal 5-joint arm with sensors, actuators, textures/materials."""
|
||||
return mujoco.MjSpec.from_string(ARM_XML)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def arm_model(arm_spec) -> mujoco.MjSpec:
|
||||
return arm_spec.compile()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def oscillator_spec() -> mujoco.MjSpec:
|
||||
"""Single-body oscillator with implicit (geom-based) inertia."""
|
||||
return mujoco.MjSpec.from_string(OSCILLATOR_XML)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def simple_timeseries() -> timeseries.TimeSeries:
|
||||
"""A TimeSeries with 5 data points and 2 columns: y = [x^2, 2*x^2]."""
|
||||
times = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
|
||||
data = np.array([
|
||||
[0.0, 0.0],
|
||||
[1.0, 2.0],
|
||||
[4.0, 8.0],
|
||||
[9.0, 18.0],
|
||||
[16.0, 32.0],
|
||||
])
|
||||
return timeseries.TimeSeries(times=times, data=data)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def box_params(box_spec) -> parameter.ParameterDict:
|
||||
"""ParameterDict for box model with modifier callbacks."""
|
||||
del box_spec
|
||||
pdict = parameter.ParameterDict()
|
||||
pdict.add(
|
||||
parameter.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]
|
||||
),
|
||||
)
|
||||
)
|
||||
pdict.add(
|
||||
parameter.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]
|
||||
),
|
||||
)
|
||||
)
|
||||
pdict.add(
|
||||
parameter.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]
|
||||
),
|
||||
)
|
||||
)
|
||||
return pdict
|
||||
@@ -0,0 +1,249 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""End-to-end integration test using the box model."""
|
||||
|
||||
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 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."""
|
||||
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)
|
||||
|
||||
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)
|
||||
state = np.squeeze(state, axis=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,
|
||||
)
|
||||
|
||||
return control_ts, sensordata, initial_state
|
||||
|
||||
|
||||
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]
|
||||
),
|
||||
)
|
||||
)
|
||||
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]
|
||||
),
|
||||
)
|
||||
)
|
||||
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]
|
||||
),
|
||||
)
|
||||
)
|
||||
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]
|
||||
),
|
||||
)
|
||||
)
|
||||
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]
|
||||
),
|
||||
)
|
||||
)
|
||||
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])),
|
||||
},
|
||||
)
|
||||
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.
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
save_results(
|
||||
experiment_results_folder=tmpdir,
|
||||
models_sequences=models_sequences,
|
||||
initial_params=params,
|
||||
opt_params=opt_params,
|
||||
opt_result=opt_result,
|
||||
residual_fn=residual_fn,
|
||||
)
|
||||
result_dir = pathlib.Path(tmpdir)
|
||||
assert (result_dir / "params_x_0.yaml").exists()
|
||||
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()
|
||||
@@ -0,0 +1,131 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""Tests for the model_modifier module."""
|
||||
|
||||
import mujoco
|
||||
from mujoco.sysid._src import model_modifier
|
||||
import numpy as np
|
||||
|
||||
|
||||
def test_apply_pgain(arm_spec):
|
||||
"""Setting a P gain correctly configures the underlying actuator parameters."""
|
||||
actuator_name = "actuator5"
|
||||
pgain_value = 74
|
||||
|
||||
modified_spec = model_modifier.apply_pgain(
|
||||
arm_spec, actuator_name, pgain_value
|
||||
)
|
||||
model = modified_spec.compile()
|
||||
|
||||
assert model.actuator(actuator_name).gainprm[0] == pgain_value
|
||||
assert model.actuator(actuator_name).biasprm[1] == -pgain_value
|
||||
|
||||
|
||||
def test_apply_dgain(arm_spec):
|
||||
"""Setting a D gain correctly configures the underlying actuator parameters."""
|
||||
actuator_name = "actuator5"
|
||||
dgain_value = 1.2
|
||||
|
||||
modified_spec = model_modifier.apply_dgain(
|
||||
arm_spec, actuator_name, dgain_value
|
||||
)
|
||||
model = modified_spec.compile()
|
||||
|
||||
assert model.actuator(actuator_name).biasprm[2] == -dgain_value
|
||||
|
||||
|
||||
def test_apply_pdgain(arm_spec):
|
||||
"""Setting P and D gains together from a single array configures both correctly."""
|
||||
actuator_name = "actuator5"
|
||||
pdgain_value = np.array([74, 1.2])
|
||||
|
||||
modified_spec = model_modifier.apply_pdgain(
|
||||
arm_spec, actuator_name, pdgain_value
|
||||
)
|
||||
model = modified_spec.compile()
|
||||
|
||||
assert model.actuator(actuator_name).gainprm[0] == pdgain_value[0]
|
||||
assert model.actuator(actuator_name).biasprm[1] == -pdgain_value[0]
|
||||
assert model.actuator(actuator_name).biasprm[2] == -pdgain_value[1]
|
||||
|
||||
|
||||
def test_apply_body_mass_explicit(arm_spec):
|
||||
"""Bodies with inertia defined in XML: changing mass proportionally scales inertia."""
|
||||
body_name = "link1"
|
||||
model = arm_spec.compile()
|
||||
original_mass = model.body(body_name).mass[0]
|
||||
original_inertia = model.body(body_name).inertia
|
||||
del model
|
||||
|
||||
scale = 3.3
|
||||
new_mass = scale * original_mass
|
||||
|
||||
modified_spec = model_modifier.apply_body_mass_ipos(
|
||||
arm_spec, body_name, mass=new_mass, rot_inertia_scale=True
|
||||
)
|
||||
model = modified_spec.compile()
|
||||
|
||||
assert model.body(body_name).mass == new_mass
|
||||
np.testing.assert_allclose(
|
||||
model.body(body_name).inertia, original_inertia * scale
|
||||
)
|
||||
|
||||
|
||||
def test_apply_body_mass_implicit(oscillator_spec):
|
||||
"""Bodies with inertia inferred from geoms: changing mass proportionally scales inertia."""
|
||||
body_name = "mass"
|
||||
model = oscillator_spec.compile()
|
||||
original_mass = model.body(body_name).mass[0]
|
||||
original_inertia = model.body(body_name).inertia
|
||||
del model
|
||||
|
||||
scale = 0.077
|
||||
new_mass = scale * original_mass
|
||||
|
||||
modified_spec = model_modifier.apply_body_mass_ipos(
|
||||
oscillator_spec, body_name, mass=new_mass, rot_inertia_scale=True
|
||||
)
|
||||
model = modified_spec.compile()
|
||||
|
||||
assert model.body(body_name).mass == new_mass
|
||||
np.testing.assert_allclose(
|
||||
model.body(body_name).inertia, original_inertia * scale
|
||||
)
|
||||
|
||||
|
||||
def test_remove_visuals(arm_spec):
|
||||
"""Stripping visuals removes all textures and materials for faster compilation."""
|
||||
cleaned_spec = model_modifier.remove_visuals(arm_spec)
|
||||
assert not cleaned_spec.textures
|
||||
assert not cleaned_spec.materials
|
||||
|
||||
|
||||
def test_apply_param_modifiers(box_spec, box_params):
|
||||
"""The full modifier pipeline applies parameter callbacks and produces an updated model."""
|
||||
spec = box_spec.copy()
|
||||
original_model = spec.compile()
|
||||
original_mass = original_model.body("box").mass[0]
|
||||
|
||||
# Change box_mass parameter.
|
||||
box_params["box_mass"].update_from_vector(np.array([5.3]))
|
||||
|
||||
modified_model = model_modifier.apply_param_modifiers(box_params, spec)
|
||||
assert modified_model.body("box").mass[0] != original_mass
|
||||
np.testing.assert_allclose(modified_model.body("box").mass[0], 5.3, atol=1e-6)
|
||||
|
||||
# Also verify apply_param_modifiers_spec returns MjSpec.
|
||||
spec2 = box_spec.copy()
|
||||
result = model_modifier.apply_param_modifiers_spec(box_params, spec2)
|
||||
assert isinstance(result, mujoco.MjSpec)
|
||||
@@ -0,0 +1,148 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""Tests for the Parameter and ParameterDict classes."""
|
||||
|
||||
from mujoco.sysid._src import parameter
|
||||
import numpy as np
|
||||
|
||||
|
||||
def test_scalar_parameter():
|
||||
"""A single-valued parameter round-trips through vector conversion, sampling, and reset."""
|
||||
param = parameter.Parameter("test", 1.0, 0.5, 2.0)
|
||||
|
||||
assert param.name == "test"
|
||||
assert param.size == 1
|
||||
assert param.shape == (1,)
|
||||
assert param.nominal == 1.0
|
||||
assert param.value == 1.0
|
||||
assert param.min_value == 0.5
|
||||
assert param.max_value == 2.0
|
||||
|
||||
np.testing.assert_array_equal(param.as_vector(), [1.0])
|
||||
param.update_from_vector(np.array([1.5]))
|
||||
np.testing.assert_array_equal(param.value, [1.5])
|
||||
np.testing.assert_array_equal(param.as_vector(), [1.5])
|
||||
|
||||
lower, upper = param.get_bounds()
|
||||
np.testing.assert_array_equal(lower, [0.5])
|
||||
np.testing.assert_array_equal(upper, [2.0])
|
||||
|
||||
param.reset()
|
||||
np.testing.assert_array_equal(param.value, [1.0])
|
||||
|
||||
rng = np.random.default_rng(42)
|
||||
sample = param.sample(rng)
|
||||
assert 0.5 <= sample[0] <= 2.0
|
||||
|
||||
|
||||
def test_vector_parameter():
|
||||
"""A multi-valued parameter preserves element-wise bounds and resets correctly."""
|
||||
param = parameter.Parameter("test_vector", [1.0, 2.0], [0.5, 1.0], [2.0, 3.0])
|
||||
|
||||
assert param.name == "test_vector"
|
||||
assert param.size == 2
|
||||
assert param.shape == (2,)
|
||||
np.testing.assert_array_equal(param.nominal, [1.0, 2.0])
|
||||
np.testing.assert_array_equal(param.value, [1.0, 2.0])
|
||||
np.testing.assert_array_equal(param.min_value, [0.5, 1.0])
|
||||
np.testing.assert_array_equal(param.max_value, [2.0, 3.0])
|
||||
|
||||
np.testing.assert_array_equal(param.as_vector(), [1.0, 2.0])
|
||||
param.update_from_vector(np.array([1.5, 2.5]))
|
||||
np.testing.assert_array_equal(param.value, [1.5, 2.5])
|
||||
|
||||
lower, upper = param.get_bounds()
|
||||
np.testing.assert_array_equal(lower, [0.5, 1.0])
|
||||
np.testing.assert_array_equal(upper, [2.0, 3.0])
|
||||
|
||||
param.reset()
|
||||
np.testing.assert_array_equal(param.value, [1.0, 2.0])
|
||||
|
||||
|
||||
def test_parameter_dict():
|
||||
"""A dict of mixed scalar/vector params flattens to one vector and reconstructs."""
|
||||
param1 = parameter.Parameter("param1", 1.0, 0.5, 2.0)
|
||||
param2 = parameter.Parameter("param2", [2.0, 3.0], [1.0, 2.0], [3.0, 4.0])
|
||||
params = parameter.ParameterDict({"param1": param1, "param2": param2})
|
||||
|
||||
assert params.size == 3 # 1 + 2
|
||||
assert len(params) == 2
|
||||
|
||||
assert params["param1"] is param1
|
||||
assert params["param2"] is param2
|
||||
|
||||
np.testing.assert_array_equal(params.as_vector(), [1.0, 2.0, 3.0])
|
||||
|
||||
params.update_from_vector(np.array([1.5, 2.5, 3.5]))
|
||||
np.testing.assert_array_equal(params["param1"].value, [1.5])
|
||||
np.testing.assert_array_equal(params["param2"].value, [2.5, 3.5])
|
||||
|
||||
lower, upper = params.get_bounds()
|
||||
np.testing.assert_array_equal(lower, [0.5, 1.0, 2.0])
|
||||
np.testing.assert_array_equal(upper, [2.0, 3.0, 4.0])
|
||||
|
||||
params.reset()
|
||||
np.testing.assert_array_equal(params["param1"].value, [1.0])
|
||||
np.testing.assert_array_equal(params["param2"].value, [2.0, 3.0])
|
||||
|
||||
rng = np.random.default_rng(42)
|
||||
sample = params.sample(rng=rng)
|
||||
assert len(sample) == 3
|
||||
|
||||
|
||||
def test_save_and_load_round_trip(tmp_path):
|
||||
"""Saving to YAML and loading back recovers modified values, nominals, and bounds."""
|
||||
param1 = parameter.Parameter("p1", 1.0, 0.0, 2.0)
|
||||
param2 = parameter.Parameter("p2", [3.0, 4.0], [1.0, 2.0], [5.0, 6.0])
|
||||
params = parameter.ParameterDict({"p1": param1, "p2": param2})
|
||||
params.update_from_vector(np.array([0.7, 3.5, 4.5]))
|
||||
|
||||
path = tmp_path / "params.yaml"
|
||||
params.save_to_disk(path)
|
||||
|
||||
loaded = parameter.ParameterDict.load_from_disk(path)
|
||||
np.testing.assert_array_equal(loaded.as_vector(), [0.7, 3.5, 4.5])
|
||||
np.testing.assert_array_equal(loaded["p1"].nominal, [1.0])
|
||||
np.testing.assert_array_equal(loaded["p2"].min_value, [1.0, 2.0])
|
||||
|
||||
|
||||
def test_randomize_stays_in_bounds():
|
||||
"""Randomized parameter values always stay within their declared bounds."""
|
||||
param1 = parameter.Parameter("a", 5.0, 2.0, 8.0)
|
||||
param2 = parameter.Parameter("b", [1.0, 2.0], [0.0, 0.0], [3.0, 3.0])
|
||||
params = parameter.ParameterDict({"a": param1, "b": param2})
|
||||
|
||||
rng = np.random.default_rng(0)
|
||||
for _ in range(10):
|
||||
params.randomize(rng=rng)
|
||||
lower, upper = params.get_bounds()
|
||||
vec = params.as_vector()
|
||||
assert np.all(vec >= lower)
|
||||
assert np.all(vec <= upper)
|
||||
|
||||
|
||||
def test_frozen_param_excluded():
|
||||
"""Freezing a parameter hides it from the optimizer: excluded from vector ops."""
|
||||
p1 = parameter.Parameter("free", 1.0, 0.0, 2.0)
|
||||
p2 = parameter.Parameter("frozen", 5.0, 3.0, 7.0, frozen=True)
|
||||
params = parameter.ParameterDict({"free": p1, "frozen": p2})
|
||||
|
||||
assert params.size == 1
|
||||
np.testing.assert_array_equal(params.as_vector(), [1.0])
|
||||
|
||||
params.update_from_vector(np.array([1.5]))
|
||||
np.testing.assert_array_equal(params["free"].value, [1.5])
|
||||
# Frozen param unchanged.
|
||||
np.testing.assert_array_equal(params["frozen"].value, [5.0])
|
||||
@@ -0,0 +1,391 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""Tests for signal_modifier and SignalTransform."""
|
||||
|
||||
from mujoco.sysid._src import parameter
|
||||
from mujoco.sysid._src import signal_modifier
|
||||
from mujoco.sysid._src import timeseries
|
||||
from mujoco.sysid._src.parameter import Parameter
|
||||
from mujoco.sysid._src.parameter import ParameterDict
|
||||
from mujoco.sysid._src.signal_transform import SignalTransform
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Helpers
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def _make_pdict(*params: Parameter) -> ParameterDict:
|
||||
pdict = ParameterDict()
|
||||
for p in params:
|
||||
pdict.add(p)
|
||||
return pdict
|
||||
|
||||
|
||||
def _make_arm_sensor_ts(arm_model):
|
||||
"""Create a synthetic TimeSeries with signal_mapping matching arm sensors."""
|
||||
n_steps = 10
|
||||
n_sensors = arm_model.nsensordata
|
||||
times = np.linspace(0, 1, n_steps)
|
||||
data = np.random.default_rng(42).standard_normal((n_steps, n_sensors))
|
||||
|
||||
mapping = {}
|
||||
for i in range(arm_model.nsensor):
|
||||
name = arm_model.sensor(i).name
|
||||
adr = arm_model.sensor_adr[i]
|
||||
dim = arm_model.sensor_dim[i]
|
||||
mapping[name] = (timeseries.SignalType.MjSensor, np.arange(adr, adr + dim))
|
||||
|
||||
return timeseries.TimeSeries(times, data, signal_mapping=mapping)
|
||||
|
||||
|
||||
def _make_resample_ts(n_steps, n_cols, seed=0):
|
||||
"""Deterministic TimeSeries for resample tests."""
|
||||
rng = np.random.default_rng(seed)
|
||||
times = np.linspace(0, 1, n_steps)
|
||||
data = rng.standard_normal((n_steps, n_cols))
|
||||
mapping = {
|
||||
f"s{i}": (timeseries.SignalType.MjSensor, np.array([i]))
|
||||
for i in range(n_cols)
|
||||
}
|
||||
return timeseries.TimeSeries(times, data, signal_mapping=mapping)
|
||||
|
||||
|
||||
def _make_transform_sensor_ts(n_steps=50, n_sensors=15, seed=0):
|
||||
"""Deterministic TimeSeries with named MjSensor columns for transforms."""
|
||||
rng = np.random.default_rng(seed)
|
||||
times = np.linspace(0, 1, n_steps)
|
||||
data = rng.standard_normal((n_steps, n_sensors))
|
||||
mapping = {
|
||||
f"joint{i + 1}_pos": (timeseries.SignalType.MjSensor, np.array([i]))
|
||||
for i in range(5)
|
||||
}
|
||||
mapping.update({
|
||||
f"joint{i - 4}_vel": (timeseries.SignalType.MjSensor, np.array([i]))
|
||||
for i in range(5, 10)
|
||||
})
|
||||
mapping.update({
|
||||
f"joint{i - 9}_torque": (timeseries.SignalType.MjSensor, np.array([i]))
|
||||
for i in range(10, 15)
|
||||
})
|
||||
return timeseries.TimeSeries(times, data, signal_mapping=mapping)
|
||||
|
||||
|
||||
def _run_both(ts, times, default_delay, sensor_delays, predicted_data):
|
||||
"""Run grouped and column-wise implementations, return both."""
|
||||
delays = signal_modifier._build_per_column_delays(
|
||||
ts, default_delay, sensor_delays, predicted_data
|
||||
)
|
||||
reference = signal_modifier._apply_resample_and_delay_columnwise(
|
||||
ts, times, delays
|
||||
)
|
||||
result = signal_modifier.apply_resample_and_delay(
|
||||
ts,
|
||||
times,
|
||||
default_delay,
|
||||
sensor_delays=sensor_delays,
|
||||
predicted_data=predicted_data,
|
||||
)
|
||||
return result.data, reference
|
||||
|
||||
|
||||
def _run_gains_biases_both(transform, ts, target_label, params):
|
||||
"""Run new and reference implementations, return both results."""
|
||||
new_result = transform._apply_gains_biases(ts, target_label, params)
|
||||
ref_result = transform._apply_gains_biases_reference(ts, target_label, params)
|
||||
return new_result, ref_result
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# signal_modifier: get_sensor_indices
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def test_get_sensor_indices(arm_model):
|
||||
"""Sensor name lookup returns the right data column indices for one or many sensors."""
|
||||
indices = signal_modifier.get_sensor_indices(arm_model, "joint1_pos")
|
||||
assert isinstance(indices, list)
|
||||
assert len(indices) == 1
|
||||
|
||||
indices = signal_modifier.get_sensor_indices(
|
||||
arm_model, ["joint1_pos", "joint2_pos"]
|
||||
)
|
||||
assert len(indices) == 2
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# signal_modifier: apply_gain / apply_bias
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def test_apply_gain(arm_model):
|
||||
"""Gain scaling affects only the named sensor's columns, leaving others untouched."""
|
||||
ts = _make_arm_sensor_ts(arm_model)
|
||||
gain = parameter.Parameter("gain", 2.0, 0.5, 3.0)
|
||||
|
||||
result = signal_modifier.apply_gain(ts, "joint1_torque", gain)
|
||||
|
||||
idx = ts.get_indices("joint1_torque")[1]
|
||||
np.testing.assert_allclose(result.data[:, idx], ts.data[:, idx] * 2.0)
|
||||
other = [i for i in range(ts.data.shape[1]) if i not in idx]
|
||||
np.testing.assert_array_equal(result.data[:, other], ts.data[:, other])
|
||||
|
||||
|
||||
def test_apply_bias(arm_model):
|
||||
"""Bias offset affects only the named sensor's columns, leaving others untouched."""
|
||||
ts = _make_arm_sensor_ts(arm_model)
|
||||
bias = parameter.Parameter("bias", 0.5, -1.0, 1.0)
|
||||
|
||||
result = signal_modifier.apply_bias(ts, "joint1_pos", bias)
|
||||
|
||||
idx = ts.get_indices("joint1_pos")[1]
|
||||
np.testing.assert_allclose(result.data[:, idx], ts.data[:, idx] + 0.5)
|
||||
other = [i for i in range(ts.data.shape[1]) if i not in idx]
|
||||
np.testing.assert_array_equal(result.data[:, other], ts.data[:, other])
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# signal_modifier: apply_delayed_ts_window
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def test_apply_delayed_ts_window(arm_model):
|
||||
"""Time-windowing crops timestamps to the overlapping region between two series."""
|
||||
ts = _make_arm_sensor_ts(arm_model)
|
||||
ts_delayed = _make_arm_sensor_ts(arm_model)
|
||||
|
||||
result = signal_modifier.apply_delayed_ts_window(
|
||||
ts, ts_delayed, min_delay=0.0, max_delay=0.0
|
||||
)
|
||||
assert result.times[0] >= ts_delayed.times[0]
|
||||
assert result.times[-1] <= ts_delayed.times[-1]
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# signal_modifier: weighted_diff / normalize_residual
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def test_weighted_diff_basic():
|
||||
"""Without weights, the residual is simply measured minus predicted."""
|
||||
predicted = np.array([[1.0, 2.0], [3.0, 4.0]])
|
||||
measured = np.array([[1.1, 2.2], [3.3, 4.4]])
|
||||
result = signal_modifier.weighted_diff(predicted, measured)
|
||||
np.testing.assert_allclose(result, measured - predicted)
|
||||
|
||||
|
||||
def test_weighted_diff_with_weights(arm_model):
|
||||
"""Sensor weights let you emphasize or de-emphasize specific channels in the residual."""
|
||||
n = arm_model.nsensordata
|
||||
predicted = np.ones((5, n))
|
||||
measured = np.ones((5, n)) * 2.0
|
||||
weights = {"joint1_pos": 0.5}
|
||||
result = signal_modifier.weighted_diff(
|
||||
predicted, measured, arm_model, weights
|
||||
)
|
||||
idx = signal_modifier.get_sensor_indices(arm_model, "joint1_pos")
|
||||
np.testing.assert_allclose(result[:, idx], 0.5)
|
||||
other = [i for i in range(n) if i not in idx]
|
||||
np.testing.assert_allclose(result[:, other], 1.0)
|
||||
|
||||
|
||||
def test_normalize_residual():
|
||||
"""Normalization makes residuals comparable across sensors with different scales."""
|
||||
residual = np.array([[2.0, 4.0], [6.0, 8.0]])
|
||||
measured = np.array([[1.0, 2.0], [3.0, 4.0]])
|
||||
result = signal_modifier.normalize_residual(residual, measured)
|
||||
norm = np.linalg.norm(measured, axis=0) / np.sqrt(2)
|
||||
np.testing.assert_allclose(result, residual / norm)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# signal_modifier: resample_and_delay grouped vs columnwise equivalence
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def test_resample_delay_mixed_delays():
|
||||
"""Optimized grouped resampling gives identical results to naive per-column resampling."""
|
||||
ts = _make_resample_ts(200, 8, seed=42)
|
||||
out_times = np.linspace(0.05, 0.95, 150)
|
||||
sensor_delays = {
|
||||
"s0": 0.01,
|
||||
"s1": 0.01,
|
||||
"s2": 0.01,
|
||||
"s3": 0.03,
|
||||
"s4": 0.03,
|
||||
}
|
||||
result, reference = _run_both(ts, out_times, 0.0, sensor_delays, True)
|
||||
np.testing.assert_array_equal(result, reference)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# SignalTransform: pattern matching
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestPatternMatching:
|
||||
"""Tests for signal transform pattern matching."""
|
||||
|
||||
def test_basic_glob(self):
|
||||
"""Wildcard patterns select the right sensors (e.g.
|
||||
|
||||
'*_pos' matches positions only).
|
||||
"""
|
||||
transform = SignalTransform()
|
||||
delay_param = Parameter("delay", [0.01], [0.0], [0.05])
|
||||
transform.delay("*_pos", delay_param)
|
||||
pdict = _make_pdict(delay_param)
|
||||
|
||||
resolved = transform._resolve_delays(
|
||||
["joint1_pos", "joint2_pos", "joint1_vel"], pdict
|
||||
)
|
||||
assert "joint1_pos" in resolved
|
||||
assert "joint2_pos" in resolved
|
||||
assert "joint1_vel" not in resolved
|
||||
assert resolved["joint1_pos"] == pytest.approx(0.01)
|
||||
|
||||
def test_last_match_wins(self):
|
||||
"""When patterns overlap, the last one registered takes priority."""
|
||||
transform = SignalTransform()
|
||||
general_delay = Parameter("delay_general", [0.01], [0.0], [0.05])
|
||||
specific_delay = Parameter("delay_specific", [0.05], [0.0], [0.1])
|
||||
transform.delay("*_torque", general_delay)
|
||||
transform.delay("joint5_torque", specific_delay)
|
||||
pdict = _make_pdict(general_delay, specific_delay)
|
||||
|
||||
resolved = transform._resolve_delays(
|
||||
["joint1_torque", "joint5_torque"], pdict
|
||||
)
|
||||
assert resolved["joint1_torque"] == pytest.approx(0.01)
|
||||
assert resolved["joint5_torque"] == pytest.approx(0.05)
|
||||
|
||||
def test_no_match(self):
|
||||
"""Patterns that don't match any sensor names produce no delay entries."""
|
||||
transform = SignalTransform()
|
||||
delay_param = Parameter("delay", [0.01], [0.0], [0.05])
|
||||
transform.delay("*_pos", delay_param)
|
||||
pdict = _make_pdict(delay_param)
|
||||
|
||||
resolved = transform._resolve_delays(["joint1_vel", "joint2_vel"], pdict)
|
||||
assert not resolved
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# SignalTransform: delay bounds
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestDelayBounds:
|
||||
"""Tests for delay bound computation."""
|
||||
|
||||
def test_single_param(self):
|
||||
"""The min/max delay window is derived from a parameter's declared bounds."""
|
||||
transform = SignalTransform()
|
||||
delay_param = Parameter("delay", [0.01], [-0.02], [0.05])
|
||||
transform.delay("*_pos", delay_param)
|
||||
|
||||
min_d, max_d = transform._compute_delay_bounds()
|
||||
assert min_d == pytest.approx(-0.02)
|
||||
assert max_d == pytest.approx(0.05)
|
||||
|
||||
def test_dedup_by_name(self):
|
||||
"""Reusing one delay param across patterns doesn't double-count its bounds."""
|
||||
transform = SignalTransform()
|
||||
delay_param = Parameter("delay", [0.01], [-0.01], [0.05])
|
||||
transform.delay("*_pos", delay_param)
|
||||
transform.delay("*_vel", delay_param)
|
||||
|
||||
min_d, max_d = transform._compute_delay_bounds()
|
||||
assert min_d == pytest.approx(-0.01)
|
||||
assert max_d == pytest.approx(0.05)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# SignalTransform: edge cases
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Tests for edge cases in signal transforms."""
|
||||
|
||||
def test_enable_sensors_stores_copy(self):
|
||||
"""The sensor list is defensively copied so callers can't mutate it after the fact."""
|
||||
transform = SignalTransform()
|
||||
sensors = ["a", "b"]
|
||||
transform.enable_sensors(sensors)
|
||||
sensors.append("c")
|
||||
assert transform._enabled_sensors == ["a", "b"]
|
||||
|
||||
def test_invalid_target(self):
|
||||
"""Typos in the target argument ('predicted'/'measured'/'both') are caught early."""
|
||||
transform = SignalTransform()
|
||||
param = Parameter("gain", [1.0], [0.5], [2.0])
|
||||
with pytest.raises(ValueError, match="target must be"):
|
||||
transform.gain("*", param, target="invalid")
|
||||
|
||||
bias_param = Parameter("bias", [0.0], [-1.0], [1.0])
|
||||
with pytest.raises(ValueError, match="target must be"):
|
||||
transform.bias("*", bias_param, target="invalid")
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# SignalTransform: _apply_gains_biases equivalence
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestApplyGainsBiasesEquivalence:
|
||||
"""Tests for gains/biases application equivalence."""
|
||||
|
||||
def test_gains_and_biases_mixed(self):
|
||||
"""Applying gains and biases together produces the same result as the reference path."""
|
||||
ts = _make_transform_sensor_ts()
|
||||
gain = Parameter("torque_scale", [1.5], [0.5], [3.0])
|
||||
bias = Parameter("torque_bias", [0.3], [-1.0], [1.0])
|
||||
pdict = _make_pdict(gain, bias)
|
||||
|
||||
transform = SignalTransform()
|
||||
transform.gain("*_torque", gain, target="both")
|
||||
transform.bias("*_torque", bias, target="both")
|
||||
|
||||
new, ref = _run_gains_biases_both(transform, ts, "predicted", pdict)
|
||||
np.testing.assert_array_equal(new.data, ref.data)
|
||||
|
||||
def test_target_filtering(self):
|
||||
"""A gain meant for measured data doesn't accidentally affect the predicted side."""
|
||||
ts = _make_transform_sensor_ts()
|
||||
gain = Parameter("gain", [2.0], [0.5], [3.0])
|
||||
pdict = _make_pdict(gain)
|
||||
|
||||
transform = SignalTransform()
|
||||
transform.gain("*_torque", gain, target="measured")
|
||||
|
||||
new, ref = _run_gains_biases_both(transform, ts, "predicted", pdict)
|
||||
np.testing.assert_array_equal(new.data, ref.data)
|
||||
np.testing.assert_array_equal(new.data, ts.data)
|
||||
|
||||
def test_original_ts_not_mutated(self):
|
||||
"""Signal transforms produce new data without mutating the input TimeSeries."""
|
||||
ts = _make_transform_sensor_ts()
|
||||
original_data = ts.data.copy()
|
||||
gain = Parameter("gain", [2.0], [0.5], [3.0])
|
||||
pdict = _make_pdict(gain)
|
||||
|
||||
transform = SignalTransform()
|
||||
transform.gain("*_torque", gain, target="predicted")
|
||||
|
||||
transform._apply_gains_biases(ts, "predicted", pdict)
|
||||
np.testing.assert_array_equal(ts.data, original_data)
|
||||
@@ -0,0 +1,331 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""Tests for the TimeSeries class and factory methods."""
|
||||
|
||||
import mujoco
|
||||
from mujoco.sysid._src.timeseries import SignalType
|
||||
from mujoco.sysid._src.timeseries import TimeSeries
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Local fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scalar_ts():
|
||||
"""y = x^2."""
|
||||
times = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
|
||||
data = np.array([0.0, 1.0, 4.0, 9.0, 16.0])
|
||||
return TimeSeries(times=times, data=data)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def multi_ts():
|
||||
"""y = [x^2, 2*x^2]."""
|
||||
times = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
|
||||
data = np.array([
|
||||
[0.0, 0.0],
|
||||
[1.0, 2.0],
|
||||
[4.0, 8.0],
|
||||
[9.0, 18.0],
|
||||
[16.0, 32.0],
|
||||
])
|
||||
return TimeSeries(times=times, data=data)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core TimeSeries tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_basics(scalar_ts, multi_ts):
|
||||
"""Basic properties: length, times array, and data array are all accessible."""
|
||||
assert len(scalar_ts) == 5
|
||||
assert len(multi_ts) == 5
|
||||
np.testing.assert_array_equal(scalar_ts.times, [0, 1, 2, 3, 4])
|
||||
np.testing.assert_array_equal(scalar_ts.data, [0, 1, 4, 9, 16])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"times, data, match",
|
||||
[
|
||||
(np.array([]), np.array([]), "Empty"),
|
||||
(np.array([[0.0], [1.0]]), np.array([0.0, 1.0]), "1D"),
|
||||
(np.array([0.0, 1.0]), np.array([0.0, 1.0, 2.0]), "Length"),
|
||||
(
|
||||
np.array([0.0, 2.0, 1.0]),
|
||||
np.array([0.0, 1.0, 2.0]),
|
||||
"strictly increasing",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_validation(times, data, match):
|
||||
"""Bad inputs (empty, non-1D times, length mismatch, non-monotonic) are rejected."""
|
||||
with pytest.raises(ValueError, match=match):
|
||||
TimeSeries(times=times, data=data)
|
||||
|
||||
|
||||
def test_zero_column_data():
|
||||
"""Zero-column data is valid (state-based models with no sensors)."""
|
||||
times = np.array([0.0, 1.0, 2.0])
|
||||
data = np.empty((3, 0))
|
||||
ts = TimeSeries(times=times, data=data)
|
||||
assert len(ts) == 3
|
||||
assert ts.data.shape == (3, 0)
|
||||
|
||||
|
||||
def test_save_and_load(scalar_ts, multi_ts, tmp_path):
|
||||
"""Saving to .npz and loading back recovers identical times and data."""
|
||||
path = tmp_path / "test.npz"
|
||||
scalar_ts.save_to_disk(path)
|
||||
loaded = TimeSeries.load_from_disk(path)
|
||||
np.testing.assert_array_equal(loaded.times, scalar_ts.times)
|
||||
np.testing.assert_array_equal(loaded.data, scalar_ts.data)
|
||||
|
||||
path2 = tmp_path / "test_multi.npz"
|
||||
multi_ts.save_to_disk(path2)
|
||||
loaded2 = TimeSeries.load_from_disk(path2)
|
||||
np.testing.assert_array_equal(loaded2.times, multi_ts.times)
|
||||
np.testing.assert_array_equal(loaded2.data, multi_ts.data)
|
||||
|
||||
|
||||
def test_save_and_load_with_signal_mapping(tmp_path):
|
||||
"""Save/load also preserves the signal_mapping (sensor name -> column index map)."""
|
||||
times = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
|
||||
data = np.array([
|
||||
[0.0, 0.0],
|
||||
[1.0, 2.0],
|
||||
[4.0, 8.0],
|
||||
[9.0, 18.0],
|
||||
[16.0, 32.0],
|
||||
])
|
||||
signal_mapping = {
|
||||
"signal1": (SignalType.MjSensor, np.array([0])),
|
||||
"signal2": (SignalType.MjSensor, np.array([1])),
|
||||
}
|
||||
ts = TimeSeries(
|
||||
times=times, data=data, signal_mapping=signal_mapping
|
||||
)
|
||||
|
||||
path = tmp_path / "test_signal_mapping.npz"
|
||||
ts.save_to_disk(path)
|
||||
loaded = TimeSeries.load_from_disk(path)
|
||||
|
||||
np.testing.assert_array_equal(loaded.times, times)
|
||||
np.testing.assert_array_equal(loaded.data, data)
|
||||
assert loaded.signal_mapping is not None
|
||||
assert loaded.signal_mapping.keys() == signal_mapping.keys()
|
||||
for key in signal_mapping:
|
||||
val_type, val_indices = signal_mapping[key]
|
||||
loaded_type, loaded_indices = loaded.signal_mapping[key]
|
||||
assert loaded_type == val_type
|
||||
np.testing.assert_array_equal(loaded_indices, val_indices)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method, expected",
|
||||
[
|
||||
("linear", 6.5),
|
||||
("cubic", 6.25),
|
||||
("quadratic", 6.25),
|
||||
("zero_order_hold", 4.0),
|
||||
("zoh", 4.0),
|
||||
],
|
||||
)
|
||||
def test_interpolate_scalar(scalar_ts, method, expected):
|
||||
"""Each interpolation method (linear, cubic, ZOH, etc.) gives the expected midpoint value."""
|
||||
result = scalar_ts.interpolate(2.5, method=method)
|
||||
assert result[0] == pytest.approx(expected, abs=1e-5)
|
||||
|
||||
|
||||
def test_interpolate_array(scalar_ts, multi_ts):
|
||||
"""Interpolating at multiple times simultaneously works for scalar and multi-column data."""
|
||||
t_values = np.array([0.5, 1.5, 2.5, 3.5])
|
||||
expected = np.array([0.5, 2.5, 6.5, 12.5])
|
||||
result = scalar_ts.interpolate(t_values, method="linear")
|
||||
np.testing.assert_allclose(result, expected, rtol=1e-5)
|
||||
|
||||
expected_multi = np.array([
|
||||
[0.5, 1.0],
|
||||
[2.5, 5.0],
|
||||
[6.5, 13.0],
|
||||
[12.5, 25.0],
|
||||
])
|
||||
result_multi = multi_ts.interpolate(t_values, method="linear")
|
||||
np.testing.assert_allclose(result_multi, expected_multi, rtol=1e-5)
|
||||
|
||||
|
||||
def test_resample_with_new_times(scalar_ts):
|
||||
"""Resampling onto a finer time grid via explicit new_times gives correct values."""
|
||||
new_times = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0])
|
||||
expected = np.array([0.0, 0.5, 1.0, 2.5, 4.0, 6.5, 9.0, 12.5, 16.0])
|
||||
resampled = scalar_ts.resample(new_times=new_times, method="linear")
|
||||
np.testing.assert_array_equal(resampled.times, new_times)
|
||||
np.testing.assert_allclose(resampled.data, expected, rtol=1e-5)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
scalar_ts.resample(new_times=np.array([0.0, 2.0, 1.0]))
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
scalar_ts.resample(new_times=np.array([[0.0], [1.0]]))
|
||||
|
||||
|
||||
def test_resample_with_target_dt(scalar_ts):
|
||||
"""Resampling by specifying a target timestep generates the right uniform grid."""
|
||||
resampled = scalar_ts.resample(target_dt=0.5, method="linear")
|
||||
expected_times = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0])
|
||||
expected_data = np.array([0.0, 0.5, 1.0, 2.5, 4.0, 6.5, 9.0, 12.5, 16.0])
|
||||
np.testing.assert_allclose(resampled.times, expected_times, rtol=1e-5)
|
||||
np.testing.assert_allclose(resampled.data, expected_data, rtol=1e-5)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
scalar_ts.resample(target_dt=-0.5)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
scalar_ts.resample()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TimeSeries factory method tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_from_model_controls_auto_resolution():
|
||||
"""Without explicit names, all model actuators are auto-discovered and mapped."""
|
||||
xml = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body>
|
||||
<joint name="j1"/>
|
||||
<geom size="0.1" mass="1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<actuator>
|
||||
<motor name="m1" joint="j1"/>
|
||||
<motor name="m2" joint="j1"/>
|
||||
</actuator>
|
||||
</mujoco>
|
||||
"""
|
||||
model = mujoco.MjModel.from_xml_string(xml)
|
||||
times = np.linspace(0, 1, 100)
|
||||
data = np.random.randn(100, 2)
|
||||
|
||||
ts = TimeSeries.from_control_names(times, data, model)
|
||||
assert ts.signal_mapping is not None
|
||||
assert "m1_ctrl" in ts.signal_mapping
|
||||
assert "m2_ctrl" in ts.signal_mapping
|
||||
assert ts.signal_mapping["m1_ctrl"][0] == SignalType.MjCtrl
|
||||
assert ts.signal_mapping["m2_ctrl"][0] == SignalType.MjCtrl
|
||||
|
||||
|
||||
def test_from_model_controls_explicit_names():
|
||||
"""Explicit actuator names are resolved; invalid or wrong-type names are rejected."""
|
||||
xml = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body>
|
||||
<joint name="j1"/>
|
||||
<geom size="0.1" mass="1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<actuator>
|
||||
<motor name="m1" joint="j1"/>
|
||||
</actuator>
|
||||
</mujoco>
|
||||
"""
|
||||
model = mujoco.MjModel.from_xml_string(xml)
|
||||
times = np.linspace(0, 1, 10)
|
||||
data = np.zeros((10, 1))
|
||||
|
||||
ts = TimeSeries.from_control_names(times, data, model, names=["m1"])
|
||||
assert ts.signal_mapping is not None
|
||||
assert "m1_ctrl" in ts.signal_mapping
|
||||
|
||||
with pytest.raises(ValueError, match="Could not resolve signal"):
|
||||
TimeSeries.from_control_names(times, data, model, names=["invalid"])
|
||||
|
||||
with pytest.raises(ValueError, match="not allowed"):
|
||||
TimeSeries.from_control_names(
|
||||
times, data, model, names=[("m1", SignalType.MjSensor)]
|
||||
)
|
||||
|
||||
|
||||
def test_from_model_auto_resolution_sensors():
|
||||
"""Without explicit names, all model sensors are auto-discovered with correct dimensions."""
|
||||
xml = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="b1">
|
||||
<geom size="0.1" mass="1"/>
|
||||
<site name="s1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<sensor>
|
||||
<accelerometer name="acc1" site="s1"/>
|
||||
<gyro name="gyro1" site="s1"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
"""
|
||||
model = mujoco.MjModel.from_xml_string(xml)
|
||||
times = np.linspace(0, 1, 10)
|
||||
data = np.zeros((10, 6))
|
||||
|
||||
ts = TimeSeries.from_names(times, data, model)
|
||||
assert ts.signal_mapping is not None
|
||||
assert "acc1" in ts.signal_mapping
|
||||
assert "gyro1" in ts.signal_mapping
|
||||
assert ts.signal_mapping["acc1"][0] == SignalType.MjSensor
|
||||
assert ts.signal_mapping["gyro1"][0] == SignalType.MjSensor
|
||||
|
||||
|
||||
def test_from_model_state_resolution():
|
||||
"""State signals (qpos, qvel) can be mapped by passing (name, SignalType) tuples."""
|
||||
xml = """
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<body name="b1">
|
||||
<joint name="j1" type="hinge"/>
|
||||
<joint name="j2" type="slide"/>
|
||||
<geom size="0.1" mass="1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
"""
|
||||
model = mujoco.MjModel.from_xml_string(xml)
|
||||
times = np.linspace(0, 1, 10)
|
||||
data = np.zeros((10, 2))
|
||||
|
||||
names = [("j1", SignalType.MjStateQPos), ("j2", SignalType.MjStateQPos)]
|
||||
ts = TimeSeries.from_names(times, data, model, names=names)
|
||||
assert ts.signal_mapping is not None
|
||||
assert "j1_qpos" in ts.signal_mapping
|
||||
assert "j2_qpos" in ts.signal_mapping
|
||||
|
||||
|
||||
def test_from_custom():
|
||||
"""Custom signal definitions (name strings and dimension tuples) are mapped correctly."""
|
||||
times = np.linspace(0, 1, 10)
|
||||
data = np.zeros((10, 3))
|
||||
signals = ["a", ("b", 2, SignalType.CustomObs)]
|
||||
|
||||
ts = TimeSeries.from_custom_map(times, data, signals)
|
||||
assert ts.signal_mapping is not None
|
||||
assert "a" in ts.signal_mapping
|
||||
assert "b" in ts.signal_mapping
|
||||
assert ts.signal_mapping["a"][1].size == 1
|
||||
assert ts.signal_mapping["b"][1].size == 2
|
||||
@@ -0,0 +1,159 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""Tests for the SystemTrajectory class."""
|
||||
|
||||
from unittest import mock
|
||||
|
||||
import mujoco
|
||||
from mujoco.sysid._src import timeseries
|
||||
from mujoco.sysid._src.trajectory import create_initial_state
|
||||
from mujoco.sysid._src.trajectory import SystemTrajectory
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model():
|
||||
model = mock.Mock(spec=mujoco.MjModel)
|
||||
model.nsensordata = 2
|
||||
model.nu = 1
|
||||
model.nq = 1
|
||||
model.nv = 1
|
||||
return model
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_trajectory(mock_model):
|
||||
"""Create a sample SystemTrajectory for testing."""
|
||||
with mock.patch.object(
|
||||
SystemTrajectory, "check_compatible", return_value=None
|
||||
):
|
||||
times = np.array([0.0, 1.0, 2.0])
|
||||
control_mapping = {"ctrl1": (timeseries.SignalType.MjCtrl, np.array([0]))}
|
||||
sensordata_mapping = {
|
||||
"sensor1": (timeseries.SignalType.MjSensor, np.array([0])),
|
||||
"sensor2": (timeseries.SignalType.MjSensor, np.array([1])),
|
||||
}
|
||||
state_mapping = {
|
||||
"qpos1": (timeseries.SignalType.MjStateQPos, np.array([0]))
|
||||
}
|
||||
|
||||
control = timeseries.TimeSeries(
|
||||
times=times,
|
||||
data=np.array([[1.0], [2.0], [3.0]]),
|
||||
signal_mapping=control_mapping,
|
||||
)
|
||||
sensordata = timeseries.TimeSeries(
|
||||
times=times,
|
||||
data=np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]),
|
||||
signal_mapping=sensordata_mapping,
|
||||
)
|
||||
state = timeseries.TimeSeries(
|
||||
times=times,
|
||||
data=np.array([[0.01], [0.02], [0.03]]),
|
||||
signal_mapping=state_mapping,
|
||||
)
|
||||
|
||||
traj = SystemTrajectory(
|
||||
model=mock_model,
|
||||
control=control,
|
||||
sensordata=sensordata,
|
||||
initial_state=np.array([0.0]),
|
||||
state=state,
|
||||
)
|
||||
yield traj, control_mapping, sensordata_mapping, state_mapping
|
||||
|
||||
|
||||
def test_save_and_load_with_signal_mapping(
|
||||
sample_trajectory, mock_model, tmp_path
|
||||
):
|
||||
"""Saving and loading a trajectory preserves all signal mappings (control, sensor, state)."""
|
||||
traj, control_mapping, sensordata_mapping, state_mapping = sample_trajectory
|
||||
|
||||
path = tmp_path / "test_traj.npz"
|
||||
traj.save_to_disk(path)
|
||||
|
||||
with mock.patch.object(
|
||||
SystemTrajectory, "check_compatible", return_value=None
|
||||
):
|
||||
loaded = SystemTrajectory.load_from_disk(path, mock_model)
|
||||
|
||||
ctrl_map = loaded.control.signal_mapping
|
||||
assert ctrl_map is not None
|
||||
assert ctrl_map.keys() == control_mapping.keys()
|
||||
for key in control_mapping:
|
||||
assert ctrl_map[key][0] == control_mapping[key][0]
|
||||
np.testing.assert_array_equal(ctrl_map[key][1], control_mapping[key][1])
|
||||
|
||||
sensor_map = loaded.sensordata.signal_mapping
|
||||
assert sensor_map is not None
|
||||
assert sensor_map.keys() == sensordata_mapping.keys()
|
||||
for key in sensordata_mapping:
|
||||
assert sensor_map[key][0] == sensordata_mapping[key][0]
|
||||
np.testing.assert_array_equal(
|
||||
sensor_map[key][1], sensordata_mapping[key][1]
|
||||
)
|
||||
|
||||
assert loaded.state is not None
|
||||
state_map = loaded.state.signal_mapping
|
||||
assert state_map is not None
|
||||
assert state_map.keys() == state_mapping.keys()
|
||||
for key in state_mapping:
|
||||
assert state_map[key][0] == state_mapping[key][0]
|
||||
np.testing.assert_array_equal(state_map[key][1], state_mapping[key][1])
|
||||
|
||||
|
||||
def test_create_initial_state(box_model):
|
||||
"""The initial MuJoCo state (qpos, qvel, act) is packed into a flat vector for rollout."""
|
||||
qpos = np.zeros(box_model.nq)
|
||||
qvel = np.zeros(box_model.nv)
|
||||
state = create_initial_state(box_model, qpos, qvel)
|
||||
expected_size = mujoco.mj_stateSize(
|
||||
box_model, mujoco.mjtState.mjSTATE_FULLPHYSICS
|
||||
)
|
||||
assert state.shape == (expected_size,)
|
||||
|
||||
|
||||
def test_create_initial_state_wrong_qpos(box_model):
|
||||
"""Wrong-sized qpos is caught early rather than causing a silent rollout bug."""
|
||||
with pytest.raises(ValueError, match="qpos"):
|
||||
create_initial_state(box_model, np.zeros(999))
|
||||
|
||||
|
||||
def test_split(sample_trajectory):
|
||||
"""A long trajectory can be split into smaller chunks for batched optimization."""
|
||||
traj, *_ = sample_trajectory
|
||||
chunks = traj.split(chunk_size=1)
|
||||
assert len(chunks) == 3
|
||||
assert len(chunks[0].sensordata) == 1
|
||||
|
||||
|
||||
def test_check_compatible_sensor_mismatch(box_model):
|
||||
"""Mismatched sensor dimensions between data and model are caught before rollout."""
|
||||
times = np.array([0.0, 0.01, 0.02])
|
||||
sensordata = timeseries.TimeSeries(times, np.ones((3, 5)))
|
||||
control = timeseries.TimeSeries(times, np.ones((3, box_model.nu)))
|
||||
initial_state = np.zeros(
|
||||
mujoco.mj_stateSize(box_model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
|
||||
)
|
||||
traj = SystemTrajectory(
|
||||
model=box_model,
|
||||
control=control,
|
||||
sensordata=sensordata,
|
||||
initial_state=initial_state,
|
||||
state=None,
|
||||
)
|
||||
with pytest.raises(ValueError, match="Sensor data dimension"):
|
||||
traj.check_compatible()
|
||||
Reference in New Issue
Block a user