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

Add docstrings and fix formatting on many things

Browse Source
This commit is contained in:
Marek Wolan
2023-10-19 09:36:23 +01:00
parent e0f8c3c5ea
commit 49e78d5291
9 changed files with 600 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,8 +1,35 @@
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Tuple, TYPE_CHECKING
"""
Manages the reward function for the agent.
Each agent is equipped with a RewardFunction, which is made up of a list of reward components. The components are
designed to calculate a reward value based on the current state of the simulation. The overall reward function is a
weighed sum of the components.
The reward function is typically specified using a config yaml file or a config dictionary. The following example shows
the structure:
```yaml
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
options:
node_ref: database_server
folder_name: database
file_name: database.db
- type: WEB_SERVER_404_PENALTY
weight: 0.5
options:
node_ref: web_server
service_ref: web_server_database_client
```
"""
from abc import abstractmethod
from typing import Dict, List, Tuple, TYPE_CHECKING
from primaite import getLogger
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.game.agent.utils import access_from_nested_dict
_LOGGER = getLogger(__name__)
@@ -11,27 +38,60 @@ if TYPE_CHECKING:
class AbstractReward:
"""Base class for reward function components."""
@abstractmethod
def calculate(self, state: Dict) -> float:
"""Calculate the reward for the current state."""
return 0.0
@classmethod
@abstractmethod
def from_config(cls, config: dict, session: "PrimaiteSession") -> "AbstractReward":
"""Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor
:type config: dict
:param session: Reference to the PrimAITE Session object
:type session: PrimaiteSession
:return: The reward component.
:rtype: AbstractReward
"""
return cls()
class DummyReward(AbstractReward):
"""Dummy reward function component which always returns 0."""
def calculate(self, state: Dict) -> float:
"""Calculate the reward for the current state."""
return 0.0
@classmethod
def from_config(cls, config: dict, session: "PrimaiteSession") -> "DummyReward":
"""Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor. Should be empty.
:type config: dict
:param session: Reference to the PrimAITE Session object
:type session: PrimaiteSession
"""
return cls()
class DatabaseFileIntegrity(AbstractReward):
"""Reward function component which rewards the agent for maintaining the integrity of a database file."""
def __init__(self, node_uuid: str, folder_name: str, file_name: str) -> None:
"""Initialise the reward component.
:param node_uuid: UUID of the node which contains the database file.
:type node_uuid: str
:param folder_name: folder which contains the database file.
:type folder_name: str
:param file_name: name of the database file.
:type file_name: str
"""
self.location_in_state = [
"network",
"nodes",
@@ -44,6 +104,11 @@ class DatabaseFileIntegrity(AbstractReward):
]
def calculate(self, state: Dict) -> float:
"""Calculate the reward for the current state.
:param state: The current state of the simulation.
:type state: Dict
"""
database_file_state = access_from_nested_dict(state, self.location_in_state)
health_status = database_file_state["health_status"]
if health_status == "corrupted":
@@ -55,6 +120,15 @@ class DatabaseFileIntegrity(AbstractReward):
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession") -> "DatabaseFileIntegrity":
"""Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor
:type config: Dict
:param session: Reference to the PrimAITE Session object
:type session: PrimaiteSession
:return: The reward component.
:rtype: DatabaseFileIntegrity
"""
node_ref = config.get("node_ref")
folder_name = config.get("folder_name")
file_name = config.get("file_name")
@@ -76,7 +150,10 @@ class DatabaseFileIntegrity(AbstractReward):
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"
(
f"{cls.__name__} could not be initialised from config because the referenced node could not be "
f"found in the simulation"
)
)
return DummyReward() # TODO: better error handling
@@ -84,10 +161,24 @@ class DatabaseFileIntegrity(AbstractReward):
class WebServer404Penalty(AbstractReward):
"""Reward function component which penalises the agent when the web server returns a 404 error."""
def __init__(self, node_uuid: str, service_uuid: str) -> None:
"""Initialise the reward component.
:param node_uuid: UUID of the node which contains the web server service.
:type node_uuid: str
:param service_uuid: UUID of the web server service.
:type service_uuid: str
"""
self.location_in_state = ["network", "nodes", node_uuid, "services", service_uuid]
def calculate(self, state: Dict) -> float:
"""Calculate the reward for the current state.
:param state: The current state of the simulation.
:type state: Dict
"""
web_service_state = access_from_nested_dict(state, self.location_in_state)
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.
@@ -100,16 +191,31 @@ class WebServer404Penalty(AbstractReward):
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession") -> "WebServer404Penalty":
"""Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor
:type config: Dict
:param session: Reference to the PrimAITE Session object
:type session: PrimaiteSession
:return: The reward component.
:rtype: WebServer404Penalty
"""
node_ref = config.get("node_ref")
service_ref = config.get("service_ref")
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."
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!
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."
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
@@ -117,6 +223,8 @@ class WebServer404Penalty(AbstractReward):
class RewardFunction:
"""Manages the reward function for the agent."""
__rew_class_identifiers: Dict[str, type[AbstractReward]] = {
"DUMMY": DummyReward,
"DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity,
@@ -124,13 +232,26 @@ class RewardFunction:
}
def __init__(self):
"""Initialise the reward function object."""
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:
"""Add a reward component to the reward function.
:param component: Instance of a reward component.
:type component: AbstractReward
:param weight: Relative weight of the reward component, defaults to 1.0
:type weight: float, optional
"""
self.reward_components.append((component, weight))
def calculate(self, state: Dict) -> float:
"""Calculate the overall reward for the current state.
:param state: The current state of the simulation.
:type state: Dict
"""
total = 0.0
for comp_and_weight in self.reward_components:
comp = comp_and_weight[0]
@@ -140,6 +261,15 @@ class RewardFunction:
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession") -> "RewardFunction":
"""Create a reward function from a config dictionary.
:param config: dict of options for the reward manager's constructor
:type config: Dict
:param session: Reference to the PrimAITE Session object
:type session: PrimaiteSession
:return: The reward manager.
:rtype: RewardFunction
"""
new = cls()
for rew_component_cfg in config["reward_components"]: