14.10 A Narrow Auto-Research Agent: ReAct, Tool-Use & Retrieval by Distillation¶

   ┌─────────────────────────────────────────────────────────────────┐
   │ 0. WORLD BUILD       stacklm.agent.corpus:                      │
   │                      build_corpus()   -> 268 short passages     │
   │                      make_task_pool() -> 200 train + 50 held-out│
   │                                          Tasks with exact gold  │
   └───────────────┬─────────────────────────────────────────────────┘
                   │  tasks + a ToolEnv over the corpus
                   ▼
   ┌─────────────────────────────────────────────────────────────────┐
   │ 1. TEACHER ROLLOUT   large model runs the REAL ReAct loop with  │
   │                      our two tools -> raw trajectory            │
   │                      (in CI: teacher is STUBBED/offline)        │
   └───────────────┬─────────────────────────────────────────────────┘
                   │  trajectory + did it solve the task?
                   ▼
   ┌─────────────────────────────────────────────────────────────────┐
   │ 2. FILTER            keep only trajectories whose final answer  │
   │                      exactly matches gold (verifiable reward)   │
   └───────────────┬─────────────────────────────────────────────────┘
                   │  successful traces
                   ▼
   ┌─────────────────────────────────────────────────────────────────┐
   │ 3. FORMAT            splice in tool special tokens, mask loss so│
   │                      the model is NOT trained on observations   │
   └───────────────┬─────────────────────────────────────────────────┘
                   │  SFT dataset
                   ▼
   ┌─────────────────────────────────────────────────────────────────┐
   │ 4. SFT Stack-100M    imitate the groove  (Ch. 14.9 loop reused) │
   └───────────────┬─────────────────────────────────────────────────┘
                   │  optional
                   ▼
   ┌─────────────────────────────────────────────────────────────────┐
   │ 5. RLVR (GRPO)       reward = final-answer exact match; sharpen │
   │                      multi-turn: loss over GENERATED tokens only│
   └─────────────────────────────────────────────────────────────────┘

# capstone/stacklm/agent/corpus.py
"""The narrow world the auto-research agent lives in.

Two artifacts, both deterministic and offline:
  build_corpus()   -> list[Passage]   268 short passages, chunked, doc_id'd
  make_task_pool() -> (train, heldout) list[Task] with EXACT gold answers

Everything is rendered from an explicit (entity, attribute, value) fact table,
so a task's gold answer is correct BY CONSTRUCTION and is always recoverable
from the corpus. Ch. 14.11's retrieval-QA probe imports these same functions.
"""

from __future__ import annotations
import random
from dataclasses import dataclass, field
from stacklm.agent.tools import Passage

# ----------------------------------------------------------------------
# 1. The fact table.  (entity, attribute, value) -- values are strings so a
#    single code path handles numeric and non-numeric facts; `_is_num`
#    decides which task templates a fact is eligible for.
#    These 12 values are byte-identical to PLAN.md 1 -- do not drift.
# ----------------------------------------------------------------------
SPEC_FACTS = [
    ("Stack-100M", "hidden size", "512"),
    ("Stack-100M", "layer count", "30"),
    ("Stack-100M", "vocabulary size", "32768"),
    ("Stack-100M", "query head count", "8"),
    ("Stack-100M", "key-value head count", "2"),
    ("Stack-100M", "head dimension", "64"),
    ("Stack-100M", "MLP inner dimension", "1408"),
    ("Stack-100M", "pretraining sequence length", "2048"),
    ("Stack-100M", "mid-training sequence length", "8192"),
    ("Stack-100M", "RoPE base", "10000"),
    ("Stack-100M", "parameter count in millions", "101"),
    ("Stack-100M", "pretraining token budget in billions", "20"),
]

# Publication years of landmark works this book cites elsewhere. These are
# real, verifiable facts -- we never invent a fact to pad the corpus.
PAPER_FACTS = [
    ("the Transformer paper", "publication year", "2017"),
    ("Adam", "publication year", "2014"),
    ("RMSNorm", "publication year", "2019"),
    ("SwiGLU", "publication year", "2020"),
    ("RoFormer", "publication year", "2021"),
    ("LoRA", "publication year", "2021"),
    ("Chinchilla", "publication year", "2022"),
    ("FlashAttention", "publication year", "2022"),
    ("ReAct", "publication year", "2022"),
    ("STaR", "publication year", "2022"),
    ("Toolformer", "publication year", "2023"),
    ("DPO", "publication year", "2023"),
    ("grouped-query attention", "publication year", "2023"),
    ("DeepSeekMath", "publication year", "2024"),
    ("Muon", "publication year", "2024"),
    ("MobileLLM", "publication year", "2024"),
]

_DISTRACTOR_ATTRS = ["hidden size", "layer count", "MLP inner dimension",
                     "vocabulary size", "head dimension"]

# Three surface forms so BM25 and the embedding backend see real lexical
# variation -- but ALL of them contain the canonical clause
# "the {attr} of {entity} is {value}", which is what makes extraction exact.
_TEMPLATES = [
    "Reference note {n}. In the Stack-100M project notes, the {attr} of {ent} "
    "is {val}. This value is fixed by the capstone specification and is not "
    "tuned per run.",
    "Configuration record {n}: the {attr} of {ent} is {val}. Downstream stages "
    "(mid-training, supervised fine-tuning, serving) all assume this setting "
    "without re-deriving it.",
    "Design log entry {n}. We record here that the {attr} of {ent} is {val}, a "
    "choice made during the architecture pass and left unchanged since.",
]


def _render(n: int, ent: str, attr: str, val: str) -> str:
    return _TEMPLATES[n % len(_TEMPLATES)].format(n=n, ent=ent, attr=attr, val=val)


def chunk_words(text: str, size: int = 80, overlap: int = 16) -> list[str]:
    """Fixed-size word chunking with overlap -- the simplest chunker that still
    respects the two rules from Ch. 9.4: (a) a chunk must be small enough that
    the answer span is not diluted by irrelevant text, (b) chunks must OVERLAP
    so a fact straddling a boundary survives in one of them. Our rendered
    passages average 28 words, so this is a no-op for them; it exists so you can
    drop real documents into the same pipeline."""
    words = text.split()
    if len(words) <= size:
        return [text]
    out, step = [], max(1, size - overlap)
    for i in range(0, len(words), step):
        piece = words[i:i + size]
        if piece:
            out.append(" ".join(piece))
        if i + size >= len(words):
            break
    return out


