cayley-puzzles

Neural distance heuristics + large-batch beam search for two CayleyPy combinatorial-puzzle competitions on Kaggle:

Solver	Puzzle	Metric	Our best (own pipeline)
IHES Picture Cube (repo root)	3×3×3 picture cube, 1003 scrambles	total move count	~23,224
Megaminx (megaminx/)	dodecahedral puzzle, 1001 scrambles	total move count	77,214

Both numbers are what our own solver pipeline produces. Some submitted leaderboard scores were lower because they min-merged publicly shared community CSVs; those community contributions are excluded from the figures above and from the PROGRESS.md narratives — see each project's progress doc for the honest, ours-only progression.

The two solvers share the same idea and the same core library: train a small neural network to estimate "distance to solved", then run a very wide GPU/TPU beam search guided by that heuristic, then post-process the resulting paths.

Layout

cayley-puzzles/
├── src/cayley/        # shared core library (puzzle ops, model, training,
│                      #   Bellman refine, beam search, post-processing, verify)
├── scripts/           # IHES Picture Cube CLI entry points (train / solve / pp / submit)
├── configs/           # IHES training configs
├── data/              # IHES puzzle_info.json (the only data file shipped)
├── tests/             # unit tests for both solvers (model-free)
├── docs/              # IHES design notes
├── PROGRESS.md        # IHES ours-only progression
│
├── megaminx/          # the Megaminx solver (reuses src/cayley by duck typing)
│   ├── src/megaminx/  # Megaminx puzzle + solver-specific modules
│   ├── scripts/       # train / solve / bridge / two-phase / data-prep
│   ├── beam_lab/      # self-contained local beam-search benchmark
│   ├── tpu/           # JAX SPMD beam driver for Cloud/Kaggle TPU
│   ├── configs/       # Megaminx training configs
│   ├── data/          # Megaminx puzzle_info.json
│   ├── docs/          # Megaminx architecture + path-shortening notes
│   ├── PROGRESS.md    # Megaminx ours-only progression
│   └── README.md
│
├── .claude/           # Claude Code skills + slash commands (secrets scrubbed)
└── .agents/           # agent-skill mirrors (secrets scrubbed)

Why this shape: the Megaminx code reuses the IHES core by duck typing — megaminx.puzzle.Megaminx is structurally identical to cayley.puzzle.PictureCube, so the cayley.* training/search/verify modules accept it unchanged. The scripts wire this up with sys.path: IHES scripts add the repo-root src/, and Megaminx scripts add both megaminx/src and the repo-root src/. So src/cayley must stay at the repo root, one level above megaminx/. No PYTHONPATH setup is needed — the scripts handle it.

Setup

python -m venv .venv
# Windows:  .venv\Scripts\activate      Linux/macOS:  source .venv/bin/activate
pip install -e ".[torch,cayleypy,dev]"   # add ".[tpu]" for the JAX TPU beam

• Python 3.10+ (the original work used 3.14 on Windows). On Windows + Python 3.14, torch.compile needs triton-windows (not triton).
• An editable install is optional — the CLI scripts add their packages to sys.path themselves, so python scripts/02_solve.py ... works from the repo root without installing.

What is not in the repo (and how to get it)

Model checkpoints, BFS/PDB lookup tables, symmetry tables, datasets, and submission CSVs are large and reproducible, so they are git-ignored. Each project's README documents how to regenerate them from the shipped scripts, or where to download them (the competition data from Kaggle; published model weights from the Kaggle Models datasets). The only data files committed are the two puzzle_info.json generator/solved-state definitions.

Running the IHES Picture Cube solver

(The IHES solver lives at the repo root. Megaminx has its own guide in megaminx/README.md.)

# 1. (once) build the BFS-depth-5 table used by post-processing
python scripts/build_bfs_table.py --depth 5 --out data/bfs_table_d5.pkl

# 2. train the base distance model (E5: K_max=26 random walks, 8000 epochs, ~70 min on a 4090)
python scripts/01_train.py --config configs/small_e5_long.yaml --output models/small_e5

# 3. (optional) Bellman-refine it (E6: self-bootstrapped targets, warm-started from E5)
python scripts/05_bellman_refine.py --config configs/e6_bellman.yaml --output models/e6

# 4. solve all 1003 puzzles with the wide khoruzhii beam
python scripts/02_solve.py \
    --checkpoint models/small_e5/epoch_7999.pt \
    --out submissions/run.csv \
    --searcher khoruzhii --beam 65536 --max-steps 50 --bf16 \
    --fallback data/kociemba_fallback.csv

# 5. post-process (pair-cancel + BFS-d5 window replacement)
python scripts/post_process_submission.py \
    --in submissions/run.csv --out submissions/run_pp.csv \
    --bfs-table data/bfs_table_d5.pkl

# 6. ensemble several runs (per-puzzle min) over the Kociemba fallback floor
python scripts/combine_submissions.py \
    --candidates submissions/run_pp.csv submissions/other_run.csv \
    --fallback data/kociemba_fallback.csv \
    --out submissions/combined.csv

Notes:

• On Windows use .venv\Scripts\python.exe instead of python.
• Always pass --fallback data/kociemba_fallback.csv. The 02_solve.py default points at a sample-quality fallback (500K+ moves); the Kociemba fallback (38,440 moves) is the real floor.
• --searcher khoruzhii is the production beam (int8 state buffers → beam up to ~10^6). Do not use --mode advanced (CayleyPy's advanced beam returns no path). torch.compile is for training only here — variable beam batch sizes trigger recompiles at inference.
• Required-but-git-ignored data: data/test.csv (the 1003 scrambles — download from the Kaggle competition), data/kociemba_fallback.csv (a classical two-phase-solver baseline used as the fallback floor), and the data/bfs_table_d5.pkl you build in step 1.

Start here

• IHES Picture Cube → the section above for commands, and PROGRESS.md for the full results story + docs/ARCHITECTURE.md for the design.
• Megaminx → megaminx/README.md for the full train → solve → beam → post-process pipeline, and megaminx/PROGRESS.md for the progression.

Method in one paragraph

A ResMLPDistance network (embedding encoder → residual-MLP trunk → scalar value head) is trained to predict distance-to-solved, first on non-backtracking random walks (an upper bound on true distance) and then refined with Bellman bootstrapping (V(s) ← 1 + min_a V_target(child)), which removes the walk-depth bias. At solve time a beam search expands every state's neighbours, scores them with the network, and keeps the best B (B up to ~10^6 on GPU, more on TPU). int8 state buffers, a Q-distilled shortlister, inference-time symmetry ensembling, and meet-in-the-middle against a BFS shell make wide beams feasible. Finished paths are tightened with BFS-window replacement, tail re-solving, and (Megaminx) residual "bridge" compression. The Megaminx solver adds an AlphaZero dual-head variant and a Kociemba-style two-phase decomposition.