From A/B to RL (2/3): Bandits and Thompson Sampling

Learn and adapt online instead of choosing once at the end.

In the first post of this series, we used a fixed randomized A/B test to compare two variants. We collected the data first, then used the posterior probability $P(B > A \mid \mathrm{data})$ to make one final decision at the end of the experiment.

Here we change the decision problem while keeping the same statistical model. We move to a class of problems commonly known as multi-armed bandits, where the learner repeatedly chooses one action from several options, observes a binary reward only for the selected action, updates its beliefs, and chooses again. This partial-feedback pattern is called bandit feedback.

The term multi-armed bandit comes from the image of a row of slot machines, sometimes called one-armed bandits. In bandit language, each option is called an arm; following reinforcement-learning terminology, we will generally call it an action.

In contrast to the fixed experiment, the learner does not wait until the end to act. Each observed reward updates its beliefs and informs the next action, so learning and decision-making happen together.

This is the second post in a three-part series:

In [1]:

The Toy Problem

Building on the previous post, we use the same website-button example, now with three possible button variants: A, B, and C. We treat these choices as actions. At each interaction, the learner chooses which button to show and observes either a click or no click.

In the fixed A/B experiment, assignments were planned in advance and we made one decision after collecting the data. Here, the learner uses each observed outcome to choose the next button. Its objective is cumulative reward: collect as many clicks as possible over time while learning which button is most effective.

In [2]:

From Posterior Beliefs to Online Decisions

We now express this setup in Bayesian and reinforcement-learning notation. At interaction step $t$, the learner selects an action $a_t$ and observes a binary reward $R_t$: $1$ for a click and $0$ for no click. Each action $a$ has a fixed but unknown click probability $\theta_a$.

Because the reward is binary, its expected value after choosing action $a$ is simply the probability of a click:

$$ \begin{aligned} \mathbb{E}[R_t \mid a_t=a] &= 1 \cdot P(R_t=1 \mid a_t=a) + 0 \cdot P(R_t=0 \mid a_t=a) \\ &= P(\text{click}\mid a) \\ &= \theta_a. \end{aligned} $$

Thus, in this example, $\theta_a$ is both the action's click probability and its expected immediate reward. It is the same quantity as the variant's CTR in the fixed A/B test.

The equation above gives the expected value of the binary reward. To describe how that reward $R_t$ is generated, we model its conditional distribution given the chosen action $a_t$ as a Bernoulli distribution:

$$ R_t \mid (a_t = a) \sim \mathrm{Bernoulli}(\theta_a). $$

Because $\theta_a$ is unknown, we need to represent our uncertainty about it before observing rewards. For each action, we use a Beta distribution as its prior:

$$ \theta_a \sim \mathrm{Beta}(1, 1). $$

We use a separate, independent prior for each possible action and model their rewards as independent given their reward probabilities. This lets us maintain one posterior for each action.

Let $(\alpha_a, \beta_a)$ denote the current Beta parameters for each action. After choosing action $a_t$ and observing reward $R_t \in \{0, 1\}$, we update only that action:

$$ \alpha_{a_t} \leftarrow \alpha_{a_t} + R_t, \qquad \beta_{a_t} \leftarrow \beta_{a_t} + (1 - R_t). $$

This is the same Beta-Bernoulli update as in the first post, applied one interaction at a time. The posterior for the selected action after one interaction becomes the prior for its next observation. For a standalone explanation of why this two-count update fits sequential binary feedback, see Beta Priors and Pólya Urns for Binary Decisions.

Online Updates and Cumulative Reward

The statistical model is the same as in the fixed A/B experiment: each action $a$ has a fixed but unknown reward probability $\theta_a$. What changes is the objective. The learner is no longer only trying to identify the best action at the end; it must also collect as much cumulative reward as possible while learning.

