Context-Action Embedding Learning for Off-Policy Evaluation in Contextual Bandits

TL;DR

This paper introduces CAEL-MIPS, a novel method that learns context-action embeddings to minimize mean squared error in off-policy evaluation for contextual bandits, improving estimation accuracy especially in large or sparse action spaces.

Contribution

It proposes a new embedding learning approach that optimizes for MSE reduction in MIPS estimators, addressing limitations of previous methods that did not consider context or MSE minimization.

Findings

CAEL-MIPS outperforms baseline methods in MSE on synthetic data.

The approach effectively leverages context-action information to improve off-policy evaluation.

Empirical results demonstrate the method's superiority in real-world datasets.

Abstract

We consider off-policy evaluation (OPE) in contextual bandits with finite action space. Inverse Propensity Score (IPS) weighting is a widely used method for OPE due to its unbiased, but it suffers from significant variance when the action space is large or when some parts of the context-action space are underexplored. Recently introduced Marginalized IPS (MIPS) estimators mitigate this issue by leveraging action embeddings. However, these embeddings do not minimize the mean squared error (MSE) of the estimators and do not consider context information. To address these limitations, we introduce Context-Action Embedding Learning for MIPS, or CAEL-MIPS, which learns context-action embeddings from offline data to minimize the MSE of the MIPS estimator. Building on the theoretical analysis of bias and variance of MIPS, we present an MSE-minimizing objective for CAEL-MIPS. In the empirical studies on a synthetic dataset and a real-world dataset, we demonstrate that our estimator outperforms baselines in terms of MSE.

Equations

$$v(\pi ):-{\int }_{\mathcal{X}}\sum _{a\in \mathcal{A}}{p}_X(x)\pi (a|x)q(x,a)dx.$$

$$\text{MSE}[\hat{v}(\pi )]$$

$$=\text{Bias}[\hat{v}(\pi ){]}^2+V[\hat{v}(\pi )],$$

$${\hat{v}}_{\text{IPS}}(\pi ):-\frac{1}{n}\sum _{i=1}^n\underset{=:w(X_i,A_i)}{\underbrace{\frac{\pi (A_i|X_i)}{\mu (A_i|X_i)}}}{R}_i,$$

$$|\text{Bias}[{\hat{v}}_{\text{IPS}}(\pi )]|=\mathbb{E}\left[\sum _{a\in \mathcal{U}(X,\mu )}\pi (a|X)q(X,a)\right],$$

$$V[{\hat{v}}_{\text{IPS}}(\pi )]=$$

$$+\mathbb{E}\left[V_{\mu }[w(X,A)q(X,A)]\right],$$

$$v(\pi ):-{\int }_{\mathcal{X}}\sum _{a\in \mathcal{A}}{\int }_{\mathcal{E}}{p}_X(x)\pi (a|x)p_E(e|x,a)q(x,a,e)dxde$$

$${\hat{v}}_{\text{MIPS}}(\pi )$$

$$=\frac{1}{n}\sum _{i=1}^n\sum _{a\in \mathcal{A}}\underset{=:\mu (a|X_i,E_i)}{\underbrace{\frac{\mu (a|X_i)p_E(E_i|X_i,a)}{p_E(E_i|X_i,\mu )}}}w(X_i,a)R_i$$

$$\text{Bias}[{\hat{v}}_{\text{MIPS}}(\pi )]=\mathbb{E}[\sum _{a<b}\mu (a|X,E)\mu (b|X,E)\cdot$$

$$(q(X,a,E)-q(X,b,E))\cdot (w(X,b)-w(X,a))].$$

$$V[{\hat{v}}_{\text{IPS}}(\pi )]-V[{\hat{v}}_{\text{MIPS}}(\pi )]$$

$$=\frac{1}{n}\mathbb{E}[{\mathbb{E}}_{p_R(\cdot |X,E)}[R^2]V_{\mu (\cdot |X,E)}[w(X,A)]],$$

$$q_{\min }\le q(x,a,e)\le q_{\max },$$

$${\sigma }_{\min }^2\le V[R|X=x,E=e]\le {\sigma }_{\max }^2,$$

$$\text{Bias}[{\hat{v}}_{\text{MIPS}}(\pi ){]}^2$$

$$\le (q_{\max }-q_{\min }{)}^2\mathbb{E}[\sum _{a<b}\mu (a|X,E)\mu (b|X,E)|w(X,b)-w(X,a)|{]}^2.$$

$${\mathcal{L}}_{\text{bias}}(\theta )$$

$$|w(X,b)-w(X,a)|{]}^2.$$

$$V[{\hat{v}}_{\text{MIPS}}(\pi )]\le V[{\hat{v}}_{\text{IPS}}(\pi )]$$

$$+\frac{1}{n}\mathbb{E}\left[\left({\sigma }_{\min }^2+\mathbb{E}[R|X,E{]}^2\right)\sum _{a\in \mathcal{A}}\mu (a|X,E{)}^2\sum _{a\in \mathcal{A}}w(X,a{)}^2\right].$$

$${\mathcal{L}}_{\text{var}}(\theta ):-$$

$$\sum _{a\in \mathcal{A}}w(a,X{)}^2]$$

$$\mathcal{L}(\theta ):-{\mathcal{L}}_R(\theta )+\alpha {\mathcal{L}}_{\text{bias}}(\theta )+\beta {\mathcal{L}}_{\text{var}}(\theta ),$$

$${\hat{\mathcal{L}}}_R(\theta ;\mathcal{D}):-\frac{1}{n}\sum _{i=1}^n(\hat{R}(X_i,f_{\theta }(X_i,A_i))-R_i{)}^2,$$

$${\hat{\mathcal{L}}}_{\text{bias}}(\theta ;\mathcal{D})$$

