File size: 3,457 Bytes
223c8ed fbd14f5 223c8ed fbd14f5 223c8ed | 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 | """
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.
The heuristic is Kociemba length, an upper bound rather than true optimal
distance, so it is inadmissible and the search is not guaranteed to return a
shortest solution. That is fine: 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):
from train_value import ValueNet
ckpt = torch.load(path, map_location=device)
model = ValueNet(ckpt["hidden"], ckpt["layers"], ckpt["heads"]).to(device)
model.load_state_dict(ckpt["state_dict"])
model.eval()
return model
@torch.no_grad()
def heuristic(model, states, device, batch_size=1024):
"""Expected distance-to-solved for each state.
Uses 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.
"""
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)
probs = torch.softmax(model(ids).float(), 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 []
|