Refactor file structure

async_agent.py

@@ -10,7 +10,10 @@ import numpy as np
 import torch.multiprocessing as mp
 from task import *
 from network import *
-from worker import *
+from async_workers.one_step_sarsa import *
+from async_workers.n_step_q import *
+from async_workers.actor_critic import *
+from async_workers.one_step_sarsa import *
 import pickle
 import os
 import time

async_workers/actor_critic.py

@@ -0,0 +1,77 @@
+#######################################################################
+# 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 numpy as np
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+
+class AdvantageActorCritic:
+    def __init__(self, config):
+        self.config = config
+        self.optimizer = config.optimizer_fn(config.learning_network.parameters())
+        self.worker_network = config.network_fn()
+        self.worker_network.load_state_dict(config.learning_network.state_dict())
+        self.task = config.task_fn()
+        self.policy = config.policy_fn()
+
+    def episode(self, deterministic=False):
+        config = self.config
+        state = self.task.reset()
+        steps = 0
+        total_reward = 0
+        pending = []
+        while not config.stop_signal.value and \
+                (not config.max_episode_length or steps < config.max_episode_length):
+            prob, log_prob, value = self.worker_network.predict(np.stack([state]))
+            action = self.policy.sample(prob.data.numpy().flatten(), deterministic)
+            next_state, reward, terminal, _ = self.task.step(action)
+
+            steps += 1
+            total_reward += reward
+
+            if deterministic:
+                if terminal:
+                    break
+                state = next_state
+                continue
+
+            pending.append([prob, log_prob, value, action, reward])
+            with config.steps_lock:
+                config.total_steps.value += 1
+
+            if terminal or len(pending) >= config.update_interval:
+                loss = 0
+                if terminal:
+                    R = torch.FloatTensor([[0]])
+                else:
+                    R = self.worker_network.critic(np.stack([next_state])).data
+                GAE = torch.FloatTensor([[0]])
+                for i in reversed(range(len(pending))):
+                    prob, log_prob, value, action, reward = pending[i]
+                    R = reward + config.discount * R
+                    advantage = Variable(R) - value
+                    GAE = config.discount * GAE + advantage.data
+                    loss += 0.5 * advantage.pow(2)
+                    loss += -log_prob.gather(1, Variable(torch.LongTensor([[action]]))) * Variable(GAE)
+                    loss += 0.01 * torch.sum(torch.mul(prob, log_prob))
+
+                pending = []
+                self.worker_network.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
+                self.optimizer.zero_grad()
+                for param, worker_param in zip(
+                        config.learning_network.parameters(), self.worker_network.parameters()):
+                    param._grad = worker_param.grad.clone()
+                self.optimizer.step()
+                self.worker_network.load_state_dict(config.learning_network.state_dict())
+                self.worker_network.reset(terminal)
+
+            if terminal:
+                break
+            state = next_state
+
+        return steps, total_reward

async_workers/n_step_q.py

@@ -0,0 +1,77 @@
+#######################################################################
+# 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 numpy as np
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+
+class NStepQLearning:
+    def __init__(self, config):
+        self.config = config
+        self.optimizer = config.optimizer_fn(config.learning_network.parameters())
+        self.worker_network = config.network_fn()
+        self.worker_network.load_state_dict(config.learning_network.state_dict())
+        self.task = config.task_fn()
+        self.policy = config.policy_fn()
+
+    def episode(self, deterministic=False):
+        config = self.config
+        state = self.task.reset()
+        steps = 0
+        total_reward = 0
+        pending = []
+        while not config.stop_signal.value and \
+                (not config.max_episode_length or steps < config.max_episode_length):
+            q = self.worker_network.predict(np.stack([state]))
+            action = self.policy.sample(q.data.numpy().flatten(), deterministic)
+            next_state, reward, terminal, _ = self.task.step(action)
+
+            steps += 1
+            total_reward += reward
+
+            if deterministic:
+                if terminal:
+                    break
+                state = next_state
+                continue
+
+            with config.steps_lock:
+                config.total_steps.value += 1
+            pending.append([q, action, reward])
+
+            if terminal or len(pending) >= config.update_interval:
+                loss = 0
+                if terminal:
+                    R = torch.FloatTensor([[0]])
+                else:
+                    R, _ = config.target_network.predict(
+                        np.stack([next_state])).data.max(1)
+
+                for i in reversed(range(len(pending))):
+                    q, action, reward = pending[i]
+                    R = reward + config.discount * R
+                    loss += 0.5 * (Variable(R) - q.gather(1, Variable(torch.LongTensor([[action]])))).pow(2)
+
+                pending = []
+                self.worker_network.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
+                self.optimizer.zero_grad()
+                for param, worker_param in zip(
+                        config.learning_network.parameters(), self.worker_network.parameters()):
+                    param._grad = worker_param.grad.clone()
+                self.optimizer.step()
+                self.worker_network.load_state_dict(config.learning_network.state_dict())
+                self.worker_network.reset(terminal)
+
+            if terminal:
+                break
+            state = next_state
+
+            if config.total_steps.value % config.target_network_update_freq == 0:
+                config.target_network.load_state_dict(config.learning_network.state_dict())
+
+        return steps, total_reward

