| """ |
| 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 [] |
|
|
