-
Notifications
You must be signed in to change notification settings - Fork 2
/
video.py
40 lines (33 loc) · 1.23 KB
/
video.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
# This piece of code was copied & modified from the following source:
#
# Title: CURL: Contrastive Unsupervised Representation Learning for Sample-Efficient Reinforcement Learning
# Author: Laskin, Michael and Srinivas, Aravind and Abbeel, Pieter
# Date: 2020
# Availability: https://github.com/MishaLaskin/curl
import imageio
import os
import numpy as np
class VideoRecorder(object):
def __init__(self, dir_name, fps=20):
self.dir_name = dir_name
self.fps = fps
self.frames = []
def init(self, enabled=True):
self.frames = []
self.enabled = self.dir_name is not None and enabled
def record(self, env):
if self.enabled:
try:
frame = env.render()
except:
print('Warning: failed to render environment, setting frame to None')
frame = None
self.frames.append(frame)
def save(self, file_name):
if self.enabled:
path = os.path.join(self.dir_name, file_name)
kargs = {'macro_block_size': None}
try:
imageio.mimsave(path, self.frames, fps=self.fps, **kargs)
except:
print('Warning: failed to save video')