oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add empty obs placeholders

Browse Source
This commit is contained in:
Marek Wolan
2023-10-10 21:01:09 +01:00
parent b53c3856dd
commit def2c3699b
9 changed files with 964 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
src/primaite/game/agent/GATE_agents.py
View File

@@ -1,57 +1,9 @@
from primaite.game.agent.interface import AbstractGATEAgent
from arcd_gate.client.gate_client import GATEClient
class GATEMan(GATEClient):
@property
def rl_framework(self) -> str:
return "SB3"
@property
def rl_framework(self) -> str:
pass
@property
def rl_algorithm(self) -> str:
pass
@property
def seed(self) -> Optional[int]:
return None
@property
def n_learn_episodes(self) -> int:
return 0
@property
def n_learn_steps(self) -> int:
return 0
@property
def n_eval_episodes(self) -> int:
return 0
@property
def n_eval_steps(self) -> int:
return 0
@property
def action_space(self) -> spaces.Space:
pass
@property
def observation_space(self) -> spaces.Space:
pass
def step(self, action: ActType) -> Tuple[np.ndarray, float, bool, bool, Dict]:
pass
def reset(self, *, seed: Optional[int] = None, options: Optional[dict[str, Any]] = None) -> Tuple[np.ndarray, Dict]:
pass
def close(self):
pass
from typing import Dict, Optional, Tuple
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.interface import AbstractGATEAgent, ObsType
from primaite.game.agent.observations import ObservationSpace
from primaite.game.agent.rewards import RewardFunction
from gymnasium.core import ActType, ObsType
class GATERLAgent(AbstractGATEAgent):
...
@@ -61,4 +13,9 @@ class GATERLAgent(AbstractGATEAgent):
# For example MultiAgentEnv in Ray allows sending a dict of observations of multiple agents, then it will reply
# with the actions for those agents.
def __init__(self, agent_name: str | None, action_space: ActionManager | None, observation_space: ObservationSpace | None, reward_function: RewardFunction | None) -> None:
super().__init__(agent_name, action_space, observation_space, reward_function)
self.most_recent_action: ActType
def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]:
return self.most_recent_action

2
src/primaite/game/agent/actions.py
View File

@@ -2,7 +2,7 @@ import itertools
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
from gym import spaces
from gymnasium import spaces
from primaite import getLogger
from primaite.simulator.sim_container import Simulation

82
src/primaite/game/agent/observations.py
View File

