From 92c9f9654fa403a64ef82e37e180e118c25a460e Mon Sep 17 00:00:00 2001 From: Sebastian Castro <4603398+sea-bass@users.noreply.github.com> Date: Sun, 7 Apr 2024 23:11:36 +0100 Subject: [PATCH] Add RRTConnect implementation (#16) --- examples/example_collision_along_path.py | 43 ++-- examples/example_differential_ik.py | 86 +------ examples/example_rrt.py | 92 ++------ src/pyroboplan/models/__init__.py | 1 + src/pyroboplan/models/panda.py | 148 ++++++++++++ .../models}/panda_description/LICENSE | 0 .../models}/panda_description/README.md | 0 .../meshes/collision/finger.stl | Bin .../meshes/collision/hand.stl | Bin .../meshes/collision/link0.stl | Bin .../meshes/collision/link1.stl | Bin .../meshes/collision/link2.stl | Bin .../meshes/collision/link3.stl | Bin .../meshes/collision/link4.stl | Bin .../meshes/collision/link5.stl | Bin .../meshes/collision/link6.stl | Bin .../meshes/collision/link7.stl | Bin .../meshes/visual/finger.dae | 0 .../panda_description/meshes/visual/hand.dae | 0 .../panda_description/meshes/visual/link0.dae | 0 .../panda_description/meshes/visual/link1.dae | 0 .../panda_description/meshes/visual/link2.dae | 0 .../panda_description/meshes/visual/link3.dae | 0 .../panda_description/meshes/visual/link4.dae | 0 .../panda_description/meshes/visual/link5.dae | 0 .../panda_description/meshes/visual/link6.dae | 0 .../panda_description/meshes/visual/link7.dae | 0 .../models}/panda_description/urdf/panda.urdf | 0 src/pyroboplan/models/utils.py | 16 ++ src/pyroboplan/planning/graph.py | 11 +- src/pyroboplan/planning/rrt.py | 217 ++++++++++++++---- test/core/test_utils.py | 64 +----- test/ik/test_differential_ik.py | 54 +---- test/planning/test_graph.py | 108 +++++++++ test/planning/test_rrt.py | 83 +++++++ 35 files changed, 590 insertions(+), 333 deletions(-) create mode 100644 src/pyroboplan/models/__init__.py create mode 100644 src/pyroboplan/models/panda.py rename {models => src/pyroboplan/models}/panda_description/LICENSE (100%) rename {models => src/pyroboplan/models}/panda_description/README.md (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/collision/finger.stl (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/collision/hand.stl (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/collision/link0.stl (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/collision/link1.stl (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/collision/link2.stl (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/collision/link3.stl (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/collision/link4.stl (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/collision/link5.stl (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/collision/link6.stl (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/collision/link7.stl (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/visual/finger.dae (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/visual/hand.dae (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/visual/link0.dae (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/visual/link1.dae (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/visual/link2.dae (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/visual/link3.dae (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/visual/link4.dae (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/visual/link5.dae (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/visual/link6.dae (100%) rename {models => src/pyroboplan/models}/panda_description/meshes/visual/link7.dae (100%) rename {models => src/pyroboplan/models}/panda_description/urdf/panda.urdf (100%) create mode 100644 src/pyroboplan/models/utils.py create mode 100644 test/planning/test_graph.py create mode 100644 test/planning/test_rrt.py diff --git a/examples/example_collision_along_path.py b/examples/example_collision_along_path.py index f65888c..801970b 100644 --- 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 index d21714f..122bb2c 100644 --- a/examples/example_differential_ik.py +++ b/examples/example_differential_ik.py @@ -1,72 +1,24 @@ -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) @@ -74,26 +26,6 @@ def prepare_scene(visual_model, collision_model): 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 index ff5e889..0f1be33 100644 --- 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 new file mode 100644 index 0000000..4c80dc4 --- /dev/null +++ b/src/pyroboplan/models/__init__.py @@ -0,0 +1 @@ +""" Contains example robot models. """ diff --git a/src/pyroboplan/models/panda.py b/src/pyroboplan/models/panda.py new file mode 100644 index 0000000..d0b4252 --- /dev/null +++ b/src/pyroboplan/models/panda.py @@ -0,0 +1,148 @@ +""" Utilities to load example Franka Emika Panda model. """ + +import hppfcl +import numpy as np +import os +import pinocchio + +from .utils import get_example_models_folder + + +def load_models(): + """ + Gets the example Panda models. + + Returns + ------- + tuple[`pinocchio.Model`] + A 3-tuple containing the model, collision geometry model, and visual geometry model. + """ + models_folder = get_example_models_folder() + package_dir = os.path.join(models_folder, "panda_description") + urdf_filename = os.path.join(package_dir, "urdf", "panda.urdf") + + return pinocchio.buildModelsFromUrdf(urdf_filename, package_dirs=models_folder) + + +def add_self_collisions(collision_model): + """ + Adds link self-collisions to the Panda collision model. + + Parameters + ---------- + collision_model : `pinocchio.Model` + The Panda collision geometry model. + """ + 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]), + ) + ) + + +def add_object_collisions(collision_model, visual_model): + """ + Adds link self-collisions to the Panda collision model. + + Parameters + ---------- + collision_model : `pinocchio.Model` + The Panda collision geometry model. + visual_model : `pinocchio.Model` + The Panda visual geometry model. + """ + # Add the collision objects + ground_plane = pinocchio.GeometryObject( + "ground_plane", + 0, + hppfcl.Box(2.0, 2.0, 0.01), + pinocchio.SE3(np.eye(3), np.array([0.0, 0.0, -0.006])), + ) + ground_plane.meshColor = np.array([0.5, 0.5, 0.5, 0.5]) + visual_model.addGeometryObject(ground_plane) + collision_model.addGeometryObject(ground_plane) + + 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) + + # Define the active collision pairs between the robot and obstacle links. + collision_names = [ + cobj.name for cobj in collision_model.geometryObjects if "panda" in cobj.name + ] + obstacle_names = [ + "ground_plane", + "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), + ) + ) diff --git a/models/panda_description/LICENSE b/src/pyroboplan/models/panda_description/LICENSE similarity index 100% rename from models/panda_description/LICENSE rename to src/pyroboplan/models/panda_description/LICENSE diff --git a/models/panda_description/README.md b/src/pyroboplan/models/panda_description/README.md similarity index 100% rename from models/panda_description/README.md rename to src/pyroboplan/models/panda_description/README.md diff --git a/models/panda_description/meshes/collision/finger.stl b/src/pyroboplan/models/panda_description/meshes/collision/finger.stl similarity index 100% rename from models/panda_description/meshes/collision/finger.stl rename to src/pyroboplan/models/panda_description/meshes/collision/finger.stl diff --git a/models/panda_description/meshes/collision/hand.stl b/src/pyroboplan/models/panda_description/meshes/collision/hand.stl similarity index 100% rename from models/panda_description/meshes/collision/hand.stl rename to src/pyroboplan/models/panda_description/meshes/collision/hand.stl diff --git a/models/panda_description/meshes/collision/link0.stl b/src/pyroboplan/models/panda_description/meshes/collision/link0.stl similarity index 100% rename from models/panda_description/meshes/collision/link0.stl rename to src/pyroboplan/models/panda_description/meshes/collision/link0.stl diff --git a/models/panda_description/meshes/collision/link1.stl b/src/pyroboplan/models/panda_description/meshes/collision/link1.stl similarity index 100% rename from models/panda_description/meshes/collision/link1.stl rename to src/pyroboplan/models/panda_description/meshes/collision/link1.stl diff --git a/models/panda_description/meshes/collision/link2.stl b/src/pyroboplan/models/panda_description/meshes/collision/link2.stl similarity index 100% rename from models/panda_description/meshes/collision/link2.stl rename to src/pyroboplan/models/panda_description/meshes/collision/link2.stl diff --git a/models/panda_description/meshes/collision/link3.stl b/src/pyroboplan/models/panda_description/meshes/collision/link3.stl similarity index 100% rename from models/panda_description/meshes/collision/link3.stl rename to src/pyroboplan/models/panda_description/meshes/collision/link3.stl diff --git a/models/panda_description/meshes/collision/link4.stl b/src/pyroboplan/models/panda_description/meshes/collision/link4.stl similarity index 100% rename from models/panda_description/meshes/collision/link4.stl rename to src/pyroboplan/models/panda_description/meshes/collision/link4.stl diff --git a/models/panda_description/meshes/collision/link5.stl b/src/pyroboplan/models/panda_description/meshes/collision/link5.stl similarity index 100% rename from models/panda_description/meshes/collision/link5.stl rename to src/pyroboplan/models/panda_description/meshes/collision/link5.stl diff --git a/models/panda_description/meshes/collision/link6.stl b/src/pyroboplan/models/panda_description/meshes/collision/link6.stl similarity index 100% rename from models/panda_description/meshes/collision/link6.stl rename to src/pyroboplan/models/panda_description/meshes/collision/link6.stl diff --git a/models/panda_description/meshes/collision/link7.stl b/src/pyroboplan/models/panda_description/meshes/collision/link7.stl similarity index 100% rename from models/panda_description/meshes/collision/link7.stl rename to src/pyroboplan/models/panda_description/meshes/collision/link7.stl diff --git a/models/panda_description/meshes/visual/finger.dae b/src/pyroboplan/models/panda_description/meshes/visual/finger.dae similarity index 100% rename from models/panda_description/meshes/visual/finger.dae rename to src/pyroboplan/models/panda_description/meshes/visual/finger.dae diff --git a/models/panda_description/meshes/visual/hand.dae b/src/pyroboplan/models/panda_description/meshes/visual/hand.dae similarity index 100% rename from models/panda_description/meshes/visual/hand.dae rename to src/pyroboplan/models/panda_description/meshes/visual/hand.dae diff --git a/models/panda_description/meshes/visual/link0.dae b/src/pyroboplan/models/panda_description/meshes/visual/link0.dae similarity index 100% rename from models/panda_description/meshes/visual/link0.dae rename to src/pyroboplan/models/panda_description/meshes/visual/link0.dae diff --git a/models/panda_description/meshes/visual/link1.dae b/src/pyroboplan/models/panda_description/meshes/visual/link1.dae similarity index 100% rename from models/panda_description/meshes/visual/link1.dae rename to src/pyroboplan/models/panda_description/meshes/visual/link1.dae diff --git a/models/panda_description/meshes/visual/link2.dae b/src/pyroboplan/models/panda_description/meshes/visual/link2.dae similarity index 100% rename from models/panda_description/meshes/visual/link2.dae rename to src/pyroboplan/models/panda_description/meshes/visual/link2.dae diff --git a/models/panda_description/meshes/visual/link3.dae b/src/pyroboplan/models/panda_description/meshes/visual/link3.dae similarity index 100% rename from models/panda_description/meshes/visual/link3.dae rename to src/pyroboplan/models/panda_description/meshes/visual/link3.dae diff --git a/models/panda_description/meshes/visual/link4.dae b/src/pyroboplan/models/panda_description/meshes/visual/link4.dae similarity index 100% rename from models/panda_description/meshes/visual/link4.dae rename to src/pyroboplan/models/panda_description/meshes/visual/link4.dae diff --git a/models/panda_description/meshes/visual/link5.dae b/src/pyroboplan/models/panda_description/meshes/visual/link5.dae similarity index 100% rename from models/panda_description/meshes/visual/link5.dae rename to src/pyroboplan/models/panda_description/meshes/visual/link5.dae diff --git a/models/panda_description/meshes/visual/link6.dae b/src/pyroboplan/models/panda_description/meshes/visual/link6.dae similarity index 100% rename from models/panda_description/meshes/visual/link6.dae rename to src/pyroboplan/models/panda_description/meshes/visual/link6.dae diff --git a/models/panda_description/meshes/visual/link7.dae b/src/pyroboplan/models/panda_description/meshes/visual/link7.dae similarity index 100% rename from models/panda_description/meshes/visual/link7.dae rename to src/pyroboplan/models/panda_description/meshes/visual/link7.dae diff --git a/models/panda_description/urdf/panda.urdf b/src/pyroboplan/models/panda_description/urdf/panda.urdf similarity index 100% rename from models/panda_description/urdf/panda.urdf rename to src/pyroboplan/models/panda_description/urdf/panda.urdf diff --git a/src/pyroboplan/models/utils.py b/src/pyroboplan/models/utils.py new file mode 100644 index 0000000..2bb9c8b --- /dev/null +++ b/src/pyroboplan/models/utils.py @@ -0,0 +1,16 @@ +""" Utilities for model loading. """ + +import importlib.resources +import pyroboplan + + +def get_example_models_folder(): + """ + Returns the full path the example models folder. + + Returns + ------- + str + The path to the `pyroboplan` example models folder. + """ + return importlib.resources.path(pyroboplan, "models").as_posix() diff --git a/src/pyroboplan/planning/graph.py b/src/pyroboplan/planning/graph.py index 2effe2d..e59e5f1 100644 --- a/src/pyroboplan/planning/graph.py +++ b/src/pyroboplan/planning/graph.py @@ -19,7 +19,7 @@ def __init__(self, q, parent=None, cost=0.0): cost : float, optional The cost associated with the node. """ - self.q = q + self.q = np.array(q) self.parent = parent self.cost = cost self.neighbors = set() @@ -77,11 +77,18 @@ def add_edge(self, nodeA, nodeB): The first node in the edge. nodeB : `pyroboplan.planning.graph.Node` The second node in the edge. + + Returns + ------- + `pyroboplan.planning.graph.Edge` + The edge that was added. """ cost = configuration_distance(nodeA.q, nodeB.q) - self.edges.add(Edge(nodeA, nodeB, cost)) + edge = Edge(nodeA, nodeB, cost) + self.edges.add(edge) nodeA.neighbors.add(nodeB) nodeB.neighbors.add(nodeA) + return edge def get_nearest_node(self, q): """ diff --git a/src/pyroboplan/planning/rrt.py b/src/pyroboplan/planning/rrt.py index 9367635..181b147 100644 --- a/src/pyroboplan/planning/rrt.py +++ b/src/pyroboplan/planning/rrt.py @@ -25,6 +25,15 @@ class RRTPlannerOptions: max_connection_dist = 0.2 """ Maximum angular distance, in radians, for connecting nodes. """ + rrt_connect = False + """ If true, enables the RRTConnect algorithm. """ + + bidirectional_rrt = False + """ + If true, uses bidirectional RRTs from both start and goal nodes. + Otherwise, only grows a tree from the start node. + """ + max_planning_time = 10.0 """ Maximum planning time, in seconds. """ @@ -38,7 +47,8 @@ class RRTPlanner: This is a sampling-based motion planner that finds collision-free paths from a start to a goal configuration. Some good resources: - * https://msl.cs.illinois.edu/~lavalle/papers/Lav98c.pdf + * Original RRT paper: https://msl.cs.illinois.edu/~lavalle/papers/Lav98c.pdf + * RRTConnect paper: https://www.cs.cmu.edu/afs/cs/academic/class/15494-s14/readings/kuffner_icra2000.pdf """ def __init__(self, model, collision_model): @@ -54,9 +64,13 @@ def __init__(self, model, collision_model): """ self.model = model self.collision_model = collision_model + self.reset() + def reset(self): + """Resets all the planning data structures.""" self.latest_path = None - self.graph = Graph() + self.start_tree = Graph() + self.goal_tree = Graph() def plan(self, q_start, q_goal, options=RRTPlannerOptions()): """ @@ -71,14 +85,18 @@ def plan(self, q_start, q_goal, options=RRTPlannerOptions()): options : `RRTPlannerOptions`, optional The options to use for planning. If not specified, default options are used. """ + self.reset() t_start = time.time() self.options = options - self.graph = Graph() + start_node = Node(q_start, parent=None) - self.graph.add_node(start_node) + self.start_tree.add_node(start_node) + goal_node = Node(q_goal, parent=None) + self.goal_tree.add_node(goal_node) goal_found = False - latest_node = start_node + latest_start_tree_node = start_node + latest_goal_tree_node = goal_node # Check start and end pose collisions. if check_collisions_at_state(self.model, self.collision_model, q_start): @@ -90,76 +108,162 @@ def plan(self, q_start, q_goal, options=RRTPlannerOptions()): # Check direct connection to goal. path_to_goal = discretize_joint_space_path( - start_node.q, q_goal, self.options.max_angle_step + q_start, q_goal, self.options.max_angle_step ) if not check_collisions_along_path( self.model, self.collision_model, path_to_goal ): - goal_node = Node(q_goal, parent=start_node) - self.graph.add_node(latest_node) - self.graph.add_edge(start_node, goal_node) + print("Start and goal can be directly connected!") goal_found = True + start_tree_phase = True while not goal_found: - # Check for timeouts + # Check for timeouts. if time.time() - t_start > options.max_planning_time: print(f"Planning timed out after {options.max_planning_time} seconds.") break + # Choose variables based on whether we're growing the start or goal tree. + tree = self.start_tree if start_tree_phase else self.goal_tree + other_tree = self.goal_tree if start_tree_phase else self.start_tree + # Sample a new configuration. if np.random.random() < self.options.goal_biasing_probability: - q_sample = q_goal + q_sample = q_goal if start_tree_phase else q_start else: q_sample = get_random_state(self.model) - nearest_node = self.graph.get_nearest_node(q_sample) - # Clip the distance between nearest and sampled nodes to max connection distance. - distance = configuration_distance(q_sample, nearest_node.q) - if distance > self.options.max_connection_dist: - scale = self.options.max_connection_dist / distance - q_sample = nearest_node.q + scale * (q_sample - nearest_node.q) - - # Add the node only if it is collision free. - if check_collisions_at_state(self.model, self.collision_model, q_sample): - continue + # Run the extend or connect operation to connect the tree to the new node. + nearest_node = tree.get_nearest_node(q_sample) + q_new = self.extend_or_connect(nearest_node, q_sample, options) + + # Only if extend/connect succeeded, add the new node to the tree. + if q_new is not None: + new_node = Node(q_new, parent=nearest_node) + tree.add_node(new_node) + tree.add_edge(nearest_node, new_node) + if start_tree_phase: + latest_start_tree_node = new_node + else: + latest_goal_tree_node = new_node - path_to_node = discretize_joint_space_path( - nearest_node.q, q_sample, self.options.max_angle_step - ) - if not check_collisions_along_path( - self.model, self.collision_model, path_to_node - ): - latest_node = Node(q_sample, parent=nearest_node) - self.graph.add_node(latest_node) - self.graph.add_edge(nearest_node, latest_node) - - # Check if latest node connects directly to goal. - path_to_goal = discretize_joint_space_path( - latest_node.q, q_goal, self.options.max_angle_step + # Check if latest node connects directly to the other tree. + # If so, planning is complete. + nearest_node_in_other_tree = other_tree.get_nearest_node(new_node.q) + path_to_other_tree = discretize_joint_space_path( + q_new, nearest_node_in_other_tree.q, self.options.max_angle_step ) if not check_collisions_along_path( - self.model, self.collision_model, path_to_goal + self.model, self.collision_model, path_to_other_tree ): - goal_node = Node(q_goal, parent=latest_node) - self.graph.add_node(goal_node) - self.graph.add_edge(latest_node, goal_node) + if start_tree_phase: + latest_goal_tree_node = nearest_node_in_other_tree + else: + latest_start_tree_node = nearest_node_in_other_tree goal_found = True + # Switch to the other tree next iteration, if bidirectional mode is enabled. + if options.bidirectional_rrt: + start_tree_phase = not start_tree_phase + # Back out the path by traversing the parents from the goal. self.latest_path = [] if goal_found: - cur_node = goal_node - path_extracted = False - while not path_extracted: - if cur_node is None: - path_extracted = True - else: - self.latest_path.append(cur_node.q) - cur_node = cur_node.parent - self.latest_path.reverse() - + self.latest_path = self.extract_path_from_trees( + latest_start_tree_node, latest_goal_tree_node + ) return self.latest_path + def extend_or_connect(self, parent_node, q_sample, options=RRTPlannerOptions()): + """ + Extends a tree towards a sampled node with steps no larger than the maximum connection distance. + + Parameters + ---------- + parent_node : `pyroboplan.planning.graph.Node` + The node from which to start extending or connecting towards the sample. + q_sample : array-like + The robot configuration sample to extend or connect towards. + options : `pyroboplan.planning.rrt.RRTPlannerOptions`, optional + The options to use for this operation. + These include whether to extend once or keep connecting (`options.rrt_connect`), + as well as the maximum angular connection distance (`options.max_connection_dist`). + """ + q_diff = q_sample - parent_node.q + q_increment = options.max_connection_dist * q_diff / np.linalg.norm(q_diff) + + terminated = False + q_out = None + q_cur = parent_node.q + while not terminated: + # Clip the distance between nearest and sampled nodes to max connection distance. + # If we have reached the sampled node, this is the final iteration. + if configuration_distance(q_cur, q_sample) > options.max_connection_dist: + q_extend = q_cur + q_increment + else: + q_extend = q_sample + terminated |= True + + # Extension is successful only if the path is collision free. + q_extend_in_collision = check_collisions_at_state( + self.model, self.collision_model, q_extend + ) + path_to_q_extend = discretize_joint_space_path( + q_cur, q_extend, self.options.max_angle_step + ) + path_to_q_extend_in_collision = check_collisions_along_path( + self.model, self.collision_model, path_to_q_extend + ) + if not q_extend_in_collision and not path_to_q_extend_in_collision: + q_cur = q_out = q_extend + else: + terminated |= True + + # If RRTConnect is disabled, only one iteration is needed. + if not options.rrt_connect: + terminated |= True + + return q_out + + def extract_path_from_trees(self, start_tree_final_node, goal_tree_final_node): + """ + Extracts the final path from the RRT trees. + + Parameters + ---------- + start_tree_final_node : `pyroboplan.planning.graph.Node` + The last node of the start tree. + goal_tree_final_node : `pyroboplan.planning.graph.Node`, optional + The last node of the goal tree. + If None, this means the goal tree is ignored. + + Return + ------ + list[array-like] + A list of robot configurations describing the path waypoints in order. + """ + path = [] + cur_node = start_tree_final_node + path_extracted = False + while not path_extracted: + if cur_node is None: + path_extracted = True + else: + path.append(cur_node.q) + cur_node = cur_node.parent + path.reverse() + + cur_node = goal_tree_final_node + path_extracted = False + while not path_extracted: + if cur_node is None: + path_extracted = True + else: + path.append(cur_node.q) + cur_node = cur_node.parent + + return path + def visualize( self, visualizer, @@ -202,15 +306,28 @@ def visualize( ) if show_tree: - for idx, edge in enumerate(self.graph.edges): + for idx, edge in enumerate(self.start_tree.edges): q_path = discretize_joint_space_path( edge.nodeA.q, edge.nodeB.q, self.options.max_angle_step ) path_tforms = extract_cartesian_poses(self.model, frame_name, q_path) visualize_path( visualizer, - f"{tree_name}/edge{idx}", + f"{tree_name}_start/edge{idx}", path_tforms, line_width=0.5, line_color=[0.9, 0.0, 0.9], ) + + for idx, edge in enumerate(self.goal_tree.edges): + q_path = discretize_joint_space_path( + edge.nodeA.q, edge.nodeB.q, self.options.max_angle_step + ) + path_tforms = extract_cartesian_poses(self.model, frame_name, q_path) + visualize_path( + visualizer, + f"{tree_name}_goal/edge{idx}", + path_tforms, + line_width=0.5, + line_color=[0.0, 0.9, 0.9], + ) diff --git a/test/core/test_utils.py b/test/core/test_utils.py index 38de5dd..ec61503 100644 --- a/test/core/test_utils.py +++ b/test/core/test_utils.py @@ -1,4 +1,3 @@ -from os.path import dirname, join, abspath import numpy as np import pinocchio import pytest @@ -12,54 +11,13 @@ get_random_state, get_random_transform, ) +from pyroboplan.models.panda import load_models, add_self_collisions # Use a fixed seed for random number generation in tests. np.random.seed(1234) -def load_panda_models(): - 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") - return pinocchio.buildModelsFromUrdf( - urdf_filename, package_dirs=pinocchio_model_dir - ) - - -def activate_panda_self_collision_pairs(collision_model): - 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 collision_pair_names: - collision_model.addCollisionPair( - pinocchio.CollisionPair( - collision_model.getGeometryId(pair[0]), - collision_model.getGeometryId(pair[1]), - ) - ) - - def test_configuration_distance(): q_start = np.array([0.0, 0.0, 0.0, 0.0, 0.0]) q_end = np.array([0.3, 0.0, -0.4, 0.0, 0.0]) @@ -67,7 +25,7 @@ def test_configuration_distance(): def test_get_random_state(): - model, _, _ = load_panda_models() + model, _, _ = load_models() for _ in range(10): state = get_random_state(model) assert np.all(state >= model.lowerPositionLimit) @@ -75,7 +33,7 @@ def test_get_random_state(): def test_get_random_state_with_scalar_padding(): - model, _, _ = load_panda_models() + model, _, _ = load_models() padding = 0.01 for _ in range(10): state = get_random_state(model, padding=padding) @@ -84,7 +42,7 @@ def test_get_random_state_with_scalar_padding(): def test_get_random_state_with_vector_padding(): - model, _, _ = load_panda_models() + model, _, _ = load_models() padding = [0.1, 0.2, 0.1, 0.2, 0.1, 0.2, 0.1, 0.01, 0.01] for _ in range(10): state = get_random_state(model, padding=padding) @@ -93,13 +51,13 @@ def test_get_random_state_with_vector_padding(): def test_get_random_transform(): - model, _, _ = load_panda_models() + model, _, _ = load_models() tform = get_random_transform(model, "panda_hand") assert isinstance(tform, pinocchio.SE3) def test_within_limits(): - model, _, _ = load_panda_models() + model, _, _ = load_models() assert check_within_limits(model, get_random_state(model)) assert not check_within_limits(model, model.upperPositionLimit + 0.2) @@ -107,7 +65,7 @@ def test_within_limits(): def test_extract_cartesian_poses(): - model, _, _ = load_panda_models() + model, _, _ = load_models() q_vec = [get_random_state(model) for _ in range(5)] tforms = extract_cartesian_poses(model, "panda_hand", q_vec) @@ -118,8 +76,8 @@ def test_extract_cartesian_poses(): def test_check_collisions_at_state(): - model, collision_model, _ = load_panda_models() - activate_panda_self_collision_pairs(collision_model) + model, collision_model, _ = load_models() + add_self_collisions(collision_model) collision_state = np.array( [ @@ -154,8 +112,8 @@ def test_check_collisions_at_state(): def test_check_collisions_along_path(): - model, collision_model, _ = load_panda_models() - activate_panda_self_collision_pairs(collision_model) + model, collision_model, _ = load_models() + add_self_collisions(collision_model) collision_state = np.array( [ diff --git a/test/ik/test_differential_ik.py b/test/ik/test_differential_ik.py index 3f8174f..edc4f85 100644 --- a/test/ik/test_differential_ik.py +++ b/test/ik/test_differential_ik.py @@ -1,55 +1,13 @@ import copy -from os.path import dirname, join, abspath import numpy as np import pinocchio from pyroboplan.ik.differential_ik import DifferentialIk, DifferentialIkOptions - - -def load_panda_models(): - 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") - return pinocchio.buildModelsFromUrdf( - urdf_filename, package_dirs=pinocchio_model_dir - ) - - -def activate_panda_self_collision_pairs(collision_model): - 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 collision_pair_names: - collision_model.addCollisionPair( - pinocchio.CollisionPair( - collision_model.getGeometryId(pair[0]), - collision_model.getGeometryId(pair[1]), - ) - ) +from pyroboplan.models.panda import load_models, add_self_collisions def test_ik_solve_trivial_ik(): - model, _, _ = load_panda_models() + model, _, _ = load_models() data = model.createData() target_frame = "panda_hand" @@ -77,7 +35,7 @@ def test_ik_solve_trivial_ik(): def test_ik_solve_ik(): - model, _, _ = load_panda_models() + model, _, _ = load_models() data = model.createData() target_frame = "panda_hand" @@ -115,7 +73,7 @@ def test_ik_solve_ik(): def test_ik_solve_impossible_ik(): - model, _, _ = load_panda_models() + model, _, _ = load_models() target_frame = "panda_hand" # Target is unreachable by the panda @@ -138,8 +96,8 @@ def test_ik_solve_impossible_ik(): def test_ik_in_collision(): - model, collision_model, _ = load_panda_models() - activate_panda_self_collision_pairs(collision_model) + model, collision_model, _ = load_models() + add_self_collisions(collision_model) target_frame = "panda_hand" # Target is reachable, but in self-collision. diff --git a/test/planning/test_graph.py b/test/planning/test_graph.py new file mode 100644 index 0000000..7813e98 --- /dev/null +++ b/test/planning/test_graph.py @@ -0,0 +1,108 @@ +import numpy as np +import pytest + +from pyroboplan.planning.graph import Node, Edge, Graph + + +def test_create_node(): + node = Node(np.array([1.0, 2.0, 3.0])) + + assert np.all(node.q == np.array([1.0, 2.0, 3.0])) + assert node.parent is None + assert node.cost == 0.0 + assert len(node.neighbors) == 0 + + +def test_create_node_nondefault_args(): + parent_node = Node([1.0, 2.0, 3.0]) + node = Node([4.0, 5.0, 6.0], parent=parent_node, cost=42.0) + + assert np.all(node.q == np.array([4.0, 5.0, 6.0])) + assert node.parent == parent_node + assert np.all(node.parent.q == np.array([1.0, 2.0, 3.0])) + assert node.cost == 42.0 + assert len(node.neighbors) == 0 + + +def test_create_edge(): + nodeA = Node([1.0, 2.0, 3.0]) + nodeB = Node([4.0, 5.0, 6.0]) + edge = Edge(nodeA, nodeB) + + assert edge.nodeA == nodeA + assert edge.nodeB == nodeB + assert edge.cost == 0.0 + + +def test_create_edge_nondefault_args(): + nodeA = Node([1.0, 2.0, 3.0]) + nodeB = Node([4.0, 5.0, 6.0]) + edge = Edge(nodeA, nodeB, cost=21.0) + + assert edge.nodeA == nodeA + assert edge.nodeB == nodeB + assert edge.cost == 21.0 + + +def test_create_empty_graph(): + graph = Graph() + assert len(graph.nodes) == 0 + assert len(graph.edges) == 0 + + +def test_add_nodes(): + nodeA = Node([1.0, 2.0, 3.0]) + nodeB = Node([4.0, 5.0, 6.0]) + graph = Graph() + + graph.add_node(nodeA) + assert len(graph.nodes) == 1 + + graph.add_node(nodeB) + assert len(graph.nodes) == 2 + + # Adding the same node again does not count as a repeat since the underlying storage is a set. + graph.add_node(nodeA) + assert len(graph.nodes) == 2 + + assert len(graph.edges) == 0 + + +def test_add_edges(): + nodeA = Node([1.0, 2.0, 3.0]) + nodeB = Node([4.0, 5.0, 6.0]) + nodeC = Node([7.0, 8.0, 9.0]) + graph = Graph() + graph.add_node(nodeA) + graph.add_node(nodeB) + graph.add_node(nodeC) + + edgeAB = graph.add_edge(nodeA, nodeB) + assert len(graph.edges) == 1 + assert nodeB in nodeA.neighbors + assert nodeA in nodeB.neighbors + assert edgeAB.cost == pytest.approx(np.linalg.norm([3.0, 3.0, 3.0])) + + edgeAC = graph.add_edge(nodeA, nodeC) + assert len(graph.edges) == 2 + assert nodeC in nodeA.neighbors + assert nodeA in nodeC.neighbors + assert edgeAC.cost == pytest.approx(np.linalg.norm([6.0, 6.0, 6.0])) + + +def test_get_nearest_node(): + nodeA = Node([1.0, 2.0, 3.0]) + nodeB = Node([4.0, 5.0, 6.0]) + + graph = Graph() + graph.add_node(nodeA) + graph.add_node(nodeB) + + nearest_node = graph.get_nearest_node([5.0, 4.0, 3.0]) + assert nearest_node == nodeB + + +def test_get_nearest_node_empty_graph(): + graph = Graph() + nearest_node = graph.get_nearest_node([1.0, 2.0, 3.0]) + assert nearest_node is None diff --git a/test/planning/test_rrt.py b/test/planning/test_rrt.py new file mode 100644 index 0000000..25aec33 --- /dev/null +++ b/test/planning/test_rrt.py @@ -0,0 +1,83 @@ +import numpy as np + +from pyroboplan.planning.rrt import RRTPlanner, RRTPlannerOptions +from pyroboplan.models.panda import ( + load_models, + add_self_collisions, + add_object_collisions, +) + + +# Use a fixed seed for random number generation in tests. +np.random.seed(1234) + + +def test_plan_trivial_rrt(): + model, collision_model, _ = load_models() + + # Initial joint states + q_start = np.array([0.0, 1.57, 0.0, 0.0, 1.57, 1.57, 0.0, 0.0, 0.0]) + q_goal = q_start + 0.01 + + # Plan with default parameters + planner = RRTPlanner(model, collision_model) + path = planner.plan(q_start, q_goal) + + assert path is not None + assert len(path) == 2 + assert np.all(path[0] == q_start) + assert np.all(path[1] == q_goal) + + +def test_plan_vanilla_rrt(): + model, collision_model, visual_model = load_models() + add_self_collisions(collision_model) + add_object_collisions(collision_model, visual_model) + + # Plan with default parameters + q_start = np.array([0.0, 1.57, 0.0, 0.0, 1.57, 1.57, 0.0, 0.0, 0.0]) + q_goal = np.array([1.57, 1.57, 0.0, 0.0, 1.57, 1.57, 0.0, 0.0, 0.0]) + + planner = RRTPlanner(model, collision_model) + path = planner.plan(q_start, q_goal) + + # The path must exist and have more than + assert path is not None + assert len(path) > 2 + assert np.all(path[0] == q_start) + assert np.all(path[-1] == q_goal) + + # Vanilla RRT should only grow the start tree, so we should check that the goal tree is just the goal node. + assert len(planner.start_tree.nodes) > 1 + assert len(planner.start_tree.edges) > 1 + assert len(planner.goal_tree.nodes) == 1 + assert len(planner.goal_tree.edges) == 0 + + +def test_plan_rrt_connect(): + model, collision_model, visual_model = load_models() + add_self_collisions(collision_model) + add_object_collisions(collision_model, visual_model) + + # Plan with RRTConnect and bidirectional RRT + q_start = np.array([0.0, 1.57, 0.0, 0.0, 1.57, 1.57, 0.0, 0.0, 0.0]) + q_goal = np.array([1.57, 1.57, 0.0, 0.0, 1.57, 1.57, 0.0, 0.0, 0.0]) + + options = RRTPlannerOptions() + options.rrt_connect = True + options.bidirectional_rrt = True + + planner = RRTPlanner(model, collision_model) + path = planner.plan(q_start, q_goal, options=options) + + # The path must exist and have more than + assert path is not None + assert len(path) > 2 + assert np.all(path[0] == q_start) + assert np.all(path[-1] == q_goal) + + # RRTConnect should have nodes and edges for both trees + assert len(planner.start_tree.nodes) > 1 + assert len(planner.start_tree.edges) > 1 + assert len(planner.goal_tree.nodes) > 1 + assert len(planner.goal_tree.edges) > 1