Major refactor to support LSTM layer

async_agent.py

@@ -20,12 +20,12 @@ class AsyncAgent:
                  network_fn,
                  optimizer_fn,
                  policy_fn,
-                 bootstrap_fn,
+                 bootstrap,
                  discount,
                  step_limit,
                  target_network_update_freq,
                  n_workers,
-                 batch_size,
+                 update_interval,
                  test_interval,
                  test_repetitions,
                  history_length,
@@ -33,16 +33,17 @@ class AsyncAgent:
         self.network_fn = network_fn
         self.learning_network = network_fn()
         self.learning_network.share_memory()
-        if bootstrap_fn != AdvantageActorCritic:
+        if bootstrap != AdvantageActorCritic:
             self.target_network = network_fn()
             self.target_network.share_memory()
             self.target_network.load_state_dict(self.learning_network.state_dict())
         else:
             self.target_network = None
-        self.bootstrap_fn = bootstrap_fn
+        self.bootstrap = bootstrap
 
         self.optimizer_fn = optimizer_fn
         self.task_fn = task_fn
+        self.task = self.task_fn()
         self.step_limit = step_limit
         self.discount = discount
         self.optimizer_fn = optimizer_fn
@@ -53,7 +54,7 @@ class AsyncAgent:
         self.total_steps = mp.Value('i', 0)
         self.stop_signal = mp.Value('i', False)
         self.n_workers = n_workers
-        self.batch_size = batch_size
+        self.update_interval = update_interval
         self.test_interval = test_interval
         self.test_repetitions = test_repetitions
         self.logger = logger
@@ -63,92 +64,69 @@ class AsyncAgent:
         state = task.reset()
         total_rewards = 0
         steps = 0
-        terminal = False
-        buffer = [state] * self.history_length
-        while not terminal and (not self.step_limit or steps < self.step_limit):
-            state = np.vstack(buffer)
-            action_values = network.predict(np.stack([state]))
-            steps += 1
-            action = np.argmax(action_values.flatten())
+        network.reset(True)
+        bootstrap = self.bootstrap(self)
+        while not self.step_limit or steps < self.step_limit:
+            action = np.argmax(bootstrap.process_state(network, state))
             state, reward, terminal, _ = task.step(action)
-            buffer.pop(0)
-            buffer.append(state)
+            steps += 1
             total_rewards += reward
             if terminal:
                 break
         return total_rewards
 
-    def async_update(self, worker_network, optimizer):
-        optimizer.zero_grad()
-        for param, worker_param in zip(self.learning_network.parameters(), worker_network.parameters()):
-            param._grad = worker_param.grad.clone().cpu()
-        optimizer.step()
-
     def worker(self, id):
         optimizer = self.optimizer_fn(self.learning_network.parameters())
         worker_network = self.network_fn()
         worker_network.load_state_dict(self.learning_network.state_dict())
+        bootstrap = self.bootstrap(self)
         task = self.task_fn()
         policy = self.policy_fn()
-        terminal = True
         episode = 0
         episode_steps = 0
-        episode_return = 0
-        episode_returns = []
-        update_target_network = False
+        episode_returns = [0]
+        state = task.reset()
+        pending_steps = 0
         while True and not self.stop_signal.value:
-            batch_states, batch_actions, batch_rewards = [], [], []
-            if terminal:
-                if episode and id == 0:
-                    episode_returns.append(episode_return)
-                    self.logger.info('episode %d, return %f, avg return %f, episode steps %d, total steps %d' % (
-                        episode, episode_return, np.mean(episode_returns[-100: ]), episode_steps,
-                    self.total_steps.value))
-                episode_steps = 0
-                episode_return = 0
-                episode += 1
-                terminal = False
-                state = task.reset()
-                buffer = [state] * self.history_length
-                state = np.vstack(buffer)
-                value = worker_network.predict(np.stack([state]))
-                action = policy.sample(value.flatten())
-            while not terminal and len(batch_states) < self.batch_size:
-                episode_steps += 1
-                with self.steps_lock:
-                    self.total_steps.value += 1
-                self.total_steps.value += 1
-                if self.total_steps.value % self.target_network_update_freq == 0:
-                    update_target_network = True
-                batch_states.append(state)
-                batch_actions.append(action)
-                state, reward, terminal, _ = task.step(action)
-                batch_rewards.append(reward)
-                episode_return += reward
-                buffer.pop(0)
-                buffer.append(state)
-                state = np.vstack(buffer)
-                value = worker_network.predict(np.stack([state]))
-                action = policy.sample(value.flatten())
-                policy.update_epsilon()
-
-            batch_rewards = self.bootstrap_fn(batch_states, batch_actions, batch_rewards,
-                                              state, action, terminal, worker_network, self.discount)
+            action = policy.sample(bootstrap.process_state(worker_network, state))
+            next_state, reward, terminal, _ = task.step(action)
+            bootstrap.process_interaction(action, reward, next_state)
 
