Add configurable seed for random numbers

a3c/train.py

@@ -1,4 +1,6 @@
 import os
+import numpy as np
+import random
 import torch
 import torch.multiprocessing as mp
 from .shared_adam import SharedAdam
@@ -6,7 +8,20 @@ from .net import Net
 from .worker import Worker
 
 
-def train(env, max_ep, model_checkpoint_dir, gamma=0., pretrained_model_path=None, save=False, min_reward=9.9, every_n_save=100):
+def _set_seed(seed: int = 100) -> None:
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+    # When running on the CuDNN backend, two further options must be set
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    # Set a fixed value for the hash seed
+    os.environ["PYTHONHASHSEED"] = str(seed)
+
+
+def train(env, max_ep, model_checkpoint_dir, gamma=0., seed=100, 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)
@@ -14,6 +29,8 @@ def train(env, max_ep, model_checkpoint_dir, gamma=0., pretrained_model_path=Non
     n_a = env.action_space.n
     words_list = env.words
     word_width = len(env.words[0])
+    # Set global seeds for randoms
+    _set_seed(seed)
     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))

main.py

@@ -16,7 +16,7 @@ def training_mode(args, env, model_checkpoint_dir):
     max_ep = args.games
     start_time = time.time()
     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)
+    global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir, args.gamma, args.seed, 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)
@@ -62,6 +62,8 @@ if __name__ == "__main__":
         "--model_name", "-m", help="If want to train from a pretrained model, the name of the pretrained model file")
     parser_train.add_argument(
         "--gamma", help="Gamma hyperparameter (discount factor) value", type=float, default=0.)
+    parser_train.add_argument(
+        "--seed", help="Seed used for random numbers generation", type=int, default=100)
     parser_train.add_argument(
         "--save", '-s', help="Save instances of the model while training", action='store_true')
     parser_train.add_argument(