Skip to content

Releases: Autonomous-Resilient-Cyber-Defence/PrimAITE

v2.0.0

15 Aug 15:16
4fb88c9
Compare
Choose a tag to compare

PrimAITE v2.0.0

✨ What's New

Command Line Interface

PrimAITE now comes with a CLI (built with Typer) that serves as the main entry point for those using PrimAITE out of the box.
PrimAITE v2 0 0 CLI

To run the default PrimAITE Session out of the box, run:

primaite session

Application Directories

To enable PrimAITE to be used as an installed Python package, and to be used as is out-of-the-box without reliance on the repository, a collection of application directories has been created. These back-end/hidden and user-facing directories are used to store things like application log files, users' config files, users' Jupyter Notebooks, PrimAITE session outputs etc. The user needs to call primaite setup after doing pip install to perform the directory setup. The directories are structured as follows:

Windows

~/
├─ AppData/
│  ├─ primaite/
│  │  ├─ 2.0.0/
│  │  │  ├─ config/
│  │  │  ├─ logs/
│  │  │  │  ├─ primaite.log
├─ primaite/
│  ├─ 2.0.0/
│  │  ├─ config/
│  │  │  ├─ example_config/
│  │  │  │  ├─ lay_down/
│  │  │  │  ├─ training/
│  │  ├─ notebooks/
│  │  │  │  ├─ primaite_demo_notebooks/
│  │  ├─ sessions/

Linux

~/
├─ .cache/
│  ├─ primaite/
│  │  ├─ 2.0.0/
│  │  │  ├─ logs/
│  │  │  │  ├─ primaite.log
├─ .config/
│  ├─ primaite/
│  │  ├─ 2.0.0/
├─ .local/
│  ├─ share/
│  │  ├─ primaite/
│  │  │  ├─ 2.0.0/
├─ primaite/
│  ├─ 2.0.0/
│  │  ├─ config/
│  │  │  ├─ example_config/
│  │  │  │  ├─ lay_down/
│  │  │  │  ├─ training/
│  │  ├─ notebooks/
│  │  │  │  ├─ primaite_demo_notebooks/
│  │  ├─ sessions/

MacOS

~/
├─ Library/
│  ├─ Application Support/
│  │  ├─ Logs/
│  │  │  ├─ primaite/
│  │  │  │  ├─ 2.0.0/
│  │  │  │  │  ├─ log/
│  │  │  │  │  │  ├─ primaite.log
│  │  ├─ Preferences/
│  │  │  ├─ primaite/
│  │  │  │  ├─ 2.0.0/
│  │  ├─ primaite/
│  │  │  ├─ 2.0.0/
├─ primaite/
│  ├─ 2.0.0/
│  │  ├─ config/
│  │  │  ├─ example_config/
│  │  │  │  ├─ lay_down/
│  │  │  │  ├─ training/
│  │  ├─ notebooks/
│  │  │  │  ├─ primaite_demo_notebooks/
│  │  ├─ sessions/

Support for Ray Rllib

PrimAITE now supports the training of PPO and A2C agents using both Stable Baselines3 and Ray RLlib. The RL framework and agent algorithm to be used for training is determined by the agent_framework and agent_identifier configurable items in the training config file. If agent_framework=RLLIB, the backend Ray RLlib RL framework can be configured to use either Tensorflow, Tensorflow 2.x, or PyTorch using the deep_learning_framework.

# Sets which agent algorithm framework will be used.
# Options are:
# "SB3" (Stable Baselines3)
# "RLLIB" (Ray RLlib)
# "CUSTOM" (Custom Agent)
agent_framework: RLLIB

# Sets which deep learning framework will be used (by RLlib ONLY).
# Default is TF (Tensorflow).
# Options are:
# "TF" (Tensorflow)
# TF2 (Tensorflow 2.X)
# TORCH (PyTorch)
deep_learning_framework: TF2

# Sets which Agent class will be used.
# Options are:
# "A2C" (Advantage Actor-Critic coupled with either SB3 or RLLIB agent_framework)
# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework)
# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type)
# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
# "RANDOM" (primaite.agents.simple.RandomAgent)
# "DUMMY" (primaite.agents.simple.DummyAgent)
agent_identifier: PPO

Random red agent

A random red agent has been provided to train the blue agent against. The random red agent will choose a random number of nodes to attack, as well as randomly choosing the actions to perform on the environment.

# Sets whether Red Agent POL and IER are randomised.
# Default is False.
# Options are:
# True
# False
random_red_agent: False

Repeatability of sessions

A seed can now be provided in the training configuration file. The seed will be used across PrimAITE so that a repeatable run is achievable. The seed needs to be an integer value and by default is set to null.

The ability to set PrimAITE to use deterministic or stochastic evaluation is also added.

# The (integer) seed to be used in random number generation
# Default is None (null)
seed: null

