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

Merge remote-tracking branch 'origin/dev' into bugfix/1587-hardcoded-agent

Browse Source
This commit is contained in:
Marek Wolan
2023-07-10 09:15:25 +01:00
parent 605a5b4cd6 47d7e9f3f6
commit bd6f9fc309
48 changed files with 232 additions and 569 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/primaite/acl/__init__.py
View File

@@ -1 +1,2 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Access Control List. Models firewall functionality."""

2
src/primaite/acl/access_control_list.py
View File

@@ -9,7 +9,7 @@ class AccessControlList:
"""Access Control List class."""
def __init__(self):
"""Init."""
"""Initialise an empty AccessControlList."""
self.acl: Dict[str, ACLRule] = {} # A dictionary of ACL Rules
def check_address_match(self, _rule: ACLRule, _source_ip_address: str, _dest_ip_address: str) -> bool:

13
src/primaite/acl/acl_rule.py
View File

@@ -7,14 +7,13 @@ class ACLRule:
def __init__(self, _permission, _source_ip, _dest_ip, _protocol, _port):
"""
Init.
Initialise an ACL Rule.
Args:
_permission: The permission (ALLOW or DENY)
_source_ip: The source IP address
_dest_ip: The destination IP address
_protocol: The rule protocol
_port: The rule port
:param _permission: The permission (ALLOW or DENY)
:param _source_ip: The source IP address
:param _dest_ip: The destination IP address
:param _protocol: The rule protocol
:param _port: The rule port
"""
self.permission = _permission
self.source_ip = _source_ip

1
src/primaite/agents/__init__.py
View File

@@ -0,0 +1 @@
"""Common interface between RL agents from different libraries and PrimAITE."""

28
src/primaite/agents/agent.py
View File

@@ -42,12 +42,21 @@ class AgentSessionABC(ABC):
"""
An ABC that manages training and/or evaluation of agents in PrimAITE.
This class cannot be directly instantiated and must be inherited from
with all implemented abstract methods implemented.
This class cannot be directly instantiated and must be inherited from with all implemented abstract methods
implemented.
"""
@abstractmethod
def __init__(self, training_config_path, lay_down_config_path):
"""
Initialise an agent session from config files.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
"""
if not isinstance(training_config_path, Path):
training_config_path = Path(training_config_path)
self._training_config_path: Final[Union[Path, str]] = training_config_path
@@ -55,7 +64,7 @@ class AgentSessionABC(ABC):
if not isinstance(lay_down_config_path, Path):
lay_down_config_path = Path(lay_down_config_path)
self._lay_down_config_path: Final[Union[Path]] = lay_down_config_path
self._lay_down_config_path: Final[Union[Path, str]] = lay_down_config_path
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path)
self.sb3_output_verbose_level = self._training_config.sb3_output_verbose_level
@@ -305,11 +314,20 @@ class HardCodedAgentSessionABC(AgentSessionABC):
"""
An Agent Session ABC for evaluation deterministic agents.
This class cannot be directly instantiated and must be inherited from
with all implemented abstract methods implemented.
This class cannot be directly instantiated and must be inherited from with all implemented abstract methods
implemented.
"""
def __init__(self, training_config_path, lay_down_config_path):
"""
Initialise a hardcoded agent session.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
"""
super().__init__(training_config_path, lay_down_config_path)
self._setup()

3
src/primaite/agents/hardcoded_acl.py
View File

@@ -439,7 +439,8 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return action
def _calculate_action_basic_view(self, obs: np.ndarray) -> int:
"""Calculate a good acl-based action for the blue agent to take.
"""
Calculate a good acl-based action for the blue agent to take.
Uses ONLY information from the current observation with NO knowledge
of previous actions taken and NO reward feedback.

12
src/primaite/agents/rllib.py
View File

@@ -44,6 +44,18 @@ class RLlibAgent(AgentSessionABC):
"""An AgentSession class that implements a Ray RLlib agent."""
def __init__(self, training_config_path, lay_down_config_path):
"""
Initialise the RLLib Agent training session.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:raises ValueError: If the training config contains an unexpected value for agent_framework (should be "RLLIB")
:raises ValueError: If the training config contains an unexpected value for agent_identifies (should be `PPO`
or `A2C`)
"""
super().__init__(training_config_path, lay_down_config_path)
if not self._training_config.agent_framework == AgentFramework.RLLIB:
msg = f"Expected RLLIB agent_framework, " f"got {self._training_config.agent_framework}"

12
src/primaite/agents/sb3.py
View File

@@ -19,6 +19,18 @@ class SB3Agent(AgentSessionABC):
"""An AgentSession class that implements a Stable Baselines3 agent."""
def __init__(self, training_config_path, lay_down_config_path):
"""
Initialise the SB3 Agent training session.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:raises ValueError: If the training config contains an unexpected value for agent_framework (should be "SB3")
:raises ValueError: If the training config contains an unexpected value for agent_identifies (should be `PPO`
or `A2C`)
"""
super().__init__(training_config_path, lay_down_config_path)
if not self._training_config.agent_framework == AgentFramework.SB3:
msg = f"Expected SB3 agent_framework, " f"got {self._training_config.agent_framework}"

3
src/primaite/agents/simple.py
View File

@@ -17,8 +17,7 @@ class DummyAgent(HardCodedAgentSessionABC):
"""
A Dummy Agent.
All action spaces setup so dummy action is always 0 regardless of action
type used.
All action spaces setup so dummy action is always 0 regardless of action type used.
"""
def _calculate_action(self, obs):

3
src/primaite/agents/utils.py
View File

@@ -67,7 +67,8 @@ def transform_action_acl_readable(action: List[str]) -> List[Union[str, int]]:
def is_valid_node_action(action: List[int]) -> bool:
"""Is the node action an actual valid action.
"""
Is the node action an actual valid action.
Only uses information about the action to determine if the action has an effect

1
src/primaite/common/__init__.py
View File

@@ -1 +1,2 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Objects which are shared between many PrimAITE modules."""

