diff --git a/overcooked_simulator/gui_2d_vis/overcooked_gui.py b/overcooked_simulator/gui_2d_vis/overcooked_gui.py
index e0a9e1aaeb121fc4469f95c060778df60fea4142..802b3b519c3eee6e75c2e546d448df86ef0010f8 100644
--- a/overcooked_simulator/gui_2d_vis/overcooked_gui.py
+++ b/overcooked_simulator/gui_2d_vis/overcooked_gui.py
@@ -59,6 +59,7 @@ class PlayerKeySet:
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.
@@ -70,6 +71,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]] = {
@@ -83,6 +85,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
@@ -168,6 +171,9 @@ class PyGameGUI:
self.current_layout_idx = 0
def setup_player_keys(self, number_players, number_key_sets=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 number_key_sets:
players = list(range(number_players))
key_set1 = PlayerKeySet(
@@ -176,6 +182,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],
@@ -183,6 +190,7 @@ class PyGameGUI:
pickup_key=pygame.K_o,
switch_key=pygame.K_p,
players=players,
+ joystick=1,
)
key_sets = [key_set1, key_set2]
@@ -224,6 +232,47 @@ 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
@@ -254,6 +303,51 @@ 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 self.CONNECT_WITH_STUDY_SERVER
+ ):
+ if event.type == pygame.JOYBUTTONDOWN:
+ key_set.next_player()
+
def set_window_size(self):
if self.fullscreen:
flags = pygame.FULLSCREEN
@@ -1294,6 +1388,7 @@ class PyGameGUI:
self.menu_state = MenuStates.PostGame
self.disconnect_websockets()
self.finished_button_press()
+ self.handle_joy_stick_input(joysticks=self.joysticks)
if self.CONNECT_WITH_STUDY_SERVER:
self.send_level_done()
@@ -1345,8 +1440,20 @@ class PyGameGUI:
self.update_screen_elements()
+ self.reset_window_size()
+
+ self.init_ui_elements()
+ self.manage_button_visibility()
+
+ self.update_selection_elements()
+
# 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.
+ self.joysticks = {}
+
while self.running:
try:
self.time_delta = clock.tick(self.FPS) / 1000
@@ -1357,6 +1464,23 @@ class PyGameGUI:
self.disconnect_websockets()
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)
+ self.joysticks[joy.get_instance_id()] = joy
+ print(f"Joystick {joy.get_instance_id()} connected")
+
+ # disconnect joystick
+ if event.type == pygame.JOYDEVICEREMOVED:
+ del self.joysticks[event.instance_id]
+ print(f"Joystick {event.instance_id} disconnected")
+ print("Number of joysticks:", pygame.joystick.get_count())
+
if event.type == pygame_gui.UI_BUTTON_PRESSED:
self.manage_button_event(event)
self.update_screen_elements()
@@ -1370,6 +1494,9 @@ class PyGameGUI:
]:
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=self.joysticks)
+
self.manager.process_events(event)
# DRAWING