Bandits, exploration, and sampling

Strip an MDP down to a single state and you get a bandit: no transitions to reason about, no discounting, just the purest form of the one dilemma that never goes away in reinforcement learning. Do I take the action that looks best right now, or do I try something else to learn whether "best" is really best? That is exploration versus exploitation, and it is not a footnote. It is the reason temperature exists as a decoding knob, it is the reason RL fine-tuning can quietly strangle itself through "entropy collapse," and it is why every serious RLVR recipe carries some term whose only job is to keep the policy from getting too sure of itself too soon. This chapter derives the bandit machinery lightly and spends most of its energy on the LLM translation, because the translation is where the money is.

Theory

The $k$-armed bandit

You face $k$ actions (arms). Each time you pull arm $a$ you receive a reward drawn from a fixed but unknown distribution with mean

$$q_{\ast }(a)=\mathbb{E}[R_t\mid A_t=a].\text{\hspace{1em}(1)}$$

This $q_{\ast }(a)$ is the action value from Chapter 5.2, collapsed to a single state so I can drop the state argument. If you knew every $q_{\ast }(a)$ you would pull $\arg {\max }_a{q}_{\ast }(a)$ forever and be done. You do not, so you estimate. Let $Q_t(a)$ be your estimate of $q_{\ast }(a)$ at time $t$. The natural estimator is the sample average of the rewards seen so far from arm $a$:

$$Q_t(a)=\frac{{\sum }_{i=1}^{t-1}{R}_{i}1[A_i=a]}{{\sum }_{i=1}^{t-1}1[A_i=a]}.\text{\hspace{1em}(2)}$$

By the law of large numbers $Q_t(a)\to q_{\ast }(a)$ as the count of pulls of $a$ grows, but only if you keep pulling $a$. That "if" is the whole problem.

The exploration-exploitation dilemma and three ways to handle it

If you always pull $\arg {\max }_a{Q}_t(a)$ (pure greedy), you can lock onto an arm that got lucky early and never discover a better one whose first few pulls were unlucky. You need to explore. Three classic mechanisms, in rising order of sophistication:

$\epsilon$-greedy. With probability $1-\epsilon$ exploit (pull the greedy arm), with probability $\epsilon$ explore (pull a uniformly random arm). Dead simple, and it works, but it explores indiscriminately: a clearly terrible arm gets pulled as often as a promising-but-uncertain one.

Upper confidence bound (UCB). Explore in proportion to how uncertain you are, not uniformly. Pick

$$A_t=\arg \underset{a}{\max }\left[Q_t(a)+c\sqrt{\frac{\ln t}{N_t(a)}}\right],\text{\hspace{1em}(4)}$$

where $N_t(a)$ is the number of pulls of $a$ so far and $c>0$ tunes the exploration bonus. Arms pulled rarely have a large bonus and get tried; arms pulled often have a small bonus and are trusted. This "optimism in the face of uncertainty" is provably near-optimal for stationary bandits.

Boltzmann / softmax. Convert estimated values into a probability distribution and sample. This is the one that matters for LLMs, so it gets its own equation:

$$\pi (a)=\frac{\exp (Q_t(a)/\tau )}{{\sum }_b\exp (Q_t(b)/\tau )}.\text{\hspace{1em}(5)}$$

The parameter $\tau >0$ is a temperature. As $\tau \to 0$ the softmax concentrates all mass on the greedy arm (pure exploitation); as $\tau \to \infty$ it flattens to uniform (pure exploration). Look at (5) and then look at how a language model picks its next token, and you should feel a jolt of recognition, because they are the same equation.

Regret, lightly

How do you score an exploration strategy? By regret: the cumulative gap between the reward you got and the reward the best arm would have given.

$$L_T=\sum _{t=1}^T(q_{\ast }(a^{\ast })-q_{\ast }(A_t)),a^{\ast }=\arg \underset{a}{\max }{q}_{\ast }(a).\text{\hspace{1em}(6)}$$

Pure greedy can suffer linear regret: if it locks onto a suboptimal arm, the per-step gap never closes and $L_T$ grows like $T$. Good strategies achieve logarithmic regret, $L_T=O(\ln T)$, and the Lai-Robbins result says you cannot do better than logarithmic in general. I will not prove that bound (it is a KL-divergence argument on how many pulls it takes to distinguish two reward distributions), but the intuition is worth stating plainly.

Temperature as the LLM's exploration knob

A language model produces a vector of logits $z=(z_1,\dots ,z_{|\mathcal{V}|})$ over the vocabulary at each step. The sampling distribution is a tempered softmax, exactly (5) with the logits playing the role of the value estimates:

$${\pi }_{\tau }(a\mid s)=\frac{\exp (z_a/\tau )}{{\sum }_b\exp (z_b/\tau )}.\text{\hspace{1em}(7)}$$

So temperature is literally Boltzmann exploration bolted onto next-token prediction. At $\tau =0$ (greedy decoding, argmax) the model exploits: it always emits its single most likely token, deterministic and safe and often bland. Crank $\tau$ up and the model explores, sampling rarer tokens, which is what you want during RL rollouts because you cannot learn from an answer you never generate. If your rollout temperature is too low, every completion for a given prompt is nearly identical, the reward signal has no variation to learn from, and gradient estimates collapse toward zero. If it is too high, completions are gibberish and rewards are uniformly low for the wrong reason. Choosing rollout temperature is choosing an exploration rate, and it is one of the few genuinely load-bearing hyperparameters in RLVR.

There is a clean derivative fact that makes the knob's behavior precise: temperature is a smooth interpolation between the uniform distribution and the point mass on the argmax.

Top-k and top-p: truncation is exploration control too

Temperature is the smooth exploration dial, but in practice LLM decoding usually stacks a second, sharper control on top of it: truncated sampling. Top-k keeps only the $k$ highest-probability tokens and renormalizes the softmax over just those, zeroing the rest. Top-p (nucleus) sampling keeps the smallest set of tokens whose cumulative probability first exceeds $p$, so the number of candidates flexes with how peaked the distribution is. Both are, in the bandit language, hard restrictions of the action set before sampling: they forbid the long tail of low-value arms entirely rather than merely making them unlikely.

This matters for RL rollouts in a way that is easy to get wrong. Truncation removes exploration in the tail. If you generate RLVR rollouts with an aggressive top-k of, say, $k=1$, you have made decoding greedy no matter what temperature says, and every completion for a prompt will be identical, giving you zero learnable variation. Conversely, leaving the full tail open at high temperature can let the policy wander into token sequences it will never actually deploy at inference time, creating a train-inference mismatch. The practical rule I follow is to make rollout sampling match, or be slightly more exploratory than, the decoding settings you will evaluate under, so that the reward signal reflects behavior you actually intend to keep. Truncation is not a rounding detail; it is a second exploration knob sitting in series with temperature.

Contextual bandits: the bridge to full RL

The pure $k$-armed bandit has one state, which is why it is the cleanest place to study exploration. A contextual bandit adds a state that you observe before choosing, but still with no transitions: each round you see a context $x$, pick an action, get a reward, and the context for the next round is drawn independently. Action values become context-dependent, $q_{\ast }(x,a)$, and the policy becomes a conditional distribution $\pi (a\mid x)$. This is exactly the sequence-level view of LLM generation from Chapter 5.1: the prompt is the context, the whole completion is the single action, the verifier is the reward, and there is no transition to a "next prompt" that your action influenced. GRPO, at its core, treats reasoning-model training as a contextual bandit, which is why the exploration machinery of this chapter, temperature and entropy, is the machinery that actually governs it, and why the value-function and multi-step credit-assignment apparatus of Chapter 5.2 can be sidestepped. The lab in Chapter 5.5 is literally a contextual bandit for that reason.

Entropy, and how RL fine-tuning collapses it

The right way to measure how much a policy explores is its entropy. For the next-token distribution,

$$H(\pi (\cdot \mid s))=-\sum _a\pi (a\mid s)\log \pi (a\mid s).\text{\hspace{1em}(8)}$$

Entropy is maximized ($\log |\mathcal{V}|$ nats) by the uniform distribution and is zero for a deterministic policy. Here is the failure mode that bites RL fine-tuning, and it is worth taking seriously because it is not hypothetical, it is the default trajectory of a naive run.

Policy-gradient training pushes probability mass toward actions that earned high reward. That is the whole idea. But nothing in the bare objective pushes back, so the policy keeps sharpening: the entropy (8) falls, the rollouts for a given prompt become more and more alike, exploration dies, and the model converges prematurely onto whatever reasoning pattern happened to work early. This is entropy collapse. Once entropy is near zero the gradient signal starves, because with no variation among sampled completions there is no advantage to learn from (Chapter 5.5 makes "advantage needs variation" precise). The model gets stuck, sometimes at a genuinely mediocre policy, and no amount of further training rescues it because it can no longer generate the alternatives it would need to discover something better.

