lalalamdbf/AgentSafeEvo

AgentSafeEvo is a reinforcement learning framework that co-evolves an attacker and a defender over multi-turn tool interactions. The attacker constructs agent tasks containing user queries, tool definitions, and simulated environment observations. The defender learns to resist malicious content while completing legitimate requirements. Both policies improve through feedback on the safety and utility of the defender's execution.

The framework covers six attack categories: direct prompt injection, indirect prompt injection, tool-description injection, misinformation, harmful tool execution, and user-query jailbreak. The training data spans 50 agent scenarios and approximately 7,000 tools.

1. 🏗️ Red-team attacker initialization. Supervised fine-tuning on benign and adversarial task-construction examples teaches the attacker to generate consistent user queries and tool-environment configurations.
2. ⚔️ Joint adversarial training. For each seed, the attacker generates multiple tasks, and the defender samples multiple tool-use trajectories for each task. Both policies are updated with Group Relative Policy Optimization (GRPO).
3. ⚖️ Learnability-guided task weighting. Within-task defender-reward standard deviation provides a proxy for learnability. Tasks with greater variation receive larger relative weights within their seed group. The same task weight scales both policies' normalized GRPO advantages.

The defender reward combines safety and legitimate task completion:

$$ R_D(\tau) = R_{\mathrm{safety}}(\tau) + R_{\mathrm{utility}}(\tau), \qquad R_A^{(i)} = -\frac{1}{K}\sum_{j=1}^{K} R_D(\tau_{i,j}). $$

Each safety and utility judgment contributes either $+1$ or $-1$. Completing a harmful objective does not count toward utility; for an entirely harmful request, safe refusal satisfies the utility criterion. The attacker receives the negative mean defender reward for its generated task. Task weighting changes the contribution to policy updates while preserving the rewards and GRPO groups.

We report mean ± standard deviation over three independent runs for Qwen3-8B. Static-AT fixes the SFT attacker and trains only the defender without task weighting.

ASR is Attack Success Rate (%), UA is Utility under Attack (%), and NRP is Net Resilient Performance. Lower ASR and harm scores are better; higher UA and NRP are better.

AgentSafeEvo improves safety across all four safety benchmarks and improves planning and multi-hop tool use. BFCL-V3 accuracy remains broadly comparable to the original model.

We are currently validating AgentSafeEvo on heterogeneous attacker–defender pairs, including a Qwen3-8B attacker paired with a Gemma4-E4B defender. Results are pending release.

Download model checkpoints from the AgentSafeEvo Hugging Face collection. Use local Hugging Face checkpoint directories in the commands below.

All three models use Qwen3-8B as their backbone. The defender's reported metrics are listed above. Model weights are distributed separately from the code.

The SFT dataset is available on Hugging Face at AgentSafeEvo-Attacker-SFT-Data. The RL training seeds and validation cases are included in the repository:

SFT records contain instruction, input, output, and system. RL seeds contain the agent scenario, intent, tool definitions, and task type. The source tasks are drawn from TOUCAN; the paper uses Gemini 3 Flash Preview for adversarial task synthesis. The prepared files support attacker initialization and co-evolution without rerunning synthesis.

Run from the repository root on Linux x86_64, using Python 3.12 and NVIDIA GPUs. The setup script installs PyTorch 2.10.0, vLLM 0.19.1, Transformers 5.5.3, and FlashAttention 2.8.3. Its FlashAttention wheel targets CUDA 12 and PyTorch 2.10 with the C++11 ABI.

git clone https://github.com/lalalamdbf/AgentSafeEvo.git
cd AgentSafeEvo

conda create -n agent_safe_evo python=3.12 -y
conda activate agent_safe_evo
bash scripts/setup.sh

The launchers explicitly load the veRL source bundled in verl/.

Note. pip check will report a NumPy dependency conflict, but this warning can be ignored.

Install LLaMA-Factory for SFT in its own environment. For evaluation:

vLLM model serving and the shared Prompt Guard 2 worker also use agent_safe_evo. Follow each benchmark's linked guide for installation; MSB's MCP attack tools use an additional .venv.