6
src/primaite/common/protocol.py
View File

@@ -7,10 +7,10 @@ class Protocol(object):
def __init__(self, _name):
"""
Init.
Initialise a protocol.
Args:
_name: The protocol name
:param _name: The name of the protocol
:type _name: str
"""
self.name = _name
self.load = 0 # bps

2
src/primaite/common/service.py
View File

@@ -9,7 +9,7 @@ class Service(object):
def __init__(self, name: str, port: str, software_state: SoftwareState):
"""
Init.
Initialise a service.
:param name: The service name.
:param port: The service port.

1
src/primaite/config/__init__.py
View File

@@ -0,0 +1 @@
"""Configuration parameters for running experiments."""

3
src/primaite/config/lay_down_config.py
View File

@@ -26,8 +26,7 @@ def load(file_path: Union[str, Path], legacy_file: bool = False) -> Dict:
Read in a lay down config yaml file.
:param file_path: The config file path.
:param legacy_file: True if the config file is legacy format, otherwise
False.
:param legacy_file: True if the config file is legacy format, otherwise False.
:return: The lay down config as a dict.
:raises ValueError: If the file_path does not exist.
"""

13
src/primaite/config/training_config.py
View File

@@ -291,14 +291,11 @@ def convert_legacy_training_config_dict(
Convert a legacy training config dict to the new format.
:param legacy_config_dict: A legacy training config dict.
:param agent_framework: The agent framework to use as legacy training
configs don't have agent_framework values.
:param agent_identifier: The red agent identifier to use as legacy
training configs don't have agent_identifier values.
:param action_type: The action space type to set as legacy training configs
don't have action_type values.
:param num_steps: The number of steps to set as legacy training configs
don't have num_steps values.
:param agent_framework: The agent framework to use as legacy training configs don't have agent_framework values.
:param agent_identifier: The red agent identifier to use as legacy training configs don't have agent_identifier
values.
:param action_type: The action space type to set as legacy training configs don't have action_type values.
:param num_steps: The number of steps to set as legacy training configs don't have num_steps values.
:return: The converted training config dict.
"""
config_dict = {

1
src/primaite/data_viz/__init__.py
View File

@@ -1,3 +1,4 @@
"""Utility to generate plots of sessions metrics after PrimAITE."""
from enum import Enum

1
src/primaite/environment/__init__.py
View File

@@ -1 +1,2 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Gym/Gymnasium environment for RL agents consisting of a simulated computer network."""

97
src/primaite/environment/observations.py
View File

@@ -25,6 +25,12 @@ class AbstractObservationComponent(ABC):
@abstractmethod
def __init__(self, env: "Primaite"):
"""
Initialise observation component.
:param env: Primaite training environment.
:type env: Primaite
"""
_LOGGER.info(f"Initialising {self} observation component")
self.env: "Primaite" = env
self.space: spaces.Space
@@ -44,11 +50,13 @@ class AbstractObservationComponent(ABC):
class NodeLinkTable(AbstractObservationComponent):
"""Table with nodes and links as rows and hardware/software status as cols.
"""
Table with nodes and links as rows and hardware/software status as cols.
This will create the observation space formatted as a table of integers.
There is one row per node, followed by one row per link.
The number of columns is 4 plus one per service. They are:
* node/link ID
* node hardware status / 0 for links
* node operating system status (if active/service) / 0 for links
@@ -67,6 +75,12 @@ class NodeLinkTable(AbstractObservationComponent):
_DATA_TYPE: type = np.int64
def __init__(self, env: "Primaite"):
"""
Initialise a NodeLinkTable observation space component.
:param env: Training environment.
:type env: Primaite
"""
super().__init__(env)
# 1. Define the shape of your observation space component
@@ -88,7 +102,8 @@ class NodeLinkTable(AbstractObservationComponent):
self.structure = self.generate_structure()
def update(self):
"""Update the observation based on current environment state.
"""
Update the observation based on current environment state.
The structure of the observation space is described in :class:`.NodeLinkTable`
"""
@@ -169,12 +184,14 @@ class NodeLinkTable(AbstractObservationComponent):
class NodeStatuses(AbstractObservationComponent):
"""Flat list of nodes' hardware, OS, file system, and service states.
"""
Flat list of nodes' hardware, OS, file system, and service states.
The MultiDiscrete observation space can be though of as a one-dimensional vector of discrete states, represented by
integers.
Each node has 3 elements plus 1 per service. It will have the following structure:
.. code-block::
[
node1 hardware state,
node1 OS state,
@@ -190,14 +207,17 @@ class NodeStatuses(AbstractObservationComponent):
node2 serviceN state (one for each service),
...
]
:param env: The environment that forms the basis of the observations
:type env: Primaite
"""
_DATA_TYPE: type = np.int64
def __init__(self, env: "Primaite"):
"""
Initialise a NodeStatuses observation component.
:param env: Training environment.
:type env: Primaite
"""
super().__init__(env)
# 1. Define the shape of your observation space component
@@ -218,7 +238,8 @@ class NodeStatuses(AbstractObservationComponent):
self.structure = self.generate_structure()
def update(self):
"""Update the observation based on current environment state.
"""
Update the observation based on current environment state.
The structure of the observation space is described in :class:`.NodeStatuses`
"""
@@ -271,28 +292,22 @@ class NodeStatuses(AbstractObservationComponent):
class LinkTrafficLevels(AbstractObservationComponent):
"""Flat list of traffic levels encoded into banded categories.
"""
Flat list of traffic levels encoded into banded categories.
For each link, total traffic or traffic per service is encoded into a categorical value.
For example, if ``quantisation_levels=5``, the traffic levels represent these values:
0 = No traffic (0% of bandwidth)
1 = No traffic (0%-33% of bandwidth)
2 = No traffic (33%-66% of bandwidth)
3 = No traffic (66%-100% of bandwidth)
4 = No traffic (100% of bandwidth)
* 0 = No traffic (0% of bandwidth)
* 1 = No traffic (0%-33% of bandwidth)
* 2 = No traffic (33%-66% of bandwidth)
* 3 = No traffic (66%-100% of bandwidth)
* 4 = No traffic (100% of bandwidth)
.. note::
The lowest category always corresponds to no traffic and the highest category to the link being at max capacity.
Any amount of traffic between 0% and 100% (exclusive) is divided evenly into the remaining categories.
:param env: The environment that forms the basis of the observations
:type env: Primaite
:param combine_service_traffic: Whether to consider total traffic on the link, or each protocol individually,
defaults to False
:type combine_service_traffic: bool, optional
:param quantisation_levels: How many bands to consider when converting the traffic amount to a categorical value ,
defaults to 5
:type quantisation_levels: int, optional
"""
_DATA_TYPE: type = np.int64
@@ -303,6 +318,18 @@ class LinkTrafficLevels(AbstractObservationComponent):
combine_service_traffic: bool = False,
quantisation_levels: int = 5,
):
"""
Initialise a LinkTrafficLevels observation component.
:param env: The environment that forms the basis of the observations
:type env: Primaite
:param combine_service_traffic: Whether to consider total traffic on the link, or each protocol individually,
defaults to False
:type combine_service_traffic: bool, optional
:param quantisation_levels: How many bands to consider when converting the traffic amount to a categorical
value, defaults to 5
:type quantisation_levels: int, optional
"""
if quantisation_levels < 3:
_msg = (
f"quantisation_levels must be 3 or more because the lowest and highest levels are "
@@ -333,7 +360,8 @@ class LinkTrafficLevels(AbstractObservationComponent):
self.structure = self.generate_structure()
def update(self):
"""Update the observation based on current environment state.
"""
Update the observation based on current environment state.
The structure of the observation space is described in :class:`.LinkTrafficLevels`
"""
@@ -374,10 +402,11 @@ class LinkTrafficLevels(AbstractObservationComponent):
class ObservationsHandler:
"""Component-based observation space handler.
"""
Component-based observation space handler.
This allows users to configure observation spaces by mixing and matching components.
Each component can also define further parameters to make them more flexible.
This allows users to configure observation spaces by mixing and matching components. Each component can also define
further parameters to make them more flexible.
"""
_REGISTRY: Final[Dict[str, type]] = {
@@ -387,6 +416,7 @@ class ObservationsHandler:
}
def __init__(self):
"""Initialise the observation handler."""
self.registered_obs_components: List[AbstractObservationComponent] = []
# internal the observation space (unflattened version of space if flatten=True)
@@ -414,7 +444,8 @@ class ObservationsHandler:
self._flat_observation = spaces.flatten(self._space, self._observation)
def register(self, obs_component: AbstractObservationComponent):
"""Add a component for this handler to track.
"""
Add a component for this handler to track.
:param obs_component: The component to add.
:type obs_component: AbstractObservationComponent
@@ -423,10 +454,11 @@ class ObservationsHandler:
self.update_space()
def deregister(self, obs_component: AbstractObservationComponent):
"""Remove a component from this handler.
"""
Remove a component from this handler.
:param obs_component: Which component to remove. It must exist within this object's
``registered_obs_components`` attribute.
``registered_obs_components`` attribute.
:type obs_component: AbstractObservationComponent
"""
self.registered_obs_components.remove(obs_component)
@@ -466,11 +498,13 @@ class ObservationsHandler:
@classmethod
def from_config(cls, env: "Primaite", obs_space_config: dict):
"""Parse a config dictinary, return a new observation handler populated with new observation component objects.
"""
Parse a config dictinary, return a new observation handler populated with new observation component objects.
The expected format for the config dictionary is:
..code-block::python
.. code-block:: python
config = {
components: [
{
@@ -510,7 +544,8 @@ class ObservationsHandler:
return handler
def describe_structure(self):
"""Create a list of names for the features of the obs space.
"""
Create a list of names for the features of the obs space.
The order of labels follows the flattened version of the space.
"""

19
src/primaite/environment/primaite_env.py
View File

@@ -73,8 +73,7 @@ class Primaite(Env):
:param training_config_path: The training config filepath.
:param lay_down_config_path: The lay down config filepath.
:param session_path: The directory path the session is writing to.
:param timestamp_str: The session timestamp in the format:
<yyyy-mm-dd>_<hh-mm-ss>.
:param timestamp_str: The session timestamp in the format: <yyyy-mm-dd>_<hh-mm- ss>.
"""
self.session_path: Final[Path] = session_path
self.timestamp_str: Final[str] = timestamp_str
@@ -660,7 +659,8 @@ class Primaite(Env):
pass
def init_observations(self) -> Tuple[spaces.Space, np.ndarray]:
"""Create the environment's observation handler.
"""
Create the environment's observation handler.
:return: The observation space, initial observation (zeroed out array with the correct shape)
:rtype: Tuple[spaces.Space, np.ndarray]
@@ -1040,7 +1040,8 @@ class Primaite(Env):
self.num_ports = len(self.ports_list)
def get_observation_info(self, observation_info):
"""Extracts observation_info.
"""
Extracts observation_info.
:param observation_info: Config item that defines which type of observation space to use
:type observation_info: str
@@ -1057,7 +1058,8 @@ class Primaite(Env):
self.action_type = ActionType[action_info["type"]]
def save_obs_config(self, obs_config: dict):
"""Cache the config for the observation space.
"""
Cache the config for the observation space.
This is necessary as the observation space can't be built while reading the config,
it must be done after all the nodes, links, and services have been initialised.
@@ -1070,10 +1072,9 @@ class Primaite(Env):
def reset_environment(self):
"""
# Resets environment.
Resets environment.
Uses config data config data in order to build the environment
configuration.
Uses config data config data in order to build the environment configuration.
"""
for item in self.lay_down_config:
if item["item_type"] == "NODE":
@@ -1157,7 +1158,6 @@ class Primaite(Env):
5: [1, 3, 1, 0],
...
}
"""
# reserve 0 action to be a nothing action
actions = {0: [1, 0, 0, 0]}
@@ -1213,7 +1213,6 @@ class Primaite(Env):
Create a dictionary mapping each possible discrete action to a more readable mutlidiscrete action.
The dictionary contains actions of both Node and ACL action types.
"""
node_action_dict = self.create_node_action_dict()
acl_action_dict = self.create_acl_action_dict()

1
src/primaite/links/__init__.py
View File

@@ -1 +1,2 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Network connections between nodes in the simulation."""

13
src/primaite/links/link.py
View File

@@ -10,14 +10,13 @@ class Link(object):
def __init__(self, _id, _bandwidth, _source_node_name, _dest_node_name, _services):
"""
Init.
Initialise a Link within the simulated network.
Args:
_id: The IER id
_bandwidth: The bandwidth of the link (bps)
_source_node_name: The name of the source node
_dest_node_name: The name of the destination node
_protocols: The protocols to add to the link
:param _id: The IER id
:param _bandwidth: The bandwidth of the link (bps)
:param _source_node_name: The name of the source node
:param _dest_node_name: The name of the destination node
:param _protocols: The protocols to add to the link
"""
self.id = _id
self.bandwidth = _bandwidth

3
src/primaite/main.py
View File

@@ -14,7 +14,8 @@ def run(
training_config_path: Union[str, Path],
lay_down_config_path: Union[str, Path],
):
"""Run the PrimAITE Session.
"""
Run the PrimAITE Session.
:param training_config_path: The training config filepath.
:param lay_down_config_path: The lay down config filepath.

1
src/primaite/nodes/__init__.py
View File

@@ -1 +1,2 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Nodes represent network hosts in the simulation."""

2
src/primaite/nodes/active_node.py
View File

@@ -26,7 +26,7 @@ class ActiveNode(Node):
config_values: TrainingConfig,
):
"""
Init.
Initialise an active node.
:param node_id: The node ID
:param name: The node name

2
src/primaite/nodes/node.py
View File

@@ -19,7 +19,7 @@ class Node:
config_values: TrainingConfig,
):
"""
Init.
Initialise a node.
:param node_id: The node id.
:param name: The name of the node.

17
src/primaite/nodes/node_state_instruction_green.py
View File

@@ -16,16 +16,15 @@ class NodeStateInstructionGreen(object):
_state,
):
"""
Init.
Initialise the Node State Instruction.
Args:
_id: The node state instruction id
_start_step: The start step of the instruction
_end_step: The end step of the instruction
_node_id: The id of the associated node
_node_pol_type: The pattern of life type
_service_name: The service name
_state: The state (node or service)
:param _id: The node state instruction id
:param _start_step: The start step of the instruction
:param _end_step: The end step of the instruction
:param _node_id: The id of the associated node
:param _node_pol_type: The pattern of life type
:param _service_name: The service name
:param _state: The state (node or service)
"""
self.id = _id
self.start_step = _start_step

25
src/primaite/nodes/node_state_instruction_red.py
View File

@@ -24,20 +24,19 @@ class NodeStateInstructionRed(object):
_pol_source_node_service_state,
):
"""
Init.
Initialise the Node State Instruction for the red agent.
Args:
_id: The node state instruction id
_start_step: The start step of the instruction
_end_step: The end step of the instruction
_target_node_id: The id of the associated node
-pol_initiator: The way the PoL is applied (DIRECT, IER or SERVICE)
_pol_type: The pattern of life type
-pol_protocol: The pattern of life protocol/service affected
_pol_state: The state (node or service)
_pol_source_node_id: The source node Id (used for initiator type SERVICE)
_pol_source_node_service: The source node service (used for initiator type SERVICE)
_pol_source_node_service_state: The source node service state (used for initiator type SERVICE)
:param _id: The node state instruction id
:param _start_step: The start step of the instruction
:param _end_step: The end step of the instruction
:param _target_node_id: The id of the associated node
:param -pol_initiator: The way the PoL is applied (DIRECT, IER or SERVICE)
:param _pol_type: The pattern of life type
:param pol_protocol: The pattern of life protocol/service affected
:param _pol_state: The state (node or service)
:param _pol_source_node_id: The source node Id (used for initiator type SERVICE)
:param _pol_source_node_service: The source node service (used for initiator type SERVICE)
:param _pol_source_node_service_state: The source node service state (used for initiator type SERVICE)
"""
self.id = _id
self.start_step = _start_step

2
src/primaite/nodes/passive_node.py
View File

@@ -18,7 +18,7 @@ class PassiveNode(Node):
config_values: TrainingConfig,
):
"""
Init.
Initialise a passive node.
:param node_id: The node id.
:param name: The name of the node.

