chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

sysid: named ic construction, bug fixes, docstrings, README, more tests

Browse Source 
Co-authored-by: Kevin Zakka <kevinarmandzakka@gmail.com>
This commit is contained in:
Levi Burner
2026-02-10 15:15:17 -05:00
parent 210cf86486
commit e89dae359e
15 changed files with 768 additions and 1630 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/sysid/tests/test_integration.py
+182 -196
View File
@@ -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()