Contextual Bandits

"For a thousand rounds I pulled the same lever because, on average, it paid best. Then someone handed me a slip of paper at the start of each round, the weather, the user's age, the time of day, and I realized: there is no single best lever. There is a best lever for this slip. I had been answering the wrong question confidently the entire time."

A Bandit That Finally Looks at the Situation Before Choosing

In Section 24.2 we met the multi-armed bandit and its exploration strategies: epsilon-greedy, upper confidence bounds, and Thompson sampling, all built on the premise that each arm has one fixed reward distribution to be estimated. That premise is exactly what real applications violate. The best news article for a retiree at breakfast is not the best article for a teenager at midnight; the best price for a price-sensitive shopper is not the best price for a brand-loyal one. The context-free bandit, forced to pick one global best arm, averages across all these situations and serves everyone the bland compromise. This section gives the bandit eyes. It observes the situation first, then chooses, and the regret it competes against is no longer "the single best arm" but "the best policy mapping contexts to arms", a far stronger and far more useful benchmark.

This matters because the contextual bandit is, by a wide margin, the most deployed form of sequential decision making in industry, more than full reinforcement learning, which it often quietly outperforms in settings without long-horizon dynamics. When your action does not change the world's next state (recommending an article does not change who arrives next), you do not need the transition machinery of an MDP, and paying for it buys you instability and sample inefficiency for nothing. The contextual bandit is the right tool precisely when the problem has rich state but trivial dynamics, and an enormous fraction of valuable decision problems have exactly that shape. We use the unified notation of Appendix A throughout: $\mathbf{x}_t$ the context, $a_t$ the chosen action, $r_t$ the observed reward, $\pi$ a policy mapping contexts to actions.

Four competencies this section installs: to write the contextual reward model and state the regret a contextual learner minimizes; to derive and implement LinUCB, the linear reward model with a ridge-regression confidence-ellipsoid exploration bonus; to describe Thompson sampling and neural contextual bandits as alternative exploration mechanisms; and to estimate the value of a policy you have not run, from logged data alone, using inverse propensity scoring and the doubly-robust estimator. These are the load-bearing skills for the deep and offline reinforcement learning of Chapter 25 and beyond.

1. Adding Context: The Situation Decides the Best Action Beginner

The defining move of the contextual bandit is a single new object revealed at the top of each round: a feature vector $\mathbf{x}_t \in \mathbb{R}^d$ describing the situation. The interaction loop becomes: the environment reveals $\mathbf{x}_t$; the learner chooses an action $a_t \in \{1, \dots, K\}$; the environment returns a reward $r_t$ whose expected value depends on both the context and the chosen action. Nothing else carries over from one round to the next. There is no state transition, no notion that the action you took nudges the context you will see next. That absence is the whole reason the contextual bandit is simpler than the MDP of Chapter 22, and the presence of $\mathbf{x}_t$ is the whole reason it is richer than the context-free bandit of Section 24.2.

Formally, we assume an expected-reward function of context and action,

with $\varepsilon_t$ a zero-mean noise term. The learner's goal is to choose actions so as to accumulate reward close to that of the best context-dependent policy. Writing the optimal policy as $\pi^\star(\mathbf{x}) = \arg\max_a \mu_a(\mathbf{x})$, the cumulative regret over $T$ rounds is

Compare this to the context-free regret of Section 24.2, which subtracts the single global best arm $\mu_{a^\star}$. The contextual benchmark $\mu_{\pi^\star(\mathbf{x}_t)}(\mathbf{x}_t)$ is per-round and per-context, and it is generally much higher than any fixed arm's mean, so beating it requires actually learning the functions $\mu_a(\cdot)$, not merely their averages. This is the precise sense in which a contextual bandit is a function-learning problem: exploration must discover not "which arm pays most" but "which arm pays most in which region of context space".

The position of the contextual bandit between the two neighbors is worth stating precisely, because it dictates which algorithms apply. From the multi-armed bandit it inherits the exploration-exploitation tension: you still must occasionally try an action you are unsure about to learn its reward. From the MDP it borrows the notion of state, the context $\mathbf{x}_t$, but throws away the transition kernel $P(\mathbf{x}_{t+1} \mid \mathbf{x}_t, a_t)$ entirely. Because there are no transitions, there is no temporal credit assignment, no discounting, no value-function bootstrapping; the reward is immediate and fully attributable to the single action that earned it. That collapse to one step is liberating: it means the rich, stable supervised-learning toolbox (linear models, gradient-boosted trees, neural networks) can estimate $\mu_a(\mathbf{x})$ directly, and the only genuinely sequential subtlety left is how to explore.

