#901 - merged dev into branch

tests/__init__.py

@@ -4,3 +4,6 @@ from typing import Final
 
 TEST_CONFIG_ROOT: Final[Path] = Path(__file__).parent / "config"
 "The tests config root directory."
+
+TEST_ASSETS_ROOT: Final[Path] = Path(__file__).parent / "assets"
+"The tests assets root directory."

tests/config/one_node_states_on_off_main_config.yaml

@@ -7,6 +7,14 @@
 # "CUSTOM" (Custom Agent)
 agent_framework: CUSTOM
 
+# Sets which deep learning framework will be used (by RLlib ONLY).
+# Default is TF (Tensorflow).
+# Options are:
+# "TF" (Tensorflow)
+# TF2 (Tensorflow 2.X)
+# TORCH (PyTorch)
+deep_learning_framework: TF2
+
 # Sets which Agent class will be used.
 # Options are:
 # "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework)
@@ -17,27 +25,66 @@ agent_framework: CUSTOM
 # "DUMMY" (primaite.agents.simple.DummyAgent)
 agent_identifier: DUMMY
 
+# Sets whether Red Agent POL and IER is randomised.
+# Options are:
+# True
+# False
+random_red_agent: False
+
+# The (integer) seed to be used in random number generation
+# Default is None (null)
+seed: null
+
+# Set whether the agent will be deterministic instead of stochastic
+# Options are:
+# True
+# False
+deterministic: False
+
+# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC.
+# Options are:
+# "BASIC" (The current observation space only)
+# "FULL" (Full environment view with actions taken and reward feedback)
+hard_coded_agent_view: FULL
+
 # Sets How the Action Space is defined:
 # "NODE"
 # "ACL"
 # "ANY" node and acl actions
 action_type: NODE
+# observation space
+observation_space:
+  # flatten: true
+  components:
+    - name: NODE_LINK_TABLE
+    # - name: NODE_STATUSES
+    # - name: LINK_TRAFFIC_LEVELS
 
 
+# Number of episodes for training to run per session
+num_train_episodes: 10
+
+# Number of time_steps for training per episode
+num_train_steps: 256
+
 # Number of episodes for evaluation to run per session
 num_eval_episodes: 1
 
 # Number of time_steps for evaluation per episode
 num_eval_steps: 15
-# Time delay between steps (for generic agents)
-time_delay: 1
 
-# Type of session to be run (TRAINING or EVALUATION)
+# Sets how often the agent will save a checkpoint (every n time episodes).
+# Set to 0 if no checkpoints are required. Default is 10
+checkpoint_every_n_episodes: 10
+
+# Time delay (milliseconds) between steps for CUSTOM agents.
+time_delay: 5
+
+# Type of session to be run. Options are:
+# "TRAIN" (Trains an agent)
+# "EVAL" (Evaluates an agent)
+# "TRAIN_EVAL" (Trains then evaluates an agent)
 session_type: EVAL
-# Determine whether to load an agent from file
-load_agent: False
-# File path and file name of agent if you're loading one in
-agent_load_file: C:\[Path]\[agent_saved_filename.zip]
 
 # Environment config values
 # The high value for the observation space
@@ -45,6 +92,13 @@ observation_space_high_value: 1000000000
 
 implicit_acl_rule: DENY
 max_number_acl_rules: 10
+# The Stable Baselines3 learn/eval output verbosity level:
+# Options are:
+# "NONE" (No Output)
+# "INFO" (Info Messages (such as devices and wrappers used))
+# "DEBUG" (All Messages)
+sb3_output_verbose_level: NONE
+
 # Reward values
 # Generic
 all_ok: 0

tests/config/test_random_red_main_config.yaml

@@ -5,7 +5,15 @@
 # "SB3" (Stable Baselines3)
 # "RLLIB" (Ray RLlib)
 # "CUSTOM" (Custom Agent)
