refactoring play function. Tests for keys to action mapping.

gym/utils/play.py

@@ -18,6 +18,27 @@
 from pygame.locals import VIDEORESIZE
 
 
+class PlayableGame:
+    def __init__(self, env):
+        self.env = env
+
+    def get_relevant_keys(self, keys_to_action=None):
+        if keys_to_action is None:
+            if hasattr(self.env, "get_keys_to_action"):
+                keys_to_action = self.env.get_keys_to_action()
+            elif hasattr(self.env.unwrapped, "get_keys_to_action"):
+                keys_to_action = self.env.unwrapped.get_keys_to_action()
+            else:
+                assert False, (
+                    self.env.spec.id
+                    + " does not have explicit key to action mapping, "
+                    + "please specify one manually"
+                )
+        relevant_keys = set(sum(map(list, keys_to_action.keys()), []))
+        return relevant_keys
+
+
+
 def display_arr(screen, arr, video_size, transpose):
     arr_min, arr_max = arr.min(), arr.max()
     arr = 255.0 * (arr - arr_min) / (arr_max - arr_min)
@@ -83,20 +104,11 @@ def callback(obs_t, obs_tp1, action, rew, done, info):
         If None, default key_to_action mapping for that env is used, if provided.
     """
     env.reset()
+    game = PlayableGame(env)
+
     rendered = env.render(mode="rgb_array")
 
-    if keys_to_action is None:
-        if hasattr(env, "get_keys_to_action"):
-            keys_to_action = env.get_keys_to_action()
-        elif hasattr(env.unwrapped, "get_keys_to_action"):
-            keys_to_action = env.unwrapped.get_keys_to_action()
-        else:
-            assert False, (
-                env.spec.id
-                + " does not have explicit key to action mapping, "
-                + "please specify one manually"
-            )
-    relevant_keys = set(sum(map(list, keys_to_action.keys()), []))
+    relevant_keys = game.get_relevant_keys(keys_to_action)
 
     video_size = [rendered.shape[1], rendered.shape[0]]
     if zoom is not None:

tests/utils/test_play.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass
+import pytest
+import numpy as np
+import gym
+
+from gym.utils.play import PlayableGame
+from gym.utils.play import play
+
+
+
+@dataclass
+class DummyEnvSpec():
+    id: str
+
+
+class DummyPlayEnv(gym.Env):
+
+    def step(self, action):
+        ...
+
+    def reset(self):
+        ...
+
+    def render(self, mode):
+        return np.zeros((1,1))
+
+
+def dummy_keys_to_action():
+    return {(ord('a'),): 0, (ord('d'),): 1}
+
+
+def test_play_relvant_keys():
+    env = DummyPlayEnv()
+    keys_to_action = {
+         (ord('a'),): 0,
+         (ord('d'),): 1
+    }
+    game = PlayableGame(env)
+    relevant_keys = game.get_relevant_keys(keys_to_action)
+    assert relevant_keys == {97, 100}
+
+
+def test_play_revant_keys_no_mapping():
+    env = DummyPlayEnv()
+    env.spec = DummyEnvSpec("DummyPlayEnv")
+    game = PlayableGame(env)
+
+    with pytest.raises(AssertionError) as info:
+        game.get_relevant_keys()
+
+
+def test_play_relevant_keys_with_env_attribute():
+    """Env has a keys_to_action attribute
+    """
+    env = DummyPlayEnv()
+    env.get_keys_to_action = dummy_keys_to_action
+    game = PlayableGame(env)
+    relevant_keys = game.get_relevant_keys()
+    assert relevant_keys == {97, 100}
+
+
+
+
+
+
+# def test_play_loop():
+#     env = DummyPlayEnv()
+#     keys_to_action = {
+#         (ord('a'),): 0,
+#         (ord('d'),): 1
+#     }
+#     play(env, keys_to_action=keys_to_action)