One distinction deserves emphasis because it separates the contextual bandit from ordinary supervised learning even though both fit a function $\mu_a(\mathbf{x})$. In supervised regression you are handed, for every training example, the label for every action: you see the reward you would have gotten under each choice. In the bandit you see only the reward of the action you actually took, the so-called bandit feedback or partial-information setting. You never observe the counterfactual: what the skipped arms would have paid. This is precisely why a contextual bandit cannot simply collect a dataset and run regression once; it must choose which arm to query, and a poor choice of queries leaves whole regions of the reward surface unobserved. Exploration is the mechanism that ensures the function estimate $\hat{\mu}_a(\cdot)$ is informed across the action-context space rather than only along the trajectory a greedy policy happened to walk. Hold that asymmetry, full feedback in supervised learning, partial feedback here, firmly, because it is the reason every algorithm in this section pairs a regressor with an explicit exploration rule.

2. LinUCB: A Linear Reward Model With a Confidence Bonus Intermediate

The most influential contextual-bandit algorithm makes one modeling assumption, that the expected reward is linear in the context, and pairs it with one exploration idea, optimism in the face of uncertainty, the same principle that powered the UCB bandit of Section 24.2. Assume each arm $a$ has an unknown parameter vector $\boldsymbol{\theta}_a \in \mathbb{R}^d$ and that

If we knew $\boldsymbol{\theta}_a$ we would just pick the arm with the largest predicted reward $\mathbf{x}_t^\top\boldsymbol{\theta}_a$. We do not, so we estimate it by ridge regression from the rounds where arm $a$ was played, and then we add an exploration bonus proportional to how uncertain that estimate is in the direction of the current context. Collect, for arm $a$, the design matrix accumulator $\mathbf{A}_a = \lambda\mathbf{I} + \sum \mathbf{x}\mathbf{x}^\top$ (over rounds where $a$ was chosen) and the response accumulator $\mathbf{b}_a = \sum r\,\mathbf{x}$. The ridge solution and the LinUCB action-selection score are

and LinUCB plays $a_t = \arg\max_a p_a(\mathbf{x}_t)$. The first term is pure exploitation, the predicted reward under the current estimate. The second term is the radius of the confidence ellipsoid around $\hat{\boldsymbol{\theta}}_a$, measured in the direction $\mathbf{x}_t$: the matrix $\mathbf{A}_a^{-1}$ is (up to noise scale) the covariance of the parameter estimate, so $\sqrt{\mathbf{x}_t^\top \mathbf{A}_a^{-1}\mathbf{x}_t}$ is the standard deviation of the predicted reward $\mathbf{x}_t^\top\hat{\boldsymbol{\theta}}_a$. The constant $\alpha$ sets how many standard deviations of optimism to add; larger $\alpha$ explores more.

The ridge-regression and confidence-ellipsoid view is exactly the online ridge of Chapter 20 seen through a decision-making lens. There, online ridge regression maintained a running inverse covariance to update a forecast as data streamed in; here, that same inverse covariance $\mathbf{A}_a^{-1}$ does double duty, it gives both the point estimate (through $\hat{\boldsymbol{\theta}}_a = \mathbf{A}_a^{-1}\mathbf{b}_a$) and the uncertainty (through the bonus). The geometry is clean: directions of context space in which arm $a$ has been pulled many times have small $\mathbf{x}^\top\mathbf{A}_a^{-1}\mathbf{x}$ (the ellipsoid is thin there, we are confident, the bonus is small), while novel directions have large $\mathbf{x}^\top\mathbf{A}_a^{-1}\mathbf{x}$ (the ellipsoid is wide, we are uncertain, the bonus is large and pulls us to explore). Optimism thus automatically steers exploration toward context directions the arm has not yet been tested in, which is precisely where learning is still possible.

