Add docstrings and fix formatting on many things

sandbox.py

@@ -1,19 +1,22 @@
-
+# flake8: noqa
-from primaite.game.session import PrimaiteSession
+import logging
 
 from primaite import _PRIMAITE_CONFIG, PRIMAITE_PATHS
-import logging
+from primaite.game.session import PrimaiteSession
-_PRIMAITE_CONFIG['log_level']=logging.DEBUG
+
+_PRIMAITE_CONFIG["log_level"] = logging.DEBUG
 print(PRIMAITE_PATHS.app_log_dir_path)
 import itertools
-from primaite.game.session import PrimaiteSession
+
-from primaite.simulator.sim_container import Simulation
-from primaite.game.agent.interface import AbstractAgent
-from primaite.simulator.network.networks import arcd_uc2_network
 import yaml
 
-with open('example_config.yaml', 'r') as file:
+from primaite.game.agent.interface import AbstractAgent
+from primaite.game.session import PrimaiteSession
+from primaite.simulator.network.networks import arcd_uc2_network
+from primaite.simulator.sim_container import Simulation
+
+with open("example_config.yaml", "r") as file:
     cfg = yaml.safe_load(file)
 sess = PrimaiteSession.from_config(cfg)
 
-sess.start_session()
+sess.start_session()

src/primaite/game/agent/GATE_agents.py

@@ -1,9 +1,13 @@
+# flake8: noqa
 from typing import Dict, Optional, Tuple
+
+from gymnasium.core import ActType, ObsType
+
 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):
     ...
@@ -13,9 +17,15 @@ 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:
+    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
+        return self.most_recent_action

src/primaite/game/agent/actions.py

@@ -695,14 +695,14 @@ class ActionManager:
         return {i: p for i, p in enumerate(all_action_possibilities)}
 
     def get_action(self, action: int) -> Tuple[str, Dict]:
-        """Produce action in CAOS format"""
+        """Produce action in CAOS format."""
         """the agent chooses an action (as an integer), this is converted into an action in CAOS format"""
         """The caos format is basically a action identifier, followed by parameters stored in a dictionary"""
         act_identifier, act_options = self.action_map[action]
         return act_identifier, act_options
 
-    def form_request(self, action_identifier: str, action_options: Dict):
+    def form_request(self, action_identifier: str, action_options: Dict) -> List[str]:
-        """Take action in CAOS format and use the execution definition to change it into PrimAITE request format"""
+        """Take action in CAOS format and use the execution definition to change it into PrimAITE request format."""
         act_obj = self.actions[action_identifier]
         return act_obj.form_request(**action_options)
 
@@ -712,7 +712,8 @@ class ActionManager:
         return spaces.Discrete(len(self.action_map))
 
     def get_node_uuid_by_idx(self, node_idx: int) -> str:
-        """Get the node UUID corresponding to the given index.
+        """
+        Get the node UUID corresponding to the given index.
 
         :param node_idx: The index of the node to retrieve.
         :type node_idx: int
@@ -722,7 +723,8 @@ class ActionManager:
         return self.node_uuids[node_idx]
 
     def get_folder_uuid_by_idx(self, node_idx: int, folder_idx: int) -> Optional[str]:
-        """Get the folder UUID corresponding to the given node and folder indices.
+        """
+        Get the folder UUID corresponding to the given node and folder indices.
 
         :param node_idx: The index of the node.
         :type node_idx: int
@@ -731,7 +733,6 @@ class ActionManager:
         :return: The UUID of the folder. Or None if the node has fewer folders than the given index.
         :rtype: Optional[str]
         """
-
         node_uuid = self.get_node_uuid_by_idx(node_idx)
         node = self.sim.network.nodes[node_uuid]
         folder_uuids = list(node.file_system.folders.keys())

src/primaite/game/agent/interface.py

@@ -1,6 +1,6 @@
 """Interface for agents."""
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Tuple, TypeAlias, Union
+from typing import Dict, List, Optional, Tuple, TypeAlias, Union
 
 import numpy as np
 
@@ -21,6 +21,18 @@ class AbstractAgent(ABC):
         observation_space: Optional[ObservationSpace],
         reward_function: Optional[RewardFunction],
     ) -> None:
+        """
+        Initialize an agent.
+
+        :param agent_name: Unique string identifier for the agent, for reporting and multi-agent purposes.
+        :type agent_name: Optional[str]
+        :param action_space: Action space for the agent.
+        :type action_space: Optional[ActionManager]
+        :param observation_space: Observation space for the agent.
+        :type observation_space: Optional[ObservationSpace]
+        :param reward_function: Reward function for the agent.
+        :type reward_function: Optional[RewardFunction]
+        """
         self.agent_name: str = agent_name or "unnamed_agent"
         self.action_space: Optional[ActionManager] = action_space
         self.observation_space: Optional[ObservationSpace] = observation_space
@@ -32,16 +44,38 @@ class AbstractAgent(ABC):
 
     def convert_state_to_obs(self, state: Dict) -> ObsType:
         """
