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.