albert-y1n/piforgeFramework abierto para red teaming de inyección de prompts basado en RL, con un entrenador compartido, aprendizaje por currículo y benchmarks como AgentDojo, InjecAgent y PIArena.
Un framework abierto para el red teaming de inyección de prompts basado en RL
💻 Código · 🤗 Modelos · 📜 Artículos
PIForge es la base de código compartida para PISmith y Climbing the Hill. PISmith aborda el problema de recompensa dispersa en el red teaming de inyección de prompts, ayudando a los atacantes RL a explorar y aprender de ataques exitosos poco frecuentes. Partiendo de PISmith, Climbing the Hill utiliza aprendizaje por currículo para abordar el problema de arranque en frío al hacer red teaming contra objetivos frontera más robustos.
| Componente | Ubicación | Propósito |
|---|---|---|
| Benchmarks | benchmarks/ | PIArena, InjecAgent, AgentDojo, AgentDyn e IPI Arena. |
| Núcleo de entrenamiento | train.py, core/, configs/ | Entrenador RL compartido y configuraciones de benchmarks. |
| Puntos de entrada | scripts/, eval/ | Un script de entrenamiento, un script de currículo y un script de evaluación. |
Ejecuta los comandos desde la raíz del repositorio. El entrenamiento requiere Python 3.10 y GPUs.
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
Los tres puntos de entrada comparten una interfaz pequeña:
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...]
Usa variables de entorno como ATTACKER_MODEL, TRAIN_SUITES, OUTPUT_DIR, LEARNING_RATE y NUM_TRAIN_EPOCHS para modificar una ejecución. DRY_RUN=1 imprime los comandos sin lanzar modelos.
Instala AgentDyn para el subconjunto github de AgentDyn, que es la suite de entrenamiento predeterminada para las recetas de AgentDojo/AgentDyn a continuación:
git clone https://github.com/SaFo-Lab/AgentDyn.git
pip install -e AgentDyn --no-deps
export OPENAI_API_KEY="your-openai-api-key"
Entrena con PISmith contra GPT-4o-mini o GPT-5-nano:
bash scripts/train.sh agentdyn gpt4o-mini
bash scripts/train.sh agentdyn gpt5-nano
Entrena hacia GPT-5.6-Luna o GPT-5.6-Terra con RL basado en currículo. Cada etapa parte del checkpoint del atacante de la etapa anterior:
bash scripts/train_curriculum.sh nano-luna
bash scripts/train_curriculum.sh nano-luna-terra
El mismo punto de entrada de objetivo único continúa el entrenamiento desde cualquier modelo atacante o checkpoint guardado:
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
Evalúa directamente un modelo publicado en 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
Para las suites de AgentDojo, elige agentdojo y establece TRAIN_SUITES o EVAL_SUITES en workspace, banking, travel o 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
Prepara el objetivo local Meta-SecAlign una vez con python merge_meta_secalign.py. Luego entrena y evalúa con el conjunto de datos de InjecAgent:
bash scripts/train.sh injecagent secalign
bash scripts/eval.sh injecagent AlbertYin/injecagent_attacker_qwen3_4b_secalign secalign 10
El script inicia el objetivo local en la GPU 0 y usa las GPUs 1,2,3 para el entrenamiento de forma predeterminada. Establece TARGET_GPU, TRAIN_GPUS o TARGET_URL para cambiar esta configuración.
La misma preparación de Meta-SecAlign se aplica a la defensa secalign. PIArena usa su propio conjunto de datos 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
Otras defensas se pueden seleccionar por su nombre de configuración, como promptguard o piguard. Establece TARGET_GPU, TRAIN_GPUS o TARGET_URL para usar GPUs diferentes o un servidor objetivo existente.
Publicamos nuestros atacantes entrenados en la colección PIForge de Hugging Face.
PIForge se publica bajo la Licencia MIT.