+            episode_returns[-1] += reward
+            episode_steps += 1
             if self.step_limit and episode_steps > self.step_limit:
                 terminal = True
+            with self.steps_lock:
+                self.total_steps.value += 1
+            pending_steps += 1
 
-            worker_network.zero_grad()
-            worker_network.gradient(np.asarray(batch_states),
-                                    worker_network.to_torch_variable(batch_actions, 'int64').unsqueeze(1),
-                                    worker_network.to_torch_variable(batch_rewards).unsqueeze(1))
-            self.async_update(worker_network, optimizer)
-            worker_network.load_state_dict(self.learning_network.state_dict())
+            if terminal or pending_steps >= self.update_interval:
+                loss = bootstrap.compute_loss(worker_network, terminal)
+                pending_steps = 0
+                worker_network.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm(worker_network.parameters(), 40)
+                optimizer.zero_grad()
+                for param, worker_param in zip(self.learning_network.parameters(), worker_network.parameters()):
+                    param._grad = worker_param.grad.clone().cpu()
+                optimizer.step()
+                worker_network.load_state_dict(self.learning_network.state_dict())
+                worker_network.reset(terminal)
 
-            if self.target_network is not None and update_target_network:
+            if terminal:
+                state = task.reset()
+                episode += 1
+                if id == 0:
+                    self.logger.info('episode %d, return %f, avg return %f, episode steps %d, total steps %d' % (
+                        episode, episode_returns[-1], np.mean(episode_returns[-100:]), episode_steps, self.total_steps.value))
+                episode_returns.append(0)
+                episode_steps = 0
+            else:
+                state = next_state
+
+            if self.target_network and self.total_steps.value % self.target_network_update_freq == 0:
                 with self.network_lock:
                     self.target_network.load_state_dict(self.learning_network.state_dict())
-                update_target_network = False
 
     def save(self, file_name):
         with open(file_name, 'wb') as f:
@@ -158,28 +136,27 @@ class AsyncAgent:
         os.environ['OMP_NUM_THREADS'] = '1'
         procs = [mp.Process(target=self.worker, args=(i, )) for i in range(self.n_workers)]
         for p in procs: p.start()
-        task = self.task_fn()
-        test_network = self.network_fn()
         test_rewards = []
         test_points = []
+        test_network = self.network_fn()
         while True:
             steps = self.total_steps.value + 1
             if steps % self.test_interval == 0:
                 with self.network_lock:
                     test_network.load_state_dict(self.learning_network.state_dict())
-                self.save('data/%s-model-%s.bin' % (self.bootstrap_fn.__name__, task.name))
+                self.save('data/%s-model-%s.bin' % (self.bootstrap.__name__, self.task.name))
                 rewards = np.zeros(self.test_repetitions)
                 for i in range(self.test_repetitions):
-                    rewards[i] = self.deterministic_episode(task, test_network)
+                    rewards[i] = self.deterministic_episode(self.task, test_network)
                 self.logger.info('total steps: %d, averaged return per episode: %f(%f)' %\
                       (steps, np.mean(rewards), np.std(rewards) / np.sqrt(self.test_repetitions)))
                 test_rewards.append(np.mean(rewards))
                 test_points.append(steps)
                 with open('data/%s-statistics-%s.bin' % (
-                    self.bootstrap_fn.__name__, task.name
+                    self.bootstrap.__name__, self.task.name
                 ), 'wb') as f:
                     pickle.dump([test_points, test_rewards], f)
-                if np.mean(rewards) > task.success_threshold:
+                if np.mean(rewards) > self.task.success_threshold:
                     self.stop_signal.value = True
                     break
         for p in procs: p.join()

atari_wrapper.py

@@ -115,17 +115,33 @@ def _process_frame84(frame):
     x_t = np.reshape(x_t, [1, 84, 84])
     return x_t.astype(np.uint8)
 
