improved handling of keyboard input for human play

examples/human_play.py

@@ -26,27 +26,43 @@
 # Thread safe variables for use with GUI
 action = Tk.IntVar()
 action.set(0)
+action_taken = Tk.BooleanVar()
+action_taken.set(False)
+action_released = Tk.BooleanVar()
+action_released.set(False)
 G = Tk.DoubleVar()
 G.set(0.0)
 is_terminate = Tk.BooleanVar()
 is_terminate.set(False)
 
-# Key handler for human player
+# Map input keys to agent actions
+key_action_map = {' ':5, 'left':1, 'up':2, 'right':3, 'down':4}
+
+# Key press handler for human player
 def on_key_event(event):
-    if event.key == "q":  # quit the game
+    if event.key == 'q':  # quit the game
         gui.quit()
     elif (event.key  == 'r'): # reset the game environment
         env.reset()
-    elif (event.key  == ' '): # shoots bullet
-        action.set(5)
-    elif (event.key == 'left'):
-        action.set(1)
-    elif (event.key == 'up'):
-        action.set(2)
-    elif (event.key  == 'right'):
-        action.set(3)
-    elif (event.key  == 'down'):
-        action.set(4)
+    elif(event.key in key_action_map):
+        key_action = key_action_map[event.key]
+        action.set(key_action)
+        # When new action is selected it has not yet been taken or released
+        action_released.set(False)
+        action_taken.set(False)
+
+# Key release handlr for human player
+def on_release_event(event):
+    if(event.key in key_action_map):
+        key_action = key_action_map[event.key]
+        a = action.get()
+        if(a==key_action):
+            # If released action has already been taken set action to no-op immediately
+            if(action_taken.get()):
+                action.set(0)
+            # Otherwise note that it has been released so that we can set back to no-op when it is taken
+            else:
+                action_released.set(True)
 
 ################################################################################################################
 # play
@@ -58,17 +74,21 @@ def play():
     if is_terminate.get() == True:
         gui.quit()
 
+    # Get players actions selection, if that action has been released, set action back to no-op for next frame
     a = action.get()
+    action_taken.set(True)
+    if(action_released.get()):
+        action.set(0)
+
     r, t = env.act(a)
     is_terminate.set(t)
     G.set(G.get()+r)
     gui.display_state(env.state())
     gui.set_message("Score: " + str(G.get()))
-    action.set(0)
     gui.update(50, play)
 
 # Hook up the key handler and initiate game play
-gui.overwrite_key_handle(on_key_event)
+gui.overwrite_key_handle(on_key_event, on_release_event)
 gui.update(0, play)
 gui.run()
 print("Final Score: "+str(G.get()))

minatar_environment/gui.py

@@ -52,7 +52,8 @@ def __init__(self, env_name, n_channels):
         self.ax = self.fig.add_subplot(111)
         self.canvas = FigureCanvasTkAgg(self.fig, master=self.root)
         self.canvas.get_tk_widget().pack(side=Tk.TOP, fill=Tk.BOTH, expand=1)
-        self.default_key_handler = self.canvas.mpl_connect('key_press_event', self.on_key_event)
+        self.key_press_handler = self.canvas.mpl_connect('key_press_event', self.on_key_event)
+        self.key_release_handler = self.canvas.mpl_connect('key_press_event', lambda x: None)
 
     # Set the message for the label on screen
     def set_message(self, str):
@@ -67,9 +68,13 @@ def display_state(self, state):
         self.canvas.draw()
 
     # Allow user to handle their own keyboard input
-    def overwrite_key_handle(self, key_handler):
-        self.canvas.mpl_disconnect(self.default_key_handler)
-        self.canvas.mpl_connect('key_press_event', key_handler)
+    def overwrite_key_handle(self, key_press_handler, key_release_handler=None):
+        self.canvas.mpl_disconnect(self.key_press_handler)
+        self.key_press_handler = self.canvas.mpl_connect('key_press_event', key_press_handler)
+        if(key_release_handler is not None):
+            self.canvas.mpl_disconnect(self.key_release_handler)
+            self.key_release_handler = self.canvas.mpl_connect('key_release_event', key_release_handler)
+
 
     # Default key handler
     def on_key_event(self, event):