NASimJax: A GPU-Accelerated Policy Learning Framework for Penetration Testing

Poster B: Monday -- 16:00 - 18:00

Raphael Simon, José Carrasquel, Elli Makdis Antoun, Wim Mees, Pieter Jules Karel Libin

Keywords: Reinforcement Learning, Penetration Testing, Unsupervised Environment Design, Contextual POMDP

Penetration testing—the practice of simulating cyberattacks to identify vulnerabilities—is a complex sequential decision-making task that is inherently partially observable and features large action spaces. Existing RL simulators for this domain are CPU-bound and fixed to narrow scenarios, making it infeasible to train policies that generalize across networks. We present NASimJax, a JAX-native framework that formulates penetration testing as a Contextual POMDP and introduces a network generation pipeline producing structurally diverse, guaranteed-solvable scenarios. The framework reaches up to 80$\times$ higher environment throughput than previous simulators, enabling experiments on larger networks and tractable hyperparameter searches. We provide PPO and PQN baselines and conduct the first systematic evaluation of unsupervised environment design for penetration testing. We find that Prioritized Level Replay and ACCEL handle dense training distributions better than Domain Randomization, and that training on sparser topologies yields an implicit curriculum that improves generalization—even to topologies denser than those seen during training. The code is available at: https://anonymous.4open.science/r/NASimJax-D7C6.