-class ProcessFrame84(gym.Wrapper):
-    def __init__(self, env=None):
-        super(ProcessFrame84, self).__init__(env)
-        self.observation_space = spaces.Box(low=0, high=255, shape=(1, 84, 84))
+def _process_frame42(frame):
+    img = np.reshape(frame, [210, 160, 3]).astype(np.float32)
+    img = img[:, :, 0] * 0.299 + img[:, :, 1] * 0.587 + img[:, :, 2] * 0.114
+    img = img[34:34 + 160, :160]
+    img = Image.fromarray(img)
+    img = img.resize((80, 80), Image.BILINEAR)
+    img = img.resize((42, 42), Image.BILINEAR)
+    resized_screen = np.array(img).reshape(1, 42, 42)
+    return resized_screen.astype(np.uint8)
+
+class ProcessFrame(gym.Wrapper):
+    def __init__(self, env=None, frame_size=84):
+        super(ProcessFrame, self).__init__(env)
+        self.observation_space = spaces.Box(low=0, high=255, shape=(1, frame_size, frame_size))
+        if frame_size == 84:
+            self.process_fn = _process_frame84
+        elif frame_size == 42:
+            self.process_fn = _process_frame42
+        else:
+            assert(False, "Unknown frame size")
 
     def _step(self, action):
         obs, reward, done, info = self.env.step(action)
-        return _process_frame84(obs), reward, done, info
+        return self.process_fn(obs), reward, done, info
 
     def _reset(self):
-        return _process_frame84(self.env.reset())
+        return self.process_fn(self.env.reset())
 
 class ClippedRewardsWrapper(gym.Wrapper):
     def _step(self, action):

bootstrap.py

@@ -4,51 +4,124 @@
 # declaration at the top                                              #
 #######################################################################
 import numpy as np
+import torch
+from torch.autograd import Variable
 
-def NStepQLearning(batch_states, batch_actions, batch_rewards,
-                   tailing_state, tailing_action, terminal, network, discount):
-    if terminal:
-        reward = 0
-    else:
-        reward = np.max(network.predict(np.stack([tailing_state])).flatten())
-    rewards = []
-    for r in reversed(batch_rewards):
-        reward = r + discount * reward
-        rewards.append(reward)
-    return rewards
+class OneStepSarsa:
+    def __init__(self, agent):
+        self.agent = agent
+        self.pending = []
 
-def OneStepQLearning(batch_states, batch_actions, batch_rewards,
-                     tailing_state, tailing_action, terminal, network, discount):
-    batch_states.append(tailing_state)
-    q_next = network.predict(np.asarray(batch_states[1:]))
-    q_next = np.max(q_next, axis=1)
-    if terminal:
-        q_next[-1] = 0
-    batch_states.pop(-1)
-    batch_rewards = np.asarray(batch_rewards) + discount * q_next
-    return batch_rewards
+    def process_state(self, network, state):
+        q = network.predict(np.stack([state]))
+        self.pending.append([q])
+        return q.data.numpy().flatten()
 
