Code_for_lecture_11.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "Dependencies:\n",
    "- System: python3, ffmpeg (for rendering animations)\n",
    "- Python: jupyter, jax, numpy, matplotlib, plotly, tqdm, hj-reachability\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 jax numpy matplotlib plotly tqdm hj-reachability\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-12/HJ%20Reachability%20--%20Pursuit%20Evasion.ipynb) and run a cell containing this command:\n",
    "```\n",
    "!pip install --upgrade hj-reachability\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "\n",
    "from IPython.display import HTML\n",
    "import matplotlib.animation as anim\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import transforms\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "import hj_reachability as hj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dynamics = hj.systems.Air3d()\n",
    "grid = hj.Grid.from_lattice_parameters_and_boundary_conditions(hj.sets.Box(lo=np.array([-6., -10., 0.]),\n",
    "                                                                           hi=np.array([20., 10., 2 * np.pi])),\n",
    "                                                               (53, 41, 50),\n",
    "                                                               periodic_dims=2)\n",
    "terminal_values = jnp.linalg.norm(grid.states[..., :2], axis=-1) - 5\n",
    "\n",
    "solver_settings = hj.SolverSettings.with_accuracy(\"very_high\",\n",
    "                                                  hamiltonian_postprocessor=hj.solver.backwards_reachable_tube)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times = np.linspace(0, -3.5, 71)\n",
    "all_values = hj.solve(solver_settings, dynamics, grid, times, terminal_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "go.Figure(data=go.Isosurface(x=grid.states[..., 0].ravel(),\n",
    "                             y=grid.states[..., 1].ravel(),\n",
    "                             z=grid.states[..., 2].ravel(),\n",
    "                             value=all_values[-1].ravel(),\n",
    "                             colorscale=\"jet\",\n",
    "                             isomin=0,\n",
    "                             surface_count=1,\n",
    "                             isomax=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vmin, vmax = all_values.min(), all_values.max()\n",
    "levels = np.linspace(round(vmin), round(vmax), round(vmax) - round(vmin) + 1)\n",
    "fig = plt.figure(figsize=(13, 8))\n",
    "plt.jet()\n",
    "\n",
    "\n",
    "def render_frame(i, colorbar=False):\n",
    "    plt.contourf(grid.coordinate_vectors[0],\n",
    "                 grid.coordinate_vectors[1],\n",
    "                 all_values[i, :, :, 30].T,\n",
    "                 vmin=vmin,\n",
    "                 vmax=vmax,\n",
    "                 levels=levels)\n",
    "    if colorbar:\n",
    "        plt.colorbar()\n",
    "        plt.title(f\"Slice at θ_rel = {float(grid.coordinate_vectors[2][30]):4.3f}\", fontsize=20)\n",
    "    plt.contour(grid.coordinate_vectors[0],\n",
    "                grid.coordinate_vectors[1],\n",
    "                all_values[-1, :, :, 30].T,\n",
    "                levels=0,\n",
    "                colors=\"black\",\n",
    "                linewidths=3)\n",
    "\n",
    "\n",
    "render_frame(0, True)\n",
    "animation = HTML(anim.FuncAnimation(fig, render_frame, all_values.shape[0], interval=50).to_html5_video())\n",
    "plt.close()\n",
    "animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Re-solve, this time just BRS (i.e., not BRT) computation.\n",
    "# This approach for pursuit-evasion policy computation is a bit hacky/memory-inefficient for now.\n",
    "solver_settings = hj.SolverSettings.with_accuracy(\"very_high\")\n",
    "dynamics = hj.systems.Air3d()\n",
    "grid = hj.Grid.from_lattice_parameters_and_boundary_conditions(hj.sets.Box(lo=np.array([-10., -10., 0.]),\n",
    "                                                                           hi=np.array([20., 10., 2 * np.pi])),\n",
    "                                                               (61, 41, 50),\n",
    "                                                               periodic_dims=2)\n",
    "times = np.linspace(0, -5, 51)\n",
    "terminal_values = jnp.linalg.norm(grid.states[..., :2], axis=-1) - 5\n",
    "\n",
    "all_values = hj.solve(solver_settings, dynamics, grid, times, terminal_values)\n",
    "all_grad_values = jax.vmap(grid.grad_values)(all_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relative_state(x):\n",
    "    xa, ya, qa, xb, yb, qb = x\n",
    "    rot_matrix = jnp.array([[jnp.cos(qa), jnp.sin(qa)], [-jnp.sin(qa), jnp.cos(qa)]])\n",
    "    return jnp.array([*(rot_matrix @ jnp.array([xb - xa, yb - ya])), jnp.mod(qb - qa, 2 * jnp.pi)])\n",
    "\n",
    "\n",
    "def joint_dynamics(x, u, relative_dynamics=dynamics):\n",
    "    return jnp.array([\n",
    "        relative_dynamics.evader_speed * jnp.cos(x[2]),\n",
    "        relative_dynamics.evader_speed * jnp.sin(x[2]),\n",
    "        u[0],\n",
    "        relative_dynamics.pursuer_speed * jnp.cos(x[5]),\n",
    "        relative_dynamics.pursuer_speed * jnp.sin(x[5]),\n",
    "        u[1],\n",
    "    ])\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def joint_step(joint_state, dt, t):\n",
    "    state = relative_state(joint_state)\n",
    "    all_state_values = jax.vmap(grid.interpolate, (0, None))(all_values, state)\n",
    "    # Find the time horizon at which the value function is minimized (excluding any immediate period\n",
    "    # where it's increasing, in which case we should be reasoning about the \"next\" min).\n",
    "    i_increasing = jnp.concatenate([np.array([True]), all_state_values[1:] - all_state_values[:-1] > 0])\n",
    "    i = jnp.argmin(jnp.where(i_increasing, np.inf, all_state_values))\n",
    "    i = jnp.where(i == 0, -1, i)  # Default to max horizon policy if min value is the current time step.\n",
    "    value = grid.interpolate(all_values[i], state)\n",
    "    grad_value = grid.interpolate(all_grad_values[i], state)\n",
    "    a_opt, b_opt = dynamics.optimal_control_and_disturbance(state, t * dt, grad_value)\n",
    "    return joint_state + joint_dynamics(joint_state, jnp.concatenate([a_opt, b_opt])) * dt, value\n",
    "\n",
    "\n",
    "def joint_trajectory(evader_state, pursuer_state, dt=1 / 30, T=5):\n",
    "    joint_states = [np.concatenate([evader_state, pursuer_state])]\n",
    "    values = []\n",
    "    for t in range(int(T / dt)):\n",
    "        joint_state, value = joint_step(joint_states[-1], dt, t)\n",
    "        joint_states.append(joint_state)\n",
    "        values.append(value)\n",
    "    return np.array(joint_states[:-1]), np.array(values)\n",
    "\n",
    "\n",
    "def animate_joint_trajectory(evader_state, pursuer_state, dt=1 / 30, T=5, animation_time_scale_factor=2):\n",
    "    joint_states, values = joint_trajectory(evader_state, pursuer_state, dt, T)\n",
    "    xmin, xmax = np.min(joint_states[:, [0, 3]]), np.max(joint_states[:, [0, 3]])\n",
    "    ymin, ymax = np.min(joint_states[:, [1, 4]]), np.max(joint_states[:, [1, 4]])\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 8))\n",
    "    ax = fig.gca()\n",
    "    ax.set_xlim(xmin - 3, xmax + 3)\n",
    "    ax.set_ylim(ymin - 3, ymax + 3)\n",
    "    ax.set_aspect(\"equal\", adjustable=\"box\")\n",
    "    triangle_pts = np.array([[-.2, -.2], [1., 0], [-.2, .2]])\n",
    "    evader = ax.add_patch(plt.Polygon(triangle_pts))\n",
    "    pursuer = ax.add_patch(plt.Polygon(triangle_pts, color=\"orange\"))\n",
    "    evader_radius = ax.add_patch(plt.Circle([0, 0], 5, alpha=0.5))\n",
    "\n",
    "    def render_frame(i):\n",
    "        evader.set_transform(\n",
    "            transforms.Affine2D().rotate(joint_states[i, 2]).translate(joint_states[i, 0], joint_states[i, 1]) +\n",
    "            ax.transData)\n",
    "        pursuer.set_transform(\n",
    "            transforms.Affine2D().rotate(joint_states[i, 5]).translate(joint_states[i, 3], joint_states[i, 4]) +\n",
    "            ax.transData)\n",
    "        evader_radius.set_center([joint_states[i, 0], joint_states[i, 1]])\n",
    "        return [evader, pursuer, evader_radius]\n",
    "\n",
    "    animation = HTML(\n",
    "        anim.FuncAnimation(fig,\n",
    "                           render_frame,\n",
    "                           joint_states.shape[0],\n",
    "                           interval=1000 * dt / animation_time_scale_factor,\n",
    "                           blit=True).to_html5_video())\n",
    "    plt.close()\n",
    "    return animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "animate_joint_trajectory(np.zeros(3), np.array([15., -6., 3*np.pi/4]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "animate_joint_trajectory(np.zeros(3), np.array([5., -6., 3*np.pi/4]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "animate_joint_trajectory(np.zeros(3), np.array([10, 0, 3.77]), T=7)"
   ]
  },
  {
   "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
}