🤖 machine learning

Learning Robust Penetration Testing Policies under Partial Observability: A systematic evaluation

This paper presents a systematic evaluation demonstrating that history aggregation techniques, such as frame-stacking and recurrent architectures, significantly enhance the convergence speed and robustness of Proximal Policy Optimization (PPO) agents in partially observable penetration testing scenarios across varying network sizes.

Original authors: Raphael Simon, Pieter Libin, Wim Mees

Published 2026-06-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Raphael Simon, Pieter Libin, Wim Mees

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 a detective trying to solve a complex mystery in a building you've never entered before. You don't have a blueprint, you don't know where the valuable items are hidden, and you can't see through the walls. Every time you try a door, you might find a key, a locked room, or nothing at all. This is the essence of penetration testing (or "pentesting"): ethical hackers trying to find security holes in computer networks.

This paper is about teaching a computer (an AI) to act like that detective, but with a twist: the building changes its layout every time the detective enters.

The Problem: The "Blind" Detective

In the real world, hackers don't get a full map of the network. They have to scan, guess, and remember what they found. In computer science terms, this is called Partial Observability. The AI only sees a tiny slice of the truth at any given moment.

Previous attempts to train AI for this task were like training a detective in a building that never changed. The AI would memorize, "If I turn left at the red door, I find the treasure." But if the building layout changed even slightly, the AI would get lost. It was "overfitting"—memorizing the specific test instead of learning the skill.

The Solution: The "Shapeshifting" Building

The authors created a new training ground called StochNASim. Think of this as a magical building that completely rebuilds itself every time you enter.

New Layout: The number of rooms (hosts) changes (sometimes 5, sometimes 8).
New Contents: The furniture (software and services) inside the rooms changes.
New Doors: The locks and keys (vulnerabilities) are different every time.

This forces the AI to stop memorizing specific paths and start learning how to think like a detective: "I need to look around, remember what I found, and then decide what to do next."

The Experiment: How to Remember?

The researchers asked: "How do we teach an AI to remember what it found in a changing building?" They tested four different "memory strategies" using a standard AI learning method called PPO (Proximal Policy Optimization):

The "No Memory" Baseline: The AI looks at the current room and guesses. It forgets everything the moment it moves. (Like a detective with amnesia).
Frame Stacking: The AI looks at the last few "snapshots" of the building to see what changed. (Like looking at a short video clip of the last few seconds).
Recurrent Networks (LSTM & TrXL): These are complex, brain-like architectures designed to remember long stories. They are like detectives with super-complex brains trying to hold the entire history of the building in their head.
Augmented Observations (The Winner): This is a clever trick. Instead of trying to "remember" in a complex way, the AI simply keeps a growing checklist. Every time it finds a piece of information (e.g., "Room 3 has a red door"), it adds it to a permanent list that it can see forever. It doesn't forget; it just accumulates facts.

The Surprising Results

The results were counter-intuitive to what many experts expected:

The Complex Brains Failed: The fancy, complex memory systems (LSTM and TrXL) struggled. They tried to be too clever, often resorting to a "brute-force" strategy: just trying every possible door until one opened. They were slow and inefficient.
The Simple Checklist Won: The Augmented Observations method (the growing checklist) was the clear winner. It learned four times faster than the other methods.
Why? The task wasn't about remembering a complex story; it was about gathering facts. The AI didn't need a complex brain to remember; it just needed a simple way to keep a list of what it had already discovered. The "checklist" approach allowed the AI to stop scanning rooms it had already checked and focus on the ones it hadn't.

The Takeaway

The paper concludes that for this specific type of cybersecurity task, simple is better.

You don't need a super-complex AI brain to be a good hacker-detective. You just need a good notepad. By giving the AI a simple way to keep a running list of what it has found, it learned to be efficient, robust, and able to handle completely new network layouts without getting confused.

The authors also proved that training in a "shapeshifting" environment (StochNASim) is crucial. If you train an AI on a static, unchanging network, it becomes a bad detective in the real world. But if you train it in a chaotic, changing environment, it learns to adapt and succeed anywhere.

In short: To teach an AI to hack networks, don't give it a super-complex memory; give it a sticky note that never gets erased. And make sure you practice in a building that changes every time you walk in.

Technical Summary: Learning Robust Penetration Testing Policies under Partial Observability

Problem Statement

Penetration testing (pentesting) is a sequential decision-making process where security professionals simulate cyberattacks to identify vulnerabilities. This task is inherently characterized by partial observability; agents do not have access to the full network state (topology, host configurations, running services) but must discover this information incrementally through reconnaissance actions. While Reinforcement Learning (RL) has been successfully applied to automate pentesting, most existing approaches formalize the problem as a fully observable Markov Decision Process (MDP). This assumption oversimplifies real-world conditions where information gathering is a core component of the task. Furthermore, current methods often suffer from overfitting, where agents memorize specific network configurations rather than learning transferable policies that generalize to unseen scenarios. The challenge lies in developing agents that can effectively aggregate historical observations to build a belief state, adapt to stochastic network dynamics, and generalize across varying network sizes and topologies.

Methodology

