Reinforcement Learning-Based Vulnerability Prioritization

Using CISA Known Exploited Vulnerabilities for Data-Driven Security Operations

This repository contains the complete implementation of our research on reinforcement learning-based vulnerability prioritization using the CISA KEV catalog, NVD CVSS scores, and EPSS data.

🎯 Overview

Modern organizations face thousands of published vulnerabilities with limited resources for remediation. This work demonstrates that reinforcement learning can learn optimal prioritization policies from real-world exploitation data, achieving:

• 98.4% classification accuracy (DQN)
• 3,587.50 average reward (3,663 points improvement over random baseline)
• 10-minute training time (production-ready)
• Balanced prioritization (52% medium, 48% immediate)

📊 Key Results

Method	Accuracy	Avg Reward	F1 (Macro)	Training Time
Random Baseline	N/A	-75.50	N/A	N/A
XGBoost	100.0%	N/A	100.0%	<1 min
DQN (Ours)	98.4%	3,587.50	65.7%	~10 min
PPO	46.9%	2,822.00	21.3%	~15 min

DQN Performance Highlights:

• ✅ 3,663 reward point improvement over random baseline
• ✅ Balanced strategy: 52% medium priority, 48% immediate
• ✅ 100% recall on high-urgency vulnerabilities
• ✅ Production-ready: 10-minute training time
• ✅ Reproducible: Fixed random seeds, public code

Feature Importance:

• 🔴 Ransomware flag: 77.9%
• 🟠 CVSS score: 19.9%
• 🟡 Days since added: 2.2%
• 🟢 EPSS: <1%

🚀 Quick Start

Installation

# Clone repository
git clone https://github.com/GitSene/RL-KEV-Vulnerability-Prioritization.git
cd RL-KEV-Vulnerability-Prioritization

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

Download and Enrich Data

# Download CISA KEV catalog
python data/download_kev.py

# Enrich with NVD CVSS and EPSS (requires API key)
python src/data_processing/kev_enrichment.py --api-key YOUR_NVD_API_KEY

Train Models

# Train DQN agent
python scripts/train_dqn.py --episodes 200

# Train PPO agent
python scripts/train_ppo.py --episodes 200

# Train XGBoost baseline
python scripts/train_xgboost.py

Evaluate

# Run complete evaluation
python scripts/evaluate_all.py

# Quick demo with trained model
python scripts/demo.py

📁 Repository Structure

RL-KEV-Vulnerability-Prioritization/
├── data/                    # Data download and storage
│   ├── download_kev.py     # KEV catalog downloader
│   ├── kev_enriched.csv    # Enriched dataset (1,464 CVEs)
│   └── .gitkeep
├── src/                     # Source code
│   ├── data_processing/    # KEV enrichment pipeline
│   │   └── kev_enrichment.py
│   ├── environment/        # RL environment (Gymnasium)
│   │   └── vuln_env.py
│   ├── agents/             # DQN and PPO implementations
│   │   ├── dqn_agent.py
│   │   └── ppo_agent.py
│   ├── baselines/          # XGBoost baseline
│   │   └── xgboost_baseline.py
│   └── evaluation/         # Evaluation scripts
│       └── evaluate.py
├── models/                  # Trained model weights
│   ├── dqn_model.pth
│   ├── ppo_model.pth
│   └── .gitkeep
├── notebooks/              # Jupyter notebooks for analysis
├── results/                # Output figures and tables
│   ├── figures/           # PNG visualizations
│   └── tables/            # CSV results
├── scripts/                # Executable training scripts
│   ├── train_dqn.py
│   ├── train_ppo.py
│   ├── train_xgboost.py
│   ├── evaluate_all.py
│   ├── reproduce_results.py
│   └── demo.py
├── tests/                  # Unit tests
├── docs/                   # Documentation
│   └── INSTALL.md
├── requirements.txt        # Python dependencies
├── .gitignore
├── LICENSE                # MIT License
├── CITATION.cff           # Citation metadata
├── CONTRIBUTING.md        # Contribution guidelines
└── README.md              # This file

🔬 Methodology

Dataset Construction

1. CISA KEV Catalog: 1,464 confirmed exploited vulnerabilities (November 2025)
2. NVD CVSS v3.1: Technical severity scores (0–10)
3. EPSS: Exploitation probability predictions (0–1)
4. Features: cvss_score, epss, epss_percentile, days_since_added, ransomware_flag

Urgency Distribution:

• Low urgency: 7 (0.5%)
• Medium urgency: 771 (52.7%)
• High urgency: 686 (46.8%)

RL Formulation

Markov Decision Process (MDP):

• State Space: 5-dimensional continuous feature vector
• Action Space: 4 discrete actions
  • a₀: Monitor only
  • a₁: Patch within 30 days
  • a₂: Patch within 7 days
  • a₃: Patch immediately
• Reward Function: Urgency alignment + SLA compliance penalties
• Environment: Custom Gymnasium-compatible implementation

Algorithms

• DQN: Value-based learning with experience replay (50k buffer)
• PPO: Policy-gradient learning with actor-critic architecture
• XGBoost: Traditional ML baseline for comparison

📈 Results Summary

DQN Performance

• Classification: 98.4% accuracy with balanced prioritization
• Recall: 100% on high-urgency vulnerabilities (no critical CVEs missed)
• Action Distribution:
  • Patch within 30 days: 52.0%
  • Patch immediately: 47.9%
  • Patch within 7 days: 0.1%
  • Monitor only: 0.0%
• Convergence: ~175 episodes (~10 minutes on CPU)

PPO Performance

• Strategy: Aggressive safety-first (100% immediate patching)
• Convergence: 8× faster than DQN (~20 episodes)
• Trade-off: Lower reward but faster learning

Feature Importance Analysis

XGBoost feature importance reveals:

1. Ransomware flag: 77.9% (dominant predictor)
2. CVSS score: 19.9%
3. Days since added: 2.2%
4. EPSS probability: <1%
5. EPSS percentile: <1%

Insight: Binary exploitation evidence (KEV membership, ransomware campaigns) dominates over probabilistic predictions (EPSS) in the KEV context.

📄 Citation

If you use this code or data in your research, please cite:

@article{habibi2025rl,
  title={Reinforcement Learning-Based Vulnerability Prioritization Using {CISA} Known Exploited Vulnerabilities},
  author={Habibi, Babek},
  journal={IEEE Transactions on Information Forensics and Security},
  year={2025},
  note={Under Review}
}

📜 License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

• CISA for maintaining the Known Exploited Vulnerabilities catalog
• NIST for the National Vulnerability Database
• FIRST for the Exploit Prediction Scoring System (EPSS)

📧 Contact

• Author: Babek Habibi
• Email: bnorouzlou19519@ucumberlands.edu
• Institution: University of the Cumberlands, Department of Computer Science
• GitHub: @GitSene

🔗 Resources

• CISA KEV Catalog
• NVD Database
• EPSS Project
• IEEE TIFS Journal

🌟 Related Work

This work builds upon and complements recent advances in vulnerability prioritization:

• Shimizu & Hashimoto (2025): Vulnerability Management Chaining - Decision tree integration of CVSS, EPSS, and KEV
• NIST LEV (2025): Likely Exploited Vulnerabilities metric

Our RL-based approach is complementary: while decision trees provide static filtering rules, reinforcement learning optimizes sequential decisions under operational constraints.

---

⭐ If you find this work useful, please consider starring the repository!

---

📊 Project Status: Under review for IEEE TIFS
🔄 Last Updated: December 2025
💻 Maintained by: @GitSene