Run a prompt-injection example from AgentDojo's workspace suite with the released defender. From the repository root, download the checkpoint and start an OpenAI-compatible server in agent_safe_evo:

conda activate agent_safe_evo
hf download lalalamdbf/AgentSafeEvo-Defender-Qwen3-8B \
  --local-dir checkpoints/defender

CUDA_VISIBLE_DEVICES=0 vllm serve checkpoints/defender \
  --served-model-name agentsafeevo \
  --host 127.0.0.1 --port 8000 \
  --tensor-parallel-size 1 \
  --enable-auto-tool-choice --tool-call-parser hermes --reasoning-parser qwen3

Once the server is ready, open a second terminal at the repository root. Install the bundled AgentDojo in its own Python 3.11 environment, then run one user-task/injection-task pair:

conda create -n agentdojo python=3.11 -y
conda activate agentdojo
python -m pip install -e evaluation/agentdojo

OPENAI_API_BASE_URL=http://127.0.0.1:8000/v1 OPENAI_API_KEY=EMPTY \
PYTHONPATH=evaluation/agentdojo/src \
python -m agentdojo.scripts.benchmark \
  --model QWEN3 --model-id agentsafeevo \
  --suite workspace --user-task user_task_0 --injection-task injection_task_0 \
  --attack important_instructions \
  --logdir outputs/quickstart/agentdojo

The example tests whether the defender completes the user's task while resisting injected instructions in tool outputs. AgentDojo prints utility and security results and saves JSON trajectories under outputs/quickstart/agentdojo/. Add --force-rerun to repeat an existing run. See the AgentDojo guide for more evaluation options.

Use the released AgentSafeEvo-Attacker-Qwen3-8B-SFT checkpoint to start joint training directly. From the AgentSafeEvo repository root, download the initialized attacker:

hf download lalalamdbf/AgentSafeEvo-Attacker-Qwen3-8B-SFT \
  --local-dir checkpoints/attacker_sft

Continue to Stage 2 using checkpoints/attacker_sft as MODEL_0_PATH.

export AGENTSAFEEVO_ROOT=/absolute/path/to/AgentSafeEvo
export HF_SFT_DATASET="lalalamdbf/AgentSafeEvo-Attacker-SFT-Data"
hf download "$HF_SFT_DATASET" redteam_sft.json \
  --repo-type dataset --local-dir data

{
  "agentsafeevo_attacker": {
    "file_name": "redteam_sft.json",
    "formatting": "alpaca",
    "columns": {
      "prompt": "instruction",
      "query": "input",
      "response": "output",
      "system": "system"
    }
  }
}

FORCE_TORCHRUN=1 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
llamafactory-cli train \
  --stage sft --do_train --finetuning_type full \
  --model_name_or_path Qwen/Qwen3-8B --template qwen3 \
  --dataset agentsafeevo_attacker --dataset_dir data \
  --cutoff_len 16384 \
  --per_device_train_batch_size 2 --gradient_accumulation_steps 4 \
  --learning_rate 3e-5 --num_train_epochs 2.0 \
  --lr_scheduler_type cosine --warmup_ratio 0.03 \
  --bf16 true --deepspeed examples/deepspeed/ds_z3_offload_config.json \
  --output_dir "$AGENTSAFEEVO_ROOT/checkpoints/attacker_sft"

From the AgentSafeEvo repository root, activate agent_safe_evo and download the original Qwen3-8B defender. Set the attacker checkpoint and an independent OpenAI-compatible reward-model endpoint, then launch training:

hf download Qwen/Qwen3-8B --local-dir checkpoints/qwen3_8b

export MODEL_0_PATH="$(pwd)/checkpoints/attacker_sft"
export MODEL_1_PATH="$(pwd)/checkpoints/qwen3_8b"
export SAFE_EVOLVE_SAFETY_BASE_URL=https://your-judge-endpoint/v1
export SAFE_EVOLVE_SAFETY_API_KEY=your-key
export CKPT_DIR=checkpoints/agentsafeevo_qwen3_8b

bash scripts/train/agent_safe_evolve/qwen3_joint_judge-q38_u1_default_s120_a1e-6_d3e-6_kl0_w0.5.sh