-def OneStepSarsa(batch_states, batch_actions, batch_rewards,
-                 tailing_state, tailing_action, terminal, network, discount):
-    batch_states.append(tailing_state)
-    batch_actions.append(tailing_action)
-    q_next = network.predict(np.asarray(batch_states[1:]))
-    q_next = q_next[np.arange(len(batch_actions[1:])), batch_actions[1:]]
-    if terminal:
-        q_next[-1] = 0
-    batch_states.pop(-1)
-    batch_actions.pop(-1)
-    batch_rewards = np.asarray(batch_rewards) + discount * q_next
-    return batch_rewards
+    def process_interaction(self, action, reward, next_state):
+        self.pending[-1].extend([action, reward, next_state])
+
+    def compute_loss(self, network, terminal):
+        loss = 0
+        valid_length = len(self.pending)
+        if not terminal:
+            valid_length -= 1
+        for i in range(valid_length):
+            q, action, reward, next_state = self.pending[i]
+            q_next = self.agent.target_network.predict(np.stack([next_state])).data
+            if i < len(self.pending) - 1:
+                next_action = self.pending[i + 1][1]
+                q_next = q_next.gather(1, torch.LongTensor([[next_action]]))
+            else:
+                q_next = torch.FloatTensor([[0]])
+            q_next = self.agent.discount * q_next + reward
+            q = q.gather(1, Variable(torch.LongTensor([[action]])))
+            loss += 0.5 * (q - Variable(q_next)).pow(2)
+        self.pending = []
+        return loss
+
+
+class OneStepQLearning:
+    def __init__(self, agent):
+        self.agent = agent
+        self.pending = []
+
+    def process_state(self, network, state):
+        q = network.predict(np.stack([state]))
+        self.pending.append([q])
+        return q.data.numpy().flatten()
+
+    def process_interaction(self, action, reward, next_state):
+        self.pending[-1].extend([action, reward, next_state])
+
+    def compute_loss(self, network, terminal):
+        loss = 0
+        for i in range(len(self.pending)):
+            q, action, reward, next_state = self.pending[i]
+            q_next, _ = self.agent.target_network.predict(np.stack([next_state])).data.max(1)
+            if terminal and i == len(self.pending) - 1:
+                q_next = torch.FloatTensor([[0]])
+            q_next = self.agent.discount * q_next + reward
+            q = q.gather(1, Variable(torch.LongTensor([[action]])))
+            loss += 0.5 * (q - Variable(q_next)).pow(2)
+        self.pending = []
+        return loss
+
+class NStepQLearning:
+    def __init__(self, agent):
+        self.agent = agent
+        self.pending = []
+
+    def process_state(self, network, state):
+        q = network.predict(np.stack([state]))
+        self.pending.append([q])
+        return q.data.numpy().flatten()
+
+    def process_interaction(self, action, reward, next_state):
+        self.pending[-1].extend([action, reward])
+        self.tailing_state = next_state
+
+    def compute_loss(self, network, terminal):
+        loss = 0
+        if terminal:
+            R = torch.FloatTensor([[0]])
+        else:
+            R, _ = self.agent.target_network.predict(
+                np.stack([self.tailing_state])).data.max(1)
+
+        for i in reversed(range(len(self.pending))):
+            q, action, reward = self.pending[i]
+            R = reward + self.agent.discount * R
+            loss += 0.5 * (Variable(R) - q.gather(1, Variable(torch.LongTensor([[action]])))).pow(2)
+        self.pending = []
+        return loss
+
+class AdvantageActorCritic:
+    def __init__(self, agent):
+        self.agent = agent
+        self.pending = []
+
+    def process_state(self, network, state):
+        prob, log_prob, value = network.predict(np.stack([state]))
+        self.pending.append([prob, log_prob, value])
+        return prob.data.numpy().flatten()
+
+    def process_interaction(self, action, reward, next_state):
+        self.pending[-1].extend([action, reward])
+        self.tailing_state = next_state
+
+    def compute_loss(self, network, terminal):
+        loss = 0
+        if terminal:
+            R = torch.FloatTensor([[0]])
+        else:
+            R = network.critic(np.stack([self.tailing_state])).data
+        for i in reversed(range(len(self.pending))):
+            prob, log_prob, value, action, reward = self.pending[i]
+            R = reward + self.agent.discount * R
+            advantage = Variable(R) - value
+            loss += 0.5 * advantage.pow(2)
+            loss += -log_prob.gather(1, Variable(torch.LongTensor([[action]]))) * Variable(advantage.data)
+            loss += 0.01 * torch.sum(torch.mul(prob, log_prob))
+        self.pending = []
+        return loss
 
-def AdvantageActorCritic(batch_states, batch_actions, batch_rewards,
-                         tailing_state, tailing_action, terminal, network, discount):
-    if terminal:
-        reward = 0
-    else:
-        reward = np.asscalar(network.critic(np.stack([tailing_state])))
-    rewards = []
-    for r in reversed(batch_rewards):
-        reward = r + discount * reward
-        rewards.append(reward)
-    return rewards

dqn_agent.py

@@ -59,7 +59,7 @@ class DQNAgent:
         total_reward = 0.0
         steps = 0
         while not self.step_limit or steps < self.step_limit:
-            value = self.learning_network.predict(np.stack([self.task.normalize_state(state)]))
+            value = self.learning_network.predict(np.stack([self.task.normalize_state(state)]), True)
             if deterministic:
                 action = np.argmax(value.flatten())
             else:
@@ -81,9 +81,9 @@ class DQNAgent:
                 states, actions, rewards, next_states, terminals = experiences
                 states = self.task.normalize_state(states)
                 next_states = self.task.normalize_state(next_states)
-                q_next = self.target_network.predict(next_states, False).detach()
+                q_next = self.target_network.predict(next_states).detach()
                 if self.double_q:
-                    _, best_actions = self.learning_network.predict(next_states, False).detach().max(1)
+                    _, best_actions = self.learning_network.predict(next_states).detach().max(1)
                     q_next = q_next.gather(1, best_actions)
                 else:
                     q_next, _ = q_next.max(1)