8
src/primaite/nodes/service_node.py
View File

@@ -27,7 +27,7 @@ class ServiceNode(ActiveNode):
config_values: TrainingConfig,
):
"""
Init.
Initialise a Service Node.
:param node_id: The node ID
:param name: The node name
@@ -77,8 +77,7 @@ class ServiceNode(ActiveNode):
Indicates whether a service is in a running state on the node.
:param protocol_name: The service (protocol)
:return: True if service (protocol) is in a running state on the
node, otherwise False.
:return: True if service (protocol) is in a running state on the node, otherwise False.
"""
for service_key, service_value in self.services.items():
if service_key == protocol_name:
@@ -93,8 +92,7 @@ class ServiceNode(ActiveNode):
Indicates whether a service is in an overwhelmed state on the node.
:param protocol_name: The service (protocol)
:return: True if service (protocol) is in an overwhelmed state on the
node, otherwise False.
:return: True if service (protocol) is in an overwhelmed state on the node, otherwise False.
"""
for service_key, service_value in self.services.items():
if service_key == protocol_name:

1
src/primaite/notebooks/__init__.py
View File

@@ -1,3 +1,4 @@
"""Contains default jupyter notebooks which demonstrate PrimAITE functionality."""
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
import importlib.util
import os

1
src/primaite/pol/__init__.py
View File

@@ -1 +1,2 @@
"""Pattern of Life- Represents the actions of users on the network."""
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

23
src/primaite/pol/ier.py
View File

@@ -23,19 +23,18 @@ class IER(object):
_running=False,
):
"""
Init.
Initialise an Information Exchange Request.
Args:
_id: The IER id
_start_step: The step when this IER should start
_end_step: The step when this IER should end
_load: The load this IER should put on a link (bps)
_protocol: The protocol of this IER
_port: The port this IER runs on
_source_node_id: The source node ID
_dest_node_id: The destination node ID
_mission_criticality: Criticality of this IER to the mission (0 none, 5 mission critical)
_running: Indicates whether the IER is currently running
:param _id: The IER id
:param _start_step: The step when this IER should start
:param _end_step: The step when this IER should end
:param _load: The load this IER should put on a link (bps)
:param _protocol: The protocol of this IER
:param _port: The port this IER runs on
:param _source_node_id: The source node ID
:param _dest_node_id: The destination node ID
:param _mission_criticality: Criticality of this IER to the mission (0 none, 5 mission critical)
:param _running: Indicates whether the IER is currently running
"""
self.id = _id
self.start_step = _start_step

8
src/primaite/primaite_session.py
View File

@@ -1,3 +1,4 @@
"""Main entry point to PrimAITE. Configure training/evaluation experiments and input/output."""
from __future__ import annotations
from pathlib import Path
@@ -21,8 +22,7 @@ class PrimaiteSession:
"""
The PrimaiteSession class.
Provides a single learning and evaluation entry point for all training
and lay down configurations.
Provides a single learning and evaluation entry point for all training and lay down configurations.
"""
def __init__(
@@ -38,12 +38,12 @@ class PrimaiteSession:
"""
if not isinstance(training_config_path, Path):
training_config_path = Path(training_config_path)
self._training_config_path: Final[Union[Path]] = training_config_path
self._training_config_path: Final[Union[Path, str]] = training_config_path
self._training_config: Final[TrainingConfig] = training_config.load(self._training_config_path)
if not isinstance(lay_down_config_path, Path):
lay_down_config_path = Path(lay_down_config_path)
self._lay_down_config_path: Final[Union[Path]] = lay_down_config_path
self._lay_down_config_path: Final[Union[Path, str]] = lay_down_config_path
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path)
self._agent_session: AgentSessionABC = None # noqa

