setuptools-based plugin for StatsWriters

ml-agents-plugin-examples/mlagents_plugin_examples/example_stats_writer.py

@@ -0,0 +1,14 @@
+from typing import Dict, List
+from mlagents.trainers.stats import StatsWriter, StatsSummary
+
+
+class ExampleStatsWriter(StatsWriter):
+    def write_stats(
+        self, category: str, values: Dict[str, StatsSummary], step: int
+    ) -> None:
+        print(f"ExampleStatsWriter category: {category} values: {values}")
+
+
+def get_example_stats_writer() -> List[StatsWriter]:
+    print("Creating a new stats writer! This is so exciting!")
+    return [ExampleStatsWriter()]

ml-agents-plugin-examples/setup.py

@@ -0,0 +1,11 @@
+from setuptools import setup
+
+setup(
+    name="mlagents_plugin_examples",
+    version="0.0.1",
+    entry_points={
+        "mlagents.stats_writer": [
+            "example=mlagents_plugin_examples.example_stats_writer:get_example_stats_writer"
+        ]
+    },
+)

ml-agents/mlagents/plugins/stats_writer.py

@@ -0,0 +1,20 @@
+from typing import List
+import importlib_metadata
+from mlagents.trainers.stats import StatsWriter
+
+
+def get_default_stats_writers() -> List[StatsWriter]:
+    # TODO move construction of default StatsWriters here
+    return []
+
+
+def register_stats_writer_plugins() -> List[StatsWriter]:
+    all_stats_writers: List[StatsWriter] = []
+    eps = importlib_metadata.entry_points()["mlagents.stats_writer"]
+
+    for ep in eps:
+        print(f"registering {ep.name}")
+        # TODO try/except around all of this
+        plugin_func = ep.load()
+        all_stats_writers += plugin_func()
+    return all_stats_writers

ml-agents/mlagents/trainers/learn.py

@@ -35,6 +35,7 @@
     add_metadata as add_timer_metadata,
 )
 from mlagents_envs import logging_util
+from mlagents.plugins.stats_writer import register_stats_writer_plugins
 
 logger = logging_util.get_logger(__name__)
 
@@ -91,14 +92,17 @@ def run_training(run_seed: int, options: RunOptions) -> None:
             )
 
         # Configure Tensorboard Writers and StatsReporter
-        tb_writer = TensorboardWriter(
-            write_path, clear_past_data=not checkpoint_settings.resume
-        )
-        gauge_write = GaugeWriter()
-        console_writer = ConsoleWriter()
-        StatsReporter.add_writer(tb_writer)
-        StatsReporter.add_writer(gauge_write)
-        StatsReporter.add_writer(console_writer)
+        stats_writers = [
+            TensorboardWriter(
+                write_path, clear_past_data=not checkpoint_settings.resume
+            ),
+            GaugeWriter(),
+            ConsoleWriter(),
+        ]
+
+        stats_writers += register_stats_writer_plugins()
+        for sw in stats_writers:
+            StatsReporter.add_writer(sw)
 
         if env_settings.env_path is None:
             port = None

ml-agents/setup.py

@@ -71,13 +71,17 @@ def run(self):
         "cattrs>=1.0.0,<1.1.0",
         "attrs>=19.3.0",
         'pypiwin32==223;platform_system=="Windows"',
+        "importlib_metadata",  # TODO for python<3.8 only
     ],
     python_requires=">=3.6.1",
     entry_points={
         "console_scripts": [
             "mlagents-learn=mlagents.trainers.learn:main",
             "mlagents-run-experiment=mlagents.trainers.run_experiment:main",
-        ]
+        ],
+        "mlagents.stats_writer": [
+            "default=mlagents.plugins.stats_writer:get_default_stats_writers"
+        ],
     },
     cmdclass={"verify": VerifyVersionCommand},
 )