Why is optimism the right exploration rule rather than a heuristic? The guarantee is that LinUCB, under the linear-reward assumption, attains cumulative regret $\tilde{O}(d\sqrt{T})$, sublinear in the horizon and graceful in the context dimension $d$. Sublinear regret means the per-round regret tends to zero: the average reward converges to that of the optimal context-dependent policy. The argument mirrors the context-free UCB bound: because the true $\boldsymbol{\theta}_a$ lies inside the confidence ellipsoid with high probability, the optimistic score $p_a$ upper-bounds the true reward, so whenever LinUCB picks a suboptimal arm it can only be because that arm's ellipsoid was wide, which means we were genuinely uncertain about it, which means we are about to learn a lot by pulling it. Exploration and information gain are aligned by construction. Code 24.3.1 in subsection five implements exactly the $\mathbf{A}_a$, $\mathbf{b}_a$ update and the score above.

3. Thompson Sampling and Neural Contextual Bandits Intermediate

LinUCB explores by adding a deterministic optimism bonus. The leading alternative explores by deliberate randomness: Thompson sampling, the posterior-sampling strategy we first met for the context-free bandit in Section 24.2, extends to the contextual setting almost verbatim. Maintain a Bayesian posterior over each arm's parameter $\boldsymbol{\theta}_a$; under the linear-Gaussian model with ridge prior, that posterior is exactly Gaussian, $\boldsymbol{\theta}_a \sim \mathcal{N}(\hat{\boldsymbol{\theta}}_a, \sigma^2 \mathbf{A}_a^{-1})$, using the same $\mathbf{A}_a$, $\mathbf{b}_a$ accumulators as LinUCB. Each round, sample a parameter $\tilde{\boldsymbol{\theta}}_a$ from each arm's posterior, then act greedily with respect to the sampled parameters:

The exploration is now stochastic: an arm whose posterior is wide will sometimes sample a high parameter and get chosen, sometimes a low one and get skipped, with frequency proportional to the posterior probability that it is actually best. This "probability matching" tends to be more robust in practice than UCB's worst-case optimism, often delivering lower empirical regret, and it parallelizes naturally because each decision needs only one posterior sample rather than a covariance-direction computation. Linear Thompson sampling enjoys the same $\tilde{O}(d\sqrt{T})$ regret order as LinUCB.

Both LinUCB and linear Thompson sampling assume the reward is linear in the context, and real reward surfaces rarely are. Neural contextual bandits replace the linear model $\mathbf{x}^\top\boldsymbol{\theta}_a$ with a neural network $f_a(\mathbf{x}; \mathbf{w})$ that can represent arbitrary reward functions, and the only real design question becomes how to do exploration when the model is no longer linear. Three families dominate. NeuralUCB and NeuralTS (Zhou et al., 2020; Zhang et al., 2021) build the confidence bonus from the network's gradient features, treating the last-layer or neural-tangent-kernel representation as the context in a LinUCB-style ellipsoid. Neural-linear methods learn a deep representation $\phi(\mathbf{x})$ with a standard network and run exact Bayesian linear regression (LinUCB or Thompson) on top of the frozen or slowly-updated features, getting nonlinearity from the representation and calibrated uncertainty from the linear head. And bootstrapped or dropout-based ensembles approximate the posterior by training several networks and sampling among them. The trade is the usual one: more expressive reward models capture richer personalization but make the uncertainty estimate that drives exploration harder to keep honest.

The choice between optimism and posterior sampling is not merely aesthetic, and it is worth knowing when each pulls ahead. LinUCB's bonus is deterministic and worst-case: it explores an arm exactly as much as its widest plausible reward justifies, which makes its behavior reproducible and its regret bound tight, but also makes it sensitive to the exploration constant $\alpha$, a value too large wastes pulls on hopeless arms and a value too small under-explores. Thompson sampling sidesteps that tuning by reading the exploration rate directly off the posterior width, so an arm is tried in proportion to the posterior probability that it is best, which adapts automatically as evidence accumulates and tends to degrade more gracefully when the reward model is slightly misspecified. The practical consequence reported across large-scale studies is that Thompson sampling is frequently the stronger default in noisy, nonstationary production systems, while LinUCB is preferred when reproducibility, an explicit confidence interval, or a provable per-round guarantee is required. Both reduce, in the linear-Gaussian case, to reading the same $\hat{\boldsymbol{\theta}}_a$ and $\mathbf{A}_a^{-1}$, so switching between them in a deployed system is often a one-line change in how the score is computed.

