notebooks/6. Optimization with JuMP.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization with JuMP\n",
    "\n",
    "JulaOpt is a collection of Julia optimization tools, and it's one of the best parts of the Julia ecosystem. Some of the packages in JuliaOpt are:\n",
    "\n",
    "![](img/juliaopt.png)\n",
    "\n",
    "For today, we're just going to show off JuMP, a modeling tool designed to make it easy to efficiently specify and solve optimizations in Julia. JuMP is similar to `yalmip` (in MATLAB), `pyomo` (in Python), or AMPL. It distinguishes itself by being fast and easy to use and taking advantage of Julia's expressiveness. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "using JuMP\n",
    "using Ipopt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating a model in JuMP is easy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$$ \\begin{alignat*}{1}\\min\\quad & 0\\\\\n",
       "\\text{Subject to} \\quad\\end{alignat*}\n",
       " $$"
      ],
      "text/plain": [
       "Feasibility problem with:\n",
       " * 0 linear constraints\n",
       " * 0 variables\n",
       "Solver is Ipopt"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Model(solver=IpoptSolver())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We add variables to the model with the `@variable` macro. Having a macro is useful because it can create variables inside the model and also, conveniently, create matching local variables in Julia:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$$ \\begin{alignat*}{1}\\min\\quad & 0\\\\\n",
       "\\text{Subject to} \\quad & x_{i} \\quad\\forall i \\in \\{1,2\\}\\\\\n",
       "\\end{alignat*}\n",
       " $$"
      ],
      "text/plain": [
       "Feasibility problem with:\n",
       " * 0 linear constraints\n",
       " * 2 variables\n",
       "Solver is Ipopt"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@variable(model, x[1:2])\n",
    "\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array{JuMP.Variable,1}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "typeof(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's add a simple objective:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$$ \\begin{alignat*}{1}\\min\\quad & x_{1}^2 + x_{2}^2\\\\\n",
       "\\text{Subject to} \\quad & x_{i} \\quad\\forall i \\in \\{1,2\\}\\\\\n",
       "\\end{alignat*}\n",
       " $$"
      ],
      "text/plain": [
       "Minimization problem with:\n",
       " * 0 linear constraints\n",
       " * 2 variables\n",
       "Solver is Ipopt"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@objective(model, Min, sum(x.^2))\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "******************************************************************************\n",
      "This program contains Ipopt, a library for large-scale nonlinear optimization.\n",
      " Ipopt is released as open source code under the Eclipse Public License (EPL).\n",
      "         For more information visit http://projects.coin-or.org/Ipopt\n",
      "******************************************************************************\n",
      "\n",
      "This is Ipopt version 3.12.4, running with linear solver mumps.\n",
      "NOTE: Other linear solvers might be more efficient (see Ipopt documentation).\n",
      "\n",
      "Number of nonzeros in equality constraint Jacobian...:        0\n",
      "Number of nonzeros in inequality constraint Jacobian.:        0\n",
      "Number of nonzeros in Lagrangian Hessian.............:        2\n",
      "\n",
      "Total number of variables............................:        2\n",
      "                     variables with only lower bounds:        0\n",
      "                variables with lower and upper bounds:        0\n",
      "                     variables with only upper bounds:        0\n",
      "Total number of equality constraints.................:        0\n",
      "Total number of inequality constraints...............:        0\n",
      "        inequality constraints with only lower bounds:        0\n",
      "   inequality constraints with lower and upper bounds:        0\n",
      "        inequality constraints with only upper bounds:        0\n",
      "\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "   0  0.0000000e+00 0.00e+00 0.00e+00  -1.0 0.00e+00    -  0.00e+00 0.00e+00   0\n",
      "\n",
      "Number of Iterations....: 0\n",
      "\n",
      "                                   (scaled)                 (unscaled)\n",
      "Objective...............:   0.0000000000000000e+00    0.0000000000000000e+00\n",
      "Dual infeasibility......:   0.0000000000000000e+00    0.0000000000000000e+00\n",
      "Constraint violation....:   0.0000000000000000e+00    0.0000000000000000e+00\n",
      "Complementarity.........:   0.0000000000000000e+00    0.0000000000000000e+00\n",
      "Overall NLP error.......:   0.0000000000000000e+00    0.0000000000000000e+00\n",
      "\n",
      "\n",
      "Number of objective function evaluations             = 1\n",
      "Number of objective gradient evaluations             = 1\n",
      "Number of equality constraint evaluations            = 0\n",
      "Number of inequality constraint evaluations          = 0\n",
      "Number of equality constraint Jacobian evaluations   = 0\n",
      "Number of inequality constraint Jacobian evaluations = 0\n",
      "Number of Lagrangian Hessian evaluations             = 0\n",
      "Total CPU secs in IPOPT (w/o function evaluations)   =      0.138\n",
      "Total CPU secs in NLP function evaluations           =      0.032\n",
      "\n",
      "EXIT: Optimal Solution Found.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       ":Optimal"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solve(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can retrieve the solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2-element Array{Float64,1}:\n",
       " 0.0\n",
       " 0.0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "getvalue(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's a pretty boring model. Let's constrain the optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[33mWARNING: \u001b[39m\u001b[22m\u001b[33mSolver does not appear to support adding constraints to an existing model. JuMP's internal model will be discarded.\u001b[39m\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$$ \\begin{alignat*}{1}\\min\\quad & x_{1}^2 + x_{2}^2\\\\\n",
       "\\text{Subject to} \\quad & x_{1} - 2 x_{2} \\geq 0\\\\\n",
       " & x_{2} \\geq 1\\\\\n",
       " & x_{i} \\quad\\forall i \\in \\{1,2\\}\\\\\n",
       "\\end{alignat*}\n",
       " $$"
      ],
      "text/plain": [
       "Minimization problem with:\n",
       " * 2 linear constraints\n",
       " * 2 variables\n",
       "Solver is Ipopt"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@constraints model begin\n",
    "    x[1] >= 2 * x[2]\n",
    "    x[2] >= 1\n",
    "end\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This is Ipopt version 3.12.4, running with linear solver mumps.\n",
      "NOTE: Other linear solvers might be more efficient (see Ipopt documentation).\n",
      "\n",
      "Number of nonzeros in equality constraint Jacobian...:        0\n",
      "Number of nonzeros in inequality constraint Jacobian.:        3\n",
      "Number of nonzeros in Lagrangian Hessian.............:        2\n",
      "\n",
      "Total number of variables............................:        2\n",
      "                     variables with only lower bounds:        0\n",
      "                variables with lower and upper bounds:        0\n",
      "                     variables with only upper bounds:        0\n",
      "Total number of equality constraints.................:        0\n",
      "Total number of inequality constraints...............:        2\n",
      "        inequality constraints with only lower bounds:        2\n",
      "   inequality constraints with lower and upper bounds:        0\n",
      "        inequality constraints with only upper bounds:        0\n",
      "\n",
      "iter    objective    inf_pr   inf_du lg(mu)  ||d||  lg(rg) alpha_du alpha_pr  ls\n",
      "   0  0.0000000e+00 1.00e+00 5.00e-01  -1.0 0.00e+00    -  0.00e+00 0.00e+00   0\n",
      "   1  5.3462199e+00 0.00e+00 1.00e-06  -1.0 2.08e+00    -  1.00e+00 1.00e+00h  1\n",
      "   2  5.0497709e+00 0.00e+00 2.00e-07  -1.7 7.02e-02    -  1.00e+00 1.00e+00f  1\n",
      "   3  5.0005613e+00 0.00e+00 1.50e-09  -3.8 1.12e-02    -  1.00e+00 1.00e+00f  1\n",
      "   4  5.0000036e+00 0.00e+00 1.84e-11  -5.7 1.27e-04    -  1.00e+00 1.00e+00f  1\n",
      "   5  4.9999999e+00 0.00e+00 2.53e-14  -8.6 8.37e-07    -  1.00e+00 1.00e+00f  1\n",
      "\n",
      "Number of Iterations....: 5\n",
      "\n",
      "                                   (scaled)                 (unscaled)\n",
      "Objective...............:   4.9999998650132769e+00    4.9999998650132769e+00\n",
      "Dual infeasibility......:   2.5313084961453569e-14    2.5313084961453569e-14\n",
      "Constraint violation....:   0.0000000000000000e+00    0.0000000000000000e+00\n",
      "Complementarity.........:   2.5066825129094524e-09    2.5066825129094524e-09\n",
      "Overall NLP error.......:   2.5066825129094524e-09    2.5066825129094524e-09\n",
      "\n",
      "\n",
      "Number of objective function evaluations             = 6\n",
      "Number of objective gradient evaluations             = 6\n",
      "Number of equality constraint evaluations            = 0\n",
      "Number of inequality constraint evaluations          = 6\n",
      "Number of equality constraint Jacobian evaluations   = 0\n",
      "Number of inequality constraint Jacobian evaluations = 6\n",
      "Number of Lagrangian Hessian evaluations             = 5\n",
      "Total CPU secs in IPOPT (w/o function evaluations)   =      0.002\n",
      "Total CPU secs in NLP function evaluations           =      0.125\n",
      "\n",
      "EXIT: Optimal Solution Found.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "2-element Array{Float64,1}:\n",
       " 2.0\n",
       " 1.0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solve(model)\n",
    "getvalue(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That was easy. But it's still not a very interesting optimization. How about something more robotics-y? \n",
    "\n",
    "## Model-Predictive Control in JuMP\n",
    "\n",
    "Let's implement a very simple model-predictive control (MPC) optimization in JuMP. Specifically, we'll write an optimization that tries to find a sequence of control inputs for a very simple model robot in order to optimize an objective function. \n",
    "\n",
    "More concretely, our model will be a brick, sliding frictionlessly in two dimensions. Our input will be the acceleration of the brick, and we'll try to minimize the brick's final velocity and distance from the origin. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       ":Optimal"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Model(solver=IpoptSolver(print_level=0))\n",
    "\n",
    "# Define our constant parameters\n",
    "Δt = 0.1\n",
    "num_time_steps = 20\n",
    "max_acceleration = 0.5\n",
    "\n",
    "# Define our decision variables\n",
    "@variables model begin\n",
    "    position[1:2, 1:num_time_steps]\n",
    "    velocity[1:2, 1:num_time_steps]\n",
    "    -max_acceleration <= acceleration[1:2, 1:num_time_steps] <= max_acceleration\n",
    "end\n",
    "\n",
    "# Add dynamics constraints\n",
    "@constraint(model, [i=2:num_time_steps, j=1:2],\n",
    "            velocity[j, i] == velocity[j, i - 1] + acceleration[j, i - 1] * Δt)\n",
    "@constraint(model, [i=2:num_time_steps, j=1:2],\n",
    "            position[j, i] == position[j, i - 1] + velocity[j, i - 1] * Δt)\n",
    "\n",
    "# Cost function: minimize final position and final velocity\n",
    "@objective(model, Min, \n",
    "    100 * sum(position[:, end].^2) + sum(velocity[:, end].^2))\n",
    "\n",
    "# Initial conditions:\n",
    "@constraint(model, position[:, 1] .== [1, 0])\n",
    "@constraint(model, velocity[:, 1] .== [0, -1])\n",
    "\n",
    "solve(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×20 Array{Float64,2}:\n",
       " -0.5  -0.5  -0.5  -0.5  -0.5  -0.5  …  -0.5  -0.5  -0.5  -0.5  0.5  0.0\n",
       "  0.5   0.5   0.5   0.5   0.5   0.5      0.5   0.5   0.5   0.5  0.5  0.0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Extract the solution from the model\n",
    "q = getvalue(position)\n",
    "v = getvalue(velocity)\n",
    "u = getvalue(acceleration)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Drawing the Result\n",
    "\n",
    "We can draw the output of the optimization using the `Plots.jl` package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Plots.GRBackend()"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Plots\n",
    "# Use the GR backend for Plots.jl, because it's fast\n",
    "gr()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[36mINFO: \u001b[39m\u001b[22m\u001b[36mSaved animation to /Users/rdeits/Downloads/agile-test/packages/v0.6/AgileRoboticsTutorial/notebooks/img/mpc1.gif\n",
      "\u001b[39m"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"img/mpc1.gif?0.4331348501284147>\" />"
      ],
      "text/plain": [
       "Plots.AnimatedGif(\"/Users/rdeits/Downloads/agile-test/packages/v0.6/AgileRoboticsTutorial/notebooks/img/mpc1.gif\")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The @animate macro creates an animated plot, which lets us draw the\n",
    "# optimized trajectory of the brick as a function of time\n",
    "anim = @animate for i = 1:num_time_steps\n",
    "    plot(q[1, :], q[2, :], xlim=(-1.1, 1.1), ylim=(-1.1, 1.1))\n",
    "    plot!([q[1, i]], [q[2, i]], marker=(:hex, 6))\n",
    "end\n",
    "\n",
    "# The gif() function saves our animated plot to an animated gif\n",
    "# Note: this may require you to have the `ffmpeg` program installed.\n",
    "# On Ubuntu 16.04, you can get this with `sudo apt-get install ffmpeg`.\n",
    "gif(anim, \"img/mpc1.gif\", fps = 30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the MPC Controller\n",
    "\n",
    "In a real application, we wouldn't just run the MPC optimization once. Instead, we might run the optimization at every time step using the robot's current state. \n",
    "\n",
    "To do that, let's wrap the MPC problem in a function called `run_mpc()` that takes the robot's current position and velocity as input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "run_mpc (generic function with 1 method)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solver = IpoptSolver(print_level=0)\n",
    "\n",
    "# run_mpc() takes the robot's current position and velocity\n",
    "# and returns an optimized trajectory of position, velocity, \n",
    "# and acceleration. \n",
    "function run_mpc(initial_position, initial_velocity)\n",
    "    \n",
    "    model = Model(solver=solver)\n",
    "\n",
    "    Δt = 0.1\n",
    "    num_time_steps = 10\n",
    "    max_acceleration = 0.5\n",
    "\n",
    "    @variables model begin\n",
    "        position[1:2, 1:num_time_steps]\n",
    "        velocity[1:2, 1:num_time_steps]\n",
    "        -max_acceleration <= acceleration[1:2, 1:num_time_steps] <= max_acceleration\n",
    "    end\n",
    "\n",
    "    # Dynamics constraints\n",
    "    @constraint(model, [i=2:num_time_steps, j=1:2],\n",
    "                velocity[j, i] == velocity[j, i - 1] + acceleration[j, i - 1] * Δt)\n",
    "    @constraint(model, [i=2:num_time_steps, j=1:2],\n",
    "                position[j, i] == position[j, i - 1] + velocity[j, i - 1] * Δt)\n",
    "\n",
    "    # Cost function: minimize final position and final velocity\n",
    "    @objective(model, Min, \n",
    "        100 * sum(position[:, end].^2) + sum(velocity[:, end].^2))\n",
    "\n",
    "    # Initial conditions:\n",
    "    @constraint(model, position[:, 1] .== initial_position)\n",
    "    @constraint(model, velocity[:, 1] .== initial_velocity)\n",
    "\n",
    "    solve(model)\n",
    "    return getvalue(position), getvalue(velocity), getvalue(acceleration)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can demonstrate this by repeatedly running the MPC program, applying its planned acceleration to the brick, and then simulating one step forward in time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[36mINFO: \u001b[39m\u001b[22m\u001b[36mSaved animation to /Users/rdeits/Downloads/agile-test/packages/v0.6/AgileRoboticsTutorial/notebooks/img/mpc2.gif\n",
      "\u001b[39m"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"img/mpc2.gif?0.28887777217280775>\" />"
      ],
      "text/plain": [
       "Plots.AnimatedGif(\"/Users/rdeits/Downloads/agile-test/packages/v0.6/AgileRoboticsTutorial/notebooks/img/mpc2.gif\")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The robot's starting position and velocity\n",
    "q = [1.0, 0.0]\n",
    "v = [0.0, -1.0]\n",
    "\n",
    "anim = @animate for i in 1:80\n",
    "    # Plot the current position\n",
    "    plot([q[1]], [q[2]], marker=(:hex, 10), xlim=(-1.1, 1.1), ylim=(-1.1, 1.1))\n",
    "    \n",
    "    # Run the MPC control optimization\n",
    "    q_plan, v_plan, u_plan = run_mpc(q, v)\n",
    "    \n",
    "    # Draw the planned future states from the MPC optimization\n",
    "    plot!(q_plan[1, :], q_plan[2, :], linewidth=5)\n",
    "    \n",
    "    # Apply the planned acceleration and simulate one step in time\n",
    "    u = u_plan[:, 1]\n",
    "    v += u * Δt\n",
    "    q += v * Δt\n",
    "end\n",
    "gif(anim, \"img/mpc2.gif\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 0.6.2",
   "language": "julia",
   "name": "julia-0.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "0.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}