Q1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cd89de4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import Pkg\n",
    "Pkg.activate(@__DIR__)\n",
    "Pkg.instantiate()\n",
    "using LinearAlgebra, Plots\n",
    "import ForwardDiff as FD\n",
    "using Test\n",
    "import Convex as cvx \n",
    "import ECOS\n",
    "using Random"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12d968e4",
   "metadata": {},
   "source": [
    "## Note: \n",
    "\n",
    "Some of the cells below will have multiple outputs (plots and animations), it can be easier to see everything if you do `Cell -> All Output -> Toggle Scrolling`, so that it simply expands the output area to match the size of the outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32f03c35",
   "metadata": {},
   "source": [
    "## Julia Warnings:\n",
    "\n",
    "1. For a function `foo(x::Vector)` with 1 input argument, it is not neccessary to do `df_dx = FD.jacobian(_x -> foo(_x), x)`. Instead you can just do `df_dx = FD.jacobian(foo, x)`. If you do the first one, it can dramatically slow down your compliation time. \n",
    "\n",
    "2. Do not define functions inside of other functions like this:\n",
    "```julia\n",
    "function foo(x)\n",
    "    # main function foo \n",
    "    \n",
    "    function body(x)\n",
    "        # function inside function (DON'T DO THIS)\n",
    "        return 2*x \n",
    "    end\n",
    "    \n",
    "    return body(x)\n",
    "end\n",
    "```\n",
    "This will also slow down your compilation time dramatically."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ce77d26",
   "metadata": {},
   "source": [
    "# Q1: Finite-Horizon LQR (55 pts)\n",
    "For this problem we are going to consider a \"double integrator\" for our dynamics model. This system has a state $x \\in \\mathbb{R}^4$, and control $u \\in \\mathbb{R}^2$, where the state describes the 2D position $p$ and velocity $v$ of an object, and the control is the acceleration $a$ of this object.  The state and control are the following:\n",
    "$$\\begin{align} x &= [p_1, p_2, v_1, v_2] \\\\ u &= [a_1, a_2] \\end{align}$$\n",
    "And the continuous time dynamics for this system are the following:\n",
    "$$\\begin{align} \\dot{x} =  \\begin{bmatrix} 0 & 0 & 1 & 0 \\\\ 0 & 0 & 0 & 1 \\\\ 0 & 0 & 0 & 0 \\\\ 0 & 0 & 0 & 0 \\end{bmatrix} x + \\begin{bmatrix} 0 & 0 \\\\ 0 & 0 \\\\ 1 & 0 \\\\ 0 & 1 \\end{bmatrix} \\end{align} u$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d976cfe8",
   "metadata": {},
   "source": [
    "## Part A: Discretize the model (5 pts)\n",
    "\n",
    "Use the matrix exponential (`exp` in Julia) to discretize the continuous time model assuming we have a zero-order hold on the control. See [this part of the first recitation](https://youtu.be/97JZi5ztc3c?si=Sxqg_uHnRXoozafM&t=2686) if you're unsure of what to do. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42a7b931",
   "metadata": {},
   "outputs": [],
   "source": [
    "# double integrator dynamics \n",
    "function double_integrator_AB(dt)::Tuple{Matrix,Matrix}\n",
    "    Ac = [0 0 1 0;\n",
    "          0 0 0 1;\n",
    "          0 0 0 0;\n",
    "          0 0 0 0.]\n",
    "    Bc = [0 0;\n",
    "          0 0;\n",
    "          1 0;\n",
    "          0 1]\n",
    "    nx, nu = size(Bc)\n",
    "        \n",
    "    # TODO: discretize this linear system using the Matrix Exponential\n",
    "    \n",
    "    A = zeros(nx,nx) # TODO \n",
    "    B = zeros(nx,nu) # TODO \n",
    "    \n",
    "    @assert size(A) == (nx,nx)\n",
    "    @assert size(B) == (nx,nu)\n",
    "    \n",
    "    return A, B \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbbfbda4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@testset \"discrete time dynamics\" begin \n",
    "    dt = 0.1 \n",
    "    A,B = double_integrator_AB(dt)\n",
    "    \n",
    "    x = [1,2,3,4.]\n",
    "    u = [-1,-3.]\n",
    "    \n",
    "    @test isapprox((A*x + B*u),[1.295, 2.385, 2.9, 3.7];atol = 1e-10) \n",
    "    \n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6154e8b",
   "metadata": {},
   "source": [
    "## Part B: Finite Horizon LQR via Convex Optimization (15 pts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "311a2439",
   "metadata": {},
   "source": [
    "We are now going to solve the finite horizon LQR problem with convex optimization. As we went over in class, this problem requires $Q \\in S_+$(Q is symmetric positive semi-definite) and $R \\in S_{++}$ (R is symmetric positive definite).  With this, the optimization problem can be stated as the following:\n",
    "$$ \\begin{align} \\min_{x_{1:N},u_{1:N-1}} \\quad & \\sum_{i=1}^{N-1} \\bigg[ \\frac{1}{2} x_i^TQx_i + \\frac{1}{2} u_i^TRu_i \\bigg] + \\frac{1}{2}x_N^TQ_fx_N\\\\ \n",
    " \\text{st} \\quad & x_1 = x_{\\text{IC}} \\\\ \n",
    " & x_{i+1} = A x_i + Bu_i \\quad \\text{for } i = 1,2,\\ldots,N-1 \n",
    " \\end{align}$$\n",
    " This problem is a convex optimization problem since the cost function is a convex quadratic and the constraints are all linear equality constraints. We will setup and solve this exact problem using the `Convex.jl` modeling package. (See 2/16 Recitation video for help with this package. [Notebook is here](https://github.com/Optimal-Control-16-745/recitations/blob/main/2_17_recitation/Convex.jl_tutorial.ipynb).)  Your job in the block below is to fill out a function `Xcvx,Ucvx = convex_trajopt(A,B,Q,R,Qf,N,x_ic)`, where you will form and solve the above optimization problem. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "11a5cfdc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "convex_trajopt (generic function with 1 method)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# utilities for converting to and from vector of vectors <-> matrix \n",
    "function mat_from_vec(X::Vector{Vector{Float64}})::Matrix\n",
    "    # convert a vector of vectors to a matrix \n",
    "    Xm = hcat(X...)\n",
    "    return Xm \n",
    "end\n",
    "function vec_from_mat(Xm::Matrix)::Vector{Vector{Float64}}\n",
    "    # convert a matrix into a vector of vectors \n",
    "    X = [Xm[:,i] for i = 1:size(Xm,2)]\n",
    "    return X \n",
    "end\n",
    "\"\"\"\n",
    "X,U = convex_trajopt(A,B,Q,R,Qf,N,x_ic; verbose = false)\n",
    "\n",
    "This function takes in a dynamics model x_{k+1} = A*x_k + B*u_k\n",
    "and LQR cost Q,R,Qf, with a horizon size N, and initial condition \n",
    "x_ic, and returns the optimal X and U's from the above optimization \n",
    "problem. You should use the `vec_from_mat` function to convert the \n",
    "solution matrices from cvx into vectors of vectors (vec_from_mat(X.value))\n",
    "\"\"\"\n",
    "function convex_trajopt(A::Matrix,      # A matrix \n",
    "                        B::Matrix,      # B matrix \n",
    "                        Q::Matrix,      # cost weight \n",
    "                        R::Matrix,      # cost weight \n",
    "                        Qf::Matrix,     # term cost weight \n",
    "                        N::Int64,       # horizon size \n",
    "                        x_ic::Vector;   # initial condition\n",
    "                        verbose = false\n",
    "                        )::Tuple{Vector{Vector{Float64}},Vector{Vector{Float64}}}\n",
    "    \n",
    "    # check sizes of everything \n",
    "    nx,nu = size(B)\n",
    "    @assert size(A) == (nx, nx)\n",
    "    @assert size(Q) == (nx, nx)\n",
    "    @assert size(R) == (nu, nu)\n",
    "    @assert size(Qf) == (nx, nx)\n",
    "    @assert length(x_ic) == nx\n",
    "    \n",
    "    # TODO: \n",
    "    \n",
    "    # create cvx variables where each column is a time step \n",
    "    # hint: x_k = X[:,k], u_k = U[:,k]\n",
    "    X = cvx.Variable(nx, N)\n",
    "    U = cvx.Variable(nu, N - 1)\n",
    "    \n",
    "    # create cost \n",
    "    # hint: you can't do x'*Q*x in Convex.jl, you must do cvx.quadform(x,Q)\n",
    "    # hint: add all of your cost terms to `cost` \n",
    "    cost = 0 \n",
    "    for k = 1:(N-1)\n",
    "        \n",
    "        # add stagewise cost \n",
    "        cost += 0\n",
    "    end\n",
    "\n",
    "    # add terminal cost \n",
    "    cost += 0\n",
    "    \n",
    "    # initialize cvx problem \n",
    "    prob = cvx.minimize(cost)\n",
    "    \n",
    "    # TODO: initial condition constraint \n",
    "    # hint: you can add constraints to our problem like this:\n",
    "    # prob.constraints += (Gz == h) \n",
    "    \n",
    "    for k = 1:(N-1)\n",
    "        # dynamics constraints \n",
    "        # prob.constraints += \n",
    "    end\n",
    "    \n",
    "    # solve problem (silent solver tells us the output)\n",
    "    cvx.solve!(prob, ECOS.Optimizer; silent_solver = !verbose)\n",
    "    \n",
    "    if prob.status != cvx.MathOptInterface.OPTIMAL\n",
    "        error(\"Convex.jl problem failed to solve for some reason\")\n",
    "    end\n",
    "        \n",
    "    # convert the solution matrices into vectors of vectors \n",
    "    X = vec_from_mat(X.value) \n",
    "    U = vec_from_mat(U.value) \n",
    "    \n",
    "    return X, U \n",
    "end     "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "652a7bb0",
   "metadata": {},
   "source": [
    "Now let's solve this problem for a given initial condition, and simulate it to see how it does:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3bd2062",
   "metadata": {},
   "outputs": [],
   "source": [
    "@testset \"LQR via Convex.jl\" begin \n",
    "    \n",
    "    # problem setup stuff \n",
    "    dt = 0.1 \n",
    "    tf = 5.0 \n",
    "    t_vec = 0:dt:tf \n",
    "    N = length(t_vec)\n",
    "    A,B = double_integrator_AB(dt)\n",
    "    nx,nu = size(B)\n",
    "    Q = diagm(ones(nx))\n",
    "    R = diagm(ones(nu))\n",
    "    Qf = 5*Q \n",
    "    \n",
    "    # initial condition \n",
    "    x_ic = [5,7,2,-1.4]\n",
    "    \n",
    "    # setup and solve our convex optimization problem (verbose = true for submission)\n",
    "    Xcvx,Ucvx = convex_trajopt(A,B,Q,R,Qf,N,x_ic; verbose = false) \n",
    "    \n",
    "    # TODO: simulate with the dynamics with control Ucvx, storing the \n",
    "    # state in Xsim \n",
    "        \n",
    "    # initial condition \n",
    "    Xsim = [zeros(nx) for i = 1:N]\n",
    "    Xsim[1] = 1*x_ic \n",
    "\n",
    "    # TODO dynamics simulation \n",
    "    \n",
    "    \n",
    "    @test length(Xsim) == N \n",
    "    @test norm(Xsim[end])>1e-13 \n",
    "    #----------------------plotting-----------------------\n",
    "    Xsim_m = mat_from_vec(Xsim)\n",
    "    \n",
    "    # plot state history \n",
    "    display(plot(t_vec,Xsim_m',label = [\"x₁\" \"x₂\" \"ẋ₁\" \"ẋ₂\"],\n",
    "                 title = \"State History\",\n",
    "                 xlabel = \"time (s)\", ylabel = \"x\"))\n",
    "    \n",
    "    # plot trajectory in x1 x2 space \n",
    "    display(plot(Xsim_m[1,:],Xsim_m[2,:],\n",
    "                 title = \"Trajectory in State Space\",\n",
    "                 ylabel = \"x₂\", xlabel = \"x₁\", label = \"\"))\n",
    "    #----------------------plotting-----------------------\n",
    "    \n",
    "    # tests \n",
    "    @test 1e-14 < maximum(norm.(Xsim .- Xcvx,Inf)) < 1e-3\n",
    "    @test isapprox(Ucvx[1], [-7.8532442316767, -4.127120137234], atol = 1e-3)\n",
    "    @test isapprox(Xcvx[end], [-0.02285990, -0.07140241, -0.21259, -0.1540299], atol = 1e-3)\n",
    "    @test 1e-14 < norm(Xcvx[end] - Xsim[end]) < 1e-3\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d14994e",
   "metadata": {},
   "source": [
    "### Bellman's Principle of Optimality\n",
    "Now we will test Bellman's Principle of optimality. This can be phrased in many different ways, but the main gist is that any section of an optimal trajectory must be optimal. Our original optimization problem was the above problem:\n",
    "$$ \\begin{align} \\min_{x_{1:N},u_{1:N-1}} \\quad & \\sum_{i=1}^{N-1} \\bigg[ \\frac{1}{2} x_i^TQx_i + \\frac{1}{2} u_i^TRu_i \\bigg] + \\frac{1}{2}x_N^TQ_fx_N\\\\ \n",
    " \\text{st} \\quad & x_1 = x_{\\text{IC}} \\\\ \n",
    " & x_{i+1} = A x_i + Bu_i \\quad \\text{for } i = 1,2,\\ldots,N-1 \n",
    " \\end{align}$$\n",
    "which has a solution $x^*_{1:N},u^*_{1:N-1}$. Now let's look at optimizing over a subsection of this trajectory. That means that instead of solving for $x_{1:N},u_{1:N-1}$, we are now solving for $x_{L:N},u_{L:N-1}$ for some new timestep $1 < L < N$. What we are going to do is take the initial condition from $x^*_L$ from our original optimization problem, and setup a new optimization problem that optimizes over $x_{L:N},u_{L:N-1}$:\n",
    "$$ \\begin{align} \\min_{x_{L:N},u_{L:N-1}} \\quad & \\sum_{i=L}^{N-1} \\bigg[ \\frac{1}{2} x_i^TQx_i + \\frac{1}{2} u_i^TRu_i \\bigg] + \\frac{1}{2}x_N^TQ_fx_N\\\\ \n",
    " \\text{st     } \\quad & x_L = x^*_L \\\\ \n",
    " & x_{i+1} = A x_i + Bu_i \\quad \\text{for } i = L,L + 1,\\ldots,N-1 \n",
    " \\end{align}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f9071f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "@testset \"Bellman's Principle of Optimality\" begin \n",
    "    \n",
    "    # problem setup \n",
    "    dt = 0.1 \n",
    "    tf = 5.0 \n",
    "    t_vec = 0:dt:tf \n",
    "    N = length(t_vec)\n",
    "    A,B = double_integrator_AB(dt)\n",
    "    nx,nu = size(B)\n",
    "    x0 = [5,7,2,-1.4] # initial condition \n",
    "    Q = diagm(ones(nx))\n",
    "    R = diagm(ones(nu))\n",
    "    Qf = 5*Q \n",
    "    \n",
    "    # solve for X_{1:N}, U_{1:N-1} with convex optimization\n",
    "    Xcvx1,Ucvx1 = convex_trajopt(A,B,Q,R,Qf,N,x0; verbose = false)\n",
    "        \n",
    "    # now let's solve a subsection of this trajectory \n",
    "    L = 18 \n",
    "    N_2 = N - L + 1\n",
    "    \n",
    "    # here is our updated initial condition from the first problem \n",
    "    x0_2 = Xcvx1[L]\n",
    "    Xcvx2,Ucvx2 = convex_trajopt(A,B,Q,R,Qf,N_2,x0_2; verbose = false)\n",
    "        \n",
    "    # test if these trajectories match for the times they share \n",
    "    U_error = Ucvx1[L:end] .- Ucvx2\n",
    "    X_error = Xcvx1[L:end] .- Xcvx2\n",
    "    @test 1e-14 < maximum(norm.(U_error)) < 1e-3\n",
    "    @test 1e-14 < maximum(norm.(X_error)) < 1e-3\n",
    "\n",
    "\n",
    "    # ---------------------------plotting ------------------------------\n",
    "    X1m = mat_from_vec(Xcvx1)\n",
    "    X2m = mat_from_vec(Xcvx2)\n",
    "    plot(X2m[1,:],X2m[2,:], label = \"optimal subtrajectory\", lw = 5, ls = :dot)\n",
    "    display(plot!(X1m[1,:],X1m[2,:],\n",
    "                 title = \"Trajectory in State Space\",\n",
    "                 ylabel = \"x₂\", xlabel = \"x₁\", label = \"full trajectory\"))\n",
    "    # ---------------------------plotting ------------------------------\n",
    "    \n",
    "    @test isapprox(Xcvx1[end], [-0.02285990, -0.07140241, -0.21259, -0.1540299], rtol = 1e-3)\n",
    "    @test 1e-14 < norm(Xcvx1[end] - Xcvx2[end],Inf) < 1e-3\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c80ae78",
   "metadata": {},
   "source": [
    "## Part C: Finite-Horizon LQR via Riccati (10 pts)\n",
    "Now we are going to solve the original finite-horizon LQR problem:\n",
    "$$ \\begin{align} \\min_{x_{1:N},u_{1:N-1}} \\quad & \\sum_{i=1}^{N-1} \\bigg[ \\frac{1}{2} x_i^TQx_i + \\frac{1}{2} u_i^TRu_i \\bigg] + \\frac{1}{2}x_N^TQ_fx_N\\\\ \n",
    " \\text{st} \\quad & x_1 = x_{\\text{IC}} \\\\ \n",
    " & x_{i+1} = A x_i + Bu_i \\quad \\text{for } i = 1,2,\\ldots,N-1 \n",
    " \\end{align}$$\n",
    " with a Riccati recursion instead of convex optimization. We describe our optimal cost-to-go function (aka the Value function) as the following:\n",
    " \n",
    " $$V_k(x) = \\frac{1}{2}x^TP_kx$$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9c43c4b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "fhlqr"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "use the Riccati recursion to calculate the cost to go quadratic matrix P and \n",
    "optimal control gain K at every time step. Return these as a vector of matrices, \n",
    "where P_k = P[k], and K_k = K[k]\n",
    "\"\"\"\n",
    "function fhlqr(A::Matrix, # A matrix \n",
    "               B::Matrix, # B matrix \n",
    "               Q::Matrix, # cost weight \n",
    "               R::Matrix, # cost weight \n",
    "               Qf::Matrix,# term cost weight \n",
    "               N::Int64   # horizon size \n",
    "               )::Tuple{Vector{Matrix{Float64}}, Vector{Matrix{Float64}}} # return two matrices \n",
    "        \n",
    "    # check sizes of everything \n",
    "    nx,nu = size(B)\n",
    "    @assert size(A) == (nx, nx)\n",
    "    @assert size(Q) == (nx, nx)\n",
    "    @assert size(R) == (nu, nu)\n",
    "    @assert size(Qf) == (nx, nx)\n",
    "        \n",
    "    # instantiate S and K \n",
    "    P = [zeros(nx,nx) for i = 1:N]\n",
    "    K = [zeros(nu,nx) for i = 1:N-1]\n",
    "    \n",
    "    # initialize S[N] with Qf \n",
    "    P[N] = deepcopy(Qf)\n",
    "    \n",
    "    # Riccati\n",
    "    for k = NaN:NaN\n",
    "        # TODO \n",
    "        \n",
    "    end\n",
    "    \n",
    "    return P, K \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3304e9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "@testset \"Convex trajopt vs LQR\" begin \n",
    "    \n",
    "    # problem stuff\n",
    "    dt = 0.1 \n",
    "    tf = 5.0 \n",
    "    t_vec = 0:dt:tf \n",
    "    N = length(t_vec)\n",
    "    A,B = double_integrator_AB(dt)\n",
    "    nx,nu = size(B)\n",
    "    x0 = [5,7,2,-1.4] # initial condition \n",
    "    Q = diagm(ones(nx))\n",
    "    R = diagm(ones(nu))\n",
    "    Qf = 5*Q \n",
    "    \n",
    "    # solve for X_{1:N}, U_{1:N-1} with convex optimization\n",
    "    Xcvx,Ucvx = convex_trajopt(A,B,Q,R,Qf,N,x0; verbose = false)\n",
    "    P, K = fhlqr(A,B,Q,R,Qf,N)\n",
    "    # now let's simulate using Ucvx \n",
    "    Xsim_cvx = [zeros(nx) for i = 1:N]\n",
    "    Xsim_cvx[1] = 1*x0 \n",
    "    Xsim_lqr = [zeros(nx) for i = 1:N]\n",
    "    Xsim_lqr[1] = 1*x0 \n",
    "    for i = 1:N-1\n",
    "        # simulate cvx control \n",
    "        Xsim_cvx[i+1] = A*Xsim_cvx[i] + B*Ucvx[i]\n",
    "        \n",
    "        # TODO: use your FHLQR control gains K to calculate u_lqr\n",
    "        # simulate lqr control \n",
    "        u_lqr = zeros(2)\n",
    "        Xsim_lqr[i+1] = A*Xsim_lqr[i] + B*u_lqr\n",
    "    end\n",
    "    \n",
    "\n",
    "    @test isapprox(Xsim_lqr[end], [-0.02286201, -0.0714058, -0.21259, -0.154030], rtol = 1e-3)\n",
    "    @test 1e-13 < norm(Xsim_lqr[end] - Xsim_cvx[end]) < 1e-3\n",
    "    @test 1e-13 < maximum(norm.(Xsim_lqr - Xsim_cvx)) < 1e-3\n",
    "\n",
    "    \n",
    "    # ------------------------plotting--------------------------\n",
    "    X1m = mat_from_vec(Xsim_cvx)\n",
    "    X2m = mat_from_vec(Xsim_lqr)\n",
    "    # plot trajectory in x1 x2 space \n",
    "    plot(X1m[1,:],X1m[2,:], label = \"cvx trajectory\", lw = 4, ls = :dot)\n",
    "    display(plot!(X2m[1,:],X2m[2,:],\n",
    "                 title = \"Trajectory in State Space\",\n",
    "                 ylabel = \"x₂\", xlabel = \"x₁\", lw = 2, label = \"lqr trajectory\"))\n",
    "    # ------------------------plotting--------------------------\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fddcbd3",
   "metadata": {},
   "source": [
    "To emphasize that these two methods for solving the optimization problem result in the same solutions, we are now going to sample initial conditions and run both solutions. You will have to fill in your LQR policy again. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f828435a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import Random \n",
    "Random.seed!(1)\n",
    "@testset \"Convex trajopt vs LQR\" begin \n",
    "    \n",
    "    # problem stuff\n",
    "    dt = 0.1 \n",
    "    tf = 5.0 \n",
    "    t_vec = 0:dt:tf \n",
    "    N = length(t_vec)\n",
    "    A,B = double_integrator_AB(dt)\n",
    "    nx,nu = size(B)\n",
    "    Q = diagm(ones(nx))\n",
    "    R = diagm(ones(nu))\n",
    "    Qf = 5*Q \n",
    "    \n",
    "    \n",
    "    plot()\n",
    "    for ic_iter = 1:20\n",
    "        x0 = [5*randn(2); 1*randn(2)]\n",
    "        # solve for X_{1:N}, U_{1:N-1} with convex optimization\n",
    "        Xcvx,Ucvx = convex_trajopt(A,B,Q,R,Qf,N,x0; verbose = false)\n",
    "        P, K = fhlqr(A,B,Q,R,Qf,N)\n",
    "        Xsim_cvx = [zeros(nx) for i = 1:N]\n",
    "        Xsim_cvx[1] = 1*x0 \n",
    "        Xsim_lqr = [zeros(nx) for i = 1:N]\n",
    "        Xsim_lqr[1] = 1*x0 \n",
    "        for i = 1:N-1\n",
    "            # simulate cvx control \n",
    "            Xsim_cvx[i+1] = A*Xsim_cvx[i] + B*Ucvx[i]\n",
    "\n",
    "            # TODO: use your FHLQR control gains K to calculate u_lqr\n",
    "            # simulate lqr control \n",
    "            u_lqr = zeros(2)\n",
    "            Xsim_lqr[i+1] = A*Xsim_lqr[i] + B*u_lqr\n",
    "        end\n",
    "\n",
    "        @test 1e-13 < norm(Xsim_lqr[end] - Xsim_cvx[end]) < 1e-3\n",
    "        @test 1e-13 < maximum(norm.(Xsim_lqr - Xsim_cvx)) < 1e-3\n",
    "        \n",
    "        # ------------------------plotting--------------------------\n",
    "        X1m = mat_from_vec(Xsim_cvx)\n",
    "        X2m = mat_from_vec(Xsim_lqr)\n",
    "        plot!(X2m[1,:],X2m[2,:], label = \"\", lw = 4, ls = :dot)\n",
    "        plot!(X1m[1,:],X1m[2,:], label = \"\", lw = 2)\n",
    "    end\n",
    "    display(plot!(title = \"cvx vs LQR\", ylabel = \"x₂\", xlabel = \"x₁\"))\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06ac015c",
   "metadata": {},
   "source": [
    "## Part D: Why LQR is so great (10 pts)\n",
    "\n",
    "Now we are going to emphasize two reasons why the feedback policy from LQR is so useful:\n",
    "\n",
    "1. It is robust to noise and model uncertainty (the Convex approach would require re-solving of the problem every time the new state differs from the expected state (this is MPC, more on this in Q3)\n",
    "2. We can drive to any achievable goal state with $u = -K(x - x_{goal})$\n",
    "\n",
    "First we are going to look at a simulation with the following white noise:\n",
    "$$x_{k+1} = Ax_k + Bu_k + \\text{noise}$$\n",
    "Where noise $ \\sim \\mathcal{N}(0,\\Sigma)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee377725",
   "metadata": {},
   "outputs": [],
   "source": [
    "@testset \"Why LQR is great reason 1\" begin \n",
    "    \n",
    "    # problem stuff\n",
    "    dt = 0.1 \n",
    "    tf = 7.0 \n",
    "    t_vec = 0:dt:tf \n",
    "    N = length(t_vec)\n",
    "    A,B = double_integrator_AB(dt)\n",
    "    nx,nu = size(B)\n",
    "    x0 = [5,7,2,-1.4] # initial condition \n",
    "    Q = diagm(ones(nx))\n",
    "    R = diagm(ones(nu))\n",
    "    Qf = 10*Q \n",
    "    \n",
    "    # solve for X_{1:N}, U_{1:N-1} with convex optimization\n",
    "    Xcvx,Ucvx = convex_trajopt(A,B,Q,R,Qf,N,x0; verbose = false)\n",
    "    P, K = fhlqr(A,B,Q,R,Qf,N)\n",
    "    # now let's simulate using Ucvx \n",
    "    Xsim_cvx = [zeros(nx) for i = 1:N]\n",
    "    Xsim_cvx[1] = 1*x0 \n",
    "    Xsim_lqr = [zeros(nx) for i = 1:N]\n",
    "    Xsim_lqr[1] = 1*x0 \n",
    "    for i = 1:N-1\n",
    "        # sampled noise to be added after each step \n",
    "        noise = [.005*randn(2);.1*randn(2)]\n",
    "        \n",
    "        # simulate cvx control \n",
    "        Xsim_cvx[i+1] = A*Xsim_cvx[i] + B*Ucvx[i] + noise\n",
    "        \n",
    "        # TODO: use your FHLQR control gains K to calculate u_lqr\n",
    "        # simulate lqr control \n",
    "        u_lqr = zeros(2)\n",
    "        Xsim_lqr[i+1] = A*Xsim_lqr[i] + B*u_lqr + noise\n",
    "    end\n",
    "    \n",
    "    # make sure our LQR achieved the goal \n",
    "    @test norm(Xsim_cvx[end]) > norm(Xsim_lqr[end])\n",
    "    @test norm(Xsim_lqr[end]) < .7\n",
    "    @test norm(Xsim_cvx[end]) > 2.0\n",
    "    \n",
    "    \n",
    "    # ------------------------plotting--------------------------\n",
    "    X1m = mat_from_vec(Xsim_cvx)\n",
    "    X2m = mat_from_vec(Xsim_lqr)\n",
    "    # plot trajectory in x1 x2 space \n",
    "    plot(X1m[1,:],X1m[2,:], label = \"CVX Trajectory (no replanning)\", lw = 4, ls = :dot)\n",
    "    display(plot!(X2m[1,:],X2m[2,:],\n",
    "                 title = \"Trajectory in State Space (Noisy Dynamics)\",\n",
    "                 ylabel = \"x₂\", xlabel = \"x₁\", lw = 2, label = \"LQR Trajectory\"))\n",
    "    ecvx = [norm(x[1:2]) for x in Xsim_cvx]\n",
    "    elqr = [norm(x[1:2]) for x in Xsim_lqr]\n",
    "    plot(t_vec, elqr, label = \"LQR Trajectory\",ylabel = \"|x - xgoal|\",\n",
    "         xlabel = \"time (s)\", title = \"Error for CVX vs LQR (Noisy Dynamics)\")\n",
    "    display(plot!(t_vec, ecvx, label = \"CVX Trajectory (no replanning)\"))\n",
    "    # ------------------------plotting--------------------------\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ea8bd00",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@testset \"Why LQR is great reason 2\" begin \n",
    "    \n",
    "    # problem stuff\n",
    "    dt = 0.1 \n",
    "    tf = 20.0 \n",
    "    t_vec = 0:dt:tf \n",
    "    N = length(t_vec)\n",
    "    A,B = double_integrator_AB(dt)\n",
    "    nx,nu = size(B)\n",
    "    x0 = [5,7,2,-1.4] # initial condition \n",
    "    Q = diagm(ones(nx))\n",
    "    R = diagm(ones(nu))\n",
    "    Qf = 10*Q \n",
    "    \n",
    "    P, K = fhlqr(A,B,Q,R,Qf,N)\n",
    "    \n",
    "    # TODO: specify a goal state with 0 velocity within a 5m radius of 0 \n",
    "    xgoal = [-3.5,-3.5,0,0]\n",
    "    @test norm(xgoal[1:2])< 5\n",
    "    @test norm(xgoal[3:4])<1e-13 # ensure 0 velocity\n",
    "\n",
    "    Xsim_lqr = [zeros(nx) for i = 1:N]\n",
    "    Xsim_lqr[1] = 1*x0 \n",
    "    \n",
    "    for i = 1:N-1\n",
    "        # TODO: use your FHLQR control gains K to calculate u_lqr\n",
    "        # simulate lqr control \n",
    "        u_lqr = zeros(2)\n",
    "        Xsim_lqr[i+1] = A*Xsim_lqr[i] + B*u_lqr\n",
    "    end\n",
    "    \n",
    "    @test norm(Xsim_lqr[end][1:2] - xgoal[1:2]) < .1 \n",
    "    \n",
    "    # ------------------------plotting--------------------------\n",
    "    Xm = mat_from_vec(Xsim_lqr)\n",
    "    plot(xgoal[1:1],xgoal[2:2],seriestype = :scatter, label = \"goal state\")\n",
    "    display(plot!(Xm[1,:],Xm[2,:],\n",
    "                 title = \"Trajectory in State Space\",\n",
    "                 ylabel = \"x₂\", xlabel = \"x₁\", lw = 2, label = \"LQR Trajectory\"))\n",
    "\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9a3329e",
   "metadata": {},
   "source": [
    "## Part E: Infinite-horizon LQR (10 pts)\n",
    "Up until this point, we have looked at finite-horizon LQR which only considers a finite number of timesteps in our trajectory.  When this problem is solved with a Riccati recursion, there is a new feedback gain matrix $K_k$ for each timestep. As the length of the trajectory increases, the first feedback gain matrix $K_1$ will begin to converge on what we call the \"infinite-horizon LQR gain\".  This is the value that $K_1$ converges to as $N \\rightarrow \\infty$.\n",
    "\n",
    "Below, we will plot the values of $P$ and $K$ throughout the horizon and observe this convergence. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edb1b05e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# half vectorization of a matrix \n",
    "function vech(A)\n",
    "    return A[tril(trues(size(A)))]\n",
    "end\n",
    "@testset \"P and K time analysis\" begin \n",
    "    \n",
    "    # problem stuff\n",
    "    dt = 0.1 \n",
    "    tf = 10.0 \n",
    "    t_vec = 0:dt:tf \n",
    "    N = length(t_vec)\n",
    "    A,B = double_integrator_AB(dt)\n",
    "    nx,nu = size(B)\n",
    "    \n",
    "    # cost terms \n",
    "    Q = diagm(ones(nx))\n",
    "    R = .5*diagm(ones(nu))\n",
    "    Qf = randn(nx,nx); Qf = Qf'*Qf + I;\n",
    "    \n",
    "    P, K = fhlqr(A,B,Q,R,Qf,N)\n",
    "    \n",
    "    Pm = hcat(vech.(P)...)\n",
    "    Km = hcat(vec.(K)...)\n",
    "    \n",
    "    # make sure these things converged \n",
    "    @test 1e-13 < norm(P[1] - P[2]) < 1e-3 \n",
    "    @test 1e-13 < norm(K[1] - K[2]) < 1e-3\n",
    "    \n",
    "    display(plot(t_vec, Pm', label = \"\",title = \"Cost-to-go Matrix (P)\", xlabel = \"time(s)\"))\n",
    "    display(plot(t_vec[1:end-1], Km', label = \"\",title = \"Gain Matrix (K)\", xlabel = \"time(s)\"))\n",
    "\n",
    "    \n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13a23f92",
   "metadata": {},
   "source": [
    "Complete this infinite horizon LQR function where you do a Riccati recursion until the cost to go matrix P converges:\n",
    "\n",
    "$$ \\|P_k - P_{k+1}\\| \\leq \\text{tol}$$\n",
    "\n",
    "And return the steady state $P$ and $K$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52b01788",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "P,K = ihlqr(A,B,Q,R)\n",
    "\n",
    "TODO: complete this infinite horizon LQR function where \n",
    "you do the Riccati recursion until the cost to go matrix \n",
    "P converges to a steady value |P_k - P_{k+1}| ≤ tol \n",
    "\"\"\"\n",
    "function ihlqr(A::Matrix,       # vector of A matrices \n",
    "               B::Matrix,       # vector of B matrices\n",
    "               Q::Matrix,       # cost matrix Q \n",
    "               R::Matrix;       # cost matrix R \n",
    "               max_iter = 1000, # max iterations for Riccati\n",
    "               tol = 1e-5       # convergence tolerance\n",
    "               )::Tuple{Matrix, Matrix} # return two matrices \n",
    "        \n",
    "    # get size of x and u from B \n",
    "    nx, nu = size(B)\n",
    "        \n",
    "    # initialize S with Q\n",
    "    P = deepcopy(Q)\n",
    "    \n",
    "    # Riccati \n",
    "    for riccati_iter = 1:max_iter \n",
    "\n",
    "        # TODO \n",
    "\n",
    "    end\n",
    "    error(\"ihlqr did not converge\")\n",
    "end\n",
    "@testset \"ihlqr test\" begin \n",
    "    # problem stuff\n",
    "    dt = 0.1 \n",
    "    A,B = double_integrator_AB(dt)\n",
    "    nx,nu = size(B)\n",
    "    \n",
    "    # we're just going to modify the system a little bit \n",
    "    # so the following graphs are still interesting\n",
    "\n",
    "    Q = diagm(ones(nx))\n",
    "    R = .5*diagm(ones(nu))\n",
    "    P, K = ihlqr(A,B,Q,R)\n",
    "    \n",
    "    # check this P is in fact a solution to the Riccati equation\n",
    "    @test typeof(P) == Matrix{Float64} \n",
    "    @test typeof(K) == Matrix{Float64}\n",
    "    @test 1e-13 < norm(Q + K'*R*K + (A - B*K)'P*(A - B*K) - P) < 1e-3\n",
    "    \n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d43bee19",
   "metadata": {},
   "source": [
    "## Part F (5 pts): One sentence short answer\n",
    "\n",
    "1. What is the difference between stage cost and terminal cost?\n",
    "\n",
    "**put one sentence answer here**\n",
    "\n",
    "2. What is a terminal cost trying to capture? (think about dynamic programming)\n",
    "\n",
    "**put one sentence answer here**\n",
    "\n",
    "3. In order to build an LQR controller for a linear system, do we need to know the initial state $x_0$?\n",
    "\n",
    "**put one sentence answer here**\n",
    "\n",
    "4. If a linear system is uncontrollable, will the finite-horizon LQR convex optimization problem have a solution?\n",
    "\n",
    "**put one sentence answer here**"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.6.7",
   "language": "julia",
   "name": "julia-1.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}