lqr.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook provides examples to go along with the [textbook](https://underactuated.csail.mit.edu/lqr.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 mpld3\n",
    "import numpy as np\n",
    "import pydot\n",
    "from IPython.display import SVG, Markdown, clear_output, display\n",
    "from pydrake.all import (AddMultibodyPlantSceneGraph, ConstantVectorSource,\n",
    "                         ControllabilityMatrix, DiagramBuilder,\n",
    "                         DiscreteAlgebraicRiccatiEquation,\n",
    "                         DiscreteTimeLinearQuadraticRegulator, IsControllable,\n",
    "                         Linearize, LinearQuadraticRegulator, LogVectorOutput,\n",
    "                         MatrixGain, MeshcatVisualizer, OutputPortSelection,\n",
    "                         Parser, Simulator, StartMeshcat, ToLatex)\n",
    "from pydrake.common.containers import namedview\n",
    "\n",
    "from underactuated import running_as_notebook\n",
    "\n",
    "if running_as_notebook:\n",
    "    mpld3.enable_notebook()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer (run this cell only once, each instance consumes a port)\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Discrete-time vs continuous-time LQR\n",
    "\n",
    "Let's compare the two solutions in our simplest of systems: the double integrator.\n",
    "\n",
    "## Continuous time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ct_demo():\n",
    "    A = np.array([[0,1],[0,0]])\n",
    "    B = np.array([[0],[1]])\n",
    "    Q = np.identity(2)\n",
    "    R = np.identity(1)\n",
    "\n",
    "    K,S = LinearQuadraticRegulator(A,B,Q,R)\n",
    "    display(Markdown(f\"$S = {ToLatex(S)}$\\n\"))\n",
    "\n",
    "    eigenvalues, eigenvectors = np.linalg.eig(S)\n",
    "    display(Markdown(f\"$eig(S) = {ToLatex(eigenvalues)}$\\n\"))\n",
    "\n",
    "ct_demo()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discrete time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dt_demo(h=1):\n",
    "    A = np.array([[1,h],[0,1]])\n",
    "    B = np.array([[0],[h]])\n",
    "    Q = h*np.identity(2)\n",
    "    R = h*np.identity(1)\n",
    "\n",
    "    K,S = DiscreteTimeLinearQuadraticRegulator(A,B,Q,R)\n",
    "    display(Markdown(f\"$S = {ToLatex(S)}$\\n\"))\n",
    "    eigenvalues, eigenvectors = np.linalg.eig(S)\n",
    "    display(Markdown(f\"$eig(S) = {ToLatex(eigenvalues)}$\\n\"))\n",
    "\n",
    "dt_demo(h=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are a few things to notice here.  First, the discrete-time solution converges to the continuous-time solution as the timestep goes to zero.  Second, the cost-to-go is always higher in the discrete-time version: you should think of the discrete time as adding an additional constraint that the control decision can only be changed once per timestep.  Adding constraints can only increase the total cost.\n",
    "\n",
    "### Side note: Algebraic solution\n",
    "\n",
    "In the DP chapter, I was able to give a nice closed-form solution for the continuous-time case:$$S = \\begin{bmatrix} \\sqrt{3} & 1 \\\\ 1 & \\sqrt{3} \\end{bmatrix}.$$  The discrete-time case is not as clean.  Even when $h=1$, the discrete-time Riccati equation, using $$S = \\begin{bmatrix} a & b \\\\ b & c \\end{bmatrix},$$ results in three equations: \\begin{gather*} b^2 = 1+c \\\\ 1 + c + bc = a + ac \\\\ c^2 = 2b + a + ac. \\end{gather*}  With a little work, you can reduce this to to a quadratic equation in $b$, with a solution, $b = \\frac{1}{4}(1 + \\sqrt{21} + \\sqrt{2(3+\\sqrt{21})}).$  Not so nice! \n",
    "\n",
    "### Side note: Exact integration\n",
    "\n",
    "One can also use the exact integral of the linear system $e^Ah$ in the discretization, instead of the Euler discretization.  It doesn't change the basic observation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LQR via Fitted Value Iteration\n",
    "\n",
    "Double integrator.  Discrete-time, infinite-horizon, discounted.  This is the \"traditional\" fitted value iteration, $J=x^TSx$ as the (linear) function approximator.  Is samples both $x$ and $u$, and takes an argmin over the samples $u$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the double integrator\n",
    "timestep = 0.1\n",
    "A = np.eye(2) + timestep*np.array([[0., 1.], [0., 0.]])\n",
    "B = timestep*np.array([[0.], [1.]])\n",
    "Q = timestep*np.eye(2)\n",
    "R = timestep*np.eye(1)\n",
    "\n",
    "def quadratic_regulator_cost(x, u):\n",
    "    return (x * (Q @ x)).sum(axis=0) + (u * (R @ u)).sum(axis=0)\n",
    "\n",
    "def FittedValueIteration(S_guess, gamma=0.9):\n",
    "\n",
    "    S_optimal = DiscreteAlgebraicRiccatiEquation(A=np.sqrt(gamma) * A,\n",
    "                                                 B=B,\n",
    "                                                 Q=Q,\n",
    "                                                 R=R / gamma)\n",
    "\n",
    "    x1s = np.linspace(-5,5,31)\n",
    "    x2s = np.linspace(-4,4,31)\n",
    "    us = np.linspace(-1,1,9)\n",
    "    Us, X1s, X2s = np.meshgrid(us, x1s, x2s, indexing='ij')\n",
    "    XwithU = np.vstack((X1s.flatten(), X2s.flatten()))\n",
    "    UwithX = Us.flatten().reshape(1,-1)\n",
    "    Nx = x1s.size * x2s.size\n",
    "    X = XwithU[:,:Nx]\n",
    "    N = X1s.size\n",
    "\n",
    "    Xnext = A @ XwithU + B @ UwithX\n",
    "    G = quadratic_regulator_cost(XwithU, UwithX)\n",
    "    Jnext = np.zeros((1,N))\n",
    "    Jd = np.zeros((1,Nx))\n",
    "\n",
    "    def cost_to_go(S, X):\n",
    "        return (X * (S @ X)).sum(axis=0)  # vectorized quadratic form\n",
    "\n",
    "    def mean_bellman_residual(S, X, J_desired):\n",
    "        N = J_desired.size\n",
    "        err = cost_to_go(S, X) - J_desired\n",
    "        loss = np.mean(err**2) # == 1/N ∑ᵢ [tr(Sxᵢxᵢ')-yᵢ]²\n",
    "        # dloss_dS = 2/N ∑ᵢ errᵢ*xᵢxᵢ' = 2/N X*Diag(err)*X'\n",
    "        dloss_dS = 2/N * X @ np.diag(err.reshape(-1,)) @ X.T\n",
    "        return loss, dloss_dS\n",
    "\n",
    "    S = S_guess\n",
    "    eta = 0.0001\n",
    "    last_loss = np.inf\n",
    "    for epoch in range(1000 if running_as_notebook else 2):\n",
    "        Jnext = cost_to_go(S, Xnext)\n",
    "        Jd[:] = np.min((G + gamma*Jnext).reshape(us.size, Nx), axis=0)\n",
    "        for i in range(10):\n",
    "            loss, dloss_dS = mean_bellman_residual(S, X, Jd)\n",
    "            S -= eta*dloss_dS\n",
    "        clear_output(wait=True)\n",
    "        print(f\"epoch {epoch}: loss = {loss}, S = {S.flatten()}\")\n",
    "        if np.linalg.norm(last_loss - loss) < 1e-8:\n",
    "            break\n",
    "        last_loss = loss\n",
    "\n",
    "    print(f\"eigenvalues of S: {np.linalg.eig(S)[0]}\")\n",
    "    print(f\"optimal S= {S_optimal.flatten()}\")\n",
    "\n",
    "\n",
    "FittedValueIteration(S_guess=np.eye(2), gamma=0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the initial guess of $S$ and the discount factor $\\gamma$ that I've used above, the algorithm converges nicely.  You'll see that it is close to the optimal $S$, but not exactly the optimal $S$.\n",
    "\n",
    "Interestingly, if you choose $\\gamma$ to be much closer to 1, the algorithm will diverge. At some point, the sampled values over $u$ cause problems -- the Bellman equation does not actually have a steady-state solution for the discretized $u$ version of the problem (the cost-to-go is infinite).\n",
    "\n",
    "To convince you that this is indeed the problem, here is a version that solves analytically for the optimal $u$ (given the estimated $S$).  It draws samples only over $x$, and converges beautifully to the true optimum, even when $\\gamma=1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def FittedValueIteration(S_guess, gamma=0.9):\n",
    "\n",
    "    S_optimal = DiscreteAlgebraicRiccatiEquation(A=np.sqrt(gamma) * A,\n",
    "                                                 B=B,\n",
    "                                                 Q=Q,\n",
    "                                                 R=R / gamma)\n",
    "\n",
    "    x1s = np.linspace(-5,5,31)\n",
    "    x2s = np.linspace(-4,4,31)\n",
    "    X1s, X2s = np.meshgrid(x1s, x2s, indexing='ij')\n",
    "    X = np.vstack((X1s.flatten(), X2s.flatten()))\n",
    "    N = X1s.size\n",
    "\n",
    "    Jd = np.zeros((1,N))\n",
    "\n",
    "    def cost_to_go(S, X):\n",
    "        return (X * (S @ X)).sum(axis=0)  # vectorized quadratic form\n",
    "\n",
    "    def policy(S, X):\n",
    "        return -gamma*np.linalg.inv(R + gamma*B.T @ S @ B) @ B.T @ S @ A @ X\n",
    "\n",
    "    def mean_bellman_residual(S, X, J_desired):\n",
    "        N = J_desired.size\n",
    "        err = cost_to_go(S, X) - J_desired\n",
    "        loss = np.mean(err**2) # == 1/N ∑ᵢ [tr(Sxᵢxᵢ')-yᵢ]²\n",
    "        # dloss_dS = 2/N ∑ᵢ errᵢ*xᵢxᵢ' = 2/N X*Diag(err)*X'\n",
    "        dloss_dS = 2/N * X @ np.diag(err.reshape(-1,)) @ X.T\n",
    "        return loss, dloss_dS\n",
    "\n",
    "    S = S_guess\n",
    "    eta = 0.0001\n",
    "    last_loss = np.inf\n",
    "    for epoch in range(1000 if running_as_notebook else 2):\n",
    "        U = policy(S, X)\n",
    "        Xnext = A @ X + B @ U\n",
    "        G = quadratic_regulator_cost(X, U)\n",
    "        Jd = G + gamma*cost_to_go(S, Xnext)\n",
    "        for i in range(10):\n",
    "            loss, dloss_dS = mean_bellman_residual(S, X, Jd)\n",
    "            S -= eta*dloss_dS\n",
    "        clear_output(wait=True)\n",
    "        print(f\"epoch {epoch}: loss = {loss}, S = {S.flatten()}\")\n",
    "        if np.linalg.norm(last_loss - loss) < 1e-8:\n",
    "            break\n",
    "        last_loss = loss\n",
    "\n",
    "    print(f\"eigenvalues of S: {np.linalg.eig(S)[0]}\")\n",
    "    print(f\"optimal S= {S_optimal.flatten()}\")\n",
    "    return S\n",
    "\n",
    "\n",
    "FittedValueIteration(S_guess=np.eye(2), gamma=1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LQR on a Manifold\n",
    "\n",
    "I'll also use these examples to illustrate a little more about how one might build a new model in drake (using multibody plant).  I've included the URDF in this notebook directly as a string so that you can see it and easily play with it.\n",
    "\n",
    "## Simulating a wheel\n",
    "\n",
    "*TODO: Consider moving this to the vehicles chapter once it takes more form*\n",
    "\n",
    "As a warm-up, let's just make sure we understand how to simulate a wheel in Drake (using the full `MultibodyPlant` contact engine).  I've done this by putting the cart on a prismatic joint (so it can't fall down) but now adding collision geometry to the wheel and the ground.\n",
    "\n",
    "The short story is that if you simply add a wheel (I use a sphere here) and put it into penetration with the ground, then the contact engine will implement the frictional contact which will cause the wheel to roll.  The one subtlety is that the contact is enforced at a point that is a somewhere inside the overlapping region of the wheel and the ground. (We can query it's exact location using the contact results output port of multibodyplant.)  This leads to an effective reduction in the radius of the wheel.\n",
    "\n",
    "I've put the wheel in this model into relatively deep penetration, just to make the point.\n",
    "\n",
    "If needed, we could have increased the friction coefficient to make sure that we are in the sticking (not sliding) regime.  But the default parameters worked just fine here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cart_with_wheel_urdf = \"\"\"\n",
    "<?xml version=\"1.0\"?>\n",
    "\n",
    "<robot xmlns=\"http://drake.mit.edu\"\n",
    " xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\"\n",
    " name=\"CartWithWheel\">\n",
    "\n",
    "  <link name=\"ground\">\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 -5\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"1000 1000 10\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0.93 .74 .4 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "    <collision>\n",
    "      <origin xyz=\"0 0 -5\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"1000 1000 10\" />\n",
    "      </geometry>\n",
    "    </collision>\n",
    "  </link>\n",
    "\n",
    "  <joint name=\"ground_weld\" type=\"fixed\">\n",
    "    <parent link=\"world\" />\n",
    "    <child link=\"ground\" />\n",
    "  </joint>\n",
    "\n",
    "  <link name=\"cart\">\n",
    "    <inertial>\n",
    "      <mass value=\"10\" />\n",
    "      <inertia ixx=\"0\" ixy=\"0\" ixz=\"0\" iyy=\"0\" iyz=\"0\" izz=\"0\" />\n",
    "    </inertial>\n",
    "    <visual>\n",
    "      <geometry>\n",
    "        <box size=\".6 .3 .3\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0 1 0 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "  </link>\n",
    "\n",
    "  <joint name=\"x\" type=\"prismatic\">\n",
    "    <parent link=\"world\" />\n",
    "    <child link=\"cart\" />\n",
    "    <origin xyz=\"0 0 0.25\" />\n",
    "    <axis xyz=\"1 0 0\" />\n",
    "  </joint>\n",
    "\n",
    "  <link name=\"wheel\">\n",
    "    <inertial>\n",
    "      <mass value=\"1\" />\n",
    "      <inertia ixx=\".02\" ixy=\"0\" ixz=\"0\" iyy=\"0.02\" iyz=\"0\" izz=\"0.02\" />\n",
    "    </inertial>\n",
    "    <visual>\n",
    "      <geometry>\n",
    "        <sphere radius=\".05\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0 0 0 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "    <collision>\n",
    "      <geometry>\n",
    "        <sphere radius=\".05\" />\n",
    "      </geometry>\n",
    "    </collision>\n",
    "  </link>\n",
    "\n",
    "  <joint name=\"wheel\" type=\"continuous\">\n",
    "    <parent link=\"cart\" />\n",
    "    <child link=\"wheel\" />\n",
    "    <!-- I've place the wheel just 0.02 into penetration -->\n",
    "    <origin xyz=\"0 0 -.22\"/>  \n",
    "    <axis xyz=\"0 1 0\" />\n",
    "  </joint>\n",
    "\n",
    "  <transmission type=\"SimpleTransmission\" name=\"wheel\">\n",
    "    <actuator name=\"wheel\" />\n",
    "    <joint name=\"wheel\" />\n",
    "  </transmission>\n",
    "  \n",
    "</robot>\n",
    "\"\"\"\n",
    "\n",
    "from pydrake.all import PidController\n",
    "\n",
    "def CartWithWheelExample():\n",
    "    builder = DiagramBuilder()\n",
    "    time_step = 0.001\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step)\n",
    "    Parser(plant).AddModelsFromString(cart_with_wheel_urdf, \"urdf\")\n",
    "    plant.Finalize()\n",
    "\n",
    "    kd_wheel = 1\n",
    "    pid = builder.AddSystem(\n",
    "        PidController(state_projection=np.array([[0, 1, 0, 0], [0, 0, 0, 1]]),\n",
    "                      kp=[0],\n",
    "                      ki=[0],\n",
    "                      kd=[kd_wheel]))\n",
    "    builder.Connect(pid.get_output_port(), plant.get_actuation_input_port())\n",
    "    builder.Connect(plant.get_state_output_port(),\n",
    "                    pid.get_input_port_estimated_state())\n",
    "    builder.ExportInput(pid.get_input_port_desired_state(), \"vd\")\n",
    "    builder.ExportOutput(plant.get_contact_results_output_port(),\n",
    "                         \"contact_results\")\n",
    "\n",
    "    logger = LogVectorOutput(plant.get_state_output_port(), builder)\n",
    "\n",
    "    # Setup visualization\n",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(xmin=-2.5, xmax=2.5, ymin=-1.0, ymax=2.5)\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    # Set up a simulator to run this diagram\n",
    "    simulator = Simulator(diagram)\n",
    "    context = simulator.get_mutable_context()\n",
    "    diagram.get_input_port().FixValue(context, [0, 10])\n",
    "\n",
    "    # Simulate\n",
    "    visualizer.StartRecording(False)\n",
    "    simulator.AdvanceTo(2 if running_as_notebook else 0.1)\n",
    "    visualizer.PublishRecording()\n",
    "    log = logger.FindLog(context)\n",
    "    fig, ax = plt.subplots(1, 2, figsize=((10, 6)))\n",
    "    ax[0].plot(log.sample_times(), log.data().T)\n",
    "    ax[0].legend(['x', 'wheel', 'xdot', 'wheel_dot'])\n",
    "    r = 0.05\n",
    "    print(f\"Wheel radius is {r}\")\n",
    "    print(f\"x(tf) / wheel(tf) = {log.data()[0,-1]/log.data()[1,-1]}\")\n",
    "    contact_results = diagram.get_output_port().Eval(context)\n",
    "    z_contact = contact_results.point_pair_contact_info(0).contact_point()[2]\n",
    "    print(\n",
    "        f\"The contact point no-slip condition is enforced at z={z_contact}\"\n",
    "    )\n",
    "    r_effective = r - 0.02 - z_contact\n",
    "    print(f\"This leads to an effective wheel radius of {r_effective}\")\n",
    "\n",
    "    F = np.array([[1, -r_effective, 0, 0], [0, 0, 1, -r_effective]])\n",
    "    ax[1].plot(log.sample_times(), (F@log.data()).T)\n",
    "    ax[1].legend(['x - r_effective*wheel', 'xdot - r_effective*wheel_dot'])\n",
    "\n",
    "CartWithWheelExample()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Balancing a 2D Segway (aka \"Ballbot\")\n",
    "\n",
    "Q: Can LQR work even better if I used absolute coordinates for the bot, instead of relative? (The wheel rotates a lot, but the bot angle should not change much)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ballbot_sliding_base_urdf = \"\"\"\n",
    "<?xml version=\"1.0\"?>\n",
    "\n",
    "<robot xmlns=\"http://drake.mit.edu\"\n",
    " xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\"\n",
    " name=\"Ballbot\">\n",
    "\n",
    "  <link name=\"ground\">\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 -5\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"1000 1000 10\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0.93 .74 .4 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "    <collision>\n",
    "      <origin xyz=\"0 0 -5\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"1000 1000 10\" />\n",
    "      </geometry>\n",
    "      <drake:proximity_properties>\n",
    "        <drake:mu_dynamic value=\"100\" />\n",
    "        <drake:mu_static value=\"100\" />\n",
    "      </drake:proximity_properties>\n",
    "    </collision>\n",
    "  </link>\n",
    "\n",
    "  <joint name=\"ground_weld\" type=\"fixed\">\n",
    "    <parent link=\"world\" />\n",
    "    <child link=\"ground\" />\n",
    "  </joint>\n",
    "\n",
    "  <link name=\"ball_x\" />\n",
    "\n",
    "  <link name=\"ball\">\n",
    "    <inertial>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <mass value=\"5\" />\n",
    "      <inertia ixx=\".02\" ixy=\"0\" ixz=\"0\" iyy=\"0.02\" iyz=\"0\" izz=\"0.02\" />\n",
    "    </inertial>\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <sphere radius=\".1\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0.25 0.52 0.96 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "    <visual> \n",
    "      <!-- add a visual cue to see when the ball is rolling -->\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"0.04 .201 0.002\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\".1 .1 .1 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "    <collision>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <sphere radius=\".1\" />\n",
    "      </geometry>\n",
    "      <drake:proximity_properties>\n",
    "        <drake:mu_dynamic value=\"100\" />\n",
    "        <drake:mu_static value=\"100\" />\n",
    "      </drake:proximity_properties>\n",
    "    </collision>\n",
    "  </link>\n",
    "\n",
    "  <joint name=\"x\" type=\"prismatic\">\n",
    "    <parent link=\"world\" />\n",
    "    <child link=\"ball_x\" />\n",
    "    <!-- height is set to (radius - 0.001) to have sufficient penetration to have a normal force that supports the required frictional force -->\n",
    "    <origin xyz=\"0 0 .099\" />\n",
    "    <axis xyz=\"1 0 0\" />\n",
    "  </joint>\n",
    "\n",
    "  <link name=\"bot\">\n",
    "    <inertial>\n",
    "      <origin xyz=\"0 0 .05\" rpy=\"0 0 0\" />\n",
    "      <mass value=\"4\" />\n",
    "      <inertia ixx=\"0.018\" ixy=\"0\" ixz=\"0\" iyy=\"0.018\" iyz=\"0\" izz=\"0.0288\" />\n",
    "    </inertial>\n",
    "    <!-- no collision geometry since Drake AutoDiffXd doesn't support cylinder on box collisions yet. -->\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 .05\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "         <cylinder length=\".1\" radius=\".12\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\".61 .63 .67 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "  </link>  \n",
    "\n",
    "  <joint name=\"theta_ball\" type=\"continuous\">\n",
    "    <parent link=\"ball_x\" />\n",
    "    <child link=\"ball\" />\n",
    "    <axis xyz=\"0 -1 0\" />   \n",
    "    <!-- use -1 to be consistent with the planar joint having z up below -->\n",
    "  </joint>\n",
    "\n",
    "  <joint name=\"theta_bot\" type=\"continuous\">\n",
    "    <parent link=\"ball\" />\n",
    "    <child link=\"bot\" />\n",
    "    <axis xyz=\"0 -1 0\" />\n",
    "  </joint>\n",
    "\n",
    "  <transmission type=\"SimpleTransmission\" name=\"ball_torque\">\n",
    "    <actuator name=\"torque\" />\n",
    "    <joint name=\"theta_bot\" />\n",
    "  </transmission>\n",
    "  \n",
    "</robot>\n",
    "\"\"\"\n",
    "\n",
    "r = 0.0995  # Effective r given the contact point (see the example above)\n",
    "n = np.sqrt(1 + r**2)\n",
    "\n",
    "F = np.array([[1, r, 0, 0, 0, 0], [0, 0, 0, 1, r, 0]])\n",
    "P = np.array([[r / n, -1 / n, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0],\n",
    "                [0, 0, 0, r / n, -1 / n, 0], [0, 0, 0, 0, 0, 1]])\n",
    "\n",
    "np.testing.assert_almost_equal(P @ P.T, np.eye(4, 4))\n",
    "np.testing.assert_almost_equal(P @ F.T, np.zeros((4, 2)))\n",
    "\n",
    "def BallbotUprightState():\n",
    "    state = (0, 0, 0, 0, 0, 0)\n",
    "    return state\n",
    "\n",
    "def MakeSlidingBaseBallbot(time_step=0.001):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step)\n",
    "    Parser(plant).AddModelsFromString(ballbot_sliding_base_urdf, \"urdf\")\n",
    "    plant.Finalize()\n",
    "\n",
    "    builder.ExportInput(plant.get_actuation_input_port(), \"torque\")\n",
    "    builder.ExportOutput(plant.get_state_output_port(), \"state\")\n",
    "    builder.ExportOutput(scene_graph.get_query_output_port(), \"query\")\n",
    "\n",
    "    return builder.Build()\n",
    "\n",
    "\n",
    "def BallbotControllability(plant):\n",
    "    context = plant.CreateDefaultContext()\n",
    "    plant.get_input_port().FixValue(context, [0])\n",
    "\n",
    "    context.SetDiscreteState(BallbotUprightState())\n",
    "\n",
    "    linear = Linearize(\n",
    "        plant,\n",
    "        context,\n",
    "        output_port_index=OutputPortSelection.kNoOutput)\n",
    "    #display(Markdown(\"$A = \" + ToLatex(linear.A()) + \"$\"))\n",
    "    #display(Markdown(\"$B = \" + ToLatex(linear.B()) + \"$\"))\n",
    "    #display(\n",
    "    #    Markdown(\"$ctrb(A,B) = \" + ToLatex(ControllabilityMatrix(linear), 6)\n",
    "    #             + \"$\"))\n",
    "    display(Markdown(f\"Is controllable? {IsControllable(linear)}\"))\n",
    "\n",
    "def BallbotManifoldLQR(plant):\n",
    "    context = plant.CreateDefaultContext()\n",
    "    plant.get_input_port().FixValue(context, [0])\n",
    "\n",
    "    Q = np.diag((10, 10, 1, 1))\n",
    "    R = np.array([1])\n",
    "\n",
    "    context.SetDiscreteState(BallbotUprightState())\n",
    "\n",
    "    linear = Linearize(\n",
    "        plant,\n",
    "        context,\n",
    "        output_port_index=OutputPortSelection.kNoOutput)\n",
    "\n",
    "    A = P @ linear.A() @ P.T\n",
    "    B = P @ linear.B()\n",
    "\n",
    "    K, S = DiscreteTimeLinearQuadraticRegulator(A, B, Q, R)\n",
    "\n",
    "    return MatrixGain(-K @ P)\n",
    "\n",
    "def BallbotSlidingBaseExample():\n",
    "    builder = DiagramBuilder()\n",
    "    time_step = 0.001\n",
    "\n",
    "    plant = builder.AddSystem(MakeSlidingBaseBallbot(time_step))\n",
    "    logger = LogVectorOutput(plant.GetOutputPort(\"state\"), builder)\n",
    "\n",
    "    # For completeness, we can check that the plant by itself is not controllable:\n",
    "    BallbotControllability(plant)\n",
    "\n",
    "    if False:  # Useful for debugging.\n",
    "        kd_wheel = 1\n",
    "        pid = builder.AddSystem(\n",
    "            PidController(state_projection=np.array([[0, 1, 0, 0, 0, 0],\n",
    "                                                    [0, 0, 0, 0, 1, 0]]),\n",
    "                        kp=[0],\n",
    "                        ki=[0],\n",
    "                        kd=[kd_wheel]))\n",
    "        builder.Connect(pid.get_output_port(), plant.get_input_port())\n",
    "        builder.Connect(plant.GetOutputPort(\"state\"),\n",
    "                        pid.get_input_port_estimated_state())\n",
    "        pid_command = builder.AddSystem(ConstantVectorSource([0, 1]))\n",
    "        builder.Connect(pid_command.get_output_port(),\n",
    "                        pid.get_input_port_desired_state())\n",
    "    else:\n",
    "        controller = builder.AddSystem(BallbotManifoldLQR(plant))\n",
    "        controller.set_name(\"LQR Controller\")\n",
    "        builder.Connect(plant.GetOutputPort(\"state\"),\n",
    "                        controller.get_input_port())\n",
    "        builder.Connect(controller.get_output_port(),\n",
    "                        plant.get_input_port())\n",
    "\n",
    "    # Setup visualization\n",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(xmin=-.5, xmax=.5, ymin=-.2, ymax=.5)\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(builder,\n",
    "                                                   plant.GetOutputPort(\"query\"),\n",
    "                                                   meshcat)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    # For reference, let's draw the diagram we've assembled:\n",
    "    #display(SVG(pydot.graph_from_dot_data(diagram.GetGraphvizString())[0].#create_svg()))\n",
    "\n",
    "    # Set up a simulator to run this diagram\n",
    "    simulator = Simulator(diagram)\n",
    "    #simulator.set_target_realtime_rate(1.0 if running_as_notebook else 0.0)\n",
    "    context = simulator.get_mutable_context()\n",
    "    x0 = BallbotUprightState() + P.T @ (0.0 * np.random.randn(4,)\n",
    "                                        + np.array([1, 0, 0, 0]))\n",
    "    assert np.allclose(F @ x0, np.zeros(2))\n",
    "    context.SetDiscreteState(x0)\n",
    "\n",
    "    # Simulate\n",
    "\n",
    "    visualizer.StartRecording(False)\n",
    "    simulator.AdvanceTo(4 if running_as_notebook else 0.1)\n",
    "    visualizer.PublishRecording()\n",
    "    log = logger.FindLog(context)\n",
    "    fig, ax = plt.subplots(1, 2, figsize=((10, 6)))\n",
    "    ax[0].plot(log.sample_times(), log.data()[:3,:].T)\n",
    "    ax[0].legend(['x', 'theta_ball', 'theta_bot'])\n",
    "    ax[1].plot(log.sample_times(), (F@log.data()).T)\n",
    "    ax[1].legend(['x - r*theta_ball', 'xdot - r*thetadot_ball'])\n",
    "\n",
    "BallbotSlidingBaseExample()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But there is subtle point to be made here.  I've modeled the wheels as a prismatic joint at the base.  Let's instead model the robot with a floating joint, and using collisions elements to simulate the interaction with the robot and the ground..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ballbot_floating_base_urdf = \"\"\"\n",
    "<?xml version=\"1.0\"?>\n",
    "\n",
    "<robot xmlns=\"http://drake.mit.edu\"\n",
    " xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\"\n",
    " name=\"Ballbot\">\n",
    "\n",
    "  <link name=\"ground\">\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 -5\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"1000 1000 10\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0.93 .74 .4 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "    <collision>\n",
    "      <origin xyz=\"0 0 -5\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"1000 1000 10\" />\n",
    "      </geometry>\n",
    "    </collision>\n",
    "  </link>\n",
    "\n",
    "  <joint name=\"ground_weld\" type=\"fixed\">\n",
    "    <parent link=\"world\" />\n",
    "    <child link=\"ground\" />\n",
    "  </joint>\n",
    "\n",
    "  <link name=\"ball\">\n",
    "    <inertial>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <mass value=\"5\" />\n",
    "      <inertia ixx=\".02\" ixy=\"0\" ixz=\"0\" iyy=\"0.02\" iyz=\"0\" izz=\"0.02\" />\n",
    "    </inertial>\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <sphere radius=\".1\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\"0.25 0.52 0.96 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "    <visual> \n",
    "      <!-- add a visual cue to see when the ball is rolling -->\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <box size=\"0.04  0.002 .201\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\".1 .1 .1 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "    <collision>\n",
    "      <origin xyz=\"0 0 0\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "        <sphere radius=\".1\" />\n",
    "      </geometry>\n",
    "    </collision>\n",
    "  </link>\n",
    "\n",
    "  <link name=\"bot\">\n",
    "    <inertial>\n",
    "      <origin xyz=\"0 0 .05\" rpy=\"0 0 0\" />\n",
    "      <mass value=\"4\" />\n",
    "      <inertia ixx=\"0.018\" ixy=\"0\" ixz=\"0\" iyy=\"0.018\" iyz=\"0\" izz=\"0.0288\" />\n",
    "    </inertial>\n",
    "    <collision>\n",
    "      <origin xyz=\"0 0 .05\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "         <cylinder length=\".1\" radius=\".12\" />\n",
    "      </geometry>\n",
    "    </collision>\n",
    "    <visual>\n",
    "      <origin xyz=\"0 0 .05\" rpy=\"0 0 0\" />\n",
    "      <geometry>\n",
    "         <cylinder length=\".1\" radius=\".12\" />\n",
    "      </geometry>\n",
    "      <material>\n",
    "        <color rgba=\".61 .63 .67 1\" />\n",
    "      </material>\n",
    "    </visual>\n",
    "  </link>\n",
    "  \n",
    "  <drake:joint name=\"floating_base\" type=\"planar\">\n",
    "    <parent link=\"world\" />\n",
    "    <child link=\"ball\" />\n",
    "    <origin rpy=\"1.57 0 0\" xyz=\"0 0 .1\" />\n",
    "  </drake:joint>\n",
    "\n",
    "  <joint name=\"theta2\" type=\"continuous\">\n",
    "    <parent link=\"ball\" />\n",
    "    <child link=\"bot\" />\n",
    "    <origin rpy=\"-1.57 0 0\" xyz=\"0 0 0\" />\n",
    "    <axis xyz=\"0 -1 0\" />\n",
    "  </joint>\n",
    "\n",
    "  <transmission type=\"SimpleTransmission\" name=\"ball_torque\">\n",
    "    <actuator name=\"torque\" />\n",
    "    <joint name=\"theta2\" />\n",
    "    <mechanicalReduction>1</mechanicalReduction>\n",
    "  </transmission>\n",
    "\n",
    "</robot>\n",
    "\"\"\"\n",
    "\n",
    "def BallbotFloatingBaseExample():\n",
    "    builder = DiagramBuilder()\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    Parser(plant).AddModelsFromString(ballbot_floating_base_urdf, \"urdf\")\n",
    "    plant.Finalize()\n",
    "\n",
    "    # Just use zero instead of a controller to start\n",
    "    command = builder.AddSystem(ConstantVectorSource([0.0]))\n",
    "    builder.Connect(command.get_output_port(), plant.get_actuation_input_port())\n",
    "\n",
    "    # Setup visualization\n",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(xmin=-.5, xmax=.5, ymin=-.2, ymax=.5)\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    # Set up a simulator to run this diagram\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.set_target_realtime_rate(1.0 if running_as_notebook else 0.0)\n",
    "    context = simulator.get_mutable_context()\n",
    "\n",
    "    State = namedview(\"State\", [\n",
    "        \"x\", \"z\", \"theta1\", \"theta2\", \"xdot\", \"zdot\", \"theta1dot\", \"theta2dot\"\n",
    "    ])\n",
    "    x0 = State(np.zeros(8))\n",
    "    x0.z = .15\n",
    "    x0.theta1 = 0.05\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "    plant.SetPositionsAndVelocities(plant_context, x0[:])\n",
    "\n",
    "    # Simulate\n",
    "    visualizer.StartRecording(False)\n",
    "    simulator.AdvanceTo(3 if running_as_notebook else 0.1)\n",
    "    visualizer.PublishRecording()\n",
    "    xf = State(plant.GetPositionsAndVelocities(plant_context))\n",
    "    print(f\"z at final time = {xf.z}\")\n",
    "\n",
    "BallbotFloatingBaseExample()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the real test... you should be able to run the controller designed for the simple model on the floating base version of the ballbot.  Give it a shot!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def BallbotFloatingBaseLqrExample():\n",
    "    builder = DiagramBuilder()\n",
    "    time_step = 0.001\n",
    "    ballbot_minimal_coordinates = builder.AddSystem(\n",
    "        MakeSlidingBaseBallbot(time_step))\n",
    "\n",
    "    controller_minimal_coordinates = builder.AddSystem(\n",
    "        BallbotManifoldLQR(ballbot_minimal_coordinates))\n",
    "    controller_minimal_coordinates.set_name(\"LQR Controller for Minimal\")\n",
    "    builder.Connect(controller_minimal_coordinates.get_output_port(),\n",
    "                    ballbot_minimal_coordinates.get_input_port())\n",
    "    builder.Connect(\n",
    "        ballbot_minimal_coordinates.GetOutputPort(\"state\"),\n",
    "        controller_minimal_coordinates.get_input_port())\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step)\n",
    "    Parser(plant).AddModelsFromString(ballbot_floating_base_urdf, \"urdf\")\n",
    "    plant.Finalize()\n",
    "\n",
    "    controller = builder.AddSystem(\n",
    "        BallbotManifoldLQR(ballbot_minimal_coordinates))\n",
    "    controller.set_name(\"LQR Controller for Floating\")\n",
    "    builder.Connect(controller.get_output_port(),\n",
    "                    plant.get_actuation_input_port())\n",
    "    #builder.ExportInput(plant.get_actuation_input_port(), \"u\")\n",
    "\n",
    "    MinimalState = namedview(\"MinimalState\", [\"x\", \"theta_ball\", \"theta_bot_m_ball\", \"xdot\", \"thetadot_bot\", \"thetadot_bot_m_ball\"])\n",
    "    FloatingState = namedview(\"FloatingState\", [\n",
    "        \"x\", \"z\", \"theta_ball\", \"theta_bot_m_ball\", \"xdot\", \"zdot\", \"thetadot_ball\", \"thetadot_bot_m_ball\"\n",
    "    ])\n",
    "\n",
    "    m_ind = MinimalState(range(6))\n",
    "    f_ind = FloatingState(range(8))\n",
    "    P_floating_to_minimal = np.delete(np.eye(8), [f_ind.z, f_ind.zdot], 0)\n",
    "    floating_state_to_minimal_state = builder.AddSystem(\n",
    "        MatrixGain(P_floating_to_minimal))\n",
    "    builder.Connect(plant.get_state_output_port(),\n",
    "                    floating_state_to_minimal_state.get_input_port())\n",
    "    builder.Connect(floating_state_to_minimal_state.get_output_port(),\n",
    "                    controller.get_input_port())\n",
    "\n",
    "    # Setup visualization\n",
    "    meshcat.Delete()\n",
    "    meshcat.Set2dRenderMode(xmin=-.5, xmax=.5, ymin=-.2, ymax=.5)\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph,\n",
    "                                                   meshcat)\n",
    "\n",
    "    minimal_state_logger = LogVectorOutput(\n",
    "        ballbot_minimal_coordinates.GetOutputPort(\"state\"), builder)\n",
    "    floating_state_logger = LogVectorOutput(plant.get_state_output_port(),\n",
    "                                            builder)\n",
    "    minimal_control_logger = LogVectorOutput(\n",
    "        controller_minimal_coordinates.get_output_port(), builder)\n",
    "    floating_control_logger = LogVectorOutput(controller.get_output_port(),\n",
    "                                              builder)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    # Set up a simulator to run this diagram\n",
    "    simulator = Simulator(diagram)\n",
    "    #simulator.set_target_realtime_rate(1.0 if running_as_notebook else 0.0)\n",
    "    context = simulator.get_mutable_context()\n",
    "    #diagram.GetInputPort(\"u\").FixValue(context, [3])\n",
    "\n",
    "    x0 = FloatingState(np.zeros(8))\n",
    "    x0.z = -0.0018  # the approximate resting depth due to soft contact\n",
    "    x0.theta_bot_m_ball = 0.05\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "    plant.SetPositionsAndVelocities(plant_context, x0[:])\n",
    "\n",
    "    ballbot_minimal_state_context = ballbot_minimal_coordinates.GetMyMutableContextFromRoot(\n",
    "        context)\n",
    "    minimal_x0 = MinimalState(P_floating_to_minimal @ x0[:])\n",
    "    ballbot_minimal_state_context.SetDiscreteState(minimal_x0[:])\n",
    "\n",
    "    values = plant.AllocateDiscreteVariables()\n",
    "    plant.CalcForcedDiscreteVariableUpdate(plant_context, values)\n",
    "    print(f\"floating_xn = {FloatingState(values.get_value())}\")\n",
    "    print(\n",
    "        f\"projected_xn = {MinimalState( P_floating_to_minimal @ values.get_value())}\"\n",
    "    )\n",
    "\n",
    "    values = ballbot_minimal_coordinates.AllocateDiscreteVariables()\n",
    "    ballbot_minimal_coordinates.CalcForcedDiscreteVariableUpdate(\n",
    "        ballbot_minimal_state_context, values)\n",
    "    print(f\"minimal_xn = {MinimalState(values.get_value())}\")\n",
    "\n",
    "    # Simulate\n",
    "    visualizer.StartRecording(False)\n",
    "    simulator.AdvanceTo(5 if running_as_notebook else 0.1)\n",
    "    visualizer.PublishRecording()\n",
    "\n",
    "    minimal_state_log = minimal_state_logger.FindLog(context)\n",
    "    floating_state_log = floating_state_logger.FindLog(context)\n",
    "    fig, ax = plt.subplots(1,2,figsize=(15,8))\n",
    "    ax[0].plot(minimal_state_log.sample_times(), minimal_state_log.data().T)\n",
    "    ax[0].set_title('minimal state')\n",
    "    ax[0].legend([\n",
    "        \"x\", \"theta_ball\", \"theta_bot_m_ball\", \"xdot\", \"thetadot_ball\",\n",
    "        \"thetadot_bot_m_ball\"\n",
    "    ])\n",
    "    if True:  # plot projection of floating state\n",
    "        ax[1].plot(floating_state_log.sample_times(),\n",
    "            (P_floating_to_minimal @ floating_state_log.data()).T)\n",
    "    else:\n",
    "        ax[1].plot(floating_state_log.sample_times(),\n",
    "                floating_state_log.data().T)\n",
    "        ax[1].legend([ \"x\", \"z\", \"theta1\", \"theta2\", \"xdot\", \"zdot\", \"theta1dot\", \"theta2dot\"])\n",
    "    ax[1].set_title('floating state')\n",
    "    ax[0].set_ylim(ax[1].get_ylim())\n",
    "\n",
    "BallbotFloatingBaseLqrExample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.8 64-bit",
   "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.10.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}