From ae2f4d472ec78d67c4268352cd773fe097572bf2 Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Wed, 31 May 2023 14:11:15 +0100 Subject: [PATCH] 1443 - reverted changes made to observation space and added config files for testing --- src/primaite/environment/primaite_env.py | 38 ------------------- .../single_action_space_lay_down_config.yaml | 29 ++++++++++++++ ...l => single_action_space_main_config.yaml} | 0 tests/test_single_action_space.py | 4 +- 4 files changed, 31 insertions(+), 40 deletions(-) create mode 100644 tests/config/single_action_space_lay_down_config.yaml rename tests/config/{single_action_space_config.yaml => single_action_space_main_config.yaml} (100%) diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index 84b485bd..49c45f3e 100644 --- a/src/primaite/environment/primaite_env.py +++ b/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: diff --git a/tests/config/single_action_space_lay_down_config.yaml b/tests/config/single_action_space_lay_down_config.yaml new file mode 100644 index 00000000..058b790b --- /dev/null +++ b/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 diff --git a/tests/config/single_action_space_config.yaml b/tests/config/single_action_space_main_config.yaml similarity index 100% rename from tests/config/single_action_space_config.yaml rename to tests/config/single_action_space_main_config.yaml diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py index 3ec1dc2e..8c87d57b 100644 --- a/tests/test_single_action_space.py +++ b/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)