Pixel atari games with DQN

README.md

@@ -6,6 +6,6 @@
 * Asynchronous N-Step Q-Learning
 * Asynchronous Advantage Actor Critic (A3C)
 
->Benchmarked by classical control tasks (CartPole, LunarLander). Atari games will make it difficult to replicate in a regular laptop without a good GPU. However it's fairly easy to adapt the components to fit Atari games.
+>Tested with both classical control tasks (CartPole, LunarLander) and Atari games.
 
 >Try it out from ```main.py```!

async_agent.py

@@ -79,7 +79,6 @@ class AsyncAgent:
                 episode_steps = 0
                 episode_return = 0
                 episode += 1
-                policy.update_epsilon()
                 terminal = False
                 state = task.reset()
                 state = state.reshape([1, -1])
@@ -97,6 +96,7 @@ class AsyncAgent:
                 state = state.reshape([1, -1])
                 value = worker_network.predict(state)
                 action = policy.sample(value.flatten())
+                policy.update_epsilon()
 
             batch_rewards = self.bootstrap_fn(batch_states, batch_actions, batch_rewards,
                                               state, action, terminal, self)

dqn_agent.py

@@ -10,7 +10,18 @@ from policy import *
 import numpy as np
 
 class DQNAgent:
-    def __init__(self, task_fn, network_fn, optimizer_fn, policy_fn, replay_fn, discount, step_limit, target_network_update_freq):
+    def __init__(self,
+                 task_fn,
+                 network_fn,
+                 optimizer_fn,
+                 policy_fn,
+                 replay_fn,
+                 discount,
+                 step_limit,
+                 target_network_update_freq,
+                 explore_steps,
+                 history_length,
+                 logger):
         self.learning_network = network_fn(optimizer_fn)
         self.target_network = network_fn(optimizer_fn)
         self.target_network.load_state_dict(self.learning_network.state_dict())
@@ -21,24 +32,37 @@ class DQNAgent:
         self.target_network_update_freq = target_network_update_freq
         self.policy = policy_fn()
         self.total_steps = 0
+        self.explore_steps = explore_steps
+        self.history_length = history_length
+        self.logger = logger
+
+    def get_state(self, history_buffer):
+        if self.history_length > 1:
+            return np.vstack(history_buffer)
+        return history_buffer[0]
 
     def episode(self):
         state = self.task.reset()
+        history_buffer = [state] * self.history_length
         total_reward = 0.0
         steps = 0
         while not self.step_limit or steps < self.step_limit:
-            value = self.learning_network.predict(np.reshape(state, (1, -1)))
+            state = self.get_state(history_buffer)
+            value = self.learning_network.predict(np.reshape(state, (1, ) + state.shape))
             action = self.policy.sample(value.flatten())
             next_state, reward, done, info = self.task.step(action)
+            history_buffer.pop(0)
+            history_buffer.append(next_state)
+            next_state = self.get_state(history_buffer)
             total_reward += reward
             self.replay.feed([state, action, reward, next_state, int(done)])
             steps += 1
             self.total_steps += 1
-            state = next_state
+            self.logger.debug('steps %d, reward %f, action %d' % (steps, reward, action))
             if done:
                 break
-            experiences = self.replay.sample()
-            if experiences is not None:
+            if self.total_steps > self.explore_steps:
+                experiences = self.replay.sample()
                 states, actions, rewards, next_states, terminals = experiences
                 targets = self.learning_network.predict(states)
                 q_next = self.target_network.predict(next_states)
@@ -46,10 +70,13 @@ class DQNAgent:
                 q_next = np.where(terminals, 0, q_next)
                 q_next = rewards + self.discount * q_next
                 targets[np.arange(len(actions)), actions] = q_next
+                self.logger.debug('start minibatch')
                 self.learning_network.learn(states, targets)
+                self.logger.debug('minibatch ended')
             if self.total_steps % self.target_network_update_freq == 0:
                 self.target_network.load_state_dict(self.learning_network.state_dict())
-        self.policy.update_epsilon()
+            if self.total_steps > self.explore_steps:
+                self.policy.update_epsilon()
         return total_reward
 
     def run(self):
@@ -60,10 +87,8 @@ class DQNAgent:
             ep += 1
             reward = self.episode()
             rewards.append(reward)
-            if len(rewards) > window_size:
-                reward = np.mean(rewards[-window_size:])
-            print 'episode %d, epsilon %f, reward %f' % (
-                ep, self.policy.epsilon, reward)
-            if reward > self.task.success_threshold:
+            avg_reward = np.mean(rewards[-window_size:])
+            self.logger.info('episode %d, epsilon %f, reward %f, avg reward %f, total steps %d' % (
+                ep, self.policy.epsilon, reward, avg_reward, self.total_steps))
+            if avg_reward > self.task.success_threshold:
                 break
