Q2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "becc9e8f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import Pkg\n",
    "Pkg.activate(@__DIR__)\n",
    "Pkg.instantiate()\n",
    "\n",
    "import MathOptInterface as MOI\n",
    "import Ipopt \n",
    "import ForwardDiff as FD \n",
    "import Convex as cvx \n",
    "import ECOS\n",
    "using LinearAlgebra\n",
    "using Plots\n",
    "using Random\n",
    "using JLD2\n",
    "using Test\n",
    "import MeshCat as mc \n",
    "using Printf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56525960",
   "metadata": {},
   "source": [
    "# Q2: iLQR (30 pts)\n",
    "\n",
    "In this problem, we are going to use iLQR to solve a trajectory optimization for a 6DOF quadrotor. This problem we will use a cost function to motivate the quadrotor to follow a specified aerobatic manuever. The continuous time dynamics of the quadrotor are detailed in `quadrotor.jl`, with the state being the following:\n",
    "\n",
    "$$ x = [r,v,{}^Np^B{},\\omega] $$ where $r\\in\\mathbb{R}^3$ is the position of the quadrotor in the world frame (N), $v\\in\\mathbb{R}^3$ is the velocity of the quadrotor in the world frame (N),  $^Np^B\\in\\mathbb{R}^3$ is the Modified Rodrigues Parameter (MRP) that is used to denote the attitude of the quadrotor, and  $\\omega\\in\\mathbb{R}^3$ is the angular velocity of the quadrotor expressed in the body frame (B).  By denoting the attitude of the quadrotor with a MRP instead of a quaternion or rotation matrix, we have to be careful to avoid any scenarios where the MRP will approach it's singularity at 360 degrees of rotation. For the manuever planned in this problem, the MRP will be sufficient. \n",
    "\n",
    "The dynamics of the quadrotor are discretized with `rk4`, resulting in the following discrete time dynamics function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c774306",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "include(joinpath(@__DIR__, \"utils\",\"quadrotor.jl\"))\n",
    "\n",
    "function discrete_dynamics(params::NamedTuple, x::Vector, u, k)\n",
    "    # discrete dynamics\n",
    "    # x - state \n",
    "    # u - control \n",
    "    # k - index of trajectory \n",
    "    # dt comes from params.model.dt \n",
    "    return rk4(params.model, quadrotor_dynamics, x, u, params.model.dt)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2ca97a2",
   "metadata": {},
   "source": [
    "## Part A: iLQR for a quadrotor (25 pts)\n",
    "\n",
    "iLQR is used to solve optimal control problems of the following form:\n",
    "$$ \\begin{align} \\min_{x_{1:N},u_{1:N-1}} \\quad & \\bigg[ \\sum_{i=1}^{N-1} \\ell(x_i,u_i)\\bigg] + \\ell_N(x_N)\\\\ \n",
    " \\text{st} \\quad & x_1 = x_{{IC}} \\\\ \n",
    " & x_{k+1} = f(x_k, u_k) \\quad \\text{for } i = 1,2,\\ldots,N-1 \\\\\n",
    " \\end{align}$$\n",
    " where $x_{IC}$ is the inital condition, $x_{k+1} = f(x_k, u_k)$ is the discrete dynamics function, $\\ell(x_i,u_i)$ is the stage cost, and $\\ell_N(x_N)$ is the terminal cost.  Since this optimization problem can be non-convex, there is no guarantee of convergence to a global optimum, or even convergene rates to a local optimum, but in practice we will see that it can work very well. \n",
    " \n",
    "For this problem, we are going to use a simple cost function consisting of the following stage cost:\n",
    "\n",
    "$$ \\ell(x_i,u_i) = \\frac{1}{2} (x_i - x_{ref,i})^TQ(x_i - x_{ref,i}) + \\frac{1}{2}(u_i - u_{ref,i})^TR(u_i - u_{ref,i}) $$\n",
    "\n",
    "And the following terminal cost:\n",
    "\n",
    "$$ \\ell_N(x_N) = \\frac{1}{2}(x_N - x_{ref,N})^TQ_f(x_N - x_{ref,N}) $$\n",
    "\n",
    "This is how we will encourange our quadrotor to track a reference trajectory $x_{ref}$. In the following sections, you will implement `iLQR` and use it inside of a `solve_quadrotor_trajectory` function. Below we have included some starter code, but you are free to use/not use any of the provided functions so long as you pass the tests. \n",
    "\n",
    "We will consider iLQR to have converged when $\\Delta J < \\text{atol}$ as calculated during the backwards pass. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dff216c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# starter code: feel free to use or not use \n",
    "\n",
    "function stage_cost(p::NamedTuple,x::Vector,u::Vector,k::Int)\n",
    "    # TODO: return stage cost at time step k \n",
    "    \n",
    "end\n",
    "function term_cost(p::NamedTuple,x)\n",
    "    # TODO: return terminal cost\n",
    "\n",
    "end\n",
    "function stage_cost_expansion(p::NamedTuple, x::Vector, u::Vector, k::Int)\n",
    "    # TODO: return stage cost expansion\n",
    "    # if the stage cost is J(x,u), you can return the following\n",
    "    # ∇ₓ²J, ∇ₓJ, ∇ᵤ²J, ∇ᵤJ\n",
    "\n",
    "end\n",
    "function term_cost_expansion(p::NamedTuple, x::Vector)\n",
    "    # TODO: return terminal cost expansion\n",
    "    # if the terminal cost is Jn(x,u), you can return the following\n",
    "    # ∇ₓ²Jn, ∇ₓJn\n",
    "\n",
    "end\n",
    "function backward_pass(params::NamedTuple,          # useful params \n",
    "                       X::Vector{Vector{Float64}},  # state trajectory \n",
    "                       U::Vector{Vector{Float64}})  # control trajectory \n",
    "    # compute the iLQR backwards pass given a dynamically feasible trajectory X and U\n",
    "    # return d, K, ΔJ  \n",
    "    \n",
    "    # outputs:\n",
    "    #     d  - Vector{Vector} feedforward control  \n",
    "    #     K  - Vector{Matrix} feedback gains \n",
    "    #     ΔJ - Float64        expected decrease in cost \n",
    "    \n",
    "    nx, nu, N = params.nx, params.nu, params.N \n",
    "    \n",
    "    # vectors of vectors/matrices for recursion \n",
    "    P = [zeros(nx,nx) for i = 1:N]   # cost to go quadratic term\n",
    "    p = [zeros(nx)    for i = 1:N]   # cost to go linear term\n",
    "    d = [zeros(nu)    for i = 1:N-1] # feedforward control\n",
    "    K = [zeros(nu,nx) for i = 1:N-1] # feedback gain\n",
    "\n",
    "    # TODO: implement backwards pass and return d, K, ΔJ \n",
    "    N = params.N\n",
    "    ΔJ = 0.0\n",
    "\n",
    "    return d, K, ΔJ\n",
    "end\n",
    "\n",
    "function trajectory_cost(params::NamedTuple,          # useful params \n",
    "                         X::Vector{Vector{Float64}},  # state trajectory \n",
    "                         U::Vector{Vector{Float64}}) # control trajectory \n",
    "    # compute the trajectory cost for trajectory X and U (assuming they are dynamically feasible)\n",
    "    N = params.N\n",
    "    \n",
    "    # TODO: add trajectory cost \n",
    "    return 0 \n",
    "end\n",
    "\n",
    "function forward_pass(params::NamedTuple,           # useful params \n",
    "                      X::Vector{Vector{Float64}},   # state trajectory \n",
    "                      U::Vector{Vector{Float64}},   # control trajectory \n",
    "                      d::Vector{Vector{Float64}},   # feedforward controls \n",
    "                      K::Vector{Matrix{Float64}};   # feedback gains\n",
    "                      max_linesearch_iters = 20)    # max iters on linesearch \n",
    "    # forward pass in iLQR with linesearch \n",
    "    # use a line search where the trajectory cost simply has to decrease (no Armijo)\n",
    "    \n",
    "    # outputs:\n",
    "    #     Xn::Vector{Vector}  updated state trajectory  \n",
    "    #     Un::Vector{Vector}  updated control trajectory \n",
    "    #     J::Float64          updated cost  \n",
    "    #     α::Float64.         step length \n",
    "\n",
    "    nx, nu, N = params.nx, params.nu, params.N \n",
    "    \n",
    "    Xn = [zeros(nx) for i = 1:N]      # new state history \n",
    "    Un = [zeros(nu) for i = 1:N-1]    # new control history \n",
    "    \n",
    "    # initial condition \n",
    "    Xn[1] = 1*X[1]\n",
    "    \n",
    "    # initial step length \n",
    "    α = 1.0\n",
    "    \n",
    "    # TODO: add forward pass \n",
    "    for i = 1:max_linesearch_iters\n",
    "\n",
    "\n",
    "    end\n",
    "\n",
    "    error(\"forward pass failed\")\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fea1b01c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "function iLQR(params::NamedTuple,         # useful params for costs/dynamics/indexing \n",
    "              x0::Vector,                 # initial condition \n",
    "              U::Vector{Vector{Float64}}; # initial controls \n",
    "              atol=1e-3,                  # convergence criteria: ΔJ < atol \n",
    "              max_iters = 250,            # max iLQR iterations \n",
    "              verbose = true)             # print logging\n",
    "    \n",
    "    # iLQR solver given an initial condition x0, initial controls U, and a \n",
    "    # dynamics function described by `discrete_dynamics`\n",
    "    \n",
    "    # return (X, U, K) where \n",
    "    # outputs:\n",
    "    #     X::Vector{Vector} - state trajectory \n",
    "    #     U::Vector{Vector} - control trajectory \n",
    "    #     K::Vector{Matrix} - feedback gains K \n",
    "\n",
    "    # first check the sizes of everything\n",
    "    @assert length(U) == params.N-1\n",
    "    @assert length(U[1]) == params.nu\n",
    "    @assert length(x0) == params.nx \n",
    "\n",
    "    nx, nu, N = params.nx, params.nu, params.N\n",
    "\n",
    "    # TODO: initial rollout\n",
    "\n",
    "\n",
    "    for ilqr_iter = 1:max_iters\n",
    "\n",
    "        \n",
    "        \n",
    "        # termination criteria \n",
    "        if ΔJ < atol \n",
    "            if verbose \n",
    "                @info \"iLQR converged\"\n",
    "            end\n",
    "            return X, U, K \n",
    "        end\n",
    "        \n",
    "        # ---------------logging -------------------\n",
    "        if verbose\n",
    "            dmax = maximum(norm.(d))\n",
    "            if rem(ilqr_iter-1,10)==0\n",
    "                @printf \"iter     J           ΔJ        |d|         α         \\n\"\n",
    "                @printf \"-------------------------------------------------\\n\"\n",
    "            end\n",
    "            @printf(\"%3d   %10.3e  %9.2e  %9.2e  %6.4f    \\n\",\n",
    "              ilqr_iter, J, ΔJ, dmax, α)\n",
    "        end\n",
    "        \n",
    "    end\n",
    "    error(\"iLQR failed\")\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f35b3036",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "function create_reference(N, dt)\n",
    "    # create reference trajectory for quadrotor \n",
    "    R = 6\n",
    "    Xref = [ [R*cos(t);R*cos(t)*sin(t);1.2 + sin(t);zeros(9)] for t = range(-pi/2,3*pi/2, length = N)]\n",
    "    for i = 1:(N-1)\n",
    "        Xref[i][4:6] = (Xref[i+1][1:3] - Xref[i][1:3])/dt\n",
    "    end\n",
    "    Xref[N][4:6] = Xref[N-1][4:6]\n",
    "    Uref = [(9.81*0.5/4)*ones(4) for i = 1:(N-1)]\n",
    "    return Xref, Uref\n",
    "end\n",
    "function solve_quadrotor_trajectory(;verbose = true)\n",
    "    \n",
    "    # problem size \n",
    "    nx = 12\n",
    "    nu = 4\n",
    "    dt = 0.05 \n",
    "    tf = 5 \n",
    "    t_vec = 0:dt:tf \n",
    "    N = length(t_vec)\n",
    "\n",
    "    # create reference trajectory \n",
    "    Xref, Uref = create_reference(N, dt)\n",
    "    \n",
    "    # tracking cost function\n",
    "    Q = 1*diagm([1*ones(3);.1*ones(3);1*ones(3);.1*ones(3)])\n",
    "    R = .1*diagm(ones(nu))\n",
    "    Qf = 10*Q \n",
    "\n",
    "    # dynamics parameters (these are estimated)\n",
    "    model = (mass=0.5,\n",
    "            J=Diagonal([0.0023, 0.0023, 0.004]),\n",
    "            gravity=[0,0,-9.81],\n",
    "            L=0.1750,\n",
    "            kf=1.0,\n",
    "            km=0.0245,dt = dt)\n",
    "\n",
    "    \n",
    "    # the params needed by iLQR \n",
    "    params = (\n",
    "        N = N, \n",
    "        nx = nx, \n",
    "        nu = nu, \n",
    "        Xref = Xref, \n",
    "        Uref = Uref, \n",
    "        Q = Q, \n",
    "        R = R, \n",
    "        Qf = Qf, \n",
    "        model = model\n",
    "    )\n",
    "\n",
    "    # initial condition \n",
    "    x0 = 1*Xref[1]\n",
    "    \n",
    "    # initial guess controls \n",
    "    U = [(uref + .0001*randn(nu)) for uref in Uref]\n",
    "    \n",
    "    # solve with iLQR\n",
    "    X, U, K = iLQR(params,x0,U;atol=1e-4,max_iters = 250,verbose = verbose)\n",
    "    \n",
    "    return X, U, K, t_vec, params\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b09022d8",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "@testset \"ilqr\" begin \n",
    "\n",
    "    # NOTE: set verbose to true here when you submit\n",
    "    Xilqr, Uilqr, Kilqr, t_vec, params =  solve_quadrotor_trajectory(verbose = true)\n",
    "    \n",
    "    # ---------------testing------------------\n",
    "    Usol = load(joinpath(@__DIR__,\"utils\",\"ilqr_U.jld2\"))[\"Usol\"]\n",
    "    @test maximum(norm.(Usol .- Uilqr,Inf)) <= 1e-2 \n",
    "    \n",
    "    # ---------------plotting------------------\n",
    "    Xm = hcat(Xilqr...)\n",
    "    Um = hcat(Uilqr...)\n",
    "    display(plot(t_vec, Xm[1:3,:]', xlabel = \"time (s)\", ylabel = \"position (m)\",\n",
    "                                   title = \"Position\", label = [\"r_x\" \"r_y\" \"r_z\"]))\n",
    "    display(plot(t_vec, Xm[4:6,:]', xlabel = \"time (s)\", ylabel = \"velocity (m/s)\",\n",
    "                                   title = \"Velocity\", label = [\"v_x\" \"v_y\" \"v_z\"]))\n",
    "    display(plot(t_vec, Xm[7:9,:]', xlabel = \"time (s)\", ylabel = \"MRP\",\n",
    "                                   title = \"Attitude (MRP)\", label = [\"p_x\" \"p_y\" \"p_z\"]))\n",
    "    display(plot(t_vec, Xm[10:12,:]', xlabel = \"time (s)\", ylabel = \"angular velocity (rad/s)\",\n",
    "                                   title = \"Angular Velocity\", label = [\"ω_x\" \"ω_y\" \"ω_z\"]))\n",
    "    display(plot(t_vec[1:end-1], Um', xlabel = \"time (s)\", ylabel = \"rotor speeds (rad/s)\",\n",
    "                                   title = \"Controls\", label = [\"u_1\" \"u_2\" \"u_3\" \"u_4\"]))\n",
    "    display(animate_quadrotor(Xilqr, params.Xref, params.model.dt))\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fca16374",
   "metadata": {},
   "source": [
    "## Part B: Tracking solution with TVLQR (5 pts)\n",
    "\n",
    "Here we will do the same thing we did in Q1 where we take a trajectory from a trajectory optimization solver, and track it with TVLQR to account for some model mismatch. In DIRCOL, we had to explicitly compute the TVLQR control gains, but in iLQR, we get these same gains out of the algorithmn as the K's. Use these to track the quadrotor through this manuever. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57b6bc43",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "@testset \"iLQR with model error\" begin \n",
    "\n",
    "    # set verbose to false when you submit \n",
    "    Xilqr, Uilqr, Kilqr, t_vec, params =  solve_quadrotor_trajectory(verbose = false)\n",
    "    \n",
    "    # real model parameters for dynamics \n",
    "    model_real = (mass=0.5,\n",
    "            J=Diagonal([0.0025, 0.002, 0.0045]),\n",
    "            gravity=[0,0,-9.81],\n",
    "            L=0.1550,\n",
    "            kf=0.9,\n",
    "            km=0.0365,dt = 0.05)\n",
    "    \n",
    "    # simulate closed loop system \n",
    "    nx, nu, N = params.nx, params.nu, params.N\n",
    "    Xsim = [zeros(nx) for i = 1:N]\n",
    "    Usim = [zeros(nx) for i = 1:(N-1)]\n",
    "    \n",
    "    # initial condition \n",
    "    Xsim[1] = 1*Xilqr[1]\n",
    "    \n",
    "    # TODO: simulate with closed loop control \n",
    "    for i = 1:(N-1) \n",
    "\n",
    "        Xsim[i+1] = rk4(model_real, quadrotor_dynamics, Xsim[i], Usim[i], model_real.dt)\n",
    "    end\n",
    "    \n",
    "    \n",
    "    # -----------------testing---------------------\n",
    "    @test 1e-6 <= norm(Xilqr[50] - Xsim[50],Inf) <= .3\n",
    "    @test 1e-6 <= norm(Xilqr[end] - Xsim[end],Inf) <= .3\n",
    "    \n",
    "    # -----------------plotting---------------------\n",
    "    Xm = hcat(Xsim...)\n",
    "    Um = hcat(Usim...)\n",
    "    Xilqrm = hcat(Xilqr...)\n",
    "    Uilqrm = hcat(Uilqr...)\n",
    "    plot(t_vec,Xilqrm[1:3,:]',ls=:dash, label = \"\",lc = [:red :green :blue])\n",
    "    display(plot!(t_vec,Xm[1:3,:]',title = \"Position (-- is iLQR reference)\",\n",
    "                 xlabel = \"time(s)\", ylabel = \"position (m)\",\n",
    "                 label = [\"r_x\" \"r_y\" \"r_z\"],lc = [:red :green :blue]))\n",
    "    \n",
    "    plot(t_vec,Xilqrm[7:9,:]',ls=:dash, label = \"\",lc = [:red :green :blue])\n",
    "    display(plot!(t_vec,Xm[7:9,:]',title = \"Attitude (-- is iLQR reference)\",\n",
    "                 xlabel = \"time(s)\", ylabel = \"Attitude (MRP)\",\n",
    "                 label = [\"p_x\" \"p_y\" \"p_z\"],lc = [:red :green :blue]))\n",
    "    \n",
    "    display(animate_quadrotor(Xilqr, params.Xref, params.model.dt))\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
}