Add the posiblity to save checkpoints of the model and the condition on which the model is saved as arguments

a3c/train.py

@@ -6,7 +6,7 @@ from .net import Net
 from .worker import Worker
 
 
-def train(env, max_ep, model_checkpoint_dir, gamma=0., pretrained_model_path=None):
+def train(env, max_ep, model_checkpoint_dir, gamma=0., pretrained_model_path=None, save=False, min_reward=9.9, every_n_save=100):
     os.environ["OMP_NUM_THREADS"] = "1"
     if not os.path.exists(model_checkpoint_dir):
         os.makedirs(model_checkpoint_dir)
@@ -23,7 +23,7 @@ def train(env, max_ep, model_checkpoint_dir, gamma=0., pretrained_model_path=Non
 
     # 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, gamma, pretrained_model_path) for i in range(mp.cpu_count())]
+                      words_list, word_width, win_ep, model_checkpoint_dir, gamma, pretrained_model_path, save, min_reward, every_n_save) for i in range(mp.cpu_count())]
     [w.start() for w in workers]
     res = []  # record episode reward to plot
     while True:
@@ -33,4 +33,6 @@ def train(env, max_ep, model_checkpoint_dir, gamma=0., pretrained_model_path=Non
         else:
             break
     [w.join() for w in workers]
+    if save:
+        torch.save(gnet.state_dict(), os.path.join(model_checkpoint_dir, f'model_{env.unwrapped.spec.id}.pth'))
     return global_ep, win_ep, gnet, res

a3c/worker.py

@@ -11,7 +11,28 @@ from .utils import v_wrap
 
 
 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, gamma, pretrained_model_path=None):
+    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,
+            gamma=0.,
+            pretrained_model_path=None,
+            save=False,
+            min_reward=9.9,
+            every_n_save=100
+    ):
         super(Worker, self).__init__()
         self.max_ep = max_ep
         self.name = 'w%02i' % name
@@ -25,6 +46,9 @@ class Worker(mp.Process):
         self.env = env.unwrapped
         self.gamma = gamma
         self.model_checkpoint_dir = model_checkpoint_dir
+        self.save = save
+        self.min_reward = min_reward
+        self.every_n_save = every_n_save
 
     def run(self):
         while self.g_ep.value < self.max_ep:
@@ -81,9 +105,9 @@ class Worker(mp.Process):
         self.lnet.load_state_dict(self.gnet.state_dict())
 
     def save_model(self):
-        if self.g_ep_r.value >= 9.9 and self.g_ep.value % 100 == 0:
+        if self.save and self.g_ep_r.value >= self.min_reward and self.g_ep.value % self.every_n_save == 0:
             torch.save(self.gnet.state_dict(), os.path.join(
-                self.model_checkpoint_dir, f'model_{ self.g_ep.value}.pth'))
+                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():

main.py

@@ -14,13 +14,8 @@ from wordle_env.wordle import WordleEnvBase
 def training_mode(args, env, model_checkpoint_dir):
     max_ep = args.games
     start_time = time.time()
-    if args.model_name:
-        pretrained_model_path = os.path.join(
-            model_checkpoint_dir, args.model_name)
-        global_ep, win_ep, gnet, res = train(
-            env, max_ep, model_checkpoint_dir, args.gamma, pretrained_model_path)
-    else:
-        global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir, args.gamma)
+    pretrained_model_path = os.path.join(model_checkpoint_dir, args.model_name) if args.model_name else args.model_name
+    global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir, args.gamma, pretrained_model_path, args.save, args.min_reward, args.every_n_save)
     print("--- %.0f seconds ---" % (time.time() - start_time))
     print_results(global_ep, win_ep, res)
     evaluate(gnet, env)
@@ -56,7 +51,13 @@ if __name__ == "__main__":
     parser_train.add_argument(
         "--model_name", "-n", help="If want to train from a pretrained model, the name of the pretrained model file")
     parser_train.add_argument(
-        "--gamma", help="Gamma hyperparameter value", type=float, default=0.)
+        "--gamma", help="Gamma hyperparameter (discount factor) value", type=float, default=0.)
+    parser_train.add_argument(
+        "--save", '-s', help="Save instances of the model while training", action='store_true')
+    parser_train.add_argument(
+        "--min_reward", help="The minimun global reward value achieved for saving the model", type=float, default=9.9)
+    parser_train.add_argument(
+        "--every_n_save", help="Check every n training steps to save the model", type=int, default=100)
     parser_train.set_defaults(func=training_mode)
 
     parser_eval = subparsers.add_parser(