Pentest-R1 是一个两阶段强化学习框架，目标是训练大语言模型（LLM）执行自主渗透测试与 CTF 挑战。它把离线强化学习（基于专家演练数据）与在线强化学习（在交互式 CTF 环境中）结合起来，以培养稳健的攻击推理能力。
用途 该项目面向安全研究与渗透测试自动化方向，用于训练和评估 LLM 在攻击推理、命令生成与多轮交互中的表现。它同时覆盖离线学习与在线交互两类训练范式。
主要能力
- 两阶段训练流程：Stage 1 为离线强化学习，Stage 2 为在线强化学习。
- Stage 1 在 500 多条真实世界专家渗透测试演练（walkthrough）数据集上训练基础 LLM。
- Stage 2 在交互式 CTF 环境中进行微调，并接收实时反馈。
- 提供奖励模型（PentestRewardModel）用于评估渗透测试动作，包含成功奖励、部分奖励与失败惩罚。
- 提供 CTF 环境封装（CTFEnvironment），基于 Docker 与 InterCode-CTF 镜像，支持启动挑战容器、执行命令与清理。
- 提供多轮推理智能体（MultiTurnAgent），根据当前观测生成下一步动作，并维护对话历史。
- 提供评估循环（evaluate_on_ctf_challenges），在 CTF 挑战集上评估模型，记录是否解出与轮数。
- 使用 GRPO（Group Relative Policy Optimization）进行训练。
使用方式 环境要求：Python 3.11.11；Docker（用于 Stage 2 与可复现环境）；NVIDIA Container Toolkit（用于 GPU 支持）；CUDA 12.4 运行时（可选，用于 GPU 加速）。
基础安装： git clone https://github.com/KHenryAegis/Pentest-R1.git cd Pentest-R1 pip install -r requirements.txt
Docker 环境（推荐用于可复现性）：
source setup-docker.sh
docker run --rm -it
--name pentest-r1
-v "$(pwd)":/root/Pentest-R1
-v ~/.cache/huggingface:/root/.cache/huggingface
-v ~/.cache/triton:/root/.cache/triton
-v ~/.cache/torch_extensions:/root/.cache/torch_extensions
-w /root/Pentest-R1
--gpus all
--net=host
pentest-r1:ubuntu22.04
Stage 1 离线强化学习： python grpo_stage1.py 关键配置包括模型名 unsloth/Meta-Llama-3.1-8B-Instruct、max_seq_length=4096、load_in_4bit=True，以及 PEFT/LoRA 配置（r=16、lora_alpha=32、target_modules 为 q_proj/k_proj/v_proj/o_proj、lora_dropout=0.05、bias=none、task_type=CAUSAL_LM）。
Stage 2 在线强化学习： cd train_ctf_env docker build -t intercode-ctf . cd .. python grpo_multi_turn_stage2.py
数据加载与预处理：使用 datasets.load_dataset 加载 JSON 格式的专家演练数据（jsonl），示例字段包括 challenge、steps（含 action、command、reasoning）与 flag。preprocess_function 将挑战与步骤格式化为 prompt 与 response。
训练配置：使用 transformers.TrainingArguments，示例参数包括 output_dir="./pentest-r1-checkpoints"、num_train_epochs=3、per_device_train_batch_size=4、gradient_accumulation_steps=4、learning_rate=5e-5、warmup_steps=100、logging_steps=10、save_steps=500、save_total_limit=3、fp16=True、report_to="wandb"、remove_unused_columns=False。
环境变量：HF_TOKEN（模型下载）、WANDB_API_KEY（日志）、CUDA_VISIBLE_DEVICES（GPU 选择）、HF_HOME、TRITON_CACHE_DIR 等缓存目录。
适用场景
- 安全研究中的自主渗透测试与 CTF 挑战求解。
- 训练与评估 LLM 的攻击推理与多轮交互能力。
- 在可复现的 Docker 环境中进行强化学习实验。
限制与注意事项
- 资料未提供收费、价格或商业授权信息。
- 资料未提供模型在具体基准上的量化效果数据。
- 资料未提供除已列平台外的其他平台支持信息。
- 训练依赖 GPU 与 Docker 环境；CUDA 显存不足时可采用梯度检查点、减小 batch size、使用 8-bit 或 4-bit 量化。
- 研究使用 unsloth==2025.5.10；若不可用，requirements.txt 中指定了最接近的版本（unsloth_zoo==2025.5.11 及指定 commit 的 unsloth）。
- 模型加载失败时可尝试不使用量化（dtype=torch.float16、load_in_4bit=False），必要时清理 HuggingFace 缓存。
- Docker 容器问题可通过查看日志、验证网络连通性、使用 --no-cache 重建镜像排查。
- 多 GPU 训练可设置 ddp_find_unused_parameters=False 与 local_rank，并使用 torchrun 启动。
参考信息：论文 arXiv:2508.07382；基础模型 Meta-Llama-3.1-8B-Instruct；微调框架 Unsloth；交互训练环境 InterCode-CTF。