Controlling a Double Integrator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "Dependencies:\n",
    "- System: python3, ffmpeg (for rendering animations)\n",
    "- Python: jupyter, numpy, matplotlib, cvxpy, ipywidgets (for interactive plots)\n",
    "\n",
    "Example setup for a Ubuntu system (Mac users, maybe `brew` instead of `sudo apt`; Windows users, learn to love [WSL](https://docs.microsoft.com/en-us/windows/wsl/install-win10)):\n",
    "```\n",
    "sudo apt install ffmpeg\n",
    "/usr/bin/python3 -m pip install --upgrade pip\n",
    "pip install --upgrade jupyter numpy matplotlib cvxpy ipywidgets\n",
    "jupyter nbextension enable --py widgetsnbextension  # necessary for interactive sliders to show up\n",
    "jupyter notebook  # from the directory of this notebook\n",
    "```\n",
    "Alternatively, view this notebook on [Google Colab](https://colab.research.google.com/github/StanfordASL/AA203-Examples/blob/master/Lecture-1/Controlling%20a%20Double%20Integrator.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.animation as animation\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "from ipywidgets import interact\n",
    "from IPython.display import HTML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define dynamics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiscreteTimeDynamics:\n",
    "\n",
    "    def __init__(self, state_dim, control_dim, time_step):\n",
    "        self.state_dim = state_dim\n",
    "        self.control_dim = control_dim\n",
    "        self.time_step = time_step\n",
    "\n",
    "    def plot_trajectory(self, x, u, figsize=(10, 10), fig=None, gs=None, return_ax=False):\n",
    "        if fig is None:\n",
    "            fig = plt.figure(figsize=figsize)\n",
    "        if gs is None:\n",
    "            gs = gridspec.GridSpec(2, 1)\n",
    "\n",
    "        state_ax = fig.add_subplot(gs[0, 0])\n",
    "        state_ax.set_ylabel('States', fontsize=24)\n",
    "        state_ax.tick_params(axis='both', labelsize=24)\n",
    "        state_ax.plot(self.time_step * np.arange(len(x)), x)\n",
    "        state_ax.legend(self.state_names, fontsize=24)\n",
    "\n",
    "        control_ax = fig.add_subplot(gs[1, 0])\n",
    "        control_ax.set_xlabel('Time', fontsize=24)\n",
    "        control_ax.set_ylabel('Controls', fontsize=24)\n",
    "        control_ax.tick_params(axis='both', labelsize=24)\n",
    "        control_ax.set_prop_cycle(color=plt.rcParams['axes.prop_cycle'].by_key()['color'][self.state_dim:])\n",
    "        control_ax.plot(self.time_step * np.arange(len(u)), u)\n",
    "        control_ax.legend(self.control_names, fontsize=24)\n",
    "\n",
    "        return (fig, state_ax, control_ax) if return_ax else fig\n",
    "\n",
    "    def animate_trajectory(self, x, u, figsize=(15, 8)):\n",
    "        fig = plt.figure(figsize=figsize)\n",
    "        gs = gridspec.GridSpec(2, 2)\n",
    "\n",
    "        _, state_ax, control_ax = self.plot_trajectory(x, u, fig=fig, gs=gs, return_ax=True)\n",
    "\n",
    "        state_dots = [state_ax.scatter([0], x[0, i]) for i in range(self.state_dim)]\n",
    "        control_dots = [control_ax.scatter([0], u[0, i]) for i in range(self.control_dim)]\n",
    "        visualization_ax = fig.add_subplot(gs[:, 1])\n",
    "        visualization_ax.tick_params(axis='both', labelsize=24)\n",
    "\n",
    "        updateables = state_dots + control_dots + self._visualize(visualization_ax, x, u)\n",
    "\n",
    "        def animate(t):\n",
    "            for i, s in enumerate(state_dots):\n",
    "                s.set_offsets([[self.time_step * t, x[t, i]]])\n",
    "            for i, s in enumerate(control_dots):\n",
    "                s.set_offsets([[self.time_step * t, u[t, i]]])\n",
    "            self._update_visualization(x[t], u[t])\n",
    "            return updateables\n",
    "\n",
    "        plt.close()\n",
    "        return HTML(\n",
    "            animation.FuncAnimation(fig, animate, interval=1000 * self.time_step, blit=True,\n",
    "                                    frames=len(u)).to_html5_video())\n",
    "\n",
    "    @property\n",
    "    def state_names(self):\n",
    "        return range(self.state_dim)\n",
    "\n",
    "    @property\n",
    "    def control_names(self):\n",
    "        return range(self.control_dim)\n",
    "\n",
    "    def _visualize(self, ax, x, u):\n",
    "        ax.text(0.2, 0.5, \"visualization not implemented\", fontsize=16)\n",
    "        return []\n",
    "\n",
    "    def _update_visualization(self, x, u):\n",
    "        pass\n",
    "\n",
    "\n",
    "class DoubleIntegatorDynamics(DiscreteTimeDynamics):\n",
    "\n",
    "    def __init__(self, spatial_dim, time_step=0.1):\n",
    "        super().__init__(2 * spatial_dim, spatial_dim, time_step)\n",
    "\n",
    "        self.spatial_dim = spatial_dim\n",
    "        self.A = np.eye(2 * spatial_dim) + time_step * np.eye(2 * spatial_dim, k=spatial_dim)\n",
    "        self.B = (time_step**2 / 2 * np.eye(2 * spatial_dim, spatial_dim) +\n",
    "                  time_step * np.eye(2 * spatial_dim, spatial_dim, k=-spatial_dim))\n",
    "\n",
    "    def __call__(self, x, u):\n",
    "        return self.A @ x + self.B @ u\n",
    "\n",
    "    @property\n",
    "    def state_names(self):\n",
    "        if self.spatial_dim == 1:\n",
    "            return ['s', 'v']\n",
    "        return [f's_{i}' for i in range(self.spatial_dim)] + [f'v_{i}' for i in range(self.spatial_dim)]\n",
    "\n",
    "    @property\n",
    "    def control_names(self):\n",
    "        if self.spatial_dim == 1:\n",
    "            return ['a']\n",
    "        return [f'a_{i}' for i in range(self.spatial_dim)]\n",
    "\n",
    "    def _visualize(self, ax, x, u):\n",
    "        plot_range = np.concatenate(\n",
    "            [x[:-1, :self.spatial_dim] + x[:-1, self.spatial_dim:], x[:-1, :self.spatial_dim] + u], 0)\n",
    "        plot_min = np.min(plot_range, 0)\n",
    "        plot_max = np.max(plot_range, 0)\n",
    "        plot_margin = (plot_max - plot_min) * 0.1\n",
    "        plot_min = plot_min - plot_margin\n",
    "        plot_max = plot_max + plot_margin\n",
    "        ax.set_xlim(plot_min[0], plot_max[0])\n",
    "        if self.spatial_dim == 1:\n",
    "            ax.set_ylim(-1, 1)\n",
    "        else:\n",
    "            ax.set_ylim(plot_min[1], plot_max[1])\n",
    "        p, v, a = x[0, :self.spatial_dim], x[0, self.spatial_dim:], u[0]\n",
    "        if self.spatial_dim == 1:\n",
    "            p, v, a = np.pad(p, [0, 1]), np.pad(v, [0, 1]), np.pad(a, [0, 1])\n",
    "        self.a_vector = plt.quiver(p[0],\n",
    "                                   p[1],\n",
    "                                   a[0],\n",
    "                                   a[1],\n",
    "                                   angles='xy',\n",
    "                                   scale_units='xy',\n",
    "                                   scale=1,\n",
    "                                   width=0.02,\n",
    "                                   color='red')\n",
    "        self.v_vector = plt.quiver(p[0],\n",
    "                                   p[1],\n",
    "                                   v[0],\n",
    "                                   v[1],\n",
    "                                   angles='xy',\n",
    "                                   scale_units='xy',\n",
    "                                   scale=1,\n",
    "                                   width=0.01,\n",
    "                                   color='black')\n",
    "        return [self.a_vector, self.v_vector]\n",
    "\n",
    "    def _update_visualization(self, x, u):\n",
    "        p, v, a = x[:self.spatial_dim], x[self.spatial_dim:], u\n",
    "        if self.spatial_dim == 1:\n",
    "            p, v, a = np.pad(p, [0, 1]), np.pad(v, [0, 1]), np.pad(a, [0, 1])\n",
    "        self.a_vector.set_offsets([[p[0], p[1]]])\n",
    "        self.v_vector.set_offsets([[p[0], p[1]]])\n",
    "        self.a_vector.set_UVC(a[0], a[1])\n",
    "        self.v_vector.set_UVC(v[0], v[1])\n",
    "\n",
    "\n",
    "def rollout_with_feedback(dynamics, controller, x0, N):\n",
    "    x = np.zeros((N + 1, dynamics.state_dim))\n",
    "    u = np.zeros((N, dynamics.control_dim))\n",
    "\n",
    "    x[0] = x0\n",
    "    for t in range(N):\n",
    "        u[t] = controller(x[t])\n",
    "        x[t + 1] = dynamics(x[t], u[t])\n",
    "\n",
    "    return x, u\n",
    "\n",
    "\n",
    "def rollout_control_sequence(dynamics, control_sequence, x0):\n",
    "    N = len(control_sequence)\n",
    "    x = np.zeros((N + 1, dynamics.state_dim))\n",
    "\n",
    "    x[0] = x0\n",
    "    for t in range(N):\n",
    "        x[t + 1] = dynamics(x[t], u[t])\n",
    "\n",
    "    return x, u"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1D double integrator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Problem setup\n",
    "dynamics = DoubleIntegatorDynamics(spatial_dim=1, time_step=0.1)\n",
    "\n",
    "time_horizon = 5\n",
    "T = int(time_horizon / dynamics.time_step)\n",
    "x0 = np.array([5, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Feedback control\n",
    "\n",
    "def linear_feedback_controller(k_p, k_d):\n",
    "\n",
    "    def controller(x):\n",
    "        return -k_p * x[0] - k_d * x[1]\n",
    "\n",
    "    return controller"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, u = rollout_with_feedback(dynamics, linear_feedback_controller(k_p=1.0, k_d=1.0), x0, T)\n",
    "dynamics.animate_trajectory(x, u)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 200\n",
    "x0 = np.array([5., 1.])\n",
    "\n",
    "def plot_trajectory(k_p=1.0, k_d=1.0):\n",
    "    x, u = rollout_with_feedback(dynamics, linear_feedback_controller(k_p, k_d), x0, T)\n",
    "    fig = dynamics.plot_trajectory(x, u)\n",
    "    fig.suptitle(f\"k_p = {k_p}, k_d = {k_d}\", fontsize=24)\n",
    "\n",
    "\n",
    "w = interact(plot_trajectory, k_p=(-4.0, 4.0), k_d=(-4.0, 4.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimal control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cvxpy as cp\n",
    "\n",
    "dynamics = DoubleIntegatorDynamics(spatial_dim=1, time_step=0.1)\n",
    "x0 = np.array([5., 0.])\n",
    "xf = np.array([0., 0.])\n",
    "T = 50\n",
    "\n",
    "x = cp.Variable((T + 1, dynamics.state_dim))\n",
    "u = cp.Variable((T, dynamics.control_dim))\n",
    "\n",
    "cost = 0\n",
    "constraints = [x[0] == x0]\n",
    "# constraints += [x[T] == xf]\n",
    "for t in range(T):\n",
    "    cost += cp.sum_squares(x[t]) + cp.sum_squares(u[t])\n",
    "    constraints += [x[t + 1] == dynamics(x[t], u[t])]\n",
    "#     constraints += [cp.norm_inf(u[t]) <= 2]\n",
    "problem = cp.Problem(cp.Minimize(cost), constraints)\n",
    "problem.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dynamics.animate_trajectory(x.value, u.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dynamics = DoubleIntegatorDynamics(spatial_dim=2, time_step=0.1)\n",
    "x0 = np.array([-5., 5., 0., 0.])\n",
    "xf = np.array([5., -5., 0., 0.])\n",
    "T = 50\n",
    "\n",
    "x = cp.Variable((T + 1, dynamics.state_dim))\n",
    "u = cp.Variable((T, dynamics.control_dim))\n",
    "\n",
    "cost = 0\n",
    "constraints = [x[0] == x0, x[T] == xf]\n",
    "for t in range(T):\n",
    "    cost += cp.sum_squares(x[t]) + cp.sum_squares(u[t])\n",
    "#     cost += cp.sum_squares(x[t, :2] - np.array([5., 5.])) + cp.sum_squares(u[t])\n",
    "    constraints += [x[t + 1] == dynamics(x[t], u[t])]\n",
    "#     constraints += [cp.norm_inf(u[t]) <= 2]\n",
    "problem = cp.Problem(cp.Minimize(cost), constraints)\n",
    "problem.solve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dynamics.animate_trajectory(x.value, u.value)"
   ]
  },
  {
   "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": 4
}