File size: 922 Bytes
05c9ac2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | from typing import Dict, List
from mlagents.trainers.settings import RunOptions
from mlagents.trainers.stats import StatsWriter, StatsSummary
class ExampleStatsWriter(StatsWriter):
"""
Example implementation of the StatsWriter abstract class.
This doesn't do anything interesting, just prints the stats that it gets.
"""
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(run_options: RunOptions) -> List[StatsWriter]:
"""
Registration function. This is referenced in setup.py and will
be called by mlagents-learn when it starts to determine the
list of StatsWriters to use.
It must return a list of StatsWriters.
"""
print("Creating a new stats writer! This is so exciting!")
return [ExampleStatsWriter()]
|