Experience replay

Deep dive

How it works

In: a stream of items (a transition, a game position with targets, a logged event) and a loss value per item from the trainer · Out: mini-batches of indices, plus importance weights if prioritized · State: the item ring, a write cursor, one priority per slot (sum-tree) · Cost: push O(1); uniform sample O(1); prioritized sample and update O(log N)

1. Write. Append at the cursor, advance modulo N; when full the oldest item is overwritten, so the buffer holds exactly the newest N.
2. Uniform sample. Draw indices uniformly over filled slots. Neighbours in time are rarely drawn together, which decorrelates the batch.
3. Priority. After training on item i with error δ, set p_i = (|δ| + ε)^α. ε keeps every item reachable; α = 0 is uniform. A new item gets the current maximum priority so it is trained on at least once.
4. Proportional draw. P(i) = p_i / Σp. A sum-tree holds partial sums, so a prefix-sum search finds the item in O(log N). Stratified: split the total mass into k equal segments, one draw per segment.
5. Rank-based draw. p_i = 1/rank(i) by |δ|: robust to outliers because only the order matters (Schaul 2015).
6. Importance weights. Prioritizing changes the data distribution, so scale each loss term by w_i = (N·P(i))^-β, divided by the batch maximum. Anneal β up to 1.
7. Sliding window. For self-play the unit is a game: keep the newest games, sample positions uniformly. Old games came from weaker policies, so window length trades diversity against staleness.

Sample reuse ratio = batch × updates per step / items written per step. DQN on Atari did one update of 32 transitions per 4 new transitions in a buffer of the last 10^6 (as restated by Schaul 2015), so each transition is drawn about 32 / 4 = 8 times before it is overwritten. Too high overfits the buffer; too low wastes data. Hold the ratio on purpose by throttling actors or trainer.

Decorrelation. Random draws across many past steps break the temporal chain and average over many behaviours, so the training distribution drifts smoothly (Mnih 2015).

Code

Ring buffer plus sum-tree prioritized sampler, std only; compiled and run, the asserts check the ring wrap-around, the analytic draw probability and the weights.

pub struct Rng(u64); // xorshift64; seed must be non-zero
impl Rng {
    pub fn next(&mut self) -> u64 {
        self.0 ^= self.0 << 13; self.0 ^= self.0 >> 7; self.0 ^= self.0 << 17; self.0
    }
    pub fn f64(&mut self) -> f64 { (self.next() >> 11) as f64 / (1u64 << 53) as f64 } // [0,1)
}

pub struct Replay<T> {
    data: Vec<T>,     // ring: slot `head` is overwritten next once full
    cap: usize,
    head: usize,
    n: usize,         // leaves of the sum-tree, cap rounded up to a power of two
    tree: Vec<f64>,   // tree[1] = total; leaf i at tree[n + i]; unused leaves stay 0
    max_p: f64,       // new items get the current max priority, so each is seen at least once
    alpha: f64,
    eps: f64,
}

impl<T> Replay<T> {
    pub fn new(cap: usize, alpha: f64, eps: f64) -> Self {
        let n = cap.next_power_of_two();
        Replay { data: Vec::with_capacity(cap), cap, head: 0, n, tree: vec![0.0; 2 * n], max_p: 1.0, alpha, eps }
    }
    pub fn len(&self) -> usize { self.data.len() }

    pub fn push(&mut self, item: T) -> usize {
        let i = self.head;
        if self.data.len() < self.cap { self.data.push(item) } else { self.data[i] = item } // overwrite oldest
        self.head = (i + 1) % self.cap;
        self.set_p(i, self.max_p);
        i
    }
    pub fn get(&self, i: usize) -> &T { &self.data[i] }

    fn set_p(&mut self, i: usize, p: f64) {
        let mut k = self.n + i;
        self.tree[k] = p;
        while k > 1 { k /= 2; self.tree[k] = self.tree[2 * k] + self.tree[2 * k + 1]; } // recompute: no float drift
    }
    /// After a train step: priority = (|td_error| + eps)^alpha.
    pub fn update(&mut self, i: usize, td_error: f64) {
        let p = (td_error.abs() + self.eps).powf(self.alpha);
        self.max_p = self.max_p.max(p);
        self.set_p(i, p);
    }
    fn find(&self, mut prefix: f64) -> usize { // leaf whose cumulative range holds `prefix`
        let mut k = 1;
        while k < self.n {
            let l = 2 * k;
            if prefix < self.tree[l] || self.tree[l + 1] == 0.0 { k = l } else { prefix -= self.tree[l]; k = l + 1 }
        }
        k - self.n
    }
    pub fn sample_uniform(&self, rng: &mut Rng) -> usize { (rng.next() % self.len() as u64) as usize }

    /// Stratified proportional sampling: one draw per equal slice of the total mass.
    /// Returns (index, importance weight), weights scaled so the batch maximum is 1.
    pub fn sample(&self, k: usize, beta: f64, rng: &mut Rng) -> Vec<(usize, f64)> {
        let total = self.tree[1];
        let seg = total / k as f64;
        let mut out: Vec<(usize, f64)> = (0..k).map(|j| {
            let i = self.find((j as f64 + rng.f64()) * seg);
            let prob = self.tree[self.n + i] / total;
            (i, (self.len() as f64 * prob).powf(-beta)) // w = (N * P(i))^-beta
        }).collect();
        let wmax = out.iter().map(|x| x.1).fold(0.0, f64::max);
        for x in out.iter_mut() { x.1 /= wmax }
        out
    }
}

