NASimJax: GPU-Accelerated Policy Learning Framework for Penetration Testing
作者: Raphael Simon, José Carrasquel, Wim Mees, Pieter Libin
分类: cs.LG, cs.CR
发布日期: 2026-03-20
💡 一句话要点
NASimJax：用于渗透测试的GPU加速策略学习框架
🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture) 支柱四：生成式动作 (Generative Motion)
关键词: 渗透测试 强化学习 GPU加速 JAX 动作空间分解
📋 核心要点
- 现有渗透测试模拟器速度慢，难以在真实网络场景中训练强化学习策略，导致策略泛化能力差。
- 提出NASimJax，基于JAX重新实现NASim，利用GPU加速，显著提升环境吞吐量，支持更大规模网络训练。
- 通过实验验证，优先级别回放和两阶段动作分解（2SAS）能有效提升策略在复杂网络中的泛化能力。
📝 摘要（中文）
渗透测试是一种模拟网络攻击以识别漏洞的实践，它是一个复杂的序列决策任务，本质上是部分可观察的，并且具有巨大的动作空间。为此领域训练强化学习（RL）策略面临一个根本瓶颈：现有的模拟器速度太慢，无法在现实的网络场景中大规模训练，导致策略无法泛化。我们提出了NASimJax，一个完整的基于JAX的Network Attack Simulator（NASim）的重新实现，实现了比原始模拟器高达100倍的环境吞吐量。通过在硬件加速器上运行整个训练流程，NASimJax能够在先前不可行的固定计算预算下，在更大的网络上进行实验。我们将自动化渗透测试形式化为Contextual POMDP，并引入了一个网络生成流程，该流程产生结构多样且保证可解的场景。这些共同为研究零样本策略泛化提供了原则性基础。我们使用该框架来研究动作空间缩放和跨最多40个主机的网络的泛化。我们发现，优先级别回放比领域随机化更好地处理密集训练分布，尤其是在更大规模上，并且在更稀疏的拓扑上进行训练会产生一个隐式课程，从而改善了分布外的泛化，即使在比训练期间看到的更密集的拓扑上也是如此。为了处理线性增长的动作空间，我们提出了一种两阶段动作分解（2SAS），它在规模上大大优于平面动作掩码。最后，我们确定了由优先级别回放的episode-reset行为和2SAS的信用分配结构之间的相互作用引起的故障模式。因此，NASimJax为推进基于RL的渗透测试提供了一个快速、灵活和现实的平台。
🔬 方法详解
问题定义：论文旨在解决渗透测试中强化学习策略训练效率低下的问题。现有的网络攻击模拟器（如NASim）速度慢，无法支持大规模、真实的渗透测试场景，导致训练出的策略泛化能力不足，难以应对实际的网络环境。
核心思路：论文的核心思路是利用JAX框架和GPU加速，对现有的NASim模拟器进行重新实现，从而大幅提升模拟器的运行速度。同时，通过改进训练策略（如优先级别回放）和动作空间处理方法（如两阶段动作分解），来提升强化学习策略的训练效率和泛化能力。
技术框架：NASimJax框架主要包含以下几个部分：1) 基于JAX的NASim模拟器：提供快速的网络攻击环境模拟；2) 网络生成流程：生成结构多样且保证可解的渗透测试场景；3) 强化学习算法：用于训练渗透测试策略，例如DQN；4) 训练策略改进：包括优先级别回放和领域随机化等；5) 动作空间处理：采用两阶段动作分解（2SAS）来处理大规模动作空间。
关键创新：论文的关键创新点在于：1) 基于JAX的NASim模拟器，实现了高达100倍的性能提升；2) 提出了两阶段动作分解（2SAS）方法，有效处理了大规模动作空间问题，显著优于传统的动作掩码方法；3) 提出了使用稀疏拓扑进行训练，从而获得隐式课程学习，提升泛化能力。
关键设计：在模拟器方面，使用JAX进行数值计算，充分利用GPU的并行计算能力。在动作空间处理方面，2SAS将动作分解为两个阶段：首先选择攻击目标，然后选择攻击方式。在训练策略方面，使用优先级别回放来关注更有价值的经验，并采用episode-reset机制来避免策略陷入局部最优。
🖼️ 关键图片
📊 实验亮点
NASimJax实现了比原始NASim模拟器高达100倍的吞吐量提升。实验表明，优先级别回放比领域随机化更适合处理密集训练分布，尤其是在大规模网络中。两阶段动作分解（2SAS）在处理大规模动作空间时，性能显著优于传统的动作掩码方法。在稀疏拓扑上训练可以提升策略在更密集拓扑上的泛化能力。
🎯 应用场景
该研究成果可应用于自动化渗透测试、网络安全态势感知、安全策略评估等领域。通过快速模拟和训练，可以帮助安全人员更有效地发现和修复网络漏洞，提升网络安全防御能力，并为未来的智能安全防御系统提供技术支撑。
📄 摘要（原文）
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.