{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "48bcf763",
   "metadata": {},
   "source": [
    "# Extension - RRT: basic finding example\n",
    "* Path finding is using sample-based (RRT) methods rather than constrained optimization to find a collision free path\n",
    "* Path finding is much easier if a final robot pose is given. We here use IK (formulated as KOMO problem) to first compute a final configuration 'qT'. (Path optimization does this jointly with optimizing the path.)\n",
    "* Then we can pass the current state and qT to a bi-directional RRT to find a collision free path.\n",
    "* Note that BotOp (the interface to sim/real) is opened only after we computed the motion. We simply pass the motion to be played by the sim/real."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e07bf36",
   "metadata": {},
   "outputs": [],
   "source": [
    "import robotic as ry\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "445a6a71",
   "metadata": {},
   "source": [
    "first a minimalistic example for testing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "584507cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "C = ry.Config()\n",
    "C.addFrame(\"base\") .setPosition([0,0,.5])\n",
    "\n",
    "C.addFrame(\"ego\", \"base\") \\\n",
    "    .setJoint(ry.JT.transXYPhi, [-1.,-1.,-3.,1.,1.,3.]) \\\n",
    "    .setRelativePosition([.2, .0, .0]) \\\n",
    "    .setShape(ry.ST.ssBox, size=[.05, .3, .05, .01]) \\\n",
    "    .setColor([0, 1., 1.]) \\\n",
    "    .setContact(1)\n",
    "\n",
    "C.addFrame(\"obstacle\") \\\n",
    "    .setPosition([.0, .0, .5]) \\\n",
    "    .setShape(ry.ST.ssBox, size=[.05, .3, .05, .01]) \\\n",
    "    .setColor([1, .5, 0]) \\\n",
    "    .setContact(1)\n",
    "\n",
    "C.view(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4d8417b",
   "metadata": {},
   "outputs": [],
   "source": [
    "q0 = [-.2, 0, 0]\n",
    "qT = [+.2, 0, 0]\n",
    "\n",
    "ry.params_clear()\n",
    "ry.params_add({'rrt/stepsize':.1, 'rrt/verbose': 3}) #verbose=3 makes it very slow, and displays result, and verbose=4 waits keypress..\n",
    "\n",
    "rrt = ry.RRT_PathFinder()\n",
    "rrt.setProblem(C)\n",
    "rrt.setStartGoal([q0], [qT])\n",
    "ret = rrt.solve()\n",
    "print(ret)\n",
    "path = ret.x\n",
    "\n",
    "ry.params_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edc4834c",
   "metadata": {},
   "outputs": [],
   "source": [
    "del rrt\n",
    "\n",
    "print('path length:', path.shape)\n",
    "# display the path\n",
    "for t in path:\n",
    "    C.setJointState(t)\n",
    "    C.view()\n",
    "    time.sleep(1./path.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96338c60",
   "metadata": {},
   "outputs": [],
   "source": [
    "#this prints all parameters used by the rrt:\n",
    "ry.params_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "059a8ee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "C = ry.Config()\n",
    "C.addFile(ry.raiPath('../rai-robotModels/scenarios/pandasTable.g'))\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "582b68ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "C.addFrame('boxR','table') \\\n",
    "    .setRelativePosition([.15,0,.1]) \\\n",
    "    .setShape(ry.ST.sphere, size=[.03]) \\\n",
    "    .setColor([1,1,0])\n",
    "C.addFrame('boxL','table') \\\n",
    "    .setRelativePosition([-.15,0,.1]) \\\n",
    "    .setShape(ry.ST.sphere, size=[.03]) \\\n",
    "    .setColor([1,.5,0])\n",
    "C.view()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eab44514",
   "metadata": {},
   "outputs": [],
   "source": [
    "# store the start configuration\n",
    "q0 = C.getJointState()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bccb7b55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute a goal configuration\n",
    "komo = ry.KOMO()\n",
    "komo.setConfig(C, True)\n",
    "komo.setTiming(1., 1, 5., 0)\n",
    "komo.addControlObjective([], 0, 1e-0)\n",
    "komo.addObjective([], ry.FS.accumulatedCollisions, [], ry.OT.eq);\n",
    "komo.addObjective([], ry.FS.jointLimits, [], ry.OT.ineq);\n",
    "komo.addObjective([], ry.FS.positionDiff, ['r_gripper', 'boxL'], ry.OT.eq, [1e1]);\n",
    "komo.addObjective([], ry.FS.positionDiff, ['l_gripper', 'boxR'], ry.OT.eq, [1e1]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "178e3d42",
   "metadata": {},
   "outputs": [],
   "source": [
    "ret = ry.NLP_Solver() \\\n",
    "    .setProblem(komo.nlp()) \\\n",
    "    .setOptions( stopTolerance=1e-2, verbose=4 ) \\\n",
    "    .solve()\n",
    "print(ret)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b727e37c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# that's the goal configuration\n",
    "qT = komo.getPath()[0]\n",
    "C.setJointState(qT)\n",
    "C.view(False, \"IK solution\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b20fc581",
   "metadata": {},
   "outputs": [],
   "source": [
    "#define a path finding problem\n",
    "rrt = ry.RRT_PathFinder()\n",
    "rrt.setProblem(C)\n",
    "rrt.setStartGoal([q0], [qT])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0cde20c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ret = rrt.solve()\n",
    "print(ret)\n",
    "path = ret.x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a615cc68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# display the path\n",
    "for t in range(0, path.shape[0]-1):\n",
    "    C.setJointState(path[t])\n",
    "    C.view()\n",
    "    time.sleep(.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51d2e8e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run the path with botop\n",
    "C.setJointState(q0)\n",
    "ry.params_add({'botsim/verbose': 1., 'physx/motorKp': 10000., 'physx/motorKd': 1000.})\n",
    "bot = ry.BotOp(C, False)\n",
    "bot.home(C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d8e62c",
   "metadata": {},
   "outputs": [],
   "source": [
    "bot.moveAutoTimed(path, 1., 1.)\n",
    "while bot.getTimeToEnd()>0:\n",
    "    bot.sync(C, .1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c02685e",
   "metadata": {},
   "outputs": [],
   "source": [
    "del bot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "929345e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "del rrt\n",
    "del C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cd2cec7",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}