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

Merge remote-tracking branch 'origin/dev' into feature/2502-placeholders-docs

Browse Source
This commit is contained in:
Marek Wolan
2024-06-13 12:52:59 +01:00
parent 4d1f84dc27 ae2bbd3e00
commit 0bd52b1c15
84 changed files with 10352 additions and 1020 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb
View File

@@ -4,13 +4,15 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"# Customising red agents\n",
"# Customising UC2 Red Agents\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook will go over some examples of how red agent behaviour can be varied by changing its configuration parameters.\n",
"\n",
"First, let's load the standard Data Manipulation config file, and see what the red agent does.\n",
"\n",
"*(For a full explanation of the Data Manipulation scenario, check out the notebook `Data-Manipulation-E2E-Demonstration.ipynb`)*"
"*(For a full explanation of the Data Manipulation scenario, check out the data manipulation scenario notebook)*"
]
},
{
@@ -22,7 +24,7 @@
"# Imports\n",
"\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.game.agent.interface import AgentActionHistoryItem\n",
"from primaite.game.agent.interface import AgentHistoryItem\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"import yaml\n",
"from pprint import pprint"
@@ -63,7 +65,7 @@
"source": [
"def friendly_output_red_action(info):\n",
" # parse the info dict form step output and write out what the red agent is doing\n",
" red_info : AgentActionHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
" red_info : AgentHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info.action\n",
" if red_action == 'DONOTHING':\n",
" red_str = 'DO NOTHING'\n",

37
src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb
View File

@@ -4,7 +4,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"# Data Manipulation Scenario\n"
"# Data Manipulation Scenario\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK"
]
},
{
@@ -59,7 +61,7 @@
"\n",
"At the start of every episode, the red agent randomly chooses either client 1 or client 2 to login to. It waits a bit then sends a DELETE query to the database from its chosen client. If the delete is successful, the database file is flagged as compromised to signal that data is not available.\n",
"\n",
"[<img src=\"_package_data/uc2_attack.png\" width=\"500\"/>](_package_data/uc2_attack.png)\n",
"![uc2_attack](./_package_data/uc2_attack.png)\n",
"\n",
"_(click image to enlarge)_"
]
@@ -79,7 +81,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"# Reinforcement learning details"
"## Reinforcement learning details"
]
},
{
@@ -180,15 +182,15 @@
"| link_id | endpoint_a | endpoint_b |\n",
"|---------|------------------|-------------------|\n",
"| 1 | router_1 | switch_1 |\n",
"| 1 | router_1 | switch_2 |\n",
"| 1 | switch_1 | domain_controller |\n",
"| 1 | switch_1 | web_server |\n",
"| 1 | switch_1 | database_server |\n",
"| 1 | switch_1 | backup_server |\n",
"| 1 | switch_1 | security_suite |\n",
"| 1 | switch_2 | client_1 |\n",
"| 1 | switch_2 | client_2 |\n",
"| 1 | switch_2 | security_suite |\n",
"| 2 | router_1 | switch_2 |\n",
"| 3 | switch_1 | domain_controller |\n",
"| 4 | switch_1 | web_server |\n",
"| 5 | switch_1 | database_server |\n",
"| 6 | switch_1 | backup_server |\n",
"| 7 | switch_1 | security_suite |\n",
"| 8 | switch_2 | client_1 |\n",
"| 9 | switch_2 | client_2 |\n",
"| 10 | switch_2 | security_suite |\n",
"\n",
"\n",
"The ACL rules in the observation space appear in the same order that they do in the actual ACL. Though, only the first 10 rules are shown, there are default rules lower down that cannot be changed by the agent. The extra rules just allow the network to function normally, by allowing pings, ARP traffic, etc.\n",
@@ -392,7 +394,7 @@
"# Imports\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.game.agent.interface import AgentActionHistoryItem\n",
"from primaite.game.agent.interface import AgentHistoryItem\n",
"import yaml\n",
"from pprint import pprint\n"
]
@@ -401,7 +403,8 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Instantiate the environment. We also disable the agent observation flattening.\n",
"Instantiate the environment. \n",
"We will also disable the agent observation flattening.\n",
"\n",
"This cell will print the observation when the network is healthy. You should be able to verify Node file and service statuses against the description above."
]
@@ -444,7 +447,7 @@
"source": [
"def friendly_output_red_action(info):\n",
" # parse the info dict form step output and write out what the red agent is doing\n",
" red_info : AgentActionHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
" red_info : AgentHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info.action\n",
" if red_action == 'DONOTHING':\n",
" red_str = 'DO NOTHING'\n",
@@ -691,7 +694,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -705,7 +708,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.8"
}
},
"nbformat": 4,

170
src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb Normal file
View File

