Pentest-R1: Towards Autonomous Penetration Testing Reasoning Optimized via Two-Stage Reinforcement Learning
Note:
- If you encounter any errors or issues, feel free to open an issue or submit a pull request.
Pentest-R1 is a two-stage reinforcement learning (RL) framework designed to substantially improve the reasoning capabilities of Large Language Models (LLMs) for autonomous penetration testing. It begins with offline RL on a curated dataset of over 500 real-world expert walkthroughs to instill core attack logic. Then, it applies online RL in an interactive Capture The Flag (CTF) environment, allowing the agent to learn robust error correction and adaptive strategies through direct environmental feedback.
Figure: The framework architecture of Pentest-R1, illustrating the two-stage training process.
Ensure your environment meets the following requirements before proceeding:
- Programming Language: Python 3.11.11
- Containerization: Docker
- Package Manager: Pip
-
Clone the Pentest-R1 repository:
git clone
-
Navigate to the project directory:
cd Pentest-R1 -
Install the required Python dependencies:
pip install -r requirements.txt
The training is divided into two main stages.
This stage trains the base LLM on the curated dataset of expert walkthroughs to learn foundational penetration testing logic.
Run the following command to start Stage 1 training:
python grpo_stage1.pyThis stage fine-tunes the agent from Stage 1 in a live, interactive CTF environment. This requires setting up the Intercode-CTF Docker environment first.
-
Build the CTF Environment: Navigate to the environment directory and build the Docker image.
cd train_ctf_env docker build -t intercode-ctf . cd ..
-
Run Stage 2 Training: Once the environment is ready, start the online RL training. This script will interact with the Docker container to provide real-time feedback to the agent.
python grpo_multi_turn_stage2.py
If you have any questions or suggestions, please open an issue on GitHub. Contributions, discussions, and improvements are always welcome!