Delete constant GAMMA and add it as an command line argument

a3c/train.py

@@ -6,7 +6,7 @@ from .net import Net
 from .worker import Worker
 
 
-def train(env, max_ep, model_checkpoint_dir, pretrained_model_path=None):
+def train(env, max_ep, model_checkpoint_dir, gamma=0., pretrained_model_path=None):
     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, pretrained_model_path=None):
 
     # 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, pretrained_model_path) for i in range(mp.cpu_count())]
+                      words_list, word_width, win_ep, model_checkpoint_dir, gamma, 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

@@ -10,11 +10,8 @@ from .net import Net
 from .utils import v_wrap
 
 
-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, 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, pretrained_model_path=None):
         super(Worker, self).__init__()
         self.max_ep = max_ep
         self.name = 'w%02i' % name
@@ -26,6 +23,7 @@ class Worker(mp.Process):
         if pretrained_model_path:
             self.lnet.load_state_dict(torch.load(pretrained_model_path))
         self.env = env.unwrapped
+        self.gamma = gamma
         self.model_checkpoint_dir = model_checkpoint_dir
 
     def run(self):
@@ -44,7 +42,7 @@ class Worker(mp.Process):
                 if done:  # update global and assign to local net
                     # sync
                     self.push_and_pull(done, s_, buffer_s,
-                                       buffer_a, buffer_r, GAMMA)
+                                       buffer_a, buffer_r)
                     goal_word = self.word_list[self.env.goal_word]
                     self.record(ep_r, goal_word,
                                 self.word_list[a], len(buffer_a))
@@ -54,7 +52,7 @@ class Worker(mp.Process):
                 s = s_
         self.res_queue.put(None)
 
-    def push_and_pull(self, done, s_, bs, ba, br, gamma):
+    def push_and_pull(self, done, s_, bs, ba, br):
         if done:
             v_s_ = 0.               # terminal
         else:
@@ -63,7 +61,7 @@ class Worker(mp.Process):
 
         buffer_v_target = []
         for r in br[::-1]:    # reverse buffer r
-            v_s_ = r + gamma * v_s_
+            v_s_ = r + self.gamma * v_s_
             buffer_v_target.append(v_s_)
         buffer_v_target.reverse()
 

main.py

@@ -18,9 +18,9 @@ def training_mode(args, env, model_checkpoint_dir):
         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, pretrained_model_path)
+            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)
+        global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir, args.gamma)
     print("--- %.0f seconds ---" % (time.time() - start_time))
     print_results(global_ep, win_ep, res)
     evaluate(gnet, env)