Code refactor

Sent Worker class to own file and unify with some utils functions
Move main methods to a3c file

a3c/discrete_A3C.py

@@ -5,48 +5,12 @@ The most simple implementation for continuous action.
 View more on my Chinese tutorial page [莫烦Python](https://morvanzhou.github.io/).
 """
 import os
+import torch
 import torch.multiprocessing as mp
-from .utils import v_wrap, push_and_pull, record, save_model
 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, model_checkpoint_dir):
-        super(Worker, self).__init__()
-        self.max_ep = max_ep
-        self.name = 'w%02i' % name
-        self.g_ep, self.g_ep_r, self.res_queue, self.winning_ep = global_ep, global_ep_r, res_queue, winning_ep
-        self.gnet, self.opt = gnet, opt
-        self.word_list = words_list
-        self.lnet = Net(N_S, N_A, words_list, word_width)           # local network
-        self.env = env.unwrapped
-        self.model_checkpoint_dir = model_checkpoint_dir
-
-    def run(self):
-        while self.g_ep.value < self.max_ep:
-            s = self.env.reset()
-            buffer_s, buffer_a, buffer_r = [], [], []
-            ep_r = 0.
-            while True:
-                a = self.lnet.choose_action(v_wrap(s[None, :]))
-                s_, r, done, _ = self.env.step(a)
-                ep_r += r
-                buffer_a.append(a)
-                buffer_s.append(s)
-                buffer_r.append(r)
-
-                if done:  # update global and assign to local net
-                    # sync
-                    push_and_pull(self.opt, self.lnet, self.gnet, done, s_, buffer_s, buffer_a, buffer_r, GAMMA)
-                    goal_word = self.word_list[self.env.goal_word]
-                    record( self.g_ep, self.g_ep_r, ep_r, self.res_queue, self.name, goal_word, self.word_list[a], len(buffer_a), self.winning_ep)
-                    save_model(self.gnet, self.model_checkpoint_dir, self.g_ep.value, self.g_ep_r.value)
-                    buffer_s, buffer_a, buffer_r = [], [], []
-                    break
-                s = s_
-        self.res_queue.put(None)
+from .utils import v_wrap
+from .worker import Worker
 
 
 def train(env, max_ep, model_checkpoint_dir):
@@ -63,7 +27,8 @@ def train(env, max_ep, model_checkpoint_dir):
     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, model_checkpoint_dir) for i in range(mp.cpu_count())]
+    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, model_checkpoint_dir) for i in range(mp.cpu_count())]
     [w.start() for w in workers]
     res = []                    # record episode reward to plot
     while True:
@@ -74,3 +39,53 @@ def train(env, max_ep, model_checkpoint_dir):
             break
     [w.join() for w in workers]
     return global_ep, win_ep, gnet, res
+
+
+def evaluate_checkpoints(dir, env):
+    n_s = env.observation_space.shape[0]
+    n_a = env.action_space.n
+    words_list = env.words
+    word_width = len(env.words[0])
+    net = Net(n_s, n_a, words_list, word_width)
+    results = {}
+    for checkpoint in os.listdir(dir):
+        checkpoint_path = os.path.join(dir, checkpoint)
+        if os.path.isfile(checkpoint_path):
+            net.load_state_dict(torch.load(checkpoint_path))
+            wins, guesses = evaluate(net, env)
+            results[checkpoint] = wins, guesses
+    return dict(sorted(results.items(), key=lambda x: (x[1][0], -x[1][1]), reverse=True))
+
+
+def evaluate(net, env):
+    n_wins = 0
+    n_guesses = 0
+    n_win_guesses = 0
+    env = env.unwrapped
+    N = env.allowable_words
+    for goal_word in env.words[:N]:
+        win, outcomes = play(net, env)
+        if win:
+            n_wins += 1
+            n_win_guesses += len(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, "
+          f"{n_guesses / N} including losses.")
+    return n_wins/N*100, n_win_guesses/n_wins
+
+
+def play(net, env):
+    state = env.reset()
+    outcomes = []
+    win = False
+    for i in range(env.max_turns):
+        action = net.choose_action(v_wrap(state[None, :]))
+        state, reward, done, _ = env.step(action)
+        outcomes.append((env.words[action], reward))
+        if done:
+            if reward >= 0:
+                win = True
+            break
+    return win, outcomes

a3c/utils.py

@@ -1,8 +1,3 @@
-"""
-Functions that use multiple times
-"""
-import os
-from torch import nn
 import torch
 import numpy as np
 
@@ -11,64 +6,3 @@ def v_wrap(np_array, dtype=np.float32):
     if np_array.dtype != dtype:
         np_array = np_array.astype(dtype)
     return torch.from_numpy(np_array)
-
-
-def set_init(layers):
-    for layer in layers:
-        nn.init.normal_(layer.weight, mean=0., std=0.1)
-        nn.init.constant_(layer.bias, 0.)
-
-
-def push_and_pull(opt, lnet, gnet, done, s_, bs, ba, br, gamma):
-    if done:
-        v_s_ = 0.               # terminal
-    else:
-        v_s_ = lnet.forward(v_wrap(s_[None, :]))[-1].data.numpy()[0, 0]
-
-    buffer_v_target = []
-    for r in br[::-1]:    # reverse buffer r
-        v_s_ = r + gamma * v_s_
-        buffer_v_target.append(v_s_)
-    buffer_v_target.reverse()
-
-    loss = lnet.loss_func(
-        v_wrap(np.vstack(bs)),
-        v_wrap(np.array(ba), dtype=np.int64) if ba[0].dtype == np.int64 else v_wrap(np.vstack(ba)),
-        v_wrap(np.array(buffer_v_target)[:, None]))
-
-    # calculate local gradients and push local parameters to global
-    opt.zero_grad()
-    loss.backward()
-    for lp, gp in zip(lnet.parameters(), gnet.parameters()):
-        gp._grad = lp.grad
-    opt.step()
-
-    # pull global parameters
-    lnet.load_state_dict(gnet.state_dict())
-
-
-def save_model(gnet, dir, episode, reward):
-    if reward >= 9 and episode % 100 == 0:
-        torch.save(gnet.state_dict(), os.path.join(dir, f'model_{episode}.pth'))
-
-
-def record(global_ep, global_ep_r, ep_r, res_queue, name, goal_word, action, action_number, winning_ep):
-    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)
-    if goal_word == action:
-        winning_ep.value += 1
-        if global_ep.value % 100 == 0:
-            print(
-                name,
-                "Ep:", global_ep.value,
-                "| Ep_r: %.0f" % global_ep_r.value,
-                "| Goal :", goal_word,
-                "| Action: ", action,
-                "| Actions: ", action_number
-            )

a3c/worker.py

@@ -0,0 +1,107 @@
+"""
+Worker class implementation of the a3c discrete algorithm
+"""
+import os
+import torch
+import numpy as np
+import torch.multiprocessing as mp
+from torch import nn
+from .net import Net
+from .utils import v_wrap
+
+
+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, model_checkpoint_dir):
+        super(Worker, self).__init__()
+        self.max_ep = max_ep
+        self.name = 'w%02i' % name
+        self.g_ep, self.g_ep_r, self.res_queue, self.winning_ep = global_ep, global_ep_r, res_queue, winning_ep
+        self.gnet, self.opt = gnet, opt
+        self.word_list = words_list
+        # local network
+        self.lnet = Net(N_S, N_A, words_list, word_width)
+        self.env = env.unwrapped
+        self.model_checkpoint_dir = model_checkpoint_dir
+
+    def run(self):
+        while self.g_ep.value < self.max_ep:
+            s = self.env.reset()
+            buffer_s, buffer_a, buffer_r = [], [], []
+            ep_r = 0.
+            while True:
+                a = self.lnet.choose_action(v_wrap(s[None, :]))
+                s_, r, done, _ = self.env.step(a)
+                ep_r += r
+                buffer_a.append(a)
+                buffer_s.append(s)
+                buffer_r.append(r)
+
+                if done:  # update global and assign to local net
+                    # sync
+                    self.push_and_pull(done, s_, buffer_s,
+                                       buffer_a, buffer_r, GAMMA)
+                    goal_word = self.word_list[self.env.goal_word]
+                    self.record(ep_r, goal_word,
+                                self.word_list[a], len(buffer_a))
+                    self.save_model()
+                    buffer_s, buffer_a, buffer_r = [], [], []
+                    break
+                s = s_
+        self.res_queue.put(None)
+
+    def push_and_pull(self, done, s_, bs, ba, br, gamma):
+        if done:
+            v_s_ = 0.               # terminal
+        else:
+            v_s_ = self.lnet.forward(v_wrap(
+                s_[None, :]))[-1].data.numpy()[0, 0]
+
+        buffer_v_target = []
+        for r in br[::-1]:    # reverse buffer r
+            v_s_ = r + gamma * v_s_
+            buffer_v_target.append(v_s_)
+        buffer_v_target.reverse()
+
+        loss = self.lnet.loss_func(
+            v_wrap(np.vstack(bs)),
+            v_wrap(np.array(ba), dtype=np.int64) if ba[0].dtype == np.int64 else v_wrap(np.vstack(ba)),
+            v_wrap(np.array(buffer_v_target)[:, None]))
+
+        # calculate local gradients and push local parameters to global
+        self.opt.zero_grad()
+        loss.backward()
+        for lp, gp in zip(self.lnet.parameters(), self.gnet.parameters()):
+            gp._grad = lp.grad
+        self.opt.step()
+
+        # pull global parameters
+        self.lnet.load_state_dict(self.gnet.state_dict())
+
+    def save_model(self):
+        if self.g_ep_r.value >= 9 and self.g_ep.value % 100 == 0:
+            torch.save(self.gnet.state_dict(), os.path.join(
+                self.model_checkpoint_dir, f'model_{ self.g_ep.value}.pth'))
+
+    def record(self, ep_r, goal_word, action, action_number):
+        with self.g_ep.get_lock():
+            self.g_ep.value += 1
+        with self.g_ep_r.get_lock():
+            if self.g_ep_r.value == 0.:
+                self.g_ep_r.value = ep_r
+            else:
+                self.g_ep_r.value = self.g_ep_r.value * 0.99 + ep_r * 0.01
+        self.res_queue.put(self.g_ep_r.value)
+        if goal_word == action:
+            self.winning_ep.value += 1
+            if self.g_ep.value % 100 == 0:
+                print(
+                    self.name,
+                    "Ep:", self.g_ep.value,
+                    "| Ep_r: %.0f" % self.g_ep_r.value,
+                    "| Goal :", goal_word,
+                    "| Action: ", action,
+                    "| Actions: ", action_number
+                )

main.py

@@ -1,62 +1,10 @@
 import sys
 import os
 import gym
-import torch
 import matplotlib.pyplot as plt
-from a3c.discrete_A3C import train
-from a3c.utils import v_wrap
-from a3c.net import Net
+from a3c.discrete_A3C import train, evaluate, evaluate_checkpoints
 from wordle_env.wordle import WordleEnvBase
 
-def evaluate_checkpoints(dir, env):
-    n_s = env.observation_space.shape[0]
-    n_a = env.action_space.n
-    words_list = env.words
-    word_width = len(env.words[0])
-    net = Net(n_s, n_a, words_list, word_width)
-    results = {}
-    print(dir)
-    for checkpoint in os.listdir(dir):
-            checkpoint_path = os.path.join(dir, checkpoint)
-            if os.path.isfile(checkpoint_path):
-                net.load_state_dict(torch.load(checkpoint_path))
-                wins, guesses = evaluate(net, env)
-                results[checkpoint] = wins, guesses
-    return dict(sorted(results.items(), key=lambda x: (x[1][0], -x[1][1]), reverse=True))
-
-
-def evaluate(net, env):
-    print("Evaluation mode")
-    n_wins = 0
-    n_guesses = 0
-    n_win_guesses = 0
-    env = env.unwrapped
-    N = env.allowable_words
-    for goal_word in env.words[:N]:
-        win, outcomes = play(net, env)
-        if win:
-            n_wins += 1
-            n_win_guesses += len(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, "
-          f"{n_guesses / N} including losses.")
-    return n_wins/N*100, n_win_guesses/n_wins
-
-def play(net, env):
-    state = env.reset()
-    outcomes = []
-    win = False
-    for i in range(env.max_turns):
-        action = net.choose_action(v_wrap(state[None, :]))
-        state, reward, done, _ = env.step(action)
-        outcomes.append((env.words[action], reward))
-        if done:
-            if reward >= 0:
-                win = True
-            break
-    return win, outcomes
 
 def print_results(global_ep, win_ep, res):
     print("Jugadas:", global_ep.value)
@@ -78,5 +26,6 @@ if __name__ == "__main__":
         print_results(global_ep, win_ep, res)
         evaluate(gnet, env)
     else:
+        print("Evaluation mode")
         results = evaluate_checkpoints(model_checkpoint_dir, env)
         print(results)