Gym CartPole 倒立摆 保存 gif 动画图像

render_mode="human" 显示画面

render_mode="rgb_array" 不显示画面

frames 动画保存，需要rgb_array模式。因此采用cv2进行渲染，解决rgb_array模式下画面显示问题。

import gym
import matplotlib.pyplot as plt
from matplotlib import animation
import cv2

# 保存gif图像
def display_frames_as_gif(frames, savepath='./pic/test.gif'):
    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=1)
    anim.save(savepath, writer='pillow', fps=30)


frames = []
# CartPole 倒立摆
env = gym.make('CartPole-v1', render_mode="rgb_array")
for episode in range(10):
    env.reset()
    print("Episode finished after {} timesteps".format(episode))
    for _ in range(100):
        screen = env.render()

        # RGB转化为BGR, cv2显示格式为BGR
        img = cv2.cvtColor(screen, cv2.COLOR_RGB2BGR)
        # 显示画面 CartPole为窗口名称
        cv2.imshow("CartPole", img)
        # 给cv2一定时间完成渲染，否则无法显示
        # cv2.waitkey(x) x的单位是ms毫秒
        cv2.waitKey(10)

        frames.append(screen)  # 加载各个时刻图像到帧

        env.step(env.action_space.sample())

env.close()

display_frames_as_gif(frames, savepath='./pic/cartpole_result.gif')  # 保存运行结果动图