save vid lab3

Author: aamini

mitdeeplearning/lab3.py

@@ -1,7 +1,7 @@
 import io
 import base64
 from IPython.display import HTML
-
+import gym
 
 def play_video(filename):
     encoded = base64.b64encode(io.open(filename, 'r+b').read())
@@ -19,3 +19,42 @@ def preprocess_pong(image):
     I[I == 109] = 0 # Remove background type 2
     I[I != 0] = 1 # Set remaining elements (paddles, ball, etc.) to 1
     return I.astype(np.float).reshape(80, 80, 1)
+
+
+def save_video_of_model(model, env_name, obs_diff=False, pp_fn=None):
+    import skvideo.io
+    from pyvirtualdisplay import Display
+    display = Display(visible=0, size=(400, 300))
+    display.start()
+
+    if pp_fn is None:
+        pp_fn = lambda x: x
+
+    env = gym.make(env_name)
+    obs = env.reset()
+    obs = pp_fn(obs)
+    prev_obs = obs
+
+    filename = env_name + ".mp4"
+    output_video = skvideo.io.FFmpegWriter(filename)
+
+    counter = 0
+    done = False
+    while not done:
+        frame = env.render(mode='rgb_array')
+        output_video.writeFrame(frame)
+
+        if obs_diff:
+            input_obs = obs - prev_obs
+        else:
+            input_obs = obs
+        action = model(np.expand_dims(input_obs, 0)).numpy().argmax()
+
+        prev_obs = obs
+        obs, reward, done, info = env.step(action)
+        obs = pp_fn(obs)
+        counter += 1
+
+  output_video.close()
+  print("Successfully saved {} frames into {}!".format(counter, filename))
+  return filename

setup.py

@@ -21,13 +21,13 @@ def get_dist(pkgname):
 setup(
   name = 'mitdeeplearning',         # How you named your package folder (MyLib)
   packages = ['mitdeeplearning'],   # Chose the same as "name"
-  version = '0.5.1',      # Start with a small number and increase it with every change you make
+  version = '0.5.2',      # Start with a small number and increase it with every change you make
   license='MIT',        # Chose a license from here: https://help.github.com/articles/licensing-a-repository
   description = 'Official software labs for MIT Introduction to Deep Learning (http://introtodeeplearning.com)',   # Give a short description about your library
   author = 'Alexander Amini',                   # Type in your name
   author_email = '[email protected]',      # Type in your E-Mail
   url = 'http://introtodeeplearning.com',   # Provide either the link to your github or to your website
-  download_url = 'https://github.com/aamini/introtodeeplearning_labs/archive/v0.5.1.tar.gz',    # I explain this later on
+  download_url = 'https://github.com/aamini/introtodeeplearning_labs/archive/v0.5.2.tar.gz',    # I explain this later on
   keywords = ['deep learning', 'neural networks', 'tensorflow', 'introduction'],   # Keywords that define your package best
   install_requires=install_deps,
   classifiers=[