#2476 Add test for episode scheduler

src/primaite/session/environment.py

@@ -1,4 +1,3 @@
-import copy
 import json
 from os import PathLike
 from typing import Any, Dict, Optional, SupportsFloat, Tuple, Union
@@ -25,10 +24,10 @@ class PrimaiteGymEnv(gymnasium.Env):
     assumptions about the agent list always having a list of length 1.
     """
 
-    def __init__(self, game_config: Union[Dict, str, PathLike]):
+    def __init__(self, env_config: Union[Dict, str, PathLike]):
         """Initialise the environment."""
         super().__init__()
-        self.episode_scheduler: EpisodeScheduler = build_scheduler(game_config)
+        self.episode_scheduler: EpisodeScheduler = build_scheduler(env_config)
         """Object that returns a config corresponding to the current episode."""
         self.io = PrimaiteIO.from_config(self.episode_scheduler(0).get("io_settings", {}))
         """Handles IO for the environment. This produces sys logs, agent logs, etc."""
@@ -140,8 +139,8 @@ class PrimaiteRayEnv(gymnasium.Env):
         :param env_config: A dictionary containing the environment configuration.
         :type env_config: Dict
         """
-        self.env = PrimaiteGymEnv(game_config=env_config)
-        self.env.episode_counter -= 1
+        self.env = PrimaiteGymEnv(env_config=env_config)
+        # self.env.episode_counter -= 1
         self.action_space = self.env.action_space
         self.observation_space = self.env.observation_space
 
@@ -157,6 +156,11 @@ class PrimaiteRayEnv(gymnasium.Env):
         """Close the simulation."""
         self.env.close()
 
+    @property
+    def game(self) -> PrimaiteGame:
+        """Pass through game from env."""
+        return self.env.game
+
 
 class PrimaiteRayMARLEnv(MultiAgentEnv):
     """Ray Environment that inherits from MultiAgentEnv to allow training MARL systems."""
@@ -168,16 +172,16 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
             which is the PrimaiteGame instance.
         :type env_config: Dict
         """
-        self.game_config: Dict = env_config
-        """PrimaiteGame definition. This can be changed between episodes to enable curriculum learning."""
-        self.io = PrimaiteIO.from_config(env_config.get("io_settings"))
+        self.episode_counter: int = 0
+        """Current episode number."""
+        self.episode_scheduler: EpisodeScheduler = build_scheduler(env_config)
+        """Object that returns a config corresponding to the current episode."""
+        self.io = PrimaiteIO.from_config(self.episode_scheduler(0).get("io_settings", {}))
         """Handles IO for the environment. This produces sys logs, agent logs, etc."""
-        self.game: PrimaiteGame = PrimaiteGame.from_config(copy.deepcopy(self.game_config))
+        self.game: PrimaiteGame = PrimaiteGame.from_config(self.episode_scheduler(self.episode_counter))
         """Reference to the primaite game"""
         self._agent_ids = list(self.game.rl_agents.keys())
         """Agent ids. This is a list of strings of agent names."""
-        self.episode_counter: int = 0
-        """Current episode number."""
 
         self.terminateds = set()
         self.truncateds = set()
@@ -203,9 +207,9 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
         if self.io.settings.save_agent_actions:
             all_agent_actions = {name: agent.action_history for name, agent in self.game.agents.items()}
             self.io.write_agent_actions(agent_actions=all_agent_actions, episode=self.episode_counter)
-        self.game: PrimaiteGame = PrimaiteGame.from_config(cfg=copy.deepcopy(self.game_config))
-        self.game.setup_for_episode(episode=self.episode_counter)
         self.episode_counter += 1
+        self.game: PrimaiteGame = PrimaiteGame.from_config(self.episode_scheduler(self.episode_counter))
+        self.game.setup_for_episode(episode=self.episode_counter)
         state = self.game.get_sim_state()
         self.game.update_agents(state)
         next_obs = self._get_obs()

src/primaite/session/episode_schedule.py

@@ -22,6 +22,7 @@ class EpisodeScheduler(pydantic.BaseModel, ABC):
     @abstractmethod
     def __call__(self, episode_num: int) -> Dict:
         """Return the config that should be used during this episode."""
+        ...
 
 
 class ConstantEpisodeScheduler(EpisodeScheduler):