AlphaZero self-play

Deep dive

How it works

Parts:

• Policy-value network — state to priors P and value v.
• PUCT tree — visit counts N, values, priors per node.
• Self-play workers — play episodes, write (s, π, z).
• Replay window — recent games, sampled uniformly.
• Trainer — minimises the loss, publishes weights.

One simulation from the root:

1. Select. At each node pick argmax_a Q(s,a) + U(s,a), U = c(s)·P(s,a)·√N(s) / (1 + N(s,a)). Q is the mean value of the child from the mover's point of view, 0 (or a first-play-urgency value) if unvisited.
2. Expand + evaluate. At a leaf, run the network once: priors P over legal moves (softmax over the masked policy logits) and value v. Terminal states use the true result.
3. Backup. Add v to every node on the path, flipping the sign at each ply, and increment N.

Self-play move: add Dirichlet noise to the root priors, P ← (1 − ε)P + ε·Dir(α), run S simulations, set π(a) ∝ N(root, a)^{1/τ} (τ = 1 early, → 0 later), play a move from π, store (s, π). At game end attach z ∈ {−1, 0, 1} from each stored position's player-to-move view.

Training loop: many self-play workers write games into a replay window (last 1M games in AlphaZero). A trainer samples uniformly from it, and publishes new weights every ~1000 steps that the workers pick up. AlphaZero always self-plays the latest network; AlphaGo Zero kept the best checkpoint and gated new ones by a 55 % win rate.

self-play workers ──(s, π, z)──► replay window ──batches──► trainer ──weights──┐
      ▲  PUCT search + net                                                     │
      └───────────────────────────────────────────────────────────────────────┘

Code

// PUCT MCTS + self-play target generation on tic-tac-toe. `evaluate` stands in for the network.
#[derive(Clone)] struct Board { c: [i8; 9], to_play: i8 }        // +1 = X, -1 = O
impl Board {
    fn winner(&self) -> i8 {
        const L: [[usize; 3]; 8] = [[0,1,2],[3,4,5],[6,7,8],[0,3,6],[1,4,7],[2,5,8],[0,4,8],[2,4,6]];
        for l in L { let s = self.c[l[0]] + self.c[l[1]] + self.c[l[2]]; if s.abs() == 3 { return s / 3; } }
        0
    }
    fn legal(&self) -> Vec<usize> { (0..9).filter(|&i| self.c[i] == 0).collect() }
    fn terminal(&self) -> bool { self.winner() != 0 || self.legal().is_empty() }
    fn play(&mut self, m: usize) { self.c[m] = self.to_play; self.to_play = -self.to_play; }
}

struct Rng(u64);
impl Rng {
    fn f(&mut self) -> f64 { self.0 ^= self.0 << 13; self.0 ^= self.0 >> 7; self.0 ^= self.0 << 17;
                             ((self.0 >> 11) as f64 + 0.5) / (1u64 << 53) as f64 }
    fn normal(&mut self) -> f64 { (-2.0 * self.f().ln()).sqrt() * (6.283185307179586 * self.f()).cos() }
    fn gamma(&mut self, a: f64) -> f64 {                          // Marsaglia-Tsang; boost for a < 1
        if a < 1.0 { return self.gamma(a + 1.0) * self.f().powf(1.0 / a); }
        let (d, c) = (a - 1.0 / 3.0, 1.0 / (9.0 * a - 3.0).sqrt());
        loop { let x = self.normal(); let v = (1.0 + c * x).powi(3);
               if v > 0.0 && self.f().ln() < 0.5 * x * x + d - d * v + d * v.ln() { return d * v; } }
    }
    fn dirichlet(&mut self, a: f64, k: usize) -> Vec<f64> {
        let g: Vec<f64> = (0..k).map(|_| self.gamma(a)).collect();
        let s: f64 = g.iter().sum(); g.iter().map(|x| x / s).collect()
    }
}

/// Stand-in network: uniform priors, value from one random playout (for the player to move).
fn evaluate(b: &Board, rng: &mut Rng) -> (Vec<(usize, f64)>, f64) {
    let moves = b.legal();
    let pri = moves.iter().map(|&m| (m, 1.0 / moves.len() as f64)).collect();
    let mut sim = b.clone();
    while !sim.terminal() { let l = sim.legal(); sim.play(l[(rng.f() * l.len() as f64) as usize]); }
    (pri, (sim.winner() * b.to_play) as f64)
}