1
src/primaite/setup/__init__.py
View File

@@ -1 +1,2 @@
"""Utilities to prepare the user's data folders."""
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

3
src/primaite/setup/reset_demo_notebooks.py
View File

@@ -15,8 +15,7 @@ def run(overwrite_existing: bool = True):
"""
Resets the demo jupyter notebooks in the users app notebooks directory.
:param overwrite_existing: A bool to toggle replacing existing edited
notebooks on or off.
:param overwrite_existing: A bool to toggle replacing existing edited notebooks on or off.
"""
notebooks_package_data_root = pkg_resources.resource_filename("primaite", "notebooks/_package_data")
for subdir, dirs, files in os.walk(notebooks_package_data_root):

3
src/primaite/setup/reset_example_configs.py
View File

@@ -14,8 +14,7 @@ def run(overwrite_existing=True):
"""
Resets the example config files in the users app config directory.
:param overwrite_existing: A bool to toggle replacing existing edited
config on or off.
:param overwrite_existing: A bool to toggle replacing existing edited config on or off.
"""
configs_package_data_root = pkg_resources.resource_filename("primaite", "config/_package_data")

1
src/primaite/transactions/__init__.py
View File

@@ -1 +1,2 @@
"""Record data of the system's state and agent's observations and actions."""
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