@@ -0,0 +1,170 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Getting information out of PrimAITE\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Imports\n",
"import yaml\n",
"from primaite import PRIMAITE_CONFIG\n",
"\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"from notebook.services.config import ConfigManager\n",
"\n",
"cm = ConfigManager().update('notebook', {'limit_output': 50}) # limit output lines to 50 - for neatness\n",
"\n",
"# create the env\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
"\n",
"env = PrimaiteGymEnv(env_config=cfg)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Visualising the Simulation Network"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The network can be visualised by running the code below."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.simulation.network.draw()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Getting the state of a simulation object\n",
"\n",
"The state of the simulation object is used to determine the observation space used by agents.\n",
"\n",
"Any object created using the ``SimComponent`` class has a ``describe_state`` method which can show the state of the object.\n",
"\n",
"An example of such an object is ``Computer`` which inherits from ``SimComponent``. In the default network configuration, ``client_1`` is a Computer object."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"client_1: Computer = env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
"client_1.describe_state()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### More specific describe_state\n",
"\n",
"As you can see, the output from the ``describe_state`` method for the ``Computer`` object includes the describe state for all its components. This can cause a large describe state output.\n",
"\n",
"As stated, the ``describe_state`` can be called on any object that inherits ``SimComponent``. This can allow you retrieve the state of a specific item."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"client_1.file_system.describe_state()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## System Logs\n",
"\n",
"Objects that inherit from the ``Node`` class will inherit the ``sys_log`` attribute.\n",
"\n",
"This is to simulate the idea that items such as Computer, Routers, Servers, etc. have a logging system used to diagnose problems."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# store config\n",
"# this is to prevent the notebook from breaking your local settings\n",
"was_enabled = PRIMAITE_CONFIG[\"developer_mode\"][\"enabled\"]\n",
"was_syslogs_enabled = PRIMAITE_CONFIG[\"developer_mode\"][\"output_sys_logs\"]\n",
"\n",
"# enable dev mode so that the default config outputs are overridden for this demo\n",
"PRIMAITE_CONFIG[\"developer_mode\"][\"enabled\"] = True\n",
"PRIMAITE_CONFIG[\"developer_mode\"][\"output_sys_logs\"] = True\n",
"\n",
"\n",
"\n",
"\n",
"# Remake the environment\n",
"env = PrimaiteGymEnv(env_config=cfg)\n",
"\n",
"# get the example computer\n",
"client_1: Computer = env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
"\n",
"# show sys logs on terminal\n",
"client_1.sys_log.show()\n",
"\n",
"\n",
"\n",
"\n",
"# restore config\n",
"PRIMAITE_CONFIG[\"developer_mode\"][\"enabled\"] = was_enabled\n",
"PRIMAITE_CONFIG[\"developer_mode\"][\"output_sys_logs\"] = was_syslogs_enabled"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.8"
}
},
"nbformat": 4,
"nbformat_minor": 2
}

4
src/primaite/notebooks/Requests-and-Responses.ipynb
View File

@@ -6,7 +6,9 @@
"source": [
"# Requests and Responses\n",
"\n",
"Agents interact with the PrimAITE simulation via the Request system.\n"
"Agents interact with the PrimAITE simulation via the Request system.\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n"
]
},
{

12
src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb
View File

@@ -4,7 +4,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## Train a Multi agent system using RLLIB\n",
"# Train a Multi agent system using RLLIB\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook will demonstrate how to use the `PrimaiteRayMARLEnv` to train a very basic system with two PPO agents."
]
@@ -25,13 +27,13 @@
"from primaite.game.game import PrimaiteGame\n",
"import yaml\n",
"\n",
"from primaite.session.environment import PrimaiteRayEnv\n",
"from primaite.session.ray_envs import PrimaiteRayEnv\n",
"from primaite import PRIMAITE_PATHS\n",
"\n",
"import ray\n",
"from ray import air, tune\n",
"from ray.rllib.algorithms.ppo import PPOConfig\n",
"from primaite.session.environment import PrimaiteRayMARLEnv\n",
"from primaite.session.ray_envs import PrimaiteRayMARLEnv\n",
"\n",
"# If you get an error saying this config file doesn't exist, you may need to run `primaite setup` in your command line\n",
"# to copy the files to your user data path.\n",
@@ -60,8 +62,8 @@
" policies={'defender_1','defender_2'}, # These names are the same as the agents defined in the example config.\n",
" policy_mapping_fn=lambda agent_id, episode, worker, **kw: agent_id,\n",
" )\n",
" .environment(env=PrimaiteRayMARLEnv, env_config=cfg)#, disable_env_checking=True)\n",
" .rollouts(num_rollout_workers=0)\n",
" .environment(env=PrimaiteRayMARLEnv, env_config=cfg)\n",
" .env_runners(num_env_runners=0)\n",
" .training(train_batch_size=128)\n",
" )\n"
]

16
src/primaite/notebooks/Training-an-RLLib-Agent.ipynb
View File

