Merge pull request #89 from g-votte/async-tb-fix2

Fix the hang in `train_agent_async` with Tensorboard

pfrl/experiments/evaluator.py

@@ -324,6 +324,15 @@ def record_tb_stats(summary_writer, agent_stats, eval_stats, env_stats, t):
     summary_writer.flush()
 
 
+def record_tb_stats_loop(outdir, queue, stop_event):
+    tb_writer = create_tb_writer(outdir)
+
+    while not (stop_event.wait(1e-6) and queue.empty()):
+        if not queue.empty():
+            agent_stats, eval_stats, env_stats, t = queue.get()
+            record_tb_stats(tb_writer, agent_stats, eval_stats, env_stats, t)
+
+
 def save_agent(agent, t, outdir, logger, suffix=""):
     dirname = os.path.join(outdir, "{}{}".format(t, suffix))
     agent.save(dirname)
@@ -476,7 +485,6 @@ class AsyncEvaluator(object):
         save_best_so_far_agent (bool): If set to True, after each evaluation,
             if the score (= mean return of evaluation episodes) exceeds
             the best-so-far score, the current agent is saved.
-        use_tensorboard (bool): Additionally log eval stats to tensorboard
     """
 
     def __init__(
@@ -489,7 +497,6 @@ def __init__(
         step_offset=0,
         save_best_so_far_agent=True,
         logger=None,
-        use_tensorboard=False,
     ):
         assert (n_steps is None) != (n_episodes is None), (
             "One of n_steps or n_episodes must be None. "
@@ -501,7 +508,6 @@ def __init__(
         self.n_episodes = n_episodes
         self.eval_interval = eval_interval
         self.outdir = outdir
-        self.use_tensorboard = use_tensorboard
         self.max_episode_len = max_episode_len
         self.step_offset = step_offset
         self.save_best_so_far_agent = save_best_so_far_agent
@@ -518,8 +524,8 @@ def __init__(
         with open(os.path.join(self.outdir, "scores.txt"), "a"):
             pass
 
-        if use_tensorboard:
-            self.tb_writer = create_tb_writer(outdir)
+        self.record_tb_stats_queue = None
+        self.record_tb_stats_thread = None
 
     @property
     def max_score(self):
@@ -563,8 +569,8 @@ def evaluate_and_update_max_score(self, t, episodes, env, agent):
         )
         record_stats(self.outdir, values)
 
-        if self.use_tensorboard:
-            record_tb_stats(self.tb_writer, agent_stats, eval_stats, env_stats, t)
+        if self.record_tb_stats_queue is not None:
+            self.record_tb_stats_queue.put([agent_stats, eval_stats, env_stats, t])
 
         with self._max_score.get_lock():
             if mean > self._max_score.value:
@@ -589,3 +595,15 @@ def evaluate_if_necessary(self, t, episodes, env, agent):
                     self.wrote_header.value = True
             return self.evaluate_and_update_max_score(t, episodes, env, agent)
         return None
+
+    def start_tensorboard_writer(self, outdir, stop_event):
+        self.record_tb_stats_queue = mp.Queue()
+        self.record_tb_stats_thread = pfrl.utils.StoppableThread(
+            target=record_tb_stats_loop,
+            args=[outdir, self.record_tb_stats_queue, stop_event],
+            stop_event=stop_event,
+        )
+        self.record_tb_stats_thread.start()
+
+    def join_tensorboard_writer(self):
+        self.record_tb_stats_thread.join()

pfrl/experiments/train_agent_async.py

@@ -253,9 +253,10 @@ def train_agent_async(
             max_episode_len=max_episode_len,
             step_offset=step_offset,
             save_best_so_far_agent=save_best_so_far_agent,
-            use_tensorboard=use_tensorboard,
             logger=logger,
         )
+        if use_tensorboard:
+            evaluator.start_tensorboard_writer(outdir, stop_event)
 
     if random_seeds is None:
         random_seeds = np.arange(processes)
@@ -312,4 +313,7 @@ def f():
 
     stop_event.set()
 
+    if evaluator is not None and use_tensorboard:
+        evaluator.join_tensorboard_writer()
+
     return agent