-
-
Save rudrasohan/1296f1e0ee6f39942476712a8a91b86d to your computer and use it in GitHub Desktop.
Save OpenAI Gym renders as GIFS
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from matplotlib import animation | |
import matplotlib.pyplot as plt | |
import gym | |
""" | |
Ensure you have imagemagick installed with | |
sudo apt-get install imagemagick | |
Open file in CLI with: | |
xgd-open <filelname> | |
""" | |
def save_frames_as_gif(frames, path='./', filename='gym_animation.gif'): | |
#Mess with this to change frame size | |
plt.figure(figsize=(frames[0].shape[1] / 72.0, frames[0].shape[0] / 72.0), 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(path + filename, writer='imagemagick', fps=60) | |
#Make gym env | |
env = gym.make('CartPole-v1') | |
#Run the env | |
observation = env.reset() | |
frames = [] | |
for t in range(1000): | |
#Render to frames buffer | |
frames.append(env.render(mode="rgb_array")) | |
action = env.action_space.sample() | |
_, _, done, _ = env.step(action) | |
if done: | |
break | |
env.close() | |
save_frames_as_gif(frames) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment