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.
Figures & tables
Figure 1: Overview of StepGen. The pipeline first constructs and rolls out an unsafe trajectory with a designated risk anchor. It then generates two prefix-aligned safe branches and an independently constructed benign tool-reuse trajectory. Finally, StepGen adds step-level supervision and retains high-quality groups through structural and semantic filtering.
Figure 2: Overview of StepGuard. StepGen constructs prefix-aligned trajectories with step-level supervision for cold-start SFT. Balance-GRPO then uses class-wise rollout feedback to reweight advantages, reducing the safe–unsafe accuracy gap during on-policy training.
Figure 3: Safe–unsafe accuracy trade-off. Points above and below the diagonal indicate over-defense and under-defense, respectively.
Model
Size
Trajectory-Level Evaluation
Step-Level Evaluation
ATBench
R-Judge
ASSE Security
Avg.
TS-Bench-Dojo
TS-Bench-Harm
Avg.
Acc.
F1
Acc.
F1
Acc.
F1
Acc.
F1
Acc.
F1
Acc.
F1
Acc.
F1
Closed-Source Reference
GPT-5.4
–
65.9
69.8
91.8
92.4
91.3
90.9
83.0
84.4
93.9
90.3
68.8
76.4
81.3
83.3
Open-Weight General Models
Qwen3.5-397B-A17B
397B
63.5
64.0
90.2
90.9
92.6
92.5
82.1
82.5
94.6
91.3
71.3
77.5
82.9
84.4
Table 1: Static safety evaluation on trajectory- and step-level benchmarks. We report accuracy and F1, averaged separately within each granularity. Bold and underlined values indicate the best and second-best open-weight models, respectively. The gray closed-source reference is excluded from ranking.
Figure 4: Runtime safety–utility trade-off. Runtime utility versus risk (ASR or malicious score) on AgentDojo, AgentDyn, and AgentHarm. Upper left is better; numerical results are reported in Table 6 .
Figure 5: Training dynamics of GRPO and Balance-GRPO. (a) Safe-minus-unsafe accuracy gap, where values closer to zero indicate better balance. (b) Macro-F1 over rollout steps.
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
AgentDojo
AgentDyn
AgentHarm
Method
ASR ↓
Utility ↑
ASR ↓
Utility ↑
Malicious ↓
Completion ↑
No Defense
25.1
89.7
21.2
73.3
22.8
70.9
MetaSecAlign-70B
1.9
84.5
9.0
55.0
–
–
TS-Guard
6.2
92.8
13.8
70.0
5.4
58.4
ShieldAgent-THU
18.1
85.3
20.0
70.0
1.3
18.7
Qwen3-4B
12.3
93.0
15.7
70.0
14.0
67.3
Appendix
Table 6: Complete guarded-agent results. Lower ASR and Malicious Score are better, while higher Utility and Completion are better. Bold denotes the best guard result in each column; No Defense is excluded from ranking.
Table 8Table 9
Full SFT Data (3,000)
Filtered SFT Data (2,814)
Benchmark
Acc.
F1
Acc.
F1
TS-Bench-Dojo
91.3
86.8
94.3
91.0
TS-Bench-Harm
73.8
78.2
71.4
74.6
ATBench
64.1
69.9
67.2
68.3
R-Judge
87.4
88.5
90.6
90.6
Macro Avg.
79.2
80.9
80.9
81.1
Appendix
Table 11: Static evaluation of independently trained SFT-only checkpoints with and without tool-overlap filtering. The two checkpoints use identical training settings and are distinct from the final SFT+Balance-GRPO model reported in the main evaluation.
Figure 7: Statistics of the final StepGuard training set. The dataset contains 3K examples for cold-start SFT and 4K additional examples for Balance-GRPO. Left: distribution over eight risk-source categories and one benign category. Right: distribution of action horizons.
Train Set
Test Set
Test n
NNmax
p99
MMS
%(<0.90)
%(<0.80)
Train MMS
Instruction Only
SFT-3K
TS-Bench-Dojo
1220
0.8114
0.8114
0.6460
100.0
96.5
0.7768
RL-4K
TS-Bench-Dojo
1220
0.8475
0.8474
0.6465
100.0
97.2
0.8341
SFT-3K
TS-Bench-Harm
416
0.8921
0.8814
0.6880
100.0
83.7
0.7768
RL-4K
TS-Bench-Harm
416
0.8839
0.8716
0.6994
100.0
75.0
0.8341
SFT-3K
ATBench
1000
0.9199
0.8850
0.7353
99.7
77.3
0.7768
Appendix
Table 12: Nearest-neighbor similarity between the StepGuard training splits and evaluation benchmarks under three input constructions. We report the maximum similarity, p99, MMS, and the percentages of test examples below the 0.90 and 0.80 thresholds.
Table 13
Figure 8: Five-field rationale for an unsafe transfer action triggered by an injected instruction. Each field captures a distinct aspect of the step-level safety decision.
Benchmark
Prompt
Prediction Target
TS-Bench-Dojo
Step
Candidate action
TS-Bench-Harm
Step
Candidate action
AT-Bench
Traj
Complete trajectory
R-Judge
Traj
Complete trajectory
ASSE-Security
Traj
Complete trajectory
Appendix
Table 15: Prompt assignment for static evaluation.
Figure 9: Guarded-agent workflow with self-reflect feedback. Unsafe tool calls are blocked before execution, without exposing the guard’s internal rationale.
False Positives
False Negatives
Failure Mode
Count
%
Failure Mode
Count
%
Benign use of sensitive tools
134
43.6
Missed malicious intent
53
19.8
Suspicious context handled safely
79
25.7
Misunderstood tool semantics or malicious tool use
47
17.5
Workflow-quality issue mistaken for a safety risk
64
20.8
Authorization or privacy violation
47
17.5
Authorized external action
10
3.3
Multi-step risk propagation
47
17.5
Broad but authorized scope
5
1.6
Context-dependent prompt injection
44
16.4
Appendix
Table 16: Failure-mode distribution over all 575 static-evaluation errors. Percentages are computed separately within the 307 false positives and 268 false negatives.
Figure 10: Correctly allowing safe behavior in the presence of malicious intent. The agent rejects the unsupported override and publishes only verified market information.
Figure 11: Correct detection of indirect prompt injection. The agent treats an instruction embedded in untrusted tool output as authoritative and proposes an environment-changing network policy.
Figure 12: False negative on an academic-integrity violation. StepGuard recognizes no direct system or security harm and therefore fails to identify the unauthorized assignment submission.