# 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 tests for mujoco.sysid.""" import logging import pathlib import tempfile import mujoco import mujoco.rollout as rollout from mujoco import sysid import numpy as np import pytest # --------------------------------------------------------------------------- # Models # --------------------------------------------------------------------------- SPRING_MASS_XML = """\ """ ARM_XML = """\ """ JOINT_NAMES = ["joint1", "joint2", "joint3", "joint4", "joint5"] TRUE_ARMATURE = {"joint1": 0.5, "joint2": 0.4, "joint3": 0.3, "joint4": 0.2, "joint5": 0.1} # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _generate_data(xml, ctrl_fn, duration): """Rollout a model and return (spec, initial_state, control_ts, sensor_ts).""" spec = mujoco.MjSpec.from_string(xml) model = spec.compile() data = mujoco.MjData(model) n_steps = int(duration / model.opt.timestep) t = np.arange(n_steps) * model.opt.timestep ctrl = ctrl_fn(t) 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] 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, ) 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 ms = sysid.ModelSequences( "spring_mass", spec, "measured", initial_state, control_ts, sensor_ts, ) residual_fn = sysid.build_residual_fn(models_sequences=[ms]) opt_params, _ = sysid.optimize( initial_params=params, residual_fn=residual_fn, optimizer="mujoco", ) 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, ) 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, ) residual_fn = sysid.build_residual_fn(models_sequences=[ms]) opt_params, opt_result = sysid.optimize( initial_params=params, residual_fn=residual_fn, optimizer="mujoco", ) for name in JOINT_NAMES: np.testing.assert_allclose( opt_params[f"{name}_armature"].value[0], TRUE_ARMATURE[name], atol=1e-4, ) # Verify save_results produces expected files. with tempfile.TemporaryDirectory() as tmpdir: sysid.save_results( experiment_results_folder=tmpdir, models_sequences=[ms], 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 / "arm.xml").exists() def _rank_1_residual_fn(x, p): # Residual depends only on x[0]: J = [[1, 0], [2, 0]] is 2x2 rank 1. del p if x.ndim == 1: r = np.array([x[0] - 1.0, 2.0 * x[0] - 2.0]) else: r = np.stack([x[0] - 1.0, 2.0 * x[0] - 2.0]) return [r], None, None def _full_rank_residual_fn(x, p): del p if x.ndim == 1: r = np.array([x[0] - 1.0, x[1] - 2.0]) else: r = np.stack([x[0] - 1.0, x[1] - 2.0]) return [r], None, None @pytest.mark.parametrize("residual_fn,expect_warning", [ (_rank_1_residual_fn, True), (_full_rank_residual_fn, False), ]) def test_check_conditioning(residual_fn, expect_warning, caplog): """check_conditioning=True warns iff cond(J^T J) is large at the start.""" params = sysid.ParameterDict() params.add(sysid.Parameter("a", 1.0, -10.0, 10.0)) params.add(sysid.Parameter("b", 1.0, -10.0, 10.0)) with caplog.at_level(logging.WARNING, logger="absl"): sysid.optimize( initial_params=params, residual_fn=residual_fn, optimizer="scipy", verbose=False, check_conditioning=True, max_iters=1, ) fired = any("cond(J^T J)" in r.message for r in caplog.records) assert fired is expect_warning