Skip to content

Graph Observations

Cyber defense is a graph problem — hosts are nodes, reachability is edges — yet most RL environments flatten the network into a vector. NetForge exposes a graph-native observation so GNN policies can operate on the network structure directly.

It lives in netforge_rl.core.graph_obs and maps trivially onto PyTorch Geometric or jraph.

Building a graph view

from netforge_rl.core.graph_obs import build_graph_observation

g = build_graph_observation(env.global_state, agent_id='blue_dmz')
# g['node_features']  -> (100, 8) float32
# g['edge_index']     -> (2, E)  int64  (COO source/target)
# g['edge_attr']      -> (E, 1)  float32 (1.0 if the edge crosses a subnet)
# g['node_mask']      -> (100,)  float32 (1.0 for real, visible hosts)
# g['n_nodes'], g['n_edges']

Node features are [privilege, online, compromised, decoy, is_domain_controller, subnet, cvss, edr_active]. Edges follow routing reachability, and cross-subnet edges are flagged in edge_attr — a direct signal for lateral-movement-aware policies.

Fog of war is respected. For a red agent_id, only hosts it has discovered are visible; blue agents see the full network. Padding hosts are always masked out.

As a wrapper

GraphObservationWrapper injects the graph into each agent's info under 'graph' without changing the observation dict, so existing policies keep working:

from netforge_rl.environment.graph_wrapper import GraphObservationWrapper
from netforge_rl.environment.parallel_env import NetForgeRLEnv

env = GraphObservationWrapper(NetForgeRLEnv({'scenario_type': 'ransomware'}))
obs, infos = env.reset(seed=0)
graph = infos['blue_dmz']['graph']

Converting to PyTorch Geometric

from netforge_rl.core.graph_obs import to_pyg

data = to_pyg(g)   # torch_geometric.data.Data(x, edge_index, edge_attr)

to_pyg imports torch / torch_geometric lazily, so the builder itself has no heavy dependencies. For jraph, pass node_features, edge_index, and edge_attr straight into a GraphsTuple.