From 361a2d1dd8cafebcabc36c1b937eef4c3c6f62c5 Mon Sep 17 00:00:00 2001
From: Sebastian Castro <sebas.a.castro@gmail.com>
Date: Fri, 5 Apr 2024 19:18:36 -0400
Subject: [PATCH] Also include self-collisions to make RRT look nicer

---
 examples/example_rrt.py        | 34 +++++++++++++++++++++++++++++++++-
 src/pyroboplan/planning/rrt.py |  2 +-
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/examples/example_rrt.py b/examples/example_rrt.py
index c8c24d6..ee98c12 100644
--- a/examples/example_rrt.py
+++ b/examples/example_rrt.py
@@ -13,6 +13,38 @@
 def prepare_scene(visual_model, collision_model):
     """Helper function to create a collision scene for this example."""
 
+    # Add self collisions
+    self_collision_pair_names = [
+        ("panda_link0_0", "panda_link2_0"),
+        ("panda_link0_0", "panda_link3_0"),
+        ("panda_link0_0", "panda_link4_0"),
+        ("panda_link0_0", "panda_link5_0"),
+        ("panda_link0_0", "panda_link6_0"),
+        ("panda_link0_0", "panda_link7_0"),
+        ("panda_link1_0", "panda_link3_0"),
+        ("panda_link1_0", "panda_link4_0"),
+        ("panda_link1_0", "panda_link5_0"),
+        ("panda_link1_0", "panda_link6_0"),
+        ("panda_link1_0", "panda_link7_0"),
+        ("panda_link2_0", "panda_link4_0"),
+        ("panda_link2_0", "panda_link5_0"),
+        ("panda_link2_0", "panda_link6_0"),
+        ("panda_link2_0", "panda_link7_0"),
+        ("panda_link3_0", "panda_link5_0"),
+        ("panda_link3_0", "panda_link6_0"),
+        ("panda_link3_0", "panda_link7_0"),
+        ("panda_link4_0", "panda_link6_0"),
+        ("panda_link4_0", "panda_link7_0"),
+        ("panda_link5_0", "panda_link7_0"),
+    ]
+    for pair in self_collision_pair_names:
+        collision_model.addCollisionPair(
+            pinocchio.CollisionPair(
+                collision_model.getGeometryId(pair[0]),
+                collision_model.getGeometryId(pair[1]),
+            )
+        )
+
     # Add collision objects
     ground_plane = pinocchio.GeometryObject(
         "ground_plane",
@@ -115,7 +147,7 @@ def prepare_scene(visual_model, collision_model):
     options.max_angle_step = 0.05
     options.max_connection_dist = 0.1
     options.goal_biasing_probability = 0.15
-    options.max_planning_time = 5.0
+    options.max_planning_time = 10.0
     options.rrt_connect = True
     options.bidirectional_rrt = True
 
diff --git a/src/pyroboplan/planning/rrt.py b/src/pyroboplan/planning/rrt.py
index 1c95727..cb73293 100644
--- a/src/pyroboplan/planning/rrt.py
+++ b/src/pyroboplan/planning/rrt.py
@@ -161,7 +161,7 @@ def plan(self, q_start, q_goal, options=RRTPlannerOptions()):
                 ):
                     goal_found = True
 
-                # Switch to the other tree
+                # Switch to the other tree next iteration, if bidirectional mode is enabled.
                 if options.bidirectional_rrt:
                     start_tree_phase = not start_tree_phase