cleaner return type from .plan()

Author: SchwartzCode

PathPlanning/TimeBasedPathPlanning/BaseClasses.py

@@ -43,8 +43,8 @@ class MultiAgentPlanner(ABC):
 
     @staticmethod
     @abstractmethod
-    def plan(grid: Grid, start_and_goal_positions: list[StartAndGoal], single_agent_planner_class: SingleAgentPlanner, verbose: bool = False) -> tuple[list[StartAndGoal], list[NodePath]]:
+    def plan(grid: Grid, start_and_goal_positions: list[StartAndGoal], single_agent_planner_class: SingleAgentPlanner, verbose: bool = False) -> dict[AgentId, NodePath]:
         """
-        Plan for all agents. Returned paths are in order corresponding to the returned list of `StartAndGoal` objects
+        Plan for all agents. Returned paths found for each agent
         """
         pass

PathPlanning/TimeBasedPathPlanning/ConflictBasedSearch.py

@@ -34,11 +34,9 @@ class ConflictBasedSearch(MultiAgentPlanner):
 
 
     @staticmethod
-    def plan(grid: Grid, start_and_goals: list[StartAndGoal], single_agent_planner_class: SingleAgentPlanner, verbose: bool = False) -> tuple[list[StartAndGoal], list[NodePath]]:
+    def plan(grid: Grid, start_and_goals: list[StartAndGoal], single_agent_planner_class: SingleAgentPlanner, verbose: bool = False) -> dict[AgentId, NodePath]:
         """
         Generate a path from the start to the goal for each agent in the `start_and_goals` list.
-        Returns the re-ordered StartAndGoal combinations, and a list of path plans. The order of the plans
-        corresponds to the order of the `start_and_goals` list.
         """
         print(f"Using single-agent planner: {single_agent_planner_class}")
 
@@ -72,7 +70,10 @@ def plan(grid: Grid, start_and_goals: list[StartAndGoal], single_agent_planner_c
                 print(f"\nFound a path with constraints after {constraint_tree.expanded_node_count()} expansions")
                 print(f"Final cost: {constraint_tree_node.cost}")
                 print(f"Number of constraints on solution: {len(constraint_tree_node.all_constraints)}")
-                return (start_and_goals, [constraint_tree_node.paths[start_and_goal.agent_id] for start_and_goal in start_and_goals])
+                final_paths = {}
+                for start_and_goal in start_and_goals:
+                    final_paths[start_and_goal.agent_id] = constraint_tree_node.paths[start_and_goal.agent_id]
+                return final_paths
 
             if not isinstance(constraint_tree_node.constraint, ForkingConstraint):
                 raise ValueError(f"Expected a ForkingConstraint, but got: {constraint_tree_node.constraint}")
@@ -89,7 +90,6 @@ def plan(grid: Grid, start_and_goals: list[StartAndGoal], single_agent_planner_c
 
                 # Deepcopy to update with applied constraint and new paths
                 applied_constraint_parent = deepcopy(constraint_tree_node)
-                # Copy paths for child node - we just need to update constrained agent's path
                 applied_constraint_parent.paths[constrained_agent.agent] = new_path
 
                 if verbose:
@@ -163,6 +163,7 @@ class Scenario(Enum):
 scenario = Scenario.HALLWAY_CROSS
 verbose = False
 show_animation = True
+use_sipp = False # Condition here mainly to appease the linter
 np.random.seed(42)  # For reproducibility
 def main():
     grid_side_length = 21
@@ -187,16 +188,16 @@ def main():
         obstacle_arrangement=obstacle_arrangement,
     )
 
+    single_agent_planner = SafeIntervalPathPlanner if use_sipp else SpaceTimeAStar
     start_time = time.time()
-    start_and_goals, paths = ConflictBasedSearch.plan(grid, start_and_goals, SafeIntervalPathPlanner, verbose)
-    # start_and_goals, paths = ConflictBasedSearch.plan(grid, start_and_goals, SpaceTimeAStar, verbose)
+    paths = ConflictBasedSearch.plan(grid, start_and_goals, single_agent_planner, verbose)
 
     runtime = time.time() - start_time
     print(f"\nPlanning took: {runtime:.5f} seconds")
 
     if verbose:
         print(f"Paths:")
-        for path in paths:
+        for path in paths.values():
             print(f"{path}\n")
 
     if not show_animation:

PathPlanning/TimeBasedPathPlanning/Plotting.py

@@ -7,6 +7,7 @@
 )
 from PathPlanning.TimeBasedPathPlanning.BaseClasses import StartAndGoal
 from PathPlanning.TimeBasedPathPlanning.Node import NodePath
+from PathPlanning.TimeBasedPathPlanning.ConstraintTree import AgentId
 
 '''
 Plot a single agent path.
@@ -50,7 +51,7 @@ def PlotNodePath(grid: Grid, start: Position, goal: Position, path: NodePath):
 '''
 Plot a series of agent paths.
 '''
-def PlotNodePaths(grid: Grid, start_and_goals: list[StartAndGoal], paths: list[NodePath]):
+def PlotNodePaths(grid: Grid, start_and_goals: list[StartAndGoal], paths: dict[AgentId, NodePath]):
     fig = plt.figure(figsize=(10, 7))
 
     ax = fig.add_subplot(
@@ -64,19 +65,21 @@ def PlotNodePaths(grid: Grid, start_and_goals: list[StartAndGoal], paths: list[N
     ax.set_yticks(np.arange(0, grid.grid_size[1], 1))
 
     # Plot start and goal positions for each agent
-    colors = [] # generated randomly in loop
+    colors = {} # generated randomly in loop. Maps agent id to color
     markers = ['D', 's', '^', 'o', 'p']  # Different markers for visual distinction
 
     # Create plots for start and goal positions
     start_and_goal_plots = []
-    for i, path in enumerate(paths):
-        marker_idx = i % len(markers)
-        agent_id = start_and_goals[i].agent_id
-        start = start_and_goals[i].start
-        goal = start_and_goals[i].goal
+    for agent_id, path in paths.items():
+        marker_idx = agent_id % len(markers)
+        start_and_goal = next((elem for elem in start_and_goals if elem.agent_id == agent_id), None)
+        if not start_and_goal:
+            raise RuntimeError(f"Failed to get start and goal for agent {agent_id}")
+        start = start_and_goal.start
+        goal = start_and_goal.goal
 
         color = np.random.rand(3,)
-        colors.append(color)
+        colors[agent_id] = color
         sg_plot, = ax.plot([], [], markers[marker_idx], c=color, ms=15, 
                             label=f"Agent {agent_id} Start/Goal")
         sg_plot.set_data([start.x, goal.x], [start.y, goal.y])
@@ -86,12 +89,11 @@ def PlotNodePaths(grid: Grid, start_and_goals: list[StartAndGoal], paths: list[N
     (obs_points,) = ax.plot([], [], "ro", ms=15, label="Obstacles")
 
     # Create plots for each agent's path
-    path_plots = []
-    for i, path in enumerate(paths):
-        agent_id = start_and_goals[i].agent_id
-        path_plot, = ax.plot([], [], "o", c=colors[i], ms=10, 
+    path_plots = {}
+    for agent_id, path in paths.items():
+        path_plot, = ax.plot([], [], "o", c=colors[agent_id], ms=10, 
                             label=f"Agent {agent_id} Path")
-        path_plots.append(path_plot)
+        path_plots[agent_id] = path_plot
 
     ax.legend(bbox_to_anchor=(1.05, 1))
 
@@ -103,32 +105,32 @@ def PlotNodePaths(grid: Grid, start_and_goals: list[StartAndGoal], paths: list[N
     )
 
     # Find the maximum time across all paths
-    max_time = max(path.goal_reached_time() for path in paths)
+    max_time = max(path.goal_reached_time() for path in paths.values())
 
     # Animation loop
-    for i in range(0, max_time + 1):
+    for t in range(0, max_time + 1):
         # Update obstacle positions
-        obs_positions = grid.get_obstacle_positions_at_time(i)
+        obs_positions = grid.get_obstacle_positions_at_time(t)
         obs_points.set_data(obs_positions[0], obs_positions[1])
 
         # Update each agent's position
-        for (j, path) in enumerate(paths):
+        for agent_id, path in paths.items():
             path_positions = []
-            if i <= path.goal_reached_time():
-                res = path.get_position(i)
+            if t <= path.goal_reached_time():
+                res = path.get_position(t)
                 if not res:
-                    print(path)
-                    print(i)
-                path_position = path.get_position(i)
+                    print(f"Error getting position for agent {agent_id} at time {t}")
+                    print(t)
+                path_position = path.get_position(t)
                 if not path_position:
-                    raise Exception(f"Path position not found for time {i}.")
+                    raise Exception(f"Path position not found for time {t}.")
 
                 # Verify position is valid
                 assert not path_position in obs_positions
                 assert not path_position in path_positions
                 path_positions.append(path_position)
 
-                path_plots[j].set_data([path_position.x], [path_position.y])
+                path_plots[agent_id].set_data([path_position.x], [path_position.y])
 
         plt.pause(0.2)
 

PathPlanning/TimeBasedPathPlanning/PriorityBasedPlanner.py

@@ -18,12 +18,13 @@
 from PathPlanning.TimeBasedPathPlanning.BaseClasses import SingleAgentPlanner
 from PathPlanning.TimeBasedPathPlanning.SafeInterval import SafeIntervalPathPlanner
 from PathPlanning.TimeBasedPathPlanning.Plotting import PlotNodePaths
+from PathPlanning.TimeBasedPathPlanning.ConstraintTree import AgentId
 import time
 
 class PriorityBasedPlanner(MultiAgentPlanner):
 
     @staticmethod
-    def plan(grid: Grid, start_and_goals: list[StartAndGoal], single_agent_planner_class: SingleAgentPlanner, verbose: bool = False) -> tuple[list[StartAndGoal], list[NodePath]]:
+    def plan(grid: Grid, start_and_goals: list[StartAndGoal], single_agent_planner_class: SingleAgentPlanner, verbose: bool = False) -> dict[AgentId, NodePath]:
         """
         Generate a path from the start to the goal for each agent in the `start_and_goals` list.
         Returns the re-ordered StartAndGoal combinations, and a list of path plans. The order of the plans
@@ -40,7 +41,7 @@ def plan(grid: Grid, start_and_goals: list[StartAndGoal], single_agent_planner_c
                     key=lambda item: item.distance_start_to_goal(),
                     reverse=True)
 
-        paths = []
+        paths = {}
         for start_and_goal in start_and_goals:
             if verbose:
                 print(f"\nPlanning for agent:  {start_and_goal}" )
@@ -54,9 +55,9 @@ def plan(grid: Grid, start_and_goals: list[StartAndGoal], single_agent_planner_c
 
             agent_index = start_and_goal.agent_id
             grid.reserve_path(path, agent_index)
-            paths.append(path)
+            paths[start_and_goal.agent_id] =path
 
-        return (start_and_goals, paths)
+        return paths
 
 verbose = False
 show_animation = True
@@ -76,7 +77,7 @@ def main():
     )
 
     start_time = time.time()
-    start_and_goals, paths = PriorityBasedPlanner.plan(grid, start_and_goals, SafeIntervalPathPlanner, verbose)
+    paths = PriorityBasedPlanner.plan(grid, start_and_goals, SafeIntervalPathPlanner, verbose)
 
     runtime = time.time() - start_time
     print(f"\nPlanning took: {runtime:.5f} seconds")