Merge remote-tracking branch 'origin/dev' into feature/901-change-functionality-acl-rules

src/primaite/environment/__init__.py

@@ -1,2 +1,2 @@
-# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
+# Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
 """Gym/Gymnasium environment for RL agents consisting of a simulated computer network."""

src/primaite/environment/observations.py

@@ -1,6 +1,8 @@
+# Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
 """Module for handling configurable observation spaces in PrimAITE."""
 import logging
 from abc import ABC, abstractmethod
+from logging import Logger
 from typing import Dict, Final, List, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
@@ -17,14 +19,15 @@ from primaite.nodes.service_node import ServiceNode
 if TYPE_CHECKING:
     from primaite.environment.primaite_env import Primaite
 
-_LOGGER = logging.getLogger(__name__)
+
+_LOGGER: Logger = logging.getLogger(__name__)
 
 
 class AbstractObservationComponent(ABC):
     """Represents a part of the PrimAITE observation space."""
 
     @abstractmethod
-    def __init__(self, env: "Primaite"):
+    def __init__(self, env: "Primaite") -> None:
         """
         Initialise observation component.
 
@@ -39,7 +42,7 @@ class AbstractObservationComponent(ABC):
         return NotImplemented
 
     @abstractmethod
-    def update(self):
+    def update(self) -> None:
         """Update the observation based on the current state of the environment."""
         self.current_observation = NotImplemented
 
@@ -74,7 +77,7 @@ class NodeLinkTable(AbstractObservationComponent):
     _MAX_VAL: int = 1_000_000_000
     _DATA_TYPE: type = np.int64
 
-    def __init__(self, env: "Primaite"):
+    def __init__(self, env: "Primaite") -> None:
         """
         Initialise a NodeLinkTable observation space component.
 
@@ -101,7 +104,7 @@ class NodeLinkTable(AbstractObservationComponent):
 
         self.structure = self.generate_structure()
 
-    def update(self):
+    def update(self) -> None:
         """
         Update the observation based on current environment state.
 
@@ -148,7 +151,7 @@ class NodeLinkTable(AbstractObservationComponent):
                 protocol_index += 1
             item_index += 1
 
-    def generate_structure(self):
+    def generate_structure(self) -> List[str]:
         """Return a list of labels for the components of the flattened observation space."""
         nodes = self.env.nodes.values()
         links = self.env.links.values()
@@ -211,7 +214,7 @@ class NodeStatuses(AbstractObservationComponent):
 
     _DATA_TYPE: type = np.int64
 
-    def __init__(self, env: "Primaite"):
+    def __init__(self, env: "Primaite") -> None:
         """
         Initialise a NodeStatuses observation component.
 
@@ -237,7 +240,7 @@ class NodeStatuses(AbstractObservationComponent):
         self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE)
         self.structure = self.generate_structure()
 
-    def update(self):
+    def update(self) -> None:
         """
         Update the observation based on current environment state.
 
@@ -268,7 +271,7 @@ class NodeStatuses(AbstractObservationComponent):
             )
         self.current_observation[:] = obs
 
-    def generate_structure(self):
+    def generate_structure(self) -> List[str]:
         """Return a list of labels for the components of the flattened observation space."""
         services = self.env.services_list
 
@@ -318,7 +321,7 @@ class LinkTrafficLevels(AbstractObservationComponent):
         env: "Primaite",
         combine_service_traffic: bool = False,
         quantisation_levels: int = 5,
-    ):
+    ) -> None:
         """
         Initialise a LinkTrafficLevels observation component.
 
@@ -360,7 +363,7 @@ class LinkTrafficLevels(AbstractObservationComponent):
 
         self.structure = self.generate_structure()
 
-    def update(self):
+    def update(self) -> None:
         """
         Update the observation based on current environment state.
 
@@ -386,7 +389,7 @@ class LinkTrafficLevels(AbstractObservationComponent):
 
         self.current_observation[:] = obs
 
-    def generate_structure(self):
+    def generate_structure(self) -> List[str]:
         """Return a list of labels for the components of the flattened observation space."""
         structure = []
         for _, link in self.env.links.items():
@@ -470,7 +473,7 @@ class AccessControlList(AbstractObservationComponent):
 
         self.structure = self.generate_structure()
 
-    def update(self):
+    def update(self) -> None:
         """Update the observation based on current environment state.
 
         The structure of the observation space is described in :class:`.AccessControlList`
@@ -550,7 +553,7 @@ class AccessControlList(AbstractObservationComponent):
 
         self.current_observation[:] = obs
 