@@ -92,7 +92,7 @@ class DQNAgent:
                 q_next = q_next * (1 - terminals)
                 q_next.add_(rewards)
                 actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
-                q = self.learning_network.predict(states, False)
+                q = self.learning_network.predict(states)
                 q = q.gather(1, actions)
                 loss = self.learning_network.criterion(q, q_next)
                 self.learning_network.zero_grad()

main.py

@@ -6,8 +6,8 @@ def dqn_cart_pole():
     config = dict()
     config['task_fn'] = lambda: CartPole()
     config['optimizer_fn'] = lambda params: torch.optim.RMSprop(params, 0.001)
-    # config['network_fn'] = lambda optimizer_fn: FullyConnectedNet([8, 50, 200, 2], optimizer_fn)
-    config['network_fn'] = lambda optimizer_fn: DuelingFullyConnectedNet([8, 50, 200, 2], optimizer_fn)
+    config['network_fn'] = lambda optimizer_fn: FullyConnectedNet([8, 50, 200, 2], optimizer_fn)
+    # config['network_fn'] = lambda optimizer_fn: DuelingFullyConnectedNet([8, 50, 200, 2], optimizer_fn)
     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
@@ -27,16 +27,16 @@ def async_cart_pole():
     config = dict()
     config['task_fn'] = lambda: CartPole()
     config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001)
-    config['network_fn'] = lambda: FullyConnectedNet([4, 50, 200, 2], gpu=False)
-    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
+    config['network_fn'] = lambda: FullyConnectedNet([4, 50, 200, 2])
+    config['policy_fn'] = lambda: GreedyPolicy(epsilon=0.5, final_step=5000, min_epsilon=0.1)
+    config['bootstrap'] = OneStepQLearning
+    # config['bootstrap'] = NStepQLearning
+    # config['bootstrap'] = OneStepSarsa
     config['discount'] = 0.99
     config['target_network_update_freq'] = 200
     config['step_limit'] = 0
     config['n_workers'] = 16
-    config['batch_size'] = 6
+    config['update_interval'] = 6
     config['test_interval'] = 4000
     config['test_repetitions'] = 50
     config['history_length'] = 1
@@ -45,19 +45,20 @@ def async_cart_pole():
     agent.run()
 
 def a3c_cart_pole():
+    update_interval = 6
     config = dict()
     config['task_fn'] = lambda: CartPole()
     config['optimizer_fn'] = lambda params: torch.optim.Adam(params, 0.001)
-    config['network_fn'] = lambda: FCActorCriticNet([4, 200, 2], gpu=False)
+    config['network_fn'] = lambda: FCActorCriticNet([4, 200, 2])
     config['policy_fn'] = SamplePolicy
-    config['bootstrap_fn'] = AdvantageActorCritic
+    config['bootstrap'] = AdvantageActorCritic
     config['discount'] = 0.99
-    config['target_network_update_freq'] = 0
+    config['target_network_update_freq'] = 200
     config['step_limit'] = 0
     config['n_workers'] = 16
-    config['batch_size'] = 6
-    config['test_interval'] = 4000
+    config['update_interval'] = update_interval
     config['history_length'] = 1
+    config['test_interval'] = 4000
     config['test_repetitions'] = 50
     config['logger'] = gym.logger
     agent = AsyncAgent(**config)
@@ -89,23 +90,25 @@ def dqn_pixel_atari(name):
 
 def async_pixel_atari(name):
     config = dict()
-    history_length = 4
+    history_length = 1
     n_actions = 6
-    config['task_fn'] = lambda: PixelAtari(name, no_op=30, frame_skip=4)
+    config['task_fn'] = lambda: PixelAtari(name, no_op=30, frame_skip=4, frame_size=42)
     config['optimizer_fn'] = lambda params: torch.optim.Adam(params, lr=0.0001)
-    config['network_fn'] = lambda: NipsConvNet(history_length, n_actions, gpu=False)
+    config['network_fn'] = lambda: OpenAIConvNet(history_length,
+                                                 n_actions,
+                                                 LSTM=False)
     config['policy_fn'] = lambda: StochasticGreedyPolicy(epsilons=[1.0, 1.0, 1.0],
                                                           final_step=1000000,
                                                           min_epsilons=[0.1, 0.01, 0.5],
                                                           probs=[0.4, 0.3, 0.3])