+        Convert a state from the simulator into an observation for the agent using the observation space.
+
         state : dict state directly from simulation.describe_state
         output : dict state according to CAOS.
         """
         return self.observation_space.observe(state)
 
     def calculate_reward_from_state(self, state: Dict) -> float:
+        """
+        Use the reward function to calculate a reward from the state.
+
+        :param state: State of the environment.
+        :type state: Dict
+        :return: Reward from the state.
+        :rtype: float
+        """
         return self.reward_function.calculate(state)
 
     @abstractmethod
     def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]:
+        """
+        Return an action to be taken in the environment.
+
+        Subclasses should implement agent logic here. It should use the observation as input to decide best next action.
+
+        :param obs: Observation of the environment.
+        :type obs: ObsType
+        :param reward: Reward from the previous action, defaults to None TODO: should this parameter even be accepted?
+        :type reward: float, optional
+        :return: Action to be taken in the environment.
+        :rtype: Tuple[str, Dict]
+        """
         # in RL agent, this method will send CAOS observation to GATE RL agent, then receive a int 0-39,
         # then use a bespoke conversion to take 1-40 int back into CAOS action
         return ("DO_NOTHING", {})
@@ -64,6 +98,15 @@ class RandomAgent(AbstractScriptedAgent):
     """Agent that ignores its observation and acts completely at random."""
 
     def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]:
+        """Randomly sample an action from the action space.
+
+        :param obs: _description_
+        :type obs: ObsType
+        :param reward: _description_, defaults to None
+        :type reward: float, optional
+        :return: _description_
+        :rtype: Tuple[str, Dict]
+        """
         return self.action_space.get_action(self.action_space.space.sample())
 
 

src/primaite/game/agent/observations.py

@@ -1,12 +1,12 @@
+"""Manages the observation space for the agent."""
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Hashable, List, Optional, Sequence, Tuple, TYPE_CHECKING
+from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
 
 from gymnasium import spaces
-from pydantic import BaseModel
 
-from primaite.simulator.sim_container import Simulation
-from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
 from primaite import getLogger
+from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
+
 _LOGGER = getLogger(__name__)
 
 if TYPE_CHECKING:
@@ -14,22 +14,25 @@ if TYPE_CHECKING:
 
 
 class AbstractObservation(ABC):
+    """Abstract class for an observation space component."""
+
     @abstractmethod
     def observe(self, state: Dict) -> Any:
-        """_summary_
+        """
+        Return an observation based on the current state of the simulation.
 
-        :param state: _description_
+        :param state: Simulation state dictionary
         :type state: Dict
-        :return: _description_
+        :return: Observation
         :rtype: Any
         """
-        ...
+        pass
 
     @property
     @abstractmethod
     def space(self) -> spaces.Space:
-        """Subclasses must define the shape that they expect"""
+        """Gymnasium space object describing the observation space."""
-        ...
+        pass
 
     @classmethod
     @abstractmethod
@@ -39,12 +42,15 @@ class AbstractObservation(ABC):
         The `session` parameter is for a the PrimaiteSession object that spawns this component. During deserialisation,
         a subclass of this class may need to translate from a 'reference' to a UUID.
         """
+        pass
 
 
 class FileObservation(AbstractObservation):
+    """Observation of a file on a node in the network."""
+
     def __init__(self, where: Optional[Tuple[str]] = None) -> None:
         """
-        _summary_
+        Initialise file observation.
 
         :param where: Store information about where in the simulation state dictionary to find the relevatn information.
             Optional. If None, this corresponds that the file does not exist and the observation will be populated with
@@ -60,6 +66,13 @@ class FileObservation(AbstractObservation):
         "Default observation is what should be returned when the file doesn't exist, e.g. after it has been deleted."
 
     def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        :return: Observation
