Add execution time of the models training and tune of the saving model parameters

a3c/worker.py

@@ -81,7 +81,7 @@ class Worker(mp.Process):
         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:
+        if self.g_ep_r.value >= 9.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'))
 

main.py

@@ -1,6 +1,7 @@
 import sys
 import os
 import gym
+import time
 import matplotlib.pyplot as plt
 from a3c.discrete_A3C import train, evaluate, evaluate_checkpoints
 from wordle_env.wordle import WordleEnvBase
@@ -22,7 +23,9 @@ if __name__ == "__main__":
     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)
+        print("--- %.0f seconds ---" % (time.time() - start_time))
         print_results(global_ep, win_ep, res)
         evaluate(gnet, env)
     else: