Kevin Zakka
2396a2541c
-- 4cabc9c52dc07ea2d64a3de6dbd58e47a2ad023c by Kevin Zakka <kevinarmandzakka@gmail.com>: Add system identification tutorial and fix template packaging - Add sysid tutorial notebook (sysid.ipynb) - Include HTML templates in package-data to fix TemplateNotFound errors when pip installing mujoco[sysid] COPYBARA_INTEGRATE_REVIEW=https://github.com/google-deepmind/mujoco/pull/3092 from kevinzakka:sysid-notebook 4cabc9c52dc07ea2d64a3de6dbd58e47a2ad023c PiperOrigin-RevId: 868900119 Change-Id: I53c4cc94b5fb6f864c5cee589af9dfad9bfc8744
1420 lines
52 KiB
Plaintext
1420 lines
52 KiB
Plaintext
{
|
|
"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
|
|
}
|