From a9db5895f8460b331bf48d594b6b49a513813baf Mon Sep 17 00:00:00 2001
From: victorlouisdg <victorlouisdg@gmail.com>
Date: Tue, 19 Mar 2024 07:29:19 +0100
Subject: [PATCH] attempt at deterministic planning

---
 airo_planner/ompl/single_arm_ompl_planner.py |  41 ++-
 airo_planner/ompl/utilities.py               |  10 +
 airo_planner/selection/goal_selection.py     |  28 ++
 airo_planner/selection/path_selection.py     |  17 ++
 notebooks/02_planning_to_tcp_poses.ipynb     | 290 ++++++++++++++++++-
 5 files changed, 371 insertions(+), 15 deletions(-)
 create mode 100644 airo_planner/selection/goal_selection.py
 create mode 100644 airo_planner/selection/path_selection.py

diff --git a/airo_planner/ompl/single_arm_ompl_planner.py b/airo_planner/ompl/single_arm_ompl_planner.py
index 7b77d66..82b1cc5 100644
--- a/airo_planner/ompl/single_arm_ompl_planner.py
+++ b/airo_planner/ompl/single_arm_ompl_planner.py
@@ -1,3 +1,5 @@
+import sys
+import time
 from typing import Callable, List
 
 import numpy as np
@@ -13,6 +15,7 @@
 
 from airo_planner import JointBoundsType, SingleArmPlanner, path_from_ompl, state_to_ompl
 from airo_planner.ompl.utilities import create_simple_setup
+from airo_planner.selection.path_selection import choose_shortest_path
 
 # TODO move this to airo_typing
 JointConfigurationsModifierType = Callable[[List[JointConfigurationType]], List[JointConfigurationType]]
@@ -35,11 +38,10 @@ def __init__(
         self,
         is_state_valid_fn: JointConfigurationCheckerType,
         joint_bounds: JointBoundsType | None = None,
-        max_planning_time: float = 5.0,
         inverse_kinematics_fn: InverseKinematicsFunctionType | None = None,
         filter_goal_configurations_fn: JointConfigurationsModifierType | None = None,
         rank_goal_configurations_fn: JointConfigurationsModifierType | None = None,
-        choose_path_fn: JointPathChooserType | None = None,
+        choose_path_fn: JointPathChooserType | None = choose_shortest_path,
         degrees_of_freedom: int = 6,
     ):
         """Instiatiate a single-arm motion planner that uses OMPL. This creates
@@ -53,7 +55,6 @@ def __init__(
         Args:
             is_state_valid_fn: A function that checks if a given joint configuration is valid.
             inverse_kinematics_fn: A function that computes the inverse kinematics of a given TCP pose.
-            max_planning_time: The maximum time allowed for planning.
         """
         self.is_state_valid_fn = is_state_valid_fn
 
