Crossprop

.gitignore

@@ -4,6 +4,9 @@ __pycache__/
 *$py.class
 
 .idea
+exp_*
+upload.py
+*.sh
 
 # C extensions
 *.so

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
+    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

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

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)

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