-agent_framework: CUSTOM
+agent_framework: SB3
+
+# Sets which deep learning framework will be used (by RLlib ONLY).
+# Default is TF (Tensorflow).
+# Options are:
+# "TF" (Tensorflow)
+# TF2 (Tensorflow 2.X)
+# TORCH (PyTorch)
+deep_learning_framework: TF2
 
 # Sets which Agent class will be used.
 # Options are:
@@ -15,7 +23,7 @@ agent_framework: CUSTOM
 # "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
 # "RANDOM" (primaite.agents.simple.RandomAgent)
 # "DUMMY" (primaite.agents.simple.DummyAgent)
-agent_identifier: DUMMY
+agent_identifier: PPO
 
 # Sets whether Red Agent POL and IER is randomised.
 # Options are:
@@ -23,92 +31,128 @@ agent_identifier: DUMMY
 # False
 random_red_agent: True
 
+# The (integer) seed to be used in random number generation
+# Default is None (null)
+seed: null
+
+# Set whether the agent will be deterministic instead of stochastic
+# Options are:
+# True
+# False
+deterministic: False
+
+# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC.
+# Options are:
+# "BASIC" (The current observation space only)
+# "FULL" (Full environment view with actions taken and reward feedback)
+hard_coded_agent_view: FULL
+
 # Sets How the Action Space is defined:
 # "NODE"
 # "ACL"
 # "ANY" node and acl actions
 action_type: NODE
+# observation space
+observation_space:
+  # flatten: true
+  components:
+    - name: NODE_LINK_TABLE
+    # - name: NODE_STATUSES
+    # - name: LINK_TRAFFIC_LEVELS
+
+
 # Number of episodes for training to run per session
-num_train_episodes: 2
+num_train_episodes: 10
 
 # Number of time_steps for training per episode
-num_train_steps: 15
+num_train_steps: 256
 
 # Number of episodes for evaluation to run per session
-num_eval_episodes: 2
+num_eval_episodes: 1
 
 # Number of time_steps for evaluation per episode
-num_eval_steps: 15
-# Time delay between steps (for generic agents)
-time_delay: 1
+num_eval_steps: 256
 
-# Type of session to be run (TRAINING or EVALUATION)
-session_type: EVAL
-# Determine whether to load an agent from file
-load_agent: False
-# File path and file name of agent if you're loading one in
-agent_load_file: C:\[Path]\[agent_saved_filename.zip]
+# Sets how often the agent will save a checkpoint (every n time episodes).
+# Set to 0 if no checkpoints are required. Default is 10
+checkpoint_every_n_episodes: 10
+
+# Time delay (milliseconds) between steps for CUSTOM agents.
+time_delay: 5
+
+# Type of session to be run. Options are:
+# "TRAIN" (Trains an agent)
+# "EVAL" (Evaluates an agent)
+# "TRAIN_EVAL" (Trains then evaluates an agent)
+session_type: TRAIN_EVAL
 
 # Environment config values
 # The high value for the observation space
 observation_space_high_value: 1000000000
 
+# The Stable Baselines3 learn/eval output verbosity level:
+# Options are:
+# "NONE" (No Output)
+# "INFO" (Info Messages (such as devices and wrappers used))
+# "DEBUG" (All Messages)
+sb3_output_verbose_level: NONE
+
 # Reward values
 # Generic
 all_ok: 0
 # Node Hardware State
-off_should_be_on: -10
-off_should_be_resetting: -5
-on_should_be_off: -2
-on_should_be_resetting: -5
-resetting_should_be_on: -5
-resetting_should_be_off: -2
-resetting: -3
+off_should_be_on: -0.001
+off_should_be_resetting: -0.0005
+on_should_be_off: -0.0002
+on_should_be_resetting: -0.0005
+resetting_should_be_on: -0.0005
+resetting_should_be_off: -0.0002
+resetting: -0.0003
 # Node Software or Service State
