mujoco

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# System Identification\n",
        "\n",
        "This notebook describes MuJoCo's [system identification](https://en.wikipedia.org/wiki/System_identification) framework. System identification optimizes parameters to make a simulation match measurements. We introduce the framework's core concepts and walk through some basic examples. More detailed exercises will be added soon.\n",
        "\n",
        "**In this notebook we will:**\n",
        "\n",
        "1. [**Formulation:**](#formulation) briefly introduce the sysid problem\n",
        "2. [**Core Concepts:**](#core-concepts) estimate the mass of a mass-spring-damper\n",
        "3. [**API:**](#api) tour core API concepts like `Parameter`, `TimeSeries`, and `ModelSequences`\n",
        "4. [**Robot arm:**](#robot-arm) identify joint armature on a 5-DOF arm, with confidence intervals, sensor plots, and an interactive HTML report\n",
        "5. [**Parameter Identifiability:**](#ambiguity) understand and diagnose when parameters are unidentifiable"
      ],
      "metadata": {
        "id": "drs7G4CAj9Nx"
      },
      "id": "drs7G4CAj9Nx"
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setup"
      ],
      "metadata": {
        "id": "8P1e2JXckIbW"
      },
      "id": "8P1e2JXckIbW"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Install MuJoCo and mediapy\n",
        "!pip install -q mujoco[sysid] --pre -f https://py.mujoco.org/\n",
        "!pip install -q mediapy"
      ],
      "metadata": {
        "id": "OAAcG7DqkK3C"
      },
      "id": "OAAcG7DqkK3C",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Set up GPU rendering\n",
        "from google.colab import files\n",
        "import distutils.util\n",
        "import os\n",
        "import subprocess\n",
        "if subprocess.run('nvidia-smi').returncode:\n",
        "  raise RuntimeError(\n",
        "      'Cannot communicate with GPU. '\n",
        "      'Make sure you are using a GPU Colab runtime. '\n",
        "      'Go to the Runtime menu and select Choose runtime type.')\n",
        "\n",
        "# Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n",
        "# This is usually installed as part of an Nvidia driver package, but the Colab\n",
        "# kernel doesn't install its driver via APT, and as a result the ICD is missing.\n",
        "# (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n",
        "NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'\n",
        "if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n",
        "  with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:\n",
        "    f.write(\"\"\"{\n",
        "    \"file_format_version\" : \"1.0.0\",\n",
        "    \"ICD\" : {\n",
        "        \"library_path\" : \"libEGL_nvidia.so.0\"\n",
        "    }\n",
        "}\n",
        "\"\"\")"
      ],
      "metadata": {
        "id": "SIrVAb4gkMin"
      },
      "id": "SIrVAb4gkMin",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Configure EGL backend\n",
        "print('Setting environment variable to use GPU rendering:')\n",
        "%env MUJOCO_GL=egl"
      ],
      "metadata": {
        "id": "QgGusTkjkNcF"
      },
      "id": "QgGusTkjkNcF",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import mujoco\n",
        "import mujoco.rollout as rollout\n",
        "from mujoco import sysid\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import mediapy as media\n",
        "from absl import logging\n",
        "import base64\n",
        "from IPython.display import IFrame\n",
        "\n",
        "logging.set_verbosity(\"INFO\")\n",
        "\n",
        "def display_report(report):\n",
        "    html_b64 = base64.b64encode(report.build().encode()).decode()\n",
        "    return IFrame(src=f\"data:text/html;base64,{html_b64}\", width=\"100%\", height=800)"
      ],
      "metadata": {
        "id": "fWs5EDSRkOfj"
      },
      "id": "fWs5EDSRkOfj",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"formulation\"></a>\n",
        "# Formulation"
      ],
      "metadata": {
        "id": "LU9bl8aSkIWp"
      },
      "id": "LU9bl8aSkIWp"
    },
    {
      "cell_type": "markdown",
      "source": [
        "This library does [**gray-box**](https://en.wikipedia.org/wiki/Grey_box_model) identification: you supply the model structure\n",
        "(rigid-body dynamics, contacts, actuators) via a MuJoCo XML, and the optimizer\n",
        "adjusts the parameters you designate as unknown.\n",
        "\n",
        "Given $K$ parameters collected in a vector $\\theta$ and $N$ sensor\n",
        "measurements $y$, we simulate the model to produce predicted outputs\n",
        "$\\bar y(\\theta)$ and minimize the weighted residual:\n",
        "\n",
        "$$\\min_\\theta \\; \\tfrac{1}{2}\\lVert W\\bigl(\\bar y(\\theta) - y\\bigr)\\rVert^2\n",
        "\\qquad \\text{s.t.}\\quad l \\preccurlyeq \\theta \\preccurlyeq u$$\n",
        "\n",
        "This is a box-constrained **nonlinear least-squares** problem. The optimizer\n",
        "uses a Gauss-Newton / Levenberg-Marquardt algorithm with finite-difference\n",
        "Jacobians. Each parameter perturbation requires an independent simulation\n",
        "rollout, and all of them execute in a single batched call to [`mujoco.rollout`](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/rollout.ipynb),\n",
        "parallelized across CPU threads.\n",
        "\n",
        "For a detailed treatment of the underlying optimizer, see the\n",
        "[Least Squares](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/least_squares.ipynb) notebook."
      ],
      "metadata": {
        "id": "pyAquNvkkIUU"
      },
      "id": "pyAquNvkkIUU"
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"core-concepts\"></a>\n",
        "# 1. Core Concepts\n",
        "\n",
        "We illustrate the framework's core concepts with an actuated [mass-spring-damper](https://en.wikipedia.org/wiki/Mass-spring-damper_model). We know the exact spring stiffness and damping, but the mass must be estimated."
      ],
      "metadata": {
        "id": "wLktJeBokISN"
      },
      "id": "wLktJeBokISN"
    },
    {
      "cell_type": "markdown",
      "source": [
        "### The model\n",
        "\n",
        "A box on a spring, driven by a force actuator, with position and velocity\n",
        "sensors. The true mass is **1.0 kg**."
      ],
      "metadata": {
        "id": "_7k6iikqkIP6"
      },
      "id": "_7k6iikqkIP6"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title SPRING_MASS_XML { vertical-output: true}\n",
        "SPRING_MASS_XML = \"\"\"\\\n",
        "<mujoco model=\"spring_mass\">\n",
        "  <option timestep=\"0.002\">\n",
        "    <flag contact=\"disable\"/>\n",
        "  </option>\n",
        "  <worldbody>\n",
        "    <body name=\"ball\" pos=\"0 0 0.1\">\n",
        "      <inertial pos=\"0 0 0\" mass=\"1.0\" diaginertia=\"0.001 0.001 0.001\"/>\n",
        "      <joint name=\"slide\" type=\"slide\" axis=\"1 0 0\"\n",
        "             stiffness=\"100\" damping=\"5.0\"/>\n",
        "      <geom type=\"box\" size=\"0.05 0.05 0.05\"/>\n",
        "    </body>\n",
        "  </worldbody>\n",
        "  <actuator>\n",
        "    <motor name=\"push\" joint=\"slide\"/>\n",
        "  </actuator>\n",
        "  <sensor>\n",
        "    <jointpos name=\"position\" joint=\"slide\"/>\n",
        "    <jointvel name=\"velocity\" joint=\"slide\"/>\n",
        "  </sensor>\n",
        "</mujoco>\n",
        "\"\"\""
      ],
      "metadata": {
        "cellView": "form",
        "id": "NvQFiFJhkXSs"
      },
      "id": "NvQFiFJhkXSs",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Generate \"measured\" data\n",
        "\n",
        "We simulate the true model to create our ground-truth sensor recordings.\n",
        "A multi-frequency excitation signal produces a non-trivial trajectory.\n",
        "<!-- ensures the system is well-excited across its dynamics. -->\n",
        "\n",
        "`create_initial_state` packs the joint positions, velocities and actuator\n",
        "activations into the\n",
        "[state vector](https://mujoco.readthedocs.io/en/stable/computation/index.html#the-state)\n",
        "that `mujoco.rollout` expects. The results are wrapped in `TimeSeries`\n",
        "objects, which pair timestamps with data columns and are covered in detail\n",
        "in [Section 2](#api)."
      ],
      "metadata": {
        "id": "KMEsYY0hkILb"
      },
      "id": "KMEsYY0hkILb"
    },
    {
      "cell_type": "code",
      "source": [
        "spec = mujoco.MjSpec.from_string(SPRING_MASS_XML)\n",
        "model = spec.compile()\n",
        "data = mujoco.MjData(model)\n",
        "\n",
        "duration = 3.0\n",
        "n_steps = int(duration / model.opt.timestep)\n",
        "t = np.arange(n_steps) * model.opt.timestep\n",
        "\n",
        "ctrl = (5.0 * np.sin(2 * np.pi * 1.5 * t)\n",
        "        + 3.0 * np.sin(2 * np.pi * 3.7 * t)).reshape(-1, 1)\n",
        "\n",
        "initial_state = sysid.create_initial_state(model, data.qpos, data.qvel, data.act)\n",
        "\n",
        "state, sensor = rollout.rollout(model, data, initial_state, ctrl[:-1])\n",
        "state = np.squeeze(state, axis=0)\n",
        "sensor = np.squeeze(sensor, axis=0)\n",
        "times = state[:, 0]"
      ],
      "metadata": {
        "id": "XXzMhCJTka8Z"
      },
      "id": "XXzMhCJTka8Z",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's visualize the result:"
      ],
      "metadata": {
        "id": "_tOP-duuYpVt"
      },
      "id": "_tOP-duuYpVt"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title { vertical-output: true}\n",
        "\n",
        "# Render the rollout.\n",
        "frames = sysid.render_rollout(\n",
        "    model, data, state[None], framerate=30, height=400, width=560\n",
        ")\n",
        "media.show_video(frames, fps=30)\n",
        "\n",
        "control_ts = sysid.TimeSeries(t, ctrl)\n",
        "sensor_ts = sysid.TimeSeries.from_names(times, sensor, model)\n",
        "\n",
        "fig, axes = plt.subplots(3, 1, figsize=(5, 4), sharex=True,\n",
        "                         gridspec_kw={\"height_ratios\": [2, 2, 1]})\n",
        "\n",
        "axes[0].plot(times, sensor[:, 0], color=\"C0\", linewidth=1.0)\n",
        "axes[0].set_ylabel(\"Position (m)\")\n",
        "axes[0].set_title(\"Measured trajectory (true mass = 1.0 kg)\")\n",
        "\n",
        "axes[1].plot(times, sensor[:, 1], color=\"C1\", linewidth=1.0)\n",
        "axes[1].set_ylabel(\"Velocity (m/s)\")\n",
        "\n",
        "axes[2].plot(t, ctrl[:, 0], color=\"0.4\", linewidth=0.8)\n",
        "axes[2].set_ylabel(\"Control (N)\")\n",
        "axes[2].set_xlabel(\"Time (s)\")\n",
        "\n",
        "for ax in axes:\n",
        "    ax.grid(True, alpha=0.3)\n",
        "\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "QwLOAPkTYsXG"
      },
      "id": "QwLOAPkTYsXG",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Construct the System Identification Problem\n",
        "\n",
        "#### Define the unknown parameter\n",
        "\n",
        "A `Parameter` wraps a scalar (or vector) value with:\n",
        "- **bounds** for box-constrained optimization\n",
        "- a **modifier** callback that stamps the current value onto an [`MjSpec`](https://mujoco.readthedocs.io/en/stable/python.html#model-editing)\n",
        "\n",
        "The `nominal` value is a fixed reference point (here we set it to the true\n",
        "value for later comparison, but in practice it would be your best prior\n",
        "guess). The mutable `value` is what the optimizer actually updates.\n",
        "\n",
        "We'll start the optimizer at **mass = 2.0 kg**, double the true value."
      ],
      "metadata": {
        "id": "8XkPEIaVkIJF"
      },
      "id": "8XkPEIaVkIJF"
    },
    {
      "cell_type": "code",
      "source": [
        "def set_mass(spec, param):\n",
        "    \"\"\"Modifier: stamp the current mass value onto the MjSpec.\"\"\"\n",
        "    spec.body(\"ball\").mass = param.value[0]\n",
        "\n",
        "params = sysid.ParameterDict()\n",
        "params.add(sysid.Parameter(\n",
        "    \"mass\",\n",
        "    nominal=1.0,\n",
        "    min_value=0.3,\n",
        "    max_value=3.0,\n",
        "    modifier=set_mass,\n",
        "))\n",
        "\n",
        "params[\"mass\"].value[:] = 2.0\n",
        "print(f\"Starting mass: {params['mass'].value[0]:.2f} kg  (true: 1.0 kg)\")"
      ],
      "metadata": {
        "id": "s6HVXGVJkfka"
      },
      "id": "s6HVXGVJkfka",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Package data\n",
        "\n",
        "`ModelSequences` bundles an `MjSpec` with one or more measured trajectories\n",
        "(initial state, controls and sensor readings)."
      ],
      "metadata": {
        "id": "KjUPZVlOkIEn"
      },
      "id": "KjUPZVlOkIEn"
    },
    {
      "cell_type": "code",
      "source": [
        "ms = sysid.ModelSequences(\n",
        "    \"spring_mass\", spec, \"measured\", initial_state, control_ts, sensor_ts,\n",
        ")"
      ],
      "metadata": {
        "id": "0KxR-RtpkiL3"
      },
      "id": "0KxR-RtpkiL3",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Optimize\n",
        "\n",
        "`build_residual_fn` creates a callable that, for a given parameter vector,\n",
        "applies the parameters to the spec, rolls out the simulation, and returns the\n",
        "difference between predicted and measured sensor data. This difference is the\n",
        "**residual** that the optimizer drives to zero."
      ],
      "metadata": {
        "id": "0aQQLaJjY09c"
      },
      "id": "0aQQLaJjY09c"
    },
    {
      "cell_type": "code",
      "source": [
        "residual_fn = sysid.build_residual_fn(models_sequences=[ms])\n",
        "\n",
        "opt_params, opt_result = sysid.optimize(\n",
        "    initial_params=params,\n",
        "    residual_fn=residual_fn,\n",
        "    optimizer=\"mujoco\",\n",
        "    verbose=True,\n",
        ")\n",
        "\n",
        "print(f\"\\nRecovered mass: {opt_params['mass'].value[0]:.4f} kg  (true: 1.0 kg)\")"
      ],
      "metadata": {
        "id": "jhQShvjLY2OG"
      },
      "id": "jhQShvjLY2OG",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Report\n",
        "\n",
        "Inspect the report to see that the optimized predictions fit the nominal (ground truth) while the initial sensor predictions do not.\n",
        "\n",
        "The report also contains measurement comparisons, parameter tables, confidence intervals, and other debugging information."
      ],
      "metadata": {
        "id": "dSRVGBIzY7is"
      },
      "id": "dSRVGBIzY7is"
    },
    {
      "cell_type": "code",
      "source": [
        "report = sysid.default_report(\n",
        "    models_sequences=[ms],\n",
        "    initial_params=params,\n",
        "    opt_params=opt_params,\n",
        "    residual_fn=residual_fn,\n",
        "    opt_result=opt_result,\n",
        "    title=\"Mass Spring Damper Identification\",\n",
        "    generate_videos=False,\n",
        ")\n",
        "display_report(report)"
      ],
      "metadata": {
        "id": "G7pcHNbSY7F4"
      },
      "id": "G7pcHNbSY7F4",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"api\"></a>\n",
        "# 2. API\n",
        "\n",
        "The warm-up introduced `Parameter`, `TimeSeries`, and `ModelSequences` in\n",
        "passing. Here we look at each one more closely."
      ],
      "metadata": {
        "id": "xQo1OQn_kjlA"
      },
      "id": "xQo1OQn_kjlA"
    },
    {
      "cell_type": "markdown",
      "source": [
        "## ParameterDict\n",
        "\n",
        "In the warm-up we created a single `Parameter` with a name, bounds and a\n",
        "modifier callback. A `ParameterDict` groups multiple parameters for the optimizer.\n",
        "\n",
        "Parameters can be **frozen**: they are excluded from the optimization vector\n",
        "but their modifiers are still applied at the frozen value."
      ],
      "metadata": {
        "id": "YvPJxqUOkjjI"
      },
      "id": "YvPJxqUOkjjI"
    },
    {
      "cell_type": "code",
      "source": [
        "# Three parameters for the spring-mass model. Damping is frozen.\n",
        "params = sysid.ParameterDict()\n",
        "params.add(sysid.Parameter(\n",
        "    \"mass\", nominal=1.0, min_value=0.3, max_value=3.0,\n",
        "    modifier=lambda s, p: setattr(s.body(\"ball\"), \"mass\", p.value[0]),\n",
        "))\n",
        "params.add(sysid.Parameter(\n",
        "    \"stiffness\", nominal=100.0, min_value=30.0, max_value=300.0,\n",
        "    modifier=lambda s, p: setattr(s.joint(\"slide\"), \"stiffness\", p.value[0]),\n",
        "))\n",
        "params.add(sysid.Parameter(\n",
        "    \"damping\", nominal=5.0, min_value=0.0, max_value=20.0, frozen=True,\n",
        "    modifier=lambda s, p: setattr(s.joint(\"slide\"), \"damping\", p.value[0]),\n",
        "))\n",
        "\n",
        "print(\"Parameter vector (excludes frozen):\", params.as_vector())\n",
        "print(\"Bounds:\", params.get_bounds())\n",
        "print(\"Frozen 'damping' still in dict:    \", params[\"damping\"].value)"
      ],
      "metadata": {
        "id": "WqFeNF9NkpRz"
      },
      "id": "WqFeNF9NkpRz",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## TimeSeries\n",
        "\n",
        "`TimeSeries` is an immutable container for timestamped signals. The\n",
        "`from_names` factory that we used in the warm-up automatically creates a\n",
        "**signal mapping**, a dict from sensor names to column indices. This mapping\n",
        "lets the library match predicted signals to measured ones by name.\n",
        "\n",
        "`TimeSeries` also supports resampling to a different timestep, which is useful\n",
        "when your measured data and simulation run at different rates."
      ],
      "metadata": {
        "id": "YlHeA2D1kjgx"
      },
      "id": "YlHeA2D1kjgx"
    },
    {
      "cell_type": "code",
      "source": [
        "# sensor_ts was created in the warm-up with from_names\n",
        "print(\"Signal mapping:\", sensor_ts.signal_mapping)\n",
        "\n",
        "# Resample to a coarser timestep (returns a new TimeSeries)\n",
        "ts_coarse = sensor_ts.resample(target_dt=0.01)\n",
        "print(f\"Original: {len(sensor_ts.times)} pts  ->  Resampled: {len(ts_coarse.times)} pts\")"
      ],
      "metadata": {
        "id": "1Aw6Io43krdq"
      },
      "id": "1Aw6Io43krdq",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## ModelSequences\n",
        "\n",
        "`ModelSequences` bundles an `MjSpec` with one or more measured trajectories\n",
        "(initial state, controls and sensor readings). It is the input to\n",
        "`build_residual_fn`.\n",
        "\n",
        "You can pass **multiple trajectories** for the same spec. This is the key\n",
        "mechanism for improving parameter identifiability, as we'll see in the next\n",
        "section."
      ],
      "metadata": {
        "id": "YP79IWujkjes"
      },
      "id": "YP79IWujkjes"
    },
    {
      "cell_type": "code",
      "source": [
        "# Single trajectory (as in the warm-up)\n",
        "ms_single = sysid.ModelSequences(\n",
        "    \"spring_mass\", spec, \"traj_0\", initial_state, control_ts, sensor_ts,\n",
        ")\n",
        "print(f\"Single trajectory: {ms_single.sequence_name}\")\n",
        "\n",
        "# Multiple trajectories for the same spec -- just pass lists\n",
        "ms_multi = sysid.ModelSequences(\n",
        "    \"spring_mass\", spec,\n",
        "    [\"traj_0\", \"traj_1\"],           # names\n",
        "    [initial_state, initial_state], # initial states\n",
        "    [control_ts, control_ts],       # controls\n",
        "    [sensor_ts, sensor_ts],         # sensor data\n",
        ")\n",
        "print(f\"Multiple trajectories: {ms_multi.sequence_name}\")"
      ],
      "metadata": {
        "id": "mGaDg60Zkuiz"
      },
      "id": "mGaDg60Zkuiz",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"robot-arm\"></a>\n",
        "# 4. Robot Arm: Identifying Joint Armature\n",
        "\n",
        "Now we tackle a realistic problem. A 5-DOF robot arm is driven by motor actuators that apply joint torques, and the armature is unknown. Armature represents reflected rotor inertia through the gear train. It is a common source of sim-to-real gap. Getting it wrong means the simulated arm accelerates too fast or too slow under the same applied torque. You can read more about it [here](https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-joint-armature)."
      ],
      "metadata": {
        "id": "yTLBbv4BZMK5"
      },
      "id": "yTLBbv4BZMK5"
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Arm model\n",
        "\n",
        "Five hinge joints with motor actuators (torque control). Joint damping\n",
        "provides passive dissipation. The true armature values decrease from base\n",
        "to tip, reflecting smaller motors on distal joints. Position sensors on every joint."
      ],
      "metadata": {
        "id": "nfGPT2f2ZPFh"
      },
      "id": "nfGPT2f2ZPFh"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ARM_XML { vertical-output: true}\n",
        "ARM_XML = \"\"\"\\\n",
        "<mujoco model=\"arm\">\n",
        "  <compiler angle=\"radian\" autolimits=\"true\"/>\n",
        "  <option integrator=\"implicitfast\" timestep=\"0.002\">\n",
        "    <flag contact=\"disable\"/>\n",
        "  </option>\n",
        "  <worldbody>\n",
        "    <body name=\"link1\" pos=\"0 0 0.1\">\n",
        "      <inertial pos=\"0 0 0.05\" mass=\"1.0\" diaginertia=\"0.01 0.01 0.005\"/>\n",
        "      <joint name=\"joint1\" type=\"hinge\" axis=\"0 0 1\" range=\"-3.14 3.14\"\n",
        "             armature=\"0.5\" damping=\"1.0\"/>\n",
        "      <geom type=\"capsule\" fromto=\"0 0 0 0 0 0.1\" size=\"0.04\"/>\n",
        "      <body name=\"link2\" pos=\"0 0 0.1\">\n",
        "        <inertial pos=\"0 0 0.05\" mass=\"0.8\" diaginertia=\"0.008 0.008 0.004\"/>\n",
        "        <joint name=\"joint2\" type=\"hinge\" axis=\"0 1 0\" range=\"-3.14 3.14\"\n",
        "               armature=\"0.4\" damping=\"0.8\"/>\n",
        "        <geom type=\"capsule\" fromto=\"0 0 0 0 0 0.1\" size=\"0.035\"/>\n",
        "        <body name=\"link3\" pos=\"0 0 0.1\">\n",
        "          <inertial pos=\"0 0 0.05\" mass=\"0.6\" diaginertia=\"0.006 0.006 0.003\"/>\n",
        "          <joint name=\"joint3\" type=\"hinge\" axis=\"0 1 0\" range=\"-3.14 3.14\"\n",
        "                 armature=\"0.3\" damping=\"0.6\"/>\n",
        "          <geom type=\"capsule\" fromto=\"0 0 0 0 0 0.1\" size=\"0.03\"/>\n",
        "          <body name=\"link4\" pos=\"0 0 0.1\">\n",
        "            <inertial pos=\"0 0 0.04\" mass=\"0.4\" diaginertia=\"0.004 0.004 0.002\"/>\n",
        "            <joint name=\"joint4\" type=\"hinge\" axis=\"0 0 1\" range=\"-3.14 3.14\"\n",
        "                   armature=\"0.2\" damping=\"0.4\"/>\n",
        "            <geom type=\"capsule\" fromto=\"0 0 0 0 0 0.08\" size=\"0.025\"/>\n",
        "            <body name=\"link5\" pos=\"0 0 0.08\">\n",
        "              <inertial pos=\"0 0 0.03\" mass=\"0.2\" diaginertia=\"0.002 0.002 0.001\"/>\n",
        "              <joint name=\"joint5\" type=\"hinge\" axis=\"0 1 0\" range=\"-3.14 3.14\"\n",
        "                     armature=\"0.1\" damping=\"0.2\"/>\n",
        "              <geom type=\"capsule\" fromto=\"0 0 0 0 0 0.06\" size=\"0.02\"/>\n",
        "            </body>\n",
        "          </body>\n",
        "        </body>\n",
        "      </body>\n",
        "    </body>\n",
        "  </worldbody>\n",
        "  <actuator>\n",
        "    <motor name=\"act1\" joint=\"joint1\"/>\n",
        "    <motor name=\"act2\" joint=\"joint2\"/>\n",
        "    <motor name=\"act3\" joint=\"joint3\"/>\n",
        "    <motor name=\"act4\" joint=\"joint4\"/>\n",
        "    <motor name=\"act5\" joint=\"joint5\"/>\n",
        "  </actuator>\n",
        "  <sensor>\n",
        "    <jointpos name=\"joint1_pos\" joint=\"joint1\"/>\n",
        "    <jointpos name=\"joint2_pos\" joint=\"joint2\"/>\n",
        "    <jointpos name=\"joint3_pos\" joint=\"joint3\"/>\n",
        "    <jointpos name=\"joint4_pos\" joint=\"joint4\"/>\n",
        "    <jointpos name=\"joint5_pos\" joint=\"joint5\"/>\n",
        "  </sensor>\n",
        "</mujoco>\n",
        "\"\"\"\n",
        "\n",
        "JOINT_NAMES = [\"joint1\", \"joint2\", \"joint3\", \"joint4\", \"joint5\"]\n",
        "TRUE_ARMATURE = {\"joint1\": 0.5, \"joint2\": 0.4, \"joint3\": 0.3,\n",
        "                 \"joint4\": 0.2, \"joint5\": 0.1}"
      ],
      "metadata": {
        "cellView": "form",
        "id": "hdLji2qelNqi"
      },
      "id": "hdLji2qelNqi",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Generate measured data\n",
        "\n",
        "We simulate the true model and add sensor noise to mimic real encoder\n",
        "readings. A multi-frequency torque excitation ensures each joint is\n",
        "well-excited. With motor actuators the control signal **is** the torque,\n",
        "so there is no feedback loop to mask the effect of armature."
      ],
      "metadata": {
        "id": "UtAAiNGIlQo-"
      },
      "id": "UtAAiNGIlQo-"
    },
    {
      "cell_type": "code",
      "source": [
        "spec = mujoco.MjSpec.from_string(ARM_XML)\n",
        "model = spec.compile()\n",
        "data = mujoco.MjData(model)\n",
        "\n",
        "duration = 2.0\n",
        "n_steps = int(duration / model.opt.timestep)\n",
        "t = np.arange(n_steps) * model.opt.timestep\n",
        "\n",
        "# Sinusoidal torques with different frequency and amplitude per joint\n",
        "ctrl = np.column_stack([\n",
        "    5.0 * np.sin(2 * np.pi * 0.5 * t),\n",
        "    4.0 * np.sin(2 * np.pi * 0.7 * t + 0.5),\n",
        "    3.0 * np.sin(2 * np.pi * 0.4 * t + 1.0),\n",
        "    2.0 * np.sin(2 * np.pi * 0.9 * t + 1.5),\n",
        "    1.0 * np.sin(2 * np.pi * 0.6 * t + 2.0),\n",
        "])\n",
        "\n",
        "initial_state = sysid.create_initial_state(model, data.qpos, data.qvel, data.act)\n",
        "state, sensor = rollout.rollout(model, data, initial_state, ctrl[:-1])\n",
        "state = np.squeeze(state, axis=0)\n",
        "sensor = np.squeeze(sensor, axis=0)\n",
        "times = state[:, 0]\n",
        "\n",
        "# Add realistic sensor noise\n",
        "rng = np.random.default_rng(seed=0)\n",
        "noise_std = np.zeros(sensor.shape[1])\n",
        "noise_std[:] = 0.6   # position noise (rad)\n",
        "sensor_noisy = sensor + rng.normal(scale=noise_std, size=sensor.shape)\n",
        "\n",
        "control_ts = sysid.TimeSeries(t, ctrl)\n",
        "sensor_ts = sysid.TimeSeries.from_names(times, sensor_noisy, model)\n",
        "\n",
        "print(f\"Sensor channels: {sensor.shape[1]}  \"\n",
        "      f\"({model.nsensor} sensors: 5 pos + 5 vel)\")\n",
        "print(f\"Timesteps: {len(times)} ({duration}s at dt={model.opt.timestep})\")\n",
        "print(f\"Noise std: {noise_std[:5][0]:.0e} rad (pos)\")"
      ],
      "metadata": {
        "id": "zTSTluXplSRm"
      },
      "id": "zTSTluXplSRm",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Define armature parameters\n",
        "\n",
        "One scalar parameter per joint. The modifier callback sets the joint's `armature` attribute on the spec. We start from a wrong initial guess of 0.01 for all joints, which is up to **50x too small** for the base joint."
      ],
      "metadata": {
        "id": "KxiNscAnlQm5"
      },
      "id": "KxiNscAnlQm5"
    },
    {
      "cell_type": "code",
      "source": [
        "INIT_ARMATURE = 0.01\n",
        "\n",
        "def make_armature_modifier(joint_name):\n",
        "    \"\"\"Create a modifier that sets armature on a specific joint.\"\"\"\n",
        "    def modifier(spec, param):\n",
        "        spec.joint(joint_name).armature = param.value[0]\n",
        "    return modifier\n",
        "\n",
        "params = sysid.ParameterDict()\n",
        "for name in JOINT_NAMES:\n",
        "    true_val = TRUE_ARMATURE[name]\n",
        "    params.add(sysid.Parameter(\n",
        "        f\"{name}_armature\",\n",
        "        nominal=true_val,\n",
        "        min_value=0.01,\n",
        "        max_value=0.6,\n",
        "        modifier=make_armature_modifier(name),\n",
        "    ))\n",
        "    # Start from wrong initial guess\n",
        "    params[f\"{name}_armature\"].value[:] = INIT_ARMATURE\n",
        "\n",
        "print(\"Initial parameter vector:\", params.as_vector())\n",
        "print(\"True values:             \",\n",
        "      np.array([TRUE_ARMATURE[n] for n in JOINT_NAMES]))"
      ],
      "metadata": {
        "id": "_98Rw6mOlUsP"
      },
      "id": "_98Rw6mOlUsP",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Optimize"
      ],
      "metadata": {
        "id": "BDJWVxGwlQlD"
      },
      "id": "BDJWVxGwlQlD"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title { vertical-output: true}\n",
        "ms = sysid.ModelSequences(\n",
        "    \"arm\", spec, \"sinusoidal\", initial_state, control_ts, sensor_ts,\n",
        ")\n",
        "\n",
        "residual_fn = sysid.build_residual_fn(models_sequences=[ms])\n",
        "\n",
        "opt_params, opt_result = sysid.optimize(\n",
        "    initial_params=params,\n",
        "    residual_fn=residual_fn,\n",
        "    optimizer=\"mujoco\",\n",
        ")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "VXQoNpmolXju"
      },
      "id": "VXQoNpmolXju",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Report\n",
        "\n",
        "**Confidence Internals.** Because the measured data contains sensor noise, the optimal residuals have variance, which can be used to estimate 95% parameter confidence intervals. These are displayed in the report \"Parameter Distribution\" section."
      ],
      "metadata": {
        "id": "RLB285HqZbNj"
      },
      "id": "RLB285HqZbNj"
    },
    {
      "cell_type": "code",
      "source": [
        "report = sysid.default_report(\n",
        "    models_sequences=[ms],\n",
        "    initial_params=params,\n",
        "    opt_params=opt_params,\n",
        "    residual_fn=residual_fn,\n",
        "    opt_result=opt_result,\n",
        "    title=\"Robot Arm Armature Identification\",\n",
        "    generate_videos=False,\n",
        ")\n",
        "display_report(report)"
      ],
      "metadata": {
        "id": "98447uDUZcgI"
      },
      "id": "98447uDUZcgI",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Rendered overlay: initial vs. optimized vs. ground truth\n",
        "\n",
        "`render_rollout` takes a list of models and a batch of state trajectories\n",
        "and renders them into a single scene, which lets us visually compare how\n",
        "different parameter values affect the motion. We render two side-by-side\n",
        "videos:\n",
        "- **Before:** initial guess (red) vs. ground truth (green)\n",
        "- **After:** optimized (blue) vs. ground truth (green)"
      ],
      "metadata": {
        "id": "bhesmm6Nln98"
      },
      "id": "bhesmm6Nln98"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title { vertical-output: true}\n",
        "def set_body_rgba(body, rgba):\n",
        "    \"\"\"Recursively set rgba on all geoms in a body tree.\"\"\"\n",
        "    for geom in body.geoms:\n",
        "        geom.rgba = rgba\n",
        "    for child in body.bodies:\n",
        "        set_body_rgba(child, rgba)\n",
        "\n",
        "def make_colored_model(base_spec, rgba, armature_values):\n",
        "    \"\"\"Copy spec, set geom colors and armature, compile.\"\"\"\n",
        "    s = base_spec.copy()\n",
        "    for name, val in armature_values.items():\n",
        "        s.joint(name).armature = val\n",
        "    set_body_rgba(s.worldbody, rgba)\n",
        "    return s.compile()\n",
        "\n",
        "true_armature = {n: TRUE_ARMATURE[n] for n in JOINT_NAMES}\n",
        "init_armature = {n: INIT_ARMATURE for n in JOINT_NAMES}\n",
        "opt_armature = {n: opt_params[f\"{n}_armature\"].value[0] for n in JOINT_NAMES}\n",
        "\n",
        "green = [0.2, 0.8, 0.2, 0.7]\n",
        "red = [1.0, 0.2, 0.2, 0.7]\n",
        "blue = [0.2, 0.4, 1.0, 0.7]\n",
        "\n",
        "truth_model = make_colored_model(spec, green, true_armature)\n",
        "init_model = make_colored_model(spec, red, init_armature)\n",
        "opt_model = make_colored_model(spec, blue, opt_armature)\n",
        "\n",
        "fps = 30\n",
        "\n",
        "# Before: initial (red) vs ground truth (green)\n",
        "models_before = [init_model, truth_model]\n",
        "datas_before = [mujoco.MjData(m) for m in models_before]\n",
        "state_before, _ = rollout.rollout(\n",
        "    models_before, datas_before, initial_state, ctrl[:-1]\n",
        ")\n",
        "frames_before = sysid.render_rollout(\n",
        "    models_before, datas_before[0], state_before,\n",
        "    framerate=fps, height=400, width=560,\n",
        ")\n",
        "\n",
        "# After: optimized (blue) vs ground truth (green)\n",
        "models_after = [opt_model, truth_model]\n",
        "datas_after = [mujoco.MjData(m) for m in models_after]\n",
        "state_after, _ = rollout.rollout(\n",
        "    models_after, datas_after, initial_state, ctrl[:-1]\n",
        ")\n",
        "frames_after = sysid.render_rollout(\n",
        "    models_after, datas_after[0], state_after,\n",
        "    framerate=fps, height=400, width=560,\n",
        ")\n",
        "\n",
        "# Side by side\n",
        "media.show_videos(\n",
        "    [frames_before, frames_after],\n",
        "    fps=fps,\n",
        "    titles=[\"Before: initial (red) vs truth (green)\",\n",
        "            \"After: optimized (blue) vs truth (green)\"],\n",
        ")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "SkpnM_zTlp5t"
      },
      "id": "SkpnM_zTlp5t",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<a name=\"ambiguity\"></a>\n",
        "# 4. Parameter Identifiability\n",
        "\n",
        "Sometimes a single experiment cannot uniquely determine all parameters. This is related to the concept of [Structural Identifiability](https://en.wikipedia.org/wiki/Structural_identifiability).\n",
        "\n",
        "Consider a cart driven by a motor:\n",
        "\n",
        "$$m\\,\\ddot x = b\\,u(t)$$\n",
        "\n",
        "where $m$ is the cart mass and $b$ is the torque constant ([`gear`](https://mujoco.readthedocs.io/en/stable/XMLreference.html#actuator-general-gear)). The response\n",
        "depends **only on the ratio** $b/m$. Doubling both $b$ and $m$ produces an\n",
        "identical trajectory. With a single recording, the individual values are\n",
        "fundamentally unidentifiable.\n",
        "\n",
        "We're going to visualize the ambiguity and resolve it by optimizing over a **second trajectory** recorded with a known mass perturbation."
      ],
      "metadata": {
        "id": "Z_LOdz7CZqlg"
      },
      "id": "Z_LOdz7CZqlg"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title CART_XML { vertical-output: true}\n",
        "CART_XML = \"\"\"\\\n",
        "<mujoco model=\"cart\">\n",
        "  <option timestep=\"0.002\">\n",
        "    <flag contact=\"disable\"/>\n",
        "  </option>\n",
        "  <worldbody>\n",
        "    <body name=\"cart\" pos=\"0 0 0.1\">\n",
        "      <inertial pos=\"0 0 0\" mass=\"2.0\" diaginertia=\"0.001 0.001 0.001\"/>\n",
        "      <joint name=\"slide\" type=\"slide\" axis=\"1 0 0\"/>\n",
        "      <geom type=\"sphere\" size=\"0.05\"/>\n",
        "    </body>\n",
        "  </worldbody>\n",
        "  <actuator>\n",
        "    <motor name=\"motor\" joint=\"slide\" gear=\"3.0\"/>\n",
        "  </actuator>\n",
        "  <sensor>\n",
        "    <jointpos name=\"position\" joint=\"slide\"/>\n",
        "    <jointvel name=\"velocity\" joint=\"slide\"/>\n",
        "  </sensor>\n",
        "</mujoco>\n",
        "\"\"\"\n",
        "\n",
        "TRUE_MASS = 2.0\n",
        "TRUE_GEAR = 3.0\n",
        "PAYLOAD_MASS = 1.0\n",
        "\n",
        "# Base cart\n",
        "spec_base = mujoco.MjSpec.from_string(CART_XML)\n",
        "\n",
        "# Payload variant: add a rigid mass to the cart body\n",
        "spec_pay = mujoco.MjSpec.from_string(CART_XML)\n",
        "payload = spec_pay.body(\"cart\").add_body()\n",
        "payload.name = \"payload\"\n",
        "payload.mass = PAYLOAD_MASS\n",
        "payload.inertia = [0.001, 0.001, 0.001]"
      ],
      "metadata": {
        "cellView": "form",
        "id": "eqgSVuaWZrp6"
      },
      "id": "eqgSVuaWZrp6",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Generate data for both configurations"
      ],
      "metadata": {
        "id": "1bfX4jpOZuLY"
      },
      "id": "1bfX4jpOZuLY"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title generate_data() { vertical-output: true}\n",
        "def generate_data(spec, duration=0.6):\n",
        "    \"\"\"Rollout a spec and return (control_ts, sensor_ts, initial_state).\"\"\"\n",
        "    model = spec.compile()\n",
        "    data = mujoco.MjData(model)\n",
        "\n",
        "    n_steps = int(duration / model.opt.timestep)\n",
        "    t = np.arange(n_steps) * model.opt.timestep\n",
        "\n",
        "    ctrl = (3.0 * np.sin(2 * np.pi * 2.0 * t)\n",
        "            + 2.0 * np.sin(2 * np.pi * 5.0 * t)).reshape(-1, 1)\n",
        "\n",
        "    initial_state = sysid.create_initial_state(\n",
        "        model, data.qpos, data.qvel, data.act\n",
        "    )\n",
        "    state, sensor = rollout.rollout(model, data, initial_state, ctrl[:-1])\n",
        "    state = np.squeeze(state, axis=0)\n",
        "    sensor = np.squeeze(sensor, axis=0)\n",
        "    times = state[:, 0]\n",
        "\n",
        "    control_ts = sysid.TimeSeries(t, ctrl)\n",
        "    sensor_ts = sysid.TimeSeries.from_names(times, sensor, model)\n",
        "    return control_ts, sensor_ts, initial_state\n",
        "\n",
        "\n",
        "ctrl_base, sens_base, state0_base = generate_data(spec_base)\n",
        "ctrl_pay, sens_pay, state0_pay = generate_data(spec_pay)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "VnFVQDLUZtfH"
      },
      "id": "VnFVQDLUZtfH",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Define parameters: mass and gear"
      ],
      "metadata": {
        "id": "L_o821XHZwZm"
      },
      "id": "L_o821XHZwZm"
    },
    {
      "cell_type": "code",
      "source": [
        "def set_cart_mass(spec, param):\n",
        "    spec.body(\"cart\").mass = param.value[0]\n",
        "\n",
        "def set_gear_ratio(spec, param):\n",
        "    spec.actuator(\"motor\").gear[0] = param.value[0]\n",
        "\n",
        "def make_params(mass_init, gear_init):\n",
        "    \"\"\"Create a ParameterDict with given starting values.\"\"\"\n",
        "    params = sysid.ParameterDict()\n",
        "    params.add(sysid.Parameter(\n",
        "        \"mass\", nominal=TRUE_MASS, min_value=0.5, max_value=5.0,\n",
        "        modifier=set_cart_mass,\n",
        "    ))\n",
        "    params.add(sysid.Parameter(\n",
        "        \"gear\", nominal=TRUE_GEAR, min_value=0.5, max_value=8.0,\n",
        "        modifier=set_gear_ratio,\n",
        "    ))\n",
        "    params[\"mass\"].value[:] = mass_init\n",
        "    params[\"gear\"].value[:] = gear_init\n",
        "    return params"
      ],
      "metadata": {
        "id": "KTomDxK3ZxQA"
      },
      "id": "KTomDxK3ZxQA",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Attempt 1: single sequence (base cart only)\n",
        "\n",
        "We start at $m = 3.5,\\; b = 5.25$, which gives the correct ratio\n",
        "$b/m = 1.5$ but wrong individual values. Since the cost is exactly zero\n",
        "everywhere along $b/m = 1.5$, the optimizer has no gradient to follow."
      ],
      "metadata": {
        "id": "jmbqF09FZzLc"
      },
      "id": "jmbqF09FZzLc"
    },
    {
      "cell_type": "code",
      "source": [
        "params_1seq = make_params(mass_init=3.5, gear_init=5.25)\n",
        "\n",
        "ms_base = sysid.ModelSequences(\n",
        "    \"cart\", spec_base, \"base_traj\", state0_base, ctrl_base, sens_base,\n",
        ")\n",
        "\n",
        "residual_fn_1seq = sysid.build_residual_fn(models_sequences=[ms_base])\n",
        "\n",
        "opt_1seq, result_1seq = sysid.optimize(\n",
        "    initial_params=params_1seq,\n",
        "    residual_fn=residual_fn_1seq,\n",
        "    optimizer=\"mujoco\",\n",
        ")\n",
        "\n",
        "print(f\"\\n--- Single sequence ---\")\n",
        "print(f\"  mass:  {opt_1seq['mass'].value[0]:.4f}  (true: {TRUE_MASS})\")\n",
        "print(f\"  gear:  {opt_1seq['gear'].value[0]:.4f}  (true: {TRUE_GEAR})\")\n",
        "print(f\"  b/m:   {opt_1seq['gear'].value[0] / opt_1seq['mass'].value[0]:.4f}\"\n",
        "      f\"  (true: {TRUE_GEAR / TRUE_MASS:.4f})\")"
      ],
      "metadata": {
        "id": "8ZSKfzE2Z0eQ"
      },
      "id": "8ZSKfzE2Z0eQ",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The optimizer converged instantly to the **wrong** individual values!\n",
        "The ratio $b/m$ is correct, but the optimizer has no way to determine $m$ and\n",
        "$b$ separately."
      ],
      "metadata": {
        "id": "yJT6PmkBZ1hy"
      },
      "id": "yJT6PmkBZ1hy"
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Attempt 2: two sequences (base + 1 kg payload)\n",
        "\n",
        "Adding a second trajectory with a **known 1.0 kg payload** gives the\n",
        "optimizer a second equation:\n",
        "- Base:    acceleration $= b / m$\n",
        "- Payload: acceleration $= b / (m + 1)$\n",
        "\n",
        "Two equations, two unknowns."
      ],
      "metadata": {
        "id": "mccbM8lmZ4O5"
      },
      "id": "mccbM8lmZ4O5"
    },
    {
      "cell_type": "code",
      "source": [
        "params_2seq = make_params(mass_init=3.5, gear_init=5.25)\n",
        "\n",
        "ms_payload = sysid.ModelSequences(\n",
        "    \"cart_payload\", spec_pay, \"payload_traj\", state0_pay, ctrl_pay, sens_pay,\n",
        ")\n",
        "\n",
        "residual_fn_2seq = sysid.build_residual_fn(\n",
        "    models_sequences=[ms_base, ms_payload],\n",
        ")\n",
        "\n",
        "opt_2seq, result_2seq = sysid.optimize(\n",
        "    initial_params=params_2seq,\n",
        "    residual_fn=residual_fn_2seq,\n",
        "    optimizer=\"mujoco\",\n",
        ")\n",
        "\n",
        "print(f\"\\n--- Two sequences (base + payload) ---\")\n",
        "print(f\"  mass:  {opt_2seq['mass'].value[0]:.4f}  (true: {TRUE_MASS})\")\n",
        "print(f\"  gear:  {opt_2seq['gear'].value[0]:.4f}  (true: {TRUE_GEAR})\")\n",
        "print(f\"  b/m:   {opt_2seq['gear'].value[0] / opt_2seq['mass'].value[0]:.4f}\"\n",
        "      f\"  (true: {TRUE_GEAR / TRUE_MASS:.4f})\")"
      ],
      "metadata": {
        "id": "lC2gdzR1Z5LS"
      },
      "id": "lC2gdzR1Z5LS",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Visualize the cost landscape\n",
        "\n",
        "Let's evaluate the cost on a grid of $(m, b)$ values to see the degeneracy."
      ],
      "metadata": {
        "id": "0guywON0Z6aS"
      },
      "id": "0guywON0Z6aS"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title compute_cost_surface() { vertical-output: true}\n",
        "mass_grid = np.linspace(0.6, 4.5, 35)\n",
        "gear_grid = np.linspace(0.6, 7.5, 35)\n",
        "\n",
        "def compute_cost_surface(residual_fn, params_template):\n",
        "    \"\"\"Evaluate cost on a (mass, gear) grid.\"\"\"\n",
        "    cost = np.zeros((len(mass_grid), len(gear_grid)))\n",
        "    p = params_template.copy()\n",
        "    for i, m in enumerate(mass_grid):\n",
        "        for j, g in enumerate(gear_grid):\n",
        "            x = np.array([m, g])\n",
        "            res, _, _ = residual_fn(x, p)\n",
        "            cost[i, j] = sum(np.sum(r**2) for r in res)\n",
        "    return cost\n",
        "\n",
        "cost_1seq = compute_cost_surface(residual_fn_1seq, params_1seq)\n",
        "cost_2seq = compute_cost_surface(residual_fn_2seq, params_2seq)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "V1l3YsFiZ7Z_"
      },
      "id": "V1l3YsFiZ7Z_",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Side-by-side cost landscapes"
      ],
      "metadata": {
        "id": "4zh6A4CTZ-NB"
      },
      "id": "4zh6A4CTZ-NB"
    },
    {
      "cell_type": "code",
      "source": [
        "#@title { vertical-output: true}\n",
        "G, M = np.meshgrid(gear_grid, mass_grid)\n",
        "\n",
        "# Compute log cost with shared color range across both panels\n",
        "log_cost_1 = np.log10(cost_1seq + 1e-12)\n",
        "log_cost_2 = np.log10(cost_2seq + 1e-12)\n",
        "vmin = min(log_cost_1.min(), log_cost_2.min())\n",
        "vmax = max(log_cost_1.max(), log_cost_2.max())\n",
        "levels = np.linspace(vmin, vmax, 30)\n",
        "\n",
        "fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True,\n",
        "                         layout=\"constrained\")\n",
        "\n",
        "for ax, log_cost, title, opt_p in [\n",
        "    (axes[0], log_cost_1, \"Single sequence\", opt_1seq),\n",
        "    (axes[1], log_cost_2, \"Two sequences (base + payload)\", opt_2seq),\n",
        "]:\n",
        "    cf = ax.contourf(G, M, log_cost, levels=levels, cmap=\"viridis\")\n",
        "\n",
        "    ax.plot(TRUE_GEAR, TRUE_MASS, \"r*\", markersize=15, label=\"True\", zorder=5)\n",
        "    ax.plot(opt_p[\"gear\"].value[0], opt_p[\"mass\"].value[0],\n",
        "            marker=\"X\", color=\"gold\", markeredgecolor=\"k\", markeredgewidth=1,\n",
        "            markersize=12, linestyle=\"none\", label=\"Optimized\", zorder=5)\n",
        "\n",
        "    m_line = np.linspace(0.5, 5.0, 200)\n",
        "    ax.plot(1.5 * m_line, m_line, \"r--\", alpha=0.5, lw=1, label=\"b/m = 1.5\")\n",
        "\n",
        "    ax.set_xlabel(\"Gear ratio (b)\")\n",
        "    ax.set_title(title)\n",
        "    ax.grid(True, alpha=0.2)\n",
        "\n",
        "axes[0].set_ylabel(\"Mass (m)\")\n",
        "\n",
        "# Shared colorbar\n",
        "fig.colorbar(cf, ax=axes, label=r\"$\\log_{10}$(cost)\", shrink=0.9)\n",
        "\n",
        "# Single legend below both plots\n",
        "handles, labels = axes[0].get_legend_handles_labels()\n",
        "fig.legend(handles, labels, loc=\"outside lower center\", ncol=3, fontsize=9)\n",
        "\n",
        "plt.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "vK2UVNnEZ_BW"
      },
      "id": "vK2UVNnEZ_BW",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Left:** With a single sequence, the cost is **exactly zero** along the entire\n",
        "$b/m = 1.5$ line. The optimizer stays at its starting point because there is no\n",
        "gradient to follow. Mass and gear are fundamentally unidentifiable.\n",
        "\n",
        "**Right:** Adding the payload trajectory collapses the valley into a localized\n",
        "minimum at the true parameter values $(m=2, b=3)$.\n",
        "\n",
        "**Takeaway:** It is easy to create models with unidentifiable parameters. When this is the case, adding diverse excitations and structural perturbations (e.g., added mass) can resolve it."
      ],
      "metadata": {
        "id": "siTyueLpaAGi"
      },
      "id": "siTyueLpaAGi"
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Conclusion\n",
        "\n",
        "We hope you enjoyed this introduction to system identification with MuJoCo.\n",
        "The library has additional features not covered here, including:\n",
        "\n",
        "- **[Physically plausible inertia parameterization](https://ieeexplore.ieee.org/document/9690029)**\n",
        "  for identifying mass, center of mass, and rotational inertia while\n",
        "  guaranteeing the result is physically valid\n",
        "- **Per-sensor weighting** to emphasize certain sensors in the cost\n",
        "- **Robust loss functions** (Huber, Cauchy, etc.) for handling outliers in\n",
        "  measured data (see the\n",
        "  [Least Squares](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/least_squares.ipynb)\n",
        "  notebook for details on non-quadratic norms)\n",
        "- **Multiple optimizer backends** (MuJoCo native, scipy, scipy with parallel\n",
        "  finite differences)\n",
        "\n",
        "For a comprehensive survey of the field, see\n",
        "[Robot Model Identification and Learning: A Modern Perspective](https://www.annualreviews.org/content/journals/10.1146/annurev-control-061523-102310)\n",
        "(Annual Review of Control, Robotics, and Autonomous Systems, 2024)."
      ],
      "metadata": {
        "id": "OggQL1VemwMV"
      },
      "id": "OggQL1VemwMV"
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.13.7"
    },
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "collapsed_sections": [
        "8P1e2JXckIbW",
        "_7k6iikqkIP6",
        "nfGPT2f2ZPFh",
        "UtAAiNGIlQo-"
      ],
      "toc_visible": true
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}