From 87f34da886c6e48b3732a271cfde289b6d8bf3d3 Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sun, 10 Apr 2022 10:26:11 +0200
Subject: [PATCH 01/21] refactoring play function. Tests for keys to action
 mapping.

---
 gym/utils/play.py        | 36 +++++++++++++-------
 tests/utils/test_play.py | 72 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 12 deletions(-)
 create mode 100644 tests/utils/test_play.py

diff --git a/gym/utils/play.py b/gym/utils/play.py
index 74667820581..75662b87f8e 100644
--- a/gym/utils/play.py
+++ b/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:
diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
new file mode 100644
index 00000000000..e53a744ce27
--- /dev/null
+++ b/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)
\ No newline at end of file

From 7982347dc0b7446875e00fa1cf45fbb598ad613d Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sun, 10 Apr 2022 11:32:56 +0200
Subject: [PATCH 02/21] Add mocking pygame events.

---
 gym/utils/play.py        | 41 ++++++++++++++++++------
 tests/utils/test_play.py | 67 +++++++++++++++++++++++++---------------
 2 files changed, 74 insertions(+), 34 deletions(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index 75662b87f8e..1199d58a5fb 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -19,10 +19,13 @@
 
 
 class PlayableGame:
-    def __init__(self, env):
+    def __init__(self, env, keys_to_action=None):
         self.env = env
+        self.relevant_keys = self.get_relevant_keys(keys_to_action)
+        self.pressed_keys = []
+        self.running = True
 
-    def get_relevant_keys(self, keys_to_action=None):
+    def get_relevant_keys(self, keys_to_action):
         if keys_to_action is None:
             if hasattr(self.env, "get_keys_to_action"):
                 keys_to_action = self.env.get_keys_to_action()
@@ -37,6 +40,30 @@ def get_relevant_keys(self, keys_to_action=None):
         relevant_keys = set(sum(map(list, keys_to_action.keys()), []))
         return relevant_keys
 
+    def get_video_size(self, zoom=None):
+        rendered = self.env.render(mode="rgb_array")
+        video_size = [rendered.shape[1], rendered.shape[0]]
+
+        if zoom is not None:
+            video_size = int(video_size[0] * zoom), int(video_size[1] * zoom)
+
+        return video_size
+
+    def process_event(self, event):
+        if event.type == pygame.KEYDOWN:
+            if event.key in self.relevant_keys:
+                self.pressed_keys.append(event.key)
+            elif event.key == 27:
+                self.running = False
+        elif event.type == pygame.KEYUP:
+            if event.key in self.relevant_keys:
+                self.pressed_keys.remove(event.key)
+        elif event.type == pygame.QUIT:
+            self.running = False
+        elif event.type == VIDEORESIZE:
+            video_size = event.size
+            screen = pygame.display.set_mode(video_size)
+            print(video_size)
 
 
 def display_arr(screen, arr, video_size, transpose):
@@ -104,15 +131,10 @@ 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")
+    game = PlayableGame(env, keys_to_action)
 
     relevant_keys = game.get_relevant_keys(keys_to_action)
-
-    video_size = [rendered.shape[1], rendered.shape[0]]
-    if zoom is not None:
-        video_size = int(video_size[0] * zoom), int(video_size[1] * zoom)
+    video_size = game.get_video_size(zoom)
 
     pressed_keys = []
     running = True
@@ -137,6 +159,7 @@ def callback(obs_t, obs_tp1, action, rew, done, info):
 
         # process pygame events
         for event in pygame.event.get():
+            # game.process_event(event)
             # test events, set key states
             if event.type == pygame.KEYDOWN:
                 if event.key in relevant_keys:
diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index e53a744ce27..cf9a5e1d8f6 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -1,66 +1,83 @@
 from dataclasses import dataclass
-import pytest
+from typing import Optional
+
 import numpy as np
+import pygame
+import pytest
+
 import gym
+from gym.utils.play import PlayableGame, play
 
-from gym.utils.play import PlayableGame
-from gym.utils.play import play
 
+@dataclass
+class MockKeyEvent:
+    type: pygame.event.Event
+    key: Optional[int]
 
 
 @dataclass
-class DummyEnvSpec():
+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 render(self, mode="rgb_array"):
+        return np.zeros((1, 1))
 
 
 def dummy_keys_to_action():
-    return {(ord('a'),): 0, (ord('d'),): 1}
+    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}
+    game = PlayableGame(env, dummy_keys_to_action())
+    assert game.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()
+        PlayableGame(env)
 
 
 def test_play_relevant_keys_with_env_attribute():
