modularized DQN

.gitignore

@@ -3,6 +3,8 @@ __pycache__/
 *.py[cod]
 *$py.class
 
+.idea
+
 # C extensions
 *.so
 

common.py

@@ -0,0 +1,13 @@
+import tensorflow as tf
+
+def fully_connected(model_name, layer_name, var_in, dim_in, dim_out,
+                    initializer, transfer):
+    with tf.variable_scope(model_name):
+        with tf.variable_scope(layer_name):
+            W = tf.get_variable("W", [dim_in, dim_out],
+                                initializer=initializer)
+            b = tf.get_variable("b", [dim_out],
+                                initializer=initializer)
+    net = tf.nn.bias_add(tf.matmul(var_in, W), b)
+    phi = transfer(net)
+    return W, b, net, phi

dqn_agent.py

@@ -0,0 +1,49 @@
+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)
+        self.task = task
+        self.step_limit = step_limit
+        self.replay = replay_fn()
+        self.discount = discount
+        self.target_network_update_freq = target_network_update_freq
+        self.policy = policy_fn()
+        self.total_steps = 0
+
+    def episode(self, sess):
+        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)) )
+            action = self.policy.sample(value.flatten())
+            next_state, reward, done, info = self.task.step(action)
+            total_reward += reward
+            self.replay.feed([state, action, reward, next_state, int(done)])
+            steps += 1
+            self.total_steps += 1
+            state = next_state
+            if done:
+                break
+            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)
+                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)
+            if self.total_steps % self.target_network_update_freq == 0:
+                sess.run(self.assign_ops)
+        self.policy.update_epsilon()
+        return total_reward

network.py

@@ -0,0 +1,31 @@
+import tensorflow as tf
+from common import *
+
+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)
+
+    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, x, target):
+        sess.run(self.train_op, feed_dict={self.x: x, self.target: target})

policy.py

@@ -0,0 +1,16 @@
+import numpy as np
+
+class GreedyPolicy:
+    def __init__(self, epsilon, decay_factor, min_epsilon):
+        self.epsilon = epsilon
+        self.decay_factor = decay_factor
+        self.min_epsilon = min_epsilon
+
+    def sample(self, state):
+        if np.random.rand() < self.epsilon:
+            return np.random.randint(0, len(state))
+        return np.argmax(state)
+
+    def update_epsilon(self):
+        if self.epsilon > self.min_epsilon:
+            self.epsilon *= self.decay_factor

replay.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+class Replay:
+    def __init__(self, memory_size, batch_size):
+        self.memory_size = memory_size
+        self.batch_size = batch_size
+
+        self.states = []
+        self.actions = []
+        self.rewards = []
+        self.next_states = []
+        self.terminals = []
+
+
+    def feed(self, experience):
+        state, action, reward, next_state, done = experience
+        self.states.append(state)
+        self.actions.append(action)
+        self.rewards.append(reward)
+        self.next_states.append(next_state)
+        self.terminals.append(done)
+        if len(self.terminals) > self.memory_size:
+            self.states.pop(0)
+            self.actions.pop(0)
+            self.rewards.pop(0)
+            self.next_states.pop(0)
+            self.terminals.pop(0)
+
+    def sample(self):
+        if len(self.terminals) >= self.batch_size:
+            sampled_indices = np.arange(len(self.terminals))
+            np.random.shuffle(sampled_indices)
+            sampled_indices = sampled_indices[: self.batch_size]
+            return [np.asarray(self.states)[sampled_indices],
+                    np.asarray(self.actions)[sampled_indices],
+                    np.asarray(self.rewards)[sampled_indices],
+                    np.asarray(self.next_states)[sampled_indices],
+                    np.asarray(self.terminals)[sampled_indices]]
+        return None
+

task.py

@@ -0,0 +1,42 @@
+import gym
+import sys
+from dqn_agent import *
+
+class BasicTask:
+    def transfer_state(self, state):
+        return state
+
+    def reset(self):
+        return self.transfer_state(self.env.reset())
+
+    def step(self, action):
+        next_state, reward, done, info = self.env.step(action)
+        next_state = self.transfer_state(next_state)
+        return next_state, reward, done, info
+
+class MountainCar(BasicTask):
+    state_space_size = 2
+    action_space_size = 3
+    name = 'MountainCar-v0'
+    def __init__(self):
+        self.env = gym.make(self.name)
+        self.env._max_episode_steps = sys.maxsize
+
+if __name__ == '__main__':
+    task = MountainCar()
+    optimizer_fn = lambda name: tf.train.GradientDescentOptimizer(name=name, learning_rate=0.01)
+    network_fn = lambda name: Network(name, task.state_space_size,
+                                      task.action_space_size, optimizer_fn, tf.random_normal_initializer())
+    policy_fn = lambda: GreedyPolicy(epsilon=0.5, decay_factor=0.95, min_epsilon=0.1)
+    replay_fn = lambda: Replay(memory_size=10000, batch_size=10)
+    agent = DQNAgent('mountain-car', task, network_fn, policy_fn, replay_fn,
+                     discount=0.99, step_limit=5000, target_network_update_freq=1000)
+    with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        ep = 0
+        while True:
+            ep += 1
+            reward = agent.episode(sess)
+            print 'episode %d: %f' % (ep, reward)
+            if reward > -110:
+                break