📊 statistics

On-Policy and Off-Policy Learning for Large Action Spaces

This thesis addresses the challenges of policy learning in contextual bandits with large action spaces by proposing structured Bayesian methods for on-policy learning to improve exploration and regret bounds, alongside novel off-policy techniques that mitigate estimation errors and control bias-variance trade-offs through optimized objectives and differentiable pessimistic approaches.

Original authors: Imad Aouali

Published 2026-07-31

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Imad Aouali

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are the captain of a massive spaceship trying to find the best route through a galaxy with millions of stars. Every time you pick a star to visit, you get a tiny, fuzzy signal telling you if it was a good choice or a bad one. This is the world of contextual bandits, a branch of artificial intelligence that helps computers make decisions when they don't know the rules of the game yet. The "context" is the situation you're in (like the weather or your mood), the "action" is what you do (like picking a star), and the "reward" is the result (like finding treasure or hitting an asteroid).

The tricky part is the sheer number of choices. If you have to guess which of a million stars is the best one, and you can only check a few at a time, you might spend your whole life exploring the wrong ones. This is the "large action space" problem. It's like trying to find a single specific needle in a haystack the size of a city, but you can only pull out one straw at a time and hope it's the needle. Scientists care about this because it's the engine behind things like recommending movies, showing you the right ads, or even designing new medicines. If the computer gets stuck guessing randomly, it wastes time and money.

This thesis tackles the problem of how to teach a computer to make smart choices when faced with millions of options, using two different strategies: learning while you go (on-policy) and learning from old logs (off-policy).

The On-Policy Adventure: Learning by Doing with a Map

First, the author looks at the "on-policy" scenario, where the computer learns by interacting with the world in real-time. Imagine you are exploring a giant library with millions of books, but you don't know which ones are good. A standard explorer would pick a book, read a page, and if it's boring, move to a completely different book, starting from scratch. This is slow and inefficient.

The paper introduces a smarter explorer using Mixed-Effect Thompson Sampling (meTS). Instead of treating every book as a unique mystery, this explorer notices that books belong to genres. It learns that "Sci-Fi" books share common traits. By grouping books into categories (like "Action," "Romance," or "Mystery"), the explorer can learn about the whole genre from just a few books. If it reads one great Sci-Fi book, it gets a hint that other Sci-Fi books might be good too. This "sharing of information" drastically speeds up learning. The math shows that instead of needing to learn about millions of individual books, the computer only needs to learn about a few dozen "genres" (latent effects) and the specific quirks of each book within those genres.

The author then takes this idea even further with Diffusion Thompson Sampling (dTS). If the first method was like grouping books by genre, this new method is like having a super-smart librarian who understands the deep, complex connections between books. Maybe a book is a mix of "Cyberpunk" and "Historical Fiction," or maybe it shares a specific writing style with a book from a different century. Using a type of AI called a "diffusion model" (the same tech behind some image generators), the computer learns a rich, deep map of how all the books relate to each other. This allows it to explore the library much faster, even if the library is huge. In simulations, these methods found the best books much quicker than older methods that treated every book as a stranger.

The Off-Policy Challenge: Learning from a Messy Diary

Next, the paper tackles the "off-policy" scenario. Imagine you can't explore the library yourself anymore. Instead, you have to learn from a messy diary left by a previous explorer who had very different tastes. Maybe that explorer only read horror movies, and now you need to find the best romance movies. This is the "off-policy" problem: learning from data collected by someone else.

The author challenges a common belief in the field: that the most important thing is to build the most accurate "reward estimator" (a crystal ball that predicts how good a choice will be). The paper argues that in huge libraries, optimization is actually the bigger problem. It's like having a perfect map (the estimator) but trying to navigate it with a broken compass (the optimization algorithm). The math shows that standard ways of using these maps often get stuck in "flat plateaus" or local traps, making it impossible to find the best path, no matter how good the map is.

To fix this, the author proposes a new approach: Policy-Weighted Log-Likelihood (PWLL). Instead of trying to predict the exact reward, this method focuses on making the optimization path smooth and easy to walk. It's like switching from a jagged, rocky mountain path to a gentle, winding road. Even if the road isn't perfectly straight, it's much easier to get to the top. In experiments with up to one million actions, this simple, smooth approach consistently beat the complex, "smart" estimators that were getting stuck.

