Newer
Older
import pygame
import numpy as np
from overcooked_simulator.overcooked_environment import Action
from overcooked_simulator.simulation_runner import Simulator
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
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
self.GET_CONTINOUS_INTERACT_AND_PICKUP = False
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. Contains the player, action type and move direction if action is a movement.
"""
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 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. (This can be
switched by the GET_CONTINOUS_INTERACT_AND_PICKUP flag)
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.
"""
BACKGROUND_LINES = (200, 200, 200)
block_size = self.counter_size//2 # Set the size of the grid block
for x in range(0, self.window_width, block_size):
for y in range(0, self.window_height, block_size):
rect = pygame.Rect(x, y, block_size, block_size)
pygame.draw.rect(self.screen, BACKGROUND_LINES, 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)))
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
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 self.GET_CONTINOUS_INTERACT_AND_PICKUP:
if event.type == pygame.KEYDOWN:
self.handle_interact_single_send(event)
self.handle_keys()
state = self.simulator.get_state()
self.draw(state)
pygame.quit()