-    """Env has a keys_to_action 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}
+    assert game.relevant_keys == {97, 100}
 
 
+def test_video_size_no_zoom():
+    env = DummyPlayEnv()
+    game = PlayableGame(env, dummy_keys_to_action())
+    video_size = game.get_video_size()
+    assert video_size == list(env.render().shape)
 
 
+def test_video_size_zoom():
+    env = DummyPlayEnv()
+    game = PlayableGame(env, dummy_keys_to_action())
+    zoom_value = 2.2
+    video_size = game.get_video_size(zoom=zoom_value)
+    assert video_size == tuple(int(shape * zoom_value) for shape in env.render().shape)
+
+
+def test_keyboard_quit_event():
+    env = DummyPlayEnv()
+    game = PlayableGame(env, dummy_keys_to_action())
+    quit_event = MockKeyEvent(pygame.KEYDOWN, 27)
+    game.process_event(quit_event)
+    assert game.running == False
 
 
 # def test_play_loop():
@@ -69,4 +86,4 @@ def test_play_relevant_keys_with_env_attribute():
 #         (ord('a'),): 0,
 #         (ord('d'),): 1
 #     }
-#     play(env, keys_to_action=keys_to_action)
\ No newline at end of file
+#     play(env, keys_to_action=keys_to_action)

From 2a73571a09a43690d929cc5c452894efd857525d Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sun, 10 Apr 2022 16:15:00 +0200
Subject: [PATCH 03/21] partial event processing in class.

---
 gym/utils/play.py        | 24 +++++--------------
 tests/utils/test_play.py | 50 ++++++++++++++++++++++++++++++++++++----
 2 files changed, 51 insertions(+), 23 deletions(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index 1199d58a5fb..8f366c5c7d2 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -1,12 +1,13 @@
 import argparse
 
-import matplotlib
 import pygame
 
 import gym
 from gym import logger
 
 try:
+    import matplotlib
+
     matplotlib.use("TkAgg")
     import matplotlib.pyplot as plt
 except ImportError as e:
@@ -133,22 +134,19 @@ def callback(obs_t, obs_tp1, action, rew, done, info):
     env.reset()
     game = PlayableGame(env, keys_to_action)
 
-    relevant_keys = game.get_relevant_keys(keys_to_action)
     video_size = game.get_video_size(zoom)
 
-    pressed_keys = []
-    running = True
     env_done = True
 
     screen = pygame.display.set_mode(video_size)
     clock = pygame.time.Clock()
 
-    while running:
+    while game.running:
         if env_done:
             env_done = False
             obs = env.reset()
         else:
-            action = keys_to_action.get(tuple(sorted(pressed_keys)), 0)
+            action = keys_to_action.get(tuple(sorted(game.pressed_keys)), 0)
             prev_obs = obs
             obs, rew, env_done, info = env.step(action)
             if callback is not None:
@@ -159,19 +157,9 @@ def callback(obs_t, obs_tp1, action, rew, done, info):
 
         # process pygame events
         for event in pygame.event.get():
-            # game.process_event(event)
+            game.process_event(event)
             # test events, set key states
-            if event.type == pygame.KEYDOWN:
-                if event.key in relevant_keys:
-                    pressed_keys.append(event.key)
-                elif event.key == 27:
-                    running = False
-            elif event.type == pygame.KEYUP:
-                if event.key in relevant_keys:
-                    pressed_keys.remove(event.key)
-            elif event.type == pygame.QUIT:
-                running = False
-            elif event.type == VIDEORESIZE:
+            if event.type == VIDEORESIZE:
                 video_size = event.size
                 screen = pygame.display.set_mode(video_size)
                 print(video_size)
diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index cf9a5e1d8f6..adff8864ca2 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -1,5 +1,5 @@
-from dataclasses import dataclass
-from typing import Optional
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
 
 import numpy as np
 import pygame
@@ -8,11 +8,15 @@
 import gym
 from gym.utils.play import PlayableGame, play
 
+RELEVANT_KEY = 100
+IRRELEVANT_KEY = 1
+
 
 @dataclass
 class MockKeyEvent:
     type: pygame.event.Event
-    key: Optional[int]
+    key: Optional[int] = field(default=None)
+    size: Optional[Tuple[int, int]] = field(default=None)
 
 
 @dataclass
@@ -75,11 +79,47 @@ def test_video_size_zoom():
 def test_keyboard_quit_event():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    quit_event = MockKeyEvent(pygame.KEYDOWN, 27)
-    game.process_event(quit_event)
+    event = MockKeyEvent(pygame.KEYDOWN, 27)
+    assert game.running == True
+    game.process_event(event)
     assert game.running == False
 
 
+def test_pygame_quit_event():
+    env = DummyPlayEnv()
+    game = PlayableGame(env, dummy_keys_to_action())
+    event = MockKeyEvent(pygame.QUIT)
+    assert game.running == True
+    game.process_event(event)
+    assert game.running == False
+
+
+def test_keyboard_relevant_keydown_event():
+    env = DummyPlayEnv()
+    game = PlayableGame(env, dummy_keys_to_action())
+    event = MockKeyEvent(pygame.KEYDOWN, RELEVANT_KEY)
+    game.process_event(event)
+    assert game.pressed_keys == [RELEVANT_KEY]
+
+
+def test_keyboard_irrelevant_keydown_event():
+    env = DummyPlayEnv()
+    game = PlayableGame(env, dummy_keys_to_action())
+    event = MockKeyEvent(pygame.KEYDOWN, IRRELEVANT_KEY)
+    game.process_event(event)
+    assert game.pressed_keys == []
+
+
+def test_keyboard_keyup_event():
+    env = DummyPlayEnv()
+    game = PlayableGame(env, dummy_keys_to_action())
+    event = MockKeyEvent(pygame.KEYDOWN, RELEVANT_KEY)
+    game.process_event(event)
+    event = MockKeyEvent(pygame.KEYUP, RELEVANT_KEY)
+    game.process_event(event)
+    assert game.pressed_keys == []
+
+
 # def test_play_loop():
 #     env = DummyPlayEnv()
 #     keys_to_action = {

From 02b81e5b3f34e661d2e1bf7022ab5fc74fc35122 Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sun, 10 Apr 2022 17:10:39 +0200
Subject: [PATCH 04/21] pre-commit.

---
 gym/utils/play.py        | 40 +++++++++++++++++-----------------------
 tests/utils/test_play.py | 25 +++++++++----------------
 2 files changed, 26 insertions(+), 39 deletions(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index 8f366c5c7d2..a06e07f24c6 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -20,13 +20,15 @@
 
 
 class PlayableGame:
-    def __init__(self, env, keys_to_action=None):
+    def __init__(self, env, keys_to_action=None, zoom=None):
         self.env = env
-        self.relevant_keys = self.get_relevant_keys(keys_to_action)
+        self.relevant_keys = self._get_relevant_keys(keys_to_action)
+        self.video_size = self._get_video_size(zoom)
+        self.screen = pygame.display.set_mode(self.video_size)
         self.pressed_keys = []
         self.running = True
 
-    def get_relevant_keys(self, keys_to_action):
+    def _get_relevant_keys(self, keys_to_action):
         if keys_to_action is None:
             if hasattr(self.env, "get_keys_to_action"):
                 keys_to_action = self.env.get_keys_to_action()
@@ -41,7 +43,7 @@ def get_relevant_keys(self, keys_to_action):
         relevant_keys = set(sum(map(list, keys_to_action.keys()), []))
         return relevant_keys
 
-    def get_video_size(self, zoom=None):
+    def _get_video_size(self, zoom=None):
         rendered = self.env.render(mode="rgb_array")
         video_size = [rendered.shape[1], rendered.shape[0]]
 
@@ -62,9 +64,8 @@ def process_event(self, event):
         elif event.type == pygame.QUIT:
             self.running = False
         elif event.type == VIDEORESIZE:
-            video_size = event.size
-            screen = pygame.display.set_mode(video_size)
-            print(video_size)
+            self.video_size = event.size
+            self.screen = pygame.display.set_mode(self.video_size)
 
 
 def display_arr(screen, arr, video_size, transpose):
@@ -132,37 +133,30 @@ 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, keys_to_action)
+    game = PlayableGame(env, keys_to_action, zoom)
 
-    video_size = game.get_video_size(zoom)
-
-    env_done = True
-
-    screen = pygame.display.set_mode(video_size)
+    done = True
     clock = pygame.time.Clock()
 
     while game.running:
-        if env_done:
-            env_done = False
+        if done:
+            done = False
             obs = env.reset()
         else:
             action = keys_to_action.get(tuple(sorted(game.pressed_keys)), 0)
             prev_obs = obs
-            obs, rew, env_done, info = env.step(action)
+            obs, rew, done, info = env.step(action)
             if callback is not None:
-                callback(prev_obs, obs, action, rew, env_done, info)
+                callback(prev_obs, obs, action, rew, done, info)
         if obs is not None:
             rendered = env.render(mode="rgb_array")
-            display_arr(screen, rendered, transpose=transpose, video_size=video_size)
+            display_arr(
+                game.screen, rendered, transpose=transpose, video_size=game.video_size
+            )
 
         # process pygame events
         for event in pygame.event.get():
             game.process_event(event)
-            # test events, set key states
-            if event.type == VIDEORESIZE:
-                video_size = event.size
-                screen = pygame.display.set_mode(video_size)
-                print(video_size)
 
         pygame.display.flip()
         clock.tick(fps)
diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index adff8864ca2..bf65193b67e 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -26,7 +26,11 @@ class DummyEnvSpec:
 
 class DummyPlayEnv(gym.Env):
     def step(self, action):
-        ...
+        obs = np.zeros((1, 1))
+        rew = 0
+        done = False
+        info = {}
+        return obs, rew, done, info
 
     def reset(self):
         ...
@@ -64,16 +68,14 @@ def test_play_relevant_keys_with_env_attribute():
 def test_video_size_no_zoom():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    video_size = game.get_video_size()
-    assert video_size == list(env.render().shape)
+    assert game.video_size == list(env.render().shape)
 
 
 def test_video_size_zoom():
     env = DummyPlayEnv()
-    game = PlayableGame(env, dummy_keys_to_action())
-    zoom_value = 2.2
-    video_size = game.get_video_size(zoom=zoom_value)
-    assert video_size == tuple(int(shape * zoom_value) for shape in env.render().shape)
+    zoom = 2.2
+    game = PlayableGame(env, dummy_keys_to_action(), zoom)
+    assert game.video_size == tuple(int(shape * zoom) for shape in env.render().shape)
 
 
 def test_keyboard_quit_event():
@@ -118,12 +120,3 @@ def test_keyboard_keyup_event():
     event = MockKeyEvent(pygame.KEYUP, RELEVANT_KEY)
     game.process_event(event)
     assert game.pressed_keys == []
-
-
-# def test_play_loop():
-#     env = DummyPlayEnv()
-#     keys_to_action = {
-#         (ord('a'),): 0,
-#         (ord('d'),): 1
-#     }
-#     play(env, keys_to_action=keys_to_action)

From 7c65be7369f7fbdc89d7195bf780b0f5ad82da3a Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sun, 10 Apr 2022 17:39:37 +0200
Subject: [PATCH 05/21] quit pygame after tests.

---
 tests/utils/test_play.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index bf65193b67e..ff39af1bcd5 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -6,7 +6,7 @@
 import pytest
 
 import gym
-from gym.utils.play import PlayableGame, play
+from gym.utils.play import PlayableGame
 
 RELEVANT_KEY = 100
 IRRELEVANT_KEY = 1
@@ -43,6 +43,12 @@ def dummy_keys_to_action():
     return {(ord("a"),): 0, (ord("d"),): 1}
 
 
+@pytest.fixture(autouse=True)
+def close_pygame():
+    yield
+    pygame.quit()
+
+
 def test_play_relvant_keys():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())

From 1b3485fb272cdf3818d609118a7a10d28f6de7ed Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sun, 10 Apr 2022 23:49:34 +0200
Subject: [PATCH 06/21] fix typos in functions names.

---
 tests/utils/test_play.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index ff39af1bcd5..02962bc8d74 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -49,13 +49,13 @@ def close_pygame():
     pygame.quit()
 
 
-def test_play_relvant_keys():
+def test_play_relevant_keys():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
     assert game.relevant_keys == {97, 100}
 
 
-def test_play_revant_keys_no_mapping():
+def test_play_relevant_keys_no_mapping():
     env = DummyPlayEnv()
     env.spec = DummyEnvSpec("DummyPlayEnv")
 

From 48d5f07f21331bcf08cf7e8cf05e39cd6d577207 Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Mon, 11 Apr 2022 23:57:10 +0200
Subject: [PATCH 07/21] Add type hint.

---
 gym/utils/play.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index a06e07f24c6..94e373bdb07 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -1,9 +1,11 @@
 import argparse
+from typing import Tuple
 
 import pygame
+from pygame.event import Event
 
 import gym
-from gym import logger
+from gym import Env, logger
 
 try:
     import matplotlib
@@ -20,7 +22,7 @@
 
 
 class PlayableGame:
-    def __init__(self, env, keys_to_action=None, zoom=None):
+    def __init__(self, env: Env, keys_to_action: dict = None, zoom: float = None):
         self.env = env
         self.relevant_keys = self._get_relevant_keys(keys_to_action)
         self.video_size = self._get_video_size(zoom)
@@ -28,7 +30,7 @@ def __init__(self, env, keys_to_action=None, zoom=None):
         self.pressed_keys = []
         self.running = True
 
-    def _get_relevant_keys(self, keys_to_action):
+    def _get_relevant_keys(self, keys_to_action: dict) -> set:
         if keys_to_action is None:
             if hasattr(self.env, "get_keys_to_action"):
                 keys_to_action = self.env.get_keys_to_action()
@@ -43,7 +45,7 @@ def _get_relevant_keys(self, keys_to_action):
         relevant_keys = set(sum(map(list, keys_to_action.keys()), []))
         return relevant_keys
 
-    def _get_video_size(self, zoom=None):
+    def _get_video_size(self, zoom: float = None) -> Tuple[int, int]:
         rendered = self.env.render(mode="rgb_array")
         video_size = [rendered.shape[1], rendered.shape[0]]
 
@@ -52,7 +54,7 @@ def _get_video_size(self, zoom=None):
 
         return video_size
 
-    def process_event(self, event):
+    def process_event(self, event: Event) -> None:
         if event.type == pygame.KEYDOWN:
             if event.key in self.relevant_keys:
                 self.pressed_keys.append(event.key)

From 3b9e08da18853e663b78c8e3cb25e1ef3be7f49f Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Wed, 13 Apr 2022 00:14:55 +0200
Subject: [PATCH 08/21] Add test for play function.

---
 tests/utils/test_play.py | 50 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 45 insertions(+), 5 deletions(-)

diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index 02962bc8d74..610a3c9447d 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -1,12 +1,14 @@
 from dataclasses import dataclass, field
-from typing import Optional, Tuple
+from typing import Callable, Optional, Tuple
 
 import numpy as np
 import pygame
 import pytest
+from pygame import KEYDOWN, QUIT, event
+from pygame.event import Event
 
 import gym
-from gym.utils.play import PlayableGame
+from gym.utils.play import PlayableGame, play
 
 RELEVANT_KEY = 100
 IRRELEVANT_KEY = 1
@@ -27,9 +29,7 @@ class DummyEnvSpec:
 class DummyPlayEnv(gym.Env):
     def step(self, action):
         obs = np.zeros((1, 1))
-        rew = 0
-        done = False
-        info = {}
+        rew, done, info = 1, False, {}
         return obs, rew, done, info
 
     def reset(self):
@@ -39,6 +39,30 @@ def render(self, mode="rgb_array"):
         return np.zeros((1, 1))
 
 
+class PlayStatus:
+    def __init__(self, callback: Callable):
+        self.data_callback = callback
+        self.cumulative_reward = 0
+
+    def callback(self, obs_t, obs_tp1, action, rew, done, info):
+        self.cumulative_reward += self.data_callback(
+            obs_t, obs_tp1, action, rew, done, info
+        )
+
+
+# set of key events to inject into the play loop as callback
+callback_events = [
+    Event(KEYDOWN, {"key": RELEVANT_KEY}),
+    Event(KEYDOWN, {"key": RELEVANT_KEY}),
+    Event(QUIT),
+]
+
+
+def callback(obs_t, obs_tp1, action, rew, done, info):
+    event.post(callback_events.pop(0))
+    return rew
+
+
 def dummy_keys_to_action():
     return {(ord("a"),): 0, (ord("d"),): 1}
 
@@ -126,3 +150,19 @@ def test_keyboard_keyup_event():
     event = MockKeyEvent(pygame.KEYUP, RELEVANT_KEY)
     game.process_event(event)
     assert game.pressed_keys == []
+
+
+def test_play_loop():
+    env = DummyPlayEnv()
+    cumulative_env_reward = 0
+    for s in range(
+        len(callback_events)
+    ):  # we run the same number of steps executed with play()
+        _, rew, _, _ = env.step(None)
+        cumulative_env_reward += rew
+
+    env_play = DummyPlayEnv()
+    status = PlayStatus(callback)
+    play(env_play, callback=status.callback, keys_to_action=dummy_keys_to_action())
+
+    assert status.cumulative_reward == cumulative_env_reward

From 9a1d4b67684785c03c187c0a7e8ab217b8f9a28f Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Wed, 13 Apr 2022 01:55:06 +0200
Subject: [PATCH 09/21] remove mockKeyEvent.

---
 tests/utils/test_play.py | 19 ++++++-------------
 1 file changed, 6 insertions(+), 13 deletions(-)

diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index 610a3c9447d..35795a59a69 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -14,13 +14,6 @@
 IRRELEVANT_KEY = 1
 
 
-@dataclass
-class MockKeyEvent:
-    type: pygame.event.Event
-    key: Optional[int] = field(default=None)
-    size: Optional[Tuple[int, int]] = field(default=None)
-
-
 @dataclass
 class DummyEnvSpec:
     id: str
@@ -111,7 +104,7 @@ def test_video_size_zoom():
 def test_keyboard_quit_event():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    event = MockKeyEvent(pygame.KEYDOWN, 27)
+    event = Event(pygame.KEYDOWN, {"key": 27})
     assert game.running == True
     game.process_event(event)
     assert game.running == False
@@ -120,7 +113,7 @@ def test_keyboard_quit_event():
 def test_pygame_quit_event():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    event = MockKeyEvent(pygame.QUIT)
+    event = Event(pygame.QUIT)
     assert game.running == True
     game.process_event(event)
     assert game.running == False
@@ -129,7 +122,7 @@ def test_pygame_quit_event():
 def test_keyboard_relevant_keydown_event():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    event = MockKeyEvent(pygame.KEYDOWN, RELEVANT_KEY)
+    event = Event(pygame.KEYDOWN, {"key": RELEVANT_KEY})
     game.process_event(event)
     assert game.pressed_keys == [RELEVANT_KEY]
 
@@ -137,7 +130,7 @@ def test_keyboard_relevant_keydown_event():
 def test_keyboard_irrelevant_keydown_event():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    event = MockKeyEvent(pygame.KEYDOWN, IRRELEVANT_KEY)
+    event = Event(pygame.KEYDOWN, {"key": IRRELEVANT_KEY})
     game.process_event(event)
     assert game.pressed_keys == []
 
@@ -145,9 +138,9 @@ def test_keyboard_irrelevant_keydown_event():
 def test_keyboard_keyup_event():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    event = MockKeyEvent(pygame.KEYDOWN, RELEVANT_KEY)
+    event = Event(pygame.KEYDOWN, {"key": RELEVANT_KEY})
     game.process_event(event)
-    event = MockKeyEvent(pygame.KEYUP, RELEVANT_KEY)
+    event = Event(pygame.KEYUP, {"key": RELEVANT_KEY})
     game.process_event(event)
     assert game.pressed_keys == []
 

From a6a7cd98247c7c6a3017b4253b14e2a5d330d214 Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Thu, 14 Apr 2022 18:38:57 +0200
Subject: [PATCH 10/21] remove unused main code.

---
 gym/utils/play.py | 17 -----------------
 1 file changed, 17 deletions(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index 94e373bdb07..66c2715c051 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -199,20 +199,3 @@ def callback(self, obs_t, obs_tp1, action, rew, done, info):
             )
             self.ax[i].set_xlim(xmin, xmax)
         plt.pause(0.000001)
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--env",
-        type=str,
-        default="MontezumaRevengeNoFrameskip-v4",
-        help="Define Environment",
-    )
-    args = parser.parse_args()
-    env = gym.make(args.env)
-    play(env, zoom=4, fps=60)
-
-
-if __name__ == "__main__":
-    main()

From 42b5ead05faf7db55644f321303287159607eaae Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Fri, 15 Apr 2022 12:33:39 +0200
Subject: [PATCH 11/21] Adding type hints.

---
 gym/utils/play.py | 41 ++++++++++++++++++++++++++++++++---------
 1 file changed, 32 insertions(+), 9 deletions(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index 66c2715c051..e8dd8e0165c 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -1,7 +1,9 @@
 import argparse
-from typing import Tuple
+from typing import Callable, Dict, Optional, Tuple
 
 import pygame
+from numpy.typing import NDArray
+from pygame import Surface
 from pygame.event import Event
 
 import gym
@@ -21,8 +23,19 @@
 from pygame.locals import VIDEORESIZE
 
 
+class MissingKeysToAction(Exception):
+    """Raised when the environment does not have
+    a default keys_to_action mapping
+    """
+
+
 class PlayableGame:
-    def __init__(self, env: Env, keys_to_action: dict = None, zoom: float = None):
+    def __init__(
+        self,
+        env: Env,
+        keys_to_action: Optional[Dict[Tuple[int], int]] = None,
+        zoom: Optional[float] = None,
+    ):
         self.env = env
         self.relevant_keys = self._get_relevant_keys(keys_to_action)
         self.video_size = self._get_video_size(zoom)
@@ -30,17 +43,18 @@ def __init__(self, env: Env, keys_to_action: dict = None, zoom: float = None):
         self.pressed_keys = []
         self.running = True
 
-    def _get_relevant_keys(self, keys_to_action: dict) -> set:
+    def _get_relevant_keys(
+        self, keys_to_action: Optional[Dict[Tuple[int], int]] = None
+    ) -> set:
         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"
+                raise MissingKeysToAction(
+                    "%s does not have explicit key to action mapping, "
+                    "please specify one manually" % self.env.spec.id
                 )
         relevant_keys = set(sum(map(list, keys_to_action.keys()), []))
         return relevant_keys
@@ -70,7 +84,9 @@ def process_event(self, event: Event) -> None:
             self.screen = pygame.display.set_mode(self.video_size)
 
 
-def display_arr(screen, arr, video_size, transpose):
+def display_arr(
+    screen: Surface, arr: NDArray, video_size: Tuple[int, int], transpose: bool
+):
     arr_min, arr_max = arr.min(), arr.max()
     arr = 255.0 * (arr - arr_min) / (arr_max - arr_min)
     pyg_img = pygame.surfarray.make_surface(arr.swapaxes(0, 1) if transpose else arr)
@@ -78,7 +94,14 @@ def display_arr(screen, arr, video_size, transpose):
     screen.blit(pyg_img, (0, 0))
 
 
-def play(env, transpose=True, fps=30, zoom=None, callback=None, keys_to_action=None):
+def play(
+    env: Env,
+    transpose: Optional[bool] = True,
+    fps: Optional[int] = 30,
+    zoom: Optional[float] = None,
+    callback: Optional[Callable] = None,
+    keys_to_action: Optional[Dict[Tuple[int], int]] = None,
+):
     """Allows one to play the game using keyboard.
 
     To simply play the game use:

