ciris_pgd-Copy1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "47244e91",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ae7f6bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydrake.multibody.plant import AddMultibodyPlantSceneGraph\n",
    "from pydrake.systems.framework import DiagramBuilder\n",
    "from pydrake.multibody.parsing import Parser\n",
    "import os\n",
    "from pydrake.all import (\n",
    "    LoadModelDirectives, ProcessModelDirectives, RevoluteJoint, \n",
    "    RationalForwardKinematics, CspaceFreePolytope, SeparatingPlaneOrder,\n",
    "    RigidTransform, RotationMatrix, Rgba,\n",
    "    AffineSubspace\n",
    ")\n",
    "import numpy as np\n",
    "# from pydrake.geometry.optimization_dev import (CspaceFreePolytope, SeparatingPlaneOrder)\n",
    "from iris_plant_visualizer import IrisPlantVisualizer\n",
    "from pydrake.geometry import Role\n",
    "from pydrake.geometry.optimization import IrisOptions, HPolyhedron, Hyperellipsoid, IrisInRationalConfigurationSpace\n",
    "from pydrake.solvers import MosekSolver, CommonSolverOption, SolverOptions, ScsSolver\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fddde847",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:425: warning: Ignoring unsupported SDFormat element in model: plugin\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:37: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:73: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:128: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:183: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:238: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:293: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:348: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:403: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:425: warning: Ignoring unsupported SDFormat element in model: plugin\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:37: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:73: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:128: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:183: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:238: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:293: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:348: warning: Ignoring unsupported SDFormat element in link: self_collide\n",
      "WARNING:drake:/home/sgrg/rlg/SUPERUROP/ciris/models/assets/oneDOF_iiwa7_with_box_collision.sdf:403: warning: Ignoring unsupported SDFormat element in link: self_collide\n"
     ]
    }
   ],
   "source": [
    "#construct our robot\n",
    "builder = DiagramBuilder()\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "parser = Parser(plant)\n",
    "\n",
    "parser.package_map().Add(\"ciris_pgd\", os.path.abspath(''))\n",
    "\n",
    "directives_file = \"/home/sgrg/rlg/SUPERUROP/ciris/models/bimanual_two_DOF.yaml\"\n",
    "directives = LoadModelDirectives(directives_file)\n",
    "models = ProcessModelDirectives(directives, plant, parser)\n",
    "\n",
    "\n",
    "# oneDOF_iiwa_file = \"/home/sgrg/rlg/SUPERUROP/ciris/assets/oneDOF_iiwa7_with_box_collision.sdf\"\n",
    "# with open(oneDOF_iiwa_file, 'r') as f:\n",
    "#     oneDOF_iiwa_string = f.read()\n",
    "# box_asset_file = \"/home/sgrg/rlg/SUPERUROP/ciris/assets/box_small.urdf\"\n",
    "# with open(box_asset_file, 'r') as f:\n",
    "#     box_asset_string = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d910e01f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plant.Finalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bf863013",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "joint index:  JointIndex(6) joint name:  iiwa_joint_2 IDX 0\n",
      "joint index:  JointIndex(14) joint name:  iiwa_joint_2 IDX 1\n"
     ]
    }
   ],
   "source": [
    "idx = 0\n",
    "q0 = [0.0, 0.0]\n",
    "val = 1.7\n",
    "q_low  = np.array([-val, -val])\n",
    "q_high = np.array([val, val])\n",
    "# set the joint limits of the plant\n",
    "# for model in models:\n",
    "for joint_index in plant.GetJointIndices():\n",
    "    joint = plant.get_mutable_joint(joint_index)\n",
    "    if isinstance(joint, RevoluteJoint):\n",
    "        print(\"joint index: \", joint_index, \"joint name: \", joint.name(), \"IDX\", idx)\n",
    "        joint.set_default_angle(q0[idx])\n",
    "        joint.set_position_limits(lower_limits= np.array([q_low[idx]]), upper_limits= np.array([q_high[idx]]))\n",
    "        idx += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f6ecc869",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ratfk = RationalForwardKinematics(plant)\n",
    "\n",
    "# the point about which we will take the stereographic projections\n",
    "# q_star = np.zeros(plant.num_positions())\n",
    "q_star = np.array([1.0,1.0])\n",
    "\n",
    "do_viz = True\n",
    "\n",
    "# The object we will use to perform our certification.\n",
    "cspace_free_polytope = CspaceFreePolytope(plant, scene_graph, SeparatingPlaneOrder.kAffine, q_star)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a7eeb7d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:drake:Meshcat listening for connections at http://localhost:7000\n",
      "INFO:drake:Meshcat listening for connections at http://localhost:7001\n",
      "INFO:drake:Meshcat listening for connections at http://localhost:7002\n"
     ]
    }
   ],
   "source": [
    "# This line builds the visualization. Change the viz_role to Role.kIllustration if you\n",
    "# want to see the plant with its illustrated geometry or to Role.kProximity if you want\n",
    "# to see the plant with the collision geometries.\n",
    "visualizer = IrisPlantVisualizer(plant, builder, scene_graph, cspace_free_polytope, rational_forward_kin=Ratfk ,viz_role=Role.kProximity, q_star=q_star)\n",
    "visualizer.visualize_collision_constraint(factor = 1.2, num_points = 100)\n",
    "visualizer.meshcat_cspace.Set2dRenderMode(RigidTransform(RotationMatrix.MakeZRotation(0), np.array([0,0,1])))\n",
    "visualizer.meshcat_nontc.Set2dRenderMode(RigidTransform(RotationMatrix.MakeZRotation(0), np.array([0,0,1])))\n",
    "visualizer.meshcat_task_space.Set2dRenderMode(RigidTransform(RotationMatrix.MakeZRotation(0), np.array([1,0,0])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3d3a4b21",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aeb38780d1694bf384aaa82476a3661d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatSlider(value=0.0, description='q0', max=1.7, min=-1.7)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "05cc3c02d2b6457992e2258e89dff326",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatSlider(value=0.0, description='q1', max=1.7, min=-1.7)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "from functools import partial\n",
    "\n",
    "sliders = []\n",
    "sliders.append(widgets.FloatSlider(min=q_low[0], max=q_high[0], value=0, description='q0'))\n",
    "sliders.append(widgets.FloatSlider(min=q_low[1], max=q_high[1], value=0, description='q1'))\n",
    "\n",
    "q = q0.copy()\n",
    "def handle_slider_change(change, idx):\n",
    "    q[idx] = change['new']\n",
    "    visualizer.show_res_q(q)\n",
    "    \n",
    "idx = 0\n",
    "for slider in sliders:\n",
    "    slider.observe(partial(handle_slider_change, idx = idx), names='value')\n",
    "    idx+=1\n",
    "\n",
    "for slider in sliders:\n",
    "    display(slider)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "30ac4007",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some seedpoints\n",
    "seed_points_q = np.array([   \n",
    "                              [0.0, 0],\n",
    "                              [0.4, 0.4],\n",
    "                              [0.9, -0.7],\n",
    "#                               [0.5,1.08],\n",
    "                              [1.3, -0.48]\n",
    "                              ])\n",
    "\n",
    "seed_points = np.array([Ratfk.ComputeSValue(seed_points_q[idx], q_star)\\\n",
    "                        for idx in range(seed_points_q.shape[0])])\n",
    "start = seed_points[0]\n",
    "end = seed_points[-1]\n",
    "\n",
    "visualizer.plot_cspace_points(seed_points, \"/iris_seed_points\", radius = 0.02)\n",
    "visualizer.plot_nontc_points(seed_points_q, \"/iris_seed_points_q\", radius = 0.02)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "54b50264",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_scale = 1e-2\n",
    "L1_ball = HPolyhedron.MakeL1Ball(2)\n",
    "Linf_ball = HPolyhedron.MakeBox(-np.ones(2), np.ones(2))\n",
    "\n",
    "template_C = np.vstack([L1_ball.A(), Linf_ball.A()])\n",
    "template_d = np.hstack([default_scale*L1_ball.b(), default_scale/np.sqrt(2)*Linf_ball.b()])\n",
    "\n",
    "\n",
    "def make_default_polytope_at_point(seed_point):\n",
    "    return HPolyhedron(template_C, template_d + template_C @ seed_point)\n",
    "\n",
    "\n",
    "# colors to plot the region.\n",
    "default_alpha = 0.2\n",
    "colors_dict = {\n",
    "    0: Rgba(0.565, 0.565, 0.565, default_alpha), # gray\n",
    "    1: Rgba(0.118, 0.533, 0.898, default_alpha), # bluish\n",
    "    2: Rgba(1,     0.757, 0.027, default_alpha), # gold\n",
    "    3: Rgba(0,     0.549, 0.024, default_alpha), # green   \n",
    "    4: Rgba(0.055, 0.914, 0.929, default_alpha), # teal \n",
    "}\n",
    "\n",
    "initial_regions = [(make_default_polytope_at_point(s), colors_dict[i]) for i, s in enumerate(seed_points)]\n",
    "# visualizer.add_group_of_regions_to_visualization(initial_regions, \"initial_regions\",\n",
    "#                             wireframe = False,\n",
    "#                            opacity = default_alpha)\n",
    "# visualizer.show_res_q(q_star)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bb389103",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 1.])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q_star"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1c1daedb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "# set up the certifier and the options for different search techniques\n",
    "solver_options = SolverOptions()\n",
    "# set this to 1 if you would like to see the solver output in terminal.\n",
    "solver_options.SetOption(CommonSolverOption.kPrintToConsole, 0)\n",
    "\n",
    "os.environ[\"MOSEKLM_LICENSE_FILE\"] = \"/home/sgrg/mosek.lic\"\n",
    "with open(os.environ[\"MOSEKLM_LICENSE_FILE\"], 'r') as f:\n",
    "    contents = f.read()\n",
    "    mosek_file_not_empty = contents != ''\n",
    "print(mosek_file_not_empty)\n",
    "\n",
    "solver_id = MosekSolver.id() if MosekSolver().available() and mosek_file_not_empty else ScsSolver.id()\n",
    "\n",
    "\n",
    "solver_id = MosekSolver.id() if MosekSolver().available() else ScsSolver.id()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0103116c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The options for when we search for a new polytope given positivity certificates.\n",
    "find_polytope_given_lagrangian_option = CspaceFreePolytope.FindPolytopeGivenLagrangianOptions()\n",
    "find_polytope_given_lagrangian_option.solver_options = solver_options\n",
    "find_polytope_given_lagrangian_option.ellipsoid_margin_cost = CspaceFreePolytope.EllipsoidMarginCost.kGeometricMean\n",
    "find_polytope_given_lagrangian_option.search_s_bounds_lagrangians = True\n",
    "find_polytope_given_lagrangian_option.ellipsoid_margin_epsilon = 1e-8\n",
    "find_polytope_given_lagrangian_option.solver_id = solver_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "07f856ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "bilinear_alternation_options = CspaceFreePolytope.BilinearAlternationOptions()\n",
    "bilinear_alternation_options.max_iter = 20 # Setting this to a high number will lead to more fill\n",
    "bilinear_alternation_options.convergence_tol = 1e-3\n",
    "bilinear_alternation_options.find_polytope_options = find_polytope_given_lagrangian_option\n",
    "\n",
    "# The options for when we search for new planes and positivity certificates given the polytopes\n",
    "# find_separation_certificate_given_polytope_options = CspaceFreePolytope.FindSeparationCertificateGivenPolytopeOptions()\n",
    "# find_separation_certificate_given_polytope_options.num_threads = -1\n",
    "# Parallelism \tparallelism {Parallelism::Max()}\n",
    "bilinear_alternation_options.find_lagrangian_options.verbose = True\n",
    "bilinear_alternation_options.find_lagrangian_options.solver_options = solver_options\n",
    "bilinear_alternation_options.find_lagrangian_options.ignore_redundant_C = False\n",
    "bilinear_alternation_options.find_lagrangian_options.solver_id = solver_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e83e89ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "binary_search_options = CspaceFreePolytope.BinarySearchOptions()\n",
    "binary_search_options.scale_min = 1\n",
    "binary_search_options.scale_max = 50\n",
    "binary_search_options.max_iter = 100\n",
    "binary_search_options.find_lagrangian_options.verbose = True\n",
    "binary_search_options.find_lagrangian_options.solver_options = solver_options\n",
    "binary_search_options.find_lagrangian_options.ignore_redundant_C = False\n",
    "binary_search_options.find_lagrangian_options.solver_id = solver_id\n",
    "# binary_search_options.find_lagrangian_options = find_separation_certificate_given_polytope_options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "75a248b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:drake:IrisInConfigurationSpace iteration 0\n",
      "INFO:drake:IrisInConfigurationSpace iteration 1\n",
      "INFO:drake:IrisInConfigurationSpace iteration 2\n",
      "INFO:drake:IrisInConfigurationSpace: Terminating because the hyperellipsoid volume change 0.004654293897454309 is below the threshold 0.02.\n",
      "INFO:drake:IrisInConfigurationSpace iteration 0\n",
      "INFO:drake:IrisInConfigurationSpace iteration 1\n",
      "INFO:drake:IrisInConfigurationSpace iteration 2\n",
      "INFO:drake:IrisInConfigurationSpace: Terminating because the hyperellipsoid volume change 0.0011874057434317353 is below the threshold 0.02.\n",
      "INFO:drake:IrisInConfigurationSpace iteration 0\n",
      "INFO:drake:IrisInConfigurationSpace iteration 1\n",
      "INFO:drake:IrisInConfigurationSpace iteration 2\n",
      "INFO:drake:IrisInConfigurationSpace: Terminating because the hyperellipsoid volume change 0.013714137862951947 is below the threshold 0.02.\n",
      "INFO:drake:IrisInConfigurationSpace iteration 0\n",
      "INFO:drake:IrisInConfigurationSpace iteration 1\n",
      "INFO:drake:IrisInConfigurationSpace iteration 2\n",
      "INFO:drake:IrisInConfigurationSpace: Terminating because the hyperellipsoid volume change 0.005382682435792319 is below the threshold 0.02.\n"
     ]
    }
   ],
   "source": [
    "iris_regions = []\n",
    "iris_ellipses = []\n",
    "\n",
    "iris_options = IrisOptions()\n",
    "iris_options.require_sample_point_is_contained = True\n",
    "iris_options.configuration_space_margin = 1e-3\n",
    "iris_options.relative_termination_threshold = 0.001\n",
    "\n",
    "context_for_iris = visualizer.task_space_diagram_context\n",
    "for i, s in enumerate(seed_points):\n",
    "    q = Ratfk.ComputeQValue(s, q_star)\n",
    "    plant.SetPositions(plant.GetMyMutableContextFromRoot(context_for_iris), q)\n",
    "    r = IrisInRationalConfigurationSpace(plant, \n",
    "                                         plant.GetMyContextFromRoot(context_for_iris),\n",
    "                                         q_star, iris_options)\n",
    "    iris_regions.append((r, colors_dict[i]))\n",
    "    iris_ellipses.append(r.MaximumVolumeInscribedEllipsoid())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "aa5ee19e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualizer.add_group_of_regions_to_visualization(iris_regions, \"/uncertified-iris\", \n",
    "#                                                             wireframe = False, opacity = 0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2d66d676",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "starting seedpoint 1/4\n",
      "starting seedpoint 2/4\n",
      "starting seedpoint 3/4\n",
      "starting seedpoint 4/4\n"
     ]
    }
   ],
   "source": [
    "binary_search_region_certificates_for_iris = dict.fromkeys([tuple(s) for s in seed_points])\n",
    "for i, (s, (initial_region, color)) in enumerate(zip(seed_points, iris_regions)):\n",
    "    print(f\"starting seedpoint {i+1}/{len(iris_regions)}\")\n",
    "    time.sleep(0.2)    \n",
    "    cert = cspace_free_polytope.BinarySearch(set(),\n",
    "                                                    initial_region.A(),\n",
    "                                                    initial_region.b(), \n",
    "                                                    initial_region.MaximumVolumeInscribedEllipsoid().center(), \n",
    "                                                    binary_search_options)\n",
    "    binary_search_region_certificates_for_iris[tuple(s)] = [(cert.certified_polytope(), cert, color)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95b8dfe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i, result in enumerate(binary_search_region_certificates_for_iris.values()):\n",
    "#     group_name = f\"/certified-iris-bin_seed_point_{i}\"\n",
    "#     visualizer.add_group_of_regions_and_certs_to_visualization(result, group_name, \n",
    "#                                                             wireframe = False, opacity = 0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ca1e2791",
   "metadata": {},
   "outputs": [],
   "source": [
    "regions = []\n",
    "# for group_name, region_and_cert_list in visualizer.region_certificate_groups.items():\n",
    "for i, result in enumerate(binary_search_region_certificates_for_iris.values()):\n",
    "#     regions.extend([l[0] for l in region_and_cert_list])\n",
    "    regions.extend(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "c545495e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(regions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "1a46fd1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "regions_visualizer = [(region[0].ReduceInequalities(), region[2]) for region in regions]\n",
    "regions = [(region[0].ReduceInequalities()) for region in regions]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a435cfb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# visualizer.add_group_of_regions_to_visualization(regions_visualizer, \"initial_regions\",\n",
    "#                             wireframe = True,\n",
    "#                            opacity = default_alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "562d1a52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<pydrake.geometry.optimization.HPolyhedron object at 0x7c409810ffb0>\n"
     ]
    }
   ],
   "source": [
    "import visualization_utils as viz_utils\n",
    "# viz_utils.plot_polytope(regions_visualizer[0][0], visualizer.meshcat_cspace, \"test\", wireframe=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "5aa31fca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pgd import *\n",
    "def get_step(values, indices, best_path, cost_func, gradient_func, step_size=0.01, backtracking=False, ndim=8):\n",
    "    # R is a delay operator\n",
    "    df_Rk = []\n",
    "    last_common = 0\n",
    "    f_Rk = 0\n",
    "    # values = np.asarray(values)\n",
    "    # t1 = time.perf_counter()\n",
    "    for v in range(len(best_path[1:-1])):\n",
    "        gradients = gradient_func(values[indices[v][0]:indices[v][1]])\n",
    "        f_Rk += cost_func(values[indices[v][0]:indices[v][1]])\n",
    "        if v==0:\n",
    "            df_Rk.extend(gradients)\n",
    "        else:\n",
    "            # df_Rk[last_common-ndim:] += gradients[:ndim]\n",
    "            for i in range(ndim):\n",
    "                df_Rk[last_common-ndim+i] += gradients[i]\n",
    "            df_Rk.extend(gradients[ndim:])\n",
    "        last_common = len(df_Rk)\n",
    "    # t2 = time.perf_counter()\n",
    "    # print(t2-t1)\n",
    "    df_Rk = np.asarray(df_Rk)\n",
    "\n",
    "#     print(len(df_Rk))\n",
    "    f_k = 0\n",
    "    t1 = 0\n",
    "    t2 = 0\n",
    "    L_norm = 0\n",
    "    t = 0\n",
    "    \n",
    "    beta = 0.5\n",
    "    # https://www.cs.cmu.edu/~ggordon/10725-F12/scribes/10725_Lecture5.pdf\n",
    "    if backtracking:\n",
    "        # print('b', end='')\n",
    "        t = 0.1\n",
    "        t1 = time.perf_counter()\n",
    "        x_k = values - t * df_Rk\n",
    "        f_k = 0 # TODO: FIX THIS\n",
    "        for v in range(len(best_path[1:-1])):\n",
    "            f_k += cost_func(x_k[indices[v][0]:indices[v][1]])\n",
    "        L_norm = df_Rk.dot(df_Rk)\n",
    "        while f_k > (f_Rk - 0.5*t*(L_norm)):\n",
    "            t = beta * t\n",
    "            x_k = values - t * df_Rk\n",
    "            f_k = 0\n",
    "            for v in range(len(best_path[1:-1])):\n",
    "                f_k += cost_func(x_k[indices[v][0]:indices[v][1]])\n",
    "            # print(f_k, f_Rk, df_Rk.dot(df_Rk))\n",
    "\n",
    "            t2 = time.perf_counter()\n",
    "        # print(t2-t1)\n",
    "\n",
    "        return x_k, f_k, t2-t1, t, L_norm\n",
    "    else:\n",
    "        return df_Rk*step_size*-1, f_k, t2-t1, t, L_norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "02baf46c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:drake:Solved GCS shortest path using Mosek with convex_relaxation=true and preprocessing=true and rounding.\n",
      "INFO:drake:Finished 1000 rounding trials with Mosek.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start vertex is  <pydrake.geometry.optimization.GraphOfConvexSets.Vertex object at 0x71ee4012ec30>   and end vertex is  <pydrake.geometry.optimization.GraphOfConvexSets.Vertex object at 0x71ee402a0eb0>\n",
      "pruned\n",
      "Cost for best path found   2.140282439151793\n"
     ]
    }
   ],
   "source": [
    "order = 3\n",
    "start = Ratfk.ComputeSValue(np.array([-0.4, 0.1]), q_star)\n",
    "goal = Ratfk.ComputeSValue(np.array([1.1, 1.0]), q_star)\n",
    "\n",
    "ORDER = order\n",
    "gcs, flows_result = generate_flows(start, goal, ORDER, regions, dim=2)\n",
    "start_vertex, end_vertex = prune(gcs, flows_result, epsilon=0.00001)\n",
    "print(\"pruned\")\n",
    "best_path, best_rounded_result = run_dfs(flows_result, gcs, start_vertex, end_vertex, iterations=5)\n",
    "# ###############################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "2c751715",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # setup for PGD\n",
    "variables, indices, h_list = generate_variable_list(best_path, ndim=2)\n",
    "init_values = [best_rounded_result.GetSolution(x) for x in variables]\n",
    "init_values.extend([best_rounded_result.GetSolution(h) for h in h_list])\n",
    "init_values = np.asarray(init_values)\n",
    "\n",
    "import numpy.linalg as la\n",
    "# # iterate until infeasible.\n",
    "# # get constraints and put into HPolyhedron - should be the same as \n",
    "prog, _ = prog_with_constraints(variables, best_path, start, goal, indices, 2, h_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "61c4ae41",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "program set\n"
     ]
    }
   ],
   "source": [
    "print(\"program set\")\n",
    "hpoly = HPolyhedron(prog)\n",
    "# is A necessary?\n",
    "A = hpoly.A()\n",
    "b = hpoly.b()\n",
    "# create affine subspace\n",
    "A_eq, b_eq = get_equality_matrices(prog)\n",
    "U, S, V = la.svd(A_eq)\n",
    "\n",
    "# nullspace is spanned by the last n-r columns of V\n",
    "# rank is almost zero entries along the diagonal of S\n",
    "r = (S>0.00001).sum()\n",
    "basis = V[r:]\n",
    "translation = la.pinv(A_eq).dot(b_eq)\n",
    "\n",
    "affine_sub = AffineSubspace(basis.T, translation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "b44ab603",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:drake:Solved GCS shortest path using Mosek with convex_relaxation=true and preprocessing=true and rounding.\n",
      "INFO:drake:Finished 1000 rounding trials with Mosek.\n"
     ]
    }
   ],
   "source": [
    "gcs, flows_result = generate_flows(start, goal, ORDER, regions, dim=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "3eccea34",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:drake:Solved GCS shortest path using Mosek with convex_relaxation=true and preprocessing=true and rounding.\n",
      "INFO:drake:Finished 1000 rounding trials with Mosek.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[-0.84228838],\n",
       "       [-0.48305507]])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# start = seed_points_q[0]\n",
    "from pydrake.all import GcsTrajectoryOptimization, Point, GraphOfConvexSetsOptions\n",
    "order = 3\n",
    "gcs = GcsTrajectoryOptimization(2)\n",
    "main_graph = gcs.AddRegions(regions, order, h_min=0.1, h_max=100, name=\"\")\n",
    "start_graph = gcs.AddRegions([Point(start)], 0)\n",
    "goal_graph = gcs.AddRegions([Point(goal)], 0)\n",
    "gcs.AddEdges(start_graph, main_graph)\n",
    "gcs.AddEdges(main_graph, goal_graph)\n",
    "\n",
    "# gcs.AddTimeCost(1000)\n",
    "gcs.AddPathLengthCost()\n",
    "gcs.AddPathContinuityConstraints(2)\n",
    "\n",
    "options = GraphOfConvexSetsOptions()\n",
    "options.max_rounding_trials = 1000\n",
    "options.max_rounded_paths = 100\n",
    "options.convex_relaxation = True\n",
    "\n",
    "traj, _ = gcs.SolvePath(start_graph, goal_graph, options)\n",
    "traj.value(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eda8e2d9",
   "metadata": {},
   "source": [
    "## Running PGD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "246e5889",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "setup\n"
     ]
    }
   ],
   "source": [
    "t1_start = perf_counter() \n",
    "init_values = [best_rounded_result.GetSolution(x) for x in variables]\n",
    "init_values.extend([best_rounded_result.GetSolution(h) for h in h_list])\n",
    "init_values = np.asarray(init_values)\n",
    " \n",
    "iters = 50\n",
    "eps = 1e-7\n",
    "qps_solved = 0\n",
    "iter_costs = []\n",
    "x_axis = []\n",
    "times = []\n",
    "t_grads = []\n",
    "t_bt = []\n",
    "t_qp = []\n",
    "bt_costs = []\n",
    "time_steps = []\n",
    "lnorms = []\n",
    "running_qp_project = []\n",
    "termin = []\n",
    "total_ts = []\n",
    "costs_ts = []\n",
    "qp, _ = prog_with_constraints(variables, best_path, start, goal, indices, 2, h_list)\n",
    "init_values, qp_c = run_qp_proj(variables, init_values, best_path, start, goal, indices, h_list, qp)\n",
    "print('setup')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "a7a7ec94",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "warm start costfn took  0.7016762330022175\n"
     ]
    }
   ],
   "source": [
    "from jax import grad\n",
    "import jax\n",
    "# define costfn\n",
    "jax.clear_caches()\n",
    "costfn = jax.jit(distance_for_vertex, backend=\"cpu\")\n",
    "t1 = perf_counter()\n",
    "costfn(init_values[indices[0][0]:indices[0][1]])\n",
    "t2 = perf_counter()\n",
    "print(\"warm start costfn took \", t2 - t1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "bdadd717",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.8280335789968376"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grad_cost = grad(distance_for_vertex)\n",
    "grad_jit = jax.jit(grad_cost, backend='cpu')\n",
    "t1 = perf_counter()\n",
    "grad_jit(init_values[indices[1][0]:indices[1][1]])#.block_until_ready()\n",
    "t2 = perf_counter()\n",
    "t2 - t1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "7454d378",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     \n",
      "......Solve time: 4.003890148997016 QPs run 6 iteration_costs [Array(3.3618877, dtype=float32), Array(3.365778, dtype=float32), Array(3.3663993, dtype=float32), Array(3.366781, dtype=float32), Array(3.366707, dtype=float32), Array(3.3661795, dtype=float32), Array(3.3661795, dtype=float32)]\n"
     ]
    }
   ],
   "source": [
    "investigative_trajs = []\n",
    "exp_avg = 0\n",
    "moving_avg = 0\n",
    "diff_termin = []\n",
    "window = 5\n",
    "i = 0\n",
    "# for i in range(1):\n",
    "while (len(diff_termin)==0 or abs(diff_termin[-1])>0.005) and i < 70:\n",
    "    total_start = perf_counter()\n",
    "    grad_time = perf_counter()\n",
    "#     print('.')\n",
    "#     bt_cost, bt_time, ts, l_norm = 0, 0, 0, 0\n",
    "    init_values[:len(variables)], bt_cost, bt_time, ts, l_norm = get_step(init_values[:len(variables)], indices, best_path, costfn, grad_jit, step_size=0.01, backtracking=True, ndim=2)\n",
    "#     init_values[:len(variables)] = get_step(init_values[:len(variables)], indices, best_path, costfn, grad_jit, step_size=0.005, backtracking=False)\n",
    "    grad_time_stop = perf_counter()\n",
    "#     print(grad_time_stop - grad_time)\n",
    "    #### NOT TIMED #####\n",
    "    t_grads.append(grad_time_stop - grad_time)\n",
    "    t_bt.append(bt_time)\n",
    "    bt_costs.append(bt_cost)\n",
    "    time_steps.append(ts)\n",
    "    lnorms.append(l_norm)\n",
    "    ## NOT TIMED FIN ##\n",
    "    qp_time = perf_counter()\n",
    "    cheap_proj = affine_sub.Projection(init_values)[1].T[0] # DOUBLE CHECK THE SHAPE OF THIS\n",
    "    if not hpoly.PointInSet(cheap_proj):\n",
    "        init_values, qp_c = run_qp_proj(variables, cheap_proj, best_path, start, goal, indices, h_list, qp, qp_c)\n",
    "        qps_solved += 1\n",
    "        running_qp_project.append(True)\n",
    "    else:\n",
    "        init_values = cheap_proj\n",
    "        running_qp_project.append(False)\n",
    "    qp_time_stop = perf_counter()\n",
    "    #### NOT TIMED #####\n",
    "    t_qp.append(qp_time_stop - qp_time)\n",
    "    cost = 0\n",
    "    cost_calc_start = perf_counter()\n",
    "    for v in range(len(best_path[1:-1])):\n",
    "        cost += costfn(init_values[indices[v][0]:indices[v][1]])\n",
    "        if cost!=cost:\n",
    "            print(v, i)\n",
    "    iter_costs.append(cost)\n",
    "    if len(iter_costs) >= window:\n",
    "        moving_avg_new = sum(iter_costs[-window:])/window\n",
    "        termin.append(moving_avg_new)\n",
    "        if moving_avg != 0:\n",
    "            diff_termin.append((moving_avg_new - moving_avg)/moving_avg)\n",
    "        moving_avg = moving_avg_new\n",
    "#     exp_avg_new = alpha * cost + (1-alpha) * exp_avg\n",
    "#     exp_avg = exp_avg_new\n",
    "#     if len(iter_costs) > 4:\n",
    "#         moving_avg_new = sum(iter_costs[-5:])/5\n",
    "#         termin.append(moving_avg_new)\n",
    "#         diff_termin.append((moving_avg_new - moving_avg)/moving_avg)\n",
    "#         moving_avg = moving_avg_new\n",
    "    cost_calc_stop = perf_counter()\n",
    "    costs_ts.append(cost_calc_stop-cost_calc_start)\n",
    "    x_axis.append(i)\n",
    "    times.append(perf_counter()-t1_start)\n",
    "    if i % 10 == 0:\n",
    "#         qp, _ = prog_with_constraints(variables, best_path, start, goal, indices, h_list)\n",
    "        print(\"     \")\n",
    "    print('.', end='')\n",
    "    ## NOT TIMED FIN ##\n",
    "    total_end = perf_counter()\n",
    "    total_ts.append(total_end-total_start)\n",
    "    i += 1\n",
    "#     if i>70:\n",
    "#         investigative_trajs.append(trajectorify_given_vars(best_path, init_values, indices))\n",
    "#     else if i==79:\n",
    "        \n",
    "t1_stop = perf_counter() \n",
    "cost = 0\n",
    "for v in range(len(best_path[1:-1])):\n",
    "    cost += costfn(init_values[indices[v][0]:indices[v][1]])\n",
    "times.append(perf_counter()-t1_start)\n",
    "iter_costs.append(cost)\n",
    "x_axis.append(i)\n",
    "\n",
    "print(\"Solve time:\", t1_stop - t1_start, \"QPs run\", qps_solved, \"iteration_costs\", iter_costs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "24eb5664",
   "metadata": {},
   "outputs": [],
   "source": [
    "traj_pgd = trajectorify_given_vars(best_path, init_values, indices, ndim=2)\n",
    "traj_og = trajectorify_given_vars(best_path, [best_rounded_result.GetSolution(x) for x in variables], indices, ndim=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "ec66e8e0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x71ed283631c0>]"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYAAAAD7CAYAAABjVUMJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAfoElEQVR4nO3dfXRd1X3m8e+jd1mWLL8I29gYYxtsPB78Mo4LKaEzQFtPkhUgtA3TzBBMWkqmZTFtpwltpp2hTJvQptNhmhYWpRAyJQXqhE5qCg2LDKWZlhcTZBNbCuFVEraxsPCLrHfpN3/cI/laSNaVLOtK9zyftbS495x97t3bh6VH+5y991FEYGZm6VOU7wqYmVl+OADMzFLKAWBmllIOADOzlHIAmJmllAPAzCylxgwASRWSXpC0S9IeSbePUOYqSbsl1UvaKenSrH21krZLapTUIOmSZPsjSfl6SW9Jqp/UlpmZ2SlprHkAkgRURUS7pFLge8CtEfFcVpnZwPGICEkXAY9GxJpk34PAP0bEfZLKgFkRcXjYd/wRcCQifncyG2dmZqMrGatAZBKiPXlbmvzEsDLtWW+rBvdLqgEuA25IyvUAPdnHJgHzc8DlY9VlwYIFsXz58rGKmZlZlpdeeum9iKgbvn3MAACQVAy8BKwC/jQinh+hzDXAl4CzgI8lm1cArcADktYnn3FrRBzPOvQjwLsR8aNRvvsm4CaAZcuWsXPnzlyqbGZmCUlvj7Q9p5vAEdEfERuApcAWSetGKPNYctnnauCOZHMJsAm4OyI2AseB24Yd+u+AvzrFd98bEZsjYnNd3QcCzMzMJmhco4CSa/fPAFtPUeZZYKWkBUAL0JLVY9hOJhAAkFQCfBJ4ZFy1NjOz05bLKKA6SbXJ60rgSqBxWJlVybV8JG0CyoBDEXEAaJa0Oil6BbA369ArgcaIaDndhpiZ2fjkcg9gMfBgch+giMwInx2SbgaIiHuAa4HrJfUCncCn4sTwoluAh5IRQG8A27I++zpOcfnHzMzOnDGHgU4nmzdvDt8ENjMbH0kvRcTm4ds9E9jMLKUcAGZmKZXTPACb+XY1H+bZV1upKi9hdkUJNRUlzC4vpboi8766ooTq8lIqSotI7uebWYFzABS43v4B/uS7r/HV7/6IgRxu95QUaSgQBgOiurwkKyhKmV2eBEgSGtkBMliutNidS7PpzgFQwN5obedXH6lnV8sRPrlpCb/z8bUAHOvq41hXH+3dfRzr6qW9u4+jXX20d514P1jmWFcvB4528aODJ8r39o+dJBWlRcwuLz0RFBUlzC4fIUCS99WDIZK8n11RwuyyEoqK3BsxO1McAAUoIvirF5q5Y8deykqK+NOf38THLlo8tL92VtlpfXZ338DJAdKVBEjW+2PdJwJkMFDeO9bBsa5ejnVnyuYyAG0wHIb+W3Fyr2RJbSUXLq5hzeIa5lSWTrhdZmnkACgw77V384Xtu3m68SCXrlrAV352PYvmVEza50uiorSYitJi6qrLJ/w5AwNBR2//6AEyFCIn3rd393Gko4eW9zuGwqWrd2DoMwfD4MLF1cl/azh33iz3IsxG4QAoIE83vMsXvrmbo119/PbH17Ltw8un7S+/oiJlLvWUl8CciX1GRHDwWDd79x+lYf9RGvYfo2H/Ub7b+O7Q/Y5ZZcWsXnQiENYurmb1oprM95qlnCeCFYCOnj7+++MNfOP5JtYsquau6zayelF1vquVN129/bz67rGTQqFh/1GOdvUNlVk2b9ZJPYW1i2tYOrfSI6CsII02Ecx/Bs1wu5oP858eqeetQ8e56bIV/PpPXUB5SXG+q5VXFaXFXLS0louW1g5tiwj2HemiYV/SWziQCYfv7H136F5EdXkJa5JQWLMocylpzaIaKsvS/e9phcsBMEP19Q/wZ8+8zl1P/4iF1eU89As/xodXLsh3taYtSSyprWRJbSVXrl04tL2jp48fHjh2Uk/hW99/h/but5Pj4Lz5VR+4t7B4ToV7CzbjOQBmoLcPHedXH6nn+02HuWrD2fzuVes8AmaCZpWVsHHZXDYumzu0bWAgaHm/M+vewlF2v3OYx1/ZP1RmTmUpa5J7C2uTUDh/4WwqSt1bsJnDATCDRAR/vbOF2/92D0VF4q7rNnDVhiX5rlbBKSoSy+bPYtn8WWxdt2ho+7Gu3qS3cJS9SY/hkReb6eztB6C4SKxYUDXUS1izuJq1i2s4q7rcvQWblhwAM0Tb8R5u++ZuvrP3XS5eMY8/+rkNLKmtzHe1UqW6opTNy+exefm8oW39A8Hbh47TsP8YjQcyvYWX3n6fb+/aN1RmXlVZ5vLRopqhcFh11mzKSjxb2vLLATAD/N8fHuTz23dzpKOXL370Qj576XnTdnhn2hQXiRV1s1lRN/ukyXZHOnqTG80nhqh+/bm36enLzFsoLRYr62Z/4N7CgtkTn1thNl4OgGmss6efLz3RwNf/+W0uWDibB7dtYe3ZNfmuluVgzqxSLl4xn4tXzB/a1tc/wFuHjg9dPmrYf5R/ev09Hnv5naEyddXlQ6EweG/hvAVVXlvJzggHwDT1g3eOcOvDL/N663E+e+l5/MZPr/YNxhmupLiIVWdVs+qsaj6x/uyh7W3He07qKTTsP8r9r783tOZSWUkR5581e6iXMBgOp7Okhxl4Iti00z8Q3PMPr/PHT73KgtnlfOVn13Pp+R7emTa9/QO83tr+gcls77X3DJVZPKcimbNw4hLSeQuqKPblQRvGE8FmgIGB4NaHX2bH7v187KLF/N7V6/xXXkqVFhexZlFmQto1G09sP3isi8asQGjYf4xnX22lL1n7oqK0iNULq7N6C5nRSDUVHiZsH+QAmEa+9EQDO3bv5/NbV/O5n1jpoYP2AWdVV3BWdQWXXVA3tK27r58fvXuit9B44Ch/v+cAD7/YPFRmcKG8tVk3nJd5obzUcwBME/d/703+/B/f5IYPL/cvfxuX8pJi1i2Zw7olJ1bViwjePdqdzFk4MRope6G8qmEL5V3ohfJSx/cApoEnXtnPf/zG9/mptQv5s0//K1/DtTOmsyezUF5jshbSYDgcy1oo79z5s7LmLGQCwgvlzWy+BzBN7XyrjVsfqWfjObXcdd1G//K3M6qyrJj159Sy/pzaoW0RwTuHO0+62dyw/yh/v/fAiYXyKkq4cFHN0GJ5Fy6uYfXCai+UN8M5APLotYPt/MLXd7K0tpL7PvMhD/O0vJDE0rmzWDp3Fj+ZtVDe8e4+fvjuyTecv/lSC8d7MktfFAmWJ0tfrF18YvVUL5Q3czgA8uTgsS5ueOAFSorE17ZtYV6VR/vY9FJVXsKmZXPZNGyhvOb3O04anrq75TCP7z6xUF7trNKThqauTZa+8B84048DIA+Od/fx2a/t5FB7D4/80sUsmz8r31Uyy0lRkTh3fhXnzq9i67oTS18czVoob3CxvL96oWnokZ3FRWJlXdVJz1pYu7iGOi+Ul1cOgCnW1z/AL3/j++zZd4T7PrP5pIeWmM1UNRWlfGj5PD40ykJ5g8Hw4ptt/J/6Ewvlza8qO+lm8/IFVRQ7EEZ0Xl3VpM/ncABMoYjgv/zND3jmh638/jX/ksvXLBz7ILMZarSF8g539JwUCo0HjvHgP59YKM9G9rVtH+Jfrz5rUj/TATCF/uS7r/Hwi83ccvkqfv7HluW7OmZ5UTurjEtWzueSlScvlPfme8dpeb+TYOYMTZ9K2fM8JosDYIr89c5m/sdTr/LJTUv4tZ+8IN/VMZtWSoqLOH9hNecvrM53VVLFa8xOgWdfbeU3v/UKHzl/AV/+5EW+6WVm04ID4Azbs+8In/vLlzh/YTV/9ulNfgqUmU0bY/42klQh6QVJuyTtkXT7CGWukrRbUr2knZIuzdpXK2m7pEZJDZIuydp3i6QfJp/7B5PXrOmh5f0Otj3wInMqS/natg9R7RUZzWwayeUeQDdweUS0SyoFvifpiYh4LqvM08C3IyIkXQQ8CqxJ9t0FPBkRPyOpDJgFIOnfAFcBF0VEt6TJvb2dZ0c6ernhgRfp7O3nm5/7MAtrKvJdJTOzk4wZAJFZLa49eVua/MSwMu1Zb6sG90uqAS4DbkjK9QCDT7T4HPDliOhO9h2caCOmm67efn7xf++k6VAHD964hQt8Y8vMpqGcLkhLKpZUDxwEnoqI50coc42kRuBx4MZk8wqgFXhA0suS7pNUley7APiIpOcl/YOkD43y3Tcll5V2tra2jq91eTAwEPz6X+/ihTfb+MrPrT9pqJuZ2XSSUwBERH9EbACWAlskrRuhzGMRsQa4Grgj2VwCbALujoiNwHHgtqx9c4GLgd8AHtUIw2Mi4t6I2BwRm+vq6obvnna+9EQDj+/ez299dM1Jz301M5tuxjUkJSIOA88AW09R5llgpaQFQAvQktVj2E4mEEj2fSsyXgAGgBn98NsH/t+Jh7r84kdW5Ls6ZmanlMsooDpJtcnrSuBKoHFYmVWDf71L2gSUAYci4gDQLGl1UvQKYG/y+m+Ay5NjLkiOee8025M3T/5gP7+7Yy8//S8W8tsfX+ux/mY27eUyCmgx8KCkYjKB8WhE7JB0M0BE3ANcC1wvqRfoBD4VJx41dgvwUDIC6A1gW7L9fuB+ST8gc2P4M1nHzCg732rj1of9UBczm1n8SMjT9HprO9fe/U/MnVXGNz/3Ya/rb2bTzmiPhPS01NNw8FgXn7k/81CXB/1QFzObYbwY3ARlP9Tl4Zv8UBczm3ncA5ig3/+7BvbsO8KffnrjSQ/YNjObKRwAE/Ry02Euu6DOD3UxsxnLATABEUFzWwfL51eNXdjMbJpyAEzA4Y5ejnX3cc48X/c3s5nLATABTW0dACxzAJjZDOYAmAAHgJkVAgfABAwGwDnzKvNcEzOziXMATEBzWwcLZpczq8zTKMxs5nIATEBTWwfL/Ne/mc1wDoAJyASAr/+b2czmABin3v4B9h3udACY2YznABinfYc7GQg8B8DMZjwHwDh5CKiZFQoHwDgNBYBX/zSzGc4BME5NbR2UFRexsLoi31UxMzstDoBxam7rYOm8Sor82Eczm+EcAOPkIaBmVigcAOPUdMgBYGaFwQEwDkc6ejna1ecAMLOC4AAYhxOLwDkAzGzmcwCMg+cAmFkhcQCMg3sAZlZIHADj0NTWwfyqMmaXexloM5v5HADj0NzW4b/+zaxgOADGwXMAzKyQOABy1Nc/wDteBtrMCogDIEf7j3TRPxAOADMrGA6AHHkEkJkVGgdAjrwMtJkVmjEDQFKFpBck7ZK0R9LtI5S5StJuSfWSdkq6NGtfraTtkholNUi6JNn+3yS9kxxTL+mjk9u0ydXU1kFpsVhU42Wgzaww5DKgvRu4PCLaJZUC35P0REQ8l1XmaeDbERGSLgIeBdYk++4CnoyIn5FUBmT/Cf3HEfGVSWjHGdfU1sHSubMo9jLQZlYgxgyAiAigPXlbmvzEsDLtWW+rBvdLqgEuA25IyvUAPadb6XzwHAAzKzQ53QOQVCypHjgIPBURz49Q5hpJjcDjwI3J5hVAK/CApJcl3SepKuuwX0kuHd0vae4o331TcllpZ2tr6ziaNrkycwAq8/b9ZmaTLacAiIj+iNgALAW2SFo3QpnHImINcDVwR7K5BNgE3B0RG4HjwG3JvruBlcAGYD/wR6N8970RsTkiNtfV1eXYrMl1pLOXwx29HgJqZgVlXKOAIuIw8Ayw9RRlngVWSloAtAAtWT2G7WQCgYh4NwmWAeDPgS3jrv0UafYqoGZWgHIZBVQnqTZ5XQlcCTQOK7NKkpLXm4Ay4FBEHACaJa1Oil4B7E3KLc76iGuAH5xeU86cZs8BMLMClMsooMXAg5KKyQTGoxGxQ9LNABFxD3AtcL2kXqAT+FRy8xjgFuChZATQG8C2ZPsfSNpA5obxW8AvTU6TJp8ngZlZIcplFNBuYOMI2+/Jen0ncOcox9cDm0fY/h/GU9F8erutg7mzSqmpKM13VczMJo1nAueg2auAmlkBcgDkoMlzAMysADkAxtDXP8A773sZaDMrPA6AMew/0kWfl4E2swLkABiD5wCYWaFyAIzBQ0DNrFA5AMbQ1NZBSZFYPMfLQJtZYXEAjKGprYMlcyspKfY/lZkVFv9WG4PnAJhZoXIAjMFzAMysUDkATuFoVy/vexloMytQDoBT8BBQMytkDoBTcACYWSFzAJyC5wCYWSFzAJxCU1sHcypLmVPpZaDNrPA4AE6hqc2LwJlZ4XIAnILnAJhZIXMAjKJ/IGh533MAzKxwOQBGceBoF739XgbazAqXA2AUTYc8BNTMCpsDYBSeA2Bmhc4BMIqmtg6Ki8TiWi8DbWaFyQEwiqa2Ds6uraDUy0CbWYHyb7dRNHkIqJkVOAfAKDwHwMwKnQNgBO3dfRw63uM5AGZW0BwAI/AIIDNLAwfACJocAGaWAg6AEbgHYGZp4AAYQVNbB9UVJV4G2swKmgNgBINDQCXluypmZmfMmAEgqULSC5J2Sdoj6fYRylwlabekekk7JV2ata9W0nZJjZIaJF0y7Nj/LCkkLZicJp2+prYOzp3vyz9mVthy6QF0A5dHxHpgA7BV0sXDyjwNrI+IDcCNwH1Z++4CnoyINcB6oGFwh6RzgJ8EmibagMk2MBC0tHV6CKiZFbwxAyAy2pO3pclPDCvTHhGD26oG90uqAS4D/iIp1xMRh7MO/WPg88M/L5/ePdZFT/+AbwCbWcHL6R6ApGJJ9cBB4KmIeH6EMtdIagQeJ9MLAFgBtAIPSHpZ0n2SqpLynwDeiYhdk9COSeNloM0sLXIKgIjoTy7vLAW2SFo3QpnHkss8VwN3JJtLgE3A3RGxETgO3CZpFvBF4HfG+m5JNyX3FXa2trbmUt3T4jkAZpYW4xoFlFy+eQbYeooyzwIrk5u6LUBLVo9hO5lAWAmcB+yS9BaZYPm+pEUjfN69EbE5IjbX1dWNp7oT0tzWQZHg7NrKM/5dZmb5lMsooDpJtcnrSuBKoHFYmVVKxkxK2gSUAYci4gDQLGl1UvQKYG9EvBIRZ0XE8ohYTiYoNiXl8yqzDHSll4E2s4JXkkOZxcCDkorJBMajEbFD0s0AEXEPcC1wvaReoBP4VNZN4VuAhySVAW8A2ya7EZPJy0CbWVqMGQARsRvYOML2e7Je3wncOcrx9cDmMb5j+Vj1mCpNbZ1ceeFZ+a6GmdkZ5+scWTp6+nivvdtzAMwsFRwAWZrbOgGPADKzdHAAZPEQUDNLEwdAFgeAmaWJAyBLc1sH1eUl1M7yMtBmVvgcAFma2jo4x8tAm1lKOACyeA6AmaWJAyAxMBA0t3WwzM8BMLOUcAAkWtu76e4b8BwAM0sNB0DCI4DMLG0cAAk/B8DM0sYBkGhq60CCJV4G2sxSwgGQaG7r4Ow5lZSV+J/EzNLBv+0SmTkA/uvfzNLDAZDwHAAzSxsHANDZ08/BY90OADNLFQcA0PJ+ZgSQ5wCYWZo4APAcADNLJwcADgAzSycHAJkAqCorZl5VWb6rYmY2ZRwAZOYAeBloM0sbBwAeAmpm6ZT6AIgIB4CZpVLqA6C1vZuu3gE/B8DMUif1ATC4CqjnAJhZ2jgAPATUzFLKAeBloM0spRwAbR0sqqmgorQ431UxM5tSqQ+AwTkAZmZpk/oA8BBQM0urVAdAV28/7x71MtBmlk6pDoDBZaAdAGaWRmMGgKQKSS9I2iVpj6TbRyhzlaTdkuol7ZR0ada+WknbJTVKapB0SbL9jqxjviPp7Mlt2tgGh4D6HoCZpVEuPYBu4PKIWA9sALZKunhYmaeB9RGxAbgRuC9r313AkxGxBlgPNCTb/zAiLkqO2QH8zkQbMVGDk8DcAzCzNCoZq0BEBNCevC1NfmJYmfast1WD+yXVAJcBNyTleoCe5PXRkY6ZSk1tnVSWFrNgtpeBNrP0yekegKRiSfXAQeCpiHh+hDLXSGoEHifTCwBYAbQCD0h6WdJ9kqqyjvk9Sc3ApxmlByDppuSy0s7W1tbxtG1MgyOAvAy0maVRTgEQEf3JpZqlwBZJ60Yo81hymedq4I5kcwmwCbg7IjYCx4Hbso75YkScAzwE/Moo331vRGyOiM11dXU5NywXngNgZmk2rlFAEXEYeAbYeooyzwIrJS0AWoCWrB7DdjKBMNw3gGvHU5fT5WWgzSztchkFVCepNnldCVwJNA4rs0rJdRRJm4Ay4FBEHACaJa1Oil4B7E3KnZ/1EZ8Y/pln2nvtPXT29rNsntcAMrN0GvMmMLAYeFBSMZnAeDQidki6GSAi7iHz1/v1knqBTuBTyc1jgFuAhySVAW8A25LtX06CYQB4G7h5shqVi6FVQP0cADNLqVxGAe0GNo6w/Z6s13cCd45yfD2weYTtU3rJZ7hmLwNtZimX2pnAgz2ApXMdAGaWTqkOgIU15V4G2sxSK9UB4Ms/ZpZmqQ0AzwEws7RLZQB09fZz4GiXewBmlmqpDIB3DncS4RFAZpZuqQyAJg8BNTNLZwB4DoCZWUoDoOlQB+UlRdRVl+e7KmZmeZPOAPAy0GZm6Q4AM7M0S10ARITnAJiZkcIAaDvew/GefvcAzCz1UhcAHgJqZpaR3gDwcwDMLOVSFwCDcwDO8TLQZpZyqQuAprYO6qrLqSzzMtBmlm6pDABf/zczS2EANLd1OgDMzEhZAPT0DbDvSKfnAJiZkbIA8DLQZmYnpCoAPAfAzOwEB4CZWUqlKgCa2zLLQJ/lZaDNzNIVAE2HMovAFRV5GWgzs3QFgOcAmJkNSU0ADC4D7QAwM8tITQAc7ujlWHef5wCYmSVSEwAeAWRmdjIHgJlZSqUuAM6ZV5nnmpiZTQ9jBoCkCkkvSNolaY+k20coc5Wk3ZLqJe2UdGnWvlpJ2yU1SmqQdEmy/Q+TbbslPSapdlJbNkxzWwcLZpczq6zkTH6NmdmMkUsPoBu4PCLWAxuArZIuHlbmaWB9RGwAbgTuy9p3F/BkRKwB1gMNyfangHURcRHwKvCbE21ELjJDQP3Xv5nZoDEDIDLak7elyU8MK9MeEYPbqgb3S6oBLgP+IinXExGHk9ffiYi+5JjngKWn15RT8xwAM7OT5XQPQFKxpHrgIPBURDw/QplrJDUCj5PpBQCsAFqBByS9LOk+SVUjfMWNwBOjfPdNyWWlna2trblU9wN6+wfYd9jPATAzy5ZTAEREf3J5ZymwRdK6Eco8llzmuRq4I9lcAmwC7o6IjcBx4Lbs4yR9EegDHhrlu++NiM0Rsbmuri6nRg2373AnA4HnAJiZZRnXKKDk8s0zwNZTlHkWWClpAdACtGT1GLaTCQQAJH0G+Djw6axLSJPOQ0DNzD4ol1FAdYMjdCRVAlcCjcPKrJKk5PUmoAw4FBEHgGZJq5OiVwB7k3JbgS8An4iIjslpzsiGAmC+A8DMbFAuYyIXAw9KKiYTGI9GxA5JNwNExD3AtcD1knqBTuBTWX/R3wI8JKkMeAPYlmz/KlAOPJVkx3MRcfMkteskTW0dlBUXsbC64kx8vJnZjDRmAETEbmDjCNvvyXp9J3DnKMfXA5tH2L5qPBU9HefNr+KajUu8DLSZWZZUzIq6bssyrtuyLN/VMDObVlKzFISZmZ3MAWBmllIOADOzlHIAmJmllAPAzCylHABmZinlADAzSykHgJlZSukMrsE26SS1Am/nux5TYAHwXr4rkUduv9vv9k+ucyPiA8spz6gASAtJOyPiA8tnpIXb7/a7/VPTfl8CMjNLKQeAmVlKOQCmp3vzXYE8c/vTze2fIr4HYGaWUu4BmJmllAPAzCylHABTRFKFpBck7ZK0R9LtI5SZI+lvs8psy9q3VdIPJb0m6baprf3pm4T2vyXpFUn1knZObe1PX47tnyvpMUm7k7Lrsval4fyfqv0z+vwPklQs6WVJO0bYJ0n/KznHu5Pnqw/uOzPnPyL8MwU/gIDZyetS4Hng4mFlfgu4M3ldB7QBZUAx8DqwInm/C1ib7zZNVfuT928BC/LdjjPc/j8E/mvyeg3wdPI6Led/xPYXwvnPasevAd8Adoyw76PAE8m/1cXA82f6/LsHMEUioz15W5r8DL8DH0C1JAGzyfwC7AO2AK9FxBsR0QM8DFw1NTWfHKfZ/hkvx/avBZ5OyjcCyyUtJD3nf7T2FwRJS4GPAfeNUuQq4OvJv9VzQK2kxZzB8+8AmEJJ968eOAg8FRHPDyvyVeBCYB/wCnBrRAwAS4DmrHItybYZ5TTaD5lfFt+R9JKkm6aqzpMph/bvAj6ZlN0CnAssJT3nf7T2QwGcf+B/Ap8HBkbZP9p5PmPn3wEwhSKiPyI2kPmfekv2Nc7ETwP1wNnABuCrkmrIdAk/8HFnrqZnxmm0H+DHI2IT8G+BX5Z02ZRUehLl0P4vA3OTX5K3AC+T6QGl5fyP1n6Y4edf0seBgxHx0qmKjbAtTrH9tDkA8iAiDgPPAFuH7doGfCvpAr4GvEnmWmgLcE5WuaVk/kqekSbQfiJiX/Lfg8BjZLrFM9Jo7Y+IoxGxLfkleT2Z+yBvkpLzf4r2F8L5/3HgE5LeInMJ53JJfzmszGjn+YydfwfAFJFUJ6k2eV0JXAk0DivWBFyRlFkIrAbeAF4Ezpd0nqQy4Drg21NU9UlxOu2XVCWpOtleBfwU8IMpqvqkyKX9kmqT8wvwC8CzEXGUlJz/0dpfCOc/In4zIpZGxHIy5++7EfHvhxX7NnB9MhroYuBIROznDJ7/ksn4EMvJYuBBScVkgvfRiNgh6WaAiLgHuAP4mqRXyHT7vhAR7wFI+hXg78mMCLg/IvbkoxGnYcLtl7QCeCxzb5gS4BsR8WReWjFxubT/QuDrkvqBvcBnk319KTn/I7YfWMjMP/8jGtb+vyMzEug1oINMj/iMnn8vBWFmllK+BGRmllIOADOzlHIAmJmllAPAzCylHABmZinlADAzSykHgJlZSv1/vGVvXRQRRRwAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.plot(times, iter_costs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b6230ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "traj_pgd.end_time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0401808b",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(init_values)\n",
    "ndim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66b6338b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d1bf85",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(best_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "750e847d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dumb init\n",
    "pts = np.squeeze(np.array([np.arctan(traj.value(it * traj.end_time() / 100)) for it in range(100)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0041d941",
   "metadata": {},
   "outputs": [],
   "source": [
    "# np.hstack([pts, 0 * np.ones((pts.shape[0], 3 - pts.shape[1]))]).T\n",
    "pts.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "4bacd7c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # visualize path given trajectory\n",
    "import visualization_utils as viz_utils\n",
    "# viz_utils.draw_traj(visualizer.meshcat_cspace, traj, 100)\n",
    "# viz_utils.draw_nontc_traj(visualizer.meshcat_nontc, traj, 100, q_star=q_star)\n",
    "# visualizer.animate_traj_s(traj, 100, 100)\n",
    "viz_utils.draw_traj(visualizer.meshcat_cspace, traj_og, 100)\n",
    "viz_utils.draw_nontc_traj(visualizer.meshcat_nontc, traj_og, 100, q_star=q_star)\n",
    "visualizer.animate_traj_s(traj_og, 100, 100)\n",
    "viz_utils.draw_traj(visualizer.meshcat_cspace, traj_pgd, 100)\n",
    "viz_utils.draw_nontc_traj(visualizer.meshcat_nontc, traj_pgd, 100, q_star=q_star)\n",
    "visualizer.animate_traj_s(traj_pgd, 100, 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f111c2bb",
   "metadata": {},
   "source": [
    "## Validating Paths?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "d6a0076e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.377200667606015 4.387048193241586\n"
     ]
    }
   ],
   "source": [
    "# take a path\n",
    "# calculate the q value at each point along the trajectory\n",
    "# Sum all of them\n",
    "def length_of_traj(traj):\n",
    "    q_prev = np.array([0.3, 1.0])\n",
    "    dist = 0\n",
    "    for i in range(50):\n",
    "        s = traj.value(traj.end_time()*i/50)\n",
    "        q = Ratfk.ComputeQValue(s, q_star)\n",
    "        dist += ((q - q_prev).dot(q - q_prev))**0.5\n",
    "        q_prev = q\n",
    "    return dist\n",
    "new_len = length_of_traj(traj_pgd)\n",
    "og_len = length_of_traj(traj_og)\n",
    "print(new_len, og_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "cc8177b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0022446814353992545"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(og_len - new_len) / og_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "f1f2617b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 1.])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q_star"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "40e6c7ea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.0"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "traj_og.end_time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "29d0238d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.0"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "traj_pgd.end_time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "b9e62178",
   "metadata": {},
   "outputs": [],
   "source": [
    "iter_dist = 0\n",
    "values = init_values\n",
    "for v in range(len(best_path[1:-1])):\n",
    "    iter_dist += distance_for_vertex(values[indices[v][0]:indices[v][1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "405ef217",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(3.3586519, dtype=float32)"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "iter_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "a1183fdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "iter_dist = 0\n",
    "values = [best_rounded_result.GetSolution(x) for x in variables]\n",
    "for v in range(len(best_path[1:-1])):\n",
    "    iter_dist += distance_for_vertex(values[indices[v][0]:indices[v][1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "8b7721ec",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(3.3597455, dtype=float32)"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "iter_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "910c39e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 1.])"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q_star"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c52a00d2",
   "metadata": {},
   "source": [
    "## 7 DOF IIWA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "929fec7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#construct our robot\n",
    "builder = DiagramBuilder()\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "parser = Parser(plant)\n",
    "\n",
    "parser.package_map().Add(\"ciris_pgd\", os.path.abspath(''))\n",
    "\n",
    "directives_file = \"/home/sgrg/rlg/SUPERUROP/ciris/models/iiwa14_sphere_collision_complex_scenario.dmd.yaml\"\n",
    "directives = LoadModelDirectives(directives_file)\n",
    "models = ProcessModelDirectives(directives, plant, parser)\n",
    "\n",
    "\n",
    "# oneDOF_iiwa_file = \"/home/sgrg/rlg/SUPERUROP/ciris/assets/oneDOF_iiwa7_with_box_collision.sdf\"\n",
    "# with open(one|DOF_iiwa_file, 'r') as f:\n",
    "#     oneDOF_iiwa_string = f.read()\n",
    "# box_asset_file = \"/home/sgrg/rlg/SUPERUROP/ciris/assets/box_small.urdf\"\n",
    "# with open(box_asset_file, 'r') as f:\n",
    "#     box_asset_string = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69731f7f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}