foxhunt/scripts/compare-per-state.py

#!/usr/bin/env python3
"""Compare raw PER state dumps between two determinism runs.

Reads the dumps produced by
`IntegratedTrainer::dump_per_state_for_debug` (env-gated on
`FOXHUNT_DETERMINISM_DEBUG=1`) and prints a per-step exact-equality
verdict for each of:

  - `sample_prng_d`      (B u32)        — device PRNG state per batch
  - `priority_tree_d`    (2 × cap f32)  — PER priority sum-tree
  - `sample_indices_d`   (B u32)        — sampled leaf indices

Hypothesis table (spec §2.F sub-investigation):
  HYPOTHESIS A (PRNG):       prng differs                 → seed propagation
                              between runs (xorshift32     bug
                              cold-init diverging)
  HYPOTHESIS B (Tree-walk):  prng EQUAL, tree EQUAL,      → fp comparison
                              indices differ              `u <= left_val`
                                                          flipping near the
                                                          fp epsilon boundary
                                                          inside rl_per_sample
  HYPOTHESIS C (Tree-rebuild): prng EQUAL, tree differs    → parallel
                              (even tiny f32 deltas)      reduction order in
                                                          rl_per_tree_rebuild
                                                          (children-before-
                                                          parents ordering)

Output is written to stdout. Exit codes:
  0  both runs identical for all three buffers at all dumped steps
  1  divergence found; verdict printed (one of A / B / C)
  2  inputs missing or malformed

Usage:
  python3 scripts/compare-per-state.py <run_a_dir> <run_b_dir> [--b-size N]

Each dir contains `step_{0,1,2}_{prng,tree,indices}.bin`.

Per `feedback_cpu_is_read_only`: this script ONLY reads device dumps;
it performs no compute that affects training.
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path
import struct


def read_u32(path: Path) -> list[int]:
    """Read a raw little-endian u32 binary file into a Python list."""
    data = path.read_bytes()
    if len(data) % 4 != 0:
        raise ValueError(f"{path}: size {len(data)} not a multiple of 4")
    n = len(data) // 4
    return list(struct.unpack(f"<{n}I", data))


def read_f32(path: Path) -> list[float]:
    """Read a raw little-endian f32 binary file into a Python list."""
    data = path.read_bytes()
    if len(data) % 4 != 0:
        raise ValueError(f"{path}: size {len(data)} not a multiple of 4")
    n = len(data) // 4
    return list(struct.unpack(f"<{n}f", data))


def compare_lists(a: list, b: list, label: str) -> tuple[bool, int, str]:
    """Compare two equal-length lists; return (equal, first_diff_idx, summary)."""
    if len(a) != len(b):
        return False, -1, f"{label}: length mismatch A={len(a)} B={len(b)}"
    diffs = []
    for i, (x, y) in enumerate(zip(a, b)):
        if x != y:
            diffs.append(i)
        if len(diffs) >= 5:
            break
    if not diffs:
        return True, -1, f"{label}: EQUAL ({len(a)} elements)"
    sample = ", ".join(
        f"[{i}]: A={a[i]!r} B={b[i]!r}" for i in diffs[:3]
    )
    return False, diffs[0], (
        f"{label}: DIVERGE at idx {diffs[0]} "
        f"(showing up to 3): {sample}"
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("run_a", type=Path)
    parser.add_argument("run_b", type=Path)
    parser.add_argument(
        "--b-size",
        type=int,
        default=128,
        help="batch size used by the run (default: 128)",
    )
    args = parser.parse_args()

    if not args.run_a.is_dir() or not args.run_b.is_dir():
        print(
            f"ERROR: run dirs missing: {args.run_a} or {args.run_b}",
            file=sys.stderr,
        )
        return 2

    # Track verdict across steps. We pick the FIRST step that shows
    # divergence and report which buffer(s) diverge there.
    overall_any_diverge = False
    first_diverge_step: int | None = None
    first_diverge_state: dict[str, bool] = {}

    for step in (0, 1, 2):
        print(f"=== step {step} ===")
        files = {
            "prng": (
                args.run_a / f"step_{step}_prng.bin",
                args.run_b / f"step_{step}_prng.bin",
                read_u32,
            ),
            "indices": (
                args.run_a / f"step_{step}_indices.bin",
                args.run_b / f"step_{step}_indices.bin",
                read_u32,
            ),
            "tree": (
                args.run_a / f"step_{step}_tree.bin",
                args.run_b / f"step_{step}_tree.bin",
                read_f32,
            ),
        }
        # Skip the step if any file missing.
        missing = [name for name, (pa, pb, _) in files.items()
                   if not (pa.is_file() and pb.is_file())]
        if missing:
            print(f"  SKIP — missing files: {missing}")
            continue

        step_state: dict[str, bool] = {}
        for name, (pa, pb, reader) in files.items():
            try:
                a = reader(pa)
                b = reader(pb)
            except ValueError as e:
                print(f"  ERROR reading {name}: {e}")
                return 2
            equal, _, summary = compare_lists(a, b, name)
            print(f"  {summary}")
            step_state[name] = equal
            if not equal:
                overall_any_diverge = True

        if (not all(step_state.values())) and first_diverge_step is None:
            first_diverge_step = step
            first_diverge_state = dict(step_state)

    print()
    if not overall_any_diverge:
        print("VERDICT: all three buffers match across all dumped steps.")
        print("PER state itself is bit-identical between runs; the")
        print("downstream `checksums.replay_sample_indices` divergence")
        print("observed by determinism-check.sh must be a SUM-OF-SQUARES")
        print("permutation artifact (the checksum is permutation-invariant).")
        print("Re-investigate: indices may differ in ORDER but match as a")
        print("multiset, indicating per-batch slot-binding non-determinism")
        print("upstream of PER.")
        return 0

    assert first_diverge_step is not None
    prng_eq = first_diverge_state.get("prng", True)
    tree_eq = first_diverge_state.get("tree", True)
    idx_eq = first_diverge_state.get("indices", True)

    print(f"VERDICT (first divergent step = {first_diverge_step}):")
    if not prng_eq:
        print("  HYPOTHESIS A — PRNG seeding non-determinism.")
        print("  sample_prng_d differs between runs at step",
              first_diverge_step,
              "→ the device-side xorshift32 state was already divergent")
        print("  before the tree walk. Apply Option A fix: explicit")
        print("  cli.seed-derived seeding in trainer construction.")
        return 1
    if prng_eq and tree_eq and not idx_eq:
        print("  HYPOTHESIS B — Tree-walk fp-comparison flip.")
        print("  prng matches, tree matches, but sample_indices diverge")
        print("  → the `u <= left_val` comparison in rl_per_sample.cu:88")
        print("  is producing different paths despite identical inputs.")
        print("  Apply Option B fix: integer priority sum-tree.")
        return 1
    if prng_eq and not tree_eq:
        print("  HYPOTHESIS C — Tree-rebuild non-determinism.")
        print("  prng matches, but priority_tree itself diverges")
        print("  → rl_per_tree_rebuild.cu accumulation order is")
        print("  non-deterministic across runs (parallel children-before-")
        print("  parents not enforced strongly enough).")
        print("  Apply Option C fix: serialize the rebuild bottom-up.")
        return 1
    # Should not be reachable (one of A/B/C must hold once any diverges).
    print("  UNEXPECTED state — review per-buffer summary above.")
    print(f"  prng_eq={prng_eq} tree_eq={tree_eq} idx_eq={idx_eq}")
    return 1


if __name__ == "__main__":
    sys.exit(main())