{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Road Runner Lite From Scratch\n",
    "\n",
    "This tutorial walks through the creation of a lite [Road Runner](https://github.com/acmerobotics/road-runner). It aims to cover all of the math and code required for a simple version. The only prerequisite is basic proficiency in single-variable calculus, vectors, and parametric curves at the level of [this Khan Academy unit](https://www.khanacademy.org/math/multivariable-calculus/multivariable-derivatives/position-vector-functions/v/position-vector-valued-functions). The full library is mostly a scaled-up version of these ideas with bells and whistles.\n",
    "\n",
    "The tutorial is written in Python using NumPy primitives instead of Kotlin. The decorator `vectorize_tail` is occasionally necessary to make functions cooperate with array inputs. You can safely ignore its definition and subsequent uses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "vectorize_tail = lambda f: np.vectorize(f, excluded={0})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trajectories\n",
    "\n",
    "Let's focus first on moving smoothly in a single dimension. Say we have a robot with position $x$ that we want to move from rest at $x_0$ to rest at $x_1$ (where $x_1 \\geq x_0$). We can cast this as the problem of finding a position function $x(t)$ that satisfies $x(0) = x_0$ and $x(T) = x_1$ (where $T \\geq 0$). This single function also determines the velocity and acceleration through its derivatives. There are many such functions, so we'll take the minimum-time one that obeys the maximum-velocity constraint $|x'(t)| \\leq v_{max}$ for all $0 \\leq t \\leq T$ (where $v_{max} > 0$). This optimal trajectory is achieved by going full throttle and setting $x'(t) = v_{max}$. We integrate to obtain\n",
    "\\begin{align*}\n",
    "    x(t) - x(0) &= \\int_0^t x'(\\tau) \\, d\\tau = \\int_0^t v_{max} \\, d\\tau = v_{max} t\\\\\n",
    "    x(t) &= x_0 + v_{max} t\n",
    "\\end{align*}\n",
    "where the start position constraint is guaranteed. The end position constraint can be satisfied by choosing the duration\n",
    "\\begin{align*}\n",
    "    x_1 &= x_0 + v_{max} T\\\\\n",
    "    T &= \\frac{x_1 - x_0}{v_{max}}.\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following snippet implements this velocity-limited trajectory generator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vel_traj_gen(x0, x1, vmax):\n",
    "    dt = (x1 - x0) / vmax\n",
    "    return x0, vmax, dt\n",
    "\n",
    "\n",
    "traj_x0, traj_v, traj_dt = vel_traj_gen(-20, 80, 30)\n",
    "t = np.linspace(0, traj_dt, 100)\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "ax[0].set_title('Velocity-Limited Trajectory Velocity')\n",
    "ax[0].set_xlabel('time [s]')\n",
    "ax[0].set_ylabel('velocity [in/s]')\n",
    "ax[0].plot(t, np.full(t.shape, traj_v)) \n",
    "\n",
    "ax[1].set_title('Velocity-Limited Trajectory Position')\n",
    "ax[1].set_xlabel('time [s]')\n",
    "ax[1].set_ylabel('position [in]')\n",
    "ax[1].plot(t, traj_x0 + traj_v * t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But can the robot faithfully execute the generated trajectory? It takes time for the robot to high velocity from rest, and we can command a more realistic plan by incorporating that limitation into our model. We need an acceleration limit $|x''(t)| \\leq a_{max}$ for all $0 \\leq t \\leq T$ (where $a_{max} > 0$) similar to the earlier velocity one. The optimal trajectory for the new setup has multiple distinct phases: acceleration, coast, and deceleration. The first and last phases go full throttle to maximize acceleration, while the second phase maintains the maximum velocity. To keep our calculations succint, we'll use the standard constant-velocity and constant-acceleration formulas. You can derive them from first principles using calculus as we did before.\n",
    "\n",
    "The first phase has acceleration $a_{max}$, and the duration $\\Delta t_1$ satisfies\n",
    "\\begin{align*}\n",
    "    v_{max} &= 0 + a_{max} \\Delta t_1\\\\\n",
    "    \\Delta t_1 &= \\frac{v_{max}}{a_{max}}.\n",
    "\\end{align*}\n",
    "The third phase has acceleration $-a_{max}$, and the duration $\\Delta t_3$ satisfies\n",
    "\\begin{align*}\n",
    "    0 &= v_{max} - a_{max} \\Delta t_3\\\\\n",
    "    \\Delta t_3 &= \\frac{v_{max}}{a_{max}},\n",
    "\\end{align*}\n",
    "which matches $\\Delta t_1$ as expected by symmetry. While the outer phases depend only on the constraints, the middle phase expands to achieve the desired total distance. It relates to the other parameters according to\n",
    "\\begin{align*}\n",
    "    x_1 &= x_0 + \\frac{1}{2} a_{max} \\Delta t_1^2 + v_{max} \\Delta t_2 + v_{max} \\Delta t_3 - \\frac{1}{2} a_{max} \\Delta t_3^2\\\\\n",
    "    &= x_0 + v_{max} \\Delta t_2 + \\frac{v_{max}^2}{a_{max}}\\\\\n",
    "    \\Delta t_2 &= \\frac{x_1 - x_0}{v_{max}} - \\frac{v_{max}}{a_{max}}.\n",
    "\\end{align*}\n",
    "But this difference expression is suspicious. What if the computed $\\Delta t_2$ is negative? The robot might not have enough runway to reach the maximum velocity and goes straight from acceleration into deceleration. Let's leverage the symmetry $\\Delta t_1 = \\Delta t_3 = \\Delta t$ to compute the new phase durations. The duration $\\Delta t$ is\n",
    "\\begin{align*}\n",
    "    x_1 &= x_0 + \\frac{1}{2} a_{max} \\Delta t_1^2 + v_{max} \\Delta t_3 - \\frac{1}{2} a_{max} \\Delta t_3^2\\\\\n",
    "    &= x_0 + \\frac{1}{2} a_{max} \\Delta t^2 + a_{max} \\Delta t^2 - \\frac{1}{2} a_{max} \\Delta t^2\\\\\n",
    "    &= x_0 + a_{max} \\Delta t^2\\\\\n",
    "    \\Delta t &= \\sqrt{\\frac{x_1 - x_0}{a_{max}}}\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With both the normal and degenerate case covered, we can implement the acceleration-limited trajectory generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accel_traj_gen(x0, x1, vmax, amax):\n",
    "    dx = x1 - x0\n",
    "    if vmax / amax < dx / vmax:\n",
    "        # normal trajectory\n",
    "        dt1 = vmax / amax\n",
    "        dt2 = dx / vmax - vmax / amax\n",
    "        dt3 = dt1\n",
    "        return x0, (\n",
    "            (amax, dt1),\n",
    "            (0, dt2),\n",
    "            (-amax, dt3)\n",
    "        )\n",
    "    else:\n",
    "        # degenerate trajectory\n",
    "        dt1 = np.sqrt(dx / amax)\n",
    "        dt2 = dt1\n",
    "        return x0, (\n",
    "            (amax, dt1), \n",
    "            (-amax, dt2)\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the trajectory to be useful, we need functions to compute instantaneous values at arbitrary points in time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@vectorize_tail\n",
    "def accel_traj_get_accel(traj, t):\n",
    "    _, phases = traj\n",
    "    for a, dt in phases:\n",
    "        if t < dt:\n",
    "            return a\n",
    "        \n",
    "        t -= dt\n",
    "    return a\n",
    "\n",
    "        \n",
    "@vectorize_tail\n",
    "def accel_traj_get_vel(traj, t):\n",
    "    _, phases = traj\n",
    "    v0 = 0\n",
    "    for a, dt in phases:\n",
    "        if t < dt:\n",
    "            return v0 + a * t\n",
    "        \n",
    "        v0 += a * dt\n",
    "        \n",
    "        t -= dt\n",
    "    return v0\n",
    "    \n",
    "    \n",
    "@vectorize_tail\n",
    "def accel_traj_get_pos(traj, t):\n",
    "    x0, phases = traj\n",
    "    v0 = 0\n",
    "    for a, dt in phases:\n",
    "        if t < dt:\n",
    "            return x0 + v0 * t + a * t**2 / 2\n",
    "        \n",
    "        x0 += v0 * dt + a * dt**2 / 2\n",
    "        v0 += a * dt\n",
    "        \n",
    "        t -= dt \n",
    "    return x0\n",
    "\n",
    "\n",
    "def accel_traj_duration(traj):\n",
    "    _, phases = traj\n",
    "    duration = 0\n",
    "    for _, dt in phases:\n",
    "        duration += dt\n",
    "    return duration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we can make some plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traj = accel_traj_gen(-20, 80, 30, 30)\n",
    "\n",
    "t = np.linspace(0, accel_traj_duration(traj), 100)\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(18, 4))\n",
    "\n",
    "ax[0].set_title('Acceleration-Limited Trajectory Acceleration')\n",
    "ax[0].set_xlabel('time [s]')\n",
    "ax[0].set_ylabel('acceleration [in/s^2]')\n",
    "ax[0].plot(t, accel_traj_get_accel(traj, t))\n",
    "\n",
    "ax[1].set_title('Acceleration-Limited Trajectory Velocity')\n",
    "ax[1].set_xlabel('time [s]')\n",
    "ax[1].set_ylabel('velocity [in/s]')\n",
    "ax[1].plot(t, accel_traj_get_vel(traj, t))\n",
    "\n",
    "ax[2].set_title('Acceleration-Limited Trajectory Position')\n",
    "ax[2].set_xlabel('time [s]')\n",
    "ax[2].set_ylabel('position [in]')\n",
    "ax[2].plot(t, accel_traj_get_pos(traj, t))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try adjusting the parameters to get a degenerate trajectory with only two segments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Splines\n",
    "\n",
    "But we want our robot to travel in more than just straight lines. We can use a sequence of polynomials called a spline to create a windy path that skirts around obstacles. Each segment of a 2D cubic spline has the form\n",
    "\\begin{align*}\n",
    "    x(u) &= a_x u^3 + b_x u^2 + c_x u + d_x\\\\\n",
    "    y(u) &= a_y u^3 + b_y u^2 + c_y u + d_y\n",
    "\\end{align*}\n",
    "where $0 \\leq u \\leq 1$. To make sure the curve hits specific start and end coordinates, we impose the constraints\n",
    "\\begin{align*}\n",
    "    x(0) &= x_0 & x(1) &= x_1\\\\\n",
    "    y(0) &= y_0 & y(1) &= y_1.\n",
    "\\end{align*}\n",
    "And now the extra flexibility of the spline enables us to specify the begin and end tangents, $(x'_0, y'_0)$ and $(x'_1, y'_1)$. These requirements give the derivative constraints\n",
    "\\begin{align*}\n",
    "    x'(0) &= x'_0 & x'(1) &= x'_1\\\\\n",
    "    y'(0) &= y'_0 & y'(1) &= y'_1.\n",
    "\\end{align*}\n",
    "By this point you may have noticed that $x$ and $y$ seem independent and can be computed separately. Let's consider the $x$ one and suppress the subscripts on the coefficients. The derivative is\n",
    "$$\n",
    "    x'(u) = 3 a u^2 + 2 b u + c.\n",
    "$$\n",
    "Evaluating all of the constraints, we obtain the system of equations\n",
    "\\begin{align*}\n",
    "    x_0 &= d\\\\\n",
    "    x'_0 &= c\\\\\n",
    "    x_1 &= a + b + c + d\\\\\n",
    "    x'_1 &= 3a + 2b + c.\n",
    "\\end{align*}\n",
    "As an exercise, you can verify that the solution is \n",
    "\\begin{align*}\n",
    "    a &= 2x_0 + x'_0 - 2x_1 + x'_1\\\\\n",
    "    b &= -3x_0 - 2x'_0 + 3x_1 - x'_1\\\\\n",
    "    c &= x'_0\\\\\n",
    "    d &= x_0.\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can write methods to fit spline parameters and give spline values and derivatives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def spline_fit(x0, dx0, x1, dx1):\n",
    "    a = 2 * x0 + dx0 - 2 * x1 + dx1\n",
    "    b = -3 * x0 - 2 * dx0 + 3 * x1 - dx1\n",
    "    c = dx0\n",
    "    d = x0\n",
    "    return a, b, c, d\n",
    "\n",
    "def spline_get(spline, u):\n",
    "    a, b, c, d = spline\n",
    "    return a * u**3 + b * u**2 + c * u + d\n",
    "\n",
    "def spline_deriv(spline, u):\n",
    "    a, b, c, d = spline\n",
    "    return 3 * a * u**2 + 2 * b * u + c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's give this a quick test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_spline = spline_fit(0, 36, 24, 30)\n",
    "y_spline = spline_fit(0, -24, 24, -9)\n",
    "u = np.linspace(0, 1, 100)\n",
    "\n",
    "plt.title('Cubic Spline')\n",
    "plt.xlabel('x [in]')\n",
    "plt.ylabel('y [in]')\n",
    "\n",
    "plt.plot(spline_get(x_spline, u), spline_get(y_spline, u))\n",
    "\n",
    "print(spline_get(x_spline, 0), spline_deriv(x_spline, 0), spline_get(x_spline, 1), spline_deriv(x_spline, 1))\n",
    "print(spline_get(y_spline, 0), spline_deriv(y_spline, 0), spline_get(y_spline, 1), spline_deriv(y_spline, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have a spline now, and we just need to marry it with a trajectory. The general idea is to compute a trajectory for the displacement along the spline and compose the two. But we only have the spline as a function of an arbitrary parameter instead of displacement. Put differently, we're missing a function $u(s)$ to compose in between.\n",
    "\n",
    "A handy integral gives the displacement as a function of $u$\n",
    "$$\n",
    "    s(u) = \\int_0^u \\sqrt{\\left( \\frac{dx}{d\\upsilon} \\right)^2 + \\left( \\frac{dy}{d\\upsilon} \\right)^2} \\, d\\upsilon.\n",
    "$$\n",
    "While we can't easily invert the integral analytically, we can write code to compute indvidual values. Recalling Riemann sums, integrals can be numerically approximated by adding up small rectangular areas. We can compute the $u$-value corresponding to a given displacement $s$ by tracking the running total and stopping when we reach $s$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "upsilon = np.linspace(0, 1, 100)\n",
    "dupsilon = upsilon[1] - upsilon[0]\n",
    "integrand = np.sqrt(\n",
    "    spline_deriv(x_spline, upsilon)**2 + \n",
    "    spline_deriv(y_spline, upsilon)**2\n",
    ")\n",
    "\n",
    "sums = np.zeros_like(upsilon)\n",
    "last_sum = 0\n",
    "for i in range(len(upsilon)):\n",
    "    sums[i] = last_sum + integrand[i] * dupsilon\n",
    "    last_sum = sums[i]\n",
    "    \n",
    "    \n",
    "@np.vectorize\n",
    "def spline_param_of_disp(s):\n",
    "    for i in range(len(sums)):\n",
    "        if s < sums[i]:\n",
    "            return upsilon[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we put everything together and generate a trajectory for the spline!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "length = sums[-1]\n",
    "traj = accel_traj_gen(0, length, 30, 30)\n",
    "t = np.linspace(0, accel_traj_duration(traj), 100)\n",
    "\n",
    "s = accel_traj_get_pos(traj, t)\n",
    "\n",
    "u = spline_param_of_disp(s)\n",
    "\n",
    "x = spline_get(x_spline, u)\n",
    "y = spline_get(y_spline, u)\n",
    "\n",
    "plt.title('Spline Trajectory Position')\n",
    "plt.xlabel('time [s]')\n",
    "plt.ylabel('position [in]')\n",
    "plt.plot(t, x, label='x')\n",
    "plt.plot(t, y, label='y')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each component of the trajectory is the composition of three functions: $(x(u(s(t))), y(u(s(t))))$. We can find the velocity by differentiating with respect to $t$. For the $x$-component, this gives\n",
    "\\begin{align*}\n",
    "    x'(t) &= \\frac{d}{dt} x(u(s(t))))\\\\\n",
    "    &= x'(u(s(t)) \\, u'(s(t)) \\, s'(t).\n",
    "\\end{align*}\n",
    "The only derivative we're missing is $u'(s)$. More chain rule abuse yields\n",
    "\\begin{align*}\n",
    "    u'(s) &= [s'(u)]^{-1}\\\\\n",
    "    &= \\left[\\frac{d}{du} \\int_0^u \\sqrt{\\left( \\frac{dx}{d\\upsilon} \\right)^2 + \\left( \\frac{dy}{d\\upsilon} \\right)^2} \\, d\\upsilon \\right]^{-1}\\\\\n",
    "    &= \\left[\\sqrt{\\left( \\frac{dx}{du} \\right)^2 + \\left( \\frac{dy}{du} \\right)^2}\\right]^{-1}.\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def spline_param_of_disp_deriv(x_spline, y_spline, u):\n",
    "    return 1.0 / np.sqrt(spline_deriv(x_spline, u)**2 + spline_deriv(y_spline, u)**2)\n",
    "\n",
    "\n",
    "dsdt = accel_traj_get_vel(traj, t)\n",
    "\n",
    "duds = spline_param_of_disp_deriv(x_spline, y_spline, u)\n",
    "\n",
    "dxdu = spline_deriv(x_spline, u)\n",
    "dydu = spline_deriv(y_spline, u)\n",
    "\n",
    "dxdt = dxdu * duds * dsdt\n",
    "dydt = dydu * duds * dsdt\n",
    "\n",
    "\n",
    "plt.title('Spline Trajectory Velocity')\n",
    "plt.xlabel('time [s]')\n",
    "plt.ylabel('velocity [in/s]')\n",
    "plt.plot(t, dxdt, label='x')\n",
    "plt.plot(t, dydt, label='y')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try computing the spline trajectory acceleration as an advanced exercise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}