-good_should_be_patching: 2
-good_should_be_compromised: 5
-good_should_be_overwhelmed: 5
-patching_should_be_good: -5
-patching_should_be_compromised: 2
-patching_should_be_overwhelmed: 2
-patching: -3
-compromised_should_be_good: -20
-compromised_should_be_patching: -20
-compromised_should_be_overwhelmed: -20
-compromised: -20
-overwhelmed_should_be_good: -20
-overwhelmed_should_be_patching: -20
-overwhelmed_should_be_compromised: -20
-overwhelmed: -20
+good_should_be_patching: 0.0002
+good_should_be_compromised: 0.0005
+good_should_be_overwhelmed: 0.0005
+patching_should_be_good: -0.0005
+patching_should_be_compromised: 0.0002
+patching_should_be_overwhelmed: 0.0002
+patching: -0.0003
+compromised_should_be_good: -0.002
+compromised_should_be_patching: -0.002
+compromised_should_be_overwhelmed: -0.002
+compromised: -0.002
+overwhelmed_should_be_good: -0.002
+overwhelmed_should_be_patching: -0.002
+overwhelmed_should_be_compromised: -0.002
+overwhelmed: -0.002
 # Node File System State
-good_should_be_repairing: 2
-good_should_be_restoring: 2
-good_should_be_corrupt: 5
-good_should_be_destroyed: 10
-repairing_should_be_good: -5
-repairing_should_be_restoring: 2
-repairing_should_be_corrupt: 2
-repairing_should_be_destroyed: 0
-repairing: -3
-restoring_should_be_good: -10
-restoring_should_be_repairing: -2
-restoring_should_be_corrupt: 1
-restoring_should_be_destroyed: 2
-restoring: -6
-corrupt_should_be_good: -10
-corrupt_should_be_repairing: -10
-corrupt_should_be_restoring: -10
-corrupt_should_be_destroyed: 2
-corrupt: -10
-destroyed_should_be_good: -20
-destroyed_should_be_repairing: -20
-destroyed_should_be_restoring: -20
-destroyed_should_be_corrupt: -20
-destroyed: -20
-scanning: -2
+good_should_be_repairing: 0.0002
+good_should_be_restoring: 0.0002
+good_should_be_corrupt: 0.0005
+good_should_be_destroyed: 0.001
+repairing_should_be_good: -0.0005
+repairing_should_be_restoring: 0.0002
+repairing_should_be_corrupt: 0.0002
+repairing_should_be_destroyed: 0.0000
+repairing: -0.0003
+restoring_should_be_good: -0.001
+restoring_should_be_repairing: -0.0002
+restoring_should_be_corrupt: 0.0001
+restoring_should_be_destroyed: 0.0002
+restoring: -0.0006
+corrupt_should_be_good: -0.001
+corrupt_should_be_repairing: -0.001
+corrupt_should_be_restoring: -0.001
+corrupt_should_be_destroyed: 0.0002
+corrupt: -0.001
+destroyed_should_be_good: -0.002
+destroyed_should_be_repairing: -0.002
+destroyed_should_be_restoring: -0.002
+destroyed_should_be_corrupt: -0.002
+destroyed: -0.002
+scanning: -0.0002
 # IER status
-red_ier_running: -5
-green_ier_blocked: -10
+red_ier_running: -0.0005
+green_ier_blocked: -0.001
 
 # Patching / Reset durations
 os_patching_duration: 5            # The time taken to patch the OS

tests/config/training_config_main_rllib.yaml

