Wasserstein Distributionally Robust Policy Evaluation and Learning for Contextual Bandits
Authors: Yi Shen
Pan Xu
Michael Zavlanos
Venue: Accepted by TMLR ()

Abstract: Off-policy evaluation and learning are concerned with assessing a given policy and learning an optimal policy from offline data without direct interaction with the environment. Often, the environment in which the data are collected differs from the environment in which the learned policy is applied. To account for the effect of different environments during learning and execution, distributionally robust optimization (DRO) methods have been developed that compute worst-case bounds on the policy values assuming that the distribution of the new environment lies within an uncertainty set. Typically, this uncertainty set is defined based on the KL divergence around the empirical distribution computed from the logging dataset. However, the KL uncertainty set fails to encompass distributions with varying support and lacks awareness of the geometry of the distribution support. As a result, KL approaches fall  short in addressing practical environment mismatches and lead to over-fitting to worst-case scenarios. To overcome these limitations, we propose a novel DRO approach that employs the Wasserstein distance instead. While Wasserstein DRO is generally computationally more expensive compared to KL DRO, we present a regularized method and a practical (biased) stochastic gradient descent method to optimize the policy efficiently. We also provide a theoretical analysis of the finite sample complexity and iteration complexity for our proposed method. We further validate our approach using a public dataset that was recorded in a randomized stoke trial.

Published in Transactions on Machine Learning Research (01/2024)
Wasserstein Distributionally Robust Policy Evaluation and Learning for Contextual Bandits
Yi Shen Duke University
yi. N.shenn77@@duke
Pan Xu Dake University
pen.zw@duke. edu
Michael M. Zavlanos
michael.zaelanos@Bal edu
Duke University
Reviewed OpenReview: https:
openreview. https:://peerreviaa
net/ forum? id= Nmp JDHWIug
Abstract
Off-policy evaluation and learning concerned with assessing given policy and learning an optimal policy from offline data without direct interaction with the environment. Often, the environment in which the data are collected differs from the environment in which the learned policy is applied. To account for the effect of different enviromments during learning and execution, distributionally robust optimization (DRO) methods have been developed that compute worst-case bounds on the policy values assuming that the distribution of the new environment lies within an uncertainty set. Typically, this uncertainty set is defined based on the KL divergence around the empirical distribution computed from the logging dataset. However, the KL uncertainty set fails to encompass distributions with varying support and lacks awareness of the geometry of the distribution support. As a result, KI approaches fall short in addressing practical environment mismatches and lead to over-fitting to worst-case scenarios. To overcome these limitations, we propose a novel DRO approach that employs the Wasserstein distance instead. While Wasserstein DRO generally com- putationally more expensive compared to KL DRO, we present regularized method and practical (biased) stochastic gradient descent method to optimize the policy efficiently. We also provide theoretical analysis of the finite sample complexity and iteration complexity for our proposed method. We further validate our approach using a public dataset that was recorded in a randomized stoke trial.
1 Introduction
Contextual bandits [Langford & Zhang 2007) are class of decision-making problems where a learner repeatedly obscrves a context, takes an action from a set of candidate actions, and receives a cost for the chosen action. This cost signal often referred to as bandit feedback since no costs for the unchosen actions are observed. The goal for the Icarner to minimize the cost by sclecting the best action for cach context. Many problems in practice can be modeled as contextual bandits, ranging from online news recommendation [Li al. 2011), advertising (Bottou et al. 2013) to personalized healthcare (Zhou et al. 2017
In many high-stakes applications, such as healtheare, it is generally unsafe and not recommended to interact with the environment directly in order to collect data and evaluate or learn desired policy. Instead, in these environments, observational data are often available, that have been collected by safe or risk-averse policies that depend on different contexts, e.g,, policies that assign placebos to high-risk patients when testing new drugs. In this paper, we study off-policy evaluation (OPE) and off off-policy learning (OPL) using observational data DDddi et al. 2014). OPE estimates the average cost. of following given policy, while OPL. aims to

---

