From 7af9d3724f41bcb4fa8e9a015e6d233598768db9 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Wed, 8 Jan 2025 14:42:35 +0000 Subject: [PATCH] #2869 - Updates to address test failures. Updated YAML configs to remove redundant start_settings --- .../_package_data/data_manipulation.yaml | 7 ++-- .../_package_data/data_manipulation_marl.yaml | 7 ++-- .../scenario_with_placeholders/reds_1.yaml | 7 ++-- .../scenario_with_placeholders/reds_2.yaml | 7 ++-- src/primaite/game/agent/interface.py | 12 +++--- .../agent/scripted_agents/abstract_tap.py | 4 +- .../scripted_agents/data_manipulation_bot.py | 4 +- .../scripted_agents/probabilistic_agent.py | 16 +++----- src/primaite/game/game.py | 1 + src/primaite/session/environment.py | 2 +- .../assets/configs/bad_primaite_session.yaml | 14 +++---- tests/assets/configs/basic_firewall.yaml | 7 ++-- .../configs/basic_switched_network.yaml | 7 ++-- tests/assets/configs/data_manipulation.yaml | 7 ++-- tests/assets/configs/dmz_network.yaml | 7 ++-- .../configs/eval_only_primaite_session.yaml | 14 +++---- tests/assets/configs/extended_config.yaml | 7 ++-- .../configs/firewall_actions_network.yaml | 7 ++-- .../assets/configs/fix_duration_one_item.yaml | 7 ++-- tests/assets/configs/multi_agent_session.yaml | 7 ++-- .../scenario_with_placeholders/reds_1.yaml | 7 ++-- .../scenario_with_placeholders/reds_2.yaml | 7 ++-- tests/assets/configs/shared_rewards.yaml | 7 ++-- .../assets/configs/software_fix_duration.yaml | 7 ++-- .../configs/test_application_install.yaml | 7 ++-- .../assets/configs/test_primaite_session.yaml | 14 +++---- tests/conftest.py | 2 +- .../_game/_agent/test_probabilistic_agent.py | 39 +++++++++++-------- 28 files changed, 111 insertions(+), 130 deletions(-) diff --git a/src/primaite/config/_package_data/data_manipulation.yaml b/src/primaite/config/_package_data/data_manipulation.yaml index 97442903..d604192e 100644 --- a/src/primaite/config/_package_data/data_manipulation.yaml +++ b/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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE diff --git a/src/primaite/config/_package_data/data_manipulation_marl.yaml b/src/primaite/config/_package_data/data_manipulation_marl.yaml index ba666781..00a34403 100644 --- a/src/primaite/config/_package_data/data_manipulation_marl.yaml +++ b/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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + variance: 5 - ref: defender_1 team: BLUE diff --git a/src/primaite/config/_package_data/scenario_with_placeholders/reds_1.yaml b/src/primaite/config/_package_data/scenario_with_placeholders/reds_1.yaml index 31675a0b..b775cb24 100644 --- a/src/primaite/config/_package_data/scenario_with_placeholders/reds_1.yaml +++ b/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 - frequency: 10 - variance: 0 + start_step: 10 + frequency: 10 + variance: 0 diff --git a/src/primaite/config/_package_data/scenario_with_placeholders/reds_2.yaml b/src/primaite/config/_package_data/scenario_with_placeholders/reds_2.yaml index c5572b89..4cae1ec6 100644 --- a/src/primaite/config/_package_data/scenario_with_placeholders/reds_2.yaml +++ b/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 - frequency: 2 - variance: 1 + start_step: 3 + frequency: 2 + variance: 1 diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 6c1c633a..1627d360 100644 --- a/src/primaite/game/agent/interface.py +++ b/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"]), - observation_manager=ObservationManager.from_config(**config["observation_space"]), - reward_function=RewardFunction.from_config(**config["reward_function"]), + action_manager=ActionManager.from_config(config["game"], config["action_manager"]), + observation_manager=ObservationManager.from_config(config["observation_manager"]), + 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]: diff --git a/src/primaite/game/agent/scripted_agents/abstract_tap.py b/src/primaite/game/agent/scripted_agents/abstract_tap.py index 9fb782d4..f0dd096d 100644 --- a/src/primaite/game/agent/scripted_agents/abstract_tap.py +++ b/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( - -self.config.agent_settings.start_settings.variance, self.config.agent_settings.start_settings.variance - ) + random_timestep_increment = random.randint(-self.config.variance, self.config.variance) self.next_execution_timestep = timestep + random_timestep_increment def _select_start_node(self) -> None: diff --git a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py index 7ec119cf..f7bf4bc5 100644 --- a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py +++ b/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) diff --git a/src/primaite/game/agent/scripted_agents/probabilistic_agent.py b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py index 159e5cd2..78f806d0 100644 --- a/src/primaite/game/agent/scripted_agents/probabilistic_agent.py +++ b/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): - """ProbabilisticAgent settings.""" + class ConfigSchema(AbstractScriptedAgent.ConfigSchema): + """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]: """ diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index db23eb14..69e294ae 100644 --- a/src/primaite/game/game.py +++ b/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 diff --git a/src/primaite/session/environment.py b/src/primaite/session/environment.py index 8e608ede..29f7c33d 100644 --- a/src/primaite/session/environment.py +++ b/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) diff --git a/tests/assets/configs/bad_primaite_session.yaml b/tests/assets/configs/bad_primaite_session.yaml index c83cadc8..6a19c2fb 100644 --- a/tests/assets/configs/bad_primaite_session.yaml +++ b/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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + 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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE diff --git a/tests/assets/configs/basic_firewall.yaml b/tests/assets/configs/basic_firewall.yaml index e37a67da..fe5e0099 100644 --- a/tests/assets/configs/basic_firewall.yaml +++ b/tests/assets/configs/basic_firewall.yaml @@ -56,10 +56,9 @@ agents: - type: DUMMY agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 + start_step: 5 + frequency: 4 + variance: 3 action_probabilities: 0: 0.4 1: 0.6 diff --git a/tests/assets/configs/basic_switched_network.yaml b/tests/assets/configs/basic_switched_network.yaml index 00ba381b..f27735d1 100644 --- a/tests/assets/configs/basic_switched_network.yaml +++ b/tests/assets/configs/basic_switched_network.yaml @@ -59,10 +59,9 @@ agents: - type: DUMMY agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 + start_step: 5 + frequency: 4 + variance: 3 action_probabilities: 0: 0.6 1: 0.4 diff --git a/tests/assets/configs/data_manipulation.yaml b/tests/assets/configs/data_manipulation.yaml index 97442903..d604192e 100644 --- a/tests/assets/configs/data_manipulation.yaml +++ b/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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE diff --git a/tests/assets/configs/dmz_network.yaml b/tests/assets/configs/dmz_network.yaml index d560efa3..41a530b0 100644 --- a/tests/assets/configs/dmz_network.yaml +++ b/tests/assets/configs/dmz_network.yaml @@ -81,10 +81,9 @@ agents: - type: DUMMY agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 + start_step: 5 + frequency: 4 + variance: 3 action_probabilities: 0: 0.4 1: 0.6 diff --git a/tests/assets/configs/eval_only_primaite_session.yaml b/tests/assets/configs/eval_only_primaite_session.yaml index 3d60eb6e..dc0acdaa 100644 --- a/tests/assets/configs/eval_only_primaite_session.yaml +++ b/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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + 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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE diff --git a/tests/assets/configs/extended_config.yaml b/tests/assets/configs/extended_config.yaml index e1a06938..fc1b72dd 100644 --- a/tests/assets/configs/extended_config.yaml +++ b/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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE diff --git a/tests/assets/configs/firewall_actions_network.yaml b/tests/assets/configs/firewall_actions_network.yaml index 2292616d..2d42e4c5 100644 --- a/tests/assets/configs/firewall_actions_network.yaml +++ b/tests/assets/configs/firewall_actions_network.yaml @@ -266,10 +266,9 @@ agents: - type: DUMMY agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 + start_step: 5 + frequency: 4 + variance: 3 diff --git a/tests/assets/configs/fix_duration_one_item.yaml b/tests/assets/configs/fix_duration_one_item.yaml index 62579e35..0252ac32 100644 --- a/tests/assets/configs/fix_duration_one_item.yaml +++ b/tests/assets/configs/fix_duration_one_item.yaml @@ -56,10 +56,9 @@ agents: - type: DUMMY agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 + start_step: 5 + frequency: 4 + variance: 3 action_probabilities: 0: 0.4 1: 0.6 diff --git a/tests/assets/configs/multi_agent_session.yaml b/tests/assets/configs/multi_agent_session.yaml index a2d64605..13cffab1 100644 --- a/tests/assets/configs/multi_agent_session.yaml +++ b/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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + variance: 5 - ref: defender_1 team: BLUE diff --git a/tests/assets/configs/scenario_with_placeholders/reds_1.yaml b/tests/assets/configs/scenario_with_placeholders/reds_1.yaml index 31675a0b..b775cb24 100644 --- a/tests/assets/configs/scenario_with_placeholders/reds_1.yaml +++ b/tests/assets/configs/scenario_with_placeholders/reds_1.yaml @@ -20,7 +20,6 @@ reds: &reds - type: DUMMY agent_settings: - start_settings: - start_step: 10 - frequency: 10 - variance: 0 + start_step: 10 + frequency: 10 + variance: 0 diff --git a/tests/assets/configs/scenario_with_placeholders/reds_2.yaml b/tests/assets/configs/scenario_with_placeholders/reds_2.yaml index c5572b89..4cae1ec6 100644 --- a/tests/assets/configs/scenario_with_placeholders/reds_2.yaml +++ b/tests/assets/configs/scenario_with_placeholders/reds_2.yaml @@ -20,7 +20,6 @@ reds: &reds - type: DUMMY agent_settings: - start_settings: - start_step: 3 - frequency: 2 - variance: 1 + start_step: 3 + frequency: 2 + variance: 1 diff --git a/tests/assets/configs/shared_rewards.yaml b/tests/assets/configs/shared_rewards.yaml index 81cb85f7..3ba480ea 100644 --- a/tests/assets/configs/shared_rewards.yaml +++ b/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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE diff --git a/tests/assets/configs/software_fix_duration.yaml b/tests/assets/configs/software_fix_duration.yaml index 3e3d6e22..98260fe3 100644 --- a/tests/assets/configs/software_fix_duration.yaml +++ b/tests/assets/configs/software_fix_duration.yaml @@ -56,10 +56,9 @@ agents: - type: DUMMY agent_settings: - start_settings: - start_step: 5 - frequency: 4 - variance: 3 + start_step: 5 + frequency: 4 + variance: 3 action_probabilities: 0: 0.4 1: 0.6 diff --git a/tests/assets/configs/test_application_install.yaml b/tests/assets/configs/test_application_install.yaml index 3a3a6890..e8b080b7 100644 --- a/tests/assets/configs/test_application_install.yaml +++ b/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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE diff --git a/tests/assets/configs/test_primaite_session.yaml b/tests/assets/configs/test_primaite_session.yaml index cf241f3c..c59bbcbf 100644 --- a/tests/assets/configs/test_primaite_session.yaml +++ b/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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + 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 - frequency: 20 - variance: 5 + start_step: 25 + frequency: 20 + variance: 5 - ref: defender team: BLUE diff --git a/tests/conftest.py b/tests/conftest.py index 0c211f49..9d18a18b 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py b/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py index 4e2378b2..69540f0a 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py +++ b/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( - actions=[ + action_space_cfg = { + "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_files_per_folder=2, - max_services_per_node=2, - max_applications_per_node=2, - max_nics_per_node=2, - max_acl_rules=10, - protocols=["TCP", "UDP", "ICMP"], - ports=["HTTP", "DNS", "ARP"], - act_map={ + "max_folders_per_node": 2, + "max_files_per_folder": 2, + "max_services_per_node": 2, + "max_applications_per_node": 2, + "max_nics_per_node": 2, + "max_acl_rules": 10, + "protocols": ["TCP", "UDP", "ICMP"], + "ports": ["HTTP", "DNS", "ARP"], + "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, - "observation_manager": observation_space, - "reward_function": reward_function, + "game": PrimaiteGame(), + "action_manager": action_space_cfg, + "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}, },