StepGuard: Learning Step-Level Guardrails with Scalable Supervision and Safety-Utility Balancing
Abstract Overview
StepGuard is a 4B step-level guard model for LLM agents that supports both pre-execution checking of candidate tool actions and post-hoc auditing of completed trajectories. The paper introduces StepGen, a synthetic data engine that creates prefix-aligned safe and unsafe trajectory branches around a designated risky step, along with benign tool-reuse examples for supervision. It also introduces Balance-GRPO, a training method that dynamically reweights optimization toward whichever class, safe or unsafe, currently has lower accuracy to reduce defense bias. Across static benchmarks, StepGuard shows strong performance at both trajectory and step granularity, and in guarded-agent experiments it substantially lowers attack success while preserving most utility, though highly adversarial AgentHarm settings remain challenging.
Novelty
The paper's main novelty is framing agent guardrails at the step level with a single model that can both block unsafe tool actions before execution and audit full trajectories afterward. It combines a prefix-aligned synthetic supervision pipeline, StepGen, with Balance-GRPO, which explicitly calibrates the safety-utility trade-off by dynamically reweighting advantages based on class-wise accuracy gaps during training.
Results
In static evaluation, StepGuard achieves the highest average accuracy among open-weight guard models, delivering 83.0 accuracy / 83.3 F1 on trajectory-level benchmarks and 84.8 accuracy / 84.1 F1 on step-level benchmarks, comparable to GPT-5.4. When deployed to guard agents on AgentDojo and AgentDyn, it reduces mean attack success rate by 77.3% relative to no defense while utility drops by only 2.8 points on average. Ablations show that Balance-GRPO reduces the safe-unsafe accuracy gap from 13.0 to 8.0 and improves guarded-agent utility by up to 6.7 points versus vanilla GRPO with only a 0.3-point ASR increase.
Key Points
- StepGuard targets action-level agent safety by evaluating candidate tool calls in context before execution and by diagnosing completed multi-step trajectories.
- StepGen provides scalable supervision through prefix-aligned safe and unsafe branches plus benign tool-reuse trajectories, helping the model separate true risk from superficial tool identity cues.
- Balance-GRPO is designed to reduce over-defense and under-defense by dynamically emphasizing the class with lower observed accuracy, improving runtime safety-utility balance.