StepSafe: Learning Timely Intervention with Context-Aware Safety Reasoning

Siyi Gu; Kunhan Gao; Chao Peter Yang; Sukanta Ganguly; Rex Ying; Arman Cohan

cs.AIOct 14, 2026

StepSafe: Learning Timely Intervention with Context-Aware Safety Reasoning

Authors: Siyi Gu, Kunhan Gao, Chao Peter Yang, Sukanta Ganguly, Rex Ying, Arman Cohan

Organizations: Yale University · NetApp Inc.

Abstract

Safety in agentic systems is inherently sequential: whether an action is safe depends on prior actions, tool outputs, information disclosures, and interactions among agents. Recent work has shifted agent safety guardrails from post-hoc trajectory-level assessment toward step-level action diagnosis and proactive risk prediction, yet existing training objectives provide limited supervision for how step evidence accumulates over time and when it becomes sufficient to justify intervention. We propose \methodlong~(ASPO): a training framework for step-level timely intervention. Pre-fix level stage alignment allows crediting each decision at its own step: it compares each response with the other responses sampled at the same prefix and rewards the action that fits the stage, continuing before the risk is visible, asking while clarification can still avert harm, and stopping before harm is irreversible. Adaptive feedback self-distillation further credits the reasoning behind each decision, with the model itself as a teacher that sees the step annotations on intervention steps and rules mined from its own false alarms on safe steps. Its training data carry a violation taxonomy and an intervention window at every step and are extended through self-evolving multi-agent red-teaming. Finally we present StepSafe, a step-level reasoning guardrail with timely intervention, that achieves the state-of-the-art peformance among open-source guardrails across five agent-safety benchmarks on safety-utility pareto curve with Youden's $J$ of 53.0 (4B) against 46.2.

Figures & tables

Figure 1: An example of realized harm to agentic system without timely intervention.

Figure 2: Overview of StepSafe . (a) Construct training signal via step-level annotation on four factored axes (12 failure, 10 fault, 10 harm, and 8 source categories) and an intervention window (signal $s$ , boundary $b$ , last preventable $h$ ) and red-teaming via a mixture of personas fills coverage gaps. (b) SAPO credits each sampled response against the others at the same visible prefix and with a class-balanced weights update. (c) AFSD enables token level supervision on unsafe trajectories via the constructed taxonomy, safe trajectories with adaptive feedback mined from false-alarms, using self distillation with an gated teacher-student gap.

Closed-source models
	ASSE-Safety			ASSE-Security			R-Judge			AgentHarm			ATBench			Overall
Guardrail	Acc	mF1 $\uparrow$	$J\uparrow$	Acc	mF1 $\uparrow$	$J\uparrow$	Acc	mF1 $\uparrow$	$J\uparrow$	Acc	mF1 $\uparrow$	$J\uparrow$	Acc	mF1 $\uparrow$	$J\uparrow$	Acc	mF1 $\uparrow$	$J\uparrow$
Claude Sonnet 5	77.5	77.5	57.0	83.8	83.8	67.7	94.0	94.0	88.2	97.2	94.6	85.5	77.4	77.2	54.7	81.9	81.9	64.0
GPT-5.6 Sol	77.0	77.0	55.0	92.2	92.1	84.3	93.8	93.8	86.7	96.6	93.6	84.8	77.1	76.9	53.9	83.3	83.3	66.6
GPT-5.6 Terra	76.2	76.1	52.8	85.9	85.9	71.8	80.3	79.6	59.1	96.6	93.8	87.6	73.0	72.9	46.1	78.8	78.8	57.6
GPT-5.6 Luna	76.4	76.4	53.5	83.0	83.0	66.0	87.1	86.9	73.3	96.6	93.6	84.8	67.5	66.7	35.2	77.9	77.9	55.7

Table 1: Benchmark Performance.

Figure 3: Safety–utility Pareto Curve. The light diagonal lines are iso-

J

guides.

	Overall					Same-task pairs
Model	Acc $\uparrow$	mF1 $\uparrow$	ASR $\downarrow$	Util $\uparrow$	$J$ $\uparrow$	Correct $\uparrow$	Flags both $\downarrow$	Passes both $\downarrow$	Inverted $\downarrow$
Qwen3-4B	64.2	64.2	33.3	64.1	28.5	36.4	55.0	7.0	1.5
+ SFT	73.0	72.9	35.4	81.8	46.2	70.9	15.3	12.5	1.2
+ SAPO	75.9	75.9	27.2	79.1	51.9	68.2	23.2	8.0	0.6
+ AFSD ( StepSafe )	76.5	76.5	27.6	80.0	53.0	75.5	16.5	7.3	0.6