$$:-\frac{1}{n^2}(\sum _{i=1}^n\sum _{a<b}\hat{\mu }(a|X_i,f_{\theta }(X_i,A_i))\hat{\mu }(b|X_i,f_{\theta }(X_i,A_i))\cdot$$

$$|w(X_i,b)-w(X_i,a)|{)}^2,$$

$${\hat{\mathcal{L}}}_{\text{var}}(\theta ;\mathcal{D}):-\frac{1}{n^2}\sum _{i=1}^n\hat{R}(X_i,f_{\theta }(X_i,A_i){)}^2\sum _{a\in \mathcal{A}}\hat{\mu }(a|X_i,f_{\theta }(X_i,A_i){)}^2\cdot$$

$$\sum _{a\in \mathcal{A}}w(a,X_i{)}^2.$$

$$H_2(X,A;\theta )=-\log \sum _{a\in \mathcal{A}}\mu (a|X,f_{\theta }(X,A){)}^2$$

$${\mathcal{L}}_{\text{bias}}(\theta )\le \frac{C^2}{4}{\left(1-\mathbb{E}\left[e^{-H_2(X,A;\theta )}\right]\right)}^2\le \frac{C^2}{4}\left(1-\frac{1}{|\mathcal{A}|}\right),$$

$${\mathcal{L}}_{\text{var}}(\theta )=\frac{1}{n}\mathbb{E}\left[\mathbb{E}[R|X,E{]}^2\cdot e^{-H_2(X,A;\theta )}\sum _{a\in \mathcal{A}}w(X,a{)}^2\right].$$

$$\theta \leftarrow \theta -\eta {\nabla }_{\theta }\hat{\mathcal{L}}(\theta ;\mathcal{B})$$

$${\hat{v}}_{\text{MIPS}}\leftarrow \frac{1}{n}\sum _{i=1}^n\sum _{a\in \mathcal{A}}{\hat{\mu }}_{\varphi }(a|X_i,f_{\theta }(X_i,A_i))w(X_i,a)R_i$$

$${\hat{v}}_{\text{DM}}(\pi )=\frac{1}{n}\sum _{i=1}^n\sum _{a\in \mathcal{A}}\pi (a|x_i)\hat{q}(x_i,a).$$

Full text

Off-policy evaluation, contextual bandits, embedding learning, inverse propensity score.

††ccs: Information systems Recommender systems††ccs: Information systems Evaluation of retrieval results

1. Introduction

The contextual bandit framework (Lattimore and Szepesvári, 2020) has been applied to various web services such as recommendation systems (Li et al., 2010; Xu et al., 2020) and digital marketing (Geng et al., 2021; Thomas et al., 2017). In this framework, the system selects an action based on contextual information and receives a reward for that action. For example, in news article recommendation, the system recommends a news article (action) based on the user’s interests (context), and receives a reward if the user clicks on the article. Since such systems are updated frequently, a key challenge is evaluating the performance of each update. The most common approach is to perform an A/B test, that is, to deploy both the updated system and the baseline in an online experiment. However, A/B testing can be costly in practice, as the updated system can lead to poor user experience or revenue loss. Moreover, A/B tests are often slow, with final evaluations taking several weeks or even months, thereby delaying system updates and progress.

In the off-policy evaluation (OPE) problem, the goal is to evaluate the performance of an algorithm or a policy using historical data. In applications such as recommendation systems and digital marketing, we often have logs of historical data, and OPE can be used to evaluate the performance of a target policy using the logged data directly (Dudík et al., 2011; Lee et al., 2022). This offline evaluation approach alleviates the need for A/B testing, enabling faster progress and improving user experience.

Most OPE estimators are based on Inverse Propensity Score (IPS) weighting (Horvitz and Thompson, 1952; Dudík et al., 2014), often called importance sampling (Sutton et al., 1998). IPS corrects the sampling distribution to the target distribution by applying IPS weights. It is unbiased under certain conditions, and has been proven to be optimal for contextual bandits (Wang et al., 2017). Despite such desired theoretical properties, estimators based on IPS suffer from high variance in practice, especially when there is a large mismatch between the target policy and the data collection policy, or when the number of actions is large.

To address the high variance problem of IPS, Saito and Joachims (2022) recently proposed the Marginalized IPS (MIPS) estimator. MIPS applies IPS weighting in an embedding space. Typically, the embedding space has a smaller dimension compared to the action space, which helps reduce the variance of IPS. Saito and Joachims (2022) assume that the embeddings are available to the estimator, and theoretically characterize how the embeddings can be used for bias-variance trade-off. For example, reducing the embedding dimension can reduce variance at the cost of increasing bias.

Cief et al. (2024) extend the work of Saito and Joachims (2022) by learning action embeddings using a linear reward model with a reward prediction loss and using the last layer of the model as embeddings. But their method learns action embeddings that are not conditioned on contexts. As a result, the action embeddings for different contexts can be similar even when the corresponding rewards are very different. For example, in clinical trials, a particular dose of adrenaline (action) can be life-saving or lethal depending on the patient’s condition (context) (Nordseth et al., 2012). Learning only action embeddings will learn the same embedding for an adrenaline dosage regardless of the patient’s condition. Moreover, it is unclear whether only minimizing the reward prediction loss is sufficient to learn good embeddings. In this work, we address these issues by proposing a novel objective to learn embeddings that are conditioned on both actions and contexts, and which directly minimizes the mean squared error (MSE) of MIPS. Kiyohara et al. (2024) also proposes a similar embedding learning objective, but for the slate bandit model, where each slate contains multiple actions. While their objective intuitively minimizes bias and variance of the estimator, it is unclear if the learned embeddings actually do so. In contrast, our objective directly estimates bias and variance, yielding MSE-minimizing embeddings.

