book/multibody/multibody.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook provides examples to go along with the [textbook](https://underactuated.csail.mit.edu/multibody.html).  I recommend having both windows open, side-by-side!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "# Note: mpld3 doesn't work on this notebook, to be fixed in matplotlib 3.4.2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Understanding the \"relaxed complementarity-free\" approximation\n",
    "\n",
    "Here is a simple interactive plot to understand the difference between the time-stepping LCP solution and the algorithm in MuJoCo.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lcp_vs_relaxed():\n",
    "    q = np.linspace(-1, 4, 51)\n",
    "    v = np.linspace(-4, 4, 51)\n",
    "    f = np.zeros((len(q), len(v)))\n",
    "    qn = np.zeros((len(q), len(v)))\n",
    "\n",
    "    [Q, V] = np.meshgrid(q, v)\n",
    "    F = 0 * Q\n",
    "    Qn = 0 * Q\n",
    "\n",
    "    h = 0.1\n",
    "    m = 1\n",
    "\n",
    "    # First plot the LCP solution:\n",
    "    for i in range(len(q)):\n",
    "        for j in range(len(v)):\n",
    "            if q[i] + h * v[j] >= 0:\n",
    "                Qn[j, i] = q[i] + h * v[j]\n",
    "                F[j, i] = 0\n",
    "            else:\n",
    "                Qn[j, i] = 0\n",
    "                F[j, i] = -m * (q[i] + h * v[j]) / h**2\n",
    "\n",
    "    fig = plt.figure()\n",
    "    ax1, ax2 = fig.subplots(1, 2, subplot_kw=dict(projection=\"3d\"))\n",
    "    ax1.plot_surface(Q, V, Qn, rstride=1, cstride=1, alpha=0.3)\n",
    "    ax2.plot_surface(Q, V, F, rstride=1, cstride=1, alpha=0.3)\n",
    "    ax1.set_xlabel(\"q[n]\")\n",
    "    ax1.set_ylabel(\"v[n]\")\n",
    "    ax1.set_title(\"q[n+1]\")\n",
    "    ax2.set_xlabel(\"q[n]\")\n",
    "    ax2.set_ylabel(\"v[n]\")\n",
    "    ax2.set_title(\"f[n]\")\n",
    "\n",
    "    # Now the \"relaxed complementarity-free\" version\n",
    "    def complementarityFree(b=0.1, e=1e-4, k=0.1):\n",
    "        for i in range(len(q)):\n",
    "            for j in range(len(v)):\n",
    "                vminus = v[j] - h * b * v[j] / (m + b)\n",
    "                vstar = v[j] - 2 * h * (1 + e) * v[j] / k - h * (1 + e) * q[i] / k**2\n",
    "                f = -(vminus - vstar) * (m + b) / (1 + e)\n",
    "                F[j, i] = np.maximum(f, 0)\n",
    "                vn = v[j] + h * (-b * v[j] + F[j, i]) / (m + b)\n",
    "                Qn[j, i] = q[i] + h * vn\n",
    "\n",
    "    complementarityFree()\n",
    "    Qn_surf = ax1.plot_surface(Q, V, Qn, rstride=1, cstride=1)\n",
    "    F_surf = ax2.plot_surface(Q, V, F, rstride=1, cstride=1)\n",
    "\n",
    "    def update(b=0.1, e=1e-3, k=0.1):\n",
    "        global Qn_surf\n",
    "        global F_surf\n",
    "        Qn_surf.remove()\n",
    "        F_surf.remove()\n",
    "        complementarityFree(b, e, k)\n",
    "        Qn_surf = ax1.plot_surface(Q, V, Qn, rstride=1, cstride=1)\n",
    "        F_surf = ax2.plot_surface(Q, V, F, rstride=1, cstride=1)\n",
    "        fig.canvas.draw()\n",
    "\n",
    "    # TODO(russt): re-create the interaction\n",
    "    if False:\n",
    "        interact(update, b=(0, 1, 0.1), e=(0, 0.1, 0.01), k=(0.001, 0.1, 0.01))\n",
    "\n",
    "\n",
    "lcp_vs_relaxed()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lcp_vs_relaxed2():\n",
    "    q = np.linspace(-1, 4, 51)\n",
    "\n",
    "    h = 0.1\n",
    "    m = 1\n",
    "\n",
    "    def lcp(q, v):\n",
    "        if q + h * v >= 0:\n",
    "            qn = q + h * v\n",
    "            f = 0\n",
    "        else:\n",
    "            qn = 0\n",
    "            f = -m * (q + h * v) / h**2\n",
    "        return qn, f\n",
    "\n",
    "    def Lcp(q, v):\n",
    "        qn = 0 * q\n",
    "        f = 0 * q\n",
    "        for i in range(len(q)):\n",
    "            qn[i], f[i] = lcp(q[i], v)\n",
    "        return qn, f\n",
    "\n",
    "    def complementarityFree(q, v, b, e, k):\n",
    "        vminus = v - h * b * v / (m + b)\n",
    "        vstar = v - 2 * h * (1 + e) * v / k - h * (1 + e) * q / k**2\n",
    "        f = -(vminus - vstar) * (m + b) / (1 + e)\n",
    "        f = np.maximum(f, 0)\n",
    "        vn = v + h * (-b * v + f) / (m + b)\n",
    "        qn = q + h * vn\n",
    "        return qn, f\n",
    "\n",
    "    def ComplementarityFree(q, v, b=0.1, e=1e-4, k=0.1):\n",
    "        qn = 0 * q\n",
    "        f = 0 * q\n",
    "        for i in range(len(q)):\n",
    "            qn[i], f[i] = complementarityFree(q[i], v, b, e, k)\n",
    "        return qn, f\n",
    "\n",
    "    fig = plt.figure()\n",
    "    ax1, ax2 = fig.subplots(1, 2)\n",
    "    qn, f = Lcp(q, 0)\n",
    "    (lcp_qn,) = ax1.plot(q, qn, label=\"lcp\")\n",
    "    (lcp_f,) = ax2.plot(q, f, label=\"lcp\")\n",
    "    ax1.set_xlabel(\"q[n]\")\n",
    "    ax1.set_ylabel(\"q[n+1]\")\n",
    "    ax2.set_xlabel(\"q[n]\")\n",
    "    ax2.set_ylabel(\"f[n]\")\n",
    "\n",
    "    qn, f = ComplementarityFree(q, 0)\n",
    "    (cfree_qn,) = ax1.plot(q, qn, label=\"cfree\")\n",
    "    (cfree_f,) = ax2.plot(q, f, label=\"cfree\")\n",
    "    ax1.legend()\n",
    "\n",
    "    def update(v=0, b=0.1, e=1e-3, k=0.1):\n",
    "        qn, f = Lcp(q, v)\n",
    "        lcp_qn.set_ydata(qn)\n",
    "        lcp_f.set_ydata(f)\n",
    "        qn, f = ComplementarityFree(q, v, b, e, k)\n",
    "        cfree_qn.set_ydata(qn)\n",
    "        cfree_f.set_ydata(f)\n",
    "        fig.canvas.draw()\n",
    "\n",
    "    if False:\n",
    "        interact(\n",
    "            update,\n",
    "            v=(-4, 4, 0.1),\n",
    "            b=(0, 1, 0.1),\n",
    "            e=(0, 1, 0.1),\n",
    "            k=(0.001, 0.1, 0.01),\n",
    "        )\n",
    "\n",
    "\n",
    "lcp_vs_relaxed2()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}