diff --git a/.gitignore b/.gitignore index 36f0f4d..5798e31 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ __pycache__/ exp_* upload.py *.sh +data # C extensions *.so diff --git a/crossprop.py b/crossprop.py deleted file mode 100644 index adf1415..0000000 --- a/crossprop.py +++ /dev/null @@ -1,115 +0,0 @@ -####################################################################### -# 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/dqn_agent.py b/dqn_agent.py index 181c120..510fabb 100644 --- a/dqn_agent.py +++ b/dqn_agent.py @@ -4,18 +4,15 @@ # declaration at the top # ####################################################################### -import tensorflow as tf from network import * from replay import * from policy import * import numpy as np class DQNAgent: - def __init__(self, name, task, network_fn, policy_fn, replay_fn, discount, step_limit, target_network_update_freq): - with tf.variable_scope(name): - self.learning_network = network_fn('learning') - self.target_network = network_fn('target') - self.assign_ops = self.target_network.get_assign_ops(self.learning_network) + def __init__(self, task, network_fn, policy_fn, replay_fn, discount, step_limit, target_network_update_freq): + self.learning_network = network_fn() + self.target_network = network_fn() self.task = task self.step_limit = step_limit self.replay = replay_fn() @@ -24,12 +21,12 @@ class DQNAgent: self.policy = policy_fn() self.total_steps = 0 - def episode(self, sess): + def episode(self): state = self.task.reset() total_reward = 0.0 steps = 0 while not self.step_limit or steps < self.step_limit: - value = self.learning_network.predict(sess, np.reshape(state, (1, -1)) ) + value = self.learning_network.predict(np.reshape(state, (1, -1)) ) action = self.policy.sample(value.flatten()) next_state, reward, done, info = self.task.step(action) total_reward += reward @@ -42,14 +39,14 @@ class DQNAgent: experiences = self.replay.sample() if experiences is not None: states, actions, rewards, next_states, terminals = experiences - targets = self.learning_network.predict(sess, states) - q_next = self.target_network.predict(sess, next_states) + targets = self.learning_network.predict(states) + q_next = self.target_network.predict(next_states) q_next = np.max(q_next, axis=1) q_next = rewards + self.discount * q_next q_next = np.where(terminals, 0, q_next) targets[np.arange(len(actions)), actions] = q_next - self.learning_network.learn(sess, states, targets) + self.learning_network.learn(states, targets) if self.total_steps % self.target_network_update_freq == 0: - sess.run(self.assign_ops) + self.target_network.sync_with(self.learning_network) self.policy.update_epsilon() return total_reward diff --git a/network.py b/network.py index 497a72c..5cf05d4 100644 --- a/network.py +++ b/network.py @@ -4,46 +4,67 @@ # declaration at the top # ####################################################################### -import tensorflow as tf -from common import * +import torch +from torch.autograd import Variable +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from SMDWrapper import SMDWrapper -class Network: - def __init__(self, name, dim_in, dim_out, optimizer_fn, initializer=tf.random_normal_initializer()): - self.x = tf.placeholder(tf.float32, shape=(None, dim_in)) - dim_hidden1 = 50 - dim_hidden2 = 200 - W1, b1, net1, phi1 = \ - fully_connected(name, 'layer1', self.x, dim_in, dim_hidden1, initializer, tf.nn.relu) - W2, b2, net2, phi2 = \ - fully_connected(name, 'layer2', phi1, dim_hidden1, dim_hidden2, initializer, tf.nn.relu) - W3, b3, net3, self.y = \ - fully_connected(name, 'layer3', phi2, dim_hidden2, 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, W3, b3] - self.train_op = optimizer_fn(name).minimize(loss=loss) +class FullyConnectedNet(nn.Module): + def __init__(self, dims, learning_rate, gpu=True): + super(FullyConnectedNet, self).__init__() + self.fc1 = nn.Linear(dims[0], dims[1]) + self.fc2 = nn.Linear(dims[1], dims[2]) + self.fc3 = nn.Linear(dims[2], dims[3]) + self.criterion = nn.MSELoss() + self.learning_rate = learning_rate + self.optimizer = torch.optim.SGD(self.parameters(), learning_rate) + self.gpu = gpu and torch.cuda.is_available() + if self.gpu: + print 'Transferring network to GPU...' + self.cuda() + print 'Network transferred.' - 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 forward(self, x): + x = torch.from_numpy(np.asarray(x, dtype='float32')) + if self.gpu: + x = x.cuda() + x = Variable(x) - def predict(self, sess, x): - y = sess.run(self.y, feed_dict={self.x: x}) + y = F.relu(self.fc1(x)) + y = F.relu(self.fc2(y)) + y = self.fc3(y) return y - def learn(self, sess, x, target): - sess.run(self.train_op, feed_dict={self.x: x, self.target: target}) + def sync_with(self, src_net): + for param_dst, param_src in zip(self.parameters(), src_net.parameters()): + param_dst.data.copy_(param_src.data) -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) + def predict(self, x): + return self.forward(x).cpu().data.numpy() + + def learn(self, x, target): + target = torch.from_numpy(target) + if self.gpu: + target = target.cuda() + target = Variable(target) + + y = self.forward(x) + loss = self.criterion(y, target) + self.zero_grad() + loss.backward() + self.optimizer.step() + + def output_transfer(self, y): + return y + +class SMDNetworkWrapper(SMDWrapper): + def __init__(self, net): + SMDWrapper.__init__(self, net) + + def sync_with(self, src_net): + self.net.sync_with(src_net.net) + + def parameters(self): + return self.net.parameters() diff --git a/task.py b/task.py index e9d8fbb..98dc6b8 100644 --- a/task.py +++ b/task.py @@ -33,9 +33,7 @@ class MountainCar(BasicTask): self.env = gym.make(self.name) self.env._max_episode_steps = sys.maxsize - 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.network_fn = lambda learning_rate=0.01: FullyConnectedNet([self.state_space_size, 50, 200, self.action_space_size], learning_rate) 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) @@ -51,45 +49,28 @@ class CartPole(BasicTask): 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.network_fn = lambda learning_rate=0.01: FullyConnectedNet([self.state_space_size, 50, 200, self.action_space_size], learning_rate) 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 - name = 'LunarLander-v2' - success_threshold = 200 - 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.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.1) - self.replay_fn = lambda: Replay(memory_size=20000, batch_size=100) - if __name__ == '__main__': - # task = MountainCar() - task = LunarLander() - agent = DQNAgent(task.name, task, task.network_fn, task.policy_fn, task.replay_fn, + task = MountainCar() + bp_network_fn = lambda learning_rate=0.001: FullyConnectedNet([task.state_space_size, 50, 200, task.action_space_size], learning_rate, gpu=False) + def smd_network_fn(learning_rate=0.001): + bp_network = bp_network_fn(learning_rate) + return SMDNetworkWrapper(bp_network) + + agent = DQNAgent(task, smd_network_fn, task.policy_fn, task.replay_fn, task.discount, task.step_limit, task.target_network_update_freq) window_size = 100 - with tf.Session() as sess: - sess.run(tf.global_variables_initializer()) - ep = 0 - rewards = [] - while True: - ep += 1 - reward = agent.episode(sess) - rewards.append(reward) - if len(rewards) > window_size: - reward = np.mean(rewards[-window_size:]) - print 'episode %d: %f' % (ep, reward) - if reward > task.success_threshold: - break \ No newline at end of file + ep = 0 + rewards = [] + while True: + ep += 1 + reward = agent.episode() + rewards.append(reward) + if len(rewards) > window_size: + reward = np.mean(rewards[-window_size:]) + print 'episode %d: %f' % (ep, reward) + if reward > task.success_threshold: + break