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

arXiv: https://arxiv.org/abs/2603.19864 · HTML: https://mlsanity.com/paper/2603.19864/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 29 of its 50 references that are in the map.
- First posted on arXiv 20 Mar 2026.

## Builds on (most cited first, 29 of 29)

- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [Solving Rubik's Cube with a Robot Hand](https://mlsanity.com/paper/1910.07113/) (2019)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [Contextual Markov Decision Processes](https://mlsanity.com/paper/1502.02259/) (2015)
- [Prioritized Level Replay](https://mlsanity.com/paper/2010.03934/) (2020)
- [Evolving Curricula with Regret-Based Environment Design](https://mlsanity.com/paper/2203.01302/) (2022)
- [A Closer Look at Invalid Action Masking in Policy Gradient Algorithms](https://mlsanity.com/paper/2006.14171/) (2020)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Mastering the Game of Stratego with Model-Free Multiagent Reinforcement Learning](https://mlsanity.com/paper/2206.15378/) (2022)
- [Replay-Guided Adversarial Environment Design](https://mlsanity.com/paper/2110.02439/) (2021)
- [Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability](https://mlsanity.com/paper/2107.06277/) (2021)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [Learning values across many orders of magnitude](https://mlsanity.com/paper/1602.07714/) (2016)
- [Simplifying Deep Temporal Difference Learning](https://mlsanity.com/paper/2407.04811/) (2024)
- [JaxMARL: Multi-Agent RL Environments and Algorithms in JAX](https://mlsanity.com/paper/2311.10090/) (2023)
- [CybORG: A Gym for the Development of Autonomous Cyber Agents](https://mlsanity.com/paper/2108.09118/) (2021)
- [Empirical Design in Reinforcement Learning](https://mlsanity.com/paper/2304.01315/) (2023)
- [Action Space Shaping in Deep Reinforcement Learning](https://mlsanity.com/paper/2004.00980/) (2020)
- [NASimEmu: Network Attack Simulator & Emulator for Training Agents Generalizing to Novel Scenarios](https://mlsanity.com/paper/2305.17246/) (2023)
- [A Method for Evaluating Hyperparameter Sensitivity in Reinforcement Learning](https://mlsanity.com/paper/2412.07165/) (2024)
- [POMDPs Make Better Hackers: Accounting for Uncertainty in Penetration Testing](https://mlsanity.com/paper/1307.8182/) (2013)
- [JaxUED: A simple and useable UED library in Jax](https://mlsanity.com/paper/2403.13091/) (2024)
- [CybORG++: An Enhanced Gym for the Development of Autonomous Cyber Agents](https://mlsanity.com/paper/2410.16324/) (2024)
- [Fleet Control using Coregionalized Gaussian Process Policy Iteration](https://mlsanity.com/paper/1911.10121/) (2019)
- [Evaluating COVID-19 vaccine allocation policies using Bayesian m-top exploration](https://mlsanity.com/paper/2301.12822/) (2023)
- [Learning Robust Penetration Testing Policies under Partial Observability: A systematic evaluation](https://mlsanity.com/paper/2509.20008/) (2025)

## Abstract (arXiv)

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× 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.