NASimJax: GPU-Accelerated Policy Learning Framework for Penetration Testing

2026-03-20

AI Summary

This work addresses the limitations of existing reinforcement learning approaches for automated penetration testing, which struggle to scale and generalize in large real-world networks due to inefficient simulators. The authors formulate the problem as a Contextual Partially Observable Markov Decision Process (Contextual POMDP) and develop an end-to-end GPU-accelerated simulator in JAX, coupled with a structured yet diverse network generation mechanism to enable zero-shot generalization studies. Key contributions include a two-stage action decomposition scheme (2SAS) that mitigates linearly growing action spaces, the discovery that training on sparse network topologies induces an implicit curriculum enhancing out-of-distribution generalization, and the identification of credit assignment failures arising from the interaction between prioritized level replay and action decomposition. Experiments demonstrate a 100× throughput improvement, enabling training on networks of up to 40 hosts; 2SAS substantially outperforms flat action masking, and the sparse-training strategy maintains strong generalization even on dense networks.

Technology Category

• Machine Learning: Reinforcement Learning
• Multiagent Systems: Adversarial Agents
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web

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. Training reinforcement learning (RL) policies for this domain faces a fundamental bottleneck: existing simulators are too slow to train on realistic network scenarios at scale, resulting in policies that fail to generalize. We present NASimJax, a complete JAX-based reimplementation of the Network Attack Simulator (NASim), achieving up to 100x higher environment throughput than the original simulator. By running the entire training pipeline on hardware accelerators, NASimJax enables experimentation on larger networks under fixed compute budgets that were previously infeasible. We formulate automated penetration testing as a Contextual POMDP and introduce a network generation pipeline that produces structurally diverse and guaranteed-solvable scenarios. Together, these provide a principled basis for studying zero-shot policy generalization. We use the framework to investigate action-space scaling and generalization across networks of up to 40 hosts. We find that Prioritized Level Replay better handles dense training distributions than Domain Randomization, particularly at larger scales, and that training on sparser topologies yields an implicit curriculum that improves out-of-distribution generalization, even on topologies denser than those seen during training. To handle linearly growing action spaces, we propose a two-stage action decomposition (2SAS) that substantially outperforms flat action masking at scale. Finally, we identify a failure mode arising from the interaction between Prioritized Level Replay's episode-reset behaviour and 2SAS's credit assignment structure. NASimJax thus provides a fast, flexible, and realistic platform for advancing RL-based penetration testing.

Problem

• penetration testing
• reinforcement learning
• simulator bottleneck
• policy generalization
• large action spaces

Innovation

• NASimJax
• GPU acceleration
• Contextual POMDP
• action decomposition
• zero-shot generalization