solve.py

"""
Solve a Rubik's Cube with tiny-cube-solver.

The network does not write moves. It reads one cube and guesses how many moves
it is from solved. This script tries every move from the current positions,
asks the network to score all the results, keeps the best `--beam` of them, and
repeats until one of them is the solved cube. Every move is applied by the code
below, and the answer is checked before it is printed.

    pip install torch transformers safetensors huggingface_hub
    python solve.py --scramble "R U F' D2 L B'"
    python solve.py --facelets UUUUUUUUURRRRRRRRRFFFFFFFFFDDDDDDDDDLLLLLLLLLBBBBBBBBB

Facelets use the same 54-character layout as the `kociemba` package: faces in
the order U R F D L B, nine stickers each, each sticker named by the face whose
centre has its colour.

Search size (`--beam`) is the one setting that matters:
  64    solves every cube scrambled with 10 moves or fewer; ~25% of full scrambles
  1024  solves ~96% of full scrambles; about 16x the work
A graphics card is strongly recommended above 10 moves.
"""
import argparse
import json
import sys
import time

import torch

# Sticker permutation for one clockwise quarter turn of each face.
QT = {
  "U": [6, 3, 0, 7, 4, 1, 8, 5, 2, 45, 46, 47, 12, 13, 14, 15, 16, 17, 9, 10, 11, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 18, 19, 20, 39, 40, 41, 42, 43, 44, 36, 37, 38, 48, 49, 50, 51, 52, 53],
  "R": [0, 1, 20, 3, 4, 23, 6, 7, 26, 15, 12, 9, 16, 13, 10, 17, 14, 11, 18, 19, 29, 21, 22, 32, 24, 25, 35, 27, 28, 51, 30, 31, 48, 33, 34, 45, 36, 37, 38, 39, 40, 41, 42, 43, 44, 8, 46, 47, 5, 49, 50, 2, 52, 53],
  "F": [0, 1, 2, 3, 4, 5, 44, 41, 38, 6, 10, 11, 7, 13, 14, 8, 16, 17, 24, 21, 18, 25, 22, 19, 26, 23, 20, 15, 12, 9, 30, 31, 32, 33, 34, 35, 36, 37, 27, 39, 40, 28, 42, 43, 29, 45, 46, 47, 48, 49, 50, 51, 52, 53],
  "D": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 24, 25, 26, 18, 19, 20, 21, 22, 23, 42, 43, 44, 33, 30, 27, 34, 31, 28, 35, 32, 29, 36, 37, 38, 39, 40, 41, 51, 52, 53, 45, 46, 47, 48, 49, 50, 15, 16, 17],
  "L": [53, 1, 2, 50, 4, 5, 47, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 0, 19, 20, 3, 22, 23, 6, 25, 26, 18, 28, 29, 21, 31, 32, 24, 34, 35, 42, 39, 36, 43, 40, 37, 44, 41, 38, 45, 46, 33, 48, 49, 30, 51, 52, 27],
  "B": [11, 14, 17, 3, 4, 5, 6, 7, 8, 9, 10, 35, 12, 13, 34, 15, 16, 33, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 36, 39, 42, 2, 37, 38, 1, 40, 41, 0, 43, 44, 51, 48, 45, 52, 49, 46, 53, 50, 47],
}
SOLVED = "".join("URFDLB"[i // 9] for i in range(54))
MOVES = [f + s for f in "URFDLB" for s in ("", "'", "2")]
_TURNS = {"": 1, "'": 3, "2": 2}

# Token ids the network was trained on: [BOS] + 54 sticker tokens + [SEP].
# The sticker ids are not contiguous; they are what the training tokenizer
# produced, and every weight in model.safetensors assumes exactly these.
BOS, SEP, PAD = 1, 2, 0
COLOR_ID = dict(zip("URFDLB", [10, 13, 16, 19, 22, 25]))


def apply_move(state: str, move: str) -> str:
    perm = QT[move[0]]
    for _ in range(_TURNS[move[1:]]):
        state = "".join(state[perm[i]] for i in range(54))
    return state


def apply_sequence(state: str, moves) -> str:
    for m in moves:
        state = apply_move(state, m)
    return state


class ValueNet(torch.nn.Module):
    """A small Llama-shaped encoder over the 54 stickers, averaged, to one number."""

    def __init__(self, hidden, layers, heads, vocab_size=28, max_positions=83):
        super().__init__()
        from transformers import LlamaConfig, LlamaModel
        cfg = LlamaConfig(
            vocab_size=vocab_size, hidden_size=hidden, intermediate_size=hidden * 4,
            num_hidden_layers=layers, num_attention_heads=heads,
            num_key_value_heads=heads, max_position_embeddings=max_positions,
            pad_token_id=PAD,
        )
        self.encoder = LlamaModel(cfg)
        self.head = torch.nn.Linear(hidden, 1)

    def forward(self, ids):
        h = self.encoder(input_ids=ids).last_hidden_state.mean(dim=1)
        return self.head(h).squeeze(-1)


def load(repo_or_dir="briscoooe/tiny-cube-solver", device="cpu"):
    """Loads the network from a local folder or from the Hugging Face Hub."""
    import os
    from safetensors.torch import load_file
    if os.path.isdir(repo_or_dir):
        cfg_path = os.path.join(repo_or_dir, "config.json")
        weights = os.path.join(repo_or_dir, "model.safetensors")
    else:
        from huggingface_hub import hf_hub_download
        cfg_path = hf_hub_download(repo_or_dir, "config.json")
        weights = hf_hub_download(repo_or_dir, "model.safetensors")
    cfg = json.load(open(cfg_path))
    net = ValueNet(cfg["hidden"], cfg["layers"], cfg["heads"],
                   cfg["vocab_size"], cfg["max_positions"])
    net.load_state_dict(load_file(weights))
    return net.to(device).eval()


@torch.no_grad()
def distance(net, states, device, batch_size=1024):
    """The network's guess of moves-to-solved for each state."""
    out = []
    for i in range(0, len(states), batch_size):
        ids = torch.tensor([[BOS] + [COLOR_ID[c] for c in s] + [SEP]
                            for s in states[i:i + batch_size]], device=device)
        out.extend(net(ids).float().clamp(min=0.0).tolist())
    return out


def solve(net, state, device="cpu", beam=64, max_steps=26):
    """Beam search over positions. Returns the solving moves, or None."""
    if state == SOLVED:
        return []
    frontier = [(state, [])]
    seen = {state}
    for _ in range(max_steps):
        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 None
        h = distance(net, children, device)
        ranked = sorted(zip(h, children, paths), key=lambda x: x[0])
        frontier = [(c, p) for _, c, p in ranked[:beam]]
    return None


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    g = p.add_mutually_exclusive_group(required=True)
    g.add_argument("--scramble", help='moves applied to a solved cube, e.g. "R U F\' D2"')
    g.add_argument("--facelets", help="54-character cube state, kociemba layout")
    p.add_argument("--beam", type=int, default=64)
    p.add_argument("--model", default="briscoooe/tiny-cube-solver",
                   help="Hub repo id or local folder")
    args = p.parse_args()

    if args.scramble:
        moves = args.scramble.split()
        bad = [m for m in moves if m not in MOVES]
        if bad:
            sys.exit(f"unknown moves: {bad}; use {' '.join(MOVES)}")
        state = apply_sequence(SOLVED, moves)
    else:
        state = args.facelets.strip()
        if len(state) != 54 or set(state) - set("URFDLB") or any(state.count(c) != 9 for c in "URFDLB"):
            sys.exit("facelets must be 54 characters, nine each of U R F D L B")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    net = load(args.model, device)
    t0 = time.time()
    moves = solve(net, state, device, beam=args.beam)
    secs = time.time() - t0
    if moves is None or apply_sequence(state, moves) != SOLVED:
        print(f"not solved (beam {args.beam}, {secs:.1f}s). Try a larger --beam.")
        sys.exit(1)
    print(" ".join(moves))
    print(f"{len(moves)} moves, {secs:.1f}s on {device}, beam {args.beam}", file=sys.stderr)


if __name__ == "__main__":
    main()