-

main.py

@@ -1,12 +1,13 @@
 from async_agent import *
 from dqn_agent import *
+import logging
 
 def async_cart_pole():
     config = dict()
     config['task_fn'] = lambda: CartPole()
     config['optimizer_fn'] = lambda params: torch.optim.SGD(params, 0.001)
     config['network_fn'] = lambda: FullyConnectedNet([4, 50, 200, 2])
-    config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, end_episode=500, min_epsilon=0.1)
+    config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=5000, min_epsilon=0.1)
     # config['bootstrap_fn'] = OneStepQLearning
     # config['bootstrap_fn'] = NStepQLearning
     config['bootstrap_fn'] = OneStepSarsa
@@ -25,7 +26,7 @@ def async_lunar_lander():
     config['task_fn'] = lambda: LunarLander()
     config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001)
     config['network_fn'] = lambda: FullyConnectedNet([8, 50, 200, 4])
-    config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, end_episode=2000, min_epsilon=0.05)
+    config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=40000, min_epsilon=0.05)
     config['bootstrap_fn'] = OneStepQLearning
     config['discount'] = 0.99
     config['target_network_update_freq'] = 200
@@ -37,30 +38,19 @@ def async_lunar_lander():
     agent = AsyncAgent(**config)
     agent.run()
 
-# Mountain Car is fairly unstable
-def dqn_mountain_car():
-    config = dict()
-    config['task_fn'] = lambda: MountainCar()
-    config['optimizer_fn'] = lambda params: torch.optim.SGD(params, 0.001)
-    config['network_fn'] = lambda optimizer_fn: FullyConnectedNet([2, 50, 200, 3], optimizer_fn)
-    config['policy_fn'] = lambda: GreedyPolicy(epsilon=0.5, end_episode=500, min_epsilon=0.1)
-    config['replay_fn'] = lambda: Replay(memory_size=10000, batch_size=10)
-    config['discount'] = 0.99
-    config['target_network_update_freq'] = 1000
-    config['step_limit'] = 5000
-    agent = DQNAgent(**config)
-    agent.run()
-
 def dqn_cart_pole():
     config = dict()
     config['task_fn'] = lambda: CartPole()
     config['optimizer_fn'] = lambda params: torch.optim.SGD(params, 0.001)
     config['network_fn'] = lambda optimizer_fn: FullyConnectedNet([4, 50, 200, 2], optimizer_fn)
-    config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, end_episode=500, min_epsilon=0.1)
+    config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=10000, min_epsilon=0.1)
     config['replay_fn'] = lambda: Replay(memory_size=10000, batch_size=10)
     config['discount'] = 0.99
     config['target_network_update_freq'] = 200
-    config['step_limit'] = 300
+    config['step_limit'] = 0
+    config['explore_steps'] = 1000
+    config['logger'] = gym.logger
+    config['history_length'] = 1
     agent = DQNAgent(**config)
     agent.run()
 
@@ -81,9 +71,28 @@ def actor_critic_cart_pole():
     agent = AsyncAgent(**config)
     agent.run()
 
+def dqn_pixel_atari(name):
+    config = dict()
+    config['task_fn'] = lambda: PixelAtari(name)
+    config['optimizer_fn'] = lambda params: torch.optim.RMSprop(params, lr=0.00025)
+    config['network_fn'] = lambda optimizer_fn: ConvNet(4, 6, optimizer_fn)
+    config['policy_fn'] = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1)
+    config['replay_fn'] = lambda: Replay(memory_size=1000000, batch_size=32)
+    config['discount'] = 0.99
+    config['target_network_update_freq'] = 10000
+    config['step_limit'] = 0
+    config['explore_steps'] = 50000
+    config['logger'] = gym.logger
+    config['history_length'] = 4
+    agent = DQNAgent(**config)
+    agent.run()
+
 if __name__ == '__main__':
+    # gym.logger.setLevel(logging.DEBUG)
+    gym.logger.setLevel(logging.INFO)
     # async_cart_pole()
     # async_lunar_lander()
     # dqn_cart_pole()
     # dqn_mountain_car()
-    actor_critic_cart_pole()
+    # actor_critic_cart_pole()
+    dqn_pixel_atari('Breakout-v0')

network.py

@@ -36,10 +36,6 @@ class FullyConnectedNet(nn.Module):
         y = self.fc3(y)
         return y
 
-    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)
-
     def predict(self, x):
         return self.forward(x).cpu().data.numpy()
 
@@ -103,4 +99,49 @@ class ActorCriticNet(nn.Module):
 
     def critic(self, x):
         phi = self.forward(x)
