diff --git a/overcooked_simulator/gui_2d_vis/overcooked_gui.py b/overcooked_simulator/gui_2d_vis/overcooked_gui.py
index 796e94e48e22cd14c664a012b8d7e4ea46523627..87a0b5c35a6907dfce594cf3310d74c1d8c3646a 100644
--- a/overcooked_simulator/gui_2d_vis/overcooked_gui.py
+++ b/overcooked_simulator/gui_2d_vis/overcooked_gui.py
@@ -274,7 +274,7 @@ class PyGameGUI:
         for key_set in self.key_sets:
             current_player_name = str(key_set.current_player)
             # if a joystick is connected for current player
-            if joysticks[key_set.joystick]:
+            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.
@@ -341,7 +341,7 @@ class PyGameGUI:
         for key_set in self.key_sets:
             current_player_name = str(key_set.current_player)
             # if a joystick is connected for current player
-            if joysticks[key_set.joystick]:
+            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")