Environment: StochNASim

The authors extend the Network Attack Simulator (NASim) to create StochNASim, a new environment designed to better reflect real-world complexity. Key modifications include:

Stochasticity: Unlike NASim, which uses fixed topologies, StochNASim regenerates the network topology, host properties (OS, services, processes), and initial states at every episode reset.
Variable Network Sizes: The environment supports networks of varying sizes (5 to 8 hosts), forcing the agent to handle different action space dimensions and observation counts within a single policy.
Partial Observability: Observations are transient and action-specific. Scanning a host reveals its attributes (e.g., OS, services) only at that moment; without explicit memory mechanisms, this information is lost in subsequent steps.
Reward Structure: The agent is rewarded for discovering hosts and compromising sensitive targets while penalized for the cost of actions (scans are cheaper than exploits).

Algorithmic Approaches

The study conducts a systematic empirical evaluation of Proximal Policy Optimization (PPO) variants designed to handle partial observability, comparing them against a vanilla PPO baseline (which lacks specific memory mechanisms). The evaluated methods include:

Frame-Stacking (PPO-FS): Concatenates the last $f_n$ observations to provide short-term temporal context.
Augmented Observations (PPO-AO): Explicitly aggregates the history of observations by maintaining a matrix where each element is the maximum of the current and previous observation values. This creates a persistent record of discovered information, mimicking how human pentesters maintain notes.
Recurrent Architectures:
- LSTM (PPO-LSTM): Uses Long Short-Term Memory networks to implicitly encode observation history.
- Transformer-XL (PPO-TrXL): Employs attention mechanisms to capture long-term dependencies in the observation sequence.

Experimental Setup

Hyperparameter Tuning: A rigorous search using Optuna (250 trials per algorithm) was conducted to ensure fair comparison and avoid tuning artifacts.
Evaluation: Policies were tested across five control seeds and evaluated on newly generated environments to measure generalization.
Analysis: Beyond quantitative metrics (cumulative reward, convergence speed), the authors performed qualitative analysis via action sequence visualization and action type distribution to understand the learned strategies.

Key Results

Performance and Convergence

PPO-AO Superiority: The PPO-AO (Augmented Observations) approach consistently outperformed all other methods. It converged four times faster than the remaining approaches (reaching optimal performance in ~1M steps vs. 4M steps for others) and achieved the highest cumulative rewards.
Complex Architectures Underperform: Contrary to the assumption that complex memory mechanisms are necessary for POMDPs, PPO-LSTM and PPO-TrXL underperformed. PPO-TrXL notably failed to learn scanning behaviors, instead adopting a "brute-force" strategy of sequentially trying all possible exploits. PPO-LSTM struggled to learn effective history representations, requiring nearly twice the steps to complete tasks compared to PPO-AO.
Baseline Limitations: Vanilla PPO, lacking any history mechanism, learned suboptimal policies characterized by high stochasticity and uncertainty, failing to generalize effectively.

Generalization and Stochasticity

Robustness: Policies trained in the stochastic StochNASim environment demonstrated strong generalization to unseen network configurations.
Overfitting in Static Environments: A critical finding was that policies trained on fixed-scenario environments (like standard NASim) failed catastrophically when tested on novel networks. These agents appeared to memorize specific observation-action sequences (lookup-table strategy) rather than learning to aggregate information, highlighting the necessity of stochastic training environments for robust policy learning.

Qualitative Behavioral Analysis

Manual inspection of learned policies revealed distinct strategic differences despite similar final rewards in some cases:

PPO-AO: Learned an optimal, human-like strategy: systematically scanning to gather intelligence, retaining that information, and then executing targeted exploits.
PPO-TrXL: Learned a brute-force approach, attempting every exploit and privilege escalation sequentially without prior reconnaissance.
PPO-FS: Showed intermediate behavior, performing some scanning but relying heavily on sequential exploitation.

Significance and Contributions

The paper claims four primary contributions:

StochNASim Environment: The adaptation of NASim into a partially observable, stochastic environment with variable network sizes, which better reflects real-world challenges and facilitates the learning of transferable policies.
Systematic Evaluation: A comprehensive empirical comparison of PPO-based approaches for handling partial observability in penetration testing, including rigorous hyperparameter tuning.
Challenge to Architectural Complexity: The demonstration that simple observation augmentation (PPO-AO) significantly outperforms complex recurrent (LSTM) and transformer (TrXL) architectures in this domain. This challenges the conventional assumption that sophisticated memory mechanisms are required for effective decision-making in partially observable security tasks.
Qualitative Policy Analysis: The provision of a detailed analysis of learned behaviors, revealing that algorithms with similar quantitative performance can learn vastly different (and sometimes suboptimal) strategies, emphasizing the need to look beyond learning curves when evaluating RL agents.

The authors conclude that for the specific domain of automated penetration testing, where the goal is information retrieval and retention, explicit history aggregation via simple augmentation is more effective and computationally efficient than learning implicit representations through complex neural architectures. Furthermore, the study underscores that stochastic training environments are crucial for preventing overfitting and ensuring policies are robust enough for real-world deployment.