PaperlayerView exact revision on arXiv

Source-linked AI summary

StepGuard: Learning Step-Level Guardrails with Scalable Supervision and Safety-Utility Balancing

Zhijie Zheng, Yu Li, Chen Qian, Yuqian Fu, Yanwei Fu, Lu Sheng, Jing Shao, Dongrui Liu

arXiv:2608.24777v1cs.AIcs.CR

TL;DR

Tool-using agents create security risks, while existing guardrails often leave step-level pre-execution monitoring underexplored. StepGuard combines pre-execution action checks with trajectory auditing, using StepGen for matched supervision and Balance-GRPO for calibration. It achieves strong guard accuracy and reduces attack success on AgentDojo and AgentDyn with a small utility loss, though open-ended and highly adversarial settings remain limitations.

  • Problem

    Existing agent guardrails often evaluate completed trajectories, leaving scalable pre-execution monitoring of step-level actions and safety–utility balance underdeveloped.

  • Method

    StepGuard checks candidate tool actions before execution and audits completed trajectories, trained with StepGen’s matched trajectories and Balance-GRPO’s accuracy-gap-based reweighting.

  • Results

    Across five static benchmarks, StepGuard achieves the best average performance among evaluated guardrail baselines, while reducing average ASR on AgentDojo and AgentDyn by 77.3% with a 2.8% utility drop.

  • Takeaways & Limitations

    The results support context-aware step-level supervision and class-wise calibration for agent guarding, with performance comparable to GPT-5.4.

  • Takeaways & Limitations

    StepGuard relies on synthetic, LLM-annotated data and benchmark-limited evaluations, is not a formal safety guarantee, and may incur false positives, false negatives, and inference cost.

Abstract

from arXiv · show

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.

1 Introduction

StepGuard addresses the difficulty of monitoring tool-using agents at the action level by combining pre-execution checks with trajectory auditing. StepGen supplies matched step-level data, while Balance-GRPO calibrates safe and unsafe decisions to improve the safety–utility trade-off.

  • The paper identifies scalable step-level supervision and safety–utility calibration as central challenges for agent guardrails.Existing guards can over-block benign actions or miss unsafe ones, sacrificing utility or safety respectively.
  • StepGuard checks candidate tool actions before execution and audits completed agent trajectories.
  • StepGen constructs prefix-aligned safe and unsafe trajectories that diverge at a designated risky step, plus benign tool-reuse trajectories.This provides scalable supervision while reducing reliance on tool identity as a safety signal.
  • Balance-GRPO reweights on-policy training toward the class with lower observed accuracy, without changing rollout prompts or raw rewards.
  • 77.3% lower mean ASR on AgentDojo and AgentDyn was achieved relative to no-guard, with mean utility decreasing by 2.8 points.StepGuard also achieved the highest average accuracy among open-weight models, comparable to GPT-5.4.

2 Related Work

Related work spans agent guardrails, synthetic safety data, and safety–utility calibration. StepGuard is positioned around step-level, context-aware supervision and dynamic class-wise balancing.

  • Agent guardrails extend ordinary LLM guardrails by using tool specifications, execution traces, and multi-step interaction histories.
  • Existing synthetic agent-safety datasets address the lack of multi-step execution context in content-level safety data, but some methods remain plan-level.
  • Balance-GRPO dynamically reweights optimization using the rollout-batch accuracy gap between safe and unsafe classes.The method is presented as a response to class-wise defense bias and does not modify prompts or raw rewards.
  • StepGen generates an unsafe trajectory, two prefix-aligned safe branches, and an independently constructed benign tool-reuse trajectory before filtering groups.

3 Task Formulation

The task formulation models an agent acting in a multi-step loop and a guard that produces structured safety diagnoses. The guard can operate either on completed trajectories or on current candidate actions before execution.

  • An LLM-based agent completes a user request through a multi-step reasoning–acting loop, proposing actions and receiving observations after execution.
  • Given an execution context X, the guard model Gθ produces a structured safety diagnosis.
  • The diagnosis includes a safe/unsafe judgment, risk-source presence, risk-source type, and optional localization of the unsafe step.Risk-source detection is distinct from unsafe judgment because a detected risk can be handled safely.
  • For post-hoc diagnosis, X is a completed trajectory; for pre-execution guarding, X is the user request, tools, history prefix, and candidate action.

4 StepGen: Prefix-Aligned Data Engine

StepGen creates prefix-aligned trajectory groups that localize the first risky action while contrasting unsafe behavior with safe alternatives and benign tool reuse. Step-level labels and structural and semantic quality checks provide scalable supervision for StepGuard.

  • Risk-Anchored Trajectory Construction: StepGen samples risk conditions and a tool subset, then constructs a plan with one designated unsafe anchor action.Plans without a unique unsafe anchor are discarded, and an environment simulator produces the unsafe base trajectory.
  • Prefix-Aligned Contrastive Branching: The generator keeps the prefix before the anchor fixed and regenerates only the remaining suffix under Refuse and Aware safe modes.Refuse declines the risky action, whereas Aware recognizes the risk and continues through a safe alternative.
  • Prefix-Aligned Contrastive Branching: Each group contains one unsafe trajectory, two prefix-aligned safe branches, and one benign trajectory reusing the same tool subset.The benign trajectory comes from a separate non-adversarial scenario, helping distinguish unsafe behavior from tool identity.
  • Step-Level Annotation and Quality Control: Every executed action receives a Safe/Unsafe label and a structured explanation, while retained groups must pass mechanical and semantic quality checks.The checks assess trajectory structure, tool-call format, parameter lineage, semantic consistency, anchor correctness, and scenario realism.

