albert-y1n/piforgeOffenes Framework für RL-basiertes Red Teaming bei Prompt-Injection, mit einem gemeinsamen Trainer, Curriculum Learning und Benchmarks wie AgentDojo, InjecAgent und PIArena.
Ein offenes Framework für RL-basiertes Prompt-Injection-Red-Teaming
PIForge ist die gemeinsame Codebasis für PISmith und Climbing the Hill. PISmith adressiert das Problem der spärlichen Belohnung beim Prompt-Injection-Red-Teaming und hilft RL-Angreifern, seltene erfolgreiche Angriffe zu erkunden und daraus zu lernen. Aufbauend auf PISmith nutzt Climbing the Hill Curriculum Learning, um das Cold-Start-Problem beim Red-Teaming robusterer Frontier-Ziele zu lösen.
| Komponente | Ort | Zweck |
|---|---|---|
| Benchmarks | benchmarks/ | PIArena, InjecAgent, AgentDojo, AgentDyn und IPI Arena. |
| Trainingskern | train.py, core/, configs/ | Gemeinsamer RL-Trainer und Benchmark-Konfigurationen. |
| Einstiegspunkte | scripts/, eval/ | Ein Trainingsskript, ein Curriculum-Skript und ein Evaluierungsskript. |
Führen Sie die Befehle aus dem Repository-Stammverzeichnis aus. Für das Training sind Python 3.10 und GPUs erforderlich.
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
Die drei Einstiegspunkte teilen sich eine kleine Schnittstelle:
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...]
Verwenden Sie Umgebungsvariablen wie ATTACKER_MODEL, TRAIN_SUITES, OUTPUT_DIR, LEARNING_RATE und NUM_TRAIN_EPOCHS, um einen Lauf zu ändern. DRY_RUN=1 gibt die Befehle aus, ohne Modelle zu starten.
Installieren Sie AgentDyn für das AgentDyn-github-Subset, das die standardmäßige Trainingssuite für die folgenden AgentDojo/AgentDyn-Rezepte ist:
git clone https://github.com/SaFo-Lab/AgentDyn.git
pip install -e AgentDyn --no-deps
export OPENAI_API_KEY="your-openai-api-key"
Trainieren Sie mit PISmith gegen GPT-4o-mini oder GPT-5-nano:
bash scripts/train.sh agentdyn gpt4o-mini
bash scripts/train.sh agentdyn gpt5-nano
Trainieren Sie mit Curriculum-RL in Richtung GPT-5.6-Luna oder GPT-5.6-Terra. Jede Stufe startet vom Angreifer-Checkpoint der vorherigen Stufe:
bash scripts/train_curriculum.sh nano-luna
bash scripts/train_curriculum.sh nano-luna-terra
Derselbe Einstiegspunkt für ein einzelnes Ziel setzt das Training von jedem Angreifermodell oder gespeicherten Checkpoint aus fort:
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
Evaluieren Sie ein veröffentlichtes Hugging-Face-Modell direkt.
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
Wählen Sie für AgentDojo-Suites agentdojo und setzen Sie TRAIN_SUITES oder EVAL_SUITES auf workspace, banking, travel oder 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
Bereiten Sie das lokale Meta-SecAlign-Ziel einmalig mit python merge_meta_secalign.py vor. Trainieren und evaluieren Sie dann mit dem InjecAgent-Datensatz:
bash scripts/train.sh injecagent secalign
bash scripts/eval.sh injecagent AlbertYin/injecagent_attacker_qwen3_4b_secalign secalign 10
Das Skript startet das lokale Ziel auf GPU 0 und verwendet standardmäßig die GPUs 1,2,3 für das Training. Setzen Sie TARGET_GPU, TRAIN_GPUS oder TARGET_URL, um dieses Setup zu ändern.
Dieselbe Meta-SecAlign-Vorbereitung gilt für die secalign-Verteidigung. PIArena verwendet einen eigenen Benchmark-Datensatz:
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
Andere Verteidigungen können über ihren Konfigurationsnamen ausgewählt werden, z. B. promptguard oder piguard. Setzen Sie TARGET_GPU, TRAIN_GPUS oder TARGET_URL, um andere GPUs oder einen vorhandenen Zielserver zu verwenden.
Wir veröffentlichen unsere trainierten Angreifer in der PIForge Hugging Face Collection.
PIForge wird unter der MIT License veröffentlicht.