fix play utility

- import pyglet window to fix mysterious issue with ui focus
  (hat tip @jeanharb)
- make matplotlib backend take effect by picking before plot import
- fix misc. issues with example usage

gym/utils/play.py

@@ -3,17 +3,19 @@
 import sys
 import time
 import matplotlib
-import matplotlib.pyplot as plt
-
-from collections import deque
-from pygame.locals import HWSURFACE, DOUBLEBUF, RESIZABLE, VIDEORESIZE
-from threading import Thread
-
 try:
     matplotlib.use('GTK3Agg')
+    import matplotlib.pyplot as plt
 except Exception:
     pass
 
+
+import pyglet.window as pw
+
+from collections import deque
+from pygame.locals import HWSURFACE, DOUBLEBUF, RESIZABLE, VIDEORESIZE
+from threading import Thread
+
 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)
@@ -180,14 +182,10 @@ def callback(self, obs_t, obs_tp1, action, rew, done, info):
 
 
 if __name__ == '__main__':
-    from rl_algs.common.atari_wrappers import wrap_deepmind
-
     def callback(obs_t, obs_tp1, action, rew, done, info):
         return [rew, obs_t.mean()]
     env_plotter = EnvPlotter(callback, 30 * 5, ["reward", "mean intensity"])
 
     env = gym.make("MontezumaRevengeNoFrameskip-v3")
-    env = wrap_deepmind(env)
-
-    play_env(env, zoom=4, callback=env_plotter.callback, fps=30)
 
+    play(env, zoom=4, callback=env_plotter.callback, fps=30)