From 7dabd13e2d55d5875ce1ee3fc6d537cf81ff67cf Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Fri, 15 Apr 2022 14:15:18 +0200
Subject: [PATCH 12/21] catch custom exception in tests.

---
 tests/utils/test_play.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index 35795a59a69..82f07494b9d 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -8,7 +8,7 @@
 from pygame.event import Event
 
 import gym
-from gym.utils.play import PlayableGame, play
+from gym.utils.play import MissingKeysToAction, PlayableGame, play
 
 RELEVANT_KEY = 100
 IRRELEVANT_KEY = 1
@@ -76,7 +76,7 @@ def test_play_relevant_keys_no_mapping():
     env = DummyPlayEnv()
     env.spec = DummyEnvSpec("DummyPlayEnv")
 
-    with pytest.raises(AssertionError) as info:
+    with pytest.raises(MissingKeysToAction) as info:
         PlayableGame(env)
 
 

From 97fa4e9c558cf6497b505739ddd33d13dcbc1539 Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Fri, 15 Apr 2022 14:54:25 +0200
Subject: [PATCH 13/21] Fix magic numbers.

---
 gym/utils/play.py        |  2 +-
 tests/utils/test_play.py | 21 +++++++++++----------
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index e8dd8e0165c..fe4c76266c9 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -72,7 +72,7 @@ def process_event(self, event: Event) -> None:
         if event.type == pygame.KEYDOWN:
             if event.key in self.relevant_keys:
                 self.pressed_keys.append(event.key)
