Multi-armed bandits

Deep dive

How it works

Parts:

• Arm statistics — pulls n_i and reward sum, or Beta counts (a_i, b_i); 16 bytes per arm.
• Policy — maps statistics (and context x, step t) to one arm.
• Feedback channel — reward arrives after the choice, possibly late or in batches.
• Decision log — arm, context, propensity (probability of that arm), reward; feeds off-policy evaluation.

Regret after T steps: R(T) = T·μ* − Σ μ_{arm_t}.

1. epsilon-greedy: with probability ε play a random arm, else the best empirical mean. Fixed ε explores forever (linear regret); ε_t ∝ 1/t gives logarithmic regret if the smallest gap is known (Auer 2002).
2. UCB1: play each arm once, then argmax x̄_i + sqrt(2 ln t / n_i). Rewards in [0, 1]; the bound is 8·ln n/Δ_i per suboptimal arm plus a constant.
3. Thompson sampling: keep a posterior per arm (Beta(a, b) for 0/1 rewards, Normal for Gaussian ones), draw one sample per arm, play the largest; update a += r, b += 1 − r. No exploration parameter (tutorial).
4. Contextual (LinUCB): per arm keep A = I + Σ xxᵀ, b = Σ r·x; score θᵀx + α·sqrt(xᵀA⁻¹x) with θ = A⁻¹b (Li 2010). Linear Thompson draws θ ~ N(A⁻¹b, v²A⁻¹) instead.
5. Log every decision with its propensity before the reward is known.
6. Drift and delay: discount old counts (n ← γn) or use a sliding window; with delayed or batched feedback keep sampling from the current posterior, which spreads traffic by itself.

        context x ─┐
 stats ──► policy ─┴─► arm a ──► world ──► reward r (maybe late)
   ▲          │ propensity p(a|x)              │
   │          ▼                                │
   │      decision log ──► off-policy eval     │
   └──────────── update(a, r) ◄────────────────┘

Code

Rust, std only: xorshift64* RNG, Gamma by Marsaglia-Tsang (2000), Beta as X/(X+Y); regret is pseudo-regret. cargo test --release passes.

