started visualization of routes in plot single runs, assets missing.

marl_factory_grid/utils/plotting/plot_single_runs.py

@@ -3,11 +3,15 @@ from os import PathLike
 from pathlib import Path
 from typing import Union
 
+import numpy as np
 import pandas as pd
 
 from marl_factory_grid.utils.helpers import IGNORED_DF_COLUMNS
 from marl_factory_grid.utils.plotting.plotting_utils import prepare_plot
 
+from marl_factory_grid.utils.renderer import Renderer
+from marl_factory_grid.utils.utility_classes import RenderEntity
+
 
 def plot_single_run(run_path: Union[str, PathLike], use_tex: bool = False, column_keys=None,
                     file_key: str = 'monitor', file_ext: str = 'pkl'):
@@ -60,3 +64,81 @@ def plot_single_run(run_path: Union[str, PathLike], use_tex: bool = False, colum
 
     prepare_plot(run_path.parent / f'{run_path.parent.name}_monitor_lineplot.png', df_melted, use_tex=use_tex)
     print('Plotting done.')
+
+
+rotation_mapping = {
+    'north': ('cardinal', 0),
+    'east': ('cardinal', 270),
+    'south': ('cardinal', 180),
+    'west': ('cardinal', 90),
+    'northeast': ('diagonal', 0),
+    'southeast': ('diagonal', 270),
+    'southwest': ('diagonal', 180),
+    'northwest': ('diagonal', 90)
+}
+
+
+def plot_routes(factory, agents):
+    renderer = Renderer(factory.map.level_shape, custom_assets_path={
+        'cardinal': 'marl_factory_grid/utils/plotting/action_assets/cardinal.png',
+        'diagonal': 'marl_factory_grid/utils/plotting/action_assets/diagonal.png',
+        'door': 'marl_factory_grid/utils/plotting/action_assets/door.png',
+        'wall': 'marl_factory_grid/environment/assets/wall.png'})
+
+    wall_positions = factory.map.walls
+
+    for index, agent in enumerate(agents):
+        action_entities = []
+        # Add walls to the action_entities list
+        for pos in wall_positions:
+            wall_entity = RenderEntity(
+                name='wall',
+                probability=1.0,
+                pos=np.array(pos),
+            )
+            action_entities.append(wall_entity)
+        current_position = agent.spawn_position
+
+        if hasattr(agent, 'action_probabilities'):
+            # Handle RL agents with action probabilities
+            top_actions = sorted(agent.action_probabilities.items(), key=lambda x: -x[1])[:4]
+        else:
+            # Handle deterministic agents by iterating through all actions in the list
+            top_actions = [(action, 1.0) for action in agent.action_list]
+
+        for action, probability in top_actions:
+            base_icon, rotation = rotation_mapping.get(action.lower(), ('north', 0))
+            icon_name = base_icon
+            new_position = action_to_coords(current_position, action.lower())
+            print(f"current position type and value: {type(current_position)}, {new_position}")
+
+            action_entity = RenderEntity(
+                name=icon_name,
+                pos=np.array(current_position),
+                probability=probability,
+                rotation=rotation
+            )
+            action_entities.append(action_entity)
+            current_position = new_position
+
+        renderer.render_action_icons(action_entities)
+
+
+def action_to_coords(current_position, action):
+    direction_mapping = {
+        'north': (0, -1),
+        'south': (0, 1),
+        'east': (1, 0),
+        'west': (-1, 0),
+        'northeast': (1, -1),
+        'northwest': (-1, -1),
+        'southeast': (1, 1),
+        'southwest': (-1, 1)
+    }
+    delta = direction_mapping.get(action)
+    if delta is not None:
+        new_position = [current_position[0] + delta[0], current_position[1] + delta[1]]
+        return new_position
+    print(f"No valid action found for {action}.")
+    return current_position
+