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
|
# Calculate the number of items that need to be included in the
|
||||||
# observation space
|
# observation space
|
||||||
"""
|
|
||||||
num_items = self.num_links + self.num_nodes
|
num_items = self.num_links + self.num_nodes
|
||||||
# Set the number of observation parameters, being # of services plus id,
|
# Set the number of observation parameters, being # of services plus id,
|
||||||
# hardware state, file system state and SoftwareState (i.e. 4)
|
# hardware state, file system state and SoftwareState (i.e. 4)
|
||||||
@@ -222,23 +221,6 @@ class Primaite(Env):
|
|||||||
shape=self.observation_shape,
|
shape=self.observation_shape,
|
||||||
dtype=np.int64,
|
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
|
# 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)
|
self.env_obs = np.zeros(self.observation_shape, dtype=np.int64)
|
||||||
@@ -696,19 +678,6 @@ class Primaite(Env):
|
|||||||
|
|
||||||
def update_environent_obs(self):
|
def update_environent_obs(self):
|
||||||
"""Updates the observation space based on the node and link status."""
|
"""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
|
item_index = 0
|
||||||
|
|
||||||
# Do nodes first
|
# Do nodes first
|
||||||
@@ -751,13 +720,6 @@ class Primaite(Env):
|
|||||||
protocol_index += 1
|
protocol_index += 1
|
||||||
item_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):
|
def load_config(self):
|
||||||
"""Loads config data in order to build the environment configuration."""
|
"""Loads config data in order to build the environment configuration."""
|
||||||
for item in self.config_data:
|
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():
|
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."""
|
"""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(
|
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
|
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)
|
print("Average Reward:", env.average_reward)
|
||||||
|
|||||||
Reference in New Issue
Block a user