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

arXiv: https://arxiv.org/abs/2607.28408 · HTML: https://mlsanity.com/paper/2607.28408/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 27 of its 158 references that are in the map.
- First posted on arXiv 30 Jul 2026.

## Builds on (most cited first, 27 of 27)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [User-friendly introduction to PAC-Bayes bounds](https://mlsanity.com/paper/2110.11216/) (2021)
- [Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks for Thompson Sampling](https://mlsanity.com/paper/1802.09127/) (2018)
- [Handbook of Convergence Theorems for (Stochastic) Gradient Methods](https://mlsanity.com/paper/2301.11235/) (2023)
- [A Primer on PAC-Bayesian Learning](https://mlsanity.com/paper/1901.05353/) (2019)
- [Off-Policy Evaluation for Large Action Spaces via Embeddings](https://mlsanity.com/paper/2202.06317/) (2022)
- [Impact of Representation Learning in Linear Bandits](https://mlsanity.com/paper/2010.06531/) (2020)
- [Meta Dynamic Pricing: Transfer Learning Across Experiments](https://mlsanity.com/paper/1902.10918/) (2019)
- [Upper Counterfactual Confidence Bounds: a New Optimism Principle for Contextual Bandits](https://mlsanity.com/paper/2007.07876/) (2020)
- [PAC-Bayes Generalisation Bounds for Heavy-Tailed Losses through Supermartingales](https://mlsanity.com/paper/2210.00928/) (2022)
- [Multi-task Representation Learning with Stochastic Linear Bandits](https://mlsanity.com/paper/2202.10066/) (2022)
- [PAC-Bayesian Offline Contextual Bandits With Guarantees](https://mlsanity.com/paper/2210.13132/) (2022)
- [Efron-Stein PAC-Bayesian Inequalities](https://mlsanity.com/paper/1909.01931/) (2019)
- [Oracle-Efficient Pessimism: Offline Policy Optimization in Contextual Bandits](https://mlsanity.com/paper/2306.07923/) (2023)
- [Towards Scalable and Robust Structured Bandits: A Meta-Learning Framework](https://mlsanity.com/paper/2202.13227/) (2022)
- [Local Policy Improvement for Recommender Systems](https://mlsanity.com/paper/2212.11431/) (2022)
- [Probabilistic Rank and Reward: A Scalable Model for Slate Recommendation](https://mlsanity.com/paper/2208.06263/) (2022)
- [VITS : Variational Inference Thompson Sampling for contextual bandits](https://mlsanity.com/paper/2307.10167/) (2023)
- [Combining Reward and Rank Signals for Slate Recommendation](https://mlsanity.com/paper/2107.12455/) (2021)
- [A Unified Approach to Translate Classical Bandit Algorithms to the Structured Bandit Setting](https://mlsanity.com/paper/1810.08164/) (2018)
- [Offline Evaluation of Reward-Optimizing Recommender Systems: The Case of Simulation](https://mlsanity.com/paper/2209.08642/) (2022)
- [Off-Policy Learning in Large Action Spaces: Optimization Matters More Than Estimation](https://mlsanity.com/paper/2509.03456/) (2025)
- [Offline Contextual Bandit with Counterfactual Sample Identification](https://mlsanity.com/paper/2509.10520/) (2025)

## Abstract (arXiv)

This thesis studies policy learning in interactive systems where an agent observes a context, selects an action from a very large set, and receives partial feedback. The main framework is contextual bandits, with two paradigms: on-policy learning, where the agent interacts sequentially with the environment and minimizes regret, and off-policy learning, where it learns from logged data collected by a logging policy. In large action spaces, both settings face major challenges: inefficient exploration, sparse data coverage, high-variance importance weights, extrapolation bias, and difficult optimization landscapes. The first part develops structured Bayesian methods for on-policy learning. We introduce meTS, a mixed-effect extension of Thompson sampling, and dTS, which leverages diffusion-inspired priors to model dependencies between actions. These methods share information across actions and yield regret guarantees depending on an effective number of actions. The second part addresses off-policy learning. We propose sDM, a structured direct method based on latent variables, show that optimization error can dominate estimation error in large action spaces, and introduce concave, efficiently optimizable policy-weighted log-likelihood objectives. Finally, we develop differentiable pessimistic methods based on exponential smoothing and PAC-Bayesian bounds to control the bias-variance trade-off of regularized importance-sampling estimators.