fn main() {
    let mut r = Replay::new(6, 0.6, 0.01); // cap 6 -> 8 leaves
    for x in 0..10u32 { r.push(x); }
    assert_eq!(r.len(), 6);
    let mut held: Vec<u32> = (0..6).map(|i| *r.get(i)).collect();
    held.sort();
    assert_eq!(held, vec![4, 5, 6, 7, 8, 9]); // oldest four overwritten
    assert_eq!(*r.get(0), 6); // slot 0 was rewritten by item 6

    for i in 0..6 { r.update(i, 0.0); }
    r.update(3, 10.0);
    let mut rng = Rng(0x9E3779B97F4A7C15);
    let (mut hits, draws) = (0usize, 20_000usize);
    for _ in 0..draws { if r.sample(1, 0.4, &mut rng)[0].0 == 3 { hits += 1 } }
    let (hi, lo) = (10.01f64.powf(0.6), 0.01f64.powf(0.6));
    let expect = hi / (hi + 5.0 * lo);
    let got = hits as f64 / draws as f64;
    assert!((got - expect).abs() < 0.02, "got {got} expect {expect}");

    let b = r.sample(64, 0.4, &mut rng);
    assert!(b.iter().any(|x| x.0 != 3));
    let wmax = b.iter().map(|x| x.1).fold(0.0, f64::max);
    assert!((wmax - 1.0).abs() < 1e-12);
    for (i, w) in &b { if *i == 3 { assert!(*w < 0.5) } } // frequent item gets a small weight
    let mut cnt = [0usize; 6];
    for _ in 0..12_000 { cnt[r.sample_uniform(&mut rng)] += 1 }
    assert!(cnt.iter().all(|&c| (1700..2300).contains(&c)));
    println!("ok: P(item 3) = {got:.3} vs {expect:.3}");
}

Complexity and constants

• Prioritized push or update is log2(N) node writes (20 for N = 1M); a draw is as many reads, so a batch of 32 is about 640 node reads. Tree memory is 2N f64 = 16 B per item (8 B with f32).
• A uniform draw is one random number. The real cost is the gather: each item is a random read, likely a cache miss, so cost is item size × batch.
• Item memory is N × item size. Storing each 84×84 Atari frame once (7,056 B) instead of four stacked per transition saves 4×; 1M frames × 7,056 B is about 7 GB.
• Updating parent sums by deltas accumulates float error; the code recomputes left + right along the path.

Efficiency

Cost: per item one write plus one tree-path update; per batch k tree descents plus k random gathers. The bound is memory latency and bandwidth for large items (frames, board planes); the trainer step usually dominates both. Levers:

• Store each frame or board once, rebuild stacks at sample time — up to 4x memory with 4-frame stacks (arithmetic above) — large observations.
• Batch priority updates and take one lock per batch — avoids k locks and k scattered writes — threaded actors and learners.
• Stratified draw — lower variance than k independent draws (Schaul 2015) — prioritized replay.
• Memory-mapped or sharded storage; actors write whole games or chunks, the sampler streams them — no per-item syscalls — windows larger than RAM.
• α = 0 skips the tree — O(1) draws — when error is a poor signal. Naive → standard → state of the art:
• Train on the newest item only · one update per item, correlated, no reuse.
• Ring buffer plus uniform batches · O(1) per op, about 8 draws per item at DQN settings.
• Sum-tree prioritized replay, stratified draws, importance weights · O(log N) per op, learns faster than uniform on Atari (Schaul 2015); replay servers such as Reverb move storage and rate limiting out of the trainer process.

Implementation notes

• Power-of-two leaf count makes the tree an implicit array: children of k are 2k and 2k+1, leaf i at n+i. Unused leaves hold 0.
• Float roundoff can land a search on an empty leaf; the code goes right only when the right child has mass.
• % len in sample_uniform has negligible bias for len far below 2^64.
• Across threads: one lock per batch, or a single-producer ring (SPSC ring buffer).
• Paper defaults: proportional α = 0.6, β 0.4 → 1; rank-based α = 0.7, β 0.5 → 1 (Schaul 2015).
• Plain windows of games need no tree; keep priorities only if you reweight.

Pitfalls

• Sampling before a warm-up count is reached gives near-duplicate batches.
• Omitting importance weights, or annealing β too early, changes the objective; the learner overfits high-error (often noisy) items.
• Priorities go stale: only sampled items are refreshed.
• Not storing done explicitly corrupts terminal targets.
• A window too long trains on policies the learner no longer resembles; too short gives a narrow, correlated set.
• Replay with an on-policy loss without correction.

Variants and alternatives

In production

• DQN on Atari: 1M-item replay, batch 32 (Q-learning and DQN).
• AlphaZero-style engines train on recent self-play games (AlphaZero self-play); Leela Chess Zero's training config selects the newest num_chunks chunks of self-play data.
• Off-policy control agents for robot manipulation and locomotion keep a buffer of recent environment steps.
• Streaming recommenders rehearse a window of logged interactions beside the live stream.

Implementations

Sources

• Long-Ji Lin, Self-improving reactive agents based on reinforcement learning, planning and teaching, Machine Learning 8, 1992, pp. 293-321. paper
• Mnih et al., Human-level control through deep reinforcement learning, Nature 518, 2015, pp. 529-533. paper
• Schaul, Quan, Antonoglou, Silver, Prioritized Experience Replay, ICLR 2016. paper
• Silver et al., A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play, Science 362, 2018, pp. 1140-1144. paper
• Leela Chess Zero, lczero-training (training pipeline, example config). code