🤖 machine learning

NASimJax: GPU-Accelerated Policy Learning Framework for Penetration Testing

This paper introduces NASimJax, a GPU-accelerated JAX-based framework that significantly accelerates reinforcement learning for penetration testing by enabling scalable training on diverse network scenarios, leading to improved zero-shot generalization through novel techniques like two-stage action decomposition and optimized curriculum learning.

Original authors: Raphael Simon, José Carrasquel, Wim Mees, Pieter Libin

Published 2026-03-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Raphael Simon, José Carrasquel, Wim Mees, Pieter Libin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to break into a highly secure bank vault. But there's a catch: the robot can't see the whole building, it doesn't know where the guards are, and it has to figure out the layout by poking around blindly. This is essentially what Penetration Testing (or "ethical hacking") is: simulating cyberattacks to find weaknesses before bad guys do.

The problem is that teaching a robot (or an AI) to do this is incredibly hard and slow. Here is a breakdown of the paper's solution, NASimJax, using simple analogies.

1. The Problem: The "Slow Motion" Simulator

Imagine you are trying to learn how to play chess by playing one game a day. By the time you've played 1,000 games, you've only spent a few months learning. Now, imagine the AI needs to play millions of games to get good.

The Old Way: Previous tools were like a single person running a simulation on a standard laptop. They were so slow that the AI could only practice on tiny, simple networks. When the AI tried to face a real, complex network, it failed because it had never seen anything that big.
The Bottleneck: The computer's brain (CPU) was talking to its muscle (GPU) too slowly. It was like a manager shouting instructions to a worker across a noisy factory floor, waiting for a reply before giving the next order.

2. The Solution: NASimJax (The "Super-Speed" Factory)

The authors built NASimJax, which is a complete rewrite of their simulation software using a special language called JAX.

The Analogy: Instead of one person shouting instructions, imagine a massive factory where 1,000 robots are working in perfect sync, all at the exact same time.
The Result: NASimJax is 100 times faster than the old version. It can simulate millions of cyber-attacks in the time it used to take to simulate thousands. This allows the AI to practice on huge, complex networks that were previously impossible to train on.

3. The Training Ground: A "Shape-Shifting" Maze

To make the AI smart enough to handle any network, they didn't just build one maze; they built a machine that generates infinite, unique mazes.

The Concept: They call this a Contextual POMDP. Think of it like a video game where the level layout, the number of enemies, and the traps change randomly every time you start a new game.
The Goal: The AI isn't memorizing a specific map. It's learning the rules of hacking so well that it can walk into a brand new, unseen building and figure out how to break in without ever having seen it before. This is called Zero-Shot Generalization.

4. The Challenge: The "Too Many Buttons" Problem

As the network gets bigger (more computers to hack), the number of possible actions the AI can take explodes.

The Old Way (Flat Action Masking): Imagine a remote control with 10,000 buttons. Most are broken or useless for the current room. The AI has to guess which one works. It's like trying to find a needle in a haystack by randomly pressing buttons.
The New Way (2SAS - Two-Stage Action Selection): The authors introduced a smarter strategy. Instead of pressing one of 10,000 buttons, the AI does it in two steps:
1. Step 1: "Which computer should I target?" (Selects a host).
2. Step 2: "What do I do to that computer?" (Selects an action).
- The Analogy: It's like a detective first deciding which suspect to interview, and then deciding what question to ask. This breaks the massive problem into two smaller, manageable ones, making learning much faster and more accurate.

5. The Secret Sauce: The "Smart Curriculum"

How do you teach an AI to hack a massive network without overwhelming it?

Domain Randomization (DR): Throwing the AI into random mazes of all sizes. It's like throwing a student into random math problems. It works okay for small problems, but gets messy with big ones.
Prioritized Level Replay (PLR): This is the "Smart Curriculum." The system watches the AI. If the AI fails at a specific type of maze, the system remembers it and makes the AI practice that specific type of maze again until it gets better.
- The Finding: They discovered that training the AI on sparser, simpler networks first (like a small office) actually made it better at hacking dense, complex networks (like a massive data center) later on. It's like learning to ride a bike on a quiet street before trying to race in a crowded city.

6. The "Oops" Moment: When Things Break

The researchers found a weird glitch when combining the "Smart Curriculum" (PLR) with the "Two-Step Strategy" (2SAS) on very large networks.

The Glitch: The system kept resetting the game before the AI could finish a long mission. Because the AI was trying to learn two steps at once (Target + Action), it got confused about who to blame for the failure. It was like a relay race where the baton was dropped, but the runners kept blaming each other for the drop, and the coach kept resetting the race before they could try again.
The Lesson: This taught them that for very complex tasks, you have to be careful about how you reset the training environment.

Summary

NASimJax is a super-fast, flexible training ground for AI hackers. By making the simulation 100x faster and teaching the AI to break down big problems into smaller steps, the researchers created a system that can learn to hack complex networks it has never seen before. It proves that if you train an AI on the right kind of "practice levels" (starting simple and getting harder), it can become a master hacker capable of handling real-world cyber threats.

1. Problem Statement

Automated penetration testing (simulating cyberattacks to find vulnerabilities) is a complex sequential decision-making task modeled as a Partially Observable Markov Decision Process (POMDP). Agents must navigate large, combinatorial action spaces and long decision horizons without full knowledge of the network topology.

Current approaches face two critical bottlenecks:

