Skip to content

Instantly share code, notes, and snippets.

@JackBosco
Forked from botforge/gym_to_gif.py
Created December 28, 2022 23:21
Show Gist options
  • Save JackBosco/9075c78b3dac43b280502a195b06f372 to your computer and use it in GitHub Desktop.
Save JackBosco/9075c78b3dac43b280502a195b06f372 to your computer and use it in GitHub Desktop.
Save OpenAI Gym renders as GIFS
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