oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#2869 - eod commit. Updates to AbstractAgent.from_config, and some minor tweaks to PrimaiteGame

Browse Source
This commit is contained in:
Charlie Crane
2024-11-20 17:51:05 +00:00
parent a3dc616126
commit 75d4ef2dfd
4 changed files with 27 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
src/primaite/game/agent/interface.py
View File

@@ -97,11 +97,12 @@ class AbstractAgent(BaseModel, ABC, identifier="Abstract_Agent"):
_registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {} _registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {}
config: "AbstractAgent.ConfigSchema"
action_manager: Optional[ActionManager] action_manager: Optional[ActionManager]
observation_manager: Optional[ObservationManager] observation_manager: Optional[ObservationManager]
reward_function: Optional[RewardFunction] reward_function: Optional[RewardFunction]
config: "AbstractAgent.ConfigSchema"
class ConfigSchema(BaseModel): class ConfigSchema(BaseModel):
""" """
Configuration Schema for AbstractAgents. Configuration Schema for AbstractAgents.
@@ -163,7 +164,14 @@ class AbstractAgent(BaseModel, ABC, identifier="Abstract_Agent"):
@classmethod @classmethod
def from_config(cls, config: Dict) -> "AbstractAgent": def from_config(cls, config: Dict) -> "AbstractAgent":
"""Creates an agent component from a configuration dictionary.""" """Creates an agent component from a configuration dictionary."""
return cls(config=cls.ConfigSchema(**config)) obj = cls(config=cls.ConfigSchema(**config))
# Pull managers out of config section for ease of use (?)
obj.observation_manager = obj.config.observation_manager
obj.action_manager = obj.config.action_manager
obj.reward_function = obj.config.reward_function
return obj
def update_observation(self, state: Dict) -> ObsType: def update_observation(self, state: Dict) -> ObsType:
""" """
@@ -172,7 +180,7 @@ class AbstractAgent(BaseModel, ABC, identifier="Abstract_Agent"):
state : dict state directly from simulation.describe_state state : dict state directly from simulation.describe_state
output : dict state according to CAOS. output : dict state according to CAOS.
""" """
return self.config.observation_manager.update(state) return self.observation_manager.update(state)
def update_reward(self, state: Dict) -> float: def update_reward(self, state: Dict) -> float:
""" """
@@ -183,7 +191,7 @@ class AbstractAgent(BaseModel, ABC, identifier="Abstract_Agent"):
:return: Reward from the state. :return: Reward from the state.
:rtype: float :rtype: float
""" """
return self.config.reward_function.update(state=state, last_action_response=self.history[-1]) return self.reward_function.update(state=state, last_action_response=self.config.history[-1])
@abstractmethod @abstractmethod
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]: def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
@@ -201,13 +209,13 @@ class AbstractAgent(BaseModel, ABC, identifier="Abstract_Agent"):
""" """
# in RL agent, this method will send CAOS observation to RL agent, then receive a int 0-39, # in RL agent, this method will send CAOS observation to RL agent, then receive a int 0-39,
# then use a bespoke conversion to take 1-40 int back into CAOS action # then use a bespoke conversion to take 1-40 int back into CAOS action
return ("DO_NOTHING", {}) return ("do_nothing", {})
def format_request(self, action: Tuple[str, Dict], options: Dict[str, int]) -> List[str]: def format_request(self, action: Tuple[str, Dict], options: Dict[str, int]) -> List[str]:
# this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator. # this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator.
# therefore the execution definition needs to be a mapping from CAOS into SIMULATOR # therefore the execution definition needs to be a mapping from CAOS into SIMULATOR
"""Format action into format expected by the simulator, and apply execution definition if applicable.""" """Format action into format expected by the simulator, and apply execution definition if applicable."""
request = self.config.action_manager.form_request(action_identifier=action, action_options=options) request = self.action_manager.form_request(action_identifier=action, action_options=options)
return request return request
def process_action_response( def process_action_response(
@@ -222,7 +230,7 @@ class AbstractAgent(BaseModel, ABC, identifier="Abstract_Agent"):
def save_reward_to_history(self) -> None: def save_reward_to_history(self) -> None:
"""Update the most recent history item with the reward value.""" """Update the most recent history item with the reward value."""
self.config.history[-1].reward = self.config.reward_function.current_reward self.config.history[-1].reward = self.reward_function.current_reward
class AbstractScriptedAgent(AbstractAgent, identifier="Abstract_Scripted_Agent"): class AbstractScriptedAgent(AbstractAgent, identifier="Abstract_Scripted_Agent"):

4
src/primaite/game/agent/scripted_agents/probabilistic_agent.py
View File

@@ -15,7 +15,7 @@ from primaite.game.agent.rewards import RewardFunction
class ProbabilisticAgent(AbstractScriptedAgent, identifier="Probabilistic_Agent"): class ProbabilisticAgent(AbstractScriptedAgent, identifier="Probabilistic_Agent"):
"""Scripted agent which randomly samples its action space with prescribed probabilities for each action.""" """Scripted agent which randomly samples its action space with prescribed probabilities for each action."""
class Settings(pydantic.BaseModel): class ConfigSchema(pydantic.BaseModel):
"""Config schema for Probabilistic agent settings.""" """Config schema for Probabilistic agent settings."""
model_config = pydantic.ConfigDict(extra="forbid") model_config = pydantic.ConfigDict(extra="forbid")
@@ -60,7 +60,7 @@ class ProbabilisticAgent(AbstractScriptedAgent, identifier="Probabilistic_Agent"
# The random number seed for np.random is dependent on whether a random number seed is set # The random number seed for np.random is dependent on whether a random number seed is set
# in the config file. If there is one it is processed by set_random_seed() in environment.py # in the config file. If there is one it is processed by set_random_seed() in environment.py
# and as a consequence the the sequence of rng_seed's used here will be repeatable. # and as a consequence the the sequence of rng_seed's used here will be repeatable.
self.settings = ProbabilisticAgent.Settings(**settings) self.settings = ProbabilisticAgent.ConfigSchema(**settings)
rng_seed = np.random.randint(0, 65535) rng_seed = np.random.randint(0, 65535)
self.rng = np.random.default_rng(rng_seed) self.rng = np.random.default_rng(rng_seed)

2
src/primaite/game/agent/scripted_agents/tap001.py
View File

@@ -16,7 +16,7 @@ class TAP001(AbstractScriptedAgent, identifier="TAP001"):
Scripted Red Agent. Capable of one action; launching the kill-chain (Ransomware Application) Scripted Red Agent. Capable of one action; launching the kill-chain (Ransomware Application)
""" """
# TODO: Link with DataManipulationAgent via a parent "TAP" agent class. # TODO: Link with DataManipulationAgent class via a parent "TAP" agent class.
config: "TAP001.ConfigSchema" config: "TAP001.ConfigSchema"

19
src/primaite/game/game.py
View File

@@ -555,17 +555,16 @@ class PrimaiteGame:
# action manager # action manager
# observation_manager # observation_manager
# reward_function # reward_function
agent_config = agent_cfg.get("agent_settings", {})
new_agent_cfg = { agent_config.update({"action_manager": action_space,
"action_manager": action_space, "observation_manager": obs_space,
"agent_name": agent_cfg["ref"], "reward_function":reward_function})
"observation_manager": obs_space,
"agent_settings": agent_cfg.get("agent_settings", {}),
"reward_function": reward_function,
}
new_agent_cfg = agent_cfg["settings"]
# new_agent_cfg.update{} # new_agent_cfg.update{}
new_agent = AbstractAgent._registry[agent_cfg["type"]].from_config(config=new_agent_cfg) new_agent = AbstractAgent._registry[agent_cfg["type"]].from_config(config=agent_config)
# If blue agent is created, add to game.rl_agents
if agent_type == "ProxyAgent":
game.rl_agents[agent_cfg["ref"]] = new_agent
if agent_type == "ProbabilisticAgent": if agent_type == "ProbabilisticAgent":
# TODO: implement non-random agents and fix this parsing # TODO: implement non-random agents and fix this parsing