Because each interaction is one-step, the reward for one interaction is just the immediate binary outcome. If we write that reward as $R_t$, then $R_t = 1$ for a click and $R_t = 0$ for no click. Across the whole run, the objective is to maximize cumulative reward:

$$ \sum_{t=1}^T R_t. $$

In this binary setup, the action value is the expected immediate reward. That expected reward is $\theta_a$, the action's hidden reward probability. Sampling a plausible $\theta_a$ is therefore the same as sampling a plausible action value.

Policy and Exploration

To act online, we need a policy: a rule for selecting the next action. In standard RL notation, a policy is often written as $\pi(a \mid s)$, a distribution over actions $a$ that can be taken from a state $s$. Because the external state remains fixed at $s_0$, we simplify the policy notation to $\pi_t(a)$ instead of $\pi_t(a \mid s_0)$. The policy changes because the learner's posterior beliefs change after each observed reward.

A bandit policy must balance the exploration-exploitation dilemma. At one extreme, a greedy policy always chooses the action that currently looks best. It uses the evidence collected so far, but may stop trying underexplored actions too early and miss one that would have been better. At the other extreme, a uniformly random policy keeps exploring every action but never uses what it has learned to improve cumulative reward.

A common compromise is an $\varepsilon$-greedy policy: most of the time choose the currently best-looking action, and with probability $\varepsilon$ explore by choosing a random action.

Here we use probability matching. We already model posterior uncertainty over each action's expected reward. In the first post, that uncertainty gave us $P(B > A \mid \mathrm{data})$. Here, the same comparison drives action selection: an action is chosen in proportion to the posterior probability that it is best. With $K$ actions, we ask this question for every action: under the current posterior, how likely is this action to have the highest expected reward? Here, $\mathrm{data}_t$ denotes all action-reward observations available before interaction $t$.

Under this probability-matching rule, we can write the policy as:

$$ \pi_t(a) = P\!\left(\theta_a > \theta_j \;\; \text{for every other action } j \mid \mathrm{data}_t\right). $$

In other words, $\pi_t(a)$ is the posterior probability that action $a$ has the highest expected reward. Because this is a one-step binary setup, the action with the highest expected reward is also the action with the highest action value.

This rule turns the posterior probabilities into a policy distribution: choose each action with probability $\pi_t(a)$ on the next interaction. This is the idea Thompson described in 1933: use the current probability that each option is best to guide allocation as more evidence arrives, instead of waiting for one final decision. Exploration comes from posterior uncertainty rather than a fixed $\varepsilon$, so an underexplored action can still be chosen when the posterior leaves a meaningful chance that it is best.

Thompson Sampling as Posterior Sampling

In the first post, we estimated $P(B > A \mid \mathrm{data})$ by sampling from the posterior. We drew many plausible expected rewards, $\tilde\theta_A$ and $\tilde\theta_B$, counted how often $\tilde\theta_B$ was higher, and used that empirical frequency as a Monte Carlo estimate.

Thompson sampling implements this probability-matching policy by sampling. For each action $i$, we draw one plausible value $\tilde{\theta}_i$ from its current posterior and choose the action with the largest draw. Repeating this procedure many times with the same posteriors would estimate $\pi_t(a)$. In the actual bandit, we perform it once per interaction and choose:

$$ a_t = \arg\max_i \tilde{\theta}_i. $$

Written in probability notation, the selected action is drawn from the policy distribution: selecting action $a$ is exactly the event that $a$ has the largest sampled expected reward:

$$ P(a_t = a \mid \mathrm{data}_t) = P\!\left(\tilde\theta_a > \tilde\theta_j \;\; \text{for every other action } j \mid \mathrm{data}_t\right) = \pi_t(a). $$

This sampling rule captures the core idea of Thompson sampling: posterior uncertainty naturally drives exploration as evidence accumulates. At interaction $t$, the algorithm works as follows:

  1. For each action $i$, sample one plausible value $\tilde{\theta}_i$ from its current posterior,
  2. interpret $\tilde{\theta}_i$ as a plausible action value, or expected immediate reward, for action $i$,
  3. choose the action with the largest sampled value, $a_t = \arg\max_i \tilde{\theta}_i$,
  4. observe the binary reward $R_t \in \{0, 1\}$ for the selected action,
  5. update only the posterior for $a_t$, using the observed reward $R_t$.

If we repeated steps 1–3 many times with the same posteriors, the action frequencies would approach $\pi_t(a)$. In the actual algorithm, we perform these steps once per interaction, observe $R_t$, and update the chosen action before the next interaction.

Each sampled $\tilde{\theta}_i$ is a plausible action value: an expected immediate reward for action $i$, not a prediction of the next binary reward $R_t$.

The simulations below make these ideas concrete. We will watch the posteriors evolve, see the induced policy concentrate on the best action, and track reward and expected regret along the way.

The figure below shows one pass through the learning loop. We start with the current posterior beliefs, sample one plausible expected reward for each action, choose the largest sample, observe one binary reward, and update only the chosen action.

In [3]:
Three-panel illustration of one Thompson-sampling interaction: sample action values, choose the largest sample, then update the chosen action after the reward.
In [4]:

We now run the interaction loop for 1,000 interactions. The code initializes one Beta posterior per action, repeats the Thompson-sampling step, and records the values used by the plots.

In [5]:
# Run Thompson sampling for repeated bandit interactions.
n_steps = 1000
steps = np.arange(1, n_steps + 1)
bandit_rng = np.random.default_rng(202604243)

# Initialize one Beta posterior per action.
alpha = np.ones(n_actions, dtype=float)
beta = np.ones(n_actions, dtype=float)

interaction_log: list[BanditInteraction] = []
posterior_history: dict[int, tuple[FloatArray, FloatArray]] = {}
checkpoints = [1, 10, 50, 200, 500, 1000]

# Run Thompson sampling as repeated agent-environment interaction.
for step in steps:
    # Sample one plausible expected reward for each action from its posterior.
    sampled_action_values = bandit_rng.beta(alpha, beta)

    # Choose the action with the highest sampled value.
    action = int(np.argmax(sampled_action_values))

    # Observe binary feedback from the selected action's hidden reward probability.
    reward = int(bandit_rng.binomial(n=1, p=true_success_probabilities[action]))

    # Update only the selected action's Beta posterior.
    alpha, beta = update_beta_bernoulli_posterior(
        alpha=alpha,
        beta=beta,
        action=action,
        reward=reward,
    )

    # Log feedback and instantaneous expected regret for later summaries.
    instant_expected_regret = float(np.max(true_success_probabilities) - true_success_probabilities[action])
    interaction_log.append(BanditInteraction(action=action, reward=reward, expected_regret=instant_expected_regret))

    # Keep posterior snapshots for the checkpoint visualizations.
    if step in checkpoints:
        posterior_history[step] = (alpha.copy(), beta.copy())
In [6]:
Best true action: C
Action counts: [ 88 200 712]
Posterior means: [0.078 0.104 0.15 ]
Total reward: 132
Total expected regret: 11.28

Posterior Evolution

Each action $a$ has a Beta posterior over its unknown reward probability $\theta_a$. The update is the same Beta-Bernoulli update used in the first post, but feedback now arrives online: after interaction $t$, only the posterior for the selected action $a_t$ changes. When $R_t=1$, that posterior shifts toward larger values; when $R_t=0$, it shifts toward smaller values.

Each panel in the figure below shows the posterior over every action's expected reward at a different point in the run. The dashed lines mark the hidden true expected rewards used to generate the simulation. After each interaction, only the selected action's posterior is updated; the others remain unchanged until they are selected again.

When reading these curves, focus on two things:

  • the location of a curve shows which expected rewards currently look plausible,
  • the width of a curve shows how uncertain we are about that action value.