Several OPE estimators use embeddings or similar notions to reduce variance at the cost of a small bias. Sachdeva et al. (2024) perform policy convolution via a similarity function between action embeddings with the implicit assumption that similar embeddings lead to similar rewards. Olivares et al. (2025) directly partition the context space into clusters and applies IPS over clusters instead of contexts, based on the assumption that the context does not affect the reward given the action and the cluster. Our work is different since our goal is to learn good embeddings for OPE, and our proposed embedding learning method can be potentially extended to these estimators.

In this paper, we propose learning context-action embeddings for the MIPS estimator by minimizing its mean squared error (MSE). In our theoretical studies, we first derive upper bounds on the bias and variance of MIPS, and propose an embedding learning objective using the upper bounds. Our proposed objective learns embeddings by balancing the bias and variance of MIPS so that the resulting estimator has a smaller MSE compared to the previous work of Cief et al. (2024). We also show information-theoretic connections to our objective, and how it grows as the number of actions increases. In our empirical studies, we conduct experiments on a synthetic dataset and a real-world bandit dataset collected from a fashion e-commerce platform (Saito et al., 2020). In both sets of experiments, we demonstrate that our method further reduces the MSE of the estimator proposed in Cief et al. (2024), which only considers reward prediction for embedding learning.

2. Preliminaries

In this section, we introduce the problem setting for OPE in contextual bandits, and MIPS estimator upon which we build our proposed method.

Notation.

Random variables are denoted with uppercase letters in standard font, such as $X,Y,Z$, etc. Expected value $\mathbb{E}$ is over the data-generating distribution unless specified otherwise. The notation $[k]$ denotes the set $\{1,\dots,k\}$.

2.1. Off-policy Evaluation in Contextual Bandits

Let $\mathcal{X}\subseteq\mathbb{R}^{d}$ be the context space, $\mathcal{A}$ be a finite set of actions, and $\mathcal{R}$ be the space of rewards. Data are collected offline by following a behavior policy $\mu\colon\mathcal{X}\to\Delta(\mathcal{A})$ in the following way. First, a context $X\sim p_{X}(\cdot)$ is sampled where $p_{X}$ is the unknown context distribution. Then an action $A\sim\mu(\cdot|X)$ is sampled. Finally, a reward $R\sim p_{R}{(\cdot|X,A)}$ is received, where $p_{R}$ is the unknown reward distribution given context $X$ and action $A$. The reward $R$ is the bandit feedback as it only reveals the value of the action taken in a context. We collect such iid samples in a dataset $\mathcal{D}=((X_{i},A_{i},R_{i}))_{i=1}^{n}$ of size $n$.

For a policy $\pi\colon\mathcal{X}\to\Delta(\mathcal{A})$, let

[TABLE]

Here $q\colon\mathcal{X}\times\mathcal{A}\to\mathbb{R}$ is the mean reward function $q(x,a)\coloneq\int rp_{R}(r|x,a)dr$.

The goal in OPE is to estimate $v(\pi)$ for a target policy $\pi$ using the offline dataset $\mathcal{D}$. The performance of an estimator $\hat{v}(\pi)$ is measured using the mean squared error or the MSE:

[TABLE]

where $\text{Bias}[\hat{v}(\pi)]$ and $\operatorname{\mathbb{V}}[\hat{v}(\pi)]$ denote the bias and variance of the estimator $\hat{v}(\pi)$ respectively.

2.2. Marginalized IPS Estimator

Perhaps the most widely used estimator for OPE is the Inverse Propensity Score (IPS) weighting, sometimes referred as importance sampling. Given a behavior policy $\mu$, a target policy $\pi$, and a dataset $\mathcal{D}=((X_{i},A_{i},R_{i}))_{i=1}^{n}$, the IPS estimator estimates $v(\pi)$ as

[TABLE]

where $w(x,a)$ is the IPS weight for $(x,a)\in\mathcal{X}\times\mathcal{A}$. Under the following coverage assumption on the action space, the IPS estimator is unbiased, i.e., $\mathbb{E}[\hat{v}_{\text{IPS}}(\pi)]=v(\pi)$.

Assumption 1** (Common support).**

$\pi(a|x)>0\implies\mu(a|x)>0$* for all $x\in\mathcal{X}$ and $a\in\mathcal{A}$.*

In practice ˜1 is often violated, making IPS biased. Under such conditions, the bias of IPS is given by (Saito and Joachims, 2022; Sachdeva et al., 2024):

[TABLE]

where $\mathcal{U}(x,\mu)=\{a\in\mathcal{A}:\mu(a|x)=0\}$ is the set of unsupported actions in context $x$ under policy $\mu$. A bigger issue is the variance of IPS (Dudík et al., 2014; Saito and Joachims, 2022), given by

[TABLE]

where $\mathbb{E}_{\mu}$ and $\operatorname{\mathbb{V}}_{\mu}$ are expectation and variance respectively under the behavior policy $\mu$. The above equation shows that the variance of IPS is high when the IPS weights $w$ are large or when $\operatorname{\mathbb{V}}[R|X,A]$ is large, i.e., the rewards are noisy.

To alleviate the issues of IPS, Saito and Joachims (2022) introduced the Marginalized IPS (MIPS) estimator. The MIPS estimator assumes access to embeddings from some set $\mathcal{E}\subseteq\mathbb{R}^{d_{\mathcal{E}}}$. The embeddings are sampled from an unknown embedding distribution $p_{E}{(e|x,a)}$. The rewards are assumed to be sampled from $p_{R}{(r|x,a,e)}$. Now the value of a policy is defined as

[TABLE]

where $q(x,a,e)\coloneq\int rp_{R}(r|x,a,e)dr$ is the reward function or expected reward given context $x$, action $a$, embedding $e$. The marginal embedding distribution induced by a policy $\pi$ is defined as $p_{E}(e|x,\pi)\coloneq\sum_{a\in\mathcal{A}}p_{E}(e|x,a)\pi(a|x)$. The goal is still to estimate $v(\pi)$ using samples. Given dataset $\mathcal{D}\coloneq((X_{i},A_{i},E_{i},R_{i}))_{i=1}^{n}$, the MIPS estimator (Saito and Joachims, 2022) is defined as

