We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Cannot use gym.utils.play.PlayPlot
```
from gym.utils.play import play, PlayPlot
def callback(obs_t, obs_tp1, action, rew, done, info):
    return [rew,]
env_plotter = PlayPlot(callback, 30 * 5, ["reward"])
env = gym.make("IceHockey-v0")
```
I got : self.fig, self.ax = plt.subplots(num_plots) NameError: name 'plt' is not defined

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `gym/utils/play.py`
Content:
```
1 import gym
2 import pygame
3 import sys
4 import time
5 import matplotlib
6 import argparse
7 try:
8     matplotlib.use('GTK3Agg')
9     import matplotlib.pyplot as plt
10 except Exception:
11     pass
12 
13 
14 import pyglet.window as pw
15 
16 from collections import deque
17 from pygame.locals import HWSURFACE, DOUBLEBUF, RESIZABLE, VIDEORESIZE
18 from threading import Thread
19 
20 parser = argparse.ArgumentParser()
21 parser.add_argument('--env', type=str, default='MontezumaRevengeNoFrameskip-v4', help='Define Environment')
22 args = parser.parse_args()
23 	
24 def display_arr(screen, arr, video_size, transpose):
25     arr_min, arr_max = arr.min(), arr.max()
26     arr = 255.0 * (arr - arr_min) / (arr_max - arr_min)
27     pyg_img = pygame.surfarray.make_surface(arr.swapaxes(0, 1) if transpose else arr)
28     pyg_img = pygame.transform.scale(pyg_img, video_size)
29     screen.blit(pyg_img, (0,0))
30 
31 def play(env, transpose=True, fps=30, zoom=None, callback=None, keys_to_action=None):
32     """Allows one to play the game using keyboard.
33 
34     To simply play the game use:
35 
36         play(gym.make("Pong-v3"))
37 
38     Above code works also if env is wrapped, so it's particularly useful in
39     verifying that the frame-level preprocessing does not render the game
40     unplayable.
41 
42     If you wish to plot real time statistics as you play, you can use
43     gym.utils.play.PlayPlot. Here's a sample code for plotting the reward
44     for last 5 second of gameplay.
45 
46         def callback(obs_t, obs_tp1, rew, done, info):
47             return [rew,]
48         env_plotter = EnvPlotter(callback, 30 * 5, ["reward"])
49 
50         env = gym.make("Pong-v3")
51         play(env, callback=env_plotter.callback)
52 
53 
54     Arguments
55     ---------
56     env: gym.Env
57         Environment to use for playing.
58     transpose: bool
59         If True the output of observation is transposed.
60         Defaults to true.
61     fps: int
62         Maximum number of steps of the environment to execute every second.
63         Defaults to 30.
64     zoom: float
65         Make screen edge this many times bigger
66     callback: lambda or None
67         Callback if a callback is provided it will be executed after
68         every step. It takes the following input:
69             obs_t: observation before performing action
70             obs_tp1: observation after performing action
71             action: action that was executed
72             rew: reward that was received
73             done: whether the environment is done or not
74             info: debug info
75     keys_to_action: dict: tuple(int) -> int or None
76         Mapping from keys pressed to action performed.
77         For example if pressed 'w' and space at the same time is supposed
78         to trigger action number 2 then key_to_action dict would look like this:
79 
80             {
81                 # ...
82                 sorted(ord('w'), ord(' ')) -> 2
83                 # ...
84             }
85         If None, default key_to_action mapping for that env is used, if provided.
86     """
87 
88     obs_s = env.observation_space
89     assert type(obs_s) == gym.spaces.box.Box
90     assert len(obs_s.shape) == 2 or (len(obs_s.shape) == 3 and obs_s.shape[2] in [1,3])
91 
92     if keys_to_action is None:
93         if hasattr(env, 'get_keys_to_action'):
94             keys_to_action = env.get_keys_to_action()
95         elif hasattr(env.unwrapped, 'get_keys_to_action'):
96             keys_to_action = env.unwrapped.get_keys_to_action()
97         else:
98             assert False, env.spec.id + " does not have explicit key to action mapping, " + \
99                           "please specify one manually"
100     relevant_keys = set(sum(map(list, keys_to_action.keys()),[]))
101 
102     if transpose:
103         video_size = env.observation_space.shape[1], env.observation_space.shape[0]
104     else:
105         video_size = env.observation_space.shape[0], env.observation_space.shape[1]
106 
107     if zoom is not None:
108         video_size = int(video_size[0] * zoom), int(video_size[1] * zoom)
109 
110     pressed_keys = []
111     running = True
112     env_done = True
113 
114     screen = pygame.display.set_mode(video_size)
115     clock = pygame.time.Clock()
116 
117 
118     while running:
119         if env_done:
120             env_done = False
121             obs = env.reset()
122         else:
123             action = keys_to_action.get(tuple(sorted(pressed_keys)), 0)
124             prev_obs = obs
125             obs, rew, env_done, info = env.step(action)
126             if callback is not None:
127                 callback(prev_obs, obs, action, rew, env_done, info)
128         if obs is not None:
129             if len(obs.shape) == 2:
130                 obs = obs[:, :, None]
131             if obs.shape[2] == 1:
132                 obs = obs.repeat(3, axis=2)
133             display_arr(screen, obs, transpose=transpose, video_size=video_size)
134 
135         # process pygame events
136         for event in pygame.event.get():
137             # test events, set key states
138             if event.type == pygame.KEYDOWN:
139                 if event.key in relevant_keys:
140                     pressed_keys.append(event.key)
141                 elif event.key == 27:
142                     running = False
143             elif event.type == pygame.KEYUP:
144                 if event.key in relevant_keys:
145                     pressed_keys.remove(event.key)
146             elif event.type == pygame.QUIT:
147                 running = False
148             elif event.type == VIDEORESIZE:
149                 video_size = event.size
150                 screen = pygame.display.set_mode(video_size)
151                 print(video_size)
152 
153         pygame.display.flip()
154         clock.tick(fps)
155     pygame.quit()
156 
157 class PlayPlot(object):
158     def __init__(self, callback, horizon_timesteps, plot_names):
159         self.data_callback = callback
160         self.horizon_timesteps = horizon_timesteps
161         self.plot_names = plot_names
162 
163         num_plots = len(self.plot_names)
164         self.fig, self.ax = plt.subplots(num_plots)
165         if num_plots == 1:
166             self.ax = [self.ax]
167         for axis, name in zip(self.ax, plot_names):
168             axis.set_title(name)
169         self.t = 0
170         self.cur_plot = [None for _ in range(num_plots)]
171         self.data     = [deque(maxlen=horizon_timesteps) for _ in range(num_plots)]
172 
173     def callback(self, obs_t, obs_tp1, action, rew, done, info):
174         points = self.data_callback(obs_t, obs_tp1, action, rew, done, info)
175         for point, data_series in zip(points, self.data):
176             data_series.append(point)
177         self.t += 1
178 
179         xmin, xmax = max(0, self.t - self.horizon_timesteps), self.t
180 
181         for i, plot in enumerate(self.cur_plot):
182             if plot is not None:
183                 plot.remove()
184             self.cur_plot[i] = self.ax[i].scatter(range(xmin, xmax), list(self.data[i]))
185             self.ax[i].set_xlim(xmin, xmax)
186         plt.pause(0.000001)
187 
188 
189 if __name__ == '__main__':
190     env = gym.make(args.env)
191     play(env, zoom=4, fps=60)
192
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/gym/utils/play.py b/gym/utils/play.py
--- a/gym/utils/play.py
+++ b/gym/utils/play.py
@@ -1,25 +1,18 @@
 import gym
 import pygame
-import sys
-import time
 import matplotlib
 import argparse
+from gym import logger
 try:
-    matplotlib.use('GTK3Agg')
+    matplotlib.use('TkAgg')
     import matplotlib.pyplot as plt
-except Exception:
-    pass
-
-
-import pyglet.window as pw
+except ImportError as e:
+    logger.warn('failed to set matplotlib backend, plotting will not work: %s' % str(e))
+    plt = None
 
 from collections import deque
-from pygame.locals import HWSURFACE, DOUBLEBUF, RESIZABLE, VIDEORESIZE
-from threading import Thread
+from pygame.locals import VIDEORESIZE
 
-parser = argparse.ArgumentParser()
-parser.add_argument('--env', type=str, default='MontezumaRevengeNoFrameskip-v4', help='Define Environment')
-args = parser.parse_args()
 	
 def display_arr(screen, arr, video_size, transpose):
     arr_min, arr_max = arr.min(), arr.max()
@@ -33,7 +26,7 @@
 
     To simply play the game use:
 
-        play(gym.make("Pong-v3"))
+        play(gym.make("Pong-v4"))
 
     Above code works also if env is wrapped, so it's particularly useful in
     verifying that the frame-level preprocessing does not render the game
@@ -43,12 +36,12 @@
     gym.utils.play.PlayPlot. Here's a sample code for plotting the reward
     for last 5 second of gameplay.
 
-        def callback(obs_t, obs_tp1, rew, done, info):
+        def callback(obs_t, obs_tp1, action, rew, done, info):
             return [rew,]
-        env_plotter = EnvPlotter(callback, 30 * 5, ["reward"])
+        plotter = PlayPlot(callback, 30 * 5, ["reward"])
 
-        env = gym.make("Pong-v3")
-        play(env, callback=env_plotter.callback)
+        env = gym.make("Pong-v4")
+        play(env, callback=plotter.callback)
 
 
     Arguments
@@ -160,6 +153,8 @@
         self.horizon_timesteps = horizon_timesteps
         self.plot_names = plot_names
 
+        assert plt is not None, "matplotlib backend failed, plotting will not work"
+
         num_plots = len(self.plot_names)
         self.fig, self.ax = plt.subplots(num_plots)
         if num_plots == 1:
@@ -181,11 +176,18 @@
         for i, plot in enumerate(self.cur_plot):
             if plot is not None:
                 plot.remove()
-            self.cur_plot[i] = self.ax[i].scatter(range(xmin, xmax), list(self.data[i]))
+            self.cur_plot[i] = self.ax[i].scatter(range(xmin, xmax), list(self.data[i]), c='blue')
             self.ax[i].set_xlim(xmin, xmax)
         plt.pause(0.000001)
 
 
-if __name__ == '__main__':
+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()

{"golden_diff": "diff --git a/gym/utils/play.py b/gym/utils/play.py\n--- a/gym/utils/play.py\n+++ b/gym/utils/play.py\n@@ -1,25 +1,18 @@\n import gym\n import pygame\n-import sys\n-import time\n import matplotlib\n import argparse\n+from gym import logger\n try:\n-    matplotlib.use('GTK3Agg')\n+    matplotlib.use('TkAgg')\n     import matplotlib.pyplot as plt\n-except Exception:\n-    pass\n-\n-\n-import pyglet.window as pw\n+except ImportError as e:\n+    logger.warn('failed to set matplotlib backend, plotting will not work: %s' % str(e))\n+    plt = None\n \n from collections import deque\n-from pygame.locals import HWSURFACE, DOUBLEBUF, RESIZABLE, VIDEORESIZE\n-from threading import Thread\n+from pygame.locals import VIDEORESIZE\n \n-parser = argparse.ArgumentParser()\n-parser.add_argument('--env', type=str, default='MontezumaRevengeNoFrameskip-v4', help='Define Environment')\n-args = parser.parse_args()\n \t\n def display_arr(screen, arr, video_size, transpose):\n     arr_min, arr_max = arr.min(), arr.max()\n@@ -33,7 +26,7 @@\n \n     To simply play the game use:\n \n-        play(gym.make(\"Pong-v3\"))\n+        play(gym.make(\"Pong-v4\"))\n \n     Above code works also if env is wrapped, so it's particularly useful in\n     verifying that the frame-level preprocessing does not render the game\n@@ -43,12 +36,12 @@\n     gym.utils.play.PlayPlot. Here's a sample code for plotting the reward\n     for last 5 second of gameplay.\n \n-        def callback(obs_t, obs_tp1, rew, done, info):\n+        def callback(obs_t, obs_tp1, action, rew, done, info):\n             return [rew,]\n-        env_plotter = EnvPlotter(callback, 30 * 5, [\"reward\"])\n+        plotter = PlayPlot(callback, 30 * 5, [\"reward\"])\n \n-        env = gym.make(\"Pong-v3\")\n-        play(env, callback=env_plotter.callback)\n+        env = gym.make(\"Pong-v4\")\n+        play(env, callback=plotter.callback)\n \n \n     Arguments\n@@ -160,6 +153,8 @@\n         self.horizon_timesteps = horizon_timesteps\n         self.plot_names = plot_names\n \n+        assert plt is not None, \"matplotlib backend failed, plotting will not work\"\n+\n         num_plots = len(self.plot_names)\n         self.fig, self.ax = plt.subplots(num_plots)\n         if num_plots == 1:\n@@ -181,11 +176,18 @@\n         for i, plot in enumerate(self.cur_plot):\n             if plot is not None:\n                 plot.remove()\n-            self.cur_plot[i] = self.ax[i].scatter(range(xmin, xmax), list(self.data[i]))\n+            self.cur_plot[i] = self.ax[i].scatter(range(xmin, xmax), list(self.data[i]), c='blue')\n             self.ax[i].set_xlim(xmin, xmax)\n         plt.pause(0.000001)\n \n \n-if __name__ == '__main__':\n+def main():\n+    parser = argparse.ArgumentParser()\n+    parser.add_argument('--env', type=str, default='MontezumaRevengeNoFrameskip-v4', help='Define Environment')\n+    args = parser.parse_args()\n     env = gym.make(args.env)\n     play(env, zoom=4, fps=60)\n+\n+\n+if __name__ == '__main__':\n+    main()\n", "issue": "Cannot use gym.utils.play.PlayPlot\n```\r\nfrom gym.utils.play import play, PlayPlot\r\ndef callback(obs_t, obs_tp1, action, rew, done, info):\r\n    return [rew,]\r\nenv_plotter = PlayPlot(callback, 30 * 5, [\"reward\"])\r\nenv = gym.make(\"IceHockey-v0\")\r\n```\r\nI got : self.fig, self.ax = plt.subplots(num_plots) NameError: name 'plt' is not defined\n", "before_files": [{"content": "import gym\nimport pygame\nimport sys\nimport time\nimport matplotlib\nimport argparse\ntry:\n    matplotlib.use('GTK3Agg')\n    import matplotlib.pyplot as plt\nexcept Exception:\n    pass\n\n\nimport pyglet.window as pw\n\nfrom collections import deque\nfrom pygame.locals import HWSURFACE, DOUBLEBUF, RESIZABLE, VIDEORESIZE\nfrom threading import Thread\n\nparser = argparse.ArgumentParser()\nparser.add_argument('--env', type=str, default='MontezumaRevengeNoFrameskip-v4', help='Define Environment')\nargs = parser.parse_args()\n\t\ndef display_arr(screen, arr, video_size, transpose):\n    arr_min, arr_max = arr.min(), arr.max()\n    arr = 255.0 * (arr - arr_min) / (arr_max - arr_min)\n    pyg_img = pygame.surfarray.make_surface(arr.swapaxes(0, 1) if transpose else arr)\n    pyg_img = pygame.transform.scale(pyg_img, video_size)\n    screen.blit(pyg_img, (0,0))\n\ndef play(env, transpose=True, fps=30, zoom=None, callback=None, keys_to_action=None):\n    \"\"\"Allows one to play the game using keyboard.\n\n    To simply play the game use:\n\n        play(gym.make(\"Pong-v3\"))\n\n    Above code works also if env is wrapped, so it's particularly useful in\n    verifying that the frame-level preprocessing does not render the game\n    unplayable.\n\n    If you wish to plot real time statistics as you play, you can use\n    gym.utils.play.PlayPlot. Here's a sample code for plotting the reward\n    for last 5 second of gameplay.\n\n        def callback(obs_t, obs_tp1, rew, done, info):\n            return [rew,]\n        env_plotter = EnvPlotter(callback, 30 * 5, [\"reward\"])\n\n        env = gym.make(\"Pong-v3\")\n        play(env, callback=env_plotter.callback)\n\n\n    Arguments\n    ---------\n    env: gym.Env\n        Environment to use for playing.\n    transpose: bool\n        If True the output of observation is transposed.\n        Defaults to true.\n    fps: int\n        Maximum number of steps of the environment to execute every second.\n        Defaults to 30.\n    zoom: float\n        Make screen edge this many times bigger\n    callback: lambda or None\n        Callback if a callback is provided it will be executed after\n        every step. It takes the following input:\n            obs_t: observation before performing action\n            obs_tp1: observation after performing action\n            action: action that was executed\n            rew: reward that was received\n            done: whether the environment is done or not\n            info: debug info\n    keys_to_action: dict: tuple(int) -> int or None\n        Mapping from keys pressed to action performed.\n        For example if pressed 'w' and space at the same time is supposed\n        to trigger action number 2 then key_to_action dict would look like this:\n\n            {\n                # ...\n                sorted(ord('w'), ord(' ')) -> 2\n                # ...\n            }\n        If None, default key_to_action mapping for that env is used, if provided.\n    \"\"\"\n\n    obs_s = env.observation_space\n    assert type(obs_s) == gym.spaces.box.Box\n    assert len(obs_s.shape) == 2 or (len(obs_s.shape) == 3 and obs_s.shape[2] in [1,3])\n\n    if keys_to_action is None:\n        if hasattr(env, 'get_keys_to_action'):\n            keys_to_action = env.get_keys_to_action()\n        elif hasattr(env.unwrapped, 'get_keys_to_action'):\n            keys_to_action = env.unwrapped.get_keys_to_action()\n        else:\n            assert False, env.spec.id + \" does not have explicit key to action mapping, \" + \\\n                          \"please specify one manually\"\n    relevant_keys = set(sum(map(list, keys_to_action.keys()),[]))\n\n    if transpose:\n        video_size = env.observation_space.shape[1], env.observation_space.shape[0]\n    else:\n        video_size = env.observation_space.shape[0], env.observation_space.shape[1]\n\n    if zoom is not None:\n        video_size = int(video_size[0] * zoom), int(video_size[1] * zoom)\n\n    pressed_keys = []\n    running = True\n    env_done = True\n\n    screen = pygame.display.set_mode(video_size)\n    clock = pygame.time.Clock()\n\n\n    while running:\n        if env_done:\n            env_done = False\n            obs = env.reset()\n        else:\n            action = keys_to_action.get(tuple(sorted(pressed_keys)), 0)\n            prev_obs = obs\n            obs, rew, env_done, info = env.step(action)\n            if callback is not None:\n                callback(prev_obs, obs, action, rew, env_done, info)\n        if obs is not None:\n            if len(obs.shape) == 2:\n                obs = obs[:, :, None]\n            if obs.shape[2] == 1:\n                obs = obs.repeat(3, axis=2)\n            display_arr(screen, obs, transpose=transpose, video_size=video_size)\n\n        # process pygame events\n        for event in pygame.event.get():\n            # test events, set key states\n            if event.type == pygame.KEYDOWN:\n                if event.key in relevant_keys:\n                    pressed_keys.append(event.key)\n                elif event.key == 27:\n                    running = False\n            elif event.type == pygame.KEYUP:\n                if event.key in relevant_keys:\n                    pressed_keys.remove(event.key)\n            elif event.type == pygame.QUIT:\n                running = False\n            elif event.type == VIDEORESIZE:\n                video_size = event.size\n                screen = pygame.display.set_mode(video_size)\n                print(video_size)\n\n        pygame.display.flip()\n        clock.tick(fps)\n    pygame.quit()\n\nclass PlayPlot(object):\n    def __init__(self, callback, horizon_timesteps, plot_names):\n        self.data_callback = callback\n        self.horizon_timesteps = horizon_timesteps\n        self.plot_names = plot_names\n\n        num_plots = len(self.plot_names)\n        self.fig, self.ax = plt.subplots(num_plots)\n        if num_plots == 1:\n            self.ax = [self.ax]\n        for axis, name in zip(self.ax, plot_names):\n            axis.set_title(name)\n        self.t = 0\n        self.cur_plot = [None for _ in range(num_plots)]\n        self.data     = [deque(maxlen=horizon_timesteps) for _ in range(num_plots)]\n\n    def callback(self, obs_t, obs_tp1, action, rew, done, info):\n        points = self.data_callback(obs_t, obs_tp1, action, rew, done, info)\n        for point, data_series in zip(points, self.data):\n            data_series.append(point)\n        self.t += 1\n\n        xmin, xmax = max(0, self.t - self.horizon_timesteps), self.t\n\n        for i, plot in enumerate(self.cur_plot):\n            if plot is not None:\n                plot.remove()\n            self.cur_plot[i] = self.ax[i].scatter(range(xmin, xmax), list(self.data[i]))\n            self.ax[i].set_xlim(xmin, xmax)\n        plt.pause(0.000001)\n\n\nif __name__ == '__main__':\n    env = gym.make(args.env)\n    play(env, zoom=4, fps=60)\n", "path": "gym/utils/play.py"}], "after_files": [{"content": "import gym\nimport pygame\nimport matplotlib\nimport argparse\nfrom gym import logger\ntry:\n    matplotlib.use('TkAgg')\n    import matplotlib.pyplot as plt\nexcept ImportError as e:\n    logger.warn('failed to set matplotlib backend, plotting will not work: %s' % str(e))\n    plt = None\n\nfrom collections import deque\nfrom pygame.locals import VIDEORESIZE\n\n\t\ndef display_arr(screen, arr, video_size, transpose):\n    arr_min, arr_max = arr.min(), arr.max()\n    arr = 255.0 * (arr - arr_min) / (arr_max - arr_min)\n    pyg_img = pygame.surfarray.make_surface(arr.swapaxes(0, 1) if transpose else arr)\n    pyg_img = pygame.transform.scale(pyg_img, video_size)\n    screen.blit(pyg_img, (0,0))\n\ndef play(env, transpose=True, fps=30, zoom=None, callback=None, keys_to_action=None):\n    \"\"\"Allows one to play the game using keyboard.\n\n    To simply play the game use:\n\n        play(gym.make(\"Pong-v4\"))\n\n    Above code works also if env is wrapped, so it's particularly useful in\n    verifying that the frame-level preprocessing does not render the game\n    unplayable.\n\n    If you wish to plot real time statistics as you play, you can use\n    gym.utils.play.PlayPlot. Here's a sample code for plotting the reward\n    for last 5 second of gameplay.\n\n        def callback(obs_t, obs_tp1, action, rew, done, info):\n            return [rew,]\n        plotter = PlayPlot(callback, 30 * 5, [\"reward\"])\n\n        env = gym.make(\"Pong-v4\")\n        play(env, callback=plotter.callback)\n\n\n    Arguments\n    ---------\n    env: gym.Env\n        Environment to use for playing.\n    transpose: bool\n        If True the output of observation is transposed.\n        Defaults to true.\n    fps: int\n        Maximum number of steps of the environment to execute every second.\n        Defaults to 30.\n    zoom: float\n        Make screen edge this many times bigger\n    callback: lambda or None\n        Callback if a callback is provided it will be executed after\n        every step. It takes the following input:\n            obs_t: observation before performing action\n            obs_tp1: observation after performing action\n            action: action that was executed\n            rew: reward that was received\n            done: whether the environment is done or not\n            info: debug info\n    keys_to_action: dict: tuple(int) -> int or None\n        Mapping from keys pressed to action performed.\n        For example if pressed 'w' and space at the same time is supposed\n        to trigger action number 2 then key_to_action dict would look like this:\n\n            {\n                # ...\n                sorted(ord('w'), ord(' ')) -> 2\n                # ...\n            }\n        If None, default key_to_action mapping for that env is used, if provided.\n    \"\"\"\n\n    obs_s = env.observation_space\n    assert type(obs_s) == gym.spaces.box.Box\n    assert len(obs_s.shape) == 2 or (len(obs_s.shape) == 3 and obs_s.shape[2] in [1,3])\n\n    if keys_to_action is None:\n        if hasattr(env, 'get_keys_to_action'):\n            keys_to_action = env.get_keys_to_action()\n        elif hasattr(env.unwrapped, 'get_keys_to_action'):\n            keys_to_action = env.unwrapped.get_keys_to_action()\n        else:\n            assert False, env.spec.id + \" does not have explicit key to action mapping, \" + \\\n                          \"please specify one manually\"\n    relevant_keys = set(sum(map(list, keys_to_action.keys()),[]))\n\n    if transpose:\n        video_size = env.observation_space.shape[1], env.observation_space.shape[0]\n    else:\n        video_size = env.observation_space.shape[0], env.observation_space.shape[1]\n\n    if zoom is not None:\n        video_size = int(video_size[0] * zoom), int(video_size[1] * zoom)\n\n    pressed_keys = []\n    running = True\n    env_done = True\n\n    screen = pygame.display.set_mode(video_size)\n    clock = pygame.time.Clock()\n\n\n    while running:\n        if env_done:\n            env_done = False\n            obs = env.reset()\n        else:\n            action = keys_to_action.get(tuple(sorted(pressed_keys)), 0)\n            prev_obs = obs\n            obs, rew, env_done, info = env.step(action)\n            if callback is not None:\n                callback(prev_obs, obs, action, rew, env_done, info)\n        if obs is not None:\n            if len(obs.shape) == 2:\n                obs = obs[:, :, None]\n            if obs.shape[2] == 1:\n                obs = obs.repeat(3, axis=2)\n            display_arr(screen, obs, transpose=transpose, video_size=video_size)\n\n        # process pygame events\n        for event in pygame.event.get():\n            # test events, set key states\n            if event.type == pygame.KEYDOWN:\n                if event.key in relevant_keys:\n                    pressed_keys.append(event.key)\n                elif event.key == 27:\n                    running = False\n            elif event.type == pygame.KEYUP:\n                if event.key in relevant_keys:\n                    pressed_keys.remove(event.key)\n            elif event.type == pygame.QUIT:\n                running = False\n            elif event.type == VIDEORESIZE:\n                video_size = event.size\n                screen = pygame.display.set_mode(video_size)\n                print(video_size)\n\n        pygame.display.flip()\n        clock.tick(fps)\n    pygame.quit()\n\nclass PlayPlot(object):\n    def __init__(self, callback, horizon_timesteps, plot_names):\n        self.data_callback = callback\n        self.horizon_timesteps = horizon_timesteps\n        self.plot_names = plot_names\n\n        assert plt is not None, \"matplotlib backend failed, plotting will not work\"\n\n        num_plots = len(self.plot_names)\n        self.fig, self.ax = plt.subplots(num_plots)\n        if num_plots == 1:\n            self.ax = [self.ax]\n        for axis, name in zip(self.ax, plot_names):\n            axis.set_title(name)\n        self.t = 0\n        self.cur_plot = [None for _ in range(num_plots)]\n        self.data     = [deque(maxlen=horizon_timesteps) for _ in range(num_plots)]\n\n    def callback(self, obs_t, obs_tp1, action, rew, done, info):\n        points = self.data_callback(obs_t, obs_tp1, action, rew, done, info)\n        for point, data_series in zip(points, self.data):\n            data_series.append(point)\n        self.t += 1\n\n        xmin, xmax = max(0, self.t - self.horizon_timesteps), self.t\n\n        for i, plot in enumerate(self.cur_plot):\n            if plot is not None:\n                plot.remove()\n            self.cur_plot[i] = self.ax[i].scatter(range(xmin, xmax), list(self.data[i]), c='blue')\n            self.ax[i].set_xlim(xmin, xmax)\n        plt.pause(0.000001)\n\n\ndef main():\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--env', type=str, default='MontezumaRevengeNoFrameskip-v4', help='Define Environment')\n    args = parser.parse_args()\n    env = gym.make(args.env)\n    play(env, zoom=4, fps=60)\n\n\nif __name__ == '__main__':\n    main()\n", "path": "gym/utils/play.py"}]}