RL-1634: Update log_stats_writer_logger

(Issue RL-1634 - Dump aggregated rollout output)
enlite-ai · Apr 5, 2024 · a5ba6cb · a5ba6cb
1 parent d78c3c3
commit a5ba6cb
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 5 deletions.
diff --git a/maze/core/log_stats/log_stats_writer_logger.py b/maze/core/log_stats/log_stats_writer_logger.py
@@ -24,11 +24,11 @@ def write(self, path: str, step: int, stats: LogStats) -> None:
 
         # print run directory
         exp_dir = os.path.abspath(".")
-        print("Output directory: {}".format(exp_dir))
+        logger.info("Output directory: {}".format(exp_dir))
 
         # print stats
-        print("{0:>5}|{1:<104}|{2:>20}".format("step", "path", "value"))
-        print("{0:>5}|{1:<104}|{2:>20}".format("="*5, "="*100, "="*20))
+        logger.info("{0:>5}|{1:<104}|{2:>20}".format("step", "path", "value"))
+        logger.info("{0:>5}|{1:<104}|{2:>20}".format("="*5, "="*104, "="*20))
 
         lines = []
         for (event, name, groups), value in stats.items():

diff --git a/maze/core/rollout/parallel_rollout_runner.py b/maze/core/rollout/parallel_rollout_runner.py
@@ -270,7 +270,6 @@ def _configure_and_launch_processes(self, parallel_worker_type: type(ParallelRol
         # Perform writer registration -- after the forks so that it is not carried over to child processes
         if self.record_event_logs:
             LogEventsWriterRegistry.register_writer(LogEventsWriterTSV(log_dir="./event_logs"))
-        register_log_stats_writer(LogStatsWriterConsole())
         register_log_stats_writer(LogStatsWriterLogger())
         self.epoch_stats_aggregator = LogStatsAggregator(LogStatsLevel.EPOCH)
         self.epoch_stats_aggregator.register_consumer(get_stats_logger("rollout_stats"))

diff --git a/maze/core/rollout/sequential_rollout_runner.py b/maze/core/rollout/sequential_rollout_runner.py
@@ -55,7 +55,6 @@ def run_with(self, env: ConfigType, wrappers: CollectionOfConfigType, agent: Con
 
         # Set up the wrappers
         # Hydra handles working directory
-        register_log_stats_writer(LogStatsWriterConsole())
         register_log_stats_writer(LogStatsWriterLogger())
         if not isinstance(env, LogStatsWrapper):
             env = LogStatsWrapper.wrap(env, logging_prefix="rollout_data")