notebooks/02_planning_to_tcp_poses.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Planning to TCP Poses 📐\n",
    "\n",
    "In the previous notebook we saw how to plan paths between joint configurations.\n",
    "In this notebook we will explore how planning to TCP poses opens doors to many interesting possibilities, such as:\n",
    "- 💥 **Filtering** e.g. pregrasp configurations to ensure moving the the grasp pose does not collide with the environment\n",
    "- 🥇 **Selecting** optimal paths, e.g. the shortest, smoothest, fastest after time parametrization or most clearance\n",
    "- 📊 **Ranking** goal joint configurations and planning to them iteratively, selecting the first path that is found. This ranking could be based on the distance of the goal joint configuration to the a desirable joint configuration, e.g the start joint configuration.\n",
    "\n",
    "**The 4-Step Process**\n",
    "\n",
    "Let's break down how we approach TCP pose planning:\n",
    "1. 🧮 **IK Solutions:** Calculate candidate goal joint configurations from the goal TCP pose.\n",
    "2. 🤔 **Rank & Prioritize (Optional):** Rank those candidates\n",
    "3. 🗺️ **Plan Paths:** Try planning paths to the candidates.\n",
    "4. 🏆 **Select the Best (Optional):** Choose the optimal path based on your criteria.\n",
    "\n",
    "\n",
    "\n",
    "<!-- **Why Planning to TCP Poses Matters**\n",
    "\n",
    "- ✨ **Flexibility:** A single TCP pose often corresponds to multiple possible joint configurations for a 6-DOF robot.\n",
    "- 💡 **Strategic choices:** This opens doors! We can filter pre-grasp configurations for collision-free grasps, optimize paths based on smoothness or speed, and more.\n",
    "\n",
    "\n",
    "However, in this notebook we will show how planning to TCP poses opens interesting opportunities.\n",
    "\n",
    "A key fact to know, is that for a given TCP pose in the workspace of a 6-DOF robot, there will be several joint configurations that can achieve that pose. \n",
    "For some robots, such as the universal robots we use at our lab, you can analytically calculate the inverse kinematics to find all the joint configurations that can achieve a given TCP pose.\n",
    "\n",
    "This finite set of possible goal configurations provide an interesting opportunity.\n",
    "For example, if the TCP pose we plan to is a pregrasp pose from which we want to move linearly to a grasp pose, we can filter the pregrasp joint configurations to only select the from which we can move to the grasp pose without colliding with the environment.\n",
    "\n",
    "Another possiblity with multiple joint configurations is that we can plan paths to each of them and then select the best path according to some criteria. For example, we might want to select the shortest path, or the smoothest path, or the path that can be executed the fastest after time parametrization.\n",
    "\n",
    "However, if any path is good enough, and we care mostly about planning time, we can rank the goal joint configurations, and plan to them iteratively, and select the first path that is found. This ranking could for example be based on the distance goal joint configuration to the a desirable joint configuration .e.g the start joint configuration.\n",
    "\n",
    "\n",
    "The way we achieve planning to TCP poses is by using a four-step process:\n",
    "1. Calculate the inverse kinematics of the goal TCP pose to find the candidate goal joint configurations.\n",
    "2. Optional: rank the goal joint configurations\n",
    "3. Attempt to plan a path to the goal joint configuration(s)\n",
    "4. Optional: select the best path -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setting up the scene & collision checker 🏗️️\n",
    "\n",
    "For more details about this, see the [`airo-drake`](https://github.com/airo-ugent/airo-drake) notebooks.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.planning import RobotDiagramBuilder, SceneGraphCollisionChecker\n",
    "from airo_drake import SingleArmScene, add_floor, add_manipulator, add_meshcat, finish_build\n",
    "\n",
    "robot_diagram_builder = RobotDiagramBuilder()\n",
    "\n",
    "meshcat = add_meshcat(robot_diagram_builder)\n",
    "arm_index, gripper_index = add_manipulator(robot_diagram_builder, \"ur3e\", \"robotiq_2f_85\", static_gripper=True)\n",
    "add_floor(robot_diagram_builder)\n",
    "robot_diagram, context = finish_build(robot_diagram_builder)\n",
    "\n",
    "scene = SingleArmScene(robot_diagram, arm_index, gripper_index, meshcat)\n",
    "\n",
    "\n",
    "collision_checker = SceneGraphCollisionChecker(\n",
    "    model=scene.robot_diagram,\n",
    "    robot_model_instances=[scene.arm_index, scene.gripper_index],\n",
    "    edge_step_size=0.125,  # Arbitrary value: we don't use the CheckEdgeCollisionFree\n",
    "    env_collision_padding=0.005,\n",
    "    self_collision_padding=0.005,\n",
    ")\n",
    "\n",
    "collision_checker.CheckConfigCollisionFree(np.zeros(6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_joints = np.deg2rad([0, -90, 90, -90, -90, 0])\n",
    "\n",
    "print(collision_checker.CheckConfigCollisionFree(start_joints))\n",
    "\n",
    "plant = scene.robot_diagram.plant()\n",
    "plant_context = plant.GetMyContextFromRoot(context)\n",
    "plant.SetPositions(plant_context, scene.arm_index, start_joints)\n",
    "scene.robot_diagram.ForcedPublish(context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Creating the motion planner 🧭\n",
    "\n",
    "To enable planning to TCP poses, we need to provide an additional argument to the `SingleArmOmplPlanner` constructor, namely an `inverse_kinematics_fn` that returns a list of joint configurations for a given TCP pose.\n",
    "\n",
    "### 2.1 `inverse_kinematics_fn` 🧮\n",
    "In this notebook we will be using the analytical inverse kinematics from the [`ur-analytic-ik`](https://github.com/Victorlouisdg/ur-analytic-ik) Python package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ur_analytic_ik import ur3e\n",
    "from airo_typing import HomogeneousMatrixType, JointConfigurationType\n",
    "from airo_drake import visualize_frame\n",
    "\n",
    "tcp_transform = np.identity(4)\n",
    "tcp_transform[2, 3] = 0.175  # 175 mm in z\n",
    "\n",
    "\n",
    "def inverse_kinematics_fn(tcp_pose: HomogeneousMatrixType) -> list[JointConfigurationType]:\n",
    "    solutions = ur3e.inverse_kinematics_with_tcp(tcp_pose, tcp_transform)\n",
    "    solutions = [solution.squeeze() for solution in solutions]\n",
    "    return solutions\n",
    "\n",
    "\n",
    "goal_pose = np.identity(4)\n",
    "\n",
    "X = np.array([-1.0, 0.0, 0.0])\n",
    "Y = np.array([0.0, 1.0, 0.0])\n",
    "Z = np.array([0.0, 0.0, -1.0])\n",
    "top_down_orientation = np.column_stack([X, Y, Z])\n",
    "goal_pose[:3, :3] = top_down_orientation\n",
    "goal_pose[:3, 3] = [0.17, 0.0, 0.0]\n",
    "\n",
    "visualize_frame(scene.meshcat, \"goal_pose\", goal_pose)\n",
    "\n",
    "solutions = inverse_kinematics_fn(goal_pose)\n",
    "\n",
    "with np.printoptions(precision=3, suppress=True):\n",
    "    for solution in solutions:\n",
    "        print(solution)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_drake import animate_joint_configurations\n",
    "\n",
    "animate_joint_configurations(scene.meshcat, scene.robot_diagram, scene.arm_index, solutions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 `SingleArmOmplPlanner` 🧭"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_planner import SingleArmOmplPlanner\n",
    "\n",
    "\n",
    "planner = SingleArmOmplPlanner(collision_checker.CheckConfigCollisionFree, inverse_kinematics_fn=inverse_kinematics_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Planning to TCP poses 🗺\n",
    "\n",
    "### 3.1 Choosing the shortest path 📏\n",
    "\n",
    "The default behavior of `SingleArmOmplPlanner`  is planning to all valid goal joint configurations and selecting the shortest path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = planner.plan_to_tcp_pose(start_joints, goal_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_drake import calculate_joint_path_length\n",
    "\n",
    "path_lengths = [calculate_joint_path_length(p) for p in planner._all_paths]\n",
    "\n",
    "print(\"Length of all paths the planner found:\")\n",
    "for i, path_length in enumerate(path_lengths):\n",
    "    print(f\"{i}: {path_length:.2f}\")\n",
    "\n",
    "print(f\"\\nLength of the chosen path: {calculate_joint_path_length(path):.2f}\")\n",
    "\n",
    "print(f\"Amount of states in the path: {len(path)}\")\n",
    "animate_joint_configurations(scene.meshcat, scene.robot_diagram, scene.arm_index, path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Ranking goal configuration desirability 🔀\n",
    "\n",
    "It's common to prefer that the robot arm stays close specific desirable configurations, such as:\n",
    "- the start joint configuration 🏁\n",
    "- a home configuration 🏡\n",
    "\n",
    "We let the planner know these goal configuration preferences by setting the:\n",
    "- `rank_goal_configurations_fn` attribute of the `SingleArmOmplPlanner`.\n",
    "\n",
    "In `airo-planner` we provide a `rank_by_distance_to_desirable_configurations()` function that we can use for this purpose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_planner import rank_by_distance_to_desirable_configurations\n",
    "\n",
    "ranked_solutions = rank_by_distance_to_desirable_configurations(solutions, [start_joints])\n",
    "\n",
    "np.printoptions(precision=3, suppress=True)\n",
    "\n",
    "print(\"Distance from start joints to the ranked IK solutions\")\n",
    "for i, solution in enumerate(ranked_solutions):\n",
    "    distance = np.linalg.norm(solution - start_joints)\n",
    "    print(f\"{distance:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "rank_fn = partial(rank_by_distance_to_desirable_configurations, desirable_configurations=[start_joints])\n",
    "\n",
    "planner.rank_goal_configurations_fn = rank_fn\n",
    "\n",
    "path = planner.plan_to_tcp_pose(start_joints, goal_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"\\nLength of the returned path: {calculate_joint_path_length(path):.2f}\")\n",
    "\n",
    "print(f\"Amount of states in the path: {len(path)}\")\n",
    "\n",
    "animate_joint_configurations(scene.meshcat, scene.robot_diagram, scene.arm_index, path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scene.meshcat.Delete(\"goal_pose\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.3 Filtering goal configurations 🧹\n",
    "\n",
    "**Filtering** is possibliy the most powerful feature of planning to TCP poses.\n",
    "By filtering the goal joint configurations, we can ensure that the robot goes to configurations that enable downstream tasks. \n",
    "For example, you want to make sure that end configuration of the path to the pregrasp pose is such that you can move linearly to the grasp pose without colliding with the environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_drake import visualize_frame\n",
    "\n",
    "\n",
    "# grasp_location = np.array([0.2, 0.5, 0.5]) # This moves the robot through a singularlity for some start configurations\n",
    "grasp_location = np.array([0.2, 0.3, 0.2])\n",
    "\n",
    "gripper_forward_direction = np.array([1, 0, 0])\n",
    "\n",
    "Z = gripper_forward_direction / np.linalg.norm(gripper_forward_direction)\n",
    "Y = np.array([0, 0, -1])  # 0, 0, 1 is also an option\n",
    "X = np.cross(Y, Z)\n",
    "\n",
    "grasp_orientation = np.column_stack([X, Y, Z])\n",
    "grasp_pose = np.identity(4)\n",
    "grasp_pose[0:3, 0:3] = grasp_orientation\n",
    "grasp_pose[0:3, 3] = grasp_location\n",
    "\n",
    "visualize_frame(scene.meshcat, \"grasp_pose\", grasp_pose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_typing import InverseKinematicsFunctionType, JointConfigurationCheckerType, JointPathType\n",
    "from loguru import logger\n",
    "from pydrake.trajectories import PiecewisePose\n",
    "from pydrake.math import RigidTransform\n",
    "from airo_drake import discretize_drake_pose_trajectory\n",
    "from airo_drake import calculate_valid_joint_paths\n",
    "from functools import partial\n",
    "from airo_planner import filter_with_distance_to_configurations, PlannerError\n",
    "\n",
    "\n",
    "def plan_pregrasp(\n",
    "    grasp_pose: HomogeneousMatrixType,\n",
    "    start_configuration: JointConfigurationType,\n",
    "    inverse_kinematics_fn: InverseKinematicsFunctionType,\n",
    "    is_state_valid_fn_pregrasp: JointConfigurationCheckerType,\n",
    "    is_state_valid_fn_grasp: JointConfigurationCheckerType,\n",
    ") -> tuple[JointPathType, JointPathType]:\n",
    "\n",
    "    # We add 1.0 so at least one pregrasp distance fails:\n",
    "    pregrasp_distances_to_try = [1.0, 0.05, 0.1, 0.15, 0.2, 0.25]\n",
    "\n",
    "    planner = SingleArmOmplPlanner(is_state_valid_fn_pregrasp, inverse_kinematics_fn=inverse_kinematics_fn)\n",
    "\n",
    "    for distance in pregrasp_distances_to_try:\n",
    "        logger.info(f\"Planning to pregrasp pose at distance {distance}.\")\n",
    "        # 1. Compute pregrasp pose\n",
    "        pregrasp_pose = grasp_pose.copy()\n",
    "        pregrasp_pose[0:3, 3] -= distance * pregrasp_pose[0:3, 2]\n",
    "\n",
    "        # 2. Compute grasp TCP path\n",
    "        rigid_transforms = [RigidTransform(pose) for pose in [pregrasp_pose, grasp_pose]]\n",
    "        times = np.linspace(0, 1, len(rigid_transforms))\n",
    "        pose_trajectory = PiecewisePose.MakeLinear(times=times, poses=rigid_transforms)\n",
    "        tcp_path = discretize_drake_pose_trajectory(pose_trajectory).poses\n",
    "\n",
    "        # 3 Compute valid grasp joint paths\n",
    "        joint_paths = calculate_valid_joint_paths(tcp_path, inverse_kinematics_fn, is_state_valid_fn_grasp)\n",
    "        grasp_path_starts = [path[0] for path in joint_paths]\n",
    "\n",
    "        # 4 plan to pregrasp tcp poses, filtering on the valid grasp joint paths\n",
    "        filter_fn = partial(filter_with_distance_to_configurations, joint_configurations_close=grasp_path_starts)\n",
    "\n",
    "        planner.filter_goal_configurations_fn = filter_fn\n",
    "\n",
    "        try:\n",
    "            pregrasp_path = planner.plan_to_tcp_pose(start_configuration, pregrasp_pose)\n",
    "        except PlannerError:\n",
    "            logger.info(f\"Failed to plan to pregrasp pose at distance {distance}, continuing to next distance.\")\n",
    "            continue\n",
    "\n",
    "        # Find the grasp path of which the start is closest to the pregrasp path end (=pregrasp end joints)\n",
    "        pregrasp_end_joints = pregrasp_path[-1]\n",
    "        distances = [np.linalg.norm(start - pregrasp_end_joints) for start in grasp_path_starts]\n",
    "        index_of_closest_start = np.argmin(distances)\n",
    "\n",
    "        assert np.isclose(distances[index_of_closest_start], 0, atol=0.01)  # sanity check\n",
    "\n",
    "        grasp_path = joint_paths[index_of_closest_start]\n",
    "\n",
    "        return pregrasp_path, grasp_path\n",
    "\n",
    "    logger.warning(\"Exhausted all pregrasp poses to try\")\n",
    "\n",
    "\n",
    "pregrasp_path, grasp_path = plan_pregrasp(\n",
    "    grasp_pose,\n",
    "    start_joints,\n",
    "    inverse_kinematics_fn,\n",
    "    collision_checker.CheckConfigCollisionFree,\n",
    "    collision_checker.CheckConfigCollisionFree,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_drake import time_parametrize_toppra\n",
    "\n",
    "pregrasp_trajectory = time_parametrize_toppra(\n",
    "    scene.robot_diagram.plant(), pregrasp_path)\n",
    "\n",
    "\n",
    "grasp_trajectory = time_parametrize_toppra(\n",
    "    scene.robot_diagram.plant(), grasp_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_drake import shift_drake_trajectory_in_time\n",
    "\n",
    "grasp_trajectory_shifted = shift_drake_trajectory_in_time(grasp_trajectory, pregrasp_trajectory.end_time())\n",
    "\n",
    "print(pregrasp_trajectory.start_time(), pregrasp_trajectory.end_time())\n",
    "print(grasp_trajectory.start_time(), grasp_trajectory.end_time())\n",
    "print(grasp_trajectory_shifted.start_time(), grasp_trajectory_shifted.end_time())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_drake import concatenate_drake_trajectories\n",
    "from airo_drake import animate_joint_trajectory\n",
    "\n",
    "joint_trajectory = concatenate_drake_trajectories([pregrasp_trajectory, grasp_trajectory])\n",
    "\n",
    "animate_joint_trajectory(scene.meshcat, scene.robot_diagram, scene.arm_index, joint_trajectory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from airo_drake import discretize_drake_joint_trajectory\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "joint_trajectory_discretized = discretize_drake_joint_trajectory(joint_trajectory, steps=1000)\n",
    "joint_velocities_discretized = [joint_trajectory.EvalDerivative(t).squeeze() for t in joint_trajectory_discretized.times]\n",
    "joint_accelerations_discretized = [joint_trajectory.EvalDerivative(t, 2).squeeze() for t in joint_trajectory_discretized.times]\n",
    "\n",
    "\n",
    "plt.figure(figsize=(20, 6))\n",
    "plt.subplot(1, 3, 1)\n",
    "plt.title(\"Joint positions\")\n",
    "for row in joint_trajectory_discretized.path.positions.T:\n",
    "    plt.plot(joint_trajectory_discretized.times, row)\n",
    "plt.ylim(-2 * np.pi, 2 * np.pi)\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.title(\"Joint velocities\")\n",
    "for row in np.array(joint_velocities_discretized).T:\n",
    "    plt.plot(row)\n",
    "plt.ylim(-4.0, 4.0)\n",
    "plt.subplot(1, 3, 3)\n",
    "plt.title(\"Joint accelerations\")\n",
    "for row in np.array(joint_accelerations_discretized).T:\n",
    "    plt.plot(row)\n",
    "plt.ylim(-10.0, 10.0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cloth-competition",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}