Merge pull request #3079 from aftersomemath:sysid-pr

PiperOrigin-RevId: 868229512
Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea

Author: copybara-github

python/build_requirements.txt

@@ -0,0 +1,60 @@
+# 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.
+# ==============================================================================
+"""Practical system identification for MuJoCo."""
+
+from mujoco.sysid._src import model_modifier
+from mujoco.sysid._src import parameter
+from mujoco.sysid._src import plotting
+from mujoco.sysid._src import signal_modifier
+from mujoco.sysid._src.io import save_results
+from mujoco.sysid._src.model_modifier import apply_body_inertia
+from mujoco.sysid._src.model_modifier import apply_dgain
+from mujoco.sysid._src.model_modifier import apply_param_modifiers
+from mujoco.sysid._src.model_modifier import apply_param_modifiers_spec
+from mujoco.sysid._src.model_modifier import apply_pdgain
+from mujoco.sysid._src.model_modifier import apply_pgain
+from mujoco.sysid._src.model_modifier import body_inertia_param
+from mujoco.sysid._src.model_modifier import remove_visuals
+from mujoco.sysid._src.optimize import calculate_intervals
+from mujoco.sysid._src.optimize import optimize
+from mujoco.sysid._src.parameter import InertiaType
+from mujoco.sysid._src.parameter import Parameter
+from mujoco.sysid._src.parameter import ParameterDict
+from mujoco.sysid._src.plotting import plot_sensor_comparison
+from mujoco.sysid._src.plotting import render_rollout
+from mujoco.sysid._src.residual import build_residual_fn
+from mujoco.sysid._src.residual import BuildModelFn
+from mujoco.sysid._src.residual import construct_ts_from_defaults
+from mujoco.sysid._src.residual import CustomRolloutFn
+from mujoco.sysid._src.residual import model_residual
+from mujoco.sysid._src.residual import ModifyResidualFn
+from mujoco.sysid._src.residual import residual
+from mujoco.sysid._src.signal_modifier import apply_bias
+from mujoco.sysid._src.signal_modifier import apply_delay
+from mujoco.sysid._src.signal_modifier import apply_delayed_ts_window
+from mujoco.sysid._src.signal_modifier import apply_gain
+from mujoco.sysid._src.signal_modifier import apply_resample_and_delay
+from mujoco.sysid._src.signal_modifier import get_sensor_indices
+from mujoco.sysid._src.signal_modifier import normalize_residual
+from mujoco.sysid._src.signal_modifier import weighted_diff
+from mujoco.sysid._src.signal_transform import SignalTransform
+from mujoco.sysid._src.timeseries import SignalType
+from mujoco.sysid._src.timeseries import TimeSeries
+from mujoco.sysid._src.trajectory import create_initial_state
+from mujoco.sysid._src.trajectory import ModelSequences
+from mujoco.sysid._src.trajectory import sysid_rollout
+from mujoco.sysid._src.trajectory import SystemTrajectory
+from mujoco.sysid.report.defaults import default_report
+from mujoco.sysid.report.defaults import default_report_matplotlib

python/mujoco/sysid/README.md

@@ -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.
+# ==============================================================================
+

python/mujoco/sysid/__init__.py

@@ -0,0 +1,82 @@
+# 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.
+# ==============================================================================
+
+"""I/O utilities for saving system identification results."""
+
+from collections.abc import Sequence
+import os
+import pathlib
+import pickle
+
+from absl import logging
+from mujoco.sysid._src import parameter
+from mujoco.sysid._src.optimize import calculate_intervals
+from mujoco.sysid._src.trajectory import ModelSequences
+import scipy.optimize as scipy_optimize
+
+
+def save_results(
+    experiment_results_folder: str | os.PathLike[str],
+    models_sequences: Sequence[ModelSequences],
+    initial_params: parameter.ParameterDict,
+    opt_params: parameter.ParameterDict,
+    opt_result: scipy_optimize.OptimizeResult,
+    residual_fn,
+):
+  """Save optimization results and confidence intervals to disk."""
+  experiment_results_folder = pathlib.Path(experiment_results_folder)
+  if not experiment_results_folder.exists():
+    experiment_results_folder.mkdir(parents=True, exist_ok=True)
+  logging.info(
+      "Experiment results will be saved to %s", experiment_results_folder
+  )
+
+  initial_params.save_to_disk(experiment_results_folder / "params_x_0.yaml")
+  opt_params.save_to_disk(experiment_results_folder / "params_x_hat.yaml")
+
+  with open(
+      os.path.join(experiment_results_folder, "results.pkl"), "wb"
+  ) as handle:
+    pickle.dump(opt_result, handle, protocol=pickle.HIGHEST_PROTOCOL)
+
+  # TODO(b/0): these intervals should be part of the params object.
+  residuals_star, _, _ = residual_fn(
+      opt_result.x, opt_params, return_pred_all=True
+  )
+  covariance, intervals = calculate_intervals(residuals_star, opt_result.jac)
+  with open(
+      os.path.join(experiment_results_folder, "confidence.pkl"), "wb"
+  ) as handle:
+    pickle.dump(
+        {"cov": covariance, "intervals": intervals},
+        handle,
+        protocol=pickle.HIGHEST_PROTOCOL,
+    )
+
+  # Dump identified models to disk.
+  for model_sequences in models_sequences:
+    model_sequences.spec.to_file(
+        (experiment_results_folder / f"{model_sequences.name}.xml").as_posix()
+    )
+
+  # Log nominal compared to initial.
+  x0 = initial_params.as_vector()
+  x_nominal = initial_params.as_nominal_vector()
+  logging.info(
+      "Initial Parameters\n%s",
+      initial_params.compare_parameters(
+          x0, opt_result.x, measured_params=x_nominal
+      ),
+  )