Q3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86bb6f30",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import Pkg\n",
    "Pkg.activate(@__DIR__)\n",
    "Pkg.instantiate()\n",
    "using LinearAlgebra, Plots\n",
    "import ForwardDiff as FD\n",
    "import MeshCat as mc \n",
    "using Test\n",
    "using Random\n",
    "import Convex as cvx \n",
    "import ECOS      # the solver we use in this hw \n",
    "# import Hypatia # other solvers you can try\n",
    "# import COSMO   # other solvers you can try \n",
    "using ProgressMeter\n",
    "include(joinpath(@__DIR__,\"utils/rendezvous.jl\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c854320",
   "metadata": {},
   "source": [
    "## Notes: \n",
    "\n",
    "1. 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.\n",
    "2. Things in space move very slowly (by design), because of this, you may want to speed up the animations when you're viewing them. You can do this in MeshCat by doing `Open Controls -> Animations -> Time Scale`, to modify the time scale. You can also play/pause/scrub from this menu as well. \n",
    "3. You can move around your view in MeshCat by `clicking + dragging`, and you can pan with `right click + dragging`, and zoom with the scroll wheel on your mouse (or trackpad specific alternatives). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f26d987e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ef2cc03",
   "metadata": {},
   "source": [
    "## Is LQR the answer for everything?\n",
    "\n",
    "Unfortunately, no. LQR is great for problems with true quadratic costs and linear dynamics, but this is a very small subset of convex trajectory optimization problems. While a quadratic cost is common in control, there are other available convex cost functions that may better motivate the desired behavior of the system. These costs can be things like an L1 norm on the control inputs ($\\|u\\|_1$), or an L2 goal error ($\\|x - x_{goal}\\|_2$). Also, control problems often have constraints like path constraints, control bounds, or terminal constraints, that can't be handled with LQR. With the addition of these constraints, the trajectory optimization problem is stil convex and easy to solve, but we can no longer just get an optimal gain $K$ and apply a feedback policy in these situations. \n",
    "\n",
    "The solution to this is Model Predictive Control (MPC). In MPC, we are setting up and solving a convex trajectory optimization at every time step, optimizing over some horizon or window into the future, and executing the first control in the solution. To see how this works, we are going to try this for a classic space control problem: the rendezvous. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57a19484",
   "metadata": {},
   "source": [
    "# Q3: Optimal Rendezvous and Docking (55 pts)\n",
    "\n",
    "In this example, we are going to use convex optimization to control the SpaceX Dragon 1 spacecraft as it docks with the International Space Station (ISS). The dynamics of the Dragon vehicle can be modeled with [Clohessy-Wiltshire equations](https://en.wikipedia.org/wiki/Clohessy%E2%80%93Wiltshire_equations), which is a linear dynamics model in continuous time. The state and control of this system are the following:\n",
    "\n",
    "$$ \\begin{align}\n",
    "x &= [r_x, r_y, r_z, v_x, v_y, v_z]^T,\\\\\n",
    "u &= [t_x, t_y, t_z]^T, \\end{align}$$\n",
    "\n",
    "where $r$ is a relative position of the Dragon spacecraft with respect to the ISS, $v$ is the relative velocity, and $t$ is the thrust on the spacecraft. The continuous time dynamics of the vehicle are the following:\n",
    "\n",
    "$$ \\begin{align}\n",
    "\\dot{x} &= \\begin{bmatrix}0  &   0 & 0  &  1 &  0 &  0 \\\\  \n",
    "         0 &    0 & 0  &  0 &  1 &  0 \\\\ \n",
    "         0 &    0 & 0 &   0 &  0 &  1\\\\\n",
    "         3n^2 &0 & 0  &  0 &  2n &0 \\\\\n",
    "         0  &   0 & 0  & -2n &0  & 0\\\\\n",
    "         0  &   0 &-n^2 & 0 &  0 &  0 \\end{bmatrix} + \\begin{bmatrix} 0 & 0 & 0 \\\\ 0 & 0 & 0 \\\\ 0 & 0 & 0 \\\\ 1 & 0 & 0 \\\\ 0 & 1 & 0 \\\\ 0 & 0 & 1 \\end{bmatrix} u,\n",
    "\\end{align}$$\n",
    "where $n = \\sqrt{\\mu/a^3}$, with $\\mu$ being the [standard gravitational parameter](https://en.wikipedia.org/wiki/Standard_gravitational_parameter), and $a$ being the semi-major axis of the orbit of the ISS. \n",
    "\n",
    "We are going to use three different techniques for solving this control problem, the first is LQR, the second is convex trajectory optimization, and the third is convex MPC where we will be able to account for unmodeled dynamics in our system (the \"sim to real\" gap). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbdd42a1",
   "metadata": {},
   "source": [
    "## Part A: Discretize the dynamics (5 pts)\n",
    "\n",
    "Use the matrix exponential to convert the linear ODE into a linear discrete time model (hint: the matrix exponential is just `exp()` in Julia when called on a matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c4e50eb",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "function create_dynamics(dt::Real)::Tuple{Matrix,Matrix}\n",
    "    mu = 3.986004418e14 # standard gravitational parameter\n",
    "    a = 6971100.0       # semi-major axis of ISS\n",
    "    n = sqrt(mu/a^3)    # mean motion\n",
    "\n",
    "    # continuous time dynamics ẋ = Ax + Bu\n",
    "    A = [0     0  0    1   0   0; \n",
    "         0     0  0    0   1   0;\n",
    "         0     0  0    0   0   1;\n",
    "         3*n^2 0  0    0   2*n 0;\n",
    "         0     0  0   -2*n 0   0;\n",
    "         0     0 -n^2  0   0   0]\n",
    "         \n",
    "    B = Matrix([zeros(3,3);0.1*I(3)])\n",
    "\n",
    "    # TODO: convert to discrete time X_{k+1} = Ad*x_k + Bd*u_k\n",
    "\n",
    "    Ad = zeros(6,6)\n",
    "    Bd = zeros(6,3)\n",
    "\n",
    "    return Ad, Bd\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a23c32b4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@testset \"discrete dynamics\" begin \n",
    "    A,B = create_dynamics(1.0)\n",
    "    \n",
    "    x = [1,3,-.3,.2,.4,-.5]\n",
    "    u = [-.1,.5,.3]\n",
    "    \n",
    "    # test these matrices \n",
    "    @test isapprox(A*x + B*u, [1.195453, 3.424786, -0.78499972, 0.190925, 0.4495759, -0.4699993], atol = 1e-3)\n",
    "    @test isapprox(det(A), 1, atol = 1e-8)\n",
    "    @test isapprox(norm(B,Inf), 0.0999999803, atol = 1e-5)\n",
    "    \n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "288de1e0",
   "metadata": {},
   "source": [
    "## Part B: LQR (10 pts)\n",
    "Now we will take a given reference trajectory `X_ref` and track it with finite-horizon LQR. Remember that finite-horizon LQR is solving this problem:\n",
    "$$ \\begin{align} \\min_{x_{1:N},u_{1:N-1}} \\quad & \\sum_{i=1}^{N-1} \\bigg[ \\frac{1}{2} (x_i - x_{ref, i})^TQ(x_i - x_{ref, i}) + \\frac{1}{2} u_i^TRu_i \\bigg] + \\frac{1}{2}(x_N- x_{ref, N})^TQ_f\n",
    "(x_N- x_{ref, 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",
    "where our policy is $u_i = -K_i(x_i - x_{ref, i})$. Use your code from the previous problem with your `fhlqr` function to generate your gain matrices. \n",
    "\n",
    "One twist we will throw into this is control constraints `u_min` and `u_max`. You should use the function `clamp.(u, u_min, u_max)` to clamp the values of your `u` to be within this range. \n",
    "\n",
    "If implemented correctly, you should see the Dragon spacecraft dock with the ISS successfuly, but only after it crashes through the ISS a little bit. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e95475a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "@testset \"LQR rendezvous\" begin \n",
    "\n",
    "    # create our discrete time model \n",
    "    dt = 1.0\n",
    "    A,B = create_dynamics(dt)\n",
    "\n",
    "    # get our sizes for state and control\n",
    "    nx,nu = size(B)\n",
    "\n",
    "    # initial and goal states\n",
    "    x0 = [-2;-4;2;0;0;.0]\n",
    "    xg = [0,-.68,3.05,0,0,0]\n",
    "\n",
    "    # bounds on U\n",
    "    u_max = 0.4*ones(3)\n",
    "    u_min = -u_max\n",
    "\n",
    "    # problem size and reference trajectory \n",
    "    N = 120\n",
    "    t_vec = 0:dt:((N-1)*dt)\n",
    "    X_ref = desired_trajectory_long(x0,xg,200,dt)[1:N]\n",
    "    \n",
    "    # TODO: FHLQR \n",
    "    Q = diagm(ones(nx))\n",
    "    R = diagm(ones(nu))\n",
    "    Qf = 10*Q\n",
    "    # TODO get K's from fhlqr \n",
    "    \n",
    "    # simulation \n",
    "    X_sim = [zeros(nx) for i = 1:N]\n",
    "    U_sim = [zeros(nu) for i = 1:N-1]\n",
    "    X_sim[1] = x0 \n",
    "    for i = 1:(N-1) \n",
    "        # TODO: put LQR control law here \n",
    "        # make sure to clamp \n",
    "        U_sim[i] = zeros(3)\n",
    "        \n",
    "        # simulate 1 step \n",
    "        X_sim[i+1] = A*X_sim[i] + B*U_sim[i]\n",
    "    end\n",
    "\n",
    "    # -------------plotting/animation---------------------------\n",
    "    Xm = mat_from_vec(X_sim)\n",
    "    Um = mat_from_vec(U_sim)\n",
    "    display(plot(t_vec,Xm[1:3,:]',title = \"Positions (LQR)\",\n",
    "                 xlabel = \"time (s)\", ylabel = \"position (m)\",\n",
    "                 label = [\"x\" \"y\" \"z\"]))\n",
    "    display(plot(t_vec,Xm[4:6,:]',title = \"Velocities (LQR)\",\n",
    "            xlabel = \"time (s)\", ylabel = \"velocity (m/s)\",\n",
    "                 label = [\"x\" \"y\" \"z\"]))\n",
    "    display(plot(t_vec[1:end-1],Um',title = \"Control (LQR)\",\n",
    "            xlabel = \"time (s)\", ylabel = \"thrust (N)\",\n",
    "                 label = [\"x\" \"y\" \"z\"]))\n",
    "\n",
    "    # feel free to toggle `show_reference`\n",
    "    display(animate_rendezvous(X_sim, X_ref, dt;show_reference = false))\n",
    "    # -------------plotting/animation---------------------------\n",
    "\n",
    "    \n",
    "    # testing \n",
    "    xs=[x[1] for x in X_sim]\n",
    "    ys=[x[2] for x in X_sim]\n",
    "    zs=[x[3] for x in X_sim]\n",
    "    @test norm(X_sim[end] - xg) < .01 # goal \n",
    "    @test (xg[2] + .1) < maximum(ys) < 0 # we should have hit the ISS \n",
    "    @test maximum(zs) >= 4 # check to see if you did the circle \n",
    "    @test minimum(zs) <= 2 # check to see if you did the circle \n",
    "    @test maximum(xs) >= 1 # check to see if you did the circle \n",
    "    @test maximum(norm.(U_sim,Inf)) <= 0.4 # control constraints satisfied \n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c517151c",
   "metadata": {},
   "source": [
    "## Part C: Convex Trajectory Optimization (15 pts)\n",
    "\n",
    "Now we are going to assume that we have a perfect model (assume there is no sim to real gap), and that we have a perfect state estimate. With this, we are going to solve our control problem as a convex trajectory optimization problem. \n",
    "\n",
    "$$ \\begin{align} \\min_{x_{1:N},u_{1:N-1}} \\quad & \\sum_{i=1}^{N-1} \\bigg[ \\frac{1}{2} (x_i - x_{ref, i})^TQ(x_i - x_{ref, i}) + \\frac{1}{2} u_i^TRu_i \\bigg] \\\\ \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",
    " & u_{min} \\leq u_i \\leq u_{max} \\quad \\text{for } i = 1,2,\\ldots,N-1 \\\\\n",
    " & x_i[2] \\leq x_{goal} [2]\\quad \\text{for } i = 1,2,\\ldots,N \\\\ \n",
    " & x_N = x_{goal}\n",
    " \\end{align}$$\n",
    " \n",
    " Where we have an LQR cost, an initial condition constraint ($x_1 = x_{\\text{IC}}$), linear dynamics constraints ($x_{i+1} = A x_i + Bu_i$), bound constraints on the control ($\\leq u_i \\leq u_{max}$), an ISS collision constraint ($x_i[2] \\leq x_{goal} [2]$), and a terminal constraint ($x_N = x_{goal}$).  This problem is convex and we will setup and solve this with `Convex.jl`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c303fe7",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Xcvx,Ucvx = convex_trajopt(A,B,X_ref,x0,xg,u_min,u_max,N)\n",
    "\n",
    "setup and solve the above optimization problem, returning \n",
    "the solutions X and U, after first converting them to \n",
    "vectors of vectors with vec_from_mat(X.value)\n",
    "\"\"\"\n",
    "function convex_trajopt(A::Matrix, # discrete dynamics A \n",
    "                        B::Matrix, # discrete dynamics B \n",
    "                        X_ref::Vector{Vector{Float64}}, # reference trajectory \n",
    "                        x0::Vector, # initial condition \n",
    "                        xg::Vector, # goal state \n",
    "                        u_min::Vector, # lower bound on u \n",
    "                        u_max::Vector, # upper bound on u\n",
    "                        N::Int64, # length of trajectory \n",
    "                        )::Tuple{Vector{Vector{Float64}}, Vector{Vector{Float64}}} # return Xcvx,Ucvx\n",
    "    \n",
    "    # get our sizes for state and control\n",
    "    nx,nu = size(B)\n",
    "    \n",
    "    @assert size(A) == (nx, nx)\n",
    "    @assert length(x0) == nx \n",
    "    @assert length(xg) == nx \n",
    "        \n",
    "    # LQR cost\n",
    "    Q = diagm(ones(nx))\n",
    "    R = diagm(ones(nu))\n",
    "\n",
    "    # variables we are solving for\n",
    "    X = cvx.Variable(nx,N)\n",
    "    U = cvx.Variable(nu,N-1)\n",
    "\n",
    "    # TODO: implement cost\n",
    "    obj = 0 \n",
    "\n",
    "    # create problem with objective\n",
    "    prob = cvx.minimize(obj)\n",
    "\n",
    "    # TODO: add constraints with prob.constraints += \n",
    "\n",
    "    cvx.solve!(prob, ECOS.Optimizer; silent_solver = true)\n",
    "\n",
    "    X = X.value\n",
    "    U = U.value\n",
    "    \n",
    "    Xcvx = vec_from_mat(X)\n",
    "    Ucvx = vec_from_mat(U)\n",
    "    \n",
    "    return Xcvx, Ucvx\n",
    "end\n",
    "    \n",
    "    \n",
    "@testset \"convex trajopt\" begin \n",
    "\n",
    "    # create our discrete time model \n",
    "    dt = 1.0\n",
    "    A,B = create_dynamics(dt)\n",
    "\n",
    "    # get our sizes for state and control\n",
    "    nx,nu = size(B)\n",
    "\n",
    "    # initial and goal states\n",
    "    x0 = [-2;-4;2;0;0;.0]\n",
    "    xg = [0,-.68,3.05,0,0,0]\n",
    "\n",
    "    # bounds on U\n",
    "    u_max = 0.4*ones(3)\n",
    "    u_min = -u_max\n",
    "\n",
    "    # problem size and reference trajectory \n",
    "    N = 100 \n",
    "    t_vec = 0:dt:((N-1)*dt)\n",
    "    X_ref = desired_trajectory(x0,xg,N,dt) \n",
    "    \n",
    "    # solve convex trajectory optimization problem \n",
    "    X_cvx, U_cvx = convex_trajopt(A,B,X_ref, x0,xg,u_min,u_max,N)\n",
    "    \n",
    "    X_sim = [zeros(nx) for i = 1:N]\n",
    "    X_sim[1] = x0 \n",
    "    for i = 1:N-1 \n",
    "        X_sim[i+1] = A*X_sim[i] + B*U_cvx[i]\n",
    "    end\n",
    "\n",
    "    # -------------plotting/animation---------------------------\n",
    "    Xm = mat_from_vec(X_sim)\n",
    "    Um = mat_from_vec(U_cvx)\n",
    "    display(plot(t_vec,Xm[1:3,:]',title = \"Positions\",\n",
    "                 xlabel = \"time (s)\", ylabel = \"position (m)\",\n",
    "                 label = [\"x\" \"y\" \"z\"]))\n",
    "    display(plot(t_vec,Xm[4:6,:]',title = \"Velocities\",\n",
    "            xlabel = \"time (s)\", ylabel = \"velocity (m/s)\",\n",
    "                 label = [\"x\" \"y\" \"z\"]))\n",
    "    display(plot(t_vec[1:end-1],Um',title = \"Control\",\n",
    "            xlabel = \"time (s)\", ylabel = \"thrust (N)\",\n",
    "                 label = [\"x\" \"y\" \"z\"]))\n",
    "\n",
    "    \n",
    "    display(animate_rendezvous(X_sim, X_ref, dt;show_reference = false))\n",
    "    # -------------plotting/animation---------------------------\n",
    "\n",
    "\n",
    "    @test maximum(norm.( X_sim .- X_cvx, Inf)) < 1e-3 \n",
    "    @test norm(X_sim[end] - xg) < 1e-3 # goal \n",
    "    xs=[x[1] for x in X_sim]\n",
    "    ys=[x[2] for x in X_sim]\n",
    "    zs=[x[3] for x in X_sim]\n",
    "    @test maximum(ys) <= (xg[2] + 1e-3)\n",
    "    @test maximum(zs) >= 4 # check to see if you did the circle \n",
    "    @test minimum(zs) <= 2 # check to see if you did the circle \n",
    "    @test maximum(xs) >= 1 # check to see if you did the circle \n",
    "    @test maximum(norm.(U_cvx,Inf)) <= 0.4 + 1e-3 # control constraints satisfied \n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7337ac55",
   "metadata": {},
   "source": [
    "## Part D (5 pts): Short answer\n",
    "\n",
    "1. List three reasons why an open loop policy wouldn't work well on a real system:\n",
    "- **put answer here**\n",
    "- **put answer here**\n",
    "- **put answer here**\n",
    "\n",
    "\n",
    "2. For convex trajectory optimization, give three examples of convex cost functions we can use:\n",
    "- **put answer here**\n",
    "- **put answer here**\n",
    "- **put answer here**\n",
    "\n",
    "\n",
    "3. List three things that convex trajectory optimization can do that LQR cannot:\n",
    "- **put answer here**\n",
    "- **put answer here**\n",
    "- **put answer here**\n",
    "\n",
    "4. Say we have the following convex trajectory optimization problem:\n",
    "$$ \\begin{align} \\min_{x_{1:N},u_{1:N-1}} \\quad & \\sum_{i=1}^{N-1} \\bigg[ \\frac{1}{2} (x_i - x_{ref, i})^TQ(x_i - x_{ref, i}) + \\frac{1}{2} u_i^TRu_i \\bigg] + \\frac{1}{2}(x_N- x_{ref, N})^TQ_f\n",
    "(x_N- x_{ref, 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",
    " & x_{min} \\leq x_i \\leq x_{max} \\quad \\text{for } i = 1,2,\\ldots,N\\\\ \n",
    " & u_{min} \\leq u_i \\leq u_{max} \\quad \\text{for } i = 1,2,\\ldots,N-1\n",
    " \\end{align}\n",
    " $$\n",
    " If the optimal solution to this problem does not violate any either the state or control bounds (the $x_{min} \\leq x_i \\leq x_{max}$ and $u_{min} \\leq u_i \\leq u_{max}$ constraints), how will it differ from the finite-horizon LQR solution?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2389dbb1",
   "metadata": {},
   "source": [
    "## Part E: Convex MPC (20 pts)\n",
    "\n",
    "In part C, we solved for the optimal rendezvous trajectory using convex optimization, and verified it by simulating it in an open loop fashion (no feedback). This was made possible because we assumed that our linear dynamics were exact, and that we had a perfect estimate of our state. In reality, there are many issues that would prevent this open loop policy from being successful.\n",
    "\n",
    "Together, these factors result in a \"sim to real\" gap between our simulated model, and the real model. Because there will always be a sim to real gap, we can't just execute open loop policies and expect them to be successful. What we can do, however, is use Model Predictive Control (MPC) that combines some of the ideas of feedback control with convex trajectory optimization. \n",
    "\n",
    "A convex MPC controller will set up and solve a convex optimization problem at each time step that incorporates the current state estimate as an initial condition. For a trajectory tracking problem like this rendezvous, we want to track $x_{ref}$, but instead of optimizing over the whole trajectory, we will only consider a sliding window of size $N_{mpc}$ (also called a horizon). If $N_{mpc} = 20$, this means our convex MPC controller is reasoning about the next 20 steps in the trajectory. This optimization problem at every timestep will start by taking the relevant reference trajectory at the current window from the current step $i$, to the end of the window $i + N_{mpc} - 1$. This slice of the reference trajectory that applies to the current MPC window will be called $\\tilde{x}_{ref} = x_{ref}[i,(i + N_{mpc} - 1)]$.\n",
    "\n",
    "\n",
    "\n",
    "$$ \\begin{align} \\min_{x_{1:N},u_{1:N-1}} \\quad & \\sum_{i=1}^{N-1} \\bigg[ \\frac{1}{2} (x_i - \\tilde{x}_{ref, i})^TQ({x}_i - \\tilde{x}_{ref, i}) + \\frac{1}{2} u_i^TRu_i \\bigg] + \\frac{1}{2}(x_N- \\tilde{x}_{ref, N})^TQ\n",
    "({x}_N- \\tilde{x}_{ref, 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",
    " & u_{min} \\leq u_i \\leq u_{max} \\quad \\text{for } i = 1,2,\\ldots,N-1 \\\\\n",
    " & x_i[2] \\leq x_{goal} [2]\\quad \\text{for } i = 1,2,\\ldots,N \n",
    " \\end{align}$$\n",
    " \n",
    "where $N$ in this case is $N_{mpc}$. This allows for the MPC controller to \"think\" about the future states in a way that the LQR controller cannot. By updating the reference trajectory window ($\\tilde{x}_{ref}$) at each step and updating the initial condition ($x_{IC}$), the MPC controller is able to \"react\" and compensate for the sim to real gap. \n",
    "\n",
    "You will now implement a function `convex_mpc` where you setup and solve this optimization problem at every timestep, and simply return $u_1$ from the solution. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b7956a8",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "`u = convex_mpc(A,B,X_ref_window,xic,xg,u_min,u_max,N_mpc)`\n",
    "\n",
    "setup and solve the above optimization problem, returning the \n",
    "first control u_1 from the solution (should be a length nu \n",
    "Vector{Float64}).  \n",
    "\"\"\"\n",
    "function convex_mpc(A::Matrix, # discrete dynamics matrix A\n",
    "                    B::Matrix, # discrete dynamics matrix B\n",
    "                    X_ref_window::Vector{Vector{Float64}}, # reference trajectory for this window \n",
    "                    xic::Vector, # current state x \n",
    "                    xg::Vector, # goal state \n",
    "                    u_min::Vector, # lower bound on u \n",
    "                    u_max::Vector, # upper bound on u \n",
    "                    N_mpc::Int64,  # length of MPC window (horizon)\n",
    "                    )::Vector{Float64} # return the first control command of the solved policy \n",
    "    \n",
    "    # get our sizes for state and control\n",
    "    nx,nu = size(B)\n",
    "    \n",
    "    # check sizes \n",
    "    @assert size(A) == (nx, nx)\n",
    "    @assert length(xic) == nx \n",
    "    @assert length(xg) == nx \n",
    "    @assert length(X_ref_window) == N_mpc \n",
    "        \n",
    "    # LQR cost\n",
    "    Q = diagm(ones(nx))\n",
    "    R = diagm(ones(nu))\n",
    "\n",
    "    # variables we are solving for\n",
    "    X = cvx.Variable(nx,N_mpc)\n",
    "    U = cvx.Variable(nu,N_mpc-1)\n",
    "\n",
    "    # TODO: implement cost function\n",
    "    obj = 0\n",
    "    \n",
    "\n",
    "    # create problem with objective\n",
    "    prob = cvx.minimize(obj)\n",
    "\n",
    "    # TODO: add constraints with prob.constraints += \n",
    "    \n",
    "\n",
    "    # solve problem \n",
    "    cvx.solve!(prob, ECOS.Optimizer; silent_solver = true)\n",
    "\n",
    "    # get X and U solutions \n",
    "    X = X.value\n",
    "    U = U.value\n",
    "    \n",
    "    # return first control U \n",
    "    return U[:,1]\n",
    "end\n",
    "        \n",
    "@testset \"convex mpc\" begin \n",
    "\n",
    "    # create our discrete time model \n",
    "    dt = 1.0\n",
    "    A,B = create_dynamics(dt)\n",
    "\n",
    "    # get our sizes for state and control\n",
    "    nx,nu = size(B)\n",
    "\n",
    "    # initial and goal states\n",
    "    x0 = [-2;-4;2;0;0;.0]\n",
    "    xg = [0,-.68,3.05,0,0,0]\n",
    "\n",
    "    # bounds on U\n",
    "    u_max = 0.4*ones(3)\n",
    "    u_min = -u_max\n",
    "\n",
    "    # problem size and reference trajectory \n",
    "    N = 100 \n",
    "    t_vec = 0:dt:((N-1)*dt)\n",
    "    X_ref = [desired_trajectory(x0,xg,N,dt)...,[xg for i = 1:N]...] \n",
    "    \n",
    "    # MPC window size \n",
    "    N_mpc = 20 \n",
    "    \n",
    "    # sim size and setup \n",
    "    N_sim = N + 20 \n",
    "    t_vec = 0:dt:((N_sim-1)*dt)\n",
    "    X_sim = [zeros(nx) for i = 1:N_sim]\n",
    "    X_sim[1] = x0 \n",
    "    U_sim = [zeros(nu) for i = 1:N_sim-1]\n",
    "    \n",
    "    # simulate \n",
    "    @showprogress \"simulating\" for i = 1:N_sim-1 \n",
    "        \n",
    "        # get state estimate\n",
    "        xi_estimate = state_estimate(X_sim[i], xg)\n",
    "        \n",
    "        # TODO: given a window of N_mpc timesteps, get current reference trajectory\n",
    "        X_ref_tilde = NaN\n",
    "        \n",
    "        # TODO: call convex mpc controller with state estimate \n",
    "        u_mpc = NaN\n",
    "        \n",
    "        # commanded control goes into thruster model where it gets modified \n",
    "        U_sim[i] = thruster_model(X_sim[i], xg, u_mpc)\n",
    "        \n",
    "        # simulate one step \n",
    "        X_sim[i+1] = A*X_sim[i] + B*U_sim[i]\n",
    "    end\n",
    "    \n",
    "    \n",
    "\n",
    "    # -------------plotting/animation---------------------------\n",
    "    Xm = mat_from_vec(X_sim)\n",
    "    Um = mat_from_vec(U_sim)\n",
    "    display(plot(t_vec,Xm[1:3,:]',title = \"Positions\",\n",
    "                 xlabel = \"time (s)\", ylabel = \"position (m)\",\n",
    "                 label = [\"x\" \"y\" \"z\"]))\n",
    "    display(plot(t_vec,Xm[4:6,:]',title = \"Velocities\",\n",
    "            xlabel = \"time (s)\", ylabel = \"velocity (m/s)\",\n",
    "                 label = [\"x\" \"y\" \"z\"]))\n",
    "    display(plot(t_vec[1:end-1],Um',title = \"Control\",\n",
    "            xlabel = \"time (s)\", ylabel = \"thrust (N)\",\n",
    "                 label = [\"x\" \"y\" \"z\"]))\n",
    "\n",
    "    \n",
    "    display(animate_rendezvous(X_sim, X_ref, dt;show_reference = false))\n",
    "    # -------------plotting/animation---------------------------\n",
    "\n",
    "    # tests \n",
    "    @test norm(X_sim[end] - xg) < 1e-3 # goal \n",
    "    xs=[x[1] for x in X_sim]\n",
    "    ys=[x[2] for x in X_sim]\n",
    "    zs=[x[3] for x in X_sim]\n",
    "    @test maximum(ys) <= (xg[2] + 1e-3)\n",
    "    @test maximum(zs) >= 4 # check to see if you did the circle \n",
    "    @test minimum(zs) <= 2 # check to see if you did the circle \n",
    "    @test maximum(xs) >= 1 # check to see if you did the circle \n",
    "    @test maximum(norm.(U_sim,Inf)) <= 0.4 + 1e-3 # control constraints satisfied \n",
    "\n",
    "end"
   ]
  }
 ],
 "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
}