def build_corpus(seed: int = 0, n_distractors: int = 120) -> list[Passage]:
    """268 short passages: 28 real facts + 120 synthetic ablation records
    (2 facts each). Deterministic given `seed`."""
    rng = random.Random(seed)
    facts = list(SPEC_FACTS) + list(PAPER_FACTS)

    # Synthetic distractors: same attribute vocabulary, different entities and
    # values, so a lazy query like "hidden size" retrieves noise, not gold.
    for i in range(n_distractors):
        ent = f"ablation run R-{i:03d}"
        for attr in rng.sample(_DISTRACTOR_ATTRS, 2):
            facts.append((ent, attr, str(rng.choice([128, 192, 256, 384, 640,
                                                     768, 896, 1024, 1536]))))

    passages: list[Passage] = []
    for n, (ent, attr, val) in enumerate(facts):
        for j, piece in enumerate(chunk_words(_render(n, ent, attr, val))):
            passages.append(Passage(doc_id=f"doc-{n:04d}-{j}", text=piece))
    return passages


# ----------------------------------------------------------------------
# 2. The task pool.  Four SHAPES, deliberately: the held-out split lets
#    Ch. 14.11 measure in-shape accuracy, and holding out a whole shape
#    (Exercise 9) measures the distribution-brittleness this chapter warns
#    about.  make_tasks() yields 262 tasks: 28 + 28 + 120 + 86.
# ----------------------------------------------------------------------
@dataclass
class Task:
    question: str
    gold: str                       # the verifiable final answer (exact match)
    kind: str                       # "lookup" | "double" | "sum" | "compare"
    # `hints` is the ANSWER KEY. The real teacher never sees it; the CI stub
    # teacher does (a test double is allowed to cheat -- see below).
    hints: dict = field(default_factory=dict)


def _is_num(v: str) -> bool:
    return v.lstrip("-").isdigit()


def make_tasks(seed: int = 0) -> list[Task]:
    rng = random.Random(seed + 1)
    facts = list(SPEC_FACTS) + list(PAPER_FACTS)
    numeric = [f for f in facts if _is_num(f[2])]      # all 28, as it happens
    tasks: list[Task] = []

    for ent, attr, val in facts:                                # SHAPE 1: lookup
        tasks.append(Task(
            question=f"What is the {attr} of {ent}? Use the corpus.",
            gold=val, kind="lookup",
            hints={"queries": [f"{ent} {attr}"], "facts": [(ent, attr)], "op": None}))

    for ent, attr, val in numeric:                              # SHAPE 2: x2
        tasks.append(Task(
            question=f"What is the {attr} of {ent}, multiplied by 2? Use the corpus.",
            gold=str(int(val) * 2), kind="double",
            hints={"queries": [f"{ent} {attr}"], "facts": [(ent, attr)], "op": "double"}))

    for _ in range(120):                                        # SHAPE 3: sum
        (e1, a1, v1), (e2, a2, v2) = rng.sample(numeric, 2)
        tasks.append(Task(
            question=(f"What do you get if you add the {a1} of {e1} to the "
                      f"{a2} of {e2}? Use the corpus."),
            gold=str(int(v1) + int(v2)), kind="sum",
            hints={"queries": [f"{e1} {a1}", f"{e2} {a2}"],
                   "facts": [(e1, a1), (e2, a2)], "op": "sum"}))

    for _ in range(120):                                        # SHAPE 4: compare
        (e1, a1, v1), (e2, a2, v2) = rng.sample(numeric, 2)
        # Two rejections, both load-bearing. Equal values leave no unique gold.
        # SAME ENTITY leaves a question that is unanswerable as posed --
        # "which is larger, the hidden size of Stack-100M or the layer count of
        # Stack-100M? Answer with the entity name" has gold "Stack-100M" down
        # EITHER branch, so it is free marks for teacher and student alike and
        # silently inflates compare-shape accuracy. Stack-100M owns 12 of the
        # 28 numeric facts, so this fires on ~17% of draws -- not a corner case.
        if e1 == e2 or int(v1) == int(v2):
            continue
        tasks.append(Task(
            question=(f"Which is larger, the {a1} of {e1} or the {a2} of {e2}? "
                      f"Answer with the entity name. Use the corpus."),
            gold=e1 if int(v1) > int(v2) else e2, kind="compare",
            hints={"queries": [f"{e1} {a1}", f"{e2} {a2}"],
                   "facts": [(e1, a1), (e2, a2)], "op": "compare",
                   "names": [e1, e2]}))
    return tasks


def make_task_pool(seed: int = 0, n_train: int = 200, n_heldout: int = 50):
    """Shuffle and split. The held-out slice is NEVER shown to the teacher and
    is what Ch. 14.11 evaluates on."""
    tasks = make_tasks(seed)
    random.Random(seed + 2).shuffle(tasks)
    assert len(tasks) >= n_train + n_heldout, "pool too small; raise the counts"
    return tasks[:n_train], tasks[n_train:n_train + n_heldout]


def qa_pairs(tasks: list[Task]) -> list[dict]:
    """Exactly the shape Ch. 14.11's `eval_retrieval_qa(..., qa_pairs)` wants."""
    return [{"question": t.question, "gold_answer": t.gold} for t in tasks]

# capstone/stacklm/agent/tools.py
"""Two narrow tools for the Stack-100M auto-research agent: an exact
calculator and a small local retriever (BM25 + an embedding-lite backend).
Everything here is deterministic, offline, pure-stdlib, and CI-safe."""

from __future__ import annotations
import ast, math, operator, re, hashlib
from collections import Counter
from dataclasses import dataclass


# ----------------------------------------------------------------------
# TOOL 1: a SAFE calculator.  We never call eval() on model output.
# Instead we parse to an AST and walk it, allowing only arithmetic nodes.
# ----------------------------------------------------------------------
_BINOPS = {
    ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul,
    ast.Div: operator.truediv, ast.Pow: operator.pow, ast.Mod: operator.mod,
    ast.FloorDiv: operator.floordiv,
}
_UNARYOPS = {ast.UAdd: operator.pos, ast.USub: operator.neg}


def _eval_node(node: ast.AST) -> float:
    if isinstance(node, ast.Expression):
        return _eval_node(node.body)
    if isinstance(node, ast.Constant):              # a literal number
        if isinstance(node.value, (int, float)):
            return node.value
        raise ValueError("non-numeric constant")
    if isinstance(node, ast.BinOp) and type(node.op) in _BINOPS:
        return _BINOPS[type(node.op)](_eval_node(node.left), _eval_node(node.right))
    if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARYOPS:
        return _UNARYOPS[type(node.op)](_eval_node(node.operand))
    raise ValueError(f"disallowed expression node: {type(node).__name__}")


def calc(expr: str) -> str:
    """Evaluate an arithmetic expression, returning a compact string.
    Errors are returned AS OBSERVATIONS (never raised) so the agent can read
    them and recover -- important for teaching self-correction."""
    try:
        tree = ast.parse(expr.strip(), mode="eval")
        val = _eval_node(tree)
    except Exception as e:                          # noqa: BLE001 (catch all)
        return f"CalcError: {e}"
    if isinstance(val, float) and val.is_integer():  # 4042.0 -> "4042"
        val = int(val)
    return str(val)


