1443 - reverted changes made to observation space and added config files for testing

src/primaite/environment/primaite_env.py

@@ -207,7 +207,6 @@ class Primaite(Env):
 
         # Calculate the number of items that need to be included in the
         # observation space
-        """
         num_items = self.num_links + self.num_nodes
         # Set the number of observation parameters, being # of services plus id,
         # hardware state, file system state and SoftwareState (i.e. 4)
@@ -222,23 +221,6 @@ class Primaite(Env):
             shape=self.observation_shape,
             dtype=np.int64,
         )
-        """
-        self.num_observation_parameters = (
-            self.num_services + self.OBSERVATION_SPACE_FIXED_PARAMETERS
-        )
-        # Define the observation space:
-        # There are:
-        #  4 Operating States (ON/OFF/RESETTING) + NONE (0)
-        #  4 OS States (GOOD/PATCHING/COMPROMISED) + NONE
-        #  5 Service States (NONE/GOOD/PATCHING/COMPROMISED/OVERWHELMED) + NONE
-        # There can be any number of services
-        # There are 5 node types No Traffic, Low Traffic, Medium Traffic, High Traffic, Overloaded/max traffic
-        self.observation_space = spaces.MultiDiscrete(
-            ([4, 4] + [5] * self.num_services) * self.num_nodes + [5] * self.num_links
-        )
-
-        # Define the observation shape
-        self.observation_shape = self.observation_space.sample().shape
 
         # This is the observation that is sent back via the rest and step functions
         self.env_obs = np.zeros(self.observation_shape, dtype=np.int64)
@@ -696,19 +678,6 @@ class Primaite(Env):
 
     def update_environent_obs(self):
         """Updates the observation space based on the node and link status."""
-        # Convert back to more readable, original format
-        reshaped_nodes = self.env_obs[: -self.num_links].reshape(
-            self.num_nodes, self.num_services + 2
-        )
-
-        # Add empty links back and add node ID back
-        s = np.zeros(
-            [reshaped_nodes.shape[0] + self.num_links, reshaped_nodes.shape[1] + 1],
-            dtype=np.int64,
-        )
-        s[:, 0] = range(1, self.num_nodes + self.num_links + 1)  # Adding ID back
-        s[: self.num_nodes, 1:] = reshaped_nodes  # put values back in
-        self.env_obs = s
         item_index = 0
 
         # Do nodes first
@@ -751,13 +720,6 @@ class Primaite(Env):
                 protocol_index += 1
             item_index += 1
 
-        # Remove ID columns, remove links and flatten to 1D array
-        node_obs = self.env_obs[: self.num_nodes, 1:].flatten()
-        # Remove ID, remove all data except link traffic status
-        link_obs = self.env_obs[self.num_nodes :, 3:].flatten()
-        # Combine nodes and links
-        self.env_obs = np.append(node_obs, link_obs)
-
     def load_config(self):
         """Loads config data in order to build the environment configuration."""
         for item in self.config_data:

tests/config/single_action_space_lay_down_config.yaml

@@ -0,0 +1,29 @@
+- itemType: ACTIONS
+  type: NODE
+- itemType: STEPS
+  steps: 15
+- itemType: PORTS
+  portsList:
+  - port: '21'
+- itemType: SERVICES
+  serviceList:
+  - name: ftp
+- itemType: NODE
+  node_id: '1'
+  name: node
+  node_class: SERVICE
+  node_type: COMPUTER
+  priority: P1
+  hardware_state: 'ON'
+  ip_address: 192.168.0.1
+  software_state: GOOD
+  file_system_state: GOOD
+  services:
+  - name: ftp
+    port: '21'
+    state: GOOD
+- itemType: POSITION
+  positions:
+  - node: '1'
+    x_pos: 309
+    y_pos: 78

tests/test_single_action_space.py

@@ -5,8 +5,8 @@ from tests.conftest import _get_primaite_env_from_config
 def test_single_action_space():
     """Test to ensure the blue agent is using the ACL action space and is carrying out both kinds of operations."""
     env = _get_primaite_env_from_config(
-        main_config_path=TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml",
+        main_config_path=TEST_CONFIG_ROOT / "single_action_space_main_config.yaml",
         lay_down_config_path=TEST_CONFIG_ROOT
-        / "one_node_states_on_off_lay_down_config.yaml",
+        / "single_action_space_lay_down_config.yaml",
     )
     print("Average Reward:", env.average_reward)