Learning When to Act or Refuse: Guarding Agentic Reasoning Models for Safe Multi-Step Tool Use
Aradhye Agarwal, Gurdit Siyan, Yash Pandya, Joykirat Singh, Akshay Nambi, Ahmed Awadallah
ICML 2026)
cs.CL
2026-03-04
MOSAIC trains agents to plan, check, then act or refuse. Qwen2.5-7B harm drops 0.18 to 0.09 with 87% refusal; Qwen3 benign completion nearly doubles 44% to 85%.
Safety training for chat models does not carry over to agents that plan, call tools, and act across many steps. A request that looks harmless in turn one can become a file read, a permission change, or a payment by turn five. Output filters and single-turn refusals never see the irreversible step.
Small models take this harder. They have shorter context and thinner world models, so they follow injected tool outputs more readily. In the paper's baselines, open SLMs either comply with injections, misuse tools, or refuse benign work. Unscaffolded GPT-4o refuses 0% of harmful AgentHarm tasks, scores 0.31 on harm, and fails 76% of direct prompt injections. Scale does not teach an agent when to act and when to stop.
Scalar rewards collapse time. Two traces can end in similar states, one refusing on step one, the other following an injected instruction for several calls before aborting. A single score treats them as close.
MOSAIC is a Microsoft Research post-training stack that forces the loop into plan, check, then act or refuse. Each turn writes a plan inside <think>, optionally writes <safetythoughts> about harm, irreversibility, permission changes, and new risk in the latest tool output, then picks one of three actions: a tool call, refusaltool with a justification (terminal), or <answer>. Refusal lives in the same action space as tools. It is not a filter applied after generation.
Whether to emit safety thoughts is learned. The gate is simply whether the model opens the tag. On AgentHarm, Qwen2.5 does this on 72% of turns, Phi-4 on 30.5%, Qwen3-4B-Thinking on 0.1%, because that model already thinks at length internally.
There are no step-level safety labels. For each prompt the trainer samples n=4 full traces. GPT-4o compares every pair (win 1, tie 0.5, loss 0), and the win sum becomes the GRPO alignment reward in [0, 3]. A format reward in [0, 2] keeps tags parseable. A soft length penalty starts only after 400 tokens per turn. Gradients skip tool-output tokens, so the policy learns what to write, not how to copy the environment.
Training data is Agent-SafetyBench: about 2,000 tasks across 349 environments, mixing harmful, benign, and dual-strategy cases. Evaluation is zero-shot on AgentHarm, Agent Security Bench, BFCL v3, and PrivacyLens.
Gains follow each model's original bias. The method does not push every model toward the same conservative point.
| Model | Setting | AgentHarm harm | Harmful refusal | ASB DPI ASR | Benign completion |
| Qwen2.5-7B | base → MOSAIC | 0.18 → 0.09 | 0.74 → 0.87 | 0.55 → 0.42 | 0.90 → 0.84 |
| Qwen3-4B-Thinking | base → MOSAIC | 0.09 → 0.08 | 0.86 → 0.89 | 0.46 → 0.29 | 0.44 → 0.85 |
| Phi-4 | base → MOSAIC | 0.06 → 0.09 | 0.94 → 0.88 | 0.19 → 0.28 | 0.78 → 0.91 |
| GPT-4o | no scaffold → MOSAIC scaffold | 0.31 → 0.07 | 0.00 → 0.92 | 0.76 → 0.21 | 0.89 → 0.93 |
Qwen2.5 is a safety hardening: harm halves, indirect-injection ASR falls 0.40→0.33, benign completion drops 6 points. That is the cleanest safety-for-utility trade in the paper.
Qwen3 was spinning in unproductive reasoning. After MOSAIC, ASB benign completion goes 44%→85% while harm barely moves. Tokens per turn fall 1172→262, about 4×, and explicit safety thoughts are almost unused.
Phi-4 started over-refusing (benign refusal 0.43). Training cuts that to 0.19 and lifts completion to 91%. Harmful refusal falls 0.94→0.88 and DPI ASR rises 0.19→0.28. Relaxing a conservative prior gives utility back and leaks some safety.
Privacy transfers. On PrivacyLens, Qwen2.5 leakage 0.32→0.26 and adjusted leakage 0.48→0.37, with helpfulness slightly up. Phi-4 leakage 0.38→0.32, helpfulness 0.87→0.76.
On purely benign BFCL v3 multi-turn tool use, Qwen2.5 base accuracy 21.0→28.5, a 35% relative gain. The safety loop did not damage ordinary tool calling.
Two ablations both matter. Dropping <safetythoughts> and keeping only <think> plus a refusal tool cuts Qwen2.5 harmful refusal 0.87→0.59 and raises harm 0.09→0.12. Replacing pairwise preferences with a scalar reward cuts refusal 0.87→0.79 and raises DPI ASR 0.42→0.51.
For GPT-4o and GPT-5 the paper applies the inference scaffold, not RL. Unscaffolded, both refuse 0% of harmful tasks. With plan-check-refuse, refusal exceeds 90% and GPT-4o harm falls 0.31→0.07. MOSAIC-trained open models beat unscaffolded frontier models; once the frontier models get the same scaffold, the remaining gap is mostly direct injection.
If you ship a tool-using agent and still rely on chat refusals, the gap is concrete: refusal has to be an action, the safety check has to be a separate block, and the reward has to rank early refusal above late abort.
The loop itself is usable without training. GPT-4o, scaffold only, drops harm from 0.31 to 0.07. Small models need GRPO to learn when to open the check. Qwen3 barely uses safety thoughts and still fixes completion and token use, so for thinking models the length penalty and the refusal tool may be doing more work than the extra reasoning block.
This is incremental engineering, not a new safety theory. The structure is copyable, the numbers sit on public agent benchmarks, and training is verl on 4×A100.
The paper has no Limitations section. The cracks that do show up are enough.
The training judge is GPT-4o and prefers the first trajectory in about 60% of pairs. Order is randomized, so the bias should average out, but the pairwise reward is noisier than the tables imply. AgentHarm is scored by GPT-4.1 and PrivacyLens by GPT-4o. Training and evaluation orbit the same family of judges. Agreement with humans is not reported.
Phi-4 shows the trade is real. Pulling back over-refusal raises DPI ASR from 0.19 to 0.28. MOSAIC retunes each model's prior. It does not find a point that is safer for every model.
Qwen3 invokes safety thoughts on 0.1% of turns. Credit for "explicit safety reasoning" does not stick on that model.
Training is Agent-SafetyBench; evaluation includes Agent Security Bench. The names and citations sit close together. The paper calls the eval OOD and does not show overlap statistics. Every benchmark is a sandbox. There are no production tools, real credentials, or real payments.
With n=4, pairwise judging is O(n²) calls. Fine for 7B rollouts, expensive for long-horizon agents whose traces already cost real money.