[Paper Review] Stochastic Bandits with Context Distributions

Johannes Kirschner, Andreas Krause|arXiv (Cornell University)|Jun 6, 2019

Advanced Bandit Algorithms ResearchDecision Sciences34 references17 citations

TL;DR

This paper introduces a novel stochastic contextual bandit model where the environment provides a distribution over contexts rather than the exact context, and the learner only observes the distribution. The authors adapt the UCB algorithm to this setting, proving an order-optimal O(d√T) high-probability regret bound for linear and kernelized reward functions, which generalizes standard contextual bandits when distributions are Dirac deltas.

ABSTRACT

We introduce a stochastic contextual bandit model where at each time step the environment chooses a distribution over a context set and samples the context from this distribution. The learner observes only the context distribution while the exact context realization remains hidden. This allows for a broad range of applications where the context is stochastic or when the learner needs to predict the context. We leverage the UCB algorithm to this setting and show that it achieves an order-optimal high-probability bound on the cumulative regret for linear and kernelized reward functions. Our results strictly generalize previous work in the sense that both our model and the algorithm reduce to the standard setting when the environment chooses only Dirac delta distributions and therefore provides the exact context to the learner. We further analyze a variant where the learner observes the realized context after choosing the action. Finally, we demonstrate the proposed method on synthetic and real-world datasets.

Motivation & Objective

• To address the challenge of contextual bandit learning when the true context is unobserved and only a distribution over contexts is available.
• To develop a learning algorithm that competes with the best action choice based on the context distribution, rather than the realized context.
• To establish high-probability regret bounds that are order-optimal in the feature dimension d and horizon T.
• To extend the framework to kernelized reward functions using kernel mean embeddings and distributional risk minimization.
• To analyze a variant where the realized context is observed after action selection, broadening applicability.

Proposed method

• The model assumes the environment selects a context distribution μt at each round t, and the context ct is drawn from μt, but only μt is observed by the learner.
• The reward function is assumed to be linear in the feature vector φx,c: f(x,c) = φx,c⊤θ, with unknown θ ∈ ℝd.
• The UCB algorithm is adapted to this setting by constructing confidence sets over θ using a least-squares estimator from observed rewards and context distributions.
• A practical variant of the algorithm is proposed that only requires sampling access to the context distributions μt, not full knowledge of them.
• For the kernelized setting, the reward function is assumed to lie in a known RKHS, and the method uses kernel mean embeddings to estimate the expected reward.
• The analysis leverages concentration inequalities and information-theoretic tools to bound the regret, showing O(d√T) high-probability regret for the linear case.

Experimental results

Research questions

• RQ1 Can a contextual bandit algorithm achieve sublinear regret when the true context is hidden and only a distribution over contexts is observed?
• RQ2 Does the UCB algorithm, adapted to context distributions, achieve order-optimal regret in the linear bandit setting?
• RQ3 How does the regret bound scale with the feature dimension d and time horizon T in the presence of context distribution uncertainty?
• RQ4 Can the framework be extended to kernelized reward functions using kernel mean embeddings and distributional risk minimization?
• RQ5 What is the performance gain of observing the realized context after action selection, compared to only observing the distribution?

Key findings

• The proposed UCB-based algorithm achieves an O(d√T) high-probability regret bound in the linear bandit setting, which is order-optimal up to logarithmic factors.
• The regret bound holds even when the environment adaptively chooses context distributions, and the learner only observes the distributions, not the realized contexts.
• The method generalizes the standard contextual bandit model: when context distributions are Dirac delta distributions, the algorithm reduces to the classical UCB.
• For the kernelized setting, the method uses kernel mean embeddings to estimate expected rewards and achieves a regret bound that depends on the kernel's eigenspectrum decay.
• The algorithm remains practical in the real-world setting where only sampling access to context distributions is available, without requiring full distributional knowledge.
• In a variant where the realized context is observed after action selection, the regret bound improves, enabling better performance in applications requiring context prediction at decision time.