A3C and main code refactor

Separate net in another fileø

a3c/discrete_A3C.py

@@ -4,81 +4,13 @@ The most simple implementation for continuous action.
 
 View more on my Chinese tutorial page [莫烦Python](https://morvanzhou.github.io/).
 """
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import gym
+import os
 import torch.multiprocessing as mp
-from .utils import v_wrap, set_init, push_and_pull, record
-import numpy as np
-
-GAMMA = 0.7
-
-class Net(nn.Module):
-    def __init__(self, s_dim, a_dim, word_list, words_width):
-        super(Net, self).__init__()
-        self.s_dim = s_dim
-        self.a_dim = a_dim
-        # n_emb = 32
-
-        word_width = 26 * words_width
-        layers = [
-            nn.Linear(s_dim, word_width),
-            nn.Tanh(),
-            # nn.Linear(128, word_width),
-            # nn.Tanh(),
-            # nn.Linear(256, n_emb),
-            # nn.Tanh(),
-        ]
-        self.v1 = nn.Sequential(*layers)
-        self.v4 = nn.Linear(word_width, 1)
-        self.actor_head = nn.Linear(word_width, word_width)
-
-        self.distribution = torch.distributions.Categorical
-        word_array = np.zeros((word_width, len(word_list)))
-        for i, word in enumerate(word_list):
-            for j, c in enumerate(word):
-                word_array[ j*26 + (ord(c) - ord('A')), i ] = 1
-        self.words = torch.Tensor(word_array)
-        # self.f_word = nn.Sequential(
-        #     nn.Linear(word_width, 64),
-        #     nn.ReLU(),
-        #     nn.Linear(64, n_emb),
-        # )
-
-    def forward(self, x):
-        # fw = self.f_word(
-        #     self.words.to(x.device.index),
-        # ).transpose(0, 1)
-        values = self.v1(x.float())
-        logits = torch.log_softmax(
-            torch.tensordot(self.actor_head(values), self.words,
-                            dims=((1,), (0,))),
-            dim=-1)
-        values = self.v4(values)
-        return logits, values
-
-    def choose_action(self, s):
-        self.eval()
-        logits, _ = self.forward(s)
-        prob = F.softmax(logits, dim=1).data
-        m = self.distribution(prob)
-        return m.sample().numpy()[0]
-
-    def loss_func(self, s, a, v_t):
-        self.train()
-        logits, values = self.forward(s)
-        td = v_t - values
-        c_loss = td.pow(2)
-        
-        probs = F.softmax(logits, dim=1)
-        m = self.distribution(probs)
-        exp_v = m.log_prob(a) * td.detach().squeeze()
-        a_loss = -exp_v
-        total_loss = (c_loss + a_loss).mean()
-        return total_loss
+from .utils import v_wrap, push_and_pull, record
+from .shared_adam import SharedAdam
+from .net import Net
 
+GAMMA = 0.65
 
 class Worker(mp.Process):
     def __init__(self, max_ep, gnet, opt, global_ep, global_ep_r, res_queue, name, env, N_S, N_A, words_list, word_width, winning_ep):
@@ -115,4 +47,27 @@ class Worker(mp.Process):
         self.res_queue.put(None)
 
 
-
+def train(env, max_ep):
+    os.environ["OMP_NUM_THREADS"] = "1"
+
+    n_s = env.observation_space.shape[0]
+    n_a = env.action_space.n
+    words_list = env.words
+    word_width = len(env.words[0])
+    gnet = Net(n_s, n_a, words_list, word_width)        # global network
+    gnet.share_memory()         # share the global parameters in multiprocessing
+    opt = SharedAdam(gnet.parameters(), lr=1e-4, betas=(0.92, 0.999))      # global optimizer
+    global_ep, global_ep_r, res_queue, win_ep = mp.Value('i', 0), mp.Value('d', 0.), mp.Queue(), mp.Value('i', 0)
+
+    # parallel training
+    workers = [Worker(max_ep, gnet, opt, global_ep, global_ep_r, res_queue, i, env, n_s, n_a, words_list, word_width, win_ep) for i in range(mp.cpu_count())]
+    [w.start() for w in workers]
+    res = []                    # record episode reward to plot
+    while True:
+        r = res_queue.get()
+        if r is not None:
+            res.append(r)
+        else:
+            break
+    [w.join() for w in workers]
+    return global_ep, win_ep, gnet, res

a3c/net.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+
+class Net(nn.Module):
+    def __init__(self, s_dim, a_dim, word_list, words_width):
+        super(Net, self).__init__()
+        self.s_dim = s_dim
+        self.a_dim = a_dim
+
+        word_width = 26 * words_width
+        layers = [
+            nn.Linear(s_dim, word_width),
+            nn.Tanh(),
+        ]
+        self.v1 = nn.Sequential(*layers)
+        self.v4 = nn.Linear(word_width, 1)
+        self.actor_head = nn.Linear(word_width, word_width)
+
+        self.distribution = torch.distributions.Categorical
+        word_array = np.zeros((word_width, len(word_list)))
+        for i, word in enumerate(word_list):
+            for j, c in enumerate(word):
+                word_array[ j*26 + (ord(c) - ord('A')), i ] = 1
+        self.words = torch.Tensor(word_array)
+
+    def forward(self, x):
+        values = self.v1(x.float())
+        logits = torch.log_softmax(
+            torch.tensordot(self.actor_head(values), self.words,
+                            dims=((1,), (0,))),
+            dim=-1)
+        values = self.v4(values)
+        return logits, values
+
+    def choose_action(self, s):
+        self.eval()
+        logits, _ = self.forward(s)
+        prob = F.softmax(logits, dim=1).data
+        m = self.distribution(prob)
+        return m.sample().numpy()[0]
+
+    def loss_func(self, s, a, v_t):
+        self.train()
+        logits, values = self.forward(s)
+        td = v_t - values
+        c_loss = td.pow(2)
+        
+        probs = F.softmax(logits, dim=1)
+        m = self.distribution(probs)
+        exp_v = m.log_prob(a) * td.detach().squeeze()
+        a_loss = -exp_v
+        total_loss = (c_loss + a_loss).mean()
+        return total_loss

a3c/shared_adam.py

@@ -1,7 +1,6 @@
 """
 Shared optimizer, the parameters in the optimizer will shared in the multiprocessors.
 """
-
 import torch
 
 
@@ -20,4 +19,3 @@ class SharedAdam(torch.optim.Adam):
                 # share in memory
                 state['exp_avg'].share_memory_()
                 state['exp_avg_sq'].share_memory_()
-

a3c/utils.py

@@ -1,7 +1,6 @@
 """
 Functions that use multiple times
 """
-
 from torch import nn
 import torch
 import numpy as np
@@ -66,4 +65,4 @@ def record(global_ep, global_ep_r, ep_r, res_queue, name, goal_word, action, act
                 "| Goal :", goal_word,
                 "| Action: ", action,
                 "| Actions: ", action_number
-            )
+            )

main.py

@@ -1,15 +1,10 @@
-import os
 import sys
 import gym
 import matplotlib.pyplot as plt
-import torch.multiprocessing as mp
-
-from a3c.discrete_A3C import Net, Worker
-from a3c.shared_adam import SharedAdam
+from a3c.discrete_A3C import train
 from a3c.utils import v_wrap
 from wordle_env.wordle import WordleEnvBase
 
-os.environ["OMP_NUM_THREADS"] = "1"
 
 def evaluate(net, env):
     print("Evaluation mode")
@@ -23,8 +18,8 @@ def evaluate(net, env):
         if win:
             n_wins += 1
             n_win_guesses += len(outcomes)
-        else:
-            print("Lost!", goal_word, outcomes)
+        # else:
+        #     print("Lost!", goal_word, outcomes)
         n_guesses += len(outcomes)
 
     print(f"Evaluation complete, won {n_wins/N*100}% and took {n_win_guesses/n_wins} guesses per win, "
@@ -44,34 +39,19 @@ def play(net, env):
             break
     return win, outcomes
 
-if __name__ == "__main__":
-    max_ep = int(sys.argv[1]) if len(sys.argv) > 1 else 100000
-    env_id = sys.argv[2] if len(sys.argv) > 2 else 'WordleEnv100FullAction-v0'
-    env = gym.make(env_id)
-    n_s = env.observation_space.shape[0]
-    n_a = env.action_space.n
-    words_list = env.words
-    word_width = len(env.words[0])
-    gnet = Net(n_s, n_a, words_list, word_width)        # global network
-    gnet.share_memory()         # share the global parameters in multiprocessing
-    opt = SharedAdam(gnet.parameters(), lr=1e-4, betas=(0.92, 0.999))      # global optimizer
-    global_ep, global_ep_r, res_queue, win_ep = mp.Value('i', 0), mp.Value('d', 0.), mp.Queue(), mp.Value('i', 0)
-
-    # parallel training
-    workers = [Worker(max_ep, gnet, opt, global_ep, global_ep_r, res_queue, i, env, n_s, n_a, words_list, word_width, win_ep) for i in range(mp.cpu_count())]
-    [w.start() for w in workers]
-    res = []                    # record episode reward to plot
-    while True:
-        r = res_queue.get()
-        if r is not None:
-            res.append(r)
-        else:
-            break
-    [w.join() for w in workers]
+def print_results(global_ep, win_ep, res):
     print("Jugadas:", global_ep.value)
     print("Ganadas:", win_ep.value)
     plt.plot(res)
     plt.ylabel('Moving average ep reward')
     plt.xlabel('Step')
     plt.show()
+
+
+if __name__ == "__main__":
+    max_ep = int(sys.argv[1]) if len(sys.argv) > 1 else 100000
+    env_id = sys.argv[2] if len(sys.argv) > 2 else 'WordleEnv100FullAction-v0'
+    env = gym.make(env_id)
+    global_ep, win_ep, gnet, res = train(env, max_ep)
+    print_results(global_ep, win_ep, res)
     evaluate(gnet, env)