trajopt.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook provides examples to go along with the [textbook](http://underactuated.csail.mit.edu/trajopt.html).  I recommend having both windows open, side-by-side!\n",
    "\n",
    "[Click here](http://underactuated.csail.mit.edu/drake.html#notebooks) for instructions on how to run the notebook on Deepnote and/or Google Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import time\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pydot\n",
    "from IPython.display import HTML, SVG, clear_output, display\n",
    "from pydrake.all import (Binding, Box, DiagramBuilder, DirectCollocation,\n",
    "                         DirectTranscription,\n",
    "                         FiniteHorizonLinearQuadraticRegulatorOptions,\n",
    "                         GraphOfConvexSets, GurobiSolver, HPolyhedron,\n",
    "                         L2NormCost, LinearSystem, LogVectorOutput,\n",
    "                         MakeFiniteHorizonLinearQuadraticRegulator,\n",
    "                         MakeFirstAvailableSolver, MathematicalProgram,\n",
    "                         MosekSolver, MultibodyPlant,\n",
    "                         MultibodyPositionToGeometryPose, Parser,\n",
    "                         PiecewisePolynomial, PlanarSceneGraphVisualizer,\n",
    "                         Point, PointCloud, Rgba, RigidTransform,\n",
    "                         RotationMatrix, SceneGraph, Simulator, Solve, Sphere,\n",
    "                         StartMeshcat, TrajectorySource, Variable, eq)\n",
    "from pydrake.examples.acrobot import AcrobotGeometry, AcrobotPlant\n",
    "from pydrake.examples.pendulum import PendulumPlant, PendulumState\n",
    "\n",
    "from underactuated import FindResource, running_as_notebook\n",
    "from underactuated.jupyter import AdvanceToAndVisualize\n",
    "from underactuated.meshcat_utils import draw_points, set_planar_viewpoint\n",
    "from underactuated.pendulum import PendulumVisualizer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer (run this cell only once, each instance consumes a port)\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trajectory optimization for the double integrator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def double_integrator():\n",
    "    # Discrete-time approximation of the double integrator.\n",
    "    dt = 0.01;\n",
    "    A = np.eye(2) + dt*np.mat('0 1; 0 0')\n",
    "    B = dt*np.mat('0; 1')\n",
    "\n",
    "    prog = MathematicalProgram()\n",
    "\n",
    "    N = 284  # Note: I had to do a manual \"line search\" to find this.\n",
    "\n",
    "    # Create decision variables\n",
    "    u = np.empty((1, N-1), dtype=Variable)\n",
    "    x = np.empty((2, N), dtype=Variable)\n",
    "    for n in range(N-1):\n",
    "        u[:,n] = prog.NewContinuousVariables(1, 'u' + str(n))\n",
    "        x[:,n] = prog.NewContinuousVariables(2, 'x' + str(n))\n",
    "    x[:,N-1] = prog.NewContinuousVariables(2, 'x' + str(N))\n",
    "\n",
    "    # Add constraints\n",
    "    x0 = [-2, 0]\n",
    "    prog.AddBoundingBoxConstraint(x0, x0, x[:,0])\n",
    "    for n in range(N-1):\n",
    "        # Will eventually be prog.AddConstraint(x[:,n+1] == A@x[:,n] + B@u[:,n])\n",
    "        # See drake issues 12841 and 8315\n",
    "        prog.AddConstraint(eq(x[:,n+1],A.dot(x[:,n]) + B.dot(u[:,n])))\n",
    "        prog.AddBoundingBoxConstraint(-1, 1, u[:,n])\n",
    "    xf = [0, 0]\n",
    "    prog.AddBoundingBoxConstraint(xf, xf, x[:, N-1])\n",
    "\n",
    "    result = Solve(prog)\n",
    "\n",
    "    x_sol = result.GetSolution(x)\n",
    "    assert(result.is_success()), \"Optimization failed\"\n",
    "\n",
    "    plt.figure()\n",
    "    plt.plot(x_sol[0,:], x_sol[1,:])\n",
    "    plt.xlabel('q')\n",
    "    plt.ylabel('qdot');\n",
    "\n",
    "double_integrator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trajectory optimization using the DirectTranscription class\n",
    "\n",
    "Because this pattern of making decision variables that are indexed over time, adding the dynamic constraints, defining the running cost and constraints, is so common, we have wrappers in drake on top of `MathematicalProgram` which handle these details for you.\n",
    "\n",
    "The optimization below is identical to the example above, but using this helper class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dirtran_example():\n",
    "    # Discrete-time approximation of the double integrator.\n",
    "    dt = 0.01;\n",
    "    A = np.eye(2) + dt*np.mat('0 1; 0 0')\n",
    "    B = dt*np.mat('0; 1')\n",
    "    C = np.eye(2)\n",
    "    D = np.zeros((2,1))\n",
    "    sys = LinearSystem(A, B, C, D, dt)\n",
    "\n",
    "    N = 284\n",
    "    x0 = [-2, 0]\n",
    "    xf = [0, 0]\n",
    "\n",
    "    dirtran = DirectTranscription(sys, sys.CreateDefaultContext(), N)\n",
    "    prog = dirtran.prog()\n",
    "    prog.AddBoundingBoxConstraint(x0, x0, dirtran.initial_state())\n",
    "    prog.AddBoundingBoxConstraint(xf, xf, dirtran.final_state())\n",
    "    dirtran.AddConstraintToAllKnotPoints(dirtran.input()[0] <= 1)\n",
    "    dirtran.AddConstraintToAllKnotPoints(dirtran.input()[0] >= -1)\n",
    "\n",
    "    result = Solve(prog)\n",
    "    x_sol = dirtran.ReconstructStateTrajectory(result)\n",
    "    assert(result.is_success()), \"Optimization failed\"\n",
    "\n",
    "    plt.figure()\n",
    "    x_values = x_sol.vector_values(x_sol.get_segment_times())\n",
    "\n",
    "    plt.plot(x_values[0,:], x_values[1,:])\n",
    "    plt.xlabel('q')\n",
    "    plt.ylabel('qdot')\n",
    "\n",
    "dirtran_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One thing that I'm very proud of (it was a lot of work!) is the fact that drake is often smart enough to introspect your system, costs, and constraints and understand whether you have formulated a convex problem or a non-convex one.  The optimization above calls a convex optimization solver.  But if you had passed in a nonlinear system instead, it would have switched to calling a solver that supports nonlinear programming."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direct Collocation for the Pendulum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pend_dircol():\n",
    "    plant = PendulumPlant()\n",
    "    context = plant.CreateDefaultContext()\n",
    "\n",
    "    N = 21\n",
    "    max_dt = 0.5\n",
    "    max_tf = N * max_dt\n",
    "    dircol = DirectCollocation(plant,\n",
    "                            context,\n",
    "                            num_time_samples=N,\n",
    "                            minimum_timestep=0.05,\n",
    "                            maximum_timestep=max_dt)\n",
    "    prog = dircol.prog()\n",
    "\n",
    "    dircol.AddEqualTimeIntervalsConstraints()\n",
    "\n",
    "    torque_limit = 3.0  # N*m.\n",
    "    u = dircol.input()\n",
    "    dircol.AddConstraintToAllKnotPoints(-torque_limit <= u[0])\n",
    "    dircol.AddConstraintToAllKnotPoints(u[0] <= torque_limit)\n",
    "\n",
    "    initial_state = PendulumState()\n",
    "    initial_state.set_theta(0.0)\n",
    "    initial_state.set_thetadot(0.0)\n",
    "    prog.AddBoundingBoxConstraint(initial_state.get_value(),\n",
    "                                  initial_state.get_value(),\n",
    "                                  dircol.initial_state())\n",
    "    # More elegant version is blocked on drake #8315:\n",
    "    # dircol.AddLinearConstraint(\n",
    "    #     dircol.initial_state() == initial_state.get_value())\n",
    "\n",
    "    final_state = PendulumState()\n",
    "    final_state.set_theta(np.pi)\n",
    "    final_state.set_thetadot(0.0)\n",
    "    prog.AddBoundingBoxConstraint(final_state.get_value(),\n",
    "                                  final_state.get_value(), dircol.final_state())\n",
    "    # dircol.AddLinearConstraint(dircol.final_state() == final_state.get_value())\n",
    "\n",
    "    R = 10  # Cost on input \"effort\".\n",
    "    dircol.AddRunningCost(R * u[0]**2)\n",
    "\n",
    "    initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(\n",
    "        [0., 4.], [initial_state.get_value(),\n",
    "                final_state.get_value()])\n",
    "    dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)\n",
    "\n",
    "    result = Solve(prog)\n",
    "    assert result.is_success()\n",
    "\n",
    "    x_trajectory = dircol.ReconstructStateTrajectory(result)\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    x_knots = np.hstack([\n",
    "        x_trajectory.value(t) for t in np.linspace(x_trajectory.start_time(),\n",
    "                                                x_trajectory.end_time(), 100)\n",
    "    ])\n",
    "\n",
    "    ax.set_xlabel('$q$')\n",
    "    ax.set_ylabel('$\\dot{q}$')\n",
    "    ax.plot(x_knots[0, :], x_knots[1, :])\n",
    "    display(plt.show())\n",
    "\n",
    "    # Animate the result.\n",
    "    vis = PendulumVisualizer(show=False)\n",
    "    ani = vis.animate(x_trajectory)\n",
    "    display(HTML(ani.to_jshtml()))\n",
    "\n",
    "pend_dircol()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direct Collocation for the Acrobot\n",
    "\n",
    "Almost identical code works to swing-up the Acrobot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dircol_acrobot():\n",
    "    plant = AcrobotPlant()\n",
    "    context = plant.CreateDefaultContext()\n",
    "\n",
    "    dircol = DirectCollocation(plant,\n",
    "                            context,\n",
    "                            num_time_samples=21,\n",
    "                            minimum_timestep=0.05,\n",
    "                            maximum_timestep=0.2)\n",
    "    prog = dircol.prog()\n",
    "\n",
    "    dircol.AddEqualTimeIntervalsConstraints()\n",
    "\n",
    "    # Add input limits.\n",
    "    torque_limit = 8.0  # N*m.\n",
    "    u = dircol.input()\n",
    "    dircol.AddConstraintToAllKnotPoints(-torque_limit <= u[0])\n",
    "    dircol.AddConstraintToAllKnotPoints(u[0] <= torque_limit)\n",
    "\n",
    "    initial_state = (0., 0., 0., 0.)\n",
    "    prog.AddBoundingBoxConstraint(initial_state, initial_state,\n",
    "                                  dircol.initial_state())\n",
    "    # More elegant version is blocked on drake #8315:\n",
    "    # dircol.AddLinearConstraint(dircol.initial_state() == initial_state)\n",
    "\n",
    "    final_state = (np.pi, 0., 0., 0.)\n",
    "    prog.AddBoundingBoxConstraint(final_state, final_state,\n",
    "                                  dircol.final_state())\n",
    "    # dircol.AddLinearConstraint(dircol.final_state() == final_state)\n",
    "\n",
    "    R = 10  # Cost on input \"effort\".\n",
    "    dircol.AddRunningCost(R * u[0]**2)\n",
    "\n",
    "    # Add a final cost equal to the total duration.\n",
    "    dircol.AddFinalCost(dircol.time())\n",
    "\n",
    "    initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(\n",
    "        [0., 4.], np.column_stack((initial_state, final_state)))  # yapf: disable\n",
    "    dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)\n",
    "\n",
    "    result = Solve(prog)\n",
    "    assert (result.is_success())\n",
    "\n",
    "    u_trajectory = dircol.ReconstructInputTrajectory(result)\n",
    "    times = np.linspace(u_trajectory.start_time(), u_trajectory.end_time(), 100)\n",
    "    u_lookup = np.vectorize(u_trajectory.value)\n",
    "    u_values = u_lookup(times)\n",
    "\n",
    "    plt.figure()\n",
    "    plt.plot(times, u_values)\n",
    "    plt.xlabel(\"time (seconds)\")\n",
    "    plt.ylabel(\"force (Newtons)\")\n",
    "    display(plt.show())\n",
    "\n",
    "    x_trajectory = dircol.ReconstructStateTrajectory(result)\n",
    "    u_trajectory = dircol.ReconstructInputTrajectory(result)\n",
    "\n",
    "    builder = DiagramBuilder()\n",
    "    source = builder.AddSystem(TrajectorySource(x_trajectory))\n",
    "    scene_graph = builder.AddSystem(SceneGraph())\n",
    "    AcrobotGeometry.AddToBuilder(builder, source.get_output_port(0),\n",
    "                                 scene_graph)\n",
    "    visualizer = builder.AddSystem(\n",
    "        PlanarSceneGraphVisualizer(scene_graph,\n",
    "                                   xlim=[-4., 4.],\n",
    "                                   ylim=[-4., 4.],\n",
    "                                   show=False))\n",
    "    builder.Connect(scene_graph.get_query_output_port(),\n",
    "                    visualizer.get_input_port(0))\n",
    "    simulator = Simulator(builder.Build())\n",
    "\n",
    "    # Simulate and animate\n",
    "    AdvanceToAndVisualize(\n",
    "        simulator, visualizer,\n",
    "        x_trajectory.end_time() if running_as_notebook else 0.1)\n",
    "\n",
    "    return x_trajectory, u_trajectory\n",
    "\n",
    "x_trajectory, u_trajectory = dircol_acrobot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trajectory stabilization with (finite-horizon) LQR\n",
    "\n",
    "The visualization above was an animation of the solution trajectory, it was not a simulation!  In fact, if you simulate only the planned torques connected to the Acrobot, then you will see that the trajectory does not follow the planned state trajectory.  The trajectory optimization only satisfies the dynamic constraints to some loose tolerance, and the (open-loop) trajectory is unstable.\n",
    "\n",
    "To do better, we must stabilizing the planned trajectory with feedback.  We do that here with a finite-horizon LQR design, which takes a time-varying linearization of the plant along the trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def finite_horizon_lqr(x_trajectory, u_trajectory):\n",
    "    options = FiniteHorizonLinearQuadraticRegulatorOptions()\n",
    "    options.x0 = x_trajectory\n",
    "    options.u0 = u_trajectory\n",
    "\n",
    "    builder = DiagramBuilder()\n",
    "    plant = builder.AddSystem(AcrobotPlant())\n",
    "    context = plant.CreateDefaultContext()\n",
    "    Q = np.diag([10.,10.,1.,1.])\n",
    "    options.Qf = Q\n",
    "    regulator = builder.AddSystem(\n",
    "        MakeFiniteHorizonLinearQuadraticRegulator(plant,\n",
    "                                                  context,\n",
    "                                                  t0=options.u0.start_time(),\n",
    "                                                  tf=options.u0.end_time(),\n",
    "                                                  Q=Q,\n",
    "                                                  R=np.eye(1),\n",
    "                                                  options=options))\n",
    "    builder.Connect(regulator.get_output_port(0), plant.get_input_port(0))\n",
    "    builder.Connect(plant.get_output_port(0), regulator.get_input_port(0))\n",
    "    input_logger = LogVectorOutput(regulator.get_output_port(0), builder)\n",
    "    state_logger = LogVectorOutput(plant.get_output_port(0), builder)\n",
    "\n",
    "    scene_graph = builder.AddSystem(SceneGraph())\n",
    "    AcrobotGeometry.AddToBuilder(builder, plant.get_output_port(0), scene_graph)\n",
    "    visualizer = builder.AddSystem(\n",
    "        PlanarSceneGraphVisualizer(scene_graph,\n",
    "                                   xlim=[-4., 4.],\n",
    "                                   ylim=[-4., 4.],\n",
    "                                   show=False))\n",
    "    builder.Connect(scene_graph.get_query_output_port(),\n",
    "                    visualizer.get_input_port(0))\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "\n",
    "    # Simulate and animate\n",
    "    AdvanceToAndVisualize(simulator, visualizer,\n",
    "                          options.u0.end_time() if running_as_notebook else 0.1)\n",
    "\n",
    "    fig, ax = plt.subplots(2,1)\n",
    "    ax[0].plot(options.u0.get_segment_times(),\n",
    "               options.u0.vector_values(options.u0.get_segment_times()).T)\n",
    "    input_log = input_logger.FindLog(simulator.get_context())\n",
    "    ax[0].plot(input_log.sample_times(), input_log.data().T)\n",
    "    ax[0].legend(('u nominal','u actual'))\n",
    "\n",
    "    ax[1].plot(options.x0.get_segment_times(),\n",
    "               options.x0.vector_values(options.x0.get_segment_times()).T, 'b')\n",
    "    state_log = state_logger.FindLog(simulator.get_context())\n",
    "    ax[1].plot(state_log.sample_times(), state_log.data().T, 'g')\n",
    "    ax[1].legend(('x nominal','x actual'))\n",
    "    display(plt.show())\n",
    "\n",
    "finite_horizon_lqr(x_trajectory, u_trajectory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direct Collocation for the Cart-Pole\n",
    "\n",
    "While the previous two examples used equations of motion that we derived and typed in manually, for the Cart-Pole we will use the Drake dynamics engine and visualization via [`MultibodyPlant`](https://drake.mit.edu/doxygen_cxx/classdrake_1_1multibody_1_1_multibody_plant.html) and [`SceneGraph`](https://drake.mit.edu/doxygen_cxx/classdrake_1_1geometry_1_1_scene_graph.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dircol_cartpole():\n",
    "    plant = MultibodyPlant(time_step=0.0)\n",
    "    scene_graph = SceneGraph()\n",
    "    plant.RegisterAsSourceForSceneGraph(scene_graph)\n",
    "    file_name = FindResource(\"models/cartpole.urdf\")\n",
    "    Parser(plant).AddModelFromFile(file_name)\n",
    "    plant.Finalize()\n",
    "\n",
    "    context = plant.CreateDefaultContext()\n",
    "    dircol = DirectCollocation(\n",
    "        plant,\n",
    "        context,\n",
    "        num_time_samples=21,\n",
    "        minimum_timestep=0.1,\n",
    "        maximum_timestep=0.4,\n",
    "        input_port_index=plant.get_actuation_input_port().get_index())\n",
    "    prog = dircol.prog()\n",
    "\n",
    "    dircol.AddEqualTimeIntervalsConstraints()\n",
    "\n",
    "    initial_state = (0., 0., 0., 0.)\n",
    "    prog.AddBoundingBoxConstraint(initial_state, initial_state,\n",
    "                                  dircol.initial_state())\n",
    "    # More elegant version is blocked by drake #8315:\n",
    "    # prog.AddLinearConstraint(dircol.initial_state() == initial_state)\n",
    "\n",
    "    final_state = (0., np.pi, 0., 0.)\n",
    "    prog.AddBoundingBoxConstraint(final_state, final_state,\n",
    "                                  dircol.final_state())\n",
    "    # prog.AddLinearConstraint(dircol.final_state() == final_state)\n",
    "\n",
    "    R = 10  # Cost on input \"effort\".\n",
    "    u = dircol.input()\n",
    "    dircol.AddRunningCost(R * u[0]**2)\n",
    "\n",
    "    # Add a final cost equal to the total duration.\n",
    "    dircol.AddFinalCost(dircol.time())\n",
    "\n",
    "    initial_x_trajectory = PiecewisePolynomial.FirstOrderHold(\n",
    "        [0., 4.], np.column_stack((initial_state, final_state)))  # yapf: disable\n",
    "    dircol.SetInitialTrajectory(PiecewisePolynomial(), initial_x_trajectory)\n",
    "\n",
    "    result = Solve(prog)\n",
    "    assert result.is_success()\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    u_trajectory = dircol.ReconstructInputTrajectory(result)\n",
    "    times = np.linspace(u_trajectory.start_time(), u_trajectory.end_time(), 100)\n",
    "    u_lookup = np.vectorize(u_trajectory.value)\n",
    "    u_values = u_lookup(times)\n",
    "\n",
    "    ax.plot(times, u_values)\n",
    "    ax.set_xlabel(\"time (seconds)\")\n",
    "    ax.set_ylabel(\"force (Newtons)\")\n",
    "    display(plt.show())\n",
    "\n",
    "    # Animate the results.\n",
    "    x_trajectory = dircol.ReconstructStateTrajectory(result)\n",
    "\n",
    "    # TODO(russt): Add some helper methods to make this workflow cleaner.\n",
    "    builder = DiagramBuilder()\n",
    "    source = builder.AddSystem(TrajectorySource(x_trajectory))\n",
    "    builder.AddSystem(scene_graph)\n",
    "    pos_to_pose = builder.AddSystem(\n",
    "        MultibodyPositionToGeometryPose(plant, input_multibody_state=True))\n",
    "    builder.Connect(source.get_output_port(0), pos_to_pose.get_input_port())\n",
    "    builder.Connect(pos_to_pose.get_output_port(),\n",
    "                    scene_graph.get_source_pose_port(plant.get_source_id()))\n",
    "\n",
    "    visualizer = builder.AddSystem(\n",
    "        PlanarSceneGraphVisualizer(scene_graph,\n",
    "                                   xlim=[-2, 2],\n",
    "                                   ylim=[-1.25, 2],\n",
    "                                   show=False))\n",
    "    builder.Connect(scene_graph.get_query_output_port(),\n",
    "                    visualizer.get_input_port(0))\n",
    "    simulator = Simulator(builder.Build())\n",
    "\n",
    "    AdvanceToAndVisualize(\n",
    "        simulator, visualizer,\n",
    "        x_trajectory.end_time() if running_as_notebook else 0.1)\n",
    "\n",
    "\n",
    "dircol_cartpole()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trajectory optimization using mixed-integer quadratic programming (MI-Convex optimization)\n",
    "\n",
    "Note: These examples don't work yet with the open-source solvers.  One needs a Mosek/Gurobi license (TODO: Link to the Drake tutorial for setting up the license on Deepnote)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mip_solver_available = (MosekSolver().available() and\n",
    "                        MosekSolver().enabled()) or (GurobiSolver().available()\n",
    "                                                     and\n",
    "                                                     GurobiSolver().enabled())\n",
    "\n",
    "A_bl = np.array([[-1, 0], [0, -1], [1, 0], [0, 1], [1, 1]])\n",
    "A_br = np.array([[1, 0], [0, -1], [-1, 0], [0, 1], [-1, 1]])\n",
    "A_tl = np.array([[-1, 0], [0, 1], [1, 0], [0, -1], [1, -1]])\n",
    "A_tr = np.array([[1, 0], [0, 1], [-1, 0], [0, -1], [-1, -1]])\n",
    "b = np.array([4, 4, 0, 0, -1])\n",
    "\n",
    "H_bl = HPolyhedron(A_bl, b)\n",
    "H_br = HPolyhedron(A_br, b)\n",
    "H_tl = HPolyhedron(A_tl, b)\n",
    "H_tr = HPolyhedron(A_tr, b)\n",
    "\n",
    "def plot_and_interact(meshcat, prog, x, start, goal):\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.DeleteAddedControls()\n",
    "    meshcat.Set2dRenderMode(xmin=-4, xmax=4, ymin=-4, ymax=4)\n",
    "    meshcat.SetObject(\"obstacle\", Box(np.sqrt(2), 1, np.sqrt(2)), Rgba(1, 0, 0, 1))\n",
    "    meshcat.SetTransform(\"obstacle\", RigidTransform(RotationMatrix.MakeYRotation(np.pi/4), [0,2,0]))\n",
    "\n",
    "    meshcat.SetObject(\"start\", Sphere(.1), Rgba(0, 0, 1, 1))\n",
    "    meshcat.SetObject(\"goal\", Sphere(.1), Rgba(0, 1, 0, 1))\n",
    "\n",
    "    x_start = start.evaluator().lower_bound()\n",
    "    x_goal = goal.evaluator().lower_bound()\n",
    "    x_sol = np.zeros((x.shape[0],3))\n",
    "    meshcat.AddSlider(\"x_start\", -4, 4, 0.25, x_start[0])\n",
    "    meshcat.AddSlider(\"y_start\", -4, -1, 0.25, x_start[1])\n",
    "    meshcat.AddSlider(\"x_goal\", -4, 4, 0.25, x_goal[0])\n",
    "    meshcat.AddSlider(\"y_goal\", 1, 4, 0.25, x_goal[1])\n",
    "    pts = PointCloud(x.shape[0])\n",
    "    pts.mutable_xyzs()[1,:] = 0\n",
    "\n",
    "    if not running_as_notebook:\n",
    "        return\n",
    "\n",
    "    meshcat.AddButton(\"Stop\")\n",
    "    while meshcat.GetButtonClicks(\"Stop\") < 1:\n",
    "        x_start = [meshcat.GetSliderValue(\"x_start\"), meshcat.GetSliderValue(\"y_start\")]\n",
    "        x_goal = [meshcat.GetSliderValue(\"x_goal\"), meshcat.GetSliderValue(\"y_goal\")]\n",
    "        start.evaluator().UpdateLowerBound(x_start)\n",
    "        start.evaluator().UpdateUpperBound(x_start)\n",
    "        goal.evaluator().UpdateLowerBound(x_goal)\n",
    "        goal.evaluator().UpdateUpperBound(x_goal)\n",
    "\n",
    "        result = Solve(prog)\n",
    "        clear_output(wait=True)\n",
    "        print(f'Feasible: {result.is_success()}')\n",
    "        print(f'Cost: {result.get_optimal_cost()}')\n",
    "\n",
    "        x_sol[:,[0,2]] = result.GetSolution(x)\n",
    "\n",
    "        meshcat.SetLine(\"x\", x_sol.T, 2.0, Rgba(0,0,1,1))\n",
    "        meshcat.SetTransform(\"start\", RigidTransform([x_start[0], 0, x_start[1]]))\n",
    "        meshcat.SetTransform(\"goal\", RigidTransform([x_goal[0], 0, x_goal[1]]))\n",
    "        pts.mutable_xyzs()[0,:] = x_sol[:,0]\n",
    "        pts.mutable_xyzs()[2,:] = x_sol[:,2]\n",
    "        meshcat.SetObject(\"x_pts\", pts, 5.0, rgba=Rgba(0, 0, 1, 1))\n",
    "\n",
    "        time.sleep(.1)\n",
    "\n",
    "    meshcat.DeleteAddedControls()\n",
    "\n",
    "def plot_and_interact_spp(meshcat, G, source, target, relaxation):\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.DeleteAddedControls()\n",
    "    meshcat.Set2dRenderMode(xmin=-4, xmax=4, ymin=-4, ymax=4)\n",
    "    meshcat.SetObject(\"obstacle\", Box(np.sqrt(2), 1, np.sqrt(2)), Rgba(1, 0, 0, 1))\n",
    "    meshcat.SetTransform(\"obstacle\", RigidTransform(RotationMatrix.MakeYRotation(np.pi/4), [0,2,0]))\n",
    "\n",
    "    meshcat.SetObject(\"start\", Sphere(.1), Rgba(0, 0, 1, 1))\n",
    "    meshcat.SetObject(\"goal\", Sphere(.1), Rgba(0, 1, 0, 1))\n",
    "\n",
    "    x_start = source.set().x()\n",
    "    x_goal = target.set().x()\n",
    "    x_sol = np.zeros((x_start.shape[0],3))\n",
    "    meshcat.AddSlider(\"x_start\", -4, 4, 0.25, x_start[0])\n",
    "    meshcat.AddSlider(\"y_start\", -4, -1, 0.25, x_start[1])\n",
    "    meshcat.AddSlider(\"x_goal\", -4, 4, 0.25, x_goal[0])\n",
    "    meshcat.AddSlider(\"y_goal\", 1, 4, 0.25, x_goal[1])\n",
    "    pts = PointCloud(2)\n",
    "    pts.mutable_xyzs()[1,:] = 0\n",
    "\n",
    "    if not relaxation and not mip_solver_available:\n",
    "        return\n",
    "\n",
    "    if not running_as_notebook:\n",
    "        return\n",
    "\n",
    "    meshcat.AddButton(\"Stop\")\n",
    "    while meshcat.GetButtonClicks(\"Stop\") < 1:\n",
    "        x_start = [meshcat.GetSliderValue(\"x_start\"), meshcat.GetSliderValue(\"y_start\")]\n",
    "        x_goal = [meshcat.GetSliderValue(\"x_goal\"), meshcat.GetSliderValue(\"y_goal\")]\n",
    "        source.set().set_x(x_start)\n",
    "        target.set().set_x(x_goal)\n",
    "\n",
    "        result = G.SolveShortestPath(source, target, relaxation)\n",
    "        clear_output(wait=True)\n",
    "        print(f'Feasible: {result.is_success()}')\n",
    "        print(f'Cost: {result.get_optimal_cost()}')\n",
    "\n",
    "        for e in G.Edges():\n",
    "            u = result.GetSolution(e.xu())\n",
    "            v = result.GetSolution(e.xv())\n",
    "            phi = result.GetSolution(e.phi())\n",
    "            meshcat.SetLineSegments(f\"edges/{e.id().get_value()}\",\n",
    "                                    [u[0], 0, u[1]], [v[0], 0, v[1]], 2.0,\n",
    "                                    Rgba(0, 0, 1, 1))\n",
    "            meshcat.SetProperty(f\"edges/{e.id().get_value()}\", \"visible\",\n",
    "                                phi > 0.5)\n",
    "            pts.mutable_xyzs()[[0,2],0] = u\n",
    "            pts.mutable_xyzs()[[0,2],1] = v\n",
    "            meshcat.SetObject(f\"edges/{e.id().get_value()}/pts\", pts, 5.0, rgba=Rgba(0, 0, 1, 1))\n",
    "\n",
    "        meshcat.SetTransform(\"start\",\n",
    "                             RigidTransform([x_start[0], 0, x_start[1]]))\n",
    "        meshcat.SetTransform(\"goal\", RigidTransform([x_goal[0], 0, x_goal[1]]))\n",
    "\n",
    "        time.sleep(.1)\n",
    "\n",
    "    meshcat.DeleteAddedControls()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Non-convex collision constraints\n",
    "\n",
    "To form the baseline for our comparison, let's first take the existing tools and apply them here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nonconvex_example():\n",
    "    meshcat.Delete()\n",
    "    N = 10     # number of timesteps\n",
    "    x_start = np.array([.2, -2.])\n",
    "    x_goal = np.array([0, 2.])\n",
    "\n",
    "    prog = MathematicalProgram()\n",
    "\n",
    "    x = prog.NewContinuousVariables(N, 2, \"x\")   # positions\n",
    "\n",
    "    # Initial and final value constraints:\n",
    "    start = prog.AddBoundingBoxConstraint(x_start, x_start, x[0])\n",
    "    goal = prog.AddBoundingBoxConstraint(x_goal, x_goal, x[N-1])\n",
    "\n",
    "    def collision(x):\n",
    "        # implemented |x0| + |x1|, except abs is not differentiable\n",
    "        d0 = np.abs(x[0])\n",
    "        if d0 < 0.1:  \n",
    "            d0 = 10*x[0]**2\n",
    "        d1 = np.abs(x[1])\n",
    "        if d1 < 0.1:\n",
    "            d1 = 10*x[1]**2\n",
    "        return [d0 + d1]\n",
    "\n",
    "    for n in range(N-1):\n",
    "        # Stage cost (sum of squared distances):\n",
    "        #prog.AddCost(dist, np.concatenate([x[n],x[n+1]]))\n",
    "        step = x[n+1] - x[n]\n",
    "        prog.AddQuadraticCost(step.dot(step))\n",
    "\n",
    "        # Collision avoidance constraints |x0| + |x1| >= 1\n",
    "        prog.AddConstraint(collision, [1], [np.inf], x[n])\n",
    "\n",
    "    prog.SetInitialGuess(\n",
    "        x,\n",
    "        np.linspace(x_start, x_goal, N) + 0.1 * np.random.rand(N, 2))\n",
    "\n",
    "    plot_and_interact(meshcat, prog, x, start, goal)\n",
    "\n",
    "nonconvex_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Big-M formulation (binaries per face)\n",
    "\n",
    "Note: If we wanted to minimize distance (instead of sum of squared distances), we could use, e.g.\n",
    "```\n",
    "  prog.AddLorentzConeConstraint(np.array([Expression(l2norm[0,n]), step[0], step[1]]))\n",
    "  prog.AddLinearCost(l2norm[0,n])\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bigM_example(convex_relaxation=False):\n",
    "    if not convex_relaxation and not mip_solver_available:\n",
    "        return\n",
    "\n",
    "    N = 10     # number of timesteps\n",
    "    bigM = 5\n",
    "    x_start = np.array([0., -2.])\n",
    "    x_goal = np.array([0., 2.])\n",
    "\n",
    "    prog = MathematicalProgram()\n",
    "\n",
    "    x = prog.NewContinuousVariables(N, 2, \"x\")   # positions\n",
    "    if convex_relaxation:\n",
    "        b = prog.NewContinuousVariables(N-1, 4, \"b\") # binaries for collisions\n",
    "        prog.AddBoundingBoxConstraint(0, 1, b)\n",
    "    else:\n",
    "        b = prog.NewBinaryVariables(N-1, 4, \"b\")     # binaries for collisions\n",
    "\n",
    "    # Initial and final value constraints:\n",
    "    start = prog.AddBoundingBoxConstraint(x_start, x_start, x[0])\n",
    "    goal = prog.AddBoundingBoxConstraint(x_goal, x_goal, x[N-1])\n",
    "\n",
    "    for n in range(N-1):\n",
    "        # Stage cost (sum of squared distances encourages equal-sized steps):\n",
    "        step = x[n+1] - x[n]\n",
    "        prog.AddQuadraticCost(step.dot(step))\n",
    "\n",
    "        # Collision avoidance constraints (each *line segment* must be in a region)\n",
    "        for k in [n,n+1]:\n",
    "            prog.AddLinearConstraint(x[k].dot([1,1]) <= -1 + (1-b[n,0])*bigM) # bottom left\n",
    "            prog.AddLinearConstraint(x[k].dot([-1,1]) <= -1 + (1-b[n,1])*bigM) # bottom right\n",
    "            prog.AddLinearConstraint(x[k].dot([1,-1]) <= -1 + (1-b[n,2])*bigM) # top left\n",
    "            prog.AddLinearConstraint(x[k].dot([-1,-1]) <= -1 + (1-b[n,3])*bigM) # top right\n",
    "        prog.AddLinearConstraint(np.sum(b[n]) == 1)\n",
    "\n",
    "    plot_and_interact(meshcat, prog, x, start, goal)\n",
    "\n",
    "bigM_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convex hull (binaries per region) using perspective functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convex_hull_example(convex_relaxation=False):\n",
    "    if not convex_relaxation and not mip_solver_available:\n",
    "        return\n",
    "\n",
    "    N = 10     # number of timesteps\n",
    "    x_start = np.array([0., -2.])\n",
    "    x_goal = np.array([0., 2.])\n",
    "\n",
    "    Region = [ H_bl, H_br, H_tl, H_tr ];\n",
    "\n",
    "    prog = MathematicalProgram()\n",
    "\n",
    "    x = prog.NewContinuousVariables(N, 2, \"x\")   # positions\n",
    "    if convex_relaxation:\n",
    "        b = prog.NewContinuousVariables(N-1, 4, \"b\") # binaries for regions\n",
    "        prog.AddBoundingBoxConstraint(0, 1, b)\n",
    "    else:\n",
    "        b = prog.NewBinaryVariables(N-1, 4, \"b\")     # binaries for regions\n",
    "\n",
    "    # Slack variables for the regions\n",
    "    x_region = [\n",
    "        prog.NewContinuousVariables(N, 2, \"x_bl\"),\n",
    "        prog.NewContinuousVariables(N, 2, \"x_br\"),\n",
    "        prog.NewContinuousVariables(N, 2, \"x_tl\"),\n",
    "        prog.NewContinuousVariables(N, 2, \"x_tr\"),\n",
    "    ]\n",
    "    for n in range(N):\n",
    "        prog.AddLinearConstraint(eq(x[n], sum([xr[n] for xr in x_region])))\n",
    "\n",
    "    # Initial and final value constraints:\n",
    "    start = prog.AddLinearConstraint(eq(x[0], x_start))\n",
    "    goal = prog.AddLinearConstraint(eq(x[N-1], x_goal))\n",
    "\n",
    "    for n in range(N-1):\n",
    "        # Stage cost (sum of squared distances encourages equal-sized steps)\n",
    "        step = x[n+1] - x[n]\n",
    "        prog.AddQuadraticCost(step.dot(step))\n",
    "\n",
    "        # Region containment constraints\n",
    "        for i in range(4):\n",
    "            Region[i].AddPointInNonnegativeScalingConstraints(prog, x_region[i][n], b[n,i])\n",
    "        prog.AddLinearConstraint(np.sum(b[n]) == 1)\n",
    "\n",
    "    plot_and_interact(meshcat, prog, x, start, goal)\n",
    "\n",
    "convex_hull_example()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To avoid cutting the corners, we can define regions for each line segment (using the Cartesian power of the set)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convex_hull_example2(convex_relaxation=False):\n",
    "    if not convex_relaxation and not mip_solver_available:\n",
    "        return\n",
    "\n",
    "    N = 10     # number of timesteps\n",
    "    x_start = np.array([0., -2.])\n",
    "    x_goal = np.array([0., 2.])\n",
    "\n",
    "    Region = [ H_bl, H_br, H_tl, H_tr ];\n",
    "    # Each *line segment* must be in a region.\n",
    "    Region2 = [ R.CartesianPower(2) for R in Region ]\n",
    "\n",
    "    prog = MathematicalProgram()\n",
    "\n",
    "    x = prog.NewContinuousVariables(N, 2, \"x\")   # positions\n",
    "    if convex_relaxation:\n",
    "        b = prog.NewContinuousVariables(N-1, 4, \"b\") # binaries for regions\n",
    "        prog.AddBoundingBoxConstraint(0, 1, b)\n",
    "    else:\n",
    "        b = prog.NewBinaryVariables(N-1, 4, \"b\")     # binaries for regions\n",
    "\n",
    "    # Slack variables for the regions\n",
    "    x_region = [\n",
    "        prog.NewContinuousVariables(N-1, 4, \"xA\"),\n",
    "        prog.NewContinuousVariables(N-1, 4, \"xB\"),\n",
    "        prog.NewContinuousVariables(N-1, 4, \"xC\"),\n",
    "        prog.NewContinuousVariables(N-1, 4, \"xD\"),\n",
    "    ]\n",
    "    for n in range(N-1):\n",
    "        prog.AddLinearConstraint(eq(x[n], sum([xr[n,:2] for xr in x_region])))\n",
    "        prog.AddLinearConstraint(eq(x[n+1], sum([xr[n,2:] for xr in x_region])))\n",
    "\n",
    "    # Initial and final value constraints:\n",
    "    start = prog.AddLinearConstraint(eq(x[0], x_start))\n",
    "    goal = prog.AddLinearConstraint(eq(x[N-1], x_goal))\n",
    "\n",
    "    for n in range(N-1):\n",
    "        # Stage cost (sum of squared distances encourages equal-sized steps)\n",
    "        step = x[n+1] - x[n]\n",
    "        prog.AddQuadraticCost(step.dot(step))\n",
    "    \n",
    "        # Region containment constraints\n",
    "        for i in range(4):\n",
    "            Region2[i].AddPointInNonnegativeScalingConstraints(prog, x_region[i][n], b[n,i])\n",
    "        prog.AddLinearConstraint(np.sum(b[n]) == 1)\n",
    "\n",
    "    plot_and_interact(meshcat, prog, x, start, goal)\n",
    "\n",
    "convex_hull_example2()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TODO: Make the point that the squared-norm is an odd cost function (for instance, changing the number of samples N changes the cost), but using the L2 norm leads to strange behavior in the formulations above.  It indicates that there is something a little wrong with the formulation.  GCS resolves it...\n",
    "\n",
    "## Graph of Convex Sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def graph_of_convex_sets_example(convex_relaxation = False, show_graph=False):\n",
    "    if not convex_relaxation and not mip_solver_available:\n",
    "        return\n",
    "\n",
    "    b = np.array([-1, -1, 4])\n",
    "    A_left = np.array([[1,1], [1, -1], [-1, 0]])\n",
    "    A_right = np.array([[-1,-1], [-1, 1], [1, 0]])\n",
    "    A_top = np.array([[-1,-1], [1, -1], [0, 1]])\n",
    "    A_bottom = np.array([[1,1], [-1, 1], [0, -1]])\n",
    "\n",
    "    H_left = HPolyhedron(A_left, b)\n",
    "    H_right = HPolyhedron(A_right, b)\n",
    "    H_top = HPolyhedron(A_top, b)\n",
    "    H_bottom = HPolyhedron(A_bottom, b)\n",
    "\n",
    "    G = GraphOfConvexSets()\n",
    "    source = G.AddVertex(Point([0.2, -2]), \"source\")\n",
    "    left =  G.AddVertex(H_left, \"left\")\n",
    "    right = G.AddVertex(H_right, \"right\")\n",
    "    top =     G.AddVertex(H_top, \"top\")\n",
    "    bottom =    G.AddVertex(H_bottom, \"bottom\")\n",
    "    target = G.AddVertex(Point([0, 2]), \"target\")\n",
    "\n",
    "    E = [\n",
    "        [source, left],\n",
    "        [source, bottom],\n",
    "        [source, right],\n",
    "        [bottom, left],\n",
    "        [bottom, right],\n",
    "        [left, top],\n",
    "        [right, top],\n",
    "        [left, target],\n",
    "        [top, target],\n",
    "        [right, target],\n",
    "    ]\n",
    "\n",
    "    # TODO: Support creating L2NormCosts from symbolic\n",
    "    # (e.g. np.linalg.norm(e.xu() - e.xv()))\n",
    "    cost = L2NormCost(A=np.block([np.eye(2), -np.eye(2)]), b=np.zeros((2,1)))\n",
    "    for u,v in E:\n",
    "        e = G.AddEdge(u, v)\n",
    "        e.AddCost(Binding[L2NormCost](cost, np.concatenate([e.xu(),\n",
    "                                                            e.xv()])))\n",
    "\n",
    "    if show_graph:\n",
    "        result = G.SolveShortestPath(source,\n",
    "                                target,\n",
    "                                convex_relaxation)\n",
    "        display(\n",
    "            SVG(\n",
    "                pydot.graph_from_dot_data(\n",
    "                    G.GetGraphvizString(result,\n",
    "                                        show_slacks=True))[0].create_svg()))\n",
    "    else:\n",
    "        plot_and_interact_spp(meshcat, G, source, target, convex_relaxation)\n",
    "\n",
    "graph_of_convex_sets_example()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_of_convex_sets_example(show_graph=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 64-bit",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}