Rrt connect

.gitignore

@@ -0,0 +1,11 @@
+gpt_rrt.py
+pybullet/gettingStarted.py
+pybullet/__pycache__/gen3lite_controller_collision_detection.cpython-313.pyc
+pybullet/__pycache__/gen3lite_controller_collision_detection.cpython-313.pyc
+pybullet/__pycache__/gen3lite_controller_collision_detection.cpython-313.pyc
+pybullet/hardest_home.npy
+.gitignore
+pybullet/davespath.npy
+.gitignore
+pybullet/__pycache__/gen3lite_controller_collision_detection.cpython-313.pyc
+pybullet/__pycache__/gen3lite_controller_collision_detection.cpython-38.pyc

pybullet/arm_rrt.py

@@ -1,161 +1,154 @@
 import numpy as np
 from tqdm import tqdm
 from scipy.spatial import cKDTree
-import matplotlib.pyplot as plt
 import random
-import math
-import gen3lite_controller_collision_detection
+import argparse
+import robots
 
-# -----------------------------
-# RRT Node
-# -----------------------------
 class Node:
-    def __init__(self, point):
+    def __init__(self, point, parent=None):
         self.point = np.array(point)
-        self.parent = None
+        self.parent = parent
+
+class Tree:
+    def __init__(self, node_list, kdtree):
+        self.node_list = node_list
+        self.kdtree = kdtree
+
+    def add(self,new_point):
+        self.node_list.append(new_point)
+
+        if len(self.node_list) % 50 == 0:
+            data = [n.point for n in self.node_list]
+            self.kdtree = cKDTree(data)
 
-# -----------------------------
-# RRT Planner
-# -----------------------------
 class RRT:
-    def __init__(
-        self,
-        start=None,
-        goal=None,
-        obstacles=None,
-        rand_area=[],
-        step_size=0.1,
-        goal_sample_rate=0.1,
-        max_iter=5000
-    ):
-        self.controller = gen3lite_controller_collision_detection.Gen3LiteArmController()
-        self.controller.createBalloonMaze()
+    def __init__(self, step_size=0.1,max_iter=20,env_name="Free"):
+        self.controller = robots.Gen3LiteArmController(env_name=env_name)
 
         self.start = Node(self.controller.getCurrentJointAngles())
-        self.goal = Node([0.1026325022237283, -0.2931188624740633, 1.2717083400432991, 0.048794139164578594, 0.07744723004754135, -0.8437927483158898, -0.024709326684397483])
+        self.goal = Node(self.controller.goal_angles)
         self.rand_ranges = self.controller.getRanges()
         self.step_size = step_size
-        self.goal_sample_rate = goal_sample_rate
         self.max_iter = max_iter
-        self.node_list = [self.start]
-        self.kdtree = cKDTree([self.start.point])
-
-    # -----------------------------
-    # Main planning loop
-    # -----------------------------
+        self.start_tree = Tree([self.start],cKDTree([self.start.point]))
+        self.goal_tree = Tree([self.goal],cKDTree([self.goal.point]))
+        self.path_to_goal = []
+
+    # This a "stub" planning function. It adds random nodes always
+    # to the root of a start and goal tree. It shows many of the 
+    # syntax elements and helper functions you could use, but 
+    # you will have to fix and extend this to make it an RRT or 
+    # RRT-Connect planner that can solve the harder environments.
     def plan(self):
 
-        for _ in tqdm(range(self.max_iter)):
-
+        # Add max_iter nodes to each start and goal tree
+        for k in tqdm(range(self.max_iter)):
             rnd_point = self.sample()
-            nearest_node = self.nearest_node(rnd_point)
-            new_node = self.steer(nearest_node, rnd_point)
+            if self.collision_free(rnd_point, self.start.point):
+                new_node = Node(rnd_point,self.start)
+                self.start_tree.add(new_node)
 
-            if self.collision_free(nearest_node.point, new_node.point):
-                self.node_list.append(new_node)
+            rnd_point = self.sample()
+            if self.collision_free(rnd_point, self.goal.point):
+                new_node = Node(rnd_point,self.goal)
+                self.goal_tree.add(new_node)
 
-                if len(self.node_list) % 50 == 0:
-                    data = [n.point for n in self.node_list]
-                    self.kdtree = cKDTree(data)
+        # Now add the same rnd_pnt to both trees, making a connection
+        # and completing the path start->goal
+        rnd_point = self.sample()   
+        if self.collision_free(rnd_point, self.start.point):
+            new_start_node = Node(rnd_point,self.start)
+            self.start_tree.add(new_start_node)
 
-                if self.reached_goal(new_node):
-                    p = self.extract_path(new_node)
-                    self.controller.execPath(p)
-                    return p
+        if self.collision_free(rnd_point, self.goal.point):
+            new_goal_node = Node(rnd_point,self.goal)
+            self.goal_tree.add(new_goal_node)
 
-        return None  # Failed
+        self.path_to_goal = self.extract_path(new_start_node,new_goal_node)
+        return True
 
