File size: 7,693 Bytes
97f78c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 | #!/usr/bin/env python3
"""Measure what the planner's field actually encodes.
Fits a ridge probe from the planner's per-cell features to the true BFS
distance and reports R^2. This is the diagnostic that caught the central
training bug: with only the sparse decision questions supervising it, the
recurrence scored R^2 = 0.21 -- it was running, it was differentiable, its
accuracy numbers looked reasonable, and it was not computing distances.
Also reports the field head's own predictions, which is what the decision
heads actually read.
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
import torch
import envs.maze as mz
from jevon.grid import bfs_field, encode_maze, pad_boards
from jevon.inference import resolve_weights
from jevon.modeling import JevonConfig, JevonModel
from jevon.planner import default_iterations
def ridge_r2(features: torch.Tensor, target: torch.Tensor) -> tuple[float, float]:
design = torch.cat([features, torch.ones(len(features), 1)], 1)
gram = design.T @ design + 1e-3 * torch.eye(design.shape[1])
weights = torch.linalg.lstsq(gram, design.T @ target.unsqueeze(1)).solution
pred = (design @ weights).squeeze(1)
ss_res = ((pred - target) ** 2).sum()
ss_tot = ((target - target.mean()) ** 2).sum()
return float(1 - ss_res / ss_tot), float((pred - target).abs().mean())
def descent_accuracy(head: torch.Tensor, truth: torch.Tensor) -> float:
"""Fraction of cells whose lowest-predicted neighbour is a correct step.
R^2 scores the field's absolute values, but greedy descent only ever
compares a cell's four neighbours, so a field can be numerically mediocre
and still route perfectly -- or fit well on average and still invert the
one comparison that matters. This is the quantity that predicts solve
rate, so it is measured rather than inferred from R^2.
"""
head = head.detach()
hits = total = 0
h, w = truth.shape
for r in range(h):
for c in range(w):
here = float(truth[r, c])
if here <= 0: # blocked, unreachable, or the goal
continue
options = [(rr, cc) for rr, cc in
((r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1))
if 0 <= rr < h and 0 <= cc < w and truth[rr, cc] >= 0]
if not options:
continue
pick = min(options, key=lambda x: float(head[x[0], x[1]]))
# Correct iff that neighbour lies on *a* shortest path.
hits += float(truth[pick]) == here - 1
total += 1
return hits / max(total, 1)
def descent_outcomes(head: torch.Tensor, truth: torch.Tensor,
cap_mult: int = 8) -> float:
"""Fraction of cells from which greedy descent actually reaches the goal.
Per-step accuracy is not the quantity a controller lives or dies by. A
field that is right 99% of the time but contains one spurious basin traps
forever; a field that is right 70% of the time with no basin still arrives,
just by a longer route. Solve rate equals *this* number, so report it next
to the per-step figure and let the gap between them speak.
"""
head = head.detach()
h, w = truth.shape
goal = min(((r, c) for r in range(h) for c in range(w)
if float(truth[r, c]) == 0.0), default=None)
if goal is None:
return 0.0
reached = total = 0
cap = cap_mult * h * w
for r in range(h):
for c in range(w):
if float(truth[r, c]) < 0:
continue
total += 1
cur, seen = (r, c), set()
for _ in range(cap):
if cur == goal:
reached += 1
break
if cur in seen:
break # a cycle: descent never escapes
seen.add(cur)
options = [(rr, cc) for rr, cc in
((cur[0] - 1, cur[1]), (cur[0] + 1, cur[1]),
(cur[0], cur[1] - 1), (cur[0], cur[1] + 1))
if 0 <= rr < h and 0 <= cc < w and truth[rr, cc] >= 0]
if not options:
break
cur = min(options, key=lambda x: float(head[x[0], x[1]]))
return reached / max(total, 1)
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--checkpoint", default="runs/jevon-final")
# None means "prefer balanced.pt", the checkpoint every shipped number
# comes from. See jevon.inference.resolve_weights.
p.add_argument("--weights", default=None)
p.add_argument("--sizes", default="11,21,31,51")
p.add_argument("--episodes", type=int, default=4)
p.add_argument("--device", default="cpu")
p.add_argument("--out", default=None)
args = p.parse_args()
root = Path(args.checkpoint)
cfg = JevonConfig(**json.loads((root / "config.json").read_text()))
model = JevonModel(cfg)
model.load_state_dict(torch.load(resolve_weights(root, args.weights),
map_location="cpu"))
model.eval().to(args.device)
report = {}
print(f"{'size':>5s} {'cells':>7s} {'probe R2':>9s} {'probe MAE':>10s} "
f"{'head R2':>8s} {'descent':>8s} {'reach':>7s} "
f"{'diameter':>9s} {'T':>6s}")
for size in [int(s) for s in args.sizes.split(",")]:
feats, heads, dists, descents, reaches = [], [], [], [], []
diameter = 0
steps = default_iterations(size)
with torch.no_grad():
for i in range(args.episodes):
topology = mz.TOPOLOGIES[i % len(mz.TOPOLOGIES)]
# Seeds ending in 9 are the held-out test split.
state = mz.generate_maze(size, topology, 70_009 + i * 10)
board, _, target = encode_maze(state)
boards, passable, _ = pad_boards([board])
cells = model.encode_boards(boards.to(args.device),
passable.to(args.device), steps, 0)
head = model.predict_field(
cells, boards.to(args.device),
passable.to(args.device), iterations=steps)[0, 0]
truth = bfs_field(board, target)
mask = truth >= 0
diameter = max(diameter, int(truth.max()))
descents.append(descent_accuracy(head.cpu(), truth))
reaches.append(descent_outcomes(head.cpu(), truth))
feats.append(cells[0, :, mask].T.cpu())
heads.append(head[mask].cpu())
dists.append(truth[mask])
features = torch.cat(feats)
truth = torch.cat(dists)
r2, mae = ridge_r2(features, truth)
head_r2, _ = ridge_r2(torch.cat(heads).unsqueeze(1), truth)
descent = sum(descents) / max(len(descents), 1)
reach = sum(reaches) / max(len(reaches), 1)
report[size] = {"cells": len(truth), "probe_r2": r2, "probe_mae": mae,
"head_r2": head_r2, "descent_acc": descent,
"reach_rate": reach,
"diameter": diameter, "iterations": steps}
print(f"{size:5d} {len(truth):7d} {r2:9.4f} {mae:10.2f} {head_r2:8.4f} "
f"{descent:8.4f} {reach:7.3f} {diameter:9d} {steps:6d}")
if args.out:
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
Path(args.out).write_text(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|