Tooling

The tool is the simulator: a ten-armed testbed you can run in a second to feel the difference between exploiting and exploring, and a temperature sweep that connects the bandit softmax directly to LLM decoding. Reuse the uv project from Chapter 5.1 (uv add numpy matplotlib if you have not already).

Lab

Two artifacts. First, an $\epsilon$-greedy vs greedy comparison on a ten-armed bandit that prints how often pure greedy locks onto the wrong arm. Second, a temperature sweep over a fixed logit vector that prints entropy as a function of $\tau$, making (7) and (8) concrete and previewing entropy collapse.

"""Two demos in one file.

(A) 10-armed bandit: greedy vs epsilon-greedy, averaged over runs.
    Shows that refusing to explore can lock onto a suboptimal arm.
(B) Temperature sweep: entropy of a tempered softmax vs tau, the
    exploration knob an LLM exposes as its sampling temperature.

Run:  uv run python labs/bandit_and_temperature.py
"""
from __future__ import annotations
import numpy as np


# ---------- (A) ten-armed testbed ----------
def run_bandit(epsilon: float, steps: int, rng: np.random.Generator):
    k = 10
    q_star = rng.normal(0.0, 1.0, size=k)      # true means
    best = int(np.argmax(q_star))
    Q = np.zeros(k)                            # estimates
    N = np.zeros(k)                            # pull counts
    optimal_pulls = 0
    for t in range(steps):
        if rng.random() < epsilon:
            a = rng.integers(k)                # explore
        else:
            a = int(np.argmax(Q))              # exploit
        reward = rng.normal(q_star[a], 1.0)    # noisy reward
        N[a] += 1
        Q[a] += (reward - Q[a]) / N[a]         # incremental mean, eq (3)
        optimal_pulls += (a == best)
    return optimal_pulls / steps


def bandit_demo(runs=500, steps=1000):
    rng = np.random.default_rng(1)
    for eps in (0.0, 0.01, 0.1):
        rate = np.mean([run_bandit(eps, steps, rng) for _ in range(runs)])
        label = "greedy" if eps == 0 else f"eps-greedy(e={eps})"
        print(f"{label:<22} optimal-arm rate over {steps} steps: {rate:.3f}")


# ---------- (B) temperature = LLM exploration knob ----------
def softmax(z, tau):
    z = z / tau
    z = z - z.max()                            # numerical stability
    e = np.exp(z)
    return e / e.sum()


def entropy(p):
    p = p[p > 0]
    return float(-(p * np.log(p)).sum())       # nats, eq (8)


def temperature_demo():
    # A plausible next-token logit vector: one clear favorite, a few
    # contenders, a long tail. Same shape an LLM produces per step.
    logits = np.array([6.0, 5.2, 4.8, 3.0, 2.5, 1.0, 0.5, 0.0, -1.0, -2.0])
    h_max = np.log(len(logits))
    print(f"\nmax possible entropy (uniform): {h_max:.3f} nats")
    for tau in (0.1, 0.5, 0.7, 1.0, 1.5, 2.0, 5.0):
        p = softmax(logits, tau)
        top = p.max()
        print(f"tau={tau:<4} entropy={entropy(p):.3f} nats   "
              f"P(argmax)={top:.3f}")


if __name__ == "__main__":
    print("=== (A) exploration on a 10-armed bandit ===")
    bandit_demo()
    print("\n=== (B) temperature sweep (LLM decoding knob) ===")
    temperature_demo()

What you should see. In part (A), pure greedy ($\epsilon =0$) lands on the optimal arm markedly less often than $\epsilon =0.1$ over 1000 steps, typically well under it, because greedy sometimes commits to an arm that got a lucky first pull and never re-checks, the linear-regret trap made visible. In part (B), entropy is tiny at $\tau =0.1$ with $P(\text{argmax})$ near 1 (this is what a collapsed policy looks like) and climbs steadily toward the uniform ceiling of $\ln 10\approx 2.303$ nats as $\tau$ grows, with $P(\text{argmax})$ falling toward $0.1$. That single monotone column of entropies is exactly the dial you turn when you set rollout temperature for RL fine-tuning: too low and you have preemptively collapsed, too high and you are sampling noise. Keeping that dial in the productive middle, and logging entropy so you notice when training pushes it down on its own, is most of what "exploration" means in practice for a reasoning model.