oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

901 - merged with dev

Browse Source
This commit is contained in:
SunilSamra
2023-07-05 11:34:15 +01:00
parent 766ee9624a 6006f022a1
commit 35045f4842
73 changed files with 4047 additions and 1376 deletions
Download Patch File Download Diff File Expand all files Collapse all files

247
tests/test_observation_space.py
View File

@@ -5,17 +5,12 @@ import time
import numpy as np
import pytest
from primaite.environment.observations import (
NodeLinkTable,
NodeStatuses,
ObservationsHandler,
)
from primaite.environment.primaite_env import Primaite
from primaite.environment.observations import NodeLinkTable, NodeStatuses, ObservationsHandler
from tests import TEST_CONFIG_ROOT
from tests.conftest import _get_primaite_env_from_config
def run_generic_set_actions(env: Primaite):
def run_generic_set_actions(env):
"""Run against a generic agent with specified blue agent actions."""
# Reset the environment at the start of the episode
env.reset()
@@ -70,61 +65,75 @@ def env(request):
yield env
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_without_obs.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
]
],
indirect=True,
)
def test_default_obs_space(env: Primaite):
def test_default_obs_space(temp_primaite_session):
"""Create environment with no obs space defined in config and check that the default obs space was created."""
env.update_environent_obs()
with temp_primaite_session as session:
session.env.update_environent_obs()
components = env.obs_handler.registered_obs_components
components = session.env.obs_handler.registered_obs_components
assert len(components) == 1
assert isinstance(components[0], NodeLinkTable)
assert len(components) == 1
assert isinstance(components[0], NodeLinkTable)
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_without_obs.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
]
],
indirect=True,
)
def test_registering_components(env: Primaite):
def test_registering_components(temp_primaite_session):
"""Test regitering and deregistering a component."""
handler = ObservationsHandler()
component = NodeStatuses(env)
handler.register(component)
assert component in handler.registered_obs_components
handler.deregister(component)
assert component not in handler.registered_obs_components
with temp_primaite_session as session:
env = session.env
handler = ObservationsHandler()
component = NodeStatuses(env)
handler.register(component)
assert component in handler.registered_obs_components
handler.deregister(component)
assert component not in handler.registered_obs_components
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_NODE_LINK_TABLE.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "obs_tests/main_config_NODE_LINK_TABLE.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
]
],
indirect=True,
)
class TestNodeLinkTable:
"""Test the NodeLinkTable observation component (in isolation)."""
def test_obs_shape(self, env: Primaite):
def test_obs_shape(self, temp_primaite_session):
"""Try creating env with box observation space."""
env.update_environent_obs()
with temp_primaite_session as session:
env = session.env
env.update_environent_obs()
# we have three nodes and two links, with two service
# therefore the box observation space will have:
# * 5 rows (3 nodes + 2 links)
# * 6 columns (four fixed and two for the services)
assert env.env_obs.shape == (5, 6)
# we have three nodes and two links, with two service
# therefore the box observation space will have:
# * 5 rows (3 nodes + 2 links)
# * 6 columns (four fixed and two for the services)
assert env.env_obs.shape == (5, 6)
def test_value(self, env: Primaite):
def test_value(self, temp_primaite_session):
"""Test that the observation is generated correctly.
The laydown has:
@@ -170,36 +179,43 @@ class TestNodeLinkTable:
* 999 (999 traffic service1)
* 0 (no traffic for service2)
"""
# act = np.asarray([0,])
obs, reward, done, info = env.step(0) # apply the 'do nothing' action
with temp_primaite_session as session:
env = session.env
# act = np.asarray([0,])
obs, reward, done, info = env.step(0) # apply the 'do nothing' action
assert np.array_equal(
obs,
[
[1, 1, 3, 1, 1, 1, 0],
[2, 1, 1, 1, 1, 4, 1],
[3, 1, 1, 1, 0, 0, 2],
[4, 0, 0, 0, 999, 0, 3],
[5, 0, 0, 0, 999, 0, 4],
],
)
assert np.array_equal(
obs,
[
[1, 1, 3, 1, 1, 1],
[2, 1, 1, 1, 1, 4],
[3, 1, 1, 1, 0, 0],
[4, 0, 0, 0, 999, 0],
[5, 0, 0, 0, 999, 0],
],
)
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_NODE_STATUSES.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "obs_tests/main_config_NODE_STATUSES.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
]
],
indirect=True,
)
class TestNodeStatuses:
"""Test the NodeStatuses observation component (in isolation)."""
def test_obs_shape(self, env: Primaite):
def test_obs_shape(self, temp_primaite_session):
"""Try creating env with NodeStatuses as the only component."""
assert env.env_obs.shape == (15,)
with temp_primaite_session as session:
env = session.env
assert env.env_obs.shape == (15,)
def test_values(self, env: Primaite):
def test_values(self, temp_primaite_session):
"""Test that the hardware and software states are encoded correctly.
The laydown has:
@@ -226,28 +242,36 @@ class TestNodeStatuses:
* service 1 = n/a (0)
* service 2 = n/a (0)
"""
obs, _, _, _ = env.step(0) # apply the 'do nothing' action
assert np.array_equal(obs, [1, 3, 1, 1, 1, 1, 1, 1, 1, 4, 1, 1, 1, 0, 0])
with temp_primaite_session as session:
env = session.env
obs, _, _, _ = env.step(0) # apply the 'do nothing' action
print(obs)
assert np.array_equal(obs, [1, 3, 1, 1, 1, 1, 1, 1, 1, 4, 1, 1, 1, 0, 0])
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
]
],
indirect=True,
)
class TestLinkTrafficLevels:
"""Test the LinkTrafficLevels observation component (in isolation)."""
def test_obs_shape(self, env: Primaite):
def test_obs_shape(self, temp_primaite_session):
"""Try creating env with MultiDiscrete observation space."""
env.update_environent_obs()
with temp_primaite_session as session:
env = session.env
env.update_environent_obs()
# we have two links and two services, so the shape should be 2 * 2
assert env.env_obs.shape == (2 * 2,)
# we have two links and two services, so the shape should be 2 * 2
assert env.env_obs.shape == (2 * 2,)
def test_values(self, env: Primaite):
def test_values(self, temp_primaite_session):
"""Test that traffic values are encoded correctly.
The laydown has:
@@ -257,36 +281,42 @@ class TestLinkTrafficLevels:
* an IER trying to send 999 bits of data over both links the whole time (via the first service)
* link bandwidth of 1000, therefore the utilisation is 99.9%
"""
obs, reward, done, info = env.step(0)
obs, reward, done, info = env.step(0)
with temp_primaite_session as session:
env = session.env
obs, reward, done, info = env.step(0)
obs, reward, done, info = env.step(0)
# the observation space has combine_service_traffic set to False, so the space has this format:
# [link1_service1, link1_service2, link2_service1, link2_service2]
# we send 999 bits of data via link1 and link2 on service 1.
# therefore the first and third elements should be 6 and all others 0
# (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%)
print(obs)
assert np.array_equal(obs, [6, 0, 6, 0])
# the observation space has combine_service_traffic set to False, so the space has this format:
# [link1_service1, link1_service2, link2_service1, link2_service2]
# we send 999 bits of data via link1 and link2 on service 1.
# therefore the first and third elements should be 6 and all others 0
# (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%)
assert np.array_equal(obs, [6, 0, 6, 0])
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_ACCESS_CONTROL_LIST.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "single_action_space_fixed_blue_actions_main_config.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown_ACL.yaml",
]
],
indirect=True,
)
class TestAccessControlList:
"""Test the AccessControlList observation component (in isolation)."""
def test_obs_shape(self, env: Primaite):
def test_obs_shape(self, temp_primaite_session):
"""Try creating env with MultiDiscrete observation space."""
env.update_environent_obs()
with temp_primaite_session as session:
env = session.env
env.update_environent_obs()
# we have two ACLs
assert env.env_obs.shape == (6 * 3)
def test_values(self, env: Primaite):
def test_values(self, temp_primaite_session):
"""Test that traffic values are encoded correctly.
The laydown has:
@@ -296,8 +326,10 @@ class TestAccessControlList:
* an IER trying to send 999 bits of data over both links the whole time (via the first service)
* link bandwidth of 1000, therefore the utilisation is 99.9%
"""
obs, reward, done, info = env.step(0)
obs, reward, done, info = env.step(0)
with temp_primaite_session as session:
env = session.env
obs, reward, done, info = env.step(0)
obs, reward, done, info = env.step(0)
# the observation space has combine_service_traffic set to False, so the space has this format:
# [link1_service1, link1_service2, link2_service1, link2_service2]
@@ -307,16 +339,13 @@ class TestAccessControlList:
print(obs)
assert np.array_equal(obs, [])
def test_observation_space_with_implicit_rule(self):
def test_observation_space_with_implicit_rule(self, temp_primaite_session):
"""Test observation space is what is expected when an agent adds ACLs during an episode."""
# Used to use env from test fixture but AtrributeError function object has no 'training_config'
env = _get_primaite_env_from_config(
training_config_path=TEST_CONFIG_ROOT
/ "single_action_space_fixed_blue_actions_main_config.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown_ACL.yaml",
)
run_generic_set_actions(env)
obs = env.env_obs
with temp_primaite_session as session:
env = session.env
run_generic_set_actions(env)
obs = env.env_obs
"""
Observation space at the end of the episode.
At the start of the episode, there is a single implicit Deny rule = 1,1,1,1,1,0
@@ -327,6 +356,4 @@ class TestAccessControlList:
"""
# assert current_obs == [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]
assert np.array_equal(
obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]
)
assert np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2])