5 StepGuard

StepGuard is trained with cold-start SFT followed by Balance-GRPO and operates by evaluating candidate tool calls before execution. Balance-GRPO reweights training toward less frequent or lower-accuracy classes to improve safety–utility calibration.

  • Training and Inference: StepGuard uses two training stages: StepGen-based cold-start SFT followed by Balance-GRPO for safe–unsafe calibration.At inference, it checks each candidate tool call against the current context before execution.
  • Cold-Start SFT: Cold-start SFT trains StepGuard on 3K StepGen demonstrations to produce safety labels, risk categories, unsafe-step identification, and explanations.Inputs are either full trajectories or contexts paired with candidate actions.
  • Balance-aware Optimization: The structured reward grants risk-category credit only when the predicted safety judgment is correct.Each rollout response is parsed into a predicted safety label and risk category before reward computation.
  • Balance-aware Optimization: Balance-GRPO gives more training weight to classes that are less frequent or have lower accuracy during on-policy optimization.Class-count correction and accuracy-aware reweighting modify normalized advantages while leaving prompt sampling and raw rewards unchanged.
  • Training Objective: The training policy uses the standard clipped GRPO objective with probability ratios comparing the current and old policies.The expectation is taken over prompts and sampled responses.

6 Experiments

StepGuard is evaluated through static safety judgments, runtime guarded-agent tests, and ablations of StepGen and Balance-GRPO. It performs strongly across trajectory and step-level evaluation, generalizes to unseen risk sources, and improves runtime safety–utility trade-offs, though AgentHarm remains challenging.

  • Evaluation setup: Experiments combine static trajectory- and step-level safety evaluation with runtime guarded-agent evaluation and ablations of StepGen and Balance-GRPO.Static results use accuracy and F1, while runtime tests measure safety risk and benign-task performance.
  • Static safety evaluation: 83.0 trajectory-level accuracy and 83.3 F1, plus 84.8 step-level accuracy and 84.1 F1, show consistent performance across granularities.Despite using a 4B backbone, StepGuard achieves the highest step-level average F1 among evaluated agent guards.
  • Runtime evaluation: AgentHarm remains difficult: malicious score falls from 22.8 to 3.4, while task completion falls from 70.9 to 52.8.No evaluated method achieves a clearly favorable trade-off on this benchmark.
  • StepGen ablation: StepGen raises average Acc/F1 from 80.4/82.0 to 83.8/83.4 with intermediate-prefix labels, while benign tool reuse raises TS-Bench-Harm F1 from 69.2 to 75.2.These components support identifying where risk emerges without relying on tool identity.
  • Generalization: 74.9/78.1 Acc./F1 on six unseen ATBench risk sources substantially exceeds the backbone’s 49.6/35.6 and approaches the full-coverage model’s 76.8/80.7.The model is trained on two risk sources while keeping training size and optimization settings fixed.
  • Balance-GRPO ablation: Balance-GRPO reduces the safe–unsafe accuracy gap from 13.0 to 8.0 while improving Acc/F1 from 81.5/81.9 to 82.2/82.1.Fixed weighting reaches a similar gap but lowers unsafe accuracy to 77.5, compared with 86.4 for Balance-GRPO.

7 Analysis

The analysis links defense bias to prediction instability and examines Balance-GRPO’s training dynamics. Less accurate action classes are less stable, while dynamic reweighting closes the accuracy gap faster and maintains higher macro-F1.

  • Prediction stability: Models are less stable on the action type they classify less accurately, with over-defense increasing safe-example flips and under-defense increasing unsafe-example flips.Defense bias therefore appears in both class-wise accuracy and prediction stability.
  • Training dynamics: Balance-GRPO reduces the safe–unsafe accuracy gap faster than standard GRPO while maintaining higher macro-F1 during training.Its dynamic weighting gives more emphasis to whichever action type currently has lower accuracy.

8 Conclusion

StepGuard combines online pre-execution checking with offline trajectory diagnosis, trained using StepGen and Balance-GRPO. Across static and guarded-agent evaluations, it performs strongly while preserving a favorable safety–utility trade-off in tool-use settings, although AgentHarm remains challenging.

  • Conclusion: StepGuard supports both online pre-execution action checking and offline trajectory diagnosis.This combines step-level guarding before tool execution with analysis of completed agent behavior.
  • Conclusion: 77.3% lower average ASR on AgentDojo and AgentDyn is obtained relative to no defense with a 2.8% utility drop.Across five static benchmarks, StepGuard has the best average performance among evaluated guardrail baselines and accuracy comparable to GPT-5.4.
  • Conclusion: AgentHarm exposes a less favorable safety–utility trade-off than the tool-use environments.The conclusion identifies highly adversarial harmful-agent settings as the remaining challenge.