@@ -64,7 +65,6 @@ def __init__(
         self.choose_path_fn = choose_path_fn
 
         # Planning parameters
-        self.max_planning_time = max_planning_time
         self._degrees_of_freedom = degrees_of_freedom
         if joint_bounds is None:
             self.joint_bounds = np.full(degrees_of_freedom, -2 * np.pi), np.full(degrees_of_freedom, 2 * np.pi)
@@ -95,7 +95,7 @@ def plan_to_joint_configuration(
         simple_setup = self._simple_setup
 
         # Run planning algorithm
-        simple_setup.solve(self.max_planning_time)
+        simple_setup.solve()
 
         if not simple_setup.haveExactSolutionPath():
             return None
@@ -126,6 +126,7 @@ def plan_to_tcp_pose(  # noqa: C901
             return None
 
         ik_solutions = self.inverse_kinematics_fn(tcp_pose_in_base)
+        ik_solutions = [solution.squeeze() for solution in ik_solutions]
         self._ik_solutions = ik_solutions  # Save for debugging
 
         if ik_solutions is None or len(ik_solutions) == 0:
@@ -152,21 +153,47 @@ def plan_to_tcp_pose(  # noqa: C901
         else:
             logger.info(f"Found {len(ik_solutions_valid)}/{len(ik_solutions_within_bounds)} valid solutions.")
 
+        # TODO filter the IK solutions -> warn when nothing is left
+
+        # TODO rank the IK solutions
+        terminate_on_first_success = self.rank_goal_configurations_fn is not None
+
+        logger.info(f"Running OMPL towards {len(ik_solutions_valid)} IK solutions.")
+        start_time = time.time()
         # Try solving to each IK solution in joint space.
         paths = []
-        for ik_solution in ik_solutions_valid:
+        for i, ik_solution in enumerate(ik_solutions_valid):
             path = self.plan_to_joint_configuration(start_configuration, ik_solution)
+
+            # My attempt to get the OMPL Info: messages to show up in the correct order between the loguru logs
+            # This however does not seems to work, so I don't know where those prints are buffered
+            sys.stdout.flush()
+
             if path is not None:
+                logger.info(f"Successfully found path to the valid IK solution with index {i}.")
+                if terminate_on_first_success is True:
+                    end_time = time.time()
+                    logger.info(f"Terminating on first success (planning time: {end_time - start_time:.2f} s).")
+                    return path
                 paths.append(path)
 
+        end_time = time.time()
+        logger.info(f"Found {len(paths)} paths towards IK solutions, (planning time: {end_time - start_time:.2f} s).")
+
         self._all_paths = paths  # Save for debugging
 
         if len(paths) == 0:
             logger.info("No paths founds towards any IK solutions, returning None.")
             return None
 
-        solution_path = paths[0]
+        solution_path = self.choose_path_fn(paths)
         return solution_path
+
+        # TODO: exhaustive mode vs iterative mode (~= greedy mode)
+        # depending on the helper functions provide to the planner, it will operate in different modes
+        # by default, it operates in an exhaustive mode, meaning that it will treat all IK solutions the same, and plan a path to them all
+        # then it will return the shortest path
+
         # path_distances = [np.linalg.norm(path[-1] - start_configuration) for path in paths]
 
         # path_desirablities = None
diff --git a/airo_planner/ompl/utilities.py b/airo_planner/ompl/utilities.py
index a64a5b8..44c37ff 100644
--- a/airo_planner/ompl/utilities.py
+++ b/airo_planner/ompl/utilities.py
@@ -46,6 +46,13 @@ def bounds_to_ompl(joint_bounds: JointBoundsType) -> ob.RealVectorBounds:
     return bounds
 
 
+def allocDeterministicStateSampler(space: ob.StateSpace) -> ob.StateSampler:
+    from loguru import logger
+
+    logger.info("Allocating RealVectorDeterministicStateSampler")
+    return ob.RealVectorDeterministicStateSampler(space)
+
+
 def create_simple_setup(
     is_state_valid_fn: JointConfigurationCheckerType, joint_bounds: JointBoundsType
 ) -> og.SimpleSetup:
@@ -65,6 +72,9 @@ def create_simple_setup(
 
     is_state_valid_ompl = function_to_ompl(is_state_valid_fn, degrees_of_freedom)
 
+    space.setStateSamplerAllocator(ob.StateSamplerAllocator(allocDeterministicStateSampler))
+    # space.allocStateSampler()
+
     # Configure the SimpleSetup object
     simple_setup = og.SimpleSetup(space)
     simple_setup.setStateValidityChecker(ob.StateValidityCheckerFn(is_state_valid_ompl))
diff --git a/airo_planner/selection/goal_selection.py b/airo_planner/selection/goal_selection.py
new file mode 100644
index 0000000..63ab45f
--- /dev/null
+++ b/airo_planner/selection/goal_selection.py
@@ -0,0 +1,28 @@
+import numpy as np
+from airo_typing import JointConfigurationType
+
+
+def rank_by_distance_to_desirable_configurations(
+    configurations: list[JointConfigurationType], desirable_configurations: list[JointConfigurationType]
+) -> list[JointConfigurationType]:
+    """Ranks joint configurations based on their distance to a set of desirable configurations.
+
+    Args:
+        configurations: The list of joint configurations to be ranked.
+        desirable_configurations: A list of desirable joint configurations.
+
+    Returns:
+        A list of joint configurations, sorted in ascending order of their
+        distance to the closest desirable configuration.
+    """
+
+    distances = []
+    for config in configurations:
+        distances_to_desirable = [
+            np.linalg.norm(config - desirable_config) for desirable_config in desirable_configurations
+        ]
+        min_distance = min(distances_to_desirable)
+        distances.append(min_distance)
+
+    ranked_configurations = [x for _, x in sorted(zip(distances, configurations))]
+    return ranked_configurations
diff --git a/airo_planner/selection/path_selection.py b/airo_planner/selection/path_selection.py
new file mode 100644
index 0000000..d58166f
--- /dev/null
+++ b/airo_planner/selection/path_selection.py
@@ -0,0 +1,17 @@
+import numpy as np
+from airo_drake import calculate_joint_path_length
+from airo_typing import JointPathType
+
+
+def choose_shortest_path(paths: list[JointPathType]) -> JointPathType:
+    """Selects the shortest path from a list of possible paths.
+
+    Args:
+        paths: A list of potential joint paths.
+
+    Returns:
+        The path from the input list that has the shortest length.
+    """
+    path_lengths = [calculate_joint_path_length(p) for p in paths]
+    shortest_idx = np.argmin(path_lengths)
+    return paths[shortest_idx]
diff --git a/notebooks/02_planning_to_tcp_poses.ipynb b/notebooks/02_planning_to_tcp_poses.ipynb
index b50fa61..3088a23 100644
--- a/notebooks/02_planning_to_tcp_poses.ipynb
+++ b/notebooks/02_planning_to_tcp_poses.ipynb
@@ -9,7 +9,7 @@
     "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 or fastest after time parametrization\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",
@@ -22,9 +22,6 @@
     "\n",
     "\n",
     "\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",
@@ -92,6 +89,22 @@
     "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": {},
@@ -100,6 +113,7 @@
     "\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."
    ]
   },
@@ -129,15 +143,15 @@
     "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.2, 0.0, 0.0]\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)\n"
+    "with np.printoptions(precision=3, suppress=True):\n",
+    "    for solution in solutions:\n",
+    "        print(solution)\n"
    ]
   },
   {
@@ -151,11 +165,271 @@
     "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": [
+    "animate_joint_configurations(scene.meshcat, scene.robot_diagram, scene.arm_index, path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**Double checking the path lengths:**"
+   ]
+  },
   {
    "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(\"Path lengths\")\n",
+    "for i, path_length in enumerate(path_lengths):\n",
+    "    print(f\"{i}: {path_length:.2f}\")\n",
+    "\n",
+    "print(f\"\\nLength of path that planner returned: {calculate_joint_path_length(path):.2f}\")"
+   ]
+  },
+  {
+   "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.selection.goal_selection 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(\"Distances to ranked 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": [
+    "planner.rank_goal_configurations_fn = rank_by_distance_to_desirable_configurations\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 path that planner returned: {calculate_joint_path_length(path):.2f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "valid_sampler = planner._simple_setup.getSpaceInformation().allocValidStateSampler()\n",
+    "valid_sampler"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ompl import base as ob\n",
+    "from ompl import util as ou\n",
+    "from loguru import logger\n",
+    "\n",
+    "class DeterministicValidStateSampler(ob.ValidStateSampler):\n",
+    "    def __init__(self, space_information: ob.SpaceInformation, random_seed: int = 1):\n",
+    "        logger.info(\"Creating DeterministicValidStateSampler\")\n",
+    "        super(DeterministicValidStateSampler, self).__init__(space_information)\n",
+    "        self.rng_ = ou.RNG(random_seed)\n",
+    "\n",
+    "\n",
+    "def allocDeterministicValidStateSampler(space_information: ob.SpaceInformation) -> ob.ValidStateSampler:\n",
+    "    logger.info(\"Allocating DeterministicValidStateSampler\")\n",
+    "    return DeterministicValidStateSampler(space_information)\n",
+    "\n",
+    "\n",
+    "space_information = planner._simple_setup.getSpaceInformation()\n",
+    "DeterministicValidStateSampler(space_information)\n",
+    "allocDeterministicValidStateSampler(space_information)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ob.RealVectorDeterministicStateSampler(planner._simple_setup.getStateSpace())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# def allocDeterministicStateSampler(space: ob.StateSpace) -> ob.StateSampler:\n",
+    "#     return ob.RealVectorDeterministicStateSampler(space)\n",
+    "\n",
+    "# planner._simple_setup.getStateSpace().setStateSamplerAllocator(ob.StateSamplerAllocator(allocDeterministicStateSampler))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from airo_planner.ompl.utilities import create_simple_setup\n",
+    "import ompl.geometric as og\n",
+    "\n",
+    "planner._simple_setup = create_simple_setup(planner.is_state_valid_fn, planner.joint_bounds)\n",
+    "# space_information = planner._simple_setup.getSpaceInformation()\n",
+    "# space_information.clearValidStateSamplerAllocator()\n",
+    "# space_information.setValidStateSamplerAllocator(ob.ValidStateSamplerAllocator(allocDeterministicValidStateSampler))\n",
+    "# space_information.allocValidStateSampler()\n",
+    "\n",
+    "# rrt_connect = og.RRTConnect(space_information)\n",
+    "# planner._simple_setup.setPlanner(rrt_connect)\n",
+    "\n",
+    "path = planner.plan_to_tcp_pose(start_joints, goal_pose)\n",
+    "calculate_joint_path_length(path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "planner._simple_setup = create_simple_setup(planner.is_state_valid_fn, planner.joint_bounds)\n",
+    "# space_information = planner._simple_setup.getSpaceInformation()\n",
+    "# space_information.clearValidStateSamplerAllocator()\n",
+    "# space_information.setValidStateSamplerAllocator(ob.ValidStateSamplerAllocator(allocDeterministicValidStateSampler))\n",
+    "# space_information.allocValidStateSampler()\n",
+    "\n",
+    "# rrt_connect = og.RRTConnect(space_information)\n",
+    "# planner._simple_setup.setPlanner(rrt_connect)\n",
+    "\n",
+    "path = planner.plan_to_tcp_pose(start_joints, goal_pose)\n",
+    "calculate_joint_path_length(path)"
+   ]
+  },
+  {
+   "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_planner.interfaces import SingleArmPlanner\n",
+    "from pydrake.trajectories import Trajectory\n",
+    "\n",
+    "# use this a filtering example: we want to filter out the grasp paths that collide with the environment\n",
+    "\n",
+    "def plan_to_grasp(planner: SingleArmPlanner, grasp_pose: HomogeneousMatrixType) -> Trajectory:\n",
+    "    pass\n",
+    "    # attempt planning the the grasp pose directly\n",
+    "    #   if all goals in collision try backing-off to a pregrasp pose\n",
+    "    #   time parametrize both paths and concatenate them\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Question: what will increase the joint distance the most? adding pi/4 to all joints or np.pi /2 to a single joint? "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
    "source": []
   }
  ],