#2450 remove the need to pass Game to observation objects
This commit is contained in:
@@ -1,7 +1,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from ipaddress import IPv4Address
|
from ipaddress import IPv4Address
|
||||||
from typing import Dict, List, Optional, TYPE_CHECKING
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -10,8 +10,6 @@ from primaite import getLogger
|
|||||||
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
||||||
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
_LOGGER = getLogger(__name__)
|
_LOGGER = getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -167,7 +165,7 @@ class ACLObservation(AbstractObservation, identifier="ACL"):
|
|||||||
)
|
)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> ACLObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> ACLObservation:
|
||||||
"""
|
"""
|
||||||
Create an ACL observation from a configuration schema.
|
Create an ACL observation from a configuration schema.
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Dict, Iterable, List, Optional, TYPE_CHECKING
|
from typing import Dict, Iterable, List, Optional
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -9,8 +9,6 @@ from primaite import getLogger
|
|||||||
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
||||||
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
_LOGGER = getLogger(__name__)
|
_LOGGER = getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -94,7 +92,7 @@ class FileObservation(AbstractObservation, identifier="FILE"):
|
|||||||
return spaces.Dict(space)
|
return spaces.Dict(space)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> FileObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> FileObservation:
|
||||||
"""
|
"""
|
||||||
Create a file observation from a configuration schema.
|
Create a file observation from a configuration schema.
|
||||||
|
|
||||||
@@ -193,7 +191,7 @@ class FolderObservation(AbstractObservation, identifier="FOLDER"):
|
|||||||
return spaces.Dict(shape)
|
return spaces.Dict(shape)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> FolderObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> FolderObservation:
|
||||||
"""
|
"""
|
||||||
Create a folder observation from a configuration schema.
|
Create a folder observation from a configuration schema.
|
||||||
|
|
||||||
@@ -211,5 +209,5 @@ class FolderObservation(AbstractObservation, identifier="FOLDER"):
|
|||||||
for file_config in config.files:
|
for file_config in config.files:
|
||||||
file_config.include_num_access = config.include_num_access
|
file_config.include_num_access = config.include_num_access
|
||||||
|
|
||||||
files = [FileObservation.from_config(config=f, game=game, parent_where=where) for f in config.files]
|
files = [FileObservation.from_config(config=f, parent_where=where) for f in config.files]
|
||||||
return cls(where=where, files=files, num_files=config.num_files, include_num_access=config.include_num_access)
|
return cls(where=where, files=files, num_files=config.num_files, include_num_access=config.include_num_access)
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Dict, List, Optional, TYPE_CHECKING
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -10,8 +10,6 @@ from primaite.game.agent.observations.acl_observation import ACLObservation
|
|||||||
from primaite.game.agent.observations.nic_observations import PortObservation
|
from primaite.game.agent.observations.nic_observations import PortObservation
|
||||||
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
_LOGGER = getLogger(__name__)
|
_LOGGER = getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -200,9 +198,7 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
|
|||||||
return space
|
return space
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> FirewallObservation:
|
||||||
cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []
|
|
||||||
) -> FirewallObservation:
|
|
||||||
"""
|
"""
|
||||||
Create a firewall observation from a configuration schema.
|
Create a firewall observation from a configuration schema.
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Dict, List, Optional, TYPE_CHECKING
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -12,8 +12,6 @@ from primaite.game.agent.observations.observations import AbstractObservation, W
|
|||||||
from primaite.game.agent.observations.software_observation import ApplicationObservation, ServiceObservation
|
from primaite.game.agent.observations.software_observation import ApplicationObservation, ServiceObservation
|
||||||
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
_LOGGER = getLogger(__name__)
|
_LOGGER = getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -201,7 +199,7 @@ class HostObservation(AbstractObservation, identifier="HOST"):
|
|||||||
return spaces.Dict(shape)
|
return spaces.Dict(shape)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> HostObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> HostObservation:
|
||||||
"""
|
"""
|
||||||
Create a host observation from a configuration schema.
|
Create a host observation from a configuration schema.
|
||||||
|
|
||||||
@@ -225,12 +223,10 @@ class HostObservation(AbstractObservation, identifier="HOST"):
|
|||||||
for nic_config in config.network_interfaces:
|
for nic_config in config.network_interfaces:
|
||||||
nic_config.include_nmne = config.include_nmne
|
nic_config.include_nmne = config.include_nmne
|
||||||
|
|
||||||
services = [ServiceObservation.from_config(config=c, game=game, parent_where=where) for c in config.services]
|
services = [ServiceObservation.from_config(config=c, parent_where=where) for c in config.services]
|
||||||
applications = [
|
applications = [ApplicationObservation.from_config(config=c, parent_where=where) for c in config.applications]
|
||||||
ApplicationObservation.from_config(config=c, game=game, parent_where=where) for c in config.applications
|
folders = [FolderObservation.from_config(config=c, parent_where=where) for c in config.folders]
|
||||||
]
|
nics = [NICObservation.from_config(config=c, parent_where=where) for c in config.network_interfaces]
|
||||||
folders = [FolderObservation.from_config(config=c, game=game, parent_where=where) for c in config.folders]
|
|
||||||
nics = [NICObservation.from_config(config=c, game=game, parent_where=where) for c in config.network_interfaces]
|
|
||||||
|
|
||||||
return cls(
|
return cls(
|
||||||
where=where,
|
where=where,
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Any, Dict, List, TYPE_CHECKING
|
from typing import Any, Dict, List
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -9,8 +9,6 @@ from primaite import getLogger
|
|||||||
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
||||||
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
_LOGGER = getLogger(__name__)
|
_LOGGER = getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -68,14 +66,12 @@ class LinkObservation(AbstractObservation, identifier="LINK"):
|
|||||||
return spaces.Dict({"PROTOCOLS": spaces.Dict({"ALL": spaces.Discrete(11)})})
|
return spaces.Dict({"PROTOCOLS": spaces.Dict({"ALL": spaces.Discrete(11)})})
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> LinkObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> LinkObservation:
|
||||||
"""
|
"""
|
||||||
Create a link observation from a configuration schema.
|
Create a link observation from a configuration schema.
|
||||||
|
|
||||||
:param config: Configuration schema containing the necessary information for the link observation.
|
:param config: Configuration schema containing the necessary information for the link observation.
|
||||||
:type config: ConfigSchema
|
:type config: ConfigSchema
|
||||||
:param game: The PrimaiteGame instance.
|
|
||||||
:type game: PrimaiteGame
|
|
||||||
:param parent_where: Where in the simulation state dictionary to find the information about this link.
|
:param parent_where: Where in the simulation state dictionary to find the information about this link.
|
||||||
A typical location might be ['network', 'links', <link_reference>].
|
A typical location might be ['network', 'links', <link_reference>].
|
||||||
:type parent_where: WhereType, optional
|
:type parent_where: WhereType, optional
|
||||||
@@ -135,14 +131,12 @@ class LinksObservation(AbstractObservation, identifier="LINKS"):
|
|||||||
return spaces.Dict({i + 1: l.space for i, l in enumerate(self.links)})
|
return spaces.Dict({i + 1: l.space for i, l in enumerate(self.links)})
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> LinksObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> LinksObservation:
|
||||||
"""
|
"""
|
||||||
Create a links observation from a configuration schema.
|
Create a links observation from a configuration schema.
|
||||||
|
|
||||||
:param config: Configuration schema containing the necessary information for the links observation.
|
:param config: Configuration schema containing the necessary information for the links observation.
|
||||||
:type config: ConfigSchema
|
:type config: ConfigSchema
|
||||||
:param game: The PrimaiteGame instance.
|
|
||||||
:type game: PrimaiteGame
|
|
||||||
:param parent_where: Where in the simulation state dictionary to find the information about these links.
|
:param parent_where: Where in the simulation state dictionary to find the information about these links.
|
||||||
A typical location might be ['network'].
|
A typical location might be ['network'].
|
||||||
:type parent_where: WhereType, optional
|
:type parent_where: WhereType, optional
|
||||||
@@ -151,5 +145,5 @@ class LinksObservation(AbstractObservation, identifier="LINKS"):
|
|||||||
"""
|
"""
|
||||||
where = parent_where + ["network"]
|
where = parent_where + ["network"]
|
||||||
link_cfgs = [LinkObservation.ConfigSchema(link_reference=ref) for ref in config.link_references]
|
link_cfgs = [LinkObservation.ConfigSchema(link_reference=ref) for ref in config.link_references]
|
||||||
links = [LinkObservation.from_config(c, game=game, parent_where=where) for c in link_cfgs]
|
links = [LinkObservation.from_config(c, parent_where=where) for c in link_cfgs]
|
||||||
return cls(where=where, links=links)
|
return cls(where=where, links=links)
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Dict, Optional, TYPE_CHECKING
|
from typing import Dict, Optional
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -8,9 +8,6 @@ from gymnasium.core import ObsType
|
|||||||
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
||||||
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
|
|
||||||
|
|
||||||
class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
|
class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
|
||||||
"""Status information about a network interface within the simulation environment."""
|
"""Status information about a network interface within the simulation environment."""
|
||||||
@@ -119,7 +116,7 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
|
|||||||
return space
|
return space
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> NICObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> NICObservation:
|
||||||
"""
|
"""
|
||||||
Create a network interface observation from a configuration schema.
|
Create a network interface observation from a configuration schema.
|
||||||
|
|
||||||
@@ -179,7 +176,7 @@ class PortObservation(AbstractObservation, identifier="PORT"):
|
|||||||
return spaces.Dict({"operating_status": spaces.Discrete(3)})
|
return spaces.Dict({"operating_status": spaces.Discrete(3)})
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> PortObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> PortObservation:
|
||||||
"""
|
"""
|
||||||
Create a port observation from a configuration schema.
|
Create a port observation from a configuration schema.
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Dict, List, Optional, TYPE_CHECKING
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -12,8 +12,6 @@ from primaite.game.agent.observations.host_observations import HostObservation
|
|||||||
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
||||||
from primaite.game.agent.observations.router_observation import RouterObservation
|
from primaite.game.agent.observations.router_observation import RouterObservation
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
_LOGGER = getLogger(__name__)
|
_LOGGER = getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -152,7 +150,7 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
|
|||||||
return space
|
return space
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> NodesObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> NodesObservation:
|
||||||
"""
|
"""
|
||||||
Create a nodes observation from a configuration schema.
|
Create a nodes observation from a configuration schema.
|
||||||
|
|
||||||
@@ -211,8 +209,8 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
|
|||||||
if firewall_config.num_rules is None:
|
if firewall_config.num_rules is None:
|
||||||
firewall_config.num_rules = config.num_rules
|
firewall_config.num_rules = config.num_rules
|
||||||
|
|
||||||
hosts = [HostObservation.from_config(config=c, game=game, parent_where=where) for c in config.hosts]
|
hosts = [HostObservation.from_config(config=c, parent_where=where) for c in config.hosts]
|
||||||
routers = [RouterObservation.from_config(config=c, game=game, parent_where=where) for c in config.routers]
|
routers = [RouterObservation.from_config(config=c, parent_where=where) for c in config.routers]
|
||||||
firewalls = [FirewallObservation.from_config(config=c, game=game, parent_where=where) for c in config.firewalls]
|
firewalls = [FirewallObservation.from_config(config=c, parent_where=where) for c in config.firewalls]
|
||||||
|
|
||||||
return cls(where=where, hosts=hosts, routers=routers, firewalls=firewalls)
|
return cls(where=where, hosts=hosts, routers=routers, firewalls=firewalls)
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -8,9 +8,6 @@ from pydantic import BaseModel, ConfigDict, model_validator, ValidationError
|
|||||||
|
|
||||||
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
|
|
||||||
|
|
||||||
class NestedObservation(AbstractObservation, identifier="CUSTOM"):
|
class NestedObservation(AbstractObservation, identifier="CUSTOM"):
|
||||||
"""Observation type that allows combining other observations into a gymnasium.spaces.Dict space."""
|
"""Observation type that allows combining other observations into a gymnasium.spaces.Dict space."""
|
||||||
@@ -76,7 +73,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
|
|||||||
return spaces.Dict({label: obs.space for label, obs in self.components.items()})
|
return spaces.Dict({label: obs.space for label, obs in self.components.items()})
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> NestedObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> NestedObservation:
|
||||||
"""
|
"""
|
||||||
Read the Nested observation config and create all defined subcomponents.
|
Read the Nested observation config and create all defined subcomponents.
|
||||||
|
|
||||||
@@ -115,7 +112,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
|
|||||||
instances = dict()
|
instances = dict()
|
||||||
for component in config.components:
|
for component in config.components:
|
||||||
obs_class = AbstractObservation._registry[component.type]
|
obs_class = AbstractObservation._registry[component.type]
|
||||||
obs_instance = obs_class.from_config(config=obs_class.ConfigSchema(**component.options), game=game)
|
obs_instance = obs_class.from_config(config=obs_class.ConfigSchema(**component.options))
|
||||||
instances[component.label] = obs_instance
|
instances[component.label] = obs_instance
|
||||||
return cls(components=instances)
|
return cls(components=instances)
|
||||||
|
|
||||||
@@ -137,9 +134,7 @@ class NullObservation(AbstractObservation, identifier="NONE"):
|
|||||||
return spaces.Discrete(1)
|
return spaces.Discrete(1)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(
|
def from_config(cls, config: NullObservation.ConfigSchema, parent_where: WhereType = []) -> NullObservation:
|
||||||
cls, config: NullObservation.ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []
|
|
||||||
) -> NullObservation:
|
|
||||||
"""Instantiate a NullObservation. Accepts parameters to comply with API."""
|
"""Instantiate a NullObservation. Accepts parameters to comply with API."""
|
||||||
return cls()
|
return cls()
|
||||||
|
|
||||||
@@ -180,7 +175,7 @@ class ObservationManager:
|
|||||||
return self.obs.space
|
return self.obs.space
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: Optional[Dict], game: "PrimaiteGame") -> "ObservationManager":
|
def from_config(cls, config: Optional[Dict]) -> "ObservationManager":
|
||||||
"""
|
"""
|
||||||
Create observation space from a config.
|
Create observation space from a config.
|
||||||
|
|
||||||
@@ -191,14 +186,12 @@ class ObservationManager:
|
|||||||
AbstractObservation
|
AbstractObservation
|
||||||
options: this must adhere to the chosen observation type's ConfigSchema nested class.
|
options: this must adhere to the chosen observation type's ConfigSchema nested class.
|
||||||
:type config: Dict
|
:type config: Dict
|
||||||
:param game: Reference to the PrimaiteGame object that spawned this observation.
|
|
||||||
:type game: PrimaiteGame
|
|
||||||
"""
|
"""
|
||||||
if config is None:
|
if config is None:
|
||||||
return cls(NullObservation())
|
return cls(NullObservation())
|
||||||
print(config)
|
print(config)
|
||||||
obs_type = config["type"]
|
obs_type = config["type"]
|
||||||
obs_class = AbstractObservation._registry[obs_type]
|
obs_class = AbstractObservation._registry[obs_type]
|
||||||
observation = obs_class.from_config(config=obs_class.ConfigSchema(**config["options"]), game=game)
|
observation = obs_class.from_config(config=obs_class.ConfigSchema(**config["options"]))
|
||||||
obs_manager = cls(observation)
|
obs_manager = cls(observation)
|
||||||
return obs_manager
|
return obs_manager
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
"""Manages the observation space for the agent."""
|
"""Manages the observation space for the agent."""
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from typing import Any, Dict, Iterable, Type, TYPE_CHECKING
|
from typing import Any, Dict, Iterable, Type
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -8,8 +8,6 @@ from pydantic import BaseModel, ConfigDict
|
|||||||
|
|
||||||
from primaite import getLogger
|
from primaite import getLogger
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
_LOGGER = getLogger(__name__)
|
_LOGGER = getLogger(__name__)
|
||||||
WhereType = Iterable[str | int] | None
|
WhereType = Iterable[str | int] | None
|
||||||
|
|
||||||
@@ -65,8 +63,6 @@ class AbstractObservation(ABC):
|
|||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def from_config(
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> "AbstractObservation":
|
||||||
cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []
|
|
||||||
) -> "AbstractObservation":
|
|
||||||
"""Create this observation space component form a serialised format."""
|
"""Create this observation space component form a serialised format."""
|
||||||
return cls()
|
return cls()
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Dict, List, Optional, TYPE_CHECKING
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -11,8 +11,6 @@ from primaite.game.agent.observations.nic_observations import PortObservation
|
|||||||
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
||||||
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
_LOGGER = getLogger(__name__)
|
_LOGGER = getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -112,7 +110,7 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
|
|||||||
return spaces.Dict(shape)
|
return spaces.Dict(shape)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []) -> RouterObservation:
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> RouterObservation:
|
||||||
"""
|
"""
|
||||||
Create a router observation from a configuration schema.
|
Create a router observation from a configuration schema.
|
||||||
|
|
||||||
@@ -142,6 +140,6 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
|
|||||||
if config.ports is None:
|
if config.ports is None:
|
||||||
config.ports = [PortObservation.ConfigSchema(port_id=i + 1) for i in range(config.num_ports)]
|
config.ports = [PortObservation.ConfigSchema(port_id=i + 1) for i in range(config.num_ports)]
|
||||||
|
|
||||||
ports = [PortObservation.from_config(config=c, game=game, parent_where=where) for c in config.ports]
|
ports = [PortObservation.from_config(config=c, parent_where=where) for c in config.ports]
|
||||||
acl = ACLObservation.from_config(config=config.acl, game=game, parent_where=where)
|
acl = ACLObservation.from_config(config=config.acl, parent_where=where)
|
||||||
return cls(where=where, ports=ports, num_ports=config.num_ports, acl=acl)
|
return cls(where=where, ports=ports, num_ports=config.num_ports, acl=acl)
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Dict, TYPE_CHECKING
|
from typing import Dict
|
||||||
|
|
||||||
from gymnasium import spaces
|
from gymnasium import spaces
|
||||||
from gymnasium.core import ObsType
|
from gymnasium.core import ObsType
|
||||||
@@ -8,9 +8,6 @@ from gymnasium.core import ObsType
|
|||||||
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
|
||||||
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from primaite.game.game import PrimaiteGame
|
|
||||||
|
|
||||||
|
|
||||||
class ServiceObservation(AbstractObservation, identifier="SERVICE"):
|
class ServiceObservation(AbstractObservation, identifier="SERVICE"):
|
||||||
"""Service observation, shows status of a service in the simulation environment."""
|
"""Service observation, shows status of a service in the simulation environment."""
|
||||||
@@ -60,9 +57,7 @@ class ServiceObservation(AbstractObservation, identifier="SERVICE"):
|
|||||||
return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(5)})
|
return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(5)})
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> ServiceObservation:
|
||||||
cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []
|
|
||||||
) -> ServiceObservation:
|
|
||||||
"""
|
"""
|
||||||
Create a service observation from a configuration schema.
|
Create a service observation from a configuration schema.
|
||||||
|
|
||||||
@@ -153,9 +148,7 @@ class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
|
|||||||
)
|
)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_config(
|
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> ApplicationObservation:
|
||||||
cls, config: ConfigSchema, game: "PrimaiteGame", parent_where: WhereType = []
|
|
||||||
) -> ApplicationObservation:
|
|
||||||
"""
|
"""
|
||||||
Create an application observation from a configuration schema.
|
Create an application observation from a configuration schema.
|
||||||
|
|
||||||
|
|||||||
@@ -396,8 +396,8 @@ class PrimaiteGame:
|
|||||||
|
|
||||||
# 2. create links between nodes
|
# 2. create links between nodes
|
||||||
for link_cfg in links_cfg:
|
for link_cfg in links_cfg:
|
||||||
node_a = net.nodes[game.ref_map_nodes[link_cfg["endpoint_a_ref"]]]
|
node_a = net.get_node_by_hostname(link_cfg["endpoint_a_ref"])
|
||||||
node_b = net.nodes[game.ref_map_nodes[link_cfg["endpoint_b_ref"]]]
|
node_b = net.get_node_by_hostname(link_cfg["endpoint_b_ref"])
|
||||||
if isinstance(node_a, Switch):
|
if isinstance(node_a, Switch):
|
||||||
endpoint_a = node_a.network_interface[link_cfg["endpoint_a_port"]]
|
endpoint_a = node_a.network_interface[link_cfg["endpoint_a_port"]]
|
||||||
else:
|
else:
|
||||||
@@ -419,7 +419,7 @@ class PrimaiteGame:
|
|||||||
reward_function_cfg = agent_cfg["reward_function"]
|
reward_function_cfg = agent_cfg["reward_function"]
|
||||||
|
|
||||||
# CREATE OBSERVATION SPACE
|
# CREATE OBSERVATION SPACE
|
||||||
obs_space = ObservationManager.from_config(observation_space_cfg, game)
|
obs_space = ObservationManager.from_config(observation_space_cfg)
|
||||||
|
|
||||||
# CREATE ACTION SPACE
|
# CREATE ACTION SPACE
|
||||||
action_space = ActionManager.from_config(game, action_space_cfg)
|
action_space = ActionManager.from_config(game, action_space_cfg)
|
||||||
|
|||||||
Reference in New Issue
Block a user