render gym by pygame and save to local
通过pygame渲染gym,并且将图片保存为视频到本地。
代码有问题,无法取消创建gym自带的渲染窗口,暂时宣布放弃这套方案。
等我找到别人成熟的方案再说。。。
前言:
gym自带的渲染虽然可以调整视角,大小什么的,但是不方便保存为视频,并且有时候我们并不需要这么完整的视角。
对于图像输入的强化学习来说,如果要分析为什么智能体学不到东西,就必须要看看智能体到底在怎么瞎折腾,因此,我们需要一个即时渲染,且能保存视频的脚本。
综合来看,pygame的即时渲染方案是最佳的,pygame库基本没什么乱七八糟的依赖,刷新一次的时间消耗非常短,比cv2和matplotlib好多了,且可以方便的加一些配套信息,比如epoch,step,reward信息,我看到好几个代码库替代渲染方案也是基于pygame。
下面贴一下当初设计的功能:
功能描述:
1.传入特定时刻的env,渲染出RGB图,可以选择,是否将其保存为一个视频流
2.需要用pygame可视化当前图
3.不需要pygame乱七八糟的功能
4.视频保存路径和当前实验log路径一致
5.视频名称需要标注好epoch
代码:
"""
功能描述:
1.传入特定时刻的env,渲染出RGB图,可以选择,是否将其保存为一个小视频
2.需要用pygame可视化当前图
3.不需要pygame乱七八糟的功能
4.视频保存路径和当前实验log路径一致
5.视频名称需要标注好epoch
"""
import pygame
import os
from pygame.locals import *
from sys import exit
import numpy as np
import cv2
import imutils
class GymRenderImageSaveVideoClass:
def __init__(self,
exp_path='',
fps=20,
render_in_gym=False,
save_flag=True,
args=None,
):
self.render_in_gym = render_in_gym
self.args = args
# pygame init
self.start_epoch = 0
self.image_width = args.image_width
self.image_height = args.image_height
self.font_size = args.image_width // 15
pygame.init()
self.break_flag = False
self.image_count = 0
self.position_font = pygame.font.SysFont("幼圆", self.font_size)
# 返回一个窗口Surface对象
self.screen = pygame.display.set_mode((self.image_width,
self.image_height),
0, 32)
self.exp_name = args.exp_name
# save video:
self.save_flag = save_flag
if self.save_flag:
self.exp_path = exp_path
self.max_store = args.max_steps-1
try:
os.mkdir(self.exp_path)
except Exception as e:
print(e)
self.fourcc = cv2.VideoWriter_fourcc(*'XVID') # 保存视频的编码
self.fps = fps
def update(self, image, epoch=0, step=0, reward=0):
# 判断是否退出
for event in pygame.event.get():
if event.type == QUIT:
self.break_flag = True
# 在窗口标题上显示参数元组
pygame.display.set_caption(self.exp_name)
robot_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# 图片可能需要旋转
robot_image = imutils.rotate(robot_image, 90, )
# 图片垂直翻转
robot_image = cv2.flip(robot_image, 0)
robot_image = cv2.resize(robot_image,
(self.image_height,
self.image_width))
my_surface = pygame.pixelcopy.make_surface(robot_image)
self.screen.blit(my_surface, (0, 0))
pygame_str = 'epoch:' + str(epoch) + '-step:' + str(step) + '-rew:' + str(np.round(reward, 3))
text1 = self.position_font.render(pygame_str, True,
(255, 0, 0), (0, 0, 0))
self.screen.blit(text1,
[10, 40],
)
pygame.display.update()
return robot_image
def step(self, epoch='', step='', reward=0.0,
env=None, specfic_name='', per_epoch=False,
img=None,
):
if self.image_count == 0 and self.save_flag:
video_name = self.args.exp_name + '_st_ep' + str(int(epoch))
video_path = os.path.join(self.exp_path, specfic_name+'_'+video_name + '.avi')
self.out = cv2.VideoWriter(video_path, self.fourcc, self.fps,
(self.image_width,
self.image_height))
if env is None:
robot_image = np.uint8(np.random.random((self.image_width, self.image_height, 3)))
else:
if img is not None:
robot_image = np.flip(img, 0)
robot_image = cv2.cvtColor(robot_image, cv2.COLOR_RGB2BGR)
else:
# TODO 如何只获取env.render()的图片,不创建窗口?
# env.unwrapped.render()
if self.render_in_gym:
env.render('human')
# robot_image = np.uint8(np.random.random((self.image_width, self.image_height, 3)))
try:
robot_image = env.render("rgb_array",
width=self.image_width,
height=self.image_height,
)
except:
robot_image = env.render("rgb_array",
)
# 判断是否退出
for event in pygame.event.get():
if event.type == QUIT:
self.break_flag = True
if self.break_flag:
pygame.quit()
import sys
print("exit")
sys.exit()
robot_image = cv2.cvtColor(robot_image, cv2.COLOR_RGB2BGR)
self.update(image=robot_image,
epoch=epoch,
step=step,
reward=reward
)
if self.save_flag:
# save img to local as video stream, and add some info to img.
post_img = cv2.resize(robot_image,
(self.image_width,
self.image_height))
fontScale = 1
text_thickness = 1
fontSize = 1
bg_color = (255, 0, 0)
fontFace = cv2.FONT_HERSHEY_SIMPLEX
cv2.putText(post_img, 'epoch:' + str(epoch),
(10, 20), fontFace, fontScale,
bg_color, text_thickness, fontSize)
cv2.putText(post_img, 'step:' + str(step),
(10, 40), fontFace, fontScale,
bg_color, text_thickness, fontSize
)
cv2.putText(post_img, 'rew:' + str(np.round(reward, 3)),
(10, 60), fontFace, fontScale,
bg_color, text_thickness, fontSize)
# cv2.imshow("post_img:", post_img)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
post_img = np.uint8(post_img)
self.out.write(post_img)
self.image_count += 1
if per_epoch:
if step == self.max_store - 1:
self.image_count = 0
self.out.release()
def main():
import time
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--image_width', type=int, default=256)
parser.add_argument('--image_height', type=int, default=256)
parser.add_argument('--batch_size', type=int, default=128)
parser.add_argument('--gamma', type=float, default=0.9)
parser.add_argument('--epochs', type=int, default=20)
parser.add_argument('--max_steps', type=int, default=20)
parser.add_argument('--exp_name', type=str, default='gym_render_by_pygame')
parser.add_argument('--output_dir', type=str, default='Gym_render_by_pygame')
args = parser.parse_args()
logger_kwargs = {'exp_name': args.exp_name,
'output_dir': args.output_dir}
dpc = GymRenderImageSaveVideoClass(exp_path=logger_kwargs['output_dir'],
save_flag=True,
render_in_gym=False,
args=args,
)
import matplotlib.pyplot as plt
st_list = []
import gym
from gym.wrappers import Monitor
def wrap_env(env):
env = Monitor(env, './video', force=True)
return env
env = gym.make("MountainCar-v0",)
env = wrap_env(env)
for i in range(args.epochs):
env.reset()
for j in range(args.max_steps):
st = time.time()
obs, r, done, info = env.step(env.action_space.sample())
dpc.step(epoch=i, step=j, reward=r,
env=env, specfic_name='train',
per_epoch=False)
# if dpc.break_flag:
# pygame.quit()
# dpc.out.release()
# cv2.destroyAllWindows()
# exit()
dpc.out.release()
cv2.destroyAllWindows()
if __name__ == "__main__":
main()