-    def generate_structure(self):
+    def generate_structure(self) -> List[str]:
         """Return a list of labels for the components of the flattened observation space."""
         structure = []
         for acl_rule in self.env.acl.acl:
@@ -593,7 +596,7 @@ class ObservationsHandler:
         "ACCESS_CONTROL_LIST": AccessControlList,
     }
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialise the observation handler."""
         self.registered_obs_components: List[AbstractObservationComponent] = []
 
@@ -606,7 +609,7 @@ class ObservationsHandler:
         # used for transactions and when flatten=true
         self._flat_observation: np.ndarray
 
-    def update_obs(self):
+    def update_obs(self) -> None:
         """Fetch fresh information about the environment."""
         current_obs = []
         for obs in self.registered_obs_components:
@@ -619,7 +622,7 @@ class ObservationsHandler:
             self._observation = tuple(current_obs)
         self._flat_observation = spaces.flatten(self._space, self._observation)
 
-    def register(self, obs_component: AbstractObservationComponent):
+    def register(self, obs_component: AbstractObservationComponent) -> None:
         """
         Add a component for this handler to track.
 
@@ -629,7 +632,7 @@ class ObservationsHandler:
         self.registered_obs_components.append(obs_component)
         self.update_space()
 
-    def deregister(self, obs_component: AbstractObservationComponent):
+    def deregister(self, obs_component: AbstractObservationComponent) -> None:
         """
         Remove a component from this handler.
 
@@ -640,7 +643,7 @@ class ObservationsHandler:
         self.registered_obs_components.remove(obs_component)
         self.update_space()
 
-    def update_space(self):
+    def update_space(self) -> None:
         """Rebuild the handler's composite observation space from its components."""
         component_spaces = []
         for obs_comp in self.registered_obs_components:
@@ -657,7 +660,7 @@ class ObservationsHandler:
             self._flat_space = spaces.Box(0, 1, (0,))
 
     @property
-    def space(self):
+    def space(self) -> spaces.Space:
         """Observation space, return the flattened version if flatten is True."""
         if len(self.registered_obs_components) > 1:
             return self._flat_space
@@ -665,7 +668,7 @@ class ObservationsHandler:
             return self._space
 
     @property
-    def current_observation(self):
+    def current_observation(self) -> Union[np.ndarray, Tuple[np.ndarray]]:
         """Current observation, return the flattened version if flatten is True."""
         if len(self.registered_obs_components) > 1:
             return self._flat_observation
@@ -673,7 +676,7 @@ class ObservationsHandler:
             return self._observation
 
     @classmethod
-    def from_config(cls, env: "Primaite", obs_space_config: dict):
+    def from_config(cls, env: "Primaite", obs_space_config: dict) -> "ObservationsHandler":
         """
         Parse a config dictinary, return a new observation handler populated with new observation component objects.
 
@@ -716,7 +719,7 @@ class ObservationsHandler:
         handler.update_obs()
         return handler
 
-    def describe_structure(self):
+    def describe_structure(self) -> List[str]:
         """
         Create a list of names for the features of the obs space.
 

src/primaite/environment/primaite_env.py

@@ -1,11 +1,12 @@
-# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
+# Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
 """Main environment module containing the PRIMmary AI Training Evironment (Primaite) class."""
 import copy
 import logging
 import uuid as uuid
+from logging import Logger
 from pathlib import Path
 from random import choice, randint, sample, uniform
-from typing import Dict, Final, Tuple, Union
+from typing import Any, Dict, Final, List, Tuple, Union
 
 import networkx as nx
 import numpy as np
@@ -20,6 +21,7 @@ from primaite.common.custom_typing import NodeUnion
 from primaite.common.enums import (
     ActionType,
     AgentFramework,
+    AgentIdentifier,
     FileSystemState,
     HardwareState,
     NodePOLInitiator,
@@ -48,7 +50,7 @@ from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_nod
 from primaite.transactions.transaction import Transaction
 from primaite.utils.session_output_writer import SessionOutputWriter
 
-_LOGGER = getLogger(__name__)
+_LOGGER: Logger = getLogger(__name__)
 
 
 class Primaite(Env):
@@ -66,7 +68,7 @@ class Primaite(Env):
         lay_down_config_path: Union[str, Path],
         session_path: Path,
         timestamp_str: str,
-    ):
+    ) -> None:
         """
         The Primaite constructor.
 
@@ -77,13 +79,14 @@ class Primaite(Env):
         """
         self.session_path: Final[Path] = session_path
         self.timestamp_str: Final[str] = timestamp_str
