#2869 - Updates to address test failures. Updated YAML configs to remove redundant start_settings

src/primaite/config/_package_data/data_manipulation.yaml

@@ -151,10 +151,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: defender
     team: BLUE

src/primaite/config/_package_data/data_manipulation_marl.yaml

@@ -150,10 +150,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: defender_1
     team: BLUE

src/primaite/config/_package_data/scenario_with_placeholders/reds_1.yaml

@@ -20,7 +20,6 @@ reds: &reds
         - type: DUMMY
 
     agent_settings:
-      start_settings:
+      start_step: 10
-        start_step: 10
+      frequency: 10
-        frequency: 10
+      variance: 0
-        variance: 0

src/primaite/config/_package_data/scenario_with_placeholders/reds_2.yaml

@@ -20,7 +20,6 @@ reds: &reds
         - type: DUMMY
 
     agent_settings:
-      start_settings:
+      start_step: 3
-        start_step: 3
+      frequency: 2
-        frequency: 2
+      variance: 1
-        variance: 1

src/primaite/game/agent/interface.py

@@ -47,13 +47,14 @@ class AbstractAgent(BaseModel):
     """Base class for scripted and RL agents."""
 
     _registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {}
-    _logger: AgentLog = AgentLog(agent_name="Abstract_Agent")
+    logger: AgentLog = AgentLog(agent_name="Abstract_Agent")
 
     history: List[AgentHistoryItem] = []
     config: "AbstractAgent.ConfigSchema"
     action_manager: "ActionManager"
     observation_manager: "ObservationManager"
     reward_function: "RewardFunction"
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
 
     class ConfigSchema(BaseModel):
         """
@@ -114,11 +115,12 @@ class AbstractAgent(BaseModel):
     @classmethod
     def from_config(cls, config: Dict) -> "AbstractAgent":
         """Creates an agent component from a configuration dictionary."""
+        print(config)
         obj = cls(
             config=cls.ConfigSchema(**config["agent_settings"]),
-            action_manager=ActionManager.from_config(**config["action_manager"]),
+            action_manager=ActionManager.from_config(config["game"], config["action_manager"]),
-            observation_manager=ObservationManager.from_config(**config["observation_space"]),
+            observation_manager=ObservationManager.from_config(config["observation_manager"]),
-            reward_function=RewardFunction.from_config(**config["reward_function"]),
+            reward_function=RewardFunction.from_config(config["reward_function"]),
         )
         return obj
 
@@ -140,7 +142,7 @@ class AbstractAgent(BaseModel):
         :return: Reward from the state.
         :rtype: float
         """
-        return self.reward_function.update(state=state, last_action_response=self.config.history[-1])
+        return self.reward_function.update(state=state, last_action_response=self.history[-1])
 
     @abstractmethod
     def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:

src/primaite/game/agent/scripted_agents/abstract_tap.py

@@ -39,9 +39,7 @@ class AbstractTAPAgent(AbstractScriptedAgent, identifier="Abstract_TAP"):
 
         :param timestep: The timestep to add variance to.
         """
-        random_timestep_increment = random.randint(
+        random_timestep_increment = random.randint(-self.config.variance, self.config.variance)
-            -self.config.agent_settings.start_settings.variance, self.config.agent_settings.start_settings.variance
-        )
         self.next_execution_timestep = timestep + random_timestep_increment
 
     def _select_start_node(self) -> None:

src/primaite/game/agent/scripted_agents/data_manipulation_bot.py

@@ -42,7 +42,7 @@ class DataManipulationAgent(AbstractTAPAgent, identifier="RedDatabaseCorruptingA
             self.logger.debug(msg="Performing do nothing action")
             return "do_nothing", {}
 
-        self._set_next_execution_timestep(timestep + self.config.agent_settings.start_settings.frequency)
+        self._set_next_execution_timestep(timestep + self.config.frequency)
         self.logger.info(msg="Performing a data manipulation attack!")
         return "node_application_execute", {
             "node_name": self.config.starting_node_name,
@@ -52,4 +52,4 @@ class DataManipulationAgent(AbstractTAPAgent, identifier="RedDatabaseCorruptingA
     def setup_agent(self) -> None:
         """Set the next execution timestep when the episode resets."""
         self._select_start_node()
-        self._set_next_execution_timestep(self.config.agent_settings.start_settings.start_step)
+        self._set_next_execution_timestep(self.config.start_step)

src/primaite/game/agent/scripted_agents/probabilistic_agent.py

@@ -6,7 +6,7 @@ import numpy as np
 import pydantic
 from gymnasium.core import ObsType
 
-from primaite.game.agent.interface import AbstractScriptedAgent, AgentSettings
+from primaite.game.agent.interface import AbstractScriptedAgent
 
 __all__ = "ProbabilisticAgent"
 
@@ -17,8 +17,10 @@ class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent")
     config: "ProbabilisticAgent.ConfigSchema"
     rng: Any = np.random.default_rng(np.random.randint(0, 65535))
 
-    class AgentSettings(AgentSettings):
+    class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
-        """ProbabilisticAgent settings."""
+        """Configuration schema for Probabilistic Agent."""
+
+        agent_name: str = "ProbabilisticAgent"
 
         action_probabilities: Dict[int, float]
         """Probability to perform each action in the action map. The sum of probabilities should sum to 1."""
@@ -42,16 +44,10 @@ class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent")
                 )
             return v
 
-    class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
-        """Configuration schema for Probabilistic Agent."""
-
-        agent_name: str = "ProbabilisticAgent"
-        agent_settings: "ProbabilisticAgent.AgentSettings"
-
     @property
     def probabilities(self) -> Dict[str, int]:
         """Convenience method to view the probabilities of the Agent."""
-        return np.asarray(list(self.config.agent_settings.action_probabilities.values()))
+        return np.asarray(list(self.config.action_probabilities.values()))
 
     def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
         """

