rl env render save

1 files changed, 40 insertions, 0 deletions

diff --git a/rl/tutorials/02_env_rendering.py b/rl/tutorials/02_env_rendering.py
new file mode 100644
index 0000000..d05a3ab
--- /dev/null
+++ b/rl/tutorials/02_env_rendering.py
@@ -0,0 +1,40 @@
+
+import gym
+import time
+import matplotlib.pyplot as plt
+from matplotlib import animation
+
+
+env_name = 'CartPole-v0'
+
+env = gym.make(env_name)
+
+state = env.reset()
+done = False
+total_score = 0
+frames = []
+while not done:
+    frames.append(env.render(mode='rgb_array'))
+    action = env.action_space.sample()
+    observation, reward, done, info = env.step(action)
+    total_score += reward
+    time.sleep(0.1)
+
+
+def display_frame_as_gif(frames):
+    plt.figure(figsize=(frames[0].shape[1]/72, frames[0].shape[0]/72), dpi=72)
+    patch = plt.imshow(frames[0])
+    plt.axis('off')
+
+    def animate(i):
+        patch.set_data(frames[i])
+
+    anim = animation.FuncAnimation(plt.gcf(), animate, frames=len(frames), interval=50)
+    anim.save('movie_cartpole.gif')
+
+
+env.close()
+
+print(len(frames), frames[0].shape)
+# display_frame_as_gif(frames)
+