struct Node { n: u32, w: f64, p: f64, kids: Vec<(usize, usize)> } // w: value sum for the player who moved here

fn puct_search(root_b: &Board, sims: usize, rng: &mut Rng) -> Vec<(usize, u32)> {
    let (c_base, c_init, eps, alpha) = (19652.0, 1.25, 0.25, 0.3);
    let mut t = vec![Node { n: 0, w: 0.0, p: 1.0, kids: vec![] }];
    for sim in 0..sims {
        let (mut b, mut path, mut node) = (root_b.clone(), vec![0usize], 0usize);
        while !t[node].kids.is_empty() {                           // 1. select
            let np = t[node].n as f64;
            let cs = ((1.0 + np + c_base) / c_base).ln() + c_init;
            let (m, child) = *t[node].kids.iter().max_by(|x, y| {
                let s = |&(_, k): &(usize, usize)| {
                    let q = if t[k].n > 0 { t[k].w / t[k].n as f64 } else { 0.0 };
                    q + cs * t[k].p * np.sqrt() / (1.0 + t[k].n as f64)
                };
                s(x).partial_cmp(&s(y)).unwrap()
            }).unwrap();
            b.play(m); node = child; path.push(node);
        }
        let v = if b.terminal() { (b.winner() * b.to_play) as f64 } else {   // 2. expand + evaluate
            let (pri, v) = evaluate(&b, rng);
            let noise = if sim == 0 { rng.dirichlet(alpha, pri.len()) } else { vec![0.0; pri.len()] };
            for (i, (m, p)) in pri.into_iter().enumerate() {
                let p = if sim == 0 { (1.0 - eps) * p + eps * noise[i] } else { p }; // root noise
                t.push(Node { n: 0, w: 0.0, p, kids: vec![] });
                let id = t.len() - 1; t[node].kids.push((m, id));
            }
            v
        };
        let mut val = -v;                                          // 3. backup: leaf was reached by the opponent
        for &k in path.iter().rev() { t[k].n += 1; t[k].w += val; val = -val; }
    }
    t[0].kids.iter().map(|&(m, k)| (m, t[k].n)).collect()
}

fn main() {
    let mut rng = Rng(11);
    let (mut b, mut samples) = (Board { c: [0; 9], to_play: 1 }, Vec::new());
    let mut ply = 0;
    while !b.terminal() {
        let visits = puct_search(&b, 400, &mut rng);
        let total: u32 = visits.iter().map(|v| v.1).sum();
        let pi: Vec<(usize, f64)> = visits.iter().map(|&(m, n)| (m, n as f64 / total as f64)).collect();
        let mv = if ply < 2 {                                      // τ = 1 for the opening plies, then argmax
            let mut r = rng.f(); pi.iter().find(|x| { r -= x.1; r <= 0.0 }).unwrap_or(&pi[0]).0
        } else { visits.iter().max_by_key(|v| v.1).unwrap().0 };
        samples.push((b.c, b.to_play, pi));
        b.play(mv); ply += 1;
    }
    let z = b.winner();
    for (c, to_play, pi) in &samples {                             // training targets (s, π, z)
        let best = pi.iter().cloned().fold((9, 0.0), |a, x| if x.1 > a.1 { x } else { a });
        println!("board {:?}  z = {:+}  top π = move {} ({:.2})", c, z * to_play, best.0, best.1);
    }
}

Complexity and constants

• Self-play dominates compute: games × plies × S network evaluations. Chess at 800 simulations and order of 100+ plies per game is order of 10⁵ evaluations per game.
• Batch network calls: run many games (or many simulations with virtual loss) per worker and evaluate leaves in GPU batches of order of 64-1024. Batch size is the main throughput lever.
• Tree memory: one node per expanded child, order of 10-50 bytes each in a compact arena. 800 simulations × 30 children is tiny, so reuse the subtree under the played move.
• Training: one gradient step per batch of positions from the window. Keep the ratio of training samples consumed per sample generated roughly constant as you add workers: it is a key hyperparameter.

