Add RRTConnect implementation (#16)

sea-bass · Apr 7, 2024 · 92c9f96 · 92c9f96
1 parent 191f847
commit 92c9f96
Show file tree

Hide file tree

Showing 35 changed files with 590 additions and 333 deletions.
diff --git a/examples/example_collision_along_path.py b/examples/example_collision_along_path.py
@@ -3,10 +3,10 @@
 from pinocchio.visualize import MeshcatVisualizer
 import meshcat.geometry as mg
 import numpy as np
-from os.path import dirname, join, abspath
 import time
 
 from pyroboplan.core.utils import extract_cartesian_poses
+from pyroboplan.models.panda import load_models, add_self_collisions
 from pyroboplan.planning.utils import discretize_joint_space_path
 from pyroboplan.visualization.meshcat_utils import visualize_frames
 
@@ -31,33 +31,12 @@ def prepare_collision_scene(collision_model):
     obstacle_1 = pinocchio.GeometryObject(
         "obstacle_1",
         0,
-        hppfcl.Box(0.25, 0.25, 0.25),
+        hppfcl.Box(0.3, 0.3, 0.3),
         pinocchio.SE3(np.eye(3), np.array([-0.5, 0.5, 0.7])),
     )
     obstacle_1.meshColor = np.array([0.0, 1.0, 0.0, 0.2])
     collision_model.addGeometryObject(obstacle_1)
 
-
-if __name__ == "__main__":
-    # Create models and data
-    pinocchio_model_dir = join(dirname(str(abspath(__file__))), "..", "models")
-    package_dir = join(pinocchio_model_dir, "panda_description")
-    urdf_filename = join(package_dir, "urdf", "panda.urdf")
-
-    model, collision_model, visual_model = pinocchio.buildModelsFromUrdf(
-        urdf_filename, package_dirs=pinocchio_model_dir
-    )
-    data = model.createData()
-
-    prepare_collision_scene(collision_model)
-
-    # Initialize visualizer
-    viz = MeshcatVisualizer(model, collision_model, visual_model, data=data)
-    viz.initViewer(open=True)
-    viz.loadViewerModel()
-    viz.displayCollisions(True)
-    viz.displayVisuals(False)
-
     # Define the active collision pairs
     collision_names = [
         cobj.name for cobj in collision_model.geometryObjects if "panda" in cobj.name
@@ -71,7 +50,23 @@ def prepare_collision_scene(collision_model):
                     collision_model.getGeometryId(obstacle_name),
                 )
             )
-    collision_data = pinocchio.GeometryData(collision_model)
+
+
+if __name__ == "__main__":
+    # Create models and data
+    model, collision_model, visual_model = load_models()
+    add_self_collisions(collision_model)
+    prepare_collision_scene(collision_model)
+
+    data = model.createData()
+    collision_data = collision_model.createData()
+
+    # Initialize visualizer
+    viz = MeshcatVisualizer(model, collision_model, visual_model, data=data)
+    viz.initViewer(open=True)
+    viz.loadViewerModel()
+    viz.displayCollisions(True)
+    viz.displayVisuals(False)
 
     # Define a joint space path
     q_start = pinocchio.randomConfiguration(model)

diff --git a/examples/example_differential_ik.py b/examples/example_differential_ik.py
@@ -1,99 +1,31 @@
-import hppfcl
-import pinocchio
 from pinocchio.visualize import MeshcatVisualizer
 import numpy as np
-from os.path import dirname, join, abspath
 
 from pyroboplan.core.utils import get_random_state, get_random_transform
 from pyroboplan.ik.differential_ik import DifferentialIk, DifferentialIkOptions
 from pyroboplan.ik.nullspace_components import joint_limit_nullspace_component
