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 (Box, DiagramBuilder, DirectCollocation,\n",
    "                         DirectTranscription,\n",
    "                         FiniteHorizonLinearQuadraticRegulatorOptions,\n",
    "                         GraphOfConvexSets, HPolyhedron, LinearSystem,\n",
    "                         LogVectorOutput,\n",
    "                         MakeFiniteHorizonLinearQuadraticRegulator,\n",
    "                         MathematicalProgram, 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 (which we can't provide on Deepnote/Colab yet)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mosek_solver = MosekSolver()\n",
    "have_mosek = mosek_solver.available() and mosek_solver.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",
    "A_left = np.array([[-1,0],[1,1],[1,-1]])\n",
    "A_right = np.array([[-1,0],[-1,1],[-1,-1]])\n",
    "b = np.array([4, -1, -1])\n",
    "H_left = HPolyhedron(A_left, b)\n",
    "H_right = HPolyhedron(A_right, b)\n",
    "\n",
    "def plot_and_interact(meshcat, prog, x, start, goal):\n",
    "\n",
    "    meshcat.Delete()\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 have_mosek:\n",
    "        print(\"This demo currently requires Mosek\")\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",
    "        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",
    "\n",
    "def plot_and_interact_spp(meshcat, G, source, target, relaxation):\n",
    "\n",
    "    meshcat.Delete()\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.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 have_mosek:\n",
    "        print(\"This demo currently requires Mosek\")\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",
    "            meshcat.SetLineSegments(f\"edges/{e.id().get_value()}\", [u[0],0,u[1]], [v[0],0,v[1]], 2.0, Rgba(0,0,1,1))\n",
    "            meshcat.SetProperty(f\"edges/{e.id().get_value()}\", \"visible\", result.GetSolution(e.phi()) > 0.2)\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\", 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()"
   ]
  },
  {
   "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": [
    "relaxation = False\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 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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convex hull (binaries per region) using perspective functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convex_relaxation = False\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)"
   ]
  },
  {
   "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": [
    "convex_relaxation = False\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph of Convex Sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "relaxation = True\n",
    "\n",
    "G = GraphOfConvexSets()\n",
    "source =       G.AddVertex(Point([0.2, -2]), \"source\")\n",
    "bottom_left =  G.AddVertex(H_bl, \"bottom_left\")\n",
    "bottom_right = G.AddVertex(H_br, \"bottom_right\")\n",
    "left =         G.AddVertex(H_left, \"left\")\n",
    "right =        G.AddVertex(H_right, \"right\")\n",
    "top_left =     G.AddVertex(H_tl, \"top_left\")\n",
    "top_right =    G.AddVertex(H_tr, \"top_right\")\n",
    "target =       G.AddVertex(Point([0,  2]), \"target\")\n",
    "\n",
    "E = [\n",
    "    [source, bottom_left], [source, bottom_right], \n",
    "    [bottom_left, left], \n",
    "    [bottom_right, right],\n",
    "    [left, top_left],\n",
    "    [right, top_right],\n",
    "    [top_left, target],\n",
    "    [top_right, target],\n",
    "]\n",
    "\n",
    "for u,v in E:\n",
    "    e = G.AddEdge(u, v)\n",
    "    step = v.x() - u.x()\n",
    "    e.AddCost(step.dot(step))\n",
    "\n",
    "plot_and_interact_spp(meshcat, G, source, target, relaxation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = G.SolveShortestPath(source, target, relaxation)\n",
    "display(\n",
    "    SVG(\n",
    "        pydot.graph_from_dot_data(G.GetGraphvizString(\n",
    "            result, show_slacks=True))[0].create_svg()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}