@@ -0,0 +1,163 @@
+# Training Config File
+
+# Sets which agent algorithm framework will be used.
+# Options are:
+# "SB3" (Stable Baselines3)
+# "RLLIB" (Ray RLlib)
+# "CUSTOM" (Custom Agent)
+agent_framework: RLLIB
+
+# Sets which deep learning framework will be used (by RLlib ONLY).
+# Default is TF (Tensorflow).
+# Options are:
+# "TF" (Tensorflow)
+# TF2 (Tensorflow 2.X)
+# TORCH (PyTorch)
+deep_learning_framework: TF2
+
+# Sets which Agent class will be used.
+# Options are:
+# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework)
+# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework)
+# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type)
+# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
+# "RANDOM" (primaite.agents.simple.RandomAgent)
+# "DUMMY" (primaite.agents.simple.DummyAgent)
+agent_identifier: PPO
+
+# Sets whether Red Agent POL and IER is randomised.
+# Options are:
+# True
+# False
+random_red_agent: False
+
+# The (integer) seed to be used in random number generation
+# Default is None (null)
+seed: null
+
+# Set whether the agent will be deterministic instead of stochastic
+# Options are:
+# True
+# False
+deterministic: False
+
+# Sets what view of the environment the deterministic hardcoded agent has. The default is BASIC.
+# Options are:
+# "BASIC" (The current observation space only)
+# "FULL" (Full environment view with actions taken and reward feedback)
+hard_coded_agent_view: FULL
+
+# Sets How the Action Space is defined:
+# "NODE"
+# "ACL"
+# "ANY" node and acl actions
+action_type: NODE
+# observation space
+observation_space:
+  # flatten: true
+  components:
+    - name: NODE_LINK_TABLE
+    # - name: NODE_STATUSES
+    # - name: LINK_TRAFFIC_LEVELS
+
+
+# Number of episodes for training to run per session
+num_train_episodes: 10
+
+# Number of time_steps for training per episode
+num_train_steps: 256
+
+# Number of episodes for evaluation to run per session
+num_eval_episodes: 1
+
+# Number of time_steps for evaluation per episode
+num_eval_steps: 256
+
+# Sets how often the agent will save a checkpoint (every n time episodes).
+# Set to 0 if no checkpoints are required. Default is 10
+checkpoint_every_n_episodes: 10
+
+# Time delay (milliseconds) between steps for CUSTOM agents.
+time_delay: 5
+
+# Type of session to be run. Options are:
+# "TRAIN" (Trains an agent)
+# "EVAL" (Evaluates an agent)
+# "TRAIN_EVAL" (Trains then evaluates an agent)
+session_type: TRAIN_EVAL
+
+# Environment config values
+# The high value for the observation space
+observation_space_high_value: 1000000000
+
+# The Stable Baselines3 learn/eval output verbosity level:
+# Options are:
+# "NONE" (No Output)
+# "INFO" (Info Messages (such as devices and wrappers used))
+# "DEBUG" (All Messages)
+sb3_output_verbose_level: NONE
+
+# Reward values
+# Generic
+all_ok: 0
+# Node Hardware State
+off_should_be_on: -0.001
+off_should_be_resetting: -0.0005
+on_should_be_off: -0.0002
+on_should_be_resetting: -0.0005
+resetting_should_be_on: -0.0005
+resetting_should_be_off: -0.0002
+resetting: -0.0003
+# Node Software or Service State
+good_should_be_patching: 0.0002
+good_should_be_compromised: 0.0005
+good_should_be_overwhelmed: 0.0005
+patching_should_be_good: -0.0005
+patching_should_be_compromised: 0.0002
+patching_should_be_overwhelmed: 0.0002
+patching: -0.0003
+compromised_should_be_good: -0.002
+compromised_should_be_patching: -0.002
+compromised_should_be_overwhelmed: -0.002
+compromised: -0.002
+overwhelmed_should_be_good: -0.002
+overwhelmed_should_be_patching: -0.002
+overwhelmed_should_be_compromised: -0.002
+overwhelmed: -0.002
+# Node File System State
+good_should_be_repairing: 0.0002
+good_should_be_restoring: 0.0002
+good_should_be_corrupt: 0.0005
+good_should_be_destroyed: 0.001
+repairing_should_be_good: -0.0005
+repairing_should_be_restoring: 0.0002
+repairing_should_be_corrupt: 0.0002
+repairing_should_be_destroyed: 0.0000
+repairing: -0.0003
+restoring_should_be_good: -0.001
+restoring_should_be_repairing: -0.0002
+restoring_should_be_corrupt: 0.0001
+restoring_should_be_destroyed: 0.0002
+restoring: -0.0006
+corrupt_should_be_good: -0.001
+corrupt_should_be_repairing: -0.001
+corrupt_should_be_restoring: -0.001
+corrupt_should_be_destroyed: 0.0002
+corrupt: -0.001
+destroyed_should_be_good: -0.002
+destroyed_should_be_repairing: -0.002
+destroyed_should_be_restoring: -0.002
+destroyed_should_be_corrupt: -0.002
+destroyed: -0.002
+scanning: -0.0002
+# IER status
+red_ier_running: -0.0005
+green_ier_blocked: -0.001
+
+# Patching / Reset durations
+os_patching_duration: 5            # The time taken to patch the OS
+node_reset_duration: 5             # The time taken to reset a node (hardware)
+service_patching_duration: 5       # The time taken to patch a service
+file_system_repairing_limit: 5      # The time take to repair the file system
+file_system_restoring_limit: 5      # The time take to restore the file system
+file_system_scanning_limit: 5       # The time taken to scan the file system