-            elif event.key == 27:
+            elif event.key == pygame.K_ESCAPE:
                 self.running = False
         elif event.type == pygame.KEYUP:
             if event.key in self.relevant_keys:
diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index 82f07494b9d..5daa778a291 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -10,7 +10,8 @@
 import gym
 from gym.utils.play import MissingKeysToAction, PlayableGame, play
 
-RELEVANT_KEY = 100
+RELEVANT_KEY_1 = ord("a")  # 97
+RELEVANT_KEY_2 = ord("d")  # 100
 IRRELEVANT_KEY = 1
 
 
@@ -45,8 +46,8 @@ def callback(self, obs_t, obs_tp1, action, rew, done, info):
 
 # set of key events to inject into the play loop as callback
 callback_events = [
-    Event(KEYDOWN, {"key": RELEVANT_KEY}),
-    Event(KEYDOWN, {"key": RELEVANT_KEY}),
+    Event(KEYDOWN, {"key": RELEVANT_KEY_1}),
+    Event(KEYDOWN, {"key": RELEVANT_KEY_1}),
     Event(QUIT),
 ]
 
@@ -57,7 +58,7 @@ def callback(obs_t, obs_tp1, action, rew, done, info):
 
 
 def dummy_keys_to_action():
-    return {(ord("a"),): 0, (ord("d"),): 1}
+    return {(RELEVANT_KEY_1,): 0, (RELEVANT_KEY_2,): 1}
 
 
 @pytest.fixture(autouse=True)