+        :rtype: Dict
+        """
         if self.where is None:
             return self.default_observation
         file_state = access_from_nested_dict(state, self.where)
@@ -69,19 +82,38 @@ class FileObservation(AbstractObservation):
 
     @property
     def space(self) -> spaces.Space:
+        """Gymnasium space object describing the observation space shape.
+
+        :return: Gymnasium space
+        :rtype: spaces.Space
+        """
         return spaces.Dict({"health_status": spaces.Discrete(6)})
 
     @classmethod
-    def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where=None):
+    def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: List[str] = None) -> "FileObservation":
+        """Create file observation from a config.
+
+        :param config: Dictionary containing the configuration for this file observation.
+        :type config: Dict
+        :param session: _description_
+        :type session: PrimaiteSession
+        :param parent_where: _description_, defaults to None
+        :type parent_where: _type_, optional
+        :return: _description_
+        :rtype: _type_
+        """
         return cls(where=parent_where + ["files", config["file_name"]])
 
 
 class ServiceObservation(AbstractObservation):
+    """Observation of a service in the network."""
+
     default_observation: spaces.Space = {"operating_status": 0, "health_status": 0}
     "Default observation is what should be returned when the service doesn't exist."
 
     def __init__(self, where: Optional[Tuple[str]] = None) -> None:
-        """
+        """Initialise service observation.
+
         :param where: Store information about where in the simulation state dictionary to find the relevant information.
             Optional. If None, this corresponds that the file does not exist and the observation will be populated with
             zeroes.
@@ -94,6 +126,13 @@ class ServiceObservation(AbstractObservation):
         self.where: Optional[Tuple[str]] = where
 
     def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        :return: Observation
+        :rtype: Dict
+        """
         if self.where is None:
             return self.default_observation
 
@@ -104,19 +143,36 @@ class ServiceObservation(AbstractObservation):
 
     @property
     def space(self) -> spaces.Space:
+        """Gymnasium space object describing the observation space shape."""
         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
+    ) -> "ServiceObservation":
+        """Create service observation from a config.
+
+        :param config: Dictionary containing the configuration for this service observation.
+        :type config: Dict
+        :param session: Reference to the PrimaiteSession object that spawned this observation.
+        :type session: PrimaiteSession
+        :param parent_where: Where in the simulation state dictionary this service's parent node is located. Optional.
+        :type parent_where: Optional[List[str]], optional
+        :return: Constructed service observation
+        :rtype: ServiceObservation
+        """
         return cls(where=parent_where + ["services", session.ref_map_services[config["service_ref"]].uuid])
 
 
 class LinkObservation(AbstractObservation):
+    """Observation of a link in the network."""
+
     default_observation: spaces.Space = {"protocols": {"all": {"load": 0}}}
     "Default observation is what should be returned when the link doesn't exist."
 
     def __init__(self, where: Optional[Tuple[str]] = None) -> None:
-        """
+        """Initialise link observation.
+
         :param where: Store information about where in the simulation state dictionary to find the relevant information.
             Optional. If None, this corresponds that the file does not exist and the observation will be populated with
             zeroes.
@@ -129,6 +185,13 @@ class LinkObservation(AbstractObservation):
         self.where: Optional[Tuple[str]] = where
 
     def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        :return: Observation
+        :rtype: Dict
+        """
         if self.where is None:
             return self.default_observation
 
@@ -147,15 +210,33 @@ class LinkObservation(AbstractObservation):
 
     @property
     def space(self) -> spaces.Space:
+        """Gymnasium space object describing the observation space shape.
+
+        :return: Gymnasium space
+        :rtype: spaces.Space
+        """
         return spaces.Dict({"protocols": spaces.Dict({"all": spaces.Dict({"load": spaces.Discrete(11)})})})
 
     @classmethod
-    def from_config(cls, config: Dict, session: "PrimaiteSession"):
+    def from_config(cls, config: Dict, session: "PrimaiteSession") -> "LinkObservation":
+        """Create link observation from a config.
+
+        :param config: Dictionary containing the configuration for this link observation.
+        :type config: Dict
+        :param session: Reference to the PrimaiteSession object that spawned this observation.
+        :type session: PrimaiteSession
+        :return: Constructed link observation
+        :rtype: LinkObservation
+        """
         return cls(where=["network", "links", session.ref_map_links[config["link_ref"]]])
 
 
 class FolderObservation(AbstractObservation):
-    def __init__(self, where: Optional[Tuple[str]] = None, files: List[FileObservation] = [], num_files_per_folder:int=2) -> None:
+    """Folder observation, including files inside of the folder."""
+
+    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.
@@ -179,7 +260,7 @@ class FolderObservation(AbstractObservation):
         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:
+        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)
@@ -191,6 +272,13 @@ class FolderObservation(AbstractObservation):
         }
 
     def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        :return: Observation
+        :rtype: Dict
+        """
         if self.where is None:
             return self.default_observation
         folder_state = access_from_nested_dict(state, self.where)
