Added pygame GUI, fixed import

overcooked_simulator/layouts/basic.layout

@@ -2,9 +2,9 @@ EEEEEEEEEEE
 ECCCCCCCCCE
 ECEEEEEEECE
 ECEEEEEEECE
-ECEEEEEEECE
-ECEEEEEEECE
-ECEEEEEEECE
+ECEEEEEEEEE
+ECEEEEEEEEE
+ECEEEEEEEEE
 ECEEEEEEECE
 ECEEEEEEECE
 ECCCCCCCCCE

overcooked_simulator/main.py

-from player import Player
+from overcooked_simulator.player import Player
 import sys
 from overcooked_simulator.simulation_runner import Simulator
 from overcooked_simulator.overcooked_environment import Environment

overcooked_simulator/pygame_gui.py

+import pygame
+import numpy as np
+from overcooked_simulator.overcooked_environment import Action
+from overcooked_simulator.simulation_runner import Simulator
+
+FPS = 60
+
+WHITE = (255, 255, 255)
+GREY = (190, 190, 190)
+BLACK = (0, 0, 0)
+COUNTERCOLOR = (240, 240, 240)
+LIGHTGREY = (220, 220, 220)
+GREEN = (0, 255, 0)
+RED = (255, 0, 0)
+BLUE = (0, 0, 255)
+YELLOW = (255, 255, 0)
+
+BACKGROUND_COLOR = GREY
+
+GET_CONTINOUS_INTERACT_AND_PICKUP = False
+
+
+class PlayerKeyset:
+    """Set of keyboard keys for controlling a player. First four keys are for movement, 5th key is for interacting
+    with counters. 6th key ist for picking up things or dropping them.
+
+    """
+    def __init__(self, keys: list[pygame.key]):
+        self.player_keys = keys
+        self.move_vectors = [[-1, 0], [1, 0], [0, -1], [0, 1]]
+        self.key_to_movement = {key: vec for (key, vec) in zip(self.player_keys[:-2], self.move_vectors)}
+        self.interact_key = self.player_keys[-2]
+        self.pickup_key = self.player_keys[-1]
+
+
+class PyGameGUI:
+    """Visualisation of the overcooked environmnent and reading keyboard inputs using pygame.
+    """
+    def __init__(self, simulator: Simulator):
+        self.simulator = simulator
+        self.counter_size = self.simulator.env.counter_side_length
+        self.WINDOW_WIDTH, self.WINDOW_HEIGHT = simulator.env.world_width, simulator.env.world_height
+
+        keys1 = [pygame.K_LEFT, pygame.K_RIGHT, pygame.K_UP, pygame.K_DOWN, pygame.K_SPACE, pygame.K_i]
+        keys2 = [pygame.K_a, pygame.K_d, pygame.K_w, pygame.K_s, pygame.K_f, pygame.K_e]
+        self.player_keysets: list[PlayerKeyset] = [PlayerKeyset(keys1), PlayerKeyset(keys2)]
+
+    def send_action(self, action: Action):
+        """Sends an action to the game environment.
+
+        Args:
+            action: The action to be sent.
+        """
+        self.simulator.enter_action(action)
+
+    def handle_keys(self):
+        """Handles keyboard inputs. Sends action for the respective players. When a key is held down, every frame
+        an action is sent in this function.
+        """
+        keys = pygame.key.get_pressed()
+        for player_idx, keyset in enumerate(self.player_keysets):
+            relevant_keys = [keys[k] for k in keyset.player_keys]
+            if any(relevant_keys[:-2]):
+                move_vec = np.zeros(2)
+                for idx, pressed in enumerate(relevant_keys[:-2]):
+                    if pressed:
+                        move_vec += keyset.move_vectors[idx]
+                if np.linalg.norm(move_vec) != 0:
+                    move_vec = move_vec / np.linalg.norm(move_vec)
+
+                action = Action(f"p{player_idx+1}", "movement", move_vec)
+                self.send_action(action)
+            if GET_CONTINOUS_INTERACT_AND_PICKUP:
+                if relevant_keys[-2]:
+                    action = Action(f"p{player_idx+1}", "interact", "interact")
+                    self.send_action(action)
+                if relevant_keys[-1]:
+                    action = Action(f"p{player_idx+1}", "pickup", "pickup")
+                    self.send_action(action)
+
+    def handle_interact_single_send(self, event):
+        """Handles key events. Here when a key is held down, only one action is sent.
+
+        Args:
+            event: Pygame event for extracting the key.
+        """
+        for player_idx, keyset in enumerate(self.player_keysets):
+            if event.key == keyset.pickup_key:
+                action = Action(f"p{player_idx + 1}", "pickup", "pickup")
+                self.send_action(action)
+            elif event.key == keyset.interact_key:
+                action = Action(f"p{player_idx + 1}", "interact", "interact")
+                self.send_action(action)
+
+    def draw_background(self):
+        """Visualizes a game background.
+        """
+        GREY = (200, 200, 200)
+        blockSize = 20  # Set the size of the grid block
+        for x in range(0, self.WINDOW_WIDTH, blockSize):
+            for y in range(0, self.WINDOW_HEIGHT, blockSize):
+                rect = pygame.Rect(x, y, blockSize, blockSize)
+                pygame.draw.rect(self.screen, GREY, rect, 1)
+
+    def draw_players(self, state):
+        """Visualizes the players as circles with an triangle for the facing diretion.
+
+        Args:
+            state: The game state returned by the environment.
+        """
+        for player in state["players"].values():
+            pos = player.pos
+            size = player.radius
+            color1 = RED if player.name == "p1" else GREEN
+            color2 = WHITE
+
+            rect = pygame.Rect(pos[0] - (size / 2), pos[1] - (size / 2), size, size)
+            pygame.draw.circle(self.screen, color2, pos, size)
+            pygame.draw.rect(self.screen, color1, rect)
+
+            facing = player.facing_direction
+
+            pygame.draw.polygon(self.screen, BLUE,
+                            ((pos[0]+(facing[1]*5), pos[1]-(facing[0]*5)), (pos[0]-(facing[1]*5), pos[1]+(facing[0]*5)), player.pos + (facing * 20)))
+
+    def draw_counters(self, state):
+        """Visualizes the counters in the environment.
+
+        Args:
+            state: The game state returned by the environment.
+        """
+        for idx, counter in enumerate(state["counters"]):
+            counter_rect_outline = pygame.Rect(counter.pos[0] - (self.counter_size / 2),
+                                       counter.pos[1] - (self.counter_size / 2), self.counter_size,
+                                       self.counter_size)
+
+            pygame.draw.rect(self.screen, COUNTERCOLOR, counter_rect_outline)
+
+    def draw(self, state):
+        """Main visualization function.
+
+        Args:
+            state: The game state returned by the environment.
+        """
+        self.screen.fill(BACKGROUND_COLOR)
+        self.draw_background()
+
+        self.draw_counters(state)
+        self.draw_players(state)
+
+        pygame.display.flip()
+
+    def start_pygame(self):
+        """Starts pygame and the gui loop. Each frame the gamestate is visualized and keyboard inputs are read.
+        """
+        pygame.init()
+        pygame.font.init()
+
+        self.screen = pygame.display.set_mode((self.WINDOW_WIDTH, self.WINDOW_HEIGHT))
+        pygame.display.set_caption("Simple Overcooked Simulator")
+        self.screen.fill(BACKGROUND_COLOR)
+
+        clock = pygame.time.Clock()
+
+        # Game loop
+        running = True
+        while running:
+
+            for event in pygame.event.get():
+                if event.type == pygame.QUIT:
+                    running = False
+                if not GET_CONTINOUS_INTERACT_AND_PICKUP:
+                    if event.type == pygame.KEYDOWN:
+                        self.handle_interact_single_send(event)
+
+            self.handle_keys()
+            clock.tick(FPS)
+            state = self.simulator.get_state()
+            self.draw(state)
+
+        pygame.quit()