colosseum.benchmark.benchmark

  1import dataclasses
  2import os
  3from enum import IntEnum
  4from typing import TYPE_CHECKING, Type, Dict
  5
  6import yaml.reader
  7
  8import colosseum
  9from colosseum import config
 10from colosseum.experiment import ExperimentConfig
 11from colosseum.utils import ensure_folder
 12from colosseum.utils.miscellanea import compare_gin_configs
 13
 14if TYPE_CHECKING:
 15    from colosseum.mdp import BaseMDP
 16
 17BENCHMARKS_DIRECTORY = (
 18    os.path.dirname(colosseum.__file__) + os.sep + "benchmark" + os.sep
 19)
 20
 21
 22@dataclasses.dataclass(frozen=True)
 23class ColosseumBenchmark:
 24    """
 25    The `ColosseumBenchmark` encodes a benchmark, which is composed of MDP configurations and an experimental
 26    configuration.
 27    """
 28
 29    name: str
 30    """The name to assign to the benchmark."""
 31    mdps_gin_configs: Dict[Type["BaseMDP"], str]
 32    """The gin config files of MDPs in the benchmark."""
 33    experiment_config: ExperimentConfig
 34    """The experiment configuration for the benchmark."""
 35
 36    def __eq__(self, other):
 37        if type(other) != ColosseumBenchmark:
 38            return False
 39        return (
 40            self.experiment_config == other.experiment_config
 41            and compare_gin_configs(self.mdps_gin_configs, other.mdps_gin_configs)
 42        )
 43
 44    def instantiate(self, benchmark_folder: str):
 45        """
 46        instantiates the benchmark locally.
 47
 48        Parameters
 49        ----------
 50        benchmark_folder : str
 51            The folder where the benchmark will be instantiated.
 52        """
 53        os.makedirs(ensure_folder(benchmark_folder) + "mdp_configs", exist_ok=True)
 54        for mdp_cl, gin_configs in self.mdps_gin_configs.items():
 55            with open(
 56                ensure_folder(benchmark_folder)
 57                + "mdp_configs"
 58                + os.sep
 59                + mdp_cl.__name__
 60                + ".gin",
 61                "w",
 62            ) as f:
 63                f.write(gin_configs)
 64
 65        self.experiment_config.store_at(benchmark_folder)
 66
 67    def get_experiments_benchmark_log_folder(self) -> str:
 68        """
 69        Returns
 70        -------
 71        str
 72            The experiment folder corresponding to the benchmark given the experiments folder configuration.
 73        """
 74        return self.get_log_folder(config.get_experiments_folder())
 75
 76    def get_hyperopt_benchmark_log_folder(self):
 77        """
 78        Returns
 79        -------
 80        str
 81            The parameters optimization folder corresponding to the benchmark given the experiments huperopt
 82            configuration.
 83        """
 84        return self.get_log_folder(config.get_hyperopt_folder())
 85
 86    def get_log_folder(self, benchmark_folder: str):
 87        """
 88        Parameters
 89        ----------
 90        benchmark_folder : str
 91            The folder where the benchmark could have been instantiated.
 92
 93        Returns
 94        -------
 95        str
 96            The folder that correspond to the benchmark if it was instantiated in the given benchmark folder.
 97        """
 98        return ensure_folder(benchmark_folder + self.name)
 99
100
101class ColosseumDefaultBenchmark(IntEnum):
102    CONTINUOUS_ERGODIC = 0
103    """The default benchmark for the continuous ergodic setting."""
104    CONTINUOUS_COMMUNICATING = 1
105    """The default benchmark for the continuous communicating setting."""
106    EPISODIC_ERGODIC = 2
107    """The default benchmark for the episodic ergodic setting."""
108    EPISODIC_COMMUNICATING = 3
109    """The default benchmark for the episodic communicating setting."""
110    EPISODIC_QUICK_TEST = 4
111    """A quick benchmark for the episodic setting."""
112    CONTINUOUS_QUICK_TEST = 5
113    """A quick benchmark for the continuous setting."""
114
115    @staticmethod
116    def get_default_experiment_config() -> ExperimentConfig:
117        """
118        Returns
119        -------
120        ExperimentConfig
121            The default experiment configuration proposed by the package for the tabular setting.
122        """
123        with open(BENCHMARKS_DIRECTORY + "experiment_config.yml", "r") as f:
124            experimental_config = yaml.load(f, yaml.Loader)
125        return ExperimentConfig(**experimental_config)
126
127    @staticmethod
128    def get_default_non_tabular_experiment_config() -> ExperimentConfig:
129        """
130        Returns
131        -------
132        ExperimentConfig
133            The default experiment configuration proposed by the package for the non-tabular setting.
134        """
135        from colosseum.emission_maps import StateInfo
136
137        default_experiment_config = (
138            ColosseumDefaultBenchmark.get_default_experiment_config()
139        )
140        default_experiment_config = dataclasses.asdict(default_experiment_config)
141        default_experiment_config["emission_map"] = StateInfo
142        return ExperimentConfig(**default_experiment_config)
143
144    def get_benchmark(
145        self,
146        postfix="",
147        experiment_config: ExperimentConfig = None,
148        non_tabular: bool = False,
149    ) -> ColosseumBenchmark:
150        """
151        creates a `ColosseumBenchmark` corresponding to the default benchmark the object encodes.
152
153        Parameters
154        ----------
155        postfix : str
156            A postfix string to add to the default name of the default benchmark.
157        experiment_config : ExperimentConfig
158            The experiment config to be substituted to the default one. By default, no substitution happens.
159        non_tabular : bool
160            If True, the default non-tabular experimental configuration is used.
161        Returns
162        -------
163        ColosseumBenchmark
164            The default benchmark the object encodes.
165        """
166
167        exp_folder = BENCHMARKS_DIRECTORY + f"benchmark_" + self.name.lower()
168        from colosseum.benchmark.utils import retrieve_benchmark
169
170        if experiment_config is None and "QUICK" not in self.name:
171            if non_tabular:
172                experiment_config = (
173                    ColosseumDefaultBenchmark.get_default_non_tabular_experiment_config()
174                )
175            else:
176                experiment_config = (
177                    ColosseumDefaultBenchmark.get_default_experiment_config()
178                )
179
180        benchmark = retrieve_benchmark(
181            exp_folder,
182            experiment_config,
183            f"{'_' if len(str(postfix)) > 0 else ''}{postfix}",
184        )
185        return benchmark