-
Notifications
You must be signed in to change notification settings - Fork 2
/
mini_network_dream.py
200 lines (149 loc) · 7.7 KB
/
mini_network_dream.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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
import numpy as np
import tensorflow as tf
import os
import param as P
from algo.ppo import Policy_net, PPOTrain
from rnn.rnn_dream import reset_graph, ConvVAE, HyperParams, DreamModel
ACTION_SPACE = 10
SIZE_1 = 64 # image latent size
SIZE_2 = 20 # non-image obs feature size
model_rnn_size = 512
model_state_space = 2 # includes C and H concatenated if 2, otherwise just H
class MiniNetwork(object):
def __init__(self, sess=None, summary_writer=tf.summary.FileWriter("logs/"), rl_training=False,
reuse=False, cluster=None, index=0, device='/gpu:0',
ppo_load_path=None, ppo_save_path=None, load_worldmodel=True, ntype='dream-model'):
self.policy_model_path_load = ppo_load_path + ntype
self.policy_model_path_save = ppo_save_path + ntype
self.ntype = ntype
self.rl_training = rl_training
self.use_norm = True
self.reuse = reuse
self.sess = sess
self.cluster = cluster
self.index = index
self.device = device
self.input_size = SIZE_1 + SIZE_2 + model_rnn_size*model_state_space
self._create_graph()
self.rl_saver = tf.train.Saver()
self.summary_writer = summary_writer
def initialize(self):
init_op = tf.global_variables_initializer()
self.sess.run(init_op)
def reset_old_network(self):
self.policy_ppo.assign_policy_parameters()
self.policy_ppo.reset_mean_returns()
self.sess.run(self.results_sum.assign(0))
self.sess.run(self.game_num.assign(0))
def _create_graph(self):
if self.reuse:
tf.get_variable_scope().reuse_variables()
assert tf.get_variable_scope().reuse
worker_device = "/job:worker/task:%d" % self.index + self.device
with tf.device(tf.train.replica_device_setter(worker_device=worker_device, cluster=self.cluster)):
self.results_sum = tf.get_variable(name="results_sum", shape=[], initializer=tf.zeros_initializer)
self.game_num = tf.get_variable(name="game_num", shape=[], initializer=tf.zeros_initializer)
self.global_steps = tf.get_variable(name="global_steps", shape=[], initializer=tf.zeros_initializer)
self.win_rate = self.results_sum / self.game_num
self.mean_win_rate = tf.summary.scalar('mean_win_rate_dis', self.results_sum / self.game_num)
self.merged = tf.summary.merge([self.mean_win_rate])
mini_scope = self.ntype
with tf.variable_scope(mini_scope):
ob_space = self.input_size
act_space_array = ACTION_SPACE
self.policy = Policy_net('policy', self.sess, ob_space, act_space_array)
self.policy_old = Policy_net('old_policy', self.sess, ob_space, act_space_array)
self.policy_ppo = PPOTrain('PPO', self.sess, self.policy, self.policy_old, lr=P.mini_lr, epoch_num=P.mini_epoch_num)
var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
self.policy_saver = tf.train.Saver(var_list=var_list)
def Update_result(self, result_list):
win = 0
for i in result_list:
if i > 0:
win += 1
self.sess.run(self.results_sum.assign_add(win))
self.sess.run(self.game_num.assign_add(len(result_list)))
def Update_summary(self, counter):
print("Update summary........")
policy_summary = self.policy_ppo.get_summary_dis()
self.summary_writer.add_summary(policy_summary, counter)
summary = self.sess.run(self.merged)
self.summary_writer.add_summary(summary, counter)
self.sess.run(self.global_steps.assign(counter))
print("Update summary finished!")
steps = int(self.sess.run(self.global_steps))
win_game = int(self.sess.run(self.results_sum))
all_game = int(self.sess.run(self.game_num))
win_rate = win_game / float(all_game) if all_game != 0 else 0.
return steps, win_rate
def get_win_rate(self):
return float(self.sess.run(self.win_rate))
def Update_policy(self, buffer):
self.policy_ppo.ppo_train_dis(buffer.observations, buffer.tech_actions,
buffer.rewards, buffer.values, buffer.values_next, buffer.gaes, buffer.returns, verbose=True)
def get_global_steps(self):
return int(self.sess.run(self.global_steps))
def save_policy(self):
self.policy_saver.save(self.sess, self.policy_model_path_save)
print("policy has been saved in", self.policy_model_path_save)
def restore_policy(self):
self.policy_saver.restore(self.sess, self.policy_model_path_load)
print("Restore policy from", self.policy_model_path_load)
model_path_name = 'tf_models'
SIZE_1 = 64 # image latent size
SIZE_2 = 20 # non-image obs feature size
model_rnn_size = 512
model_num_mixture = 5
model_restart_factor = 10.
def default_hps():
return HyperParams(num_steps=2000, # train model for 2000 steps.
max_seq_len=300, # train on sequences of 300
seq_width=SIZE_1, # width of our data (64)
rnn_size=model_rnn_size, # number of rnn cells
batch_size=100, # minibatch sizes
grad_clip=1.0,
num_mixture=model_num_mixture, # number of mixtures in MDN
restart_factor=model_restart_factor, # factor of importance for restart=1 rare case for loss.
learning_rate=0.001,
decay_rate=1.0,
min_learning_rate=0.00001,
use_layer_norm=0, # set this to 1 to get more stable results (less chance of NaN), but slower
use_recurrent_dropout=0,
recurrent_dropout_prob=0.90,
use_input_dropout=0,
input_dropout_prob=0.90,
use_output_dropout=0,
output_dropout_prob=0.90,
is_training=1)
hps_model = default_hps()
hps_sample = hps_model._replace(batch_size=1, max_seq_len=2, use_recurrent_dropout=0, is_training=0)
class SecondNetwork(object):
def __init__(self, sess=None, rl_training=False, index=0,
reuse=False, cluster=None, device='/gpu:0',
load_model=True, net_path_name=model_path_name, ntype='assist-model'):
self.index = index
#reset_graph()
self.vae = ConvVAE(batch_size=1, gpu_mode=False, is_training=False, reuse=True)
self.rnn = DreamModel(hps_sample, gpu_mode=False, reuse=True)
if load_model:
self.vae.load_json(os.path.join(net_path_name, 'vae.json'))
self.rnn.load_json(os.path.join(net_path_name, 'rnn.json'))
self.outwidth = SIZE_1 + SIZE_2
def rnn_init_state(self):
return self.rnn.sess.run(self.rnn.initial_state)
def rnn_next_state(self, feature, action, reward, prev_state):
prev_feature = np.zeros((1, 1, self.outwidth))
prev_feature[0][0] = feature
prev_action = np.zeros((1, 1))
prev_action[0] = action
prev_reward = np.ones((1, 1))
prev_reward[0] = reward
feed = {self.rnn.input_z: prev_feature[:,:,:self.rnn.hps.seq_width],
self.rnn.input_obs: prev_feature[:,:,self.rnn.hps.seq_width:],
self.rnn.input_reward: prev_reward,
self.rnn.input_action: prev_action,
self.rnn.initial_state: prev_state
}
return self.rnn.sess.run(self.rnn.final_state, feed)
def rnn_output(self, hidden_state, feature):
return np.concatenate([feature, np.concatenate((hidden_state.c, hidden_state.h), axis=1)[0]])