@@ -69,7 +70,7 @@ def close_pygame():
 def test_play_relevant_keys():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    assert game.relevant_keys == {97, 100}
+    assert game.relevant_keys == {RELEVANT_KEY_1, RELEVANT_KEY_2}
 
 
 def test_play_relevant_keys_no_mapping():
@@ -104,7 +105,7 @@ def test_video_size_zoom():
 def test_keyboard_quit_event():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    event = Event(pygame.KEYDOWN, {"key": 27})
+    event = Event(pygame.KEYDOWN, {"key": pygame.K_ESCAPE})
     assert game.running == True
     game.process_event(event)
     assert game.running == False
@@ -122,9 +123,9 @@ def test_pygame_quit_event():
 def test_keyboard_relevant_keydown_event():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    event = Event(pygame.KEYDOWN, {"key": RELEVANT_KEY})
+    event = Event(pygame.KEYDOWN, {"key": RELEVANT_KEY_1})
     game.process_event(event)
-    assert game.pressed_keys == [RELEVANT_KEY]
+    assert game.pressed_keys == [RELEVANT_KEY_1]
 
 
 def test_keyboard_irrelevant_keydown_event():
@@ -138,9 +139,9 @@ def test_keyboard_irrelevant_keydown_event():
 def test_keyboard_keyup_event():
     env = DummyPlayEnv()
     game = PlayableGame(env, dummy_keys_to_action())
