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

Add documentation

Browse Source
This commit is contained in:
Marek Wolan
2023-10-19 01:56:40 +01:00
parent 027addc485
commit e0f8c3c5ea
6 changed files with 445 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
src/primaite/game/agent/rewards.py
View File

@@ -1,8 +1,9 @@
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Tuple, TYPE_CHECKING
from primaite import getLogger
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
if TYPE_CHECKING:
@@ -10,14 +11,13 @@ if TYPE_CHECKING:
class AbstractReward:
@abstractmethod
def calculate(self, state: Dict) -> float:
return 0.0
@classmethod
@abstractmethod
def from_config(cls, config:dict, session:"PrimaiteSession") -> "AbstractReward":
def from_config(cls, config: dict, session: "PrimaiteSession") -> "AbstractReward":
return cls()
@@ -26,16 +26,26 @@ class DummyReward(AbstractReward):
return 0.0
@classmethod
def from_config(cls, config: dict, session:"PrimaiteSession") -> "DummyReward":
def from_config(cls, config: dict, session: "PrimaiteSession") -> "DummyReward":
return cls()
class DatabaseFileIntegrity(AbstractReward):
def __init__(self, node_uuid:str, folder_name:str, file_name:str) -> None:
self.location_in_state = ["network", "nodes", node_uuid, "file_system", "folders",folder_name, "files", file_name]
def __init__(self, node_uuid: str, folder_name: str, file_name: str) -> None:
self.location_in_state = [
"network",
"nodes",
node_uuid,
"file_system",
"folders",
folder_name,
"files",
file_name,
]
def calculate(self, state: Dict) -> float:
database_file_state = access_from_nested_dict(state, self.location_in_state)
health_status = database_file_state['health_status']
health_status = database_file_state["health_status"]
if health_status == "corrupted":
return -1
elif health_status == "good":
@@ -48,29 +58,39 @@ class DatabaseFileIntegrity(AbstractReward):
node_ref = config.get("node_ref")
folder_name = config.get("folder_name")
file_name = config.get("file_name")
if not (node_ref):
_LOGGER.error(f"{cls.__name__} could not be initialised from config because node_ref parameter was not specified")
return DummyReward() #TODO: better error handling
if not node_ref:
_LOGGER.error(
f"{cls.__name__} could not be initialised from config because node_ref parameter was not specified"
)
return DummyReward() # TODO: better error handling
if not folder_name:
_LOGGER.error(f"{cls.__name__} could not be initialised from config because folder_name parameter was not specified")
return DummyReward() # TODO: better error handling
_LOGGER.error(
f"{cls.__name__} could not be initialised from config because folder_name parameter was not specified"
)
return DummyReward() # TODO: better error handling
if not file_name:
_LOGGER.error(f"{cls.__name__} could not be initialised from config because file_name parameter was not specified")
return DummyReward() # TODO: better error handling
_LOGGER.error(
f"{cls.__name__} could not be initialised from config because file_name parameter was not specified"
)
return DummyReward() # TODO: better error handling
node_uuid = session.ref_map_nodes[node_ref]
if not node_uuid:
_LOGGER.error(f"{cls.__name__} could not be initialised from config because the referenced node could not be found in the simulation")
return DummyReward() # TODO: better error handling
_LOGGER.error(
f"{cls.__name__} could not be initialised from config because the referenced node could not be found in the simulation"
)
return DummyReward() # TODO: better error handling
return cls(node_uuid=node_uuid, folder_name=folder_name, file_name=file_name)
return cls(node_uuid = node_uuid, folder_name=folder_name, file_name=file_name)
class WebServer404Penalty(AbstractReward):
def __init__(self, node_uuid:str, service_uuid:str) -> None:
self.location_in_state = ['network','nodes', node_uuid, 'services', service_uuid]
def __init__(self, node_uuid: str, service_uuid: str) -> None:
self.location_in_state = ["network", "nodes", node_uuid, "services", service_uuid]
def calculate(self, state: Dict) -> float:
web_service_state = access_from_nested_dict(state, self.location_in_state)
most_recent_return_code = web_service_state['most_recent_return_code']
most_recent_return_code = web_service_state["most_recent_return_code"]
# TODO: reward needs to use the current web state. Observation should return web state at the time of last scan.
if most_recent_return_code == 200:
return 1
elif most_recent_return_code == 404:
@@ -85,13 +105,13 @@ class WebServer404Penalty(AbstractReward):
if not (node_ref and service_ref):
msg = f"{cls.__name__} could not be initialised from config because node_ref and service_ref were not found in reward config."
_LOGGER.warn(msg)
return DummyReward() #TODO: should we error out with incorrect inputs? Probably!
return DummyReward() # TODO: should we error out with incorrect inputs? Probably!
node_uuid = session.ref_map_nodes[node_ref]
service_uuid = session.ref_map_services[service_ref].uuid
if not (node_uuid and service_uuid):
msg = f"{cls.__name__} could not be initialised because node {node_ref} and service {service_ref} were not found in the simulator."
_LOGGER.warn(msg)
return DummyReward() # TODO: consider erroring here as well
return DummyReward() # TODO: consider erroring here as well
return cls(node_uuid=node_uuid, service_uuid=service_uuid)
@@ -101,13 +121,13 @@ class RewardFunction:
"DUMMY": DummyReward,
"DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity,
"WEB_SERVER_404_PENALTY": WebServer404Penalty,
}
}
def __init__(self):
self.reward_components: List[Tuple[AbstractReward, float]] = []
"attribute reward_components keeps track of reward components and the weights assigned to each."
def regsiter_component(self, component:AbstractReward, weight:float=1.0) -> None:
def regsiter_component(self, component: AbstractReward, weight: float = 1.0) -> None:
self.reward_components.append((component, weight))
def calculate(self, state: Dict) -> float:
@@ -124,8 +144,8 @@ class RewardFunction:
for rew_component_cfg in config["reward_components"]:
rew_type = rew_component_cfg["type"]
weight = rew_component_cfg.get("weight",1.0)
weight = rew_component_cfg.get("weight", 1.0)
rew_class = cls.__rew_class_identifiers[rew_type]
rew_instance = rew_class.from_config(config=rew_component_cfg.get('options',{}), session=session)
rew_instance = rew_class.from_config(config=rew_component_cfg.get("options", {}), session=session)
new.regsiter_component(component=rew_instance, weight=weight)
return new