1443 - reverted changes made to observation space and added config files for testing
This commit is contained in:
@@ -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:
|
||||
|
||||
29
tests/config/single_action_space_lay_down_config.yaml
Normal file
29
tests/config/single_action_space_lay_down_config.yaml
Normal 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
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user