async_workers/one_step_q.py

@@ -0,0 +1,75 @@
+#######################################################################
+# 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 numpy as np
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+
+class OneStepQLearning:
+    def __init__(self, config):
+        self.config = config
+        self.optimizer = config.optimizer_fn(config.learning_network.parameters())
+        self.worker_network = config.network_fn()
+        self.worker_network.load_state_dict(config.learning_network.state_dict())
+        self.task = config.task_fn()
+        self.policy = config.policy_fn()
+
+    def episode(self, deterministic=False):
+        config = self.config
+        state = self.task.reset()
+        steps = 0
+        total_reward = 0
+        pending = []
+        while not config.stop_signal.value and \
+                (not config.max_episode_length or steps < config.max_episode_length):
+            q = self.worker_network.predict(np.stack([state]))
+            action = self.policy.sample(q.data.numpy().flatten(), deterministic)
+            next_state, reward, terminal, _ = self.task.step(action)
+
+            steps += 1
+            total_reward += reward
+
+            if deterministic:
+                if terminal:
+                    break
+                state = next_state
+                continue
+
+            with config.steps_lock:
+                config.total_steps.value += 1
+            pending.append([q, action, reward, next_state])
+
+            if terminal or len(pending) >= config.update_interval:
+                loss = 0
+                for i in range(len(pending)):
+                    q, action, reward, next_state = pending[i]
+                    q_next, _ = config.target_network.predict(np.stack([next_state])).data.max(1)
+                    if terminal and i == len(pending) - 1:
+                        q_next = torch.FloatTensor([[0]])
+                    q_next = config.discount * q_next + reward
+                    q = q.gather(1, Variable(torch.LongTensor([[action]])))
+                    loss += 0.5 * (q - Variable(q_next)).pow(2)
+
+                pending = []
+                self.worker_network.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
+                self.optimizer.zero_grad()
+                for param, worker_param in zip(
+                        config.learning_network.parameters(), self.worker_network.parameters()):
+                    param._grad = worker_param.grad.clone()
+                self.optimizer.step()
+                self.worker_network.load_state_dict(config.learning_network.state_dict())
+                self.worker_network.reset(terminal)
+
+            if terminal:
+                break
+            state = next_state
+
+            if config.total_steps.value % config.target_network_update_freq == 0:
+                config.target_network.load_state_dict(config.learning_network.state_dict())
+
+        return steps, total_reward

async_workers/one_step_sarsa.py

