Adds visualization and simple policy step to getting started docs on custom MARL env

docs/tutorials/custom_environment/6-adding-visualization.md

@@ -0,0 +1,40 @@
+---
+title: "Tutorial: Adding visualization with pygame and a simple random policy"
+---
+
+# Tutorial: Adding visualization with pygame and a simple random policy
+
+## Introduction
+
+The standard library for adding graphics to your environments is [pygame](https://www.pygame.org/).
+
+Here we present a simple visualization tool to see the guard and the prisoner in action.
+
+They both have a very simple policy: they're just taking random actions. Occasionaly you can see that the prisoner is able to scape or tha guard is able to caught them. You can modify the amount of steps you want to run the simulation for and also write your own policy.
+
+You can add the following file in the root of your directory. If you name it `visualization.py`, this is how it will look:
+
+    Custom-Environment
+    ├── custom-environment
+        └── env
+            └── custom_environment.py
+        └── custom_environment_v0.py
+    ├── README.md
+    └── requirements.txt
+    └── visualization.py
+
+## Setup and execution
+
+To run the visualization make sure you have `pygame` installed for the graphics and `numpy` to generate the random policy.
+
+You can add both of these packages with pip or conda. `pip install pygame numpy` or `conda install pygame numpy`.
+
+Now you can run the code with `python3 visualization.py`.
+
+## Code
+
+```{eval-rst}
+.. literalinclude:: ../../../tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py
+   :language: python
+   :caption: /custom-environment/env/custom_environment.py
+```

docs/tutorials/custom_environment/index.md

@@ -14,6 +14,8 @@ These tutorials walk you though the full process of creating a custom environmen
 
 4. [Testing Your Environment](/tutorials/custom_environment/4-testing-your-environment.md)
 
+5. [Adding visualization with pygame and a simple random policy](/tutorials/custom_environment/6-adding-visualization.md)
+
 For a simpler example environment, including both [AEC](/api/aec/) and [Parallel](/api/aec/) implementations, see our [Environment Creation](/content/environment_creation/) documentation.
 
 
@@ -25,4 +27,5 @@ For a simpler example environment, including both [AEC](/api/aec/) and [Parallel
 2-environment-logic
 3-action-masking
 4-testing-your-environment
+6-adding-visualization
 ```

tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py

@@ -0,0 +1,109 @@
+import random
+
+import numpy as np
+import pygame
+from custom_environment import CustomActionMaskedEnvironment, CustomEnvironment
+
+
+class PygameVisualizer:
+    def __init__(self, env, width=700, height=700):
+        pygame.init()
+        self.env = env
+        self.screen = pygame.display.set_mode((width, height))
+        self.cell_size = width // 7
+        self.colors = {
+            "prisoner": (255, 0, 0),  # Red
+            "guard": (0, 0, 255),  # Blue
+            "escape": (0, 255, 0),  # Green
+            "background": (255, 255, 255),  # White
+        }
+
+    def draw_grid(self):
+        self.screen.fill(self.colors["background"])
+        for y in range(7):
+            for x in range(7):
+                rect = pygame.Rect(
+                    x * self.cell_size,
+                    y * self.cell_size,
+                    self.cell_size,
+                    self.cell_size,
+                )
+                pygame.draw.rect(self.screen, (0, 0, 0), rect, 1)
+
+        # Draw the agents
+        p_x, p_y = self.env.prisoner_x, self.env.prisoner_y
+        g_x, g_y = self.env.guard_x, self.env.guard_y
+        e_x, e_y = self.env.escape_x, self.env.escape_y
+
+        pygame.draw.rect(
+            self.screen,
+            self.colors["prisoner"],
+            (
+                p_x * self.cell_size,
+                p_y * self.cell_size,
+                self.cell_size,
+                self.cell_size,
+            ),
+        )
+        pygame.draw.rect(
+            self.screen,
+            self.colors["guard"],
+            (
+                g_x * self.cell_size,
+                g_y * self.cell_size,
+                self.cell_size,
+                self.cell_size,
+            ),
+        )
+        pygame.draw.rect(
+            self.screen,
+            self.colors["escape"],
+            (
+                e_x * self.cell_size,
+                e_y * self.cell_size,
+                self.cell_size,
+                self.cell_size,
+            ),
+        )
+
+    def run(self, steps):
+        running = True
+        while steps > 0:
+            # Both agents take random actions
+            actions = {
+                agent: self.env.action_space(agent).sample()
+                for agent in self.env.agents
+            }
+
+            # Step the environment
+            observations, rewards, terminations, truncations, infos = self.env.step(
+                actions
+            )
+
+            # Draw the grid and agents
+            self.draw_grid()
+
+            pygame.display.flip()
+            pygame.time.delay(100)
+
+            steps -= 1
+
+            # Check for termination conditions
+            if any(terminations.values()):
+                if rewards["prisoner"] == 1:
+                    print("Prisoner escaped!")
+                elif rewards["prisoner"] == -1:
+                    print("Prisoner caught by the guard!")
+                return
+
+        print("This game has ended because the specified number of steps was reached.")
+        return
+
+
+# Initialize your environment
+env = CustomActionMaskedEnvironment()
+env.reset()
+
+visualizer = PygameVisualizer(env)
+# You can change the number of steps
+visualizer.run(100)