[Paper Review] Conformal Off-Policy Prediction in Contextual Bandits

Muhammad Faaiz Taufiq, Jean-François Ton+3 more|arXiv (Cornell University)|Jun 9, 2022

Advanced Bandit Algorithms Research4 citations

TL;DR

This paper proposes Conformal Off-Policy Prediction (COPP), a novel method for constructing finite-sample valid predictive intervals for outcomes in contextual bandits under a target policy using observational data. COPP leverages conformal prediction to deliver covariate-adaptive, distribution-free intervals with marginal coverage guarantees, outperforming existing methods in coverage accuracy and interval width across synthetic and real-world benchmarks.

ABSTRACT

Most off-policy evaluation methods for contextual bandits have focused on the expected outcome of a policy, which is estimated via methods that at best provide only asymptotic guarantees. However, in many applications, the expectation may not be the best measure of performance as it does not capture the variability of the outcome. In addition, particularly in safety-critical settings, stronger guarantees than asymptotic correctness may be required. To address these limitations, we consider a novel application of conformal prediction to contextual bandits. Given data collected under a behavioral policy, we propose \emph{conformal off-policy prediction} (COPP), which can output reliable predictive intervals for the outcome under a new target policy. We provide theoretical finite-sample guarantees without making any additional assumptions beyond the standard contextual bandit setup, and empirically demonstrate the utility of COPP compared with existing methods on synthetic and real-world data.

Motivation & Objective

• To address the limitations of traditional off-policy evaluation methods that focus only on expected outcomes and lack finite-sample validity.
• To develop a method that provides reliable predictive intervals capturing outcome variability, especially in small-sample or risk-sensitive settings.
• To ensure coverage guarantees that are adaptive to covariates $X$, avoiding overly conservative intervals common in existing approaches.
• To extend conformal prediction to stochastic policies and continuous action spaces in the contextual bandit setting.
• To empirically validate that COPP maintains target coverage levels while producing tighter intervals than competing methods.

Proposed method

• COPP applies conformal prediction to off-policy outcome prediction by constructing prediction sets $\hat{C}(x)$ that are guaranteed to contain the true outcome $Y$ with pre-specified probability.
• It uses inverse probability weighting to reweight observed data under the behavioral policy $\pi^b$ to estimate the distribution of outcomes under the target policy $\pi^*$.
• The method defines a score function $s(x,y)$ based on the estimated conditional probability $\hat{P}^{\pi^b}(y|x)$ to rank outcomes for conformal prediction.
• Prediction sets are formed by including outcomes with the lowest scores up to the $1-\alpha$ quantile of the empirical distribution of scores on a calibration set.
• For discrete outcomes, the method computes the quantile directly over the outcome space $\mathcal{Y}$, avoiding grid-based approximations.
• The approach ensures finite-sample marginal coverage: $\mathbb{P}(Y \in \hat{C}(X)) \geq 1 - \alpha - o(n^{-1})$, with asymptotic conditional coverage under mild regularity conditions.

Experimental results

Research questions

• RQ1 Can conformal prediction be adapted to provide finite-sample, covariate-adaptive predictive intervals for off-policy outcomes in contextual bandits?
• RQ2 How does COPP’s coverage performance compare to existing off-policy evaluation methods like WIS and SBA in terms of accuracy and interval width?
• RQ3 Does COPP maintain reliable coverage across varying degrees of policy shift and heterogeneous data distributions?
• RQ4 What are the conditions under which COPP achieves asymptotic conditional coverage, and how do they relate to model misspecification?
• RQ5 Can COPP be extended to handle label-conditioned coverage guarantees in discrete outcome settings?

Key findings

• COPP achieves coverage levels that are consistently close to the nominal $90\%$ level across all tested target policies, while WIS and SBA methods become increasingly conservative with larger policy shifts.
• In experiments with $n = 5000$ calibration points, COPP maintained $90.1\%$ average coverage, significantly closer to the target than WIS ($94.3\%$) and SBA ($95.1\%$).
• COPP produces narrower predictive intervals than WIS and SBA, especially in discrete outcome settings where WIS sets often include all labels due to lack of $X$-adaptivity.
• The method demonstrates robustness to policy shift and heteroscedasticity, avoiding the over-conservativeness seen in non-adaptive methods.
• Label-conditioned coverage is achieved via class-balanced conformal prediction, with empirical results showing $\geq 90\%$ conditional coverage per label, unlike standard COPP which may fail under such constraints.
• Empirical results confirm that COPP’s coverage improves with increasing calibration data size, converging toward the target level, while competing methods remain overly conservative.