Resolve merge conflicts

src/primaite/session/environment.py

@@ -0,0 +1,162 @@
+from typing import Any, Dict, Final, Optional, SupportsFloat, Tuple
+
+import gymnasium
+from gymnasium.core import ActType, ObsType
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+
+from primaite.game.agent.interface import ProxyAgent
+from primaite.game.game import PrimaiteGame
+
+
+class PrimaiteGymEnv(gymnasium.Env):
+    """
+    Thin wrapper env to provide agents with a gymnasium API.
+
+    This is always a single agent environment since gymnasium is a single agent API. Therefore, we can make some
+    assumptions about the agent list always having a list of length 1.
+    """
+
+    def __init__(self, game: PrimaiteGame):
+        """Initialise the environment."""
+        super().__init__()
+        self.game: "PrimaiteGame" = game
+        self.agent: ProxyAgent = self.game.rl_agents[0]
+
+    def step(self, action: ActType) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict[str, Any]]:
+        """Perform a step in the environment."""
+        # make ProxyAgent store the action chosen my the RL policy
+        self.agent.store_action(action)
+        # apply_agent_actions accesses the action we just stored
+        self.game.apply_agent_actions()
+        self.game.advance_timestep()
+        state = self.game.get_sim_state()
+        self.game.update_agents(state)
+
+        next_obs = self._get_obs()
+        reward = self.agent.reward_function.current_reward
+        terminated = False
+        truncated = self.game.calculate_truncated()
+        info = {}
+
+        return next_obs, reward, terminated, truncated, info
+
+    def reset(self, seed: Optional[int] = None) -> Tuple[ObsType, Dict[str, Any]]:
+        """Reset the environment."""
+        self.game.reset()
+        state = self.game.get_sim_state()
+        self.game.update_agents(state)
+        next_obs = self._get_obs()
+        info = {}
+        return next_obs, info
+
+    @property
+    def action_space(self) -> gymnasium.Space:
+        """Return the action space of the environment."""
+        return self.agent.action_manager.space
+
+    @property
+    def observation_space(self) -> gymnasium.Space:
+        """Return the observation space of the environment."""
+        return gymnasium.spaces.flatten_space(self.agent.observation_manager.space)
+
+    def _get_obs(self) -> ObsType:
+        """Return the current observation."""
+        unflat_space = self.agent.observation_manager.space
+        unflat_obs = self.agent.observation_manager.current_observation
+        return gymnasium.spaces.flatten(unflat_space, unflat_obs)
+
+
+class PrimaiteRayEnv(gymnasium.Env):
+    """Ray wrapper that accepts a single `env_config` parameter in init function for compatibility with Ray."""
+
+    def __init__(self, env_config: Dict[str, PrimaiteGame]) -> None:
+        """Initialise the environment.
+
+        :param env_config: A dictionary containing the environment configuration. It must contain a single key, `game`
+            which is the PrimaiteGame instance.
+        :type env_config: Dict[str, PrimaiteGame]
+        """
+        self.env = PrimaiteGymEnv(game=env_config["game"])
+        self.action_space = self.env.action_space
+        self.observation_space = self.env.observation_space
+
+    def reset(self, *, seed: int = None, options: dict = None) -> Tuple[ObsType, Dict]:
+        """Reset the environment."""
+        return self.env.reset(seed=seed)
+
+    def step(self, action: ActType) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict]:
+        """Perform a step in the environment."""
+        return self.env.step(action)
+
+
+class PrimaiteRayMARLEnv(MultiAgentEnv):
+    """Ray Environment that inherits from MultiAgentEnv to allow training MARL systems."""
+
+    def __init__(self, env_config: Optional[Dict] = None) -> None:
+        """Initialise the environment.
+
+        :param env_config: A dictionary containing the environment configuration. It must contain a single key, `game`
+            which is the PrimaiteGame instance.
+        :type env_config: Dict[str, PrimaiteGame]
+        """
+        self.game: PrimaiteGame = env_config["game"]
+        """Reference to the primaite game"""
+        self.agents: Final[Dict[str, ProxyAgent]] = {agent.agent_name: agent for agent in self.game.rl_agents}
+        """List of all possible agents in the environment. This list should not change!"""
+        self._agent_ids = list(self.agents.keys())
+
+        self.terminateds = set()
+        self.truncateds = set()
+        self.observation_space = gymnasium.spaces.Dict(
+            {name: agent.observation_manager.space for name, agent in self.agents.items()}
+        )
+        self.action_space = gymnasium.spaces.Dict(
+            {name: agent.action_manager.space for name, agent in self.agents.items()}
+        )
+        super().__init__()
+
+    def reset(self, *, seed: int = None, options: dict = None) -> Tuple[ObsType, Dict]:
+        """Reset the environment."""
+        self.game.reset()
+        state = self.game.get_sim_state()
+        self.game.update_agents(state)
+        next_obs = self._get_obs()
+        info = {}
+        return next_obs, info
+
+    def step(
+        self, actions: Dict[str, ActType]
+    ) -> Tuple[Dict[str, ObsType], Dict[str, SupportsFloat], Dict[str, bool], Dict[str, bool], Dict]:
+        """Perform a step in the environment. Adherent to Ray MultiAgentEnv step API.
+
+        :param actions: Dict of actions. The key is agent identifier and the value is a gymnasium action instance.
+        :type actions: Dict[str, ActType]
+        :return: Observations, rewards, terminateds, truncateds, and info. Each one is a dictionary keyed by agent
+            identifier.
+        :rtype: Tuple[Dict[str,ObsType], Dict[str, SupportsFloat], Dict[str,bool], Dict[str,bool], Dict]
+        """
+        # 1. Perform actions
+        for agent_name, action in actions.items():
+            self.agents[agent_name].store_action(action)
+        self.game.apply_agent_actions()
+
+        # 2. Advance timestep
+        self.game.advance_timestep()
+
+        # 3. Get next observations
+        state = self.game.get_sim_state()
+        self.game.update_agents(state)
+        next_obs = self._get_obs()
+
+        # 4. Get rewards
+        rewards = {name: agent.reward_function.current_reward for name, agent in self.agents.items()}
+        terminateds = {name: False for name, _ in self.agents.items()}
+        truncateds = {name: self.game.calculate_truncated() for name, _ in self.agents.items()}
+        infos = {}
+        terminateds["__all__"] = len(self.terminateds) == len(self.agents)
+        truncateds["__all__"] = self.game.calculate_truncated()
+        return next_obs, rewards, terminateds, truncateds, infos
+
+    def _get_obs(self) -> Dict[str, ObsType]:
+        """Return the current observation."""
+        return {name: agent.observation_manager.current_observation for name, agent in self.agents.items()}

