LinUCB Contextual Bandit

Background

LinUCB is a contextual bandit for personalised recommendations. Each arm learns a linear model $\hat\theta_a = A_a^{-1} b_a$ over the context vector; arm selection adds an upper confidence bound bonus proportional to the model's uncertainty in the current context direction. Classic baseline behind ad ranking and news personalisation.

Problem statement

Implement LinUCB(n_arms, d, alpha=1.0) with:

• choose(x) — return int arm index with the highest UCB score:

$$\text{UCB}_a(x) \;=\; \hat\theta_a^\top x \;+\; \alpha \sqrt{x^\top A_a^{-1} x}$$

• update(arm, x, reward) — Bayesian-style update: $A_a \mathrel{+}= xx^\top$, $b_a \mathrel{+}= r\,x$.

Input

• n_arms — int >= 1, number of arms.
• d — int >= 1, context dimension.
• alpha — float > 0, exploration weight (default $1.0$).
• choose(x) — x is a length-$d$ array.
• update(arm, x, reward) — arm in $[0, n_{\text{arms}})$, x length-$d$, reward scalar.

Output

• choose returns an int arm index.
• update returns None (mutates state).

Examples

Example 1 — cold start: all arms tied → arm 0

b = LinUCB(3, d=2)
b.choose([1.0, 0.0])  # -> 0  (tie-break by lower index)

Example 2 — positive reward shifts preference

b = LinUCB(2, d=2, alpha=0.01)
for _ in range(6): b.update(1, [1.0, 0.0], 10.0)
b.choose([1.0, 0.0])  # -> 1

Constraints

• A initialised to identity for each arm; b initialised to zero.
• alpha controls exploration vs exploitation; higher = more exploration.
• Same input from a deterministic state returns the same arm.

Notes

• Regret. LinUCB achieves $\tilde O(d\sqrt{T})$ regret under linear reward — provably better than ε-greedy for context dimensions that matter.
• Disjoint vs hybrid. This is the disjoint variant (separate $\theta$ per arm). Hybrid LinUCB shares a common $\theta_0$ across arms plus per-arm $\theta_a$.