Add posibility to save and load models

Also add an evaluation task to evaluate saved models

.gitignore

@@ -113,4 +113,7 @@ GitHub.sublime-settings
 !.vscode/tasks.json 
 !.vscode/launch.json 
 !.vscode/extensions.json 
-.history
+.history
+
+# PyTorch model files
+*.pth

a3c/discrete_A3C.py

@@ -6,14 +6,14 @@ View more on my Chinese tutorial page [莫烦Python](https://morvanzhou.github.i
 """
 import os
 import torch.multiprocessing as mp
-from .utils import v_wrap, push_and_pull, record
+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):
+    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
@@ -22,6 +22,7 @@ class Worker(mp.Process):
         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:
@@ -40,16 +41,18 @@ class Worker(mp.Process):
                     # 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)
+                    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)
 
 
-def train(env, max_ep):
+def train(env, max_ep, model_checkpoint_dir):
     os.environ["OMP_NUM_THREADS"] = "1"
-
+    if not os.path.exists(model_checkpoint_dir):
+        os.makedirs(model_checkpoint_dir)
     n_s = env.observation_space.shape[0]
     n_a = env.action_space.n
     words_list = env.words
@@ -60,7 +63,7 @@ def train(env, max_ep):
     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())]
+    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:

a3c/utils.py

@@ -1,6 +1,7 @@
 """
 Functions that use multiple times
 """
+import os
 from torch import nn
 import torch
 import numpy as np
@@ -46,6 +47,11 @@ def push_and_pull(opt, lnet, gnet, done, s_, bs, ba, br, gamma):
     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

main.py

@@ -1,10 +1,29 @@
 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 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")
@@ -21,9 +40,9 @@ def evaluate(net, env):
         # 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()
@@ -51,7 +70,13 @@ def print_results(global_ep, win_ep, res):
 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'
+    evaluation = True if len(sys.argv) > 3 and sys.argv[3] == 'evaluation' else False
     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)
+    model_checkpoint_dir = os.path.join('checkpoints', env.unwrapped.spec.id)
+    if not evaluation:
+        global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir)
+        print_results(global_ep, win_ep, res)
+        evaluate(gnet, env)
+    else:
+        results = evaluate_checkpoints(model_checkpoint_dir, env)
+        print(results)