# Set whether the agent evaluation will be deterministic instead of stochastic
# Default is False (stochastic).
# Options are:
# True
# False
deterministic: False

Session loading

PrimAITE can now load previously run sessions for SB3 Agents (SB3 agents only, see Known Issues. This can be done via:

CLI

primaite session --load "<PREVIOUS_SESSION_DIRECTORY>"

Python

from primaite.main import run

run(session_path=<PREVIOUS_SESSION_DIRECTORY>)

The output for the loaded session will be in the target directory i.e. the PREVIOUS_SESSION_DIRECTORY. While most of the outputs won't be overwritten, the agent zip file will be overwritten.

Agent Session Classes

An AgentSessionABC class with SB3Agent and RLlib subclasses and a HardCodedAgentSessionABC class with various hard-coded agent subclasses have been created. The Agent Session classes act as a wrapper around various RL and hard-coded agents. They help to standardise how agents are trained in PrimAITE using a common interface. They also provide a suite of standardised session outputs.

Standardised Session Output

When a session is run, a session output sub-directory is created in the user's app sessions directory
(~/primaite/sessions). The sub-directory is formatted as such: ~/primaite/sessions/<yyyy-mm-dd>/<yyyy-mm-dd>_<hh-mm-dd>/. This session directory is populated with four types of outputs:

  • Session Metadata
  • Results
  • Diagrams
  • Saved agents (training checkpoints and a final trained agent)

Example Session Directory Structure

~/
└── primaite/
    └── 2.0.0/
        └── sessions/
            └── 2023-07-18/
                └── 2023-07-18_11-06-04/
                    ├── evaluation/
                    │   ├── all_transactions_2023-07-18_11-06-04.csv
                    │   ├── average_reward_per_episode_2023-07-18_11-06-04.csv
                    │   └── average_reward_per_episode_2023-07-18_11-06-04.png
                    ├── learning/
                    │   ├── all_transactions_2023-07-18_11-06-04.csv
                    │   ├── average_reward_per_episode_2023-07-18_11-06-04.csv
                    │   ├── average_reward_per_episode_2023-07-18_11-06-04.png
                    │   ├── checkpoints/
                    │   │   └── sb3ppo_5.zip
                    │   ├── SB3_PPO.zip
                    │   └── tensorboard_logs/
                    │       ├── PPO_1/
                    │       │   └── events.out.tfevents.1689674765.METD-9PMRFB3.42960.0
                    │       ├── PPO_2/
                    │       │   └── events.out.tfevents.1689674766.METD-9PMRFB3.42960.1
                    │       ├── PPO_3/
                    │       │   └── events.out.tfevents.1689674766.METD-9PMRFB3.42960.2
                    │       ├── PPO_4/
                    │       │   └── events.out.tfevents.1689674767.METD-9PMRFB3.42960.3
                    │       └── PPO_5/
                    │           └── events.out.tfevents.1689674767.METD-9PMRFB3.42960.4
                    ├── network_2023-07-18_11-06-04.png
                    └── session_metadata.json

PrimaiteSession Class

The PrimaiteSession class acts as a single wrapper around the Agent Session classes. It is both an entry point and a broker for the individual Agent Session classes.

Action Space

Discrete Action Space

The NODE and ACL action spaces have been changed from multi-discrete to discrete action spaces.

NODE and ACL action spaces are both dictionaries where a single number reflects an entire action an agent can take.

The below code block is an example of a dictionary entry for the NODE action space:

{
    1: [1, 1, 1, 0], 
    
}

The below code block is an example of a dictionary entry for the ACL action space:

{
    1: [1, 0, 1, 2, 1, 0, 3],
}

Combined Action Spaces

A new ANY action space option has been introduced. This allows the agent to do both NODE actions and ACL actions in the same episode (e.g., scan a node in Step 1 and create an ACL rule in Step 2).

The below code block is an example of a dictionary entry for the ANY action space:

{
    0: [1, 0, 0, 0], 
    1: [1, 1, 1, 0], 
    2: [1, 0, 1, 2, 1, 0, 3]
}

is_valid_acl_action_extra, is_valid_node_action and is_valid_acl_action in the primaite.agents.utils module help to slim down the dictionary to contain the relevant actions only.

For example, a node action to do PATCHING on a node's hardware state CANNOT happen so it is not added to the dictionary.

transform_action_node_readable and transform_action_acl_readable in the primaite.agents.utils module converts the enumerated node action into a more readable form. The readable form is used by functions such as is_valid_node_action to determine if the action is valid or not.

An example using transform_action_node_readable:

{ 
    # Converts node action into readable form
    [1, 3, 1, 0] -> [1, 'SERVICE', 'PATCHING', 0]
}

ACL Action Space

Previously, the ACL action space was made up of 6 items: [Decision, Permission, Source, Dest, Protocol, Port].

Now the agent can specifically choose ...

Read more