浏览代码

Added keyboard handling to standalone.py

Maxime Chevalier-Boisvert 7 年之前
父节点
当前提交
a7a56cdfe2
共有 2 个文件被更改,包括 46 次插入5 次删除
  1. 26 0
      gym_minigrid/envs/rendering.py
  2. 20 5
      standalone.py

+ 26 - 0
gym_minigrid/envs/rendering.py

@@ -35,12 +35,38 @@ class Window(QMainWindow):
 
         self.closed = False
 
+        # Callback for keyboard events
+        self.keyDownCb = None
+
     def closeEvent(self, event):
         self.closed = True
 
     def setPixmap(self, pixmap):
         self.imgLabel.setPixmap(pixmap)
 
+    def setKeyDownCb(self, callback):
+        self.keyDownCb = callback
+
+    def keyPressEvent(self, e):
+        if self.keyDownCb == None:
+            return
+
+        keyName = None
+        if e.key() == Qt.Key_Left:
+            keyName = 'LEFT'
+        elif e.key() == Qt.Key_Right:
+            keyName = 'RIGHT'
+        elif e.key() == Qt.Key_Up:
+            keyName = 'UP'
+        elif e.key() == Qt.Key_Down:
+            keyName = 'DOWN'
+        elif e.key() == Qt.Key_Space:
+            keyName = 'SPACE'
+
+        if keyName == None:
+            return
+        self.keyDownCb(keyName)
+
 class Renderer:
     def __init__(self, width, height, ownWindow=False):
         self.width = width

+ 20 - 5
standalone.py

@@ -4,8 +4,10 @@ from __future__ import division, print_function
 
 import numpy
 import gym
+import time
 
 import gym_minigrid
+from gym_minigrid.envs import MiniGridEnv
 
 def main():
 
@@ -15,11 +17,18 @@ def main():
     # Create a window to render into
     renderer = env.render('human')
 
-    while True:
-
-        env.render('human')
-
+    def keyDownCb(keyName):
         action = 0
+        if keyName == 'LEFT':
+            action = MiniGridEnv.ACTION_LEFT
+        elif keyName == 'RIGHT':
+            action = MiniGridEnv.ACTION_RIGHT
+        elif keyName == 'UP':
+            action = MiniGridEnv.ACTION_FORWARD
+        elif keyName == 'SPACE':
+            action = MiniGridEnv.ACTION_TOGGLE
+        else:
+            "unknown key"
 
         obs, reward, done, info = env.step(action)
 
@@ -29,8 +38,14 @@ def main():
             print('done!')
             env.reset()
 
+    renderer.window.setKeyDownCb(keyDownCb)
+
+    while True:
+        env.render('human')
+        time.sleep(0.01)
+
         # If the window was closed
-        if not renderer.window:
+        if renderer.window == None:
             break
 
 if __name__ == "__main__":