diff --git a/overcooked_simulator/pygame_gui/pygame_gui.py b/overcooked_simulator/pygame_gui/pygame_gui.py
index 758440e6efec9bf4fc43ff3a1b399156d9c5f551..0705681286f2e6e20b94b1ee73eae39300956ba5 100644
--- a/overcooked_simulator/pygame_gui/pygame_gui.py
+++ b/overcooked_simulator/pygame_gui/pygame_gui.py
@@ -1,31 +1,25 @@
+import colorsys
import math
import numpy as np
import numpy.typing as npt
import pygame
import yaml
+from scipy.spatial import KDTree
from overcooked_simulator import ROOT_DIR
from overcooked_simulator.game_items import (
ProgressibleItem,
- Plate,
Item,
CookingEquipment,
Meal,
)
from overcooked_simulator.overcooked_environment import Action
-from overcooked_simulator.pygame_gui.game_colors import colors
+from overcooked_simulator.pygame_gui.game_colors import BLUE
+from overcooked_simulator.pygame_gui.game_colors import colors, Color
from overcooked_simulator.simulation_runner import Simulator
-COUNTER_COLOR = colors["whitesmoke"]
-GREEN = colors["green"]
-RED = colors["crimson"]
-BLUE = colors["blue"]
-WHITE = colors["white"]
-BACKGROUND_COLOR = colors["sgigray76"]
-BACKGROUND_LINES_COLOR = colors["gray79"]
-
-PLAYER_DEBUG_VIZ = False
+PLAYER_DEBUG_VIZ = True
class PlayerKeySet:
@@ -89,6 +83,27 @@ class PyGameGUI:
self.images_path = ROOT_DIR / "pygame_gui" / "images"
+ self.player_colors = self.create_player_colors()
+
+ def create_player_colors(self) -> list[Color]:
+ number_player = len(self.simulator.env.players)
+ hue_values = np.linspace(0, 1, number_player+1)
+
+ colors_vec = np.array([col for col in colors.values()])
+
+ tree = KDTree(colors_vec)
+
+ color_names = list(colors.keys())
+
+ player_colors = []
+ for hue in hue_values:
+ rgb = colorsys.hsv_to_rgb(hue, 1, 1)
+ query_color = np.array([int(c * 255) for c in rgb])
+ _, index = tree.query(query_color, k=1)
+ player_colors.append(color_names[index])
+
+ return player_colors
+
def send_action(self, action: Action):
"""Sends an action to the game environment.
@@ -142,7 +157,12 @@ class PyGameGUI:
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_COLOR, rect, 1)
+ pygame.draw.rect(
+ self.screen,
+ self.visualization_config["Kitchen"]["background_lines"],
+ rect,
+ 1,
+ )
def draw_players(self, state):
"""Visualizes the players as circles with a triangle for the facing direction.
@@ -151,11 +171,11 @@ class PyGameGUI:
Args:
state: The game state returned by the environment.
"""
- for player in state["players"].values():
+ for p_idx, player in enumerate(state["players"].values()):
if PLAYER_DEBUG_VIZ:
pos = player.pos
size = player.radius
- color1 = RED if player.name == "p1" else GREEN
+ color1 = self.player_colors[p_idx]
color2 = colors["white"]
pygame.draw.circle(self.screen, color2, pos, size)
@@ -163,7 +183,7 @@ class PyGameGUI:
pygame.draw.circle(
self.screen, BLUE, pos, player.interaction_range, width=1
)
- pygame.draw.circle(self.screen, color1, pos, size // 2)
+ pygame.draw.circle(self.screen, colors[color1], pos, size // 2)
pos = player.pos
facing = player.facing_direction
@@ -261,8 +281,8 @@ class PyGameGUI:
if isinstance(item, ProgressibleItem) and not item.finished:
self.draw_progress_bar(pos, item.progressed_steps, item.steps_needed)
- if isinstance(item, Plate) and item.holds:
- self.draw_item(pos, item.holds)
+ # if isinstance(item, Plate) and item.holds:
+ # self.draw_item(pos, item.holds)
if isinstance(item, CookingEquipment) and item.content:
self.draw_item(pos, item.content)
@@ -286,11 +306,11 @@ class PyGameGUI:
elif item.parts:
match len(item.parts):
case 1:
- pygame.draw.circle(self.screen, RED, pos, 4)
+ pygame.draw.circle(self.screen, colors["red"], pos, 4)
case 2:
- pygame.draw.circle(self.screen, RED, pos, 8)
+ pygame.draw.circle(self.screen, colors["red"], pos, 8)
case 3:
- pygame.draw.circle(self.screen, RED, pos, 12)
+ pygame.draw.circle(self.screen, colors["red"], pos, 12)
else:
for i, o in enumerate(item.parts):
self.draw_item(np.abs([pos[0], pos[1] - (i * 5)]), o)
@@ -306,7 +326,7 @@ class PyGameGUI:
progress_width,
bar_height,
)
- pygame.draw.rect(self.screen, GREEN, progress_bar)
+ pygame.draw.rect(self.screen, colors["green1"], progress_bar)
def draw_counter(self, counter):
"""Visualization of a counter at its position. If it is occupied by an item, it is also shown.
@@ -317,14 +337,7 @@ class PyGameGUI:
counter: The counter to visualize.
"""
- 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, COUNTER_COLOR, counter_rect_outline)
-
+ self.draw_thing(counter.pos, self.visualization_config["Counter"]["parts"])
if str(counter) in self.visualization_config:
self.draw_thing(
counter.pos, self.visualization_config[str(counter)]["parts"]
@@ -361,7 +374,9 @@ class PyGameGUI:
Args:
state: The game state returned by the environment.
"""
- self.screen.fill(BACKGROUND_COLOR)
+ self.screen.fill(
+ colors[self.visualization_config["Kitchen"]["ground_tiles_color"]]
+ )
self.draw_background()
self.draw_counters(state)
@@ -376,7 +391,9 @@ class PyGameGUI:
self.screen = pygame.display.set_mode((self.window_width, self.window_height))
pygame.display.set_caption("Simple Overcooked Simulator")
- self.screen.fill(BACKGROUND_COLOR)
+ self.screen.fill(
+ colors[self.visualization_config["Kitchen"]["ground_tiles_color"]]
+ )
clock = pygame.time.Clock()
diff --git a/overcooked_simulator/pygame_gui/visualization.yaml b/overcooked_simulator/pygame_gui/visualization.yaml
index f264639349521930155ffb0f2ca6cd8f912723db..54fe70b4443164fd517cc493bea32246f53b09a8 100644
--- a/overcooked_simulator/pygame_gui/visualization.yaml
+++ b/overcooked_simulator/pygame_gui/visualization.yaml
@@ -1,3 +1,7 @@
+Kitchen:
+ ground_tiles_color: sgigray76
+ background_lines: gray79
+
Counter:
parts:
- type: rect