Feat: A3C

Author: choru-k

README.md

@@ -20,16 +20,16 @@ So you can run this example in your computer(maybe it take just only 1~2 minitue
 - [x] Actor Critic [[10]](#reference)
 - [x] Advantage Actor Critic
 - [x] GAE(Generalized Advantage Estimation) [[12]](#reference)
-- [ ] TNPG [[20]](#reference)
+- [x] TNPG [[20]](#reference)
 - [ ] TRPO [[13]](#reference)
 - [ ] PPO [[14]](#reference)
-- [ ] ACER [[21]](#reference)
 
 ## Parallel
 - [x] Asynchronous Q-learning [[11]](#reference)
-- [ ] A3C (Asynchronous Advange Actor Critice) [[11]](#reference)
+- [x] A3C (Asynchronous Advange Actor Critice) [[11]](#reference)
 
 ## Will
+- [ ] ACER [[21]](#reference)
 - [ ] APE-X [[15]](#reference)
 - [ ] R2D2 [[16]](#reference)
 - [ ] RND [[17]](#reference)

parallel/1-Async-Q-Learning/worker.py

@@ -46,11 +46,9 @@ def run(self):
         running_score = 0
         epsilon = 1.0
         steps = 0
-        total_step = 0
         while self.global_ep.value < max_episode:
             if self.global_ep_r.value > goal_score:
                 break
-            total_step +=1
             done = False
 
             score = 0

parallel/2-A3C/config.py

@@ -0,0 +1,10 @@
+import torch
+
+env_name = 'CartPole-v1'
+gamma = 0.99
+lr = 0.0001
+goal_score = 200
+log_interval = 10
+n_step = 10
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+max_episode = 30000

parallel/2-A3C/memory.py

@@ -1,9 +1,6 @@
 import random
 from collections import namedtuple
 
-# Taken from
-# https://github.com/pytorch/tutorials/blob/master/Reinforcement%20(Q-)Learning%20with%20PyTorch.ipynb
-
 Transition = namedtuple('Transition', ('state', 'next_state', 'action', 'reward', 'mask'))
 
 

parallel/2-A3C/model.py

@@ -2,6 +2,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from config import gamma
+
 def set_init(layers):
     for layer in layers:
         nn.init.normal_(layer.weight, mean=0., std=0.1)
@@ -13,26 +15,27 @@ def __init__(self, num_inputs, num_outputs):
         self.num_inputs = num_inputs
         self.num_outputs = num_outputs
 
-        self.fc1 = nn.Linear(num_inputs, 128)
-        self.fc2 = nn.Linear(128, 128)
+        self.fc = nn.Linear(num_inputs, 128)
         self.fc_actor = nn.Linear(128, num_outputs)
-
-        self.fc3 = nn.Linear(num_inputs, 128)
-        self.fc4 = nn.Linear(128, 128)
         self.fc_critic = nn.Linear(128, 1)
 
-        set_init([self.fc1, self.fc2, self.fc_actor, self.fc3, self.fc4, self.fc_critic])
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform(m.weight)
 
     def forward(self, input):
-        x = F.relu(self.fc1(input))
-        x = F.relu(self.fc2(x))
+        x = F.relu(self.fc(input))
         policy = F.softmax(self.fc_actor(x))
-
-        y = F.relu(self.fc3(input))
-        y = F.relu(self.fc4(y))
-        value = self.fc_critic(y)
+        value = self.fc_critic(x)
         return policy, value
 
+    def get_action(self, input):
+        policy, _ = self.forward(input)
+        policy = policy[0].data.numpy()
+
+        action = np.random.choice(self.num_outputs, 1, p=policy)[0]
+        return action
+
 
 class GlobalModel(Model):
     def __init__(self, num_inputs, num_outputs):
@@ -43,34 +46,41 @@ class LocalModel(Model):
     def __init__(self, num_inputs, num_outputs):
         super(LocalModel, self).__init__(num_inputs, num_outputs)
 
-    def push_to_global_model(self, batch, global_model, global_optimizer, args):
+    def push_to_global_model(self, batch, global_model, global_optimizer):
         states = torch.stack(batch.state)
         next_states = torch.stack(batch.next_state)
-        actions = torch.Tensor(batch.action).long()
+        actions = torch.stack(batch.action)
         rewards = torch.Tensor(batch.reward)
         masks = torch.Tensor(batch.mask)
 
-        policy, value = self.forward(states[0])
+        policy, value = self.forward(states)
+        policy = policy.view(-1, self.num_outputs)
+        value = value.view(-1)
+
         _, last_value = self.forward(next_states[-1])
 
-        running_returns = last_value[0]
+        running_return = last_value[0].data
+        running_returns = torch.zeros(rewards.size())
         for t in reversed(range(0, len(rewards))):
-            running_returns = rewards[t] + args.gamma * running_returns * masks[t]
+            running_return = rewards[t] + gamma * running_return * masks[t]
+            running_returns[t] = running_return
 
-        pred = running_returns
-        td_error = pred - value[0]
 
-        log_policy = torch.log(policy[0] + 1e-5)[actions[0]]
-        loss1 = - log_policy * td_error.item()
-        loss2 = F.mse_loss(value[0], pred.detach())
-        entropy = torch.log(policy + 1e-5) * policy
-        loss = loss1 + loss2 - 0.01 * entropy.sum()
+        td_error = running_returns - value.detach()
+        log_policy = (torch.log(policy + 1e-10) * actions).sum(dim=1, keepdim=True)
+        loss_policy = - log_policy * td_error
+        loss_value = torch.pow(td_error, 2)
+        entropy = (torch.log(policy + 1e-10) * policy).sum(dim=1, keepdim=True)
+
+        loss = (loss_policy + loss_value - 0.01 * entropy).mean()
 
         global_optimizer.zero_grad()
         loss.backward()
         for lp, gp in zip(self.parameters(), global_model.parameters()):
             gp._grad = lp.grad
         global_optimizer.step()
 
+        return loss
+
     def pull_from_global_model(self, global_model):
         self.load_state_dict(global_model.state_dict())

parallel/2-A3C/shared_adam.py

@@ -1,7 +1,3 @@
-"""
-Shared optimizer, the parameters in the optimizer will shared in the multiprocessors.
-"""
-
 import torch
 class SharedAdam(torch.optim.Adam): # extend a pytorch optimizer so it shares grads across processes
     def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):

parallel/2-A3C/train.py

@@ -1,49 +1,41 @@
-import os
-import sys
 import gym
-import argparse
-import numpy as np
+import torch
 
 from model import Model
 from worker import Worker
 from shared_adam import SharedAdam
 from tensorboardX import SummaryWriter
 import torch.multiprocessing as mp
 
-parser = argparse.ArgumentParser()
-parser.add_argument('--env_name', type=str, default="CartPole-v1", help='')
-parser.add_argument('--load_model', type=str, default=None)
-parser.add_argument('--save_path', default='./save_model/', help='')
-parser.add_argument('--render', default=False, action="store_true")
-parser.add_argument('--gamma', default=0.9, help='')
-parser.add_argument('--goal_score', default=400, help='')
-parser.add_argument('--log_interval', default=10, help='')
-parser.add_argument('--logdir', type=str, default='./logs',
-                    help='tensorboardx logs directory')
-parser.add_argument('--MAX_EP', default=10000)
-args = parser.parse_args()
+from config import env_name, lr
 
-if __name__ == "__main__":
-    env = gym.make(args.env_name)
-    global_model = Model(env.observation_space.shape[0], env.action_space.n)
+def main():
+    env = gym.make(env_name)
+    env.seed(500)
+    torch.manual_seed(500)
+
+    num_inputs = env.observation_space.shape[0]
+    num_actions = env.action_space.n
+    global_model = Model(num_inputs, num_actions)
     global_model.share_memory()
-    global_optimizer = SharedAdam(global_model.parameters(), lr=0.0001)
+    global_optimizer = SharedAdam(global_model.parameters(), lr=lr)
     global_ep, global_ep_r, res_queue = mp.Value('i', 0), mp.Value('d', 0.), mp.Queue()
 
-    # mp.cpu_count()
-    workers = [Worker(global_model, global_optimizer, global_ep, global_ep_r, res_queue, i, args) for i in range(mp.cpu_count())]
+    writer = SummaryWriter('logs')
+
+    workers = [Worker(global_model, global_optimizer, global_ep, global_ep_r, res_queue, i) for i in range(mp.cpu_count())]
     [w.start() for w in workers]
     res = []
     while True:
         r = res_queue.get()
         if r is not None:
             res.append(r)
+            [ep, ep_r, loss] = r
+            writer.add_scalar('log/score', float(ep_r), ep)
+            writer.add_scalar('log/loss', float(loss), ep)
         else:
             break
     [w.join() for w in workers]
 
-    # import matplotlib.pyplot as plt
-    # plt.plot(res)
-    # plt.ylabel('Moving average ep reward')
-    # plt.xlabel('Step')
-    # plt.show()
+if __name__=="__main__":
+    main()

parallel/2-A3C/worker.py

@@ -4,28 +4,13 @@
 import numpy as np
 from model import LocalModel
 from memory import Memory
-
-def record(global_ep, global_ep_r, ep_r, res_queue, name):
-    with global_ep.get_lock():
-        global_ep.value += 1
-    with global_ep_r.get_lock():
-        if global_ep_r.value == 0.:
-            global_ep_r.value = ep_r
-        else:
-            global_ep_r.value = global_ep_r.value * 0.99 + ep_r * 0.01
-    res_queue.put(global_ep_r.value)
-    print(
-        name,
-        "Ep:", global_ep.value,
-        "| Ep_r:", global_ep_r.value,
-    )
+from config import env_name, n_step, max_episode, log_interval
 
 class Worker(mp.Process):
-    def __init__(self, global_model, global_optimizer, global_ep, global_ep_r, res_queue, name, args):
+    def __init__(self, global_model, global_optimizer, global_ep, global_ep_r, res_queue, name):
         super(Worker, self).__init__()
-        self.args = args
 
-        self.env = gym.make(self.args.env_name)
+        self.env = gym.make(env_name)
         self.env.seed(500)
 
         self.name = 'w%i' % name
@@ -34,16 +19,28 @@ def __init__(self, global_model, global_optimizer, global_ep, global_ep_r, res_q
         self.local_model = LocalModel(self.env.observation_space.shape[0], self.env.action_space.n)
         self.num_actions = self.env.action_space.n
 
+    def record(self, score, loss):
+        with self.global_ep.get_lock():
+            self.global_ep.value += 1
+        with self.global_ep_r.get_lock():
+            if self.global_ep_r.value == 0.:
+                self.global_ep_r.value = score
+            else:
+                self.global_ep_r.value = 0.99 * self.global_ep_r.value + 0.01 * score
+        if self.global_ep.value % log_interval == 0:
+            print('{} , {} episode | score: {:.2f}'.format(
+                self.name, self.global_ep.value, self.global_ep_r.value))
+
+        self.res_queue.put([self.global_ep.value, self.global_ep_r.value, loss])
 
     def get_action(self, policy, num_actions):
         policy = policy.data.numpy()[0]
         action = np.random.choice(num_actions, 1, p=policy)[0]
         return action
 
     def run(self):
-        self.local_model.train()
-        total_step = 1
-        while self.global_ep.value < self.args.MAX_EP:
+
+        while self.global_ep.value < max_episode:
             self.local_model.pull_from_global_model(self.global_model)
             done = False
             score = 0
@@ -52,10 +49,9 @@ def run(self):
             state = self.env.reset()
             state = torch.Tensor(state)
             state = state.unsqueeze(0)
-            memory = Memory(100)
+            memory = Memory(n_step)
 
             while True:
-                self.local_model.eval()
                 policy, value = self.local_model(state)
                 action = self.get_action(policy, self.num_actions)
 
@@ -65,21 +61,22 @@ def run(self):
 
                 mask = 0 if done else 1
                 reward = reward if not done or score == 499 else -1
-                score += reward
+                action_one_hot = torch.zeros(2)
+                action_one_hot[action] = 1
+                memory.push(state, next_state, action_one_hot, reward, mask)
 
-                memory.push(state, next_state, action, reward, mask)
+                score += reward
+                state = next_state
 
-                if len(memory) == 10 or done:
+                if len(memory) == n_step or done:
                     batch = memory.sample()
-                    self.local_model.push_to_global_model(batch, self.global_model, self.global_optimizer, self.args)
+                    loss = self.local_model.push_to_global_model(batch, self.global_model, self.global_optimizer)
                     self.local_model.pull_from_global_model(self.global_model)
-                    memory = Memory(100)
+                    memory = Memory(n_step)
 
                     if done:
-                        record(self.global_ep, self.global_ep_r, score, self.res_queue, self.name)
+                        running_score = self.record(score, loss)
                         break
 
 
-                total_step += 1
-                state = next_state
         self.res_queue.put(None)