@@ -0,0 +1,82 @@
+#######################################################################
+# 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 numpy as np
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+
+class OneStepSarsa:
+    def __init__(self, config):
+        self.config = config
+        self.optimizer = config.optimizer_fn(config.learning_network.parameters())
+        self.worker_network = config.network_fn()
+        self.worker_network.load_state_dict(config.learning_network.state_dict())
+        self.task = config.task_fn()
+        self.policy = config.policy_fn()
+
+    def episode(self, deterministic=False):
+        config = self.config
+        state = self.task.reset()
+        q = self.worker_network.predict(np.stack([state]))
+        action = self.policy.sample(q.data.numpy().flatten(), deterministic)
+        steps = 0
+        total_reward = 0
+        pending = []
+        while not config.stop_signal.value and \
+                (not config.max_episode_length or steps < config.max_episode_length):
+            next_state, reward, terminal, _ = self.task.step(action)
+            next_q = self.worker_network.predict(np.stack([next_state]))
+            next_action = self.policy.sample(next_q.data.numpy().flatten(), deterministic)
+            pending.append([q, action, reward, next_state, next_action])
+
+            steps += 1
+            total_reward += reward
+
+            if deterministic:
+                if terminal:
+                    break
+                state = next_state
+                action = next_action
+                continue
+
+            with config.steps_lock:
+                config.total_steps.value += 1
+
+            if terminal or len(pending) >= config.update_interval:
+                loss = 0
+                for i in range(len(pending)):
+                    q, action, reward, next_state, next_action = pending[i]
+                    q_next = config.target_network.predict(np.stack([next_state])).data
+                    if terminal and i == len(pending) - 1:
+                        q_next = torch.FloatTensor([[0]])
+                    else:
+                        q_next = q_next.gather(1, torch.LongTensor([[next_action]]))
+                    q_next = config.discount * q_next + reward
+                    q = q.gather(1, Variable(torch.LongTensor([[action]])))
+                    loss += 0.5 * (q - Variable(q_next)).pow(2)
+
+                pending = []
+                self.worker_network.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
+                self.optimizer.zero_grad()
+                for param, worker_param in zip(
+                        config.learning_network.parameters(), self.worker_network.parameters()):
+                    param._grad = worker_param.grad.clone()
+                self.optimizer.step()
+                self.worker_network.load_state_dict(config.learning_network.state_dict())
+                self.worker_network.reset(terminal)
+
+            if terminal:
+                break
+            else:
+                q = next_q
+                action = next_action
+
+            if config.total_steps.value % config.target_network_update_freq == 0:
+                config.target_network.load_state_dict(config.learning_network.state_dict())
+
+        return steps, total_reward

main.py

@@ -189,11 +189,11 @@ if __name__ == '__main__':
 
     # dqn_cart_pole()
     # async_cart_pole()
-    # a3c_cart_pole()
+    a3c_cart_pole()
 
     # dqn_pixel_atari('PongNoFrameskip-v3')
     # async_pixel_atari('PongNoFrameskip-v3')
-    a3c_pixel_atari('PongNoFrameskip-v3')
+    # a3c_pixel_atari('PongNoFrameskip-v3')
 
     # dqn_pixel_atari('BreakoutNoFrameskip-v3')
     # async_pixel_atari('BreakoutNoFrameskip-v3')

worker.py