@@ -1,10 +1,12 @@
from abc import ABC, abstractmethod
from typing import Any, Dict, Hashable, List, Optional, Sequence, Tuple, TYPE_CHECKING
from gym import spaces
from gymnasium import spaces
from pydantic import BaseModel
from primaite.simulator.sim_container import Simulation
from primaite import getLogger
_LOGGER = getLogger(__name__)
if TYPE_CHECKING:
from primaite.game.session import PrimaiteSession
@@ -181,7 +183,7 @@ class LinkObservation(AbstractObservation):
class FolderObservation(AbstractObservation):
def __init__(self, where: Optional[Tuple[str]] = None, files: List[FileObservation] = []) -> None:
def __init__(self, where: Optional[Tuple[str]] = None, files: List[FileObservation] = [], num_files_per_folder:int=2) -> None:
"""Initialise folder Observation, including files inside of the folder.
:param where: Where in the simulation state dictionary to find the relevant information for this folder.
@@ -203,6 +205,13 @@ class FolderObservation(AbstractObservation):
self.where: Optional[Tuple[str]] = where
self.files: List[FileObservation] = files
while len(self.files) < num_files_per_folder:
self.files.append(FileObservation())
while len(self.files)> num_files_per_folder:
truncated_file = self.files.pop()
msg = f"Too many files in folde observation. Truncating file {truncated_file}"
_LOGGER.warn(msg)
raise UserWarning(msg)
self.default_observation = {
"health_status": 0,
@@ -235,13 +244,13 @@ 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]], num_files_per_folder:int=2):
where = parent_where + ["folders", config["folder_name"]]
file_configs = config["files"]
files = [FileObservation.from_config(config=f, session=session, parent_where=where) for f in file_configs]
return cls(where=where, files=files)
return cls(where=where, files=files, num_files_per_folder=num_files_per_folder)
class NicObservation(AbstractObservation):
@@ -277,7 +286,10 @@ class NodeObservation(AbstractObservation):
folders: List[FolderObservation] = [],
nics: List[NicObservation] = [],
logon_status: bool = False,
) -> None:
num_services_per_node: int = 2,
num_folders_per_node: int = 2,
num_files_per_folder: int = 2
) -> None:
"""
Configurable observation for a node in the simulation.
@@ -302,7 +314,24 @@ class NodeObservation(AbstractObservation):
self.where: Optional[Tuple[str]] = where
self.services: List[ServiceObservation] = services
while len(self.services)<num_services_per_node:
# add an empty service observation without `where` parameter that will always return default (blank) observations
self.services.append(ServiceObservation())
while len(self.services)>num_services_per_node:
truncated_service = self.services.pop()
msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}"
_LOGGER.warn(msg)
raise UserWarning(msg)
# truncate service list
self.folders: List[FolderObservation] = folders
while len(self.folders) < num_folders_per_node:
# add an empty folder observation without `where` parameter that will always return default (blank) observations
self.folders.append(FolderObservation())
while len(self.folders) > num_folders_per_node:
truncated_folder = self.folders.pop()
msg = f"Too many folders in Node observation for node. Truncating service {truncated_folder.where[-1]}"
self.nics: List[NicObservation] = nics
self.logon_status: bool = logon_status
@@ -349,7 +378,13 @@ class NodeObservation(AbstractObservation):
@classmethod
def from_config(
cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]] = None
cls,
config: Dict,
session: "PrimaiteSession",
parent_where: Optional[List[str]] = None,
num_services_per_node: int = 2,
num_folders_per_node: int = 2,
num_files_per_folder: int = 2,
) -> "NodeObservation":
node_uuid = session.ref_map_nodes[config["node_ref"]]
if parent_where is None:
@@ -360,12 +395,21 @@ class NodeObservation(AbstractObservation):
svc_configs = config.get("services", {})
services = [ServiceObservation.from_config(config=c, session=session, parent_where=where) for c in svc_configs]
folder_configs = config.get("folders", {})
folders = [FolderObservation.from_config(config=c, session=session, parent_where=where) for c in folder_configs]
folders = [FolderObservation.from_config(config=c, session=session, parent_where=where, num_files_per_folder=num_files_per_folder) 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}] 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)
return cls(where=where, services=services, folders=folders, nics=nics, logon_status=logon_status)
return cls(
where=where,
services=services,
folders=folders,
nics=nics,
logon_status=logon_status,
num_services_per_node = num_services_per_node,
num_folders_per_node = num_folders_per_node,
num_files_per_folder = num_files_per_folder,
)
class AclObservation(AbstractObservation):
@@ -467,11 +511,12 @@ class AclObservation(AbstractObservation):
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession") -> "AclObservation":
node_ip_to_idx = {}
for node_idx, node_cfg in enumerate(config["node_order"]):
n_ref = node_cfg["node_ref"]
n_obj = session.simulation.network.nodes[session.ref_map_nodes[n_ref]]
for nic_uuid, nic_obj in n_obj.nics.items():
node_ip_to_idx[nic_obj.ip_address] = node_idx + 2
for ip_idx, ip_map_config in enumerate(config["ip_address_order"]):
node_ref = ip_map_config["node_ref"]
nic_num = ip_map_config["nic_num"]
node_obj = session.simulation.network.nodes[session.ref_map_nodes[node_ref]]
nic_obj = node_obj.ethernet_port[nic_num]
node_ip_to_idx[nic_obj.ip_address] = ip_idx + 2
router_uuid = session.ref_map_nodes[config["router_node_ref"]]
return cls(
@@ -552,7 +597,16 @@ class UC2BlueObservation(AbstractObservation):
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession"):
node_configs = config["nodes"]
nodes = [NodeObservation.from_config(config=n, session=session) for n in node_configs]
num_services_per_node = config["num_services_per_node"]
num_folders_per_node = config["num_folders_per_node"]
num_files_per_folder = config["num_files_per_folder"]
nodes = [NodeObservation.from_config(
config=n,
session=session,
num_services_per_node= num_services_per_node,
num_folders_per_node=num_folders_per_node,
num_files_per_folder=num_files_per_folder,
) for n in node_configs]
link_configs = config["links"]
links = [LinkObservation.from_config(config=l, session=session) for l in link_configs]

49
src/primaite/game/session.py
View File

@@ -7,12 +7,16 @@
from ipaddress import IPv4Address
from typing import Any, Dict, List, Optional, Tuple
from gymnasium.vector.utils import spaces
from gymnasium import spaces
from gymnasium.spaces.utils import flatten, flatten_space, unflatten
from gymnasium.core import ObsType, ActType
import numpy as np
from pydantic import BaseModel
from primaite import getLogger
from primaite.game.agent.GATE_agents import GATERLAgent
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.interface import AbstractAgent, RandomAgent
from primaite.game.agent.observations import (
@@ -54,7 +58,7 @@ _LOGGER = getLogger(__name__)
class PrimaiteGATEClient(GATEClient):
def __init__(self, parent_session:"PrimaiteSession", service_port: int = 50000):
super().__init__(service_port=service_port)
self.parent_session:"PrimaiteSession"
self.parent_session:"PrimaiteSession" = parent_session
@property
def rl_framework(self) -> str:
@@ -86,21 +90,33 @@ class PrimaiteGATEClient(GATEClient):
@property
def action_space(self) -> spaces.Space:
return self.parent_session.rl_agent.action_space
return self.parent_session.rl_agent.action_space.space
@property
def observation_space(self) -> spaces.Space:
return self.parent_session.rl_agent.observation_space
print("YEEY0")
print(flatten_space(spaces.Dict({})))
print("YEEY1")
# print(self.parent_session.rl_agent.observation_space.space)
return flatten_space(self.parent_session.rl_agent.observation_space.space)
def step(self, action: ActType) -> Tuple[ndarray, float, bool, bool, Dict]:
def step(self, action: ActType) -> Tuple[ObsType, float, bool, bool, Dict]:
self.parent_session.rl_agent.most_recent_action = action
self.parent_session.step()
#TODO: not sure how to go about this.
obs = self.parent_session.rl_agent.observation_space.observe()
obs = flatten(self.parent_session.rl_agent.observation_space.space, obs)
rew = self.parent_session.rl_agent.reward_function.calculate()
term = False
trunc = False
info = {}
return obs, rew, term, trunc, info
def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None) -> Tuple[ndarray, Dict]:
...
self.parent_session.reset()
def close(self):
...
self.parent_session.close()
class PrimaiteSessionOptions(BaseModel):
ports: List[str]
@@ -131,15 +147,11 @@ class PrimaiteSession:
self.ref_map_nodes: Dict[str, Node] = {}
self.ref_map_services: Dict[str, Service] = {}
self.ref_map_links: Dict[str, Link] = {}
self.gate_client: PrimaiteGATEClient = PrimaiteGATEClient(self)
def start_session(self, opts="TODO..."):
"""Commence the session, this gives the gate client control over the simulation/agent loop."""
...
def eval(self, opts="TODO..."):
...
self.gate_client.start()
def step(self):
_LOGGER.debug(f"Stepping primaite session. Step counter: {self.step_counter}")
@@ -172,12 +184,18 @@ class PrimaiteSession:
# 10. primaite session receives the action from the agents and asks the simulation to apply each
_LOGGER.debug(f"Sending request to simulation: {agent_request}")
self.simulation.apply_action(agent_request)
self.simulation.apply_request(agent_request)
_LOGGER.debug(f"Initiating simulation step {self.step_counter}")
self.simulation.apply_timestep(self.step_counter)
self.step_counter += 1
def reset(self):
pass
def close(self):
pass
@classmethod
def from_config(cls, cfg: dict) -> "PrimaiteSession":
sess = cls()
@@ -351,6 +369,7 @@ class PrimaiteSession:
reward_function=rew_function,
)
sess.agents.append(new_agent)
sess.rl_agent = new_agent
elif agent_type == "RedDatabaseCorruptingAgent":
new_agent = RandomAgent(
agent_name=agent_cfg["ref"],