pub struct Rng(u64);
impl Rng {
    pub fn new(seed: u64) -> Self { Rng(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1) }
    pub fn next_u64(&mut self) -> u64 { // xorshift64*
        let mut x = self.0;
        x ^= x >> 12; x ^= x << 25; x ^= x >> 27;
        self.0 = x;
        x.wrapping_mul(0x2545_F491_4F6C_DD1D)
    }
    pub fn f64(&mut self) -> f64 { (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64 } // [0,1)
    fn normal(&mut self) -> f64 { // Box-Muller; 1-u keeps ln() finite
        let (u, v) = (1.0 - self.f64(), self.f64());
        (-2.0 * u.ln()).sqrt() * (std::f64::consts::TAU * v).cos()
    }
    fn gamma(&mut self, shape: f64) -> f64 { // Marsaglia-Tsang, shape >= 1
        let d = shape - 1.0 / 3.0;
        let c = 1.0 / (9.0 * d).sqrt();
        loop {
            let x = self.normal();
            let v = (1.0 + c * x).powi(3);
            if v <= 0.0 { continue; }
            if (1.0 - self.f64()).ln() < 0.5 * x * x + d - d * v + d * v.ln() { return d * v; }
        }
    }
    pub fn beta(&mut self, a: f64, b: f64) -> f64 { let (x, y) = (self.gamma(a), self.gamma(b)); x / (x + y) }
}

pub trait Policy {
    fn select(&mut self, t: u64, rng: &mut Rng) -> usize;
    fn update(&mut self, arm: usize, reward: f64);
}
fn argmax(v: impl Iterator<Item = f64>) -> usize {
    v.enumerate().fold((0, f64::NEG_INFINITY), |b, (i, x)| if x > b.1 { (i, x) } else { b }).0
}

pub struct Random(pub usize);
impl Policy for Random {
    fn select(&mut self, _: u64, rng: &mut Rng) -> usize { (rng.next_u64() % self.0 as u64) as usize }
    fn update(&mut self, _: usize, _: f64) {}
}
/// epsilon-greedy (eps > 0) and UCB1 (eps = None) share pull counts and reward sums.
pub struct Counts { n: Vec<f64>, sum: Vec<f64>, eps: Option<f64> }
impl Counts { pub fn new(k: usize, eps: Option<f64>) -> Self { Counts { n: vec![0.0; k], sum: vec![0.0; k], eps } } }
impl Policy for Counts {
    fn select(&mut self, t: u64, rng: &mut Rng) -> usize {
        if let Some(i) = self.n.iter().position(|&n| n == 0.0) { return i; } // try each arm once
        let mean = |i: usize| self.sum[i] / self.n[i];
        match self.eps {
            Some(e) if rng.f64() < e => (rng.next_u64() % self.n.len() as u64) as usize,
            Some(_) => argmax((0..self.n.len()).map(mean)),
            None => argmax((0..self.n.len()).map(|i| mean(i) + (2.0 * (t as f64).ln() / self.n[i]).sqrt())),
        }
    }
    fn update(&mut self, arm: usize, r: f64) { self.n[arm] += 1.0; self.sum[arm] += r; }
}
pub struct Thompson { a: Vec<f64>, b: Vec<f64> } // Beta(1,1) priors
impl Thompson { pub fn new(k: usize) -> Self { Thompson { a: vec![1.0; k], b: vec![1.0; k] } } }
impl Policy for Thompson {
    fn select(&mut self, _: u64, rng: &mut Rng) -> usize { argmax((0..self.a.len()).map(|i| rng.beta(self.a[i], self.b[i]))) }
    fn update(&mut self, arm: usize, r: f64) { self.a[arm] += r; self.b[arm] += 1.0 - r; }
}

/// Mean pseudo-regret (sum of gaps of the arms played) at steps t1 and t2 over `seeds` runs.
pub fn regret(make: &dyn Fn() -> Box<dyn Policy>, means: &[f64], t1: u64, t2: u64, seeds: u64) -> (f64, f64) {
    let best = means.iter().cloned().fold(0.0, f64::max);
    let (mut r1, mut r2) = (0.0, 0.0);
    for seed in 0..seeds {
        let (mut rng, mut pol, mut reg) = (Rng::new(seed + 1), make(), 0.0);
        for t in 1..=t2 {
            let arm = pol.select(t, &mut rng);
            pol.update(arm, if rng.f64() < means[arm] { 1.0 } else { 0.0 });
            reg += best - means[arm];
            if t == t1 { r1 += reg; }
        }
        r2 += reg;
    }
    (r1 / seeds as f64, r2 / seeds as f64)
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn regret_ordering_and_sublinear_growth() {
        let means = [0.10, 0.20, 0.30, 0.40, 0.55];
        let k = means.len();
        let run = |m: &dyn Fn() -> Box<dyn Policy>| regret(m, &means, 10_000, 20_000, 50);
        let rand = run(&|| Box::new(Random(k)));
        let eps = run(&|| Box::new(Counts::new(k, Some(0.1))));
        let ucb = run(&|| Box::new(Counts::new(k, None)));
        let ts = run(&|| Box::new(Thompson::new(k)));
        for r in [ucb, ts] {
            assert!(r.0 < eps.0 && r.0 < rand.0, "{r:?} {eps:?} {rand:?}");
            assert!(r.1 < 1.5 * r.0, "R(2T) must be well under 2x R(T): {r:?}");
        }
        assert!(eps.1 > 1.7 * eps.0); // fixed epsilon keeps paying: linear
    }
}

Measured, 50 seeds, regret at T = 10,000 (20,000). Easy env [.10,.20,.30,.40,.55]: random 2400 (4800), epsilon-0.1 270 (510), UCB1 200 (238), Thompson 39 (43). Hard env, 10 arms from .10 to .40 (gaps down to .02): random 1729, epsilon-0.1 285, UCB1 524, Thompson 116. UCB1 is still exploring at T = 10,000 there and loses to epsilon-greedy; its growth is slowing (R(2T)/R(T) 1.33 vs 1.69).

Complexity and constants

• Regret: UCB1 ≤ 8·Σ ln n/Δ_i + O(ΣΔ_j) (Auer 2002); Thompson O(ln T) (Agrawal-Goyal); fixed-ε greedy Θ(εT).
• Lower bound: every consistent policy has regret at least about Σ Δ_i / KL(μ_i, μ*) · ln T on Bernoulli arms (Lai-Robbins 1985; book Lattimore-Szepesvári).
• Per decision O(K) scores; Thompson draws two Gamma variates per arm. LinUCB: O(d²) per arm with rank-1 updates of A⁻¹, O(d³) if re-inverted.
• Cost drivers: K, gap sizes, feedback delay, posterior sampling.

Efficiency

• Cost: one decision, K = 10, Apple M4 Pro, release build, one thread, measured here on a busy machine (order of): epsilon-greedy about 16-22 ns, UCB1 about 22-28 ns, Thompson (Beta via Gamma) about 450-550 ns, each including one RNG draw for the simulated reward. State 16 bytes per arm. Bound: compute for Thompson; in a service the log write and network dominate.
• Levers:
  • Gaussian approximation of the Beta posterior — about 130 ns instead of about 500 ns per decision (measured here; regret 37 vs 43 at 2T) — when counts are large.
  • Sample once per batch and split requests by the sampled win probabilities — policy cost per batch, not per request (by construction) — when feedback arrives in batches anyway.
  • • Naive → standard → state of the art:
  • fixed uniform split · linear regret, 2400 at T = 10,000 here
  • epsilon-greedy 0.1 · linear, 270 here
  • UCB1 / Thompson · logarithmic, 200 / 39 here; contextual versions with logged propensities and off-policy evaluation in production (Vowpal Wabbit)

Implementation notes

• Log the propensity at decision time. For Thompson estimate it by Monte Carlo, or keep a small uniform floor (1-5 %): inverse-propensity weights explode when an arm's probability is near 0.
• One seeded generator per worker (PRNGs); non-binary rewards: scale to [0, 1] for UCB1 or use a Gaussian posterior.
• The Gamma code needs shape ≥ 1; for priors below 1 use G(a) = G(a+1)·U^(1/a).
• Priors: Beta(1,1) is flat; centre it near the baseline rate to calm early swings.

Pitfalls

• Delayed conversions: a pending pull counted as failure biases against recent arms; count it once its window closed.
• Drift (weekday effects, novelty): plain averages lag; discount or window.
• Arm means under adaptive sampling are biased: do not read them as an A/B result.

Variants and alternatives

In production

Yahoo! Front Page news (LinUCB paper: 33 million events, 12.5 % click lift over a context-free bandit, offline replay); Vowpal Wabbit contextual bandits (--cb_explore, --cb_explore_adf; epsilon, bagging, cover, softmax).

Implementations

Sources

• Auer, Cesa-Bianchi, Fischer, Finite-time Analysis of the Multiarmed Bandit Problem, Machine Learning 47, 2002. paper
• Lai, Robbins, Asymptotically efficient adaptive allocation rules, Advances in Applied Mathematics 6, 1985. paper
• Chapelle, Li, An Empirical Evaluation of Thompson Sampling, NIPS 2011. paper
• Agrawal, Goyal, Analysis of Thompson Sampling for the multi-armed bandit problem, COLT 2012. paper
• Li, Chu, Langford, Schapire, A Contextual-Bandit Approach to Personalized News Article Recommendation, WWW 2010. paper
• Russo, Van Roy, Kazerouni, Osband, Wen, A Tutorial on Thompson Sampling, Foundations and Trends in ML 11(1), 2018. paper
• Lattimore, Szepesvári, Bandit Algorithms, Cambridge University Press 2020. book