-    event = Event(pygame.KEYDOWN, {"key": RELEVANT_KEY})
+    event = Event(pygame.KEYDOWN, {"key": RELEVANT_KEY_1})
     game.process_event(event)
-    event = Event(pygame.KEYUP, {"key": RELEVANT_KEY})
+    event = Event(pygame.KEYUP, {"key": RELEVANT_KEY_1})
     game.process_event(event)
     assert game.pressed_keys == []
 

From 808d0a146db8139f7518ebde665c980f4b3cb7f5 Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Fri, 15 Apr 2022 16:06:01 +0200
Subject: [PATCH 14/21] Add test with an actual environment.

---
 gym/utils/play.py        |  5 ++-
 tests/utils/test_play.py | 83 ++++++++++++++++++++++++++++++++--------
 2 files changed, 70 insertions(+), 18 deletions(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index fe4c76266c9..bf09eed6b94 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -101,6 +101,7 @@ def play(
     zoom: Optional[float] = None,
     callback: Optional[Callable] = None,
     keys_to_action: Optional[Dict[Tuple[int], int]] = None,
+    seed: Optional[int] = None,
 ):
     """Allows one to play the game using keyboard.
 
@@ -157,7 +158,7 @@ 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()
+    env.reset(seed=seed)
     game = PlayableGame(env, keys_to_action, zoom)
 
     done = True
@@ -166,7 +167,7 @@ def callback(obs_t, obs_tp1, action, rew, done, info):
     while game.running:
         if done:
             done = False
-            obs = env.reset()
+            obs = env.reset(seed=seed)
         else:
             action = keys_to_action.get(tuple(sorted(game.pressed_keys)), 0)
             prev_obs = obs
diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index 5daa778a291..65abd9bc75d 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -4,7 +4,7 @@
 import numpy as np
 import pygame
 import pytest
-from pygame import KEYDOWN, QUIT, event
+from pygame import KEYDOWN, KEYUP, QUIT, event
 from pygame.event import Event
 
 import gym
@@ -26,7 +26,7 @@ def step(self, action):
         rew, done, info = 1, False, {}
         return obs, rew, done, info
 
-    def reset(self):
+    def reset(self, seed=None):
         ...
 
     def render(self, mode="rgb_array"):
@@ -37,24 +37,14 @@ class PlayStatus:
     def __init__(self, callback: Callable):
         self.data_callback = callback
         self.cumulative_reward = 0
+        self.last_observation = None
 
     def callback(self, obs_t, obs_tp1, action, rew, done, info):
-        self.cumulative_reward += self.data_callback(
+        _, obs_tp1, _, rew, _, _ = self.data_callback(
             obs_t, obs_tp1, action, rew, done, info
         )
-
-
-# set of key events to inject into the play loop as callback
-callback_events = [
-    Event(KEYDOWN, {"key": RELEVANT_KEY_1}),
-    Event(KEYDOWN, {"key": RELEVANT_KEY_1}),
-    Event(QUIT),
-]
-
-
-def callback(obs_t, obs_tp1, action, rew, done, info):
-    event.post(callback_events.pop(0))
-    return rew
+        self.cumulative_reward += rew
+        self.last_observation = obs_tp1
 
 
 def dummy_keys_to_action():
@@ -147,6 +137,17 @@ def test_keyboard_keyup_event():
 
 
 def test_play_loop():
+    # set of key events to inject into the play loop as callback
+    callback_events = [
+        Event(KEYDOWN, {"key": RELEVANT_KEY_1}),
+        Event(KEYDOWN, {"key": RELEVANT_KEY_1}),
+        Event(QUIT),
+    ]
+
+    def callback(obs_t, obs_tp1, action, rew, done, info):
+        event.post(callback_events.pop(0))
+        return obs_t, obs_tp1, action, rew, done, info
+
     env = DummyPlayEnv()
     cumulative_env_reward = 0
     for s in range(
@@ -160,3 +161,53 @@ def test_play_loop():
     play(env_play, callback=status.callback, keys_to_action=dummy_keys_to_action())
 
     assert status.cumulative_reward == cumulative_env_reward
+
+
+def test_play_loop_real_env():
+    SEED = 42
+    ENV = "CartPole-v1"
+
+    # set of key events to inject into the play loop as callback
+    callback_events = [
+        Event(KEYDOWN, {"key": RELEVANT_KEY_1}),
+        Event(KEYUP, {"key": RELEVANT_KEY_1}),
+        Event(KEYDOWN, {"key": RELEVANT_KEY_2}),
+        Event(KEYUP, {"key": RELEVANT_KEY_2}),
+        Event(KEYDOWN, {"key": RELEVANT_KEY_1}),
+        Event(KEYUP, {"key": RELEVANT_KEY_1}),
+        Event(KEYDOWN, {"key": RELEVANT_KEY_1}),
+        Event(KEYUP, {"key": RELEVANT_KEY_1}),
+        Event(KEYDOWN, {"key": RELEVANT_KEY_2}),
+        Event(KEYUP, {"key": RELEVANT_KEY_2}),
+        Event(QUIT),
+    ]
+    keydown_events = [k for k in callback_events if k.type == KEYDOWN]
+
+    def callback(obs_t, obs_tp1, action, rew, done, info):
+        pygame_event = callback_events.pop(0)
+        event.post(pygame_event)
+
+        # after releasing a key, post new events until
+        # we have one keydown
+        while pygame_event.type == KEYUP:
+            pygame_event = callback_events.pop(0)
+            event.post(pygame_event)
+
+        return obs_t, obs_tp1, action, rew, done, info
+
+    env = gym.make(ENV)
+    env.reset(seed=SEED)
+    keys_to_action = dummy_keys_to_action()
+
+    # first action is 0 because at the first iteration
+    # we have no input in the game
+    env.step(0)
+    for e in keydown_events:  # we run the same number of steps executed with play()
+        action = keys_to_action[(e.key,)]
+        obs, _, _, _ = env.step(action)
+
+    env_play = gym.make(ENV)
+    status = PlayStatus(callback)
+    play(env_play, callback=status.callback, keys_to_action=keys_to_action, seed=SEED)
+
+    assert (status.last_observation == obs).all()

From bf4d2d895977fa94138a582f50571bad5c6fcda5 Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Fri, 15 Apr 2022 16:08:03 +0200
Subject: [PATCH 15/21] fix comment.

---
 tests/utils/test_play.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/utils/test_play.py b/tests/utils/test_play.py
index 65abd9bc75d..0f0ee1e46eb 100644
--- a/tests/utils/test_play.py
+++ b/tests/utils/test_play.py
@@ -200,9 +200,9 @@ def callback(obs_t, obs_tp1, action, rew, done, info):
     keys_to_action = dummy_keys_to_action()
 
     # first action is 0 because at the first iteration
-    # we have no input in the game
+    # we can not inject a callback event into play()
     env.step(0)
-    for e in keydown_events:  # we run the same number of steps executed with play()
+    for e in keydown_events:
         action = keys_to_action[(e.key,)]
         obs, _, _, _ = env.step(action)
 

From 52b37ccbcd563f0f17975251500aafd121924147 Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sat, 16 Apr 2022 00:07:23 +0200
Subject: [PATCH 16/21] Add TODO memo on env.render.

---
 gym/utils/play.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index bf09eed6b94..151fe48b338 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -60,6 +60,7 @@ def _get_relevant_keys(
         return relevant_keys
 
     def _get_video_size(self, zoom: float = None) -> Tuple[int, int]:
+        # TODO: this needs to be updated when the render API API change goes through
         rendered = self.env.render(mode="rgb_array")
         video_size = [rendered.shape[1], rendered.shape[0]]
 
@@ -175,6 +176,7 @@ def callback(obs_t, obs_tp1, action, rew, done, info):
             if callback is not None:
                 callback(prev_obs, obs, action, rew, done, info)
         if obs is not None:
+            # TODO: this needs to be updated when the render API API change goes through
             rendered = env.render(mode="rgb_array")
             display_arr(
                 game.screen, rendered, transpose=transpose, video_size=game.video_size

From 5a8f6a4c9b5131819c458330da075af1673b00fb Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sat, 16 Apr 2022 00:26:32 +0200
Subject: [PATCH 17/21] change map with list comprehension.

---
 gym/utils/play.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index 151fe48b338..db3e75eeb01 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -56,7 +56,7 @@ def _get_relevant_keys(
                     "%s does not have explicit key to action mapping, "
                     "please specify one manually" % self.env.spec.id
                 )
-        relevant_keys = set(sum(map(list, keys_to_action.keys()), []))
+        relevant_keys = set(sum((list(k) for k in keys_to_action.keys()), []))
         return relevant_keys
 
     def _get_video_size(self, zoom: float = None) -> Tuple[int, int]:

From b5aae1b8182b7ca73a8cf997f2a6a25046452440 Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sat, 16 Apr 2022 00:27:10 +0200
Subject: [PATCH 18/21] remove unused imports.

---
 gym/utils/play.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index db3e75eeb01..d0549eb0304 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -1,4 +1,3 @@
-import argparse
 from typing import Callable, Dict, Optional, Tuple
 
 import pygame

From e9d333b988367865996607fbaec15033f2ca970c Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sat, 16 Apr 2022 11:52:50 +0200
Subject: [PATCH 19/21] Add type hint.

---
 gym/utils/play.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index d0549eb0304..31022a62405 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -58,7 +58,7 @@ def _get_relevant_keys(
         relevant_keys = set(sum((list(k) for k in keys_to_action.keys()), []))
         return relevant_keys
 
-    def _get_video_size(self, zoom: float = None) -> Tuple[int, int]:
+    def _get_video_size(self, zoom: Optional[float] = None) -> Tuple[int, int]:
         # TODO: this needs to be updated when the render API API change goes through
         rendered = self.env.render(mode="rgb_array")
         video_size = [rendered.shape[1], rendered.shape[0]]

From b0081c2a05a0d12413d5ce0bfbb030567bd9821f Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sat, 16 Apr 2022 12:14:58 +0200
Subject: [PATCH 20/21] typo.

---
 gym/utils/play.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index 31022a62405..ea351f39c03 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -59,7 +59,7 @@ def _get_relevant_keys(
         return relevant_keys
 
     def _get_video_size(self, zoom: Optional[float] = None) -> Tuple[int, int]:
-        # TODO: this needs to be updated when the render API API change goes through
+        # TODO: this needs to be updated when the render API change goes through
         rendered = self.env.render(mode="rgb_array")
         video_size = [rendered.shape[1], rendered.shape[0]]
 
@@ -175,7 +175,7 @@ def callback(obs_t, obs_tp1, action, rew, done, info):
             if callback is not None:
                 callback(prev_obs, obs, action, rew, done, info)
         if obs is not None:
-            # TODO: this needs to be updated when the render API API change goes through
+            # TODO: this needs to be updated when the render API change goes through
             rendered = env.render(mode="rgb_array")
             display_arr(
                 game.screen, rendered, transpose=transpose, video_size=game.video_size

From a02323c7d3cd59d12f38d608cffe391c32c86694 Mon Sep 17 00:00:00 2001
From: Gianluca De Cola <gianluca.deco@gmail.com>
Date: Sun, 17 Apr 2022 11:01:48 +0200
Subject: [PATCH 21/21] docstring.

---
 gym/utils/play.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/gym/utils/play.py b/gym/utils/play.py
index ea351f39c03..f04ccae4a62 100644
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -157,6 +157,8 @@ def callback(obs_t, obs_tp1, action, rew, done, info):
                 # ...
             }
         If None, default key_to_action mapping for that env is used, if provided.
+    seed: bool or None
+        Random seed used when resetting the environment. If None, no seed is used.
     """
     env.reset(seed=seed)
     game = PlayableGame(env, keys_to_action, zoom)