Collision filtering (#4)

* Fix #3 * Remove accidentally added notebook checkpoint * Collision checking and filtering with SceneGraphCollisionChecker * Add collision filtering and checking functions to airo-drake and call these from the notebook * Variable naming * More Pythonic version via list comprehension * filtered is now also allowed to be a positional argument and is now True by default
airo-ugent · Apr 22, 2024 · 37b9c61 · 37b9c61
1 parent c3fd5ab
commit 37b9c61
Show file tree

Hide file tree

Showing 4 changed files with 187 additions and 5 deletions.
diff --git a/airo_drake/__init__.py b/airo_drake/__init__.py
@@ -38,6 +38,10 @@
     animate_joint_configurations,
     animate_joint_trajectory,
 )
+from airo_drake.collision.collision_checking import (
+    filter_collisions_between_all_body_pairs,
+    list_collisions_between_bodies,
+)
 
 
 __all__ = [
@@ -69,4 +73,6 @@
     "calculate_path_ik_solutions",
     "calculate_joint_paths",
     "calculate_valid_joint_paths",
+    "filter_collisions_between_all_body_pairs",
+    "list_collisions_between_bodies",
 ]
diff --git a/airo_drake/collision/__init__.py b/airo_drake/collision/__init__.py
diff --git a/airo_drake/collision/collision_checking.py b/airo_drake/collision/collision_checking.py
@@ -0,0 +1,52 @@
+import itertools
+from typing import List, Tuple
+
+from airo_typing import JointConfigurationType
+from pydrake.all import BodyIndex
+from pydrake.planning import RobotCollisionType, SceneGraphCollisionChecker
+
+
+def filter_collisions_between_all_body_pairs(
+    collision_checker: SceneGraphCollisionChecker,
+    body_indices_one: List[BodyIndex],
+    body_indices_two: List[BodyIndex],
+    filtered: bool = True,
+) -> None:
+    """Enable or disable collision filtering between all pairs of the bodies listed in body_indices_one and body_indices_two.
+    When collision filtering is enabled for two bodies, they are allowed to collide.
+
+    Args:
+        collision_checker: The collision checker instance to alter.
+        body_indices_one: A list of body indices.
+        body_indices_two: A list of body indices.
+        filtered: Whether to filter collisions between these bodies."""
+
+    body_combinations = itertools.product(body_indices_one, body_indices_two)
+    for body_index_1, body_index_2 in body_combinations:
+        collision_checker.SetCollisionFilteredBetween(body_index_1, body_index_2, filtered)
+
+
+def list_collisions_between_bodies(
+    collision_checker: SceneGraphCollisionChecker, q: JointConfigurationType
+) -> List[Tuple[BodyIndex, BodyIndex, bool]]:
+    """List all collisions between bodies for a joint configuration.
+    The return type of this function is a list that contains all colliding bodies (if any) and a boolean value that indicates whether the collision is a self-collision.
+    When the boolean value is false, the collision is an environment collision.
+
+    Args:
+        collision_checker: The collision checker instance to use.
+        q: The joint configuration of the robot to check.
+
+    Returns:
+        A list of tuples of colliding body indices and a boolean value that is True when the collision is a self collision."""
+    robot_clearance = collision_checker.CalcRobotClearance(q, 1.0)
+    indices_one = robot_clearance.robot_indices()
+    indices_two = robot_clearance.other_indices()
+    collision_types = robot_clearance.collision_types()
+    distances = robot_clearance.distances()
+
+    return [
+        (tup[0], tup[1], tup[2] != RobotCollisionType.kEnvironmentCollision)
+        for tup in zip(indices_one, indices_two, collision_types, distances)
+        if tup[3] <= 0.0
+    ]
diff --git a/notebooks/03_collision_checking.ipynb b/notebooks/03_collision_checking.ipynb
@@ -35,14 +35,15 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "from pydrake.planning import RobotDiagramBuilder\n",
-    "from airo_drake import SingleArmScene, add_floor, add_manipulator, add_meshcat, finish_build\n",
+    "from pydrake.planning import RobotDiagramBuilder, RobotCollisionType\n",
+    "from airo_drake import SingleArmScene, add_floor, add_manipulator, add_meshcat, finish_build, list_collisions_between_bodies, filter_collisions_between_all_body_pairs\n",
+    "from airo_typing import JointConfigurationType\n",
     "\n",
     "robot_diagram_builder = RobotDiagramBuilder()\n",
     "\n",
     "meshcat = add_meshcat(robot_diagram_builder)\n",
     "arm_index, gripper_index = add_manipulator(robot_diagram_builder, \"ur5e\", \"robotiq_2f_85\", static_gripper=True)\n",
-    "add_floor(robot_diagram_builder)\n",
+    "floor_index = 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",
@@ -160,6 +161,30 @@
     "scene.robot_diagram.ForcedPublish(context)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can also check which bodies are colliding."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def print_robot_collisions(q: JointConfigurationType):\n",
+    "    colliding_bodies = list_collisions_between_bodies(collision_checker, q)\n",
+    "\n",
+    "    for colliding_body_1, colliding_body_2, is_self_collision in colliding_bodies:\n",
+    "            bodyA_name = plant.get_body(colliding_body_1).scoped_name()\n",
+    "            bodyB_name = plant.get_body(colliding_body_2).scoped_name()\n",
+    "            print(f\"Body {bodyA_name} is colliding with body {bodyB_name} (is self collision: {is_self_collision})\")\n",
+    "\n",
+    "print_robot_collisions(joints_table_collision)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -707,6 +732,105 @@
     "animate_joint_trajectory(scene.meshcat, scene.robot_diagram, scene.arm_index, fastest_trajectory, context)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Allowing collisions between bodies 🏳️\n",
+    "\n",
+    "It is possible to allow collisions between certain bodies. Just for this example, let's allow collisions between the ur5e and the floor."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "arm_body_indices = plant.GetBodyIndices(arm_index)\n",
+    "floor_body_indices = plant.GetBodyIndices(floor_index)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's go to a configuration where only the arm and the floor collide."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "joints_table_collision = np.deg2rad([0, 45, 0, 0, 0, 0])\n",
+    "print_robot_collisions(joints_table_collision)\n",
+    "collision_checker.CheckConfigCollisionFree(joints_table_collision)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plant.SetPositions(plant_context, scene.arm_index, joints_table_collision)\n",
+    "scene.robot_diagram.ForcedPublish(context)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "filter_collisions_between_all_body_pairs(collision_checker, arm_body_indices, floor_body_indices, filtered=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We should now see no collisions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print_robot_collisions(joints_table_collision)\n",
+    "collision_checker.CheckConfigCollisionFree(joints_table_collision)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can re-enable collisions equally as easily."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "filter_collisions_between_all_body_pairs(collision_checker, arm_body_indices, floor_body_indices, filtered=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print_robot_collisions(joints_table_collision)\n",
+    "collision_checker.CheckConfigCollisionFree(joints_table_collision)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -717,7 +841,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "cloth-competition",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -735,5 +859,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }