agent-decision-3b

A small, fast, calibrated decision model for AI agents. Given the tools, the user's request and the agent's proposed tool calls, it answers typed questions such as are these calls correct?, what is wrong with them? and should the agent call a tool, ask the user, or decline?, returning a probability for every allowed answer in one forward pass, with no text generation.

Use it anywhere an agent's step needs a fast, calibrated check: as a hook that vets tool calls before they run (like the grounding check in grounded_ai's Claude Code plugin), as a guardrail that stops an agent from acting on a confidently wrong call, for labelling or filtering agent traces, or as the cheap first stage of an evaluation cascade where only low-confidence decisions go to an LLM judge.

This repository holds a LoRA adapter for LiquidAI/d1-3B (3.1B) and calibration.json, the temperature applied to the option scores at readout. It is the 3B sibling of grounded-ai/agent-decision-9b, trained on the same rows with the same recipe.

• Dataset builder and showcase notebooks: github.com/grounded-ai/agent-decision
• Test data: grounded-ai/agent-decision-eval

Results

Test set: 1,139 rows from APIs and items never seen in training, sampled to a fixed number of rows per kind of example so rare kinds are measured as precisely as common ones. "Base" is the released d1-3B on the same rows, zero-shot. Accuracy / expected calibration error, x100.

Question	Base d1-3B	agent-decision-3b	agent-decision-9b
Are the tool calls correct and complete? (yes/no)	86.2 / 3.7	98.3 / 1.0	98.1 / 1.0
What is wrong with the calls? (correct / wrong tool / wrong arguments / incomplete)	75.3 / 8.6	96.4 / 2.1	96.4 / 1.9
Should the assistant call a tool, ask, or decline?	69.7 / 9.4	85.4 / 3.9	88.2 / 3.8
All	79.5 / 1.8	94.8 / 1.1	95.3 / 1.4

As a cascade's first stage: 86% of decisions come out at 90% confidence or more, and 99.1% of those are right.

Are the calls correct? by kind of call (accuracy x100):

Kind	Base d1-3B	agent-decision-3b
Correct calls	97.3	96.7
A parameter the tool does not define	64.0	99.3
A value the user gave is missing	91.0	99.0
A value differs from what the user asked	92.4	97.2
The wrong tool	96.7	100
One of several requested calls is missing	70.7	97.6

Call, ask or decline, by correct answer: call 94.4 → 95.2, decline 60.3 → 90.5, ask 52.3 → 68.5.

Unseen presentation (the same test rows rendered as markdown, a layout never used in training): 93.1% overall (base 80.5%). The weak spot is "ask the user": 52.3% in this layout (base 51.4%).

Usage

Needs transformers>=5.14 and a GPU with bfloat16 (an L4 is plenty; the model is about 6 GB). d1 ships its own model code, so it loads with trust_remote_code=True.

import json
import torch
from huggingface_hub import snapshot_download
from peft import PeftModel
from transformers import AutoModel

model = AutoModel.from_pretrained("LiquidAI/d1-3B", revision="051bcc464b01b9f92942b364d9586b0ef5912432",
                                  trust_remote_code=True, dtype=torch.bfloat16).to("cuda")
ours = snapshot_download("grounded-ai/agent-decision-3b")
model = PeftModel.from_pretrained(model, ours).merge_and_unload()


class Temperature:
    """Divides each option's score by the fitted temperature before the softmax."""
    def __init__(self, t):
        self.t = t

    def apply(self, question, scores):
        return [s / self.t for s in scores]


model.engine.calibration = Temperature(json.load(open(f"{ours}/calibration.json"))["temperature"])

tools = [{"name": "loginuser", "description": "Log in a user.",
          "parameters": {"username": {"type": "str"}, "password": {"type": "str"}}}]
state = {"tools": tools, "user": "Log in a user named 'guest' with the password 'guestpass'.",
         "assistant_tool_calls": [{"name": "loginuser", "arguments": {"username": "guest789", "password": "guestpass"}}]}
print(model.system_one(state, {
    "call_ok": {"type": "noul", "instructions": "Are the assistant's tool calls correct and complete for the user's request?"},
    "call_kind": {"type": "choice", "instructions": "What, if anything, is wrong with the proposed tool calls?",
                  "criteria": {"correct": "the calls are right and do everything the user asked",
                               "wrong_tool": "a call uses a tool that cannot do what was asked",
                               "wrong_arguments": "the right tool is called with arguments that do not match the request or the tool",
                               "incomplete": "the calls are right but part of the request is not done"}},
}))

Questions are free-form: any noul (yes/no), choice or score question can be asked about any state, as with d1. The fine-tune targets the questions above; other questions behave like the base model.

Training data

All labels are exact, and nothing is truncated:

• Salesforce/xlam-function-calling-60k (CC-BY-4.0). Half the rows are the execution-verified calls. The rest carry exactly one edit that makes them wrong by construction: a tool swapped for one whose schema rejects the arguments, a parameter the tool does not define, a value the user typed removed or changed to one they did not type, or one of several distinct calls dropped. Splits are by API name, so test APIs are never seen in training.
• nvidia/When2Call (CC-BY-4.0), multiple-choice split: call a tool, ask for missing information, or decline.

52,552 training rows, rendered as JSON or as a transcript; markdown is held out for testing. The prompts are rendered by d1's own prompt code, so training sees exactly what the model reads at inference.

Training

d1 answers by reading the language model's logits at the answer slot: each option's score is its best token logit there, and the answer is the softmax over the option scores. Training optimises exactly that: cross-entropy over the option scores, one position per question, rather than next-token loss over a response.

LoRA rank 16, alpha 32, no dropout, on the attention, short-convolution and MLP projections of the language model; the vision tower, the projector and the tied embeddings / LM head stay frozen (24.5M trainable parameters). AdamW, learning rate 1e-4, constant schedule, effective batch 32 rows (4 length-grouped micro-batches of 8), gradient clipping 1.0, bf16, gradient checkpointing. Early stopping on validation log loss stopped training at step 175; the weights from step 125 were kept. One temperature (1.158) fitted on a held-out calibration split, stored in calibration.json (d1's LM head is tied to its embeddings, so it is applied at readout rather than folded into weights). One L4 GPU, about 70 minutes of training. run_manifest.json holds the exact configuration.

Limitations

• Wrong calls in training and test are exact edits of correct calls. They are certain to be wrong, but may be easier to spot than the mistakes real agents make; treat the results above as an upper bound until measured on real agent errors.
• A call with a parameter the tool does not define counts as wrong. Some real APIs would ignore it; for a guardrail we chose strict.
• "Ask the user" is the weakest decision (68.5%, and 52.3% in an unfamiliar layout); the 9B does better there (81.1%).
• Confidence for a choice is its top probability, not normalized for the number of options.

License

This adapter is a derivative of LiquidAI/d1-3B and is released under the LFM Open License v1.0, like the base model: commercial use is licensed only for organisations with annual revenue under US$10 million (qualified non-profits are exempt for non-commercial and research use). Read the license before use. The training data is CC-BY-4.0: Salesforce xLAM function-calling dataset and NVIDIA When2Call.