diff --git a/overcooked_simulator/gui_2d_vis/overcooked_gui.py b/overcooked_simulator/gui_2d_vis/overcooked_gui.py
index a537861748079b3b6aaefc828b4b7942d41483ba..ad2cf169e30f1b48179a2302ca68b99d20765189 100644
--- a/overcooked_simulator/gui_2d_vis/overcooked_gui.py
+++ b/overcooked_simulator/gui_2d_vis/overcooked_gui.py
@@ -1,6 +1,5 @@
import argparse
import dataclasses
-import glob
import json
import logging
import random
@@ -52,12 +51,13 @@ class PlayerKeySet:
"""
def __init__(
- self,
- move_keys: list[pygame.key],
- interact_key: pygame.key,
- pickup_key: pygame.key,
- switch_key: pygame.key,
- players: list[int],
+ self,
+ move_keys: list[pygame.key],
+ interact_key: pygame.key,
+ pickup_key: pygame.key,
+ switch_key: pygame.key,
+ players: list[int],
+ joystick: int
):
"""Creates a player key set which contains information about which keyboard keys control the player.
@@ -69,6 +69,7 @@ class PlayerKeySet:
pickup_key: The key to pick items up or put them down.
switch_key: The key for switching through controllable players.
players: The player indices which this keyset can control.
+ joystick: number of joystick (later check if available)
"""
self.move_vectors: list[list[int]] = [[-1, 0], [1, 0], [0, -1], [0, 1]]
self.key_to_movement: dict[pygame.key, list[int]] = {
@@ -82,6 +83,7 @@ class PlayerKeySet:
self.current_player: int = players[0] if players else 0
self.current_idx = 0
self.other_keyset: list[PlayerKeySet] = []
+ self.joystick = joystick
def set_controlled_players(self, controlled_players: list[int]) -> None:
self.controlled_players = controlled_players
@@ -102,10 +104,10 @@ class PyGameGUI:
"""Visualisation of the overcooked environment and reading keyboard inputs using pygame."""
def __init__(
- self,
- url: str,
- port: int,
- manager_ids: list[str],
+ self,
+ url: str,
+ port: int,
+ manager_ids: list[str],
):
pygame.init()
pygame.display.set_icon(
@@ -160,10 +162,10 @@ class PyGameGUI:
self.kitchen_height = state["kitchen"]["height"]
self.kitchen_aspect_ratio = self.kitchen_height / self.kitchen_width
game_width = self.visualization_config["GameWindow"]["min_width"] - (
- 2 * self.screen_margin
+ 2 * self.screen_margin
)
game_height = self.visualization_config["GameWindow"]["min_height"] - (
- 2 * self.screen_margin
+ 2 * self.screen_margin
)
if self.kitchen_width > game_width:
@@ -200,6 +202,9 @@ class PyGameGUI:
self.game_height -= residual_y
def setup_player_keys(self, n=1, disjunct=False):
+ # First four keys are for movement. Order: Down, Up, Left, Right.
+ # 5th key is for interacting with counters.
+ # 6th key ist for picking up things or dropping them.
if n:
players = list(range(self.number_humans_to_be_added))
key_set1 = PlayerKeySet(
@@ -208,6 +213,7 @@ class PyGameGUI:
pickup_key=pygame.K_e,
switch_key=pygame.K_SPACE,
players=players,
+ joystick=0
)
key_set2 = PlayerKeySet(
move_keys=[pygame.K_LEFT, pygame.K_RIGHT, pygame.K_UP, pygame.K_DOWN],
@@ -215,6 +221,7 @@ class PyGameGUI:
pickup_key=pygame.K_o,
switch_key=pygame.K_p,
players=players,
+ joystick=1
)
key_sets = [key_set1, key_set2]
@@ -256,6 +263,41 @@ class PyGameGUI:
)
self.send_action(action)
+ def handle_joy_stick_input(self, joysticks):
+ """Handles joystick inputs for movement every frame
+ Args:
+ joysticks: list of joysticks
+ """
+ # Axis 0: joy stick left: -1 = left, ~0 = center, 1 = right
+ # Axis 1: joy stick left: -1 = up, ~0 = center, 1 = down
+ # see control stuff here (at the end of the page): https://www.pygame.org/docs/ref/joystick.html
+ for key_set in self.key_sets:
+ current_player_name = str(key_set.current_player)
+ # if a joystick is connected for current player
+ if key_set.joystick in joysticks:
+ # Usually axis run in pairs, up/down for one, and left/right for the other. Triggers count as axes.
+ # You may want to take into account some tolerance to handle jitter, and
+ # joystick drift may keep the joystick from centering at 0 or using the full range of position values.
+ tolerance_threshold = 0.2
+ # axis 0 = joy stick left --> left & right
+ axis_left_right = joysticks[key_set.joystick].get_axis(0)
+ axis_up_down = joysticks[key_set.joystick].get_axis(1)
+ if abs(axis_left_right) > tolerance_threshold or abs(axis_up_down) > tolerance_threshold:
+ move_vec = np.zeros(2)
+ if abs(axis_left_right) > tolerance_threshold:
+ move_vec[0] += axis_left_right
+ # axis 1 = joy stick right --> up & down
+ if abs(axis_up_down) > tolerance_threshold:
+ move_vec[1] += axis_up_down
+
+ if np.linalg.norm(move_vec) != 0:
+ move_vec = move_vec / np.linalg.norm(move_vec)
+
+ action = Action(
+ current_player_name, ActionType.MOVEMENT, move_vec, duration=self.time_delta
+ )
+ self.send_action(action)
+
def handle_key_event(self, event):
"""Handles key events for the pickup and interaction keys. Pickup is a single action,
for interaction keydown and keyup is necessary, because the player has to be able to hold
@@ -286,6 +328,42 @@ class PyGameGUI:
if event.type == pygame.KEYDOWN:
key_set.next_player()
+ def handle_joy_stick_event(self, event, joysticks):
+ """Handles joy stick events for the pickup and interaction keys. Pickup is a single action,
+ for interaction buttondown and buttonup is necessary, because the player has to be able to hold
+ the button down.
+
+ Args:
+ event: Pygame event for extracting the button action.
+ joysticks: list of joysticks
+ """
+
+ for key_set in self.key_sets:
+ current_player_name = str(key_set.current_player)
+ # if a joystick is connected for current player
+ if key_set.joystick in joysticks:
+ # pickup = Button A <-> 0
+ if joysticks[key_set.joystick].get_button(0) and event.type == pygame.JOYBUTTONDOWN:
+ action = Action(current_player_name, ActionType.PUT, "pickup")
+ self.send_action(action)
+
+ # interact = Button X <-> 2
+ if joysticks[key_set.joystick].get_button(2) and event.type == pygame.JOYBUTTONDOWN:
+ action = Action(
+ current_player_name, ActionType.INTERACT, InterActionData.START
+ )
+ self.send_action(action)
+ # stop interaction if last pressed button was X <-> 2
+ if event.button == 2 and event.type == pygame.JOYBUTTONUP:
+ action = Action(
+ current_player_name, ActionType.INTERACT, InterActionData.STOP
+ )
+ self.send_action(action)
+ # switch button Y <-> 3
+ if joysticks[key_set.joystick].get_button(3) and not CONNECT_WITH_STUDY_SERVER:
+ if event.type == pygame.JOYBUTTONDOWN:
+ key_set.next_player()
+
def init_ui_elements(self):
self.manager = pygame_gui.UIManager((self.window_width, self.window_height))
self.manager.get_theme().load_theme(ROOT_DIR / "gui_2d_vis" / "gui_theme.json")
@@ -745,7 +823,7 @@ class PyGameGUI:
self.player_info = {self.player_info["player_id"]: self.player_info}
else:
environment_config_path = (
- ROOT_DIR / "game_content" / "environment_config.yaml"
+ ROOT_DIR / "game_content" / "environment_config.yaml"
)
layout_path = self.layout_file_paths[self.layout_selection.selected_option]
item_info_path = ROOT_DIR / "game_content" / "item_info.yaml"
@@ -789,7 +867,7 @@ class PyGameGUI:
)
)
assert (
- json.loads(websocket.recv())["status"] == 200
+ json.loads(websocket.recv())["status"] == 200
), "not accepted player"
self.websockets[player_id] = websocket
else:
@@ -865,13 +943,13 @@ class PyGameGUI:
self.menu_state = MenuStates.Game
self.number_players = (
- self.number_humans_to_be_added + self.number_bots_to_be_added
+ self.number_humans_to_be_added + self.number_bots_to_be_added
)
self.vis.create_player_colors(self.number_players)
if self.split_players:
assert (
- self.number_humans_to_be_added > 1
+ self.number_humans_to_be_added > 1
), "Not enough players for key configuration."
num_key_set = 2 if self.multiple_keysets else 1
self.key_sets = self.setup_player_keys(
@@ -1035,6 +1113,11 @@ class PyGameGUI:
# Game loop
self.running = True
+ # This dict can be left as-is, since pygame will generate a
+ # pygame.JOYDEVICEADDED event for every joystick connected
+ # at the start of the program.
+ joysticks = {}
+
while self.running:
try:
self.time_delta = clock.tick(self.FPS) / 1000
@@ -1044,6 +1127,20 @@ class PyGameGUI:
if event.type == pygame.QUIT:
self.running = False
+ # connect joystick
+ if pygame.joystick.get_count() > 0 and event.type == pygame.JOYDEVICEADDED:
+ # This event will be generated when the program starts for every
+ # joystick, filling up the list without needing to create them manually.
+ joy = pygame.joystick.Joystick(event.device_index)
+ joysticks[joy.get_instance_id()] = joy
+ print(f"Joystick {joy.get_instance_id()} connected")
+
+ # disconnect joystick
+ if event.type == pygame.JOYDEVICEREMOVED:
+ del joysticks[event.instance_id]
+ print(f"Joystick {event.instance_id} disconnected")
+ print("Number of joysticks:", pygame.joystick.get_count())
+
# elif event.type == pygame.VIDEORESIZE:
# # scrsize = event.size
# self.window_width_windowed = event.w
@@ -1057,8 +1154,8 @@ class PyGameGUI:
match event.ui_element:
case self.start_button:
if not (
- self.number_humans_to_be_added
- + self.number_bots_to_be_added
+ self.number_humans_to_be_added
+ + self.number_bots_to_be_added
):
continue
self.start_button_press()
@@ -1122,12 +1219,14 @@ class PyGameGUI:
self.manage_button_visibility()
if (
- event.type in [pygame.KEYDOWN, pygame.KEYUP]
- and self.menu_state == MenuStates.Game
+ event.type in [pygame.KEYDOWN, pygame.KEYUP]
+ and self.menu_state == MenuStates.Game
):
- pass
self.handle_key_event(event)
+ if event.type in [pygame.JOYBUTTONDOWN, pygame.JOYBUTTONUP] and self.menu_state == MenuStates.Game:
+ self.handle_joy_stick_event(event, joysticks=joysticks)
+
self.manager.process_events(event)
# drawing:
@@ -1144,6 +1243,7 @@ class PyGameGUI:
state = self.request_state()
self.handle_keys()
+ self.handle_joy_stick_input(joysticks=joysticks)
if state["ended"]:
self.finished_button_press()