POTEC: Off-Policy Contextual Bandits for Large Action Spaces via Policy Decomposition
Authors: Yuta Saito
Jihan Yao
Thorsten Joachims
Venue: ICLR 2025 Spotlight (2025)
Keywords: Off-Policy Learning
Contextual Bandits
Large Action Space
Importance Weighting
Clustering
Policy Gradient
Primary area: reinforcement learning
TL;DR: We propose a new two-stage off-policy learning algorithm called POTEC, which unifies regression- and policy-based approaches via a policy decomposition and is particularly effective in large action spaces.

Abstract: We study off-policy learning (OPL) of contextual bandit policies in large discrete action spaces where existing methods -- most of which rely crucially on reward-regression models or importance-weighted policy gradients -- fail due to excessive bias or variance. To overcome these issues in OPL, we propose a novel two-stage algorithm, called Policy Optimization via Two-Stage Policy Decomposition (POTEC). It leverages clustering in the action space and learns two different policies via policy- and regression-based approaches, respectively. In particular, we derive a novel low-variance gradient estimator that enables to learn a first-stage policy for cluster selection efficiently via a policy-based approach. To select a specific action within the cluster sampled by the first-stage policy, POTEC uses a second-stage policy derived from a regression-based approach within each cluster. We show that a local correctness condition, which only requires that the regression model preserves the relative expected reward differences of the actions within each cluster, ensures that our policy-gradient estimator is unbiased and the second-stage policy is optimal. We also show that POTEC provides a strict generalization of policy- and regression-based approaches and their associated assumptions. Comprehensive experiments demonstrate that POTEC provides substantial improvements in OPL effectiveness particularly in large and structured action spaces.