The paper also introduces a new way to handle the "noise" in the old diary. When the previous explorer rarely visited certain sections, the data is unreliable. The author suggests using Exponential Smoothing combined with "principled pessimism." Think of this as being a cautious explorer who trusts the diary but adds a safety buffer. If the diary says a path is great but the data is shaky, the explorer assumes it might be slightly worse than reported to avoid disaster. The paper proves mathematically that this method keeps the explorer safe while still allowing them to learn effectively, and it works well even when the data is sparse.

The Big Picture

In short, this thesis shows that when you have millions of choices, you can't just brute-force your way through. You need to find the hidden structures (like genres or deep connections) to share what you learn, and you need to make sure your learning path is smooth enough to actually find the solution. Whether you are learning in real-time or digging through old logs, the key is to be smart about how you group information and how you navigate the math. The results, tested on both fake data and real-world movie recommendation datasets, suggest that these new methods are a significant step forward in making AI decision-making scalable and efficient.

Technical Summary: On-Policy and Off-Policy Learning for Large Action Spaces

Problem Statement

The thesis addresses the fundamental challenge of learning in interactive systems, specifically contextual bandits, where the action space $K$ is extremely large (ranging from thousands to millions). In such settings, standard algorithms for both on-policy (online) and off-policy (offline) learning face severe scalability bottlenecks:

On-Policy: Standard exploration strategies like Thompson Sampling (TS) or Upper Confidence Bound (UCB) typically rely on disjoint models where each action has an independent parameter. In large action spaces, learning each parameter independently is statistically inefficient, leading to high cumulative regret. Conversely, joint models that share parameters across all actions are often computationally intractable due to the $O(K^2)$ or $O(K^3)$ complexity of maintaining and updating a full covariance matrix.
Off-Policy: Learning from static logged data collected by a logging policy $\pi_0$ $π_{0}$ faces three main issues:
1. Statistical Inefficiency: Direct methods (DM) suffer from high bias when actions are rarely observed.
2. Variance: Inverse Propensity Scoring (IPS) estimators exhibit exploding variance when the logging probability $\pi_0(a|x)$ is small.
3. Optimization Intractability: The thesis argues that in large action spaces, the optimization landscape of standard IPS-based objectives becomes highly non-concave with flat plateaus and exponentially many local maxima, causing gradient-based optimizers to fail regardless of the statistical quality of the estimator.

Methodology and Contributions

The thesis is divided into two parts, proposing novel methods for each learning paradigm.

Part I: On-Policy Learning (Scaling Thompson Sampling)

The core insight is to retain the flexibility of disjoint reward models while introducing structured Bayesian priors that enable information sharing across actions.

1. Mixed-Effect Thompson Sampling (meTS) - Chapter 3

Methodology: The authors introduce a hierarchical Bayesian model where action parameters $\theta_a$ are coupled through $L$ shared latent effect parameters $\Psi = (\psi_\ell)_{\ell \in [L]}$ . The generative process assumes $\theta_a | \Psi \sim p_0(\cdot|\Psi)$ .
Algorithm: meTS performs hierarchical sampling: first sampling the global effects $\Psi_t$ from their posterior, then sampling individual action parameters $\theta_{t,a}$ conditioned on $\Psi_t$ .
Theoretical Results: For linear-Gaussian models, the authors derive exact closed-form posterior updates. They prove a Bayesian regret bound of $\tilde{O}(\sqrt{Td K_{eff}})$ , where $K_{eff}$ is an effective number of actions. When the number of effects $L \ll K$ , $K_{eff} \ll K$ , yielding a multiplicative improvement over standard TS.
Complexity: The method reduces memory complexity from $O(K^2d^2)$ to $O((L^2 + K)d^2)$ and runtime from $O(K^3d^3)$ to $O((L^3 + K)d^3)$ .
Extensions: The framework is extended to Generalized Linear Models (GLMs) using tractable Laplace approximations.

2. Diffusion Thompson Sampling (dTS) - Chapter 4

Methodology: To capture more complex, non-linear dependencies between actions, the authors replace the single-layer latent structure with a deep hierarchy of latent variables governed by pre-trained diffusion models. The prior is defined as a Markov chain of latent variables $\psi_L \to \dots \to \psi_1 \to \theta_a$ connected by non-linear link functions $f_\ell$ .
Algorithm: dTS employs a recursive inference procedure. It approximates the intractable posterior by maintaining a "posterior diffusion model" where the link functions and covariances are updated via Gaussian-like updates (combining prior precision and data precision) without backpropagating through the neural network.
Theoretical Results: In a linear-Gaussian setting, the regret bound scales with the effective number of actions and the sum of variances across the $L$ layers. The complexity scales linearly with $L$ , i.e., $O((L+K)d^2)$ , improving upon meTS.
Significance: This demonstrates that deep generative models can be used not just for representation, but to structure the uncertainty required for efficient exploration.