tests/conftest.py

@@ -5,7 +5,7 @@ import shutil
 import tempfile
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Dict, Union
+from typing import Any, Dict, Tuple, Union
 from unittest.mock import patch
 
 import pytest
@@ -13,7 +13,7 @@ import pytest
 from primaite import getLogger
 from primaite.environment.primaite_env import Primaite
 from primaite.primaite_session import PrimaiteSession
-from primaite.utils.session_output_reader import av_rewards_dict
+from primaite.utils.session_output_reader import all_transactions_dict, av_rewards_dict
 from tests.mock_and_patch.get_session_path_mock import get_temp_session_path
 
 ACTION_SPACE_NODE_VALUES = 1
@@ -37,16 +37,26 @@ class TempPrimaiteSession(PrimaiteSession):
         super().__init__(training_config_path, lay_down_config_path)
         self.setup()
 
-    def learn_av_reward_per_episode(self) -> Dict[int, float]:
+    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_csv(self) -> Dict[int, float]:
+    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:
@@ -113,7 +123,7 @@ def temp_primaite_session(request):
     """
     training_config_path = request.param[0]
     lay_down_config_path = request.param[1]
-    with patch("primaite.agents.agent.get_session_path", get_temp_session_path) as mck:
+    with patch("primaite.agents.agent_abc.get_session_path", get_temp_session_path) as mck:
         mck.session_timestamp = datetime.now()
 
         return TempPrimaiteSession(training_config_path, lay_down_config_path)

tests/test_reward.py

@@ -48,5 +48,5 @@ def test_rewards_are_being_penalised_at_each_step_function(
     """
     with temp_primaite_session as session:
         session.evaluate()
-        ev_rewards = session.eval_av_reward_per_episode_csv()
+        ev_rewards = session.eval_av_reward_per_episode_dict()
         assert ev_rewards[1] == -8.0

tests/test_rllib_agent.py

@@ -0,0 +1,23 @@
+import pytest
+
+from primaite import getLogger
+from primaite.config.lay_down_config import dos_very_basic_config_path
+from tests import TEST_CONFIG_ROOT
+
+_LOGGER = getLogger(__name__)
+
+
+@pytest.mark.parametrize(
+    "temp_primaite_session",
+    [[TEST_CONFIG_ROOT / "training_config_main_rllib.yaml", dos_very_basic_config_path()]],
+    indirect=True,
+)
+def test_primaite_session(temp_primaite_session):
+    """Test the training_config_main_rllib.yaml training config file."""
+    with temp_primaite_session as session:
+        session_path = session.session_path
+        assert session_path.exists()
+        session.learn()
+
+        assert len(session.learn_av_reward_per_episode_dict().keys()) == 10
+        assert len(session.learn_all_transactions_dict().keys()) == 10 * 256

tests/test_seeding_and_deterministic_session.py

