jupyter_animate.py

from IPython.display import display
from JSAnimation.IPython_display import display_animation
from matplotlib import animation

def run_env_steps(env, num_steps=1000):
    env.reset()
    frames = []
    for _ in range(num_steps):
        # Render into buffer. 
        # You will still see the window.
        frames.append(env.render(mode = 'rgb_array'))
        action = env.action_space.sample()
        observation, reward, done, info = env.step(action)
        if done:
            break
    env.render(close=True)
    return frames

def display_frames_as_gif(frames):
    """
    Displays a list of frames as a gif, with controls
    """
    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)
    display(display_animation(anim, default_mode='loop'))