| |
| """ |
| Observability: dump the full recurrent trajectory for one position. |
| |
| Emits a human-readable trace and/or machine-readable JSON containing |
| h0..hn, Δh_t, memory, routing weights, halting probabilities, per-step move |
| latents and candidate move scores. |
| |
| python inspect_trace.py --checkpoint runs/x/latest.pt --fen "..." --steps 8 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
|
|
| import chess |
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
|
|
| from tinychess.batching import encode_batch |
| from tinychess.checkpoint import load_model_only |
|
|
|
|
| @torch.no_grad() |
| def trace_position(model, board, steps=8, adaptive=False, topk=5): |
| pb = encode_batch([board], [board.peek() if board.move_stack else None]) |
| out = model(**pb.model_args(), steps=steps, adaptive=adaptive, trace=True) |
| traj = out["trajectory"] |
| coll = out.get("collect", {}) or {} |
| moves = pb.real_moves[0] |
|
|
| rec = { |
| "fen": board.fen(), |
| "side_to_move": "white" if board.turn else "black", |
| "n_legal": int(pb.n_legal[0]), |
| "params": model.param_report()["total"], |
| "active_params": model.param_report()["active"], |
| "adaptive": adaptive, |
| "n_steps_used": float(out["n_steps"][0]), |
| "value": float(out["value"][0, 0]), |
| "steps": [], |
| } |
|
|
| prev = None |
| for t, h in enumerate(traj): |
| h0 = h[0] |
| s = {"t": t, |
| "h_norm": float(h0.norm()), |
| "h_mean": float(h0.mean()), |
| "h_std": float(h0.std()), |
| "active_units": int((h0.abs() > 1e-3).float().mean(0).gt(0.5).sum()), |
| } |
| if prev is not None: |
| d = h0 - prev |
| s["delta_norm"] = float(d.norm()) |
| s["delta_rel"] = float(d.norm() / prev.norm().clamp_min(1e-6)) |
| s["most_changed_squares"] = [ |
| chess.square_name(i) for i in d.norm(dim=-1).topk(4).indices.tolist()] |
| prev = h0 |
| i = t - 1 |
| if 0 <= i < len(coll.get("router", [])): |
| w = coll["router"][i][0] |
| s["route"] = {op: round(float(v), 3) |
| for op, v in zip(model.core.ops, w.tolist())} |
| if 0 <= i < len(coll.get("halt", [])): |
| s["p_halt"] = round(float(coll["halt"][i][0]), 4) |
| if 0 <= i < len(coll.get("mem_attn", [])): |
| ma = coll["mem_attn"][i][0] |
| s["memory_attention"] = [round(float(v), 3) for v in ma.mean(0).tolist()] |
| rec["steps"].append(s) |
|
|
| if out.get("memory") is not None: |
| m = out["memory"][0] |
| rec["memory_final"] = {"n_slots": int(m.shape[0]), |
| "slot_norms": [round(float(v), 3) for v in m.norm(dim=-1)]} |
|
|
| |
| if "refine_logits" in out: |
| rec["refinement"] = [] |
| for i, lg in enumerate(out["refine_logits"]): |
| if lg is None: |
| continue |
| p = torch.softmax(lg[0], -1) |
| k = min(topk, len(moves)) |
| top = p.topk(k) |
| rec["refinement"].append({ |
| "z_index": i, |
| "top": [{"move": moves[j].uci(), "p": round(float(v), 4)} |
| for v, j in zip(top.values, top.indices)]}) |
|
|
| p = torch.softmax(out["logits"][0], -1) |
| k = min(topk, len(moves)) |
| top = p.topk(k) |
| rec["final_policy"] = [{"move": moves[j].uci(), "p": round(float(v), 4)} |
| for v, j in zip(top.values, top.indices)] |
| rec["chosen_move"] = moves[int(p.argmax())].uci() |
| rec["policy_entropy"] = float(-(p * p.clamp_min(1e-12).log()).sum()) |
| return rec |
|
|
|
|
| def human(rec) -> str: |
| L = [] |
| L.append("=" * 72) |
| L.append(f"Position: {rec['fen']}") |
| L.append(f" side to move: {rec['side_to_move']} legal moves: {rec['n_legal']}") |
| L.append(f" params: {rec['params']:,} (active {rec['active_params']:,}) " |
| f"steps used: {rec['n_steps_used']:.2f} value: {rec['value']:+.3f}") |
| L.append("=" * 72) |
| for s in rec["steps"]: |
| L.append(f"Step {s['t']}") |
| L.append(f" |h|={s['h_norm']:.3f} mean={s['h_mean']:+.4f} std={s['h_std']:.4f}" |
| f" active units={s['active_units']}") |
| if "delta_norm" in s: |
| L.append(f" Δh={s['delta_norm']:.4f} (rel {s['delta_rel']:.3f}) " |
| f"changed: {', '.join(s['most_changed_squares'])}") |
| if "route" in s: |
| L.append(" route: " + " ".join(f"{k}={v:.2f}" for k, v in s["route"].items())) |
| if "p_halt" in s: |
| L.append(f" p_halt: {s['p_halt']:.3f}") |
| if "memory_attention" in s: |
| L.append(f" memory attn: {s['memory_attention']}") |
| if "memory_final" in rec: |
| L.append(f"\nMemory ({rec['memory_final']['n_slots']} slots) norms: " |
| f"{rec['memory_final']['slot_norms']}") |
| if "refinement" in rec: |
| L.append("\nMove-latent refinement:") |
| for r in rec["refinement"]: |
| tops = " ".join(f"{d['move']}:{d['p']:.3f}" for d in r["top"]) |
| L.append(f" z{r['z_index']}: {tops}") |
| L.append("\nFinal policy:") |
| for d in rec["final_policy"]: |
| L.append(f" {d['move']} p={d['p']:.4f}") |
| L.append(f"\nFinal move: {rec['chosen_move']} entropy={rec['policy_entropy']:.3f}") |
| return "\n".join(L) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--checkpoint", required=True) |
| ap.add_argument("--fen", default=chess.STARTING_FEN) |
| ap.add_argument("--steps", type=int, default=8) |
| ap.add_argument("--adaptive", action="store_true") |
| ap.add_argument("--format", choices=["human", "json", "both"], default="human") |
| ap.add_argument("--out", default=None) |
| a = ap.parse_args() |
|
|
| model, cfg = load_model_only(a.checkpoint) |
| rec = trace_position(model, chess.Board(a.fen), a.steps, a.adaptive) |
| if a.format in ("human", "both"): |
| print(human(rec)) |
| if a.format in ("json", "both"): |
| js = json.dumps(rec, indent=2) |
| if a.out: |
| open(a.out, "w").write(js) |
| print(f"\n[saved json] {a.out}") |
| elif a.format == "json": |
| print(js) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