Frequently chosen actions accumulate evidence quickly, so their posteriors become narrower. Less frequently chosen actions remain more diffuse. Thompson sampling uses this difference in uncertainty: uncertain actions can still be explored, while well-performing and well-measured actions are exploited more often.

In [7]:
Beta posterior densities for three bandit actions at selected interaction checkpoints, with dashed lines marking hidden true action values.

At each checkpoint, we pause the simulation after the indicated number of interactions. Using the posterior distributions at that point, we estimate the action-selection policy for the next interaction. The plot below shows the probability assigned to each action, or equivalently, how likely each action is to have the highest expected reward under the current posterior distributions.

In [8]:
Estimated probability that each of three actions is selected at several checkpoints, showing how posterior uncertainty changes the policy.

Reward, Expected Regret, and Action Allocation

With distinct action values, Thompson sampling usually explores several actions at first, then allocates more interactions to actions that continue to look promising. This adaptive allocation is the main behavioral difference from a fixed experiment.

We can evaluate the run in three ways:

  • Cumulative reward: the total reward collected through interaction $T$, $\sum_{t=1}^{T} R_t$. This is the quantity the bandit tries to maximize.
  • Cumulative expected regret: the expected reward lost by not always choosing the best action. If $\theta^* = \max_a \theta_a$, choosing action $a_t$ at interaction $t$ incurs expected regret $\theta^* - \theta_{a_t}$. Summing over interactions gives cumulative expected regret.
  • Action allocation: how often the learner chooses each action, both overall and across successive blocks of interactions.

Because this is a simulation, we know the hidden action values and can compute expected regret. We can also compare Thompson sampling with two reference strategies: uniform allocation and an oracle that always chooses the best action. Their reward curves are expected reference values, not realized alternative runs. In a real deployment, the hidden action values and expected regret would not be directly observable.

In [9]:
Three-panel bandit evaluation showing cumulative reward, cumulative expected regret, and how often each action was chosen.
Allocation share for each bandit action over successive blocks of interactions.

One practical caveat is that bandit data are collected adaptively, so they should not be treated like data from a fixed randomized experiment. The policy that chooses actions affects which outcomes are observed. To make causal or off-policy claims from logged bandit data, we must account for that policy, for example by logging each selected action's selection probability and ensuring that relevant alternatives retain a nonzero chance of being selected.

Final Learned Picture

In this run, the posterior means roughly recover the ordering of the hidden true expected rewards. But estimating that ordering is only part of the story. Unlike the fixed A/B test, the bandit did not wait until the end to make a single choice; it used its evolving beliefs to steer later interactions toward actions that appeared more promising.

The plot below compares the hidden true expected rewards with the final posterior means and 90% credible intervals. The labels show how often each action was selected, making the exploration-exploitation pattern visible.

In [10]:
Final posterior means and 90% credible intervals for the three action values, compared with their hidden true expected rewards; labels show selection counts.

Summary

In the first post, we used the posterior after the experiment to make one final decision about which variant to ship. Here, the learner updates its posterior after each interaction and uses it to choose the next action.

At each interaction:

  • the policy chooses one action using the current posterior beliefs,
  • the environment returns an immediate binary reward for that action only,
  • the learner updates only the chosen action's Beta posterior,
  • the updated posterior beliefs inform the next action choice.

This is Thompson sampling in action: the learner draws one plausible expected reward from each action's posterior, chooses the action with the largest draw, observes the binary reward, and updates the selected action. These draws represent plausible average rewards, not predictions of the next binary outcome.

This is still a simple RL setting: the external state is fixed, rewards are immediate, and there is no delayed credit assignment. The next post adds both changing state and delayed feedback.

References

Further Reading

In [11]:
Python: 3.13.12
numpy: 2.4.6
scipy: 1.17.1
matplotlib: 3.11.0
seaborn: 0.13.2

This post is generated from an IPython notebook file. Link to the full IPython notebook file