#2869 - Updates to address test failures. Updated YAML configs to remove redundant start_settings
This commit is contained in:
@@ -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]:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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]:
|
||||
"""
|
||||
|
||||
@@ -538,6 +538,7 @@ class PrimaiteGame:
|
||||
"observation_manager": observation_space_cfg,
|
||||
"reward_function": reward_function_cfg,
|
||||
"agent_settings": agent_settings,
|
||||
"game": game,
|
||||
}
|
||||
|
||||
# CREATE AGENT
|
||||
|
||||
Reference in New Issue
Block a user