From b4b6c16872622e79176849f6d43ac1bc8863c434 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 15 Jan 2025 10:08:14 +0000 Subject: [PATCH] #2869 - Make observation and action managers use config schemas --- src/primaite/game/agent/actions/manager.py | 49 ++++++------- src/primaite/game/agent/interface.py | 52 +++++++------- .../agent/observations/observation_manager.py | 72 ++++++++++++++++--- 3 files changed, 115 insertions(+), 58 deletions(-) diff --git a/src/primaite/game/agent/actions/manager.py b/src/primaite/game/agent/actions/manager.py index 400d30e4..0f7db2f3 100644 --- a/src/primaite/game/agent/actions/manager.py +++ b/src/primaite/game/agent/actions/manager.py @@ -13,9 +13,10 @@ agents: from __future__ import annotations -from typing import Dict, Optional, Tuple +from typing import Dict, Tuple from gymnasium import spaces +from pydantic import BaseModel, ConfigDict, Field, field_validator # from primaite.game.game import PrimaiteGame # TODO: Breaks things from primaite.game.agent.actions.abstract import AbstractAction @@ -38,35 +39,35 @@ class DoNothingAction(AbstractAction, identifier="do_nothing"): return ["do_nothing"] -class ActionManager: +class ActionManager(BaseModel): """Class which manages the action space for an agent.""" - def __init__(self, act_map: Optional[Dict[int, Dict]] = None) -> None: - """Init method for ActionManager. + class ConfigSchema(BaseModel): + """Config Schema for ActionManager.""" - :param act_map: Action map which maps integers to actions. Used for restricting the set of possible actions. - :type act_map: Optional[Dict[int, Dict]] - """ - self.action_map: Dict[int, Tuple[str, Dict]] = {} - """ - Action mapping that converts an integer to a specific action and parameter choice. + model_config = ConfigDict(extra="forbid") + action_map: Dict[int, Tuple[str, Dict]] = {} + """Mapping between integer action choices and CAOS actions.""" - For example : - {0: ("node_service_scan", {node_name:"client_1", service_name:"WebBrowser"})} - """ - # allows restricting set of possible actions - TODO: Refactor to be a list? - if act_map is None: - # raise RuntimeError("Action map must be specified in the config file.") - pass - else: - self.action_map = {i: (a["action"], a["options"]) for i, a in act_map.items()} - # make sure all numbers between 0 and N are represented as dict keys in action map - assert all([i in self.action_map.keys() for i in range(len(self.action_map))]) + @field_validator("action_map", mode="after") + def consecutive_action_nums(cls, v: Dict) -> Dict: + """Make sure all numbers between 0 and N are represented as dict keys in action map.""" + assert all([i in v.keys() for i in range(len(v))]) + + config: ActionManager.ConfigSchema = Field(default_factory=lambda: ActionManager.ConfigSchema()) + + @property + def action_map(self) -> Dict[int, Tuple[str, Dict]]: + """Convenience method for accessing the action map.""" + return self.config.action_map def get_action(self, action: int) -> Tuple[str, Dict]: - """Produce action in CAOS format.""" - """the agent chooses an action (as an integer), this is converted into an action in CAOS format""" - """The CAOS format is basically a action identifier, followed by parameters stored in a dictionary""" + """ + Produce action in CAOS format. + + The agent chooses an action (as an integer), this is converted into an action in CAOS format + The CAOS format is basically an action identifier, followed by parameters stored in a dictionary. + """ act_identifier, act_options = self.action_map[action] return act_identifier, act_options diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 9d8f3f63..0b55c1db 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -2,7 +2,7 @@ """Interface for agents.""" from __future__ import annotations -from abc import abstractmethod +from abc import ABC, abstractmethod from typing import Any, ClassVar, Dict, List, Optional, Tuple, Type, TYPE_CHECKING from gymnasium.core import ActType, ObsType @@ -43,35 +43,31 @@ class AgentHistoryItem(BaseModel): reward_info: Dict[str, Any] = {} -class AbstractAgent(BaseModel): +class AbstractAgent(BaseModel, ABC): """Base class for scripted and RL agents.""" - _registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {} - logger: AgentLog = AgentLog(agent_name="Abstract_Agent") - - history: List[AgentHistoryItem] = [] - config: "AbstractAgent.ConfigSchema" = Field(default_factory=lambda: AbstractAgent.ConfigSchema()) - action_manager: "ActionManager" - observation_manager: "ObservationManager" - reward_function: "RewardFunction" model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) - class ConfigSchema(BaseModel): - """ - Configuration Schema for AbstractAgents. - - :param type: Type of agent being generated. - :type type: str - :param observation_space: Observation space for the agent. - :type observation_space: Optional[ObservationSpace] - :param reward_function: Reward function for the agent. - :type reward_function: Optional[RewardFunction] - :param agent_settings: Configurable Options for Abstracted Agents. - :type agent_settings: Optional[AgentSettings] - """ + class ConfigSchema(BaseModel, ABC): + """Configuration Schema for AbstractAgents.""" model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) type: str = "AbstractAgent" + action_space: ActionManager.ConfigSchema = Field(default_factory=lambda: ActionManager.ConfigSchema()) + observation_space: ObservationManager.ConfigSchema = Field( + default_factory=lambda: ObservationManager.ConfigSchema() + ) + reward_function: RewardFunction.ConfigSchema = Field(default_factory=lambda: RewardFunction.ConfigSchema()) + + config: "AbstractAgent.ConfigSchema" = Field(default_factory=lambda: AbstractAgent.ConfigSchema()) + + logger: AgentLog = AgentLog(agent_name="Abstract_Agent") + history: List[AgentHistoryItem] = [] + action_manager: "ActionManager" + observation_manager: "ObservationManager" + reward_function: "RewardFunction" + + _registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {} def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) @@ -81,6 +77,14 @@ class AbstractAgent(BaseModel): raise ValueError(f"Cannot create a new agent under reserved name {identifier}") cls._registry[identifier] = cls + def __init__(self, config: ConfigSchema, **kwargs): + kwargs["action_manager"] = kwargs.get("action_manager") or ActionManager.from_config(config.action_space) + kwargs["observation_manager"] = kwargs.get("observation_manager") or ObservationManager( + config.observation_space + ) + kwargs["reward_function"] = kwargs.get("reward_function") or RewardFunction.from_config(config.reward_function) + super().__init__(config=config, **kwargs) + @property def flatten_obs(self) -> bool: """Return agent flatten_obs param.""" @@ -94,7 +98,7 @@ class AbstractAgent(BaseModel): obj = cls( config=cls.ConfigSchema(**config["agent_settings"]), action_manager=ActionManager.from_config(config["action_space"]), - observation_manager=ObservationManager.from_config(config["observation_space"]), + observation_manager=ObservationManager(config["observation_space"]), reward_function=RewardFunction.from_config(config["reward_function"]), ) return obj diff --git a/src/primaite/game/agent/observations/observation_manager.py b/src/primaite/game/agent/observations/observation_manager.py index 71a60433..6964ce2c 100644 --- a/src/primaite/game/agent/observations/observation_manager.py +++ b/src/primaite/game/agent/observations/observation_manager.py @@ -1,11 +1,12 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK from __future__ import annotations +from functools import cached_property from typing import Any, Dict, List, Optional from gymnasium import spaces from gymnasium.core import ObsType -from pydantic import BaseModel, ConfigDict, model_validator, ValidationError +from pydantic import BaseModel, computed_field, ConfigDict, Field, model_validator, ValidationError from primaite.game.agent.observations.observations import AbstractObservation, WhereType @@ -140,7 +141,7 @@ class NullObservation(AbstractObservation, identifier="NONE"): return cls() -class ObservationManager: +class ObservationManager(BaseModel): """ Manage the observations of an Agent. @@ -150,15 +151,66 @@ class ObservationManager: 3. Formatting this information so an agent can use it to make decisions. """ - def __init__(self, obs: AbstractObservation) -> None: - """Initialise observation space. + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) - :param observation: Observation object - :type observation: AbstractObservation - """ - self.obs: AbstractObservation = obs - self.current_observation: ObsType - """Cached copy of the observation at the time it was most recently calculated.""" + class ConfigSchema(BaseModel): + """Config Schema for Observation Manager.""" + + model_config = ConfigDict(extra="forbid") + type: str = "NONE" + """Identifier name for the top-level observation.""" + options: AbstractObservation.ConfigSchema = Field( + default_factory=lambda: NullObservation.ConfigSchema(), validate_default=True + ) + """Options to pass into the top-level observation during creation.""" + + @model_validator(mode="before") + @classmethod + def resolve_obs_options_type(cls, data: Any) -> Any: + """ + When constructing the model from a dict, resolve the correct observation class based on `type` field. + + Workaround: The `options` field is statically typed as AbstractObservation. Therefore, it falls over when + passing in data that adheres to a subclass schema rather than the plain AbstractObservation schema. There is + a way to do this properly using discriminated union, but most advice on the internet assumes that the full + list of types between which to discriminate is known ahead-of-time. That is not the case for us, because of + our plugin architecture. + + We may be able to revisit and implement a better solution when needed using the following resources as + research starting points: + https://docs.pydantic.dev/latest/concepts/unions/#discriminated-unions + https://github.com/pydantic/pydantic/issues/7366 + https://github.com/pydantic/pydantic/issues/7462 + https://github.com/pydantic/pydantic/pull/7983 + """ + if not isinstance(data, dict): + return data + + # (TODO: duplicate default definition between here and the actual model) + obs_type = data["type"] if "type" in data else "NONE" + obs_class = AbstractObservation._registry[obs_type] + + # if no options are passed in, try to create a default schema. Only works if there are no mandatory fields + if "options" not in data: + data["options"] = obs_class.ConfigSchema() + + # if options passed as a dict, convert to a schema + elif isinstance(data["options"], dict): + data["options"] = obs_class.ConfigSchema(**data["options"]) + + return data + + config: ConfigSchema = Field(default_factory=lambda: ObservationManager.ConfigSchema()) + + current_observation: ObsType = 0 + + @computed_field + @cached_property + def obs(self) -> AbstractObservation: + """Create the main observation component for the observation manager from the config.""" + obs_class = AbstractObservation._registry[self.config.type] + obs_instance = obs_class.from_config(config=self.config.options) + return obs_instance def update(self, state: Dict) -> Dict: """