Add possibility to train from a pretrained model

a3c/discrete_A3C.py

@@ -13,7 +13,7 @@ from .utils import v_wrap
 from .worker import Worker
 
 
-def train(env, max_ep, model_checkpoint_dir):
+def train(env, max_ep, model_checkpoint_dir, pretrained_model_path=None):
     os.environ["OMP_NUM_THREADS"] = "1"
     if not os.path.exists(model_checkpoint_dir):
         os.makedirs(model_checkpoint_dir)
@@ -22,13 +22,15 @@ def train(env, max_ep, model_checkpoint_dir):
     words_list = env.words
     word_width = len(env.words[0])
     gnet = Net(n_s, n_a, words_list, word_width)        # global network
+    if pretrained_model_path:
+        gnet.load_state_dict(torch.load(pretrained_model_path))
     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, model_checkpoint_dir) for i in range(mp.cpu_count())]
+                      words_list, word_width, win_ep, model_checkpoint_dir, pretrained_model_path) for i in range(mp.cpu_count())]
     [w.start() for w in workers]
     res = []                    # record episode reward to plot
     while True:

a3c/worker.py

@@ -14,7 +14,7 @@ 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):
+    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, pretrained_model_path=None):
         super(Worker, self).__init__()
         self.max_ep = max_ep
         self.name = 'w%02i' % name
@@ -23,6 +23,8 @@ class Worker(mp.Process):
         self.word_list = words_list
         # local network
         self.lnet = Net(N_S, N_A, words_list, word_width)
+        if pretrained_model_path:
+            self.lnet.load_state_dict(torch.load(pretrained_model_path))
         self.env = env.unwrapped
         self.model_checkpoint_dir = model_checkpoint_dir
 

main.py

@@ -20,11 +20,16 @@ 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
+    pretrained = True if len(sys.argv) > 3 and sys.argv[3] == 'pretrained' else False
     env = gym.make(env_id)
     model_checkpoint_dir = os.path.join('checkpoints', env.unwrapped.spec.id)
     if not evaluation:
         start_time = time.time()
-        global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir)
+        if pretrained:
+            pretrained_model_path = os.path.join(model_checkpoint_dir, sys.argv[4]) if len(sys.argv) > 4 else ''
+            global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir, pretrained_model_path)
+        else:
+            global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir)
         print("--- %.0f seconds ---" % (time.time() - start_time))
         print_results(global_ep, win_ep, res)
         evaluate(gnet, env)