File size: 6,443 Bytes
9eb2e79 95c986f 9eb2e79 d32fcdc 9eb2e79 95c986f 9eb2e79 95c986f 9eb2e79 d32fcdc 9eb2e79 d32fcdc 9eb2e79 d32fcdc 9eb2e79 d32fcdc 9eb2e79 | 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 | """
Scores the value function + state search on the canonical benchmark scrambles.
Separate from eval_canonical.py because this is a different system: that one
decodes a policy (greedy, beam, or chunked), this one searches the state space
with a learned distance estimate as the heuristic and the engine as the
transition function. The policy is not involved at all.
Run on a GPU. Each solve evaluates `beam * 18` child states per step for up to
`--max-depth` steps -- tens of thousands of network calls, which is nothing on a
GPU and hopeless on a few CPU cores.
python eval_value_search.py --value <path-or-hub> --beam 64 --limit-per-depth 25
"""
import argparse, json, os, sys, time
from collections import defaultdict
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
from gen_data import SOLVED, apply_sequence
FIXTURE = Path(__file__).parent / "fixtures" / "canonical_states.json"
def main():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--value", required=True,
help="Path to a value checkpoint, or a Hub repo id to fetch it "
"from. Either flavour works -- supervised (train_value.py) "
"or value iteration (train_value_iteration.py); search.py "
"detects which from the checkpoint.")
p.add_argument("--value-file", default="value.pt",
help="Filename within the Hub repo, when --value is a repo id.")
p.add_argument("--beam", default="64",
help="Beam width, or a comma-separated list to sweep. A sweep "
"reuses one box for every width, which matters because "
"the deep buckets are where width decides the outcome and "
"a single guess at it wastes the whole run.")
p.add_argument("--lam", default="0.0",
help="Weight on path length: states are ranked by h + lam*g. "
"0 is pure greedy on the heuristic (fastest, and longest "
"solutions), 1.0 is standard A*. Accepts a comma-separated "
"list to sweep.")
p.add_argument("--max-depth", type=int, default=26)
p.add_argument("--limit-per-depth", type=int, default=25)
p.add_argument("--depths", default="")
p.add_argument("--records", default="",
help="Also write one JSON line per cube (depth, index, beam, "
"moves, solved, seconds, states scored) to this path, and "
"print each as a RECORD line. The totals alone cannot be "
"uploaded: they do not say which cube failed or with what "
"moves, and the first value-iteration sweep kept only them.")
args = p.parse_args()
import torch
from search import load_value_model, solve_value_beam
path = args.value
if not os.path.exists(path):
from huggingface_hub import hf_hub_download
path = hf_hub_download(args.value, args.value_file,
token=os.environ.get("HF_TOKEN"))
device = "cuda" if torch.cuda.is_available() else "cpu"
model = load_value_model(path, device)
print(f"value model ({model.value_kind}) on {device} | beam {args.beam} | "
f"lam {args.lam}", flush=True)
states = json.loads(FIXTURE.read_text())
if args.depths:
keep = {int(d) for d in args.depths.split(",")}
states = [s for s in states if s["depth"] in keep]
if args.limit_per_depth:
seen = defaultdict(int)
kept = []
for s in states:
if seen[s["depth"]] < args.limit_per_depth:
kept.append(s)
seen[s["depth"]] += 1
states = kept
rec_fh = open(args.records, "a") if args.records else None
grid = [(b, l) for b in [int(x) for x in str(args.beam).split(",")]
for l in [float(x) for x in str(args.lam).split(",")]]
for beam, lam in grid:
print(f"\n=== beam {beam} lam {lam} ===", flush=True)
by_depth = defaultdict(lambda: [0, 0, 0, 0.0])
print(f"{'depth':>6} {'solved':>8} {'n':>4} {'rate':>7} {'moves':>7} {'s/solve':>8}",
flush=True)
for s_ in states:
t0 = time.time()
stats = {}
moves = solve_value_beam(model, s_["facelets"], device, beam=beam,
max_depth=args.max_depth, lam=lam, stats=stats)
ok = bool(moves) and apply_sequence(s_["facelets"], moves) == SOLVED
if rec_fh:
rec = {"depth": s_["depth"], "index": s_["index"],
"facelets": s_["facelets"], "beam": beam, "lam": lam,
"max_depth": args.max_depth, "moves": moves, "solved": ok,
"seconds": round(time.time() - t0, 3),
"scored": stats.get("scored"), "steps": stats.get("steps"),
"device": torch.cuda.get_device_name(0) if device == "cuda" else "cpu",
# Fixed per cube, so the upload derives the same row id
# from the same record however many times it is re-run.
"finished_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}
rec_fh.write(json.dumps(rec) + "\n")
rec_fh.flush()
print("RECORD " + json.dumps(rec), flush=True)
b = by_depth[s_["depth"]]
b[0] += ok
b[1] += 1
b[2] += len(moves) if ok else 0
b[3] += time.time() - t0
# Printed per depth as each depth's last scramble lands, rather than
# only at the end: a wide beam on the deep buckets can run for a long
# time, and a box that is killed or times out mid-sweep should still
# have reported every bucket it did finish.
if b[1] == sum(1 for x in states if x["depth"] == s_["depth"]):
ok_, n_, mv_, secs_ = b
print(f"{s_['depth']:>6} {ok_:>8} {n_:>4} {ok_/n_:>6.1%} "
f"{mv_/max(ok_,1):>7.1f} {secs_/n_:>8.1f}", flush=True)
tot_ok = sum(v[0] for v in by_depth.values())
tot_n = sum(v[1] for v in by_depth.values())
print(f"{'ALL':>6} {tot_ok:>8} {tot_n:>4} {tot_ok/tot_n:>6.1%}", flush=True)
if __name__ == "__main__":
main()
|