# ----------------------------------------------------------------------
# TOOL 2: a tiny retriever over an in-memory corpus of passages.
# ----------------------------------------------------------------------
_TOK = re.compile(r"[a-z0-9]+")


def _tokenize(text: str) -> list[str]:
    return _TOK.findall(text.lower())


@dataclass
class Passage:
    doc_id: str
    text: str


class BM25Retriever:
    """Textbook Okapi BM25 (Robertson & Zaragoza). Lexical, exact, fast, and
    shockingly hard to beat on keyword-heavy narrow corpora."""

    def __init__(self, passages: list[Passage], k1: float = 1.5, b: float = 0.75):
        self.passages = passages
        self.k1, self.b = k1, b
        self.docs = [_tokenize(p.text) for p in passages]
        self.doc_len = [len(d) for d in self.docs]
        self.avgdl = sum(self.doc_len) / max(1, len(self.docs))
        df: Counter[str] = Counter()
        for d in self.docs:
            for t in set(d):
                df[t] += 1
        N = len(self.docs)
        # BM25 idf with +0.5 smoothing; the +1.0 inside the log keeps it
        # non-negative even for terms appearing in more than half the corpus.
        self.idf = {t: math.log((N - n + 0.5) / (n + 0.5) + 1.0)
                    for t, n in df.items()}
        self.tf = [Counter(d) for d in self.docs]

    def score(self, query_terms: list[str], i: int) -> float:
        s, dl = 0.0, self.doc_len[i]
        for t in query_terms:
            if t not in self.idf:
                continue
            f = self.tf[i][t]
            if f == 0:
                continue
            denom = f + self.k1 * (1 - self.b + self.b * dl / self.avgdl)
            s += self.idf[t] * (f * (self.k1 + 1)) / denom
        return s

    def search(self, query: str, k: int = 3) -> list[tuple[Passage, float]]:
        """Top-k by BM25. NOTE: hits with score <= 0 are dropped, so this can
        return FEWER than k passages -- including zero, which the environment
        renders as the observation 'NoResults'. That is deliberate: a small
        model must be trained on what a failed search actually looks like."""
        q = _tokenize(query)
        scored = [(self.passages[i], self.score(q, i)) for i in range(len(self.docs))]
        scored.sort(key=lambda x: x[1], reverse=True)
        return [(p, sc) for p, sc in scored[:k] if sc > 0.0]


class HashEmbedRetriever:
    """'Embedding-lite': hash each token into a fixed-width vector via the
    hashing trick (Weinberger et al., 2009), L2-normalise, cosine-rank.
    No trained encoder, no network -- a stand-in for a real dense retriever
    (Ch. 9.1) that is still meaningfully robust to paraphrase."""

    def __init__(self, passages: list[Passage], dim: int = 512):
        self.passages, self.dim = passages, dim
        self.vecs = [self._embed(p.text) for p in passages]

    def _embed(self, text: str) -> list[float]:
        v = [0.0] * self.dim
        for t in _tokenize(text):
            h = int(hashlib.md5(t.encode()).hexdigest(), 16)
            # Signed hashing needs the sign bit to be INDEPENDENT of the bucket,
            # or collisions stop cancelling. Taking the sign from a low bit of
            # `h` is a real bug when `dim` is a power of two: `h % dim` is
            # exactly the low bits, so every token in a bucket gets the same
            # sign and signed hashing silently degrades to plain counting.
            # Splitting h into quotient and remainder keeps the two independent.
            idx, rest = h % self.dim, h // self.dim
            v[idx] += 1.0 if rest & 1 else -1.0
        norm = math.sqrt(sum(x * x for x in v)) or 1.0
        return [x / norm for x in v]

    def search(self, query: str, k: int = 3) -> list[tuple[Passage, float]]:
        qv = self._embed(query)
        scored = [(self.passages[i], sum(a * b for a, b in zip(qv, self.vecs[i])))
                  for i in range(len(self.passages))]
        scored.sort(key=lambda x: x[1], reverse=True)
        return [(p, sc) for p, sc in scored[:k] if sc > 0.0]

# capstone/stacklm/agent/retriever_dense.py  (NOT imported by CI: needs network + deps)
"""A real dense retriever with the same interface as HashEmbedRetriever.
pip install sentence-transformers faiss-cpu"""

from __future__ import annotations
from stacklm.agent.tools import Passage


class DenseRetriever:
    def __init__(self, passages: list[Passage],
                 model_name: str = "google/embeddinggemma-300m"):
        from sentence_transformers import SentenceTransformer
        import faiss
        self.passages = passages
        self.encoder = SentenceTransformer(model_name)
        # normalize_embeddings=True + inner product == cosine similarity.
        emb = self.encoder.encode([p.text for p in passages],
                                  normalize_embeddings=True,
                                  convert_to_numpy=True).astype("float32")
        self.index = faiss.IndexFlatIP(emb.shape[1])   # exact search; 268 docs
        self.index.add(emb)

    def search(self, query: str, k: int = 3) -> list[tuple[Passage, float]]:
        q = self.encoder.encode([query], normalize_embeddings=True,
                                convert_to_numpy=True).astype("float32")
        scores, idx = self.index.search(q, k)
        return [(self.passages[i], float(s))
                for i, s in zip(idx[0], scores[0]) if i >= 0]

<|system|>You are a narrow research assistant. Think, then either call a
tool or give a final answer. Tools: search(query,k), calc(expr).<|end|>
<|user|>What is the publication year of RoFormer, multiplied by 2? Use the corpus.<|end|>
<|assistant|>Thought: I need the publication year of RoFormer. I'll search the corpus.
<|tool_call|>{"tool": "search", "args": {"query": "RoFormer publication year", "k": 2}}<|end|>
<|tool_result|>[1] Reference note 16. In the Stack-100M project notes, the publication year of RoFormer is 2021. ... [2] ...<|end|>
<|assistant|>Thought: Now I compute 2021*2 with the calculator.
<|tool_call|>{"tool": "calc", "args": {"expr": "2021*2"}}<|end|>
<|tool_result|>4042<|end|>
<|assistant|>Thought: I have everything I need.
Answer: 4042<|end|>

# capstone/stacklm/agent/react.py
"""ReAct wire-format constants, (de)serialisation, and the single tokenizer
entry point. The special-token strings and the non-standard JSON spacing in
`render_call` MUST stay byte-identical across distillation, SFT, and serving
or the model's memorized groove breaks."""

from __future__ import annotations
import json, re
from dataclasses import dataclass

# The reserved special tokens (must match stacklm/tokenizer specials, Ch. 14.3).
BOS, EOS, PAD, END = "<|bos|>", "<|eos|>", "<|pad|>", "<|end|>"
SYS, USER, ASST = "<|system|>", "<|user|>", "<|assistant|>"
TOOL_CALL, TOOL_RESULT = "<|tool_call|>", "<|tool_result|>"