9 Limitations

The paper identifies limitations in StepGen’s synthetic supervision, benchmark coverage, and StepGuard’s lack of formal guarantees. Deployment can still involve classification errors, inference cost, and policy-management requirements.

  • StepGen data may inherit coverage limits, biases, or annotation errors from synthetic generation, the teacher model, and its risk taxonomy.
  • Current evaluations do not cover fully open-ended tool ecosystems, longer-horizon workflows, multi-agent interactions, or adaptive adversaries.
  • StepGuard is a pre-execution guardrail rather than a formal safety guarantee, so false positives and false negatives may still occur.
  • Deployment adds inference cost and requires policies for revising blocked actions.

10 Ethics Statement

The ethics statement restricts access to sensitive synthetic safety data and recommends safeguards around deployment. The accompanying evaluation reports runtime overhead and human-verified generalization, while clarifying how repeated results should be interpreted.

  • 10 Ethics Statement: Sensitive examples and executable unsafe workflows are restricted to authorized researchers under ethical-use and data-handling requirements.
  • 10 Ethics Statement: StepGuard should be deployed with access controls, monitoring, and human oversight because it provides no formal safety guarantee.
  • Evaluation scope: Guarded-agent comparisons use the same agent backbone, while repeated StepGuard results reflect evaluation variability from one fixed checkpoint rather than independent training seeds.
  • Runtime overhead: 599.9 ms per guard call and 2.53 seconds per task represent StepGuard’s runtime overhead under AgentDojo’s all-calls-inspected protocol.
  • Runtime overhead: Guard inference accounts for only 7.24% of total AgentDojo task time under the reported serving and evaluation settings.

A.3 Human-Verified Held-Out Evaluation

The held-out evaluation tests transfer beyond StepGen using 100 human-verified execution traces. Balance-GRPO improves both overall and class-specific safety metrics, while additional analyses examine rollout balance, harmful-intent coverage, overlap, and training-set composition.

  • Human-verified held-out evaluation: 100 held-out traces comprise 25 benign, 25 attack-exposed without unsafe execution, and 50 involving unsafe execution.
  • Human-verified held-out evaluation: Balance-GRPO raises Acc/F1 from 75.0/76.6 to 82.0/83.0 over vanilla GRPO on the human-verified traces.
  • Human-verified held-out evaluation: Safe accuracy increases from 68.0 to 76.0 and Unsafe accuracy from 82.0 to 88.0, without sacrificing one class to reduce the class-wise gap.
  • Balance-GRPO robustness: Class-count weighting reduces rollout-ratio sensitivity, shrinking the F1 difference from 1.92 points under vanilla GRPO to 0.74 points.
  • Harmful-intent supervision: Adding 1K harmful-intent examples improves TS-Bench-Harm F1 by 7.8 points while decreasing TS-Bench-Dojo F1 by 0.2 points.
  • Overlap and data analysis: Tool-overlap filtering changes macro-averaged Acc/F1 from 79.2/80.9 to 80.9/81.1 without systematic degradation, weakening overlap as the primary explanation.
  • Training-set composition: The final training set contains 7K examples, balanced across step-level versus trajectory-level records and safe versus unsafe examples.
  • Overlap and data analysis: Instruction-level nearest-neighbor similarity remains limited: p99 does not exceed 0.885, and at least 99.7% of test examples fall below 0.90.

H Error Analysis

Error analysis finds distinct over-defense and under-defense patterns across 575 static-evaluation errors. StepGuard often overweights suspicious context or sensitive tools, while missing risks requiring intent, authorization, semantics, and multi-step reasoning.

  • The analysis covers 575 errors: 307 false positives and 268 false negatives.
  • False positives: Benign uses of sensitive tools account for 43.6% of false positives, while safely handled suspicious inputs account for another 25.7%.
  • False positives: False positives indicate excessive reliance on sensitive tool identities or suspicious context without fully accounting for authorization and executed actions.
  • False negatives: Missed malicious intent is the most frequent false-negative mode at 19.8%; tool, authorization or privacy, and multi-step propagation errors each account for 17.5%.
  • False negatives: Context-dependent prompt injection contributes 16.4% of false negatives, which require connecting intent, authorization, tool effects, and downstream consequences.
  • Improvement directions: The analysis recommends authorization-grounded safe examples and stronger tool-semantic, authorization-tracking, and multi-step causal reasoning.
  • Representative cases: Figure examples show correct separation of malicious intent from safe actions and localization of an indirect-injection violation to the consequential network-change step.
  • Representative cases: A representative academic-integrity case is a false negative because identity misuse lacks direct prompt-injection, leakage, or system-damage cues.
Loading 2608.24777v1…