[change] Move console printing to StatsWriter class

ml-agents/mlagents/trainers/learn.py

@@ -21,6 +21,7 @@
     CSVWriter,
     StatsReporter,
     GaugeWriter,
+    ConsoleWriter,
 )
 from mlagents_envs.environment import UnityEnvironment
 from mlagents.trainers.sampler_class import SamplerManager
@@ -270,9 +271,11 @@ def run_training(run_seed: int, options: RunOptions) -> None:
         )
         tb_writer = TensorboardWriter(summaries_dir)
         gauge_write = GaugeWriter()
+        console_writer = ConsoleWriter()
         StatsReporter.add_writer(tb_writer)
         StatsReporter.add_writer(csv_writer)
         StatsReporter.add_writer(gauge_write)
+        StatsReporter.add_writer(console_writer)
 
         if options.env_path is None:
             port = UnityEnvironment.DEFAULT_EDITOR_PORT

ml-agents/mlagents/trainers/stats.py

@@ -4,10 +4,14 @@
 import abc
 import csv
 import os
+import time
+import logging
 
 from mlagents.tf_utils import tf
 from mlagents_envs.timers import set_gauge
 
+logger = logging.getLogger("mlagents.trainers")
+
 
 class StatsSummary(NamedTuple):
     mean: float
@@ -58,6 +62,45 @@ def write_text(self, category: str, text: str, step: int) -> None:
         pass
 
 
+class ConsoleWriter(StatsWriter):
+    def __init__(self):
+        self.training_start_time = time.time()
+
+    def write_stats(
+        self, category: str, values: Dict[str, StatsSummary], step: int
+    ) -> None:
+        is_training = "Not Training."
+        if "Is Training" in values:
+            stats_summary = stats_summary = values["Is Training"]
+            if stats_summary.mean > 0.0:
+                is_training = "Training."
+        if "Environment/Cumulative Reward" in values:
+            stats_summary = values["Environment/Cumulative Reward"]
+            logger.info(
+                "{}: Step: {}. "
+                "Time Elapsed: {:0.3f} s "
+                "Mean "
+                "Reward: {:0.3f}"
+                ". Std of Reward: {:0.3f}. {}".format(
+                    category,
+                    step,
+                    time.time() - self.training_start_time,
+                    stats_summary.mean,
+                    stats_summary.std,
+                    is_training,
+                )
+            )
+        else:
+            logger.info(
+                "{}: Step: {}. No episode was completed since last summary. {}".format(
+                    category, step, is_training
+                )
+            )
+
+    def write_text(self, category: str, text: str, step: int) -> None:
+        pass
+
+
 class TensorboardWriter(StatsWriter):
     def __init__(self, base_dir: str):
         """

ml-agents/mlagents/trainers/tests/test_stats.py

@@ -2,6 +2,7 @@
 import os
 import pytest
 import tempfile
+import unittest
 import csv
 
 from mlagents.trainers.stats import (
@@ -10,6 +11,7 @@
     CSVWriter,
     StatsSummary,
     GaugeWriter,
+    ConsoleWriter,
 )
 
 
@@ -138,3 +140,33 @@ def test_gauge_stat_writer_sanitize():
         GaugeWriter.sanitize_string("Very/Very/Very Nested Stat")
         == "Very.Very.VeryNestedStat"
     )
+
+
+class ConsoleWriterTest(unittest.TestCase):
+    def test_console_writer(self):
+        # Test write_stats
+        with self.assertLogs("mlagents.trainers", level="INFO") as cm:
+            category = "category1"
+            console_writer = ConsoleWriter()
+            statssummary1 = StatsSummary(mean=1.0, std=1.0, num=1)
+            console_writer.write_stats(
+                category,
+                {
+                    "Environment/Cumulative Reward": statssummary1,
+                    "Is Training": statssummary1,
+                },
+                10,
+            )
+            statssummary2 = StatsSummary(mean=0.0, std=0.0, num=1)
+            console_writer.write_stats(
+                category,
+                {
+                    "Environment/Cumulative Reward": statssummary1,
+                    "Is Training": statssummary2,
+                },
+                10,
+            )
+        self.assertIn(
+            "Mean Reward: 1.000. Std of Reward: 1.000. Training.", cm.output[0]
+        )
+        self.assertIn("Not Training.", cm.output[1])

ml-agents/mlagents/trainers/trainer/trainer.py

@@ -9,7 +9,6 @@
 
 from collections import deque
 
-from mlagents_envs.timers import set_gauge
 from mlagents.model_serialization import export_policy_model, SerializationSettings
 from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.stats import StatsReporter
@@ -193,33 +192,7 @@ def _write_summary(self, step: int) -> None:
         """
         Saves training statistics to Tensorboard.
         """
-        is_training = "Training." if self.should_still_train else "Not Training."
-        stats_summary = self.stats_reporter.get_stats_summaries(
-            "Environment/Cumulative Reward"
-        )
-        if stats_summary.num > 0:
-            logger.info(
-                "{}: {}: Step: {}. "
-                "Time Elapsed: {:0.3f} s "
-                "Mean "
-                "Reward: {:0.3f}"
-                ". Std of Reward: {:0.3f}. {}".format(
-                    self.run_id,
-                    self.brain_name,
-                    step,
-                    time.time() - self.training_start_time,
-                    stats_summary.mean,
-                    stats_summary.std,
-                    is_training,
-                )
-            )
-            set_gauge(f"{self.brain_name}.mean_reward", stats_summary.mean)
-        else:
-            logger.info(
-                " {}: {}: Step: {}. No episode was completed since last summary. {}".format(
-                    self.run_id, self.brain_name, step, is_training
-                )
-            )
+        self.stats_reporter.add_stat("Is Training", float(self.should_still_train))
         self.stats_reporter.write_stats(int(step))
 
     @abc.abstractmethod