Cyber Threat Mitigation with Knowledge-Infused Reinforcement Learning and LLM-Guided Policies
Md. Shamim Towhid, Shahrear Iqbal, Euclides Carlos Pinto Neto, Nashid Shahriar, Scott Buffett, Madeena Sultana, Adrian Taylor · 2025
As cyber threats continue to evolve, there is a need for autonomous cyber defense (ACD) strategies capable of fast and context-aware responses. Reinforcement learning (RL) has shown promise for automating cyber defense by exploring and learning effective countermeasures, yet it often struggles with sparse reward signals and insufficient context to handle diverse attack scenarios. Furthermore, the convergence time taken by an RL agent is often high, which makes it difficult to train the RL agent in online settings. To address these challenges, we propose a large language model (LLM)-enhanced RL method that builds and queries a knowledge graph (KG) derived from agent-environment interactions. We leverage the pre-trained knowledge of an LLM on different cybersecurity frameworks and use the LLM to analyze a part of the KG to generate appropriate actions for the RL agent. We infuse the knowledge extracted from the LLM into the RL agent’s training loop in two ways. First, the state vector of the RL agent is augmented with the most effective action and its corresponding reward, as determined from the KG. Second, the suggested action from the LLM is used as a reference policy. In addition, we introduce a regularization term in the loss function to make the RL policy close to the reference policy. To validate our approach, we develop a custom RL environment guided by the MITRE ATT&CK framework, enabling the agent to generate tailored mitigation strategies for detected cyber attacks. Experimental results show that our proposed approach significantly outperforms the baseline RL by over $75 \%$ in terms of taking better mitigation actions.