Koopman RL 5-5.ipynb

{
 "metadata": {
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "source": [
    "import observables\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import integrate\n",
    "from control import lqr"
   ],
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "source": [
    "# Slow Manifold Example (KOOC/KRONIC)\n",
    "\n",
    "In order to test our iteration methodology, we wanted to explore a simple system where the Koopman operator is known\n",
    "\n",
    "In this section, we study a deterministic dynamical system where we can explicitly find a finite dimensional linear embedding for the Koopman operator. This example comes from Steve's KRONIC paper. In particular, we consider a controlled system with quadratic nonlinearity that gives rise to a slow manifold:\n",
    "\n",
    "\\begin{align}\n",
    "    \\frac{d}{dt}\n",
    "    \\begin{bmatrix}\n",
    "        x_1\\\\\n",
    "        x_2\n",
    "    \\end{bmatrix} = \n",
    "    \\begin{bmatrix}\n",
    "        \\mu x_1\n",
    "        \\\\\n",
    "        \\lambda (x_2 - x_1^2)\n",
    "    \\end{bmatrix}\n",
    "    + \\mathbf{D} \\ u\n",
    "\\end{align}\n",
    "Considering a transformation to a Koopman invariant subspace, $(y_1,y_2,y_3) = (x_1,x_2,x_1^2)$, we can express the system dynamics as\n",
    "\\begin{align}\n",
    "    \\frac{d}{dt}\n",
    "    \\begin{bmatrix}\n",
    "        y_1\\\\\n",
    "        y_2\\\\\n",
    "        y_3\n",
    "    \\end{bmatrix} = \n",
    "    \\begin{bmatrix}\n",
    "        \\mu & 0 & 0\n",
    "        \\\\\n",
    "        0 & \\lambda & -\\lambda\n",
    "        \\\\\n",
    "        0 & 0 & 2\\mu\n",
    "    \\end{bmatrix}\n",
    "    \\begin{bmatrix}\n",
    "        y_1\\\\\n",
    "        y_2\\\\\n",
    "        y_3\n",
    "    \\end{bmatrix}\n",
    "    +\n",
    "    \\underbrace{\\begin{bmatrix}\n",
    "        1 & 0\\\\\n",
    "        0 & 1\\\\\n",
    "        2y_1 & 0\n",
    "    \\end{bmatrix}\n",
    "    \\mathbf{D}}_{D_y} \\ u\n",
    "\\end{align}"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "source": [
    "## Definitions:"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu = -0.1\n",
    "lamb = 1\n",
    "tspan = np.arange(0, 50, 0.01)\n",
    "x0 = np.array(\n",
    "    [[-5],\n",
    "     [5]]\n",
    ")\n",
    "\n",
    "A = np.array(\n",
    "    [[mu, 0   ],\n",
    "     [0,  lamb]]\n",
    ")\n",
    "B = np.array(\n",
    "    [[0],\n",
    "     [1]]\n",
    ")\n",
    "Q = np.identity(2)\n",
    "R = 1\n",
    "\n",
    "A2 = np.array(\n",
    "    [[mu, 0,    0    ],\n",
    "     [0,  lamb, -lamb],\n",
    "     [0,  0,    2*mu ]]\n",
    ")\n",
    "B2 = np.array(\n",
    "    [[0],\n",
    "     [1],\n",
    "     [0]]\n",
    ")\n",
    "Q2 = np.array(\n",
    "    [[1, 0, 0],\n",
    "     [0, 1, 0],\n",
    "     [0, 0, 0]]\n",
    ")"
   ]
  },
  {
   "source": [
    "## Calculate LQR controller\n",
    "C is supposed to be \\[0.0, 2.4142\\] and C2 is supposed to be \\[0.0, 2.4142, -1.4956\\]"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C = lqr(A, B, Q, R)[0][0]\n",
    "vf = lambda tau, x: A @ x + np.array([0, -lamb * x[0]**2]) + B.T * (-C @ x)\n",
    "xLQR = integrate.solve_ivp(vf, (0,50), x0[:,0], first_step=0.01, max_step=0.01)\n",
    "xLQR = xLQR.y[:,:-2]\n",
    "\n",
    "C2 = lqr(A2, B2, Q2, R)[0][0]\n",
    "vf2 = lambda tau, x: A @ x + np.array([0, -lamb * x[0]**2]) + B.T * ((-C2[0:2] @ x) - (C2[2] * x[0]**2))\n",
    "xKOOC = integrate.solve_ivp(vf2, (0,50), x0[:,0], first_step=0.01, max_step=0.01)\n",
    "xKOOC = xKOOC.y[:,:-2]\n",
    "\n",
    "JLQR = np.cumsum(xLQR[0]**2 + xLQR[1]**2 + (C @ np.conjugate(xLQR))**2)\n",
    "JKOOC = np.cumsum(xKOOC[0]**2 + xKOOC[1]**2 + (C @ np.conjugate(xKOOC))**2)"
   ]
  },
  {
   "source": [
    "## Plots"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('dark_background')\n",
    "\n",
    "fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(18, 6))\n",
    "\n",
    "ax1.plot(xLQR[0], xLQR[1], 'w-')\n",
    "ax1.plot(xKOOC[0], xKOOC[1], 'r--')\n",
    "ax1.grid(b=True, which='major', color='#666666', linestyle='-')\n",
    "ax1.set_xlabel('x_1')\n",
    "ax1.set_ylabel('x_2')\n",
    "\n",
    "ax2.set_xlim(0, 50)\n",
    "ax2.plot(tspan, xLQR.T, 'w-')\n",
    "ax2.plot(tspan, xKOOC.T, 'r--')\n",
    "ax2.grid(b=True, which='major', color='#666666', linestyle='-')\n",
    "ax2.set_xlabel('t')\n",
    "ax2.set_ylabel('x_k')\n",
    "\n",
    "ax3.set_xlim(0, 50)\n",
    "ax3.plot(tspan, JLQR, 'w-')\n",
    "ax3.plot(tspan, JKOOC, 'r--')\n",
    "ax3.grid(b=True, which='major', color='#666666', linestyle='-')\n",
    "ax3.set_xlabel('t')\n",
    "ax3.set_ylabel('J')\n",
    "\n",
    "plt.subplots_adjust(wspace=0.5)\n",
    "plt.show()"
   ]
  },
  {
   "source": [
    "## Koopman MDP Approach\n",
    "\n",
    "Let $\\mathbf{F}\\ y = x$,\n",
    "\\begin{align}\n",
    "    \\mathbf{F} = \\begin{bmatrix}\n",
    "        1 & 0 & 0\\\\\n",
    "        0 & 1 & 0\n",
    "    \\end{bmatrix}\n",
    "    \\implies\n",
    "    \\frac{\\text{d}x}{\\text{d}t} = \\mathbf{E}\\ \\frac{\\text{d}y}{\\text{d}t}\n",
    "\\end{align}\n",
    "Note: $\\exists$ linear transformation $y \\to x$ but not $x \\to y$\n",
    "\n",
    "Now looking at the continuation term of the HJB equation:\n",
    "\\begin{align}\n",
    "    \\mathcal{L} V = \\dot{V} = \\frac{\\text{d}V}{\\text{d}x}\\ \\frac{\\text{d}x}{\\text{d}t} = \\nabla_x V\\ \\mathbf{F}\\ \\frac{\\text{d}y}{\\text{d}t}\n",
    "\\end{align}\n",
    "In order to characterize $\\dot{V}$, we must then derive an expression for $\\frac{\\text{d}V}{\\text{d}x}$. Assuming that $V$ is a linear combination of the Koopman invariant subspace $y$:\n",
    "\\begin{align}\n",
    "    V = \\mathbf{B}^\\top\\ y \\implies \\nabla_x V &= \\nabla_x \\mathbf{B}^\\top\\ y\\\\\n",
    "    &= \\nabla_x (y^\\top\\ \\mathbf{B})\\\\\n",
    "    &= J_y (x_1, x_2)\\ \\mathbf{B}\\\\\n",
    "    &= \\begin{bmatrix}\n",
    "        \\frac{\\partial y_1}{\\partial x_1} & \\frac{\\partial y_2}{\\partial x_1} & \\frac{\\partial y_3}{\\partial x_1}\\\\\n",
    "        \\frac{\\partial y_1}{\\partial x_2} & \\frac{\\partial y_2}{\\partial x_2} & \\frac{\\partial y_3}{\\partial x_2}\n",
    "    \\end{bmatrix}\n",
    "    \\ \\mathbf{B}\n",
    "\\end{align}\n",
    "\n",
    "If it is correct to think statically (holding time fixed), then the Jacobian matrix is not hard to calculate and is given by\n",
    "\\begin{align*}\n",
    "    J_y (x_1, x_2) =\n",
    "    \\begin{bmatrix}\n",
    "        \\frac{\\partial y_1}{\\partial x_1} & \\frac{\\partial y_2}{\\partial x_1} & \\frac{\\partial y_3}{\\partial x_1}\\\\\n",
    "        \\frac{\\partial y_1}{\\partial x_2} & \\frac{\\partial y_2}{\\partial x_2} & \\frac{\\partial y_3}{\\partial x_2}\n",
    "    \\end{bmatrix}\n",
    "    &= \\begin{bmatrix}\n",
    "        \\mu & 0 & 2x_1\\\\\n",
    "        0 & 1 & 0\n",
    "    \\end{bmatrix}\n",
    "\\end{align*}\n",
    "\n",
    "Because LQR has no discounting and it is framed as a minimization problem, the HJB becomes\n",
    "\\begin{align}\n",
    "    0 &=  \\inf_{\\pi_t \\in \\mathcal{P}(U)} \\underbrace{\n",
    "    \\int_U \\left( r(x,u) + \\lambda \\ln{\\pi_t(u)} + (\\mathcal{L}v)(x,u)\\right) \\pi_t(u) \\text{d}u}_{H\\left(x,\\pi(x),\\mathcal{L}V\\right)} \\label{HJB_LQR}\n",
    "\\end{align}\n",
    "\n",
    "where $r(x,u) = x^\\top Q x + u^\\top Ru$ and the function $H$ denotes the Hamiltonian to be minimized. The optimal policy is then\n",
    "\\begin{align}\n",
    "    \\pi(x) = \\text{argmin}_{\\pi} H\\left(x,\\pi(x),\\mathcal{L}V\\right)\n",
    "\\end{align}\n",
    "the solution of which is\n",
    "\\begin{align}\n",
    "    \\pi^*(u|x) &= \\frac{\\exp\\left(-\\frac{1}{\\lambda}(r(x,u) + (\\mathcal{L}v)(x,u))\\right)}{\\int_U \\exp\\left(-\\frac{1}{\\lambda}(r(x,u) + (\\mathcal{L}v)(x,u))\\right)du}\\label{opt_pol_LQR}\n",
    "\\end{align}\n",
    "We can use the DP relationship to update the value function as follows\n",
    "\\begin{align}\n",
    "    V_{j+1}(x(t)) = \\mathbb{E}_{u\\sim\\pi}\\left(\\int_{t}^{t+\\tau} r(x(s),u)ds\\right) + V_j(x(t+\\tau))\n",
    "\\end{align}\n",
    "Expressing this in terms of the Koopman operator\n",
    "\\begin{align}\n",
    "    V_{j+1}(x(t)) = \\mathbb{E}_{u\\sim\\pi}\\left(\\int_{t}^{t+\\tau} r(x(s),u)ds\\right) + \\mathcal{K}_\\tau V_j(x(t))\n",
    "\\end{align}\n",
    "Using the relationship between the Koopman operator and its generator $\\mathcal{K}_{\\tau} = \\exp(\\tau\\mathcal{L})$\n",
    "\\begin{align}\n",
    "    V_{j+1}(x(t)) &= \\mathbb{E}_{u\\sim\\pi}\\left(\\int_{t}^{t+\\tau} r(x(s),u)ds\\right) + \\exp(\\tau\\mathcal{L}) V_j(x(t))\n",
    "    \\\\\n",
    "    &= \\int_U\\left(\\int_{t}^{t+\\tau} r(x(s),u)ds\\right)\\pi(u)du + \\exp(\\tau\\mathcal{L}) V_j(x(t))\n",
    "    \\\\\n",
    "     &= \\int_{t}^{t+\\tau}\\left( \\int_U r(x(s),u)\\pi(u)du\\right)ds + \\exp(\\tau\\mathcal{L}) V_j(x(t))\n",
    "\\end{align}\n",
    "\n",
    "Specializing the reward to the special case of LQR $r(x,u) = x^\\top Q x + u^\\top Ru$ \n",
    "\\begin{align}\n",
    "    V_{j+1}(x(t))\n",
    "     &= \\int_{t}^{t+\\tau}x(t)^\\top Q x(t)ds + \\tau\\int_U  u^\\top Ru \\;\\pi(u)du + \\exp(\\tau\\mathcal{L}) V_j(x(t))\n",
    "\\end{align}\n",
    "Furthermore, if we take the control to be univariate\n",
    "\\begin{align}\n",
    "    V_{j+1}(x(t))\n",
    "     &= \\int_{t}^{t+\\tau}x(t)^\\top Q x(t)ds + \\tau \\mathbb{E}_{u\\sim\\pi}(Ru^2) + \\exp(\\tau\\mathcal{L}) V_j(x(t)) \\label{BellmanEq}\n",
    "\\end{align}\n",
    "Next, we outline an algorithm for value function iteration in this case\n",
    "\n",
    "\\begin{enumerate}\n",
    "    \\item Simulate states by taking (for now) the optimal policy of $u$ from the KOOC approach as given and using the dynamics of x. This gives us samples $\\bf{x_1,...,x_m}$ and $u_1,...,u_{m-1}$. In addition, we also have observations $\\bf{y_1,...,y_m}$.\n",
    "    \n",
    "    \\item Guess initial value function and evaluate it at sample points\n",
    "    \n",
    "    \\item Use least squares to project onto dictionary space (which is richer than the minimal function space that spans the dynamics). Denote this matrix by $\\hat B^\\top$.\n",
    "    \n",
    "    \\item Get the unforced Koopman operator of ${\\bf y}=(x_1,x_2,x_1^2)$\n",
    "    \n",
    "    \\item Calculate $\\hat{\\mathcal{L}} V = \\dot{V} = \\frac{\\text{d}V}{\\text{d}x}\\ \\frac{\\text{d}x}{\\text{d}t} = \\nabla_x \\psi(x)^\\top \\hat B\\ \\mathbf{F}\\ \\frac{\\text{d}y}{\\text{d}t}$\n",
    "    \n",
    "    \\item Find opt policy $\\pi^*$ using \\eqref{opt_pol_LQR}\n",
    "    \n",
    "    \\item Calculate next value function by plugging in optimal policy and the current value function into the Bellman equation \\eqref{BellmanEq}\n",
    "\\end{enumerate}\n",
    "Unfortunately, dealing with the numerical issues with integration is proving to be difficult, especially when evaluating the optimal policy on each iteration and integrating wrt it."
   ],
   "cell_type": "markdown",
   "metadata": {}
  }
 ]
}