Code_for_lecture_6.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "Dependencies:\n",
    "- System: python3\n",
    "- Python: jupyter, numpy, matplotlib, scipy, jax, cvxpy\n",
    "\n",
    "Example setup for a Ubuntu system (Mac users, maybe `brew` instead of `sudo apt`; Windows users, learn to love [WSL](https://docs.microsoft.com/en-us/windows/wsl/install-win10)):\n",
    "```\n",
    "/usr/bin/python3 -m pip install --upgrade pip\n",
    "pip install --upgrade jupyter numpy matplotlib scipy jax cvxpy\n",
    "jupyter notebook  # from the directory of this notebook\n",
    "```\n",
    "Alternatively, view this notebook on [Google Colab](https://colab.research.google.com/github/StanfordASL/AA203-Examples/blob/master/Lecture-11/Zermelo's%20Problem.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cvxpy as cp\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize import minimize, Bounds\n",
    "\n",
    "# Constants used in plots below\n",
    "t_f, N = 10, 20\n",
    "h = t_f / N"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simultaneous Method (state and control decision variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def zermelo_simultaneous(u_max=1.0, z0=None, verbose=True):\n",
    "\n",
    "    # Problem parameters\n",
    "    t_f = 10\n",
    "    M = 10\n",
    "    ℓ = 5\n",
    "    v = 1\n",
    "    flow = lambda y: 0.35 / (ℓ**2 / 4) * y * (ℓ - y)\n",
    "\n",
    "    # Discretization\n",
    "    N = 20\n",
    "    h = t_f / N\n",
    "\n",
    "    # Decision variables\n",
    "    # z = np.concatenate([x, y, u])\n",
    "    get_x = lambda z: z[:N + 1]\n",
    "    get_y = lambda z: z[N + 1:-N]\n",
    "    get_u = lambda z: z[-N:]\n",
    "\n",
    "    # Set up problem and `minimize`\n",
    "    eps = 1e-3  # `minimize` can be a bit finicky about constraints, so slightly loosen state/control box constraints\n",
    "    bounds = Bounds(\n",
    "        np.concatenate([np.zeros(N + 1), np.zeros(N + 1), -u_max * np.ones(N)]) - eps,\n",
    "        np.concatenate([M * np.ones(N + 1), ℓ * np.ones(N + 1), u_max * np.ones(N)]) + eps)\n",
    "\n",
    "    cost = lambda z: h * np.sum(np.square(get_u(z)))\n",
    "\n",
    "    def constraints(z):\n",
    "        return np.concatenate([\n",
    "            get_x(z)[1:] - get_x(z)[:-1] - h * (v * np.cos(get_u(z)) + flow(get_y(z)[:-1])),  # x dynamics constraint\n",
    "            get_y(z)[1:] - get_y(z)[:-1] - h * v * np.sin(get_u(z)),                          # y dynamics constraint\n",
    "            get_x(z)[[0, -1]] - np.array([0., M]),                                            # x initial/terminal constraint\n",
    "            get_y(z)[[0, -1]] - np.array([0., ℓ])                                             # y initial/terminal constraint\n",
    "        ])\n",
    "\n",
    "    if z0 is None:\n",
    "        z0 = np.concatenate([np.ones(N + 1), np.ones(N + 1), 0.5 * u_max * np.ones(N)])\n",
    "        # z0 = np.concatenate([np.linspace(0, M, N + 1), np.linspace(0, ℓ, N + 1), 0.5 * u_max * np.zeros(N)])\n",
    "    result = minimize(cost,\n",
    "                      z0,\n",
    "                      bounds=bounds,\n",
    "                      constraints={\n",
    "                          'type': 'eq',\n",
    "                          'fun': constraints\n",
    "                      },\n",
    "                      options={'maxiter': 1000})\n",
    "    if verbose:\n",
    "        print(result)\n",
    "\n",
    "    return get_x(result.x), get_y(result.x), get_u(result.x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_opt, y_opt, u_opt = zermelo_simultaneous(u_max=1.)\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(x_opt, y_opt)\n",
    "plt.title(\"Optimal Trajectory\", fontsize=24)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Optimal Control\", fontsize=24)\n",
    "plt.plot(np.arange(N) * h, u_opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "x_opt, y_opt, u_opt = zermelo_simultaneous(1.0, verbose=False)\n",
    "x_opt, y_opt, u_opt = zermelo_simultaneous(0.75, np.concatenate([x_opt, y_opt, u_opt]))\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(x_opt, y_opt)\n",
    "plt.title(\"Optimal Trajectory\", fontsize=24)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Optimal Control\", fontsize=24)\n",
    "plt.plot(np.arange(N) * h, u_opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Direct Shooting Method (control decision variables only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def zermelo_shooting(u_max=1.0, u0=None, verbose=True):\n",
    "\n",
    "    # Problem parameters\n",
    "    t_f = 10\n",
    "    M = 10\n",
    "    ℓ = 5\n",
    "    v = 1\n",
    "    flow = lambda y: 0.35 / (ℓ**2 / 4) * y * (ℓ - y)\n",
    "\n",
    "    # Discretization\n",
    "    N = 20\n",
    "    h = t_f / N\n",
    "    dynamics = lambda x, u: x + h * np.array([v * np.cos(u) + flow(x[..., 1]), v * np.sin(u)])\n",
    "\n",
    "    # Set up problem and `minimize`\n",
    "    eps = 1e-3\n",
    "    bounds = Bounds(-u_max * np.ones(N), u_max * np.ones(N))  # control constraints\n",
    "\n",
    "    cost = lambda u: h * np.sum(np.square(u))\n",
    "\n",
    "    def constraints(u):\n",
    "        state = np.zeros(2)\n",
    "        constraint_list = []\n",
    "        for ui in u:\n",
    "            state = dynamics(state, ui)\n",
    "            constraint_list.append(state)                     # state >= 0 (box constraint with below)\n",
    "            constraint_list.append(np.array([M, ℓ]) - state)  # state <= [M, ℓ]\n",
    "        constraint_list.append(state - np.array([M, ℓ]))      # terminal state >= [M, ℓ] as well (enforcing equality)\n",
    "        return np.concatenate(constraint_list)\n",
    "\n",
    "    if u0 is None:\n",
    "        u0 = 0.5 * u_max * np.ones(N)\n",
    "    result = minimize(cost,\n",
    "                      u0,\n",
    "                      bounds=bounds,\n",
    "                      constraints={\n",
    "                          'type': 'ineq',\n",
    "                          'fun': constraints\n",
    "                      },\n",
    "                      options={'maxiter': 1000})\n",
    "    if verbose:\n",
    "        print(result)\n",
    "\n",
    "    u_opt = result.x\n",
    "    states = [np.zeros(2)]\n",
    "    for ui in u_opt:\n",
    "        states.append(dynamics(states[-1], ui))\n",
    "    return [state[0] for state in states], [state[1] for state in states], u_opt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_opt, y_opt, u_opt = zermelo_shooting(1.0)\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(x_opt, y_opt)\n",
    "plt.title(\"Optimal Trajectory\", fontsize=24)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Optimal Control\", fontsize=24)\n",
    "plt.plot(np.arange(N) * h, u_opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_opt, y_opt, u_opt = zermelo_shooting(0.75)\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(x_opt, y_opt)\n",
    "plt.title(\"Optimal Trajectory\", fontsize=24)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Optimal Control\", fontsize=24)\n",
    "plt.plot(np.arange(N) * h, u_opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sequential Convex Programming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def zermelo_scp(u_max=1.0, verbose=True):\n",
    "\n",
    "    # Problem parameters\n",
    "    t_f = 10\n",
    "    M = 10\n",
    "    ℓ = 5\n",
    "    v = 1\n",
    "    flow = lambda y: 0.35 / (ℓ**2 / 4) * y * (ℓ - y)\n",
    "\n",
    "    # Discretization\n",
    "    N = 20\n",
    "    h = t_f / N\n",
    "    dynamics = lambda x, u, np=jnp: x + h * np.array([v * np.cos(u[0]) + flow(x[1]), v * np.sin(u[0])])\n",
    "\n",
    "    # Linearization\n",
    "    dynamics_jacobian = jax.jit(jax.jacfwd(dynamics, (0, 1)))\n",
    "\n",
    "    def linearized_dynamics(x_nom, u_nom, x, u):\n",
    "        # Slightly different interface compared to your homework to force you guys to think\n",
    "        f_x, f_u = map(np.array, dynamics_jacobian(x_nom, u_nom))\n",
    "        return dynamics(x_nom, u_nom, np=np) + f_x @ (x - x_nom) + f_u @ (u - u_nom)\n",
    "\n",
    "    # Set up problem and apply sequential convex programming\n",
    "    x0 = np.zeros(2)\n",
    "    xf = np.array([M, ℓ])\n",
    "    x_nom = np.linspace(x0, xf, N + 1)\n",
    "    u_nom = 0.6 * u_max * np.ones((N, 1))\n",
    "\n",
    "    max_absolute_difference = np.inf\n",
    "    for i in range(100):\n",
    "        x = cp.Variable((N + 1, 2))\n",
    "        u = cp.Variable((N, 1))\n",
    "        s = cp.Variable()\n",
    "        cost = cp.sum(h * cp.square(u))\n",
    "        constraints = [x[i + 1] == linearized_dynamics(x_nom[i], u_nom[i], x[i], u[i]) for i in range(N)]\n",
    "        constraints += [x[0] == x0, x[-1] == xf]                   # initial/terminal constraints\n",
    "        constraints += [x >= x0[np.newaxis], x <= xf[np.newaxis]]  # state box constraints\n",
    "        constraints += [cp.abs(u) <= u_max + s]                    # control constraints\n",
    "        constraints += [s >= 0]                                    # slack variable must be positive\n",
    "        # problem = cp.Problem(cp.Minimize(cost + cp.square(s)), constraints)\n",
    "        problem = cp.Problem(cp.Minimize(cost + s), constraints)\n",
    "\n",
    "        cost_value = problem.solve()\n",
    "        max_absolute_difference = max(np.max(np.abs(x_nom - x.value)), np.max(np.abs(u_nom - u.value)))\n",
    "        if verbose:\n",
    "            print(f\"Iteration {i:2}, cost = {cost_value:7.4f}, slack variable = {s.value:7.4f}, \"\n",
    "                  f\"max_absolute_difference = {max_absolute_difference}\")\n",
    "        converged = max_absolute_difference < 1e-3\n",
    "        x_nom = x.value\n",
    "        u_nom = u.value\n",
    "        if converged:\n",
    "            break\n",
    "\n",
    "    return x_nom[:, 0], x_nom[:, 1], u_nom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_opt, y_opt, u_opt = zermelo_scp()\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(x_opt, y_opt)\n",
    "plt.title(\"Optimal Trajectory\", fontsize=24)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Optimal Control\", fontsize=24)\n",
    "plt.plot(np.arange(N) * h, u_opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_opt, y_opt, u_opt = zermelo_scp(0.75)\n",
    "plt.figure(figsize=(15, 8))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(x_opt, y_opt)\n",
    "plt.title(\"Optimal Trajectory\", fontsize=24)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(\"Optimal Control\", fontsize=24)\n",
    "plt.plot(np.arange(N) * h, u_opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "AA203",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "e95adbe7dca8d9d43e38406dc053eb0871452db61fa85ecb56111a071905356c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}