src/primaite/game/game.py

@@ -538,6 +538,7 @@ class PrimaiteGame:
                 "observation_manager": observation_space_cfg,
                 "reward_function": reward_function_cfg,
                 "agent_settings": agent_settings,
+                "game": game,
             }
 
             # CREATE AGENT

src/primaite/session/environment.py

@@ -89,7 +89,7 @@ class PrimaiteGymEnv(gymnasium.Env):
         :return: Action mask
         :rtype: List[bool]
         """
-        if not self.agent.action_masking:
+        if not self.agent.config.action_masking:
             return np.asarray([True] * len(self.agent.action_manager.action_map))
         else:
             return self.game.action_mask(self._agent_name)

tests/assets/configs/bad_primaite_session.yaml

@@ -31,10 +31,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: data_manipulation_attacker
     team: RED
@@ -63,10 +62,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: defender
     team: BLUE

tests/assets/configs/basic_firewall.yaml

@@ -56,10 +56,9 @@ agents:
         - type: DUMMY
 
     agent_settings:
-      start_settings:
+      start_step: 5
-        start_step: 5
+      frequency: 4
-        frequency: 4
+      variance: 3
-        variance: 3
       action_probabilities:
         0: 0.4
         1: 0.6

tests/assets/configs/basic_switched_network.yaml

@@ -59,10 +59,9 @@ agents:
         - type: DUMMY
 
     agent_settings:
-      start_settings:
+      start_step: 5
-        start_step: 5
+      frequency: 4
-        frequency: 4
+      variance: 3
-        variance: 3
       action_probabilities:
         0: 0.6
         1: 0.4

tests/assets/configs/data_manipulation.yaml

@@ -151,10 +151,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: defender
     team: BLUE

tests/assets/configs/dmz_network.yaml

@@ -81,10 +81,9 @@ agents:
         - type: DUMMY
 
     agent_settings:
-      start_settings:
+      start_step: 5
-        start_step: 5
+      frequency: 4
-        frequency: 4
+      variance: 3
-        variance: 3
       action_probabilities:
         0: 0.4
         1: 0.6

tests/assets/configs/eval_only_primaite_session.yaml

@@ -35,10 +35,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: data_manipulation_attacker
     team: RED
@@ -75,10 +74,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: defender
     team: BLUE

tests/assets/configs/extended_config.yaml

@@ -151,10 +151,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: defender
     team: BLUE

tests/assets/configs/firewall_actions_network.yaml

@@ -266,10 +266,9 @@ agents:
         - type: DUMMY
 
     agent_settings:
-      start_settings:
+      start_step: 5
-        start_step: 5
+      frequency: 4
-        frequency: 4
+      variance: 3
-        variance: 3
 
 
 

tests/assets/configs/fix_duration_one_item.yaml

@@ -56,10 +56,9 @@ agents:
         - type: DUMMY
 
     agent_settings:
-      start_settings:
+      start_step: 5
-        start_step: 5
+      frequency: 4
-        frequency: 4
+      variance: 3
-        variance: 3
       action_probabilities:
         0: 0.4
         1: 0.6

tests/assets/configs/multi_agent_session.yaml

@@ -150,10 +150,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: defender_1
     team: BLUE

tests/assets/configs/scenario_with_placeholders/reds_1.yaml

@@ -20,7 +20,6 @@ reds: &reds
         - type: DUMMY
 
     agent_settings:
-      start_settings:
+      start_step: 10
-        start_step: 10
+      frequency: 10
-        frequency: 10
+      variance: 0
-        variance: 0

tests/assets/configs/scenario_with_placeholders/reds_2.yaml

@@ -20,7 +20,6 @@ reds: &reds
         - type: DUMMY
 
     agent_settings:
-      start_settings:
+      start_step: 3
-        start_step: 3
+      frequency: 2
-        frequency: 2
+      variance: 1
-        variance: 1

tests/assets/configs/shared_rewards.yaml

@@ -146,10 +146,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: defender
     team: BLUE

tests/assets/configs/software_fix_duration.yaml

@@ -56,10 +56,9 @@ agents:
         - type: DUMMY
 
     agent_settings:
-      start_settings:
+      start_step: 5
-        start_step: 5
+      frequency: 4
-        frequency: 4
+      variance: 3
-        variance: 3
       action_probabilities:
         0: 0.4
         1: 0.6

tests/assets/configs/test_application_install.yaml

@@ -150,10 +150,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: defender
     team: BLUE

tests/assets/configs/test_primaite_session.yaml

@@ -43,10 +43,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
       action_probabilities:
         0: 1.0
 
@@ -86,10 +85,9 @@ agents:
         - type: DUMMY
 
     agent_settings: # options specific to this particular agent type, basically args of __init__(self)
-      start_settings:
+      start_step: 25
-        start_step: 25
+      frequency: 20
-        frequency: 20
+      variance: 5
-        variance: 5
 
   - ref: defender
     team: BLUE

tests/conftest.py

@@ -1,5 +1,5 @@
 # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import pytest
 import yaml

tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py

@@ -3,6 +3,7 @@ from primaite.game.agent.actions import ActionManager
 from primaite.game.agent.observations.observation_manager import NestedObservation, ObservationManager
 from primaite.game.agent.rewards import RewardFunction
 from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent
+from primaite.game.game import PrimaiteGame
 
 
 def test_probabilistic_agent():
@@ -25,36 +26,41 @@ def test_probabilistic_agent():
     MIN_NODE_FILE_DELETE = 5750
     MAX_NODE_FILE_DELETE = 6250
 
-    action_space = ActionManager(
+    action_space_cfg = {
-        actions=[
+        "action_list": [
             {"type": "DONOTHING"},
             {"type": "NODE_APPLICATION_EXECUTE"},
             {"type": "NODE_FILE_DELETE"},
         ],
-        nodes=[
+        "nodes": [
             {
                 "node_name": "client_1",
                 "applications": [{"application_name": "WebBrowser"}],
                 "folders": [{"folder_name": "downloads", "files": [{"file_name": "cat.png"}]}],
             },
         ],
-        max_folders_per_node=2,
+        "max_folders_per_node": 2,
-        max_files_per_folder=2,
+        "max_files_per_folder": 2,
-        max_services_per_node=2,
+        "max_services_per_node": 2,
-        max_applications_per_node=2,
+        "max_applications_per_node": 2,
-        max_nics_per_node=2,
+        "max_nics_per_node": 2,
-        max_acl_rules=10,
+        "max_acl_rules": 10,
-        protocols=["TCP", "UDP", "ICMP"],
+        "protocols": ["TCP", "UDP", "ICMP"],
-        ports=["HTTP", "DNS", "ARP"],
+        "ports": ["HTTP", "DNS", "ARP"],
-        act_map={
+        "act_map": {
             0: {"action": "DONOTHING", "options": {}},
             1: {"action": "NODE_APPLICATION_EXECUTE", "options": {"node_id": 0, "application_id": 0}},
             2: {"action": "NODE_FILE_DELETE", "options": {"node_id": 0, "folder_id": 0, "file_id": 0}},
         },
-    )
+        "options": {},
+    }
     observation_space = ObservationManager(NestedObservation(components={}))
     reward_function = RewardFunction()
 
+    observation_space_cfg = None
+
+    reward_function_cfg = {}
+
     # pa = ProbabilisticAgent(
     #     agent_name="test_agent",
     #     action_space=action_space,
@@ -67,9 +73,10 @@ def test_probabilistic_agent():
 
     pa_config = {
         "agent_name": "test_agent",
-        "action_manager": action_space,
+        "game": PrimaiteGame(),
-        "observation_manager": observation_space,
+        "action_manager": action_space_cfg,
-        "reward_function": reward_function,
+        "observation_manager": observation_space_cfg,
+        "reward_function": reward_function_cfg,
         "agent_settings": {
             "action_probabilities": {0: P_DO_NOTHING, 1: P_NODE_APPLICATION_EXECUTE, 2: P_NODE_FILE_DELETE},
         },