# Passed to tok.encode(..., allowed_special=ALL_SPECIAL) EVERYWHERE we encode a
# rendered transcript. Forgetting this is the bug the pitfall box warns about:
# without it the BPE shreds "<|tool_call|>" into a dozen ordinary tokens.
ALL_SPECIAL = frozenset({BOS, EOS, PAD, END, SYS, USER, ASST,
                         TOOL_CALL, TOOL_RESULT})

ANSWER_RE = re.compile(r"Answer:\s*(.+?)\s*$", re.DOTALL)


def enc(tok, text: str) -> list[int]:
    """The ONLY way this package tokenizes wire-format text."""
    return tok.encode(text, allowed_special=ALL_SPECIAL)


@dataclass
class Action:
    kind: str                    # "tool" or "final"
    thought: str = ""
    tool: str | None = None
    args: dict | None = None
    answer: str | None = None


def parse_assistant_step(text: str) -> Action:
    """Parse ONE assistant emission (thought + optional tool call OR answer).
    `text` is what the model generated up to (and excluding) <|end|>."""
    thought = ""
    m = re.search(r"Thought:\s*(.*?)(?:\n|$)", text)
    if m:
        thought = m.group(1).strip()

    if TOOL_CALL in text:
        payload = text.split(TOOL_CALL, 1)[1].split(END, 1)[0].strip()
        try:
            obj = json.loads(payload)
            return Action("tool", thought, obj["tool"], obj.get("args", {}))
        except Exception:
            # Malformed call: surface it so the harness returns an error obs.
            return Action("tool", thought, tool="__malformed__", args={"raw": payload})

    m = ANSWER_RE.search(text)
    if m:
        return Action("final", thought, answer=m.group(1).strip())

    # No tool call, no answer: the classic small-model derailment. Treat the
    # whole emission as the (probably wrong) final answer so the loop halts.
    return Action("final", thought, answer=text.strip())


def render_tool_result(obs: str) -> str:
    return f"{TOOL_RESULT}{obs}{END}"


def render_call(tool: str, args: dict) -> str:
    # Canonical key order (tool, args) and canonical spacing -> ONE grammar for
    # the model to learn, and one grammar for the acceptor below to enforce.
    body = json.dumps({"tool": tool, "args": args}, separators=(", ", ": "))
    return f"{TOOL_CALL}{body}{END}"

call        ::= '{"tool": "' toolbody '}}'
toolbody    ::= 'search", "args": {"query": "' query '", "k": ' digit
              | 'calc", "args": {"expr": "' expr '"'
query       ::= [A-Za-z0-9 ._,()-]{1,64}
expr        ::= [0-9+\-*/.() ]{1,40}
digit       ::= [1-9]

# Illustrative: the tool-call schema you would hand a served model.
# vLLM has renamed this parameter across releases (guided_json ->
# guided_decoding -> structured_outputs); the semantics have been stable.
# Check the version you install.
TOOL_CALL_SCHEMA = {
    "type": "object",
    "properties": {
        "tool": {"type": "string", "enum": ["search", "calc"]},
        "args": {"type": "object"},
    },
    "required": ["tool", "args"],
}

# capstone/stacklm/agent/grammar.py
"""A character-level acceptor for the ONE JSON shape our agent may emit, and a
token-level logit mask built from it. This is a miniature, readable version of
what XGrammar / Outlines / llguidance do inside vLLM and SGLang."""

from __future__ import annotations
from dataclasses import dataclass

_DIGITS = frozenset("123456789")
_QUERY  = frozenset("abcdefghijklmnopqrstuvwxyz"
                    "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 ._,()-")
_EXPR   = frozenset("0123456789+-*/.() ")


@dataclass(frozen=True)
class Hole:
    chars: frozenset
    min_len: int
    max_len: int


# A template is a tuple of literal strings and Holes, and matches `render_call`
# byte for byte (that is what the CI test below checks). Both templates share
# the prefix '{"tool": "' and diverge at 's' vs 'c'; the acceptor tracks a SET
# of live states, so the branch costs nothing and needs no backtracking.
SEARCH_T = ('{"tool": "search", "args": {"query": "', Hole(_QUERY, 1, 64),
            '", "k": ', Hole(_DIGITS, 1, 1), '}}')
CALC_T   = ('{"tool": "calc", "args": {"expr": "', Hole(_EXPR, 1, 40), '"}}')
TEMPLATES = (SEARCH_T, CALC_T)

# A state is (template_index, item_index, chars_consumed_within_that_item).
State = tuple


def start_states() -> frozenset:
    return frozenset((ti, 0, 0) for ti in range(len(TEMPLATES)))


def _succ(state: State, ch: str):
    """Yield every state reachable from `state` by consuming character `ch`."""
    ti, ii, n = state
    tmpl = TEMPLATES[ti]
    if ii >= len(tmpl):                       # already complete: nothing follows
        return
    item = tmpl[ii]
    if isinstance(item, str):                 # literal: exactly one legal char
        if item[n] == ch:
            yield (ti, ii, n + 1) if n + 1 < len(item) else (ti, ii + 1, 0)
        return
    # a Hole: either extend it...
    if ch in item.chars and n < item.max_len:
        yield (ti, ii, n + 1)
    # ...or (if long enough) close it and hand `ch` to the next item.
    if n >= item.min_len:
        yield from _succ((ti, ii + 1, 0), ch)


def step(states: frozenset, ch: str) -> frozenset:
    out: set = set()
    for s in states:
        out.update(_succ(s, ch))
    return frozenset(out)


def accepts(states: frozenset) -> bool:
    """True iff some live state sits at the end of its template -- i.e. the
    JSON object is complete and <|end|> is now legal."""
    return any(ii >= len(TEMPLATES[ti]) for ti, ii, _ in states)


# ---- lifting the CHARACTER acceptor to a TOKEN mask ---------------------
_MASK_CACHE: dict[frozenset, tuple[list[int], dict[int, frozenset]]] = {}


def build_vocab_strings(tok) -> list[str]:
    """Decode every vocabulary id once. Special tokens decode to their literal
    '<|...|>' form, whose '<' is illegal in every grammar position, so they are
    excluded automatically -- except <|end|>, which the decoder re-admits by
    hand exactly when accepts(states) is True."""
    return [tok.decode([i]) for i in range(tok.vocab_size)]


def token_transitions(states: frozenset, vocab_strings: list[str]):
    """Return (allowed_ids, next_states). A token is allowed iff feeding its
    decoded characters one at a time leaves at least one state alive.

    Cost: O(V * avg_token_len) per DISTINCT state set -- ~10^5 char steps at
    V=32768. We memoize on the state set, which collapses the cost to near zero
    in practice because the same states recur constantly. THIS is the
    inefficiency XGrammar removes properly, by precompiling the vocabulary into
    a trie and caching per-state adaptive masks."""
    hit = _MASK_CACHE.get(states)
    if hit is not None:
        return hit
    allowed, nxt = [], {}
    for tid, s in enumerate(vocab_strings):
        if not s:
            continue
        cur = states
        for ch in s:
            cur = step(cur, ch)
            if not cur:
                break
        if cur:
            allowed.append(tid)
            nxt[tid] = cur
    _MASK_CACHE[states] = (allowed, nxt)
    return allowed, nxt

