Save OpenAI Gym renders as GIFS

gym_to_gif.py

from matplotlib import animation
import matplotlib.pyplot as plt
import gym 
import collections
import os

"""
Ensure you have imagemagick installed with 
sudo apt-get install imagemagick

Open file in CLI with:
xgd-open <filelname>
"""
#the list of frames is a named tuple with the matrix, action, and reward at each index
Frame = collections.namedtuple('Frame', ['image', 'action', 'reward'])

def save_frames_as_gif(frames, path='./', filename='gym_animation.gif'):

    #Mess with this to change frame size
    plt.figure(figsize=(frames[0].image.shape[1] / 72.0, frames[0].image.shape[0] / 72.0), dpi=72)
    
    #allows for the extra display
    ax = plt.gcf().add_subplot()

    plt.axis('off')
    
    patch = plt.imshow(frames[0].image)
    
    #configure extra display
    at = ax.text(x=0.05, y=0.5, s=f"Action: {frames[0].action}\nReward: {frames[0].reward}", fontsize=18 \
        , bbox={'facecolor': 'grey', 'alpha': 0.5, 'pad': 10}, fontweight='black', fontfamily='serif')

    #added extra display to animate method
    def animate(i):
        patch.set_data(frames[i].image)
        at.set_text(s=f"Action: {frames[i].action}\nReward: {frames[i].reward}")
                          
    anim = animation.FuncAnimation(plt.gcf(), animate, frames = len(frames), interval=100)
    anim.save(path + filename, writer='imagemagick', fps=30)
    

def main():
    #Make gym env
    env = gym.make('CartPole-v1')

    #Run the env
    observation = env.reset()
    frames = []
    printAction = ['←', '→']
    for t in range(1000):
        #Render to frames buffer
        action = env.action_space.sample()
        frames.append(Frame(env.render(mode='rgb_array'), printAction[int(action)], t))
        _, _, done, _ = env.step(action)
        if done:
            break
    env.close()
    
    #handles filename collisions
    path='.' + os.sep
    filename='gym_animation.gif'
    n = 0
    while filename in os.listdir(path):
        n += 1
        filename = 'gym_animation' + str(n) + '.gif'
    save_frames_as_gif(frames, filename=filename, path=path)

if __name__ == '__main__':
    main()

I added a display that gives the current action and reward at each step