VaidikML0508

qwen2.5-3b-investigation-grpo

README

License: apache-2.0

What this model does

On every turn it should emit exactly one tool call:

call_action — run an unlocked investigation action (gather evidence, unlock the DAG).
close_case — accuse a culprit and cite evidence ids you actually observed.

Success (four-way solve) requires: closed case + correct culprit + all required evidence discovered + citations ⊆ observed evidence.

Results

Animation — Base vs SFT vs SFT+RL

Side-by-side rollout from the interactive case showcase (The Adventure of the Speckled Band, hop-3). Frames advance turn-by-turn for each agent.

Base vs SFT vs SFT+RL investigation compare

Upload compare_medium_speckled_band.gif next to this README on the Hub so the image renders.

Table with columns: Agent, Solved, Closed, Culprit, Evidence complete, Cite OK, Recall, Steps
Agent	Solved	Closed	Culprit	Evidence complete	Cite OK	Recall	Steps
Base	✗	✗	✗	✗	✗	0%	16 (invalid tool format)
SFT	✗	✗	✗	✓	✗	100%	16 (explores; broken `citations_json`)
SFT+RL	✓	✓	✓	✓

Takeaway: Base cannot emit valid <tool_call>s. SFT learns to investigate the unlock DAG but fails to close with valid citations. SFT+RL gathers required evidence and closes with a correct, observable citation set.

Harder case (hop-4 showcase)

The Boscombe Valley Mystery (8 required evidence ids):

Table with columns: Agent, Solved, Closed, Recall, Steps
Agent	Solved	Closed	Recall	Steps
Base	✗	✗	0%	16
SFT	✗	✗	0%	16
SFT+RL	✗	✗	50%	16

RL improves exploration on harder DAGs but does not yet fully solve every case — capacity and early-close / cite discipline remain active work.

Held-out eval (n=8 cases)

Same protocol for base / SFT / SFT+RL:

Table with columns: Checkpoint, Solve, Closed, Culprit, Mean recall, Valid actions
Checkpoint	Solve	Closed	Culprit	Mean recall	Valid actions
Base	0.00	0.00	0.00	0.00	0.00
SFT	0.00	0.13	0.13	0.33	0.41
SFT+RL	0.00*

*Aggregate solve on this slice was still near zero when models often closed before full evidence; the Speckled Band showcase shows a clean full solve under the investigation gym. Prefer case-level rollouts + evidence_complete / cite_ok / early_close diagnostics alongside solve rate.

Install

bash
pip install -U transformers peft accelerate bitsandbytes torch
# optional (private repos / gated base):
pip install -U huggingface_hub
huggingface-cli login

GPU with ~8–16GB VRAM is enough for 4-bit inference.

Load the adapter

Replace ADAPTER_ID with the RL or SFT repo above.

python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel

BASE_MODEL = "Qwen/Qwen2.5-3B-Instruct"
ADAPTER_ID = "VaidikML0508/mystery-rl-adapter22"  # or mystery-sft-adapter22

bnb = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

base = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    quantization_config=bnb,
    device_map="auto",
    trust_remote_code=True,
)
model = PeftModel.from_pretrained(base, ADAPTER_ID)
model.eval()

Without bitsandbytes (bf16 / fp16, more VRAM):

python
base = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)
model = PeftModel.from_pretrained(base, ADAPTER_ID)

Tool format (required)

The policy was trained on text tool calls (not native HF tools= JSON). Output must look like:

text
<tool_call>{"name":"call_action","arguments":{"action_name":"inspect_study","arguments_json":"{}"}}</tool_call>

text
<tool_call>{"name":"close_case","arguments":{"culprit":"Jordan Lee","citations_json":"[\"ev_ledger\",\"ev_key\"]"}}</tool_call>

Rules used in training:

Exactly one <tool_call>...</tool_call> per turn, then stop.
Only call currently unlocked action names from the user message.
Do not close until required evidence is gathered (RL phase rejects / discourages early close).
citations_json is a stringified JSON list of evidence ids you observed.

Quick inference (single turn)

Minimal smoke test: one generate step from a chat prompt.

python
import re
import torch

SYSTEM = (
    "You are an investigation agent. On EVERY turn output exactly ONE tool call, then stop.\n"
    "Format (exact):\n"
    '<tool_call>{"name":"call_action","arguments":{"action_name":"ACTION","arguments_json":"{}"}}</tool_call>\n'
    "or\n"
    '<tool_call>{"name":"close_case","arguments":{"culprit":"NAME","citations_json":"[\\"ev1\\",\\"ev2\\"]"}}</tool_call>\n'
    "Rules:\n"
    "- Only use currently unlocked action names from the user message.\n"
    "- Do not repeat the same action if it already returned evidence.\n"
    "- After you have required evidence, call close_case with citations you observed.\n"
)

