Description

This record contains code associated with the paper Structural Generalization in Autonomous Cyber Incident Response with Message-Passing Neural Networks and Reinforcement Learning.

It is used to train an agent with Reinforcement Learning (RL) agent that uses a Graph Neural Network (GNN) to embed the state. The agent is trained and tested on the CAGE 2 environment. The task of the agent is to defend the network modeled by the CAGE 2 challenge from an attacking agent. The GNN is used to facilitate agent generalization across multiple network topologies, which is tested with the policy_eval script after training.

Dependencies

Dependencies are managed by uv. PyTorch is not included in the dependency list to allow users to install GPU or CPU variants as they please.

Usage

The scripts train_mlp and train_gnn will respectively train MLP and GNN models on all the scenarios in the scenarios folder. This may take a bit of time. Settings can only be changed by editing the scripts directly, as no CLI interface has been implemented.

If you want to evaluate trained models, run policy_eval which by default will run two evaluations, one with different GNN layers and one with all models on all scenarios. This can take a long time to complete. To run a shorter evaluation, edit the corresponding values in the script.

The scenario files which define the network structures/graphs are in scenarios/.

Code Sources

The modules CybORG, gnn and stable_baselines3 contains modified and unmodified code from external sources. The copyright to this code belongs to the respective authors.

• Oracle SAGE, SR-DRL code adapted for SB3. gnn/ module.
  • Oracle Sage is itself derived from SR-DRL, GNN RL policy.
• Stable Baselines 3, PPO implementation. stable_baselines3 module.
• CybORG, CybORG/CAGE2 environment. CybORG/ module.