4. Off-Policy Evaluation: Scoring a Policy You Have Not Run Advanced

In almost every production setting you cannot freely experiment. Deploying an untested recommendation policy to live users to see whether it is better is slow, expensive, and risky; a bad policy loses revenue and trust before you can measure it. Yet you have a mountain of logged data: contexts, the actions your current (logging) policy took, the probabilities with which it took them, and the rewards it received. Off-policy evaluation (OPE) is the set of techniques that estimate the value of a new policy $\pi$ from that logged data alone, without ever running $\pi$. It is the workhorse that lets teams compare dozens of candidate policies offline and deploy only the winner.

The logging policy $\pi_0$ recorded, for each round, the context $\mathbf{x}_t$, the action $a_t$ it sampled, the propensity $p_t = \pi_0(a_t \mid \mathbf{x}_t)$ (the probability it chose that action), and the reward $r_t$. The target policy's value is $V(\pi) = \mathbb{E}_{\mathbf{x}}\big[\sum_a \pi(a \mid \mathbf{x})\,\mu_a(\mathbf{x})\big]$, the expected reward if $\pi$ were deployed. The inverse propensity scoring (IPS) estimator reweights each logged reward by how much more (or less) likely the target policy was to take the same action than the logging policy:

The importance weight $\pi(a_t\mid\mathbf{x}_t)/\pi_0(a_t\mid\mathbf{x}_t)$ corrects for the fact that the data was collected under $\pi_0$, not $\pi$. IPS is unbiased whenever the logging policy had nonzero probability of every action the target policy might take (the "full support" or overlap condition), which is exactly why production logging policies must keep exploring rather than collapse to a deterministic greedy choice: a deterministic logger leaves you blind to every action it never tried. The cost of IPS is high variance: when $\pi_0$ took an action rarely but $\pi$ favors it, the weight $\pi/\pi_0$ is huge, and a single such round can dominate the estimate.

The doubly-robust (DR) estimator tames that variance by combining IPS with a learned reward model $\hat{\mu}_a(\mathbf{x})$ (a "direct method" regressor fit to the logged rewards). It uses the model as a baseline and applies the expensive importance-weighting only to the model's residual:

The name "doubly robust" is the payoff: the estimator is consistent if either the propensities $\pi_0$ are correct or the reward model $\hat{\mu}$ is correct, you need only one of the two to be right. When the reward model is good its residual is small, so the high-variance importance weight multiplies a tiny number and the variance collapses; when the model is bad the IPS correction term rescues unbiasedness. DR is the default OPE estimator in modern contextual-bandit pipelines for exactly this reason, and it is the same control-variate idea that appears throughout statistics: subtract a correlated estimate, importance-weight only what is left.

5. Worked Example: LinUCB From Scratch, Then a Library Advanced

We now make the algorithm executable on a synthetic contextual bandit where we control the ground truth, so we can measure regret exactly. The plan is a clean three-way comparison: LinUCB (context-aware, explores), a greedy linear model (context-aware, never explores), and a context-free UCB bandit (explores, but ignores context). The synthetic environment has $K$ arms, each with a true hidden parameter $\boldsymbol{\theta}_a^\star$; each round draws a random context and the reward of arm $a$ is $\mathbf{x}_t^\top\boldsymbol{\theta}_a^\star$ plus noise. Code 24.3.1 defines the environment and the from-scratch LinUCB agent.

import numpy as np

rng = np.random.default_rng(0)
K, d, T = 4, 5, 3000               # arms, context dimension, rounds
alpha, lam = 1.0, 1.0             # exploration constant; ridge regularizer

# Ground-truth: each arm has a hidden linear reward parameter theta*_a.
theta_star = rng.normal(0, 1.0, size=(K, d))

def draw_context():
    x = rng.normal(0, 1.0, size=d)
    return x / np.linalg.norm(x)   # unit-norm contexts keep reward scales comparable

def reward(a, x):
    return float(x @ theta_star[a] + rng.normal(0, 0.1))   # linear reward + noise

class LinUCB:
    """From-scratch LinUCB: a ridge model + confidence-ellipsoid bonus per arm."""
    def __init__(self, K, d, alpha, lam):
        self.alpha = alpha
        self.A = [lam * np.eye(d) for _ in range(K)]   # A_a = lam I + sum x x^T
        self.b = [np.zeros(d) for _ in range(K)]        # b_a = sum r x

    def select(self, x):
        scores = []
        for a in range(K):
            Ainv = np.linalg.inv(self.A[a])
            theta = Ainv @ self.b[a]                     # ridge estimate theta_hat_a
            exploit = x @ theta                          # predicted reward
            explore = self.alpha * np.sqrt(x @ Ainv @ x) # confidence width (the bonus)
            scores.append(exploit + explore)             # the LinUCB score p_a(x)
        return int(np.argmax(scores))

    def update(self, a, x, r):
        self.A[a] += np.outer(x, x)                      # rank-one design update
        self.b[a] += r * x                               # response accumulation

Now the two baselines and the comparison loop. The greedy agent reuses the same ridge accumulators but drops the exploration bonus (it sets $\alpha = 0$, always trusting its current estimate), and the context-free UCB agent ignores $\mathbf{x}_t$ entirely, tracking only a scalar mean and count per arm. We run all three on the same context stream and accumulate regret against the per-round optimal arm. Code 24.3.2 is the driver.

class ContextFreeUCB:
    """Ignores context: classic UCB1 on per-arm scalar means (Section 24.2)."""
    def __init__(self, K):
        self.n = np.zeros(K); self.mean = np.zeros(K); self.t = 0
    def select(self, x):
        self.t += 1
        if (self.n == 0).any():
            return int(np.argmin(self.n))               # pull each arm once first
        bonus = np.sqrt(2 * np.log(self.t) / self.n)    # UCB1 exploration bonus
        return int(np.argmax(self.mean + bonus))
    def update(self, a, x, r):
        self.n[a] += 1
        self.mean[a] += (r - self.mean[a]) / self.n[a]  # incremental mean

def run(agent, contextual=True):
    total_regret = 0.0
    for _ in range(T):
        x = draw_context()
        best = max(x @ theta_star[a] for a in range(K))  # oracle: best arm THIS context
        a = agent.select(x)
        r = reward(a, x)
        agent.update(a, x, r)
        total_regret += best - (x @ theta_star[a])       # instantaneous regret
    return total_regret

linucb  = LinUCB(K, d, alpha=1.0, lam=lam)
greedy  = LinUCB(K, d, alpha=0.0, lam=lam)               # alpha=0 => pure exploitation
cfucb   = ContextFreeUCB(K)
print("LinUCB        cumulative regret = %.1f" % run(linucb))
print("Greedy linear cumulative regret = %.1f" % run(greedy))
print("Context-free  cumulative regret = %.1f" % run(cfucb))

LinUCB        cumulative regret = 41.7
Greedy linear cumulative regret = 188.4
Context-free  cumulative regret = 612.9

The ordering is the lesson of the whole section in three numbers. The context-free bandit pays by far the most regret: with four arms whose best choice flips depending on $\mathbf{x}_t$, committing to any single arm is wrong most of the time, and exploration cannot fix a model that has no place to put the context. The greedy linear model fixes the modeling gap but creates an exploration gap: by always trusting its current ridge estimate it under-explores arms it initially misjudged, never collecting the data that would correct them, and pays for it. Only LinUCB has both a context-aware model and a principled reason to gather data about under-tested arm-context directions, and only LinUCB gets sublinear regret. Now the library equivalent.

# pip install obp
from obp.policy import LinUCB
from obp.ope import OffPolicyEvaluation, InverseProbabilityWeighting, DoublyRobust

policy = LinUCB(dim=d, n_actions=K, epsilon=0.0)        # the agent: ~1 line vs ~25
# ... fit policy / collect logged feedback into `bandit_feedback` ...

ope = OffPolicyEvaluation(                                # score a NEW policy offline
    bandit_feedback=bandit_feedback,
    ope_estimators=[InverseProbabilityWeighting(), DoublyRobust()],
)
value_estimates = ope.estimate_policy_values(            # IPS and DR in one call
    action_dist=new_policy_action_dist,                  # pi(a|x) of the candidate
    estimated_rewards_by_reg_model=reward_model_preds,   # mu_hat for the DR term
)

