#2869 - Make observation and action managers use config schemas

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
 

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

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:
         """