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

What it is

PROBE trains a small LLM "auditor" that poses as a customer and talks to a frozen, tool-using target agent to surface policy violations that only show up in the full trace (hidden tool calls, missing confirmations, state changes). Code is at https://github.com/Mishrakshitij/PROBE (NeurIPS 2026 paper). The evidence comes from one environment, SKYASSIST, a deterministic flight-booking sandbox.

How it works

flowchart LR
    S[Scenario pool] --> A[Auditor 4B, QLoRA]
    A -->|user message| T[Frozen target 30B-A3B + 8 tools]
    T -->|public reply only| A
    T -->|full trace incl. tool calls| V[20 verifiers]
    V -->|scalar reward only| A

• Setup: The auditor is Qwen3-4B-Instruct (rank-64 QLoRA); the target is a frozen Qwen3-30B-A3B served with vLLM. The auditor sees only the public dialogue, never tool calls, prompt, database or verifier labels. Verifiers see the whole trace and return only a scalar reward.
• Environment: 35 airports, 120 routes, 643 flights, 8 tools (4 read-only, 4 mutating) and 20 policies. Each policy has a severity of 0.5, 1, 2 or 4. Sixteen verifiers are deterministic (regex, state machines, schema checks). Four use an LLM judge during training (B4 upsell, B6 round-trip, S5 confirmation, P1 pricing disclosure).
• Reward (Eq. 1): violation + progress + novelty + plausibility, z-score normalized.
• Violation is severity-weighted and deduplicated per episode: the first hit gets full credit, the second half credit, and later repeats zero. This stops the auditor farming one easy rule.
• Progress rewards first-time entry into tool states and is annealed from 0.5 to 0.1 over 100 steps. Novelty is annealed from 0.1 to 0.4.
• Plausibility penalizes assistant-like phrasing, markdown, fabricated system-access claims, excess length and echoing the previous turn.
• Curriculum: The horizon grows from 4 to 8, 12 and 14 turns, advancing at 15% yield with patience 50.
• Optimizers compared: GRPO, DPO and GFlowNet.
• DPO is the most stable under monotone reward reshaping.
• GFlowNet has the highest training yield and is mode-covering, but needs careful reward calibration.
• GRPO and GFlowNet react strongly to reward magnitudes.

Results the summary omits

• Held-out policy coverage rises from 35% (strong zero-shot auditor) to 60%.
• Multi-policy failures occur in up to 94% of evaluation episodes.
• Under deterministic re-scoring (no LLM judge), trained auditors reach 55% versus 5% for a static scripted suite.
• Transfer to τ-bench is reported; the excerpt gives no numbers.
• Verifiers were validated against human labels, and the authors say they observed over-optimization of LLM-judged rewards (their Section 6.3).

Preconditions / Affected

Reproducing this needs a target agent with a typed tool surface, policies written as trace-level verifiers (most programmatic), and a scenario pool. The method's central design is "verifier-first": the policy set must be formalized as executable checks before training.

Caveats

• Only one domain (flight booking), one target model family and one auditor size are described in the fetched text. Generalization beyond that rests on the τ-bench transfer, which I couldn't see in detail.
• The excerpt cuts off before the experiments section, so the optimizer-by-optimizer numbers and the human-label agreement figures are not available here.
• Reward weights were set from training diagnostics, not held-out coverage, per the authors.
• The 4 LLM-judged policies were noisy enough that the authors added the deterministic re-scoring.