-        self._training_config_path = training_config_path
-        self._lay_down_config_path = lay_down_config_path
+        self._training_config_path: Union[str, Path] = training_config_path
+        self._lay_down_config_path: Union[str, Path] = lay_down_config_path
 
         self.training_config: TrainingConfig = training_config.load(training_config_path)
         _LOGGER.info(f"Using: {str(self.training_config)}")
 
         # Number of steps in an episode
+        self.episode_steps: int
         if self.training_config.session_type == SessionType.TRAIN:
             self.episode_steps = self.training_config.num_train_steps
         elif self.training_config.session_type == SessionType.EVAL:
@@ -94,7 +97,7 @@ class Primaite(Env):
         super(Primaite, self).__init__()
 
         # The agent in use
-        self.agent_identifier = self.training_config.agent_identifier
+        self.agent_identifier: AgentIdentifier = self.training_config.agent_identifier
 
         # Create a dictionary to hold all the nodes
         self.nodes: Dict[str, NodeUnion] = {}
@@ -113,42 +116,42 @@ class Primaite(Env):
         self.green_iers_reference: Dict[str, IER] = {}
 
         # Create a dictionary to hold all the node PoLs (this will come from an external source)
-        self.node_pol = {}
+        self.node_pol: Dict[str, NodeStateInstructionGreen] = {}
 
         # Create a dictionary to hold all the red agent IERs (this will come from an external source)
-        self.red_iers = {}
+        self.red_iers: Dict[str, IER] = {}
 
         # Create a dictionary to hold all the red agent node PoLs (this will come from an external source)
-        self.red_node_pol = {}
+        self.red_node_pol: Dict[str, NodeStateInstructionRed] = {}
 
         # Create the Access Control List
