Add an FPS counter

Author: ikostrikov2

main.py

@@ -49,7 +49,6 @@
 
     # uncomment when it's fixed in pytorch
     # torch.manual_seed(args.seed)
-
     env = create_atari_env(args.env_name)
     shared_model = ActorCritic(
         env.observation_space.shape[0], env.action_space)
@@ -63,12 +62,15 @@
 
     processes = []
 
-    p = mp.Process(target=test, args=(args.num_processes, args, shared_model))
+    counter = mp.Value('i', 0)
+    lock = mp.Lock()
+
+    p = mp.Process(target=test, args=(args.num_processes, args, shared_model, counter))
     p.start()
     processes.append(p)
 
     for rank in range(0, args.num_processes):
-        p = mp.Process(target=train, args=(rank, args, shared_model, optimizer))
+        p = mp.Process(target=train, args=(rank, args, shared_model, counter, lock, optimizer))
         p.start()
         processes.append(p)
     for p in processes:

test.py

@@ -9,7 +9,7 @@
 from model import ActorCritic
 
 
-def test(rank, args, shared_model):
+def test(rank, args, shared_model, counter):
     torch.manual_seed(args.seed + rank)
 
     env = create_atari_env(args.env_name)
@@ -55,9 +55,10 @@ def test(rank, args, shared_model):
             done = True
 
         if done:
-            print("Time {}, episode reward {}, episode length {}".format(
+            print("Time {}, num steps {}, FPS {:.0f}, episode reward {}, episode length {}".format(
                 time.strftime("%Hh %Mm %Ss",
                               time.gmtime(time.time() - start_time)),
+                counter.value, counter.value / (time.time() - start_time),
                 reward_sum, episode_length))
             reward_sum = 0
             episode_length = 0

train.py

@@ -15,7 +15,7 @@ def ensure_shared_grads(model, shared_model):
         shared_param._grad = param.grad
 
 
-def train(rank, args, shared_model, optimizer=None):
+def train(rank, args, shared_model, counter, lock, optimizer=None):
     torch.manual_seed(args.seed + rank)
 
     env = create_atari_env(args.env_name)
@@ -64,6 +64,9 @@ def train(rank, args, shared_model, optimizer=None):
             done = done or episode_length >= args.max_episode_length
             reward = max(min(reward, 1), -1)
 
+            with lock:
+                counter.value += 1
+
             if done:
                 episode_length = 0
                 state = env.reset()