Add basic RRT implementation (#10)

examples/example_collision_along_path.py

@@ -126,4 +126,4 @@ def prepare_collision_scene(collision_model):
         )
 
         viz.display(q)
-        time.sleep(0.2)
+        time.sleep(0.1)

examples/example_rrt.py

@@ -0,0 +1,122 @@
+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.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 = pinocchio.buildModelsFromUrdf(
+        urdf_filename, package_dirs=pinocchio_model_dir
+    )
+    data = model.createData()
+
+    prepare_scene(visual_model, collision_model)
+
+    # 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)
+    viz.display(q_start)
+    time.sleep(1.0)
+
+    # Search for a path
+    options = RRTPlannerOptions()
+    options.max_angle_step = 0.05
+    options.max_connection_dist = 0.25
+    options.goal_biasing_probability = 0.15
+
+    planner = RRTPlanner(model, collision_model)
+    path = planner.plan(q_start, q_end, options=options)
+    planner.visualize(viz, "panda_hand", show_tree=True)
+
+    # Animate the path
+    input("Press 'Enter' to animate the path.")
+    for idx in range(1, len(path)):
+        segment_start = path[idx - 1]
+        segment_end = path[idx]
+        q_path = discretize_joint_space_path(
+            segment_start, segment_end, options.max_angle_step
+        )
+        for q in q_path:
+            viz.display(q)
+            time.sleep(0.05)

src/pyroboplan/core/utils.py

@@ -5,6 +5,85 @@
 import pinocchio
 
 
+def check_collisions_at_state(model, collision_model, q):
+    """
+    Checks whether a specified joint configuration is collision-free.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to generate a random state.
+        collision_model : `pinocchio.Model`
+            The model to use for collision checking.
+        q : array-like
+            The joint configuration of the model.
+
+    Returns
+    -------
+        bool
+            True is there are any collisions, otherwise False.
+    """
+    data = model.createData()
+    collision_data = collision_model.createData()
+    stop_at_first_collision = True  # For faster computation
+
+    pinocchio.computeCollisions(
+        model, data, collision_model, collision_data, q, stop_at_first_collision
+    )
+    return np.any([cr.isCollision() for cr in collision_data.collisionResults])
+
+
+def check_collisions_along_path(model, collision_model, q_path):
+    """
+    Checks whether a path consisting of multiple joint configurations is collision-free.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to generate a random state.
+        collision_model : `pinocchio.Model`
+            The model to use for collision checking.
+        q_path : list[array-like]
+            A list of joint configurations describing the path.
+
+    Returns
+    -------
+        bool
+            True is there are any collisions, otherwise False.
+    """
+    data = model.createData()
+    collision_data = collision_model.createData()
+    stop_at_first_collision = True  # For faster computation
+
+    for q in q_path:
+        pinocchio.computeCollisions(
+            model, data, collision_model, collision_data, q, stop_at_first_collision
+        )
+        if np.any([cr.isCollision() for cr in collision_data.collisionResults]):
+            return True
+
+    return False
+
+
+def configuration_distance(q_start, q_end):
+    """
+    Returns the distance between two joint configurations.
+
+    Parameters
+    ----------
+        q_start : array-like
+            The start joint configuration.
+        q_end : array-like
+            The end joint configuration.
+
+    Returns
+    -------
+        float
+            The distance between the two joint configurations.
+    """
+    return np.linalg.norm(q_end - q_start)
+
+
 def get_random_state(model, padding=0.0):
     """
     Returns a random state that is within the model's joint limits.
@@ -26,6 +105,37 @@ def get_random_state(model, padding=0.0):
     )
 
 
+def get_random_collision_free_state(model, collision_model, padding=0.0, max_tries=100):
+    """
+    Returns a random state that is within the model's joint limits and is collision-free according to the collision model.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to generate a random state.
+        collision_model : `pinocchio.Model`
+            The model to use for collision checking.
+        padding : float or array-like, optional
+            The padding to use around the sampled joint limits.
+        max_tries : int, optional
+            The maximum number of tries for sampling a collision-free state.
+
+    Returns
+    -------
+        array-like
+            A set of randomly generated collision-free joint variables, or None if one cannot be found.
+    """
+    num_tries = 0
+    while num_tries < max_tries:
+        state = get_random_state(model, padding=padding)
+        if not check_collisions_at_state(model, collision_model, state):
+            return state
+        num_tries += 1
+
+    print(f"Could not generate collision-free state after {max_tries} tries.")
+    return None
+
+
 def get_random_transform(model, target_frame, joint_padding=0.0):
     """
     Returns a random transform for a target frame that is within the model's joint limits.

src/pyroboplan/planning/graph.py

@@ -0,0 +1,109 @@
+import numpy as np
+
+from ..core.utils import configuration_distance
+
+
+class Node:
+    """Describes an individual node in a graph."""
+
+    def __init__(self, q, parent=None, cost=0.0):
+        """
+        Creates a graph node instance.
+
+        Parameters
+        ----------
+            q : array-like
+                The joint configuration associated with the node.
+            parent : `pyroboplan.planning.graph.Node`, optional.
+                The parent of the node, which is either None or another Node object.
+            cost : float, optional
+                The cost associated with the node.
+        """
+        self.q = q
+        self.parent = parent
+        self.cost = cost
+        self.neighbors = set()
+
+
+class Edge:
+    """Describes an individual edge in a graph."""
+
+    def __init__(self, nodeA, nodeB, cost=0.0):
+        """
+        Creates a graph edge instance.
+
+        Parameters
+        ----------
+            nodeA : `pyroboplan.planning.graph.Node`
+                The first node in the edge.
+            nodeB : `pyroboplan.planning.graph.Node`
+                The second node in the edge.
+            cost : float, optional
+                The cost associated with the edge.
+        """
+        self.nodeA = nodeA
+        self.nodeB = nodeB
+        self.cost = cost
+
+
+class Graph:
+    """Describes a graph of robot configuration states for motion planning."""
+
+    def __init__(self):
+        """
+        Creates a graph instance with no edges or nodes.
+        """
+        self.nodes = set()
+        self.edges = set()
+
+    def add_node(self, node):
+        """
+        Adds a node to the graph.
+
+        Parameters
+        ----------
+            node : `pyroboplan.planning.graph.Node`
+                The node to add to the graph.
+        """
+        self.nodes.add(node)
+
+    def add_edge(self, nodeA, nodeB):
+        """
+        Adds an edge to the graph.
+
+        Parameters
+        ----------
+            nodeA : `pyroboplan.planning.graph.Node`
+                The first node in the edge.
+            nodeB : `pyroboplan.planning.graph.Node`
+                The second node in the edge.
+        """
+        cost = configuration_distance(nodeA.q, nodeB.q)
+        self.edges.add(Edge(nodeA, nodeB, cost))
+        nodeA.neighbors.add(nodeB)
+        nodeB.neighbors.add(nodeA)
+
+    def get_nearest_node(self, q):
+        """
+        Gets the nearest node to a specifed robot configuration.
+
+        Parameters
+        ----------
+            q : array-like
+                The robot configuration to use in this query.
+
+        Returns
+        -------
+            `pyroboplan.planning.graph.Node` or None
+                The nearest node to the input configuration, or None if the graph is empty.
+        """
+        nearest_node = None
+        min_dist = np.inf
+
+        for node in self.nodes:
+            dist = configuration_distance(q, node.q)
+            if dist < min_dist:
+                min_dist = dist
+                nearest_node = node
+
+        return nearest_node