-        self.acl = AccessControlList(
+        self.acl: AccessControlList = AccessControlList(
             self.training_config.implicit_acl_rule,
             self.training_config.max_number_acl_rules,
         )
         # Sets limit for number of ACL rules in environment
-        self.max_number_acl_rules = self.training_config.max_number_acl_rules
+        self.max_number_acl_rules: int = self.training_config.max_number_acl_rules
 
         # Create a list of services (enums)
-        self.services_list = []
+        self.services_list: List[str] = []
 
         # Create a list of ports
-        self.ports_list = []
+        self.ports_list: List[str] = []
 
         # Create graph (network)
-        self.network = nx.MultiGraph()
+        self.network: nx.Graph = nx.MultiGraph()
 
         # Create a graph (network) reference
-        self.network_reference = nx.MultiGraph()
+        self.network_reference: nx.Graph = nx.MultiGraph()
 
         # Create step count
-        self.step_count = 0
+        self.step_count: int = 0
 
         self.total_step_count: int = 0
         """The total number of time steps completed."""
 
         # Create step info dictionary
-        self.step_info = {}
+        self.step_info: Dict[Any] = {}
 
         # Total reward
         self.total_reward: float = 0
@@ -157,22 +160,23 @@ class Primaite(Env):
         self.average_reward: float = 0
 
         # Episode count
-        self.episode_count = 0
+        self.episode_count: int = 0
 
         # Number of nodes - gets a value by examining the nodes dictionary after it's been populated
-        self.num_nodes = 0
+        self.num_nodes: int = 0
 
         # Number of links - gets a value by examining the links dictionary after it's been populated
-        self.num_links = 0
+        self.num_links: int = 0
 
         # Number of services - gets a value when config is loaded
-        self.num_services = 0
+        self.num_services: int = 0
 
         # Number of ports - gets a value when config is loaded
-        self.num_ports = 0
+        self.num_ports: int = 0
 
         # The action type
-        self.action_type = 0
+        # TODO: confirm type
+        self.action_type: int = 0
 
         # TODO fix up with TrainingConfig
         # stores the observation config from the yaml, default is NODE_LINK_TABLE
@@ -184,7 +188,7 @@ class Primaite(Env):
         # It will be initialised later.
         self.obs_handler: ObservationsHandler
 
-        self._obs_space_description = None
+        self._obs_space_description: List[str] = None
         "The env observation space description for transactions writing"
 
         # Open the config file and build the environment laydown
@@ -216,9 +220,13 @@ class Primaite(Env):
             _LOGGER.error("Could not save network diagram", exc_info=True)
 
         # Initiate observation space
+        self.observation_space: spaces.Space
+        self.env_obs: np.ndarray
         self.observation_space, self.env_obs = self.init_observations()
 
         # Define Action Space - depends on action space type (Node or ACL)
+        self.action_dict: Dict[int, List[int]]
+        self.action_space: spaces.Space
         if self.training_config.action_type == ActionType.NODE:
             _LOGGER.debug("Action space type NODE selected")
             # Terms (for node action space):
@@ -246,8 +254,12 @@ class Primaite(Env):
         else:
             _LOGGER.error(f"Invalid action type selected: {self.training_config.action_type}")
 
-        self.episode_av_reward_writer = SessionOutputWriter(self, transaction_writer=False, learning_session=True)
-        self.transaction_writer = SessionOutputWriter(self, transaction_writer=True, learning_session=True)
+        self.episode_av_reward_writer: SessionOutputWriter = SessionOutputWriter(
+            self, transaction_writer=False, learning_session=True
+        )
+        self.transaction_writer: SessionOutputWriter = SessionOutputWriter(
+            self, transaction_writer=True, learning_session=True
+        )
 
     @property
     def actual_episode_count(self) -> int:
@@ -256,7 +268,7 @@ class Primaite(Env):
             return self.episode_count - 1
         return self.episode_count
 
-    def set_as_eval(self):
+    def set_as_eval(self) -> None:
         """Set the writers to write to eval directories."""
         self.episode_av_reward_writer = SessionOutputWriter(self, transaction_writer=False, learning_session=False)
         self.transaction_writer = SessionOutputWriter(self, transaction_writer=True, learning_session=False)
@@ -265,12 +277,12 @@ class Primaite(Env):
         self.total_step_count = 0
         self.episode_steps = self.training_config.num_eval_steps
 
-    def _write_av_reward_per_episode(self):
+    def _write_av_reward_per_episode(self) -> None:
         if self.actual_episode_count > 0:
             csv_data = self.actual_episode_count, self.average_reward
             self.episode_av_reward_writer.write(csv_data)
 
-    def reset(self):
+    def reset(self) -> np.ndarray:
         """
         AI Gym Reset function.
 
@@ -304,7 +316,7 @@ class Primaite(Env):
 
         return self.env_obs
 
-    def step(self, action):
+    def step(self, action: int) -> Tuple[np.ndarray, float, bool, Dict]:
         """
         AI Gym Step function.
 
@@ -423,7 +435,7 @@ class Primaite(Env):
         # Return
         return self.env_obs, reward, done, self.step_info
 
-    def close(self):
+    def close(self) -> None:
         """Override parent close and close writers."""
         # Close files if last episode/step
         # if self.can_finish:
@@ -432,18 +444,18 @@ class Primaite(Env):
         self.transaction_writer.close()
         self.episode_av_reward_writer.close()
 
-    def init_acl(self):
+    def init_acl(self) -> None:
         """Initialise the Access Control List."""
         self.acl.remove_all_rules()
 
-    def output_link_status(self):
+    def output_link_status(self) -> None:
         """Output the link status of all links to the console."""
         for link_key, link_value in self.links.items():
             _LOGGER.debug("Link ID: " + link_value.get_id())
             for protocol in link_value.protocol_list:
                 print("    Protocol: " + protocol.get_name().name + ", Load: " + str(protocol.get_load()))
 
-    def interpret_action_and_apply(self, _action):
+    def interpret_action_and_apply(self, _action: int) -> None:
         """
         Applies agent actions to the nodes and Access Control List.
 
@@ -462,7 +474,7 @@ class Primaite(Env):
         else:
             logging.error("Invalid action type found")
 
-    def apply_actions_to_nodes(self, _action):
+    def apply_actions_to_nodes(self, _action: int) -> None:
         """
         Applies agent actions to the nodes.
 
@@ -550,7 +562,7 @@ class Primaite(Env):
         else:
             return
 
-    def apply_actions_to_acl(self, _action):
+    def apply_actions_to_acl(self, _action: int) -> None:
         """
         Applies agent actions to the Access Control List [TO DO].
 
@@ -630,7 +642,7 @@ class Primaite(Env):
             else:
                 return
 
-    def apply_time_based_updates(self):
+    def apply_time_based_updates(self) -> None:
         """
         Updates anything that needs to count down and then change state.
 
@@ -686,12 +698,12 @@ class Primaite(Env):
 
         return self.obs_handler.space, self.obs_handler.current_observation
 
-    def update_environent_obs(self):
+    def update_environent_obs(self) -> None:
         """Updates the observation space based on the node and link status."""
         self.obs_handler.update_obs()
         self.env_obs = self.obs_handler.current_observation
 
-    def load_lay_down_config(self):
+    def load_lay_down_config(self) -> None:
         """Loads config data in order to build the environment configuration."""
         for item in self.lay_down_config:
             if item["item_type"] == "NODE":
@@ -729,7 +741,7 @@ class Primaite(Env):
         _LOGGER.info("Environment configuration loaded")
         print("Environment configuration loaded")
 
-    def create_node(self, item):
+    def create_node(self, item: Dict) -> None:
         """
         Creates a node from config data.
 
@@ -810,7 +822,7 @@ class Primaite(Env):
         # Add node to network (reference)
         self.network_reference.add_nodes_from([node_ref])
 
-    def create_link(self, item: Dict):
+    def create_link(self, item: Dict) -> None:
         """
         Creates a link from config data.
 
@@ -854,7 +866,7 @@ class Primaite(Env):
             self.services_list,
         )
 