{
  "name": "search",
  "description": "Search the local corpus and return the top-k passages.",
  "inputSchema": {
    "type": "object",
    "properties": {
      "query": {"type": "string", "description": "Keyword query."},
      "k": {"type": "integer", "minimum": 1, "maximum": 5, "default": 2}
    },
    "required": ["query"]
  }
}

# capstone/stacklm/agent/mcp_adapter.py
"""Render an MCP tool list into (a) the terse system-prompt tool line the 100M
model was distilled on, and (b) the JSON Schemas a constrained decoder wants."""

def tools_line(mcp_tools: list[dict]) -> str:
    """[{'name': 'search', 'inputSchema': {...}}, ...] -> 'search(query,k), calc(expr)'"""
    parts = []
    for t in mcp_tools:
        props = t.get("inputSchema", {}).get("properties", {})
        parts.append(f"{t['name']}({','.join(props)})")
    return ", ".join(parts)


def schemas(mcp_tools: list[dict]) -> dict[str, dict]:
    return {t["name"]: t.get("inputSchema", {}) for t in mcp_tools}

# Illustrative ONLY -- the real teacher call. In CI this path is never taken.
import anthropic  # or any strong-model client

def teacher_step(client, messages, tools_desc):
    resp = client.messages.create(
        model="a-strong-agentic-model",       # the large teacher
        system=tools_desc,                    # tool schemas + 2 gold exemplars
        messages=messages,
        max_tokens=512, temperature=0.7,      # some temperature -> diverse traces
    )
    return resp.content[0].text

# capstone/stacklm/agent/distill.py
"""Generate ReAct trajectories, filter to successes, export SFT data.
The teacher is an INJECTED callable (prompt -> one assistant step) so CI can
pass a hermetic stub while the full run passes a real large-model client. No
network import lives anywhere on this path."""

from __future__ import annotations
from stacklm.agent.react import (
    Action, parse_assistant_step, render_tool_result,
    SYS, USER, ASST, END,
)
from stacklm.agent.tools import calc, BM25Retriever, Passage
from stacklm.agent.corpus import Task          # Task lives with the pool


class ToolEnv:
    """The environment: dispatch a parsed Action to a real tool."""

    def __init__(self, corpus: list[Passage]):
        self.retriever = BM25Retriever(corpus)

    def run_tool(self, act: Action) -> str:
        if act.tool == "calc":
            return calc(str(act.args.get("expr", "")))
        if act.tool == "search":
            hits = self.retriever.search(act.args.get("query", ""),
                                         int(act.args.get("k", 2)))
            if not hits:
                return "NoResults"
            return " ".join(f"[{i+1}] {p.text}" for i, (p, _) in enumerate(hits))
        return f"ToolError: unknown or malformed tool '{act.tool}'"


SYSTEM_PROMPT = ("You are a narrow research assistant. Think, then either "
                 "call a tool or give a final answer. Tools: search(query,k), "
                 "calc(expr).")


def normalize(s: str | None) -> str:
    """Verifier-side normalisation so '4042', ' 4042 ', '4042.0' all match."""
    if s is None:
        return ""
    s = s.strip().rstrip(".")
    try:
        f = float(s)
        return str(int(f)) if f.is_integer() else str(f)
    except ValueError:
        return " ".join(s.lower().split())


def rollout(task: Task, teacher, env: ToolEnv, max_steps: int = 6):
    """Run ONE ReAct trajectory with `teacher`. Returns (transcript, solved).
    The transcript is the full wire format including observations."""
    transcript = f"{SYS}{SYSTEM_PROMPT}{END}{USER}{task.question}{END}"
    solved = False
    for _ in range(max_steps):
        gen = teacher(transcript + ASST)           # teacher emits one step
        step_text = gen.split(END, 1)[0]           # up to the first <|end|>
        act = parse_assistant_step(step_text)
        transcript += f"{ASST}{step_text}{END}"    # record it verbatim
        if act.kind == "final":
            solved = (normalize(act.answer) == normalize(task.gold))
            break
        obs = env.run_tool(act)                    # REAL tool call
        transcript += render_tool_result(obs)      # inject observation
    return transcript, solved


def distill(tasks, teacher, env, samples_per_task: int = 4):
    """Rejection sampling: multiple rollouts per task, keep only solved ones,
    deduplicated per task."""
    kept = []
    for task in tasks:
        seen = set()
        for _ in range(samples_per_task):
            transcript, solved = rollout(task, teacher, env)
            if solved and transcript not in seen:
                seen.add(transcript)
                kept.append({"task": task.question, "text": transcript,
                             "kind": task.kind})
    return kept

# capstone/stacklm/agent/stub_teacher.py
"""A deterministic 'teacher' that solves the generated auto-research tasks by a
fixed policy driven by Task.hints (the answer key). It emits exactly the wire
format a real large model would emit, but offline and in microseconds.

It is a TEST DOUBLE: it cheats by construction. Its only contract is that every
downstream code path -- dispatch, observation rendering, the success filter,
the SFT exporter -- is exercised end to end on every CI run."""

from __future__ import annotations
import re
from stacklm.agent.react import TOOL_RESULT, END, USER, render_call


def _observations(prompt: str) -> list[str]:
    """Every tool result seen so far, in order."""
    return [chunk.split(END, 1)[0]
            for chunk in prompt.split(TOOL_RESULT)[1:]]


def _extract(observations: list[str], ent: str, attr: str) -> str | None:
    """Pull a value out of an observation using the corpus's canonical clause
    'the {attr} of {ent} is {value}'. Returns None if the search missed --
    which is exactly what makes the rejection filter non-trivial.

    Note the character class EXCLUDES the comma: one of the three passage
    templates continues the sentence with ', a choice made during...', and a
    greedy '[^.]+' would swallow that whole clause into the 'value'. Verifiers
    fail silently in exactly this way; write the test that catches it."""
    pat = re.compile(rf"the {re.escape(attr)} of {re.escape(ent)} is ([^.,;]+)[.,;]")
    for obs in observations:
        m = pat.search(obs)
        if m:
            return m.group(1).strip()
    return None