Published in Transactions on Machine Learning Research (01/2024)
Distribution
0.200
Distribution
0.10
0:175
0.150
Robabiity
Probebiitty
0:125
0.06
0.100
0.01
0.075
0.050
0.02
0.025
0.00
0.000
40
45
55
50
10
Risk index
Patient age distribution
(b) Risk index distribution
Two artificial datasets that represent patients' contextual distributions, where æ axis is the context support and y axis is the probability. In (a), the distributions P and Q represent the patients' ages and are only different at the ages 55 and 60. The KI. divergence is KT( = +oo since the two distributions have different supports, where the KL divergence between two discrete distributions P and Q is defined as In (b), the distributions Q1 and Q2 represent the patients' risk index, the higher the worse. They are equally distant from P under the KT. divergence, i.c., (T((QI] = KL(Q2||P) 0.21. However, the distribution Q1 represents a less challenging environment compared to the nominal distribution P the possibilities of encountering patients with higher risks in QI are smaller than both P and Q2. Q2 represents similar environment P and is indeed closer to P than Q1 under the Wasserstein distance, i..., W(P,0) = 2.07 and W(P, Q2) = 1.42. See equation 2.2 for the Wasscrstein distance definition.
Figure 1:
KL(Q|[P) = E=Ex
find the optimal policy that achieves the minimum cost. Generally, both OPE and OPL assume that the environment in which the observational are collected (training set. and the enviromment. in which the trained policy is deployed (testing set) are identical. This means that the distributions of both contexts and costs do not change. However, in practice, this assumption often violated. For instance, in clinical trials (Finlayson et al. 2021: Imai & Ratkovic, 2013). treatments that are tested and found effective using Randomized Clinical Trials (RCT) are not necessarily equally effective in local hospitals since the patient populations can be significantly different; RCTs use stringent inclusion and exclusion critcria to seleet the patient population while local hospitals reflect the local care environment, determined, e.g., by local social determinants of health. These differences between the training and testing environments are termed distribution shifts.
To address the effect of distribution shifts on OPE and OPL, distributionally robust optimization (DRO) methods have been recently developed (Si el al. 2020: Kallus et al., 2022 Mu el al. 2022), that aim to obtain worst-case bounds on the true OPE and OPL costs by assuming that the testing distributions lie within uncertainty set around the training distributions. This way, over-optimism in unknown environments is avoided. Si et al. (2020) proposed importance sampling methods to solve both OPE and OPL problems, assurming that the behavior policy used to collect the observational data is known. They also provided an asymptotic convergence result, demonstrating that the proposed estimator converges to the true values at a rate of Op(nn1/2), where n is the number of data samples. Kallus et al. (2022) further proposed doubly robust methods that relax the assumption of knowing the behavior policy and achieve the same convergence rate, even when the unknown behavior policy the previous work (Si et al., 2020; Kallus al. 2022) is that potential distribution shifts appear across both context and cost distributions. Mu et al. (2022) argue that this assumption can lead to conservative performance, and the DRO objective may degenerate to non-robust solution, especially when features are continuous and the loss function is binary (Hu et al. 2018). To address this issue, Mu et al. (2022) proposed a factored DRO formulation that treats distributions shifts in contexts and costs separately. This approach provides non asympttpiooic sample complexity results.
non-parametrically estimated. A common assumption in
All previous works consider the distribution shifts under Ballback-Leibler divergence (KL), i..., the size of the uncertainty set is measured using the KT, divergence. DRO problems based on the KT. divergence enjoy
2

---

Published in Transactions on Machine Learning Research (01/2024)
simple dual formulation with scalar dual variable and strong duality holds under mild conditions; see Hu & Hong (2013) for details. However, a KT. uncertainty set requires a strong assumption on the testing distributions, namely, that they are absolutely contimuous with respect to the training (nominal) distribution. This assumption excludes distributions with different supports compared to the training distribution and can be often violated in practice; see Figure [(a] for an example. In addition, using the KL. divergence it is not straightforward
obtain an initial estimate of the radius of the uncertainty set., One approach to determine an initial radius of the uncertainty set is to divide training set into two subsets and compute the distance between them. This approach is not possible using the KL divergence due to potential differences in the supports of these subsets, that will result to the radius being infinite. Finally, it is important to note that KL divergence cannot capture the geometric differences between distributions. For example, it is possible that two distinct lesting distributions have equal KL distances from the training distribution, but one of them represents a more significant distribution shift compared to the other. A visual illustration of this scenario can be observed in Figure 1(b] From a computational perspective, when the training set is large, stochastic gradient descent. (SGD) appears to be a promising method to solve KL DRO problems, especially OPL problems since as it allows to iteratively search for the best policy. Yet, as mentioned by Namkoong & Duchi (2016), SGD generally fails to converge because the variance of the objective and its subgradients explodes as the dual variable approaches O. The above limitations motivate the study of DRO problems using different measures to define the uncertainty set.
Two popular measures, alternative to the KL divergence, that have been used to characterize the uncertainty sets in DRO literature are uncertainty sets based on moment constraints (Delage & Ye. 2010; Zymler et al., Yang et. al. 2023) and uncertainty sets defined by the Wasserstein distance (Moha Esfahani & Kuhn 2018 Zhao & Guan 2018; Gao & Kleywegt, 2022). Moment constraints generally roquire the testing distribution to have similar moments to the training distribution, such as means and variances (first and second moments). Finite order moment constraints can be extended to infinite orders and have been explored in the context of Kernel DRO (Zhu et al.] 2021), where the kernel function defines a. metric between distributions, known as the maximum mean discrepancy (MMD) (Gretton et al., 2012). However, it has been observed that in optimal transport problems (Feydy et al., 2019). the gradient vanishes when evaluating the distance between two distributions under the MMD distance, leading to slow convergence rates in practice. Detailed examples and discussions can be found in [Fyydy (2020, Chapter 3.2). In contrast, the Wasserstein distance has been shown to result in fast and stable convergence in many optimal transport and machine learning problems (Frogner et al., 2015 Arjovsky et al.] 2017]. The Wasserstein distance considers the underlying geometry of the distributions, including their supports, and therefore allows to compare testing and training distributions with different supports, something that KL methods are not able to do. While Wasserstein DRO has been used to quantify uncertainty in various learning tasks including classification and regression (Kuhn et al. 2 its application to OPE and OPL problems is mostly unexplored yet with a few recent exceptions (Kido, 2022 Adjaho & Christensen 2022). Though Kido 2022) measures the distribution shifts by Wasserstein distance, this work assumee that the context distribution of the testing set is absolutely continuous with respect to the context distribution of the training set, i.c. the same assumption as in methods that rely on KL (Si et al., 2020 Mu et al., 2022). This assumption simplifies the dual problem but excludes potential shifts of different supports. The approach in Adjaho & Christensen (2022) considers OPL problems with binary actions and allows for distribution shifts of different supports; simple and interpretable policies are studied due to the assumptions that actions are binary, e.g., a threshold policy. However, finite sample analysis is not the focus of this study.
2013
In this paper, we focus on OPE and OPL for contextual bandit problems under distribution shifts. Specifi- cally, we assume that shifts occur in both the context and the cost distributions for each context-action pair and employ Wasserstein DRO to develop worst-case bounds on policy values. For OPE problems, we first calculate upper bounds on cost values that account for cost shifts in each context-action pair. These bounds are then used to determine upper bounds on the policy values that account for context shifts. We provide a finite sample analysis of our proposed method and demonstrate that the policy value obtained from a dataset of size n converges to the true value at a rate of (n-12)) Compared to OPE problems that rely KL DRO (Si et al.) 2020: Kallus et al., 2022 Mu et al. 2022 the use of Wasserstein DRO introduces an inner maximization sub- aro-problem, that increases computational cost. To overcome this challenge, we propose a regularized DRO method that approximates the inner maximization sub-problem using smoothing. For
3

---

Published in Transactions on Machine Learning Research (01/2024)
OPL problems, we demonstrate that both Wasserstein DRO and regularized Wasserstein DRO methods converge to the optimal policy at the rate of Op(n-1/2). The regularized DRO method improves the query complexity of the Wasserstein DRO method in finding a 8-accurate solution by an order of O(8-1) when the OPL problem is convex. Since the query complexity of both Wasserstein DRO and regularized Wasserstein DRO methods depends on the size of the support of the distributions, we also propose a (biased) SGD method that can find the optimal policy with query complexity independent of the distribution support size. It worth noting that previous works (Si et al. 2020 Kallus et al., 2022 Mu al.) 2022) do not provide explicit query complexities. In our study, we validate the proposed methods using randomized controlled trial dataset that study the effects of drug treatment on acute ischemic stroke, further demonstrating their effectiveness.
2 Problem definition and preliminaries
Denote X C R" the set of contexts and A = fai....,,0 the set of discrcte actions. Let Dn == [(x, A,, Y.)) be the observational dataset of size n. We assume the datasel is generated as follows. The contexts {Xi]i= are independent and identically distributed (i.id.) samples from a fixed context dis- tribution po. Actions {Ai]==z are selected according to context-dependent behavior policy no: i.e, A, is sampled from the distribution To(-(X,), After taking the action A,, a random cost Y, is announced and it is an i.i.d sample from a fixed cost. distribution P.....
Assumption 2.1 We make the following assumplions on the behavior policy To, the cost functions and the contezts.
Sufficient exploration: To(A|X) 2 Mmo > 0 for all contert action poirs;
Discrete and bounded costs: Vx E X. E A, Bza is discrete, finite, and knoun. Moreover, S (e, S Ymax: VY E For simplicily, we assume Ez,a = for all (x.a) E X A;
Discrete conterts with positive mass: the contect set X is a finite and known discrete set and the probabilily mass of each context is bounded from below, i.e., there exists an pro such that P(X z) Mao for all E X.
The first two assumptions in Assumption 2.1 are standard in non-robust OPE and OPL problems (Dudik et al., 2014; Wang et al., 2017) and are adopted in the KL DRO problems (Si et al. 2020 Kallus et al. 2022 Mu et al. 2022). The last assumption together with the first assumption in Assumption 2.1 ensure that the observational dataset contains sufficient samples across all context and action pairs, which is also assumed in Mu et al. (2022). We further make the following assumption on the random costs Y.
Assumption 2.2 For each contert action (z, paár, are given measurable function y2,a (E) : - and E C Ih" is known. The observational dataset is provided in the form of D, := {(X,, Ai, E,, Fo where & is an i.i.d sample from fixed distribution po Ex.,A. and the cost Y; = (X..A((S)) is then calculated
from the given function Va.a-
Assumption 2.2 holds in many applications. For example, in clinical trials, E can represent a patient's 's health condition after the treatment, e.g., blood pressure, which is random and can be monitored; a cost function is defined according to this health condition. Alternatively, in news recommendation systems, & can represent a user's behavior after receiving batch of news, c.g., clicking or not, time spent on the page. A cost function can be designed to encourage specific user behaviors, e.g., s(click) -1 and ynoo click) = 1.
If the random cost Y is observed directly from the environment, one can set (e) = & which is identity function mapping the random cost to itself or set Vz,a({) Ynax for all & that are not observed in the observational dataset. We note, however, that our analysis here applies to general cost functions that only need to satisfy Assumption 2.2 We note also that Assumption [2.2] subsumes the special case when VC) = v(E) for all (x, a) pairs and given cost function v(E), ie.,, the case when the cost function only depends on the random observation &, as the examples discussed above. Thus, knowing the cost function does not imply knowing the relations between different (x.,) a) pairs.
4

---

Published in Transactions on Machine Learning Research (01/2024)
We assume distribution shifts happen both in the distribution of contexts and the distributions of the cost functions. Specifically, denote by Pr the distribution of the contexts in the testing environment and assume that P& is contained in a uncertainty set U(Ezi P2) that is centered around the training distribution Pe i.c., Ps E [(cx; = {P D(P, P2) & & where ez > and D is a distance metric to be specified later. Similarly, for each context and action pair (r,a), we assume the distribution of the cost in the testing environment satisfies Pz,a E(ce: Pi,a) where 6e > 0.
Due to the possible distributional shifts in Pe and P 2.a: want to be conservative when evaluating learning new policy in the testing environment. To do so, focus on the worst-case value of the new policy obtained over all possible distributions within the uncertainty sets.
Definition 2.3 For given Co.fe > 0 and policy n, define the distributionally robust policy value V(*) as:
P.... M] EEtreE
V(T)
sup PP)
Ex (awr(lla)
sup
(2.1)
Remark 2.4 The formulation in equation 2.1 implies that given any two contert and action pairs (x,, a), (x',a') and their truining cost distributions Pe..r aa the testing cost distributions Pa.a Pr are independent. Thus, the inner mazimization problem is not coupled across conterts or actions and can be solved individually for each context and action pair.
In this paper, we employ the Wasserstein distance as a distance metric between two probability distributions P and Q, defined as
D(P,Q) == inf (B(e.c)aa (c(5, () 01 P..02 = =
(2.2)
where <) is a continuous function that captures the cost of moving the point to Ç. In what follows, we assume <) (E <)² equation 2.2 is the coupling distribution with its first marginal on and second marginal d2 being P and Q, respectively. The optimal coupling distribution a provides the best transportation plan that moves P to Q with respect to the cost function c(:;).
We now present some preliminary results on the DRO using the Wasserstein distance. Specifically, given a nominal distribution pe with support X, a loss function f(r), and a uncertainty parameter e, the (primal) Wasserstein DRO problem defined as
(P) = sup EgeP [f(x)], Pell(e;
(2.3)
where (P inf, (c(E, () o :0 <) and M represents the set of all probability measures. Observe that solving the primal DRO problem is challenging since the optimization variable is a measure function and the expectation is taken with respect to an unobserved distribution P. This difficulty can be avoided if the DRO problem 2.3 is studied in its dual form, where the optimization variable becomes a scalar. In fact, Zhao & Guan (2018) showed that, under mild conditions, strong duality holds, i.e., zero duality gap; see Proposition 2 thereof for the details. The dual Wasserstein DRO problem can be formulated as
(D) = inf X20
eA + Px sup (f(() - A(2 M
(2.4)
CEx
Indeed, in this dual problem, the optimization variable is a scalar and the expectation is taken with respect to the given nominal distribution p0. Recent literature policy evaluation has studied this dual DRO problem (Si et al. 2020: Mu et al. 2022). but using the KI. divergence to characterize the uncertainty set Hu & Hong 2013). In this case, the dual KL DRO problem bocomes:
(DKL.) inf {ex + XIn (exp A20
(2.5)
As discussed before, compared to the Wasserstein DRO problem (D), the KL DRO problem, (DKL) tends up-weight the samples with high costs and may over-fit to outliers since it is not aware of the geometry of the
5

---

Published in Transactions on Machine Learning Research (01/2024)
distribution. See Appendix E for an example that compares KL DRO and Wasserstein DRO. Nevertheless, the Wasserstein dual problem (D) is computationally more challenging as it involves an inner maximization problem.
The following lemma shows that if f()) bounded, the optimal solution to the dual problem (D) is finite. This avoids limit arguments the proofs that involve the dual variable X.
Lemma 2.5 Consider the dual problem in equation 2-4
1f0 f(x) S mmax for all E X, then the oplämal solution equation 2.4 satisfies X [0, fmax/ and the optimal value of equation 2.4 attained at X satisfies D fmax.
The proof is provided in Appendix
B
3 Wasserstein DRO for policy evaluation
with the uncertainty set described using the In this section, we focus on the OPE problem equation Wasserstein distance metric.
Proposition 3.1 Let Assumption 2. hold. The policy evaluation problem equation 2.1 is equivalent to the
following two-step dual problem.
V(T)= inf { End + Ben sup X20
(Eawa(Is) - }
(3.1)
<EX
{ {ead pe sup (Vr, - E (3.2)
where m(x, a) inf X20
(E
Proof The result follows by applying the duality as equation 2.4 to both contexts and costs.
Recall that Assumption 2.2 poses extra condition on the cost functions. From the dual problem perspec- tive, we observe that this assumption is necessary for addressing distribution shifts in Wasserstein DRO the dual problem equation 3.2 contains an inner maximization problem, which requires the knowledge of the function values for all realizations. Additional assumptions on the cost functions are required when applying KL DRO to OPE and OPL problems, where the probability mass function P(Y = y|æ,a] a) is lower bounded for y across all context and action pairs (Si et al., 2020 Kallus ct al.] 2022 Mu al. 2022). Similarly, this lower bound assumption ensures that the dual problem under KL well defined.
Note that, in general, it is intractable to directly solve equation 3.1 and equation 3.2 since these problems involve an inner maxinization problem that is hard to solve. The inner maximization problem can be solved efficiently assuming, e.g., that X' and E are discrete. Besides, the distributions Pe and P0 in equation 3.1 and equation 3.2 are unknown and can only be estimated from the finite sample dataset Dn..
Given the observational dataset Da, denote po and pd and {Po bre., respectively. We define the empirical policy evaluation problem as follows:
the empirical distribution estimates of
((T) inf { (Eg) + A20
Ao sup (Euwn(1s) (mi(S, a)) X(x -xxx-))) -))) CEX
(3.3)
where mix a) inf {ed B X
Pg
sup CEEu.
(Sn,a(S) - ( [) Vr EX,
(3.4)
The goal of the OPE problem to estimate the policy value V(z). As we can only estimate the empirical policy value (T) using the offline data in practice, in the following theorem, we provide finite sample analysis of the gap between (Tr and V(r)).

---

Published in Transactions on Machine Learning Research (01/2024)
Theorem 3.2 Given an evoluation policy n, let Assumptions 2.1 and 2. 2 hold and select the total number of samples as 2 2 2log(4))/]A///))/( (Hoo )2. Then, with probabilily 1 5, the following inequality holds,
2 log (2/0) (())) < maax (v
X
log (A|/8)
2
+
+
n
n
1ppooaa Paro
log (11)) =0
X + E
O
(3.5)
n
Theorem 3.2 shows that the policy evaluation error decreases at a rate of n1122 Note that Mu et al. (2022) consider a similar problem using the KL divergence to characterize the uncertainty set. In this case, they show that the policy evaluation error satisfies (V(T) V(=) ((V(log(n) log ((2()(//0)/(<<r)),, where 6 is the KL uncertainty set radius.
The proof is provided in Appendix
Our error bound in Theorem 3.2 contains an extra term O ([v( | that is due the use of Wasser- stein distance and is directly related to the total variation of the empirical discrete distribution compared to the true discrete distribution. This error incrcases if the size of the supports of X and E increases as the Wasserstein distance is a geometry-await metric. Comparing the first term in equation 3.5 to the KL error bounds provided in Mu et al.] (2022), we observe that our result. improves the error bound by an order of Vlog(n). We believe that the extra vlog(n) term in the error bound in Mu et al. (2022) is caused by the corresponding analysis rather than the different characterizations of the uncertainty set; the KL error bound can be improved by using a similar analysis in Lemma 5 in Zhou et al. (0221] 1 which studies distributionally robust reinforcement lcarning problem using KL uncertainty set.
Note that the KL error bound additionally depends on the radius of the uncertainty set c,, which is due to the fact that the dual variable A multiplies the log- eegeexeccttion term as in equation 2.5 This KL error bound increases when the radius of the uncortainty set decreases. Although this error bound is always upper bounded by Ymax, it becomes uninformative when the distributional shift is small. This is because O( ((V(log(n) log (2z1[A1//)/((Enn)) > Ymax when Eg - O. On the other hand, this KL error bound decreases the radius of the uncertainty set increases. This is because when Co is large, the optimal dual variable X in KL dual problem is small and thus the value of Aln Po- (exp(f (x) converges to the highest cost. In contrast, our proposed Wasserstein error bound in equation 3.5 does not depend on the uncertainty radius.
Concurrently with the method proposed here, Xu et al. (023)) studied a DRO problem for a robust reinforce- ment learning problem also using the Wasserstein distance. However, in Xu et al. (2023) the cost function is known and deterministic while here we consider stochastic costs with distributional shifts. Besides, the error bound provided in Xu et al. (2023) is radius-dependent and will not be informative when the radius is small, similar to the KL formulation in Mu et al. (2022) discussed above. We believe that this dependency is due to the covering number type of arguments in the analysis, which is designed for improving state de- pendency in the sample complexity results. A reinforcement lcarning OPE problem using Wasscrstein DRO has recently been studied in Wang et al. (2021), assuming that the behavior policy that collects the training set is known or can be estimated accurately. This known behavior policy is directly used in the primal and dual formulations. More importantly, to achieve zero duality gap, extra assumptions on this behavior policy are required as statod in their Theorem in Wang et al. (2021). These extra assumptions cannot be verified for the problem considered here that only assumes knowledge of the finite training set..
In practice, Ee can be selected proportional to the sample size to capture uncertainties caused by finite samples [Ramdas et al. 2017), e.g., c = ((nn1/2). Then, the policy evaluation error decreases at the rate of n 1114 under KI. divergence while the rate under Wasserstein distance in equation 3.5 remains the same of n 11/2.
4 Regularized Wasserstein DRO for policy evaluation
The empirical policy evaluation problem under the Wasserstein distance can be reformulated as a lincar program (LP) and thus can be solved using commercial linear programming solvers, e.g., Gurobi (Gurobi Optimization, LLC 2023). However, as often discussed in much of the optimal transport literature (Cuturi
7

---

Published in Transactions on Machine Learning Research (01/2024)
2013 Genevay et al. 2016 Weed 2018 ) calculating the Wasserstein distance between two distributions is computationally expensive when the support set is large, e.g., when |X| or | is large. In this case, entropy-reullarized methods have been proposed to reduce the computational complexity of optimal trans- port problems. Inspired by these methods, we simplify the dual problem in equation 2.4 by replacing the maximization operation with a smooth "sofumax" function, egg the log-sum-exp function. Specifically, we define the smoothed dual problem as:
(4.1) where 17 > 0 is a hyper-paranee parameter that controls the distance between the "softmax" and the maximum. The primal problem of the smoothed dual problem (Dy) is indeed DRO problem but with entropy-regulariieed objective function. See Appendix D for the duality connection.
(D,) = inf Eewpg A20
[ log (xxxX
Although the smoothed dual problem (Dn) and dual problem DKL in equation 2.5 may look similor each other, they are in fact fundomentally different. Specifically, in DKL, the dual variable X directly affects the exponential function inside the expectation, e.g., a small up-weights regions of X with high costs f(x). On the other hand, in (Dn), oi is a pre-specified tuning parameter and the dual variable X is direclly related to the geomelry of the domain X, e.9., for a given x, large X only up-weighls regions of X with high costs that are close to x. In addition, (D₂) simplifies the theoretical analysis since the "softmax" function in (D₂) is not be unbounded when X goes to O. To avoid this technical challenge, Mu et al. (2022) instead assume that the X in DKL is bounded from below.
Remark 4.1
We define the empirical policy evaluation problem using the regularized dual problem as follows:
(T( inf (EEy + E Pa [1/n log / -(x-)))))(((-)))))))) X20
(4.2)
where m,, a) inf & + Bg [1/7 log (Eçez... (0 ()))] (
(4.3)
In the following theorem, we provide a finite sample analysis on the gap between V., and V for the empirical policy evaluation problem equation 4.2 and equation 4.3 that uses the regularized dual method.
Given an evaluation policy n. let Assumptions 2.1 and 2.2 2 hold and select the total number Theorem 4.2 of samples as n 2 2 log (1/44)/1////8)////maaa
Then, with probability 1 6, the following inequality holds,
(x) < Ymax
( 2lg (2/8) /n+ vxxlnn+ V 44o (41x)) + (2peo/(Mpsoola))
+ +log(121)/m
Theorem 4.2 implies that the entropy regularized OPE problem can find the policy value (T) almost efficiently as the un-regularized OPE problem in equation 3.3 and equation 3.4. The error term log()>)))/ is caused by the approximation of inner maximization by the "softmax".
5 Wasserstein DRO for policy learning
In this section, we consider the OPL problem using the Wasserstein distance to characterize the uncertainty. Let II be the policy space. Define the optimal distributionally robust policy म* arg min V(T) with optimal value V* VTT*). Similarly, define the empirical optimal distributionally robust policy ** arg min (T)) with optimal value V = v(*). In the following theorem, we provide finite sample analysis of the gap between V and V=
Theorem 5.1 Given a learning policy space II, let Assumptions 2.1 and 2.2 hold and select the lotal number of samples n > 2 [lgg(4/\(]|A ( ). Then, with probability 1 8, the following inequality holds,
V - -V < Yiaox (v 2 2log (2/8)/n + vixl/n + V 44log (41x)1111/8) (mppoH
+ (2
8

---

Published in Transactions on Machine Learning Research (01/2024)
The proof is provided in Appendix C.1 Note that the results of Theorem 5.1 can be easily extended to the entropy-regnlarized problem described in Section B
Theorem 5.1 assumes that we can solve the optimization problem exactly without specifying how. In practice, if we parameterize the policy π by 0 E +, we can use first-order methods to find the optimal policy. Specifically, denote [(0,5) = Bowrr((-c)) [ Then, we can define the OPL, problem for the Wasserstein and regularized Wasserstein DRO formulations follows:
((T))= = inf
{B
sup
(1(0, - A(x - 5)2 + end)]}
(5.1)
0,220
çEX
(T((T") = inf
E
1 log
Σ SEA
1
er (100))) + teaj..
(5.2)
0,220
17
x
In general, the regularized problem in equation 5.2 is preferred to equation 5.1 in solving for the optimal solution. For example, consider a specific where P is a point mass at zu- Then, equation 5.1 equivalent to inf( L(O, A), where L(O,C) = maxcelxi A)) and fc(0, [(0, () - X( - C)2 + end. Therefore, in this case, L(0, A) is the point-wise maximum of the set of functions fo indexed by Ç. Note that L(O, A) is not guaranteed to be smooth even if all functions {fe]ce()) x are smooth. Thus, if further assume that each fc(0,A) A) is convex in (0, A) and solve the problem using first-order methods, e.g., subgradient methods (Nesterov et al., 2018), it will require 3( iterations to find a 6-accurate optimal solution and each iteration requires O(xx)) queries to evaluate fc and Vcc. However, as discussed in Nesterov (2005). if the "softmax" is used to approximate the inner maximization problem, the smoothed problem L,, X) is smooth when all functions (Sckce|z) are smooth, e.g., when Vfc is ((0*t)-Lipschitz continuous, and it only takes Ö(8-1) iterations to find a d-accurate optimal solution to the original problem L(0, by selecting y = Olog (8). For each iteration, it requires the same O(X') of queries to calculate the gradient. Since querying the functions Se is computational costly in practice, solving equation 5.2 reduces the query complexity from 0(1x15=2) to ((1x18=1) when fc are convex and smooth and f is a &point mass.
Note that although the regularized OPL problem equation 5.2 improves the query complexity by an order of 0(5-1), the lincar dependency on | remains. If X is large, c.g., if | > 8-1, the computational advantage of solving the regularized DRO problem equation 5.2 compared to the un-regularized DRO problem equation 5.1 is marginal. To further improve the query complexity in |x), we propose biased stochastic gradient descent (SGD) method to approximate the gradient of equation 5.2 Specifically, at each time step, we sample one point from Oo and sample multiple points uniformly from the set X. We then construct a biased gradient estimate using those samples. Note that unlike the SGD methods typically used in the machine learning literature, the finite sample stochastic gradient estimate is biased due to the nonlincar logarithm function equation 5.2. The full algorithm presented in Algorithm 1
Theorem 5.2 (Iteration and query complerity of Algorithm D.. Assume that the policy space θ is compact and there exists fonite number B > O such that for any x, E X, || vll2 S B. i) Suppose that ((0,c) T
is convex in 0 and Li-smooth in 0 for any < X X. Then, the number of itcrations required to find a 5-accurate solution is in the order ((8-2) iterations to find 5-accurate optimal solution and the total query complexity is O(8-2). ii) Suppose that i(0, Ç) is Li-smooth in 0 for any Ç E X. Then, the number of iterations required to fond d-aceurate slationary point is in the order O(8-1) iterations and the total query complezity is 0(8-6).
Theorem 5.2 provides query complexity of Algorithm foor both convex and nonconvex smooth functions. The regularized DRO problem is a special case of so-called conditional stochastic optimization (CSO) problems (Hu al. 2020 The proof of Theorem 5.2 is directly adapted from Theorem 3.2 and 3.3 in Hu et al. (2020) and is provided in Appendix C.2 As finding the optimal query complexity not the focus of this paper, we
defer readers to Carmon et al. (2021) b Hu et al. 20200 Wang et al. (2021) for related discussions.

---

Published in Transactions on Machine Learning Research (01/2024)
Algorithm 1 Policy learning using biased stochastic gradient descent
Require: Robust cost estimates ri(s,a) a) for each (s,a) pair; context uncertainty radius 6z; parameterized policy wo: where 0 E θ number of iterations T;; inner batch size my and step size Yei initial policy parameter Bo and dual variable Ao; smoothing parameter f); context support set X and nominal context distribution pa.
1: for i = 0....,
do
2: Sample a single context ia from P and my points {<:)z=) uniformly from the support set X.
3: Evaluate the function values {o((; 02, Ae)) where 0i, Ar) c=((0::))--((e=-))))
4: Calculate the biased stochastic gradients Ve, (o,,, and (0, (e0(0,,,)//000 Ao0 according to equation 5.3 equa-
tion 54
DØ g(G; Or Ao vo.((0..5)
Vo, (o,,,
(5.3)
Σj g(G;; 02, 21)
av, (V,
E me 0(()) 82,, Ae Ae) 6
Ez.
(5.4)
-タ=1 g(Ss: Ne)
5: Update the policy parameter and the dual variable:
Dc+1= Ille (0. 17 Vo, (0c, A4+1= = (
6: end for
6 Policy evaluation and learning for the international stroke trial
In this section we validate our regularized Wasserstein DRO methods for both OPE and OPL problems on the International Stroke Trial (IST) (Group et al. 1997p dataset.
IST is a randomized controlled trial that aims to study the effects of carly administration of aspirin and/or heparin on the clinical course of acute ischaemic stroke. The IST datasel (Sandercock et al. 2011) includes 19,435 patients with logged contextual information such as age, gender, and level of conscionsness before treatment admissions, as well as follow-up results on day 14, including the occurrence of recurrent stroke, pulmonary embolism, and death. We define two actions from the IST dataset: the treatment action of pre- scribing both aspirin and heparin (high or medium doses) (a1) and the control action of of not administering any treatment, neither aspirin or heparin (02). The behavioral policy so(a- = a1) = 0.3. The cost function is calculated based on the recorded follow-up events day 14.
To introduce distribution shifts, we split the dataset into a training set and a testing set, and we introduce a selection bias into the training set. Specifically, we randomly remove 50% of the patients in the training set who are not fully conscious. This creates difference in the context distributtin between the training set and the testing set, with the patients in the testing set being more likely to be unconscious before treatment than those in the training set. In this simulation, we only consider context shifts and assume that the cost functions do not change between the training and testing sets. Due to the limited number of data points in the IST dataset, we further use decision trees and binning methods to reduce the number of distinct contexts. The detailed experimental setup are provided in Appendix F
We first consider the OPE problem for a random policy using the training set. Although any policy can be evaluated using the proposed methods, we select the random policy since it the same as the behavioral policy, and thus the underlying truth can be estimated more accurately from the testing set, i.e., the impor- tance ratio is always 1 (Dudik et al. 2014h. One main challenge in DRO is the selection of the radius, as the final robust value depends this choice. When similar datasets that include both the testing and training sets are available, the uncertainty set radius can be selected by calculating the distance between the testing and the training sets; however, this distance is not always well-defined under KL divergence. Specifically, for a large number of contexts, it is unlikely that all contexts have been observed for any two given finite datascts. Therefore, the two datascts are likely to have different supports, resulting an unbounded KI. radius. Since the Wasserstein distance well-defined regardless of possible differences in the supports of the context distributions in the two datasets, we can split the training set into two parts and calculate the
10

---

Published in Transactions on Machine Learning Research (01/2024)
Table 1: Policy evaluation and learning results
Method Uncertainty set radius (ew)
0.03
0.05
0.1
OPE: Factor KL DRO (EKL)
0.30 (0.01)
0.374 (0.1)
0.74 (1.0)
OPE: Wasserstein DRO (LP on sub-support)
0.53
0.61
0.79
OPE: Regulated Wasserstein DRO (BSGD)
0.59
0.76
1.05
OPL: Factor KL DRO (EKL)
0.28 (0.01)
0.368 (0.1)
0.69 (1.0)
OPL: Regulated Wasserstein DRO (BSGD)
0.54
0.62
0.92
Expectation under P (training set, random policy) Expectation under Q (testing set, random policy)
0.28
0.38
Wasserstein distance between them, which yields the uncertainty set éw 0.03 in this simulation. This distance captures the data uncertainty due to finite samples and provides a lower bound on the uncertainty set radius. To further capture distribution shifts, should select the radius larger than this distance.
The policy evaluation results are presented in Table I Solving the OPE problem using the Wasserstein DRO method is challenging since the support size is large. Therefore, we approximate the solution by only considering observed contexts in training set instead of the full set X. However, this Wasserstein DRO using the sub-support only provides a lower bound to the full support DRO method, as the inner maximum problem is approximated by maximizing over smaller subset for given radius. On the other hand, regularized Wasserstein DRO does not face the issue of large support set, as it directly approximates the solution through sampling. It is important to observe that both Wasserstein DRO methods provide upper bounds to the true expectation under the testing distribution Q as in Table I Since we consider separate uncertainty sets on both the contexts and the rewards, we adopt the Factor KT, DRO method developed in Mu et al. 2022 to benchmark the performance of Wasserstein DRO, where similar separate uncertainty sets are considered. As no specific numerical algorithm is provided in Mu al. (022) to solve the KL DRO problem, we simply use a gradient descent method to find the optimal solution. Note that the DRO problem under KL divergence is convex, gradient descent is guaranteed to converge to the global optimal point. The gradient derivations of KL DRO can be found in the appendix of Mu et al. (2022) As estimating the uncertainty set radius for KL divergence is not straightforward, we test different radii when using KL DRO, However, it is worth noting that the KL DRO method fails to provide valid upper bounds when EKL. = 0.01 or 0.1.
Finally, we consider the OPT, problem by parameterizing a context-dependent policy. Specifically, if patient is not fully conscious, there is a probability 01 of taking action a1, and if a patient is fully conscious, there is a probability 82 of taking action 01. The optimal policy values are shown in Table I where the Factor KT. DRO is solved by gradient descent on both the policy parameters and the KT dual variable. Note that the optimal policies achieve lower costs compared to the random policy for both methods under all three uncertainty sets, indicating that the random policy is not robustly optimal. The optimal policy parameters of the regulated Wasserstein DRO converge to 0₃ = 0.8 and 0₂ = 1 when éw = 0.03, which suggest that the early administration of aspirin and/or heparin effective on the clinical course of acute ischaemic stoke at the population level. Note that the optimal robust policy is not necessary deterministic the value function in equation 5.1 is not linear in the policy parameters. Further results on policy learning can be found in
Appendix F
7 Conclusion
In this work, we proposed a distributionally robust optimization (DRO) framework using Wasscrstein dis- fance (Wasserstein DRO) to address distribution shifts in OPE and OPL problems. We provided a finite sample analysis and showed that the policy value obtained from a dataset of size n converges to the true values at a rate of On under the proposod framework. As Wasscrstein DRO computationally challenging when the support set of the underlying distribution large, we further formulated a regularized Wasserstein DRO problem and provided corresponding biased stochastic gradient descent algorithm that is
11

---

Published in Transactions on Machine Learning Research (01/2024)
independent of the distribution support size and can numerically solve the regularized problem; the proposed algorithm achieves better iteration and query complexity compared to solving Wasserstein DRO using linear programs while maintaining the same convergence rate. We validated our proposed methods using medical stroke trial dataset.
Acknowledgments
This work is supported in part by AFOSR under award #FA9550-19-1-0169 and by NSF under award CNS- 1932011. Pan Xu is supported by the Whitchead Scholars Program and the Department of Biostatistics and Bioinformaties at Duke University.
References
arXiv preprint Christopher Adjaho and Timothy Christensen. Externally valid treatment choice. arXi::220500556,, 2022.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214 223. PMLR, 2017.
Waiss Azizian, Franck Tutzeler, and Jérôme Malick. Regularization for wasserstein distributionally robust optimization. arviv preprint arXiv:2205.088. 2022.
Jose Blanchet and Yang Kang. Semi-supervised learning based on distributionallly robust optimization. Data Analysis and Applications 3: Computational, Classification, Financial, Statisticol and Stochastic Methods, 5:1-33, 2020.
Léon Bottou, Jonas Peters, Joaquin Quihonero-Candella Denis Charles, D Max Chickering, Elon Portugaly, Dipankar Ray, Patrice Simard, and Fd Snelson. Counterfactual reasoning and learning systems: The example of computational advertising. Journal of Machine Learning Research, 14(11), 2013.
Yair Carmon, Arun Jambulapati, Yujia Jin, and Aaron Sidford. Thinking inside the ball: Near-optimal PMLR, 2021.
minimization of the maximal loss. In Conference on Learning Theory, pp.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Erick Delage and Yinyu Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations research, 58(3):595 612, 2010.
Miroslav Dudik, Dumitru Erhan, John Langford, and Lihong Li. Doubly robust policy evaluation and optimization. Statistical Science, 29(4):485-51) 511, 2014.
Jean Feydy. Geomelric dala analysis, beyond convolutions. PhD thesis, Université Paris-Saclay Gif-sur- Yvette, France, 2020.
Jean Feydy, Thibault Séjourné, François- Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2681-2690. PMLR, 2019.
-aaier
Samuel G Finlayson, Adarsh Subbaswamy, Karandeep Singh, John Bowers, Annabel Kupke, Jonathan Zit- train, Isaac S Kohane, and Suchi Saria. The clinician and dataset shift in artificial intelligence. New England Journal of Medicine, 385(3):283 286, 2021.
Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio. Learning with a wasserstein loss. Advances in neural information processing systems, 28, 2015.
Rui Gao and Anton Kleywegt. Distributionally robust stochastic optimization with wasserstein distance. Mathematics of Operalions Research, 2022.
12

---

Published in Transactions on Machine Learning Research (01/2024)
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large-scale optimal transport. Advances in neural information processing systems, 29, 2016.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. kernel two-sample test. The Journal of Machine Learning Research, 13(1):723 773, 2012.
International Stroke Trial Collaborative Group et al. The international stroke trial (ist): a randomised trial of aspirin, subcutancous hoparin, both, or noither among 19 435 patients with acute ischaemic stroke. The Lancel, 349(9065):1569 1581, 1997.
Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2023. URL
com
Weihua Hu, Gang Niu, Issci Sato, and Masashi Sugiyama. Docs distributionally robust supervised learning give robust classifiers? International Conference on Machine Learning, pp. 2029 2037. PMLR, 2018.
Yilan Hu, Siqi Zhang, Xin Chen, and Niao He. Biased stochastic first-order methods for conditional stochastic optimization and applications in meta learning. Advances in Newral Information Processing Systems, 33: 2759-2753 2020.
Zhaolin Hu and L Jeff Hong. Kullback-leibler divergence constrained distributionally robust optimization. Available at Optimization Online, 1(2):9, 2013.
Kosuke Imai and Marc Ratkovic. Estimating treatment effect heterogeneity randomized program evalua-
tion. Annals of Applied Statislics, 2013.
Nathan Kallus, Xiaojie Mao, Kaiwen Wang, and Zhengyuan Zhou. Doubly robust distributionally robust off- off-policy evaluation and learning. In International Conference on Machine Learning, pp. 10598-10633 PMLR, 2022.
Daido Kido. Distributionally robust policy learning with wasserstein distance. arXiv preprint arXiv:2205.0460 2022.
Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Bateeeezzde---baadee Wasser- stein distributionally robust optimization: Theory and applications in machine learning. In Operalions research & management science in the age of analytics, pp. 130- 166. Informs, 2019.
John Langford and Tong Zhang. The epoch- greedy algorithm for contextual multi-armer bandits. Advances in neural information processing systems, 20(1):96-1, 2007.
Lihong Li, Wei Chu, John Langford, and Xuanhui Wang. Unbiased offline evaluation of contextuul-bandit- based news article recommendation algorithms. In Proceedings of the fourth ACM international conference on Web search and data mining, pp. 297-306, 2011.
Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the wasserstein metric: Performance guarantees and tractable reformulations. Mothematical Programming. 171((12):115-166, 2018.
Tong Mu, Yash Chandak, Tatsunori Hashimoto, and Emma Brunskill. Factored dro: Factored distribu- tionally robust policies for contextual bandits. Advances in Neural Information Processing Systems, 35: 8318-8331, 2022.
Hongscok Namkoong and John C Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. Advances in neural information processing systems, 29, 2016.
Yu Nesterov. Smooth minimization of non-smooth functions. Mathematical programmingg 103:127-152. 2005.
Yurii Nesterov et al. Lectures on conver optimization, volume 137. Springer, 2018.
Aaditya Ramdas, Nicolás García Trillos, and Marco Cuturi. On wasscrstein two-sample testing and related families of nonparametric tests. Entropy, 19(2):47, 2017.
13

---

Published in Transactions on Machine Learning Research (01/2024)
Peter AG Sandercock, Maciej Niewada, and Anna Czlonkowska. The international stroke trial database.
Trials, 12(1):1-7, 2011.
Nian Si, Fan Zhang, Zhengyuan Zhou, and Jose Blanchet. Distributionally robust policy evaluation and learning in offline contextual bandits. In International Conference on Machine Learning, pp. PMLR, 2020.
Jie Wang, Rui Gao, and Yao Xie. arXiv:2109.1 11926, 2021.
Sinkhorn distributionally robust optimization. arXiv preprint
Yu-Xiang Wang, Alekh Agarwal, and Miroslav Dudık. Optimal and adaptive off-policy evaluation in con- textual bandits. In International Conference on Machine Learning, pp. 3589 3597. PMLR, 2017.
Jonathan Weed. An explicit analysis of the entropic penalty in linear programming. In Conference On Learning Theory, pp. 1841-1855. PMLR, 2018.
Zaiyan Xu, Kishan Panaganti, and Dileep Kalathil. Improved sample complexity bounds for distributionally robust reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pp. 9728-9754 PMLR, 2023.
Zhouhao Yang, Yihong Guo, Pan Xu, Anqi Liu, and Animashree Anandkumar. Distributionally robust policy gradient for oflline contextual bandits. In International Conference on Artificial Intelligence and Statistics, pp. 6443 6462. PMLR, 2023.
Chaoyue Zhao and Yongpei Guan. Data-driven risk-averse stochastic optimization with wasserstein metric. Operations Research Letters, 46(2):262 216-26776 2018
Xin Zhou, Nicole Mayer-Hambless Umer Khan, and Michael R. Kosorok. Residual weighted learning for estimating individualized treatment rules. Journal of the American. Statistical Association, 112(517):169- 187, 2017.
Zhengqing Zhou, Zhengyuan Zhou, Qinxun Bai, Linhai Qiu, Jose Blanchet, and Peter Glynn. Finito-sample regret bound for distributionally robust offline tabular reinforcement learning. In International Conference on Artificiall Intelligence and Statislics, pp. 3331-3333 PMLR, 2021.
Jia-Jie Zhu, Wittawat Jitkrittum, Moritz Dichl, and Bernhard Schölkopf. Kernel distributionally robust optimization: Generalized duality theorem and stochastic approximation. In International Conference Artificial Intelligence and Statistics, pp. 280-288. PMLR, 2021.
Steve Zymler, Daniel Kuhn, and Berç Rustem. Distributionally robust joint chance constraints with second- order moment information. Mathematical Programming, 137, 2013.
14

---

Published in Transactions on Machine Learning Research (01/2024)
Appendix
A Proofs for Off-Policy Evaluation
A.1 Proof Theorem 3.2
In what follows, we first provide an error decomposition lemmaa.
Lemma A.1 Let Assumptions 2.2 and 2.1 hold. Denote p(z,) = Po = xx) and p(x,) = pa x.). For any given policy T, define 1(S) B (a (m((s
a)] and i(C) = Bawa((l)) [)]] We have that
((()) V(T)I Mmax
Σ
l((xa)) - -p(xz))
sup CE" (((())
where I(C) represents the distribulionally robust cost of following policy at the conlert Ç.
Lemma A.1 indicates that the gap between (T) and V(T) upper bounded by the error from the contexts distribution estimate and the empirical distributionally robust cost function estimate. Since both crrors are cansed by finite samples, the errors will decrease if we increase the size of the observational dataset.
Lemma A.2 Let P be a discrete distribution of size m. For each [m], denote p(i) the true probability and p(i) a Eー= Ifsy = i} the empirical probability, where {sy}ī=1 are n i.i.d. samples drawing from P. Denote f(p) = Ei (f(x;) - p(x:)] the random distance between and P after observing n. samples. For a given å E (0, 1), we have that, with probobility at least 1, 8,
2 2lo (1/8) f(P) S V
m n
(A.1)
T2
Lemma A.3 Consider Bernoulli distribution that lakes the value with probubilily p 0. probability 1 - 8, one observes at least uo samples of value out of n trials as long as n 2 2log(1/8) z
With
F²
Lemma A.4 Let Assumption 2. hold. Given d E (0,1), with probability 1 8, if we choose n 2 2 (1 we have V(x, a) E X x A, it holds that (Mry Mxg
|m(x, 1 (n(x,a)-mis mi( a)l Ymar (v Alog (21x1111185)
2 MpPolao
(A.2)
Proof [Proof of Theorem 3.2
According to Lemma
A.1
we have that
(V) V(T) < Mmxx
Σ
lp(x.) - p(ai))
+ sup SEa
(())))
|2| [p(x;) - p(zx))). which relates to the total variation between the em- E
We first bound the term Ymax pirical estimate and the true discrete distribution. For a given E (0, 1), Lemma A.2 implics that, with probability least 5, we have that.
Ymax (
210g (1/8) n
X' n
Y ymax
Σ (p(xe) p(ax))
(A.3)
where is the number of samples.
15

---

Published in Transactions on Machine Learning Research (01/2024)
We now bound the term supeex (i()) - u(s)). which relates to the difference between m(ç, a) and mi(S. a)
due to the following fact.
sup cex
sup CEX
( [M((. E [m{{,
I sup (m(s,)) ÇER,uEA
(A.4)
[1 then [m(x, a) m(x a)l s
Lemma
A.4
implies that, with probability 1 - 8, if we choose n 2 (Mog
Ymax
4log (2) (2
[A)66)
((2,a) E X x A. Since holds for all context actions pairs, we have nitotaao
2/3]
that
(Ifu(s.)-m(() m(S. S Ymax (AAlog(22R111 (21x1
2 m1p Maoo
sup ÇEX,aEA
(A.5)
7 71mooMo
Taking union bound of the events when both inequalities equation A.3 and equation A.5 hold, have that with probability --6, 6, by selecting n 2 log((4 ( the following inequality holds
(Mrzg Azo
(m) [(+) S Ymax (V log (2/5)
X n1
4 4log (4120))A)10)
2 71/1 Mono
11 Maro
A.2 Proof of Theorem 4.2
Lemma A.5 Given set of wumbers (xi,.... ..n. and let Tmax maxx _ ,nn). we define the Log-
Sum-Eap (LSE) function LSE(#1," [..) log ( Σn exp(πx;)). Then, the following inequalities hold:
log(n)
Tmax
LSE (21, ... "..,"",,,, ",,)), S Tmax.
Σi= exp(nx;) < exp(n)max). The result follows by taking the logarithms to the incqualitics and multiplying by 1/n.
Proof Observe that 1 exp(()max) <
Lemma A.5 shows that the DRO problem under the regularized Wasserstein distance can be good approx- imation to the original DRO problem if select the smoothing parameter 7) properly.
A.5 implies that 1 Loo (Ecex Y ( A(=
Proof Proof of Theorem 4.2
Lemma
suPcex (J(C) - Mxx - 2)2) for any given function f(C). As result, similar to the proof as in Theorem 3.2 for
any given policy म, both Vr (T) and (a, a), V(I, a) E X x are upper bounded by Ymax according to Lemma
2.5 In addition, if we denote X and X the optimal solution to V and VT we have A E [0, Ymax/ Define I(C) Be- (m((,a)) and i(C) (F (M((,a)) Without loss of generality, for given policy
16

---

Published in Transactions on Machine Learning Research (01/2024)
", assume (TT) 2 V(z) and we have that
(T) [TT)
- 11 ))
|
e\ +
(z) -
log
Σ SEX |
1
(
ex*
Σ
p(z)) SEX
sup {1(() - X (2; c)
|21
1.21
ex*
+
(1rr) log
Σ SEX
1
en ((5) (1((())) (xc-)))
ex*
p(
sup {1(5) - >² - CEX
-5² 5)²}
|
17
im
x
sup {1(() - X* (2: 5)²} )
Σ
((1.)) -
- log
Σ <E>
((((()
p((x))
11
SEA
where the inequality holds since X is the optimal solution to (T) while A is not. Then, it follows that
(x)( - [++]
12
Σ (x((c)--(000-9)) cex
- -p(xa) _ log Σ
e ((06)))(00011)))
p(zi) log
12
|
(EE
+p(x;)
1
log
1
p(x;) sup
{i(() >² - {)
|
77
SEX
CEX
+p(x;) sup SEX
(i(S) - -21(x; - {) p(ax) SEt
sup (((() >' - c)²]
|
|
< Y Ymax
Σ 2=1
p(x.) p(x:)]+ log(1=()
+
Σ 1=1
p(a;)
sup CEX { {((() { -(2))
11
The last inequality holds since log Ecex t// t ( (
maxçex ((()) X²xx;
() S (((() S Ymax- In addition, according to Lemma A.5 we have that 1 1 log (Eçex 1xT (((())(()---)))))
supçex {i()) > x1x= c)
log(1>() Observe that the last
inequality is the same as Lemma A.1 except for the extra error term lox(Ix)) ; we complete the proof by
applying the results from Theorem 3.2 To the last inequality.
B Proofs of Supporting Lemmas
B.1 Proof of Lemma 2.5
Proof Denote D* ex* + po pa (supeex (S(1) - Xxx Xxx-) For a given point æ and X", we have
that supçex (f(c) - X³ (z - 2) > sup (f(() (xx Taking the expectation of x with respect to the distribution pa and adding eX* on both sides, we have ed" + Ee-- [sup;ex (f(() - Ax () >
ex* <x (sup (f(c) (= - 2) +Ex [f(x))], which implies D* 2 ex po [f(x)) 2 ex".
By the definition of the dual problem, we have D < c..0+ + Egaro [supcex (f(s) - 0 (x- --]) = fmax-
The result follows from the fact that 0 < ex* S D* < Smax.
17

---

Published in Transactions on Machine Learning Research (01/2024)
B.2 Proof of Lemma A.1
Proof Denote X and A the optimal solution to V and V. Applying Lemma 2.5 to equation 3.2 and
equation 3.4. we have m(s, a) and m(s, a) are upper bounded by Ymax since the costs are upper bounded
by Ymax. This further implies that ((x:) and ((x) are upper bounded by Ymax- Applying Lemma 2.5 to
equation 3.1 and equation 3.3 we have X², x E [0, e2)x/ Without loss of generality, for given policy
H,, assume (TT) 2 ((T)
and we have [V(T) [TT)
|x
X
= (c} +
((x,) sup (x) {)p}) (eX' Σ Cex
p(az) sup {((() (1; (x;-<)p)
ÇEX
|x
S (EX*
+
(1xs))
sup {i(() (x, ()p))))-00 -ce2* +
Σ
p(zz)
sup {((() - X² ()²)) ÇEt
SEtt
i=1
i=1
=
p(xx))
sup ÇEX
{i(( \² - -)p] -x;
p(ax)
sup {'(() X Xa () CEX
The first inequality follows from the fact X is the optimal solution to V while X is not. Then, we have
[2| Σ
S ((x,))
sup {((() - X² - SEA
-)²} (=,--)))
p(x;) sup {i()) - -xaxx -0x0--))) c)²} CEX
-c)
+p(x:)
sup CEX
{i()) - \² - (=x==)))) -s)²}
p(ax;) sup {(({) x²; - -spp] CEX
|
x
((xs) - p(a;))
sup x.EE
x0x (x= ()²} {i()) x2
+
Σ
p(x:)
sup CEX
i=1
|2
< YYma
Σ (m(xs) p(zs)]
+
Σ i=1
p(as) sup SEX
(((s)-((())
The second inequality holds since sup, f(x) sup, g(z) s sup, fxx - g(x)]. The last inequality holds
due to the fact that ?(C) S Ymax and X 2 o. Note that if (T) < V(T),, we will replace the term
supa,çex {i(s) - xx } by the term sup,Çex ((()) '0(x - <)² in first incquality and the result still holds.
B.3 Proof of Lemma A.2
Proof Observe that changing any single sample can at most change the value of f by 1/n. Applying McDiarmid's inequality to the function f, we have that
2.2 (((()) - E((2 20se
= e 2nc²
which implies with probability at least 1- 8, f(p) B[f(p) S
2 log (1/8)
by setting / 2log (1/8) Thus,
n to bound f(p), we only need to bound the first frst-moment term [p)) Specifically, we have that
"
1 p(ax)) - p(xx)),
E(f(p)] = Σ [lp(2x:) -(xp)p(()1)--0))))0
<
E ((<x:))p(<:)Ppp
1=1
18

---

Published in Transactions on Machine Learning Research (01/2024)
where the first inequality follows from Jensen's inequality. The second incquality holds due to the fact that np((zz) follows from the binomial distribution and E [lí(x.) - p(x:)) 1/n2 Var(nf(1:)) 1 (n
p(z:)). Since - p(z;)= - 1, we have that
Σ
ploe (1 p(x))) S Σ Vpla)
1 -===-
m
where the second incquality follows from Cauchy-Schin Schwarz inoquality. Putting everything together, we have f(p) S / (1/66
m
B.4 Proof of Lemma
Proof Let the empirical estimate of p after n trials, i.e., p= Di= 1ss = 1}, where {sy}ā=1 are iid. samples drawing from P. We first show that the probability p S is low. We have that
P P 2) P p ())
2) 5 S exp (-2m (
npp 2
exp
where the inequality follows from Hoeffding's inequality. For a given > 0, by choosing n 2 log we have that with probability at least 1 - 5, P 2 P. This implies that we observe at least uy samples of value 1 out of n trials.
/²
B.5 Proof of Lemma A.4
Proof Without loss of generality, we assume m(xxa) 2 m(x, a) and let and X the optimal solution to m and mn we have that
mm(a,a) a) m(xxa))
ex* + Pg py sup CE
(S, 2-(6
-8:(5)-------0- ( -(²)
ex Be sup (Uz,((S) X2 5) CEE
<
ex* + Be 6o sup (y (S) - X- - 6)9)
ex* + Eg sup (Va. 27 <)²)
CE
<EE
E
S Y mmax
(pE))) p(E:)l-
Observe that EE1 Ip(E)) P(E:) directly depends the number of samples. Denote n( a) the random variable representing the number of (x,a) pairs observed in the dataset. Let Pr.o be the probability of observing each pair (x, a). Lemma A.3 implies that we select n 2log(1/8) we will observe at least nps.s
pe.
samples of (s, a) pairs with high probability. In addition, according to Lemma A.2 with high probability, have
log (1/8) n(xx, a)
Σ Ip(E.) - p(E)) =
+
n(I, a)
Take union bound on the following two events
{ (n(2,a)
11pr.0 ) and Σ (OEE)) p(Ex)) S
2 log (1/5) n(s, a)
n(x, a
19

---

Published in Transactions on Machine Learning Research (01/2024)
have that, with probability - 6, if select n 2
2/66)
then
Pi.a
4 llg (2/6)
2
(M(E)) - p(E))) &
v
nPz,a
npr,a
We complete the proof by taking the union bounds over all state action pairs and using the fact that Pa,a is lower bounded by PraMaas which is implied by Assumption 2.1
C Proofs for Off-Policy Learning
C.1 Proof of Theorem 5.1
Proof When V 2 V*, obtain that V ---ve V*I
V(**) V(**) I V(T)) - V(T'). Then,
the result follows by applying Theorem 3.2 to the policy x². When V (+, have that V ---
V* = -((")-((()) & - The result follows by applying Theorem 3.2 to the policy A*.
C.2 Proof of Theorem 5.2
Due to page limit, the sample size and the learning rate are not specified in Theorem 5.2 The theorem below provides a complete version of Theorem 5.2
Theorem C.1 (Query complerity of Algorithm Assume that the policy space θ is compact and conver and there erist a fivite number B such that for any I, X, || ylla S B. i) Suppose that i(0, () Ei_ (o(a:)()m(((.a)) ai) is convex in and Li-smooth in 0 for any Ç E X. For a given accuracy 6 > 0, if we set m, = 0 and " = O then it requires (6-²)) iterations to find a d-accurate optimal solution. The total query complerity is 0(8 ti) Suppose that ((0, () Σ (e(a))()m a.) is Lp-smooth in 0 for any X. For a given accuracy d > O, if we set m, = ( and 7o 1/22,, then it requires O(8-1) iterations to find d-accurate stationary point. The total query complerity is O 6).
Before proving Theorem C.1 we first introduce some notations and study our problem under the framework of conditional stochastic optimization (Hu al. 2020).
Notations A function f(-) : R + R is said to be L-Lipschitz continuous X' if for any 2, E X,
[f(z) - f(y) < L| y1l2; f(-) is said to be S-Lipschitz smooth on X' if for any z, E X, f(x) f(y) -
(f(y) (x-y) < =1|l-v1l3 f(-) is said to be p-convex on X' if for any x, E X, f()) (x- 2
- v12..
Conditional stochastic optimization Consider the following conditional stochastic optimization prob-
lem:
min F(u) = E, [Sa (Eqi Isc (u, x)))]
(C.1)
where U C Rd (x,x)) is function that depends on random variables and Ç. For fixed and given samples {,, 1 define the finite sample function
f(u ((4x,155)7
fx m
gC, (u,
20

---

Published in Transactions on Machine Learning Research (01/2024)
In addition, the biased gradient of F()) is defined as:
Vullu; (-))==))= Vulz
-
Σ
m
gcs (u,
Σ j=1
ugg (2,
m
Recall the regularized Wasserstein DRO problem is defined as:
onlless-Mes))))
+ tead)
VT(TT)= = inf 8,220
E
log
Σ SEN
(C.2)
1
X
We convert the regularized Wasscrstein DRO problem into the conditional stochastic optimization problem. Denole = A), for fixed sample from 2 we define
ge(u, x) -M( e fz(-) log()))
17
Thus we have (T(TT) inf, (u) inf stein DRO problem is indeed a conditional stochastic optimization problem. Its biased gradient with respect to 0, A can be computed as follows:
{ (Be-Py [Se (Bcvmmiform(x)) (u)(u,))))] The regularized Wasser-
Em j=1 Σ ((00))))-N
2,,
Σ ((((((,)))-((0-- -
(3)
V,
Σj (10.6))) + Ex.
As proven Theorem 3.2 and 3.3 Hu al. (2020), prove Theorem C.1 suffices show that the following properties hold for our problem:
Bounded variance: of == sup E (1, x)) Ee (uc m)11
+00
fu; v. (()))nn) is convex in if f(0,C) is convex in and is Lysmooth in (couvex case); fa and
Рay are Lipschitz contimuous and smooth if i(0,ç) is only Ly-smooth in 0 (non-convex case)
3. Bounded second moment: there exist M > 0 such that E [IV. ( (+) s M2 for any
u.
1. We first verify that regularized DRO problem satisfies the bounded variance propricty. Note that for any 2, (, as O s Yman and 0 < < Ymaz/ez, we have that
(0 Ymua B0+ S ge(2,, u) s ( -0te Yaus
Ott,
which implics that of ( mm B²* /en)) /en) )²
S
AnnMmss +00.
2. We now verify the convex case of the second propricty. For given and Sj, denote h(0, 2; C3) ((0,)) A(x - 2. As i(0,,)) is convex in by assumption and A(x - G2 linear X, h(0, X; G) is convex in and X jointly. Consequently, for a given x E X and { {G)))> the function Flu;z, x, (G)))))) - Le=((e,,))-Ma--0
+ ε is convex in (0, X) since log EØ= 1 exp (b,0,);,,Gi)) preserves convexity when h(0, Xx (z) is convex in (0, A).
") log T
For the non-convex case, we need to show both fz and 9G, are smooth. As derived in the first property, gc is bounded. Thus, the domain of the function fa is restricted to the set nY fez 2n)ymas) and thus
21

---

Published in Transactions on Machine Learning Research (01/2024)
Table Comparison between KL DRO and Wasserstein DRO
Method
0.8ef
c
1.2c]
0.8c.
<
1.2c5
KL DRO
47.27 41.53
49.84
52.20
49.86
53.07
56.09
Wasserstein DRO
43.36
45.04
41.82
43.72
45.49
Expectation under Q
35.58
f.(+) 1 log(-) is enYmmx / /n-Lipschitz continuous and e B /es /n22 Lipschits smooth. We now show 7 that gc(,,z) smooth. For any given two points u and u', we have that
[ugg(u, - (ugg(u, x| & Vegc(u, x) Vegg(u', (=))( (V - (u =l) g6 (u' Vol00 O) + (1); - ge(u' x)) m/ (a ()
(C.3)
where the equality follows from the definition of gc functions for the regularized DRO problem. As i is Lipschitz smooth in 8, we further bound equation C.3 as follows
equation
C.3
2 mix || <) Vol(0' (0] O) + 17 (B - 9c xx))ll
u mms 1 || + B max l
- ex
S ne 2 1 | + BB 2n a X & 22n( + 1) (Lp + 1)]lu - w'Il-
As a result, 9c (u, is (B + 2n - (T+ + T)Iiisccitz smooth in a for any x,
3. The bounded second moment property hold due to the fact that both || G/Il and || - Gill are bounded as θ is compact and || - G B.
As all three propertics hold for the regularized Wasserstein DRO problem, we complete the proof.
D Connection of regularized Wasserstein DRO to entropy regularized primal problems
Lemma D.1 (Strong duality for cost-regularized Wasserstein DRO Take e, 17 > 0. Then we have
Azizian et al. 2022 Theorem 3.1)).
log (E<moo(12)) [er((co- 11]]
sup Bgwr(f(E))- (oloo) inf Po
cA + EgnPo
A20
17
where Po); {P E M inf, (B(,,))-ll((,CC)) pp:00.0pp0p probability measure for each x E X
p0,02 =P s c} and oo(-1z) is pre-specified
The dual problem in Lemma D.1 coincides with our dual formulation equation 4.1 if 0-(-12) is chosen as a uniform distribution over X.
Remark D.2
Note that the entropg-regularized term in the primal problem added the objective function.
It is possible to add the entropy-renlarrzed term on the uncertainty set. well, which is referred to
Sinkhorn distance in the optimal transport literature Cuturi 2013). Distributionally robust optimization
using Sinkhorn dislance has vecently been explored in Wong et al. (2021). Our construction is molivated
by directly smoothing the dual problem while Sinkhorn DRO starts from regularizing the primal problem. A
similar idea of smoothing dual problems was studied for semi-supervised learning problems; see Blanchet B
Kang 2020) for details.
E Comparison between KL DRO and Wasserstein DRO
We provide a simple DRO example illustrate the differences when using different uncertainty measures. Let X' {1 be the support set that spanning uniformly from to 10. The true training distribution
22

---

Published in Transactions on Machine Learning Research (01/2024)
0.08
0.06
Probabiltty ity
0.04
0.02
0.00
10
Figure Distributions
P, the empirical training distribution P and testing distribution Q are depicted in Figure Note that only the empirical distribution P can be observed. We define the cost function f X' -+ TR = x² We then set the radius for both KL (Gkr) and Wasserstein (Gv) by dircctly calculating the distance between P and Q. The robust values are shown in Table 2 W(P, Q) 1.85 for each method, respectively. Since in practice we don't observe and camnot compute 1.24
With slightly abuse of notation, cj is defined as KL(QIIP) 0.52 and
the true distance c², we also provide results when is not accurate, c.g., when c 0.8c The expectation under Q method represents the true expected cost when x is distributed according to the testing distribution Q. Both DRO methods provide valid upper bounds to the true expected costs. As mentioned in Scction 2 KL DRO tends to over-fit the largest cost and is not aware of the geometry of the support. Here, we modify the distribution P by setting ds11 = and the cost. f(xs1) = = and keep the probability P(X- = #51) the same as before, i.e., only the last point of the support: modified. We call the modified distribution P. This new point T51 can be seen as an outlier to the original support set. We then define E KL( 0.52 and <j
Q) 1.92 for each method, respectively. Observe that the KT, divergence does not change since the probability P(x x51) remains the same while the Wasserstein distance increases as we increase one value of the support. As shown in Table 2 the Wasserstein DRO method is robust to the outlier; the robust values obtained by the KL DRO method increase substantially though the true expect cost remains the same.
In practice, the datasot is not free of outliers and the outliers ineur larger costs, it better to use the Wasserstein DRO method since it considers the geometry of the dataset and does not simply over-fit to high costs.
F Additional Details of the Stroke Trial Application
The selected contextual variables, actions and follow-up events are described in Table 3 The cost function is defined according to the random follow-up events:
cost = I(ISC14) + I(NK14) + [(HTI14) + I(PE14) + +(DVVT4) I(NCB14) +3 +3- I.(DDDED
After removing patients according to the selection bias, the training set contains 4,340 patients. However, since there are 9,936 distinct contexts, the number of data points for most contexts is limited. To address this issue, we train two decision (ree models one for each action, to learn the average cost of each context and action pair by grouping similar contexts together. We further group the contextual variables AGE and RSBP into bins with size of 10 reduce the support sizes. The docision trees parameters are cross-valis voss-validatod and shown in Table g The same decision trees are used as inputs to all KL and Wasserstein DRO OPE and OPT, methods that listed in Table 1
23

---

Published in Transactions on Machine Learning Research (01/2024)
Table Dataset summary
Variables
Descriptions
RCONSC
Conscious state at randomization (F-fully alert, D-drowsy, U-unconscious)
SEX
M male; female
AGE
Age in years
RSLEEP
Symptoms noted waking (Y/N)
RATRIAL
Atrial fibrillation (Y/N)
RSBP
Systolic blood pressure at randomisation (mmHg)
STYPE
Stroke subtype
Action 1
If aspirin is allocated and heparin is allocated either 12500 5000 unites If both aspirin and heparin are not allocated
Action 2
DDEAD
Dead on discharge form within 14 days (Y/N)
ISC14
Indicator of ischaemic stroke within days (Y/N)
NK14
Indicator of indeterminate stroke within 14 days (Y/N)
HTI14
Indicator of haemorrhagic transformation within 14 days (Y/N)
PE14
Indicator of pulmonary embolism within 14 days (Y/N)
DVT14
Indicator of deep vein thrombosis on discharge form (Y/N)
NCB14
Indicator of any non-cerebral bleed within 14 days (Y/N)
Table 4: Decision trees parameters
Parameters
Action 1
Action 2
max depth
4
min samples leaf score function
2
mean squared error mean squared error
We calculate the initial Wasserstein distance 6w = 0.03 by solving a linear program. Since the number of optimization variables in the linear program very large (approximately 10 variables), affecting our ability to solve it, we restrict linear program the subset of the support X, specifically, we only consider the observed contexts in the training set. We calculate the Wasserstein distance between two random splits of the training set restricted to subset of the context variables. Note that this linear program restricted to subset of the context variables provides upper bound to the solution of the linear program with full context, as the search space is smaller and the linear program is a minimization problem.
The dual variables convergence results for OPE and OPL problems using Wasscrstein DRO are presented in Figure ((a) and Figure 4(b) and the convergence results of using KL DRO are presented Figure 3(b) and Figure 5(b] Additionally, the policy parameters learning curves of the OPL problems of two different methods are provided in 4(a) and 5(a) respectively.
24

---

Published in Transactions on Machine Learning Research (01/2024)
80
25-
70
10
10
250
500
750
1000
1250
1500
1750
2000
250
500
750
1000
1250
1500
1750
2000
Time
Tire
(a) Convergence curve of the dual variable of the regulated (b) Convergence curve of the dual variable of the Factor Wasserstein DRO (BSGD) (ew = 0.03).
KL DRO (gradient descent) (fKL = 0.1).
Figure 3: Policy evaluation convergence curves of the dual variables. The solid line and shades are averages and standard deviations over 20 runs.
0.8
0.6
50
40
30
20
0.0
250
500
750
1000
1250
1500
1750
2000
0
250
500
750
1000
1250
1500
1750
2000
Time
(a) Convergence curve of the policy parameters
(b) Convergence curve of the dual variable
Figure 4: Policy learning convergence results of the regulated Wasserstein DRO (BSGD) (ew = 0.03). The solid lines and shades are averages and standard deviations over 20 runs.
25
044
20
0.3
15
250
500
750
1000
1230
1500
1750
2000
0
250
500
750
1000
1250
1500
1750
2000
Time
Time
(a) Couvergence curve of the policy parameters
(b) Convergence curve of the dual variable
Figure 5: Policy learning convergence results of the Factor KL DRO (gradient descent) (EKL = 0.1). There is no variance since the gradient descent method is deterministic.
25


Reviews:
[{"review_id": "rnoDuEfPn6", "reviewer": "Reviewer_oaWR", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The strength of this study lies in the proposal of the distributionally robust policy value and algorithms appropriately designed for the policy value. The authors effectively highlight the issue of distributional overlap in DRO and introduce a new policy value that addresses this concern. Subsequently, they propose OPE and OPL algorithms and reveal the theoretical properties. The paper makes insightful contributions, and I am inclined to support its acceptance.\n\nDespite the valuable contributions, I have two principal concerns:\n\n1. On the necessity of considering DRO in OPE and OPL.\n\nI am curious when the problem setting is useful. For example, in OPE and OPL, we usually assume that $E[Y|X]$ is identifiable. Once $E[Y|X]$ is identified, it is at least robust for distributional change in $X$. Furthermore, if we do not know $E[Y|X]$ and do not assume common support, it seems that OPE and OPL become infeasible. I mean that the necessity of DRO in OPE and OPL significantly depends on the underlying assumptions. I would like the authors to clarify these points, especially what assumptions are required.  \n\n2. On the use of the Wasserstein distance. \n\nAs the authors discuss, the use of the Wasserstein distance is necessary when there is no overlap support, however, without any assumptions on either the functional form of the costs. In fact, the authors assume some restrictions or knowledge on the functional form. I believe that there is a tradeoff between weakening the overlap assumptions and imposing additional assumptions on the related random variables.  \n\nWhile the definition of policy value and the methods proposed are persuasive, the study's framework seems somewhat contrived within the context of OPE and OPL. It appears that the problem setting was adapted to incorporate the Wasserstein distance within OPE and OPL rather than applying the Wasserstein distance to address issues within the existing problem setting. The presentation would be improved by clarifying what assumptions are the same and what assumptions are different compared with the original OPE and OPL problems.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This study considers distributionally robust optimization (DRO) in off-policy evaluation (OPE) and off-policy learning (OPL). First, in Definition 2.1, the authors define the distributionally robust policy value $V(\\pi)$ as\n$$V(\\pi) := \\sup_{P_x}E_{x\\sim P_x}[E_{a\\sim\\pi(\\cdot|x)}[\\sup_{P_{x,a}}E_{P_{x,a}}[Y]] ],$$\nwhere the supremums are taken over perturbations within small deviations computed by the Wasserstein distance. \nThen, they propose OPE and OPL algorithms based on the policy value with strong theoretical guarantees. Theorem 3.4 and 5.1 show seemingly tight upper bounds for OPE and OPL, respectively.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I felt that this manuscript is a bit hard to read because the problem setting is unclear in the sense of OPE and OPL. In OPE and OPL, the difference between the potential outcomes and observations plays an important role. In this sense, I believe that observations should be described in the problem-setting section (Section 2). However, they are described in Assumption 3.1 with a new variable $\\xi$. It would be more desirable if these definitions were mentioned in the problem-setting section. \n\nFurthermore, the overall assumptions are incrementally added as the section goes on, which makes this study hard to understand. Some assumptions are more related to the problem setting itself, as well as the difference between the potential outcomes and observations. For example, the knowledge about the functional form of $y(\\cdot)$ and the existence of $xi$ should be clarified in Section 2.\n\nAdditionally, the authors should clarify which variables are known or unknown before the OPE and OPL. For example, is $\\pi$ known to us before the OPE and OPL?\n\nMinor comments.\n1. page 5: set(Hu & Hong, 2013) -> set (Hu & Hong, 2013) (space).", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xwC3fvH2Zu", "reviewer": "Reviewer_Z8aE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths \n\n1. OPE and OPL under distribution shifts are important problems. The proposed DRO with Wasserstein distance is motivated to improve over previous works of DRO with KL divergence. The proposed Wasserstein DRO framework is intuitive. \n\n2. Motivated by the high computation cost of Wasserstein DRO, the authors further proposed a regularized Wasserstein DRO problem and a biased stochastic gradient descent algorithm. \n\n3. Besides the proposed new methods, theoretical analysis is another major contribution of the paper. Both convergence rate and sample complexity guarantees are analyzed for the proposed solutions. \n\n\nWeaknesses and questions:\n\n1. Assumptions. I appreciate that the authors clearly stated the assumptions. Assumption 3.3 is standard in OPE literature while assumption 3.1 is new to me. The necessity and benefit of this assumption are unclear: what would be the applications where assumption 3.1 is necessary? How does this assumption change the method/solution? Are new theoretical tools introduced to handle this assumption? I would also suggest the authors to further clarify the technical novelty compared to existing works in general.\n\n2. Experiments. The authors are suggested to include details of KL DRO, and explain why not to compare with  KL DRO-based methods such as Si et al., 2020, Kallus et al., 2022, Mu et al., 2022. Regarding OPL, it is unclear why there is no baseline to compare with.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studied off-policy evaluation (OPE) and off-policy learning (OPL) for contextual bandits under distribution shifts. The authors proposed  distributionally robust optimization (DRO) with Wasserstein distance, and its regularized version with a biased stochastic gradient descent algorithm for policy learning. The authors analyzed convergence rate and sample complexity guarantees. Empirical results on a stroke trial dataset validated the effectiveness of proposed Wasserstein DRO methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see weaknesses and questions.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "v5MLFh4kOK", "reviewer": "Reviewer_8pbY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \nThe paper is well-written despite a few typos, see minor remarks\nMany contributions, including theoretical guarantees and empirical evaluation\nThe premise, that KL is not well suited for certain distribution shift tasks, is relevant and should be better known. \n\nMajor remarks:\n\nIn algorithm 1, is it not a problem that you use the same batch of sample to estimate the two gradients? I think you are introducing correlations and you should sample 2 separate batches. Can you please explain why you are allowed to do that? \n\nMinor remarks: \nSome typos are left so I recommend the authors to double check everything. I highlight a few below. \n“ in practices”: incorrect\n\n“ as shown in (Zhao & Guan, 2018) Proposition 2,”: awkward use of citation\n\n“ The goal of the OPE problem is to bound the difference between Vˆ (π) and V”: no the goal of the OPE problem is to estimate the value of the policy offline. *Your* goal is to bound the difference between the estimator and V. \n\n“Although the radius can be estimated when similar datasets that include both the testing and training sets are available, using the KL divergence, as discussed before, to calculate this radius poses additional challenges”: This sentence is upside down and must be rewritten", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes to use the Wasserstein distance instead of the KL in a Distributionally Robust formulation of Off-policy evaluation and off-policy learning for contextual bandits. \nThe paper is well organised: it starts from the direct definition of the problem  (2.1) and progressively highlights the computational issues and limitations that arise as we try to solve it. Through various clever ideas of regularizations and relaxation, an efficient estimator and algorithm is be provided with theoretical guarantees (Th 4.2 for the estimator, Th 5.1 for algorithm 1). Finally, an empirical study is done on medical data and showcases the benefits of the proposed method.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "A pass on the text to clean up some typos and convoluted sentences.\n\nDepending on the answer to my main question (major remarks), a modification may be needed in Algorithm 1.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: I agree with the reviewers that the paper is providing a very well-rounded contribution to the problems of policy evaluation/optimization in contextual bandits under distribution shifts. Especially, this work seems to properly address some prior works limitations, which either model the distribution shift through KL divergence or do not provide comparable theoretical results for Wasserstein DRO.

While one reviewer is concerned by the relevance of the problem setting, I personally think it is sufficiently motivated and that the contribution will interest a sub-community of TMLR.
Thus, I am recommending to accept the paper as is. 

I am further proposing a "featured certification" in view of its technical quality, which has been mostly acknowledged by all the reviewers.
Decision: Accept (poster)