| """ |
| 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, stats=None): |
| """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. |
| |
| `stats`, when given a dict, receives `scored` (states the network evaluated) |
| and `steps` (levels expanded) -- the cost of the solve, which a replay page |
| can show beside the time. |
| """ |
| if stats is not None: |
| stats["scored"], stats["steps"] = 0, 0 |
| 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) |
| if stats is not None: |
| stats["scored"] += len(children) |
| stats["steps"] += 1 |
| 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 [] |
|
|
