-
Notifications
You must be signed in to change notification settings - Fork 19
/
test.py
86 lines (70 loc) · 2.48 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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import math
import os
import sys
import pickle
import numpy as np
import random
import torch
import torch.nn.functional as F
import torch.optim as optim
import env_wrapper
from model import ActorCritic
from torch.autograd import Variable
from torchvision import datasets, transforms
import time
from collections import deque
def test(rank, args, shared_model):
torch.manual_seed(args.seed + rank)
env = env_wrapper.create_doom(args.record, outdir=args.outdir)
model = ActorCritic(env.observation_space.shape[0], env.action_space)
model.eval()
state = env.reset()
state = torch.from_numpy(state)
reward_sum = 0
done = True
start_time = time.time()
# a quick hack to prevent the agent from stucking
actions = deque(maxlen=2100)
episode_length = 0
result = []
while True:
episode_length += 1
# Sync with the shared model
if done:
model.load_state_dict(shared_model.state_dict())
cx = Variable(torch.zeros(1, 256), volatile=True)
hx = Variable(torch.zeros(1, 256), volatile=True)
else:
cx = Variable(cx.data, volatile=True)
hx = Variable(hx.data, volatile=True)
value, logit, (hx, cx) = model(
(Variable(state.unsqueeze(0), volatile=True), (hx, cx)),
icm = False
)
prob = F.softmax(logit)
action = prob.max(1)[1].data.numpy()
state, reward, done, _ = env.step(action[0, 0])
state = torch.from_numpy(state)
done = done or episode_length >= args.max_episode_length
reward_sum += reward
# a quick hack to prevent the agent from stucking
actions.append(action[0, 0])
if actions.count(actions[0]) == actions.maxlen:
done = True
if done:
end_time = time.time()
print("Time {}, episode reward {}, episode length {}".format(
time.strftime("%Hh %Mm %Ss",
time.gmtime(end_time - start_time)),
reward_sum, episode_length))
result.append((reward_sum, end_time - start_time))
f = open('output/result.pickle','w')
pickle.dump(result, f)
f.close()
torch.save(model.state_dict(), 'output/{}.pth'.format((end_time - start_time)))
reward_sum = 0
episode_length = 0
actions.clear()
state = env.reset()
state = torch.from_numpy(state)
time.sleep(60)