@@ -208,6 +296,11 @@ class FolderObservation(AbstractObservation):
 
     @property
     def space(self) -> spaces.Space:
+        """Gymnasium space object describing the observation space shape.
+
+        :return: Gymnasium space
+        :rtype: spaces.Space
+        """
         return spaces.Dict(
             {
                 "health_status": spaces.Discrete(6),
@@ -216,7 +309,26 @@ class FolderObservation(AbstractObservation):
         )
 
     @classmethod
-    def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]], num_files_per_folder:int=2):
+    def from_config(
+        cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]], num_files_per_folder: int = 2
+    ) -> "FolderObservation":
+        """Create folder observation from a config. Also creates child file observations.
+
+        :param config: Dictionary containing the configuration for this folder observation. Includes the name of the
+            folder and the files inside of it.
+        :type config: Dict
+        :param session: Reference to the PrimaiteSession object that spawned this observation.
+        :type session: PrimaiteSession
+        :param parent_where: Where in the simulation state dictionary to find the information about this folder's
+            parent node. A typical location for a node ``where`` can be:
+            ['network','nodes',<node_uuid>,'file_system']
+        :type parent_where: Optional[List[str]]
+        :param num_files_per_folder: How many spaces for files are in this folder observation (to preserve static
+            observation size) , defaults to 2
+        :type num_files_per_folder: int, optional
+        :return: Constructed folder observation
+        :rtype: FolderObservation
+        """
         where = parent_where + ["folders", config["folder_name"]]
 
         file_configs = config["files"]
@@ -226,13 +338,30 @@ class FolderObservation(AbstractObservation):
 
 
 class NicObservation(AbstractObservation):
+    """Observation of a Network Interface Card (NIC) in the network."""
+
     default_observation: spaces.Space = {"nic_status": 0}
 
     def __init__(self, where: Optional[Tuple[str]] = None) -> None:
+        """Initialise NIC observation.
+
+        :param where: Where in the simulation state dictionary to find the relevant information for this NIC. A typical
+            example may look like this:
+            ['network','nodes',<node_uuid>,'NICs',<nic_uuid>]
+            If None, this denotes that the NIC does not exist and the observation will be populated with zeroes.
+        :type where: Optional[Tuple[str]], optional
+        """
         super().__init__()
         self.where: Optional[Tuple[str]] = where
 
     def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        :return: Observation
+        :rtype: Dict
+        """
         if self.where is None:
             return self.default_observation
         nic_state = access_from_nested_dict(state, self.where)
@@ -243,14 +372,31 @@ class NicObservation(AbstractObservation):
 
     @property
     def space(self) -> spaces.Space:
+        """Gymnasium space object describing the observation space shape."""
         return spaces.Dict({"nic_status": spaces.Discrete(3)})
 
     @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]]
+    ) -> "NicObservation":
+        """Create NIC observation from a config.
+
+        :param config: Dictionary containing the configuration for this NIC observation.
+        :type config: Dict
+        :param session: Reference to the PrimaiteSession object that spawned this observation.
+        :type session: PrimaiteSession
+        :param parent_where: Where in the simulation state dictionary to find the information about this NIC's parent
+            node. A typical location for a node ``where`` can be: ['network','nodes',<node_uuid>]
+        :type parent_where: Optional[List[str]]
+        :return: Constructed NIC observation
+        :rtype: NicObservation
+        """
         return cls(where=parent_where + ["NICs", config["nic_uuid"]])
 
 
 class NodeObservation(AbstractObservation):
+    """Observation of a node in the network. Includes services, folders and NICs."""
+
     def __init__(
         self,
         where: Optional[Tuple[str]] = None,
@@ -260,8 +406,8 @@ class NodeObservation(AbstractObservation):
         logon_status: bool = False,
         num_services_per_node: int = 2,
         num_folders_per_node: int = 2,
-        num_files_per_folder: int = 2
+        num_files_per_folder: int = 2,
-        ) -> None:
+    ) -> None:
         """
         Configurable observation for a node in the simulation.
 
@@ -271,7 +417,8 @@ class NodeObservation(AbstractObservation):
         :type where: List[str], optional
         :param services: Mapping between position in observation space and service UUID, defaults to {}
         :type services: Dict[int,str], optional
-        :param max_services: Max number of services that can be presented in observation space for this node, defaults to 2
+        :param max_services: Max number of services that can be presented in observation space for this node
+            , defaults to 2
         :type max_services: int, optional
         :param folders: Mapping between position in observation space and folder name, defaults to {}
         :type folders: Dict[int,str], optional