-
-
-def prepare_scene(visual_model, collision_model):
-    """Helper function to create a collision scene for this example."""
-
-    # Add collision objects
-    obstacle_sphere_1 = pinocchio.GeometryObject(
-        "obstacle_sphere_1",
-        0,
-        hppfcl.Sphere(0.2),
-        pinocchio.SE3(np.eye(3), np.array([0.0, 0.1, 1.1])),
-    )
-    obstacle_sphere_1.meshColor = np.array([0.0, 1.0, 0.0, 0.5])
-    visual_model.addGeometryObject(obstacle_sphere_1)
-    collision_model.addGeometryObject(obstacle_sphere_1)
-
-    obstacle_sphere_2 = pinocchio.GeometryObject(
-        "obstacle_sphere_2",
-        0,
-        hppfcl.Sphere(0.25),
-        pinocchio.SE3(np.eye(3), np.array([0.5, 0.5, 0.5])),
-    )
-    obstacle_sphere_2.meshColor = np.array([1.0, 1.0, 0.0, 0.5])
-    visual_model.addGeometryObject(obstacle_sphere_2)
-    collision_model.addGeometryObject(obstacle_sphere_2)
-
-    obstacle_box_1 = pinocchio.GeometryObject(
-        "obstacle_box_1",
-        0,
-        hppfcl.Box(0.25, 0.25, 0.25),
-        pinocchio.SE3(np.eye(3), np.array([-0.5, 0.5, 0.7])),
-    )
-    obstacle_box_1.meshColor = np.array([1.0, 0.0, 0.0, 0.5])
-    visual_model.addGeometryObject(obstacle_box_1)
-    collision_model.addGeometryObject(obstacle_box_1)
-
-    obstacle_box_2 = pinocchio.GeometryObject(
-        "obstacle_box_2",
-        0,
-        hppfcl.Box(0.33, 0.33, 0.33),
-        pinocchio.SE3(np.eye(3), np.array([-0.5, -0.5, 0.75])),
-    )
-    obstacle_box_2.meshColor = np.array([0.0, 0.0, 1.0, 0.5])
-    visual_model.addGeometryObject(obstacle_box_2)
-    collision_model.addGeometryObject(obstacle_box_2)
+from pyroboplan.models.panda import (
+    load_models,
+    add_self_collisions,
+    add_object_collisions,
+)
 
 
 if __name__ == "__main__":
     # Create models and data
-    pinocchio_model_dir = join(dirname(str(abspath(__file__))), "..", "models")
-    package_dir = join(pinocchio_model_dir, "panda_description")
-    urdf_filename = join(package_dir, "urdf", "panda.urdf")
+    model, collision_model, visual_model = load_models()
+    add_self_collisions(collision_model)
+    add_object_collisions(collision_model, visual_model)
 
-    model, collision_model, visual_model = pinocchio.buildModelsFromUrdf(
-        urdf_filename, package_dirs=pinocchio_model_dir
-    )
     data = model.createData()
-
-    # Comment out to exclude collision objects.
-    prepare_scene(visual_model, collision_model)
+    collision_data = collision_model.createData()
 
     # Initialize visualizer
     viz = MeshcatVisualizer(model, collision_model, visual_model, data=data)
     viz.initViewer(open=True)
     viz.loadViewerModel()
     np.set_printoptions(precision=3)
 
