albert-y1n/piforgeEstrutura aberta para red teaming de injeção de prompt baseado em RL, com um trainer compartilhado, aprendizado por currículo e benchmarks como AgentDojo, InjecAgent e PIArena.
Um Framework Aberto para Red Teaming de Injeção de Prompt baseado em RL
💻 Código · 🤗 Modelos · 📜 Artigos
PIForge é a base de código compartilhada para PISmith e Climbing the Hill. O PISmith aborda o problema de recompensa esparsa em red teaming de injeção de prompt, ajudando atacantes de RL a explorar e aprender com ataques bem-sucedidos raros. Com base no PISmith, o Climbing the Hill usa aprendizado por currículo para resolver o problema de inicialização a frio ao fazer red teaming em alvos de fronteira mais robustos.
| Componente | Onde | Propósito |
|---|---|---|
| Benchmarks | benchmarks/ | PIArena, InjecAgent, AgentDojo, AgentDyn e IPI Arena. |
| Núcleo de treinamento | train.py, core/, configs/ | Treinador RL compartilhado e configurações de benchmark. |
| Pontos de entrada | scripts/, eval/ | Um script de treinamento, um script de currículo e um script de avaliação. |
Execute os comandos a partir da raiz do repositório. O treinamento requer Python 3.10 e 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
Os três pontos de entrada compartilham uma pequena 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...]
Use variáveis de ambiente como ATTACKER_MODEL, TRAIN_SUITES, OUTPUT_DIR, LEARNING_RATE e NUM_TRAIN_EPOCHS para alterar uma execução. DRY_RUN=1 imprime os comandos sem iniciar os modelos.
Instale o AgentDyn para o subconjunto github do AgentDyn, que é a suíte de treinamento padrão para as receitas AgentDojo/AgentDyn abaixo:
git clone https://github.com/SaFo-Lab/AgentDyn.git
pip install -e AgentDyn --no-deps
export OPENAI_API_KEY="your-openai-api-key"
Treine com o PISmith contra GPT-4o-mini ou GPT-5-nano:
bash scripts/train.sh agentdyn gpt4o-mini
bash scripts/train.sh agentdyn gpt5-nano
Treine em direção ao GPT-5.6-Luna ou GPT-5.6-Terra com RL por currículo. Cada estágio começa a partir do checkpoint do atacante do estágio anterior:
bash scripts/train_curriculum.sh nano-luna
bash scripts/train_curriculum.sh nano-luna-terra
O mesmo ponto de entrada de alvo único continua o treinamento a partir de qualquer modelo atacante ou checkpoint salvo:
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
Avalie um modelo lançado no Hugging Face diretamente.
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 as suítes do AgentDojo, escolha agentdojo e defina TRAIN_SUITES ou EVAL_SUITES como 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
Prepare o alvo local Meta-SecAlign uma vez com python merge_meta_secalign.py. Em seguida, treine e avalie com o conjunto de dados InjecAgent:
bash scripts/train.sh injecagent secalign
bash scripts/eval.sh injecagent AlbertYin/injecagent_attacker_qwen3_4b_secalign secalign 10
O script inicia o alvo local na GPU 0 e usa as GPUs 1,2,3 para treinamento por padrão. Defina TARGET_GPU, TRAIN_GPUS ou TARGET_URL para alterar essa configuração.
A mesma preparação do Meta-SecAlign se aplica à defesa secalign. O PIArena usa seu próprio conjunto de dados 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
Outras defesas podem ser selecionadas pelo nome de sua configuração, como promptguard ou piguard. Defina TARGET_GPU, TRAIN_GPUS ou TARGET_URL para usar GPUs diferentes ou um servidor de alvo existente.
Lançamos nossos atacantes treinados na coleção PIForge no Hugging Face.
PIForge é distribuído sob a Licença MIT.