demos/planning_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "This file is part of DGPMP2-ND    \n",
    "\n",
    "Copyright (C) 2024 ArtiMinds Robotics GmbH\n",
    "\n",
    "This program is free software: you can redistribute it and/or modify\n",
    "it under the terms of the GNU General Public License as published by\n",
    "the Free Software Foundation, either version 3 of the License, or\n",
    "(at your option) any later version.\n",
    "\n",
    "This program is distributed in the hope that it will be useful,\n",
    "but WITHOUT ANY WARRANTY; without even the implied warranty of\n",
    "MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the\n",
    "GNU General Public License for more details.\n",
    "\n",
    "You should have received a copy of the GNU General Public License\n",
    "along with this program.  If not, see <https://www.gnu.org/licenses/>.\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "from pathlib import Path\n",
    "import open3d as o3d\n",
    "from open3d.geometry import TriangleMesh\n",
    "import torch\n",
    "\n",
    "from diff_gpmp2.differentiable_robot_model.diff_robot_model import DiffRobotModel\n",
    "from diff_gpmp2.utils.sdf import CollisionObject, SignedDistanceField\n",
    "\n",
    "NOTEBOOK_DIR = Path(os.path.abspath(os.path.dirname(\"__file__\")))\n",
    "\n",
    "\n",
    "boxes = [\n",
    "    CollisionObject(torch.tensor([-0.5, -0.1, 0.1, 0, 0, 0]), TriangleMesh.create_box(width=0.2, height=0.2, depth=0.2)),\n",
    "    CollisionObject(torch.tensor([-0.6, -0.2, 0.2, 0, 0, 0]), TriangleMesh.create_box(width=0.1, height=0.1, depth=0.1)),\n",
    "    CollisionObject(torch.tensor([-0.2, -0.5, 0.05, 0, 0, 0]), TriangleMesh.create_box(width=0.1, height=0.1, depth=0.1)),\n",
    "    CollisionObject(torch.tensor([0, 0.5, 0.05, 0, 0, 0]), TriangleMesh.create_box(width=0.1, height=0.1, depth=0.1))\n",
    "]\n",
    "robot = DiffRobotModel(\n",
    "    urdf_path=(NOTEBOOK_DIR / 'urdf' / 'ur5' / 'robot.urdf').as_posix(),\n",
    "    good_initial_conf=torch.deg2rad(torch.tensor([23, -86, 94, -98, -90, 113]))\n",
    ")\n",
    "sdf = SignedDistanceField.from_obstacles(\n",
    "    obstacles=boxes,\n",
    "    collision_padding=0.01,\n",
    "    voxel_size=0.01)\n",
    "\n",
    "start_conf = torch.deg2rad(torch.tensor([-46, -79, 112, -123, -90, 43]))\n",
    "start_vel = torch.zeros(6)\n",
    "start = torch.cat((start_conf, start_vel))\n",
    "goal_conf = torch.deg2rad(torch.tensor([26, -76, 108, -122, -90, 116]))\n",
    "goal_vel = torch.zeros(6)\n",
    "goal = torch.cat((goal_conf, goal_vel))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diff_gpmp2.gpmp2.diff_gpmp2_planner import DiffGPMP2Planner\n",
    "from diff_gpmp2.gpmp2.factors import Factors\n",
    "from diff_gpmp2.gpmp2.gp.gp_factor import GPFactor\n",
    "from diff_gpmp2.gpmp2.gp.prior_factor import PriorFactor\n",
    "from diff_gpmp2.gpmp2.gp.robot_limits_factor import RobotLimitsFactor\n",
    "from diff_gpmp2.gpmp2.obstacle.obstacle_factor import ObstacleFactor\n",
    "from diff_gpmp2.gpmp2.plan_layer import OptimParams\n",
    "from diff_gpmp2.utils.planner_utils import straight_line_traj\n",
    "\n",
    "device = torch.device('cpu')\n",
    "\n",
    "# Factors\n",
    "batch_size = 1\n",
    "state_dim = robot.n_dofs * 2  # position & velocity\n",
    "wksp_dim = 3\n",
    "total_time_sec = 10\n",
    "total_time_step = 20\n",
    "\n",
    "gp_prior = GPFactor(\n",
    "    name=\"GP\",\n",
    "    weight=torch.tensor(1),\n",
    "    dof=robot.n_dofs,\n",
    "    device=device,\n",
    "    threshold=10.0\n",
    ")\n",
    "gp_prior.set_total_time_sec(total_time_sec)\n",
    "start_prior = PriorFactor(\n",
    "    name=\"Start\",\n",
    "    meanb=start,\n",
    "    K=torch.tensor(0.01),\n",
    "    state_dim=state_dim,\n",
    "    thb_point_idx=0,\n",
    "    device=device,\n",
    "    threshold=1e-3,\n",
    ")\n",
    "goal_prior = PriorFactor(\n",
    "    name=\"Goal\",\n",
    "    meanb=goal,\n",
    "    K=torch.tensor(0.01),\n",
    "    state_dim=state_dim,\n",
    "    thb_point_idx=-1,\n",
    "    device=device,\n",
    "    threshold=1e-3,\n",
    ")\n",
    "obs_factor = ObstacleFactor(\n",
    "    name=\"Obst.\",\n",
    "    state_dim=state_dim,\n",
    "    weight=torch.tensor(0.05),\n",
    "    threshold=0,\n",
    "    eps=0.001,\n",
    "    sdf=sdf,\n",
    "    robot_model=robot,\n",
    "    device=device,\n",
    ")\n",
    "lims_factor = RobotLimitsFactor(\n",
    "    name=\"Limits\",\n",
    "    state_dim=state_dim,\n",
    "    robot_model=robot,\n",
    "    K=torch.tensor(0.01),\n",
    "    device=device,\n",
    "    threshold=0\n",
    ")\n",
    "\n",
    "factors = Factors(\n",
    "    state_dim=state_dim,\n",
    "    batch_size=batch_size,\n",
    "    device=device,\n",
    "    factors=[start_prior, goal_prior, gp_prior, obs_factor, lims_factor],\n",
    ")\n",
    "\n",
    "optim_params = OptimParams(method=\"gauss_newton\", reg=0.1, init_lr=1, max_iters=100,\n",
    "                           error_patience=15)\n",
    "\n",
    "planner = DiffGPMP2Planner(\n",
    "    factors=factors,\n",
    "    optim_params=optim_params,\n",
    "    state_dim=state_dim,\n",
    "    batch_size=batch_size,\n",
    "    device=torch.device('cpu')\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tempfile import TemporaryDirectory\n",
    "import cv2\n",
    "\n",
    "th_init = straight_line_traj(\n",
    "    start_conf, goal_conf, total_time_sec, total_time_step, robot.n_dofs, device\n",
    ")\n",
    "\n",
    "obs_factor.plot_trajectory(th_init.unsqueeze(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Start': 5.960464477539063e-08, 'Goal': 1.1920928955078125e-07, 'GP': 0.12716001272201538, 'Obst.': 0.4467200040817261, 'Limits': 0.0}\n",
      "{'Start': 3.516674041748047e-06, 'Goal': 2.9772520065307617e-05, 'GP': 1.1861517429351807, 'Obst.': 0.3933199942111969, 'Limits': 0.0}\n",
      "{'Start': 6.020069122314453e-06, 'Goal': 3.978610038757324e-05, 'GP': 1.8765300512313843, 'Obst.': 0.39169999957084656, 'Limits': 0.0}\n",
      "{'Start': 6.556510925292969e-06, 'Goal': 4.360079765319824e-05, 'GP': 1.6250630617141724, 'Obst.': 0.3894599974155426, 'Limits': 0.0}\n",
      "{'Start': 9.775161743164062e-06, 'Goal': 4.690885543823242e-05, 'GP': 1.9955742359161377, 'Obst.': 0.3844600021839142, 'Limits': 0.0}\n",
      "{'Start': 1.2278556823730469e-05, 'Goal': 4.926323890686035e-05, 'GP': 2.253143548965454, 'Obst.': 0.37386003136634827, 'Limits': 0.0}\n",
      "{'Start': 9.417533874511719e-06, 'Goal': 5.1140785217285156e-05, 'GP': 2.1540846824645996, 'Obst.': 0.3797999918460846, 'Limits': 0.0}\n",
      "{'Start': 9.894371032714844e-06, 'Goal': 5.3256750106811523e-05, 'GP': 2.436846971511841, 'Obst.': 0.36480000615119934, 'Limits': 0.0}\n",
      "{'Start': 1.0013580322265625e-05, 'Goal': 5.346536636352539e-05, 'GP': 2.464782238006592, 'Obst.': 0.3643600046634674, 'Limits': 0.0}\n",
      "{'Start': 1.0073184967041016e-05, 'Goal': 5.1349401473999023e-05, 'GP': 2.846646785736084, 'Obst.': 0.32392001152038574, 'Limits': 0.0}\n",
      "{'Start': 9.894371032714844e-06, 'Goal': 2.6047229766845703e-05, 'GP': 2.6042513847351074, 'Obst.': 0.3446199893951416, 'Limits': 0.0}\n",
      "{'Start': 8.046627044677734e-06, 'Goal': 4.5239925384521484e-05, 'GP': 2.580972671508789, 'Obst.': 0.3173600137233734, 'Limits': 0.0}\n",
      "{'Start': 5.0067901611328125e-06, 'Goal': 2.6881694793701172e-05, 'GP': 2.358196973800659, 'Obst.': 0.3175399899482727, 'Limits': 0.0}\n",
      "{'Start': 4.708766937255859e-06, 'Goal': 2.422928810119629e-05, 'GP': 2.076734781265259, 'Obst.': 0.31391996145248413, 'Limits': 0.0}\n",
      "{'Start': 5.245208740234375e-06, 'Goal': 5.7071447372436523e-05, 'GP': 2.0428903102874756, 'Obst.': 0.29339998960494995, 'Limits': 0.0}\n",
      "{'Start': 6.318092346191406e-06, 'Goal': 2.9087066650390625e-05, 'GP': 1.98263680934906, 'Obst.': 0.2740800082683563, 'Limits': 0.0}\n",
      "{'Start': 6.854534149169922e-06, 'Goal': 3.546476364135742e-05, 'GP': 2.017223596572876, 'Obst.': 0.2777000069618225, 'Limits': 0.0}\n",
      "{'Start': 7.331371307373047e-06, 'Goal': 3.74913215637207e-05, 'GP': 2.109099864959717, 'Obst.': 0.26322001218795776, 'Limits': 0.0}\n",
      "{'Start': 8.046627044677734e-06, 'Goal': 2.98917293548584e-05, 'GP': 2.1502487659454346, 'Obst.': 0.21974000334739685, 'Limits': 0.0}\n",
      "{'Start': 6.377696990966797e-06, 'Goal': 2.1755695343017578e-05, 'GP': 1.8549418449401855, 'Obst.': 0.2088799923658371, 'Limits': 0.0}\n",
      "{'Start': 5.602836608886719e-06, 'Goal': 1.806020736694336e-05, 'GP': 1.7961905002593994, 'Obst.': 0.1955999881029129, 'Limits': 0.0}\n",
      "{'Start': 5.125999450683594e-06, 'Goal': 1.817941665649414e-05, 'GP': 1.68352472782135, 'Obst.': 0.18836000561714172, 'Limits': 0.0}\n",
      "{'Start': 5.125999450683594e-06, 'Goal': 3.5703182220458984e-05, 'GP': 1.7985186576843262, 'Obst.': 0.1690399944782257, 'Limits': 0.0}\n",
      "{'Start': 4.708766937255859e-06, 'Goal': 4.32431697845459e-05, 'GP': 1.634299874305725, 'Obst.': 0.16541999578475952, 'Limits': 0.0}\n",
      "{'Start': 4.470348358154297e-06, 'Goal': 4.6253204345703125e-05, 'GP': 1.788252830505371, 'Obst.': 0.12316000461578369, 'Limits': 0.0}\n",
      "{'Start': 4.112720489501953e-06, 'Goal': 3.8057565689086914e-05, 'GP': 1.5769553184509277, 'Obst.': 0.1424800008535385, 'Limits': 0.0}\n",
      "{'Start': 4.172325134277344e-06, 'Goal': 4.017353057861328e-05, 'GP': 1.6046658754348755, 'Obst.': 0.1328199952840805, 'Limits': 0.0}\n",
      "{'Start': 4.0531158447265625e-06, 'Goal': 3.790855407714844e-05, 'GP': 1.5696766376495361, 'Obst.': 0.11349999904632568, 'Limits': 0.0}\n",
      "{'Start': 3.7550926208496094e-06, 'Goal': 3.787875175476074e-05, 'GP': 1.5510854721069336, 'Obst.': 0.11349999904632568, 'Limits': 0.0}\n",
      "{'Start': 3.874301910400391e-06, 'Goal': 3.731250762939453e-05, 'GP': 1.5353924036026, 'Obst.': 0.10383999347686768, 'Limits': 0.0}\n",
      "{'Start': 4.231929779052734e-06, 'Goal': 3.6656856536865234e-05, 'GP': 1.5191969871520996, 'Obst.': 0.10745999962091446, 'Limits': 0.0}\n",
      "{'Start': 4.410743713378906e-06, 'Goal': 3.6329030990600586e-05, 'GP': 1.5220057964324951, 'Obst.': 0.07848000526428223, 'Limits': 0.0}\n",
      "{'Start': 4.291534423828125e-06, 'Goal': 3.5375356674194336e-05, 'GP': 1.4681055545806885, 'Obst.': 0.07485999166965485, 'Limits': 0.0}\n",
      "{'Start': 4.410743713378906e-06, 'Goal': 3.460049629211426e-05, 'GP': 1.4940447807312012, 'Obst.': 0.06881999969482422, 'Limits': 0.0}\n",
      "{'Start': 4.470348358154297e-06, 'Goal': 3.4362077713012695e-05, 'GP': 1.50822913646698, 'Obst.': 0.0494999997317791, 'Limits': 0.0}\n",
      "{'Start': 4.410743713378906e-06, 'Goal': 3.471970558166504e-05, 'GP': 1.5291274785995483, 'Obst.': 0.033799998462200165, 'Limits': 0.0}\n",
      "{'Start': 3.814697265625e-06, 'Goal': 3.457069396972656e-05, 'GP': 1.4188724756240845, 'Obst.': 0.03500000014901161, 'Limits': 0.0}\n",
      "{'Start': 3.3974647521972656e-06, 'Goal': 3.382563591003418e-05, 'GP': 1.3767887353897095, 'Obst.': 0.03017999976873398, 'Limits': 0.0}\n",
      "{'Start': 3.2782554626464844e-06, 'Goal': 3.3408403396606445e-05, 'GP': 1.3540339469909668, 'Obst.': 0.033799998462200165, 'Limits': 0.0}\n",
      "{'Start': 3.5762786865234375e-06, 'Goal': 3.316998481750488e-05, 'GP': 1.3340615034103394, 'Obst.': 0.016899999231100082, 'Limits': 0.0}\n",
      "{'Start': 3.159046173095703e-06, 'Goal': 3.272294998168945e-05, 'GP': 1.3003944158554077, 'Obst.': 0.022940000519156456, 'Limits': 0.0}\n",
      "{'Start': 3.2782554626464844e-06, 'Goal': 3.24249267578125e-05, 'GP': 1.2765458822250366, 'Obst.': 0.013280000537633896, 'Limits': 0.0}\n",
      "{'Start': 2.9802322387695312e-06, 'Goal': 3.2007694244384766e-05, 'GP': 1.253401279449463, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.9802322387695312e-06, 'Goal': 3.159046173095703e-05, 'GP': 1.2317312955856323, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 3.0994415283203125e-06, 'Goal': 3.11732292175293e-05, 'GP': 1.210630178451538, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 3.039836883544922e-06, 'Goal': 3.075599670410156e-05, 'GP': 1.1900683641433716, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.9206275939941406e-06, 'Goal': 3.0249357223510742e-05, 'GP': 1.1700384616851807, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.7418136596679688e-06, 'Goal': 2.9712915420532227e-05, 'GP': 1.1505581140518188, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.7418136596679688e-06, 'Goal': 2.9325485229492188e-05, 'GP': 1.1316444873809814, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.7418136596679688e-06, 'Goal': 2.893805503845215e-05, 'GP': 1.1133087873458862, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.682209014892578e-06, 'Goal': 2.866983413696289e-05, 'GP': 1.0955548286437988, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.682209014892578e-06, 'Goal': 2.8401613235473633e-05, 'GP': 1.078379511833191, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.6226043701171875e-06, 'Goal': 2.804398536682129e-05, 'GP': 1.0617791414260864, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.6226043701171875e-06, 'Goal': 2.7894973754882812e-05, 'GP': 1.0457446575164795, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.5033950805664062e-06, 'Goal': 2.7507543563842773e-05, 'GP': 1.0302729606628418, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.562999725341797e-06, 'Goal': 2.7358531951904297e-05, 'GP': 1.0153508186340332, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.4437904357910156e-06, 'Goal': 2.7120113372802734e-05, 'GP': 1.000967025756836, 'Obst.': 0.00965999998152256, 'Limits': 0.0}\n",
      "{'Start': 2.384185791015625e-06, 'Goal': 2.6971101760864258e-05, 'GP': 0.9871139526367188, 'Obst.': 0.0, 'Limits': 0.0}\n"
     ]
    }
   ],
   "source": [
    "th_final, err_per_iter, k = planner(th_init, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_factor.plot_trajectory(th_final.unsqueeze(0))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dgpmp2-nd",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}