scripts/readme_table.py

#!/usr/bin/env python3
"""Rewrite the README's headline table from the evaluation JSON.

Every other number in this project is generated: RESULTS.md comes from
`make_results.py`, the arcade report from `report.py`. The README headline was
the exception, and it drifted -- it carried figures from a `runs/arm-*` probe
long after the shipped checkpoint had changed, which is the specific failure
this repository criticises NanoJev for.

The table lives between two HTML comment markers and is replaced wholesale.
`--check` re-renders and compares instead of writing, so a stale README fails
the test suite rather than shipping.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]


def begin(block: str) -> str:
    return f"<!-- BEGIN:{block} (generated by scripts/readme_table.py) -->"


def end(block: str) -> str:
    return f"<!-- END:{block} -->"


def block(name: str, body) -> list[str]:
    """One generated block, markers included.

    Shared with the test that checks the README is in sync. It was duplicated
    there, and when marker emission moved out of the renderers the copy in the
    test was left comparing marker-less output against a README that has
    markers -- an assertion that could never pass, and did not fail only
    because the test skips when the run directory is incomplete.

    Trailing blanks are trimmed off the body so a renderer that ends its table
    with one does not stack it with the blank added here. Spacing is decided
    in this one place, not by each renderer.
    """
    body = list(body)
    while body and not body[0].strip():
        body.pop(0)
    while body and not body[-1].strip():
        body.pop()
    return [begin(name), ""] + body + ["", end(name)]


def read(path: Path):
    try:
        return json.loads(path.read_text())
    except (OSError, json.JSONDecodeError):
        return None


def maze_row(summaries: dict, key: str) -> dict | None:
    data = summaries.get(key)
    return (data.get("summary", {}) or {}).get("maze") if data else None


def snake_row(summaries: dict, key: str) -> dict | None:
    data = summaries.get(key)
    return (data.get("summary", {}) or {}).get("snake") if data else None


def pick_split(ev: dict) -> tuple[dict, str]:
    """The split the headline quotes, and its name.

    `test` first, and the preference is not cosmetic. `best.pt` and
    `balanced.pt` are selected on `val`, once per eval across the whole run,
    so a val figure is a maximum over ~18 noisy estimates of the same 74
    snake questions and is biased upward by construction. `test` is scored
    once, by a checkpoint that never competed on it.

    evaluate.py keys its output by `Path(split).stem`, so the keys are bare
    names. This used to look up "data/frozen/test.jsonl", matched nothing,
    and silently rendered the headline with no accuracy table at all -- the
    lookup and the writer had never been run against each other, because the
    test fixture made up its own key format. Path-style keys are still
    accepted in case that writer changes back, and an unrecognised shape now
    says so rather than rendering an empty section.
    """
    for name in ("test", "val"):
        for key in (name, f"data/frozen/{name}.jsonl"):
            if isinstance(ev.get(key), dict) and ev[key].get("groups"):
                return ev[key], name
    return {}, ""


def render(run: Path) -> list[str]:
    play = run / "play"
    summaries = {p.stem.replace("_summary", ""): read(p)
                 for p in sorted(play.glob("*_summary.json"))}
    ev = read(run / "eval.json") or {}
    out = []

    # Accuracy next to its control, because a number without one is what the
    # NanoJev log reports and the whole point of this repo is not to.
    block, source = pick_split(ev)
    groups = block.get("groups", {})
    acc = [(g, groups[g]) for g in ("maze/choice", "snake/choice") if g in groups]
    if acc:
        out += [f"**Decision accuracy** (held-out `{source}` split, next to "
                "the uniform control). "
                "`boards` is the number of distinct boards behind those "
                "questions; it tracks `n` because the generator draws a fresh "
                "board per state, which is what makes `n` an honest "
                "denominator rather than an assumed one:", "",
                "| question | n | boards | accuracy | uniform control |",
                "| --- | ---: | ---: | ---: | ---: |"]
        for name, row in acc:
            b = row.get("boards")
            out.append(f"| `{name}` | {row['n']} | {'--' if b is None else b} | "
                       f"**{row['accuracy']:.4f}** | "
                       f"{row['uniform_control']:.4f} |")
        out.append("")

    rows = [("`model` (network alone)", "maze_model"),
            ("`model-field` (architectural)", "maze_model_field"),
            ("`random` (floor)", "maze_random"),
            ("`reference` (BFS)", "maze_reference")]
    body = [(lbl, maze_row(summaries, k)) for lbl, k in rows]
    body = [(lbl, m) for lbl, m in body if m]
    if body:
        # Episode count read from the summary, not typed: this block exists
        # because typed numbers drift.
        n = max(m.get("episodes", 0) for _, m in body)
        out += [f"**Maze**, {n} episodes across topologies and sizes:", "",
                "| controller | solve rate | efficiency | cells visited |",
                "| --- | ---: | ---: | ---: |"]
        for lbl, m in body:
            # Bold the model rows only. Bolding a control's solve rate reads
            # as a boast about the floor.
            sr = (f"**{m['solve_rate']:.2f}**" if "model" in lbl
                  else f"{m['solve_rate']:.2f}")
            out.append(f"| {lbl} | {sr} | {m['mean_efficiency']:.3f} | "
                       f"{m.get('mean_distinct_cells', float('nan')):.1f} |")
        out.append("")

    srows = [("`model`", "snake_model"), ("`random` (floor)", "snake_random"),
             ("`reference` (oracle)", "snake_reference")]
    sbody = [(lbl, snake_row(summaries, k)) for lbl, k in srows]
    sbody = [(lbl, s) for lbl, s in sbody if s]
    if sbody:
        n = max(v.get("episodes", 0) for _, v in sbody)
        out += [f"**Snake**, {n} episodes:", "",
                "| controller | mean food | max food | survival |",
                "| --- | ---: | ---: | ---: |"]
        for lbl, v in sbody:
            food = (f"**{v['mean_food']:.2f}**" if "model" in lbl
                    else f"{v['mean_food']:.2f}")
            out.append(f"| {lbl} | {food} | {v['max_food']} | "
                       f"{v['survival_rate']:.2f} |")
        out.append("")

    if ev and not groups:
        out += ["_`eval.json` has no recognised split; expected a `test` or "
                "`val` key from `scripts/evaluate.py`._", ""]
    if not out:
        out.append("_No evaluation artifacts yet; run `scripts/benchmark.sh`._")
        out.append("")
    return out


def render_memory(run: Path) -> list[str]:
    """`model+memory` beside `random+memory`, per board size.

    The one table in the README that has to be per-size, because the finding
    *is* the size dependence: at 11 the memory alone matches the model on
    arrival, and at 21 it beats it. A whole-benchmark average hides that by
    mixing the size where the control wins into the sizes where it does not.
    """
    play = run / "play"
    summaries = {p.stem.replace("_summary", ""): read(p)
                 for p in sorted(play.glob("*_summary.json"))}
    pairs = [("`model+memory`", "maze_model_memory"),
             ("`random+memory`", "maze_random_memory")]
    have = {lbl: maze_row(summaries, k) for lbl, k in pairs}
    if not all(have.values()):
        return [
                "_No `+memory` play artifacts yet; run `scripts/benchmark.sh`._"]
    sizes = sorted({int(z) for m in have.values()
                    for z in (m.get("by_size") or {})})
    out = ["| controller | size | solve | efficiency |",
           "| --- | ---: | ---: | ---: |"]
    for size in sizes:
        for lbl, m in have.items():
            row = (m.get("by_size") or {}).get(str(size))
            if not row:
                continue
            # Bold whichever of the pair wins the column, rather than always
            # bolding the model. At size 21 the control wins solve rate, and a
            # table that bolded the model there would be advertising a loss.
            other = (have["`random+memory`"] if lbl == "`model+memory`"
                     else have["`model+memory`"])
            orow = (other.get("by_size") or {}).get(str(size)) or {}
            sr, eff = row["solve_rate"], row.get("mean_efficiency", 0.0)
            sr_s = (f"**{sr:.2f}**" if sr > orow.get("solve_rate", sr)
                    else f"{sr:.2f}")
            eff_s = (f"**{eff:.3f}**" if eff > orow.get("mean_efficiency", eff)
                     else f"{eff:.3f}")
            out.append(f"| {lbl} | {size} | {sr_s} | {eff_s} |")
    n = max(((m.get("by_size") or {}).get(str(sizes[0])) or {}).get("episodes", 0)
            for m in have.values())
    out += ["",
            f"{n} episodes per size per topology, from the shipped checkpoint."]
    return out


# Block name -> renderer. Adding a generated table to the README means adding
# a marker pair and one entry here; nothing else in main() changes.
def render_volatility(run: Path) -> list[str]:
    """How much `val` still moves between evals, split into halves of training.

    This paragraph used to quote the numbers by hand, and said the movement
    "stops" because four consecutive evals had come back byte-identical. Two
    evals later one gave back six questions and the sentence was false. The
    claim was structurally unstable: anything phrased about the last N evals
    goes stale the moment training continues past them.

    Halves are stable under continued training in a way a trailing window is
    not, and they are the comparison the paragraph is actually making -- that
    the movement shrinks by a large factor without reaching zero. Both columns
    are derived here so the prose can point at the table instead of repeating
    it.
    """
    history = read(run / "history.json")
    if not history or len(history) < 4:
        return ["_No training history to report volatility from._"]
    rows = [(r["step"], r.get("snake_action"), r.get("maze_action"))
            for r in history if r.get("snake_action") is not None]
    if len(rows) < 4:
        return ["_Training history carries no `snake_action` to report._"]
    # Questions per split, recovered from the accuracy grid rather than
    # hard-coded: an accuracy on n questions is always a multiple of 1/n, so
    # the smallest gap between distinct observed values bounds n. Falls back
    # to reporting moves alone if the values are too few to infer it.
    def questions(values: list[float]) -> int | None:
        uniq = sorted(set(round(v, 6) for v in values))
        gaps = [b - a for a, b in zip(uniq, uniq[1:]) if b - a > 1e-9]
        if not gaps:
            return None
        n = round(1.0 / min(gaps))
        return n if 1 <= n <= 10_000 else None

    n_snake = questions([r[1] for r in rows])
    mid = len(rows) // 2
    out = ["| half of the run | evals | largest eval-to-eval move in "
           "`snake_action` |" + (" in questions |" if n_snake else ""),
           "| --- | ---: | ---: |" + (" ---: |" if n_snake else "")]
    for label, chunk in (("first", rows[:mid + 1]), ("second", rows[mid:])):
        # The halves overlap by one eval on purpose: a move is a property of
        # a *pair*, so splitting the list cleanly would drop the move across
        # the seam and quietly understate whichever half it belonged to.
        moves = [abs(b[1] - a[1]) for a, b in zip(chunk, chunk[1:])]
        biggest = max(moves) if moves else 0.0
        row = (f"| {label} ({chunk[0][0]}-{chunk[-1][0]}) | {len(chunk)} | "
               f"`{biggest:.3f}` |")
        if n_snake:
            row += f" {round(biggest * n_snake)} of {n_snake} |"
        out.append(row)
    return out


def render_shipped(run: Path) -> list[str]:
    """The two checkpoints the shipped run actually wrote, side by side.

    The section above makes its case with a snapshot taken mid-training,
    which a reader cannot check against anything: those checkpoints were
    overwritten by the next eval. These two are on disk. If the argument for
    selecting on `min(maze_action, snake_action)` is right, it should be
    visible in the files the repository ships, not only in an anecdote.

    Says so plainly when the two selectors agree. That is a real outcome --
    it means loss and gameplay pointed at the same step -- and rendering an
    identical pair of columns as though it were a contrast would be worse
    than saying nothing.
    """
    best, balanced = read(run / "best.json"), read(run / "balanced.json")
    if not best or not balanced:
        return ["_No selection metadata yet; train a run to write "
                "`best.json` and `balanced.json`._"]
    if best.get("step") == balanced.get("step"):
        return [f"_On this run both selectors chose step "
                f"{best.get('step')}: lowest eval loss and best "
                f"`min(maze_action, snake_action)` landed on the same "
                f"checkpoint, so there is nothing to contrast._"]
    rows = [("eval loss", "loss", 4, False),
            ("`acc_boolean`", "acc_boolean", 5, True),
            ("`acc_score`", "acc_score", 5, True),
            ("`acc_choice`", "acc_choice", 4, True),
            ("`maze_action`", "maze_action", 4, True),
            ("`snake_action`", "snake_action", 4, True)]
    out = [f"| | `best.pt` (step {best['step']}) | "
           f"`balanced.pt` (step {balanced['step']}) |",
           "| --- | ---: | ---: |"]
    for label, key, places, higher_is_better in rows:
        a, b = best.get(key), balanced.get(key)
        if a is None or b is None:
            continue
        # Bold the better cell, with the direction depending on the metric:
        # loss is the one row where smaller wins, and bolding the larger
        # number there would silently invert the whole table's argument.
        win_a = a > b if higher_is_better else a < b
        fa, fb = f"{a:.{places}f}", f"{b:.{places}f}"
        if abs(a - b) > 10 ** -(places + 1):
            fa, fb = (f"**{fa}**", fb) if win_a else (fa, f"**{fb}**")
        out.append(f"| {label} | {fa} | {fb} |")
    return out


BLOCKS = {"headline": render, "memory": render_memory,
          "volatility": render_volatility, "shipped": render_shipped}


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--run", default="runs/jevon-final")
    p.add_argument("--readme", default=str(ROOT / "README.md"))
    p.add_argument("--check", action="store_true",
                   help="exit 1 if the README block is stale instead of writing")
    args = p.parse_args()

    readme = Path(args.readme)
    text = updated = readme.read_text()
    for name, renderer in BLOCKS.items():
        b, e = begin(name), end(name)
        if b not in updated or e not in updated:
            print(f"markers for {name!r} not found in {readme}; add:\n{b}\n{e}",
                  file=sys.stderr)
            raise SystemExit(2)
        head, rest = updated.split(b, 1)
        _, tail = rest.split(e, 1)
        # The markers are re-emitted here, not by the renderers. They used to
        # be each renderer's job, which meant a new renderer that forgot them
        # deleted its own markers on the first write and made every later run
        # die with "markers not found" -- pointing at the README, which was
        # fine until this script ate it.
        updated = head + "\n".join(block(name, renderer(Path(args.run)))) + tail
    if args.check:
        if updated != text:
            print("README generated tables are stale; run "
                  "scripts/readme_table.py", file=sys.stderr)
            raise SystemExit(1)
        print(f"README generated tables are in sync ({len(BLOCKS)} blocks)")
        return
    readme.write_text(updated)
    print(f"updated {readme} ({len(BLOCKS)} blocks)")


if __name__ == "__main__":
    main()