{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TaskTree Tutorial\n",
    "\n",
    "In this tutorial we will walk through the capabilities of task trees in pycram.\n",
    "\n",
    "First we have to import the necessary functionality from pycram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-28T09:45:06.870087Z",
     "start_time": "2023-04-28T09:45:06.870015Z"
    }
   },
   "outputs": [],
   "source": [
    "from pycram.bullet_world import BulletWorld\n",
    "from pycram.robot_descriptions import robot_description\n",
    "import pycram.task\n",
    "from pycram.enums import Arms, ObjectType\n",
    "from pycram.designators.action_designator import *\n",
    "from pycram.designators.location_designator import *\n",
    "from pycram.process_module import simulated_robot\n",
    "from pycram.designators.object_designator import *\n",
    "from pycram.pose import Pose\n",
    "from pycram.enums import ObjectType\n",
    "import anytree\n",
    "import pycram.plan_failures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we will create a bullet world with a PR2 in a kitchen containing milk and cereal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-27T17:32:49.441717Z",
     "start_time": "2023-04-27T17:32:49.441589Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Unknown tag \"material\" in /robot[@name='plane']/link[@name='planeLink']/collision[1]\n",
      "Unknown tag \"contact\" in /robot[@name='plane']/link[@name='planeLink']\n",
      "Unknown tag \"material\" in /robot[@name='plane']/link[@name='planeLink']/collision[1]\n",
      "Unknown tag \"contact\" in /robot[@name='plane']/link[@name='planeLink']\n",
      "Unknown attribute \"type\" in /robot[@name='pr2']/link[@name='base_laser_link']\n",
      "Unknown attribute \"type\" in /robot[@name='pr2']/link[@name='wide_stereo_optical_frame']\n",
      "Unknown attribute \"type\" in /robot[@name='pr2']/link[@name='narrow_stereo_optical_frame']\n",
      "Unknown attribute \"type\" in /robot[@name='pr2']/link[@name='laser_tilt_link']\n",
      "Unknown attribute \"type\" in /robot[@name='pr2']/link[@name='base_laser_link']\n",
      "Unknown attribute \"type\" in /robot[@name='pr2']/link[@name='wide_stereo_optical_frame']\n",
      "Unknown attribute \"type\" in /robot[@name='pr2']/link[@name='narrow_stereo_optical_frame']\n",
      "Unknown attribute \"type\" in /robot[@name='pr2']/link[@name='laser_tilt_link']\n",
      "Scalar element defined multiple times: limit\n",
      "Scalar element defined multiple times: limit\n"
     ]
    }
   ],
   "source": [
    "world = BulletWorld()\n",
    "pr2 = Object(\"pr2\", ObjectType.ROBOT, \"pr2.urdf\")\n",
    "kitchen = Object(\"kitchen\", ObjectType.ENVIRONMENT, \"kitchen.urdf\")\n",
    "milk = Object(\"milk\", ObjectType.MILK, \"milk.stl\", pose=Pose([1.3, 1, 0.9]))\n",
    "cereal = Object(\"cereal\", ObjectType.BREAKFAST_CEREAL, \"breakfast_cereal.stl\", pose=Pose([1.3, 0.7, 0.95]))\n",
    "milk_desig = ObjectDesignatorDescription(names=[\"milk\"])\n",
    "cereal_desig = ObjectDesignatorDescription(names=[\"cereal\"])\n",
    "robot_desig = ObjectDesignatorDescription(names=[\"pr2\"]).resolve()\n",
    "kitchen_desig = ObjectDesignatorDescription(names=[\"kitchen\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we create a plan where the robot parks his arms, walks to the kitchen counter and picks the cereal and places it on the table. Then we execute the plan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-27T17:32:58.124468Z",
     "start_time": "2023-04-27T17:32:58.124325Z"
    }
   },
   "outputs": [],
   "source": [
    "@pycram.task.with_tree\n",
    "def plan():\n",
    "    with simulated_robot:\n",
    "        ParkArmsAction.Action(Arms.BOTH).perform()\n",
    "        MoveTorsoAction([0.3]).resolve().perform()\n",
    "        pickup_pose = CostmapLocation(target=cereal_desig.resolve(), reachable_for=robot_desig).resolve()\n",
    "        pickup_arm = pickup_pose.reachable_arms[0]\n",
    "        NavigateAction(target_locations=[pickup_pose.pose]).resolve().perform()\n",
    "        PickUpAction(object_designator_description=cereal_desig, arms=[pickup_arm], grasps=[\"front\"]).resolve().perform()\n",
    "        ParkArmsAction([Arms.BOTH]).resolve().perform()\n",
    "\n",
    "        place_island = SemanticCostmapLocation(\"kitchen_island_surface\", kitchen_desig.resolve(),\n",
    "                                           cereal_desig.resolve()).resolve()\n",
    "\n",
    "        place_stand = CostmapLocation(place_island.pose, reachable_for=robot_desig, reachable_arm=pickup_arm).resolve()\n",
    "\n",
    "        NavigateAction(target_locations=[place_stand.pose]).resolve().perform()\n",
    "\n",
    "        PlaceAction(cereal_desig, target_locations=[place_island.pose], arms=[pickup_arm]).resolve().perform()\n",
    "\n",
    "        ParkArmsAction([Arms.BOTH]).resolve().perform()\n",
    "\n",
    "        ParkArmsAction.Action(Arms.BOTH).perform()\n",
    "\n",
    "plan()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we get the task tree from its module and render it. Rendering can be done with any render method described in the anytree package. We will use ascii rendering here for ease of displaying."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-27T17:33:02.572319Z",
     "start_time": "2023-04-27T17:33:02.572253Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no_operation()\n",
      "└── plan()\n",
      "    ├── perform(ParkArmsAction, )\n",
      "    ├── perform(MoveTorsoAction, )\n",
      "    ├── perform(NavigateAction, )\n",
      "    ├── perform(PickUpAction, )\n",
      "    ├── perform(ParkArmsAction, )\n",
      "    ├── perform(NavigateAction, )\n",
      "    ├── perform(PlaceAction, )\n",
      "    ├── perform(ParkArmsAction, )\n",
      "    └── perform(ParkArmsAction, )\n"
     ]
    }
   ],
   "source": [
    "tt = pycram.task.task_tree\n",
    "print(anytree.RenderTree(tt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we see every task in the plan got recorded correctly. It is noticeable that the tree begins with a NoOperation node. This is done because several, not connected, plans that get executed after each other should still appear in the task tree. Hence, a NoOperation node is the root of any tree. If we re-execute the plan we would see them appear in the same tree even though they are not connected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-28T09:45:59.592864Z",
     "start_time": "2023-04-28T09:45:58.928492Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no_operation()\n",
      "├── plan()\n",
      "│   ├── perform(ParkArmsAction, )\n",
      "│   ├── perform(MoveTorsoAction, )\n",
      "│   ├── perform(NavigateAction, )\n",
      "│   ├── perform(PickUpAction, )\n",
      "│   ├── perform(ParkArmsAction, )\n",
      "│   ├── perform(NavigateAction, )\n",
      "│   ├── perform(PlaceAction, )\n",
      "│   ├── perform(ParkArmsAction, )\n",
      "│   └── perform(ParkArmsAction, )\n",
      "└── plan()\n",
      "    ├── perform(ParkArmsAction, )\n",
      "    ├── perform(MoveTorsoAction, )\n",
      "    ├── perform(NavigateAction, )\n",
      "    ├── perform(PickUpAction, )\n",
      "    ├── perform(ParkArmsAction, )\n",
      "    ├── perform(NavigateAction, )\n",
      "    ├── perform(PlaceAction, )\n",
      "    ├── perform(ParkArmsAction, )\n",
      "    └── perform(ParkArmsAction, )\n"
     ]
    }
   ],
   "source": [
    "world.reset_bullet_world()\n",
    "plan()\n",
    "print(anytree.RenderTree(tt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Projecting a plan in a new environment with its own task tree that only exists while the projected plan is running can be done with the ``with`` keyword. When this is done, both the bullet world and task tree are saved and new, freshly reset objects are available. At the end of a with block the old state is restored. The root for such things is then called ``simulation()``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-28T09:46:37.041421Z",
     "start_time": "2023-04-28T09:46:36.999813Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "simulation()\n",
      "no_operation()\n",
      "├── plan()\n",
      "│   ├── perform(ParkArmsAction, )\n",
      "│   ├── perform(MoveTorsoAction, )\n",
      "│   ├── perform(NavigateAction, )\n",
      "│   ├── perform(PickUpAction, )\n",
      "│   ├── perform(ParkArmsAction, )\n",
      "│   ├── perform(NavigateAction, )\n",
      "│   ├── perform(PlaceAction, )\n",
      "│   ├── perform(ParkArmsAction, )\n",
      "│   └── perform(ParkArmsAction, )\n",
      "└── plan()\n",
      "    ├── perform(ParkArmsAction, )\n",
      "    ├── perform(MoveTorsoAction, )\n",
      "    ├── perform(NavigateAction, )\n",
      "    ├── perform(PickUpAction, )\n",
      "    ├── perform(ParkArmsAction, )\n",
      "    ├── perform(NavigateAction, )\n",
      "    ├── perform(PlaceAction, )\n",
      "    ├── perform(ParkArmsAction, )\n",
      "    └── perform(ParkArmsAction, )\n"
     ]
    }
   ],
   "source": [
    "with pycram.task.SimulatedTaskTree() as stt:\n",
    "    print(anytree.RenderTree(pycram.task.task_tree))\n",
    "print(anytree.RenderTree(pycram.task.task_tree))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Task tree can be manipulated with ordinary anytree manipulation. If we for example want to discard the second plan, we would write"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-28T09:46:45.637721Z",
     "start_time": "2023-04-28T09:46:45.635535Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no_operation()\n",
      "+-- plan()\n",
      "    |-- perform(ParkArmsAction, )\n",
      "    |-- perform(MoveTorsoAction, )\n",
      "    |-- perform(NavigateAction, )\n",
      "    |-- perform(PickUpAction, )\n",
      "    |-- perform(ParkArmsAction, )\n",
      "    |-- perform(NavigateAction, )\n",
      "    |-- perform(PlaceAction, )\n",
      "    |-- perform(ParkArmsAction, )\n",
      "    +-- perform(ParkArmsAction, )\n"
     ]
    }
   ],
   "source": [
    "tt.root.children = (tt.root.children[0],)\n",
    "print(anytree.RenderTree(tt, style=anytree.render.AsciiStyle()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now re-execute this (modified) plan by executing the leaf in pre-ordering iteration using the anytree functionality. This will not append the re-execution to the task tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-28T09:46:48.259716Z",
     "start_time": "2023-04-28T09:46:48.233474Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no_operation()\n",
      "+-- plan()\n",
      "    |-- perform(ParkArmsAction, )\n",
      "    |-- perform(MoveTorsoAction, )\n",
      "    |-- perform(NavigateAction, )\n",
      "    |-- perform(PickUpAction, )\n",
      "    |-- perform(ParkArmsAction, )\n",
      "    |-- perform(NavigateAction, )\n",
      "    |-- perform(PlaceAction, )\n",
      "    |-- perform(ParkArmsAction, )\n",
      "    +-- perform(ParkArmsAction, )\n"
     ]
    }
   ],
   "source": [
    "world.reset_bullet_world()\n",
    "with simulated_robot:\n",
    "    [node.code.execute() for node in tt.root.leaves]\n",
    "print(anytree.RenderTree(pycram.task.task_tree, style=anytree.render.AsciiStyle()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nodes in the task tree contain additional information about the status and time of a task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-28T09:46:52.656020Z",
     "start_time": "2023-04-28T09:46:52.653083Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Code: plan() \n",
      " start_time: 2023-11-08 14:01:12.716969 \n",
      " Status: TaskStatus.SUCCEEDED \n",
      " end_time: 2023-11-08 14:01:24.675108 \n",
      " \n"
     ]
    }
   ],
   "source": [
    "print(pycram.task.task_tree.children[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The task tree can also be reset to an empty one by invoking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-28T09:46:55.385481Z",
     "start_time": "2023-04-28T09:46:55.381316Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no_operation()\n"
     ]
    }
   ],
   "source": [
    "pycram.task.reset_tree()\n",
    "print(anytree.RenderTree(pycram.task.task_tree, style=anytree.render.AsciiStyle()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If a plan fails using the PlanFailure exception, the plan will not stop. Instead, the error will be logged and saved in the task tree as a failed subtask. First let's create a simple failing plan and execute it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-28T09:46:57.445908Z",
     "start_time": "2023-04-28T09:46:57.443154Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Oopsie!\n"
     ]
    }
   ],
   "source": [
    "@pycram.task.with_tree\n",
    "def failing_plan():\n",
    "    raise pycram.plan_failures.PlanFailure(\"Oopsie!\")\n",
    "\n",
    "try:\n",
    "    failing_plan()\n",
    "except pycram.plan_failures.PlanFailure as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now investigate the nodes of the tree, and we will see that the tree indeed contains a failed task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-28T09:47:00.292660Z",
     "start_time": "2023-04-28T09:47:00.287579Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no_operation()\n",
      "+-- failing_plan()\n",
      "Code: failing_plan() \n",
      " start_time: 2023-11-08 14:02:00.257042 \n",
      " Status: TaskStatus.FAILED \n",
      " end_time: 2023-11-08 14:02:00.257507 \n",
      " \n"
     ]
    }
   ],
   "source": [
    "print(anytree.RenderTree(pycram.task.task_tree, style=anytree.render.AsciiStyle()))\n",
    "print(pycram.task.task_tree.children[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-04-28T09:47:03.581366Z",
     "start_time": "2023-04-28T09:47:03.447150Z"
    }
   },
   "outputs": [],
   "source": [
    "world.exit()"
   ]
  }
 ],
 "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": 1
}