-    def create_green_ier(self, item):
+    def create_green_ier(self, item: Dict) -> None:
         """
         Creates a green IER from config data.
 
@@ -895,7 +907,7 @@ class Primaite(Env):
             ier_mission_criticality,
         )
 
-    def create_red_ier(self, item):
+    def create_red_ier(self, item: Dict) -> None:
         """
         Creates a red IER from config data.
 
@@ -925,7 +937,7 @@ class Primaite(Env):
             ier_mission_criticality,
         )
 
-    def create_green_pol(self, item):
+    def create_green_pol(self, item: Dict) -> None:
         """
         Creates a green PoL object from config data.
 
@@ -959,7 +971,7 @@ class Primaite(Env):
             pol_state,
         )
 
-    def create_red_pol(self, item):
+    def create_red_pol(self, item: Dict) -> None:
         """
         Creates a red PoL object from config data.
 
@@ -1000,7 +1012,7 @@ class Primaite(Env):
             pol_source_node_service_state,
         )
 
-    def create_acl_rule(self, item):
+    def create_acl_rule(self, item: Dict) -> None:
         """
         Creates an ACL rule from config data.
 
@@ -1023,7 +1035,8 @@ class Primaite(Env):
             acl_rule_position,
         )
 
-    def create_services_list(self, services):
+    # TODO: confirm typehint using runtime
+    def create_services_list(self, services: Dict) -> None:
         """
         Creates a list of services (enum) from config data.
 
@@ -1039,7 +1052,7 @@ class Primaite(Env):
         # Set the number of services
         self.num_services = len(self.services_list)
 
-    def create_ports_list(self, ports):
+    def create_ports_list(self, ports: Dict) -> None:
         """
         Creates a list of ports from config data.
 
@@ -1055,7 +1068,8 @@ class Primaite(Env):
         # Set the number of ports
         self.num_ports = len(self.ports_list)
 
-    def get_observation_info(self, observation_info):
+    # TODO: this is not used anymore, write a ticket to delete it
+    def get_observation_info(self, observation_info: Dict) -> None:
         """
         Extracts observation_info.
 
@@ -1064,7 +1078,8 @@ class Primaite(Env):
         """
         self.observation_type = ObservationType[observation_info["type"]]
 
-    def get_action_info(self, action_info):
+    # TODO: this is not used anymore, write a ticket to delete it.
+    def get_action_info(self, action_info: Dict) -> None:
         """
         Extracts action_info.
 
@@ -1073,7 +1088,7 @@ class Primaite(Env):
         """
         self.action_type = ActionType[action_info["type"]]
 
-    def save_obs_config(self, obs_config: dict):
+    def save_obs_config(self, obs_config: dict) -> None:
         """
         Cache the config for the observation space.
 
@@ -1086,7 +1101,7 @@ class Primaite(Env):
         """
         self.obs_config = obs_config
 
-    def reset_environment(self):
+    def reset_environment(self) -> None:
         """
         Resets environment.
 
@@ -1111,7 +1126,7 @@ class Primaite(Env):
         for ier_key, ier_value in self.red_iers.items():
             ier_value.set_is_running(False)
 
-    def reset_node(self, item):
+    def reset_node(self, item: Dict) -> None:
         """
         Resets the statuses of a node.
 
@@ -1159,7 +1174,7 @@ class Primaite(Env):
             # Bad formatting
             pass
 
-    def create_node_action_dict(self):
+    def create_node_action_dict(self) -> Dict[int, List[int]]:
         """
         Creates a dictionary mapping each possible discrete action to more readable multidiscrete action.
 