-    # config['bootstrap_fn'] = OneStepQLearning
-    # config['bootstrap_fn'] = NStepQLearning
-    config['bootstrap_fn'] = OneStepSarsa
+    # config['bootstrap'] = OneStepQLearning
+    config['bootstrap'] = NStepQLearning
+    # config['bootstrap'] = OneStepSarsa
     config['discount'] = 0.99
     config['target_network_update_freq'] = 10000
     config['step_limit'] = 10000
     config['n_workers'] = 16
-    config['batch_size'] = 32
+    config['update_interval'] = 32
     config['test_interval'] = 50000
     config['test_repetitions'] = 1
     config['history_length'] = history_length
@@ -117,16 +120,18 @@ def a3c_pixel_atari(name):
     config = dict()
     history_length = 1
     n_actions = 6
-    config['task_fn'] = lambda: PixelAtari(name, no_op=30, frame_skip=4)
+    config['task_fn'] = lambda: PixelAtari(name, no_op=30, frame_skip=4, frame_size=42)
     config['optimizer_fn'] = lambda params: torch.optim.Adam(params, lr=0.0001)
-    config['network_fn'] = lambda: ConvActorCriticNet(history_length, n_actions, gpu=False)
+    config['network_fn'] = lambda: OpenAIConvActorCriticNet(history_length,
+                                                            n_actions,
+                                                            LSTM=True)
     config['policy_fn'] = SamplePolicy
-    config['bootstrap_fn'] = AdvantageActorCritic
+    config['bootstrap'] = AdvantageActorCritic
     config['discount'] = 0.99
     config['target_network_update_freq'] = 0
     config['step_limit'] = 10000
     config['n_workers'] = 16
-    config['batch_size'] = 20
+    config['update_interval'] = 20
     config['test_interval'] = 50000
     config['test_repetitions'] = 1
     config['history_length'] = history_length
@@ -138,12 +143,14 @@ if __name__ == '__main__':
     # gym.logger.setLevel(logging.DEBUG)
     gym.logger.setLevel(logging.INFO)
 
-    # async_cart_pole()
     # dqn_cart_pole()
-    # dqn_pixel_atari('BreakoutNoFrameskip-v3')
-    # async_pixel_atari('BreakoutNoFrameskip-v3')
-    # a3c_pixel_atari('BreakoutNoFrameskip-v3')
+    # async_cart_pole()
     # a3c_cart_pole()
+
     # dqn_pixel_atari('PongNoFrameskip-v3')
-    async_pixel_atari('PongNoFrameskip-v3')
+    # async_pixel_atari('PongNoFrameskip-v3')
     # a3c_pixel_atari('PongNoFrameskip-v3')
+
+    # dqn_pixel_atari('BreakoutNoFrameskip-v3')
+    async_pixel_atari('BreakoutNoFrameskip-v3')
+    # a3c_pixel_atari('BreakoutNoFrameskip-v3')

network.py

@@ -12,63 +12,57 @@ import numpy as np
 
 # Base class for all kinds of network
 class BasicNet:
-    def __init__(self, optimizer_fn, gpu):
+    def __init__(self, optimizer_fn, gpu, LSTM=False):
         if optimizer_fn is not None:
             self.optimizer = optimizer_fn(self.parameters())
         self.gpu = gpu and torch.cuda.is_available()
+        self.LSTM = LSTM
         if self.gpu:
             print 'Transferring network to GPU...'
             self.cuda()
             print 'Network transferred.'
 
     def to_torch_variable(self, x, dtype='float32'):
-        x = torch.from_numpy(np.asarray(x, dtype=dtype))
+        if not isinstance(x, torch.FloatTensor):
+            x = torch.from_numpy(np.asarray(x, dtype=dtype))
         if self.gpu:
             x = x.cuda()
         return Variable(x)
 
+    def reset(self, terminal):
+        if not self.LSTM:
+            return
+        if terminal:
+            self.h.data.zero_()
+            self.c.data.zero_()
+        self.h = Variable(self.h.data)
+        self.c = Variable(self.c.data)
+
 # Base class for value based methods
 class VanillaNet(BasicNet):
-    def predict(self, x, to_numpy=True):
+    def predict(self, x, to_numpy=False):
         y = self.forward(x)
         if to_numpy:
             y = y.cpu().data.numpy()
         return y
 
-    def gradient(self, x, actions, targets):
-        y = self.forward(x)
-        y = y.gather(1, actions)
-        loss = self.criterion(y, targets)
-        loss.backward()
-
 # Base class for actor critic method
 class ActorCriticNet(BasicNet):
     def predict(self, x):
