StepGuard: Learning Step-Level Guardrails with Scal...

StepGuard: Learning Step-Level Guardrails with Scal... | AI Research

Key Takeaways

StepGuard is a guard model designed to improve the safety of LLM-based agents by auditing their actions at the step level.
LLM-based agents can interact with external environments through tool invocation, but this capability also introduces security risks such as file modification, information leakage, and unauthorized actions.
Existing guardrails often evaluate completed trajectories, leaving pre-execution monitoring of step-level actions underexplored.
We propose StepGuard, a step-level guard model that can audit completed agent trajectories and check tool actions before they are executed.
To train StepGuard, we introduce StepGen, an automatic data engine that generates safe and unsafe trajectories with the same context but different actions at the risky step.

Paper AbstractExpand

LLM-based agents can interact with external environments through tool invocation, but this capability also introduces security risks such as file modification, information leakage, and unauthorized actions. Existing guardrails often evaluate completed trajectories, leaving pre-execution monitoring of step-level actions underexplored. We propose StepGuard, a step-level guard model that can audit completed agent trajectories and check tool actions before they are executed. To train StepGuard, we introduce StepGen, an automatic data engine that generates safe and unsafe trajectories with the same context but different actions at the risky step. To further reduce over-defense and under-defense, we propose Balance-GRPO, which dynamically balances learning between safe and unsafe actions based on their observed accuracy. Experiments show that StepGuard achieves the highest average accuracy among open-weight guard models, with performance comparable to GPT-5.4. When used to guard agents on AgentDojo and AgentDyn, StepGuard reduces mean attack success rate by 77.3% relative to the no-guard setting, while mean utility drops by only 2.8 percentage points.

StepGuard is a guard model designed to improve the safety of LLM-based agents by auditing their actions at the step level. It monitors agent behavior before tools are executed and evaluates completed trajectories to prevent risks such as unauthorized file modifications, information leakage, and harmful tool usage.

Addressing the Safety-Utility Gap

A primary challenge in deploying agent guardrails is the "defense bias," where models either block too many benign actions (over-defense) or fail to catch enough unsafe ones (under-defense). Existing methods often struggle to balance these outcomes. StepGuard addresses this by using a training approach that dynamically adjusts to the accuracy gap between safe and unsafe actions, ensuring that the guard remains effective without unnecessarily hindering the agent's ability to complete tasks.

Data Generation with StepGen

To overcome the scarcity of high-quality, step-level safety data, the researchers developed StepGen, an automatic data engine. StepGen creates "prefix-aligned" trajectories, meaning it generates both safe and unsafe versions of a task that share the same starting context. By branching from a common point, the model learns to distinguish between safe and unsafe actions within the same scenario. The engine also generates benign trajectories that use similar tools to ensure the guard does not mistakenly flag a tool as unsafe simply because it is being used.

Training with Balance-GRPO

StepGuard is trained in two stages. First, it undergoes supervised fine-tuning on data generated by StepGen. Second, it uses a method called Balance-GRPO, which extends standard reinforcement learning techniques to reduce defense bias. During training, Balance-GRPO monitors the accuracy of the model on both safe and unsafe examples. If the model performs worse on one class, the training process automatically increases the weight of those examples, forcing the model to improve its performance on the weaker category.

Performance and Results

Experiments show that StepGuard achieves the highest average accuracy among open-weight guard models, with performance comparable to the closed-source GPT-5.4. When deployed in agent environments like AgentDojo and AgentDyn, StepGuard reduced the mean attack success rate by 77.3% compared to using no guard at all. Notably, this safety improvement came with a minimal impact on task utility, which dropped by only 2.8 percentage points. Ablation studies confirmed that the Balance-GRPO method was effective in reducing the accuracy gap between safe and unsafe classifications while maintaining high performance.