@@ -4,7 +4,10 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## Train a Single agent system using RLLib\n",
"# Train a Single agent system using RLLib\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook will demonstrate how to use PrimaiteRayEnv to train a basic PPO agent."
]
},
@@ -18,8 +21,7 @@
"import yaml\n",
"from primaite.config.load import data_manipulation_config_path\n",
"\n",
"from primaite.session.environment import PrimaiteRayEnv\n",
"from ray.rllib.algorithms import ppo\n",
"from primaite.session.ray_envs import PrimaiteRayEnv\n",
"from ray import air, tune\n",
"import ray\n",
"from ray.rllib.algorithms.ppo import PPOConfig\n",
@@ -52,8 +54,8 @@
"\n",
"config = (\n",
" PPOConfig()\n",
" .environment(env=PrimaiteRayEnv, env_config=env_config, disable_env_checking=True)\n",
" .rollouts(num_rollout_workers=0)\n",
" .environment(env=PrimaiteRayEnv, env_config=env_config)\n",
" .env_runners(num_env_runners=0)\n",
" .training(train_batch_size=128)\n",
")\n"
]
@@ -74,7 +76,7 @@
"tune.Tuner(\n",
" \"PPO\",\n",
" run_config=air.RunConfig(\n",
" stop={\"timesteps_total\": 5 * 128}\n",
" stop={\"timesteps_total\": 512}\n",
" ),\n",
" param_space=config\n",
").fit()\n"
@@ -97,7 +99,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,

11
src/primaite/notebooks/Training-an-SB3-Agent.ipynb
View File

@@ -6,6 +6,8 @@
"source": [
"# Training an SB3 Agent\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook will demonstrate how to use primaite to create and train a PPO agent, using a pre-defined configuration file."
]
},
@@ -43,7 +45,10 @@
"outputs": [],
"source": [
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)"
" cfg = yaml.safe_load(f)\n",
"for agent in cfg['agents']:\n",
" if agent['ref'] == 'defender':\n",
" agent['agent_settings']['flatten_obs']=True"
]
},
{
@@ -163,7 +168,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -177,7 +182,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.10"
"version": "3.10.8"
}
},
"nbformat": 4,

54
src/primaite/notebooks/Using-Episode-Schedules.ipynb
View File

@@ -6,6 +6,8 @@
"source": [
"# Using Episode Schedules\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"\n",
"PrimAITE supports the ability to use different variations on a scenario at different episodes. This can be used to increase \n",
"domain randomisation to prevent overfitting, or to set up curriculum learning to train agents to perform more complicated tasks.\n",
"\n",
@@ -13,50 +15,6 @@
"directory with several config files that work together."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Defining variations in the config file.\n",
"\n",
"### Base scenario\n",
"The base scenario is essentially the same as a fixed YAML configuration, but it can contain placeholders that are \n",
"populated with episode-specific data at runtime. The base scenario contains any network, agent, or settings that\n",
"remain fixed for the entire training/evaluation session.\n",
"\n",
"The placeholders are defined as YAML Aliases and they are denoted by an asterisk (`*placeholder`).\n",
"\n",
"### Variations\n",
"For each variation that could be used in a placeholder, there is a separate yaml file that contains the data that should populate the placeholder.\n",
"\n",
"The data that fills the placeholder is defined as a YAML Anchor in a separate file, denoted by an ampersand (`&anchor`).\n",
"\n",
"[Learn more about YAML Aliases and Anchors here.](https://www.educative.io/blog/advanced-yaml-syntax-cheatsheet#:~:text=YAML%20Anchors%20and%20Alias)\n",
"\n",
"### Schedule\n",
"Users must define which combination of scenario variations should be loaded in each episode. This takes the form of a\n",
"YAML file with a relative path to the base scenario and a list of paths to be loaded in during each episode.\n",
"\n",
"It takes the following format:\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"```yaml\n",
"base_scenario: base.yaml\n",
"schedule:\n",
" 0: # list of variations to load in at episode 0 (before the first call to env.reset() happens)\n",
" - laydown_1.yaml\n",
" - attack_1.yaml\n",
" 1: # list of variations to load in at episode 1 (after the first env.reset() call)\n",
" - laydown_2.yaml\n",
" - attack_2.yaml\n",
"```\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
@@ -298,8 +256,8 @@
"table = PrettyTable()\n",
"table.field_names = [\"step\", \"Green Action\", \"Red Action\"]\n",
"for i in range(21):\n",
" green_action = env.game.agents['green_A'].action_history[i].action\n",
" red_action = env.game.agents['red_A'].action_history[i].action\n",
" green_action = env.game.agents['green_A'].history[i].action\n",
" red_action = env.game.agents['red_A'].history[i].action\n",
" table.add_row([i, green_action, red_action])\n",
"print(table)"
]
@@ -329,8 +287,8 @@
"table = PrettyTable()\n",
"table.field_names = [\"step\", \"Green Action\", \"Red Action\"]\n",
"for i in range(21):\n",
" green_action = env.game.agents['green_B'].action_history[i].action\n",
" red_action = env.game.agents['red_B'].action_history[i].action\n",
" green_action = env.game.agents['green_B'].history[i].action\n",
" red_action = env.game.agents['red_B'].history[i].action\n",
" table.add_row([i, green_action, red_action])\n",
"print(table)"
]

9
src/primaite/notebooks/multi-processing.ipynb
View File

@@ -4,8 +4,11 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## Simple multi-processing demo using SubprocVecEnv from SB3\n",
"Based on a code example provided by Rachael Proctor."
"# Simple multi-processing demonstration\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook uses SubprocVecEnv from SB3."
]
},
{
@@ -140,7 +143,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,