Read what the three code blocks establish together. Code 24.3.1 implemented LinUCB's ridge accumulators and confidence-ellipsoid bonus exactly as derived in subsection two. Code 24.3.2 ran it against a greedy context-aware baseline and a context-blind bandit on identical data, and Output 24.3.1 showed LinUCB's regret a factor of four below greedy and fifteen below context-free, isolating the two ingredients (context and calibrated exploration) that each baseline lacks one of. The library shortcut then collapsed the agent and the entire off-policy-evaluation apparatus to a few lines. The payoff: a contextual bandit is not a black box, you can build LinUCB by hand, watch context and exploration each earn their regret reduction, and reach for a library only once you understand what it computes.

6. Neural Contextual Bandits: NeuralUCB and NeuralTS Advanced

LinUCB achieves sublinear regret under one assumption: the expected reward is linear in the context. That assumption is violated the moment the reward surface curves, and it curves in nearly every real application. A news article's appeal to a user is not a dot product of five features; a patient's drug response is not linear in their biomarkers; an ad's click probability is not linear in the pixel histogram of the creative. Neural contextual bandits discard the linearity constraint and replace the linear reward model with a neural network $f_\theta(\mathbf{x}, a)$, gaining the full approximation power of deep learning. The challenge that follows immediately is that the LinUCB confidence bonus $\alpha\sqrt{\mathbf{x}^\top \mathbf{A}_a^{-1} \mathbf{x}}$ was derived for a linear model and has no obvious analogue for a 50-layer network. NeuralUCB and NeuralTS, the two landmark algorithms from Zhou et al. (NeurIPS 2020) and Zhang et al. (2021), solve this by identifying the right object in the neural network to play the role of the linear feature: the gradient of the network output with respect to its parameters.

6.1 The Neural Tangent Kernel as a Feature Map

Fix a neural network with parameters $\theta$ and output $f_\theta(\mathbf{x}, a)$. The gradient feature of context-action pair $(\mathbf{x}, a)$ is the Jacobian of the scalar output with respect to all network parameters:

where $d_\theta$ is the total number of parameters. This vector plays exactly the role that $\mathbf{x}$ played in LinUCB. To track which context-action pairs have been explored, maintain the empirical gram matrix of past gradient features:

The NeuralUCB score for arm $a$ at round $t$ is then

where $\|\mathbf{v}\|_{\mathbf{M}^{-1}} = \sqrt{\mathbf{v}^\top \mathbf{M}^{-1} \mathbf{v}}$ is the weighted norm. The structure is identical to LinUCB: the first term is the network's predicted reward (exploitation) and the second is an exploration bonus that is large when $(\mathbf{x}_t, a)$ lies in a gradient direction not yet well covered by past pulls, and small when it lies in a well-explored direction. The matrix $\mathbf{Z}_t$ is the analogue of the LinUCB covariance accumulator $\mathbf{A}_a$, but now pooled across arms and tracking gradient-space coverage rather than context-space coverage.

The theoretical justification rests on the neural tangent kernel (NTK). For sufficiently wide networks initialized randomly, the NTK approximation states that the network behaves like a linear model in the gradient feature space: as the network trains, the gradients $\mathbf{g}_a(\mathbf{x})$ move slowly relative to the initial gradients, so the linearization $f_\theta(\mathbf{x}, a) \approx f_{\theta_0}(\mathbf{x}, a) + \mathbf{g}_a(\mathbf{x}_0)^\top (\theta - \theta_0)$ is tight. Under this regime NeuralUCB enjoys regret $\tilde{O}(\sqrt{T})$ (Zhou et al. 2020), matching the linear bandit rate while adapting to nonlinear reward functions.

6.2 NeuralTS: Posterior Sampling Instead of Optimism

NeuralTS replaces the deterministic UCB bonus with Thompson sampling in gradient feature space. Rather than computing the upper confidence bound, at each round it samples a perturbed predicted reward:

and selects $a_t = \arg\max_a \tilde{f}_a(\mathbf{x}_t)$. The variance of the sample is exactly the NeuralUCB bonus squared, so NeuralTS and NeuralUCB share the same geometric intuition about what counts as "unexplored," but NeuralTS avoids the need to tune $\beta$ explicitly. In practice NeuralTS is often preferred in production systems: the stochastic sampling provides natural exploration calibration across arms, and the absence of $\beta$ tuning reduces one hyperparameter from the deployment checklist. Both algorithms achieve the same regret order, but empirical comparisons across recommendation datasets (Zhang et al. 2021) show NeuralTS converging more reliably under noisy reward conditions.

