#2476 Get episode schedule working

src/primaite/session/episode_schedule.py

@@ -0,0 +1,127 @@
+import copy
+from abc import ABC, abstractmethod
+from os import PathLike
+from pathlib import Path
+from typing import Dict, List, Mapping, Sequence, Union
+
+import pydantic
+
+from primaite import getLogger
+
+_LOGGER = getLogger(__name__)
+import warnings
+from itertools import chain
+
+import yaml
+
+
+class EpisodeScheduler(pydantic.BaseModel, ABC):
+    """
+    Episode schedulers provide functionality to select different scenarios and game setups for each episode.
+
+    This is useful when implementing advanced RL concepts like curriculum learning and domain randomisation.
+    """
+
+    @abstractmethod
+    def __call__(self, episode_num: int) -> Dict:
+        """Return the config that should be used during this episode."""
+
+
+class ConstantEpisodeScheduler(EpisodeScheduler):
+    """
+    The constant episode schedule simply provides the same game setup every time.
+    """
+
+    config: Dict
+
+    def __call__(self, episode_num: int) -> Dict:
+        """Return the same config every time."""
+        return copy.deepcopy(self.config)
+
+
+class EpisodeListScheduler(EpisodeScheduler):
+    """The episode list u"""
+
+    schedule: Mapping[int, List[str]]
+    """Mapping from episode number to list of filenames"""
+    episode_data: Mapping[str, str]
+    """Mapping from filename to yaml string."""
+    base_scenario: str
+    """yaml string containing the base scenario."""
+
+    _exceeded_episode_list: bool = False
+    """
+    Flag that's set to true when attempting to keep generating episodes after schedule runs out.
+
+    When this happens, we loop back to the beginning, but a warning is raised.
+    """
+
+    # TODO: be careful about off-by-one errors with episode number- should it start at 0 or 1?
+    def __call__(self, episode_num: int) -> Dict:
+        if episode_num > len(self.schedule):
+            if not self._exceeded_episode_list:
+                self._exceeded_episode_list = True
+                _LOGGER.warn(
+                    f"Running episode {episode_num} but the schedule only defines"
+                    f"{len(self.schedule)} episodes. Looping back to the beginning"
+                )
+                # not sure if we should be using a traditional warning, or a _LOGGER.warning
+            episode_num = episode_num % len(self.schedule)
+
+        filenames_to_join = self.schedule[episode_num]
+        yaml_data_to_join = [self.episode_data[fn] for fn in filenames_to_join] + [self.base_scenario]
+        joined_yaml = "\n".join(yaml_data_to_join)
+        parsed_cfg = yaml.safe_load(joined_yaml)
+
+        # Unfortunately, using placeholders like this is slightly hacky, so we have to flatten the list of agents
+        flat_agents_list = []
+        for a in parsed_cfg["agents"]:
+            if isinstance(a, Sequence):
+                flat_agents_list.extend(a)
+            else:
+                flat_agents_list.append(a)
+        parsed_cfg["agents"] = flat_agents_list
+
+        return parsed_cfg
+
+
+def build_scheduler(config: Union[str, Path, Dict]) -> EpisodeScheduler:
+    """
+    Convenience method to build an EpisodeScheduler with a dict, file path, or folder path.
+
+    If a path to a folder is provided, it will be treated as a list of game scenarios.
+    Otherwise, if a dict or a single file is provided, it will be treated as a constant game scenario.
+    """
+    # If we get a dict, return a constant episode schedule that repeats that one config forever
+    if isinstance(config, Dict):
+        return ConstantEpisodeScheduler(config=config)
+
+    # Cast string to Path
+    if isinstance(config, str):
+        config = Path(config)
+
+    if not config.exists():
+        raise FileNotFoundError(f"Provided config path {config} could not be found.")
+
+    if config.is_file():
+        with open(config, "r") as f:
+            cfg_data = yaml.safe_load(f)
+        return ConstantEpisodeScheduler(config=cfg_data)
+
+    if not config.is_dir():
+        raise RuntimeError("Something went wrong while building Primaite config.")
+
+    root = config
+    schedule_path = root / "schedule.yaml"
+
+    with open(schedule_path, "r") as f:
+        schedule = yaml.safe_load(f)
+
+    base_scenario_path = root / schedule["base_scenario"]
+    files_to_load = set(chain.from_iterable(schedule["schedule"].values()))
+
+    episode_data = {fp: (root / fp).read_text() for fp in files_to_load}
+
+    return EpisodeListScheduler(
+        schedule=schedule["schedule"], episode_data=episode_data, base_scenario=base_scenario_path.read_text()
+    )