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

Abstract

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. A recurrent PPO variant confirms that these distributional findings are not an artifact of feed-forward policies. The code is available at: https://github.com/raphsimon/NASimJax.