-    # Define the active collision pairs
-    collision_names = [
-        cobj.name for cobj in collision_model.geometryObjects if "panda" in cobj.name
-    ]
-    obstacle_names = [
-        "obstacle_box_1",
-        "obstacle_box_2",
-        "obstacle_sphere_1",
-        "obstacle_sphere_2",
-    ]
-    for obstacle_name in obstacle_names:
-        for collision_name in collision_names:
-            collision_model.addCollisionPair(
-                pinocchio.CollisionPair(
-                    collision_model.getGeometryId(collision_name),
-                    collision_model.getGeometryId(obstacle_name),
-                )
-            )
-    collision_data = pinocchio.GeometryData(collision_model)
-
     # Set up the IK solver
     ik = DifferentialIk(
         model, data=data, collision_model=collision_model, visualizer=viz

diff --git a/examples/example_rrt.py b/examples/example_rrt.py
@@ -1,98 +1,30 @@
-import hppfcl
-import pinocchio
 from pinocchio.visualize import MeshcatVisualizer
-import numpy as np
-from os.path import dirname, join, abspath
 import time
 
 from pyroboplan.core.utils import get_random_collision_free_state
+from pyroboplan.models.panda import (
+    load_models,
+    add_self_collisions,
+    add_object_collisions,
+)
 from pyroboplan.planning.rrt import RRTPlanner, RRTPlannerOptions
 from pyroboplan.planning.utils import discretize_joint_space_path
 
 
-def prepare_scene(visual_model, collision_model):
-    """Helper function to create a collision scene for this example."""
-
-    # Add collision objects
-    obstacle_sphere_1 = pinocchio.GeometryObject(
-        "obstacle_sphere_1",
-        0,
-        hppfcl.Sphere(0.2),
-        pinocchio.SE3(np.eye(3), np.array([0.0, 0.1, 1.1])),
-    )
-    obstacle_sphere_1.meshColor = np.array([0.0, 1.0, 0.0, 0.5])
-    visual_model.addGeometryObject(obstacle_sphere_1)
-    collision_model.addGeometryObject(obstacle_sphere_1)
-
-    obstacle_sphere_2 = pinocchio.GeometryObject(
-        "obstacle_sphere_2",
-        0,
-        hppfcl.Sphere(0.25),
-        pinocchio.SE3(np.eye(3), np.array([0.5, 0.5, 0.5])),
-    )
-    obstacle_sphere_2.meshColor = np.array([1.0, 1.0, 0.0, 0.5])
-    visual_model.addGeometryObject(obstacle_sphere_2)
-    collision_model.addGeometryObject(obstacle_sphere_2)
-
-    obstacle_box_1 = pinocchio.GeometryObject(
-        "obstacle_box_1",
-        0,
-        hppfcl.Box(0.25, 0.25, 0.25),
-        pinocchio.SE3(np.eye(3), np.array([-0.5, 0.5, 0.7])),
-    )
-    obstacle_box_1.meshColor = np.array([1.0, 0.0, 0.0, 0.5])
-    visual_model.addGeometryObject(obstacle_box_1)
-    collision_model.addGeometryObject(obstacle_box_1)
-
-    obstacle_box_2 = pinocchio.GeometryObject(
-        "obstacle_box_2",
-        0,
-        hppfcl.Box(0.33, 0.33, 0.33),
-        pinocchio.SE3(np.eye(3), np.array([-0.5, -0.5, 0.75])),
-    )
-    obstacle_box_2.meshColor = np.array([0.0, 0.0, 1.0, 0.5])
-    visual_model.addGeometryObject(obstacle_box_2)
-    collision_model.addGeometryObject(obstacle_box_2)
-
-
 if __name__ == "__main__":
     # Create models and data
-    pinocchio_model_dir = join(dirname(str(abspath(__file__))), "..", "models")
-    package_dir = join(pinocchio_model_dir, "panda_description")
-    urdf_filename = join(package_dir, "urdf", "panda.urdf")
+    model, collision_model, visual_model = load_models()
+    add_self_collisions(collision_model)
+    add_object_collisions(collision_model, visual_model)
 
-    model, collision_model, visual_model = pinocchio.buildModelsFromUrdf(
-        urdf_filename, package_dirs=pinocchio_model_dir
-    )
     data = model.createData()
-
-    prepare_scene(visual_model, collision_model)
+    collision_data = collision_model.createData()
 
     # Initialize visualizer
     viz = MeshcatVisualizer(model, collision_model, visual_model, data=data)
     viz.initViewer(open=True)
     viz.loadViewerModel()
 
-    # Define the active collision pairs
-    collision_names = [
-        cobj.name for cobj in collision_model.geometryObjects if "panda" in cobj.name
-    ]
-    obstacle_names = [
-        "obstacle_box_1",
-        "obstacle_box_2",
-        "obstacle_sphere_1",
-        "obstacle_sphere_2",
-    ]
-    for obstacle_name in obstacle_names:
-        for collision_name in collision_names:
-            collision_model.addCollisionPair(
-                pinocchio.CollisionPair(
-                    collision_model.getGeometryId(collision_name),
-                    collision_model.getGeometryId(obstacle_name),
-                )
-            )
-    collision_data = pinocchio.GeometryData(collision_model)
-
     # Define the start and end configurations
     q_start = get_random_collision_free_state(model, collision_model)
     q_end = get_random_collision_free_state(model, collision_model)
@@ -102,9 +34,11 @@ def prepare_scene(visual_model, collision_model):
     # Search for a path
     options = RRTPlannerOptions()
     options.max_angle_step = 0.05
-    options.max_connection_dist = 0.5
+    options.max_connection_dist = 0.25
     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
 
     planner = RRTPlanner(model, collision_model)
     path = planner.plan(q_start, q_end, options=options)

diff --git a/src/pyroboplan/models/__init__.py b/src/pyroboplan/models/__init__.py
@@ -0,0 +1 @@
+""" Contains example robot models. """