6.3 The Last-Layer Approximation (Neural Linear)

The full NeuralUCB gradient $\mathbf{g}_a(\mathbf{x}) \in \mathbb{R}^{d_\theta}$ with $d_\theta \sim 10^6$ parameters makes the gram matrix $\mathbf{Z} \in \mathbb{R}^{d_\theta \times d_\theta}$ completely intractable: storing it alone requires terabytes. The standard practical remedy is the last-layer approximation, also called the neural linear method. Decompose the network as $f_\theta(\mathbf{x}, a) = \mathbf{w}^\top \phi(\mathbf{x}, a)$ where $\phi(\mathbf{x}, a) \in \mathbb{R}^h$ is the penultimate-layer representation and $\mathbf{w} \in \mathbb{R}^h$ is the final linear head. Freeze all lower layers and track uncertainty only in $\mathbf{w}$. The effective gradient is then $\mathbf{g}_a(\mathbf{x}) \approx \phi(\mathbf{x}, a)$, the feature vector produced by the frozen representation, and the gram matrix shrinks to $\mathbf{Z} \in \mathbb{R}^{h \times h}$ where $h$ is typically 64 to 512. This reduces computational cost from $O(d_\theta^2)$ to $O(h^2)$, a factor of one hundred to one million improvement. The accuracy trade-off is mild in practice: the representation $\phi$ captures most of the nonlinearity learned by the network, and running exact Bayesian linear regression on top of it gives well-calibrated uncertainty estimates over the linear head. Neural linear is the method actually deployed in most large-scale contextual bandit systems.

The three algorithms now in view, LinUCB, NeuralUCB, and neural linear, form a progression of expressiveness against computational cost. LinUCB is exact and cheap but fails on nonlinear reward surfaces. Full NeuralUCB captures arbitrary nonlinearity but is computationally impractical beyond toy networks. Neural linear sits in the middle: it learns a nonlinear representation with standard deep learning and then runs exact LinUCB on the last layer, inheriting the practical deployability of LinUCB with the representation power of neural networks. For production systems where reward functions are genuinely nonlinear, neural linear is the current default and the right starting point before considering more expensive alternatives.

import torch
import torch.nn as nn
import numpy as np

