Separate obs functions and provide docstrings

src/primaite/environment/primaite_env.py

@@ -641,19 +641,29 @@ class Primaite(Env):
             else:
                 pass
 
-    def init_observations(self) -> Tuple[spaces.Space, np.ndarray]:
+    def _init_box_observations(self) -> Tuple[spaces.Space, np.ndarray]:
-        """Build the observation space based on network laydown and provide initial obs.
+        """Initialise the observation space with the BOX option chosen.
 
-        This method uses the object's `num_links`, `num_nodes`, `num_services`,
+        This will create the observation space formatted as a table of integers.
-        `OBSERVATION_SPACE_FIXED_PARAMETERS`, `OBSERVATION_SPACE_HIGH_VALUE`, and `observation_type`
+        There is one row per node, followed by one row per link.
-        attributes to figure out the correct shape and format for the observation space.
+        Columns are as follows:
+            * node/link ID
+            * node hardware status / 0 for links
+            * node operating system status (if active/service) / 0 for links
+            * node file system status (active/service only) / 0 for links
+            * node service1 status / traffic load from that service for links
+            * node service2 status / traffic load from that service for links
+            * ...
+            * node serviceN status / traffic load from that service for links
 
-        :return: Gym observation space
+        For example if the environment has 5 nodes, 7 links, and 3 services, the observation space shape will be
-        :rtype: gym.spaces.Space
+        ``(12, 7)``
+
+        :return: Box gym observation
+        :rtype: gym.spaces.Box
         :return: Initial observation with all entires set to 0
         :rtype: numpy.Array
         """
-        if self.observation_type == ObservationType.BOX:
         _LOGGER.info("Observation space type BOX selected")
 
         # 1. Determine observation shape from laydown
@@ -672,7 +682,33 @@ class Primaite(Env):
         )
         initial_observation = np.zeros(observation_shape, dtype=np.int64)
 
-        elif self.observation_type == ObservationType.MULTIDISCRETE:
+        return observation_space, initial_observation
+
+    def _init_multidiscrete_observations(self) -> Tuple[spaces.Space, np.ndarray]:
+        """Initialise the observation space with the MULTIDISCRETE option chosen.
+
+        This will create the observation space with node observations followed by link observations.
+        Each node has 3 elements in the observation space plus 1 per service, more specifically:
+            * hardware state
+            * operating system state
+            * file system state
+            * service states (one per service)
+        Each link has one element in the observation space, corresponding to the traffic load,
+        it can take the following values:
+            0 = No traffic (0% of bandwidth)
+            1 = No traffic (0%-33% of bandwidth)
+            2 = No traffic (33%-66% of bandwidth)
+            3 = No traffic (66%-100% of bandwidth)
+            4 = No traffic (100% of bandwidth)
+
+        For example if the environment has 5 nodes, 7 links, and 3 services, the observation space shape will be
+        ``(37,)``
+
+        :return: MultiDiscrete gym observation
+        :rtype: gym.spaces.MultiDiscrete
+        :return: Initial observation with all entires set to 0
+        :rtype: numpy.Array
+        """
         _LOGGER.info("Observation space MULTIDISCRETE selected")
 
         # 1. Determine observation shape from laydown
@@ -691,17 +727,47 @@ class Primaite(Env):
         # 2. Create observation space & zeroed out sample from space.
         observation_space = spaces.MultiDiscrete(observation_shape)
         initial_observation = np.zeros(len(observation_shape), dtype=np.int64)
-        else:
-            raise ValueError(
-                f"Observation type must be {ObservationType.BOX} or {ObservationType.MULTIDISCRETE}"
-                f", got {self.observation_type} instead"
-            )
 
         return observation_space, initial_observation
 
-    def update_environent_obs(self):
+    def init_observations(self) -> Tuple[spaces.Space, np.ndarray]:
-        """Updates the observation space based on the node and link status."""
+        """Build the observation space based on network laydown and provide initial obs.
+
+        This method uses the object's `num_links`, `num_nodes`, `num_services`,
+        `OBSERVATION_SPACE_FIXED_PARAMETERS`, `OBSERVATION_SPACE_HIGH_VALUE`, and `observation_type`
+        attributes to figure out the correct shape and format for the observation space.
+
+        :raises ValueError: If the env's `observation_type` attribute is not set to a valid `enums.ObservationType`
+        :return: Gym observation space
+        :rtype: gym.spaces.Space
+        :return: Initial observation with all entires set to 0
+        :rtype: numpy.Array
+        """
         if self.observation_type == ObservationType.BOX:
+            observation_space, initial_observation = self._init_box_observations()
+            return observation_space, initial_observation
+        elif self.observation_type == ObservationType.MULTIDISCRETE:
+            (
+                observation_space,
+                initial_observation,
+            ) = self._init_multidiscrete_observations()
+            return observation_space, initial_observation
+        else:
+            errmsg = (
+                f"Observation type must be {ObservationType.BOX} or {ObservationType.MULTIDISCRETE}"
+                f", got {self.observation_type} instead"
+            )
+            _LOGGER.error(errmsg)
+            raise ValueError(errmsg)
+
+    def _update_env_obs_box(self):
+        """Update the environment's observation state based on the current status of nodes and links.
+
+        This function can only be called if the observation space setting is set to BOX.
+
+        :raises AssertionError: If this function is called when the environment has the incorrect ``observation_type``
+        """
+        assert self.observation_type == ObservationType.BOX
         item_index = 0
 
         # Do nodes first
@@ -744,7 +810,14 @@ class Primaite(Env):
                 protocol_index += 1
             item_index += 1
 
-        elif self.observation_type == ObservationType.MULTIDISCRETE:
+    def _update_env_obs_multidiscrete(self):
+        """Update the environment's observation state based on the current status of nodes and links.
+
+        This function can only be called if the observation space setting is set to MULTIDISCRETE.
+
+        :raises AssertionError: If this function is called when the environment has the incorrect ``observation_type``
+        """
+        assert self.observation_type == ObservationType.MULTIDISCRETE
         obs = []
         # 1. Set nodes
         # Each node has the following variables in the observation space:
@@ -806,6 +879,13 @@ class Primaite(Env):
 
         self.env_obs = np.asarray(obs)
 
+    def update_environent_obs(self):
+        """Updates the observation space based on the node and link status."""
+        if self.observation_type == ObservationType.BOX:
+            self._update_env_obs_box()
+        elif self.observation_type == ObservationType.MULTIDISCRETE:
+            self._update_env_obs_multidiscrete()
+
     def load_config(self):
         """Loads config data in order to build the environment configuration."""
         for item in self.config_data: