#1595:

- Added ability to load sessions via PrimaiteSession
- PrimaiteSession loading test
- Added a NotImplemented RLlib loading for now
- Added the ability to load sessions for hardcoded agents
- Moved Session metadata parsing to utils

src/primaite/agents/agent_abc.py

@@ -7,14 +7,13 @@ from pathlib import Path
 from typing import Dict, Optional, Union
 from uuid import uuid4
 
-import yaml
-
 import primaite
 from primaite import getLogger, SESSIONS_DIR
 from primaite.config import lay_down_config, training_config
 from primaite.config.training_config import TrainingConfig
 from primaite.data_viz.session_plots import plot_av_reward_per_episode
 from primaite.environment.primaite_env import Primaite
+from primaite.utils.session_metadata_parser import parse_session_metadata
 
 _LOGGER = getLogger(__name__)
 
@@ -253,47 +252,21 @@ class AgentSessionABC(ABC):
 
     def load(self, path: Union[str, Path]):
         """Load an agent from file."""
-        if not isinstance(path, Path):
-            path = Path(path)
+        md_dict, training_config_path, laydown_config_path = parse_session_metadata(path)
 
-        if path.exists():
-            # Unpack the session_metadata.json file
-            md_file = path / "session_metadata.json"
-            with open(md_file, "r") as file:
-                md_dict = json.load(file)
+        # set training config path
+        self._training_config_path: Union[Path, str] = training_config_path
+        self._training_config: TrainingConfig = training_config.load(self._training_config_path)
+        self._lay_down_config_path: Union[Path, str] = laydown_config_path
+        self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path)
+        self.sb3_output_verbose_level = self._training_config.sb3_output_verbose_level
 
-            # Create a temp directory and dump the training and lay down
-            # configs into it
-            temp_dir = path / ".temp"
-            temp_dir.mkdir(exist_ok=True)
+        # set random UUID for session
+        self._uuid = md_dict["uuid"]
 
-            temp_tc = temp_dir / "tc.yaml"
-            with open(temp_tc, "w") as file:
-                yaml.dump(md_dict["env"]["training_config"], file)
-
-            temp_ldc = temp_dir / "ldc.yaml"
-            with open(temp_ldc, "w") as file:
-                yaml.dump(md_dict["env"]["lay_down_config"], file)
-
-            # set training config path
-            self._training_config_path: Union[Path, str] = temp_tc
-            self._training_config: TrainingConfig = training_config.load(self._training_config_path)
-            self._lay_down_config_path: Union[Path, str] = temp_ldc
-            self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path)
-            self.sb3_output_verbose_level = self._training_config.sb3_output_verbose_level
-
-            # set random UUID for session
-            self._uuid = md_dict["uuid"]
-
-            # set the session path
-            self.session_path = path
-            "The Session path"
-
-        else:
-            # Session path does not exist
-            msg = f"Failed to load PrimAITE Session, path does not exist: {path}"
-            _LOGGER.error(msg)
-            raise FileNotFoundError(msg)
+        # set the session path
+        self.session_path = path
+        "The Session path"
 
     @property
     def _saved_agent_path(self) -> Path:

src/primaite/agents/hardcoded_abc.py

@@ -1,5 +1,7 @@
 import time
 from abc import abstractmethod
+from pathlib import Path
+from typing import Optional, Union
 
 from primaite import getLogger
 from primaite.agents.agent_abc import AgentSessionABC
@@ -16,7 +18,12 @@ class HardCodedAgentSessionABC(AgentSessionABC):
     implemented.
     """
 
-    def __init__(self, training_config_path, lay_down_config_path):
+    def __init__(
+        self,
+        training_config_path: Optional[Union[str, Path]] = "",
+        lay_down_config_path: Optional[Union[str, Path]] = "",
+        session_path: Optional[Union[str, Path]] = None,
+    ):
         """
         Initialise a hardcoded agent session.
 
@@ -26,7 +33,7 @@ class HardCodedAgentSessionABC(AgentSessionABC):
         :param lay_down_config_path: YAML file containing configurable items for generating network laydown.
         :type lay_down_config_path: Union[path, str]
         """
-        super().__init__(training_config_path, lay_down_config_path)
+        super().__init__(training_config_path, lay_down_config_path, session_path)
         self._setup()
 
     def _setup(self):

src/primaite/agents/rllib.py

@@ -4,7 +4,7 @@ import json
 import shutil
 from datetime import datetime
 from pathlib import Path
-from typing import Union
+from typing import Optional, Union
 from uuid import uuid4
 
 from ray.rllib.algorithms import Algorithm
@@ -43,7 +43,12 @@ def _custom_log_creator(session_path: Path):
 class RLlibAgent(AgentSessionABC):
     """An AgentSession class that implements a Ray RLlib agent."""
 
-    def __init__(self, training_config_path, lay_down_config_path):
+    def __init__(
+        self,
+        training_config_path: Optional[Union[str, Path]] = "",
+        lay_down_config_path: Optional[Union[str, Path]] = "",
+        session_path: Optional[Union[str, Path]] = None,
+    ):
         """
         Initialise the RLLib Agent training session.
 
@@ -56,6 +61,13 @@ class RLlibAgent(AgentSessionABC):
         :raises ValueError: If the training config contains an unexpected value for agent_identifies (should be `PPO`
             or `A2C`)
         """
+        # TODO: implement RLlib agent loading
+        if session_path is not None:
+            msg = "RLlib agent loading has not been implemented yet"
+            _LOGGER.error(msg)
+            print(msg)
+            raise NotImplementedError
+
         super().__init__(training_config_path, lay_down_config_path)
         if not self._training_config.agent_framework == AgentFramework.RLLIB:
             msg = f"Expected RLLIB agent_framework, " f"got {self._training_config.agent_framework}"