Simulation Speed: Existing simulators (e.g., NASim, CyberBattleSim) are CPU-bound Python implementations. They cannot generate experience fast enough to utilize modern GPU accelerators, creating a throughput bottleneck that prevents training on realistic, large-scale networks.
Generalization: Policies often overfit to specific training scenarios. Existing simulators lack the structural diversity and guaranteed solvability required to train agents that generalize to unseen network topologies (Zero-Shot Policy Transfer).

2. Methodology: NASimJax

The authors present NASimJax, a complete reimplementation of the Network Attack Simulator (NASim) using JAX, a Python library for high-performance numerical computing and accelerator-oriented array computation.

Core Architecture & Design

Contextual POMDP Formulation: The environment is formalized as a Contextual MDP (CMDP). Each episode is conditioned on a "context" (a specific network instance) sampled from a distribution $p(c)$ . The agent does not observe the context directly but must infer network properties through interaction.
Network Generation Pipeline: A new generator creates structurally diverse, realistic, and guaranteed-solvable scenarios.
- Parameters: Controls include host count ( $N_h$ ), subnet count ( $N_s$ ), topology density ( $t_d$ ), service/process density, and sensitive host density.
- Constraints: Ensures every subnet has at least one service and every sensitive host has exploitable paths, preventing "dead ends."
- Curriculum via Density: By varying topology density, the system naturally generates a curriculum of varying difficulty (fewer reachable hosts = shorter horizons).
State & Observation Space:
- State: Encoded as batched adjacency matrices and host properties (OS, services, access levels) using boolean flags and one-hot encoding to minimize memory footprint.
- Observation: Partial; the agent only sees the outcome of its last action (e.g., scan results, exploit success/failure) and a history of past observations.
Action Space: Discrete actions include reconnaissance (scans) and exploitation (exploits, privilege escalation). The action space grows linearly with the number of hosts ( $|A| \propto |H|$ ), creating a massive search space for large networks.

Algorithmic Innovations

Two-Stage Action Selection (2SAS): To handle linearly growing action spaces, the authors decompose the decision process:
- Stage 1: Select a target host (masked to exclude unreachable/undiscovered hosts).
- Stage 2: Select an action for that specific host.
- Benefit: Reduces the effective branching factor from $|H| \times |Actions|$ to $\max(|H|, |Actions|)$ , significantly lowering exploration difficulty.
Reward Scaling: Rewards are scaled by the theoretical maximum return ( $N_s \cdot V_h$ ) to normalize the learning signal across networks of different sizes and densities, stabilizing value function estimation.
Unsupervised Environment Design (UED): The framework evaluates Domain Randomization (DR) and Prioritized Level Replay (PLR) to improve generalization.

3. Key Contributions

100× Speedup: NASimJax achieves up to 1.6 million steps per second on a single entry-level GPU (NVIDIA RTX A4000), a 100× improvement over the original CPU-based NASim. This is achieved via JIT compilation and vectorized environment stepping (vmap).
Scalable Action Decomposition: The introduction of 2SAS demonstrates that factoring action spaces is critical for large-scale networks, outperforming standard action masking significantly as network size increases.
Principled Generalization Framework: By combining CMDP formulation with a controllable network generator, the paper provides a rigorous testbed for studying zero-shot policy transfer.
Identification of Failure Modes: The authors identify a specific failure mode where PLR's episode-reset behavior interacts destructively with 2SAS's credit assignment in large, dense networks, leading to training collapse.

4. Experimental Results

Experiments were conducted on networks ranging from 16 to 40 hosts.

Performance Scaling:
- 16 Hosts: 2SAS offered no significant advantage over flat masking (problem was too simple).
- 26 Hosts: 2SAS achieved an 82% solve rate vs. 66% for masking.
- 40 Hosts: The gap widened dramatically. 2SAS achieved 42% solve rate, while flat masking dropped to 14%.
Zero-Shot Policy Transfer (ZSPT):
- Training on Sparse Networks: Training on low-density topologies (fewer active hosts) consistently yielded better generalization to denser, unseen networks than training directly on the target distribution. This acts as an implicit curriculum.
- PLR vs. DR: Prioritized Level Replay (PLR) generally outperformed Domain Randomization (DR) at larger scales by constructing a natural curriculum.
The PLR-2SAS Failure: At 40 hosts with high density ( $t_d=0.15$ $t_{d} = 0.15$ ), the combination of PLR and 2SAS collapsed (0% solve rate).
- Cause: PLR resets environments frequently. In dense 40-host networks, episodes are too long to complete in a single rollout. 2SAS cannot distinguish between a bad host choice and a bad action choice when the episode fails, leading to mutual degradation of both policy heads. PLR-2SAS failed to populate the replay buffer with successful trajectories to recover.

5. Significance

Enabling Large-Scale RL Research: NASimJax removes the computational barrier that previously prevented training RL agents on realistic, large-scale network scenarios. It makes hyperparameter searches and extensive experimentation feasible.
Practical Security Insights: The finding that training on sparse networks improves generalization to dense networks suggests a cost-effective strategy for security teams: train on cheaper, simpler simulations to build robust policies before deploying them on complex targets.
Methodological Guidance: The paper highlights the importance of credit assignment structures in decomposed action spaces and warns against naive combinations of curriculum learning (PLR) and complex action decomposition (2SAS) in long-horizon tasks.
Open Source Platform: By providing a flexible, JAX-native framework, NASimJax allows the community to rapidly prototype new algorithms for cybersecurity, moving beyond fixed benchmarks to dynamic, research-driven environments.