Table 3: A1: Ablation on ASPO.

Table 7

Data	Acc $\uparrow$	mF1 $\uparrow$	ASR $\downarrow$	Util $\uparrow$	$J\uparrow$
E1: full ( StepSafe -4B)	76.5	76.5	27.6	80.0	53.0
E2: w/o red-team	75.8	75.8	28.3	78.4	51.7

Table 6: A4: Ablation on Red teaming

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Axis	Category	Definition
Failure (12)	authorization	the action exceeds the actor, resource, or tool scope
	consent	an applicable required agreement is missing
	information handling	restricted data is collected, disclosed, or used
	instruction trust	untrusted instructions are promoted to authority
	evidence reasoning	an unsupported factual inference drives a consequential decision
	representation	misleading claims about identity, facts, or status

Appendix

Table 7: Categories of the factored violation taxonomy and the definitions given to the annotator.

Unsafe trajectories, each with a benign twin	512
failure-driven (monitor errors)	328
coverage-driven (taxonomy gaps)	184
Domain	Pairs
Enterprise operations	74
Public services	68
Finance and commerce	56

Appendix

Table 8: The red-teamed trajectory pairs by origin and domain.

Figure 6: Taxonomy coverage of the unsafe training trajectories without and with the red-teamed trajectories. Each arrow runs from a category’s share without red-teaming to its share with it. Shaded rows grow by at least

1.5\times

, with the factor marked; “new” marks a category absent without red-teaming. Categories are sorted by their share without red-teaming.

Setting	Supervised warm start	Reinforcement learning
Initialization	Qwen3-4B, thinking enabled	the warm-start SFT (2 epochs)
Data	2,402 step-level examples from 1,291 trajectories	3,184 trajectories with 22,000 annotated steps
Format	one visible prefix per example, single turn	replay of every agent action with the full visible history
Parameters updated	all (no adapters)	all (no adapters)
Optimizer	AdamW, $\beta=(0.9,0.95)$ , weight decay 0.01	AdamW, weight decay 0.01
Learning rate	$10^{-5}$ , cosine, 5% warmup	$10^{-6}$ , constant after linear warmup over 10% of a 2,985-step schedule

Appendix

Table 9: Training settings of the two stages.

Arm	Acc $\uparrow$	mF1 $\uparrow$	ASR $\downarrow$	Util $\uparrow$	$J$ $\uparrow$
StepSafe -4B	$76.4\pm 0.5$	$76.4\pm 0.5$	$28.0\pm 0.6$	$79.7\pm 0.8$	$52.9\pm 0.8$

Appendix

Table 10: Mean

\pm

standard deviation over three seeds.

Untrained backbone
Backbone	Acc $\uparrow$	mF1 $\uparrow$	ASR $\downarrow$	Util $\uparrow$	$J$ $\uparrow$
Qwen3-4B	64.2	64.2	33.3	64.1	28.5
Qwen3-8B	68.0	67.9	22.7	63.6	35.9
StepSafe
StepSafe -4B	76.5	76.5	27.6	80.0	53.0
StepSafe -8B	77.2	77.2	21.7	74.8	54.4

Appendix

Table 11: Model scale. Untrained backbones use our prompt with thinking on. StepSafe rows share data and hyperparameters.

Monitor	Schedule	Acc $\uparrow$	mF1 $\uparrow$	ASR $\downarrow$	Util $\uparrow$	$J$ $\uparrow$
StepGuard ( Zheng et al., 2026 )	—	73.0	73.0	32.7	78.5	46.2
StepSafe	default	76.5	76.5	27.6	80.0	53.0
StepSafe	utility-leaning: Continue $1\to 1.25$	74.1	74.0	32.4	84.3	48.3

Appendix

Table 12: Operating points of StepSafe under the intervention schedule reward, against StepGuard.

Training	Traj.	Steps	Acc $\uparrow$	mF1 $\uparrow$	ASR $\downarrow$	Util $\uparrow$	$J$ $\uparrow$
SFT, key steps balanced	1.3k	2.4k	73.0	72.9	35.4	81.8	46.2
SFT, all steps, multi-turn	3.2k	22.0k	72.3	72.3	31.5	70.7	44.6
SFT, all steps, step-level	3.2k	22.0k	62.0	59.8	67.1	86.3	24.5
SFT, all trajectories, class-balanced	3.2k	6.5k	70.0	69.8	38.7	79.2	40.1
StepSafe -4B	3.2k	22.0k	76.5	76.5	27.6	80.0	53.0

Appendix

Table 13: SFT on all annotated steps versus RL, on Qwen3-4B. Traj. and Steps count the supervised trajectories and steps; for StepSafe , those of the RL stage.

	Unanimous wrong $\to$ corrected			Unanimous correct	Mixed
