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',render_mode="rgb_array")
#Run the env
observation = env.reset()
frames = []
for t in range(1000):
#Render to frames buffer
frames.append(env.render())
action = env.action_space.sample()
_, _, done, _ = env.step(action)
if done:
break
env.close()
save_frames_as_gif(frames)
Note: Please not, render_mode para should be setted in the gym.mak()
Besides, if you want to save the frame from env.render() , and make it as gif format.
You should make sure that the render_mode is "rgb_array" . if you set the render_mode as "human", you can see the result when you run the code. But the render() will return none.
You can not get the array of the frame and cannot generate the gif format file.