save the whole RRT tree

guided_mrmp/main.py

@@ -49,15 +49,15 @@ def plan_decoupled_path(env, start, goal, solver="RRT*",
     """
     if solver == "RRT":
         rrt = RRT(env, start, goal, step_length, goal_sample_rate, num_samples)
-        path = rrt.run()
+        path, tree = rrt.run()
     elif solver == "RRT*":
         rrtstar = RRTStar(env, start, goal, step_length, goal_sample_rate, num_samples,r)
-        path = rrtstar.run()
+        path, tree = rrtstar.run()
     else:
         print(f"Solver {solver} is not yet implemented. Choose something else.")
         return None
 
-    return list(reversed(path))
+    return list(reversed(path)), tree
 
 def initialize_robots(starts, goals, dynamics_models, radii, target_v, env):
     """
@@ -72,7 +72,7 @@ def initialize_robots(starts, goals, dynamics_models, radii, target_v, env):
 
     for i, (start, goal, dynamics_model, radius, color) in enumerate(zip(starts, goals, dynamics_models, radii, colors)):
         # print(f"planning path for robot {i} from {start} to {goal} ")
-        rrtpath = plan_decoupled_path(env, (start[0],start[1]), (goal[0],goal[1]))
+        rrtpath, tree = plan_decoupled_path(env, (start[0],start[1]), (goal[0],goal[1]))
         xs = []
         ys = []
         for node in rrtpath:
@@ -86,7 +86,7 @@ def initialize_robots(starts, goals, dynamics_models, radii, target_v, env):
         start_heading = np.arctan2(goal[1] - start[1], goal[0] - start[0])
         start = [start[0], start[1], start_heading]
         
-        r = Robot(i,color,radius,start,goal,dynamics_model,target_v,rrtpath,waypoints)
+        r = Robot(i,color,radius,start,goal,dynamics_model,target_v,rrtpath,waypoints, tree)
         robots.append(r)
 
     return robots
@@ -148,8 +148,8 @@ if __name__ == "__main__":
     robots = initialize_robots(robot_starts, robot_goals, dynamics_models, robot_radii, target_v, env)
 
     # Load the libraries
-    # libs = initialize_libraries()
-    libs = [None, None, None]
+    libs = initialize_libraries()
+    # libs = [None, None, None]
 
     # Create the Guided MRMP policy
     T = settings['prediction_horizon']