Code_for_lecture_8.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "western-breeding",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "Dependencies:\n",
    "- System: python3\n",
    "- Python: jupyter, numpy, matplotlib, jax\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",
    "/usr/bin/python3 -m pip install --upgrade pip\n",
    "pip install --upgrade jupyter numpy matplotlib jax\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-10/LQR%20Variants.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "stretch-singer",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt; plt.rcParams.update({'font.size': 20})\n",
    "from ipywidgets import interact, interactive\n",
    "\n",
    "# `NamedTuple`s are used (more accurately, abused) in this notebook to minimize dependencies;\n",
    "# much better JAX-compatible choices to fit the archetype of \"parameterized function\" would be\n",
    "# `flax.struct.dataclass` or `equinox.Module`.\n",
    "from typing import Callable, NamedTuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "prescribed-sharp",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Notation from https://github.com/StanfordASL/AA203-Notes/blob/master/notes.pdf;\n",
    "# notation from https://asl.stanford.edu/aa203/pdfs/lecture/lectures_9_and_10.pdf in comments.\n",
    "\n",
    "\n",
    "class LinearDynamics(NamedTuple):\n",
    "    f_x: jnp.array  # A\n",
    "    f_u: jnp.array  # B\n",
    "\n",
    "    def __call__(self, x, u, k=None):\n",
    "        f_x, f_u = self\n",
    "        return f_x @ x + f_u @ u if k is None else self[k](x, u)\n",
    "\n",
    "    def __getitem__(self, key):\n",
    "        return jax.tree.map(lambda x: x[key], self)\n",
    "\n",
    "\n",
    "class AffinePolicy(NamedTuple):\n",
    "    l: jnp.array  # l\n",
    "    l_x: jnp.array  # L\n",
    "\n",
    "    def __call__(self, x, k=None):\n",
    "        l, l_x = self\n",
    "        return l + l_x @ x if k is None else self[k](x)\n",
    "\n",
    "    def __getitem__(self, key):\n",
    "        return jax.tree.map(lambda x: x[key], self)\n",
    "\n",
    "\n",
    "class QuadraticCost(NamedTuple):\n",
    "    c: jnp.array  # c\n",
    "    c_x: jnp.array  # q\n",
    "    c_u: jnp.array  # r\n",
    "    c_xx: jnp.array  # Q\n",
    "    c_uu: jnp.array  # R\n",
    "    c_ux: jnp.array  # H.T\n",
    "\n",
    "    @classmethod\n",
    "    def from_pure_quadratic(cls, c_xx, c_uu, c_ux):\n",
    "        return cls(\n",
    "            jnp.zeros((c_xx.shape[:-2])),\n",
    "            jnp.zeros(c_xx.shape[:-1]),\n",
    "            jnp.zeros(c_uu.shape[:-1]),\n",
    "            c_xx,\n",
    "            c_uu,\n",
    "            c_ux,\n",
    "        )\n",
    "\n",
    "    def __call__(self, x, u, k=None):\n",
    "        c, c_x, c_u, c_xx, c_uu, c_ux = self\n",
    "        return c + c_x @ x + c_u @ u + x @ c_xx @ x / 2 + u @ c_uu @ u / 2 + u @ c_ux @ x if k is None else self[k](x)\n",
    "\n",
    "    def __getitem__(self, key):\n",
    "        return jax.tree.map(lambda x: x[key], self)\n",
    "\n",
    "\n",
    "class QuadraticStateCost(NamedTuple):\n",
    "    v: jnp.array  # p (scalar)\n",
    "    v_x: jnp.array  # p (vector)\n",
    "    v_xx: jnp.array  # P\n",
    "\n",
    "    @classmethod\n",
    "    def from_pure_quadratic(cls, v_xx):\n",
    "        return cls(\n",
    "            jnp.zeros(v_xx.shape[:-2]),\n",
    "            jnp.zeros(v_xx.shape[:-1]),\n",
    "            v_xx,\n",
    "        )\n",
    "\n",
    "    def __call__(self, x, k=None):\n",
    "        v, v_x, v_xx = self\n",
    "        return v + v_x @ x + x @ v_xx @ x / 2 if k is None else self[k](x)\n",
    "\n",
    "    def __getitem__(self, key):\n",
    "        return jax.tree.map(lambda x: x[key], self)\n",
    "\n",
    "\n",
    "def rollout_state_feedback_policy(dynamics, policy, x0, step_range, x_nom=None, u_nom=None):\n",
    "\n",
    "    def scan_fn(x, k):\n",
    "        u = policy(x, k) if x_nom is None else u_nom[k] + policy(x - x_nom[k], k)\n",
    "        x1 = dynamics(x, u, k)\n",
    "        return (x1, (x1, u))\n",
    "\n",
    "    xs, us = jax.lax.scan(scan_fn, x0, step_range)[1]\n",
    "    return jnp.concatenate([x0[None], xs]), us"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "behind-hungary",
   "metadata": {},
   "source": [
    "## [Discrete-time, finite-horizon] Linear Quadratic Regulator (LQR)\n",
    "Minimizing a quadratic regulation cost:\n",
    "\n",
    "$\n",
    "J_k(\\mathbf{x}_k) = \\frac{1}{2}Q_N \\|\\mathbf{x}_N\\|^2 + \\frac{1}{2}\\sum_{i=k}^{N-1} \\left(\\|\\mathbf{x}_k\\|^2 + \\|\\mathbf{u}_k\\|^2\\right)\n",
    "$\n",
    "\n",
    "subject to double-integrator dynamics with a time discretization of $\\Delta t$:\n",
    "\n",
    "$\\begin{bmatrix}p_{k+1} \\\\ v_{k+1}\\end{bmatrix} =\n",
    "\\begin{bmatrix}1 & \\Delta t \\\\ 0 & 1\\end{bmatrix} \\begin{bmatrix}p_k \\\\ v_k\\end{bmatrix} +\n",
    "\\begin{bmatrix}\\frac{1}{2}\\Delta t^2 \\\\ \\Delta t\\end{bmatrix} \\begin{bmatrix}u_k\\end{bmatrix}.\n",
    "$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "administrative-denver",
   "metadata": {},
   "outputs": [],
   "source": [
    "def riccati_step(\n",
    "    current_step_dynamics: LinearDynamics,\n",
    "    current_step_cost: QuadraticCost,\n",
    "    next_state_value: QuadraticStateCost,\n",
    "):\n",
    "    f_x, f_u = current_step_dynamics\n",
    "    c, c_x, c_u, c_xx, c_uu, c_ux = current_step_cost\n",
    "    v, v_x, v_xx = next_state_value\n",
    "\n",
    "    q = c + v\n",
    "    q_x = c_x + f_x.T @ v_x\n",
    "    q_u = c_u + f_u.T @ v_x\n",
    "    q_xx = c_xx + f_x.T @ v_xx @ f_x\n",
    "    q_uu = c_uu + f_u.T @ v_xx @ f_u\n",
    "    q_ux = c_ux + f_u.T @ v_xx @ f_x\n",
    "\n",
    "    l = -jnp.linalg.solve(q_uu, q_u)\n",
    "    l_x = -jnp.linalg.solve(q_uu, q_ux)\n",
    "\n",
    "    current_state_value = QuadraticStateCost(\n",
    "        q - l.T @ q_uu @ l / 2,\n",
    "        q_x - l_x.T @ q_uu @ l,\n",
    "        q_xx - l_x.T @ q_uu @ l_x,\n",
    "    )\n",
    "    current_step_optimal_policy = AffinePolicy(l, l_x)\n",
    "    return current_state_value, current_step_optimal_policy\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def linear_quadratic_regulator(Qs, Rs, Hs, As, Bs):\n",
    "    final_state_value = QuadraticStateCost.from_pure_quadratic(Qs[-1])\n",
    "    dynamics = LinearDynamics(As, Bs)\n",
    "    cost = QuadraticCost.from_pure_quadratic(Qs[:-1], Rs, Hs)\n",
    "\n",
    "    def scan_fn(next_state_value, current_step_dynamics_cost):\n",
    "        current_step_dynamics, current_step_cost = current_step_dynamics_cost\n",
    "        current_state_value, current_step_optimal_policy = riccati_step(\n",
    "            current_step_dynamics,\n",
    "            current_step_cost,\n",
    "            next_state_value,\n",
    "        )\n",
    "        return current_state_value, (current_state_value, current_step_optimal_policy)\n",
    "\n",
    "    value_functions, optimal_policy = jax.lax.scan(scan_fn, final_state_value, (dynamics, cost), reverse=True)[1]\n",
    "    return jax.tree.map(lambda x, y: jnp.concatenate([x, y]), value_functions, final_state_value[None]), optimal_policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "civic-airport",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10\n",
    "dt = 0.5\n",
    "Q_N = 100  # Change to 100!\n",
    "\n",
    "Qs = np.broadcast_to(np.eye(2), (N + 1, 2, 2)).copy()\n",
    "Qs[-1] *= Q_N\n",
    "Rs = np.broadcast_to(np.eye(1), (N, 1, 1))\n",
    "Hs = np.broadcast_to(np.zeros((1, 2)), (N, 1, 2))\n",
    "As = np.broadcast_to(np.array([[1, dt], [0, 1]]), (N, 2, 2))\n",
    "Bs = np.broadcast_to(np.array([[dt**2 / 2], [dt]]), (N, 2, 1))\n",
    "value_functions, optimal_policy = linear_quadratic_regulator(Qs, Rs, Hs, As, Bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fancy-portable",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "X0, X1 = np.meshgrid(np.linspace(-5, 5, 100), np.linspace(-5, 5, 100))\n",
    "X = np.stack([X0, X1], -1)\n",
    "all_values = jax.vmap(lambda v: jax.vmap(jax.vmap(v))(X))(value_functions)\n",
    "\n",
    "dynamics = LinearDynamics(jnp.array(As), jnp.array(Bs))\n",
    "x0s = np.random.uniform(-3, 3, (20, 2))\n",
    "example_trajectories = [\n",
    "    jax.vmap(lambda x0: rollout_state_feedback_policy(\n",
    "        dynamics,\n",
    "        optimal_policy,\n",
    "        x0,\n",
    "        np.arange(k, N),\n",
    "    ))(x0s)[0] for k in range(N)\n",
    "]\n",
    "\n",
    "\n",
    "@interact(k=(0, N))\n",
    "def plot(k=N):\n",
    "    plt.figure(figsize=(10, 8))\n",
    "    plt.contour(X0, X1, all_values[k], levels=40)\n",
    "    plt.colorbar()\n",
    "    plt.xlim(-5, 5)\n",
    "    plt.ylim(-5, 5)\n",
    "    plt.xlabel(\"Position $p$\")\n",
    "    plt.ylabel(\"Velocity $v$\")\n",
    "    plt.title(f\"Value Function at Step {k}\")\n",
    "\n",
    "    if k < N:\n",
    "        plt.scatter(*x0s.T, color=\"red\")\n",
    "        plt.plot(*example_trajectories[k].T, color=\"grey\", alpha=0.5)\n",
    "        plt.plot(*example_trajectories[k][:, :2].T, color=\"red\")\n",
    "    else:\n",
    "        plt.scatter(*x0s.T, color=\"red\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "twelve-roots",
   "metadata": {},
   "source": [
    "## Tracking LQR for Linear Dynamics\n",
    "Minimizing a quadratic tracking cost:\n",
    "\n",
    "$\n",
    "J_k(\\mathbf{\\delta x}_k) = \\frac{1}{2}Q_N \\|\\mathbf{\\delta x}_N\\|^2 + \\frac{1}{2}\\sum_{i=k}^{N-1} \\left(\\|\\mathbf{\\delta x}_k\\|^2 + \\|\\mathbf{\\delta u}_k\\|^2\\right)\n",
    "$\n",
    "\n",
    "subject to single-integrator dynamics with a time discretization of $\\Delta t$:\n",
    "\n",
    "$\\begin{bmatrix}p_{k+1} \\\\ v_{k+1}\\end{bmatrix} =\n",
    "\\begin{bmatrix}1 & \\Delta t \\\\ 0 & 1\\end{bmatrix} \\begin{bmatrix}p_k \\\\ v_k\\end{bmatrix} +\n",
    "\\begin{bmatrix}\\frac{1}{2}\\Delta t^2 \\\\ \\Delta t\\end{bmatrix} \\begin{bmatrix}u_k\\end{bmatrix}\n",
    "$\n",
    "\n",
    "tracking a sinusoidal nominal trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "raised-asset",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_nom, u_nom = rollout_state_feedback_policy(dynamics, lambda x, k: -x[:1], np.array([0., 2.]), np.arange(N))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf9dbf15",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(x_nom)\n",
    "plt.plot(u_nom, label=r'$\\bar u$')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "8eac85da",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10\n",
    "dt = 0.5\n",
    "Q_N = 100  # Change to 100!\n",
    "\n",
    "Qs = np.broadcast_to(np.eye(2), (N + 1, 2, 2)).copy()\n",
    "Qs[-1] *= Q_N\n",
    "Rs = np.broadcast_to(np.eye(1), (N, 1, 1))\n",
    "Hs = np.broadcast_to(np.zeros((1, 2)), (N, 1, 2))\n",
    "As = np.broadcast_to(np.array([[1, dt], [0, 1]]), (N, 2, 2))\n",
    "Bs = np.broadcast_to(np.array([[dt**2 / 2], [dt]]), (N, 2, 1))\n",
    "value_functions, optimal_policy = linear_quadratic_regulator(Qs, Rs, Hs, As, Bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "infinite-respect",
   "metadata": {},
   "outputs": [],
   "source": [
    "X0, X1 = np.meshgrid(np.linspace(-5, 5, 100), np.linspace(-5, 5, 100))\n",
    "X = np.stack([X0, X1], -1)\n",
    "all_values = jax.vmap(lambda v, x: jax.vmap(jax.vmap(v))(X - x))(value_functions, x_nom)\n",
    "\n",
    "dynamics = LinearDynamics(jnp.array(As), jnp.array(Bs))\n",
    "x0s = np.random.uniform(-4, 4, (20, 2))\n",
    "example_trajectories = [\n",
    "    # Note that this is the exact same policy as above!\n",
    "    jax.vmap(lambda x0: rollout_state_feedback_policy(\n",
    "        dynamics,\n",
    "        optimal_policy,\n",
    "        x0,\n",
    "        np.arange(k, N),\n",
    "        x_nom,\n",
    "        u_nom,\n",
    "    ))(x0s)[0] for k in range(N)\n",
    "]\n",
    "\n",
    "\n",
    "@interact(k=(0, N))\n",
    "def plot(k=N):\n",
    "    plt.figure(figsize=(10, 8))\n",
    "    plt.contour(X0, X1, all_values[k], levels=40)\n",
    "    plt.colorbar()\n",
    "    plt.xlim(-5, 5)\n",
    "    plt.ylim(-5, 5)\n",
    "    plt.xlabel(\"Position $p$\")\n",
    "    plt.ylabel(\"Velocity $v$\")\n",
    "    plt.title(f\"Value Function at Step {k}\")\n",
    "\n",
    "    if k < N:\n",
    "        plt.scatter(*x0s.T, color=\"red\")\n",
    "        plt.plot(*example_trajectories[k].T, color=\"grey\", alpha=0.5)\n",
    "        plt.plot(*example_trajectories[k][:, :2].T, color=\"red\")\n",
    "\n",
    "    plt.plot(*x_nom.T, color=\"black\", linewidth=4)\n",
    "    plt.scatter(*x_nom[k], color=\"green\", s=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "automatic-sydney",
   "metadata": {},
   "source": [
    "## Tracking LQR for Nonlinear Dynamics\n",
    "Minimizing a quadratic tracking cost:\n",
    "\n",
    "$\n",
    "J_k(\\mathbf{\\delta x}_k) = \\frac{1}{2}Q_N \\|\\mathbf{\\delta x}_N\\|^2 + \\frac{1}{2}\\sum_{i=k}^{N-1} \\left(\\|\\mathbf{\\delta x}_k\\|^2 + \\|\\mathbf{\\delta u}_k\\|^2\\right)\n",
    "$\n",
    "\n",
    "subject to nonlinear dynamics:\n",
    "\n",
    "$\\begin{bmatrix}p_{k+1} \\\\ v_{k+1}\\end{bmatrix} =\n",
    "\\begin{bmatrix}1 & \\Delta t \\\\ 0 & 1\\end{bmatrix} \\begin{bmatrix}p_k \\\\ v_k\\end{bmatrix} +\n",
    "\\begin{bmatrix}\\frac{1}{2}\\Delta t^2 p_k \\\\ \\Delta t\\end{bmatrix} \\begin{bmatrix}u_k\\end{bmatrix}\n",
    "$\n",
    "\n",
    "tracking a sinusoidal nominal trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "injured-alcohol",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simple_nonlinear_dynamics(x, u, k):\n",
    "    p, v = x\n",
    "    a, = u\n",
    "    return jnp.array([p + v * dt + p * a * dt**2 / 2, v + a * dt])\n",
    "\n",
    "x_nom, u_nom = rollout_state_feedback_policy(simple_nonlinear_dynamics, lambda x, k: -x[:1], np.array([0., 2.]), np.arange(N))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "noble-bookmark",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10\n",
    "dt = 0.5\n",
    "Q_N = 1  # Change to 100!\n",
    "\n",
    "Qs = np.broadcast_to(np.eye(2), (N + 1, 2, 2)).copy()\n",
    "Qs[-1] *= Q_N\n",
    "Rs = np.broadcast_to(np.eye(1), (N, 1, 1))\n",
    "Hs = np.broadcast_to(np.zeros((1, 2)), (N, 1, 2))\n",
    "As, Bs = jax.vmap(jax.jacobian(simple_nonlinear_dynamics, (0, 1)))(x_nom[:-1], u_nom, np.arange(10))\n",
    "value_functions, optimal_policy = linear_quadratic_regulator(Qs, Rs, Hs, As, Bs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pointed-property",
   "metadata": {},
   "outputs": [],
   "source": [
    "X0, X1 = np.meshgrid(np.linspace(-8, 5, 100), np.linspace(-5, 5, 100))\n",
    "X = np.stack([X0, X1], -1)\n",
    "all_values = jax.vmap(lambda v, x: jax.vmap(jax.vmap(v))(X - x))(value_functions, x_nom)\n",
    "\n",
    "dynamics = LinearDynamics(As, Bs)\n",
    "x0s = np.random.uniform(-4, 4, (20, 2))\n",
    "example_trajectories = [\n",
    "    jax.vmap(lambda x0: rollout_state_feedback_policy(\n",
    "        simple_nonlinear_dynamics,\n",
    "        optimal_policy,\n",
    "        x0,\n",
    "        np.arange(k, N),\n",
    "        x_nom,\n",
    "        u_nom,\n",
    "    ))(x0s)[0] for k in range(N)\n",
    "]\n",
    "\n",
    "\n",
    "@interact(k=(0, N))\n",
    "def plot(k=0):\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    plt.contour(X0, X1, all_values[k], levels=40)\n",
    "    plt.colorbar()\n",
    "    plt.xlim(-8, 5)\n",
    "    plt.ylim(-5, 5)\n",
    "    plt.xlabel(\"Position $p$\")\n",
    "    plt.ylabel(\"Velocity $v$\")\n",
    "    plt.title(f\"(Locally Accurate) Value Function at Step {k}\")\n",
    "\n",
    "    if k < N:\n",
    "        plt.scatter(*x0s.T, color=\"red\")\n",
    "        plt.plot(*example_trajectories[k].T, color=\"grey\", alpha=0.5)\n",
    "        plt.plot(*example_trajectories[k][:, :2].T, color=\"red\")\n",
    "\n",
    "    plt.plot(*x_nom.T, color=\"black\", linewidth=4)\n",
    "    plt.scatter(*x_nom[k], color=\"green\", s=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "casual-exhibition",
   "metadata": {},
   "source": [
    "## Iterative LQR\n",
    "Minimizing a nonconvex cost:\n",
    "\n",
    "$\n",
    "    J_k(\\mathbf{x}_k) = 10 \\|\\mathbf{x}_N - [-4, 0]^T\\|^2 + \\sum_{i=k}^{N-1} \\left(\\sin^2(v_k\\pi) + \\|\\mathbf{u}_k\\|^2\\right)\n",
    "$\n",
    "\n",
    "subject to nonlinear dynamics:\n",
    "\n",
    "$\\begin{bmatrix}p_{k+1} \\\\ v_{k+1}\\end{bmatrix} =\n",
    "\\begin{bmatrix}1 & \\Delta t \\\\ 0 & 1\\end{bmatrix} \\begin{bmatrix}p_k \\\\ v_k\\end{bmatrix} +\n",
    "\\begin{bmatrix}\\frac{1}{2}\\Delta t^2 p_k \\\\ \\Delta t\\end{bmatrix} \\begin{bmatrix}u_k\\end{bmatrix}.\n",
    "$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "graphic-edward",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ensure_positive_definite(a, eps=1e-3):\n",
    "    w, v = jnp.linalg.eigh(a)\n",
    "    return (v * jnp.maximum(w, eps)) @ v.T\n",
    "\n",
    "\n",
    "class TotalCost(NamedTuple):\n",
    "    running_cost: Callable\n",
    "    terminal_cost: Callable\n",
    "\n",
    "    def __call__(self, xs, us):\n",
    "        step_range = jnp.arange(us.shape[0])\n",
    "        return jnp.sum(jax.vmap(self.running_cost)(xs[:-1], us, step_range)) + self.terminal_cost(xs[-1])\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def iterative_linear_quadratic_regulator(dynamics, total_cost, x0, u_guess, maxiter=100, atol=1e-3):\n",
    "    running_cost, terminal_cost = total_cost\n",
    "    n, (N, m) = x0.shape[-1], u_guess.shape\n",
    "    step_range = jnp.arange(N)\n",
    "\n",
    "    xs_iterates, us_iterates = jnp.zeros((maxiter, N + 1, n)), jnp.zeros((maxiter, N, m))\n",
    "    xs, us = rollout_state_feedback_policy(dynamics, lambda x, k: u_guess[k], x0, step_range)\n",
    "    xs_iterates, us_iterates = xs_iterates.at[0].set(xs), us_iterates.at[0].set(us)\n",
    "    j_curr = total_cost(xs, us)\n",
    "    value_functions_iterates = QuadraticStateCost.from_pure_quadratic(jnp.zeros((maxiter, N + 1, n, n)))\n",
    "\n",
    "    def continuation_criterion(loop_vars):\n",
    "        i, _, _, j_curr, j_prev, _ = loop_vars\n",
    "        return (j_curr < j_prev - atol) & (i < maxiter)\n",
    "\n",
    "    def ilqr_iteration(loop_vars):\n",
    "        i, xs_iterates, us_iterates, j_curr, j_prev, value_functions_iterates = loop_vars\n",
    "        xs, us = xs_iterates[i], us_iterates[i]\n",
    "\n",
    "        f_x, f_u = jax.vmap(jax.jacobian(dynamics, (0, 1)))(xs[:-1], us, step_range)\n",
    "        c = jax.vmap(running_cost)(xs[:-1], us, step_range)\n",
    "        c_x, c_u = jax.vmap(jax.grad(running_cost, (0, 1)))(xs[:-1], us, step_range)\n",
    "        (c_xx, c_xu), (c_ux, c_uu) = jax.vmap(jax.hessian(running_cost, (0, 1)))(xs[:-1], us, step_range)\n",
    "        v, v_x, v_xx = terminal_cost(xs[-1]), jax.grad(terminal_cost)(xs[-1]), jax.hessian(terminal_cost)(xs[-1])\n",
    "\n",
    "        # Ensure quadratic cost terms are positive definite.\n",
    "        c_zz = jnp.block([[c_xx, c_xu], [c_ux, c_uu]])\n",
    "        c_zz = jax.vmap(ensure_positive_definite)(c_zz)\n",
    "        c_xx, c_uu, c_ux = c_zz[:, :n, :n], c_zz[:, -m:, -m:], c_zz[:, -m:, :n]\n",
    "        v_xx = ensure_positive_definite(v_xx)\n",
    "\n",
    "        linearized_dynamics = LinearDynamics(f_x, f_u)\n",
    "        quadratized_running_cost = QuadraticCost(c, c_x, c_u, c_xx, c_uu, c_ux)\n",
    "        quadratized_terminal_cost = QuadraticStateCost(v, v_x, v_xx)\n",
    "\n",
    "        def scan_fn(next_state_value, current_step_dynamics_cost):\n",
    "            current_step_dynamics, current_step_cost = current_step_dynamics_cost\n",
    "            current_state_value, current_step_policy = riccati_step(\n",
    "                current_step_dynamics,\n",
    "                current_step_cost,\n",
    "                next_state_value,\n",
    "            )\n",
    "            return current_state_value, (current_state_value, current_step_policy)\n",
    "\n",
    "        value_functions, policy = jax.lax.scan(scan_fn,\n",
    "                                               quadratized_terminal_cost,\n",
    "                                               (linearized_dynamics, quadratized_running_cost),\n",
    "                                               reverse=True)[1]\n",
    "        value_functions_iterates = jax.tree.map(lambda x, xi, xiN: x.at[i].set(jnp.concatenate([xi, xiN[None]])),\n",
    "                                                value_functions_iterates, value_functions, quadratized_terminal_cost)\n",
    "\n",
    "        def rollout_linesearch_policy(alpha):\n",
    "            # Note that we roll out the true `dynamics`, not the `linearized_dynamics`!\n",
    "            l, l_x = policy\n",
    "            return rollout_state_feedback_policy(dynamics, AffinePolicy(alpha * l, l_x), x0, step_range, xs, us)\n",
    "\n",
    "        # Backtracking line search (step sizes evaluated in parallel).\n",
    "        all_xs, all_us = jax.vmap(rollout_linesearch_policy)(0.5**jnp.arange(16))\n",
    "        js = jax.vmap(total_cost)(all_xs, all_us)\n",
    "        a = jnp.argmin(js)\n",
    "        j = js[a]\n",
    "        xs_iterates = xs_iterates.at[i + 1].set(jnp.where(j < j_curr, all_xs[a], xs))\n",
    "        us_iterates = us_iterates.at[i + 1].set(jnp.where(j < j_curr, all_us[a], us))\n",
    "        return i + 1, xs_iterates, us_iterates, jnp.minimum(j, j_curr), j_curr, value_functions_iterates\n",
    "\n",
    "    i, xs_iterates, us_iterates, j_curr, j_prev, value_functions_iterates = jax.lax.while_loop(\n",
    "        continuation_criterion, ilqr_iteration,\n",
    "        (0, xs_iterates, us_iterates, j_curr, jnp.inf, value_functions_iterates))\n",
    "\n",
    "    return {\n",
    "        \"optimal_trajectory\": (xs_iterates[i], us_iterates[i]),\n",
    "        \"optimal_cost\": j_curr,\n",
    "        \"num_iterations\": i,\n",
    "        \"trajectory_iterates\": (xs_iterates, us_iterates),\n",
    "        \"value_functions_iterates\": value_functions_iterates\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "political-reputation",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleNonlinearDynamics(NamedTuple):\n",
    "    dt: float = 0.5\n",
    "\n",
    "    def __call__(self, x, u, k):\n",
    "        p, v = x\n",
    "        a, = u\n",
    "        return jnp.array([p + v * self.dt + p * a * self.dt**2 / 2, v + a * self.dt])\n",
    "\n",
    "\n",
    "class SimpleRunningCost(NamedTuple):\n",
    "    gain: float = 1.0\n",
    "\n",
    "    def __call__(self, x, u, k):\n",
    "        return self.gain * (jnp.sin(x[1] * np.pi)**2 + u[0]**2)\n",
    "\n",
    "\n",
    "class SimpleTerminalCost(NamedTuple):\n",
    "    gain: float = 10.0\n",
    "    target: jnp.array = jnp.array([-4., 0.])\n",
    "\n",
    "    def __call__(self, x):\n",
    "        return self.gain * jnp.sum(jnp.square(x - self.target))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "missing-inclusion",
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact(terminal_cost_gain=(1, 100))\n",
    "def plot(terminal_cost_gain=10.0):\n",
    "    N = 10\n",
    "    x0 = jnp.array([4., 0.])\n",
    "    u_guess = jnp.zeros((N, 1))\n",
    "    solution = iterative_linear_quadratic_regulator(\n",
    "        SimpleNonlinearDynamics(),\n",
    "        TotalCost(SimpleRunningCost(), SimpleTerminalCost(gain=terminal_cost_gain)),\n",
    "        x0,\n",
    "        u_guess,\n",
    "    )\n",
    "    i = solution[\"num_iterations\"]\n",
    "\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    plt.xlim(-5, 5)\n",
    "    plt.ylim(-3, 1)\n",
    "    plt.xlabel(\"Position $p$\")\n",
    "    plt.ylabel(\"Velocity $v$\")\n",
    "    for xs, c in zip(solution[\"trajectory_iterates\"][0][:i], np.linspace([1, 0, 0, 1], [0, 1, 0, 1], i)):\n",
    "        plt.plot(*xs.T, color=c)\n",
    "    plt.plot(*solution[\"optimal_trajectory\"][0].T, color=\"black\", linewidth=4, marker=\"o\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "specific-retention",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10\n",
    "x0 = jnp.array([4., 0.])\n",
    "u_guess = jnp.zeros((N, 1))\n",
    "solution = iterative_linear_quadratic_regulator(\n",
    "    SimpleNonlinearDynamics(),\n",
    "    TotalCost(SimpleRunningCost(gain=0.0), SimpleTerminalCost()),\n",
    "    x0,\n",
    "    u_guess,\n",
    ")\n",
    "i = solution[\"num_iterations\"]\n",
    "\n",
    "X0, X1 = np.meshgrid(np.linspace(-5, 5, 100), np.linspace(-3, 1, 100))\n",
    "X = np.stack([X0, X1], -1)\n",
    "all_values = jax.vmap(jax.vmap(lambda v, x: jax.vmap(jax.vmap(v))(X - x)))(solution[\"value_functions_iterates\"][:i],\n",
    "                                                                           solution[\"trajectory_iterates\"][0][:i])\n",
    "\n",
    "\n",
    "@interact(j=(0, int(i) - 1), k=(0, N))\n",
    "def plot(j=0, k=0):\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    plt.contour(X0, X1, all_values[j, k], levels=40)\n",
    "    plt.colorbar()\n",
    "    plt.xlim(-5, 5)\n",
    "    plt.ylim(-3, 1)\n",
    "    plt.xlabel(\"Position $p$\")\n",
    "    plt.ylabel(\"Velocity $v$\")\n",
    "    plt.title(f\"(Locally Accurate) Value Function at iLQR Iterate {j}, Step {k}\")\n",
    "\n",
    "    plt.plot(*solution[\"trajectory_iterates\"][0][j].T, color=\"black\", linewidth=4, marker=\"o\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "AA203",
   "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.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "cf32138a9007cd0f1b0fcd1f4c0b99886618bc4804f4ff5e19b02064504ee0d6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}