PISmith: Reinforcement Learning-based Red Teaming for Prompt Injection Defenses(COLM 2026)

This is an official implementation of PISmith: Reinforcement Learning-based Red Teaming for Prompt Injection Defenses

Environment Setup

PISmith has been tested using Python 3.10 and CUDA Version: 12.9

1. Create a Python 3.10 conda environment

conda create -n PISmith python=3.10 -y
conda activate PISmith

2. Install dependencies

pip install -r requirements.txt

3. (Optional) Prepare the Meta-SecAlign model checkpoint

For experiments targeting the secalign defense, run the provided merge script to download and merge the base model with the SecAlign adapter:

python merge_meta_secalign.py

This downloads meta-llama/Llama-3.1-8B-Instruct and facebook/Meta-SecAlign-8B from HuggingFace, merges them, and saves the result to checkpoints/Meta-SecAlign-8B-merged/.

Scripts

PIArena

PIArena supports training and evaluation against a range of prompt injection defenses. Use the defense argument to select the target defense.

Supported defenses: secalign, none, promptguard, promptarmor, sandwich, instructional, datasentinel, piguard, datafilter

Training

bash scripts/train_piarena.sh <defense> [train_gpus] [target_gpu] [target_port]

Examples:

# Train against SecAlign defense
bash scripts/train_piarena.sh secalign

# Train against no defense (plain LLM)
bash scripts/train_piarena.sh none

Evaluation

bash scripts/eval_piarena.sh <checkpoint> <defense> [target_port] [target_gpu] [attacker_gpu] [attacker_port] [num_samples]

Examples:

# Evaluate against SecAlign (default settings)
bash scripts/eval_piarena.sh checkpoints/piarena/checkpoint-500 secalign

# Evaluate against no piguard, pass@10
bash scripts/eval_piarena.sh checkpoints/piarena_none/checkpoint-500 piguard

AgentDojo

Supports GPT-4o-mini, GPT-4o, GPT-5-nano, and local vLLM targets.

bash scripts/train_agentdojo.sh [target_type] [suites] [train_gpus]

# Default: GPT-4o-mini target on the firsr 7 injected task of workspace suite
bash scripts/train_agentdojo.sh

# Train on all suites (workspace, banking, travel, slack)
bash scripts/train_agentdojo.sh gpt4o-mini all

Evaluation:

bash scripts/eval_agentdojo.sh <checkpoint> [target_type] [eval_suites] [num_samples] [target_defense]

# Example
bash scripts/eval_agentdojo.sh checkpoints/agentdojo/checkpoint-500 gpt4o-mini

AgentDyn

AgentDyn is built on top of AgentDojo. Install it separately before running AgentDyn experiments:

git clone https://github.com/SaFo-Lab/AgentDyn.git
cd AgentDyn
pip install -e . --no-deps

Training supports AgentDyn suites such as github, dailylife, and shopping.

bash scripts/train_agentdyn.sh [target_type] [suites] [train_gpus]

# Example
bash scripts/train_agentdyn.sh gpt5-nano github "0,1,2,3"

Evaluation reports both pass@k and sample-level average ASR, and supports data-parallel attacker vLLM serving:

ATTACKER_GPUS=0,1,2,3 ATTACKER_DP_SIZE=4 \
bash scripts/eval_agentdyn.sh checkpoints/agentdyn/checkpoint-500 gpt5-nano "github,dailylife,shopping" 5

InjecAgent

Supports a local vLLM target, GPT-4o-mini, or multi-target mode.

bash scripts/train_injecagent.sh [target_type] [train_gpus] [target_gpu] [target_port]

# Default: local vLLM target (Meta-SecAlign-8B)
bash scripts/train_injecagent.sh

# GPT-4o-mini API target
bash scripts/train_injecagent.sh gpt4o-mini

Evaluation:

bash scripts/eval_injecagent.sh <checkpoint> [target_type] [target_gpu] [target_port] [eval_gpu] [num_samples]

# Example
bash scripts/eval_injecagent.sh checkpoints/injecagent/checkpoint-500

Experiment Results

Main Results (vs. Meta-SecAlign-8B, 13 Benchmarks)

PISmith is evaluated against 7 baselines spanning static, search-based, and RL-based attack categories. All RL-based methods report ASR@10 / ASR@1; static and search-based methods report ASR@1.

Utility–Robustness Trade-off (8 Defenses, Qwen3-4B-Instruct-2507)

PISmith ASR@1 averaged over 13 benchmarks. Utility measures task accuracy without attack.

It remains challenging for state-of-the-art defenses to simultaneously achieve high utility (≥0.70) and low ASR (≤0.60), revealing a fundamental utility–robustness trade-off.

Agentic Settings

InjecAgent

AgentDojo (best static baseline vs. PISmith)