Published as a conference paper at ICLR 2025
POTEC: OFF-POLICY CONTEXTUAL BANDITS FOR LARGE ACTION SPACES VIA POLICY DECOMPOSITION
Yuta Saito
Jihan Yao
Thorsten Joachims Cornell University
Cornell University ys552@cornel1.ee
University of Washington
j 3ihany2ccs. washingtoon. edu
(lBcs.cornell edu
ABSTRACT
We study off- off-policy learning (OPL) of contextual bandit policies in large discrete action spaces where existing methods most of which rely crucially on reward- regression models or importacce-weighted policy gradients fail due to excessive bias or variance. To overcome these issues in OPL, we propose novel two- stage algorithm, called Policy Optimization via Two-Stage Policy Decomposition (POTEC). It leverages clustering in the action space and learns two different policies via policy- and regression- based approaches, respectively. In particular, we derive novel low-variance gradient estimator that enables to learn first-stage policy for cluster selection efficiently via policy-based approach. To select a specific action within the cluster sampled by the first-stage policy, POTEC uses second-stage policy derived from a regression-basee approach within each cluster. We show that local correctness condition. which only requires that the regression model preserves the relative expected reward differences of the actions within each cluster, ensures that our seliy-trradiet estimator unbiased and the second-star stage policy is optimal. We also show that POTEC provides a strict generalization of policy- and regression- bbased approaches and their associated assumptions. Comprehensive experiments demonstrate that POTEC provides substantial improvements in OPL effectiveness particularly in large and structured action spaces.
INTRODUCTION
Many interactive systems (e.g., language models, recommender systems) are increasingly controlled by automated decision-making policies that learn from historical interactions. These interactions con- sist of the context (e.g., prompt, user profile), the action chosen by the logging policy (e.g., sentence, recommended product), and the resulting reward (e.g., click). Using such logged interactions, goal is to train a new policy that improves the expected reward. This off-policy learning (OPL) task is of great practical relevance, as it enables to improve system effectiveness without the risky, slow, and potentially unethical use of online exploration (Saito & Joachims 2021 Gao et al. 2022).
A highly effective approach OPL is policy learning by estimating the policy gradient, which has resulted in a number of practical OPL methods for small action spaces (Swaminathan & Joachims
2015a Joachims al., 2018 Su et al. [019 2020aa Metelli et al. 2021). Unfortunately, this policy-based approach can deteriorate dramatically for large action spaces, which are prevalent in many potential applications of OPL where there exist millions of items (e.g., recommendations of movies, songs, products). In particular, in such large-scale environments, existing policy- based methods, which are mostly based on importance-weighee
policy gradients, can collapse due to extremely large variance (Saito & Joachims 2022; Cief et al. 20244 Sachdevv et al. 2024). While regression- based approaches, which learn the expected reward function and choose the action with the highest predicted reward, could potentially circumvent the variance issuc, they are known to suffer from high bias due to model misspecification (Farajtabar et al. 2018 VVlooshi et al., 2019 Sachdeva al. 2020 Saito et al. 2021a) and thus do not provide readily available solution either.
To overcome this bias and variance dilemma of OPL arising particularly in large action spaces, we develop novel two-stage OPL algorithm called Policy Optimization via Sto-Stage Policy Decomposition (POTEC). POTEC operates under a novel policy decomposition framework, wherein the typical overall policy (marginal action distribution) is decomposed into first-stage and second-

---

Published as a conference paper at ICLR 2025
stage policies via an action cluster space. The first-stage policy identifies promising action clusters (cluster distribution), while the second-stage policy selects the optimal action within a specific cluster sampled from the first-stage policy (conditional action distribution). A key feature of POTEC is its distinct learning approaches for the policy at each stage. The first-stage policy is learned using a policy-basee approach with novel policy gradient estimator, called the POTEC gradient estimator. It relies on importance weighting in the action cluster space to estimate the value of clusters while using a pairwise reward model to deal with the effect of individual actions within each cluster. We show that our gradient estimator is unbiased under local correctness (Saito et al. (023), requiring only that the regression model accurately preserves the relative reward differences within each action cluster. We also show that we can leverage the regression model used in the POTEC gradient estimator to construct a second-stage policy through a regression- based approach.
Compared to standard policy-based methods, the POTEC gradient estimator for the first-stage policy exhibits significantly lower variance large action spaces. This because it applies importance weighting to only the action cluster space. which is considerably more compact than the original action space, Furthermore, POTEC is more resilient to estimation bias than typical regression-ban approaches, since our first-stage policy is based policy only needs to learn the relative value differences between actions, which is less demanding than conventional absolute reward regression. Moreover, we show that POTEC and local correctness provide a full spectrum of OPI. approaches whose endpoints are policy- and regression- -based methods and their associated reward-modeling conditions. Experiments on synthetic and real-world recommendation datasets demonstrate that POTEC can provide substantially more effective OPL than conventional methods particularly when the per-action training data size small.
unbiased policy gradient and the second-stage
Related Work. There is a recent line of work on off-policy evaluation (OPE) for large action
spaces (Satto & Joachims 2022
Peng et al.
2023
Saito al., [023)] Sachdeva et al. 2024 Cief In particular, Saito & Joachims (2022) tackle
et al. [Quaa]Aouali et al.) 2024: Taufig et al.
2024
the problem of OPE for large action spaces and propose the marginalized IPS (MIPS) estimator that leverages pre-defined action embeddings. |Png et al. (2023), [aachdeaa et al. (2024), and Cief et al. (2024) propose methods to learn such a structure in the action space from logged data and further improve MIPS. Compared to these works, which focus solely on the OPE problem, we focus on the OPL problem in large action spaces, where OPE methods cannot be directly applied.
Among these, we particularly distinguish our contributions from those Saito al. (2023), which develop the OffCEM estimator to handle OPE in large action spaces. Although our algorithm is inspired by the reward function decomposition in OffCEM, its application to our concept of policy decomposition and the associated two-stage OPL. algorithm are novel. Moreover, we address the off- off-policy learning (OPL) problem and develop corresponding algorithm, whereas (Saito al.] (2023) focus solely on the off-policy evaluation problem in their analysis and experiments. Thus, our work is the first to formulate and propose methods specific to OPL for large discrete action spaces and offers several unique contributions from both methodological and empirical perspectives.
There also exist several related methods aimed at improving sample efficiency in large action spaces for reinforcement learning (RL). Chandak et al. (2019) propose method to learn action representations to enhance the sample efficiency of on-policy RI. However, the setting of Chandak et al. (2019) is not on offline policy learning, and thus their proposed method is not considered as bascline in our paper. Additionally, the supervised representation learning procedure in their work relies on RL-specific structures (i.e., state transitions), making inapplicable to our contextual bandit setup. Similarly, Gu et al. (2022) study offline RL in large action spaces and propose a method to learn latent representations in the action space. Their method leverages a data-distributional metric to learn action embeddings. This metric is based on the MDP structure, and how to apply it to the offline contextual bandit problem is not discussed, making it incomparable in our experiments.
Finally, Ban & He (202T) and Zhu et al. 2022) study online bandits in large action spaces where active exploration is allowed, whereas we focus on the offine setup with no additional exploration.
2 OFF-POLICY LEARNING FOR CONTEXTUAL BANDITS
We formulate OPL under the general contextual bandit process, where a decision maker repeatedly observes context X' drawn i.i.d. from unknown distribution p(x:). Given context æ,
2

---

Published as a conference paper at ICLR 2025
potentially stochastic policy x(a (z) chooses action from finite action space denoted as A. The reward r E [0, 7'max] is then sampled from some unknown conditional distribution p(r 2, a). We define the value of policy π as measure of its effectiveness:
V (T) EE (ptee)((llshprrt...a) [r] (p(a))(ala) (g( a)],
where we use q(z, a) == Eir x, al denote the expected reward function given and a.
Our goal is to learn a new policy Te, parameterized by 0, to maximize the policy value as
arg max (TT). dE@
The logged data we can use for performing OPL takes the form D == {(z,, a, ri)l= which contains n independent observations drawn from the logging policy zn-
Below, we describe two typical approaches to OPL, namely the policy-based and regression- based approaches, and summarize their limitations, particularly in large action spaces.
The policy-based approach learns the policy parameter via iterative gradient ascent as 8z+1 + 0g + (TV(T)). Since we not know the true gradient
VoV (mo) (p(=)= [g(2;, Vo log (o(a [
we need estimate it from the logged data. A common way do so is to apply importance weighting:
V Virs(mo: =
w(xi, (1)
where w(x,a) == (a x/) (o x)) is the vanilla importance weight defined with respect to the action space A and so [z, a) = Vo log (a |a| x) is the policy score function.
Eq. I is unbiased (i.e., E o (T D) = Vo VoV(To) under the following condition.
Condition 2.1. (Full Support) The logging policy To is said to have full support if ro(a x)) > 0, a(x,a)= X A.
For large action spaces, unfortunately, this requirement of full support is problematic for two reasons. First, violating the requirement can introduce substantial bias (Sachdeva et al., 2020 Felicion et al., 2022). Second, fulfilling the requirement for large action spaces leads to excessive variance, since To(a x) becomes extremely small (Saito & Joachims [022 Sachdeva et al. 2024). At first glance, doubly-robust (DR) estimation may appear helpful for dealing with the variance issue.
V To D) = [ulaa a:)(ri ((x;, 0:) (s)))s((i; as) + ((21,))))
(2)
DR uses reward function estimator q(x. a) while maintaining unbiasedness under Condition 2.1 and its variance is often lower than that of Eq. (T). However, unless the rewards are close to deterministic and the reward estimates ((x, a) are close perfect, its variance can still be extremely large due to importance weighting in the action space, which leads to inefficient OPL in large action spaces (Saito) & &Jocchims 2022. Peng et al.) 2023 (Sachdeva et al.) (023).. The issue of the IPS and DR policy gradients can be seen by calculating their variance (for a particular parameter E Rd) as
ntr (CovD (VOVOR D)]
Σ
((x (I.
+ Ep(x) "oolall)luu a a) a)s (x, a))] + + Vp(z) [mtaae) (z, a)]
(3)
where (2 a) == V/r x, a) and Ag a) = g(zz a) g(0, a). sp (I, a) is the j-th dimension of the score function. Note that the variance of IPS can be obtained by setting q(z a) 0. The variance reduction of DR comes from the second term where A a) is smaller than q2x a) if qzz a) is accurate. However, can also that the variance contributed by the first term be extremely
3

---

Published as a conference paper at ICLR 2025
Action Clustering
1st- Ast-Stage Policy 1et
Stage Policy a
action
cluster
C ~ T
(- x)
ne 2, c)
best action in this cluster
promising
cluster
Figure 1: The Two-Stage Off-Pollcc Learning Procedure of Our POTEC Algorithm, which first forms action clustering ca. and then identifies a promising cluster by the 1st-stage policy To lst and finally picks the best action in the cluster by the 2nd-stage policy "w
2nd
large, regardless of the accuracy of ((x, a). when the reward is noisy and the weights w(x a) become large, which occurs when Te assigns large probabilities to actions that are less likely under To-
The regression- approach employs an off- the-shelf supervised machine learning method, such neural networks, estimate the reward function q(z,, a), for example, by solving
= arg min
Σ
c(r. (o).)o(z))))
(u.0,r)ED
for some loss function é such as the squared loss. Then, it transforms the estimated reward function Go (2z, a) into a decision-making rule Ta(a zz). This approach avoids the use of importance weighting and is therefore relatively robust to high variance compared to the policy-based approach, even in large action spaces (Jeunen & Goethals] (021). However, it is widely acknowledged that this approach may fail significantly due to bias issues resulting from the difficulty in accurately estimating the expected reward for every action in A (Farajtabar et al. 2018 Voloshin et al. 2019)).
THE POTEC ALGORITHM
The following proposes new OPL algorithm, named POTEC, that eircumvents the challenges of pasicy-based and regression basid approaches for large action spaces. As depicted in Figure POTEC leverages the following novel decomposition of an overall policy.
overall T8
(a (z)
Σ GEC
T 1st
(c ())
2nd (a xx, c),
where the marginal action-selection (overall) policy (aggerall) is decomposed into the cluster- selection (first-stage) policy (ITT*) and conditional action-sclection (second-stage) policy (x² parametrized by 0 and vh, respectively. This policy decomposition defined via some pre-defined clustering structure in the action space, where ca C represents the cluster to which action belongs (typically |C] < |A). Action clusters are learnable, by applying clustering algorithms, such KMeans, in the action feature space to the averaged estimated rewards for each action, ((a) = (1/n) Σi= ((z,, a). as the embedding of a, which is a heuristic but performed effectively in our experiments. Note that although consider context-independer and deterministic action clusters in the main text, our framework can be extended to more general types of action clustering (i.e., context- and stochastic).
Leveraging this decomposition, POTEC (i) trains the sst-stage policy Tha a parameterized distri- bution over the cluster space C, using policy-based approach, and (ii) trains the 2nd-stage policy 2nd n2nd a parameterized distribution over the action space A conditional on a cluster sampled by the 1 st-stage policy, using regression-based approach. The underlying intuition that we can apply a policy-based approach to identify promising action clusters with low bias and variance, as the cluster space is much smaller than the original action space. We then apply a regression- policy identify promising actions within each cluster, minimizing variance. The resulting overall policy is more robust to reward modeling errors than typical regression-base approaches because we apply the regression- based policy only within each cluster[]
2nd-stage
'ne might wonder why not use regression- based policy for the first stage and policy-based approach for the second stage. While this idea may also seem reasonable first, produce an unbiased policy gradient
4

---

Published as a conference paper at ICLR 2025
When performing inference for incoming context z, first sample cluster from the Ist-stage policy as C ~ T8
1 z)). We then apply the 2nd-stage policy to choose the action in the cluster as
a ~ 2nd (-1 æ, c). This procedure is equivalent to producing an action as a ~ T (- [x).
T
The following describes how train 1st- and 2taes policies to improve the overall policy.
3.1 TRAINING THE IST-STAGE POLICY s
First, we develop a training procedure for the 1st-stage policy given a (pre-trained) 2nd-stage policy. Then, the theoretical analysis of the proposed training procedure will naturally tell us how we should construct the 2nd-stage policy (which will be described in the next subsection).
As mentioned earlier, given a (pre-trained) 2nd-stage policy के we consider training the Ist-stage policy !st parameterized by 0, via a policy- policy-based approach as below.
8+41 + 0g + VoV (18,0 )
(5)
This performs gradient ascent of 0 with the aim of improving the value of the overall policy T8,0 og The true policy gradient in Eq. (5) is given as follows (derived in Appendix D).
V (OV (Togerall) E Bp(s))3""((ax) [q" (I, (z,c)so(2, (2,)) c)
(6)
where we use 7" 2nd (ar, c) = (o (1(,) a)] to denote the value of cluster e under a 2nd-stage
(1(
policy and (0 c) = Vo log m] (cl (c| (>) to denote the policy score function of the Ist-stage policy.
Hence, given a 2nd-stage policy, our objective is to estimate the policy gradient in Eq. (6) to train 1st-stage policy. We achieve this via the following POTEC gradient estimator,
(NVrOTE((Tgg (overall ;D)
wo woii,, Cri) [(,,, (0)) Ca.) + P (ett ( (oi, ())
where wxx, == mp "c (o x) is the cluster importance weight. Specifically, the first term of Eq. 7 estimates the value of cluster via cluster importance weighting and the second term deals with the value of individual actions via the regression model f. Since our policy gradient estimator applies importance weighting with respect to only the action cluster space, it provides a substantial reduction in variance compared to typical estimators such as IPS and DR. We will discuss how should optimize the regression model based on the following analysis of our gradient estimator.
First, we characterize the bias of the POTEC gradient estimator under the following full cluster support condition (which is less restrictive than Condition 2.1p.
Condition 3.1. (Full Cluster Support) The logging policy mo has full cluster support if To(e| x)) > O, V(z,e) e) E X x C.
In the following theorem, we denote with A, a, b) = g(x,a) g(x,, b) the difference in the expected rewards between the pair of actions and b given æ. which we call the relative value difference of the actions. A a, b) := f(z, a) f(x, b) is an estimate of the relative value difference between and based on a regression model f.
Theorem 3.2. (Bias Analysis) If Condition 3.7 is true, the POTEC gradient estimator in Eq. 7 has the following bias for some given regression model f(z,a), a).
Bias(Ve Vo PPOTEC(TTT overali, D)) = E (p(()))""((m
Σ
z Ir, c)m 2)m xx,c) c)
acbc
(w(a,b) -y,a,,, w(z, a)) se(z,c)). (8)
in the second stage using policy-based approach, we would need to apply the typical importance weight w(z, a). resulting in the same variance issues typical methods. Hence, focus the proposed approach of using a policy-hased method in the first stage and regression-haseet approach in the second stage, leading substantial variance reduction while performing unhiased gradient estimation.
"This implies that the optimal cluster that should he chosen by the 1st-stage policy can be different given different 2nd-stage policics. Appendix D.neaniiixD.Tlaaaoo.. elaborrtes on this via numerical example.
5

---

Published as a conference paper at ICLR 2025
where a be A.
The proof is given in Appendix D.2 Theorem 3.2 shows that the bias of the POTEC gradient estimator is characterized by the accuracy of the regression model f with respect to the relative value difference, which is quantified by A, a. b) A, When f preserves the relative value difference of the actions within each cluster accurately, the second factor in Eq. (8) becomes small and so does the bias of the POTEC gradient estimator. This also suggests that, in an ideal case when the local correctness condition (Saito et al. 2023) holds true, the POTEC gradient estimator becomes unbiased (Ev[Vo (POTRO ( D) VaV V (moyeral)))
Condition 3.3. (Local Correctness) regression model and action clustering satisfy local correctness if ( a, b) = و (I, a, b) for all z E X' and a, b E s.t. Cb.
This implies that, to directly minimize the bias of the POTEC gradient estimator, should idcally optimize the regression model fo X'x A- + R via the following pairwise regression procedure (Saito et al. 2023) so that we can achieve a small [ b) - A a,b)l by preserving the relative value difference of the actions within each cluster.
min (r.
Σ
(9)
where Dpoir is a dataset augmented for performing pairwise regression, which is defined as (Ia, (a)) b, D
D Pair := (2,
I Ta Tb. Ca Co }. As suggested in Theorem 3.2 fo a)
characterizes the bias of the POTEC gradient estimator. Even if Eq. (9) is infeasible due insufficient pairwise data, we can still perform a conventional regression for the expected absolute reward to directly optimize the parameterized function (1, a) via min, (x.a))) ((r; J (z., a)) and then use fo in Eq. (7). Even for such conventionally trained regression model far POTEC still has advantages over existing policy gradient estimators, such as IPS and DR, due to its substantially reduced variance demonstrated in our experiments.
Next, the following shows the variance of the POTEC gradient estimator, and its variance reduction compared to typical importance weighting in the action space.
Proposition 3.4. (Variance Analysis) Under Conditions 3.7 and 3.3 for particular parameter 0 E Rd, the POTEC gradient estimator has the following variance.
n tr (Covpl WoWD\VoPPOTCC (ovrrat:))) overall :DI)
(
Σ 1=1
Ep(r)oo(aln)
[Iu( (a)s ( (2. (z,cc)))oPP
2nd + Ep(a) [V (Vaocale) [w(z (a) (or, a)s (x. ca)]] + Vp(a) [Eau"ccte)) (x
(x, (2, (x,))) }.
where A (r, = ((x,0) f(x, a) is the error of [xx a) against a).
Proposition 3.5. (Variance Reduction) The difference in the variance of the cluster and vanilla importance weights can be represented as follows.
V ((x))mtaal)] [wol a)] V (l [w [w(z, c)] = (E(())xatel)) [V (wll
Proposition 3.4 shows that the variance of the POTEC gradient estimator depends only on w(2z, xz, c) rather than w(x, a), implying reduced variance compared to IPS and DR. In addition, Proposition 3.5 characterizes the reduction in variance provided by cluster importance weighting of POTEC. It is worth noting that the variance reduction characterized by the variance of the vanilla importance weight (p(a))o(eaa)) [Vrotala,ec) [w[x,, a)]]. which suggests that cluster importance weighting provides increasingly larger variance reduction when typical weighting has a larger variance.
2nd 3.2 TRAINING THE 2ND-STAGE POLICY T
We have thus far developed policy-based approach for learning an effective cluster selection (Ist- stage) policy via the POTEC gradient estimator. The remaining objective is identify the optimal

---

Published as a conference paper at ICLR 2025
C
The POTEC Algorithm Spedtrum OPL Approaches
Regression-based
Pollcy-based
Num. of Clusters
Strong Condition
Condition
Local Correctness Spectrum of Reward Modeling Conditions
Figure 2:
The POTEC algorithm and local correctness condition generalize policy- and regression- based approaches and their respective conditions about the reward function (g(æ, a)) estimation.
Algorithm 1 The POTEC Algorithm
Input:
logged bandit data D, conventionally trained regression model q(x, a).
Output:
Ist-stage (policy-based) policy x 1st and 2nd-stage tregression-basedd policy x 1: Perform action clustering by applying clustering algorithm (such as KMeans) to the averaged estimated rewards for each action, q(a) == (1/n) D ((z;,a), as the embedding a.
2: Perform pairwise regression and obtain (x, a) as in Eq. (9). which works as the 2nd-stage policy in Eq. (10) and also regression model to help train the 1st-stage policy via the POTEC gradient estimator
Perform policy-based learning of the Ist-stage policy based on the POTEC estimator in Eq. 7
actions, given cluster selected by the Ist-stage policy. In essence, regression model fo from the previous section to establish the 2nd-stage policy =2nd because f is already optimized towards estimating the relative value differences of actions within each action cluster (i.e., local correctness). Specifically,
simply use the pairwise
construct 2nd stage policy based on fg as
2nd (a xz, : 1 (a = arg max to (o(x,a') (otherwise)
T
(10)
which implies that the 2nd-stage policy selects the action with the highest value of the pairwise regression function f within the already sampled cluster c. This action selection procedure justified since we have learned the function fo so that it can estimate the relative value difference of the actions given a cluster in bias minimization (Eq. (9)). ). In ideal scenario where Condition 3.3 holds true, our 2nd-stage policy achieves optimal action selection. In our experiments, we will demonstrate that our overall policy overall outperforms existing approaches by a considerable margin even with learned 2nd-stage policy that may not perfectly satisfy local correctness.
3.3 THE OVERALL POTEC ALGORITHM AND ITS INTERPRETATION
Algorithm describes the overall procedure of POTEC. It first performs action clustering based on an estimated reward function q. then performs pairwise regression (if feasible) and obtains the regression function fur which forms the 2nd-stage policy (as in Eq. (10). Note that pairwise regression is not feasible, we can instead use the conventional regressor as f <- Then, we train the 1st-stage policy TT based on the POTEC gradient estimator, which is based on cluster importance weighting and a learned regression model (x, a), which governs the bias of the estimator.
It worth mentioning that POTEC and its associated local correctness condition generalize typical OPL approaches, i.c., based and piaression-bassee as depicted in Figure 2] That is, when there only one action cluster (CC 1), the 2nd-stage policy POTEC needs choose the best action in reduction to the regression- approach. Moreover,
the entire action space, which can be seen in this case, the local correctness condition becomes relatively stringent (since all actions are grouped into the same cluster), which is akin to the typical condition of the regression-based approach, i.e., globally accurate estimation of the reward function. On the other hand, when the cluster space is equivalent to the original action space (C A), the 1st-stage policy selects action from the original action space, akin to the policy-base approach. this scenario, local correctness imposes no specific requirements, each action cluster contains only one unique action. This absence of requirements
7

---

Published as a conference paper at ICLR 2025
aligns with the policy-based approach, which does not necessitate specific conditions for reward function estimation to produce an unbiased gradient. Thus, POTEC and local correctness encompass the full spectrum of existing OPL approaches and respective reward-modelin conditions (Figure 2 As a strict generalization, POTEC generally outperforms both approaches with a good (even not perfect) selection of the number of clusters, the following section empirically demonstrates.
is also worth noting that the number of clusters and clustering algorithm are hyperparameters of POTEC, so can tune them via the cross-validation procedure. However, it is important to perform careful OPE on validation data to evaluate the performance of each hyperparameter with low bias. To address the variance issue in large action spaces during validation. we would recommend using OPE methods from Sachdeva et al. (2024); Saito et al. (2023); Peng et al. (2023).
4 EMPIRICAL EVALUATION
We first evaluate POTEC on synthetic data with ground-truth cluster information to compare the effectiveness of POTEC w/ and w/o true cluster information and w/ and w/o pairwise regression. We then assess the real-world applicability of POTEC on public recommendation dataset.
4.1 SYNTHETIC EXPERIMENT
First, we create synthetic datasets to compare the policy learning algorithms based on their ground- truth policy values. Specifically, we first sample 5-dimensional context vectors and action features from the standard normal distribution. We then form ground-truth action clusters based on the action features. We synthesize the expected reward function as g( a) g ca) + he, (zr, a), where g(2, c) and (1, a) define the values of the cluster and individual actions, respectively, as detailed in Appendix B We then sample an action a based on the logging policy To. which is defined as below.
exp(B (g(x, a) + nw,a))
To(a|z) Ev'e exp(3 (g(x, a') + Na,a)))
(1I)
where 3 is set to 5 as default and the noise aa, is sampled from normal distribution.
After sampling action a from the logging policy, we finally sample the reward r from a normal distribution with mean q(2, a) and standard deviation o, 1.0. Repeating the above procedure times generate logged training data of the form D {(1;, ag r:)
Baselines: We compare POTEC with the regression-basee method (Reg-based), IPS-PG (Eq. (I). and DR-PG (Eq. ). We use neural network with 3 hidden layers to parameterize the policy Te. ((z, a) for DR-PG and Reg- baseda and fo for POTEC. We apply the variance reduction technique proposed by Lopez et al. (202T) to IPS-PG and DR-PG. In the synthetic experiment, we compare POTEC w/ true clusters that define the true reward function g(z, a) and POTEC w/ learned clusters obtained by performing KMeans to the averaged estimated rewards for each action
Results: Figure 3 shows the policy values (normalized by the value of the logging policy V(To)) of the OPL methods on test data obtained from 30 simulations with varying seeds. Note that we employ default experiment parameters of 4, 000, |AI 500, and [CI = 30, and the shaded regions in the plots represent 95% confidence intervals of the test policy values estimated via bootstrap.
First, in most situations, POTEC provides significant improvements in policy value over the baseline methods. Specifically, in Figure (i), we observe that POTEC performs increasingly better with larger sample sizes, while the baseline methods remain inferior compared to POTEC. We also see that the policy value of POTEC is particularly higher than the baseline methods when the training data size is small, such as n < 2, 000, which suggests that POTEC is efficient even with small per-action sample size (i.e., nt/lA). while the baseline methods require even more data for each action to be effective. is also interesting to compare the policy value of POTEC (w/ true clusters) and POTEC (w/ learned clusters), where observe that POTEC (w/ true clusters) consistently performs best, while POTEC (w/ learned clusters) achieves competitive performance with POTEC (w/ true clusters) and outperforms the baselines by large margin across range of training data sizes. This suggests
'hhe number of learned action clusters (i.c., the parameter K of KMeans) is set equal CC..
8

---

Published as a conference paper at ICLR 2025
Reg- based
IPS-PG
POTEC/w/ Cluster?
POTECOW/ Learned Cluster)
(iit)
Normalized Policy
1.2
1.0
2000
4000
200
500
1000
2000
5000
100
230
training cista size
number of actions [.A]
number (true) clusters [C]
Figure 3: Comparing the test policy value (normalized by (=ao) of the OPL methods, with varying
(i) training data sizes, (ii) numbers of actions, and (iii) numbers of (true) clusters, on Synthetic Data.
DR PG
POTEC(Absolute Regression)
POTEC(Pairwise Regression)
Normalized Policy Value
(ii)
(iit)
1.3
03
2000
training data size
number of actions [4]
number of (true) clusters
[c]
Figure 4: Comparing the test policy value of POTEC w/ and w/o pairwise regression, with varying (i)
training data sizes, (ii) numbers of actions, and (iii) numbers (true) clusters, on Synthetic Data.
the real- world applicability of POTEC, even with the feasible and simple clustering procedure based on ((z a). Next, in Figure (ii), we vary the number of actions ((.A)) to investigate the robustness of the methods to increasing action spaces. We observe that POTEC performs consistently well, even as the action space grows, as long as the underlying cluster space C does not expand. In contrast, the performance of the baseline methods clearly deteriorates as the number of actions increases, and they perform even worse than the logging policy when |A| 2 2, 000. We can also see that POTEC (w/ learned clusters) performs substantially better than the baseline methods, particularly when the action space becomes large, even with heuristically learned action clusters. Finally, Figure (iii) evaluates POTEC as we increase the number of (true) clusters while keeping the number of actions fixed. The figure shows that the advantage of POTEC is greatest when the cluster effect can be captured by a small number of underlying clusters, However, even for data with [C] = 200 clusters, POTEC remains competitive with the baselines
In Figure we present the results of an ablation study for (i) varying training data sizes, (ii) varying numbers of actions, and (iii) varying numbers of (true) clusters, where we compare POTEC with pairwise regression [(r a), as described in Eq. (9) and with the standard absolute regression q(ar, a). We also include DR-PG, the best method among the baselines, as reference. From the figure, it clear that both versions of POTEC significantly outperform the best baseline (DR-PG) in most experimental settings, which empirically highlights the advantage of using POTEC, even with conventionally trained regression model q(x, a). However, is also evident that performing pairwise regression consistently improves POTEC, making it beneficial to apply when feasible.
4.2 REAL-WORLD EXPERIMENT ON KUAIREC
To assess the real-world applicability of POTEC, now evaluate it the KuaiRec dataset (Gao) ct al. 2022). a publicly available recommendation dataset collected on a short video platform, where 1,411 users have viewed all 3,317 videos and provide watch ratio (play duration divided by the vidco duration) as reward feedback. This full feedback nature of the dataset enables an OPL experiment without synthesizing the reward function q(x.a)) a).
'Not that this experiment varies the ground-truth underlying cluster structure, and as result, the policy values of the baseline methods may not remain constant.

---

Published as a conference paper at ICLR 2025
Reg-based
IPS-PG
DR-PG
POTEC(Ours)
Normalized Palicy Value
(1)
(ii)
(iil)
1.6
144
1.4
1.2
1.0
0.0
0.
0.3
0.5
5000
10000
0.5 0.3
training data Size
number of clusters [C]
ratio of deficient actions [[al/1-]]
Figure 5: Comparing the test policy value (normalized by (TTo)), with varying (i) training data sizes, (ii) numbers of (learndd) clusters, and (iii) ratios of deficient actions, on the KuaiRee dataset.
To perform an OPI. experiment and generate logged training data, we first sample user index u from uniform distribution, and for each user, we sample an action based the logging policy wo defined in Eq. (II). After sampling action a, we sample the reward r from a normal distribution with mean (z,,, a) and standard deviation a, 1.0 where the expected reward q(I,,, a) is the watch ratio recorded in the dataset. Repeating the above procedure n times generates 'D {[Tua Qg. r:) Pt)
We evaluate POTEC against IPS-PG, DR-PG, and Reg-based methods under varying training data sizes n, numbers of (learndd) clusters [C], and ratios of deficient actions (24 := fa E <o(a|z) = 0///|A/ (Sachdeva et al. 2020 For POTEC, we perform pairwise reward regression using Random Forest to obtain f(x a). Note that, in the real-data, there is ground-truth action clusters, so we implement only a feasible version of POTEC with learned action clusters.
Results:
Figure 5 presents the real-world experiment results for various configurations. First, Figure () compares the methods across different training data sizes, showing that POTEC is the most sample- efficient and performs significantly better than the baselines, particularly when the training dataset is small. Second, Figure (ii) compares POTEC across varying numbers of clusters [C]. a key hyperparameter of POTEC The figure shows that POTEC is most sample-efficient with a moderate number of clusters, such as |C] 50 or [C] = 100. It is also noteworthy that POTEC does not underperform the baseline or logging policies even in the worst cases, with overly small ((C) = 10) or large (CC) 1, 000) cluster spaces. Finally, Figure 5 (iii) compares the methods as we vary the violations of the full support condition (Condition 2.11 where we gradually expand the set of deficient actions Uo. When the logging policy To becomes nearly deterministic and offers limited exploration, (ol grows, making OPL more challenging. Indeed, in Figure (iii we observe that the baseline policy-based methods, particularly DR-PG, struggle the number of deficient actions [o4]] increases, which produces greater bias in policy gradient estimation. However, POTEC continues to perform significantly better than the baselines and logging policy, even with many deficient actions, because the POTEC gradient estimator requires only full cluster support (Condition I.T.. which is milder than full support and thus results in smaller bias in the gradient estimation of the Tst-stage policy.
CONCLUSION AND FUTURE WORK
This work introduces a novel two-stage OPL procedure called POTEC, which is particularly advanta- geous in large action spaces. POTEC learns the first-stage cluster-selectine selection policy via a new policy gradient estimator that is unbiased under local correctness and has substantially lower variance. The second-stage action-selection policy is learned through pairwise reward regression within each cluster, offering greater robustness to estimation bias compared to traditional regression-based approaches.
Our findings suggest valuable directions for future research. For instance, even though we have empirically demonstrated that POTEC generally outperforms existing OPL methods, we now rely on a heuristic action clustering, which could be seen as a limitation. It would thus be valuable to explore more principled method, such as context-dependert and iterative procedure to better satisfy local correctness. Extending POTEC to offline reinforcement learning and large language models beyond contextual bandits is also an intriguing future direction.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Shipra Agrawal and Navin Goyal. Thompson sampling for contextual bandits with linear payoff's. In International Conference on Machine Learning, pp. 127-135. PMLR, 2013.
Iaad Aouali, Victor-Emmanuel Brunel, David Rohde, and Anna Korba. Bayesian off-policy evalua- tion and learning for large action spaces. arXiv preprint arXiv:2402.146644 2024.
Susan Athey, Raj Chetty, Guido Imbens, and Hyunseung Kang. The surrogate index: Combining short-term proxies estimate long-term treatment effects more rapidly and precisely. Technical report, National Bureau of Economic Research, 2019.
Susan Athey, Raj Chetty, and Guido Imbens. Combining experimental and observational data to estimate treatment effects on long term outcomes. arXiv preprint arXiv:2006. 09676, 2020.
Yikun Ban and Jingrui He. Local clustering in contextual Web Conference 2021, pp. 2335-2346, 2021.
bandits. In Proceedings of the
K. Bhatia, K. Dahiya, H. Jain, P. Kar, A. Mittal, Y. Prabhu, and M. Varma. The extreme classi- fication repository: Multi-label datasets and code, 2016. URL httt::://mmiivaaatie org/
Yash Chandak, Georgios Theocharous, James Kostas, Scott Jordan, and Philip Thomas. Learning ac- tion representations for reinforcement learning. In International Conference on Machine Learning, pp. 941-950. PMLR, 2019.
Jiafeng Chen and David Ritzwoller. Semiparametric estimation of long-term treatment effects. arXiv preprint arXiv:2107.14400 2021.
Matej Cief, Jacek Golebiowski, Philipp Schmidt, Ziawasch Abedjan, and Artur Bekasov. Learning action embeddings for off-po evaluation. In European Conference on Information Retrieval, pp. 102-112. Springer. 2024.
Miroslav Dudík, Dumitru Erhan, John Langford, and Lihong Li. Doubly robust policy evaluation and optimization. Statistical Science, 5114 2014.
Mehrdad Farajtabar, Yinlam Chow, and Mohammad Ghavamzadeh. More robust doubly robust off-policy evaluation. In Proceedings of the 35th International Conference on Machine Learning, volume 80, pp. 1447-1456. PMLR, 2018.
Nicolo Felicioni, Maurizio Ferrari Dacrema, Marcello Restelli, and Paolo Cremonesi. Off-policy eval- uation with deficient support using side information. Advances Neural Information Processing Systems, 35, 2022.
Chongming Gao, Shijun Li, Wenqiang Lei, Jiawei Chen, Biao Li, Peng Jiang, Xiangnan He, Jiaxin Mao, and Tat-Seng Chua. Kuairee: A fully-observed dataset and insights for evaluating recom- mender systems. In Proceedings of the 31st ACM International Conference on Information & Knowledge Management, pp. 5500-550--00 2022.
Pengjie Gu, Mengchen Zhao, Chen Chen, Dong Li, Jianye Hao, and Bo An. Learning pseudometric- based action representations for offline reinforcement learning. In International Conference on Machine Learning, pp. 7902-7918. PMLR, 2022.
Olivier Jeunen and Bart Goethals. Pessimistic reward models for off off-policy learning in recom- mendation. In Proceedings of the 15th ACM Conference on Recommender Systems, pp. 63-74, 2021.
Nan Jiang and Lihong Li. Doubly robust off-policy value evaluation for reinforcement learning. In Proceedings of the 33rd International Conference on Machine Learning, volume 48, pp. 652-661. PMLR. 2016.
Thorsten Joachims, Adith Swaminathan, and Maarten de Rijke. Deep learning with logged bandit feedback. In International Conference Learning Representations, 2018.

---

Published as a conference paper at ICLR 2025
Nathan Kallus and Masatoshi Uchara. Double reinforcement learning for efficient off- off-policy evalua- tion in markov decision processes. J. Mach. Learn. Res., 21:167-1,
Nathan Kallus, Yuta Saito, and Masatoshi Uehara. Optimal off-policy evaluation from multiple logging policies. In Proceedings of the 38th International Conference on Machine Learning, volume 139, pp. 5247-5256, PMLR, 2021.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 11Xx::1412.6980, 2014.
Haruka Kiyohara, Yuta Saito, Tatsuya Matsuhiro, Yusuke Narita, Nobuyuki Shimizu, and Yasuo Yamamoto. Doubly robust off off-policy evaluation for ranking policies under the cascade behavior model. In Proceedings of the 15th International Conference on Web Search and Data Mining, 2022.
Haruka Kiyohara, Masatoshi Uchara, Yusuke Narita, Nobuyuki Shimizu, Yasuo Yamamoto, and Yuta Saito. Off-policy evaluation of ranking policies under diverse user behavior. In Proceedings the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 1154-1163, 2023.
Tor Lattimore and Csaha Szepesvári. Bandit algorithms. Cambridge University Press, 2020.
Jaron JR Lee, David Arbour, and Georgios Theocharous. Off-policy evaluation in embedded spaces.
arXiv preprint arXiv:2203.02887, 2022.
Lihong Li. Wei Chu, John Langford, and Robert Schapire. A contextual-bandit approach to personalized news article recommendation. In Proceedings the 19th international conference World wide web, pp. 661-670, 2010.
Dawen Liang and Nikos Vlassis. Local policy improvement for recommender systems. arXiv preprint arXiv:2212.11431, 2022.
Qiang Liu, Lihong Li. Ziyang Tang, and Dengyong Zhou. Breaking the curse of horizon: infinite- horizon off- off-policy estimation. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 5361-5371, 2018.
Yao Liu, Pierre-Luc Bacon, and Emma Brunskill. Understanding the curse horizon in off-policy evaluation via conditional importance sampling. In International Conference Machine Learning, pp. 6184-6193. PMLR, 2020.
Ben London and Ted Sandler. Bayesian counterfactual risk minimization. In International Conference on Machine Learning, pp.
PMLR, 2019.
Romain Lopez, Inderjit S Dhillon, and Michael I Jordan. Learning from extreme bandit feedback. Proc. Association for the Advancement of Artificial Intelligence, 2021.
Yifei Ma. Yu-X Xo-Xiang Wang, and Balakrishnan Narayanaswamy, Imitation-regullaiieed offline learning. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2956-2965. PMLR, 2019.
Alberto Maria Metelli, Alessio Russo, and Marcello Restelli. Subgaussian and differentiable impor- tance sampling for off- off-policy evaluation and learning. Advances in Neural Information Processing Systems, 34, 2021.
Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, Jake Vanderplas, Alexandre Passos, David Cournapeau, Matthieu Brucher, Matthieu Perrot, and Edouard Duchesnay. Scikit-leam: Machine learning in python. Journal of Machine Learning Research, 12:2825-288 28300, 2011.
Jie Peng. Hao Zou, Jiashuo Liu, Shaoming Li, Yibao Jiang, Jian Pei, and Peng Cui. Offline policy evaluation large action spaces via outcome-orienter action grouping. In Proceedings of the ACM Web Conference 2023, pp. 1220-1230, 2023.
12

---

Published as a conference paper at ICLR 2025
Noveen Sachdeva, Yi Su, and Thorsten Joachims. Off-policy bandits with deficient support. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 965-975, 2020.
Noveen Sachdeva, Lequn Wang, Dawen Liang, Nathan Kallus, and Julian McAuley, Off-policy evaluation for large action spaces via policy convolution. arXiv preprint arXiv:2310. 2023.
Noveen Sachdeva, Lequn Wang, Dawen Liang, Nathan Kallus, and Julian McAuley. Off-policy evaluation for large action spaces via policy convolution. In Proceedings of the ACM on Web Conference 2024, pp. 3576-3585, 2024.
Yuta Saito and Thorsten Joachims. Counterfactual learning and evaluation for recommender systems: Foundations, implementations, and recent advances. In Proceedings of the 15th ACM Conference Recommender Systems, pp. 828-88 830, 2021.
Yuta Saito and Thorsten Joachims. Off-policy evaluation for large action spaces via embeddings. In International Conference Machine Learning, pp- 19089-19122. PMLR, 2022.
Yuta Saito, Shunsuke Aihara, Megumi Matsutani, and Yusuke Narita. Open bandit dataset and pipeline: Towards realistic and reproducible off-policy evaluation. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2021a.
Yuta Saito, Takuma Udagawa, Haruka Kiyohara, Kazuki Mogi. Yusuke Narita, and Kei Tateno. Evaluating the robustness of off-policy evaluation. In Proceedings of the 15th ACM Conference on Recommender Systems, pp. 143,123, 2021b.
Yuta Saito, Ren Qingyang, and Thorsten Joachims. Off-policy evaluation for large action spaces via conjunct effect modeling. International Conference on Machine Learning, pp. 29734-29759, PMLR, 2023.
Yi Su, Lequn Wang, Michele Santacatterina, and Thorsten Joachims. Cab: Continuous adaptive blending for policy evaluation and learning. In International Conference on Machine Learning. volume 84, pp. 6005-6014, 2019.
Yi Su, Maria Dimakopoulou, Akshay Krishnamurthy, and Miroslav Dudík. Doubly robust off-policy evaluation with shrinkage. In Proceedings of the 37th International Conference on Machine Learning, volume 119, pp. 9167-9176. PMLR, 2020a.
Yi Su, Pavithra Srinath, and Akshay Krishnamurthy. Adaptive estimator selection for off-policy evaluation. In International Conference on Machine Learning, pp. 9196-9205. PMLR, 2020b.
Adith Swaminathan and Thorsten Joachims. Batch learning from logged bandit feedback through counterfactual risk minimization. The Journal of Machine Learning Research, 16(1):1331-1755, 2015a.
Adith Swaminathan and Thorsten Joachims. Counterfactual risk minimization: Learning from logged bandit feedback. International Conference on Machine Learning, pp. 814-823, PMLR, 2015b.
Adith Swaminathan and Thorsten Joachims. The nelf-normalized estimator for counterfactual learning. Advances in Neural Information Processing Systems, 28, 2015c.
Muhammad Faaiz Taufiq, Arnaud Doucet, Rob Cornish, and Tean-Francois Ton. Marginal density ratio for off-policy evaluation in contextual bandits. Advances in Neural Information Processing Systems, 36, 2024.
Philip Thomas and Emma Brunskill. Data-efficient off-policy policy evaluation for reinforcement learning. In Proceedings of the 33rd International Conference Machine Learning, volume 48, pp. 2139-21483 PMLR, 2016.
Takuma Udagawa, Haruka Kiyohara, Yusuke Narita, Yuta Saito, and Kei Tateno. Policy- addatie estimator selection for off-policy evaluation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, 2023.
13

---

Published as a conference paper at ICLR 2025
Cameron Voloshin, Hoang M Le, Nan Jiang, and Yisong Yue. Empirical study of off- off-policy policy evaluation for reinforcement learning. arXiv preprint wwXiv:191.00685/1 2019.
Yu-Xiang Wang, Alekh Agarwal, and Miroslav Dudik. Optimal and adaptive off- evaluation in contextual bandits. In International Conference on Machine Learning, pp. 3589-3597. PMLR, 2017.
Tengyang Xie, Yifei Ma, and Yu- Yu-Xiang Wang. Towards optimal off-policy evaluation for rein- forcement learning with marginalized importance sampling. In Advances in Neural Information Processing Systems, pp. 9665-9675, 2019.
Yinglun Zhu, Dylan J Foster, John Langford, and Paul Mineiro. Contextual bandits with large action spaces: Made practical. In International Conference on Machine Learning, pp. 27428-27453 27453. PMLR, 2022.
14

---

Published as a conference paper at ICLR 2025
A
MORE RELATED WORK
Off-Pollicy Evaluation:
Off-policy evaluation of counterfactual policies has recently garnered
significant interest in both contextual bandits (Dudik et al. [2014] Farajtabar et al. 2018 Kallus et al..
2021
Kiyohara et al., 2022 2023 Metelli et al. 2021] 1
Saito & Joachims, 2021; 1 Su et al. 2020a
2019 Wang al. 2017 ) and reinforcement learning (RL) (Jiang & Li 2016, Kallus & Uehara 2020
Liu et al. 2018 2020] Thomas & Brunskill 2016 Xie et al. 2019). The literature encompasses three
mai approaches. The first approach, named Direct Method (DM), is defined as:
VDM((T D,9) = Σ Bx(alz))0(01, a)]
Σ
x(a 2,)(((,,)))
REAA
where ((2; a) estimates q(z, a) based on logged bandit data. This approach exhibits lower variance than IPS and has been utilized to address violations of full support (Sachdeva et al. 2020). where IPS can be severely biased. However, DM is often vulnerable to reward function misspecification. This issue is problematic, as the extent of misspecification cannot be easily detected and evaluated for real-world data due to non- -ineariyy or partial observability of the environment (Farajtabar et al., (Sachdacee al., 2020] Voloshin al. 2019). The second approach is IPS, which estimates the value of by re-weighting the observed rewards as
(a: 22) To(a; x3))
VIPs (T; D)
Σ (1)
where w(xr,a) a) == (aa] (x) Toaa (x) is called the (vanilla) importance weight. Under some identífica- tion assumptions such as no interference, full support, and no unobserved confounders. IPS provides unbiased and consistent estimation of the value of new policies. However, this approach has a critical drawback: it can suffer from high bias and variance in the presence of numerous actions. First, high bias occur when the logging policy fails to provide full support (Condition 2.I}, which is likely in larger action spaces (Sachdeva et al. 2020 Saito & Joachims 2022). Furthermore, its variance can be particularly excessive for large action spaces, the importance weights are prone to taking extremely large values. It is possible to apply weight clipping Su et al. (2020a; 2019); [S Swaminathan & Joachims (2015b) and self- nelf-normalization (Swaminathan & Joachims 2015c) to somewhat alleviate
the variance issue, however, they introduce additional bias in return. DR, which is given as follows, a third approach that can be considered hybrid of the previous two approaches, achieving lower bias than DM and lower variance than IPS (Dudik et al. 2014
Farajtabar et al. 2018).
(T; D.
(w(rn, [(0; a:)) + (n(a)=)) a)
Several recent studies have extended DR to further improve its finite sample accuracy (Su et al. 2020a Wang et al., 2017 Metelli et al. 202T] or its robustness to model misspecification (Farajtabar et al. 2018 Kallus et al.
2021). Although there is number of extensions of DR in both bandits (as
described above) and RL. (Jiang Li 2016 Kallus & Ueharal [020)] Thomas & Brunskill] (016)),
these variants of DR still face the critical variance issue in large action spaces due to the same reasons
as IPS Saito & Joachims 2022); Saito et al. (003)).
To address the fundamental issues of typical OPE estimators for large action spaces, (Satoo & Joachims| 2022) proposed a new framework and estimator called Marginalized IPS (MIPS). This approach leverages auxiliary information about the actions, called action embeddings or action features, which are available in many potential applications of OPE such as recommender systems, and provide useful structure in the action space. More specifically, MIPS is defined as:
a
T(e: xi) To xT) ri
VAIPS (T; D) =
wox ei)ri
n
where the logged dataset D (0 ayy Cy 1i)l now contains action embeddings for each data =(e E, p(e Ple [.a)=(a|m)) point and w(x,e) := <)(e|æ] S. ple (.a)=o(a| is the marginal importance weight. This weight is defined with respect to the marginal distributions of the action embeddings induced by the
(r,, ~ p())to((| x)p(e a)p(r a, e) where p(e ax a) is an action embedding distribution.
15

---

Published as a conference paper at ICLR 2025
target and logging policies. This enhanced weighting scheme results in significantly lower variance compared to IPS and DR in larger action spaces, while maintaining unbiasedness under the no direct effect assumption. This assumption necessitates that the given action embeddings be informative enough to mediate every causal effect of the actions on the rewards (i.e., 0 ... r 2c, e). A similar condition regarding the causal structure has been utilized to address the deficient support problem in OPE (Felicioni et al., 2022 Lee et al. [2022 Peng et al. 20233 Sachdeva et al. 2023)) and to conduct causal inference of long-term outcomes through short-term proxies (Athey et al. 2020 2019 Chen & Ritzwoller 2021). However, MIPS may still exhibit high variance, similarly to IPS, when the provided action embeddings are high-dimensional and fine-grained. Additionally, it may generate substantial bias if the no direct effect condition is violated and action embeddings fail to explain much of the causal effects of the actions. This bias issue is particularly expected when performing action feature selection on high-dimensional action embeddings to reduce variance Su et al. (220b);
Udagawa et al. (2023).
To circumvent the bias-variance dilemma of MIPS, (Saito et al. 2023) proposed more general formulation and a refined estimator. Specifically, instead of relying on the often demanding no direct effect condition, (Saito al., 2023) introduced the conjunct effect model (CEM) of the reward function. The CEM is a useful decomposition of the expected reward function into what is called the cluster effect and residual effect. Building on the CEM, we can employ model-free estimation utilizing cluster importance weights to estimate the cluster effect without bias, and apply model-based estimation using the pairwise regression procedure to estimate the residual effect with low variance as
VORCEM(T: D) ==
1 T1
Ww Mr. ax)) + Exte
where w(x,e) = The first term of OffCEM estimates the cluster effect through cluster importance weighting, while the second term addresses the residual effect using the regression model f, which is ideally learned via a two-step procedure similar to POTEC. As result, the OffCEM estimator is likely to achieve significantly lower variance than IPS, DR, and MIPS in scenarios with many actions or high- dimensional action embeddings. while often reducing the bias of MIPS since OffCEM does not ignore the residual effect. Our OPL algorithm is inspired by this CEM formulation, and suggests training two distinct policies via policy-based (model-free) and regression-bast based (model-based) approaches, respectively.
s(cl=) x) To(c|a x)) Enea [fc,=c]zo(aa z) is referred to as the cluster importance weight.
EnEa ((a=))
Off-Policy Learning: The contextual bandit framework has become popular and practical ap- proach for online learning and decision-making under uncertainty (Lattimore & Szepesvári] 2020). with many efficient algorithms proposed for exploring the (potentially large or even infinite) action spaces [Arawall & Goyal y; Li et al. (2010). However, a pressing need exists for offline procedure that optimizes maaision-making without requiring risky and time-consuming active ex- ploration. As a result, there has been significant recent interest in developing efficient off-policy learning methods in the contextual bandit setting (Sachdeva et al. (2020): Saito & Joachims (021)). Fortunately, many real-world interactive systems can often leverage logged interaction data to learn an improved policy fully offline, which allows us to safely enhance the performance of the current system Joachims et al. (0118) London & Sandler (2019); Sachdeva al. (2020): Saito & Joachims (2021 Swaminathan & Joachims (2015a b).
As already described in Section there are two main families of approaches in OPL: regression- and posicy-based methods. The regression-basse
approach relies on reduction to supervised learning. where a regression estimate is trained to predict the rewards from the logged data JatanJeunen & Goethals (202T): |acchdeva et al. (2020). To derive a policy, the action with the highest predicted reward is chosen deterministically. y, or a distribution can be formed based on the estimated rewards as well. A drawback of this straightforward approach is the bias that arises from the misspecification of the re- gression model. On the other hand, the policy-based approach aims to update the parameterized policy To by performing gradient ascent iterations of the form: 0,+1 + 0, + (Ta) at each step t during policy learning. Since the true policy gradient (oV(TT) E (z) [q(2, a) Ve log a x x)) is unknown, it must be estimated from the logged data using OPE techniques, such as IPS (Eq. (
and DR (Eq. (2). However, these estimators necessitate the assumption that the logging policy has full support for every policy in the policy space. This assumption is frequently violated in large
16

---

Published as a conference paper at ICLR 2025
action spaces, leading to significant bias in gradient estimation. Moreover, existing policy gradient estimators heavily rely on the vanilla importance weight with respect to the original (potentially large) action space, resulting in critical variance issues and insufficient off oif-policy learning for large action spaces. One possible approach to address the variance issue in OPE is to apply conservative imitation regularization Jeunen & Goethals (202T): Liang & Vlassis (2022); Ma et al.) (2019); Swaminathan & Joachims (2015b), which penalize policies that diverge from the logging policy. However, in large action spaces, these regularization techniques often yield a policy that is too close to the logging policy. To tackle the challenges associated with OPE in large action spaces, (Lopez al. (02I) recently proposed the following selective IPS (sIPS) estimator to estimate the policy gradient.
To(a; Di Oj ((x))) Ys Vo log ( (x))
V VOPAIS(TT: D)
(12)
To(a; (i)
where ((x) fa E Al Alq( > 0} is the set of relevant actions called the action selector. The idea to reduce the variance in importance weighting by focusing only on relevant actions assuming that there are many irrelevant actions that have (almost) zero expected rewards in real applications. However, we argue that the variance reduction effect of sIPS is often limited, as it still relies on the logging policy in the denominator. Furthermore, reliable method for identifying the action selector has not yet been provided.
To address the limitations of existing approaches, utilize the CEM from Saito et al. (2023) and proposed the POTEC algorithm, which is the first OPL framework to unify regression- sagression-bbsed and policy-based approaches. This algorithm trains separate policies using regression-basee and policy-based approaches, respectively typical policy- and regression-bastt importance weighting when training the Ist-stage policy and regression- approach when training the 2nd-stage policy, which should yield significantly lower variance compared to existing policy-based methods that apply importance weighting over the original action space. Furthermore, our algorithm likely to be more robust to reward function misspecification than the regression- based approach, as it relies on provably unbiased policy gradient in the Ist-stage and aims to estimate only the relative value difference in the 2nd-stage. This is arguably a simpler task compared to the absolute value regression of the conventional regression-based approach.
6 In particular, our POTEC algorithm is expected to outperform approaches in large action spaces. First, utilize cluster
Note that the context reinforcement learning (RL), there are some related ideas and methods to improve sample- -eapee-efficincyy in large action spaces. For example, Chandak et al. (2019) propose method to learn action representation to improve sample-efficiency of on- on-policy RL. However, the focus of Chandak et al. 2 is not offline policy learning, and thus its proposed method is not considered as a baseline in our paper. In addition, the supervised representation learning procedure of this paper uses the structure specific to RL (i.e., state transition), it cannot be applied to contextual bandit setup. In addition, Gu et al. (2022) study offline RL in large action spaces and propose method learn latent representation the action space. However, the proposed method Gu et al. (2022) leverages the data nattadistributionall metric to learn action embeddings to deal with large action spaces in offline RL, but the metric is based on the MDP structure, and how apply the method to the offline contextual bandit problem was not discussed and it is non-trivial.
Table 1: Examples of locally correct regression models
Go
02
43
0g
0:
4g
Go
02
Oltoo a)
0
1
o a) ((
0
1
O a)
0
1
q(Za, a)
4
1
3
2
4
2
q a)
4
1
3
2
3
0
0
fe(2o. a)
50
47
-30
-31
fa(#o.
4
1
3
2
3
a.
A A(ra,
3
1
'Note that DR in Eq. 2 should be classified as a policy-based approach since its aim is to accurately estimate the true policy gradient, even though it employs a regression-based reward function estimator for variance reduction from IPS.
17

---

Published as a conference paper at ICLR 2025
B EXAMPLES: LOCALLY CORRECT REGRESSION MODELS
This section provides some examples regression model f that satisfies Condition 3.3 (local correct- ness). Suppose that there is only single context X' {mo) and four actions A = (0n a1 02 03).
The expected reward function q(z, a) and clustering function ((x, a) are given as follows.
g( uo) 4, (( a1) 1, g a2) = 3, ((z0, da) = 2, ((20, ao) = O, 6( a1) = 0, p(z,, a2) = 1, @ 03) =
Then, Table provides three locally correct regression models (fi to fa). More specifically, these example models succeed in preserving the relative value difference of the actions within each action cluster (c = for ap, a1 and = 1 for 42, a3). In fact, we can see that Ag do, a1)
Aj, fi (Fu, ao, 01)
A fo (To, do, a1) A fg (Io, au, 01) and (T 02, az) A (Zo, 02 as)
A fa (Zo, 02 03) = A (2o, 02, 03) = 1 where 6(20, 0o) = 6 a1) and o(zo, a2) = b(20, az).
C GENERALIZATION OF OUR FRAMEWORK AND POTEC ALGORITHM
In this section, we describe the generalization of our framework and algorithm the situation under the presence of some predefined action representation X' x A - & C IRd which is often available in practice and can be used to better parameterize the policy. Under the presence of such action representations, can first generalize the CEM as follows.
q(z,a) a) glz, cla, (xx, + h(2, (x)
(13)
cluster effect
residuul effect
where X' provides discretization the action representation space &, Note also that the residual effect depends on the representation of the action (ar, a) rather than the atomic actions as simpler version presented in the main text.
Leveraging this general version of the CEM in Eq. estimator Eq. the following two ways.
13 we can generalize our POTEC gradient
Implementation Option 1:
This option trains parameterized distribution over the action repre- sentation space & as the lst-stage policy via the following version of the POTEC gradient estimator.
Vo PrOTEC(TO (overall,
oveerlll D) = m Σ
(u(x,,))(r - f(x,, (( az)) V log ( ax))
+ E Eewagee (f and
(I Vo log
log Tole m3)1 m.)]] (14)
where c c(1,, B(2,, (,)),,)))) fm (z,c) =E, Egggnll [f(z a)] and
1st
Secota.e)=e T 1st (e(x))
w(x, c) = mg
Jese(a,e)=e 70 1
This general version of the POTEC gradient estimator is unbiased under local correctness (i.e, (, a, 6) = 42 a, b), Va, a, such that c(x, (xx, a)) = c(2,, ((x, b))). Since the Ist-stage policy learned the action representation space, it can naturally exploit the smoothness E.
If we follow this implementation, in the inference time, for an incoming context we first sample point in the action representation space & from the Ist-stage policy as e ~ (=]x(= which implics promising region in E. Note that, in general, E È will not match with any already observed action representation { a.))am Then, the second-stage क which constructed from the pairwise regression model hy X + 1R, identifies the best action within the promising region as
hnp (xxxa))))
arg max
a':o(z, ((.==))
where (a' E Al ((x, d')) = c(x,e)) is the set of actions whose representation lies in the promising region induced by nz T 1z).
18

---

Published as a conference paper at ICLR 2025
Implementation Option 2: This option first leams parameterized distribution over the action space A as the sst-stage policy using (22 a) as its imput via the following version of the POTEC gradient estimator.
Va TANPOTEC(TTN) ( D) :- 1 Σ
w(a ci)r( - fax p(r,, a.))) log (a(a; | (( a:))
+ Bawnya ( Vo my (2
(15)
where G c(ar, pari, a:)); f
(2r,c) (
-Raguul (f a)] and
T 1 (c x) w(r, c)
Emuo(a Tp (a +(2
To st (c)x)
Emue( TD Tb (a x; )(x,a))
This version also unbiased under local correctness (i.e, Ag a, b) ه ,(6 a, b), Vx, a, b such that c(z, ((x, a)) = c(z, ((x, The Ist-stage policy also simply leverages the action representation as its input
we follow this implementation, in the inference time, for an incoming context a. we first sample point in the action space A from the 1st-stage policy implies promising region in E. Then, the second-stage m2nd which is constructed from the pairwise regression model hy X 3 - R. identifies the best action within the promising region as
x0 ( xx; (xx, a)), which merely
(a, D(x, d')),
arg max
a':a(a, (=,(,,,zz,)))))) (rxaa))
where {a' E A c(z, ((x, d')) = c(x, ((r, a))} is the set of actions whose representation lies in the promising region induced by no ( x;; (x),)))
The empirical comparison of the above two options highly depends on each application. For example, Implementation Option may perform better when the action representation space is low-dimensional while it may suffer when is high-dimensional. Therefore, under the presence of some action representation (xx, a), we would encourage the practitioners identify the best implementations for their particular application in data-driven fashion, for example, by performing a careful cross-validation.
C.1 THE ONE-STAGE VARIANT OF POTEC
It is worth noting that there exists one-stage variant of POTEC, opposed to the two-stage variant, which is our primary proposal. More specifically, the one-stage variant directly trains a parameterized overall policy in the action space, No æ)). via the POTEC gradient estimator as follows:
V VoVPOTECI (To; D)
Σ w(r,, Cas (r: f(r,, ()))ss(();))))))))))))))))))( 0:) + Es [f(e a) a) ()) all.
where s0(a, a) = V Logn((( x).. Although the one-stage variant is categorized as policy-based approach, as it trains the overall policy directly via policy gradient, it still achieves significant variance reduction compared to IPS-PG and DR-PG and remains unbiased under local correctness. However, the one-stage variant could be considered suboptimal utilization of the local correctness condition since, given a locally correct regression model, we should be able to optimally choose the action within cluster in Eq. (10) and thus do not need learn the overall policy solely through policy gradient. Nevertheless, the one-stage variant may be valuable practice, as it do not need to maintain and execute multiple policies. We provide an empirical comparison of the one-stage and two-stage variants of POTEC in Appendix D
Forr example, we can define a parameterized policy as
exp(fo(2, a())a))) En'E cxp(Selz, ((x
ne 1: (1r, a))
where fo + TR: is some parameterized function having action representation (x,, a) its input.
19