[TABLE]

where $\rho(X,E)$ is the marginal IPS weight for a pair of context and embedding $(X,E)$ in the offline data, and with some overload of notation, $\mu(a|X,E)$ is the posterior distribution of action $a\in\mathcal{A}$ given $(X,E)$ in the offline data (Saito and Joachims, 2022). In practice, the estimator $\hat{v}_{\text{MIPS}}(\pi)$ is computed following Equation˜2, where the posterior $\mu(a|X,E)$ is estimated via a logistic regression model (Saito et al., 2020). MIPS requires a common support assumption on embeddings similar to ˜1 for having finite marginal IPS weights.

Assumption 2** (Common Embedding Support, Assumption 3.1 in (Saito and Joachims, 2022)).**

Given policies $\pi$ and $\mu$, $p(e|x,\pi)>0\implies p(e|x,\mu)>0$ for all $e\in\mathcal{E}$ and $x\in\mathcal{X}$, where $p_{E}(e|x,\pi)$ is the marginal embedding distribution given policy $\pi$ and $x\in\mathcal{X}$ defined as $p(e|x,\pi)\coloneq\sum_{a\in\mathcal{A}}p_{E}(e|x,a)\pi(a|x)$.

Note that the embedding space is generally smaller than the action space to reduce the variance of the MIPS estimator. Therefore, this assumption is weaker than the common support assumption (˜1). Under the common embedding support assumption (˜2), Saito et al. (2020) showed that the bias of MIPS can be decomposed as

[TABLE]

The MIPS estimator is unbiased under the following assumption.

Assumption 3** (No Direct Effect, Assumption 3.2 in (Saito and Joachims, 2022)).**

Action $A$ is said to have no direct effect on reward $R$ if they are conditionally independent given the context $X$ and the embedding $E$: $A\perp R\,|\,X,E$.

If, in addition to ˜2, ˜3 also holds, then the bias in Equation˜3 becomes zero.” due to $q(X,a,E)=q(X,b,E)$. Saito and Joachims (2022) also showed that under Assumptions 1, 2, and 3, the variance reduction of MIPS compared to IPS is given by

[TABLE]

where $p_{R}(\cdot|X,E)$ is the reward distribution that depends only on $X$ and $E$ and not on $A$.

From Equation˜3, note that the bias is small when the posterior probabilities $\mu(a|x,e)$ are close one or zero. However, in such cases, Equation˜4 shows that the variance reduction is small, or the variance of MIPS estimator increases closer to the variance of IPS estimator, since the term $\operatorname{\mathbb{V}}_{\mu(\cdot\,|\,X,E)}[w(X,A)]$ is small due to reduced stochasticity of the posterior probability.

However, the MIPS estimator does not have control over its bias and variance as it assumes that the embeddings are given. In real-world applications, embeddings are not given in many cases. Even if the embeddings are provided, they might not be optimal embeddings that minimize the MSE of the MIPS estimator, as they were not designed to do so.

To learn embeddings based on offline data, Cief et al. (2024) proposed an embedding learning method for MIPS that learns action embeddings by minimizing a reward prediction loss. However, embeddings that minimize the reward prediction loss do not necessarily minimize the MSE of the MIPS estimator. Specifically, their embedding learning objective does not account for how concentrated or spread out the posterior $\mu(a\,|\,x,e)$ is, where the posterior controls the bias and variance of MIPS estimator as shown in Equations˜3 and 4. Moreover, the work of Cief et al. (2024) only learns action embeddings, while the embeddings assigned to actions should be different depending on the context, as an action may have a different reward in different contexts.

In the following section, we introduce our proposed method Context-Action Embedding Learning for MIPS (CAEL-MIPS) that learns MSE-minimizing embeddings that depend on both contexts and actions.

3. Context-Action Embedding Learning

CAEL-MIPS learns context-action embeddings by minimizing bias and variance of MIPS. To this end, we derive upper bounds for bias and variance and propose an embedding learning objective that combines the reward prediction error from Cief et al. (2024) with these upper bounds. The proofs for all the results are deferred to Appendix A.

3.1. Embedding Learning through MIPS MSE Minimization

As the bias and variance of MIPS cannot be directly evaluated from offline data, we derive their upper bounds.

3.1.1. Upper bound of MIPS Bias

The bias of MIPS (Equation˜3) cannot be evaluated using offline data as we cannot estimate the mean reward $q(x,a,e)=\int rp_{R}(r|x,a,e)$ since we do not have reward sampled from $p_{R}(r|x,a,e)$. Therefore, we derive an upper bound of bias that can be evaluated with offline data under the following mild assumption.

Assumption 4**.**

The mean and the variance of the rewards are bounded:

[TABLE]

for all $(x,a,e)\in\mathcal{X}\times\mathcal{A}\times\mathcal{E}$.

Given that the mean and the variance are bounded in most applications, e.g. click through rate of a webpage is between 0 and 1, ˜4 is generally satisfied in real-world scenarios. The following theorem establishes an upper bound on the bias of the MIPS estimator.

Proposition 0**.**

Under Assumptions 2 and 4, the upper bound of MIPS estimator’s bias is given by

[TABLE]

To evaluate the upper bound in Proposition˜1 with offline data, we use a function approximator $f_{\theta}$ to obtain context-action embeddings $E\coloneq f_{\theta}(X,A)$. Using the embeddings, we estimate the posterior $\mu(\cdot\,|\,X,f_{\theta}(X,A))$ as in the MIPS estimator via logistic regression. Finally, we use the estimated posterior and the IPS weights $w(\cdot\,,\cdot)$ to evaluate the upper bound in Proposition˜1. Since the exact bounds on the mean rewards, $q_{\max}$ and $q_{\min}$, are typically unknown, we introduce a hyperparameter $\alpha$ in our final objective (see Equation˜7). To learn bias-minimizing embeddings from offline data, we formulate the bias upper bound as an objective $\mathcal{L}_{\text{bias}}(\theta)$ for CAEL-MIPS:

