File size: 5,096 Bytes
9eb2e79 95c986f 9eb2e79 95c986f 9eb2e79 95c986f 9eb2e79 95c986f 9eb2e79 95c986f 9eb2e79 95c986f 9eb2e79 95c986f 9eb2e79 95c986f 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 | """
Search over cube states guided by a learned distance estimate.
The policy beam search in serve.py explores sequences the *policy* ranks highly.
This explores states the *value function* thinks are closest to solved, which is
a different and stronger thing: it can prefer a move the policy never considered,
because the ranking comes from an independent estimate rather than from the
policy's own confidence.
This is DeepCubeA's shape in miniature -- batch-weighted search with the engine
as the transition function and a network supplying the heuristic. The weighting
`h + lambda * g` trades solution length against search effort: lambda=0 is pure
greedy on the heuristic (fast, longer solutions), higher lambda behaves more like
uniform-cost search.
Two kinds of value model can drive this, and `load_value_model` tells them apart
by what the checkpoint carries:
- **supervised** (`train_value.py`): a distribution over distances, trained on
Kociemba lengths. Those are upper bounds rather than true optimal distances, so
the heuristic is inadmissible and no shortest-solution guarantee holds.
- **value iteration** (`train_value_iteration.py`): a scalar, bootstrapped from
the solved state with no solver involved.
Either way the engine verifies that whatever comes back actually solves, and
finding *a* solution at depth 20 is the open problem, not finding the shortest.
"""
import sys
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).parent))
import cube_tokenizer as T
from gen_data import SOLVED, MOVES, apply_move
def load_value_model(path, device):
"""Loads either value-model flavour, detected from the checkpoint's own keys.
The two trainers save under different keys ("state_dict" vs "model"), which is
the only reliable discriminator: both carry identical hidden/layers/heads, so
shape alone cannot distinguish a 27-way distribution head from a scalar one
without guessing. The flavour is recorded on the model so `heuristic` reads
each correctly -- feeding a scalar head through softmax would silently return
a constant and turn the search into an untargeted walk.
"""
ckpt = torch.load(path, map_location=device)
if "state_dict" in ckpt:
from train_value import ValueNet
model = ValueNet(ckpt["hidden"], ckpt["layers"], ckpt["heads"]).to(device)
model.load_state_dict(ckpt["state_dict"])
kind = "distribution"
elif "model" in ckpt:
from train_value_iteration import ValueNet as ValueNetScalar
model = ValueNetScalar(ckpt["hidden"], ckpt["layers"], ckpt["heads"]).to(device)
model.load_state_dict(ckpt["model"])
kind = "scalar"
else:
raise ValueError(f"{path}: no 'state_dict' or 'model' key; not a value "
f"checkpoint (keys: {sorted(ckpt)[:8]})")
model.eval()
model.value_kind = kind
return model
@torch.no_grad()
def heuristic(model, states, device, batch_size=1024):
"""Estimated distance-to-solved for each state.
For a distribution head, the expectation over the predicted distribution
rather than the argmax: the extra resolution matters when ranking siblings
that all round to the same integer, which is most of the frontier. A scalar
head already has that resolution and is read directly.
"""
kind = getattr(model, "value_kind", "distribution")
out = []
for i in range(0, len(states), batch_size):
chunk = states[i:i + batch_size]
ids = torch.tensor([T.encode_state(s) for s in chunk], device=device)
raw = model(ids).float()
if kind == "scalar":
out.extend(raw.clamp(min=0.0).tolist())
else:
probs = torch.softmax(raw, dim=-1)
values = torch.arange(probs.shape[-1], device=device, dtype=torch.float32)
out.extend((probs * values).sum(-1).tolist())
return out
def solve_value_beam(model, state, device, beam=64, max_depth=26, lam=0.0,
batch_size=1024):
"""Beam search over states, ranked by `h + lam * g`.
Returns the move list that solves, or [] if the beam is exhausted. Visited
states are tracked globally: revisiting one cannot help, and on a group this
symmetric the frontier collapses onto duplicates quickly without it.
"""
if state == SOLVED:
return []
frontier = [(state, [])]
seen = {state}
for _ in range(max_depth):
children, paths = [], []
for s, path in frontier:
for mv in MOVES:
nxt = apply_move(s, mv)
if nxt in seen:
continue
if nxt == SOLVED:
return path + [mv]
seen.add(nxt)
children.append(nxt)
paths.append(path + [mv])
if not children:
return []
h = heuristic(model, children, device, batch_size)
scored = sorted(zip(h, children, paths), key=lambda x: x[0] + lam * len(x[2]))
frontier = [(c, p) for _, c, p in scored[:beam]]
return []
|