A Vulnerability-Driven Gym for Training Autonomous Cyber Agents with Reinforcement Learning

IEEE International Conference on Ubiquitous Intelligence and Computing(UIC) (UIC) 2024.9.26,

Weixia Cai, Huashan Chen, Han Miao, Feng Liu, Yong Zhang, Xiaojia Yang

Abstract

Autonomous agents allow frequent and regular pen etration testing to be performed, which is increasingly necessary as there is a growing need for skilled and expensive specialists to examine and improve the cybersecurity of organizations. Besides, given the success of reinforcement learning in various domains, it is promising to explore the application of its methods to the development of intelligent and autonomous cyber agents. Currently, the cyber agents are trained only on limited scenarios and have weak transfer capabilities to new environments. This study is intended to train red agents that can conduct automated attacks in unexplored new cyber scenarios. To achieve this, we propose a vulnerability-driven approach called VDGym, which is a four-step cyber gym generator. By combining two augmentation methods, VDGym is able to generate diverse cyber scenarios. To verify the validity of our method, we train our agents on plenty of RL training scenarios randomly generated by VDGym and test the trained model on existing well-designed scenarios without reinforcement. The experimental results show that the proposed method can effectively help train agents and generalize better to new scenarios.