AnnaMats/ppo-Pyramids-Training

1from typing import Dict, List2from mlagents.trainers.settings import RunOptions3from mlagents.trainers.stats import StatsWriter, StatsSummary4 5 6class ExampleStatsWriter(StatsWriter):7    """8    Example implementation of the StatsWriter abstract class.9    This doesn't do anything interesting, just prints the stats that it gets.10    """11 12    def write_stats(13        self, category: str, values: Dict[str, StatsSummary], step: int14    ) -> None:15        print(f"ExampleStatsWriter category: {category} values: {values}")16 17 18def get_example_stats_writer(run_options: RunOptions) -> List[StatsWriter]:19    """20    Registration function. This is referenced in setup.py and will21    be called by mlagents-learn when it starts to determine the22    list of StatsWriters to use.23 24    It must return a list of StatsWriters.25    """26    print("Creating a new stats writer! This is so exciting!")27    return [ExampleStatsWriter()]28