[TABLE]

From Section˜3.1.1, we can see that the bias loss term approaches zero as the stochasticity of the posterior $\mu(a\,|\,X,f_{\theta}(X,A))$ is reduced ($\mu(a\,|\,X,f_{\theta}(X,A))\to 0$ or 1). The embedding learning function $f_{\theta}(X,A)$ that minimizes this loss would induce the posterior probability closer to 1 or 0 when the difference of IPS weights $|w(X,b)-w(X,a)|$ is large due to large discrepancy between $\pi$ and $\mu$ on context $X$ sampled in the offline data.

3.1.2. Upper bound of MIPS Variance

The variance reduction of MIPS compared to IPS is given by Equation˜4. In order to evaluate the variance of MIPS using offline data without calculating the variance of the IPS weights, we establish an upper bound on the variance of MIPS in the next proposition.

Proposition 0**.**

When Assumptions 1, 2, 3, and 4 hold, then

[TABLE]

Similar to Proposition˜1, we evaluate the upper bound in Proposition˜2 by parametrizing $E=f_{\theta}(X,A)$ and estimating the posterior $\mu(\cdot\,|\,X,f_{\theta}(X,A))$ using offline data. Using $E$ and $X$, we use a reward prediction model to get $E[R|X,E]^{2}$. Finally, with the estimated posterior, the IPS weights, and $E[R|X,E]^{2}$, we can evaluate the upper bound in Proposition˜2. Furthermore, since the first term $\operatorname{\mathbb{V}}[\hat{v}_{\text{IPS}}(\pi)]$ in the upper bound in Proposition˜2 does not depend on the embeddings, we can ignore it while minimizing the variance upper bound to learn embeddings. We introduce another hyperparameter $\beta$ in the final objective (see Equation˜7) since $\sigma_{\min}^{2}$ is not typically known in advance. We define the objective based on variance as

[TABLE]

Note that $\mathcal{L}_{\text{var}}(\theta)\to 0$ when $\sum_{a}\mu(a\,|\,X,f_{\theta}(X,A))^{2}$ is minimized. This sum of squared probabilities is minimized when the posterior $\mu(\cdot\,|\,X,f_{\theta}(X,A))$ is a uniform distribution over the actions (see Proposition˜4). Learning the embedding function $f_{\theta}$ to minimize the upper bound of variance $\mathcal{L}_{\text{var}}(\theta)$ (Section˜3.1.2) induces the posterior to be more uniform on the context $X$ where $\sum_{a\in\mathcal{A}}w(a,X)^{2}$ is high. In other words, when there is high variance due to a large discrepancy between the behavior and target policies (Saito and Joachims, 2022; Dudík et al., 2014), the embedding function is learned to smooth out the induced posterior distribution. Also note that $\mathcal{L}_{\text{var}}(\theta)\to 0$ as $n\to\infty$. This is expected as the variance of an estimator goes to zero as the sample size goes to infinity. So minimizing Section˜3.1.2 is more useful for small sample sizes.

3.1.3. Embedding Learning Objective

Propositions˜1 and 2 show the upper bounds of bias and variance of MIPS. On one hand, by learning the embeddings to minimize the upper bound of bias, $\mu(\cdot|X,f_{\theta}(X,A))$ induced from the learned embeddings concentrates around some action as discussed earlier. On the other hand, minimizing upper bound of variance pushes the posterior distribution towards a uniform distribution. By balancing this trade-off, we propose an embedding learning objective that combines the upper bounds of bias and variance, as well as the reward prediction loss introduced by Cief et al. (2024). The reward prediction loss is used since we would like to learn embeddings that capture most of the reward information from contexts and actions to further reduce bias. It is also empirically shown to have good performance in previous works (Cief et al., 2024; Kiyohara et al., 2024). Our final objective function is

[TABLE]

where $\mathcal{L}_{R}(\theta)$ is the MSE of reward prediction and $\alpha$ and $\beta$ are positive coefficients. The first term $\mathcal{L}_{R}(\theta)$ in Equation˜7 aims to reduce bias since predicting the rewards well based on embeddings controls bias as seen in Equation˜3 and Proposition˜1. The second and the third term aim to reduce bias and variance respectively, as governed by the posterior distribution $\mu$ and discussed in Propositions˜1 and 2. The hyperparameters $\alpha$ and $\beta$ control the contribution of $\mathcal{L}_{\text{bias}}(\theta)$ and $\mathcal{L}_{\text{var}}(\theta)$ to the total loss. Larger values of $\alpha$ and $\beta$ induce larger bias and variance reduction respectively. We estimate Equation˜7 using offline data $\mathcal{D}$ by replacing expectations with empirical averages to get $\hat{\mathcal{L}}(\theta;\mathcal{D})\coloneq\hat{\mathcal{L}}_{R}(\theta;\mathcal{D})+\alpha\hat{\mathcal{L}}_{\text{bias}}(\theta;\mathcal{D})+\beta\hat{\mathcal{L}}_{\text{var}}(\theta;\mathcal{D})$. Recall that $f_{\theta}(x,a)$ outputs the embedding of $(x,a)\in\mathcal{X}\times\mathcal{A}$, which is used for reward prediction $\hat{R}$. The reward prediction loss, the bias loss, and the variance loss that consist the embedding learning loss $\hat{\mathcal{L}}(\theta;\mathcal{D})$ are:

[TABLE]

[TABLE]

and

[TABLE]

3.1.4. Discussion on the Bias and Variance Upper Bounds

