itm-kitware/alignment-cage-challenge-2CybORG CAGE 2 environment updated for the paper "Alignment of Autonomous Cyber Defense Agents via Adaptive Policy Shaping"
Code for our paper Alignment of Autonomous Cyber Defense Agents via Adaptive Policy Shaping. This repository extends CybORG CAGE Challenge 2 to study how autonomous defenders can be steered towards operational priorities while retaining their ability to defend the network.
The original CAGE 2 challenge evaluates blue defenders against built-in red agents using the CAGE reward. This modified version adds operational objectives, modified attacker policies, and alignment methods for blue defenders. The simulator and experiment scripts remain in the CybORG/ layout. The original CAGE 2 guide retains the original challenge description.
The modified network includes authentication (Auth), database (Database), and front-end (Front) servers in addition to the user, enterprise, operational, and defender hosts. These servers represent operational priorities expressed through confidentiality (C), integrity (I), and availability (A):
| Server type | Operational priorities |
|---|---|
Authentication (Auth) | Confidentiality, integrity, availability |
Database (Database) | Confidentiality, availability |
Front-end (Front) | Integrity, availability |
Resilience scoring is adapted from Cadet et al. Quantitative Resilience Modeling for Autonomous Cyber Defense. The implementation builds on the associated CyberResilience code.
Our resilience metric records CIA drops at each step and combines them into a normalized resilience drop, with default weights of 0.6 for confidentiality, 0.2 for integrity, and 0.2 for availability. It also scores candidate blue actions for the attribute-based policy and alignment analysis. The generated scenarios assign CIA-critical servers zero ordinary CAGE host reward values, allowing operational alignment and general network defense performance to be evaluated separately.
Additional internal changes made to the CAGE 2 codebase are detailed in experiment README.
The original Meander and B-line attackers are extended to target the added operational servers and support host names across the generated network sizes. Evaluation uses five adversaries:
| Paper adversary | Python class | Impact targets |
|---|---|---|
| M-CIA | RedMeanderAgent_Combination_CIA | CIA-critical Auth, Database, and Front servers |
| M-CAGE | RedMeanderAgent_Combination_CAGE | Other eligible compromised hosts contributing to CAGE performance |
| Bline-C | B_lineAgent_Resilience_C | Auth or Database servers |
| Bline-I | B_lineAgent_Resilience_I | Auth or Front servers |
| Bline-A | B_lineAgent_Resilience_A | Auth, Database, or Front servers |
Training uses RedMeanderAgent_Combination, which selects between CIA and CAGE impact targets with equal probability at each impact-selection decision. The B-line variants select a target server type once per episode and resolve matching hosts from observations. See Meander_Resilience.py and B_line_resilience.py.
Use Python 3.10 or newer and install uv. From the repository root:
uv sync
cd CybORG
uv sync creates the project environment and installs CybORG and its dependencies. Run scripts with uv run ... and shell launchers with uv run bash ...; manual environment activation is not required.
| Location | Contents |
|---|---|
CybORG/CybORG/ | Simulator, agents, rewards, metrics, scenarios, and upstream tests |
CybORG/experiment_configs/ | Test-time, training-time, adaptive, scaling, and red-agent evaluation configurations |
CybORG/train_agents.py | Train baseline or reward-shaped blue agents |
CybORG/test_experiment.py | Evaluate baseline, fixed, and adaptive policies |
CybORG/test_experiment_training_time.py | Evaluate agents aligned during training |
CybORG/alignment_analysis.py | Compute alignment scores and analyze reward trade-offs |
CybORG/CybORG/Shared/Scenarios/scenario_generation.py | Generate scenarios with configurable host populations and reward settings |
Use train_agents.py to train a PPO defender against the mixed meandering attacker. PPOAgent trains the baseline; AlignedPPOAgent adds a reward term for the selected CIA attribute. The training-time bonus is called lambda in the paper and alpha in the CLI only; it is separate from the test-time interpolation weight.
To train one baseline checkpoint for base:
uv run python train_agents.py \
--scenario scalability_experiments/base \
--agent_type PPOAgent \
--red_agent RedMeanderAgent_Combination \
--alpha 0 --seed 0 --game_steps 200 --training_iter 500000 \
--agent_export checkpoints/PPO_base_0_seed0
To train a defender aligned to confidentiality with a bonus of 1.5:
uv run python train_agents.py \
--scenario scalability_experiments/base \
--agent_type AlignedPPOAgent \
--red_agent RedMeanderAgent_Combination \
--alignment_attribute C --alpha 1.5 \
--seed 0 --game_steps 200 --training_iter 500000 \
--agent_export checkpoints/PPO_base_1.5_C_seed0
For the complete training sweeps:
uv run bash experiment_configs/train_agents/train_base.sh
uv run bash experiment_configs/train_agents/train_trainingtime_alignment.sh
These launchers cover base, hosts_2, hosts_3, and hosts_4 with training seeds 0 1 2 3 4. Checkpoints are saved under CybORG/checkpoints/ as PPO_{scenario}_0_seed{seed}.zip for baselines and PPO_{scenario}_{bonus}_{attribute}_seed{seed}.zip for reward-shaped agents. After saving a checkpoint, train_agents.py also runs a short evaluation. Logging is disabled by default; add --wandb to a Python training command to enable.
Evaluation scripts load trained checkpoints and run the agents against the configured red adversary. The supplied configurations use 100 episodes, an evaluation seed starting at 42, and five training seeds. Episode horizons are 200, 300, 400, and 500 steps for base, hosts_2, hosts_3, and hosts_4, respectively.
test_experiment.py evaluates Blue Monitor and PPO baselines alongside fixed or adaptive test-time alignment:
# Baselines and fixed test-time interpolation
uv run python test_experiment.py --config experiment_configs/base.yaml
uv run python alignment_analysis.py --config experiment_configs/base.yaml
# Adaptive test-time policy selection
uv run python test_experiment.py --config experiment_configs/adaptive_alpha_config/base/cia.yaml
uv run python alignment_analysis.py --config experiment_configs/adaptive_alpha_config/base/cia.yaml
test_experiment_training_time.py evaluates the reward-shaped checkpoints: