Make GIF of OpenAI gym environment

make_gym_gif.py

from jutility import util, plotting
import gymnasium as gym

def main():
    table = util.Table(
        util.CountColumn("c"),
        util.TimeColumn("t"),
        print_interval=util.TimeInterval(1),
    )

    env = gym.make("HalfCheetah-v4", render_mode="rgb_array")
    gif = plotting.Gif()
    state, info = env.reset(seed=0)
    env.action_space.seed(0)
    env.observation_space.seed(0)

    while True:
        table.update()

        action = env.action_space.sample()

        state, reward, terminated, truncated, info = env.step(action)
        gif.add_rgb_array_frame(env.render())
        if terminated or truncated:
            break

    env.close()
    gif.save(output_name="make_gym_gif_output", dir_name=".")

if __name__ == "__main__":
    with util.Timer("main function"):
        main()

make_gym_gif_output.gif