def make_stub_teacher(tasks):
    by_q = {t.question: t for t in tasks}

    def teacher(prompt: str) -> str:
        question = prompt.split(USER, 1)[1].split(END, 1)[0]
        task = by_q[question]                      # the cheat: the answer key
        obs = _observations(prompt)
        queries, facts = task.hints["queries"], task.hints["facts"]
        op = task.hints["op"]

        # PHASE 1 -- one search per required fact, in order.
        if len(obs) < len(queries):
            i = len(obs)
            return (f"Thought: I need the {facts[i][1]} of {facts[i][0]}. "
                    f"I'll search the corpus.\n"
                    f"{render_call('search', {'query': queries[i], 'k': 2})}")

        vals = [_extract(obs, e, a) for e, a in facts]
        if any(v is None for v in vals):           # retrieval genuinely failed
            return "Thought: The corpus did not give me the fact.\nAnswer: unknown"

        # PHASE 2 -- one calculator call, if the task shape needs arithmetic.
        n_calc = len(obs) - len(queries)           # observations past the searches
        if n_calc == 0 and op is not None:
            if op == "double":
                expr = f"{vals[0]}*2"
            elif op == "sum":
                expr = f"{vals[0]}+{vals[1]}"
            else:                                  # "compare": sign of difference
                expr = f"{vals[0]}-{vals[1]}"
            return (f"Thought: Now I compute {expr} with the calculator.\n"
                    f"{render_call('calc', {'expr': expr})}")

        # PHASE 3 -- synthesize.
        if op is None:
            answer = vals[0]
        elif op == "compare":
            diff = obs[-1].strip()
            names = task.hints["names"]
            answer = names[0] if not diff.startswith("-") else names[1]
        else:
            answer = obs[-1].strip()               # the calculator's output
        return f"Thought: I have everything I need.\nAnswer: {answer}"

    return teacher

# capstone/tests/test_agent_pipeline.py
"""Hermetic, CPU-only, no network. Runs in a couple of seconds.
`tok` is the session-scoped toy StackTokenizer fixture from conftest.py."""

from stacklm.agent.corpus import build_corpus, make_task_pool
from stacklm.agent.distill import ToolEnv, distill, normalize
from stacklm.agent.stub_teacher import make_stub_teacher
from stacklm.agent.sft_format import build_example, _segment, IGNORE
from stacklm.agent.react import TOOL_CALL, TOOL_RESULT, END, render_call, enc
from stacklm.agent import grammar as G


def test_distill_yields_correct_transcripts():
    corpus = build_corpus()
    train, held = make_task_pool()
    assert len(corpus) == 268 and len(train) == 200 and len(held) == 50
    env = ToolEnv(corpus)
    teacher = make_stub_teacher(train)
    kept = distill(train[:40], teacher, env, samples_per_task=2)

    # (1) The pipeline produces data at all -- the bug this test exists for.
    #     The stub reads the answer key, so its yield is ~100%; a real teacher
    #     lands nearer 0.6. Assert on the YIELD, not merely on "no exception":
    #     a stub that silently solves nothing returns [] and every downstream
    #     stage becomes a no-op while CI stays green.
    assert len(kept) >= 35, f"teacher yield collapsed: {len(kept)}/40"
    # (2) Every kept transcript really ends in the gold answer.
    gold = {t.question: t.gold for t in train}
    for row in kept:
        final = row["text"].rsplit("Answer:", 1)[1].split(END, 1)[0]
        assert normalize(final) == normalize(gold[row["task"]])
    # (3) All four task shapes survive the filter (no shape is silently lost).
    assert {r["kind"] for r in kept} == {"lookup", "double", "sum", "compare"}


def test_special_tokens_are_single_ids_and_masking_is_nonvacuous(tok):
    # The pitfall box's assertion, as an actual test.
    for s in (TOOL_CALL, TOOL_RESULT, END):
        assert len(enc(tok, s)) == 1, f"{s} is not a single special id"
    corpus = build_corpus()
    train, _ = make_task_pool()
    env, teacher = ToolEnv(corpus), make_stub_teacher(train)
    kept = distill(train[:4], teacher, env, samples_per_task=1)
    # _segment must PARTITION the transcript -- no dropped and no duplicated
    # spans. The bug this catches: a tail span appended twice when a transcript
    # does not end in a marker, which silently doubles those tokens in the SFT
    # example AND in the RL generated-token mask.
    segs = _segment(kept[0]["text"])
    assert "".join(s for s, _ in segs) == kept[0]["text"]
    ids, labels = build_example(kept[0]["text"], tok)
    n_sup = int((labels != IGNORE).sum())
    assert 0 < n_sup < len(labels), "loss mask is empty or masks nothing"


def test_grammar_accepts_what_we_render_and_rejects_what_we_fear():
    # Every search query the task pool can produce, plus both calc shapes.
    train, _ = make_task_pool()
    bodies = [render_call("calc", {"expr": "2021*2"}),
              render_call("calc", {"expr": "2019+10000"})]
    bodies += [render_call("search", {"query": q, "k": 2})
               for t in train for q in t.hints["queries"]]
    for body in bodies:
        json_part = body.split(TOOL_CALL, 1)[1].split(END, 1)[0]
        st = G.start_states()
        for ch in json_part:
            st = G.step(st, ch)
            assert st, f"grammar rejected our own output at {ch!r}"
        assert G.accepts(st)
    # A single dropped quote must be rejected, and rejected EARLY.
    st = G.start_states()
    for ch in '{"tool": search':
        st = G.step(st, ch)
    assert not st

# capstone/stacklm/agent/sft_format.py
"""Turn a distilled transcript into (input_ids, labels) with the correct loss
mask: supervise ASSISTANT emissions (thoughts, tool CALLS, final answer) only;
mask the system/user prompt and every tool RESULT span."""

from __future__ import annotations
import numpy as np
from stacklm.agent.react import ASST, END, TOOL_RESULT, USER, SYS, enc

IGNORE = -100   # standard CrossEntropyLoss ignore_index (matches Ch. 14.9)


def _segment(t: str):
    """Split a transcript into (text, supervised) spans that PARTITION it
    exactly -- concatenating the pieces must reproduce `t` byte for byte.
    Assistant content between <|assistant|> and the NEXT role/result marker is
    supervised; everything else (system, user, and <|tool_result|>...<|end|>)
    is masked. <|tool_call|> is deliberately NOT a segmentation marker, so tool
    calls stay INSIDE the supervised span -- we very much want them emitted.

    This function is load-bearing twice over: it defines the SFT loss mask
    here, and it defines the RL generated-token mask in `rlvr.py`."""
    out, i, supervised = [], 0, False
    markers = [ASST, USER, SYS, TOOL_RESULT, END]
    while i < len(t):
        j, m = min([(t.find(x, i), x) for x in markers if t.find(x, i) != -1],
                   default=(len(t), None), key=lambda z: z[0])
        if j > i:
            out.append((t[i:j], supervised))
        if m is None:
            break            # trailing text, already appended -- do NOT re-add
        if m == ASST:
            supervised = True                 # begin supervising assistant span
            out.append((m, False))            # the marker token itself: masked
        elif m in (USER, SYS, TOOL_RESULT):
            supervised = False                # stop supervising
            out.append((m, False))
        else:                                 # END
            out.append((m, supervised))       # <|end|> IS supervised if in asst
            supervised = False                # ...then close the span
        i = j + len(m)
    return out