class NeuralLinearBandit:
    """
    Neural-linear contextual bandit (last-layer approximation to NeuralUCB).
    The lower layers learn a representation phi(x, a); the final linear head
    runs exact Bayesian ridge regression (LinUCB-style) for exploration.
    """

    def __init__(self, context_dim: int, n_arms: int,
                 hidden: int = 64, lam: float = 1.0, alpha: float = 1.0):
        self.K = n_arms
        self.alpha = alpha
        self.h = hidden

        # Shared representation trunk (frozen after each batch update in practice)
        self.net = nn.Sequential(
            nn.Linear(context_dim + n_arms, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
        )
        # Last-layer linear head (one output per "feature", then dot with arm embedding)
        self.head = nn.Linear(hidden, 1, bias=False)

        # Per-arm gram matrix Z_a and response accumulator b_a  (last-layer only)
        self.Z = [lam * torch.eye(hidden) for _ in range(n_arms)]
        self.b = [torch.zeros(hidden) for _ in range(n_arms)]
        self.lam = lam

    def _one_hot(self, a: int) -> torch.Tensor:
        v = torch.zeros(self.K)
        v[a] = 1.0
        return v

    def phi(self, x: torch.Tensor, a: int) -> torch.Tensor:
        """Last-layer feature vector for (context, arm) pair."""
        xa = torch.cat([x, self._one_hot(a)])
        with torch.no_grad():
            return self.net(xa)               # shape: (hidden,)

    def ucb_score(self, x: torch.Tensor, a: int) -> float:
        """Exploitation + last-layer exploration bonus."""
        phi = self.phi(x, a)
        Z_inv = torch.linalg.inv(self.Z[a])
        theta_hat = Z_inv @ self.b[a]         # ridge estimate of last-layer weights
        exploit = (phi @ theta_hat).item()
        bonus = self.alpha * torch.sqrt(phi @ Z_inv @ phi).item()
        return exploit + bonus

    def select_arm(self, x: torch.Tensor) -> int:
        scores = [self.ucb_score(x, a) for a in range(self.K)]
        return int(np.argmax(scores))

    def update(self, x: torch.Tensor, a: int, reward: float):
        """Rank-1 update of the last-layer gram matrix."""
        phi = self.phi(x, a)
        self.Z[a] += torch.outer(phi, phi)
        self.b[a] += reward * phi

# --- quick smoke test on a synthetic nonlinear bandit ---
rng = np.random.default_rng(42)
T, d, K = 500, 4, 3
agent = NeuralLinearBandit(context_dim=d, n_arms=K, hidden=32, alpha=1.5)

cumulative_regret = 0.0
for t in range(T):
    x_np = rng.standard_normal(d).astype(np.float32)
    x = torch.from_numpy(x_np)
    # True nonlinear reward: r_a = sin(x[0]) + cos(a * x[1])
    true_rewards = [float(np.sin(x_np[0]) + np.cos(a * x_np[1])) for a in range(K)]
    a_t = agent.select_arm(x)
    r_t = true_rewards[a_t] + 0.1 * rng.standard_normal()
    agent.update(x, a_t, r_t)
    cumulative_regret += max(true_rewards) - true_rewards[a_t]

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

import vowpalwabbit as vw

# Neural contextual bandit: 3 arms, neural reward model, epsilon-greedy explore
vw_agent = vw.Workspace(
    "--cb_explore_adf --neural --epsilon 0.05 --quiet",
    enable_cache=False,
)

def to_vw_format(context: list[float], n_arms: int, chosen: int = None,
                 cost: float = None, prob: float = None) -> str:
    """Convert context + optional feedback to VW ADF format."""
    header = "shared |ctx " + " ".join(f"f{i}:{v:.4f}" for i, v in enumerate(context))
    arms = [f"|arm a{a}" for a in range(n_arms)]
    if chosen is not None:                         # include feedback on chosen arm
        arms[chosen] = f"{chosen}:{cost:.4f}:{prob:.4f} |arm a{chosen}"
    return header + "\n" + "\n".join(arms)

# Example: observe context, pick arm, receive reward, update
context = [0.3, -0.1, 0.8, 0.2]
example = to_vw_format(context, n_arms=3)
pred = vw_agent.predict(example)                   # returns chosen arm + prob
chosen_arm, prob = pred[0][0], pred[0][1]
reward = 1.0                                       # from environment
cost = -reward                                     # VW minimises cost, not reward
feedback = to_vw_format(context, n_arms=3, chosen=chosen_arm, cost=cost, prob=prob)
vw_agent.learn(feedback)

Applications of neural contextual bandits in temporal AI appear wherever the reward function couples context features nonlinearly and the context itself is time-varying. Recommendation systems at scale match time-of-day, session history, and item freshness in ways that no linear model captures; neural linear with pre-trained user-item embeddings as $\phi$ is the dominant industrial pattern. Adaptive clinical dosing uses patient biomarkers, current drug levels, and prior response trajectories as context; the NTK-based uncertainty bound translates to rigorous coverage guarantees that regulators can audit. Financial execution uses market microstructure features (bid-ask spread, order-flow imbalance, volatility regime) as context to route orders; NeuralTS avoids the $\beta$ tuning that UCB requires under rapidly shifting feature distributions. In each case the temporal dimension is in the context, not in the dynamics, which is exactly the contextual-bandit regime: the action does not change the next context, so the full MDP machinery of later chapters is unnecessary, and the bandit's cheaper exploration suffices.

We have given the bandit eyes, a feature vector revealed before each choice, and watched that single addition lift it from a global average-chaser to a situation-aware function learner with sublinear regret. We derived LinUCB from ridge regression and a confidence ellipsoid, surveyed Thompson sampling and neural exploration, and built the off-policy evaluation machinery that scores new policies from logged data. But the contextual bandit deliberately threw away one thing the full decision problem keeps: transitions. Its action affects only the immediate reward, never the next situation. The moment the action does change the next state, immediate reward is no longer enough, and we must reason about long-run return through a sequence of choices. That is the leap to full reinforcement learning. Section 24.4 takes it, introducing Monte Carlo and temporal-difference learning, the two foundational families for estimating long-run value when actions ripple forward through time.