world-models

How It Works

Two phases. Training is a single supervised pass: Stockfish labels positions, the student fits those labels. Inference is Monte Carlo Tree Search using the trained network as its prior. The network alone never picks a move — MCTS does, and the network just guides where MCTS looks.

1 · Training (one-shot, no self-play)

1. Stockfish self-plays games at search-depth 15. At every position, before the chosen move is played, Stockfish reports its top-8 candidate moves with centipawn scores.
2. Those scores become a soft distribution π over the 8 candidates (softmax at T = 1 pawn). The game outcome z ∈ {−1, 0, +1} is recorded from each side-to-move POV.
3. The student trains on (state, π, z) triples — cross-entropy on the policy head against π, MSE on the value head against z. No self-play, no MCTS at training time.

2 · The Network

20 residual blocks × 256 channels, ~24M parameters. Input: a 19-plane encoding of the current board (piece positions, side-to-move, castling rights, repetitions). Output: a 4,672-dim policy logits vector and a scalar value in [−1, +1]. Same architecture AlphaZero used — only the way we got the weights differs.

3 · Playing a Move (Network + MCTS, end-to-end)

The trained network is a prior over moves and a value estimator for leaf positions; the actual move comes from Monte Carlo Tree Search built on top of those two signals. One move at the root costs N simulations (800 for the routine eval, 4,000 for the search-saturated eval). Each simulation is:

1. Selection. Starting at the root (the current position), walk down the search tree by repeatedly picking the child that maximises Q + U: average value seen below this move + an exploration bonus proportional to the network's prior P and inversely proportional to how many times we've already visited this child. The network's policy is the bias; the running statistics correct it as evidence accumulates.
2. Expansion + evaluation. When you hit a node that hasn't been expanded yet, encode that position to 19 planes, run one network forward pass, and store the policy logits as priors for its children + the value estimate as the value of the leaf.
3. Backup. Propagate that leaf value back up the path you took, incrementing each parent's visit count and updating the running Q with the new estimate. Sign-flip at every ply because chess alternates side-to-move.

Repeat for N sims. Then pick the move at the root with the highest visit count (not highest Q; visit counts are MCTS's improved policy estimate). At higher sims/move, MCTS spends more rollouts refining the prior into a sharper visit distribution — that's why the same checkpoint scores 1,807 Elo at 800 sims and 2,084 Elo at 4,000 sims (see the search ablation).

Implementation: the encode is encode_board, the network forward is AlphaZeroNet.forward, and the MCTS loop is run_mcts / run_mcts_batched. The selection formula is straight from the AlphaGo Zero paper: U(s,a) = c_puct · P(s,a) · √Σ_b N(s,b) / (1 + N(s,a)).

→ Full method, including the loss formula and Elo bisection