USER = """Case brief: A rare manuscript vanished from the university archive overnight.
Suspects: Avery Quinn, Jordan Lee, Sam Ortiz
Places: archive, loading dock, dormitory

Unlocked actions: list_archive_log, interview_night_guard
Discovered evidence: []

Call exactly one tool."""

messages = [
    {"role": "system", "content": SYSTEM},
    {"role": "user", "content": USER},
]

prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=160,
        do_sample=True,
        temperature=0.3,
        top_p=0.9,
        pad_token_id=tokenizer.pad_token_id,
    )

gen = tokenizer.decode(out[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True)
print(gen)

# Optional: extract first tool_call JSON
m = re.search(r"<tool_call>\s*(\{.*?\})\s*(?:</tool_call>|$)", gen, flags=re.DOTALL)
print("parsed:", m.group(1) if m else None)

Multi-turn investigation loop (inference recipe)

Wire generations to your own environment (or case gym). Pseudocode:

python
import json
import re
from typing import Any

def extract_tool_call(text: str) -> dict[str, Any] | None:
    m = re.search(
        r"<tool_call>\s*(\{.*?\})\s*(?:</tool_call>|(?=<tool_call>)|$)",
        text,
        flags=re.DOTALL,
    )
    if not m:
        return None
    try:
        return json.loads(m.group(1))
    except json.JSONDecodeError:
        return None


def step_env(tool: dict[str, Any]) -> str:
    """Replace with your InvestigationGym / API."""
    name = tool.get("name")
    args = tool.get("arguments") or {}
    if name == "call_action":
        return env.call_action(args["action_name"], args.get("arguments_json", "{}"))
    if name == "close_case":
        return env.close_case(args["culprit"], args["citations_json"])
    return f"ERROR: unknown tool {name}"


messages = [
    {"role": "system", "content": SYSTEM},
    {"role": "user", "content": initial_case_prompt},  # brief + unlocked + discovered
]

for turn in range(12):
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.inference_mode():
        out = model.generate(
            **inputs,
            max_new_tokens=160,
            do_sample=True,
            temperature=0.3,
            top_p=0.9,
            pad_token_id=tokenizer.pad_token_id,
        )
    gen = tokenizer.decode(out[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True)

    # Keep only the first tool_call block if the model continues
    if "<tool_call>" in gen:
        i = gen.find("<tool_call>")
        rest = gen[i:]
        j = rest.find("</tool_call>")
        gen = rest[: j + len("</tool_call>")] if j != -1 else rest

    messages.append({"role": "assistant", "content": gen})
    tool = extract_tool_call(gen)
    if not tool:
        messages.append({
            "role": "user",
            "content": "ERROR: no valid tool_call. Output exactly one <tool_call>{...}</tool_call>.",
        })
        continue

    obs = step_env(tool)
    # Include updated unlocked / discovered lists in the user turn when you have them
    messages.append({
        "role": "user",
        "content": (
            f"Tool result ({tool.get('name')}):\n{obs}\n"
            "Continue with exactly one next tool_call "
            "(new evidence action, or close_case if ready)."
        ),
    })

    if tool.get("name") == "close_case" and not str(obs).startswith("ERROR"):
        break

Tip: Prefer temperature≈0.2–0.4 for eval; raise slightly for exploration demos.

Intended inputs

User turns should include:

Public case brief + suspects (no gold solution)
Currently unlocked action names
Evidence discovered so far
After each tool: observation text + refreshed unlocked/discovered

Without an environment that returns observations and unlocks, the adapter cannot run a full investigation — it only proposes the next tool call.

Training sketch (for card readers)

Build validated investigation cases (unlock DAG, required evidence, no solution leak in the public brief).
SFT on scripted full-solve trajectories (call_action* → gold close_case).
GRPO-style multi-turn RL in the gym with outcome rewards tied to the four-way solve; early close blocked until evidence is complete.

This repo is the adapter only. Load it on top of the base Instruct model as shown above.

Limitations

Small 3B policy: long / deep DAGs may still fail to fully solve.
Must use the <tool_call> text protocol; native tools= chat APIs are not the training format.
Needs a compatible case environment at inference for multi-turn solves.
Mystery domain English cases; not a general chat model replacement.

Citation

If you use this adapter, please cite the base model and note QLoRA SFT → GRPO investigation training on a sealed-case gym with four-way solve (closed + culprit + evidence complete + faithful citations).

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Container

Run this model inference with full control and performance in your environment.