-        return self.fc_critic(phi).cpu().data.numpy()
+        return self.fc_critic(phi).cpu().data.numpy()
+
+class ConvNet(nn.Module):
+    def __init__(self, in_channels, n_actions, optimizer_fn=None, gpu=True):
+        super(ConvNet, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
+        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
+        self.fc4 = nn.Linear(7 * 7 * 64, 512)
+        self.fc5 = nn.Linear(512, n_actions)
+
+        self.criterion = nn.MSELoss()
+        if optimizer_fn is not None:
+            self.optimizer = optimizer_fn(self.parameters())
+
+        self.gpu = gpu and torch.cuda.is_available()
+        if self.gpu:
+            print 'Transferring network to GPU...'
+            self.cuda()
+            print 'Network transferred.'
+
+    def to_torch_variable(self, x):
+        x = torch.from_numpy(np.asarray(x, dtype='float32'))
+        if self.gpu:
+            x = x.cuda()
+        return Variable(x)
+
+    def forward(self, x):
+        y = F.relu(self.conv1(x))
+        y = F.relu(self.conv2(y))
+        y = F.relu(self.conv3(y))
+        y = y.view(y.size(0), -1)
+        y = F.relu(self.fc4(y))
+        return self.fc5(y)
+
+    def predict(self, x):
+        return self.forward(self.to_torch_variable(x)).cpu().data.numpy()
+
+    def learn(self, x, target):
+        x = self.to_torch_variable(x)
+        target = self.to_torch_variable(target)
+        y = self.forward(x)
+        loss = self.criterion(y, target)
+        self.zero_grad()
+        loss.backward()
+        self.optimizer.step()

policy.py

@@ -7,11 +7,11 @@
 import numpy as np
 
 class GreedyPolicy:
-    def __init__(self, epsilon, end_episode, min_epsilon):
+    def __init__(self, epsilon, final_step, min_epsilon):
         self.init_epsilon = self.epsilon = epsilon
-        self.current_episode = 0
+        self.current_steps = 0
         self.min_epsilon = min_epsilon
-        self.end_episode = end_episode
+        self.final_step = final_step
 
     def sample(self, action_value):
         if np.random.rand() < self.epsilon:
@@ -19,9 +19,9 @@ class GreedyPolicy:
         return np.argmax(action_value)
 
     def update_epsilon(self):
-        self.epsilon = self.init_epsilon - float(self.current_episode) / self.end_episode * (self.init_epsilon - self.min_epsilon)
+        self.epsilon = self.init_epsilon - float(self.current_steps) / self.final_step * (self.init_epsilon - self.min_epsilon)
         self.epsilon = max(self.epsilon, self.min_epsilon)
-        self.current_episode += 1
+        self.current_steps += 1
 
 class SamplePolicy:
     def sample(self, action_value):

replay.py

@@ -37,10 +37,21 @@ class Replay:
             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]]
+            sampled_states = []
+            sampled_actions = []
+            sampled_rewards = []
+            sampled_next_states = []
+            sampled_terminals = []
+            for ind in sampled_indices:
+                sampled_states.append(self.states[ind])
+                sampled_actions.append(self.actions[ind])
+                sampled_rewards.append(self.rewards[ind])
+                sampled_next_states.append(self.next_states[ind])
+                sampled_terminals.append(self.terminals[ind])
+            return [np.asarray(sampled_states),
+                    np.asarray(sampled_actions),
+                    np.asarray(sampled_rewards),
+                    np.asarray(sampled_next_states),
+                    np.asarray(sampled_terminals)]
         return None
 

task.py

@@ -5,6 +5,8 @@
 #######################################################################
 import gym
 import sys
+import numpy as np
+import cv2
 
 class BasicTask:
     def transfer_state(self, state):
@@ -16,7 +18,7 @@ class BasicTask:
     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
+        return next_state, np.sign(reward), done, info
 
 class MountainCar(BasicTask):
     name = 'MountainCar-v0'
@@ -38,4 +40,17 @@ class LunarLander(BasicTask):
     success_threshold = 200
 
     def __init__(self):
-        self.env = gym.make(self.name)
+        self.env = gym.make(self.name)
+
+class PixelAtari(BasicTask):
+    width = 84
+    height = 84
+    success_threshold = 1000
+
+    def __init__(self, name):
+        self.env = gym.make(name)
+
+    def transfer_state(self, state):
+        img = (state[:, :, 0] * 0.299 + state[:, :, 1] * 0.587 + state[:, :, 2] * 0.114) / 255.0
+        img = cv2.resize(img, (self.width, self.height))
+        return np.reshape(img, (1, self.width, self.height))