@@ -286,10 +433,10 @@ class NodeObservation(AbstractObservation):
         self.where: Optional[Tuple[str]] = where
 
         self.services: List[ServiceObservation] = services
-        while len(self.services)<num_services_per_node:
+        while len(self.services) < num_services_per_node:
-            # add an empty service observation without `where` parameter that will always return default (blank) observations
+            # add empty service observation without `where` parameter so it always returns default (blank) observation
             self.services.append(ServiceObservation())
-        while len(self.services)>num_services_per_node:
+        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)
@@ -297,8 +444,8 @@ class NodeObservation(AbstractObservation):
             # truncate service list
 
         self.folders: List[FolderObservation] = folders
+        # add empty folder observation without `where` parameter that will always return default (blank) observations
         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()
@@ -317,6 +464,13 @@ class NodeObservation(AbstractObservation):
             self.default_observation["logon_status"] = 0
 
     def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        :return: Observation
+        :rtype: Dict
+        """
         if self.where is None:
             return self.default_observation
 
@@ -337,6 +491,7 @@ class NodeObservation(AbstractObservation):
 
     @property
     def space(self) -> spaces.Space:
+        """Gymnasium space object describing the observation space shape."""
         space_shape = {
             "SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}),
             "FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}),
@@ -358,6 +513,27 @@ class NodeObservation(AbstractObservation):
         num_folders_per_node: int = 2,
         num_files_per_folder: int = 2,
     ) -> "NodeObservation":
+        """Create node observation from a config. Also creates child service, folder and NIC observations.
+
+        :param config: Dictionary containing the configuration for this node observation.
+        :type config: Dict
+        :param session: Reference to the PrimaiteSession object that spawned this observation.
+        :type session: PrimaiteSession
+        :param parent_where: Where in the simulation state dictionary to find the information about this node's parent
+            network. A typical location for it would be: ['network',]
+        :type parent_where: Optional[List[str]]
+        :param num_services_per_node: How many spaces for services are in this node observation (to preserve static
+            observation size) , defaults to 2
+        :type num_services_per_node: int, optional
+        :param num_folders_per_node: How many spaces for folders are in this node observation (to preserve static
+            observation size) , defaults to 2
+        :type num_folders_per_node: int, optional
+        :param num_files_per_folder: How many spaces for files are in the folder observations (to preserve static
+            observation size) , defaults to 2
+        :type num_files_per_folder: int, optional
+        :return: Constructed node observation
+        :rtype: NodeObservation
+        """
         node_uuid = session.ref_map_nodes[config["node_ref"]]
         if parent_where is None:
             where = ["network", "nodes", node_uuid]
@@ -367,7 +543,12 @@ 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, num_files_per_folder=num_files_per_folder) 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]
@@ -378,13 +559,15 @@ class NodeObservation(AbstractObservation):
             folders=folders,
             nics=nics,
             logon_status=logon_status,
-            num_services_per_node = num_services_per_node,
+            num_services_per_node=num_services_per_node,
-            num_folders_per_node = num_folders_per_node,
+            num_folders_per_node=num_folders_per_node,
-            num_files_per_folder = num_files_per_folder,
+            num_files_per_folder=num_files_per_folder,
-            )
+        )
 
 
 class AclObservation(AbstractObservation):
+    """Observation of an Access Control List (ACL) in the network."""
+
     # TODO: should where be optional, and we can use where=None to pad the observation space?
     # definitely the current approach does not support tracking files that aren't specified by name, for example
     # if a file is created at runtime, we have currently got no way of telling the observation space to track it.
@@ -397,6 +580,21 @@ class AclObservation(AbstractObservation):
         where: Optional[Tuple[str]] = None,
         num_rules: int = 10,
     ) -> None:
+        """Initialise ACL observation.
+
+        :param node_ip_to_id: Mapping between IP address and ID.
+        :type node_ip_to_id: Dict[str, int]
+        :param ports: List of ports which are part of the game that define the ordering when converting to an ID
+        :type ports: List[int]
+        :param protocols: List of protocols which are part of the game, defines ordering when converting to an ID
+        :type protocols: list[str]
+        :param where: Where in the simulation state dictionary to find the relevant information for this ACL. A typical
+            example may look like this:
+            ['network','nodes',<router_uuid>,'acl','acl']
+        :type where: Optional[Tuple[str]], optional
+        :param num_rules: , defaults to 10
+        :type num_rules: int, optional
+        """
         super().__init__()
         self.where: Optional[Tuple[str]] = where
         self.num_rules: int = num_rules
@@ -423,6 +621,13 @@ class AclObservation(AbstractObservation):
         }
 
     def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        :return: Observation
+        :rtype: Dict
+        """
         if self.where is None:
             return self.default_observation
         acl_state: Dict = access_from_nested_dict(state, self.where)
