#917 -Finished integrating all agents to either train (policy agents) or evaluate (hard-coded agents). Still some fixing up to do, tidying up, loading etc. also docs. But this is all now working.

src/primaite/main.py

@@ -1,137 +1,15 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
-The main PrimAITE session runner module.
-
-TODO: This will eventually be refactored out into a proper Session class.
-TODO: The passing about of session_path and timestamp_str is temporary and
-    will be cleaned up once we move to a proper Session class.
-"""
+"""The main PrimAITE session runner module."""
 import argparse
-import json
-import time
-from datetime import datetime
 from pathlib import Path
-from typing import Final, Union
-from uuid import uuid4
+from typing import Union
 
-from stable_baselines3 import A2C, PPO
-from stable_baselines3.common.evaluation import evaluate_policy
-from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
-from stable_baselines3.ppo import MlpPolicy as PPOMlp
-
-from primaite import SESSIONS_DIR, getLogger
-from primaite.config.training_config import TrainingConfig
-from primaite.environment.primaite_env import Primaite
+from primaite import getLogger
 from primaite.primaite_session import PrimaiteSession
-from primaite.transactions.transactions_to_file import \
-    write_transaction_to_file
 
 _LOGGER = getLogger(__name__)
 
 
-def run_generic(env: Primaite, config_values: TrainingConfig):
-    """
-    Run against a generic agent.
-
-    :param env: An instance of
-        :class:`~primaite.environment.primaite_env.Primaite`.
-    :param config_values: An instance of
-        :class:`~primaite.config.training_config.TrainingConfig`.
-    """
-    for episode in range(0, config_values.num_episodes):
-        env.reset()
-        for step in range(0, config_values.num_steps):
-            # Send the observation space to the agent to get an action
-            # TEMP - random action for now
-            # action = env.blue_agent_action(obs)
-            action = env.action_space.sample()
-
-            # Run the simulation step on the live environment
-            obs, reward, done, info = env.step(action)
-
-            # Break if done is True
-            if done:
-                break
-
-            # Introduce a delay between steps
-            time.sleep(config_values.time_delay / 1000)
-
-        # Reset the environment at the end of the episode
-
-    env.close()
-
-
-def run_stable_baselines3_ppo(
-    env: Primaite, config_values: TrainingConfig, session_path: Path, timestamp_str: str
-):
-    """
-    Run against a stable_baselines3 PPO agent.
-
-    :param env: An instance of
-        :class:`~primaite.environment.primaite_env.Primaite`.
-    :param config_values: An instance of
-        :class:`~primaite.config.training_config.TrainingConfig`.
-    :param session_path: The directory path the session is writing to.
-    :param timestamp_str: The session timestamp in the format:
-        <yyyy-mm-dd>_<hh-mm-ss>.
-    """
-    if config_values.load_agent:
-        try:
-            agent = PPO.load(
-                config_values.agent_load_file,
-                env,
-                verbose=0,
-                n_steps=config_values.num_steps,
-            )
-        except Exception:
-            print(
-                "ERROR: Could not load agent at location: "
-                + config_values.agent_load_file
-            )
-            _LOGGER.error("Could not load agent")
-            _LOGGER.error("Exception occured", exc_info=True)
-    else:
-        agent = PPO(PPOMlp, env, verbose=0, n_steps=config_values.num_steps)
-
-    if config_values.session_type == "TRAINING":
-        # We're in a training session
-        print("Starting training session...")
-        _LOGGER.debug("Starting training session...")
-        for episode in range(config_values.num_episodes):
-            agent.learn(total_timesteps=config_values.num_steps)
-        _save_agent(agent, session_path, timestamp_str)
-    else:
-        # Default to being in an evaluation session
-        print("Starting evaluation session...")
-        _LOGGER.debug("Starting evaluation session...")
-        evaluate_policy(agent, env, n_eval_episodes=config_values.num_episodes)
-
-    env.close()
-
-
-
-
-def _save_agent(agent: OnPolicyAlgorithm, session_path: Path, timestamp_str: str):
-    """
-    Persist an agent.
-
-    Only works for stable baselines3 agents at present.
-
-    :param session_path: The directory path the session is writing to.
-    :param timestamp_str: The session timestamp in the format:
-        <yyyy-mm-dd>_<hh-mm-ss>.
-    """
-    if not isinstance(agent, OnPolicyAlgorithm):
-        msg = f"Can only save {OnPolicyAlgorithm} agents, got {type(agent)}."
-        _LOGGER.error(msg)
-    else:
-        filepath = session_path / f"agent_saved_{timestamp_str}"
-        agent.save(filepath)
-        _LOGGER.debug(f"Trained agent saved as: {filepath}")
-
-
-
-
 def run(training_config_path: Union[str, Path], lay_down_config_path: Union[str, Path]):
     """Run the PrimAITE Session.