From 6745fd292513a5cd38807944101706aa9233d398 Mon Sep 17 00:00:00 2001 From: marimeireles Date: Thu, 8 Feb 2024 13:58:29 +0700 Subject: [PATCH 1/2] Adds visualization and simple policy step to getting started docs on custom MARL env --- .../6-adding-visualization.md | 40 +++++++ docs/tutorials/custom_environment/index.md | 3 + ...torial6_adding_visualization_and_policy.py | 109 ++++++++++++++++++ 3 files changed, 152 insertions(+) create mode 100644 docs/tutorials/custom_environment/6-adding-visualization.md create mode 100644 tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py diff --git a/docs/tutorials/custom_environment/6-adding-visualization.md b/docs/tutorials/custom_environment/6-adding-visualization.md new file mode 100644 index 000000000..fe341d991 --- /dev/null +++ b/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 +``` \ No newline at end of file diff --git a/docs/tutorials/custom_environment/index.md b/docs/tutorials/custom_environment/index.md index 77c9a8b5e..6a8fffc07 100644 --- a/docs/tutorials/custom_environment/index.md +++ b/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 ``` diff --git a/tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py b/tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py new file mode 100644 index 000000000..50885aadb --- /dev/null +++ b/tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py @@ -0,0 +1,109 @@ +import pygame +import numpy as np +import random + +from custom_environment import CustomEnvironment, CustomActionMaskedEnvironment + + +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) From b94006066dfd1f2e98118da037b730382513cb1a Mon Sep 17 00:00:00 2001 From: marimeireles Date: Thu, 8 Feb 2024 14:10:44 +0700 Subject: [PATCH 2/2] Run pre-commit --- docs/tutorials/custom_environment/6-adding-visualization.md | 2 +- .../tutorial6_adding_visualization_and_policy.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/tutorials/custom_environment/6-adding-visualization.md b/docs/tutorials/custom_environment/6-adding-visualization.md index fe341d991..34e4fe50c 100644 --- a/docs/tutorials/custom_environment/6-adding-visualization.md +++ b/docs/tutorials/custom_environment/6-adding-visualization.md @@ -37,4 +37,4 @@ Now you can run the code with `python3 visualization.py`. .. literalinclude:: ../../../tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py :language: python :caption: /custom-environment/env/custom_environment.py -``` \ No newline at end of file +``` diff --git a/tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py b/tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py index 50885aadb..b1bbc490d 100644 --- a/tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py +++ b/tutorials/CustomEnvironment/tutorial6_adding_visualization_and_policy.py @@ -1,8 +1,8 @@ -import pygame -import numpy as np import random -from custom_environment import CustomEnvironment, CustomActionMaskedEnvironment +import numpy as np +import pygame +from custom_environment import CustomActionMaskedEnvironment, CustomEnvironment class PygameVisualizer: