PIForge

Открытый фреймворк для RL-based Red Teaming промпт-инъекций

💻 Код · 🤗 Модели · 📜 Статьи

---

PIForge — это общая кодовая база для PISmith и Climbing the Hill. PISmith решает проблему разреженного вознаграждения в red teaming промпт-инъекций, помогая RL-атакующим исследовать и учиться на редких успешных атаках. Развивая PISmith, Climbing the Hill использует curriculum learning для решения проблемы холодного старта при red teaming более устойчивых frontier-целей.

✨ Новости

• 2026.09 — Мы выпускаем: Climbing the Hill: Prompt Injection Red-Teaming Against Frontier Models with Curriculum Reinforcement Learning — метод curriculum RL, решающий проблему холодного старта в red-teaming промпт-инъекций и достигающий 93.8%/45.0% ASR@10 против GPT-5.6-Luna/GPT-5.6-Terra (против 0% у предыдущих RL-методов).
• 2026.07 — PISmith был принят на COLM 2026.

Что здесь есть

Компонент	Где	Назначение
Бенчмарки	benchmarks/	PIArena, InjecAgent, AgentDojo, AgentDyn и IPI Arena.
Ядро обучения	train.py, core/, configs/	Общий RL-тренер и конфигурации бенчмарков.
Точки входа	scripts/, eval/	Один скрипт обучения, один скрипт curriculum и один скрипт оценки.

Установка

Запускайте команды из корня репозитория. Для обучения требуются 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, которое является набором обучения по умолчанию для приведённых ниже рецептов AgentDojo/AgentDyn:

git clone https://github.com/SaFo-Lab/AgentDyn.git
pip install -e AgentDyn --no-deps
export OPENAI_API_KEY="your-openai-api-key"

Обучение с PISmith против GPT-4o-mini или GPT-5-nano:

bash scripts/train.sh agentdyn gpt4o-mini
bash scripts/train.sh agentdyn gpt5-nano

Обучение в направлении GPT-5.6-Luna или GPT-5.6-Terra с curriculum RL. Каждый этап начинается с чекпоинта атакующего предыдущего этапа:

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

InjecAgent

Подготовьте локальную цель Meta-SecAlign один раз с помощью python merge_meta_secalign.py. Затем обучите и оцените с набором данных 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, чтобы изменить эту конфигурацию.

PIArena

Та же подготовка 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. Установите TARGET_GPU, TRAIN_GPUS или TARGET_URL, чтобы использовать другие GPU или существующий сервер цели.

Выпущенные атакующие

Мы публикуем наши обученные модели-атакующие в коллекции PIForge на Hugging Face.

Статьи

• PISmith: Reinforcement Learning-based Red Teaming for Prompt Injection Defenses
• Climbing the Hill: Prompt Injection Red-Teaming Against Frontier Models with Curriculum Reinforcement Learning

Лицензия

PIForge выпускается под лицензией MIT.