XYQV Splines.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import sympy\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mean_control_effort_coefficients(x0, dx0, xf, dxf):\n",
    "    \"\"\"Returns `(c4, c3, c2)` corresponding to `c4 * tf**-4 + c3 * tf**-3 + c2 * tf**-2`.\"\"\"\n",
    "    return (12 * (x0 - xf)**2, 12 * (dx0 + dxf) * (x0 - xf), 4 * dx0**2 + 4 * dx0 * dxf + 4 * dxf**2)\n",
    "\n",
    "\n",
    "def cubic_spline_coefficients(x0, dx0, xf, dxf, tf):\n",
    "    return (x0, dx0, -2 * dx0 / tf - dxf / tf - 3 * x0 / tf**2 + 3 * xf / tf**2,\n",
    "            dx0 / tf**2 + dxf / tf**2 + 2 * x0 / tf**3 - 2 * xf / tf**3)\n",
    "\n",
    "\n",
    "def compute_interpolating_spline(state_0, state_f, tf=None):\n",
    "    x0, y0, q0, v0 = state_0\n",
    "    xf, yf, qf, vf = state_f\n",
    "    dx0, dy0 = v0 * jnp.cos(q0), v0 * jnp.sin(q0)\n",
    "    dxf, dyf = vf * jnp.cos(qf), vf * jnp.sin(qf)\n",
    "    if tf is None:\n",
    "        c4, c3, c2 = [\n",
    "            a + b for a, b in zip(mean_control_effort_coefficients(x0, dx0, xf, dxf),\n",
    "                                  mean_control_effort_coefficients(y0, dy0, yf, dyf))\n",
    "        ]\n",
    "        c, b, a = -4 * c4, -3 * c3, -2 * c2\n",
    "        d = b**2 - 4 * a * c\n",
    "        tf = (-b + jnp.sqrt(d)) / (2 * a)\n",
    "        tf = jnp.where((d < 0) | (tf < 0), jnp.maximum(jnp.hypot(x0 - xf, y0 - yf), 1), tf)\n",
    "    return (\n",
    "        jnp.array(cubic_spline_coefficients(x0, dx0, xf, dxf, tf)),\n",
    "        jnp.array(cubic_spline_coefficients(y0, dy0, yf, dyf, tf)),\n",
    "        tf,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Free Final Time\n",
    "Minimizing mean control effort happens to have a closed-form solution for the local optimum (if it exists); the global optimum is of course at `tf → ∞`. Alternatively we could consider a mixed time/control effort penalty and optimize numerically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_0 = np.array([0, 0, 0, 10.])\n",
    "state_f = np.array([\n",
    "    np.array([x, y, q, v]) for x in [24, 30, 36] for y in [-4, 0, 4] for q in [-np.pi / 12, 0, np.pi / 12]\n",
    "    for v in [8, 10, 12]\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_spline_xyvat(x_coefficients, y_coefficients, tf, N=30):\n",
    "    t = jnp.linspace(0, tf, N)\n",
    "    tp = t[:, None]**np.arange(4)\n",
    "    dtp = t[:, None]**np.array([0, 0, 1, 2]) * np.arange(4)\n",
    "    ddtp = t[:, None]**np.array([0, 0, 0, 1]) * np.array([0, 0, 2, 6])\n",
    "    return (\n",
    "        tp @ x_coefficients,\n",
    "        tp @ y_coefficients,\n",
    "        jnp.hypot(dtp @ x_coefficients, dtp @ y_coefficients),\n",
    "        jnp.hypot(ddtp @ x_coefficients, ddtp @ y_coefficients),\n",
    "        t,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xc, yc, tf = jax.vmap(compute_interpolating_spline, in_axes=(None, 0))(state_0, state_f)\n",
    "x, y, v, a, t = jax.vmap(compute_spline_xyvat)(xc, yc, tf)\n",
    "plt.figure(figsize=(20, 10))\n",
    "plt.plot(x.T, y.T);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fixed Final Time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xc, yc, tf = jax.vmap(lambda s0, s1: compute_interpolating_spline(s0, s1, 3), in_axes=(None, 0))(state_0, state_f)\n",
    "x, y, v, a, t = jax.vmap(compute_spline_xyvat)(xc, yc, tf)\n",
    "plt.figure(figsize=(20, 10))\n",
    "plt.plot(x.T, y.T);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Derivations\n",
    "Precomputing these expressions saves a bit of computation over doing linear system solves online."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t, tf = sympy.symbols(\"t, t_f\")\n",
    "x0, dx0, xf, dxf = sympy.symbols(\"x_0, \\dot{x}_0, x_f, \\dot{x}_f\")\n",
    "symbols_to_variables = [(tf, \"tf\"), (x0, \"x0\"), (dx0, \"dx0\"), (xf, \"xf\"), (dxf, \"dxf\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "monomials = sympy.Array(t**np.arange(4))\n",
    "coefficients = sympy.Matrix([\n",
    "    monomials.subs(t, 0),\n",
    "    sympy.diff(monomials, t).subs(t, 0),\n",
    "    monomials.subs(t, tf),\n",
    "    sympy.diff(monomials, t).subs(t, tf)\n",
    "]).LUsolve(sympy.Matrix([x0, dx0, xf, dxf])).expand()\n",
    "coefficients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_control_effort = sympy.integrate(sympy.diff(coefficients.dot(monomials), t, 2)**2 / tf, (t, 0, tf)).expand()\n",
    "mean_control_effort"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sympy.pycode(np.array(coefficients.subs(symbols_to_variables)).ravel()), sympy.pycode([\n",
    "    c.subs(symbols_to_variables)\n",
    "    for c in reversed(sympy.factor(sympy.Poly((tf**4 * mean_control_effort), tf).all_coeffs()))\n",
    "])"
   ]
  },
  {
   "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
}