def build_example(transcript: str, tok, max_len: int = 1024):
    """-> (input_ids, target) or None if the transcript does not fit.

    `tok` is the Stack-100M byte-level BPE tokenizer (Ch. 14.3). We encode
    through `enc`, which passes allowed_special so every '<|...|>' marker maps
    to its single reserved id instead of being shredded into byte pieces.

    We DROP over-length transcripts rather than truncating them. `ids[:max_len]`
    looks harmless and is not: it cuts the tail off the longest, most
    compositional traces -- precisely the ones whose `Answer:` span lives at the
    end -- so the model is trained on trajectories that never terminate, and it
    learns not to terminate. Log the drop rate; if it is non-trivial, raise
    `max_len` (mid-training gave us 8192 tokens of context) rather than quietly
    mutilating your best data."""
    ids, labels = [], []
    for text, supervised in _segment(transcript):
        piece = enc(tok, text)
        ids.extend(piece)
        labels.extend(piece if supervised else [IGNORE] * len(piece))
    if len(ids) > max_len:
        return None
    # Teacher forcing: predict token t+1 from tokens <= t, so shift the labels.
    input_ids = np.array(ids[:-1], dtype=np.int64)
    target    = np.array(labels[1:], dtype=np.int64)
    return input_ids, target


def build_dataset(kept: list[dict], tok, max_len: int = 1024):
    """Export the whole distilled set, reporting what was dropped."""
    out = [ex for ex in (build_example(r["text"], tok, max_len) for r in kept)
           if ex is not None]
    dropped = len(kept) - len(out)
    if dropped:
        print(f"[sft_format] dropped {dropped}/{len(kept)} over-length traces "
              f"(max_len={max_len}); raise max_len if this is more than a few %")
    return out

# capstone/stacklm/agent/loop.py
"""The runtime auto-research loop: Stack-100M interleaves thought, tool calls,
and observations, then synthesizes a final grounded answer. Guards against the
known small-model failure modes, and carries token ids so the RL stage never
has to re-tokenize a rollout."""

from __future__ import annotations
from dataclasses import dataclass, field
import torch
from stacklm.agent.react import (parse_assistant_step, render_tool_result,
                                 SYS, USER, ASST, END, TOOL_CALL, enc)
from stacklm.agent.distill import ToolEnv, SYSTEM_PROMPT
from stacklm.agent import grammar as G


@dataclass
class Trace:
    """A finished rollout, in two views.

    `steps` is for humans, for logging, and for Ch. 14.11's failure taxonomy;
    it iterates as (role, text) pairs so `for role, text in trace` works.
    `ids` / `gen` are for RL: the exact token sequence the forward pass must
    see, and a 0/1 flag per token that is 1 iff the POLICY sampled it."""
    steps: list[tuple[str, str]] = field(default_factory=list)
    ids: list[int] = field(default_factory=list)
    gen: list[int] = field(default_factory=list)

    def __iter__(self):
        return iter(self.steps)

    def __len__(self):
        return len(self.steps)

    def add(self, ids: list[int], generated: int) -> None:
        self.ids.extend(ids)
        self.gen.extend([generated] * len(ids))


def run_agent(model, tok, question: str, env: ToolEnv,
              max_steps: int = 6, max_new: int = 160,
              temperature: float = 0.0, constrain: bool = True):
    """-> (answer, Trace).

    temperature=0.0 -> greedy (serve time). temperature>0 -> sampling, which is
    what the GRPO rollouts below REQUIRE: a greedy policy in a deterministic
    environment produces G identical trajectories, a zero-variance group, and
    therefore a zero gradient."""
    asst_id, end_id = tok.special_token_id(ASST), tok.special_token_id(END)
    tr = Trace()
    tr.add(enc(tok, f"{SYS}{SYSTEM_PROMPT}{END}{USER}{question}{END}"), 0)
    seen_calls = set()                                # loop guard

    for _ in range(max_steps):
        tr.add([asst_id], 0)                          # the HARNESS opens the turn
        new_ids, closed = generate(model, tok, tr.ids, max_new=max_new,
                                   temperature=temperature, constrain=constrain)
        tr.add(new_ids, 1)                            # <- the policy's own tokens
        # The harness always closes the turn so the wire format stays legal, but
        # the token counts as OURS only if the policy chose to emit it. If we
        # hit max_new, flagging it 1 would reinforce a token nobody sampled.
        tr.add([end_id], 1 if closed else 0)

        step_text = tok.decode(new_ids)
        tr.steps.append(("assistant", step_text))
        act = parse_assistant_step(step_text)
        if act.kind == "final":
            return act.answer, tr

        # --- GUARDS ---------------------------------------------------
        call_key = (act.tool, str(act.args))
        if call_key in seen_calls:                    # exact repeat -> break loop
            obs = "RepeatedCall: you already ran this; use the prior result."
        elif act.tool == "__malformed__":             # bad JSON -> teach recovery
            obs = "FormatError: emit a valid JSON tool call."
        else:
            obs = env.run_tool(act)
            seen_calls.add(call_key)
        # --------------------------------------------------------------
        tr.add(enc(tok, render_tool_result(obs)), 0)  # the ENVIRONMENT wrote this
        tr.steps.append(("observation", obs))

    # Out of steps: force one synthesis attempt. The forcing text is the
    # HARNESS's, not the policy's, so it is flagged 0 -- otherwise RL would
    # reinforce a string the model never chose to say.
    tr.add(enc(tok, f"{ASST}Thought: I must answer now.\nAnswer:"), 0)
    new_ids, closed = generate(model, tok, tr.ids, max_new=40,
                               temperature=temperature, constrain=False)
    tr.add(new_ids, 1)
    tr.add([end_id], 1 if closed else 0)
    answer = tok.decode(new_ids).strip()
    tr.steps.append(("assistant", f"Thought: I must answer now.\nAnswer: {answer}"))
    return answer, tr