We give some intuitions on the bias and variance upper bounds via the entropy of the posterior, and show their growth with respect to the size of the action space. Let

[TABLE]

be the Rényi entropy of order 2 of the posterior $\mu(a|X,f_{\theta}(X,A))$. The next proposition relates the bias upper bound with the entropy of the posterior and the size of the action set.

Proposition 0**.**

Let $C>0$ and $|w(x,a)-w(x,b)|<C$ for all $x\in\mathcal{X}$, and for all $a,b\in\mathcal{A}$. Then

[TABLE]

where $H_{2}(X,a;\theta)$ is the Rényi entropy of order 2 of the posterior $\mu(\cdot|X,f_{\theta}(X,A))$. Further, $\mathcal{L}_{\text{bias}}(\theta)\leq(C^{2}/4)\mathbb{E}[H_{2}(X,A;\theta)]$.

Proposition˜3 states that minimizing the entropy $H_{2}$ corresponds to minimizing the bias upper bound. The entropy is minimized when $\mu(\cdot\,|\,X,f_{\theta}(X,A))$ is concentrated around some action. This means that when the context $X$ and the learned embedding $f_{\theta}(X,A)$ can fully recover an action, then the bias upper bound is zero. This verifies our earlier intuition about $\mathcal{L}_{\text{bias}}(\theta)$. From Proposition˜3, we also observe that as the number of actions increases, the bias upper bound increases and is further upper bounded by $(1-1/|\mathcal{A}|)$. Later in Section˜4.1, we validate this observation empirically. Next, we show that minimizing the variance upper bound is equivalent to maximizing $H_{2}$ and show the convergence rate of the variance upper bound with respect to the dataset size and the number of actions.

Proposition 0**.**

Let $H_{2}(X,A;\theta)$ be the Rényi entropy of order 2. Then

[TABLE]

Further, we have that $\mathcal{L}_{\text{var}}(\theta)\leq\frac{1}{n}\sum_{a}\mathbb{E}[\mathbb{E}[R|X,E]^{2}w(X,a)^{2}]=O(|\mathcal{A}|/n)$ for bounded IPS weights.

From Proposition˜4, we observe that minimizing the variance loss is equivalent to maximizing $H_{2}$, weighted by the sum of squared IPS weights. The entropy $H_{2}$ is maximized when the posterior $\mu(\cdot\,|\,X,f_{\theta}(X,A))$ is uniform implying that variance minimization encourages embeddings that yield a uniform posterior over actions. The weighting emphasizes regions where the squared IPS weights are large, that is, where the overlap between the behavior and target policies is small. Additionally, we observe that the variance loss increases linearly with the number of actions, which we verify empirically in Section˜4.1.

3.2. Algorithm

The pseudocode for CAEL-MIPS is presented in Algorithm˜1. CAEL-MIPS uses a three-layer feedforward neural network $f_{\theta}$ to parameterize the embeddings and the rewards. The hidden layers use batchnorm, ReLU activations, and a dropout rate of 0.2. The network takes $(x,a)\in\mathcal{X}\times\mathcal{A}$ as input and outputs the embedding $\hat{e}$ for the context $x$ and action $a$, that is, $\hat{e}\coloneq f_{\theta}(x,a)$. Following Cief et al. (2024), the network also outputs a reward prediction using the embeddings given by $\hat{r}\coloneq\hat{e}^{\top}x$.

To compute $\hat{v}_{\text{MIPS}}(\pi)$, Algorithm˜1 first learns the context-action embeddings $(\hat{E}_{i})_{i=1}^{n}$ by performing gradient updates on the loss defined in Equation˜7 (lines 3-8). After learning the embeddings, Algorithm˜1 estimates $\hat{\mu}_{\phi}$ using supervised learning on $((\hat{E}_{i},A_{i}))_{i=1}^{n}$ (line 10). We use logistic regression from scikit-learn (Pedregosa et al., 2011) to estimate $\hat{\mu}_{\phi}$ following Saito et al. (2020) and Cief et al. (2024). Finally, on line 11, Algorithm˜1 computes the MIPS estimator using $\hat{\mu}_{\phi}$ as in Equation˜2.

4. Experiments

For experiments, CAEL-MIPS is tested on both a synthetic dataset and a real-world dataset. We prepared the synthetic dataset to test OPE methods on various settings and to empirically validate the theoretical findings in Section˜3. We also test CAEL-MIPS on a real-world bandit dataset (Saito and Joachims, 2022) to demonstrate the applicability of CAEL-MIPS in practical scenarios and to show that CAEL-MIPS outperforms baselines in most experimental runs. We use $\alpha=10$ and $\beta=0.1$ for CAEL-MIPS in all the experiments. The baselines used in our experiments include AEL-MIPS (Cief et al., 2024), IPS, and DM. AEL-MIPS is included as a baseline, as it learns embeddings for MIPS, similar to CAEL-MIPS, but only learns embeddings for actions based on the reward prediction loss (Equation˜8). We also compare our proposed model with DM, as both CAEL-MIPS and AEL-MIPS utilize the reward prediction loss for learning embeddings. For DM, we estimate the policy value using the approximated expected reward learned through the reward prediction loss:

[TABLE]

where $\hat{q}$ is the learned reward model. Lastly, we compare with IPS to see if MIPS methods with embedding learning reduce the MSE of the OPE estimation.

4.1. Synthetic Data

We conduct a synthetic data experiment to assess the capacity of CAEL-MIPS on reducing the MSE of MIPS under diverse experimental conditions. The subsequent description outlines the sampling distributions used for rewards, actions, and contexts in the synthetic data, along with a description of the target policy used in the experiment.

Rewards

