diff --git a/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml b/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml index dfd200f3..8c83bf79 100644 --- a/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml +++ b/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml @@ -55,50 +55,50 @@ agents: action_space: action_list: - - type: DONOTHING - - type: NODE_SHUTDOWN - - type: NODE_STARTUP - - type: HOST_NIC_ENABLE - - type: HOST_NIC_DISABLE + - type: do_nothing + - type: node_shutdown + - type: node_startup + - type: host_nic_enable + - type: host_nic_enable action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: NODE_SHUTDOWN + action: node_shutdown options: - node_id: 0 + node_name: client_1 2: - action: NODE_SHUTDOWN + action: node_shutdown options: - node_id: 1 + node_name: server 3: - action: NODE_STARTUP + action: node_startup options: - node_id: 0 + node_name: client_1 4: - action: NODE_STARTUP + action: node_startup options: - node_id: 1 + node_name: server 5: - action: HOST_NIC_DISABLE + action: host_nic_disable options: - node_id: 0 + node_name: client_1 nic_id: 0 6: - action: HOST_NIC_DISABLE + action: host_nic_disable options: - node_id: 1 + node_name: server nic_id: 0 7: - action: HOST_NIC_ENABLE + action: host_nic_enable options: - node_id: 0 + node_name: client_1 nic_id: 0 8: - action: HOST_NIC_ENABLE + action: host_nic_enable options: - node_id: 1 + node_name: server nic_id: 0 options: nodes: diff --git a/src/primaite/game/agent/actions/__init__.py b/src/primaite/game/agent/actions/__init__.py index 428c6c58..3540b128 100644 --- a/src/primaite/game/agent/actions/__init__.py +++ b/src/primaite/game/agent/actions/__init__.py @@ -1,6 +1,7 @@ # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK from primaite.game.agent.actions import ( + abstract, acl, application, config, @@ -17,6 +18,7 @@ from primaite.game.agent.actions.manager import ActionManager __all__ = ( "acl", + "abstract", "application", "config", "file", diff --git a/src/primaite/game/agent/actions/abstract.py b/src/primaite/game/agent/actions/abstract.py new file mode 100644 index 00000000..5250e532 --- /dev/null +++ b/src/primaite/game/agent/actions/abstract.py @@ -0,0 +1,48 @@ +# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK +from __future__ import annotations +from abc import ABC +from typing import Any, ClassVar, Dict, Type + +from pydantic import BaseModel, ConfigDict + +from primaite.interface.request import RequestFormat + +class AbstractAction(BaseModel): + """Base class for actions.""" + + # notes: + # we actually don't need to hold any state in actions, so there's no need to define any __init__ logic. + # all the init methods in the old actions are just used for holding a verb and shape, which are not really used. + # the config schema should be used to the actual parameters for formatting the action itself. + # (therefore there's no need for creating action instances, just the action class contains logic for converting + # CAOS actions to requests for simulator. Similar to the network node adder, that class also doesn't need to be + # instantiated.) + class ConfigSchema(BaseModel, ABC): + """Base configuration schema for Actions.""" + + model_config = ConfigDict(extra="forbid") + type: str + + _registry: ClassVar[Dict[str, Type[AbstractAction]]] = {} + + def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if identifier in cls._registry: + raise ValueError(f"Cannot create new action under reserved name {identifier}") + cls._registry[identifier] = cls + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + return [] + + @classmethod + def from_config(cls, config: Dict) -> "AbstractAction": + """Create an action component from a config dictionary""" + + type_id = config.get("type") + + if type_id in cls._registry: + return cls(type=type_id, model_config=config) + else: + return [] \ No newline at end of file diff --git a/src/primaite/game/agent/actions/application.py b/src/primaite/game/agent/actions/application.py index 3a254d57..e0496dc7 100644 --- a/src/primaite/game/agent/actions/application.py +++ b/src/primaite/game/agent/actions/application.py @@ -1,7 +1,7 @@ # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK from typing import ClassVar -from primaite.game.agent.actions.manager import AbstractAction +from primaite.game.agent.actions.abstract import AbstractAction from primaite.interface.request import RequestFormat __all__ = ( diff --git a/src/primaite/game/agent/actions/manager.py b/src/primaite/game/agent/actions/manager.py index 7677b39a..98ffbd00 100644 --- a/src/primaite/game/agent/actions/manager.py +++ b/src/primaite/game/agent/actions/manager.py @@ -14,48 +14,17 @@ agents: from __future__ import annotations import itertools -from abc import ABC -from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Type +from typing import Dict, List, Literal, Optional, Tuple from gymnasium import spaces -from pydantic import BaseModel, ConfigDict # from primaite.game.game import PrimaiteGame # TODO: Breaks things +from primaite.game.agent.actions.abstract import AbstractAction from primaite.interface.request import RequestFormat # TODO: Make sure that actions are backwards compatible where the old YAML format is used. -class AbstractAction(BaseModel): - """Base class for actions.""" - - # notes: - # we actually don't need to hold any state in actions, so there's no need to define any __init__ logic. - # all the init methods in the old actions are just used for holding a verb and shape, which are not really used. - # the config schema should be used to the actual parameters for formatting the action itself. - # (therefore there's no need for creating action instances, just the action class contains logic for converting - # CAOS actions to requests for simulator. Similar to the network node adder, that class also doesn't need to be - # instantiated.) - class ConfigSchema(BaseModel, ABC): # TODO: not sure if this better named something like `Options` - """Base configuration schema for Actions.""" - - model_config = ConfigDict(extra="forbid") - type: str - - _registry: ClassVar[Dict[str, Type[AbstractAction]]] = {} - - def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None: - super().__init_subclass__(**kwargs) - if identifier in cls._registry: - raise ValueError(f"Cannot create new action under reserved name {identifier}") - cls._registry[identifier] = cls - - @classmethod - def form_request(self, config: ConfigSchema) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - return [] - - class DoNothingAction(AbstractAction, identifier="do_nothing"): """Do Nothing Action.""" @@ -204,8 +173,8 @@ class ActionManager: # where `type` decides which AbstractAction subclass should be used # and `options` is an optional dict of options to pass to the init method of the action class act_type = act_spec.get("type") - act_options = act_spec.get("options", {}) - # self.actions[act_type] = self.act_class_identifiers[act_type](self, **global_action_args, **act_options) + # act_options = act_spec.get("options", {}) # Don't need this anymore I think? + self.actions[act_type] = AbstractAction._registry[act_type] self.action_map: Dict[int, Tuple[str, Dict]] = {} """ @@ -235,10 +204,10 @@ class ActionManager: :return: An action map maps consecutive integers to a combination of Action type and parameter choices. An example output could be: - {0: ("DONOTHING", {'dummy': 0}), - 1: ("NODE_OS_SCAN", {'node_id': 0}), - 2: ("NODE_OS_SCAN", {'node_id': 1}), - 3: ("NODE_FOLDER_SCAN", {'node_id:0, folder_id:0}), + {0: ("do_nothing", {'dummy': 0}), + 1: ("node_os_scan", {'node_name': computer}), + 2: ("node_os_scan", {'node_name': server}), + 3: ("node_folder_scan", {'node_name:computer, folder_name:downloads}), ... #etc... } :rtype: Dict[int, Tuple[AbstractAction, Dict]] @@ -269,7 +238,7 @@ class ActionManager: def form_request(self, action_identifier: str, action_options: Dict) -> RequestFormat: """Take action in CAOS format and use the execution definition to change it into PrimAITE request format.""" act_obj = self.actions[action_identifier] - return act_obj.form_request(**action_options) + return act_obj.form_request(action_options) @property def space(self) -> spaces.Space: diff --git a/src/primaite/notebooks/Action-masking.ipynb b/src/primaite/notebooks/Action-masking.ipynb index d22e171d..858b4bb6 100644 --- a/src/primaite/notebooks/Action-masking.ipynb +++ b/src/primaite/notebooks/Action-masking.ipynb @@ -19,8 +19,7 @@ "source": [ "from primaite.session.environment import PrimaiteGymEnv\n", "from primaite.config.load import data_manipulation_config_path\n", - "from prettytable import PrettyTable\n", - "UDP=\"UDP\"" + "from prettytable import PrettyTable" ] }, {