@torch.no_grad()
def generate(model, tok, prompt_ids: list[int], max_new: int,
             temperature: float = 0.0, constrain: bool = True):
    """Decode ONE assistant emission. -> (sampled_ids, closed_with_end).

    Takes and returns IDS, never text: the caller keeps them verbatim.

    `constrain=True` turns on grammar-constrained decoding for the tool-call
    JSON: the constraint activates the moment the model commits to <|tool_call|>
    and releases when the JSON is complete. The choice between calling a tool
    and answering stays FREE -- we constrain syntax, never the decision."""
    stop_id = tok.special_token_id(END)
    call_id = tok.special_token_id(TOOL_CALL)
    vocab_strings = G.build_vocab_strings(tok) if constrain else None
    cap = model.cfg.max_seq_len                       # the RoPE cache length
    dev = next(model.parameters()).device

    ids = torch.tensor([prompt_ids[-cap:]], dtype=torch.long, device=dev)
    out, states, nxt = [], None, None
    for _ in range(max_new):
        logits, _ = model(ids[:, -cap:])              # Stack100M -> (logits, loss)
        logits = logits[:, -1, :].float()             # [1, vocab]

        if states is not None:                        # inside a tool call
            allowed, nxt = G.token_transitions(states, vocab_strings)
            mask = torch.full_like(logits, float("-inf"))
            mask[0, allowed] = 0.0
            if G.accepts(states):                     # JSON complete -> may stop
                mask[0, stop_id] = 0.0
            logits = logits + mask

        if temperature > 0.0:
            probs = torch.softmax(logits / temperature, dim=-1)
            tid = int(torch.multinomial(probs, 1))
        else:
            tid = int(logits.argmax(-1))              # greedy

        if tid == stop_id:                            # the model closed the step
            return out, True
        out.append(tid)
        ids = torch.cat([ids, torch.tensor([[tid]], device=dev)], dim=1)

        if states is not None:
            states = nxt.get(tid)                     # advance the grammar
        elif constrain and tid == call_id:
            states = G.start_states()                 # commit -> constrain
    return out, False                                 # ran out of budget

# capstone/stacklm/agent/roles.py
"""Split one distilled trace into role-specific SFT sets, so each role's
adapter sees ONLY the decisions it will be asked to make at serve time."""

from stacklm.agent.react import ASST, END


def split_roles(transcript: str):
    """-> (query_writer_examples, synthesizer_examples), each a list of
    (prefix, target) pairs in the same wire format. The prefix accumulates the
    FULL history including observations, so each role still conditions on
    everything that happened -- only the target differs."""
    qw, syn, prefix = [], [], ""
    for chunk in transcript.split(ASST):
        if not prefix:                       # the system+user preamble
            prefix = chunk
            continue
        target, rest = chunk.split(END, 1)   # rest = observation, if any
        role = qw if '"tool": "search"' in target else syn
        role.append((prefix + ASST, target + END))
        prefix += ASST + target + END + rest
    return qw, syn

# capstone/stacklm/agent/rlvr.py
"""Narrow multi-turn RLVR with GRPO. One 'rollout' = one full ReAct trajectory;
reward = final-answer exact match minus a step penalty. The policy-gradient
loss covers GENERATED tokens only -- observations condition the forward pass
but never enter the loss."""

from __future__ import annotations
import numpy as np, torch
from stacklm.agent.loop import run_agent
from stacklm.agent.distill import normalize
from stacklm.agent.react import PAD
from stacklm.post.grpo import token_logprobs          # Ch. 14.9, unchanged


def trajectory_reward(answer, gold, n_steps, max_steps, lam=0.05):
    correct = 1.0 if normalize(answer) == normalize(gold) else 0.0
    return correct - lam * (n_steps / max_steps)


def grpo_advantages(rewards):
    r = np.asarray(rewards, dtype=np.float64)
    # group-relative: subtract the group mean, scale by the group std (GRPO)
    return (r - r.mean()) / (r.std() + 1e-6)


def grpo_agent_step(policy, ref, opt, tok, env, task, *, group_size=8,
                    max_steps=6, temperature=0.7, clip_eps=0.2, kl_beta=0.02,
                    lam=0.05, device="cuda"):
    """ONE GRPO iteration on ONE task. Returns a stats dict."""
    # ---- 1. rollouts. SAMPLED, not greedy (see the aside above). ----------
    traces, rewards = [], []
    for _ in range(group_size):
        answer, tr = run_agent(policy, tok, task.question, env,
                               max_steps=max_steps, temperature=temperature)
        n_steps = sum(1 for role, _ in tr if role == "assistant")
        rewards.append(trajectory_reward(answer, task.gold, n_steps, max_steps, lam))
        traces.append(tr)

    adv_np = grpo_advantages(rewards)
    if float(np.abs(adv_np).max()) < 1e-6:
        # Degenerate group: every rollout scored the same. No gradient exists;
        # skipping is strictly better than taking a zero step (and it is the
        # signal to fix your task difficulty, not your learning rate).
        return {"reward": float(np.mean(rewards)), "skipped": True}

    # ---- 2. pad the rollouts. NO re-tokenization: the ids came from the
    #         sampler, and `tr.gen` came with them. ------------------------
    G, T = group_size, max(len(tr.ids) for tr in traces)
    pad_id = tok.special_token_id(PAD)
    seqs  = torch.full((G, T), pad_id, dtype=torch.long, device=device)
    gmask = torch.zeros((G, T), dtype=torch.float32, device=device)
    pmask = torch.zeros((G, T), dtype=torch.float32, device=device)  # real, non-pad
    for i, tr in enumerate(traces):
        n = len(tr.ids)
        seqs[i, :n]  = torch.tensor(tr.ids, dtype=torch.long, device=device)
        gmask[i, :n] = torch.tensor(tr.gen, dtype=torch.float32, device=device)
        pmask[i, :n] = 1.0
    # Right-padding is exact under CAUSAL attention: no real token can attend
    # to a pad, and every pad position is masked out of the loss below.
    adv = torch.tensor(adv_np, dtype=torch.float32, device=device)

    # ---- 3. old log-probs: the policy that GENERATED these rollouts -------
    with torch.no_grad():
        old_lp = token_logprobs(policy, seqs)          # (G, T-1)
        ref_lp = token_logprobs(ref, seqs)             # frozen KL anchor
    m    = gmask[:, 1:]                                # align to prediction targets
    real = pmask[:, 1:]                                # non-pad positions

    # ---- 4. one clipped-surrogate step over GENERATED tokens only --------
    new_lp = token_logprobs(policy, seqs)
    ratio  = torch.exp(new_lp - old_lp)                # == 1 on the first step
    a      = adv.unsqueeze(1)                          # (G,1) broadcasts over t
    surr   = torch.min(ratio * a,
                       torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps) * a)
    logr   = ref_lp - new_lp                           # k3 KL estimator (Schulman)
    kl     = torch.exp(logr) - logr - 1.0
    # Token-mean over the WHOLE batch: every generated token counts once,
    # regardless of which trajectory it came from (Dr. GRPO length correction).
    denom  = m.sum().clamp(min=1.0)
    loss   = -((surr - kl_beta * kl) * m).sum() / denom

    loss.backward()
    torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
    opt.step(); opt.zero_grad(set_to_none=True)
    return {"reward": float(np.mean(rewards)), "loss": float(loss),
            # `real_tokens`, NOT m.numel(): the latter counts right-padding, so
            # the ratio below would move with trajectory-length variance rather
            # than with the thing we want to monitor.
            "gen_tokens": int(m.sum()), "real_tokens": int(real.sum()),
            "skipped": False}