albert-y1n/piforgeFramework ouvert pour le red teaming d'injection de prompts basé sur le RL, avec un entraîneur partagé, un apprentissage par curriculum et des benchmarks comme AgentDojo, InjecAgent et PIArena.
Un framework ouvert pour le red teaming d'injection de prompt basé sur le RL
💻 Code · 🤗 Modèles · 📜 Articles
PIForge est la base de code partagée pour PISmith et Climbing the Hill. PISmith traite le problème de récompense sparse dans le red teaming d'injection de prompt, aidant les attaquants RL à explorer et apprendre à partir d'attaques réussies rares. S'appuyant sur PISmith, Climbing the Hill utilise l'apprentissage par curriculum pour résoudre le problème de démarrage à froid lors du red teaming de cibles frontières plus robustes.
| Composant | Emplacement | Objectif |
|---|---|---|
| Benchmarks | benchmarks/ | PIArena, InjecAgent, AgentDojo, AgentDyn et IPI Arena. |
| Cœur d'entraînement | train.py, core/, configs/ | Entraîneur RL partagé et configurations de benchmarks. |
| Points d'entrée | scripts/, eval/ | Un script d'entraînement, un script de curriculum et un script d'évaluation. |
Exécutez les commandes depuis la racine du dépôt. L'entraînement nécessite Python 3.10 et des 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
Les trois points d'entrée partagent une petite interface :
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...]
Utilisez des variables d'environnement telles que ATTACKER_MODEL, TRAIN_SUITES, OUTPUT_DIR, LEARNING_RATE et NUM_TRAIN_EPOCHS pour modifier une exécution. DRY_RUN=1 affiche les commandes sans lancer les modèles.
Installez AgentDyn pour le sous-ensemble github d'AgentDyn, qui est la suite d'entraînement par défaut pour les recettes AgentDojo/AgentDyn ci-dessous :
git clone https://github.com/SaFo-Lab/AgentDyn.git
pip install -e AgentDyn --no-deps
export OPENAI_API_KEY="your-openai-api-key"
Entraînez avec PISmith contre GPT-4o-mini ou GPT-5-nano :
bash scripts/train.sh agentdyn gpt4o-mini
bash scripts/train.sh agentdyn gpt5-nano
Entraînez vers GPT-5.6-Luna ou GPT-5.6-Terra avec le RL par curriculum. Chaque étape démarre à partir du checkpoint d'attaquant de l'étape précédente :
bash scripts/train_curriculum.sh nano-luna
bash scripts/train_curriculum.sh nano-luna-terra
Le même point d'entrée mono-cible poursuit l'entraînement à partir de n'importe quel modèle attaquant ou checkpoint sauvegardé :
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
Évaluez directement un modèle Hugging Face publié.
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
Pour les suites AgentDojo, choisissez agentdojo et définissez TRAIN_SUITES ou EVAL_SUITES sur workspace, banking, travel ou 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
Préparez une fois la cible locale Meta-SecAlign avec python merge_meta_secalign.py. Ensuite, entraînez et évaluez avec le jeu de données InjecAgent :
bash scripts/train.sh injecagent secalign
bash scripts/eval.sh injecagent AlbertYin/injecagent_attacker_qwen3_4b_secalign secalign 10
Le script démarre la cible locale sur le GPU 0 et utilise les GPU 1,2,3 pour l'entraînement par défaut. Définissez TARGET_GPU, TRAIN_GPUS ou TARGET_URL pour modifier cette configuration.
La même préparation Meta-SecAlign s'applique à la défense secalign. PIArena utilise son propre jeu de données de benchmark :
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
D'autres défenses peuvent être sélectionnées par leur nom de configuration, comme promptguard ou piguard. Définissez TARGET_GPU, TRAIN_GPUS ou TARGET_URL pour utiliser d'autres GPU ou un serveur cible existant.
Nous publions nos attaquants entraînés dans la collection Hugging Face PIForge.
PIForge est publié sous la licence MIT.