From 3f13ebc711fd7f96f9ef0077244b8d53307923f1 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Wed, 3 May 2017 19:48:03 -0600 Subject: [PATCH] Crossprop --- .gitignore | 3 ++ common.py | 39 ++++++++++++++++- crossprop.py | 115 +++++++++++++++++++++++++++++++++++++++++++++++++++ network.py | 12 ++++++ task.py | 22 +++++++++- 5 files changed, 188 insertions(+), 3 deletions(-) create mode 100644 crossprop.py diff --git a/.gitignore b/.gitignore index a8caa7d..36f0f4d 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,9 @@ __pycache__/ *$py.class .idea +exp_* +upload.py +*.sh # C extensions *.so diff --git a/common.py b/common.py index d2a2dcf..24cec97 100644 --- a/common.py +++ b/common.py @@ -16,4 +16,41 @@ def fully_connected(model_name, layer_name, var_in, dim_in, dim_out, initializer=initializer) net = tf.nn.bias_add(tf.matmul(var_in, W), b) phi = transfer(net) - return W, b, net, phi \ No newline at end of file + return W, b, net, phi + +class Relu: + def __init__(self): + self.gate_fun = tf.nn.relu + self.gate_fun_gradient = \ + lambda phi, net: tf.where(net >= 0, tf.ones(tf.shape(net)), tf.zeros(tf.shape(net))) + + +class Tanh: + def __init__(self): + self.gate_fun = tf.tanh + self.gate_fun_gradient = \ + lambda phi, net: tf.subtract(1.0, tf.pow(phi, 2)) + +class Identity: + def __init__(self): + self.gate_fun = tf.identity + self.gate_fun_gradient = \ + lambda phi, net: tf.ones(tf.shape(phi)) + +def crossprop_layer(model_name, layer_name, var_in, dim_in, dim_hidden, dim_out, gate_fun, initializer): + with tf.variable_scope(model_name): + with tf.variable_scope(layer_name): + U = tf.get_variable('U', [dim_in, dim_hidden], + initializer=initializer) + b_hidden = tf.get_variable('b_hidden', [dim_hidden], + initializer=initializer) + W = tf.get_variable('W', [dim_hidden, dim_out], + initializer=initializer) + b_out = tf.get_variable('b_out', [dim_out], + initializer=initializer) + net = tf.matmul(var_in, U) + net = tf.nn.bias_add(net, b_hidden) + phi = gate_fun(net) + y = tf.matmul(phi, W) + y = tf.nn.bias_add(y, b_out) + return U, b_hidden, net, phi, W, b_out, y \ No newline at end of file diff --git a/crossprop.py b/crossprop.py new file mode 100644 index 0000000..adf1415 --- /dev/null +++ b/crossprop.py @@ -0,0 +1,115 @@ +####################################################################### +# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) # +# Permission given to modify the code as long as you keep this # +# declaration at the top # +####################################################################### + +import tensorflow as tf +from common import * +import numpy as np + +class CrossProp: + def __init__(self, name, dim_in, dim_hidden, dim_out, learning_rate, gate=Relu(), + initializer=tf.random_normal_initializer(), bottom_layer=None, + optimizer=None, output_layer='MSE'): + + self.learning_rate = learning_rate + self.lam = 0 + self.output_layer = output_layer + if optimizer is None: + optimizer = tf.train.GradientDescentOptimizer(learning_rate=self.learning_rate) + + if bottom_layer is None: + self.x = tf.placeholder(tf.float32, shape=(None, dim_in)) + var_in = self.x + trainable_vars = [] + else: + self.x = bottom_layer.x + var_in = bottom_layer.var_out + trainable_vars = bottom_layer.trainable_vars + self.bottom_layer = bottom_layer + + self.h = tf.placeholder(tf.float32, shape=(dim_hidden, dim_out)) + + self.target = tf.placeholder(tf.float32, shape=(None, dim_out)) + + U, b_hidden, net, phi, W, b_out, y =\ + crossprop_layer(name, 'crossprop_layer', var_in, dim_in, dim_hidden, dim_out, gate.gate_fun, initializer) + if self.output_layer == 'CE': + self.pred = tf.nn.softmax(y) + ce_loss = tf.nn.softmax_cross_entropy_with_logits(logits=y, labels=self.target) + self.loss = tf.reduce_mean(ce_loss) + self.total_loss = tf.reduce_sum(ce_loss) + correct_prediction = tf.equal(tf.argmax(self.target, 1), tf.argmax(self.pred, 1)) + self.correct_labels = tf.reduce_sum(tf.cast(correct_prediction, "float")) + delta = tf.subtract(self.pred, self.target) + elif self.output_layer == 'MSE': + se_loss = 0.5 * tf.squared_difference(y, self.target) + self.loss = tf.reduce_mean(se_loss) + self.total_loss = tf.reduce_sum(se_loss) + delta = y - self.target + self.correct_labels = tf.constant(0) + else: + assert False + + trainable_vars.extend([W, b_out]) + + h_decay = tf.subtract(1.0, tf.scalar_mul(learning_rate, tf.pow(phi, 2))) + h_decay = tf.reshape(tf.tile(h_decay, [1, tf.shape(self.h)[1]]), [-1, tf.shape(self.h)[1], tf.shape(self.h)[0]]) + h_decay = tf.transpose(h_decay, [0, 2, 1]) + self.h_decay = tf.reduce_sum(h_decay, axis=0) + + h_delta = tf.reshape(tf.tile(delta, [1, tf.shape(self.h)[0]]), [-1, tf.shape(self.h)[0], tf.shape(self.h)[1]]) + self.h_delta = tf.reduce_sum(h_delta, axis=0) + + new_grads = [] + phi_phi_grad = tf.multiply(phi, gate.gate_fun_gradient(phi, net)) + weight = tf.transpose(tf.matmul(self.h, tf.transpose(delta))) + phi_phi_grad = tf.multiply(phi_phi_grad, weight) + new_u_grad = tf.matmul(tf.transpose(var_in), phi_phi_grad) + new_u_grad = tf.scalar_mul(1.0 / tf.cast(tf.shape(var_in)[0], tf.float32), new_u_grad) + + phi_error = tf.matmul(delta, tf.transpose(W)) + net_error = tf.multiply(phi_error, gate.gate_fun_gradient(phi, net)) + bp_u_grad = tf.matmul(tf.transpose(var_in), net_error) + bp_u_grad = tf.scalar_mul(1.0 / tf.cast(tf.shape(var_in)[0], tf.float32), bp_u_grad) + + new_u_grad = (1 - self.lam) * new_u_grad + self.lam * bp_u_grad + new_grads.append(new_u_grad) + + new_b_hidden_grad = tf.reduce_mean(phi_phi_grad, axis=0) + bp_b_hidden_grad = tf.reduce_mean(net_error, axis=0) + new_b_hidden_grad = (1 - self.lam) * new_b_hidden_grad + self.lam * bp_b_hidden_grad + new_grads.append(new_b_hidden_grad) + + old_grads = optimizer.compute_gradients(self.loss, var_list=[U, b_hidden]) + for i, (grad, var) in enumerate(old_grads): + old_grads[i] = (new_grads[i], var) + other_grads = optimizer.compute_gradients(self.loss, var_list=trainable_vars) + + self.all_gradients = old_grads + other_grads + self.train_op = optimizer.apply_gradients(self.all_gradients) + self.h_var = np.zeros((dim_hidden, dim_out)) + self.variables = [W, b_out, U, b_hidden] + self.y = y + + def get_assign_ops(self, src_network): + assign_ops = [] + for dst_var, src_var in zip(self.variables, src_network.variables): + assign_ops.append(dst_var.assign(src_var)) + return assign_ops + + def predict(self, sess, x): + y = sess.run(self.y, feed_dict={self.x: x}) + return y + + def learn(self, sess, train_x, train_y): + _, h_decay_var, h_delta_var = \ + sess.run([self.train_op, self.h_decay, self.h_delta], + feed_dict={ + self.x: train_x, + self.target: train_y, + self.h: self.h_var + }) + batch_size = float(train_x.shape[0]) + self.h_var = np.multiply(h_decay_var / batch_size, self.h_var) - self.learning_rate * h_delta_var / batch_size diff --git a/network.py b/network.py index 9ad55cc..497a72c 100644 --- a/network.py +++ b/network.py @@ -35,3 +35,15 @@ class Network: def learn(self, sess, x, target): sess.run(self.train_op, feed_dict={self.x: x, self.target: target}) + +class SimpleNetwork(Network): + def __init__(self, name, dim_in, dim_out, dim_hidden, optimizer_fn, initializer=tf.random_normal_initializer()): + self.x = tf.placeholder(tf.float32, shape=(None, dim_in)) + W1, b1, net1, phi1 = \ + fully_connected(name, 'layer1', self.x, dim_in, dim_hidden, initializer, tf.nn.relu) + W2, b2, net2, self.y = \ + fully_connected(name, 'layer2', phi1, dim_hidden, dim_out, initializer, tf.identity) + self.target = tf.placeholder(tf.float32, shape=(None, dim_out)) + loss = 0.5 * tf.reduce_mean(tf.squared_difference(self.y, self.target)) + self.variables = [W1, b1, W2, b2] + self.train_op = optimizer_fn(name).minimize(loss=loss) diff --git a/task.py b/task.py index 4234229..e9d8fbb 100644 --- a/task.py +++ b/task.py @@ -39,6 +39,24 @@ class MountainCar(BasicTask): self.policy_fn = lambda: GreedyPolicy(epsilon=0.5, decay_factor=0.95, min_epsilon=0.1) self.replay_fn = lambda: Replay(memory_size=10000, batch_size=10) +class CartPole(BasicTask): + state_space_size = 4 + action_space_size = 2 + name = 'CartPole-v0' + success_threshold = 195 + discount = 0.99 + step_limit = 5000 + target_network_update_freq = 200 + + def __init__(self): + self.env = gym.make(self.name) + + optimizer_fn = lambda name: tf.train.GradientDescentOptimizer(name=name, learning_rate=0.01) + self.network_fn = lambda name: Network(name, self.state_space_size, + self.action_space_size, optimizer_fn, tf.random_normal_initializer()) + self.policy_fn = lambda: GreedyPolicy(epsilon=0.5, decay_factor=0.95, min_epsilon=0.1) + self.replay_fn = lambda: Replay(memory_size=10000, batch_size=10) + class LunarLander(BasicTask): state_space_size = 8 action_space_size = 4 @@ -53,7 +71,7 @@ class LunarLander(BasicTask): optimizer_fn = lambda name: tf.train.GradientDescentOptimizer(name=name, learning_rate=0.001) self.network_fn = lambda name: Network(name, self.state_space_size, self.action_space_size, optimizer_fn, tf.random_normal_initializer()) - self.policy_fn = lambda: GreedyPolicy(epsilon=0.5, decay_factor=0.99, min_epsilon=0) + self.policy_fn = lambda: GreedyPolicy(epsilon=0.5, decay_factor=0.99, min_epsilon=0.1) self.replay_fn = lambda: Replay(memory_size=20000, batch_size=100) if __name__ == '__main__': @@ -74,4 +92,4 @@ if __name__ == '__main__': reward = np.mean(rewards[-window_size:]) print 'episode %d: %f' % (ep, reward) if reward > task.success_threshold: - break + break \ No newline at end of file