#1632 - Added full benchmarking script that included plots and a LaTeX report. Ran the v2.0.0rc1 benchmark. Tidied a few other things up.

src/primaite/__init__.py

@@ -152,4 +152,4 @@ def getLogger(name: str) -> Logger:  # noqa
 
 
 with open(Path(__file__).parent.resolve() / "VERSION", "r") as file:
-    __version__ = file.readline()
+    __version__ = file.readline().strip()

src/primaite/config/training_config.py

@@ -246,6 +246,7 @@ class TrainingConfig:
         return data
 
     def __str__(self) -> str:
+        obs_str = ",".join([c["name"] for c in self.observation_space["components"]])
         tc = f"{self.agent_framework}, "
         if self.agent_framework is AgentFramework.RLLIB:
             tc += f"{self.deep_learning_framework}, "
@@ -253,7 +254,7 @@ class TrainingConfig:
         if self.agent_identifier is AgentIdentifier.HARDCODED:
             tc += f"{self.hard_coded_agent_view}, "
         tc += f"{self.action_type}, "
-        tc += f"observation_space={self.observation_space}, "
+        tc += f"observation_space={obs_str}, "
         if self.session_type is SessionType.TRAIN:
             tc += f"{self.num_train_episodes} episodes @ "
             tc += f"{self.num_train_steps} steps"

src/primaite/data_viz/session_plots.py

@@ -10,7 +10,7 @@ from plotly.graph_objs import Figure
 from primaite import _PLATFORM_DIRS
 
 
-def _get_plotly_config() -> Dict:
+def get_plotly_config() -> Dict:
     """Get the plotly config from primaite_config.yaml."""
     user_config_path = _PLATFORM_DIRS.user_config_path / "primaite_config.yaml"
     with open(user_config_path, "r") as file:
@@ -41,7 +41,7 @@ def plot_av_reward_per_episode(
         if subtitle:
             title = subtitle
 
-    config = _get_plotly_config()
+    config = get_plotly_config()
     layout = go.Layout(
         autosize=config["size"]["auto_size"],
         width=config["size"]["width"],

src/primaite/nodes/node_state_instruction_red.py

@@ -1,6 +1,5 @@
 # Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
 """Defines node behaviour for Green PoL."""
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Union
 
 from primaite.common.enums import NodePOLType
@@ -9,8 +8,7 @@ if TYPE_CHECKING:
     from primaite.common.enums import FileSystemState, HardwareState, NodePOLInitiator, SoftwareState
 
 
-@dataclass()
-class NodeStateInstructionRed(object):
+class NodeStateInstructionRed:
     """The Node State Instruction class."""
 
     def __init__(

src/primaite/pol/red_agent_pol.py

@@ -250,6 +250,11 @@ def apply_red_agent_node_pol(
             # continue --------------------------
             target_node: NodeUnion = nodes[target_node_id]
 
+            # check if the initiator type is a str, and if so, cast it as
+            # NodePOLInitiator
+            if isinstance(initiator, str):
+                initiator = NodePOLInitiator[initiator]
+
             # Based the action taken on the initiator type
             if initiator == NodePOLInitiator.DIRECT:
                 # No conditions required, just apply the change

src/primaite/primaite_session.py

@@ -2,8 +2,9 @@
 """Main entry point to PrimAITE. Configure training/evaluation experiments and input/output."""
 from __future__ import annotations
 
+import json
 from pathlib import Path
-from typing import Any, Dict, Final, Optional, Union
+from typing import Any, Dict, Final, Optional, Tuple, Union
 
 from primaite import getLogger
 from primaite.agents.agent_abc import AgentSessionABC
@@ -16,6 +17,7 @@ from primaite.common.enums import ActionType, AgentFramework, AgentIdentifier, S
 from primaite.config import lay_down_config, training_config
 from primaite.config.training_config import TrainingConfig
 from primaite.utils.session_metadata_parser import parse_session_metadata
+from primaite.utils.session_output_reader import all_transactions_dict, av_rewards_dict
 
 _LOGGER = getLogger(__name__)
 
@@ -186,3 +188,28 @@ class PrimaiteSession:
     def close(self) -> None:
         """Closes the agent."""
         self._agent_session.close()
+
+    def learn_av_reward_per_episode_dict(self) -> Dict[int, float]:
+        """Get the learn av reward per episode from file."""
+        csv_file = f"average_reward_per_episode_{self.timestamp_str}.csv"
+        return av_rewards_dict(self.learning_path / csv_file)
+
+    def eval_av_reward_per_episode_dict(self) -> Dict[int, float]:
+        """Get the eval av reward per episode from file."""
+        csv_file = f"average_reward_per_episode_{self.timestamp_str}.csv"
+        return av_rewards_dict(self.evaluation_path / csv_file)
+
+    def learn_all_transactions_dict(self) -> Dict[Tuple[int, int], Dict[str, Any]]:
+        """Get the learn all transactions from file."""
+        csv_file = f"all_transactions_{self.timestamp_str}.csv"
+        return all_transactions_dict(self.learning_path / csv_file)
+
+    def eval_all_transactions_dict(self) -> Dict[Tuple[int, int], Dict[str, Any]]:
+        """Get the eval all transactions from file."""
+        csv_file = f"all_transactions_{self.timestamp_str}.csv"
+        return all_transactions_dict(self.evaluation_path / csv_file)
+
+    def metadata_file_as_dict(self) -> Dict[str, Any]:
+        """Read the session_metadata.json file and return as a dict."""
+        with open(self.session_path / "session_metadata.json", "r") as file:
+            return json.load(file)

src/primaite/utils/session_output_reader.py

@@ -18,7 +18,7 @@ def av_rewards_dict(av_rewards_csv_file: Union[str, Path]) -> Dict[int, float]:
     """
     df_dict = pl.read_csv(av_rewards_csv_file).to_dict()
 
-    return {v: df_dict["Average Reward"][i] for i, v in enumerate(df_dict["Episode"])}
+    return {int(v): df_dict["Average Reward"][i] for i, v in enumerate(df_dict["Episode"])}
 
 
 def all_transactions_dict(all_transactions_csv_file: Union[str, Path]) -> Dict[Tuple[int, int], Dict[str, Any]]: