RLベースのプロンプトインジェクションレッドチーミングのためのオープンフレームワーク
PIForgeは PISmith と Climbing the Hill の共有コードベースです。PISmithはプロンプトインジェクションレッドチーミングにおける報酬の疎性問題に対処し、RL攻撃者が稀に成功する攻撃を探索し学習できるようにします。PISmithを基盤として、Climbing the Hillはカリキュラム学習を用いて、より堅牢なフロンティアターゲットに対するレッドチーミング時のコールドスタート問題に対処します。
| コンポーネント | 場所 | 目的 |
|---|---|---|
| ベンチマーク | benchmarks/ | PIArena、InjecAgent、AgentDojo、AgentDyn、IPI Arena。 |
| トレーニングコア | train.py、core/、configs/ | 共有RLトレーナーとベンチマーク設定。 |
| エントリーポイント | scripts/、eval/ | トレーニングスクリプト1つ、カリキュラムスクリプト1つ、評価スクリプト1つ。 |
コマンドはリポジトリのルートから実行してください。トレーニングには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
3つのエントリーポイントは小さなインターフェースを共有しています：
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
ローカルの 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 を設定します。
同じ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 の下で公開されています。