Contextual Bandits

9 min read

Learn

Quick Reading

Estimated 9 mins

Prereq

Intermediate

Basic ML concepts helpful

Interactive

Static Playbook

Static guide & reference tables

Adding context to the exploration-exploitation problem

A classic multi-armed bandit picks the same 'best arm' for every user, ignoring everything about who they are. But the best ad, the best article, or the best price usually depends heavily on context: a user's device, location, browsing history, time of day. A contextual bandit observes a context vector $x_t$ before choosing an action, and the reward distribution is now a function of both the context and the action: $r_t \sim P(r | x_t, a_t)$. The goal becomes learning a policy $\pi(x) \rightarrow a$ that picks the best action *for each context*, not a single best action overall — this is the core building block behind most modern recommendation, ads, and personalization systems.

Formalizing the problem

At each round $t$: (1) the environment reveals a context $x_t \in \mathbb{R}^d$; (2) the policy chooses an action $a_t$ from a set of $K$ arms, possibly using $x_t$; (3) the environment reveals a reward $r_t$ for the chosen action only — rewards for the arms *not* chosen remain unobserved. This partial feedback (only ever seeing the outcome of your own choices, never the counterfactual) is what separates bandit learning from supervised learning, and it's the same fundamental problem causal inference deals with: you only observe $Y(a_t)$, never $Y(a)$ for $a \neq a_t$.

Regret is defined analogously to the non-contextual case, but now relative to the best possible *context-dependent* policy: $R_T = \sum_t \left[ \max_a E[r|x_t,a] - E[r|x_t,a_t] \right]$.

LinUCB: linear contextual bandits

The most widely deployed contextual bandit algorithm is LinUCB (Li, Chu, Langford, Schapire 2010), which assumes the expected reward of arm $a$ is linear in the context: $E[r | x, a] = x^\top \theta_a$ for an unknown weight vector $\theta_a$ per arm. LinUCB maintains a ridge-regression estimate of $\theta_a$ from past data for that arm, and — mirroring UCB1's optimism principle — selects the arm maximizing the predicted reward plus a confidence bonus derived from the covariance of observed contexts:

$a_t = \arg\max_a \left( x_t^\top \hat{\theta}_a + \alpha \sqrt{x_t^\top A_a^{-1} x_t} \right)$

where $A_a = D_a^\top D_a + I$ accumulates the outer products of contexts seen by arm $a$, and $\alpha$ controls how aggressively to explore. This closed-form update is what makes LinUCB cheap enough to run online at scale.

linucb.py

python

import numpy as np

class LinUCB:
    def __init__(self, n_arms, n_features, alpha=1.0):
        self.alpha = alpha
        self.A = [np.eye(n_features) for _ in range(n_arms)]   # ridge regression design matrix
        self.b = [np.zeros(n_features) for _ in range(n_arms)]  # accumulated reward-weighted context

    def select_arm(self, context):
        scores = []
        for A_a, b_a in zip(self.A, self.b):
            A_inv = np.linalg.inv(A_a)
            theta_hat = A_inv @ b_a
            mean = context @ theta_hat
            bonus = self.alpha * np.sqrt(context @ A_inv @ context)
            scores.append(mean + bonus)
        return int(np.argmax(scores))

    def update(self, arm, context, reward):
        self.A[arm] += np.outer(context, context)
        self.b[arm] += reward * context


rng = np.random.default_rng(3)
n_features, n_arms, T = 5, 4, 3000
true_theta = rng.normal(size=(n_arms, n_features))

bandit = LinUCB(n_arms, n_features, alpha=1.5)
cumulative_regret = 0.0
for t in range(T):
    context = rng.normal(size=n_features)
    true_rewards = context @ true_theta.T
    best_arm = int(np.argmax(true_rewards))

    chosen_arm = bandit.select_arm(context)
    observed_reward = true_rewards[chosen_arm] + rng.normal(0, 0.1)
    bandit.update(chosen_arm, context, observed_reward)

    cumulative_regret += true_rewards[best_arm] - true_rewards[chosen_arm]

print(f"Cumulative regret after {T} rounds: {cumulative_regret:.2f}")

Contextual Thompson Sampling generalizes naturally

Just as Thompson Sampling for plain bandits samples from a posterior over each arm's mean, contextual Thompson Sampling samples a weight vector $\theta_a$ from a posterior over the linear model's coefficients (typically a multivariate Gaussian with covariance derived the same way as LinUCB's $A_a^{-1}$), then picks the arm maximizing $x_t^\top \theta_a^{sampled}$. In practice it often performs comparably to or better than LinUCB and is easier to extend to richer reward models.

Beyond linear: neural and tree-based contextual bandits

Linear reward models are limiting when the true relationship between context and reward is nonlinear — which is the common case in recommendation and ads. Neural bandits replace the linear layer with a neural network trained online or in mini-batches, using approaches like NeuralUCB or NeuralLinear (a neural network for feature extraction with a Bayesian linear layer on top for the exploration bonus) to retain tractable uncertainty estimates. Tree-based approaches (bandit forests) similarly extend UCB-style exploration bonuses to non-linear base learners. The tradeoff is that nonlinear models make the uncertainty quantification needed for principled exploration much harder — this is an active area of bandit research.

Off-policy evaluation is notoriously fragile

Because bandit data is never i.i.d. — the policy that logged the data actively chose which arms to show — evaluating a *new* candidate policy against historical logs (off-policy evaluation) requires correcting for that selection bias, typically via inverse propensity scoring (IPS): weighting each logged example by $1/\pi_{old}(a|x)$. IPS estimates have notoriously high variance when the old and new policies disagree substantially, and clipping or doubly-robust estimators are standard mitigations. Teams that naively evaluate a new bandit policy by replaying old logs without this correction routinely get badly biased estimates of how the new policy would actually perform.

What's Next?

Contextual bandits still treat each decision as a single-step choice: pick an action, get an immediate reward, move on. Many real systems care about consequences that unfold over a sequence of decisions — a recommendation today affecting what a user wants tomorrow — which is the setting full reinforcement learning is built to handle.

I build these systems professionally.

Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.

Work With Me →