PIForge

一个用于基于强化学习的提示注入红队测试的开放框架

💻 代码 · 🤗 模型 · 📜 论文

---

PIForge 是 PISmith 和 Climbing the Hill 的共享代码库。PISmith 解决了提示注入红队测试中的稀疏奖励问题，帮助强化学习攻击者探索并从罕见的成功攻击中学习。在 PISmith 的基础上，Climbing the Hill 使用课程学习来解决对更鲁棒的前沿目标进行红队测试时的冷启动问题。

✨ 新闻

• 2026.09 — 我们发布了：Climbing the Hill: Prompt Injection Red-Teaming Against Frontier Models with Curriculum Reinforcement Learning，一种课程强化学习方法，解决了提示注入红队测试中的冷启动问题，针对 GPT-5.6-Luna/GPT-5.6-Terra 达到 93.8%/45.0% ASR@10（而此前的强化学习方法为 0%）。
• 2026.07 — PISmith 被 COLM 2026 接收。

内容概览

组件	位置	用途
基准测试	benchmarks/	PIArena、InjecAgent、AgentDojo、AgentDyn 和 IPI Arena。
训练核心	train.py、core/、configs/	共享的强化学习训练器和基准配置。
入口点	scripts/、eval/	一个训练脚本、一个课程脚本和一个评估脚本。

环境配置

从仓库根目录运行命令。训练需要 Python 3.10 和 GPU。

git clone https://github.com/albert-y1n/PIForge.git
cd PIForge
conda create -n piforge python=3.10 -y
conda activate piforge
pip install -r requirements.txt

这三个入口点共享一个简单的接口：

bash scripts/train.sh <benchmark> <target> [train_gpus] [train.py overrides...]
bash scripts/train_curriculum.sh [curriculum]
bash scripts/eval.sh <benchmark> <attacker path or HF model ID> <target> [num_samples] [eval.py overrides...]

使用诸如 ATTACKER_MODEL、TRAIN_SUITES、OUTPUT_DIR、LEARNING_RATE 和 NUM_TRAIN_EPOCHS 等环境变量来更改运行配置。DRY_RUN=1 会打印命令而不启动模型。

AgentDojo / AgentDyn

为 AgentDyn 的 github 子集安装 AgentDyn，这是下面 AgentDojo/AgentDyn 配方默认的训练套件：

git clone https://github.com/SaFo-Lab/AgentDyn.git
pip install -e AgentDyn --no-deps
export OPENAI_API_KEY="your-openai-api-key"

使用 PISmith 针对 GPT-4o-mini 或 GPT-5-nano 进行训练：

bash scripts/train.sh agentdyn gpt4o-mini
bash scripts/train.sh agentdyn gpt5-nano

使用课程强化学习针对 GPT-5.6-Luna 或 GPT-5.6-Terra 进行训练。每个阶段都从前一阶段的攻击者检查点开始：

bash scripts/train_curriculum.sh nano-luna
bash scripts/train_curriculum.sh nano-luna-terra

同一个单目标入口点可以从任何攻击者模型或保存的检查点继续训练：

ATTACKER_MODEL=checkpoints/agentdyn_curriculum/nano-luna-terra/stage_3_gpt-5.6-terra/checkpoint-XXX \
OUTPUT_DIR=checkpoints/muse_spark \
  bash scripts/train.sh agentdyn muse-spark-1.2

直接评估已发布的 Hugging Face 模型。

bash scripts/eval.sh agentdyn AlbertYin/agentdojo_attacker_qwen3_4b_5_nano gpt5-nano 10
bash scripts/eval.sh agentdyn AlbertYin/agentdojo_attacker_qwen3_4b_5.6_terra gpt-5.6-terra 10

对于 AgentDojo 套件，选择 agentdojo 并将 TRAIN_SUITES 或 EVAL_SUITES 设置为 workspace、banking、travel 或 slack：

TRAIN_SUITES=workspace bash scripts/train.sh agentdojo gpt4o-mini
EVAL_SUITES=workspace bash scripts/eval.sh agentdojo AlbertYin/agentdojo_attacker_qwen3_4b_4o_mini gpt4o-mini 10

InjecAgent

使用 python merge_meta_secalign.py 一次性准备本地 Meta-SecAlign 目标。然后使用 InjecAgent 数据集进行训练和评估：

bash scripts/train.sh injecagent secalign
bash scripts/eval.sh injecagent AlbertYin/injecagent_attacker_qwen3_4b_secalign secalign 10

该脚本默认在 GPU 0 上启动本地目标，并使用 GPU 1、2、3 进行训练。设置 TARGET_GPU、TRAIN_GPUS 或 TARGET_URL 来更改此配置。

PIArena

同样的 Meta-SecAlign 准备工作适用于 secalign 防御。PIArena 使用其自己的基准数据集：

bash scripts/train.sh piarena secalign
bash scripts/train.sh piarena none
bash scripts/eval.sh piarena AlbertYin/piarena_attacker_qwen3_4b_secalign secalign 10

其他防御可以通过其配置名称来选择，例如 promptguard 或 piguard。设置 TARGET_GPU、TRAIN_GPUS 或 TARGET_URL 以使用不同的 GPU 或现有的目标服务器。

已发布的攻击者

我们在 PIForge Hugging Face 集合中发布了我们训练好的攻击者。

论文

• PISmith: Reinforcement Learning-based Red Teaming for Prompt Injection Defenses
• Climbing the Hill: Prompt Injection Red-Teaming Against Frontier Models with Curriculum Reinforcement Learning

许可证

PIForge 在 MIT 许可证下发布。