-
Notifications
You must be signed in to change notification settings - Fork 2
/
test.py
69 lines (51 loc) · 1.61 KB
/
test.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
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
from __future__ import print_function
import os
import gym
import numpy as np
import torch
from train import data_transform, actions_set, Net, DATA_DIR, MODEL_FILE
from pyglet.window import key
def play(model):
"""
Let the agent play
:param model: the network
"""
env = gym.make('CarRacing-v0')
# use ESC to exit
global exit_test
exit_test = False
def key_press(k, mod):
global exit_test
if k == key.ESCAPE: exit_test = True
# initialize environment
state = env.reset()
env.unwrapped.viewer.window.on_key_press = key_press
while 1:
env.render()
state = np.moveaxis(state, 2, 0) # channel first image
# numpy to tensor
state = torch.from_numpy(np.flip(state, axis=0).copy())
state = data_transform(state) # apply transformations
state = state.unsqueeze(0) # add additional dimension
# forward
with torch.set_grad_enabled(False):
outputs = model(state)
normalized = torch.nn.functional.softmax(outputs, dim=1)
# translate from net output to env action
max_action = np.argmax(normalized.cpu().numpy()[0])
action = actions_set[max_action]
# adjust brake power
if action[2] != 0:
action[2] = 0.3
state, _, terminal, _ = env.step(action) # one step
# if terminal:
# env.close()
# return
if exit_test:
env.close()
return
if __name__ == '__main__':
m = Net()
m.load_state_dict(torch.load(os.path.join(DATA_DIR, MODEL_FILE)))
m.eval()
play(m)