laya: tools/laya_finetune, the routing head's CPU fine-tune over --dump-encoder dumps, with before/after on the game (Refs #1313)

CMakeLists.txt

@@ -823,6 +823,12 @@ add_executable(laya_eval_test tests/laya_eval_test.cpp)
 target_link_libraries(laya_eval_test PRIVATE onebit_laya nlohmann_json::nlohmann_json)
 add_test(NAME laya_eval COMMAND laya_eval_test)
 add_dependencies(onebit laya_eval_test)
+# engine#1313 item 3: the head-only fine-tune over --dump-encoder dumps, on a synthetic checkpoint and dump (no model,
+# no device). It needs torch + numpy, which the engine build does not, so a box without them skips (77).
+add_python_test(laya_finetune ${CMAKE_SOURCE_DIR}/tests/laya_finetune_test.py)
+if(Python3_Interpreter_FOUND)
+    set_tests_properties(laya_finetune PROPERTIES SKIP_RETURN_CODE 77)
+endif()
 # Crafted Laya checkpoints the reader and the scorer must refuse (written at run time); runs in CI.
 add_executable(laya_bounds_test tests/laya_bounds_test.cpp)
 target_link_libraries(laya_bounds_test PRIVATE onebit_laya nlohmann_json::nlohmann_json)

docs/laya-models.md

@@ -113,8 +113,8 @@ numbers a demo HUD reads; every one was measured by serve or reported by the bac
   `args`. Linux only.
 - Laya's zero-shot answer to a new question is a starting point, not a measured router: the base
   checkpoints are near chance on decision types they were not tuned for (Laya's own README). The game
-  table's accuracy is measured below, and a fine-tune on the game's own logged decisions is the plan for
-  raising it (engine #1313).
+  table's accuracy is measured below, and a head-only fine-tune on the game's own logged decisions
+  (`tools/laya_finetune/`, below) raises its confidence on them (engine #1313).
 
 ## Measuring and fine-tuning the router
 
@@ -173,8 +173,51 @@ Measured on this box (`typed-decisions/`, the 20 game steps of the measurement b
 | a torch copy of the head over the dump | reproduces the engine's probabilities within 2.4e-7 (logits 1.3e-6) on all 40 |
 | option order | mean \|p(reflex) in one order - the other\| 0.065 (max 0.171); 1 of 20 decisions flips with the order |
 
+### Fine-tuning the head (`tools/laya_finetune/`)
+
+```sh
+# 1. labelled cases -> the head's input, every option order (see the join above for labels)
+1bit route --laya-models table.json --laya-model DIR --dump-encoder train.jsonl --out dump-train/ --rotate
+1bit route --laya-models table.json --laya-model DIR --dump-encoder dev.jsonl   --out dump-dev/   --rotate
+# 2. head-only fine-tune on the CPU (torch + numpy): the encoder is frozen, its output is in the dump
+tools/laya_finetune/train_head.py --model DIR --train dump-train --dev dump-dev --out NEW_DIR [--epochs 4] [--lr 2e-5]
+# 3. before / after on held-out cases, by the path serve takes, and the same numbers from the per-case files
+1bit route --laya-models table.json --laya-model DIR     --eval test.jsonl --per-case before.jsonl
+1bit route --laya-models table.json --laya-model NEW_DIR --eval test.jsonl --per-case after.jsonl
+tools/laya_finetune/metrics.py before.jsonl after.jsonl
+```
+
+`train_head.py` trains `type_emb`, the two head layers and `scorer` (`head.py` is the head of `laya/scorer.cpp` in
+torch) on Laya's proper reward (log score + 0.5 spherical score), class-weighted, with early stopping on dev log loss.
+It then refits the `choice:k` temperature on dev, and keeps the checkpoint's own temperature when the refit lands on
+the grid's edge: a separable dev set says nothing about how sure to be. `NEW_DIR` is a copy of `DIR` with only the head
+tensors' bytes replaced (same names, dtypes, shapes and offsets), the temperature in `rl_agent_config.json`, and
+`finetune.json` (data, seed, epochs, dev before and after, command line). It refuses to overwrite an existing `NEW_DIR`.
+Serve takes it unchanged: `--laya-model NEW_DIR`.
+
+Measured on the game (engine #1313, the private `laya-plays` harness): labels are the step kind each logged decision
+was made for. Episodes were held out whole: `test` = the hand-labelled steps, `test2` and `test3` = later episodes
+the fine-tune never saw.
+
+| set | rows | correct, base / fine-tuned | log loss, base / fine-tuned |
+|---|---|---|---|
+| test | 39 | 33 / 33 | 0.543 / 1.395 |
+| test2 | 110 | 107 / 106 | 0.478 / 0.294 |
+| test3 | 56 | 55 / 55 | 0.475 / 0.188 |
+
+The fine-tune keeps accuracy and makes the router much surer on the game's own steps (log loss 0.48 -> 0.19-0.29). On the
+hand-written `test` set it is overconfident on its misses, so read the table row by row, not as one number. A
+second label tried in the harness, "plan" for a reflex step whose game action failed, taught the head nothing: 0 of 10
+held-out failed-reflex rows moved to plan, because those failures come from the plan and are not visible in the
+reflex step's own state. That label needs a planner's view of the same state, not the reflex prompt.
+
 ## Tests
 
+`tests/laya_finetune_test.py` (ctest `laya_finetune`, no checkpoint; skips without torch): `train_head.py` on a
+synthetic checkpoint (the head's real tensor names at d 32) and a dump whose label sits on the right option's marker
+row. Dev log loss must fall and accuracy must not; every non-head byte of the new `model.safetensors` must be
+identical, the head must change, the temperature must land in its bucket, an existing `--out` must be refused; and
+`metrics.py` must agree with a hand-computed `--per-case` file.
 `tests/laya_models_test.cpp` (ctest `laya_models`, no checkpoint): parsing, the question, the pick and
 its fallback, aliases, shared backends, the size gate, the refusals, and the shipped game table.
 `tests/laya_eval_test.cpp` (ctest `laya_eval`, no checkpoint): the eval's arithmetic on hand-computed cases.

tests/laya_finetune_test.py

@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+# Copyright 2026 bong-water-water-bong
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""tools/laya_finetune on synthetic data (engine #1313): no checkpoint, no game, no device.
+
+A tiny checkpoint with the real head's tensor names (d 32, the 16-head layout of laya/scorer.cpp) and one encoder
+tensor the fine-tune must not touch, plus `--dump-encoder`-shaped dumps whose label is carried by the marker row of
+the right option. Checked:
+  1. train_head.py lowers dev log loss and does not lower dev accuracy;
+  2. every byte of the new model.safetensors outside the head tensors is identical (header, encoder tensor), the
+     head tensors changed, and the temperature landed in rl_agent_config.json's bucket;
+  3. it refuses to overwrite an existing --out;
+  4. metrics.py on a hand-computed --per-case file.
+Exit 77 (ctest skip) when torch or numpy is missing: the tool needs them, the engine build does not.
+"""
+import json
+import math
+import os
+import shutil
+import struct
+import subprocess
+import sys
+import tempfile
+
+try:
+    import numpy as np
+    import torch  # noqa: F401
+except ImportError:
+    print("SKIP: torch/numpy not installed")
+    sys.exit(77)
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+TOOL = os.path.join(ROOT, "tools", "laya_finetune")
+sys.path.insert(0, TOOL)
+import metrics  # noqa: E402
+
+D, F_, DH = 32, 64, 32
+fails = 0
+
+
+def check(ok, what):
+    global fails
+    print(("ok   " if ok else "FAIL ") + what)
+    fails += 0 if ok else 1
+
+
+def write_safetensors(path, tensors):
+    header, blobs, off = {}, [], 0
+    for name, a in tensors.items():
+        raw = a.astype(np.float32).tobytes()
+        header[name] = {"dtype": "F32", "shape": list(a.shape), "data_offsets": [off, off + len(raw)]}
+        blobs.append(raw)
+        off += len(raw)
+    h = json.dumps(header).encode()
+    h += b" " * (-len(h) % 8)
+    with open(path, "wb") as f:
+        f.write(struct.pack("<Q", len(h)) + h + b"".join(blobs))
+
+
+def make_checkpoint(d, rng):
+    t = {"encoder.final_norm.weight": rng.standard_normal(D), "type_emb.weight": 0.02 * rng.standard_normal((3, D))}
+    for i in range(2):
+        p = f"head.layers.{i}."
+        t[p + "norm1.weight"], t[p + "norm1.bias"] = np.ones(D), np.zeros(D)
+        t[p + "self_attn.in_proj_weight"] = 0.05 * rng.standard_normal((3 * D, D))
+        t[p + "self_attn.in_proj_bias"] = np.zeros(3 * D)
+        t[p + "self_attn.out_proj.weight"] = 0.05 * rng.standard_normal((D, D))
+        t[p + "self_attn.out_proj.bias"] = np.zeros(D)
+        t[p + "norm2.weight"], t[p + "norm2.bias"] = np.ones(D), np.zeros(D)
+        t[p + "linear1.weight"], t[p + "linear1.bias"] = 0.05 * rng.standard_normal((F_, D)), np.zeros(F_)
+        t[p + "linear2.weight"], t[p + "linear2.bias"] = 0.05 * rng.standard_normal((D, F_)), np.zeros(D)
+    t["scorer.0.weight"], t["scorer.0.bias"] = np.ones(D), np.zeros(D)
+    t["scorer.1.weight"], t["scorer.1.bias"] = 0.1 * rng.standard_normal((DH, D)), np.zeros(DH)
+    t["scorer.3.weight"], t["scorer.3.bias"] = 0.1 * rng.standard_normal((1, DH)), np.zeros(1)
+    t["temperature"] = np.ones(3)
+    os.makedirs(d)
+    write_safetensors(f"{d}/model.safetensors", t)
+    json.dump({"temperature_by_options": {}}, open(f"{d}/rl_agent_config.json", "w"))
+
+
+def make_dump(d, n, rng):
+    """n states, each written in both option orders (the --rotate shape); the right option's marker row carries +3 on
+    feature 0 and the other one -3, plus noise, so the class is learnable from the head's input alone."""
+    os.makedirs(d)
+    rows, cases, L = [], [], 6
+    for i in range(n):
+        y = i % 2  # 0 = reflex, 1 = plan
+        for rot in (0, 1):
+            options = ["reflex", "plan"] if rot == 0 else ["plan", "reflex"]
+            label = options.index(["reflex", "plan"][y])
+            h = 0.5 * rng.standard_normal((L, D))
+            markers = [2, 4]
+            for k, m in enumerate(markers):
+                h[m, 0] += 3.0 if k == label else -3.0
+            cases.append({"src": f"s{i}", "rotation": rot, "label": label, "options": options, "qtype": 0,
+                          "offset": len(rows) * L, "len": L, "markers": markers})
+            rows.append(h)
+    np.concatenate(rows).astype(np.float32).tofile(f"{d}/h.f32")
+    with open(f"{d}/cases.jsonl", "w") as f:
+        for c in cases:
+            f.write(json.dumps(c) + "\n")
+    json.dump({"d": D, "rows": len(rows) * L, "cases": len(cases), "rotate": True}, open(f"{d}/dump.json", "w"))
+
+
+def tensor_ranges(path):
+    with open(path, "rb") as f:
+        n = struct.unpack("<Q", f.read(8))[0]
+        return 8 + n, json.loads(f.read(n))
+
+
+def main():
+    work = tempfile.mkdtemp(prefix="laya-finetune-", dir=os.environ.get("ONEBIT_TEST_SCRATCH")
+                            or os.path.expanduser("~/.cache/scratch-tests"))
+    try:
+        rng = np.random.default_rng(1313)
+        ck, tr, dv, out = (os.path.join(work, x) for x in ("ckpt", "train", "dev", "out"))
+        make_checkpoint(ck, rng)
+        make_dump(tr, 24, rng)
+        make_dump(dv, 8, rng)
+        cmd = [sys.executable, os.path.join(TOOL, "train_head.py"), "--model", ck, "--train", tr, "--dev", dv,
+               "--out", out, "--epochs", "6", "--lr", "2e-3", "--threads", "2"]
+        r = subprocess.run(cmd, capture_output=True, text=True)
+        check(r.returncode == 0, "train_head.py exits 0" + ("" if r.returncode == 0 else ": " + r.stderr[-400:]))
+        if r.returncode != 0:
+            return
+        meta = json.load(open(f"{out}/finetune.json"))
+        b, a = meta["dev_before"], meta["dev_after"]
+        check(a["nll"] < b["nll"], f"dev log loss falls: {b['nll']:.4f} -> {a['nll']:.4f}")
+        check(a["acc"] >= b["acc"], f"dev accuracy does not fall: {b['acc']:.3f} -> {a['acc']:.3f}")
+        check(a["acc"] >= 0.9, f"the synthetic boundary is learned: dev accuracy {a['acc']:.3f} >= 0.9")
+        check(meta["train_sequences"] == 48 and meta["class_counts"] == {"reflex": 24, "plan": 24},
+              "finetune.json records the data it trained on")
+
+        base0, hdr0 = tensor_ranges(f"{ck}/model.safetensors")
+        base1, hdr1 = tensor_ranges(f"{out}/model.safetensors")
+        raw0, raw1 = open(f"{ck}/model.safetensors", "rb").read(), open(f"{out}/model.safetensors", "rb").read()
+        check(hdr0 == hdr1 and base0 == base1 and len(raw0) == len(raw1), "header, offsets and size are unchanged")
+        changed, frozen_same = 0, True
+        for name, info in hdr0.items():
+            lo, hi = info["data_offsets"]
+            same = raw0[base0 + lo:base0 + hi] == raw1[base1 + lo:base1 + hi]
+            if name.startswith(("head.", "type_emb.", "scorer.")):
+                changed += 0 if same else 1
+            elif not same:
+                frozen_same = False
+        check(frozen_same, "every non-head tensor (encoder, temperature) is byte-identical")
+        check(changed > 0, f"the head tensors were rewritten ({changed} changed)")
+        cfg = json.load(open(f"{out}/rl_agent_config.json"))
+        check("choice:2" in cfg.get("temperature_by_options", {}), "the temperature is written to the choice:2 bucket")
+
+        r2 = subprocess.run(cmd, capture_output=True, text=True)
+        check(r2.returncode != 0 and "refusing to overwrite" in (r2.stdout + r2.stderr), "an existing --out is refused")
+
+        # metrics.py, by hand: 3 cases, 2 right; p(true) 0.9, 0.6, 0.3
+        pc = os.path.join(work, "per.jsonl")
+        with open(pc, "w") as f:
+            for lab, cho, p, ms in (("reflex", "reflex", 0.9, 100), ("plan", "plan", 0.6, 200), ("plan", "reflex", 0.3, 300)):
+                other = "plan" if lab == "reflex" else "reflex"
+                f.write(json.dumps({"label": lab, "chosen": cho, "probabilities": {lab: p, other: 1 - p},
+                                    "decision_ms": ms}) + "\n")
+        m = metrics.m(pc)
+        ll = -(math.log(0.9) + math.log(0.6) + math.log(0.3)) / 3
+        # top p 0.9 (right), 0.6 (right), 0.7 (wrong) -> bins 9, 6, 7, each |p - correct|: 0.1, 0.4, 0.7
+        ece = (0.1 + 0.4 + 0.7) / 3
+        check(m["correct"] == 2 and m["acc"] == round(2 / 3, 4), "metrics: accuracy 2/3")
+        check(m["log_loss"] == round(ll, 4), f"metrics: log loss {m['log_loss']} == {round(ll, 4)}")
+        check(m["ece"] == round(ece, 4), f"metrics: ECE {m['ece']} == {round(ece, 4)}")
+        check(m["ms_p50"] == 200 and m["min_p_true"] == 0.3, "metrics: decision p50 and min p(true)")
+    finally:
+        shutil.rmtree(work, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    os.makedirs(os.path.expanduser("~/.cache/scratch-tests"), exist_ok=True)
+    main()
+    print("all laya_finetune checks passed" if fails == 0 else f"{fails} laya_finetune check(s) FAILED")
+    sys.exit(1 if fails else 0)

tools/laya_finetune/head.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+# Copyright 2026 bong-water-water-bong
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Laya's decision head in plain torch, over the encoder states `1bit route --dump-encoder` writes.
+
+It is the head of laya/scorer.cpp (and of DecisionModel upstream): h + type_emb[qtype], two pre-norm
+TransformerEncoderLayers (d 1024, 16 heads, ff 4096, relu, key padding), the option markers, and the scorer
+LayerNorm -> Linear -> GELU(erf) -> Linear(1); probabilities are softmax(logit / T) with T from
+temperature_by_options. The encoder is frozen: only these tensors are trained.
+"""
+import json
+import struct
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+HEAD_PREFIXES = ("head.", "type_emb.", "scorer.")
+DT = {"F32": np.float32, "F16": np.float16, "BF16": None}
+
+
+def read_safetensors(path, want=lambda name: True):
+    with open(path, "rb") as f:
+        n = struct.unpack("<Q", f.read(8))[0]
+        header = json.loads(f.read(n))
+        base = 8 + n
+        out = {}
+        for name, info in header.items():
+            if name == "__metadata__" or not want(name):
+                continue
+            a, b = info["data_offsets"]
+            f.seek(base + a)
+            raw = f.read(b - a)
+            if info["dtype"] == "BF16":
+                t = torch.frombuffer(bytearray(raw), dtype=torch.bfloat16).float()
+            else:
+                t = torch.from_numpy(np.frombuffer(raw, dtype=DT[info["dtype"]]).astype(np.float32).copy())
+            out[name] = t.reshape(info["shape"])
+        return header, out
+
+
+def load_head(model_dir):
+    _, w = read_safetensors(f"{model_dir}/model.safetensors",
+                            lambda n: n.startswith(HEAD_PREFIXES) or n == "temperature")
+    cfg = json.load(open(f"{model_dir}/rl_agent_config.json"))
+    return {k: v.clone().requires_grad_(k != "temperature") for k, v in w.items()}, cfg
+
+
+def bucket(qtype, k):
+    sz = "2" if k <= 2 else "3-5" if k <= 5 else "6-10" if k <= 10 else "11+"
+    return ("choice:", "score:", "noul:")[qtype] + sz
+
+
+def temperature(w, cfg, qtype, k):
+    return float(cfg.get("temperature_by_options", {}).get(bucket(qtype, k), w["temperature"][qtype]))
+
+
+def forward(w, h, pad, markers, qtype, nhead=16):
+    """h [B, L, D] (encoder after final norm), pad [B, L] True = padding, markers [B, K] long, qtype [B] long.
+    Returns the logits [B, K] before temperature."""
+    x = h + w["type_emb.weight"][qtype][:, None, :]
+    B, L, D = x.shape
+    for i in range(2):
+        p = f"head.layers.{i}."
+        y = F.layer_norm(x, (D,), w[p + "norm1.weight"], w[p + "norm1.bias"], 1e-5)
+        a, _ = F.multi_head_attention_forward(
+            y.transpose(0, 1), y.transpose(0, 1), y.transpose(0, 1), D, nhead,
+            w[p + "self_attn.in_proj_weight"], w[p + "self_attn.in_proj_bias"], None, None, False, 0.0,
+            w[p + "self_attn.out_proj.weight"], w[p + "self_attn.out_proj.bias"], training=False,
+            key_padding_mask=pad, need_weights=False)
+        x = x + a.transpose(0, 1)
+        y = F.layer_norm(x, (D,), w[p + "norm2.weight"], w[p + "norm2.bias"], 1e-5)
+        y = F.linear(F.relu(F.linear(y, w[p + "linear1.weight"], w[p + "linear1.bias"])),
+                     w[p + "linear2.weight"], w[p + "linear2.bias"])
+        x = x + y
+    m = torch.gather(x, 1, markers[:, :, None].expand(-1, -1, D))
+    m = F.layer_norm(m, (D,), w["scorer.0.weight"], w["scorer.0.bias"], 1e-5)
+    m = F.gelu(F.linear(m, w["scorer.1.weight"], w["scorer.1.bias"]))
+    return F.linear(m, w["scorer.3.weight"], w["scorer.3.bias"])[..., 0]
+
+
+class Dump:
+    """A `1bit route --dump-encoder --out DIR` directory."""
+
+    def __init__(self, d):
+        self.dir = d
+        self.info = json.load(open(f"{d}/dump.json"))
+        self.cases = [json.loads(x) for x in open(f"{d}/cases.jsonl") if x.strip()]
+        D = self.info["d"]
+        self.h = np.memmap(f"{d}/h.f32", dtype=np.float32, mode="r", shape=(self.info["rows"], D))
+
+    def batch(self, idx):
+        cs = [self.cases[i] for i in idx]
+        L = max(c["len"] for c in cs)
+        K = max(len(c["markers"]) for c in cs)
+        D = self.info["d"]
+        h = torch.zeros(len(cs), L, D)
+        pad = torch.ones(len(cs), L, dtype=torch.bool)
+        mk = torch.zeros(len(cs), K, dtype=torch.long)
+        for b, c in enumerate(cs):
+            h[b, :c["len"]] = torch.from_numpy(np.array(self.h[c["offset"]:c["offset"] + c["len"]]))
+            pad[b, :c["len"]] = False
+            mk[b, :len(c["markers"])] = torch.tensor(c["markers"])
+        qt = torch.tensor([c["qtype"] for c in cs])
+        y = torch.tensor([c["label"] for c in cs])
+        return h, pad, mk, qt, y, cs