-        phi = self.forward(x)
-        return F.softmax(self.fc_actor(phi)).cpu().data.numpy()
-
-    def gradient(self, x, actions, rewards):
-        phi = self.forward(x)
-        logit = self.fc_actor(phi)
-        prob = F.softmax(logit)
-        log_prob_ = F.log_softmax(logit)
-        state_value = self.fc_critic(phi)
-        log_prob = log_prob_.gather(1, actions)
-        advantage = (rewards - state_value).detach()
-        policy_loss = -torch.sum(log_prob * advantage)
-        value_loss = 0.5 * torch.sum(torch.pow(rewards - state_value, 2))
-        entropy = -torch.sum(torch.mul(prob, log_prob_))
-        loss = policy_loss + value_loss - self.xentropy_weight * entropy
-        loss.backward()
-        nn.utils.clip_grad_norm(self.parameters(), self.grad_threshold)
+        phi = self.forward(x, True)
+        pre_prob = self.fc_actor(phi)
+        prob = F.softmax(pre_prob)
+        log_prob = F.log_softmax(pre_prob)
+        value = self.fc_critic(phi)
+        return prob, log_prob, value
 
     def critic(self, x):
-        phi = self.forward(x)
-        return self.fc_critic(phi).cpu().data.numpy()
+        phi = self.forward(x, False)
+        return self.fc_critic(phi)
 
 # Base class for dueling architecture
 class DuelingNet(BasicNet):
-    def predict(self, x, to_numpy=True):
+    def predict(self, x, to_numpy=False):
         phi = self.forward(x)
         value = self.fc_value(phi)
         advantange = self.fc_advantage(phi)
@@ -77,6 +71,8 @@ class DuelingNet(BasicNet):
             return q.cpu().data.numpy()
         return q
 
+# Starting of several network instances
+
 # Network for CartPole with value based methods
 class FullyConnectedNet(nn.Module, VanillaNet):
     def __init__(self, dims, optimizer_fn=None, gpu=True):
@@ -178,21 +174,30 @@ class DuelingConvNet(nn.Module, DuelingNet):
 class FCActorCriticNet(nn.Module, ActorCriticNet):
     def __init__(self,
                  dims,
-                 xentropy_weight=0.01,
-                 grad_threshold=40,
-                 gpu=True):
+                 LSTM=False):
         super(FCActorCriticNet, self).__init__()
-        self.fc1 = nn.Linear(dims[0], dims[1])
+        if LSTM:
+            self.layer1 = nn.LSTMCell(dims[0], dims[1])
+        else:
+            self.layer1 = nn.Linear(dims[0], dims[1])
         self.fc_actor = nn.Linear(dims[1], dims[2])
         self.fc_critic = nn.Linear(dims[1], 1)
-        self.xentropy_weight = xentropy_weight
-        self.grad_threshold = grad_threshold
-        BasicNet.__init__(self, optimizer_fn=None, gpu=gpu)
+        BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM)
+        if LSTM:
+            self.h = self.to_torch_variable(np.zeros((1, dims[1])))
+            self.c = self.to_torch_variable(np.zeros((1, dims[1])))
 
-    def forward(self, x):
+    def forward(self, x, update_LSTM=True):
         x = self.to_torch_variable(x)
         x = x.view(x.size(0), -1)
-        phi = self.fc1(x)
+        if self.LSTM:
+            h, c = self.layer1(x, (self.h, self.c))
+            if update_LSTM:
+                self.h = h
+                self.c = c
+            phi = h
+        else:
+            phi = self.layer1(x)
         return phi
 
 # Network for pixel Atari game with actor critic
@@ -222,3 +227,87 @@ class ConvActorCriticNet(nn.Module, ActorCriticNet):
         y = y.view(y.size(0), -1)
         return F.elu(self.fc4(y))
 