model_0 is the attacker; model_1 is the defender. The judge is configured as qwen3.8-27b and supplies independent safety and utility judgments. Set safe_evolve.components.outcome_reward_model.kwargs.model if your endpoint uses a different served name.

The provided launcher defaults to 120 steps. Checkpoints and validation run every 30 steps. The default allocation is eight GPUs, four per role; set GPUS_PER_MODEL to adjust it. W&B logging defaults to offline mode.

Additional Hydra overrides can be appended to the launcher:

bash scripts/train/agent_safe_evolve/qwen3_joint_judge-q38_u1_default_s120_a1e-6_d3e-6_kl0_w0.5.sh \
  data.train_files=data/agent_safe_evolve/train.jsonl \
  data.val_files=data/agent_safe_evolve/val.jsonl \
  safe_evolve.components.outcome_reward_model.kwargs.model=your-judge-model

The base configuration and joint configuration define the remaining options.

Use the same model paths, data, and judge settings with one of the following scripts. Set a separate CKPT_DIR for each experiment.

Use the evolved defender checkpoint for benchmark evaluation. Each guide provides its own environment setup, data preparation, model configuration, run commands, and scoring procedure.

For API-based evaluation, serve the defender through an OpenAI-compatible endpoint with structured message.tool_calls, as in the quick start. Configure the endpoint and actual served model name in each benchmark. Agent-SafetyBench trajectory scoring and AgentHarm's refusal and semantic judges default to gpt-5.6-luna. See the Agent-SafetyBench scoring guide and AgentHarm judge setup; both can share the same judge endpoint. AgentHarm requires a pinned external Inspect Evals checkout, MSB requires MCP tool setup, and ALFWorld and ToolHop require their benchmark data.

Spotlighting and Prompt Guard 2 baselines are available through the benchmark adapters. Prompt Guard 2 loads in a persistent subprocess using the shared agent_safe_evo environment. See the evaluation overview and shared Prompt Guard 2 setup.

Use the same model, prompts, sampling settings, task set, and tool configuration when comparing methods.

The direct evaluator generates and caches one fixed task per seed for each attacker, then evaluates defenders on the same tasks. It uses two GPU workers in successive attacker and defender phases, plus the independent judge endpoint configured for training.

After a 120-step training run, evaluate the initial models and saved checkpoints:

bash evaluation/ad_eval/run.sh \
  --attacker-base "$MODEL_0_PATH" \
  --defender-base "$MODEL_1_PATH" \
  --checkpoint-root "$CKPT_DIR" \
  --val-file data/agent_safe_evolve/val.jsonl \
  --steps 30 60 90 120 \
  --risk-type malicious \
  --defender-tool-parser hermes \
  --judge-model qwen3.8-27b \
  --output-dir outputs/ad_eval/agentsafeevo_qwen3_8b

Each evaluated checkpoint must contain exported Hugging Face weights and fit on one GPU. Append --dry-run to validate paths and print serving commands before launching. For a small live run, add --limit 2 and use a separate output directory.

Outputs include asr_matrix.csv, pair_metrics.csv, summary.json, cached attacks, and per-pair trajectories. The direct evaluation guide covers checkpoint discovery, resuming, report rebuilding, and evaluating external defender APIs with cached attacks through run_external.sh.

If you find AgentSafeEvo useful, please cite our paper:

@misc{wang2026agentsafeevoadversarialcoevolutionllm,
  title={AgentSafeEvo: Adversarial Co-Evolution for LLM Agent Safety},
  author={Chenxu Wang and Xiaoyu Wen and Han Qi and Ziyuan Zhou and Cheng Jin and Chaochao Lu and Qiaosheng Zhang},
  year={2026},
  eprint={2610.19422},
  archivePrefix={arXiv},
  primaryClass={cs.AI},
  url={https://arxiv.org/abs/2610.19422},
}

The training pipeline builds on PettingLLMs and veRL, with vLLM for rollout generation and LLaMA-Factory for attacker SFT. We thank the TOUCAN dataset and benchmark authors for their tools, datasets, and evaluation infrastructure.

The project code is released under the MIT License. Bundled veRL retains its Apache-2.0 license and notice. Benchmarks, datasets, and model checkpoints retain their respective licenses and attribution; see their own documentation.