The reward distribution for the synthetic data is $P_{R}=\mathcal{N}(q(x,a),1)$, where $q(x,a)=10e^{-(x_{1}-a_{1})^{2}}$, and $x_{1}$ and $a_{1}$ are the first components of the context and action vectors respectively. Since AEL-MIPS employs a linear reward model for embedding learning, we construct the reward function $q(x,a)$ to be a nonlinear function of context–action pairs to examine whether CAEL-MIPS, which does not utilize a linear reward function, can achieve lower MSE than AEL-MIPS.

Actions

The action space consists of $k=500$ actions. Each action $a$ has a vector representation where its first component is $a/k$, and the remaining $k-1$ components are sampled uniformly at random from $[0,1]$. The actions in the dataset were sampled from a uniform random behavior policy unless specifed otherwise.

Contexts

The contexts in the synthetic data were sampled uniformly from the context space $\mathcal{X}=[0,1]^{5}$.

Target Policy

The target policy $\pi$ is $\epsilon$-greedy with respect to $q(x,a)$ with $\epsilon=0.2$.

Using the synthetic data, we first compare CAEL-MIPS and baselines in terms of bias, variance, and MSE (Figure˜1). CAEL-MIPS achieves the lowest MSE among the OPE estimators. CAEL-MIPS has the bias of OPE estimation significantly reduced while inducing a small increase in the variance compared to AEL-MIPS. This is due to the additional MSE-minimizing loss terms in the learning objective of CAEL-MIPS in addition to the reward prediction loss used in AEL-MIPS. The use of a nonlinear reward model further contributes to the bias reduction, as AEL-MIPS assumes a linear reward model for embedding learning. Lastly, the context-action dependent embeddings of CAEL-MIPS also contribute to a reduction in bias compared to AEL-MIPS. Both CAEL-MIPS and AEL-MIPS show significantly reduced MSE compared to IPS due to substantially reduced variance through learned embeddings. DM, which uses a reward function learned with the reward prediction loss for OPE estimation, exhibits the highest bias and the lowest variance among OPE methods. The bias arises from the reward model overfitting to uniformly sampled, uninformative action components, while the low variance results from the deterministic mapping between predicted rewards and context–action pairs.

Next, we evaluate the performance of CAEL-MIPS and baselines for various experimental conditions and see if the performance ordering observed in Figure˜1 persists (Figure˜3). In Figures˜3(a), 3(b) and 3(c), we show MSE of OPE estimators, as well as bias and variance that compose the MSE on a varying dataset size. The figures were drawn with 100 repeated trials for each dataset size. In Figure˜3(a), we can observe that the performance ordering of OPE methods that was observed in Figure˜1 persists on varying dataset size. CAEL-MIPS outperforms other OPE estimators in terms of MSE. CAEL-MIPS shows significantly reduced bias while inducing a small variance compared to AEL-MIPS. However, as the dataset size becomes larger, CAEL-MIPS reduces both bias and variance due to the MSE-minimizing loss terms in its embedding learning objective and the variance of CAEL-MIPS becomes similar to the variance of AEL-MIPS. Both IPS and DM shows large MSE due to large variance and bias respectively as in Figure˜1.

Figures˜3(d), 3(e) and 3(f) show the MSE, bias, and variance of OPE estimators with respect to the number of actions. We test OPE estimators on increasing variance due to increaseing number of actions (Sachdeva et al., 2024) and see if the performance ordering observed in Figure˜1 persists. The figures were drawn with 30 repeated trials for each number of actions (50, 100, 1000, 1500). In Figure˜3(f), IPS, AEL-MIPS, and CAEL-MIPS show increase in variance as the number of actions increases while DM is not significantly affected by the increase in the number of actions. For IPS that has variance as a dominant factor in its MSE, the increase in the variance is reflected heavily on its MSE. In Figure˜3(e) we observe that CAEL-MIPS has the lowest bias among all the estimators due to the MSE-minimization loss terms. The bias grows with the number of actions, and our empirical results match the rate given by Proposition˜3. Figure˜3(f) shows how the variance increases as we increase the number of actions.

Figure˜3(g) shows MSE of OPE esetimators on varying behavior policy $\mu(a|x)$. The figure was drawn with 30 repeated trials for each value of $\gamma\in\{-1.0,-0.5,0.0,0.5,1.0\}$ that controls the softmax behavior policy $\mu(a|x)$:

[TABLE]

When $\gamma=0$, the policy is a uniform random distribution, while large values of $\gamma$ makes the policy tend towards greedy, and negative $\gamma$ values makes the policy suboptimal. We observe that for negative $\gamma$ values, the MSE is large as the mismatch between the target policy and the behavior policy is large. We also observe that the MSE of CAEL-MIPS is always smaller than that of AEL-MIPS and the performance ordering observed in Figure˜1 generally persists in this setting.

Figure˜3(h) shows the MSE as a function of $\epsilon$ which controls the randomness of the $\epsilon$-greedy target policy with uniform random behavior policy. The figure was drawn with 30 repeated trials for each value of $\epsilon$. When $\epsilon=0$, the target policy is greedy, and is the optimal policy. Since the behavior policy is uniform random, as the target policy changes from greedy to uniform random by increasing $\epsilon$ from zero to one, target policy approaches uniform random behavior policy and MSE approaches zero for all OPE estimators. However, the performance ordering of the OPE estimators observed in Figure˜1 is still observed in this experiment setting and CAEL-MIPS outperforms the baseline OPE estimators across varying $\epsilon$.

Figure˜3(i) shows the MSE on varying reward distribution. The figure was drawn with 30 repeated trials for each value of standard deviation that controls the reward distribution. Similar to the other empirical study made on varying experiment environment parameters, we can see that the performance ordering observed in Figure˜1 is still observed across varying standard deviation of the reward distribution.

4.2. Real-World Data