Model	All $\uparrow$	Misses $\uparrow$	False alarms $\uparrow$	$\to$ kept $\uparrow$	$\to$ correct $\uparrow$
	(204)	(92)	(112)	(2,254)	(1,653)
+ SAPO	15.1	18.6	12.3	98.7	80.6
+ AFSD ( StepSafe )	16.9	19.7	14.5	99.2	81.2

Appendix

Table 14: Decisions on training prefixes, grouped by the eight sampled verdicts of the SFT initialization. Each cell is the share of eight sampled decisions of the trained monitor that are correct; the number of prefixes is in parentheses. Unanimous errors split into missed violations (

d^{\star}_{t}=1

) and false alarms (

d^{\star}_{t}=0

Arm	Reward	$J$ $\uparrow$	ASR $\downarrow$	Util $\uparrow$	BIR $\uparrow$	Stop share $\uparrow$
	Binary	41.0	27.8	70.3	26.5	62.9
C4	Stage schedule	41.2	26.4	70.1	30.7	71.7

Appendix

Table 15: SAPO with the binary reward and with the stage schedule, both trained from Qwen3-4B without SFT. BIR is the share of unsafe trajectories whose first Stop falls exactly at the annotated boundary, and Stop share the share with any Stop .

Figure 7: Validation ASR (left, lower is better) and utility (middle) over the first 250 RL steps. Markers are individual evaluations, lines a centered five-point moving average, and bands

\pm 1

standard deviation of that window; the shaded region is the window summarized on the right. Right:

J

after step 150 as mean

\pm

standard deviation, with the mean absolute change between consecutive evaluations (

|\Delta J|

Estimator	Comparison group	Credit signal	Credit stays with the decision	Credit follows the stage	Stratified class budget
Trajectory-level GRPO ( Shao et al., 2024 )	replays of a task	terminal return	–	–	–
RLOO ( Ahmadian et al., 2024 )	replays of a task	terminal return, leave-one-out	–	–	–
Step-level GRPO	all decisions of a task	per-step reward	–	–	–
GiGPO ( Feng et al., 2026 )	anchor state	return-to-go	–	–	–
SAPO , binary reward	prefix	$R_{\text{binary}}$ , leave-one-out	✓	–	✓
SAPO	prefix	$\sigma$ , leave-one-out	✓	✓	✓

Appendix

Table 16: Credit estimators for a step-level monitor; rows correspond to Table 5 .

	SAPO (decision credit)	AFSD (token credit)
Unit of credit	response	token
Weight	signed, uniform within a response	varies by token
Reference	other responses at the same prefix	the model with privileged context $\rho_{t}$
Target from	stage schedule $\sigma$ (landmarks)	annotated verdict, description, rubric
Silent when	scores tie at the prefix	teacher and student agree

Appendix

Table 17: Roles of the two terms of Eq. 16 .

AI Agent Safety

AI agent safety studies how to prevent and detect harmful actions by autonomous systems that use tools, memory, and multi-step plans.

See all 326 papers

StepGuard: Learning Step-Level Guardrails with Scalable Supervision and Safety-Utility Balancing
Aug 25, 2026Zhijie Zheng, Yu Li, Chen Qian +5LLM Guardrails AI Agent Safety
Shanghai Artificial Intelligence Laboratory · Beihang University · Fudan University +2
PASTABench: Proactive Assessment of Sequential Trajectories for Agent Safety
Sep 23, 2026Jiapeng Sun, Yujin Zhou, Han Zhu +4AI Agent Safety Benchmarks Inference-Time Intervention
The Hong Kong University of Science and Technology, Hong Kong, China · Peking University, Beijing, China
PROACT-Agent: Progressive Runtime Oversight and Active Circuit-breaking for Real-Time Safety
Sep 28, 2026Ding Jia, Wei Liu, Xianglong Du +5LLM Guardrails Runtime Enforcement for AI Agents
State Key Laboratory of Intelligent Vehicle Safety Technology, Changan Automobile · Independent Researcher

Aug 25, 2026J/K move · Enter open · S save

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

Zhijie Zheng, Yu Li, Chen Qian, Yuqian Fu+4

Shanghai Artificial Intelligence Laboratory · Beihang University · Fudan University+2

LLM Guardrails AI Agent Safety LLM Agent Safety Tool Access Control for LLM Agents RL for LLM Agents

Open paper AskPDF arXiv

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.

Abstract

AI Agent Safety

Explore similar work

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

PASTABench: Proactive Assessment of Sequential Trajectories for Agent Safety

PROACT-Agent: Progressive Runtime Oversight and Active Circuit-breaking for Real-Time Safety