albert-y1n/piforgeRL 기반 프롬프트 인젝션 레드팀을 위한 오픈 프레임워크로, 공유 트레이너, 커리큘럼 학습, 그리고 AgentDojo, InjecAgent, PIArena와 같은 벤치마크를 제공합니다.
RL 기반 프롬프트 인젝션 레드팀을 위한 오픈 프레임워크
PIForge는 **PISmith**와 **Climbing the Hill**의 공용 코드베이스입니다. PISmith는 프롬프트 인젝션 레드팀에서의 희소 보상 문제를 해결하여, RL 공격자가 드물게 성공하는 공격을 탐색하고 학습할 수 있도록 돕습니다. PISmith를 기반으로, Climbing the Hill은 커리큘럼 학습을 사용하여 더 견고한 프론티어 타깃을 레드팀할 때 발생하는 콜드 스타트 문제를 해결합니다.
| Component | Where | Purpose |
|---|---|---|
| Benchmarks | benchmarks/ | PIArena, InjecAgent, AgentDojo, AgentDyn, and IPI Arena. |
| Training core | train.py, core/, configs/ | Shared RL trainer and benchmark configurations. |
| Entry points | scripts/, eval/ | One training script, one curriculum script, and one evaluation script. |
저장소 루트에서 명령을 실행하세요. 학습에는 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을 설치하세요:
git clone https://github.com/SaFo-Lab/AgentDyn.git
pip install -e AgentDyn --no-deps
export OPENAI_API_KEY="your-openai-api-key"
GPT-4o-mini 또는 GPT-5-nano에 대해 PISmith로 학습하세요:
bash scripts/train.sh agentdyn gpt4o-mini
bash scripts/train.sh agentdyn gpt5-nano
커리큘럼 RL로 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
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을 설정하세요.
동일한 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와 같은 구성 이름으로 선택할 수 있습니다. 다른 GPU 또는 기존 타깃 서버를 사용하려면 TARGET_GPU, TRAIN_GPUS, TARGET_URL을 설정하세요.
학습된 공격자를 PIForge Hugging Face 컬렉션에서 공개합니다.
PIForge는 MIT License에 따라 공개됩니다.