Fix certain observation bugs
This commit is contained in:
Marek Wolan
@@ -3,7 +3,6 @@ from typing import Any, Dict, Hashable, List, Optional, TYPE_CHECKING
|
||||
|
||||
from gym import spaces
|
||||
from pydantic import BaseModel
|
||||
from primaite.game.session import PrimaiteSession
|
||||
|
||||
from primaite.simulator.sim_container import Simulation
|
||||
if TYPE_CHECKING:
|
||||
@@ -56,8 +55,8 @@ class AbstractObservation(ABC):
|
||||
"""Subclasses must define the shape that they expect"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def from_config(cls, config:Dict, session:"PrimaiteSession"):
|
||||
"""Create this observation space component form a serialised format.
|
||||
|
||||
@@ -132,7 +131,7 @@ class ServiceObservation(AbstractObservation):
|
||||
return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(6)})
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict, session: PrimaiteSession, parent_where:Optional[List[str]]=None):
|
||||
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]=None):
|
||||
return cls(where=parent_where+["services",session.ref_map_services[config['service_ref']]])
|
||||
|
||||
|
||||
@@ -235,7 +234,7 @@ class FolderObservation(AbstractObservation):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict, session: PrimaiteSession, parent_where:Optional[List[str]]):
|
||||
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]):
|
||||
where = parent_where + ["folders", config['folder_name']]
|
||||
|
||||
file_configs = config["files"]
|
||||
@@ -324,7 +323,6 @@ class NodeObservation(AbstractObservation):
|
||||
return self.default_observation
|
||||
|
||||
obs = {}
|
||||
|
||||
obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)}
|
||||
obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)}
|
||||
obs["operating_status"] = node_state["operating_state"]
|
||||
@@ -349,7 +347,7 @@ class NodeObservation(AbstractObservation):
|
||||
return spaces.Dict(space_shape)
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]= None):
|
||||
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]= None) -> "NodeObservation":
|
||||
node_uuid = session.ref_map_nodes[config['node_ref']]
|
||||
if parent_where is None:
|
||||
where = ["network", "nodes", node_uuid]
|
||||
@@ -361,12 +359,10 @@ class NodeObservation(AbstractObservation):
|
||||
folder_configs = config.get('folders', {})
|
||||
folders = [FolderObservation.from_config(config=c,session=session, parent_where=where) for c in folder_configs]
|
||||
nic_uuids = session.simulation.network.nodes[node_uuid].nics.keys()
|
||||
nic_configs = [{'nic_uuid':n for n in nic_uuids }]
|
||||
nic_configs = [{'nic_uuid':n for n in nic_uuids }] if nic_uuids else []
|
||||
nics = [NicObservation.from_config(config=c, session=session, parent_where=where) for c in nic_configs]
|
||||
logon_status = config.get('logon_status',False)
|
||||
cls(where=where, services=services, folders=folders, nics=nics, logon_status=logon_status)
|
||||
return super().from_config(config, session)
|
||||
|
||||
return cls(where=where, services=services, folders=folders, nics=nics, logon_status=logon_status)
|
||||
|
||||
|
||||
class AclObservation(AbstractObservation):
|
||||
@@ -486,7 +482,7 @@ class NullObservation(AbstractObservation):
|
||||
return spaces.Dict({})
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, cfg:Dict) -> "NullObservation":
|
||||
def from_config(cls, config:Dict, session:Optional["PrimaiteSession"]=None) -> "NullObservation":
|
||||
return cls()
|
||||
|
||||
class ICSObservation(NullObservation): pass
|
||||
@@ -539,15 +535,15 @@ class UC2BlueObservation(AbstractObservation):
|
||||
})
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config:Dict, sess:"PrimaiteSession"):
|
||||
def from_config(cls, config:Dict, session:"PrimaiteSession"):
|
||||
node_configs = config["nodes"]
|
||||
nodes = [NodeObservation.from_config(n) for n in node_configs]
|
||||
nodes = [NodeObservation.from_config(config=n, session=session) for n in node_configs]
|
||||
|
||||
link_configs = config["links"]
|
||||
links = [LinkObservation.from_config(l) for l in link_configs]
|
||||
links = [LinkObservation.from_config(config=l, session=session) for l in link_configs]
|
||||
|
||||
acl_config = config["acl"]
|
||||
acl = AclObservation.from_config(acl_config)
|
||||
acl = AclObservation.from_config(config=acl_config, session=session)
|
||||
|
||||
ics_config = config["ics"]
|
||||
ics = ICSObservation.from_config(ics_config)
|
||||
@@ -561,19 +557,30 @@ class UC2RedObservation(AbstractObservation):
|
||||
self.where:Optional[List[str]] = where
|
||||
self.nodes: List[NodeObservation] = nodes
|
||||
|
||||
self.default_observation=...#TODO
|
||||
self.default_observation : Dict = {
|
||||
"NODES": {i+1: n.default_observation for i,n in enumerate(self.nodes)},
|
||||
}
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
return super().observe(state)
|
||||
def observe(self, state: Dict) -> Dict:
|
||||
if self.where is None:
|
||||
return self.default_observation
|
||||
|
||||
obs = {}
|
||||
obs['NODES'] = {i+1: node.observe(state) for i, node in enumerate(self.nodes)}
|
||||
return obs
|
||||
|
||||
@property
|
||||
def space(self) -> spaces.Space:
|
||||
... #TODO
|
||||
return spaces.Dict({
|
||||
"NODES": spaces.Dict({i+1: node.space for i, node in enumerate(self.nodes)}),
|
||||
})
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict, sim:Simulation):
|
||||
def from_config(cls, config: Dict, session: "PrimaiteSession"):
|
||||
node_configs = config["nodes"]
|
||||
nodes = [NodeObservation.from_config(config=cfg, session=session) for cfg in node_configs]
|
||||
return cls(nodes=nodes, where=["network"])
|
||||
|
||||
... #TODO
|
||||
|
||||
class UC2GreenObservation(NullObservation): pass
|
||||
|
||||
@@ -605,10 +612,10 @@ class ObservationSpace:
|
||||
@classmethod
|
||||
def from_config(cls, config:Dict, session:"PrimaiteSession") -> "ObservationSpace":
|
||||
if config['type'] == "UC2BlueObservation":
|
||||
return cls(UC2BlueObservation(config['options']))
|
||||
return cls(UC2BlueObservation.from_config(config.get('options',{}), session=session))
|
||||
elif config['type'] == "UC2RedObservation":
|
||||
return cls(UC2RedObservation(config['options']))
|
||||
return cls(UC2RedObservation.from_config(config.get('options',{}), session=session))
|
||||
elif config['type'] == "UC2GreenObservation":
|
||||
return cls(UC2GreenObservation(config["options"]))
|
||||
return cls(UC2GreenObservation.from_config(config.get("options",{}), session=session))
|
||||
else:
|
||||
raise ValueError("Observation space type invalid")
|
||||
|
||||
@@ -21,6 +21,7 @@ from primaite.game.agent.observations import (
|
||||
NicObservation,
|
||||
NodeObservation,
|
||||
NullObservation,
|
||||
ObservationSpace,
|
||||
ServiceObservation,
|
||||
UC2BlueObservation,
|
||||
UC2GreenObservation,
|
||||
@@ -177,7 +178,7 @@ class PrimaiteSession:
|
||||
if service_type in service_types_mapping:
|
||||
new_node.software_manager.install(service_types_mapping[service_type])
|
||||
new_service = new_node.software_manager.software[service_type]
|
||||
ref_map_services[service_ref] = new_service
|
||||
sess.ref_map_services[service_ref] = new_service
|
||||
else:
|
||||
print(f"service type not found {service_type}")
|
||||
# service-dependent options
|
||||
@@ -198,12 +199,12 @@ class PrimaiteSession:
|
||||
|
||||
net.add_node(new_node)
|
||||
new_node.power_on()
|
||||
ref_map_nodes[node_ref] = new_node.uuid
|
||||
sess.ref_map_nodes[node_ref] = new_node.uuid
|
||||
|
||||
# 2. create links between nodes
|
||||
for link_cfg in links_cfg:
|
||||
node_a = net.nodes[ref_map_nodes[link_cfg["endpoint_a_ref"]]]
|
||||
node_b = net.nodes[ref_map_nodes[link_cfg["endpoint_b_ref"]]]
|
||||
node_a = net.nodes[sess.ref_map_nodes[link_cfg["endpoint_a_ref"]]]
|
||||
node_b = net.nodes[sess.ref_map_nodes[link_cfg["endpoint_b_ref"]]]
|
||||
if isinstance(node_a, Switch):
|
||||
endpoint_a = node_a.switch_ports[link_cfg["endpoint_a_port"]]
|
||||
else:
|
||||
@@ -213,7 +214,7 @@ class PrimaiteSession:
|
||||
else:
|
||||
endpoint_b = node_b.ethernet_port[link_cfg["endpoint_b_port"]]
|
||||
new_link = net.connect(endpoint_a=endpoint_a, endpoint_b=endpoint_b)
|
||||
ref_map_links[link_cfg["ref"]] = new_link.uuid
|
||||
sess.ref_map_links[link_cfg["ref"]] = new_link.uuid
|
||||
|
||||
# 3. create agents
|
||||
game_cfg = cfg["game_config"]
|
||||
@@ -229,108 +230,112 @@ class PrimaiteSession:
|
||||
reward_function_cfg = agent_cfg["reward_function"]
|
||||
|
||||
# CREATE OBSERVATION SPACE
|
||||
if observation_space_cfg is None:
|
||||
obs_space = NullObservation()
|
||||
elif observation_space_cfg["type"] == "UC2BlueObservation":
|
||||
node_obs_list = []
|
||||
link_obs_list = []
|
||||
obs_space=ObservationSpace.from_config(observation_space_cfg, sess)
|
||||
|
||||
# node ip to index maps ip addresses to node id, as there are potentially multiple nics on a node, there are multiple ip addresses
|
||||
node_ip_to_index = {}
|
||||
for node_idx, node_cfg in enumerate(nodes_cfg):
|
||||
n_ref = node_cfg["ref"]
|
||||
n_obj = net.nodes[ref_map_nodes[n_ref]]
|
||||
for nic_uuid, nic_obj in n_obj.nics.items():
|
||||
node_ip_to_index[nic_obj.ip_address] = node_idx + 2
|
||||
"""
|
||||
# if observation_space_cfg is None:
|
||||
# obs_space = NullObservation()
|
||||
# elif observation_space_cfg["type"] == "UC2BlueObservation":
|
||||
# node_obs_list = []
|
||||
# link_obs_list = []
|
||||
|
||||
for node_obs_cfg in observation_space_cfg["options"]["nodes"]:
|
||||
node_ref = node_obs_cfg["node_ref"]
|
||||
folder_obs_list = []
|
||||
service_obs_list = []
|
||||
if "services" in node_obs_cfg:
|
||||
for service_obs_cfg in node_obs_cfg["services"]:
|
||||
service_obs_list.append(
|
||||
ServiceObservation(
|
||||
where=[
|
||||
"network",
|
||||
"nodes",
|
||||
ref_map_nodes[node_ref],
|
||||
"services",
|
||||
ref_map_services[service_obs_cfg["service_ref"]],
|
||||
]
|
||||
)
|
||||
)
|
||||
if "folders" in node_obs_cfg:
|
||||
for folder_obs_cfg in node_obs_cfg["folders"]:
|
||||
file_obs_list = []
|
||||
if "files" in folder_obs_cfg:
|
||||
for file_obs_cfg in folder_obs_cfg["files"]:
|
||||
file_obs_list.append(
|
||||
FileObservation(
|
||||
where=[
|
||||
"network",
|
||||
"nodes",
|
||||
ref_map_nodes[node_ref],
|
||||
"folders",
|
||||
folder_obs_cfg["folder_name"],
|
||||
"files",
|
||||
file_obs_cfg["file_name"],
|
||||
]
|
||||
)
|
||||
)
|
||||
folder_obs_list.append(
|
||||
FolderObservation(
|
||||
where=[
|
||||
"network",
|
||||
"nodes",
|
||||
ref_map_nodes[node_ref],
|
||||
"folders",
|
||||
folder_obs_cfg["folder_name"],
|
||||
],
|
||||
files=file_obs_list,
|
||||
)
|
||||
)
|
||||
nic_obs_list = []
|
||||
for nic_uuid in net.nodes[ref_map_nodes[node_obs_cfg["node_ref"]]].nics.keys():
|
||||
nic_obs_list.append(
|
||||
NicObservation(where=["network", "nodes", ref_map_nodes[node_ref], "NICs", nic_uuid])
|
||||
)
|
||||
node_obs_list.append(
|
||||
NodeObservation(
|
||||
where=["network", "nodes", ref_map_nodes[node_ref]],
|
||||
services=service_obs_list,
|
||||
folders=folder_obs_list,
|
||||
nics=nic_obs_list,
|
||||
logon_status=False,
|
||||
)
|
||||
)
|
||||
for link_obs_cfg in observation_space_cfg["options"]["links"]:
|
||||
link_ref = link_obs_cfg["link_ref"]
|
||||
link_obs_list.append(LinkObservation(where=["network", "links", ref_map_links[link_ref]]))
|
||||
# # node ip to index maps ip addresses to node id, as there are potentially multiple nics on a node, there are multiple ip addresses
|
||||
# node_ip_to_index = {}
|
||||
# for node_idx, node_cfg in enumerate(nodes_cfg):
|
||||
# n_ref = node_cfg["ref"]
|
||||
# n_obj = net.nodes[ref_map_nodes[n_ref]]
|
||||
# for nic_uuid, nic_obj in n_obj.nics.items():
|
||||
# node_ip_to_index[nic_obj.ip_address] = node_idx + 2
|
||||
|
||||
acl_obs = AclObservation(
|
||||
node_ip_to_id=node_ip_to_index,
|
||||
ports=game_cfg["ports"],
|
||||
protocols=game_cfg["ports"],
|
||||
where=["network", "nodes", observation_space_cfg["options"]["acl"]["router_node_ref"]],
|
||||
)
|
||||
obs_space = UC2BlueObservation(
|
||||
nodes=node_obs_list, links=link_obs_list, acl=acl_obs, ics=ICSObservation()
|
||||
)
|
||||
elif observation_space_cfg["type"] == "UC2RedObservation":
|
||||
obs_space = UC2RedObservation.from_config(observation_space_cfg["options"], sim=sim)
|
||||
elif observation_space_cfg["type"] == "UC2GreenObservation":
|
||||
obs_space = UC2GreenObservation.from_config(observation_space_cfg.get('options',{}))
|
||||
else:
|
||||
print("observation space config not specified correctly.")
|
||||
obs_space = NullObservation()
|
||||
# for node_obs_cfg in observation_space_cfg["options"]["nodes"]:
|
||||
# node_ref = node_obs_cfg["node_ref"]
|
||||
# folder_obs_list = []
|
||||
# service_obs_list = []
|
||||
# if "services" in node_obs_cfg:
|
||||
# for service_obs_cfg in node_obs_cfg["services"]:
|
||||
# service_obs_list.append(
|
||||
# ServiceObservation(
|
||||
# where=[
|
||||
# "network",
|
||||
# "nodes",
|
||||
# ref_map_nodes[node_ref],
|
||||
# "services",
|
||||
# ref_map_services[service_obs_cfg["service_ref"]],
|
||||
# ]
|
||||
# )
|
||||
# )
|
||||
# if "folders" in node_obs_cfg:
|
||||
# for folder_obs_cfg in node_obs_cfg["folders"]:
|
||||
# file_obs_list = []
|
||||
# if "files" in folder_obs_cfg:
|
||||
# for file_obs_cfg in folder_obs_cfg["files"]:
|
||||
# file_obs_list.append(
|
||||
# FileObservation(
|
||||
# where=[
|
||||
# "network",
|
||||
# "nodes",
|
||||
# ref_map_nodes[node_ref],
|
||||
# "folders",
|
||||
# folder_obs_cfg["folder_name"],
|
||||
# "files",
|
||||
# file_obs_cfg["file_name"],
|
||||
# ]
|
||||
# )
|
||||
# )
|
||||
# folder_obs_list.append(
|
||||
# FolderObservation(
|
||||
# where=[
|
||||
# "network",
|
||||
# "nodes",
|
||||
# ref_map_nodes[node_ref],
|
||||
# "folders",
|
||||
# folder_obs_cfg["folder_name"],
|
||||
# ],
|
||||
# files=file_obs_list,
|
||||
# )
|
||||
# )
|
||||
# nic_obs_list = []
|
||||
# for nic_uuid in net.nodes[ref_map_nodes[node_obs_cfg["node_ref"]]].nics.keys():
|
||||
# nic_obs_list.append(
|
||||
# NicObservation(where=["network", "nodes", ref_map_nodes[node_ref], "NICs", nic_uuid])
|
||||
# )
|
||||
# node_obs_list.append(
|
||||
# NodeObservation(
|
||||
# where=["network", "nodes", ref_map_nodes[node_ref]],
|
||||
# services=service_obs_list,
|
||||
# folders=folder_obs_list,
|
||||
# nics=nic_obs_list,
|
||||
# logon_status=False,
|
||||
# )
|
||||
# )
|
||||
# for link_obs_cfg in observation_space_cfg["options"]["links"]:
|
||||
# link_ref = link_obs_cfg["link_ref"]
|
||||
# link_obs_list.append(LinkObservation(where=["network", "links", ref_map_links[link_ref]]))
|
||||
|
||||
# acl_obs = AclObservation(
|
||||
# node_ip_to_id=node_ip_to_index,
|
||||
# ports=game_cfg["ports"],
|
||||
# protocols=game_cfg["ports"],
|
||||
# where=["network", "nodes", observation_space_cfg["options"]["acl"]["router_node_ref"]],
|
||||
# )
|
||||
# obs_space = UC2BlueObservation(
|
||||
# nodes=node_obs_list, links=link_obs_list, acl=acl_obs, ics=ICSObservation()
|
||||
# )
|
||||
# elif observation_space_cfg["type"] == "UC2RedObservation":
|
||||
# obs_space = UC2RedObservation.from_config(observation_space_cfg["options"], sim=sim)
|
||||
# elif observation_space_cfg["type"] == "UC2GreenObservation":
|
||||
# obs_space = UC2GreenObservation.from_config(observation_space_cfg.get('options',{}))
|
||||
# else:
|
||||
# print("observation space config not specified correctly.")
|
||||
# obs_space = NullObservation()
|
||||
"""
|
||||
|
||||
# CREATE ACTION SPACE
|
||||
action_space_cfg['options']['node_uuids'] = []
|
||||
# if a list of nodes is defined, convert them from node references to node UUIDs
|
||||
for action_node_option in action_space_cfg.get('options',{}).pop('nodes', {}):
|
||||
if 'node_ref' in action_node_option:
|
||||
node_uuid = ref_map_nodes[action_node_option['node_ref']]
|
||||
node_uuid = sess.ref_map_nodes[action_node_option['node_ref']]
|
||||
action_space_cfg['options']['node_uuids'].append(node_uuid)
|
||||
# Each action space can potentially have a different list of nodes that it can apply to. Therefore,
|
||||
# we will pass node_uuids as a part of the action space config.
|
||||
@@ -342,7 +347,7 @@ class PrimaiteSession:
|
||||
if 'options' in action_config:
|
||||
if 'target_router_ref' in action_config['options']:
|
||||
_target = action_config['options']['target_router_ref']
|
||||
action_config['options']['target_router_uuid'] = ref_map_nodes[_target]
|
||||
action_config['options']['target_router_uuid'] = sess.ref_map_nodes[_target]
|
||||
|
||||
action_space = ActionManager.from_config(sess, action_space_cfg)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user