StepGuard is a 4B Guard model that inspects each agent tool call before execution and rebalances training with Balance-GRPO to close the gap between blocking too much and blocking too little, cutting mean Automatic Speech Recognition by 77.3% against no-guard with only a 2.8-point utility drop.
You’ve shipped an LLM agent that calls tools: sending emails, moving money, editing repos. A single wrong tool call can leak data or execute an injected instruction hidden in a webpage the agent just read. Today’s defenses tend to fail in one of two ways. Content-style guards like LlamaGuard3-8B score isolated prompts and miss risks that only appear inside a multi-step trajectory. Agent-specific guards like TS-Guard work at the step level but often over-block benign tool use, tanking task completion. The paper’s target is the class-wise imbalance itself: a guard that treats safe and unsafe examples asymmetrically during training will keep that asymmetry at inference.
Two pieces: a data engine and a training tweak.
The data engine, StepGen, addresses the fact that real unsafe agent executions are rare. It picks a risk type and a tool subset, then has a planner LLM write a scenario where that risk appears at exactly one step, the anchor. It rolls out an unsafe trajectory, then re-rolls two safe alternatives from the same prefix: one that refuses at the anchor, one that recognizes the risk and continues via a safe path. It also generates a separate benign trajectory that reuses the same tools in a non-adversarial scenario, so the guard doesn’t learn that certain tool names are just always dangerous. A rule validator plus an LLM auditor filters groups; only ~31% survive.
Training has two stages: standard supervised fine-tuning on 3K examples from Qwen3-4B-Instruct, then reinforcement learning with Balance-GRPO. The base Group Relative Policy Optimization (GRPO) algorithm normalizes rewards within a group of sampled responses per prompt. Balance-GRPO adds two multipliers to each normalized advantage: one for class-count imbalance, one that grows for whichever class (safe or unsafe) currently has lower rollout accuracy. In plain terms: when the model is over-blocking benign calls, safe-class examples get bigger gradient updates. When it under-blocks, unsafe examples do.
for batch in rollouts:
acc_safe, acc_unsafe = measure_accuracy(batch)
gap = acc_safe - acc_unsafe # signed
for example in batch:
A = normalize_reward(example)
c = class_count_weight(example.label, batch)
if example.label == "safe":
w = 1 + lam * max(-gap, 0)
else:
w = 1 + lam * max(gap, 0)
update_policy(example, advantage = c * w * A)
The prevailing move when a safety-tuned model over-refuses is to reweight the loss with fixed class weights or upsample the weaker class. This paper argues the imbalance is a moving target during training: which class is losing changes rollout to rollout. Measure the safe-versus-unsafe accuracy gap live in each RL batch and let it drive the advantage weighting, so the harder class gets more pressure exactly when it needs it. The evidence that matters is the ablation against fixed-weight and upsampling baselines, not the headline benchmark numbers.
The load-bearing result is the ablation isolating Balance-GRPO from simpler balancing. Starting from an over-defensive SFT checkpoint (safe accuracy 69.3, unsafe 91.1, gap 21.8), vanilla Group Relative Policy Optimization (GRPO) closes the gap to 13.0. Balance-GRPO closes it to 8.0 while keeping unsafe accuracy at 86.4. Fixed 1.5/0.5 weights reach a similar gap (7.9) but drop unsafe accuracy to 77.5, meaning they trade safety for balance. Balance-GRPO doesn’t.
•
Static safety judgment across five benchmarks (ATBench, R-Judge, ASSE-Security, TS-Bench-Dojo, TS-Bench-Harm): StepGuard averages 83.0/84.8 accuracy at trajectory/step level, best among open-weight guards at 4B, comparable to GPT-5.4.
•
Deployed as a runtime guard on AgentDojo (Qwen3.6-35B agent), Automatic Speech Recognition drops from 25.1 to 1.2 with utility 90.7 vs 89.7. On AgentDyn, ASR drops from 21.2 to 9.3 with utility 66.7 vs 73.3.
•
Held-out generalization: trained on only 2 of 8 risk sources, StepGuard hits 74.9/78.1 on the 6 unseen sources, close to the full-coverage model at 76.8/80.7.
•
AgentHarm is the honest weak spot: malicious score drops 22.8 → 3.4, but benign completion also drops 70.9 → 52.8. The paper concedes no guard achieves a clean trade-off on this benchmark.
Reach for this when you’re shipping a tool-using agent and want a pre-execution circuit breaker that isn’t a full model retrain. StepGuard runs as a separate 4B model inspecting each candidate tool call in context. On AgentDojo it adds ~600ms per call and ~2.5s per task, about 7% overhead, and returns a safe/unsafe label plus a risk category. Unsafe calls get blocked and replaced with a self-reflection message to the agent, which can then try a different action.
Code and weights are released at GitHub and HuggingFace. The Balance-GRPO idea (dynamic class reweighting based on live rollout accuracy) is portable to other RL settings with class imbalance, not just safety guards. The synthetic data recipe (prefix-aligned safe/unsafe branches around a designated anchor) is reusable if you need contrastive step-level supervision in your own domain.
When a classifier is losing to one class, reweight during training based on how badly it’s losing right now, not on a fixed schedule set upfront. Balance-GRPO’s contribution isn’t the safety domain; it’s the observation that class imbalance in an RL objective is a signal you can measure and respond to in the same loop.
•
The training data comes from LLM-generated synthetic trajectories audited by another LLM. Coverage biases and label errors from the teacher chain propagate, and the paper doesn’t audit StepGen labels against human judgment directly, only downstream performance on held-out traces.
•
The strong runtime results are on prompt-injection-style benchmarks (AgentDojo, AgentDyn). On explicitly harmful-agent tasks (AgentHarm) the safety-utility trade-off is materially worse: ~18 points of benign completion lost to reach a low malicious score. Different failure regime, weaker fix.
•
Error analysis shows two boundary failures: StepGuard over-blocks legitimate use of sensitive tools (44% of false positives) and under-detects institutional violations like submitting work under someone else’s identity when no technical harm is involved. A guard that reasons about tool effects is not a guard that reasons about policy.