-    # -----------------------------
-    # Sampling
-    # -----------------------------
     def sample(self):
-        if random.random() < self.goal_sample_rate:
-            return self.goal.point
         point = []
         for i in range(0,len(self.rand_ranges[0])):
             point.append(random.uniform(self.rand_ranges[0][i], self.rand_ranges[1][i]))
         return np.array(point)
 
-    # -----------------------------
-    # Nearest node
-    # -----------------------------
-    def nearest_node(self, point):
-         _, idx = self.kdtree.query(point)
-         return self.node_list[idx]
+    def nearest_node(self, point, tree):
+         _, idx = tree.kdtree.query(point)
+         return tree.node_list[idx]
 
-    # -----------------------------
-    # Steer
-    # -----------------------------
     def steer(self, from_node, to_point):
         direction = to_point - from_node.point
         distance = np.linalg.norm(direction)
-        direction = direction / distance
 
-        new_point = from_node.point + self.step_size * direction
+        if distance < self.step_size:
+            new_point = to_point
+        else:
+            direction = direction / distance
+            new_point = from_node.point + self.step_size * direction
+
         new_node = Node(new_point)
         new_node.parent = from_node
         return new_node
 
-    # -----------------------------
-    # Collision checking
-    # -----------------------------
     def collision_free(self, p1, p2):
         return self.controller.collision_free(p1,p2)
 
-    # -----------------------------
-    # Goal check
-    # -----------------------------
-    def reached_goal(self, node):
-        return np.linalg.norm(node.point - self.goal.point) < self.step_size
-
-    # -----------------------------
-    # Path extraction
-    # -----------------------------
-    def extract_path(self, node):
-        path = [self.goal.point]
-        while node is not None:
-            path.append(node.point)
-            node = node.parent
-        return path[::-1]
-
-    # -----------------------------
-    # Visualization
-    # -----------------------------
-    def draw(self, path=None):
-        plt.figure()
-        for node in self.node_list:
-            if node.parent:
-                plt.plot(
-                    [node.point[0], node.parent.point[0]],
-                    [node.point[1], node.parent.point[1]],
-                    "-g"
-                )
-
-        for (ox, oy, r) in self.obstacles:
-            circle = plt.Circle((ox, oy), r, color="r")
-            plt.gca().add_patch(circle)
-
-        plt.plot(self.start.point[0], self.start.point[1], "bo", label="Start")
-        plt.plot(self.goal.point[0], self.goal.point[1], "ro", label="Goal")
-
-        if path:
-            px, py = zip(*path)
-            plt.plot(px, py, "-b", linewidth=2, label="Path")
-
-        plt.axis("equal")
-        plt.grid(True)
-        plt.legend()
-        plt.show()
-
-
-# -----------------------------
-# Example Usage
-# -----------------------------
-if __name__ == "__main__":
+    def reached_goal(self, node,goal=None):
+        if goal is None:
+            goal = self.goal
+
+        return np.linalg.norm(node.point - goal.point) < self.step_size and self.collision_free(node.point,goal.point)
+
+    def extract_path(self, start_node,goal_node):
+        # Build the start-tree path, leaf to root (backwards) 
+        start_tree_path = []
+        while start_node is not None:
+            start_tree_path.append(start_node.point)
+            start_node = start_node.parent
+
+        # Build the goal-tree path, leaf to root (forwards)
+        goal_tree_path = []
+        while goal_node is not None:
+            goal_tree_path.append(goal_node.point)
+            goal_node = goal_node.parent 
+
+        # First add the start path, reversing
+        overall_path = start_tree_path[::-1]
+        # Add the goal path
+        overall_path.extend(goal_tree_path)
+        return overall_path
 
-    rrt = RRT()
+if __name__ == "__main__":
 
-    path = rrt.plan()
-    print("Planning complete.")
-    print(path)
-    #rrt.draw(path)
+    parser = argparse.ArgumentParser(
+                    prog='arm_rrt',
+                    description='Plans and executes paths for arms around obstacles.')
+    parser.add_argument('--filename',default='rrt_path.npy')           
+    parser.add_argument('-e', '--environment',default='middle')
+    parser.add_argument('-p', '--plan',action='store_true')
+    parser.add_argument('-r', '--run',action='store_true')
+    args = parser.parse_args()
+
+    rrt = RRT(env_name=args.environment)
+
+    if args.plan:
+        success = rrt.plan()
+
+        if success:
+            print("Tree planning reached the goal.")
+            np.save(args.filename,rrt.path_to_goal)
+        else:
+            print("Failed to find a path to the goal.")
+
+        rrt.controller.visTreesAndPaths( 
+            [rrt.start_tree,rrt.goal_tree],             
+            [rrt.path_to_goal],
+            rgbas_in=[[0.5,0.0,0.5,1.0],[0.902,0.106,0.714,1.0]]
+        )
+
+    if args.run:
+        path_to_goal = np.load(args.filename)
+        rrt.controller.execPath(path_to_goal)
+