In this section, we test CAEL-MIPS and baselines on a complex real-world dataset to assess the scalability of our method in practical use cases. We test the OPE estimators on Open Bandit Dataset, a real-world bandit dataset collected on a fashion e-commerce platform ZOZOTOWN (Saito et al., 2020). The dataset was collected from the A/B test of two policies: a uniformly random behavior policy and a Thompson sampling target policy. We use $n=10000$ samples. The context dimension is $d=20$, the total number of actions is $|\mathcal{A}|=240$, and the reward is binary, which indicates whether the user clicked the recommended item or not. To learn the embeddings, we use the same three-layer feedforward neural network as discussed in Section˜3.2. We use one-hot action representations as input to the neural network. For DM, we use the same network as CAEL-MIPS, and use it as a reward model to estimate $v(\pi)$ using Equation˜11. We repeat the experiments 30 times to estimate the squared errors of all the estimators.

We plot the empirical cumulative distribution function (CDF) of the squared errors relative to IPS in Figure˜2 following the work of Saito and Joachims (2022) and Cief et al. (2024). To compute the squared errors relative to IPS, the squared errors of all the estimators are normalized by dividing them by the squared error of IPS. Then empirical CDF of the relative squared errors is computed. From Figure˜2, we observe that CAEL-MIPS outperforms IPS in more than 75% of the runs. Furthermore, while the relative errors of AEL-MIPS and DM never reach below 0.1 and 0.01, respectively, CAEL-MIPS’s relative error reaches below 0.1 for about 40% of the runs and below 0.01 for about 20% of the runs.

5. Discussion

In this work, we proposed CAEL-MIPS, a MIPS-based estimator that learns context-action embeddings to perform OPE in contextual bandits. We derived upper bounds for bias and variance of MIPS, and proposed an embedding learning objective function using those bounds and the reward prediction loss. By minimizing the MSE upperbound in the embedding learning objective, the learned embeddings directly reduce the MSE of MIPS estimator. In the theoretical analysis, we showed that our embedding learning objective function is related to the entropy of the posterior distribution $\mu(a|x,e)$. In our experiments, we empirically validated that CAEL-MIPS outperforms the baselines on various experiment settings on synthetic dataset due to its context-action dependent embeddings learned through the learning objective that directly minimizes MSE. We further demonstrated the scalability of CAEL-MIPS to practical use cases by testing it on a real-world bandit dataset, where it consistently outperformed the other baselines.

Our work opens several directions for future research. Due to the unknown degrees of tightness of the upper bounds on bias and variance of MIPS, our method introduced constant weights for bias and variance upperbounds used in the embedding learning objective, which may be difficult to tune. A potential future work is to choose these hyperparameters using methods such as SLOPE (Su et al., 2020; Tucker and Lee, 2021). Furthermore, computing the bias upper bound in Algorithm˜1 requires $O(|\mathcal{A}|^{2})$ operations. Reducing this computation would make our method more scalable to large action spaces (Lopez et al., 2021). In addition, Algorithm˜1 relies on multiple calls to a logistic regression subroutine to estimate the posterior $\hat{\mu}(\cdot|x,e)$. A potential future work would be to use a generative network to model the posterior. Finally, extending our method to policy optimization and the full reinforcement learning setting are also promising directions.

Appendix A Proofs

Proof of Proposition˜3.

We first show the upper bound on $\mathcal{L}_{\text{bias}}(\theta)$, that is, $\mathcal{L}_{\text{bias}}(\theta)\leq(C^{2}/4)(1-\mathbb{E}[e^{H_{2}(X,A;\theta)}])^{2}$. First recall from Section˜3.1.1 that

[TABLE]

Since $|w(x,a)-w(x,b)|<C$ for all $(x,a,b)\in\mathcal{X}\times\mathcal{A}^{2}$, we have

[TABLE]

It remains to bound $\sum_{a<b}\mu(i|X,f_{\theta}(X,A))\mu(j|X,f_{\theta}(X,A))$. Recall the identity $2\sum_{a<b}a_{i}a_{j}=(\sum_{i}a_{i})^{2}-\sum_{i}a_{i}^{2}$. Applying it to the posterior $\mu(\cdot|x,f_{\theta}(x,a))$ for all $(x,a)\in\mathcal{X}\times\mathcal{A}$, we have

[TABLE]

which proves the result. The second inequality follows from the fact that the entropy is maximized for the uniform distribution of the actions. Now, from Taylor series expansion, we have $e^{-x}\geq 1-x$ for all $x\in\mathbb{R}$. Applying it to $e^{-H_{2}(X,A;\theta)}$, we obtain

[TABLE]

proving the last inequality. ∎

Proof of Proposition˜4.

Recall that

[TABLE]

The proposition follows by noting that from the definition of $H_{2}$, we have $\sum_{a}\mu(a|X,f_{\theta}(X,A))=e^{-H_{2}(X,A;\theta)}$. The upper bound can be obtained by noting that $H_{2}(X,A;\theta)\geq 0$, implying that $e^{-H_{2}(X,A;\theta)}\leq 1$, the boundedness of IPS weights, and the fact that $\mathbb{E}[R|X,E]^{2}$ does not depend on $a$. ∎

Proof of Proposition˜1.

Our goal is to bound $|\text{Bias}[\hat{v}_{\text{MIPS}}(\pi)]$. From Equation˜3, recall that the bias of MIPS is given by (Saito and Joachims, 2022):

[TABLE]

Now, we have that

[TABLE]

where we used Jensen’s inequality in $(\clubsuit)$, triangle inequality in $(\spadesuit)$, and ˜4 in $(\diamondsuit)$. ∎

Proof of Proposition˜2.

Here, we upper bound the variance of MIPS. Recall from Appendix˜A that

[TABLE]

where the last inequality follows from ˜4. We now upper bound the RHS of Appendix˜A as follows.

[TABLE]

where we use Cauchy-Schwarz in the final step. Combining Appendix˜A and Appendix˜A, we get

[TABLE]

∎