| 38 | |
| 39 | |
| 40 | class Agent: |
| 41 | def __init__(self, Network, ob_space, ac_space, nenvs, nsteps, nstack, |
| 42 | ent_coef=0.01, vf_coef=0.5, max_grad_norm=0.5, lr=7e-4, |
| 43 | alpha=0.99, epsilon=1e-5, total_timesteps=int(80e6)): |
| 44 | config = tf.ConfigProto(intra_op_parallelism_threads=nenvs, |
| 45 | inter_op_parallelism_threads=nenvs) |
| 46 | config.gpu_options.allow_growth = True |
| 47 | sess = tf.Session(config=config) |
| 48 | nbatch = nenvs * nsteps |
| 49 | |
| 50 | A = tf.placeholder(tf.int32, [nbatch]) |
| 51 | ADV = tf.placeholder(tf.float32, [nbatch]) |
| 52 | R = tf.placeholder(tf.float32, [nbatch]) |
| 53 | LR = tf.placeholder(tf.float32, []) |
| 54 | |
| 55 | step_model = Network(sess, ob_space, ac_space, nenvs, 1, nstack, reuse=False) |
| 56 | train_model = Network(sess, ob_space, ac_space, nenvs, nsteps, nstack, reuse=True) |
| 57 | |
| 58 | neglogpac = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=train_model.pi, labels=A) |
| 59 | pg_loss = tf.reduce_mean(ADV * neglogpac) |
| 60 | vf_loss = tf.reduce_mean(tf.squared_difference(tf.squeeze(train_model.vf), R) / 2.0) |
| 61 | entropy = tf.reduce_mean(cat_entropy(train_model.pi)) |
| 62 | loss = pg_loss - entropy * ent_coef + vf_loss * vf_coef |
| 63 | |
| 64 | params = find_trainable_variables("model") |
| 65 | grads = tf.gradients(loss, params) |
| 66 | if max_grad_norm is not None: |
| 67 | grads, grad_norm = tf.clip_by_global_norm(grads, max_grad_norm) |
| 68 | grads_and_params = list(zip(grads, params)) |
| 69 | trainer = tf.train.RMSPropOptimizer(learning_rate=LR, decay=alpha, epsilon=epsilon) |
| 70 | _train = trainer.apply_gradients(grads_and_params) |
| 71 | |
| 72 | def train(states, rewards, actions, values): |
| 73 | advs = rewards - values |
| 74 | feed_dict = {train_model.X: states, A: actions, ADV: advs, R: rewards, LR: lr} |
| 75 | policy_loss, value_loss, policy_entropy, _ = sess.run( |
| 76 | [pg_loss, vf_loss, entropy, _train], |
| 77 | feed_dict |
| 78 | ) |
| 79 | return policy_loss, value_loss, policy_entropy |
| 80 | |
| 81 | def save(save_path): |
| 82 | ps = sess.run(params) |
| 83 | joblib.dump(ps, save_path) |
| 84 | |
| 85 | def load(load_path): |
| 86 | loaded_params = joblib.load(load_path) |
| 87 | restores = [] |
| 88 | for p, loaded_p in zip(params, loaded_params): |
| 89 | restores.append(p.assign(loaded_p)) |
| 90 | ps = sess.run(restores) |
| 91 | |
| 92 | self.train = train |
| 93 | self.train_model = train_model |
| 94 | self.step_model = step_model |
| 95 | self.step = step_model.step |
| 96 | self.value = step_model.value |
| 97 | self.save = save |