3
src/primaite/transactions/transaction.py
View File

@@ -86,8 +86,7 @@ def _turn_obs_space_to_array(obs_space, obs_assets, obs_features) -> List[str]:
Turns observation space into a string array so it can be saved to csv.
:param obs_space: The observation space
:param obs_assets: The number of assets (i.e. nodes or links) in the
observation space
:param obs_assets: The number of assets (i.e. nodes or links) in the observation space
:param obs_features: The number of features associated with the asset
:return: The observation space as an array of strings
"""

1
src/primaite/utils/__init__.py
View File

@@ -0,0 +1 @@
"""Utilities for PrimAITE."""

3
src/primaite/utils/session_output_reader.py
View File

@@ -10,8 +10,7 @@ def av_rewards_dict(av_rewards_csv_file: Union[str, Path]) -> Dict[int, float]:
"""
Read an average rewards per episode csv file and return as a dict.
The dictionary keys are the episode number, and the values are the mean
reward that episode.
The dictionary keys are the episode number, and the values are the mean reward that episode.
:param av_rewards_csv_file: The average rewards per episode csv file path.
:return: The average rewards per episode cdv as a dict.

15
src/primaite/utils/session_output_writer.py
View File

@@ -29,6 +29,18 @@ class SessionOutputWriter:
transaction_writer: bool = False,
learning_session: bool = True,
):
"""
Initialise the Session Output Writer.
:param env: PrimAITE gym environment.
:type env: Primaite
:param transaction_writer: If `true`, this will output a full account of every transaction taken by the agent.
If `false` it will output the average reward per episode, defaults to False
:type transaction_writer: bool, optional
:param learning_session: Set to `true` to indicate that the current session is a training session. This
determines the name of the folder which contains the final output csv. Defaults to True
:type learning_session: bool, optional
"""
self._env = env
self.transaction_writer = transaction_writer
self.learning_session = learning_session
@@ -68,8 +80,7 @@ class SessionOutputWriter:
"""
Write a row of session data.
:param data: The row of data to write. Can be a Tuple or an instance
of Transaction.
:param data: The row of data to write. Can be a Tuple or an instance of Transaction.
"""
if isinstance(data, Transaction):
header, data = data.as_csv_data()