NASimJax: ペネトレーションテスト向けのGPU加速ポリシー学習フレームワーク

要約: ペネトレーションテスト、脆弱性を特定するためにサイバー攻撃を模倣する実践は、固有に部分的に観測可能であり、大規模なアクション空間を特徴とする複雑な逐次意思決定タスクである。強化学習（RL）ポリシーの訓練はこの領域で根本的なボトルネックに直面している。既存のシミュレータは現実的なネットワークシナリオを大規模に訓練するには遅すぎ、その結果、一般化に失敗するポリシーが生じる。我々は NASimJax を提示する。Network Attack Simulator（NASim）の完全な JAX ベース再実装で、元のシミュレータより最大 100 倍の環境スループットを達成する。訓練パイプライン全体をハードウェアアクセラレータ上で実行することにより、NASimJax は以前は不可能だった固定された計算予算の下で、より大規模なネットワークでの実験を可能にする。自動化されたペネトレーションテストを Contextual POMDP として定式化し、構造的に多様で解決可能性が保証されたシナリオを生成するネットワーク生成パイプラインを導入する。これらを合わせて、ゼロショットポリシーの一般化を研究するための原理的な基盤を提供する。フレームワークを用いて、アクション空間のスケーリングと最大 40 ホストのネットワークでの一般化を調査する。Prioritized Level Replay は、Domain Randomization よりも密な訓練分布をより適切に扱うことが特に大規模なスケールであることを示し、疎なトポロジーでの訓練は訓練中に見られたトポロジより密度が高いトポロジ上でも分布外一般化を改善する暗黙のカリキュラムを生み出すことを発見する。線形に増大するアクション空間を扱うため、2 段階のアクション分解（2SAS）を提案し、スケールで平坦なアクションマスキングを大幅に上回る。最後に、Prioritized Level Replay のエピソードリセット挙動と 2SAS のクレジット割り当て構造の相互作用から生じる故障モードを特定する。NASimJax は、RL ベースのペネトレーションテストを進展させるための、迅速で柔軟かつ現実的なプラットフォームを提供する。