{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c017196f",
   "metadata": {},
   "source": [
    "# PVTOL LQR + EQF example\n",
    "RMM, 14 Feb 2022\n",
    "\n",
    "This notebook illustrates the implementation of an extended Kalman filter and the use of the estimated state for LQR feedback."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "544525ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "import control as ct"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "859834cf",
   "metadata": {},
   "source": [
    "## System definition\n",
    "The dynamics of the system\n",
    "with disturbances on the $x$ and $y$ variables is given by\n",
    "$$\n",
    "  \\begin{aligned}\n",
    "    m \\ddot x &= F_1 \\cos\\theta - F_2 \\sin\\theta - c \\dot x + d_x, \\\\\n",
    "    m \\ddot y &= F_1 \\sin\\theta + F_2 \\cos\\theta - c \\dot y - m g + d_y, \\\\\n",
    "    J \\ddot \\theta &= r F_1.\n",
    "  \\end{aligned}\n",
    "$$\n",
    "The measured values of the system are the position and orientation,\n",
    "with added noise $n_x$, $n_y$, and $n_\\theta$:\n",
    "$$\n",
    "  \\vec y = \\begin{bmatrix} x \\\\ y \\\\ \\theta \\end{bmatrix} + \n",
    "  \\begin{bmatrix} n_x \\\\ n_y \\\\ n_z \\end{bmatrix}.\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffafed74",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pvtol = nominal system (no disturbances or noise)\n",
    "# noisy_pvtol = pvtol w/ process disturbances and sensor noise\n",
    "from pvtol import pvtol, pvtol_noisy, plot_results\n",
    "\n",
    "# Find the equiblirum point corresponding to the origin\n",
    "xe, ue = ct.find_eqpt(\n",
    "    pvtol, np.zeros(pvtol.nstates),\n",
    "    np.zeros(pvtol.ninputs), [0, 0, 0, 0, 0, 0],\n",
    "    iu=range(2, pvtol.ninputs), iy=[0, 1])\n",
    "\n",
    "x0, u0 = ct.find_eqpt(\n",
    "    pvtol, np.zeros(pvtol.nstates),\n",
    "    np.zeros(pvtol.ninputs), np.array([2, 1, 0, 0, 0, 0]),\n",
    "    iu=range(2, pvtol.ninputs), iy=[0, 1])\n",
    "\n",
    "# Extract the linearization for use in LQR design\n",
    "pvtol_lin = pvtol.linearize(xe, ue)\n",
    "A, B = pvtol_lin.A, pvtol_lin.B\n",
    "\n",
    "print(pvtol, \"\\n\")\n",
    "print(pvtol_noisy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b63bf5b",
   "metadata": {},
   "source": [
    "We now define the properties of the noise and disturbances.  To make things (a bit more) interesting, we include some cross terms between the noise in $\\theta$ and the noise in $x$ and $y$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e1ee7c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disturbance and noise intensities\n",
    "Qv = np.diag([1e-2, 1e-2])\n",
    "Qw = np.array([[2e-4, 0, 1e-5], [0, 2e-4, 1e-5], [1e-5, 1e-5, 1e-4]])\n",
    "Qwinv = np.linalg.inv(Qw)\n",
    "\n",
    "# Initial state covariance\n",
    "P0 = np.eye(pvtol.nstates)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4c52c73",
   "metadata": {},
   "source": [
    "## Control system design\n",
    "\n",
    "To design the control system, we first construct an estimator for the state (given the commanded inputs and measured outputs.  Since this is a nonlinear system, we use the update law for the nominal system to compute the state update.  We also make use of the linearization around the current state for the covariance update (using the function `pvtol.A(x, u)`, which is defined in `pvtol.py`, making this an extended Kalman filter)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3647bf15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the disturbance input and measured output matrices\n",
    "F = np.array([[0, 0], [0, 0], [0, 0], [1/pvtol.params['m'], 0], [0, 1/pvtol.params['m']], [0, 0]])\n",
    "C = np.eye(3, 6)\n",
    "\n",
    "# Estimator update law\n",
    "def estimator_update(t, x, u, params):\n",
    "    # Extract the states of the estimator\n",
    "    xhat = x[0:pvtol.nstates]\n",
    "    P = x[pvtol.nstates:].reshape(pvtol.nstates, pvtol.nstates)\n",
    "\n",
    "    # Extract the inputs to the estimator\n",
    "    y = u[0:3]                  # just grab the first three outputs\n",
    "    u = u[6:8]                  # get the inputs that were applied as well\n",
    "\n",
    "    # Compute the linearization at the current state\n",
    "    A = pvtol.A(xhat, u)        # A matrix depends on current state\n",
    "    # A = pvtol.A(xe, ue)       # Fixed A matrix (for testing/comparison)\n",
    "    \n",
    "    # Compute the optimal again\n",
    "    L = P @ C.T @ Qwinv\n",
    "\n",
    "    # Update the state estimate\n",
    "    xhatdot = pvtol.updfcn(t, xhat, u, params) - L @ (C @ xhat - y)\n",
    "\n",
    "    # Update the covariance\n",
    "    Pdot = A @ P + P @ A.T - P @ C.T @ Qwinv @ C @ P + F @ Qv @ F.T\n",
    "\n",
    "    # Return the derivative\n",
    "    return np.hstack([xhatdot, Pdot.reshape(-1)])\n",
    "\n",
    "def estimator_output(t, x, u, params):\n",
    "    # Return the estimator states\n",
    "    return x[0:pvtol.nstates]\n",
    "\n",
    "estimator = ct.NonlinearIOSystem(\n",
    "    estimator_update, estimator_output,\n",
    "    states=pvtol.nstates + pvtol.nstates**2,\n",
    "    inputs= pvtol_noisy.output_labels \\\n",
    "        + pvtol_noisy.input_labels[0:pvtol.ninputs],\n",
    "    outputs=[f'xh{i}' for i in range(pvtol.nstates)],\n",
    ")\n",
    "print(estimator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba3d2640",
   "metadata": {},
   "source": [
    "For the controller, we will use an LQR feedback with physically motivated weights (see OBC, Example 3.5):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9787db61",
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# LQR design w/ physically motivated weighting\n",
    "#\n",
    "# Shoot for 1 cm error in x, 10 cm error in y.  Try to keep the angle\n",
    "# less than 5 degrees in making the adjustments.  Penalize side forces\n",
    "# due to loss in efficiency.\n",
    "#\n",
    "\n",
    "Qx = np.diag([100, 10, (180/np.pi) / 5, 0, 0, 0])\n",
    "Qu = np.diag([10, 1])\n",
    "K, _, _ = ct.lqr(A, B, Qx, Qu)\n",
    "\n",
    "#\n",
    "# Control system construction: combine LQR w/ EKF\n",
    "#\n",
    "# Use the linearization around the origin to design the optimal gains\n",
    "# to see how they compare to the final value of P for the EKF\n",
    "#\n",
    "\n",
    "# Construct the state feedback controller with estimated state as input\n",
    "statefbk, _ = ct.create_statefbk_iosystem(pvtol, K, estimator=estimator)\n",
    "print(statefbk, \"\\n\")\n",
    "\n",
    "# Reconstruct the control system with the noisy version of the process\n",
    "# Create a closed loop system around the controller\n",
    "clsys = ct.interconnect(\n",
    "    [pvtol_noisy, statefbk, estimator],\n",
    "    inplist = statefbk.input_labels[0:pvtol.ninputs + pvtol.nstates] + \\\n",
    "        pvtol_noisy.input_labels[pvtol.ninputs:],\n",
    "    inputs = statefbk.input_labels[0:pvtol.ninputs + pvtol.nstates] + \\\n",
    "        pvtol_noisy.input_labels[pvtol.ninputs:],\n",
    "    outlist = pvtol.output_labels + statefbk.output_labels + estimator.output_labels,\n",
    "    outputs = pvtol.output_labels + statefbk.output_labels + estimator.output_labels\n",
    ")\n",
    "print(clsys)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f527f16",
   "metadata": {},
   "source": [
    "Note that we have to construct the closed loop system manually since we need to allow the disturbance and noise inputs to be sent to the closed loop system and `create_statefbk_iosystem` does not support this (to be fixed in an upcoming release)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bf558a0",
   "metadata": {},
   "source": [
    "## Simulations\n",
    "\n",
    "Finally, we can simulate the system to see how it all works.  We start by creating the noise for the system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2583a0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the time vector for the simulation\n",
    "Tf = 10\n",
    "timepts = np.linspace(0, Tf, 1000)\n",
    "\n",
    "# Create representative process disturbance and sensor noise vectors\n",
    "np.random.seed(117)           # avoid figures changing from run to run\n",
    "V = ct.white_noise(timepts, Qv) # smaller disturbances and noise then design\n",
    "W = ct.white_noise(timepts, Qw)\n",
    "plt.plot(timepts, V[0], label=\"V[0]\")\n",
    "plt.plot(timepts, W[0], label=\"W[0]\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d944709",
   "metadata": {},
   "source": [
    "### LQR with EKF\n",
    "\n",
    "We can now feed the desired trajectory plus the noise and disturbances into the system and see how well the controller with a state estimator does in holding the system at an equilibrium point:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad7a9750",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Put together the input for the system\n",
    "U = [xe, ue, V, W]\n",
    "X0 = [x0, xe, P0.reshape(-1)]\n",
    "\n",
    "# Initial condition response\n",
    "resp = ct.input_output_response(clsys, timepts, U, X0)\n",
    "\n",
    "# Plot the response\n",
    "plot_results(timepts, resp.states, resp.outputs[pvtol.nstates:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86f10064",
   "metadata": {},
   "source": [
    "To see how well the estimtator did, we can compare the estimated position with the actual position:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5f24119",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Response of the first two states, including internal estimates\n",
    "h1, = plt.plot(resp.time, resp.outputs[0], 'b-', linewidth=0.75)\n",
    "h2, = plt.plot(resp.time, resp.outputs[1], 'r-', linewidth=0.75)\n",
    "\n",
    "# Add on the internal estimator states\n",
    "xh0 = clsys.find_output('xh0')\n",
    "xh1 = clsys.find_output('xh1')\n",
    "h3, = plt.plot(resp.time, resp.outputs[xh0], 'k--')\n",
    "h4, = plt.plot(resp.time, resp.outputs[xh1], 'k--')\n",
    "\n",
    "plt.plot([0, 10], [0, 0], 'k--', linewidth=0.5)\n",
    "plt.ylabel(\"Position $x$, $y$ [m]\")\n",
    "plt.xlabel(\"Time $t$ [s]\")\n",
    "plt.legend(\n",
    "    [h1, h2, h3, h4], ['$x$', '$y$', '$\\hat{x}$', '$\\hat{y}$'], \n",
    "    loc='upper right', frameon=False, ncol=2);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7139202f",
   "metadata": {},
   "source": [
    "Note the rapid convergence of the estimate to the proper value, since we are directly measuring the position variables.  If we look at the full set of states, we see that other variables have different convergence properties:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78a61e74",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 3)\n",
    "var = ['x', 'y', r'\\theta', r'\\dot x', r'\\dot y', r'\\dot \\theta']\n",
    "for i in [0, 1]:\n",
    "    for j in [0, 1, 2]:\n",
    "        k = i * 3 + j\n",
    "        axs[i, j].plot(resp.time, resp.outputs[k], label=f'${var[k]}$')\n",
    "        axs[i, j].plot(resp.time, resp.outputs[xh0+k], label=f'$\\hat {var[k]}$')\n",
    "        axs[i, j].legend()\n",
    "        if i == 1:\n",
    "            axs[i, j].set_xlabel(\"Time $t$ [s]\")\n",
    "        if j == 0:\n",
    "            axs[i, j].set_ylabel(\"State\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2039578e",
   "metadata": {},
   "source": [
    "Note the lag in tracking changes in the $\\dot x$ and $\\dot y$ states (varies from simulation to simulation, depending on the specific noise signal)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c0d5c99",
   "metadata": {},
   "source": [
    "### Full state feedback\n",
    "\n",
    "To see how the inclusion of the estimator affects the system performance, we compare it with the case where we are able to directly measure the state of the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b6a1f1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the full state feedback solution\n",
    "lqr_ctrl, _ = ct.create_statefbk_iosystem(pvtol, K)\n",
    "\n",
    "lqr_clsys = ct.interconnect(\n",
    "    [pvtol_noisy, lqr_ctrl],\n",
    "    inplist = lqr_ctrl.input_labels[0:pvtol.ninputs + pvtol.nstates] + \\\n",
    "        pvtol_noisy.input_labels[pvtol.ninputs:],\n",
    "    inputs = lqr_ctrl.input_labels[0:pvtol.ninputs + pvtol.nstates] + \\\n",
    "        pvtol_noisy.input_labels[pvtol.ninputs:],\n",
    "    outlist = pvtol.output_labels + lqr_ctrl.output_labels,\n",
    "    outputs = pvtol.output_labels + lqr_ctrl.output_labels\n",
    ")\n",
    "\n",
    "# Put together the input for the system (turn off sensor noise)\n",
    "U = [xe, ue, V, W*0]\n",
    "\n",
    "# Run a simulation with full state feedback\n",
    "lqr_resp = ct.input_output_response(lqr_clsys, timepts, U, x0)\n",
    "\n",
    "# Compare the results\n",
    "plt.plot(resp.states[0], resp.states[1], 'b-', label=\"Extended KF\")\n",
    "plt.plot(lqr_resp.states[0], lqr_resp.states[1], 'r-', label=\"Full state\")\n",
    "\n",
    "plt.xlabel('$x$ [m]')\n",
    "plt.ylabel('$y$ [m]')\n",
    "plt.axis('equal')\n",
    "plt.legend(frameon=False);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c0083cb",
   "metadata": {},
   "source": [
    "Things to try:\n",
    "* Compute a feasable trajectory and stabilize around that instead of the origin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "777053a4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}