【Python】保存gym截圖

如果想做基於圖像cnn的深度強化學習，需要拿到gym的截圖，下面是兩種截圖方法。

1. 利用render結果生成圖像：

import gym
import warnings
import os
from PIL import Image

warnings.filterwarnings("ignore")

if os.path.exists('img')==False:
    os.mkdir('img')

env = gym.make("CartPole-v1")

num = 0
for i in range(3000):
    state = env.reset()

    while True:
        num +=1
        screen = env.render(mode='rgb_array')
        img = Image.fromarray(screen)

        filename = './img/screen_image'+str(num)+'.png'
        print('save img at '+filename)
        img.save(filename)

        action = env.action_space.sample()
        next_state, reward, done, _ = env.step(action)

        if done:
            break

2.  針對atari仿真，沒有render，可以用env.env.ale.saveScreenPNG生成圖像：

import gym
import warnings
import time
import os

warnings.filterwarnings("ignore")

if os.path.exists('img')==False:
    os.mkdir('img')

env = gym.make('Pong-v0', render_mode='human')

obs = env.reset()
for i in range(1000):

    env.env.ale.saveScreenPNG('./img/screen_image'+str(i)+'.png')

    action = env.action_space.sample()
    obs, reward, done, info = env.step(action)
    time.sleep(0.01)
env.close()