---

Published as a conference paper at ICLR 2025
Table 2: Dependence of the cluster value on the 2nd-stage policy (q (
(x, c))
Go
Q1
Q/
Q3
40
41
Q2
43
c(zo, a) G(Io. a) 2nd Ri
0
c(ro, a)
0
1
4
2
5
0
q( a)
4
2
5
(a xr, c)
1
0
1
0
π 2nd
(a 20
0.5
0.5
0.5
0.5
2nd
(2)
2nd q
(x,
3
2.5
D OMITTED PROOFS
D.I DERIVATION OF EQ. 6
VoV V T0,
ovurall)
Bp(x)
Σ
q(E, a) V overall
(a x))
LaEA
-p((x)
Σ q(E, a) Σ Vex (c) अ) 2nd (a x,
LaEA
= -Ep(x)
Σ
VoTT lat (c(())
2nd (alaaa)zzmuaa
Lee
aEA
-E(())
Σ
m [x) 1og (c)2))²) and (x)
LEEC
=pp()))""(claa)
(z. (z))))((
where q- 2nd (x,c) c) (alad(a a. (((xaa) and s0(x,,) V log m (c| The above
policy gradient suggests increasing the choice probability of a cluster that is promising under the given 2nd-stage policy n. annd where the effectiveness of a cluster under the 2nd-stage policy is
quantified by y - (ar, This implies that the optimal cluster can be different given different 2nd-stage policies. A toy example in Table Ssoos that the value of cluster can indeed be very different given different 2nd-stage policies. More specifically, the left table shows the case with the optimal 2nd-stage policy that can identify the best action within each cluster. Then, we can see that the optimal cluster = 1, since the maximum expected reward in the actions of this cluster is larger. In contrast, the right table shows the case with uniform 2nd-stage policy. Under such a 2nd-stage policy, the optimal cluster then becomes e = O, since the average expected reward of the actions in 0 is larger than that of c= 1.
Below we prove the theorems presented the main text based on the following general version of the POTEC gradient estimator.
V TOVOTECCTTQV Toverall, overall D) = n
w(r, ()(r (.r. Ca) ) + Enger (f 2nd (Is c)so cill
where p(11 1x,, 0;) is a stochastic and context-dependen clustering. The POTEC gradient estimator defined in Eq. (7) can be considered special case with a deterministic and context- independent clustering function e A C.
Note that we w(x,c) ox(a|z,o)//( [w(2, (2, a)] and w(x,a) =(alz) <(a,clz) in the following. Ma(a x) mofa.c
D.2 PROOF OF THEOREM 3.2 AND COROLLARY ??
Proof. To derive the bias of the POTEC gradient estimator, we calculate the difference between its expectation and the true policy gradient given in Eq. (6) below.
20

---

Published as a conference paper at ICLR 2025
Bias( Vo (= (agerrt :D)
-yoomdaaacceearrkaauuuu o) (tt + Beto) " (2, ((,,,))((1))))
- -Bp(s)ss"((t))
and 1
(1,)))o(((,c))
= Ppp
Σ LaEA
To(a|1)4 (x, a)
Σ DEC
p(c a)w(s, () e) + Rp())
Σ प a (c(x)f(0 (c. ())
LOEC
-pp(
Σ
1st (c|x))"0 (I,, ())
LeEC
Tbt a () ( w(x, c))
Bp(s)) Σ To(a x) 4. (
To x))
DEC
x) To st (c x
T (c|z))
+ Ep(())
Σ LEEC
40
1st (c|2)
1 (c|
fm (I, (1,))80(2,C))
Ep(x)
Σ Leed
(cla
x)
1st
(c|x)
9 (2,
to
= =Pp(a)m("(ela)))
w(r, (0)cc) c)
Σ
40
(a (al 2r, c) ()
aEd
w(z, c)fa Bad (I,
Ep(e)s ( (z,c)
w(2, c))
Σ
To
(a x, c) g) (I)
acA
- -Ep())s
50(x, C)
T
To
( x.) T
1st (cl
x) T (a x)
2nd (al xr c)
To
2nd (a| x. (.,,
s0(I,c) c)
Σ
w(x a) 2nd (a x, c) Σ
2nd
aEA
beA
- -Bp(s)e
80(2, c)
Σ a 2nd (a x, c) (
aEA
(E
co(
Σ
w(x, a)d I,,)
2 (b () - - (e,
aEAl
bEd
where A f (x, = g( a) - f(x,a). By applying Lemma B.1 of (Saito & Joachims 2022) to
the last line (setting ((a) w(, g(a) (an I,), h(a) = A(, we obtain the following expression of the bias.
(cla)
so s0 c)
x Σ 2nd (a x, o (b|z, æ c) (A ( (x, (w(z, b) - w(z, a))
a<b
In particular, in the simpler case of deterministic and context-indepentt clustering as in the main text, we can simplify the expression of the bias as below.
2nd 2nd(a (a x, c) To
2nd (b x c)
- (w(z, 6) w(x, a)) s))
(p(a)m (clo)
Lactbe, =c
= =Bp(a)m("((la)))
Σ
2nd (a x, c) 2nd (b| c) (2, - ,, (1(2,,) b) - wlo, (0(
Lastica =oo=
where, we used Ti and (a c) = sate He and A f a) Ag,p(z,b) = (, Aslo, b).
21

---

Published as a conference paper at ICLR 2025
D.3 PROOF OF PROPOSITION 3.4
Proof. We apply the law of total variance several times to obtain the variance of the j-th element of the POTEC gradient estimator for particular parameter IRd in the following.
pehmaldleppiddeeaapprpooo w(( c) ()) (x, c) + Bapa(eea)) (z. (00 (z,
(x d1]
(prria, (r
- a) (0) (I, (I Big
(x, c']]]
Bp(rlan) { (r ( (z,c) B. (x.c)s (z, cx]]
+ (oleemoaaao))
(w0(1, (10t,,,))000,,0)))))0)000
(2,c) (od (z, dd]]
(z,d)ss
(100 (1
(p(ass((a)a))
Vololaaa
w(I,
Bn (o, (x, (x,'')
+ (x (c, (1,6)s (c, (z,d))]
(e(emm(aaa))
(10)
(p( (olelan)
a)) (z. o)]
Vo(e)o((aa))
(W(r,))
c)
B (x, [d]]
(wolx,c),)
+ (p(z)s((aaa) Volcion
wlx, 1) ( a) s) (2) o0]
+ Ep(a))
V (o(()l))
E Ep(c
w w(2,
(I, c)
+ Exzo((11))
(z.
+ VV(e)
E (mo(al=)
E Bp(ck 10 ((1
(x, c)
+ +Batn
(a, (x,)))s0)
(w(z, (ww(z,))) c)s (z, c)) (z, a) + (p(x))((aaa)) VP(clona) (w(,, c) (z, a a)s²²² (z. d]]]
[0 c) g (z, a)sg (x, (z,) ]]]
+ Pp()) (m " (z, c)s (x, 0 (z, [r,c]]
+ Ep(x)
Bp(c)z,0) VTo(alz)
where we rely on local correctness in the last line to use
Emp(olz) (p(elaaa) [w(z c)4 (x, a)) (x, e)) + Enga(een (z,'')o (x, = Eejr((ok)) (
and (2, (2,,))o((xx)))
In particular, in the case of deterministic and context-independent clustering, the variance can be simplified follows.
V "pco)molala)ptrllaa))) (((r,c.)(r--0 [utt, f(x, a))s (x, ca) + (ng (c, () (c, c.] -Bp(r)mo(aaa) (w((x,c))s (1) (z, ca (x, oz) az)
(z, ca) ]] + V)()) (agu F 2nd (x, c)s (z, oI] (x,)))
+ Ep(s) Woo(olz)
(z,
D.4 PROOF OF PROPOSITION 3.5
Proof. Since (al)))) [w c)] = (ap a.a) [w( a)] 1, the difference in the variance of the cluster and vanilla importance weights is attributed the difference in their second
22

---

Published as a conference paper at ICLR 2025
Table Hyperparameter search spaces used in the experiments. the hyperparameter for weight decay. is the learning rate. B is the batch size.
Datasets
Methods
[x(x) in Eq. (12
IPS-PG
{10 2, 10 1)
(10 5
(64, 128, 256)
(0.1 0.3
Synthetic
DR-PG
{10-2,101 100-,,00-.110 10-
{10-3, 5 10 10-4)
(64, 128, 256}
{0.1]A], 0.5
POCEM
10
5 10
128
IPS-PG
10 110
10 10
1,024
A
Real-World
DR-PG
[10 100",10-2
[10 10
1,024
A
POCEM
10
10 3
1, 021
moment, which is calculated below.
V
ootebrooeeeepoccceossllle a)] V (tceoooaaeppcccas//llcc..
Bp(s))o(aa))))ll...))) [w2 (z. a)] bx (x,
(a(a)mm(l))))(cl....)) 1 (I, a) (Eraaela.))((xt,
:. wa, Bx
Σ <o(alz)
Σ
p(c|z, a) (w (z, (1mla a))))
Bp(x)
LEC
To(a|æ) Σ xu(cla)mo(alz,e) (w² (wxxt a)])
<o(c2²) no (a
Bpp
To(a(()
To (a(aa)
LaEA
CEC
(c(a)ro(cla))
E E lw² (xx (x, - (wo
(p(a)mo(caa)
V (w
E ADDITIONAL EXPERIMENT SETUPS AND RESULTS
E SYNTHETIC EXPERIMENT
Detailed Setup. This section describes how define the synthetic reward function and perform hyperparameter tuning in detail. Recall that, in the synthetic experiment, we synthesized the expected reward function as
q a) g(x, ca) + ha, (x, a),
(16)
where we use the following functions as g (cluster effect) and (residual effect), respectively.
g(x, ca) Ibase (x, + [uT 2
Ta) 1.5}
d=1
3
10 ra) > 3.0} + uS na < 1.0),
+ U2 U2I
I.a < 0.5) - uzl {()
3=3
d=2
d=
he. (I, x Mc, one hot, + TT + T one hota;
where Fd is the d-th dimension of the context vector function from OpenBandiiPipeline as Abase (-..) and u....., 21 sampled from uniform distribution with range [-3, 3]. Mc. O Oz,c,, and 6a,c, parameter matrices or vectors sampled from uniform distribution with range 1, separately for each given action cluster Ca.
We
Based on the above reward function qa, a), we synthesized the logging policy To as exp(B q(x, a) + n2,a) En'e exp(3 g(x, a') + 72.a)
ma (r)
(17)
23

---

Published as a conference paper at ICLR 2025
Table 4: Dataset Statistics
Dataset
Asrain
Theat
[A]
EUR-Lex 4K
15,449
3,865
3,956
Wikil0-31K
14,146
6,616
30,938
where 3 is parameter that controls the optimality of the logging policy, ,3 is set to 5 as default and the noise Me,, is sampled from a normal distribution.
To summarize, first sample context and define the expected reward q(x,, a) as in Eq. (6). We then sample discrete action a from To based on Eq. (17) where action a is associated with cluster co- The reward is then sampled from a normal distribution with mean g(2; a.). Iterating this procedure times generates logged data D with n independent copies (x, O, r).
We tuned the weight decay hyperparameter, learning rate, batch size, and the number of irrelevant actions for variance reduction for the baseline methods (i.e., IPS- and DR (R-PG) using the test policy value, while advantage to the baselines. For all methods, we used Adam (Kingma & Ba 2014) as the optimizer and used neural networks with hidden layers to parameterize the policy.
fixed set of hyperparameters for POTEC shown in Table 3 giving unfair
E.2 EXPERIMENT ON EXTREME CLASSIFICATION DATA
In addition to synthetic and real-worle recommendation data, we performed OPL experiments on two extreme classification datasets provided by Bhatia et al. (2016).
Setup. Following previous studies (Dudik et al. 2014: Saito et al. 2021b Su et al. 2020a Wang
ct al. 2017), we transform the extreme classification datasets to contextual bandit feedback data with many actions. In classification dataset {( a:) J= we have some feature vector I; E X' and ground-truth label a; A, which will considered an action.
We consider stochastic continuous rewards where we define the expected reward function follows.
q(z, a)
if is positive label na otherwise
1
7]a
(18)
To is a noise parameter sampled separately for each action a from uniform distribution with range [0, 0.1]. After defining the expected reward function, we sample the reward from a normal distribution as .N(g(z,)) a), o²) with standard deviation a 0.05 for each data.
We define the logging policy To by applying softmax to an estimated reward function q(x, a) as
exp(3 q(z, a))
(a (x) = En'EA cxp(3 ((x,aa))
(19)
where we use 3 10 for both datasets. We obtain q(τ; a) by learning matrix factorization model where we use the test data recorded in the original datasets for obtaining a logging policy while we use the training data for performing OPL to make them independent.
Results. We evaluate POTEC against IPS-PG, DR-PG, and Reg-based under varying numbers of clusters evaluate POTEC's robustness the choice of this key byper-parameter. We optimize the hyperparameters of POTEC and the baselines based on the ground-truth policy value in the validation set, and the effectiveness of the OPL methods evaluated on the test set. For POTEC, we evaluate it with two types of clustering methods to investigate its robustness to the ways the clustering is performed. The first method through learning an action embedding via Lipschitz regularization (Lip) recently proposed for improving OPE in large action spaces (Peng et al. (023). The second method is to apply Agglomerative clustering (AC) implemented in scikit-learn (Pedregosa et al.. more accurate clustering by leveraging
2011) to the full-information labels, which provides the true reward correlation. Note that we perform a conventional reward regression rather than the two-step regression for POTEC due insufficient pairwise data in these specific datasets.
Figures 6 and 7 report the test policy value (normalized by ((TO)) of the OPL methods with varying numbers of clusters Eurlex-4K and Wiki10-31K, using two types of logging policies. For these
24

---

Published as a conference paper at ICLR 2025
Reg-based
IPS-PG
DR-PG
POTEC (w/ AC)
POTEC (w/ Lip)
weak logging
strong logging
Test Normalized Value
1.7
1.7
1.6
1.5
1.5
1.5
14
1.3
1.3
200
200
number of clusters
number of clusters
Figure 6: Comparing the test policy value the OPL methods (normalized by V (TT)) on the Eurlex-
4K dataset with weak and strong logging policies, respectively.
Reg-based
IPS-PG
DR-PG
POTEC (w/ AC)
POTEC (w/ Lip)
weak logging
strong logging
Test Normalized Value
202
30
1000
200
10
50
100
200
number of clusters
number of clusters
Figure 7: Comparing the test policy value of the OPL methods (normalized by V ((TO)) on the Wikil0- 31K dataset with weak and strong logging policies, respectively.
Note: The results are averaged over different sets of synthetic logged data replicated with different random seeds. The shaded regions in the plots represent the 95% confidence intervals of the policy value estimated with bootstrap.
experiments, we trained "weak logging" policy with (two times) fewer samples than the "strong logging" policy. We optimized the hyperparameters of POTEC and the basclines based on the ground-truth policy value in the validation set, and the effectiveness of the OPL methods evaluated on the test set. It should be noted that the baseline methods do not depend on action clusters, which results in flat lines in the figures.
The figures demonstrate that POTEC, with both clustering methods (Lipschitz regularization; Lip and Agglomerative clustering; AC), typically outperforms all baseline methods across range of numbers of clusters, indicating its potential for real-world applications even with action clustering learned only from observable logged data (i.e., POTEC w/ Lip). We can also see that POTEC with a more accurate clustering (i.c., POTEC w/ AC) slightly outperforms POTEC w/ Lip. implying an even better potential of POTEC with more refined clustering procedure. The regression- competitively with POTEC only for strong logging policy on the Wiki10-31K dataset, but we can see, all other scenarios, POTEC typically performs the best. We also compared the one-stage and two-stage variants of POTEC on the real-world datasets, but we did not find significant difference between them for both types of clustering.
method performs
25


Reviews:
[{"review_id": "4Cpd2x3U2u", "reviewer": "Reviewer_J9Lw", "summary": "This paper introduces POTEC, a novel off-policy learning (OPL) method designed to handle large action spaces, a common challenge in real-world recommendation systems. The algorithm decomposes the policy into two stages: a first-stage policy that selects promising action clusters, and a second-stage policy that chooses specific actions within the selected cluster. The paper provides theoretical analysis showing that POTEC achieves unbiased gradient estimation under a local correctness condition. It also shows its cluster-based importance weighting leads to substantially lower variance compared to traditional methods. The authors present a complete learning framework, including how to optimize both policies through a combination of policy-based and regression-based approaches. Through extensive experiments on both synthetic and real-world recommendation data, POTEC demonstrates superior performance compared to existing methods such as IPS-PG and DR-PG, particularly when dealing with sparse action observations in logged data and large action spaces.", "questions": "How test data is constructed in empirical study (both synthetic and real world experiment)? For example, is there overlap of context between test data and training data, if so how much overlap?\n\nWhat will happen if full cluster support condition is not met in logging policy? In the real world, how do you check if this condition is met?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper introduces a novel two-stage approach to address the challenges of large action spaces in off-policy learning.\n\nIt demonstrates high quality through rigorous theoretical proof of the unbiasedness of the POTEC gradient estimator under local correctness condition, and shows how cluster-level importance weighting effectively reduces variance.\n\nThe paper provides clear theoretical justification for the two-stage decomposition and progresses logically from the problem formulation to the solution.\n\nthe paper addresses fundamental challenges in large-scale off-policy learning by providing a unified framework with theoretical guarantees for the joint optimization of two-stage policies.", "weaknesses": "The paper did not specify how different clustering algorithm could impact the performance of POTEC from both theoretical point of view and empirical study.\n\nThere could be more details provided on how the test data is constructed in empirical study.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oLCztoJ0cv", "reviewer": "Reviewer_nd8T", "summary": "The authors derive a novel theoretical framework for off policy contextual bandits that strictly generalizes the complementary approaches of policy based and value based methods via a two stage decomposition involving both. It is argued that this provides a non-trivial advantage in large and structured action spaces. The two stage decomposition involves first using a policy based parametrization to identify an action cluster, followed by a regression based approach within each cluster to discriminate between the fine grained individual actions.", "questions": "- Clustering the action space based on the scalar embedding of averaged marginal rewards seems somewhat surprising and not justified even informally -- can you comment on this? Have you considered a random bucketing strategy as a baseline for the first stage, and if so, how much better does proposed method do over this baseline?\n\n\n- Is there a typo in Equation (3)? -- should the third variance term on RHS be $\\hat{q} (x, a) s_\\theta (x, a)$ ? For easier verification, please consider adding a quick derivation of this to the appendix.\n\n\n- Is there a fundamental reason for the theory to only work with the policy at the base stage and value at the second stage? A footnote on page 4 is discussing this, but the implication is unclear for whether this is a fundamental difference. The point about the importance weights being an issue doesn't seem convincing since the same reasoning applies to the first stage in the current proposal.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Novel general approach for combining policy and value based learning methods in a hierarchical fashion so as to avoid bias and variance simultaneously.\n- The paper is very well written and a pleasure to read\n- Good discussion of related literature and its context for the current work.", "weaknesses": "- A good bit of the technical novelty in the current paper is shared with  prior work [Saito et. al. 2023] which the authors duly point out. Having said, that the prior work applies a similar idea for off policy estimation while the current work deals with policy optimization.\n\n- No satisfactory general recipe to perform the right kind of clustering that can best exploit such a two stage decomposition.\n\n- The main issue with a regression based approach is non-zero bias due to model capacity as the authors point out,  and in practice this error seems more plausible to have its source in an inability to fit within cluster, if at all, rather than across clusters. However, this is precisely what is assumed for a good performance in the second stage. If it is assumed to do sufficiently well in the second stage, wouldn't this somehow make the problem more likely to achieve a low error for a simple supervised learning global fit?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qWQwA19Dhz", "reviewer": "Reviewer_5MGC", "summary": "This work studies the problem of off-policy learning for contextual bandit policies in large discrete action spaces. The authors propose an approach that decomposes the overall policy into two components: a cluster-selection policy and a conditional action-selection policy. To optimize the cluster-selection policy, they derive a policy gradient, while a pairwise regression procedure is used to learn the conditional action-selection policy. The authors provide a thorough analysis of the bias and variance of the proposed policy gradient estimator, which makes the work more solid and convincing. The authors perform comprehensive numerical studies to demonstrate the performance of the proposed algorithm compared to several baselines.", "questions": "1.\tThe selection of the number of clusters is a crucial aspect of the proposed method, as illustrated in the numerical studies. The author could provide more details on how to select the number of clusters in the real-world applications.\n2.\tPlease provide more details on how to incorporate context-dependent clustering into the algorithm.\n3.\tThe authors could compare the proposed method with Saito et al.,2023, highlighting performance differences and potential advantages/disadvantages.\n4.\tPlease provide more details on data generating mechanism in the real-world experiment.\n5.\tThe authors could consider mentioning potential limitations of the proposed method.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The authors study an important problem of off-policy learning for contextual bandit policies in large discrete action spaces.\n2. The theoretical analysis of the bias and variance for policy gradient estimator makes the work more solid and convincing.\n3. The authors perform comprehensive numerical studies to demonstrate the performance of the proposed algorithm", "weaknesses": "The clarity of the writing could be further improved by providing more details on certain aspects, such as practicability of the algorithm and experiment settings. See Questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cF3PI5qJkO", "reviewer": "Reviewer_p75W", "summary": "This paper provides a solution to off-policy learning for large action spaces through a two-stage policy. The action space is first clustered with action embedding (reward). Then, the first-stage policy selects an action cluster, and the second-stage policy selects an action within the cluster. The two-stage policy is updated with the policy gradient method, providing a bias and variance analysis of the gradient. They show that the bias of the gradient is controlled by the relative difference in the reward estimation and can be further reduced with a pairwise regression if possible. Also, they show that the variance of the gradient is significantly smaller than the gradient of the original action space (which is straightforward). Empirical studies demonstrate its effectiveness.", "questions": "1, How do you choose the number of cluster C? \n\n2, how do you get the embedding of action a? Is it from the reward estimation? (more details of the clustering is beneficial)\n\n3, do you assume the logging policy $\\pi_0$ known? What if it is not known?\n\n4, Is the logging dataset sampled based on the ground-truth cluster? If I understand correctly when creating the logging dataset, you first cluster the action based on the reward of $(x,a_i)$, and then you have a reward for both cluster r(x,c_a) and r(x,a|c). What if the reward is not in this hierarchy way? For example, if the reward function is just defined as r(x,a) instead of the reward of cluster + reward of the action. Or, in other words, what if the logging policy is not defined as Eq.(11)? I believe that this algorithm might benefit from the logging policy and reward definition while it might not work that well in a more general logging policy. Because the logging policy itself is somehow a two-stage policy.\n\nI am willing to further raise my score if these questions are addressed.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1, propose a novel new two-stage policy for OPL for bandits with large action space and the intuition/motivation is well-justified.\n\n2, provide a variance and bias analysis of the gradient of the proposed algorithm and connect with the algorithm. Also, the proposed gradient has a smaller variance than the IPS method. \n\n3, the paper is well-written, solid, easy to follow.", "weaknesses": "1, The paper did not provide very detailed description of the clustering algorithm. More details will be preferred. \n\n2, The paper did not discuss how the number of the cluster 'C' can affect the performance of the algorithm.\n\n3, They did not experiment with other logging policy settings. \n\n4, Others, see questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)