@@ -1199,7 +1214,7 @@ class Primaite(Env):
 
         return actions
 
-    def create_acl_action_dict(self):
+    def create_acl_action_dict(self) -> Dict[int, List[int]]:
         """Creates a dictionary mapping each possible discrete action to more readable multidiscrete action."""
         # Terms (for ACL action space):
         # [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule)
@@ -1240,7 +1255,7 @@ class Primaite(Env):
 
         return actions
 
-    def create_node_and_acl_action_dict(self):
+    def create_node_and_acl_action_dict(self) -> Dict[int, List[int]]:
         """
         Create a dictionary mapping each possible discrete action to a more readable mutlidiscrete action.
 
@@ -1257,7 +1272,7 @@ class Primaite(Env):
         combined_action_dict = {**acl_action_dict, **new_node_action_dict}
         return combined_action_dict
 
-    def _create_random_red_agent(self):
+    def _create_random_red_agent(self) -> None:
         """Decide on random red agent for the episode to be called in env.reset()."""
         # Reset the current red iers and red node pol
         self.red_iers = {}

src/primaite/environment/reward.py

@@ -1,25 +1,31 @@
-# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
+# Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
 """Implements reward function."""
-from typing import Dict
+from logging import Logger
+from typing import Dict, TYPE_CHECKING, Union
 
 from primaite import getLogger
+from primaite.common.custom_typing import NodeUnion
 from primaite.common.enums import FileSystemState, HardwareState, SoftwareState
 from primaite.common.service import Service
 from primaite.nodes.active_node import ActiveNode
 from primaite.nodes.service_node import ServiceNode
 
-_LOGGER = getLogger(__name__)
+if TYPE_CHECKING:
+    from primaite.config.training_config import TrainingConfig
+    from primaite.pol.ier import IER
+
+_LOGGER: Logger = getLogger(__name__)
 
 
 def calculate_reward_function(
-    initial_nodes,
-    final_nodes,
-    reference_nodes,
-    green_iers,
-    green_iers_reference,
-    red_iers,
-    step_count,
-    config_values,
+    initial_nodes: Dict[str, NodeUnion],
+    final_nodes: Dict[str, NodeUnion],
+    reference_nodes: Dict[str, NodeUnion],
+    green_iers: Dict[str, "IER"],
+    green_iers_reference: Dict[str, "IER"],
+    red_iers: Dict[str, "IER"],
+    step_count: int,
+    config_values: "TrainingConfig",
 ) -> float:
     """
     Compares the states of the initial and final nodes/links to get a reward.
@@ -93,7 +99,9 @@ def calculate_reward_function(
     return reward_value
 
 
-def score_node_operating_state(final_node, initial_node, reference_node, config_values) -> float:
+def score_node_operating_state(
+    final_node: NodeUnion, initial_node: NodeUnion, reference_node: NodeUnion, config_values: "TrainingConfig"
+) -> float:
     """
     Calculates score relating to the hardware state of a node.
 
@@ -142,7 +150,12 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
     return score
 
 
-def score_node_os_state(final_node, initial_node, reference_node, config_values) -> float:
+def score_node_os_state(
+    final_node: Union[ActiveNode, ServiceNode],
+    initial_node: Union[ActiveNode, ServiceNode],
+    reference_node: Union[ActiveNode, ServiceNode],
+    config_values: "TrainingConfig",
+) -> float:
     """
     Calculates score relating to the Software State of a node.
 
@@ -193,7 +206,9 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
     return score
 
 
-def score_node_service_state(final_node, initial_node, reference_node, config_values) -> float:
+def score_node_service_state(
+    final_node: ServiceNode, initial_node: ServiceNode, reference_node: ServiceNode, config_values: "TrainingConfig"
+) -> float:
     """
     Calculates score relating to the service state(s) of a node.
 
@@ -265,7 +280,12 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
     return score
 
 
-def score_node_file_system(final_node, initial_node, reference_node, config_values) -> float:
+def score_node_file_system(
+    final_node: Union[ActiveNode, ServiceNode],
+    initial_node: Union[ActiveNode, ServiceNode],
+    reference_node: Union[ActiveNode, ServiceNode],
+    config_values: "TrainingConfig",
+) -> float:
     """
     Calculates score relating to the file system state of a node.