tiny-cube-solver: value network, solve.py and model card

README.md

@@ -0,0 +1,113 @@
+---
+library_name: pytorch
+tags:
+  - rubiks-cube
+  - search
+  - value-iteration
+  - reinforcement-learning
+pipeline_tag: other
+---
+
+# tiny-cube-solver
+
+A 33.6M-parameter network that guesses how many moves a Rubik's Cube is from
+solved, plus a search that uses the guess to find the moves. It was trained from
+nothing: no solver, no example solutions, no human data.
+
+- Solves **96%** of fully scrambled cubes (25 cubes, 20-move scrambles)
+- Solves **every** cube scrambled with 10 moves or fewer
+- Built as a control for [RubiksBench](https://brianbriscoe.dev/lab/rubiksbench),
+  where the best frontier language models solve 50 to 80% of scrambles when they
+  write the whole answer in one go
+
+Write-up: https://brianbriscoe.dev/lab/tiny-cube-solver (coming soon)
+
+## Try it
+
+```sh
+pip install torch transformers safetensors huggingface_hub
+python solve.py --scramble "R U F' D2 L B'"
+```
+
+or with a cube state in the `kociemba` layout (faces U R F D L B, nine stickers
+each):
+
+```sh
+python solve.py --facelets <54 characters>
+```
+
+`solve.py` has no other dependencies. It downloads `config.json` and
+`model.safetensors` from this repo the first time.
+
+## How it works
+
+The network does not write moves. It reads the 54 stickers of one cube and
+returns one number: its guess of the distance to solved.
+
+`solve.py` then searches:
+
+1. Apply each of the 18 moves to every position it is keeping.
+2. Ask the network to score all the new positions.
+3. Keep the best `--beam` of them, and repeat.
+4. Stop when one of them is the solved cube.
+
+Every move is applied by ordinary code, and the answer is checked before it is
+printed. The network only decides which positions are worth keeping.
+
+### Training: the model makes its own answers
+
+The model was trained by approximate value iteration (the method behind
+[DeepCubeA](https://www.nature.com/articles/s42256-019-0070-z)). It starts from
+one free fact: the solved cube is 0 moves from solved. A cube is then 1 more
+than its best neighbour. The network learns from its own previous guesses, and
+the correct answers spread outward from the solved cube, one move further each
+round. No solver is involved at any point.
+
+An earlier version with the identical network was trained on distances from the
+Kociemba solver instead. It solved only 16% of full scrambles. Kociemba returns
+good solutions, not the shortest, and 47% of random cubes got the label "21", so
+there was nothing to learn from above distance 15. Changing only the labels took
+full scrambles from 16% to 96%.
+
+## Results
+
+Canonical RubiksBench scrambles (the same cubes every language model on the
+board is tested on), 25 cubes per scramble length:
+
+| Scramble length | beam 64 | beam 1024 |
+|---|---|---|
+| 3 | 100% | 100% |
+| 6 | 100% | 100% |
+| 10 | 100% | 100% |
+| 15 | 28% | 100% |
+| 20 | 24% | 96% |
+
+Past about 15 moves, every scramble is fully mixed, so 25, 50 or 100 moves are
+no harder than 20. Beam 1024 takes about 30 seconds per full scramble on one
+consumer graphics card. Beam 64 on a processor with no graphics card took about 5 seconds for a
+3-move scramble and about 50 seconds for a 6-move scramble; it is not practical
+for full scrambles.
+
+## Model
+
+- A Llama-shaped encoder, trained from scratch: 8 layers, width 512, 8 heads.
+  The architecture only; no Llama weights.
+- Input: `[BOS]` + 54 sticker tokens + `[SEP]`. Output: the mean of the final
+  layer, through one linear layer, to one number.
+- 33,577,985 parameters, float32.
+
+## Things that give much lower numbers
+
+| What you did | What you get at 20 moves |
+|---|---|
+| Beam 64 instead of 1024 | about 24% |
+| Always take the best-scored move, no search | 34% even at 10 moves |
+| Judge the network by its average error | about 1.5 moves, which looks poor; the search only needs the correct ranking of the 18 neighbours |
+
+## Limits
+
+- It is not fast. The Kociemba two-phase solver solves any cube in milliseconds.
+  This model exists to show that a small network can learn the task at all.
+- Solutions are not guaranteed to be shortest. Full scrambles average about 21
+  moves.
+- 3x3x3 only.

config.json

@@ -0,0 +1,8 @@
+{
+  "hidden": 512,
+  "layers": 8,
+  "heads": 8,
+  "vocab_size": 28,
+  "max_positions": 83,
+  "training_step": 27000
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16ed8f2bef7283f875e287ac3f434f706c72133e75a5222f40d29e185f001519
+size 134320420

solve.py

@@ -0,0 +1,172 @@
+"""
+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()