@@ -1,284 +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 numpy as np
-import torch
-from torch.autograd import Variable
-import torch.nn as nn
-
-class AdvantageActorCritic:
-    def __init__(self, config):
-        self.config = config
-        self.optimizer = config.optimizer_fn(config.learning_network.parameters())
-        self.worker_network = config.network_fn()
-        self.worker_network.load_state_dict(config.learning_network.state_dict())
-        self.task = config.task_fn()
-        self.policy = config.policy_fn()
-
-    def episode(self, deterministic=False):
-        config = self.config
-        state = self.task.reset()
-        steps = 0
-        total_reward = 0
-        pending = []
-        while not config.stop_signal.value and \
-                (not config.max_episode_length or steps < config.max_episode_length):
-            prob, log_prob, value = self.worker_network.predict(np.stack([state]))
-            action = self.policy.sample(prob.data.numpy().flatten(), deterministic)
-            next_state, reward, terminal, _ = self.task.step(action)
-
-            steps += 1
-            total_reward += reward
-
-            if deterministic:
-                if terminal:
-                    break
-                state = next_state
-                continue
-
-            pending.append([prob, log_prob, value, action, reward])
-            with config.steps_lock:
-                config.total_steps.value += 1
-
-            if terminal or len(pending) >= config.update_interval:
-                loss = 0
-                if terminal:
-                    R = torch.FloatTensor([[0]])
-                else:
-                    R = self.worker_network.critic(np.stack([next_state])).data
-                GAE = torch.FloatTensor([[0]])
-                for i in reversed(range(len(pending))):
-                    prob, log_prob, value, action, reward = pending[i]
-                    R = reward + config.discount * R
-                    advantage = Variable(R) - value
-                    GAE = config.discount * GAE + advantage.data
-                    loss += 0.5 * advantage.pow(2)
-                    loss += -log_prob.gather(1, Variable(torch.LongTensor([[action]]))) * Variable(GAE)
-                    loss += 0.01 * torch.sum(torch.mul(prob, log_prob))
-
-                pending = []
-                self.worker_network.zero_grad()
-                loss.backward()
-                nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
-                self.optimizer.zero_grad()
-                for param, worker_param in zip(
-                        config.learning_network.parameters(), self.worker_network.parameters()):
-                    param._grad = worker_param.grad.clone()
-                self.optimizer.step()
-                self.worker_network.load_state_dict(config.learning_network.state_dict())
-                self.worker_network.reset(terminal)
-
-            if terminal:
-                break
-            state = next_state
-
-        return steps, total_reward
-
-class NStepQLearning:
-    def __init__(self, config):
-        self.config = config
-        self.optimizer = config.optimizer_fn(config.learning_network.parameters())
-        self.worker_network = config.network_fn()
-        self.worker_network.load_state_dict(config.learning_network.state_dict())
-        self.task = config.task_fn()
-        self.policy = config.policy_fn()
-
-    def episode(self, deterministic=False):
-        config = self.config
-        state = self.task.reset()
-        steps = 0
-        total_reward = 0
-        pending = []
-        while not config.stop_signal.value and \
-                (not config.max_episode_length or steps < config.max_episode_length):
-            q = self.worker_network.predict(np.stack([state]))
-            action = self.policy.sample(q.data.numpy().flatten(), deterministic)
-            next_state, reward, terminal, _ = self.task.step(action)
-
-            steps += 1
-            total_reward += reward
-
-            if deterministic:
-                if terminal:
-                    break
-                state = next_state
-                continue
-
-            with config.steps_lock:
-                config.total_steps.value += 1
-            pending.append([q, action, reward])
-
-            if terminal or len(pending) >= config.update_interval:
-                loss = 0
-                if terminal:
-                    R = torch.FloatTensor([[0]])
-                else:
-                    R, _ = config.target_network.predict(
-                        np.stack([next_state])).data.max(1)
-
-                for i in reversed(range(len(pending))):
-                    q, action, reward = pending[i]
-                    R = reward + config.discount * R
-                    loss += 0.5 * (Variable(R) - q.gather(1, Variable(torch.LongTensor([[action]])))).pow(2)
-
-                pending = []
-                self.worker_network.zero_grad()
-                loss.backward()
-                nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
-                self.optimizer.zero_grad()
-                for param, worker_param in zip(
-                        config.learning_network.parameters(), self.worker_network.parameters()):
-                    param._grad = worker_param.grad.clone()
-                self.optimizer.step()
-                self.worker_network.load_state_dict(config.learning_network.state_dict())
-                self.worker_network.reset(terminal)
-
-            if terminal:
-                break
-            state = next_state
-
-            if config.total_steps.value % config.target_network_update_freq == 0:
-                config.target_network.load_state_dict(config.learning_network.state_dict())
-
-        return steps, total_reward
-
-class OneStepQLearning:
-    def __init__(self, config):
-        self.config = config
-        self.optimizer = config.optimizer_fn(config.learning_network.parameters())
-        self.worker_network = config.network_fn()
-        self.worker_network.load_state_dict(config.learning_network.state_dict())
-        self.task = config.task_fn()
-        self.policy = config.policy_fn()
-
-    def episode(self, deterministic=False):
-        config = self.config
-        state = self.task.reset()
-        steps = 0
-        total_reward = 0
-        pending = []
-        while not config.stop_signal.value and \
-                (not config.max_episode_length or steps < config.max_episode_length):
-            q = self.worker_network.predict(np.stack([state]))
-            action = self.policy.sample(q.data.numpy().flatten(), deterministic)
-            next_state, reward, terminal, _ = self.task.step(action)
-
-            steps += 1
-            total_reward += reward
-
-            if deterministic:
-                if terminal:
-                    break
-                state = next_state
-                continue
-
-            with config.steps_lock:
-                config.total_steps.value += 1
-            pending.append([q, action, reward, next_state])
-
-            if terminal or len(pending) >= config.update_interval:
-                loss = 0
-                for i in range(len(pending)):
-                    q, action, reward, next_state = pending[i]
-                    q_next, _ = config.target_network.predict(np.stack([next_state])).data.max(1)
-                    if terminal and i == len(pending) - 1:
-                        q_next = torch.FloatTensor([[0]])
-                    q_next = config.discount * q_next + reward
-                    q = q.gather(1, Variable(torch.LongTensor([[action]])))
-                    loss += 0.5 * (q - Variable(q_next)).pow(2)
-
-                pending = []
-                self.worker_network.zero_grad()
-                loss.backward()
-                nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
-                self.optimizer.zero_grad()
-                for param, worker_param in zip(
-                        config.learning_network.parameters(), self.worker_network.parameters()):
-                    param._grad = worker_param.grad.clone()
-                self.optimizer.step()
-                self.worker_network.load_state_dict(config.learning_network.state_dict())
-                self.worker_network.reset(terminal)
-
-            if terminal:
-                break
-            state = next_state
-
-            if config.total_steps.value % config.target_network_update_freq == 0:
-                config.target_network.load_state_dict(config.learning_network.state_dict())
-
-        return steps, total_reward
-
-class OneStepSarsa:
-    def __init__(self, config):
-        self.config = config
-        self.optimizer = config.optimizer_fn(config.learning_network.parameters())
-        self.worker_network = config.network_fn()
-        self.worker_network.load_state_dict(config.learning_network.state_dict())
-        self.task = config.task_fn()
-        self.policy = config.policy_fn()
-
-    def episode(self, deterministic=False):
-        config = self.config
-        state = self.task.reset()
-        q = self.worker_network.predict(np.stack([state]))
-        action = self.policy.sample(q.data.numpy().flatten(), deterministic)
-        steps = 0
-        total_reward = 0
-        pending = []
-        while not config.stop_signal.value and \
-                (not config.max_episode_length or steps < config.max_episode_length):
-            next_state, reward, terminal, _ = self.task.step(action)
-            next_q = self.worker_network.predict(np.stack([next_state]))
-            next_action = self.policy.sample(next_q.data.numpy().flatten(), deterministic)
-            pending.append([q, action, reward, next_state, next_action])
-
-            steps += 1
-            total_reward += reward
-
-            if deterministic:
-                if terminal:
-                    break
-                state = next_state
-                action = next_action
-                continue
-
-            with config.steps_lock:
-                config.total_steps.value += 1
-
-            if terminal or len(pending) >= config.update_interval:
-                loss = 0
-                for i in range(len(pending)):
-                    q, action, reward, next_state, next_action = pending[i]
-                    q_next = config.target_network.predict(np.stack([next_state])).data
-                    if terminal and i == len(pending) - 1:
-                        q_next = torch.FloatTensor([[0]])
-                    else:
-                        q_next = q_next.gather(1, torch.LongTensor([[next_action]]))
-                    q_next = config.discount * q_next + reward
-                    q = q.gather(1, Variable(torch.LongTensor([[action]])))
-                    loss += 0.5 * (q - Variable(q_next)).pow(2)
-
-                pending = []
-                self.worker_network.zero_grad()
-                loss.backward()
-                nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
-                self.optimizer.zero_grad()
-                for param, worker_param in zip(
-                        config.learning_network.parameters(), self.worker_network.parameters()):
-                    param._grad = worker_param.grad.clone()
-                self.optimizer.step()
-                self.worker_network.load_state_dict(config.learning_network.state_dict())
-                self.worker_network.reset(terminal)
-
-            if terminal:
-                break
-            else:
-                q = next_q
-                action = next_action
-
-            if config.total_steps.value % config.target_network_update_freq == 0:
-                config.target_network.load_state_dict(config.learning_network.state_dict())
-
-        return steps, total_reward