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

This commit is contained in:
SunilSamra
2023-05-31 14:11:15 +01:00
parent a8cc50a495
commit 81e9ddca9b
4 changed files with 31 additions and 40 deletions

View File

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

View File

@@ -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

View File

@@ -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)