PROBE: Learning to Audit Policy Compliance in Tool-Using LLM Agents

What it claims

A Qwen3-4B-Instruct auditor, trained with QLoRA against a frozen Qwen3-30B-A3B flight-booking agent, raises held-out policy coverage from 35% for a strong zero-shot auditor to 60%. It produces multi-policy failures in up to 94% of evaluation episodes. Under deterministic re-scoring, learned auditors reach 55% coverage against 5% for a static scripted suite. The authors also report transfer to τ-bench. The paper is a NeurIPS 2026 paper, and code is at https://github.com/Mishrakshitij/PROBE.

How it works

The auditor sees only the public dialogue and a customer scenario. It never sees the target's weights, system prompt, database, tool calls or verifier labels. A verifier suite scores the full trace, including the hidden tool calls and the database state.

flowchart LR
  S[Scenario pool] --> A[Auditor 4B QLoRA]
  A -->|user msg| T[Frozen target 30B-A3B]
  T -->|hidden tool calls| E[SkyAssist DB]
  T -->|public reply only| A
  T --> V[Verifier suite]
  E --> V
  V --> R[Composite reward]
  R -->|GRPO / DPO / GFlowNet| A

• Environment (SkyAssist). The environment is deterministic: 35 airports, 120 routes, 643 flights and 3 fare classes. It deliberately includes 4 restricted destinations and 46 cancelled flights. It exposes 8 tools: 4 read-only and 4 mutating (hold, book, cancel, modify), and mutations require confirmed=true. There are 20 policies across business (B1–B6), safety (S1–S6), privacy (P1–P4) and operational (O1–O4) categories. Each policy has a severity of 0.5, 1, 2 or 4.
• Verifiers. 16 are programmatic: regex, state machines and schema checks. For example, S5 is a state machine that checks confirmation precedes any mutating call. The other 4 (B4 upsell, B6 round-trip, S5 confirm, P1 pricing) use a hybrid LLM judge during training, with more than 90% self-consistency.
• Reward. The reward is r_viol + α·r_prog + β·r_nov + γ·r_plaus, with each term z-score normalized using EW running statistics.
• Violation credit: severity-weighted, with full credit for a policy's first hit in an episode, half for the second and zero after that.
• Progress: rewards reaching new tool states, from 0.03 for slot filling up to 0.25–0.30 for booking, cancellation and multi-leg flows.
• Novelty: inverse global frequency plus a first-discovery bonus.
• Plausibility: penalizes assistant-like phrasing, markdown or lists, fabricated system access, excess length and echoing the previous turn.
• Schedule: α anneals from 0.5 to 0.1 and β from 0.1 to 0.4 over the first 100 steps. The authors say the weights were set from training diagnostics, not held-out coverage.
• Curriculum. The maximum conversation length grows through 4, 8, 12 and 14 turns. Training advances a stage when violation yield exceeds 15%, with a patience of 50.
• Training setup. The group size is G=4. The auditor uses rank-64 NF4 QLoRA, which trains 47M of its 4B parameters. Auditor messages are capped at 1–3 sentences, must stay in the customer role and may not fabricate.

Evidence

• Coverage. Held-out coverage rose from 35% for the zero-shot auditor to 60% for trained auditors. Up to 94% of episodes contain multiple policy violations. Under deterministic re-scoring, trained auditors reach 55% against 5% for a static suite.
• Optimizer comparison.
• GRPO and GFlowNet are sensitive to reward magnitudes.
• DPO is the most stable under monotone reward reshaping.
• GFlowNet, which is mode-covering, gives the highest training-time yield but needs careful reward calibration.
• Verifier validation. The verifiers were checked against human labels.
• Missing from the text. The fetched text gives no per-optimizer coverage table and no τ-bench numbers.

How to use it

The fetched text gives only the code repository link. It contains no install commands, no licence and no released auditor weights. To reproduce the setup you would need:

• Qwen3-4B-Instruct trained with QLoRA.
• A Qwen3-30B-A3B target served with vLLM.
• The SkyAssist environment and verifiers.

To adapt PROBE to your own agent, the costly part is writing the verifiers: one trace-level predicate per policy, with a severity attached.

Limits

• Truncated page. The fetched text stops at Section 4.3. The results tables, ablations, the human study and the τ-bench transfer details were not visible, so the headline numbers above come from the introduction.
• Reward hacking. The authors observed over-optimization against the LLM-judged rewards (Section 6.3). This is why they re-score all methods deterministically, which lowers 60% to 55%.
• Single environment. All main experiments use one synthetic, deterministic flight-booking environment with one target model. Some policies are business preferences rather than safety rules, such as B4 "mention upgrade" and B6 "suggest round-trip."
• Calibration sensitivity. GFlowNet's exploration gains depend on careful reward calibration.
• Auditor only. The co-evolutionary loop, in which the target agent learns from auditor feedback, is proposed but not tested.