+class OpenAIConvActorCriticNet(nn.Module, ActorCriticNet):
+    def __init__(self,
+                 in_channels,
+                 n_actions,
+                 LSTM=False):
+        super(OpenAIConvActorCriticNet, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1)
+        self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+        self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+
+        self.LSTM = LSTM
+        hidden_units = 256
+
+        if LSTM:
+            self.layer5 = nn.LSTMCell(32 * 3 * 3, hidden_units)
+        else:
+            self.layer5 = nn.Linear(32 * 3 * 3, hidden_units)
+
+        self.fc_actor = nn.Linear(hidden_units, n_actions)
+        self.fc_critic = nn.Linear(hidden_units, 1)
+        BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM)
+        if LSTM:
+            self.h = self.to_torch_variable(np.zeros((1, hidden_units)))
+            self.c = self.to_torch_variable(np.zeros((1, hidden_units)))
+
+    def forward(self, x, update_LSTM=True):
+        x = self.to_torch_variable(x)
+        y = F.elu(self.conv1(x))
+        y = F.elu(self.conv2(y))
+        y = F.elu(self.conv3(y))
+        y = F.elu(self.conv4(y))
+        y = y.view(y.size(0), -1)
+        if self.LSTM:
+            h, c = self.layer5(y, (self.h, self.c))
+            if update_LSTM:
+                self.h = h
+                self.c = c
+            phi = h
+        else:
+            phi = F.elu(self.layer5(y))
+        return phi
+
+class OpenAIConvNet(nn.Module, VanillaNet):
+    def __init__(self,
+                 in_channels,
+                 n_actions,
+                 LSTM=False):
+        super(OpenAIConvNet, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1)
+        self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+        self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+
+        self.LSTM = LSTM
+        hidden_units = 256
+
+        if LSTM:
+            self.layer5 = nn.LSTMCell(32 * 3 * 3, hidden_units)
+        else:
+            self.layer5 = nn.Linear(32 * 3 * 3, hidden_units)
+
+        self.fc6 = nn.Linear(hidden_units, n_actions)
+        BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM)
+        if LSTM:
+            self.h = self.to_torch_variable(np.zeros((1, hidden_units)))
+            self.c = self.to_torch_variable(np.zeros((1, hidden_units)))
+
+    def forward(self, x, update_LSTM=True):
+        x = self.to_torch_variable(x)
+        y = F.elu(self.conv1(x))
+        y = F.elu(self.conv2(y))
+        y = F.elu(self.conv3(y))
+        y = F.elu(self.conv4(y))
+        y = y.view(y.size(0), -1)
+        if self.LSTM:
+            h, c = self.layer5(y, (self.h, self.c))
+            if update_LSTM:
+                self.h = h
+                self.c = c
+            phi = h
+        else:
+            phi = F.elu(self.layer5(y))
+        return self.fc6(phi)

policy.py

@@ -13,7 +13,9 @@ class GreedyPolicy:
         self.min_epsilon = min_epsilon
         self.final_step = final_step
 
-    def sample(self, action_value):
+    def sample(self, action_value, deterministic=False):
+        if deterministic:
+            return np.argmax(action_value)
         if np.random.rand() < self.epsilon:
             return np.random.randint(0, len(action_value))
         return np.argmax(action_value)
@@ -30,15 +32,17 @@ class StochasticGreedyPolicy:
         for epsilon, min_epsilon in zip(epsilons, min_epsilons):
             self.policies.append(GreedyPolicy(epsilon, final_step, min_epsilon))
 
-    def sample(self, action_value):
-        return np.random.choice(self.policies, p=self.probs).sample(action_value)
+    def sample(self, action_value, deterministic=False):
+        return np.random.choice(self.policies, p=self.probs).sample(action_value, deterministic)
 
     def update_epsilon(self):
         for policy in self.policies:
             policy.update_epsilon()
 
 class SamplePolicy:
-    def sample(self, action_value):
+    def sample(self, action_value, deterministic=False):
+        if deterministic:
+            return np.argmax(action_value)
         return np.random.choice(np.arange(len(action_value)), p=action_value)
     def update_epsilon(self):
         pass

task.py

@@ -53,12 +53,12 @@ class LunarLander(BasicTask):
         self.env = gym.make(self.name)
 
 class PixelAtari(BasicTask):
-    success_threshold = 1000
-
-    def __init__(self, name, no_op, frame_skip, normalized_state=True):
+    def __init__(self, name, no_op, frame_skip, normalized_state=True,
+                 frame_size=84, success_threshold=1000):
         BasicTask.__init__(self)
         self.normalized_state = normalized_state
         self.name = name
+        self.success_threshold = success_threshold
         env = gym.make(name)
         assert 'NoFrameskip' in env.spec.id
         env = EpisodicLifeEnv(env)
@@ -66,7 +66,7 @@ class PixelAtari(BasicTask):
         env = MaxAndSkipEnv(env, skip=frame_skip)
         if 'FIRE' in env.unwrapped.get_action_meanings():
             env = FireResetEnv(env)
-        env = ProcessFrame84(env)
+        env = ProcessFrame(env, frame_size)
         self.env = ClippedRewardsWrapper(env)
 
     def normalize_state(self, state):