examples/mujoco_sysid_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g-cUBZu4hnyE"
   },
   "source": [
    "\n",
    "<a href=\"https://colab.research.google.com/github/lvjonok/mujoco-sysid/blob/master/examples/mujoco_sysid_demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\" width=\"150\" align=\"center\"/></a>\n",
    "<h1><center>System Identification and Linear Parametrization in Robotic Systems</center></h1>\n",
    "\n",
    "This notebook offers a brief introduction to system identification (Sys-ID) in robotic systems. It is inspired by recent advancements in MuJoCo utilities, specifically the [Levenberg-Marquardt nonlinear least squares method](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/least_squares.ipynb).\n",
    "\n",
    "We focus primarily on mechanical systems where the **model structure is known**, including the number of state variables and the configuration of the kinematic tree. Although the dynamics of these systems can be inherently complex, the general forms of the equations are known and have already been implemented in MuJoCo. In this context, the task of identification essentially involves estimating the parameters within a structured model.\n",
    "\n",
    "### Contents of This Notebook\n",
    "\n",
    "This notebook includes the following examples:\n",
    "* [Estimation of cart-pole inertial parameters through random forcing and LQR stabilization of the identified system](#cell-bwnDFHRjc836)\n",
    "* [Identification of end-effector load for the Franka Emika Panda and compensation using inverse dynamics](#cell-puimYTxi6OqK)\n",
    "* [Determination of mass, center of mass, and spatial inertia for a Skydio X2 Quadrotor following trajectory via LTV LQR](#cell-cKG867aUc84I)\n",
    "\n",
    "In addition to the examples, we provide theoretical background and introduce two utility functions that can greatly facilitate system identification and adaptive control in robotic systems:\n",
    "\n",
    "* `mj_bodyRegressor`: Calculates the regressor matrix for a single rigid body in the MuJoCo model.\n",
    "* `mj_jointRegressor`: Calculates the regressor matrix for the entire inverse dynamics in the MuJoCo model.\n",
    "\n",
    "We hope these examples and utilities will be useful for all MuJoCo users and assist in resolving their system identification challenges.\n",
    "\n",
    "For further questions and suggestions, please do not hesitate to create an issue or start a discussion [on GitHub](https://github.com/lvjonok/mujoco-sysid).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F8vVVm3ipYL0"
   },
   "source": [
    "\n",
    "###### Imports and config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jCWbI_ZaiuVm"
   },
   "outputs": [],
   "source": [
    "!pip install mujoco\n",
    "\n",
    "# Set up GPU rendering.\n",
    "from google.colab import files\n",
    "import distutils.util\n",
    "import os\n",
    "import subprocess\n",
    "\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",
    "\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",
    "            \"\"\"{\n",
    "    \"file_format_version\" : \"1.0.0\",\n",
    "    \"ICD\" : {\n",
    "        \"library_path\" : \"libEGL_nvidia.so.0\"\n",
    "    }\n",
    "}\n",
    "\"\"\"\n",
    "        )\n",
    "\n",
    "# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n",
    "print(\"Setting environment variable to use GPU rendering:\")\n",
    "%env MUJOCO_GL=egl\n",
    "\n",
    "# Check if installation was succesful.\n",
    "try:\n",
    "    print(\"Checking that the installation succeeded:\")\n",
    "    import mujoco\n",
    "\n",
    "    mujoco.MjModel.from_xml_string(\"<mujoco/>\")\n",
    "except Exception as e:\n",
    "    raise e from RuntimeError(\n",
    "        \"Something went wrong during installation. Check the shell output above \"\n",
    "        \"for more information.\\n\"\n",
    "        \"If using a hosted Colab runtime, make sure you enable GPU acceleration \"\n",
    "        'by going to the Runtime menu and selecting \"Choose runtime type\".'\n",
    "    )\n",
    "\n",
    "print(\"Installation successful.\")\n",
    "\n",
    "# Other imports and helper functions\n",
    "import time\n",
    "import itertools\n",
    "import numpy as np\n",
    "\n",
    "# Graphics and plotting.\n",
    "print(\"Installing mediapy:\")\n",
    "!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
    "!pip install -q mediapy\n",
    "import mediapy as media\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# More legible printing from numpy.\n",
    "np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
    "\n",
    "from IPython.display import clear_output\n",
    "\n",
    "clear_output()\n",
    "\n",
    "# Fix random seed\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UM-DG_N8BGKG"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip3 install dm_control robot_descriptions matplotlib scipy\n",
    "\n",
    "\n",
    "import copy\n",
    "import mujoco\n",
    "from dm_control import mjcf\n",
    "import numpy as np\n",
    "import mediapy as media\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bwnDFHRjc836"
   },
   "source": [
    "## Cart Pole Sys-ID and LQR stablization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SbjvqRr0tkcw"
   },
   "source": [
    "In the following introduction example we will consider the cart pole system, with unknown pendulum length and mass, as well as cart mass. The goal is to show how one can deduce these parameters given the system trajectory and measurements."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pO4Pnjo_uOJu"
   },
   "source": [
    "First let us download and load the model of cart pole:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EDYhZ-IMvWtc",
    "outputId": "c351b55f-a924-4533-f648-9c0e4ed45167"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/lvjonok/mujoco-sysid.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "doAtzzt3c838"
   },
   "outputs": [],
   "source": [
    "CART_POLE_MJCF = \"mujoco-sysid/data/models/cart_pole/model.xml\"\n",
    "\n",
    "sim_model = mujoco.MjModel.from_xml_path(CART_POLE_MJCF)\n",
    "sim_data = mujoco.MjData(sim_model)\n",
    "renderer = mujoco.Renderer(sim_model, height=480, width=640)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8bxaxGtwc839"
   },
   "source": [
    "### LQR stabilization\n",
    "\n",
    "As a first test we will perform the LQR stabilization assuming the all parameters are known, the LQR is done by computing linearization similarlly to **LQR tutorial** <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"90\"/></a>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AwXU5yHHc839"
   },
   "outputs": [],
   "source": [
    "import scipy\n",
    "\n",
    "\n",
    "def compute_gains(model, data, configuration, Q, R):\n",
    "    # 1. Set configuration and find control that stabilizes it (ctrl0)\n",
    "    newdata = mujoco.MjData(model)\n",
    "    newdata = copy.copy(data)\n",
    "\n",
    "    mujoco.mj_resetData(model, newdata)\n",
    "    newdata.qpos = configuration\n",
    "    # compute the control\n",
    "    mujoco.mj_forward(model, newdata)\n",
    "    newdata.qacc = 0\n",
    "    mujoco.mj_inverse(model, newdata)\n",
    "\n",
    "    # define control and configuration to linearize around\n",
    "    print(newdata.qfrc_actuator)\n",
    "    ctrl0 = newdata.qfrc_inverse.copy() @ np.linalg.pinv(newdata.actuator_moment)\n",
    "    qpos0 = newdata.qpos.copy()\n",
    "\n",
    "    # 2. Compute LQR gains given weights\n",
    "    mujoco.mj_resetData(model, newdata)\n",
    "    newdata.ctrl = ctrl0\n",
    "    newdata.qpos = qpos0\n",
    "\n",
    "    # 3. Allocate the A and B matrices, compute them.\n",
    "    A = np.zeros((2 * model.nv, 2 * model.nv))\n",
    "    B = np.zeros((2 * model.nv, model.nu))\n",
    "    epsilon = 1e-6\n",
    "    flg_centered = True\n",
    "    mujoco.mjd_transitionFD(model, newdata, epsilon, flg_centered, A, B, None, None)\n",
    "\n",
    "    # Solve discrete Riccati equation.\n",
    "    P = scipy.linalg.solve_discrete_are(A, B, Q, R)\n",
    "\n",
    "    # Compute the feedback gain matrix K.\n",
    "    K = np.linalg.inv(R + B.T @ P @ B) @ B.T @ P @ A\n",
    "\n",
    "    return ctrl0, K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TCJ2hiyjwljn"
   },
   "source": [
    "Perform the simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "NrCqEGykc83-",
    "outputId": "8d6744b3-c9b8-41ac-ad20-f45a5ab599be"
   },
   "outputs": [],
   "source": [
    "# Parameters.\n",
    "DURATION = 3  # seconds\n",
    "BALANCE_STD = 0.2  # actuator units\n",
    "\n",
    "dq = np.zeros(sim_model.nv)\n",
    "\n",
    "qpos0 = np.array([0, 2.05])\n",
    "target = np.array([0, np.pi])\n",
    "\n",
    "Q = np.diag([1, 1, 1, 1]) * 1e2\n",
    "R = np.eye(sim_model.nu)\n",
    "\n",
    "# ctrl0, K = compute_gains(uncertain_model, uncertain_data, target)\n",
    "ctrl0, K = compute_gains(sim_model, sim_data, target, Q, R)\n",
    "\n",
    "# Reset data, set initial pose.\n",
    "mujoco.mj_resetData(sim_model, sim_data)\n",
    "sim_data.qpos = qpos0\n",
    "\n",
    "qhist = []\n",
    "\n",
    "while sim_data.time < DURATION:\n",
    "    # Get state difference dx.\n",
    "    mujoco.mj_differentiatePos(sim_model, dq, 1, target, sim_data.qpos)\n",
    "    dx = np.hstack((dq, sim_data.qvel)).T\n",
    "\n",
    "    # LQR control law.\n",
    "    sim_data.ctrl = ctrl0 - K @ dx + np.random.randn(sim_model.nu) * BALANCE_STD\n",
    "\n",
    "    # Step the simulation.\n",
    "    mujoco.mj_step(sim_model, sim_data)\n",
    "\n",
    "    # Save history.\n",
    "    qhist.append(sim_data.qpos.copy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 452
    },
    "id": "crpuusvlc83_",
    "outputId": "76385904-b223-4d75-f9ef-477eb726ab0e"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "qhist = np.array(qhist)\n",
    "plt.title(\"Stabilization of ideal model\")\n",
    "plt.plot(qhist[:, 0], label=\"x\")\n",
    "plt.plot(qhist[:, 1], label=\"theta\")\n",
    "\n",
    "plt.hlines(0, 0, len(qhist), color=\"black\", linestyle=\"--\")\n",
    "plt.hlines(np.pi, 0, len(qhist), color=\"black\", linestyle=\"--\")\n",
    "\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o0tD3Ntoc83_"
   },
   "source": [
    "### Introducing Uncertainty\n",
    "\n",
    "let us alter dynamical parameters of the cart-pole as if we were unsure of them and made a some guess\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PUWYDlRnc83_",
    "outputId": "f45b2a15-5ba5-4eb5-f1fe-5cc3e2c34299"
   },
   "outputs": [],
   "source": [
    "# create a model with uncertainty, with the same structure as the original model\n",
    "uncertain_model = mjcf.from_path(CART_POLE_MJCF)\n",
    "# change the mass of the cart\n",
    "cart = uncertain_model.find(\"body\", \"cart\")\n",
    "cart.inertial.mass = 1  # from 0.5 to 1.0\n",
    "\n",
    "# change the mass of the pole\n",
    "pole = uncertain_model.find(\"body\", \"pole\")\n",
    "pole.inertial.mass = 0.5  # from 0.1 to 0.5\n",
    "# change the length of the pole\n",
    "pole.inertial.pos[2] = -0.9  # from -0.5 to -0.9\n",
    "\n",
    "# create a model with uncertainty\n",
    "uncertain_model = mjcf.Physics.from_mjcf_model(uncertain_model).model._model\n",
    "uncertain_data = mujoco.MjData(uncertain_model)\n",
    "\n",
    "ctrl0, K = compute_gains(uncertain_model, uncertain_data, target, Q, R)\n",
    "\n",
    "# Reset data, set initial pose.\n",
    "mujoco.mj_resetData(sim_model, sim_data)\n",
    "sim_data.qpos = qpos0\n",
    "\n",
    "qhist_uncertain = []\n",
    "\n",
    "while sim_data.time < DURATION:\n",
    "    # Get state difference dx.\n",
    "    mujoco.mj_differentiatePos(sim_model, dq, 1, target, sim_data.qpos)\n",
    "    dx = np.hstack((dq, sim_data.qvel)).T\n",
    "\n",
    "    # LQR control law.\n",
    "    sim_data.ctrl = ctrl0 - K @ dx + np.random.randn(sim_model.nu) * BALANCE_STD\n",
    "\n",
    "    # Step the simulation.\n",
    "    mujoco.mj_step(sim_model, sim_data)\n",
    "\n",
    "    # Save history.\n",
    "    qhist_uncertain.append(sim_data.qpos.copy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 452
    },
    "id": "FlJ3dEWyc84A",
    "outputId": "6cc4d74a-e87e-4fb8-ae60-dda300b91efd"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "qhist_uncertain = np.array(qhist_uncertain)\n",
    "plt.title(\"Stabilization of uncertain model\")\n",
    "plt.plot(qhist[:, 0], label=\"x\")\n",
    "plt.plot(qhist[:, 1], label=\"$\\\\theta$\")\n",
    "plt.plot(qhist_uncertain[:, 0], label=\"$x_{uncertain}$\")\n",
    "plt.plot(qhist_uncertain[:, 1], label=\"$\\\\theta_{uncertain}$\")\n",
    "\n",
    "plt.hlines(0, 0, len(qhist), color=\"black\", linestyle=\"--\")\n",
    "plt.hlines(np.pi, 0, len(qhist), color=\"black\", linestyle=\"--\")\n",
    "\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0v0Q6Nj0c84A"
   },
   "source": [
    "It is obvious that even though cart pole with uncertain model can still balance the pole, the controller performance is far from that provided by original LQR controller.\n",
    "\n",
    "In order to tackle this problem, we will try to identify the actual parameters of the cart pole system."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DF76GCRdw73o"
   },
   "source": [
    "### Cart Pole Identification\n",
    "\n",
    "Now let us build the identification routine, we will perfrom the several steps: this by gather data from the uncertain model and use it to identify the system dynamics. Then, we can design a controller based on the identified model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mpvuczk_c84A"
   },
   "source": [
    "#### Data gathering via random forcing\n",
    "\n",
    "We begin with gathering the data by randomly applying force to the cart:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Aox8d0bQc84A"
   },
   "outputs": [],
   "source": [
    "# random forcing and measure data\n",
    "\n",
    "# Reset data, set initial pose.\n",
    "mujoco.mj_resetData(sim_model, sim_data)\n",
    "\n",
    "identification_data = {\n",
    "    \"qhist\": [],\n",
    "    \"vhist\": [],\n",
    "    \"dvhist\": [],\n",
    "    \"uhist\": [],\n",
    "}\n",
    "\n",
    "for i in range(3000):\n",
    "    identification_data[\"qhist\"].append(sim_data.qpos.copy())\n",
    "    identification_data[\"vhist\"].append(sim_data.qvel.copy())\n",
    "\n",
    "    u = np.random.randn(1) * 10\n",
    "    identification_data[\"uhist\"].append(u)\n",
    "\n",
    "    sim_data.ctrl = u\n",
    "    mujoco.mj_step(sim_model, sim_data)\n",
    "    identification_data[\"dvhist\"].append(sim_data.qacc.copy())\n",
    "\n",
    "identification_data[\"qhist\"] = np.array(identification_data[\"qhist\"])\n",
    "identification_data[\"vhist\"] = np.array(identification_data[\"vhist\"])\n",
    "identification_data[\"dvhist\"] = np.array(identification_data[\"dvhist\"])\n",
    "identification_data[\"uhist\"] = np.array(identification_data[\"uhist\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fZk7DW0qc84B"
   },
   "source": [
    "#### Formulate regressors of system dynamics\n",
    "\n",
    "Once data is gathered we can formulate the sys-id problem, to do so we will dig into structure of problem at hand by recalling that dynamics of cart pole system is given by:\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\begin{cases}\n",
    "\\left(M+m\\right){\\ddot {q}_1}-m L \\ddot{q}_2 \\cos q_2 +m L \\dot{q}_2^{2}\\sin q_2 = u \\\\\n",
    "L \\ddot{q}_2- g\\sin q_2 =\\ddot{q}_1 \\cos q_2 \\\\\n",
    "\\end{cases}\n",
    "\\end{equation}\n",
    "$$\n",
    "\n",
    "where $q_2$ is angle of the pendulum measured from the upper equilibrium and $q_1$ is position of cart.\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CvhMzDUwzZSW"
   },
   "source": [
    "The notable property of any rigid body mechanical system  (cart-pole in particaluar that) the inverse dynamics [can be parametrized linearly to the set of constant inertial parameters](https://journals.sagepub.com/doi/abs/10.1177/027836498600500306?journalCode=ijra), indeed one may introduce the parameters $\\boldsymbol{\\theta} = [M+m, mL, L]$ and derive linearly parametrized inverse dynamics in so called regressor form:\n",
    "\n",
    "$$\n",
    "\\begin{bmatrix}\n",
    "\\ddot{q}_1 & - \\ddot{q}_2 \\cos \\theta + \\dot{q}_2^2 \\sin q_2 & 0 \\\\\n",
    "0 & 0 & \\ddot{q}_2 \\\\\n",
    "\\end{bmatrix} \\begin{bmatrix}\n",
    "M + m \\\\\n",
    "mL \\\\\n",
    "L\n",
    "\\end{bmatrix} = \\begin{bmatrix}\n",
    "u \\\\\n",
    "g \\sin q_2 + \\ddot{q}_1 \\cos q_2 \\\\\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "For given measurements of $\\mathbf{q}(t), \\dot{\\mathbf{q}}(t), \\ddot{\\mathbf{q}}(t), \\mathbf{u}(t)$ the above is the formulated as simple linear equation:\n",
    ":\n",
    "$$\n",
    "\\mathbf{A}(t)\\boldsymbol{\\theta} = \\mathbf{b}(t)\n",
    "$$\n",
    "\n",
    "Now given the sequance of time samples we may build the $\\mathbf{A}, \\mathbf{b}$ and solve associated linear least squares problem (posibly with bounds on parameters):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IuyzQtvOc84B"
   },
   "outputs": [],
   "source": [
    "# formulate regressor\n",
    "A = []\n",
    "b = []\n",
    "\n",
    "g = 9.81\n",
    "\n",
    "for i in range(1, len(qhist)):\n",
    "    q = identification_data[\"qhist\"][i]\n",
    "    v = identification_data[\"vhist\"][i]\n",
    "    dv = identification_data[\"dvhist\"][i]\n",
    "    u = identification_data[\"uhist\"][i]\n",
    "\n",
    "    regressor = [\n",
    "        [dv[0], -dv[1] * np.cos(q[1]) + v[1] ** 2 * np.sin(q[1]), 0],\n",
    "        [0, 0, dv[1]],\n",
    "    ]\n",
    "\n",
    "    A.extend(regressor)\n",
    "    b.extend(\n",
    "        [\n",
    "            u[0],\n",
    "            dv[0] * np.cos(q[1]) + g * np.sin(q[1]),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "A = np.array(A)\n",
    "b = np.array(b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Tswi8dHRc84B"
   },
   "source": [
    "#### Solving Least Squares problem:\n",
    "\n",
    "Once the regressor matrix is gathered we may solve the simple leas squares problem to get the parameter estimates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ehF7C7Qyc84B",
    "outputId": "3b3785e4-7231-477c-e515-f59c59e38100"
   },
   "outputs": [],
   "source": [
    "import mujoco.minimize as minimize\n",
    "\n",
    "\n",
    "def residual(x):\n",
    "    # Ensure x is 2D for consistent processing\n",
    "    if x.ndim == 1:\n",
    "        x = x[:, np.newaxis]\n",
    "\n",
    "    # Compute residuals using matrix operations\n",
    "    return A @ x - b[:, np.newaxis]\n",
    "\n",
    "\n",
    "result = minimize.least_squares(\n",
    "    np.array((0.0, 0.0, 0.0)),\n",
    "    residual,\n",
    ")\n",
    "\n",
    "identified_parameters = result[0]\n",
    "\n",
    "newL = identified_parameters[2]\n",
    "new_poleM = identified_parameters[1] / newL\n",
    "new_cartM = identified_parameters[0] - new_poleM\n",
    "\n",
    "identified_parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AxfCCcubc84C"
   },
   "source": [
    "### LQR on Identified System\n",
    "\n",
    "Ones parameters are identified we will build a new model and perform an updated LQR on identified system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mTyGfT2ic84C",
    "outputId": "ba09a8b0-7563-4413-9c0c-e5ae5dc113d0"
   },
   "outputs": [],
   "source": [
    "# create a model with identified parameters, with the same structure as the original model\n",
    "identified_model = mjcf.from_path(CART_POLE_MJCF)\n",
    "# change the mass of the cart\n",
    "cart = identified_model.find(\"body\", \"cart\")\n",
    "cart.inertial.mass = new_cartM\n",
    "\n",
    "# change the mass of the pole\n",
    "pole = identified_model.find(\"body\", \"pole\")\n",
    "pole.inertial.mass = new_poleM\n",
    "# change the length of the pole\n",
    "pole.inertial.pos[2] = -newL\n",
    "\n",
    "# create a model with uncertainty\n",
    "identified_model = mjcf.Physics.from_mjcf_model(identified_model).model._model\n",
    "identified_data = mujoco.MjData(identified_model)\n",
    "\n",
    "ctrl0, K = compute_gains(identified_model, identified_data, target, Q, R)\n",
    "\n",
    "# Reset data, set initial pose.\n",
    "mujoco.mj_resetData(sim_model, sim_data)\n",
    "sim_data.qpos = qpos0\n",
    "\n",
    "qhist_identified = []\n",
    "\n",
    "while sim_data.time < DURATION:\n",
    "    # Get state difference dx.\n",
    "    mujoco.mj_differentiatePos(sim_model, dq, 1, target, sim_data.qpos)\n",
    "    dx = np.hstack((dq, sim_data.qvel)).T\n",
    "\n",
    "    # LQR control law.\n",
    "    sim_data.ctrl = ctrl0 - K @ dx + np.random.randn(sim_model.nu) * BALANCE_STD\n",
    "\n",
    "    # Step the simulation.\n",
    "    mujoco.mj_step(sim_model, sim_data)\n",
    "\n",
    "    # Save history.\n",
    "    qhist_identified.append(sim_data.qpos.copy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 452
    },
    "id": "cnUTua7ic84C",
    "outputId": "e75ea305-a6ed-4b29-9a24-3496def018df"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "qhist_identified = np.array(qhist_identified)\n",
    "plt.title(\"Stabilization of identified model\")\n",
    "plt.plot(qhist[:, 0], label=\"x\", linestyle=\"--\")\n",
    "plt.plot(qhist[:, 1], label=\"$\\\\theta$\", linestyle=\"--\")\n",
    "plt.plot(qhist_identified[:, 0], label=\"$\\hat{x}$\")\n",
    "plt.plot(qhist_identified[:, 1], label=r\"$\\hat{\\theta}$\")\n",
    "plt.plot(qhist_uncertain[:, 0], label=\"$x_{uncertain}$\", alpha=0.5)\n",
    "plt.plot(qhist_uncertain[:, 1], label=r\"$\\theta_{uncertain}$\", alpha=0.5)\n",
    "\n",
    "plt.hlines(0, 0, len(qhist), color=\"black\", linestyle=\"--\")\n",
    "plt.hlines(np.pi, 0, len(qhist), color=\"black\", linestyle=\"--\")\n",
    "\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LzuQNsPgDdaj"
   },
   "source": [
    "as one can see the estimated feedback gains provided by identified parameters result in the response that is very similar to the original on thus seggesting the good estimation accuracy.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "173uHJU3FN3G"
   },
   "source": [
    "## The Building Blocks of System Identification\n",
    "Despite the simplicity of the cart-pole system, it serves as a useful starting point, as it encompasses the main building blocks that are consistently used in most system identification applications. These building blocks include:\n",
    "\n",
    "1. **Model Parametrization**: The first step involves constructing a parametrized model of the system, particularly a linear one. This parametrization captures the essential characteristics and dynamics of the system in a concise and manageable form.\n",
    "\n",
    "2. **Data Acquisition**: To identify the system parameters, it is necessary to gather data by applying specific forcing or excitation to the system. In the case of the cart-pole example, random forcing was used to generate informative data that captures the system's behavior under various conditions.\n",
    "\n",
    "3. **Least Squares Estimation**: With the obtained data, the next step is to formulate and solve a least squares problem. This involves minimizing the difference between the predicted outputs of the parametrized model and the actual measured outputs from the system. By solving the least squares problem, the optimal estimates of the system parameters can be determined.\n",
    "\n",
    "4. **Parameter Verification**: Once the system parameters have been estimated, it is crucial to verify their validity and performance in a particular scenario. This verification step ensures that the identified parameters accurately represent the system's behavior and can be used for control or other purposes.\n",
    "\n",
    "These building blocks form the foundation of system identification and are applicable to a wide range of systems, from simple toy problems to complex real-world applications for control, simulation, or analysis purposes.\n",
    "\n",
    "Having explored the basics of system identification through the cart-pole example, let us now move on to a more practical and challenging application: the load extimation on Panda manipulator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "puimYTxi6OqK"
   },
   "source": [
    "## Panda Load Estimation and Compensation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DR_fSJESc84C"
   },
   "source": [
    "\n",
    "\n",
    "Let's delve into a more hands-on challenge involving the [Panda manipulator](https://github.com/google-deepmind/mujoco_menagerie/tree/main/franka_emika_panda#readme). Suppose the manipulator has securely grasped an object, the mass and offset of which from the end-effector are unknown. The goal is to ascertain these parameters by guiding the robot through a sequence of predetermined configurations and monitoring the corresponding changes in torque.\n",
    "\n",
    "Once these parameters are precisely determined, we can proceed to compensate for the external load by updating the control policy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lgN5VmMgYcVx"
   },
   "source": [
    "#### Actuator Type\n",
    "\n",
    "Before progressing, it is crucial to set up all actuators for torque control. This configuration can be efficiently achieved by modifying the [actuator type](https://github.com/lvjonok/mujoco-actuators-types):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WPgB00GwfQzz"
   },
   "outputs": [],
   "source": [
    "from robot_descriptions.panda_mj_description import MJCF_PATH\n",
    "\n",
    "# get path to panda without gripper\n",
    "MJCF_PATH = MJCF_PATH[: -len(\"panda.xml\")] + \"panda_nohand.xml\"\n",
    "\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_-HjrpASiijQ"
   },
   "outputs": [],
   "source": [
    "class ActuatorMotor:\n",
    "    def __init__(self) -> None:\n",
    "        self.dyn = np.array([1, 0, 0])\n",
    "        self.gain = np.array([1, 0, 0])\n",
    "        self.bias = np.array([0, 0, 0])\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return f\"ActuatorMotor(dyn={self.dyn}, gain={self.gain}, bias={self.bias})\"\n",
    "\n",
    "\n",
    "def update_actuator(model, actuator_id, actuator):\n",
    "    \"\"\"\n",
    "    Update actuator in model\n",
    "    model - mujoco.MjModel\n",
    "    actuator_id - int or str (name) (for reference see, named access to model elements)\n",
    "    actuator - ActuatorMotor, ActuatorPosition, ActuatorVelocity\n",
    "    \"\"\"\n",
    "\n",
    "    model.actuator(actuator_id).dynprm[:3] = actuator.dyn\n",
    "    model.actuator(actuator_id).gainprm[:3] = actuator.gain\n",
    "    model.actuator(actuator_id).biasprm[:3] = actuator.bias\n",
    "    model.actuator(actuator_id).ctrlrange = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FcGr2dlBA-ZG"
   },
   "source": [
    "### Modifying the Model and Adding Load\n",
    "\n",
    "To facilitate our discussion, we will work with two distinct models. The `old_model` will represent the configuration without any load, whereas the `sim_model` will include a mass attached to the end effector (`link_7`).\n",
    "\n",
    "Additionally, to avoid complications during simulations, particularly as we will be working with randomly selected desired points in joint space, we will disable all self-collisions. This precaution helps prevent excessive joint torques that can arise from self-collisions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZRnlXSFWA9nA"
   },
   "outputs": [],
   "source": [
    "# create a model with external load\n",
    "mjcf_model = mjcf.from_path(MJCF_PATH, escape_separators=True)\n",
    "# get the body of end-effector\n",
    "ee_body = mjcf_model.find(\"body\", \"link7\")\n",
    "# create new body with offset from the end-effector\n",
    "rc = np.array([0.1, 0.05, -0.1])\n",
    "load_mass = ee_body.add(\"body\", name=\"load_mass\", pos=f\"{rc[0]} {rc[1]} {rc[2]}\")\n",
    "load_mass.add(\"geom\", name=\"load_mass_geom\", type=\"sphere\", size=\"0.05\", rgba=\"0.8 0.2 0.2 1\")\n",
    "\n",
    "# add inertial properties to the new body\n",
    "M = 2.0\n",
    "I = 1e-6\n",
    "load_mass.add(\"inertial\", mass=f\"{M}\", pos=\"0 0 0\", diaginertia=f\"{I} {I} {I}\")\n",
    "\n",
    "# disable collision for all the geoms\n",
    "for geom in mjcf_model.find_all(\"geom\"):\n",
    "    geom.conaffinity = 0\n",
    "    geom.contype = 0\n",
    "\n",
    "physics = mjcf.Physics.from_mjcf_model(mjcf_model)\n",
    "sim_model = physics.model._model\n",
    "sim_data = mujoco.MjData(sim_model)\n",
    "renderer = mujoco.Renderer(sim_model, height=480, width=640)\n",
    "\n",
    "# create an old model without the load mass\n",
    "old_mjcf_model = mjcf.from_path(MJCF_PATH, escape_separators=True)\n",
    "# disable collisions for all the geoms\n",
    "for geom in old_mjcf_model.find_all(\"geom\"):\n",
    "    geom.conaffinity = 0\n",
    "    geom.contype = 0\n",
    "\n",
    "old_model = mjcf.Physics.from_mjcf_model(old_mjcf_model).model._model\n",
    "old_data = mujoco.MjData(old_model)\n",
    "body_offset = np.where(old_model.body_dofadr == 0)[0][0]\n",
    "\n",
    "# update actuator type\n",
    "for actuator_id in range(sim_model.nu):\n",
    "    actuator = ActuatorMotor()\n",
    "    update_actuator(sim_model, actuator_id, actuator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A_ib_qo8BIXq"
   },
   "source": [
    "To better understand the changes and dynamics, we will visually represent the `sim_model` that includes the external load attached to the end effector (`link_7`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 503
    },
    "id": "3ZpMsi6BBGeb",
    "outputId": "b7b5be41-96bd-4905-c274-88362895c05f"
   },
   "outputs": [],
   "source": [
    "# Reset data, set initial pose.\n",
    "mujoco.mj_resetData(sim_model, sim_data)\n",
    "\n",
    "# Step the simulation.\n",
    "mujoco.mj_step(sim_model, sim_data)\n",
    "\n",
    "renderer.update_scene(sim_data)\n",
    "media.show_image(renderer.render())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iPBiq2Bzc84C"
   },
   "source": [
    "### Parametrizing the Model with External Load\n",
    "\n",
    "We begin with deriving external force and torque produced by the load on end effector local frame, the load is given by:\n",
    "$$\n",
    "m a_g = f_{ee}\n",
    "$$\n",
    "while the torque\n",
    "$$\n",
    "ma_g\\times r = \\tau_{ee}\n",
    "$$\n",
    "\n",
    "Here, $a_g$ is the acceleration due to gravity, specifically $a_g = R^T \\begin{bmatrix} 0 & 0 & 9.81 \\end{bmatrix}^T$ where $R$ is the rotation matrix of the end-effector.\n",
    "\n",
    "Thus external load can be parametrized using four parameters: mass $m$, and the first moment, i.e mass multiplied by its CoM position: $mr_x, mr_y, mr_z$. The linear parametrization of the static force exerted by the load can be represented as:\n",
    "\n",
    "$$\n",
    "\\begin{bmatrix}\n",
    "a_1 & 0 & 0 & 0 \\\\\n",
    "a_2 & 0 & 0 & 0 \\\\\n",
    "a_3 & 0 & 0 & 0 \\\\\n",
    "0 & 0 & a_3 & -a_2 \\\\\n",
    "0 & -a_3 & 0 & a_1 \\\\\n",
    "0 & a_2 & -a_1 & 0\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "m \\\\\n",
    "mr_x \\\\\n",
    "mr_y \\\\\n",
    "mr_z\n",
    "\\end{bmatrix} = \\begin{bmatrix}\n",
    "f_{ee} \\\\\n",
    "\\tau_{ee} \\\\\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "This matrix equation can be compactly expressed in matrix-vector notation as:\n",
    "\n",
    "$$\n",
    "\\mathbf{Y_b} \\theta = \\bar{f}\n",
    "$$\n",
    "\n",
    "where $\\mathbf{Y_b}$ is the matrix composed of gravitational components and cross-product terms, and $\\theta$ encapsulates the inertial parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "glA9gvmkDred"
   },
   "source": [
    "The implementation of the above regressor yields:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HpOo8Jz2DTkY"
   },
   "outputs": [],
   "source": [
    "def static_body_regressor(\n",
    "    mj_model: mujoco.MjModel, mj_data: mujoco.MjData, body_id: int, body_offset: int = 0\n",
    ") -> np.ndarray:\n",
    "    \"\"\"Calculates the static body regressor matrix for a given body.\n",
    "\n",
    "    Args:\n",
    "        mj_model (mujoco.MjModel): MuJoCo model.\n",
    "        mj_data (mujoco.MjData): MuJoCo data.\n",
    "        body_id (int): ID of the body.\n",
    "        body_offset (int, optional): Offset for the body ID. Defaults to 0.\n",
    "\n",
    "    Returns:\n",
    "        np.ndarray: Static body regressor matrix of shape (6, 4).\n",
    "    \"\"\"\n",
    "\n",
    "    # mujoco.mj_kinematics(mj_model, mj_data)\n",
    "    rotation = mj_data.xmat[body_id + body_offset].reshape(3, 3).copy()\n",
    "\n",
    "    a1, a2, a3 = rotation.T @ np.array([0, 0, 9.81])\n",
    "\n",
    "    static_load_regressor = np.array(\n",
    "        [\n",
    "            [a1, 0, 0, 0],\n",
    "            [a2, 0, 0, 0],\n",
    "            [a3, 0, 0, 0],\n",
    "            [0, 0, a3, -a2],\n",
    "            [0, -a3, 0, a1],\n",
    "            [0, a2, -a1, 0],\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    return static_load_regressor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZNhZy-YFD-sA"
   },
   "source": [
    "\n",
    "Transitioning to a focus on practical application, it is beneficial to reformulate the load model in terms of motor torques (commonly derived from motor currents). This reorientation leads us back to the inverse dynamics:\n",
    "\n",
    "$$\n",
    "M(q) \\ddot{q} + C(q, \\dot{q}) \\dot{q} + g(q) = \\tau - J^T(q) \\bar{f}\n",
    "$$\n",
    "\n",
    "Here, $M(q)$ is the mass matrix, $C(q, \\dot{q})$ is the Coriolis matrix, $G(q)$ is the gravitational force vector, and $\\tau$ represents the joint torques. In the absence of external loads and other applied wrenches, the last term would be omitted.\n",
    "\n",
    "In static conditions ($\\ddot{q} = 0, \\dot{q} = 0$) and substituting $\\bar{f} = \\mathbf{Y_b} \\theta$, we derive:\n",
    "\n",
    "$$\n",
    "g(q) = \\tau - J^T(q) \\mathbf{Y}_b(q) \\theta\n",
    "$$\n",
    "\n",
    "The term $\\mathbf{Y}_s(q) = J^T(q)\\mathbf{Y}_b(q) \\theta$, henceforth referred to as the joint regressor (for statics), illustrates how the end-effector load, parametrized linearly in terms of the load's mass and CoM, contributes to joint torques in static conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9a5oA1VQIonE"
   },
   "outputs": [],
   "source": [
    "def static_joint_regressor(\n",
    "    mj_model: mujoco.MjModel, mj_data: mujoco.MjData, body_id: int, active_joints: int = 7, body_offset: int = 0\n",
    ") -> np.ndarray:\n",
    "    \"\"\"Calculates the static joint regressor matrix for a given body.\n",
    "\n",
    "    Args:\n",
    "        mj_model (mujoco.MjModel): MuJoCo model.\n",
    "        mj_data (mujoco.MjData): MuJoCo data.\n",
    "        body_id (int): ID of the body.\n",
    "        active_joints (int, optional): Number of active joints. Defaults to 7.\n",
    "        body_offset (int, optional): Offset for the body ID. Defaults to 0.\n",
    "\n",
    "    Returns:\n",
    "        np.ndarray: Static joint regressor matrix of shape (active_joints, 4).\n",
    "    \"\"\"\n",
    "    accel = np.zeros(6)\n",
    "\n",
    "    mujoco.mj_rnePostConstraint(mj_model, mj_data)\n",
    "    mujoco.mj_objectAcceleration(mj_model, mj_data, 2, body_id, accel, 1)\n",
    "\n",
    "    static_load_regressor = static_body_regressor(mj_model, mj_data, body_id, body_offset)\n",
    "    rotation = mj_data.xmat[body_id + body_offset].reshape(3, 3).copy()\n",
    "\n",
    "    jac_lin = np.zeros((3, mj_model.nv))\n",
    "    jac_rot = np.zeros((3, mj_model.nv))\n",
    "    mujoco.mj_jacBody(mj_model, mj_data, jac_lin, jac_rot, body_id + body_offset)\n",
    "    ee_local_jacobian = np.vstack(\n",
    "        (\n",
    "            rotation.T @ jac_lin[:, :active_joints],\n",
    "            rotation.T @ jac_rot[:, :active_joints],\n",
    "        )\n",
    "    )\n",
    "    joint_regressor = ee_local_jacobian.T @ static_load_regressor\n",
    "\n",
    "    return joint_regressor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o7ZQaGOxNYzH"
   },
   "source": [
    "For practical implementation, see the code available in this [GitHub repository](https://github.com/lvjonok/mujoco-sysid/blob/master/mujoco_sysid/modeling.py), which demonstrates how to integrate more general case of full inertial parameters using the Mujoco `mj_data`. We will undertake similar exercises in subsequent examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "28N3FHLRc84D"
   },
   "source": [
    "### Load Estimation\n",
    "\n",
    "With the foundational concepts established, we are now equipped to initiate the main load estimation routine. The methodology will be structured as follows:\n",
    "\n",
    "1. **Joint Space Regulation:** Initially, we will control the manipulator to achieve a sequence of random points in joint space. This control will be executed using the inverse dynamics of an uncertain model.\n",
    "\n",
    "2. **Data Collection during Steady State:** After allowing sufficient time for the transients to settle at each target point, we will record the applied torques. Concurrently, we will calculate the joint regressors, which are crucial for the next steps of load estimation.\n",
    "\n",
    "3. **Least Squares Estimation:** With the collected data, we will employ a least squares method to estimate the load's mass and center of mass (CoM).\n",
    "\n",
    "4. **Model Compensation and Verification:** After estimating the load parameters, we will update the control policy to compensate for the identified load and verify the accuracy of our estimates.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rR9jMNH1UTPB"
   },
   "source": [
    "First we generate random points in joint space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "psLrrVOhc84E"
   },
   "outputs": [],
   "source": [
    "# find limits except the last two fingers\n",
    "lower, upper = np.zeros(sim_model.nq), np.zeros(sim_model.nq)\n",
    "for jnt_idx in range(sim_model.nq):\n",
    "    lower[jnt_idx] = sim_model.joint(jnt_idx).range[0]\n",
    "    upper[jnt_idx] = sim_model.joint(jnt_idx).range[1]\n",
    "\n",
    "num_of_phases = 3\n",
    "waypoints = [np.random.uniform(lower, upper) for _ in range(num_of_phases)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O25tbZOMUd09"
   },
   "source": [
    "#### Controller Implementation on Approximate Dynamics and Data Collection\n",
    "\n",
    "To ensure stability at each targeted joint position, we will employ a basic inverse dynamics controller defined by the following equation:\n",
    "\n",
    "$$\n",
    "u = M_0(q) (K_p\\tilde{q} - K_d\\dot{q} ) + C_0(q, \\dot{q}) \\dot{q} + g_0(q)\n",
    "$$\n",
    "\n",
    "It is important to note that this controller does not account for the load, as it relies solely on the nominal model parameters (`old_model`), indicated by the subscripts $M_0, C_0, g_0$. Consequently, the controller might not achieve zero tracking error. Additionally, the control input $u$ is applied to the actual uncertain model (`sim_model`), meaning the generalized angles and velocities reflect the real system conditions, designated as `old_data = sim_data`.\n",
    "\n",
    "Whenever the manipulator approximately reaches the desired point and the transient dynamics have settled, we will capture the regressor matrix along with the torque differences, which are formulated as:\n",
    "\n",
    "$$\n",
    "u - g_0(q) =  \\mathbf{Y}(q) \\theta\n",
    "$$\n",
    "\n",
    "This relationship facilitates the segregation of the dynamic effects solely influenced by the load:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VSg9aozYc84E"
   },
   "outputs": [],
   "source": [
    "# define PD controller for inverse dynamics\n",
    "kp = 50\n",
    "kd = 10\n",
    "target_q = waypoints[0]\n",
    "\n",
    "frames = []\n",
    "framerate = 25\n",
    "\n",
    "phase_time = 5.0\n",
    "phases = 0\n",
    "old_data.time = 0\n",
    "sim_data.time = 0\n",
    "\n",
    "load_joint_regressor = np.zeros(((sim_model.nq) * num_of_phases, 4))\n",
    "load_joint_force = np.zeros((sim_model.nq) * num_of_phases)\n",
    "active_joints = sim_model.nq\n",
    "\n",
    "hist_q = []\n",
    "target_hist_q = []\n",
    "\n",
    "while sim_data.time < (num_of_phases - 1) * phase_time:\n",
    "    # PD controller\n",
    "    mujoco.mj_step(sim_model, sim_data)\n",
    "\n",
    "    hist_q.append(sim_data.qpos.copy())\n",
    "    target_hist_q.append(target_q)\n",
    "\n",
    "    target_acc = kp * (target_q - sim_data.qpos) + kd * (np.zeros(sim_model.nv) - sim_data.qvel)\n",
    "\n",
    "    # use old data to solve inverse dynamics\n",
    "    old_data.qpos = sim_data.qpos.copy()\n",
    "    old_data.qvel = sim_data.qvel.copy()\n",
    "    old_data.qacc = target_acc\n",
    "    mujoco.mj_inverse(old_model, old_data)\n",
    "    sol = old_data.qfrc_inverse\n",
    "\n",
    "    sim_data.ctrl = sol\n",
    "\n",
    "    if sim_data.time > phase_time * phases and phases < num_of_phases:\n",
    "        # after each phase we add the regressor and the joint forces\n",
    "        # and change the target q\n",
    "        regressor = static_joint_regressor(old_model, old_data, old_data.body(ee_body.name).id).copy()\n",
    "\n",
    "        load_joint_regressor[7 * phases : 7 * (phases + 1), :] = regressor\n",
    "\n",
    "        load_qforce = sol[:7] - old_data.qfrc_bias[:7]\n",
    "        load_joint_force[7 * phases : 7 * (phases + 1)] = load_qforce.copy()\n",
    "        actual_load_qforce = regressor @ np.array([M, *(M * rc)])\n",
    "\n",
    "        target_q = waypoints[phases]\n",
    "        phases += 1\n",
    "\n",
    "hist_q = np.array(hist_q)\n",
    "target_hist_q = np.array(target_hist_q)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hojQJelEc84E"
   },
   "source": [
    "The plot below illustrates that the errors are not zero, indicating that the robot struggles to accurately reach the desired waypoints. This deviation from the expected positions is a predictable outcome, given the controller's inability to account for additional external forces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 468
    },
    "id": "0kCNltcsc84E",
    "outputId": "dc1b66c4-eb86-4406-86d2-9d74685cc435"
   },
   "outputs": [],
   "source": [
    "each_joint_color = [\"r\", \"g\", \"b\", \"c\", \"m\", \"y\", \"k\"]\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.title(\"Joint positions\")\n",
    "for jidx in range(7):\n",
    "    plt.plot(hist_q[:, jidx], color=each_joint_color[jidx])\n",
    "    plt.plot(target_hist_q[:, jidx], linestyle=\"--\", color=each_joint_color[jidx])\n",
    "\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3z2WG9XAYCQ-"
   },
   "source": [
    "#### Building and solving least sqaures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RnyWLGPvc84E"
   },
   "source": [
    "Having collected the necessary data, we can now proceed to structure this information into a system of equations suitable for least squares analysis. Each data point collected gives us an equation based on the deviation from the control input due to the unmodeled load:\n",
    "\n",
    "$$\n",
    "u - g_0(q) = \\Delta u_i = \\mathbf{Y}_i(q_i) \\theta\n",
    "$$\n",
    "\n",
    "Here, the subscript $i$ signifies the specific data point recorded. These individual equations are then compiled into larger matrices:\n",
    "\n",
    "$$\n",
    "\\begin{bmatrix}\n",
    "\\mathbf{Y}_1 \\\\\n",
    "\\mathbf{Y}_2 \\\\\n",
    "\\vdots \\\\\n",
    "\\mathbf{Y}_N\n",
    "\\end{bmatrix} \\theta = \\begin{bmatrix}\n",
    "\\Delta u_1 \\\\\n",
    "\\Delta u_2 \\\\\n",
    "\\vdots \\\\\n",
    "\\Delta u_N\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "This sets the stage for the least squares estimation, where the tall matrix on the left contains the stacked regressor matrices $\\mathbf{Y}_i$ for each point $i$, and the vector on the right consists of the corresponding torque deviations $\\Delta u_i$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o2XIYfHPYrn8"
   },
   "outputs": [],
   "source": [
    "def manip_load_ls(x):\n",
    "    # Ensure x is 2D for consistent processing\n",
    "    if x.ndim == 1:\n",
    "        x = x[:, np.newaxis]\n",
    "\n",
    "    # Compute residuals using matrix operations\n",
    "    return load_joint_regressor[7:] @ x - load_joint_force[7:][:, np.newaxis]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4yp4LME6Yzhn"
   },
   "source": [
    "This initial point typically  (first 7 elements) is subject to sudden changes in acceleration and velocities due to the initial response of the PD controller. These abrupt dynamics can introduce significant noise and anomalies, potentially leading to inaccurate estimations if included in our model, thus we ignore it above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hRMFhUznc84G"
   },
   "source": [
    "Now, we move forward to solve the associated least squares problem and will also print out the load parameters for reference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "K3YPv2xkYauz",
    "outputId": "681a80f6-4b36-4ec0-e7bd-a7422a7d4be7"
   },
   "outputs": [],
   "source": [
    "M, rc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UV05OvIbc84G",
    "outputId": "3b42e062-c3c1-484b-e458-76afd55ccf51"
   },
   "outputs": [],
   "source": [
    "import mujoco.minimize as minimize\n",
    "\n",
    "result = minimize.least_squares(\n",
    "    np.ones(4),  # Initial guess for the parameter values\n",
    "    manip_load_ls,\n",
    ")\n",
    "load_estimates = result[0]\n",
    "print(load_estimates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rjLKs6sRal4C"
   },
   "source": [
    "From these estimates, we can parse the actual center of mass position:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JvzbHUctalXF",
    "outputId": "a1a91ec0-c774-4471-f0c6-ba54fb298f21"
   },
   "outputs": [],
   "source": [
    "mass = load_estimates[0]\n",
    "coordinate = load_estimates[1:] / mass\n",
    "print(mass, coordinate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w2GllP2pawOy"
   },
   "source": [
    "As observed, the estimates are remarkably close to the actual external load. It's worth noting that this problem is fundamentally a linear least squares, which can be resolved using the simple pseudoinverse method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8_Sbl1oQFciB",
    "outputId": "da5041d3-f468-4910-98d2-885bd11bd97a"
   },
   "outputs": [],
   "source": [
    "load_pinv_estimates = np.linalg.pinv(load_joint_regressor[7:, :]) @ load_joint_force[7:]\n",
    "mass_pinv = load_pinv_estimates[0]\n",
    "coordinate_pinv = load_pinv_estimates[1:] / mass\n",
    "mass_pinv, coordinate_pinv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nmdY3J5HXr8C"
   },
   "source": [
    "However, if one wants to account for bounds on the parameters, nonlinear dependencies such as friction models, or physical consistency, it is advisable to utilize a more general framework provided by minimize.least_squares."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Hr_uVYvSc84G"
   },
   "source": [
    "#### Verifying and Utilizing the Estimated Load\n",
    "To verify and utilize the estimated load parameters effectively, we can modify our control policy to include the estimated parameters.\n",
    "\n",
    "This adjustment involves adding an additional term to the inverse controller:\n",
    "$$\n",
    "u = M_0(q) (K_p\\tilde{q} - K_d\\dot{q} ) + C_0(q, \\dot{q}) \\dot{q} + g_0(q) + \\mathbf{Y}(q) \\hat{\\theta}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TrhN0X4Hc84H"
   },
   "outputs": [],
   "source": [
    "# define PD controller for inverse dynamics\n",
    "kp = 50\n",
    "kd = 10\n",
    "target_q = waypoints[0]\n",
    "\n",
    "frames = []\n",
    "framerate = 25\n",
    "\n",
    "phase_time = 5.0\n",
    "phases = 0\n",
    "old_data.time = 0\n",
    "sim_data.time = 0\n",
    "\n",
    "load_joint_regressor = np.zeros(((sim_model.nq) * num_of_phases, 4))\n",
    "load_joint_force = np.zeros((sim_model.nq) * num_of_phases)\n",
    "active_joints = sim_model.nq\n",
    "\n",
    "hist_q_new = []\n",
    "target_hist_q_new = []\n",
    "\n",
    "while sim_data.time < (num_of_phases - 1) * phase_time:\n",
    "    # PD controller\n",
    "    mujoco.mj_step(sim_model, sim_data)\n",
    "\n",
    "    hist_q_new.append(sim_data.qpos.copy())\n",
    "    target_hist_q_new.append(target_q)\n",
    "\n",
    "    target_acc = kp * (target_q - sim_data.qpos) + kd * (-sim_data.qvel)\n",
    "\n",
    "    # use old data to solve inverse dynamics\n",
    "    old_data.qpos = sim_data.qpos.copy()\n",
    "    old_data.qvel = sim_data.qvel.copy()\n",
    "    old_data.qacc = target_acc\n",
    "    mujoco.mj_inverse(old_model, old_data)\n",
    "    sol = old_data.qfrc_inverse\n",
    "\n",
    "    actual_load_qforce = (\n",
    "        static_joint_regressor(old_model, old_data, old_data.body(ee_body.name).id).copy() @ load_estimates\n",
    "    )\n",
    "\n",
    "    sim_data.ctrl = sol + actual_load_qforce\n",
    "\n",
    "    if sim_data.time > phase_time * phases and phases < num_of_phases:\n",
    "        target_q = waypoints[phases]\n",
    "        phases += 1\n",
    "\n",
    "hist_q_new = np.array(hist_q_new)\n",
    "target_hist_q_new = np.array(target_hist_q_new)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WQVZYC1Mc84H"
   },
   "source": [
    "now let us plot the results of tracking (the updated model is given by solid lines):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 468
    },
    "id": "yCE6yjJBc84H",
    "outputId": "b9ab952d-b38a-4174-c747-9f43a69d3f38"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 5))\n",
    "plt.title(\"Joint positions after estimating load parameters\")\n",
    "for jidx in range(7):\n",
    "    plt.plot(hist_q_new[:, jidx], color=each_joint_color[jidx])\n",
    "    plt.plot(target_hist_q_new[:, jidx], linestyle=\"--\", color=each_joint_color[jidx])\n",
    "    plt.plot(hist_q[:, jidx], color=each_joint_color[jidx], alpha=0.5)\n",
    "\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 391
    },
    "id": "atw0o43vc84H",
    "outputId": "17cf73fa-868b-4323-a42c-b2ae5bbc9cc9"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(10, 4))\n",
    "\n",
    "ax[0].set_title(\"Joints errors before estimating load parameters\")\n",
    "ax[0].plot(hist_q - target_hist_q, alpha=0.5)\n",
    "ax[0].set_ylim(-1, 1)\n",
    "ax[0].grid()\n",
    "\n",
    "ax[1].set_title(\"Errors after compensation\")\n",
    "ax[1].plot(hist_q_new - target_hist_q_new, alpha=0.5)\n",
    "ax[1].set_ylim(-1, 1)\n",
    "ax[1].grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EjIzJfPrfkf6"
   },
   "source": [
    "As we can see, the updated control law, which incorporates the estimated load parameters, significantly improves the accuracy of joint tracking. This improvement is evident from the decreased error in joint positions when compared to the target positions.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ogBVBbxe22BG"
   },
   "source": [
    "Up to now, our discussions have primarily focused on the static estimation of load parameters based on known inputs in joints, with a particular emphasis on regressor forms, which include both body and joint regressors. These concepts can be effectively extended to dynamic scenarios by utilizing adaptive control techniques, such as those found in [Slotine's adaptive controller](https://journals.sagepub.com/doi/abs/10.1177/027836498700600303) or its [recent modifications](https://ieeexplore.ieee.org/abstract/document/8593727). These techniques are designed to continuously adjust and optimize control parameters in real-time based on observed system behaviors.\n",
    "\n",
    "However, before moving on to such advanced topics, let us first develop next regressors that will incorporate inertia parameters, starting with a simple single rigid body, specifically a model of a simple quadrotor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cKG867aUc84I"
   },
   "source": [
    "## Quadrotor Inertia Identification\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wcNQIpGE9V22"
   },
   "source": [
    "In the previous example, we considered estimating mass-related quantities, which were done basically through statics. However, some of the inertial parameters are not possible to extract under such circumstances; they inherently require active maneuvers.\n",
    "\n",
    "This is the topic of the next example. Here, we will take into account the full system dynamics, including the inertia tensor:\n",
    "\n",
    "$$\n",
    "\\mathbf{I} =\\left[\\begin{array}{lll}\n",
    "I_{x x} & I_{x y} & I_{x z} \\\\\n",
    "I_{x y} & I_{y y} & I_{y z} \\\\\n",
    "I_{x z} & I_{y z} & I_{z z}\n",
    "\\end{array}\\right]\n",
    "$$\n",
    "\n",
    "Note that the inertia tensor is represented in the body-fixed frame. Since the inertia tensor is symmetric, there are only $6$ unique parameters coming from this matrix. Adding mass and CoM position thus yields the following set of inertial parameters:\n",
    "\n",
    "$$\n",
    "\\mathbf{\\theta} = \\begin{bmatrix} m, mr_x, mr_y, mr_z, I_{xx}, I_{xy}, I_{yy}, I_{xz}, I_{yz}, I_{zz} \\end{bmatrix}^T\n",
    "$$\n",
    "\n",
    "Although these ten parameters are sufficient, they may not be the best possible set to represent the inertial properties of rigid bodies that are [tightly coupled with density distribution](https://arxiv.org/pdf/1701.04395). However, we will use these parameters as they are the most intuitive and easy to work with.\n",
    "\n",
    "Our goal will be to estimate these ten parameters based on measurement data provided by a controlled Skydio X2 quadrotor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FPBNo7BYJ7fe"
   },
   "outputs": [],
   "source": [
    "from robot_descriptions.skydio_x2_mj_description import MJCF_PATH\n",
    "\n",
    "model = mujoco.MjModel.from_xml_path(MJCF_PATH[: -len(\"x2.xml\")] + \"scene.xml\")\n",
    "data = mujoco.MjData(model)\n",
    "renderer = mujoco.Renderer(model, height=480, width=640)\n",
    "\n",
    "for act_id in range(4):\n",
    "    model.actuator(act_id).ctrlrange = np.array([-1e4, 1e4])\n",
    "\n",
    "model.opt.timestep = 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 503
    },
    "id": "RYlhrU6RO3WT",
    "outputId": "ca41cf5f-e4e2-4dce-bd99-2512ab1aa31f"
   },
   "outputs": [],
   "source": [
    "mujoco.mj_step(model, data)\n",
    "renderer.update_scene(data)\n",
    "\n",
    "media.show_image(renderer.render())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mTBAAO6pCcb2"
   },
   "source": [
    "One may easily parse the inertial parameters $\\theta$ from it's sparce spatial form M (here it is trivial since $M$ is fixed in local frame):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1_Oxr7hVCU10",
    "outputId": "1bac7d9d-380b-4f5a-9854-19adcdbf9f90"
   },
   "outputs": [],
   "source": [
    "sparceM = np.zeros((6, 6))\n",
    "\n",
    "mujoco.mj_step(model, data)\n",
    "mujoco.mj_fullM(model, sparceM, data.qM)\n",
    "\n",
    "mSrc = sparceM[3:6, :3]\n",
    "mrc = np.array([mSrc[2, 1], mSrc[0, 2], mSrc[1, 0]])\n",
    "Ib = sparceM[3:6, 3:6]\n",
    "vechIb = np.array([Ib[0, 0], Ib[0, 1], Ib[1, 1], Ib[0, 2], Ib[1, 2], Ib[2, 2]])\n",
    "skydio_parameters = np.array([sparceM[0, 0], *mrc, *vechIb])\n",
    "\n",
    "with np.printoptions(precision=3, suppress=True):\n",
    "    print(skydio_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A5_MtscU1WB6"
   },
   "source": [
    "Now, let us treat these parameters as the true baseline and attempt to estimate them from one maneuver executed by the quadrotor. This process will involve several key steps:\n",
    "\n",
    "* Simulate **the system response** under a specific controller (in this case, we will obtain an offline trajectory and perform Linear Quadratic Regulator (LQR) tracking). This step is designed to simulate the data coming from the actual quadrotor control system.\n",
    "\n",
    "* Introduce the **full-body regressor** - a matrix that relates the inertial parameters to the measured data in a simple linear form.\n",
    "\n",
    "* Formulate and solve the least squares problem based on the body regressor and compare the parameter estimates with true values."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XeqyP_o8_UAU"
   },
   "source": [
    "#### LTV LQR Controller\n",
    "\n",
    "Firstly, we will simulate the response of an actual control system by exploiting the LTV LQR technique. You may skip this step since it is not related to the problem at hand - estimating the model inertia. However, this example is useful on its own since LTV LQR is a widely used controller in practice. Below, we will introduce utility functions and some brief descriptions of the theory required to roughly understand time-varying LQR. For notation and more insight, please follow the original **LQR tutorial** <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"90\"/></a>.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fmMj_-K6IzVh"
   },
   "source": [
    "First we will load precomputed in advance desired trajectory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "B9Ve6XKXI3uc",
    "outputId": "15a4e3bf-6d98-44c1-b913-76ebc1de7bf2"
   },
   "outputs": [],
   "source": [
    "# load trajectory\n",
    "traj = np.load(\"mujoco-sysid/data/trajectories/skydio.npz\")\n",
    "tf = traj[\"t\"]\n",
    "\n",
    "q_data = traj[\"q\"].T\n",
    "v_data = traj[\"v\"].T\n",
    "u_data = traj[\"u\"].T\n",
    "N_DATA = q_data.shape[1]\n",
    "t_data = np.linspace(0, tf, N_DATA)\n",
    "\n",
    "q_traj = scipy.interpolate.interp1d(t_data, q_data, axis=1, kind=\"linear\")\n",
    "v_traj = scipy.interpolate.interp1d(t_data, v_data, axis=1, kind=\"linear\")\n",
    "u_traj = scipy.interpolate.interp1d(t_data[:-1], u_data, axis=1, kind=\"linear\")\n",
    "N_TRAJ = int(tf / model.opt.timestep)\n",
    "traj[\"q\"].shape, traj[\"v\"].shape, traj[\"u\"].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9UgvxLpOIux8"
   },
   "source": [
    "\n",
    "\n",
    "**System Model**\n",
    "\n",
    "For an LTV system, the state-space representation at each discrete time index $t$ is given by:\n",
    "\n",
    "$$x_{t+1} = A_t x_t + B_t u_t$$\n",
    "\n",
    "where $x_t$ is the state vector, $u_t$ is the control input, $A_t$ is the time-varying state transition matrix, and $B_t$ is the time-varying control input matrix.\n",
    "\n",
    "**Linearization and Jacobian Matrices**\n",
    "\n",
    "The matrices $A_t$ and $B_t$ represent the linearization of the system dynamics around a trajectory point at time $t$. They can be obtained by evaluating the Jacobian matrices of the nonlinear system function $f$ with respect to the state $x$ and control input $u$ at each point along the trajectory:\n",
    "\n",
    "$$A_t = \\frac{\\partial f}{\\partial x}\\bigg|_{x_t, u_t}, \\quad B_t = \\frac{\\partial f}{\\partial u}\\bigg|_{x_t, u_t}$$\n",
    "\n",
    "This is why $A_t$ and $B_t$ are time-varying, as they depend on the specific trajectory point $(x_t, u_t)$ at each time step $t$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eUieY9T1c84I"
   },
   "outputs": [],
   "source": [
    "def linearize(q: np.ndarray, v: np.ndarray, u: np.ndarray) -> tuple[np.ndarray, np.ndarray]:\n",
    "    \"\"\"Finds the linearization over a point.\n",
    "\n",
    "    Args:\n",
    "        q (np.ndarray): Generalized coordinates (positions).\n",
    "        v (np.ndarray): Generalized velocities.\n",
    "        u (np.ndarray): Control inputs.\n",
    "\n",
    "    Returns:\n",
    "        tuple[np.ndarray, np.ndarray]: A tuple containing:\n",
    "            - A (np.ndarray): State transition matrix of shape (2 * model.nv, 2 * model.nv).\n",
    "            - B (np.ndarray): Control input matrix of shape (2 * model.nv, model.nu).\n",
    "    \"\"\"\n",
    "    newdata = mujoco.MjData(model)\n",
    "    newdata = copy.copy(data)\n",
    "\n",
    "    mujoco.mj_resetData(model, newdata)\n",
    "    newdata.ctrl = u\n",
    "    newdata.qpos = q\n",
    "    newdata.qvel = v\n",
    "\n",
    "    # 4. Allocate the A and B matrices, compute them.\n",
    "    A = np.zeros((2 * model.nv, 2 * model.nv))\n",
    "    B = np.zeros((2 * model.nv, model.nu))\n",
    "    epsilon = 1e-6  # model.opt.timestep\n",
    "    flg_centered = False\n",
    "    mujoco.mjd_transitionFD(model, newdata, epsilon, flg_centered, A, B, None, None)\n",
    "\n",
    "    return A, B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f0hAtCLP8qqA"
   },
   "source": [
    "\n",
    "**Cost Function**\n",
    "\n",
    "The objective is to minimize a time-varying quadratic cost function over a finite horizon:\n",
    "\n",
    "$$J = \\sum_{t=0}^{T-1} (x_t^\\top Q_t x_t + u_t^\\top R_t u_t) + x_T^\\top P_T x_T$$\n",
    "\n",
    "Here, $Q_t$ and $R_t$ are symmetric positive definite matrices that may vary with time, and $P_T$ is the terminal cost matrix.\n",
    "\n",
    "**Optimal Control Law**\n",
    "\n",
    "The control law that minimizes the cost function $J$ is given by:\n",
    "\n",
    "$$u_t^* = -K_t x_t$$\n",
    "\n",
    "where $K_t$ are the optimal feedback gain matrices derived from the solution of the Riccati equation.\n",
    "\n",
    "**Riccati Equation (LTV)**\n",
    "\n",
    "The gain matrices $K_t$ and cost-to-go matrices $P_t$ are computed using a backward recursion through the Riccati equation:\n",
    "\n",
    "$$P_t = Q + K_t^\\top R K_t + (A_t - B_t K_t)^\\top P_{t+1} (A_t - B_t K_t)$$\n",
    "\n",
    "where the LQR gain $K_t$ is computed as:\n",
    "\n",
    "$$K_t = (R + B_t^\\top P_{t+1} B_t)^{-1} B_t^\\top P_{t+1} A_t$$\n",
    "\n",
    "starting from the terminal condition $P_{N-1} = Q$, this terminal condition may be calculated from steady state solution.\n",
    "\n",
    "**Discrete Algebraic Riccati Equation (LTI)**\n",
    "\n",
    "For an infinite horizon problem with constant system matrices $A$ and $B$, and cost matrices $Q$ and $R$, the steady-state solution $P$ can be determined by solving the Discrete Algebraic Riccati Equation (DARE):\n",
    "\n",
    "$$P = Q + A^\\top P A - A^\\top P B (R + B^\\top P B)^{-1} B^\\top P A$$\n",
    "\n",
    "The LQR gain $K$ is then computed as:\n",
    "\n",
    "$$K = (B^\\top P B + R)^{-1} B^\\top P A$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pEvpqbiAc84I"
   },
   "outputs": [],
   "source": [
    "from scipy.linalg import solve_discrete_are as dare\n",
    "\n",
    "\n",
    "def dlqr(A: np.ndarray, B: np.ndarray, Q: np.ndarray, R: np.ndarray) -> tuple[np.ndarray, np.ndarray]:\n",
    "    \"\"\"Discrete time LTI LQR.\n",
    "\n",
    "    Args:\n",
    "        A (np.ndarray): State transition matrix of shape (n, n).\n",
    "        B (np.ndarray): Control input matrix of shape (n, m).\n",
    "        Q (np.ndarray): State cost matrix of shape (n, n).\n",
    "        R (np.ndarray): Control cost matrix of shape (m, m).\n",
    "\n",
    "    Returns:\n",
    "        tuple[np.ndarray, np.ndarray]: A tuple containing:\n",
    "            - K (np.ndarray): LQR gain matrix of shape (m, n).\n",
    "            - P (np.ndarray): Cost-to-go matrix of shape (n, n).\n",
    "    \"\"\"\n",
    "    # Solve discrete Ricatti equation (DARE)\n",
    "    P = dare(A, B, Q, R)\n",
    "    # Compute the LQR gain\n",
    "    K = np.linalg.inv(B.T @ P @ B + R) @ (B.T @ P @ A)\n",
    "    return K, P\n",
    "\n",
    "\n",
    "def dlqr_ltv(A: np.ndarray, B: np.ndarray, Q: np.ndarray, R: np.ndarray) -> tuple[np.ndarray, np.ndarray]:\n",
    "    \"\"\"Discrete time LTV LQR controller.\n",
    "\n",
    "    Args:\n",
    "        A (np.ndarray): State transition matrices of shape (N, n, n).\n",
    "        B (np.ndarray): Control input matrices of shape (N, n, m).\n",
    "        Q (np.ndarray): State cost matrix of shape (n, n).\n",
    "        R (np.ndarray): Control cost matrix of shape (m, m).\n",
    "\n",
    "    Returns:\n",
    "        tuple[np.ndarray, np.ndarray]: A tuple containing:\n",
    "            - K (np.ndarray): LQR gains of shape (N, m, n).\n",
    "            - P (np.ndarray): Cost-to-go matrices of shape (N, n, n).\n",
    "    \"\"\"\n",
    "    N = np.shape(A)[0]\n",
    "    n = np.shape(A)[1]\n",
    "    m = np.shape(B)[2]\n",
    "\n",
    "    # Matrix to store cost-to-go\n",
    "    P = np.zeros((N, n, n))\n",
    "    P[-1] = Q\n",
    "    # Matrix to store gains\n",
    "    K = np.zeros((N, m, n))\n",
    "\n",
    "    for k in range(N - 2, -1, -1):\n",
    "        # Compute the LQR gain\n",
    "        invM_k = np.linalg.inv(R + B[k].T @ P[k + 1] @ B[k])\n",
    "        K[k] = invM_k @ B[k].T @ P[k + 1] @ A[k]\n",
    "        Ac_k = A[k] - B[k] @ K[k]  # Update with difference Riccati equation\n",
    "        # Compute cost-to-go matrix\n",
    "        P[k] = Q + K[k].T @ R @ K[k] + Ac_k.T @ P[k + 1] @ Ac_k\n",
    "\n",
    "    return K, P"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7IlZ6q1180QJ"
   },
   "source": [
    "\n",
    "**Feedback Gains and policy**\n",
    "\n",
    "The feedback gains $K_t$ are computed as:\n",
    "\n",
    "$$K_t = (R_t + B_t^\\top P_{t+1} B_t)^{-1} B_t^\\top P_{t+1} A_t$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A4esbnTs89nY"
   },
   "source": [
    "Now we may calculate the feedback gains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2iZk9iWGc84J"
   },
   "outputs": [],
   "source": [
    "# linearize around trajectory\n",
    "A_traj = []\n",
    "B_traj = []\n",
    "\n",
    "traj_t = traj[\"t\"]\n",
    "\n",
    "traj_dt = model.opt.timestep\n",
    "\n",
    "for i in range(N_DATA):\n",
    "    t = i * traj_dt\n",
    "    A, B = linearize(q_traj(t), v_traj(t), u_traj(t))\n",
    "    A_traj.append(A)\n",
    "    B_traj.append(B)\n",
    "\n",
    "Q = np.eye(2 * model.nv)\n",
    "R = np.eye(4) * 0.3\n",
    "\n",
    "# find gains for stabilization at the end\n",
    "A_stab, B_stab = linearize(\n",
    "    q_data[:, -1],\n",
    "    np.zeros(model.nv),\n",
    "    np.ones(4) * 1.325 * 9.81 / 4,\n",
    ")\n",
    "# find gains for stabilization at the end\n",
    "K_stab, P_stab = dlqr(A_stab, B_stab, Q, R)\n",
    "K, _ = dlqr_ltv(A_traj, B_traj, Q, R)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U133st8m7HQE"
   },
   "source": [
    "Now we are ready to simulate the controller behaviour with following feedback policy:\n",
    "\n",
    "$$u = u_t -K_t (x - x_t)$$\n",
    "\n",
    "where $x$ is the current state\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 503
    },
    "id": "AFdbYCELKRZl",
    "outputId": "8f8a8f73-d57d-442d-8f71-0ae2545ff85a"
   },
   "outputs": [],
   "source": [
    "FRAMERATE = 20  # Hz\n",
    "\n",
    "# store different in tangent space between configurations\n",
    "dq = np.zeros(model.nv)\n",
    "\n",
    "# Reset data, set initial pose.\n",
    "mujoco.mj_resetData(model, data)\n",
    "data.qpos = q_traj(0)\n",
    "data.qvel = np.zeros(6)\n",
    "last_time = 0\n",
    "phase = 0\n",
    "\n",
    "hist_qpos = []\n",
    "hist_target_qpos = []\n",
    "\n",
    "frames = []\n",
    "k = 0\n",
    "\n",
    "ltv_log = {\n",
    "    \"q\": [],\n",
    "    \"v\": [],\n",
    "    \"dv\": [],\n",
    "    \"u\": [],\n",
    "}\n",
    "\n",
    "while data.time < tf + 3.0:\n",
    "    hist_qpos.append(data.qpos.copy())\n",
    "    t = data.time\n",
    "\n",
    "    k = int(t / (tf / N_DATA))\n",
    "\n",
    "    if t < tf - 0.1:\n",
    "        q_target = q_traj(t)\n",
    "        v_target = v_traj(t)\n",
    "        u_target = u_traj(t)\n",
    "\n",
    "        # Get state difference dx.\n",
    "        mujoco.mj_differentiatePos(model, dq, 1, q_target, data.qpos.copy())\n",
    "\n",
    "        dv = data.qvel.copy() - v_target\n",
    "        dx = np.hstack((dq, dv)).T\n",
    "\n",
    "        # LQR control law.\n",
    "        data.ctrl = u_target - K[k] @ dx\n",
    "    else:\n",
    "        # stabilize around the last configuration\n",
    "        q_target = q_data[:, -1]\n",
    "        v_target = np.zeros(model.nv)\n",
    "        u_target = np.ones(4) * 1.325 * 9.81 / 4\n",
    "\n",
    "        # Get state difference dx.\n",
    "        mujoco.mj_differentiatePos(model, dq, 1, q_target, data.qpos.copy())\n",
    "\n",
    "        dv = data.qvel.copy() - v_target\n",
    "        dx = np.hstack((dq, dv)).T\n",
    "\n",
    "        # LQR control law.\n",
    "        data.ctrl = u_target - K_stab @ dx\n",
    "\n",
    "    mujoco.mj_step(model, data)\n",
    "\n",
    "    hist_target_qpos.append(q_target.copy())\n",
    "\n",
    "    # extend the log\n",
    "    ltv_log[\"q\"].append(data.qpos.copy())\n",
    "    ltv_log[\"v\"].append(data.qvel.copy())\n",
    "    ltv_log[\"dv\"].append(data.qacc.copy())\n",
    "    ltv_log[\"u\"].append(data.ctrl.copy())\n",
    "\n",
    "    if len(frames) < data.time * framerate:\n",
    "        renderer.update_scene(data, camera=0)\n",
    "        pixels = renderer.render()\n",
    "        frames.append(pixels)\n",
    "\n",
    "\n",
    "ltv_log[\"q\"] = np.array(ltv_log[\"q\"])\n",
    "ltv_log[\"v\"] = np.array(ltv_log[\"v\"])\n",
    "ltv_log[\"dv\"] = np.array(ltv_log[\"dv\"])\n",
    "ltv_log[\"u\"] = np.array(ltv_log[\"u\"])\n",
    "\n",
    "media.show_video(frames, fps=framerate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 545
    },
    "id": "rRRpouTCN67v",
    "outputId": "b0369f82-cc27-4b1b-c293-2a299bfc4425"
   },
   "outputs": [],
   "source": [
    "hist_qpos = np.array(hist_qpos)\n",
    "hist_target_qpos = np.array(hist_target_qpos)\n",
    "\n",
    "plt.figure(figsize=(12, 6))\n",
    "plt.title(\"Position tracking\")\n",
    "plt.plot(hist_qpos[:, :3], label=[\"x\", \"y\", \"z\"])\n",
    "plt.plot(hist_target_qpos[:, :3], linestyle=\"--\")\n",
    "plt.ylim([-5, 5])\n",
    "\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gLu8m9S3c84J"
   },
   "source": [
    "#### Model Parametrization and Body Regressor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WttohcXVzVMC"
   },
   "source": [
    "To effectively parametrize the model, we refer to the Newton-Euler formulation in spatial form, which is a pivotal concept in rigid body dynamics. If you need to recap this, refer to Roy Featherstone's book, [\"Rigid Body Dynamics Algorithms\"](https://link.springer.com/book/10.1007/978-1-4899-7560-7).\n",
    "\n",
    "In a nutshell, the dynamics of a single rigid body is expressed as:\n",
    "\n",
    "$$\n",
    "M \\mathrm{a}_{g} + \\bar{v} \\hat{\\times} M \\bar{v} = \\bar{f}\n",
    "$$\n",
    "\n",
    "where:\n",
    "- $M \\in \\mathbb{R}^{6 \\times 6}$ is the spatial inertia matrix of the body,\n",
    "$$\n",
    "M=\\left[\\begin{array}{cc}\n",
    "m 1_{3} & -m S(r_c)\\\\\n",
    "m S(r_c) & I  \n",
    "\\end{array}\\right]\n",
    "$$\n",
    "\n",
    "- $\\bar{v} \\in \\mathbb{R}^6$ represents the spatial velocity,\n",
    "- $\\bar{f} \\in \\mathbb{R}^6$ is the external wrench applied,\n",
    "- $g \\in \\mathbb{R}^{3}$ is the constant vector of gravity acceleration expressed with the orientation of the inertial frame $A$.\n",
    "\n",
    "- $\\mathrm{a}^{g} \\in \\mathbb{R}^{6}$ denotes the proper body acceleration, i.e., the difference between the body **spatial acceleration** and the gravity acceleration:  \n",
    "\n",
    "$$\n",
    "\\mathrm{a}_{g}=\\dot{\\bar{\\mathrm{v}}}-\\left[\\begin{array}{c}\n",
    "{ }^{A} R_{B}^{\\top} g \\\\\n",
    "0_{3}  \n",
    "\\end{array}\\right]\n",
    "$$\n",
    "\n",
    "- $\\mathrm{v} \\hat{\\times}$ is the $6 \\mathrm{D}$ cross product operator, defined as:\n",
    "\n",
    "$$\n",
    "\\left[\\begin{array}{l} v \\\\  \n",
    "\\omega\n",
    "\\end{array}\\right] \\hat{\\times}=\\left[\\begin{array}{cc}\n",
    "S(\\omega) & 0_{3 \\times 3} \\\\\n",
    "S(v) & S(\\omega)  \n",
    "\\end{array}\\right]  \n",
    "$$  \n",
    "\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "\n",
    "As one can see in the Inverse Dynamics (ID) for a single rigid body, there are only linear operations (multiplication with matrices) with respect to the spatial inertia matrix $M$. Thus, the force given by ID is **linearly parametrizable**:\n",
    "\n",
    "$$\n",
    "M \\mathrm{a}_{g} + \\bar{v} \\hat{\\times} M \\bar{v} = \\mathbf{Y}(\\bar{v}, \\mathrm{a}_{g})\\boldsymbol{\\theta} = \\bar{f}\n",
    "$$\n",
    "\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "\n",
    "where $\\mathbf{Y}(\\bar{v}, \\mathrm{a}^{g}) \\in \\mathbb{R}^{6\\times10}$ is the so-called **body regressor**.\n",
    "\n",
    "The easiest way to obtain this matrix is to take the partial derivative of force $\\bar{f}$ with respect to parameters $\\boldsymbol{\\theta}$, or just reparametrize cross products in a certain way. To shorten this, we provide handy notes on this topic in a Colab notebook <a href=\"https://colab.research.google.com/drive/1xFte2FT0nQ0ePs02BoOx4CmLLw5U-OUZ#scrollTo=MJEQtlpX0-Ya\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"90\"/></a>. Consider studying it; later on, we will add more theoretical information there.\n",
    "\n",
    "All in all, one may derive this linear parametrization and obtain the following equations for the body regressor:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-sa_JJajc84J"
   },
   "outputs": [],
   "source": [
    "def mj_bodyRegressor(mj_model: mujoco.MjModel, mj_data: mujoco.MjData, body_id: int) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Computes the body regressor matrix for a specific body in a MuJoCo model.\n",
    "\n",
    "    Args:\n",
    "        mj_model (mujoco.MjModel): MuJoCo model.\n",
    "        mj_data (mujoco.MjData): MuJoCo data.\n",
    "        body_id (int): ID of the body for which to compute the regressor matrix.\n",
    "\n",
    "    Returns:\n",
    "        np.ndarray: Body regressor matrix of shape (6, 10).\n",
    "    \"\"\"\n",
    "    velocity = np.zeros(6)\n",
    "    accel = np.zeros(6)\n",
    "    _cross = np.zeros(3)\n",
    "\n",
    "    mujoco.mj_objectVelocity(mj_model, mj_data, 2, body_id, velocity, 1)\n",
    "    mujoco.mj_rnePostConstraint(mj_model, mj_data)\n",
    "    mujoco.mj_objectAcceleration(mj_model, mj_data, 2, body_id, accel, 1)\n",
    "\n",
    "    v, w = velocity[3:], velocity[:3]\n",
    "    dv, dw = accel[3:], accel[:3]  # dv - classical acceleration, already contains g\n",
    "    mujoco.mju_cross(_cross, w, v)\n",
    "    dv -= _cross\n",
    "\n",
    "    v1, v2, v3 = v\n",
    "    v4, v5, v6 = w\n",
    "\n",
    "    a1, a2, a3 = dv\n",
    "    a4, a5, a6 = dw\n",
    "\n",
    "    # fmt: off\n",
    "    Y = np.array(\n",
    "        [\n",
    "            [a1 - v2 * v6 + v3 * v5, -(v5**2) - v6**2, -a6 + v4 * v5, a5 + v4 * v6, 0, 0, 0, 0, 0, 0],\n",
    "            [a2 + v1 * v6 - v3 * v4, a6 + v4 * v5, -(v4**2) - v6**2, -a4 + v5 * v6, 0, 0, 0, 0, 0, 0],\n",
    "            [a3 - v1 * v5 + v2 * v4, -a5 + v4 * v6, a4 + v5 * v6, -(v4**2) - v5**2, 0, 0, 0, 0, 0, 0],\n",
    "            [0, 0, a3 - v1 * v5 + v2 * v4, -a2 - v1 * v6 + v3 * v4, a4, a5 - v4 * v6, -v5 * v6, a6 + v4 * v5, v5**2 - v6**2, v5 * v6],\n",
    "            [0, -a3 + v1 * v5 - v2 * v4, 0, a1 - v2 * v6 + v3 * v5, v4 * v6, a4 + v5 * v6, a5, -(v4**2) + v6**2, a6 - v4 * v5, -v4 * v6],\n",
    "            [0, a2 + v1 * v6 - v3 * v4, -a1 + v2 * v6 - v3 * v5, 0, -v4 * v5, v4**2 - v5**2, v4 * v5, a4 - v5 * v6, a5 + v4 * v6, a6],\n",
    "        ]\n",
    "    )\n",
    "    # fmt: on\n",
    "\n",
    "    return Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7ouVghAa7q6r"
   },
   "source": [
    "#### Parameter Estimation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-cwKHl7N8LTV"
   },
   "source": [
    "In fact now we are ready to do the parameter estimation itself, first we will use body regressor to build the stacked measurements and inputs:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f-n9BzqxDUGH"
   },
   "outputs": [],
   "source": [
    "LHS = []\n",
    "RHS = []\n",
    "\n",
    "for i in range(len(ltv_log[\"q\"])):\n",
    "    # Regressor computation\n",
    "    datacopy = copy.copy(data)\n",
    "    datacopy.qpos = ltv_log[\"q\"][i]\n",
    "    datacopy.qvel = ltv_log[\"v\"][i]\n",
    "    datacopy.qacc = ltv_log[\"dv\"][i]\n",
    "    datacopy.ctrl = ltv_log[\"u\"][i]\n",
    "\n",
    "    mujoco.mj_step(model, datacopy)\n",
    "    mujoco.mj_inverse(model, datacopy)\n",
    "\n",
    "    # Fill matrices\n",
    "    regressor = mj_bodyRegressor(model, datacopy, 1)\n",
    "    LHS.extend(regressor)\n",
    "\n",
    "    force = datacopy.qfrc_actuator.copy()\n",
    "    force[:3] = datacopy.xmat[1].reshape(3, 3).T @ force[:3]\n",
    "    RHS.extend(force)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "otfEsbHF_hR5"
   },
   "source": [
    "Once LHS and RHS terms are calculated we may proceed with least squares:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BKVdnnGRc84K",
    "outputId": "99f9cf3a-6da1-47b8-896f-0f38ccf7913e"
   },
   "outputs": [],
   "source": [
    "from mujoco.minimize import least_squares\n",
    "\n",
    "LHS = np.array(LHS)\n",
    "RHS = np.array(RHS)\n",
    "\n",
    "\n",
    "def skydio_ls(x):\n",
    "    # Ensure x is 2D for consistent processing\n",
    "    if x.ndim == 1:\n",
    "        x = x[:, np.newaxis]\n",
    "\n",
    "    # Compute residuals using matrix operations\n",
    "    return LHS @ x - RHS[:, np.newaxis]\n",
    "\n",
    "\n",
    "result_skydio, _ = least_squares(\n",
    "    np.zeros(10),\n",
    "    skydio_ls,\n",
    ")\n",
    "\n",
    "with np.printoptions(precision=3, suppress=True):\n",
    "    print(\"Identified parameters\", result_skydio)\n",
    "    print(\"Actual parameters\", skydio_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2LtuOF-mA5IU"
   },
   "source": [
    "Great! We've successfully extracted all inertial parameters from a single quadrotor trajectory, which is a promising result. However, it's important to acknowledge that we made several notable simplifications and assumptions. The most significant assumption is that the actual system trajectory, along with its first and second derivatives and the input signal, is known. In practice, the accelerations and velocities are obtained after differentiating positions that have been smoothed by non-causal filters, such as the [Savitzky-Golay](https://en.wikipedia.org/wiki/Savitzky%E2%80%93Golay_filter)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gCwdAJcd4K7Z"
   },
   "source": [
    "## Conclusion: Stay Tuned!\n",
    "\n",
    "This just the beginning of exploration into system identification, it's important to emphasize that the **linear parametrization of robot dynamics**, which we covered in this tutorial, forms the cornerstone of system identification and adaptive control methodologies. To support your learning and application of these concepts, we have provided two practical utility functions within [our GitHub repository](https://github.com/lvjonok/mujoco-sysid/blob/master/mujoco_sysid/modeling.py):\n",
    "\n",
    "- `mj_bodyRegressor(mj_model, mj_data, body_id)`: This function computes the regressor matrix $\\mathbf{Y}_b \\in \\mathbb{R}^{6 \\times 10}$ for a single rigid body in the MuJoCo model:\n",
    "\n",
    "  $$\n",
    "  M \\mathrm{a}_g + \\bar{v} \\hat{\\times} M \\bar{v} = \\mathbf{Y}_b(\\bar{v}, \\mathrm{a}_g)\\boldsymbol{\\theta} = \\bar{f}\n",
    "  $$\n",
    "\n",
    "- `mj_jointRegressor(mj_model, mj_data)`: This function calculates the regressor $\\mathbf{Y} \\in \\mathbb{R}^{nv \\times 10nb}$ (where \\(nb\\) is the number of bodies) for the whole inverse dynamics in the MuJoCo model:\n",
    "\n",
    "  $$\n",
    "  M(q) \\dot{v} + C(q, v) v + g(q) = \\mathbf{Y}_j(q, v, \\dot{v})\\boldsymbol{\\theta}\n",
    "  $$\n",
    "\n",
    "While the latter function was not explored within this tutorial, it serves as the primary tool in system identification (sys-id) for articulated multi-body robotic systems (will be exploited in future examples). We are hope that these utilities will be valuable in your system identification projects.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Gb-jExCSKnZR"
   },
   "source": [
    "It's also important to note that the examples provided above are simplified scenarios, and they may not fully reflect the complexity of real-world practical applications. For instance, in these examples, we assume that all signals (including acceleration) are perfectly known and that there are no unmodeled effects. However, in practice, this is not always the case. To enhance your understanding and prepare for real-world challenges, we highly encourage you to modify the above examples by introducing noise and considering partial measurements of some signals.\n",
    "\n",
    "The field of system identification, including its application in robotic systems, is currently experiencing a resurgence and gaining significant attention. This renewed interest can be attributed to the high practical relevance of system identification and its potential to significantly reduce the **sim-to-real gap** while enhancing the fidelity of your simulations, which may be beneficial, for instance, in reinforcement learning.\n",
    "\n",
    "Here are some useful links and possible directions for further exploration:\n",
    "\n",
    "1. [Different regressor models: energy, power, momentum](https://ieeexplore.ieee.org/document/574537) - these offer alternative perspectives on regressor forms for identification.\n",
    "2. [Online parameter estimation and adaptive control in regressor form](https://ieeexplore.ieee.org/abstract/document/8593727) - real-time adaptation.\n",
    "3. [Physical consistency](https://arxiv.org/abs/1701.04395), [Riemannian regularization](https://ieeexplore.ieee.org/document/8922724), and [inertia parametrization](https://ieeexplore.ieee.org/document/9690029) - recent research that exploits the geometric manifold of physically consistent inertial parameters.\n",
    "\n",
    "In the near future, we plan to further investigate these concepts and provide new examples that showcase their application and benefits. We are excited about the potential of system identification in robotics and its ability to bridge the gap between simulation and reality.\n",
    "\n",
    "So, stay tuned for more in-depth tutorials and practical examples!"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}