-
Notifications
You must be signed in to change notification settings - Fork 0
/
gymrecord
95 lines (80 loc) · 3.11 KB
/
gymrecord
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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
#!/usr/bin/env python
import play
import gym
from collections import namedtuple
from datetime import datetime
import pickle
import threading
from time import sleep
import argparse
import sys
import numpy as np
Transition = namedtuple('Transition',
('state', 'action', 'reward', 'next_state', 'done'))
class PlayClass(threading.Thread):
def __init__(self, env, fps = 60):
super(PlayClass, self).__init__()
self.env = env
self.fps = fps
def run(self):
play.play(self.env, fps = self.fps, zoom = 4)
class Record(gym.Wrapper):
def __init__(self, env, memory, args, skipframes = 3):
gym.Wrapper.__init__(self, env)
self.memory_lock = threading.Lock()
self.memory = memory
self.args = args
self.skipframes = skipframes
self.current_i = skipframes
def reset(self):
return self.env.reset()
def step(self, action):
self.memory_lock.acquire()
state = self.env.env._get_obs()
next_state, reward, done, info = self.env.step(action)
if self.current_i <= 0:
self.memory.append(Transition(state, action, reward, next_state, done))
self.current_i = self.skipframes
else: self.current_i -= 1
self.memory_lock.release()
return next_state, reward, done, info
def log_transitions(self):
self.memory_lock.acquire()
if len(self.memory) > 0:
basename = self.args['logdir'] + "/{}.{}".format(self.args['environment_name'], datetime.now().strftime("%Y-%m-%d-%H-%M-%s"))
print("Base Filename: ", basename)
state, action, reward, next_state, done = zip(*self.memory)
np.save(basename + "-state.npy", np.array(state), allow_pickle = False)
np.save(basename + "-action.npy", np.array(action), allow_pickle = False)
np.save(basename + "-reward.npy", np.array(reward), allow_pickle = False)
np.save(basename + "-nextstate.npy", np.array(next_state), allow_pickle = False)
np.save(basename + "-done.npy", np.array(done), allow_pickle = False)
self.memory.clear()
self.memory_lock.release()
## Parsing arguments
parser = argparse.ArgumentParser(description="Play and log the environment")
parser.add_argument("--environment_name", type=str, help="The environment name in OpenAI gym to play.")
parser.add_argument("--logdir", type=str, help="Directory to log video and (state, action, reward, next_state, done) in.")
parser.add_argument("--skip", type=int, help="Number of frames to skip logging.")
parser.add_argument("--fps", type=int, help="Number of frames per second")
args = vars(parser.parse_args())
if args['environment_name'] is None or args['logdir'] is None:
parser.print_help()
sys.exit(1)
if args['skip'] is None:
args['skip'] = 3
if args['fps'] is None:
args['fps'] = 30
## Starting the game
memory = []
env = Record(gym.make(args['environment_name']), memory, args, skipframes = args['skip'])
env = gym.wrappers.Monitor(env, args['logdir'], force=True)
playThread = PlayClass(env, args['fps'])
playThread.start()
## Logging portion
while playThread.is_alive():
playThread.join(60)
print("Logging....", end = " ")
env.log_transitions()
# Save what's remaining after process died
env.log_transitions()