added plotting probability maps

marl_factory_grid/utils/renderer.py

@@ -1,7 +1,7 @@
 import sys
 
 from pathlib import Path
-from collections import deque
+from collections import deque, defaultdict
 from itertools import product
 
 import numpy as np
@@ -235,7 +235,7 @@ class Renderer:
         return np.transpose(rgb_obs, (2, 0, 1))
         # return torch.from_numpy(rgb_obs).permute(2, 0, 1)
 
-    def render_action_icons(self, action_entities):
+    def render_single_action_icons(self, action_entities):
         """
         Renders action icons based on the entities' specified actions' name, position, rotation and probability.
         Renders probabilities unequal 0.
@@ -244,7 +244,6 @@ class Renderer:
         :type action_entities: List[RenderEntity]
         """
         self.fill_bg()
-        font = pygame.font.Font(None, 24)  # Initialize the font once for all text rendering
 
         for action_entity in action_entities:
             if not isinstance(action_entity.pos, np.ndarray) or action_entity.pos.ndim != 1:
@@ -266,13 +265,74 @@ class Renderer:
 
             # Render the probability next to the icon if it exists
             if hasattr(action_entity, 'probability') and action_entity.probability != 0:
-                prob_text = font.render(f"{action_entity.probability:.2f}", True, (255, 0, 0))
+                prob_text = self.font.render(f"{action_entity.probability:.2f}", True, (255, 0, 0))
                 prob_text_rect = prob_text.get_rect(top=img_rect.bottom, left=img_rect.left)
                 self.screen.blit(prob_text, prob_text_rect)
 
         pygame.display.flip()  # Update the display with all new blits
         self.save_screen("route_graph")
 
+    def render_multi_action_icons(self, action_entities):
+        """
+        Renders multiple action icons at the same position without overlap and arranges them based on direction, except for
+        walls which cover the entire grid cell.
+        """
+        self.fill_bg()
+        font = pygame.font.Font(None, 18)
+
+        # prepare position dict to iterate over
+        position_dict = defaultdict(list)
+        for entity in action_entities:
+            position_dict[tuple(entity.pos)].append(entity)
+
+        for position, entities in position_dict.items():
+            num_entities = len(entities)
+            entity_size = self.cell_size // 2  # Adjust size to fit multiple entities for non-wall entities
+
+            # Define offsets for each direction based on a quadrant layout within the cell
+            offsets = {
+                0: (-entity_size // 2, -entity_size // 2),  # North
+                90: (-entity_size // 2, entity_size // 2),  # East
+                180: (entity_size // 2, entity_size // 2),  # South
+                270: (entity_size // 2, -entity_size // 2)  # West
+            }
+
+            # Sort entities based on direction to ensure consistent positioning
+            entities.sort(key=lambda x: x.rotation)
+
+            for entity in entities:
+                img = self.assets[entity.name.lower()]
+                if img is None:
+                    print(f"Error: No asset available for '{entity.name}'. Skipping rendering this entity.")
+                    continue
+
+                img = pygame.transform.rotate(img, entity.rotation)
+
+                # Check if the entity is a wall and adjust the size and position accordingly
+                if entity.name == 'wall':
+                    img = pygame.transform.scale(img, (self.cell_size, self.cell_size))
+                    img_rect = img.get_rect(center=(position[0] * self.cell_size + self.cell_size // 2,
+                                                    position[1] * self.cell_size + self.cell_size // 2))
+                else:
+                    img = pygame.transform.scale(img, (entity_size, entity_size))  # Scale down the image for arrows
+                    offset = offsets.get(entity.rotation, (0, 0))
+                    img_rect = img.get_rect(center=(
+                        position[0] * self.cell_size + self.cell_size // 2 + offset[0],
+                        position[1] * self.cell_size + self.cell_size // 2 + offset[1]
+                    ))
+
+                self.screen.blit(img, img_rect)
+
+                # Render the probability next to the icon if it exists and is non-zero
+                if entity.probability > 0 and entity.name != 'wall':
+                    formatted_probability = f"{entity.probability:.4f}"
+                    prob_text = font.render(formatted_probability, True, (0, 0, 0))  # Black color for readability
+                    prob_text_rect = prob_text.get_rect(center=img_rect.center)  # Center text on the arrow
+                    self.screen.blit(prob_text, prob_text_rect)
+
+        pygame.display.flip()  # Update the display
+        self.save_screen("multi_action_graph")
+
     def save_screen(self, filename):
         """
         Saves the current screen to a PNG file, appending a counter to ensure uniqueness.