Efficiency

• Cost: one simulation is a tree descent (order of 100 ns) plus one network evaluation, which dominates. A game of 100+ plies at 800 simulations is order of 10⁵ evaluations, so self-play on accelerators bounds the whole run. AlphaZero searched about 80,000 positions per second in chess against about 70 million for Stockfish (Silver 2018).
• Levers:
  • Batch leaf evaluations across games and with virtual loss — batches of order of 64-1024 keep the accelerator busy (this page) — always.
  • Reuse the subtree under the played move — keeps the visits already spent there (this page) — every move.
  • Fewer simulations on most moves (playout cap randomization, within KataGo's bundle of changes) — the bundle matched ELF OpenGo with about 50x less compute (Wu 2019) — large self-play runs.
  • Gumbel root selection — policy improves with as few as 2 simulations (Danihelka 2022) — small search budgets.
• Naive → standard → state of the art:
  • One leaf per network call, one game at a time · accelerator mostly idle, latency per simulation
  • Many parallel games, batched leaves, virtual loss, subtree reuse · throughput set by batch inference
  • KataGo-style reduced-simulation self-play plus Gumbel search · about 50x less compute than ELF OpenGo (Wu 2019)

Implementation notes

• Policy head over a fixed move encoding (AlphaZero chess: 8×8×73 = 4672 move planes). Mask illegal moves before the softmax.
• Values: AlphaZero's pseudocode keeps values in [0, 1] with unvisited children at 0 (a pessimistic first-play urgency). Lc0 and KataGo use [−1, 1] or win/draw/loss heads and tune the FPU (parent Q minus a reduction).
• Exploit symmetries for data augmentation where the game has them (Go: 8 dihedral symmetries). Chess and shogi have none that are exact.
• Resignation saves compute. AlphaGo Zero set the threshold so false resignations stayed below 5 %, measured by disabling resignation in 10 % of games.
• Playout cap randomization (KataGo): most moves use a cheap search and are not recorded; a fraction get the full search and become training targets.

Pitfalls

• Sign errors in backup (value from the wrong player's view): the search plays for the opponent. Test on positions with a forced win in one.
• No root noise or no temperature: self-play collapses to a single line and the value head overfits to it.
• Training on visit counts that include the noise-induced exploration of forced-bad moves: KataGo prunes forced playouts from the policy target.
• Too small a replay window: the network overfits recent games and forgets. Too large: stale data from weak networks.
• Draw-heavy games with z = 0 everywhere give no value signal. Add WDL heads or auxiliary targets.

Variants and alternatives

In production

AlphaZero (chess, shogi, Go), Leela Chess Zero (Lc0), KataGo, Leela Zero, ELF OpenGo; AlphaTensor (new matrix-multiplication algorithms), AlphaDev (faster small-input sort routines); research frameworks OpenSpiel (alpha_zero) and DeepMind mctx (JAX, batched PUCT and Gumbel search).

Implementations

Sources

• Fawzi et al., Discovering faster matrix multiplication algorithms with reinforcement learning, Nature 2022. paper
• Mankowitz et al., Faster sorting algorithms discovered using deep reinforcement learning, Nature 2023. paper
• Silver et al., Mastering the Game of Go without Human Knowledge, Nature 2017. paper
• Silver et al., A General Reinforcement Learning Algorithm that Masters Chess, Shogi, and Go through Self-Play, Science 2018 (with published pseudocode). paper
• Schrittwieser et al., Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model, Nature 2020. paper
• Danihelka, Guez, Schrittwieser, Silver, Policy Improvement by Planning with Gumbel, ICLR 2022. paper · code
• Wu, Accelerating Self-Play Learning in Go, 2019 (KataGo). paper · code
• Anthony, Tian, Barber, Thinking Fast and Slow with Deep Learning and Tree Search, NeurIPS 2017. paper
• Rosin, Multi-armed Bandits with Episode Context, Annals of Mathematics and Artificial Intelligence 2011 (PUCB). paper
• Marsaglia and Tsang, A Simple Method for Generating Gamma Variables, ACM TOMS 2000. paper