@@ -457,6 +662,11 @@ class AclObservation(AbstractObservation):
 
     @property
     def space(self) -> spaces.Space:
+        """Gymnasium space object describing the observation space shape.
+
+        :return: Gymnasium space
+        :rtype: spaces.Space
+        """
         return spaces.Dict(
             {
                 "RULES": spaces.Dict(
@@ -482,6 +692,15 @@ class AclObservation(AbstractObservation):
 
     @classmethod
     def from_config(cls, config: Dict, session: "PrimaiteSession") -> "AclObservation":
+        """Generate ACL observation from a config.
+
+        :param config: Dictionary containing the configuration for this ACL observation.
+        :type config: Dict
+        :param session: Reference to the PrimaiteSession object that spawned this observation.
+        :type session: PrimaiteSession
+        :return: Observation object
+        :rtype: AclObservation
+        """
         node_ip_to_idx = {}
         for ip_idx, ip_map_config in enumerate(config["ip_address_order"]):
             node_ref = ip_map_config["node_ref"]
@@ -500,26 +719,44 @@ class AclObservation(AbstractObservation):
 
 
 class NullObservation(AbstractObservation):
+    """Null observation, returns a single 0 value for the observation space."""
+
     def __init__(self, where: Optional[List[str]] = None):
+        """Initialise null observation."""
         self.default_observation: Dict = {}
 
     def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation."""
         return 0
 
     @property
     def space(self) -> spaces.Space:
+        """Gymnasium space object describing the observation space shape."""
         return spaces.Discrete(1)
 
     @classmethod
     def from_config(cls, config: Dict, session: Optional["PrimaiteSession"] = None) -> "NullObservation":
+        """
+        Create null observation from a config.
+
+        The parameters are ignored, they are here to match the signature of the other observation classes.
+        """
         return cls()
 
 
 class ICSObservation(NullObservation):
+    """ICS observation placeholder, currently not implemented so always returns a single 0."""
+
     pass
 
 
 class UC2BlueObservation(AbstractObservation):
+    """Container for all observations used by the blue agent in UC2.
+
+    TODO: there's no real need for a UC2 blue container class, we should be able to simply use the observation handler
+        for the purpose of compiling several observation components.
+    """
+
     def __init__(
         self,
         nodes: List[NodeObservation],
@@ -528,6 +765,20 @@ class UC2BlueObservation(AbstractObservation):
         ics: ICSObservation,
         where: Optional[List[str]] = None,
     ) -> None:
+        """Initialise UC2 blue observation.
+
+        :param nodes: List of node observations
+        :type nodes: List[NodeObservation]
+        :param links: List of link observations
+        :type links: List[LinkObservation]
+        :param acl: The Access Control List observation
+        :type acl: AclObservation
+        :param ics: The ICS observation
+        :type ics: ICSObservation
+        :param where: Where in the simulation state dict to find information. Not used in this particular observation
+            because it only compiles other observations and doesn't contribute any new information, defaults to None
+        :type where: Optional[List[str]], optional
+        """
         super().__init__()
         self.where: Optional[Tuple[str]] = where
 
@@ -544,6 +795,13 @@ class UC2BlueObservation(AbstractObservation):
         }
 
     def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        :return: Observation
+        :rtype: Dict
+        """
         if self.where is None:
             return self.default_observation
 
@@ -557,6 +815,12 @@ class UC2BlueObservation(AbstractObservation):
 
     @property
     def space(self) -> spaces.Space:
+        """
+        Gymnasium space object describing the observation space shape.
+
+        :return: Space
+        :rtype: spaces.Space
+        """
         return spaces.Dict(
             {
                 "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}),
@@ -567,21 +831,34 @@ class UC2BlueObservation(AbstractObservation):
         )
 
     @classmethod
-    def from_config(cls, config: Dict, session: "PrimaiteSession"):
+    def from_config(cls, config: Dict, session: "PrimaiteSession") -> "UC2BlueObservation":
+        """Create UC2 blue observation from a config.
+
+        :param config: Dictionary containing the configuration for this UC2 blue observation. This includes the nodes,
+            links, ACL and ICS observations.
+        :type config: Dict
+        :param session: Reference to the PrimaiteSession object that spawned this observation.
+        :type session: PrimaiteSession
+        :return: Constructed UC2 blue observation
+        :rtype: UC2BlueObservation
+        """
         node_configs = config["nodes"]
         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(
+        nodes = [
-            config=n,
+            NodeObservation.from_config(
-            session=session,
+                config=n,
-            num_services_per_node= num_services_per_node,
+                session=session,
-            num_folders_per_node=num_folders_per_node,
+                num_services_per_node=num_services_per_node,
-            num_files_per_folder=num_files_per_folder,
+                num_folders_per_node=num_folders_per_node,
-            ) for n in node_configs]
+                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]
+        links = [LinkObservation.from_config(config=link, session=session) for link in link_configs]
 
         acl_config = config["acl"]
         acl = AclObservation.from_config(config=acl_config, session=session)
@@ -593,6 +870,8 @@ class UC2BlueObservation(AbstractObservation):
 
 
 class UC2RedObservation(AbstractObservation):
+    """Container for all observations used by the red agent in UC2."""
+
     def __init__(self, nodes: List[NodeObservation], where: Optional[List[str]] = None) -> None:
         super().__init__()
         self.where: Optional[List[str]] = where
@@ -603,6 +882,7 @@ class UC2RedObservation(AbstractObservation):
         }
 
     def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation."""
         if self.where is None:
             return self.default_observation
 
@@ -612,6 +892,7 @@ class UC2RedObservation(AbstractObservation):
 
     @property
     def space(self) -> spaces.Space:
+        """Gymnasium space object describing the observation space shape."""
         return spaces.Dict(
             {
                 "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}),
@@ -619,43 +900,74 @@ class UC2RedObservation(AbstractObservation):
         )
 
     @classmethod
-    def from_config(cls, config: Dict, session: "PrimaiteSession"):
+    def from_config(cls, config: Dict, session: "PrimaiteSession") -> "UC2RedObservation":
+        """
+        Create UC2 red observation from a config.
+
+        :param config: Dictionary containing the configuration for this UC2 red observation.
+        :type config: Dict
+        :param session: Reference to the PrimaiteSession object that spawned this observation.
+        :type 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"])
 
 
 class UC2GreenObservation(NullObservation):
+    """Green agent observation. As the green agent's actions don't depend on the observation, this is empty."""
+
     pass
 
 
 class ObservationSpace:
     """
-    Manage the observations of an Actor.
+    Manage the observations of an Agent.
 
     The observation space has the purpose of:
       1. Reading the outputted state from the PrimAITE Simulation.
       2. Selecting parts of the simulation state that are requested by the simulation config
-      3. Formatting this information so an actor can use it to make decisions.
+      3. Formatting this information so an agent can use it to make decisions.
     """
 
-    ...
+    # TODO: Dear code reader: This class currently doesn't do much except hold an observation object. It will be changed
+    # to have more of it's own behaviour, and it will replace UC2BlueObservation and UC2RedObservation during the next
+    # refactor.
 
-    # what this class does:
-    # keep a list of observations
-    # create observations for an actor from the config
     def __init__(self, observation: AbstractObservation) -> None:
+        """Initialise observation space.
+
+        :param observation: Observation object
+        :type observation: AbstractObservation
+        """
         self.obs: AbstractObservation = observation
 
-    def observe(self, state) -> Dict:
+    def observe(self, state: Dict) -> Dict:
+        """
+        Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        """
         return self.obs.observe(state)
 
     @property
     def space(self) -> None:
+        """Gymnasium space object describing the observation space shape."""
         return self.obs.space
 
     @classmethod
     def from_config(cls, config: Dict, session: "PrimaiteSession") -> "ObservationSpace":
+        """Create observation space from a config.
+
+        :param config: Dictionary containing the configuration for this observation space.
+            It should contain the key 'type' which selects which observation class to use (from a choice of:
+            UC2BlueObservation, UC2RedObservation, UC2GreenObservation)
+            The other key is 'options' which are passed to the constructor of the selected observation class.
+        :type config: Dict
+        :param session: Reference to the PrimaiteSession object that spawned this observation.
+        :type session: PrimaiteSession
+        """
         if config["type"] == "UC2BlueObservation":
             return cls(UC2BlueObservation.from_config(config.get("options", {}), session=session))
         elif config["type"] == "UC2RedObservation":

src/primaite/game/agent/rewards.py

@@ -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"]:

src/primaite/game/agent/scripted_agents.py

@@ -1,9 +1,14 @@
+"""Agents with predefined behaviours."""
 from primaite.game.agent.interface import AbstractScriptedAgent
 
 
 class GreenWebBrowsingAgent(AbstractScriptedAgent):
-    ...
+    """Scripted agent which attempts to send web requests to a target node."""
+
+    raise NotImplementedError
 
 
 class RedDatabaseCorruptingAgent(AbstractScriptedAgent):
-    ...
+    """Scripted agent which attempts to corrupt the database of the target node."""
+
+    raise NotImplementedError

src/primaite/game/agent/utils.py

@@ -1,4 +1,4 @@
-from typing import Dict, Sequence, Hashable, Any
+from typing import Any, Dict, Hashable, Sequence
 
 NOT_PRESENT_IN_STATE = object()
 """
@@ -6,6 +6,7 @@ Need an object to return when the sim state does not contain a requested value.
 the thing requested in the state could equal None. This NOT_PRESENT_IN_STATE is a sentinel for this purpose.
 """
 
+
 def access_from_nested_dict(dictionary: Dict, keys: Sequence[Hashable]) -> Any:
     """
     Access an item from a deeply dictionary with a list of keys.
@@ -26,4 +27,4 @@ def access_from_nested_dict(dictionary: Dict, keys: Sequence[Hashable]) -> Any:
     k = key_list.pop(0)
     if k not in dictionary:
         return NOT_PRESENT_IN_STATE
-    return access_from_nested_dict(dictionary[k], key_list)
+    return access_from_nested_dict(dictionary[k], key_list)

src/primaite/game/session.py

@@ -1,6 +1,6 @@
 """PrimAITE session - the main entry point to training agents on PrimAITE."""
 from ipaddress import IPv4Address
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Never, Optional, Tuple
 
 from arcd_gate.client.gate_client import ActType, GATEClient
 from gymnasium import spaces
@@ -32,13 +32,17 @@ _LOGGER = getLogger(__name__)
 
 
 class PrimaiteGATEClient(GATEClient):
+    """Lightweight wrapper around the GATEClient class that allows PrimAITE to message GATE."""
+
     def __init__(self, parent_session: "PrimaiteSession", service_port: int = 50000):
-        """Create a new GATE client for PrimAITE.
+        """
+        Create a new GATE client for PrimAITE.
 
         :param parent_session: The parent session object.
         :type parent_session: PrimaiteSession
         :param service_port: The port on which the GATE service is running.
-        :type service_port: int, optional"""
+        :type service_port: int, optional
+        """
         super().__init__(service_port=service_port)
         self.parent_session: "PrimaiteSession" = parent_session
 
@@ -133,9 +137,11 @@ class PrimaiteGATEClient(GATEClient):
 
 
 class PrimaiteSessionOptions(BaseModel):
-    """Global options which are applicable to all of the agents in the game.
+    """
+    Global options which are applicable to all of the agents in the game.
 
-    Currently this is used to restrict which ports and protocols exist in the world of the simulation."""
+    Currently this is used to restrict which ports and protocols exist in the world of the simulation.
+    """
 
     ports: List[str]
     protocols: List[str]
@@ -154,9 +160,7 @@ class TrainingOptions(BaseModel):
 
 
 class PrimaiteSession:
-    """
+    """The main entrypoint for PrimAITE sessions, this manages a simulation, agents, and connections to ARCD GATE."""
-    The main entrypoint for PrimAITE sessions, this coordinates a simulation, agents, and connections to ARCD GATE.
-    """
 
     def __init__(self):
         self.simulation: Simulation = Simulation()
@@ -183,7 +187,7 @@ class PrimaiteSession:
         self.gate_client: PrimaiteGATEClient = PrimaiteGATEClient(self)
         """Reference to a GATE Client object, which will send data to GATE service for training RL agent."""
 
-    def start_session(self, opts="TODO..."):
+    def start_session(self) -> Never:
         """Commence the training session, this gives the GATE client control over the simulation/agent loop."""
         self.gate_client.start()
 
@@ -221,7 +225,7 @@ class PrimaiteSession:
             #    to discrete(40) is only necessary for purposes of RL learning, therefore that bit of
             #    code should live inside of the GATE agent subclass)
             # gets action in CAOS format
-            _LOGGER.debug(f"Getting agent action")
+            _LOGGER.debug("Getting agent action")
             agent_action, action_options = agent.get_action(agent_obs, agent_reward)
             # 9. CAOS action is converted into request (extra information might be needed to enrich
             # the request, this is what the execution definition is there for)
@@ -236,11 +240,11 @@ class PrimaiteSession:
         self.simulation.apply_timestep(self.step_counter)
         self.step_counter += 1
 
-    def reset(self):
+    def reset(self) -> None:
         """Reset the session, this will reset the simulation."""
         return NotImplemented
 
-    def close(self):
+    def close(self) -> None:
         """Close the session, this will stop the gate client and close the simulation."""
         return NotImplemented