Part II: Off-Policy Learning (Addressing Estimation and Optimization)

The thesis challenges the prevailing "estimation-centric" paradigm and proposes three complementary approaches.

1. Structured Direct Method (sDM) - Chapter 6

Methodology: Analogous to meTS, sDM introduces a shared latent vector $\psi$ to couple action parameters in the offline setting. It computes the posterior over latent effects and conditional posteriors for action parameters.
Results: The authors prove that sDM achieves $O(1/\sqrt{n})$ convergence in Bayesian suboptimality without requiring the restrictive "full logging support" assumption. The bound depends on the alignment between the logging policy and the optimal policy.
Insight: Under the Bayesian suboptimality criterion considered, greedy policies are shown to be optimal and can outperform pessimistic policies, contrasting with frequentist settings where pessimism is often necessary.

2. Optimization Matters More Than Estimation - Chapter 7

Methodology: The authors demonstrate that in large action spaces, optimization error dominates estimation error. Standard IPS-based objectives induce non-concave landscapes that trap optimizers.
Proposal: Instead of estimating value and then optimizing, the thesis advocates for Policy-Weighted Log-Likelihood (PWLL) objectives:
$\hat{U}_g(\pi) = \frac{1}{n} \sum_{i=1}^n g(R_i, \pi_0(A_i|X_i)) \log \pi(A_i|X_i)$
Results: For linear softmax policies, PWLL objectives are concave (and strongly concave with $\ell_2$ regularization), guaranteeing a unique global maximum. Experiments on datasets with up to one million actions show that these simple, stable objectives consistently outperform sophisticated value-estimator-based methods.

3. Principled Pessimism for Exponential Smoothing - Chapter 8

Methodology: To address the variance of IPS, the authors introduce differentiable Exponential Smoothing (ES) estimators that interpolate between no regularization and strong variance reduction.
Theory: A two-sided PAC-Bayesian generalization bound is derived for these regularized estimators. This bound decomposes the error into divergence, regularization bias, and remaining variance.
Application: The bound yields a fully differentiable, pessimistic learning objective that can be optimized via stochastic gradient ascent. The framework is unified to cover other regularizers like clipping and implicit exploration.
Extension: The authors also introduce Logarithmic Smoothing (LS), a soft, differentiable alternative to clipping that achieves tighter concentration bounds.

Experimental Results

The proposed methods are validated on both synthetic and real-world datasets:

On-Policy: In synthetic linear and logistic bandits, meTS and dTS significantly outperform baselines (LinTS, GLM-TS, HierTS), with the performance gap widening as the action space $K$ increases. In MovieLens experiments (recommender systems), meTS achieves the lowest regret even when the true reward model is misspecified, demonstrating robustness.
Off-Policy:
- sDM outperforms standard Direct Methods, with gains increasing as $K$ grows.
- PWLL objectives (Chapter 7) consistently outperform state-of-the-art IPS-based methods on datasets with up to one million actions, showing superior robustness to hyperparameter tuning.
- The PAC-Bayesian pessimistic objectives (Chapter 8) achieve better rewards than baselines (including London et al. and Sakhi et al.) across various logging policy qualities, particularly when the logging policy is not uniform.

Significance and Claims

The thesis claims to bridge the gap between theoretical bandit algorithms and the constraints of modern industrial systems (e.g., recommender systems with millions of items). Its central argument is that scaling to large action spaces requires:

Structuring Uncertainty: Using hierarchical and diffusion-based priors to share statistical strength across actions (On-Policy).
Re-evaluating Objectives: Recognizing that in large spaces, the tractability of the optimization landscape is as critical as the statistical accuracy of the estimator (Off-Policy).
Differentiable Pessimism: Developing theoretically grounded, yet computationally feasible, pessimistic objectives that control variance and bias without relying on intractable constants.

The work asserts that by moving away from independent action modeling and standard value estimation, contextual bandits can be effectively deployed in environments where action spaces routinely exceed thousands or millions.