src/primaite/session/io.py

@@ -0,0 +1,63 @@
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict
+
+from primaite import PRIMAITE_PATHS
+from primaite.simulator import SIM_OUTPUT
+
+
+class SessionIOSettings(BaseModel):
+    """Schema for session IO settings."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    save_final_model: bool = True
+    """Whether to save the final model right at the end of training."""
+    save_checkpoints: bool = False
+    """Whether to save a checkpoint model every `checkpoint_interval` episodes"""
+    checkpoint_interval: int = 10
+    """How often to save a checkpoint model (if save_checkpoints is True)."""
+    save_logs: bool = True
+    """Whether to save logs"""
+    save_transactions: bool = True
+    """Whether to save transactions, If true, the session path will have a transactions folder."""
+    save_tensorboard_logs: bool = False
+    """Whether to save tensorboard logs. If true, the session path will have a tenorboard_logs folder."""
+
+
+class SessionIO:
+    """
+    Class for managing session IO.
+
+    Currently it's handling path generation, but could expand to handle loading, transaction, tensorboard, and so on.
+    """
+
+    def __init__(self, settings: SessionIOSettings = SessionIOSettings()) -> None:
+        self.settings: SessionIOSettings = settings
+        self.session_path: Path = self.generate_session_path()
+
+        # set global SIM_OUTPUT path
+        SIM_OUTPUT.path = self.session_path / "simulation_output"
+
+        # warning TODO: must be careful not to re-initialise sessionIO because it will create a new path each time it's
+        # possible refactor needed
+
+    def generate_session_path(self, timestamp: Optional[datetime] = None) -> Path:
+        """Create a folder for the session and return the path to it."""
+        if timestamp is None:
+            timestamp = datetime.now()
+        date_str = timestamp.strftime("%Y-%m-%d")
+        time_str = timestamp.strftime("%H-%M-%S")
+        session_path = PRIMAITE_PATHS.user_sessions_path / date_str / time_str
+        session_path.mkdir(exist_ok=True, parents=True)
+        return session_path
+
+    def generate_model_save_path(self, agent_name: str) -> Path:
+        """Return the path where the final model will be saved (excluding filename extension)."""
+        return self.session_path / "checkpoints" / f"{agent_name}_final"
+
+    def generate_checkpoint_save_path(self, agent_name: str, episode: int) -> Path:
+        """Return the path where the checkpoint model will be saved (excluding filename extension)."""
+        return self.session_path / "checkpoints" / f"{agent_name}_checkpoint_{episode}.pt"

src/primaite/session/policy/__init__.py

@@ -0,0 +1,4 @@
+from primaite.session.policy.rllib import RaySingleAgentPolicy
+from primaite.session.policy.sb3 import SB3Policy
+
+__all__ = ["SB3Policy", "RaySingleAgentPolicy"]

src/primaite/session/policy/policy.py

@@ -0,0 +1,82 @@
+"""Base class and common logic for RL policies."""
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Dict, Type, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from primaite.session.session import PrimaiteSession, TrainingOptions
+
+
+class PolicyABC(ABC):
+    """Base class for reinforcement learning agents."""
+
+    _registry: Dict[str, Type["PolicyABC"]] = {}
+    """
+    Registry of policy types, keyed by name.
+
+    Automatically populated when PolicyABC subclasses are defined. Used for defining from_config.
+    """
+
+    def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None:
+        """
+        Register a policy subclass.
+
+        :param name: Identifier used by from_config to create an instance of the policy.
+        :type name: str
+        :raises ValueError: When attempting to create a policy with a duplicate name.
+        """
+        super().__init_subclass__(**kwargs)
+        if identifier in cls._registry:
+            raise ValueError(f"Duplicate policy name {identifier}")
+        cls._registry[identifier] = cls
+        return
+
+    @abstractmethod
+    def __init__(self, session: "PrimaiteSession") -> None:
+        """
+        Initialize a reinforcement learning policy.
+
+        :param session: The session context.
+        :type session: PrimaiteSession
+        :param agents: The agents to train.
+        :type agents: List[RLAgent]
+        """
+        self.session: "PrimaiteSession" = session
+        """Reference to the session."""
+
+    @abstractmethod
+    def learn(self, n_episodes: int, timesteps_per_episode: int) -> None:
+        """Train the agent."""
+        pass
+
+    @abstractmethod
+    def eval(self, n_episodes: int, timesteps_per_episode: int, deterministic: bool) -> None:
+        """Evaluate the agent."""
+        pass
+
+    @abstractmethod
+    def save(self, save_path: Path) -> None:
+        """Save the agent."""
+        pass
+
+    @abstractmethod
+    def load(self) -> None:
+        """Load agent from a file."""
+        pass
+
+    def close(self) -> None:
+        """Close the agent."""
+        pass
+
+    @classmethod
+    def from_config(cls, config: "TrainingOptions", session: "PrimaiteSession") -> "PolicyABC":
+        """
+        Create an RL policy from a config by calling the relevant subclass's from_config method.
+
+        Subclasses should not call super().from_config(), they should just handle creation form config.
+        """
+        # Assume that basically the contents of training_config are passed into here.
+        # I should really define a config schema class using pydantic.
+
+        PolicyType = cls._registry[config.rl_framework]
+        return PolicyType.from_config(config=config, session=session)

src/primaite/session/policy/rllib.py

@@ -0,0 +1,106 @@
+from pathlib import Path
+from typing import Literal, Optional, TYPE_CHECKING
+
+from primaite.session.environment import PrimaiteRayEnv, PrimaiteRayMARLEnv
+from primaite.session.policy.policy import PolicyABC
+
+if TYPE_CHECKING:
+    from primaite.session.session import PrimaiteSession, TrainingOptions
+
+import ray
+from ray import air, tune
+from ray.rllib.algorithms import ppo
+from ray.rllib.algorithms.ppo import PPOConfig
+
+
+class RaySingleAgentPolicy(PolicyABC, identifier="RLLIB_single_agent"):
+    """Single agent RL policy using Ray RLLib."""
+
+    def __init__(self, session: "PrimaiteSession", algorithm: Literal["PPO", "A2C"], seed: Optional[int] = None):
+        super().__init__(session=session)
+
+        config = {
+            "env": PrimaiteRayEnv,
+            "env_config": {"game": session.game},
+            "disable_env_checking": True,
+            "num_rollout_workers": 0,
+        }
+
+        ray.shutdown()
+        ray.init()
+
+        self._algo = ppo.PPO(config=config)
+
+    def learn(self, n_episodes: int, timesteps_per_episode: int) -> None:
+        """Train the agent."""
+        for ep in range(n_episodes):
+            self._algo.train()
+
+    def eval(self, n_episodes: int, deterministic: bool) -> None:
+        """Evaluate the agent."""
+        for ep in range(n_episodes):
+            obs, info = self.session.env.reset()
+            for step in range(self.session.game.options.max_episode_length):
+                action = self._algo.compute_single_action(observation=obs, explore=False)
+                obs, rew, term, trunc, info = self.session.env.step(action)
+
+    def save(self, save_path: Path) -> None:
+        """Save the policy to a file."""
+        self._algo.save(save_path)
+
+    def load(self, model_path: Path) -> None:
+        """Load policy parameters from a file."""
+        raise NotImplementedError
+
+    @classmethod
+    def from_config(cls, config: "TrainingOptions", session: "PrimaiteSession") -> "RaySingleAgentPolicy":
+        """Create a policy from a config."""
+        return cls(session=session, algorithm=config.rl_algorithm, seed=config.seed)
+
+
+class RayMultiAgentPolicy(PolicyABC, identifier="RLLIB_multi_agent"):
+    """Mutli agent RL policy using Ray RLLib."""
+
+    def __init__(self, session: "PrimaiteSession", algorithm: Literal["PPO"], seed: Optional[int] = None):
+        """Initialise multi agent policy wrapper."""
+        super().__init__(session=session)
+
+        self.config = (
+            PPOConfig()
+            .environment(env=PrimaiteRayMARLEnv, env_config={"game": session.game})
+            .rollouts(num_rollout_workers=0)
+            .multi_agent(
+                policies={agent.agent_name for agent in session.game.rl_agents},
+                policy_mapping_fn=lambda agent_id, episode, worker, **kw: agent_id,
+            )
+            .training(train_batch_size=128)
+        )
+
+    def learn(self, n_episodes: int, timesteps_per_episode: int) -> None:
+        """Train the agent."""
+        checkpoint_freq = self.session.io_manager.settings.checkpoint_interval
+        tune.Tuner(
+            "PPO",
+            run_config=air.RunConfig(
+                stop={"training_iteration": n_episodes * timesteps_per_episode},
+                checkpoint_config=air.CheckpointConfig(checkpoint_frequency=checkpoint_freq),
+            ),
+            param_space=self.config,
+        ).fit()
+
+    def load(self, model_path: Path) -> None:
+        """Load policy parameters from a file."""
+        return NotImplemented
+
+    def eval(self, n_episodes: int, deterministic: bool) -> None:
+        """Evaluate trained policy."""
+        return NotImplemented
+
+    def save(self, save_path: Path) -> None:
+        """Save policy parameters to a file."""
+        return NotImplemented
+
+    @classmethod
+    def from_config(cls, config: "TrainingOptions", session: "PrimaiteSession") -> "RayMultiAgentPolicy":
+        """Create policy from config."""
+        return cls(session=session, algorithm=config.rl_algorithm, seed=config.seed)

src/primaite/session/policy/sb3.py

@@ -0,0 +1,80 @@
+"""Stable baselines 3 policy."""
+from pathlib import Path
+from typing import Literal, Optional, Type, TYPE_CHECKING, Union
+
+from stable_baselines3 import A2C, PPO
+from stable_baselines3.a2c import MlpPolicy as A2C_MLP
+from stable_baselines3.common.callbacks import CheckpointCallback
+from stable_baselines3.common.evaluation import evaluate_policy
+from stable_baselines3.ppo import MlpPolicy as PPO_MLP
+
+from primaite.session.policy.policy import PolicyABC
+
+if TYPE_CHECKING:
+    from primaite.session.session import PrimaiteSession, TrainingOptions
+
+
+class SB3Policy(PolicyABC, identifier="SB3"):
+    """Single agent RL policy using stable baselines 3."""
+
+    def __init__(self, session: "PrimaiteSession", algorithm: Literal["PPO", "A2C"], seed: Optional[int] = None):
+        """Initialize a stable baselines 3 policy."""
+        super().__init__(session=session)
+
+        self._agent_class: Type[Union[PPO, A2C]]
+        if algorithm == "PPO":
+            self._agent_class = PPO
+            policy = PPO_MLP
+        elif algorithm == "A2C":
+            self._agent_class = A2C
+            policy = A2C_MLP
+        else:
+            raise ValueError(f"Unknown algorithm `{algorithm}` for stable_baselines3 policy")
+        self._agent = self._agent_class(
+            policy=policy,
+            env=self.session.env,
+            n_steps=128,  # this is not the number of steps in an episode, but the number of steps in a batch
+            seed=seed,
+        )
+
+    def learn(self, n_episodes: int, timesteps_per_episode: int) -> None:
+        """Train the agent."""
+        if self.session.io_manager.settings.save_checkpoints:
+            checkpoint_callback = CheckpointCallback(
+                save_freq=timesteps_per_episode * self.session.io_manager.settings.checkpoint_interval,
+                save_path=self.session.io_manager.generate_model_save_path("sb3"),
+                name_prefix="sb3_model",
+            )
+        else:
+            checkpoint_callback = None
+        self._agent.learn(total_timesteps=n_episodes * timesteps_per_episode, callback=checkpoint_callback)
+
+    def eval(self, n_episodes: int, deterministic: bool) -> None:
+        """Evaluate the agent."""
+        reward_data = evaluate_policy(
+            self._agent,
+            self.session.env,
+            n_eval_episodes=n_episodes,
+            deterministic=deterministic,
+            return_episode_rewards=True,
+        )
+        print(reward_data)
+
+    def save(self, save_path: Path) -> None:
+        """
+        Save the current policy parameters.
+
+        Warning: The recommended way to save model checkpoints is to use a callback within the `learn()` method. Please
+        refer to https://stable-baselines3.readthedocs.io/en/master/guide/callbacks.html for more information.
+        Therefore, this method is only used to save the final model.
+        """
+        self._agent.save(save_path)
+
+    def load(self, model_path: Path) -> None:
+        """Load agent from a checkpoint."""
+        self._agent = self._agent_class.load(model_path, env=self.session.env)
+
+    @classmethod
+    def from_config(cls, config: "TrainingOptions", session: "PrimaiteSession") -> "SB3Policy":
+        """Create an agent from config file."""
+        return cls(session=session, algorithm=config.rl_algorithm, seed=config.seed)

src/primaite/session/session.py

@@ -0,0 +1,113 @@
+from enum import Enum
+from pathlib import Path
+from typing import Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, ConfigDict
+
+from primaite.game.game import PrimaiteGame
+from primaite.session.environment import PrimaiteGymEnv, PrimaiteRayEnv, PrimaiteRayMARLEnv
+from primaite.session.io import SessionIO, SessionIOSettings
+
+# from primaite.game.game import PrimaiteGame
+from primaite.session.policy.policy import PolicyABC
+
+
+class TrainingOptions(BaseModel):
+    """Options for training the RL agent."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    rl_framework: Literal["SB3", "RLLIB_single_agent", "RLLIB_multi_agent"]
+    rl_algorithm: Literal["PPO", "A2C"]
+    n_learn_episodes: int
+    n_eval_episodes: Optional[int] = None
+    max_steps_per_episode: int
+    # checkpoint_freq: Optional[int] = None
+    deterministic_eval: bool
+    seed: Optional[int]
+    n_agents: int
+    agent_references: List[str]
+
+
+class SessionMode(Enum):
+    """Helper to keep track of the current session mode."""
+
+    TRAIN = "train"
+    EVAL = "eval"
+    MANUAL = "manual"
+
+
+class PrimaiteSession:
+    """The main entrypoint for PrimAITE sessions, this manages a simulation, policy training, and environments."""
+
+    def __init__(self, game: PrimaiteGame):
+        """Initialise PrimaiteSession object."""
+        self.training_options: TrainingOptions
+        """Options specific to agent training."""
+
+        self.mode: SessionMode = SessionMode.MANUAL
+        """Current session mode."""
+
+        self.env: Union[PrimaiteGymEnv, PrimaiteRayEnv, PrimaiteRayMARLEnv]
+        """The environment that the RL algorithm can consume."""
+
+        self.policy: PolicyABC
+        """The reinforcement learning policy."""
+
+        self.io_manager = SessionIO()
+        """IO manager for the session."""
+
+        self.game: PrimaiteGame = game
+        """Primaite Game object for managing main simulation loop and agents."""
+
+    def start_session(self) -> None:
+        """Commence the training/eval session."""
+        self.mode = SessionMode.TRAIN
+        n_learn_episodes = self.training_options.n_learn_episodes
+        n_eval_episodes = self.training_options.n_eval_episodes
+        max_steps_per_episode = self.training_options.max_steps_per_episode
+
+        deterministic_eval = self.training_options.deterministic_eval
+        self.policy.learn(
+            n_episodes=n_learn_episodes,
+            timesteps_per_episode=max_steps_per_episode,
+        )
+        self.save_models()
+
+        self.mode = SessionMode.EVAL
+        if n_eval_episodes > 0:
+            self.policy.eval(n_episodes=n_eval_episodes, deterministic=deterministic_eval)
+
+        self.mode = SessionMode.MANUAL
+
+    def save_models(self) -> None:
+        """Save the RL models."""
+        save_path = self.io_manager.generate_model_save_path("temp_model_name")
+        self.policy.save(save_path)
+
+    @classmethod
+    def from_config(cls, cfg: Dict, agent_load_path: Optional[str] = None) -> "PrimaiteSession":
+        """Create a PrimaiteSession object from a config dictionary."""
+        game = PrimaiteGame.from_config(cfg)
+
+        sess = cls(game=game)
+
+        sess.training_options = TrainingOptions(**cfg["training_config"])
+
+        # READ IO SETTINGS (this sets the global session path as well) # TODO: GLOBAL SIDE EFFECTS...
+        io_settings = cfg.get("io_settings", {})
+        sess.io_manager.settings = SessionIOSettings(**io_settings)
+
+        # CREATE ENVIRONMENT
+        if sess.training_options.rl_framework == "RLLIB_single_agent":
+            sess.env = PrimaiteRayEnv(env_config={"game": game})
+        elif sess.training_options.rl_framework == "RLLIB_multi_agent":
+            sess.env = PrimaiteRayMARLEnv(env_config={"game": game})
+        elif sess.training_options.rl_framework == "SB3":
+            sess.env = PrimaiteGymEnv(game=game)
+
+        sess.policy = PolicyABC.from_config(sess.training_options, session=sess)
+        if agent_load_path:
+            sess.policy.load(Path(agent_load_path))
+
+        return sess