@@ -43,11 +43,9 @@ def test_seeded_learning(temp_primaite_session):
             session._training_config.seed == 67890
         ), "Expected output is based upon a agent that was trained with seed 67890"
         session.learn()
+        actual_mean_reward_per_episode = session.learn_av_reward_per_episode_dict()
 
-        print("\n")
-        print(session.learn_av_reward_per_episode())
-
-        assert expected_mean_reward_per_episode == session.learn_av_reward_per_episode()
+    assert actual_mean_reward_per_episode == expected_mean_reward_per_episode
 
 
 @pytest.mark.skip(reason="Inconsistent results. Needs someone with RL knowledge to investigate further.")
@@ -62,5 +60,5 @@ def test_deterministic_evaluation(temp_primaite_session):
         # do stuff
         session.learn()
         session.evaluate()
-        eval_mean_reward = session.eval_av_reward_per_episode_csv()
+        eval_mean_reward = session.eval_av_reward_per_episode_dict()
         assert len(set(eval_mean_reward.values())) == 1

tests/test_session_loading.py

@@ -0,0 +1,188 @@
+import os.path
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Union
+from uuid import uuid4
+
+from primaite import getLogger
+from primaite.agents.sb3 import SB3Agent
+from primaite.common.enums import AgentFramework, AgentIdentifier
+from primaite.main import run
+from primaite.primaite_session import PrimaiteSession
+from primaite.utils.session_output_reader import av_rewards_dict
+from tests import TEST_ASSETS_ROOT
+
+_LOGGER = getLogger(__name__)
+
+
+def copy_session_asset(asset_path: Union[str, Path]) -> str:
+    """Copies the asset into a temporary test folder."""
+    if asset_path is None:
+        raise Exception("No path provided")
+
+    if isinstance(asset_path, Path):
+        asset_path = str(os.path.normpath(asset_path))
+
+    copy_path = str(Path(tempfile.gettempdir()) / "primaite" / str(uuid4()))
+
+    # copy the asset into a temp path
+    try:
+        shutil.copytree(asset_path, copy_path)
+    except Exception as e:
+        msg = f"Unable to copy directory: {asset_path}"
+        _LOGGER.error(msg, e)
+        print(msg, e)
+
+    _LOGGER.debug(f"Copied test asset to: {copy_path}")
+
+    # return the copied assets path
+    return copy_path
+
+
+def test_load_sb3_session():
+    """Test that loading an SB3 agent works."""
+    expected_learn_mean_reward_per_episode = {
+        10: 0,
+        11: -0.008037109374999995,
+        12: -0.007978515624999988,
+        13: -0.008191406249999991,
+        14: -0.00817578124999999,
+        15: -0.008085937499999998,
+        16: -0.007837890624999982,
+        17: -0.007798828124999992,
+        18: -0.007777343749999998,
+        19: -0.007958984374999988,
+        20: -0.0077499999999999835,
+    }
+
+    test_path = copy_session_asset(TEST_ASSETS_ROOT / "example_sb3_agent_session")
+
+    loaded_agent = SB3Agent(session_path=test_path)
+
+    # loaded agent should have the same UUID as the previous agent
+    assert loaded_agent.uuid == "301874d3-2e14-43c2-ba7f-e2b03ad05dde"
+    assert loaded_agent._training_config.agent_framework == AgentFramework.SB3.name
+    assert loaded_agent._training_config.agent_identifier == AgentIdentifier.PPO.name
+    assert loaded_agent._training_config.deterministic
+    assert loaded_agent._training_config.seed == 12345
+    assert str(loaded_agent.session_path) == str(test_path)
+
+    # run another learn session
+    loaded_agent.learn()
+
+    learn_mean_rewards = av_rewards_dict(
+        loaded_agent.learning_path / f"average_reward_per_episode_{loaded_agent.timestamp_str}.csv"
+    )
+
+    # run is seeded so should have the expected learn value
+    assert learn_mean_rewards == expected_learn_mean_reward_per_episode
+
+    # run an evaluation
+    loaded_agent.evaluate()
+
+    # load the evaluation average reward csv file
+    eval_mean_reward = av_rewards_dict(
+        loaded_agent.evaluation_path / f"average_reward_per_episode_{loaded_agent.timestamp_str}.csv"
+    )
+
+    # the agent config ran the evaluation in deterministic mode, so should have the same reward value
+    assert len(set(eval_mean_reward.values())) == 1
+
+    # the evaluation should be the same as a previous run
+    assert next(iter(set(eval_mean_reward.values()))) == -0.009896484374999988
+
+    # delete the test directory
+    shutil.rmtree(test_path)
+
+
+def test_load_primaite_session():
+    """Test that loading a Primaite session works."""
+    expected_learn_mean_reward_per_episode = {
+        10: 0,
+        11: -0.008037109374999995,
+        12: -0.007978515624999988,
+        13: -0.008191406249999991,
+        14: -0.00817578124999999,
+        15: -0.008085937499999998,
+        16: -0.007837890624999982,
+        17: -0.007798828124999992,
+        18: -0.007777343749999998,
+        19: -0.007958984374999988,
+        20: -0.0077499999999999835,
+    }
+
+    test_path = copy_session_asset(TEST_ASSETS_ROOT / "example_sb3_agent_session")
+
+    # create loaded session
+    session = PrimaiteSession(session_path=test_path)
+
+    # run setup on session
+    session.setup()
+
+    # make sure that the session was loaded correctly
+    assert session._agent_session.uuid == "301874d3-2e14-43c2-ba7f-e2b03ad05dde"
+    assert session._agent_session._training_config.agent_framework == AgentFramework.SB3.name
+    assert session._agent_session._training_config.agent_identifier == AgentIdentifier.PPO.name
+    assert session._agent_session._training_config.deterministic
+    assert session._agent_session._training_config.seed == 12345
+    assert str(session._agent_session.session_path) == str(test_path)
+
+    # run another learn session
+    session.learn()
+
+    learn_mean_rewards = av_rewards_dict(
+        session.learning_path / f"average_reward_per_episode_{session.timestamp_str}.csv"
+    )
+
+    # run is seeded so should have the expected learn value
+    assert learn_mean_rewards == expected_learn_mean_reward_per_episode
+
+    # run an evaluation
+    session.evaluate()
+
+    # load the evaluation average reward csv file
+    eval_mean_reward = av_rewards_dict(
+        session.evaluation_path / f"average_reward_per_episode_{session.timestamp_str}.csv"
+    )
+
+    # the agent config ran the evaluation in deterministic mode, so should have the same reward value
+    assert len(set(eval_mean_reward.values())) == 1
+
+    # the evaluation should be the same as a previous run
+    assert next(iter(set(eval_mean_reward.values()))) == -0.009896484374999988
+
+    # delete the test directory
+    shutil.rmtree(test_path)
+
+
+def test_run_loading():
+    """Test loading session via main.run."""
+    expected_learn_mean_reward_per_episode = {
+        10: 0,
+        11: -0.008037109374999995,
+        12: -0.007978515624999988,
+        13: -0.008191406249999991,
+        14: -0.00817578124999999,
+        15: -0.008085937499999998,
+        16: -0.007837890624999982,
+        17: -0.007798828124999992,
+        18: -0.007777343749999998,
+        19: -0.007958984374999988,
+        20: -0.0077499999999999835,
+    }
+
+    test_path = copy_session_asset(TEST_ASSETS_ROOT / "example_sb3_agent_session")
+
+    # create loaded session
+    run(session_path=test_path)
+
+    learn_mean_rewards = av_rewards_dict(
+        next(Path(test_path).rglob("**/learning/average_reward_per_episode_*.csv"), None)
+    )
+
+    # run is seeded so should have the expected learn value
+    assert learn_mean_rewards == expected_learn_mean_reward_per_episode
+
+    # delete the test directory
+    shutil.rmtree(test_path)