multibody.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook provides examples to go along with the [textbook](http://underactuated.csail.mit.edu/multibody.html).  I recommend having both windows open, side-by-side!\n",
    "\n",
    "[Click here](http://underactuated.csail.mit.edu/drake.html#notebooks) for instructions on how to run the notebook on Deepnote and/or Google Colab."
   ]
  },
  {
   "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=.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=.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, .1, 0.01), k=(0.001, .1, 0.01))\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=.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=.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(update, v=(-4, 4, 0.1), b=(0, 1, 0.1), e=(0, 1, 0.1), k=(0.001, .1, 0.01))\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
}