Factored DRO: Factored Distributionally Robust Policies for Contextual Bandits
Anonymous Author(s)
AfﬁliationAddress email
Abstract
While there has been extensive work on learning from ofﬂine data for contextual multi-armed bandit settings, existing methods typically assume there is noenvironment shift: that the learned policy will operate in the same environmentalprocess as that of data collection.
However, this assumption may overly limit theuse of these methods for many practical situations where there may be distributionshifts.
In this work we propose Factored Distributionally Robust Optimization(Factored-DRO), which is able to separately handle distribution shifts in the contextdistribution and shifts in the reward generating process.
Prior work that eitherignores potential shifts in the context, or considers them jointly, can lead to performance that is too conservative, especially under certain forms of reward feedback.
Our Factored-DRO objective mitigates this by considering the shifts separately,and our proposed estimators are consistent and converge asymptotically.
We alsointroduce a practical algorithm and demonstrate promising empirical results inenvironments based on real-world datasets, such as voting outcomes and sceneclassiﬁcation.
Introduction
In many real-world applications we often have access to historical data about past decisions andtheir associated outcomes.
This data can be leveraged for off-policy learning to optimize for adecision making rule, called a policy that is more effective for the given application than thatused for data collection.
Off-policy learning from historical data for contextual multi-armed banditsettings has been well studied, due in part to the myriad potential applications from personalizedrecommendation
[Strehl et al., 2010] to optimizing medical treatment [Kim et al., 2011].
Priormethods often focus on the stationary setting where the learned policy will be deployed in a settingwith the same underlying data generation process that produces the observed contexts and rewardsupon intervention [Precup, 2000, Liu et al., 2018, Kallus and Uehara, 2020a,b, Liu et al., 2020].
However, in many settings, this assumption is not true and the distributions may shift.
As a motivating example, consider the problem of increasing political participation during electionsusing various interventions to encourage voting.
In the past, different strategies have been implemented, ranging from reminding people of their own voting history, to reminding them of the votinghistory of their neighborhood, etc.
[Gerber et al., 2008].
Leveraging this historical data, an importantquestion is: what strategy should be used during the next election to maximize political participation ?
Addressing this question requires careful consideration of two potential sources of distribution shift:
(a) How much has the population demographic in the area changed since the past data was collected?
(e
g. the percentage of single person households may have changed due to industrial developments).
and (b) How much has the effectiveness of each strategy on
each demographic group
changed (e g.as technology develops,
the effectiveness of mail- based interventions may change).
Such challengesare ubiquitous in real-world problems
For example, in automated healthcare, data from patients of one locality may be used to develop treatments for other localities where the effectiveness oftreatments may differ [Finlayson et al., 2021].
Similarly, online recommendation systems need toaccount for both the changing interests and the demographic shifts of its user base [Koh et al., 2021].
Developing a general procedure for addressing these challenges can provide a reliable solution forsuch applications with high stakes in terms of societal behavior, public health, and monetary assets.
Without considering potential shifts, there is a high risk of using suboptimal strategies, however itis often difﬁcult to predict exactly the shifts that will occur.
Distributionally robust optimizationformulations such as ours attempt to address this uncertainty over potential shifts by learning policiesthat attain good performance on all potential distributions within a user-deﬁned set of distributions
tobe considered (denoted the uncertainty set ).
Prior methods [Si et al., 2020, Mo et al., 2021, Kalluset al., 2022] that aim to provide robustness against distribution shifts do not separate the differentsources of shifts, and consider them in a joint formulation.
These methods may be disadvantageousas they lack the ﬂexibility to separate the sources.
Consequently, some of these methods mayproduce degenerate results under binary reward feedback, an issue theoretically proven to occur inthe DRO literature in supervised learning [Hu et al., 2018].
Due to this, methods that better usestructure to mitigate this issue have been an active area of research in the supervised learning DROliterature [Sagawa et al., 2019, Duchi et al., 2019].
We leverage this insight that better usage ofstructure is important, as well as the insight that it is common to have expert opinion about the amountof distribution shift that may occur between the two sources, to obtain better results.
In this work we address such challenges under the framework of contextual bandits.
We propose our
Factored Distributionally Robust Optimization (FDRO) algorithm that uses past data to search for apolicy that has the best worst-case performance, when the context (e.g., population demographic)or/and the reward model (e.g., inﬂuence of strategies on voting behavior) can be different from theones during data collection.
Towards this goal, we make three primary contributions: I. A factored distributionally robust optimization formulation for contextual bandits that providesend-users the ﬂexibility to separately specify different amounts of potential shifts in the twomain sources of uncertainty: the context distribution and the reward distribution.
Our factored formulation can overcome some failure cases of non-factored prior methods Si et al.
[2020].II.
We establish theoretically that the proposed estimators provide asymptotically consistent evalua- tion of any policy’s worst-case performance under distribution shifts.
Additionally, the estimators provide asymptotic convergence to a policy with the best worst-case performance.
III. Building on the theoretical insights, we develop a practical algorithm, FDRO, that uses approximations based on moment generating functions.
We show that FDRO provides promising empiricalresults on several domains based on real-world challenges.
Related Works
Both off-policy learning for bandits and reinforcement learning, as well as distributionally robustoptimization (DRO) for supervised learning and stochastic optimization have rich histories of research.
To brieﬂy summarize, methods for off-policy evaluation/learning predominantly assume that thedeployment environment for the learned policy will be the same as the environment the historicaldata was collected in [Precup, 2000, Jiang and Li, 2016, Thomas and Brunskill, 2016a,b, Gelada and
Bellemare, 2019, Liu et al., 2018, Imani et al., 2018, Yang et al., 2020, Jiang and Huang, 2020, Kallusand Uehara, 2020a,b, Liu et al., 2020].
In contrast, DRO methods for tackling distributional shiftshave predominantly focused on settings like supervised learning and do not consider decision makingaspects present in (contextual) bandits [Bertsimas and Sim, 2004, Hu and Hong, 2013, Namkoongand Duchi, 2016, Shapiro, 2017, Duchi et al., 2016, Duchi and Namkoong, 2018, Hashimoto et al.
,2018, Duchi et al., 2019, Sagawa et al., 2019, Oren et al., 2019, Michel et al., 2022].
In addition to the prior work already previously mentioned in Section 1, there have been a few recentpapers investigating DRO for use in bandits and reinforcement learning.
Some works use DRO toaccount for the mismatch between the true distributions (of the context and rewards) and the observedempirical distributions [Swaminathan and Joachims, 2015a,b, Faury et al., 2020, Sakhi et al., 2020].
These methods account for the potential errors of ﬁnite-data estimation, and, under inﬁnite data, theyconverge to the best training distribution policy.
However we are concerned with the setting when thetest distribution may be different, and the best training-distribution policy may not be optimal.
There has also been other work that consider settings when different forms of information are availablea priori.
For example, knowledge about the distribution of covariates during evaluation [Kato et al.
,2020], bounds on the ratio of train-eval covariate distribution [Hatt et al., 2021, Giguere et al.
, 2021],bounds on train-eval performance differences [Chandak et al., 2021], causal directed acyclic graph[Subbaswamy et al., 2019, Singh et al., 2021], or when there is data through time that can be used tomodel potential non-stationarity
[Thomas et al., 2017, Chandak et al., 2020, Xie et al., 2020, Qin and
Russo, 2021].
While relevant, they look at complementary problems.
Preliminaries
In this work, we consider the batch setting where the goals is to learn from an existing dataset D .
Weconsider a setting with discrete actions, where A denotes the set of k discrete actions, A .
Let X ⇢ R ddenote the set of contexts, X .
A policy ⇡ denotes a mapping from a context X to a distribution over Afrom which an action A can be sampled.
For a given context X and action A we observe a stochasticreward R .
We assume we the dataset, D consists of n independently collected datapoints , each a(context, action, reward) triplet:
D =
{ X i , A i , R i } ni =1 .
For this dataset, we assume the actions areselected with some data collection policy ⇡ 0 where A i ⇠ ⇡ 0 ( ·| X i ) .
We assume there is an underlying 105 context distribution P x 0 , and for each context-action pair there is an underlying reward distribution P r 0 | X,A .
We additionally assume the observed contexts, X i are independent and identially distributed 107 (i.i.d.) draws from P x 0 , and the observed rewards are i.i.d.
draws from P r| X
i ,A i .
In this work we consider using D to learn a policy ⇡ that will achieve the highest worst-case-rewardwhen the deployment environment may be shifted from the one of data collection.
Given a ﬁxedpolicy, notice that there are two ways the environment can shift and change the distribution of data weobserve.
First the context distribution can shift so the inputs we see are different.
Second the rewardgeneration model can shift, changing the reward distribution of each (context, action) pair.
We assumethe exact shift is not known, so we would like to learn a policy that would do well on all environmentswithin a user speciﬁed set of distributions of the context distribution P x and reward distributions P R | X,A .
Formally, given D we aim to ﬁnd the policy that optimizes the following distributionally 116 robust objective:
[Equation]
In (1), the uncertainty sets P r 0 | X,A and P x 0 describes the set of possible distributions P r and P x , andthe optimization set P r 0 | X,A is conditioned on random variables X ⇠ P x and A ⇠ ⇡ ( ·| X ) .
We term 119 our objective Factored-DRO as we have two layers of distributionally robust optimization, ﬁrst overreward shifts P R | X,A , and another over context shifts P x .
For theoretical analysis, we make the following assumptions common in prior literature [Si et al.
,2020, Swaminathan and Joachims, 2015a, Kallus and Uehara, 2020a].
Assumption 1.
Assumptions on ⇡ 0 and the reward distribution
(i) Unconfoundedness: The rewards are independent with A conditioned on X: ( R X,A 1 , R X,A 2 ... ) ?? 125 A
| X
(ii) Bounded reward support: 0  R  R M for all rewards for a ﬁnite R M
(iii) Coverage/Sufﬁcient exploration: ⇡
( A | X )  ✏ ⇡ for all ( X, A ) 2 X ⇥
A for some ✏ ⇡ >
(iv) Positive density of rewards: Let f 0 ,R | X,A be the pdf (continuous case) or p R
| X,A be the pmf 129 (discrete case) of the reward distribution P R | X,A .
Then assume f 0 ,R | X,A  ✏ r (continous) or 130 p
R | X,A  ✏ r (discrete) on the support of of R for all ( X, A ) 2 X ⇥
A for some ✏ r > 0 .
We make an additional assumption, which we do not require in our practical algorithm (Section 5):
Assumption 2.
Discrete contexts with sufﬁcient mass: The context set
X is a ﬁnite discrete set andthe probability mass of each context is lower bounded with P x 0 ( x = X )  ✏ X for some ✏ X >
Note about model misspeciﬁcation: Context shifts are an issue when we have model misspeciﬁcation and cannot model the best action for each context.
In these cases, the learned model encodesthe biases of the population distribution.
Model misspeciﬁcation can occur when the model is under-speciﬁed, for example, to reduce variance (bias-variance tradeoff), for interpretability, or dueto unobserved variables.
Misspeciﬁcation also occurs when the model is over-speciﬁed, where itmay overﬁt to irrelevant features(ex. neural networks [Sagawa et al., 2019]).
The real world is oftencomplex and difﬁcult to model correctly, so throughout this work we assume the misspeciﬁed setting.
Prior Work with Joint Shifts: Recently Si et al.
[2020] proposed an algorithm for solving a
DRO formulation for contextual bandits that does not separate the context and reward shift.
Theyconsider the optimization problem: ⇡ ⇤ = argmax ⇡ inf P 2 P P 0 E P [ R ] , where P encapsulates the jointdistribution of contexts and rewards and P P
= { P : D kl ( P || P 0 )
  } .
Degenerate Outcomes with binary rewards under the joint objective: Unfortunately, the jointformulation has a key limitation in the common setting of binary reward outcomes (ex, clicked ornot-clicked[Chu et al., 2011], voted or did-not-votedGerber et al.
[2008], etc).
In supervised learning(SL), similar DRO formulations have been theoretically proven to fail with certain loss functions.
Hu et al.
[2018] prove that when features are continous, and 0-1 classiﬁcation loss is used, the DROobjective is equivalent to the standard, non-robust optimization of the loss function on the trainingdistribution.
Intuitively, DRO can be seen as changing the weights of the datapoints based on theincurred loss, and increasing the radius parameters increases the weights on datapoints with largerincurred losses.
However, under 0-1 loss, DRO upweights all missclassiﬁed examples equally, whichdoes not change the decision boundary.
Consequently, in SL, methods that can mitigate this issue byusing structure, such as Duchi et al.
[2019] or Sagawa et al.
[2019], have been proposed.
The sameissue holds in the binary reward case for the joint formulation and we observe this issue empirically(see Section 6).
1 Intuitively, the lower reward points for each context action pair are all upweightedequally so the learned policy is equivalent to that optimized for the training distribution.
Our methodwhich trades-off need to set an additional hyperparameter with better structure usage to separate outthe shifts mitigates this issue.
Factored DRO for Policy Evaluation and Learning
In our work, we consider uncertainty sets deﬁned on the KL radius from the data collectiondistributions.
Due to the joint decomposition, we require the user to specify two KL-radiushyperparameters,  c and  x , for rewards and contexts respectively as opposed to a single hyperparameter required in prior work.
The uncertainty sets deﬁned in Equation 1 are deﬁned as
P r 0 | X,A = P r : D kl ( P r , P r 0 |
X
,A ) <
c and P x 0 =
P x : D kl ( P x , P x 0 ) <
x .
We can decompose our DRO objective from Equation 1 into two steps, ﬁrst a policy evaluation stepwhich given a policy evaluates the worst case reward policy reward.
Notice the policy evaluationstep itself can be separated into two steps: a step considering the potential reward generation shift( Q ( ⇡ , X, A ) ) and a step over potential context shift to arrive at the desired value ( V ( ⇡ ) ):
[Equation]
Second, we have a policy optimization step over the worst case policy rewards:
[Equation]
In this section, we present our estimators for policy evaluation and learning and asymptotic convergence guarantees.
We give proof sketches of lemmas and theorems throughout and the full proofs forthose not given in the text can be found in the appendix. Policy Evaluation: We start by discussing the policy evaluation step.
Equation 2 as written is adifﬁcult constrained optimization over probability distribution space.
Fortunately, we can use strongduality to derive an equivalent simpler unconstrained convex optimization over a single scalar:
Theorem 1 (Strong Duality) .
The optimization problem in equation 2 is equivalent to solving:
[Equation]
where Q X,A =  min ↵ c
[Equation] Proof Sketch This is shown utilizing Theorem 1 of Hu and Hong [2013].
We will use this easier to optimize dual form to derive estimators for policy evaluation.
We will ﬁrstconsider the potential shift of the reward distribution of each (context, action) pair and calculate anestimate of Q X,A of Equation 5.
We will then use these values to perform the second DRO layer overcontext shifts to estimate V ( ⇡ ) (Equation 4).4.1.1 First DRO, Reward Shift:
For theoretical analysis, we consider the setting when the contextand action sets are both ﬁnite and discrete (Assumptions 1.2 and 2).
Therefore we have a ﬁnitenumber of unique ( X, A ) pairs, each which occurs with a non-zero probability.
For a given ( X, A )pair, let R X,A correspond to the set of the observed reward R i every time ( X, A ) is observed in
D : R X,A
= { R i : ( X i , A
i , R i ) 2 D
if X i
=
X,
A i = A } .
Additionally, let m xa =
|R X,A | 189 denote the number of datapoints in each ( X, A ) bin and let  c denote the user speciﬁed KL-radiusof the ambiguity set.
For every unique ( X, A ) we propose the following estimator for worst casecontext-action reward for Q X,A in equation 5: [Equation]
Using Theorem 1 of Si et al.
[2020], we can show the above estimator converges with O p (1 / p n ) :
Lemma 1 (Convergence of Reward Shift Estimates) .
Suppose Assumption 1 is enforced, then:p
[Equation]
Q
[Equation]
,A ) d !
N (0 ,  2 ( ↵ ⇤ c ))
if ↵ ⇤ c > 0 and p n ( ˆ Q X,A  Q X,A ) p ! 0 if ↵ ⇤ c = 0 , where  2 ( ↵ c )
=
↵ 2 c
{+E [ exp (  R/ ↵ c )] 2 E [( exp (  R/ ↵ c )  E [ exp (  R/ ↵ c )]) 2 ] and ↵ ⇤ c is the value of ↵ cthat achieves the min for equation 
+}
Note that the objective inside the minimization in equation 6 is convex in ↵ c for ↵ c > 0 , thereforemany optimization methods (ex. gradient descent, bisection search) are able to achieve the globaloptimum when ↵ c 6 = 0 .
We give the gradients and a proof of convexity in the appendix.
The ˆ Q X,A ’sdo not depend on the ⇡ we are trying to evaluate, therefore they only need to be computed once andcan be used throughout the rest of the policy evaluation and optimization
procedure.4.1.2 Second DRO, Context Shift: Let  x denote the user-speciﬁed KL-radius for context distributionshift, then with the computed ˆ Q X,A for each ( X, A ) pair, we can complete the policy evaluationestimation and estimate V ( ⇡ ) from equation 2.
[Equation]
We can show ˆ V ( ⇡ ) converges with o p (1) :
Theorem 2 (Convergence in probability of policy evaluation) .
Under Assumptions 1 and 2 ,
[Equation]
every positive
[Equation]
Proof Sketch: We ﬁrst demonstrate pointwise convergence of the optimization objective. We thenutilize a result from stochastic equicontinuity [Newey, 1991] to obtain uniform convergence, afterdemonstrating a required Lipschitz condition for stochastic sequences is [-also convex and easy to optimize for↵ x > 0 .4.
-]satisﬁed.4.2 Policy Learning: We now consider the optimization over policy of equation 3.
Deﬁne the optimalrobust V ⇤ = min ⇡ V ( ⇡ ) and the set of ✏ -optimal policies ⇧ ⇤ ✏ = { ⇡ : V ( ⇡ )  V ⇤  ✏ } .
Assume wehave a procedure for ﬁnding the optimal ˆ ⇡ ⇤ that satisﬁes ˆ ⇡ ⇤ = argmax ⇡
ˆ V ( ⇡ ) , we can show ˆ ⇡ ⇤ , willconverge to an item in ⇧ ⇤ ✏ {+! 
+}
with o p (1) .
Theorem 3 (Convergence in Probability of Policy Learning) .
Under Assumptions 1 and 2 ,
[Equation]
Our estimator for solving for the worst case reward shift value, ˆ Q X,A for each context action pairin our ﬁrst DRO step (section 4) relies on having a sufﬁcient number of datapoints per each ( X, A )to approximate E R ⇠ P R | X,A [exp(
R/ ↵ c )] = 1 /m xa
P j exp(  R j / ↵ c ) .
However, in many settingscontexts may not be discrete, and, in these cases, each continuous context may only be observed once.
One approach can be to discretize the context space, however selecting the correct discretizationis challenging.
Too coarse a discretization introduces looseness and leads to overly conservativeresults, but too ﬁne a discretization may not leave enough items in each bin to accurately estimate
E R ⇠ P R | X,A [exp(  R/ ↵ c )] .
Building on the theoretical insights, instead of binning we use function 223 approximation that may not have the same convergence guarantees but as we show in Section 6 it canbe quite accurate and lead to better results than the baseline methods.
We develop an empiricallystable method for estimating E R ⇠ P R | X,A [exp(  R/ ↵ c )] using a moment generating function Taylor 226 series approximation.
The key feature of our method is that it does not require ﬁtting the exponentialdependence on R and ↵ c .
Approximating E R ⇠ P R | X,A [exp(  R/ ↵ c )] : We ﬁrst notice that E P R | X,A [exp(  R/ ↵ c )] is a momentgenerating function (MGF) of random variable R ⇠ P R | X,A with t = 1 / ↵ c .
We can use the Taylorseries expansion, M R ( t ) =
P 1 i {+=
+}
t i i !
E [(  R ) i ] , to estimate an order m approximation.
We do this byﬁtting m function approximators, F j for j = 1 , ...m to approximate the ﬁrst m moments.
We train F j to predict (  R i ) j from ( X i , A i ) , when trained on all the data, F j should model the j th moment
E P R | Xi,A
[(  R ) j ] .
Let D X represent the sequence of contexts in the dataset, then for every context
X i 2 D X and every a 2 A and any ↵ c , we use our approximators F j to estimate the desired value: [Equation] [Equation] any
[Equation]
This approximation achieves better regularity as we are no longer ﬁtting the dependence on ↵ c .
Practically, we ﬁnd order m = 5 can often be quite accurate.
Other Methods: Another method is to train a function approximator to predict exp(  R i / ↵ c ) giventhe corresponding values of X i , A i and multiple values of ↵ c value sampled from the relevant range.
Then the trained model should estimate E R ⇠ P R | X,A [exp(  R/ ↵ c )] .
However, we found this method 240 to be unstable at small values of ↵ c due to the exponential dependence, so we develop a method thatdoes not ﬁt this dependence.
Additionally, note that only estimating the ﬁrst moment, E [ R ] and usingexp(  E [ R ] / ↵ c ) is a biased estimator of E [exp(  R/ ↵ c )] due to Jensen’s inequality.
Complete Algorithm:
With our estimator ˆ E P R | Xi,A [exp(  R/ ↵ c )] we proceed with our method byﬁrst calculating n ⇥ k values of ˆ Q X i ,a for each context in the dataset and each action in the action set.
To do this, we iterate through ( X i , a ) 2 ( D X ⇥ A ) and solve the following convex problem in ↵ c :
[Equation]
We then carry out the procedure detailed in Section 4 onwards with these ˆ Q X i ,a .
Algorithm 1 in
Appendix A.1 provides our full algorithm for policy evaluation and optimization, that considers boththe discrete binned case (Section 4) and the function approximation case just discussed.
Experiments
We test and analyze our algorithm empirically in three settings, including two derived from real data.
While our method works for any policy that can be optimized, we consider learning the parameters ✓of a stochastic softmax linear contextual bandit policy of the form:
[Equation]
where
[Equation]
( X, A [Equation]
A
[Equation]
This policy ⇡ ✓ is differentiable so we optimize ✓ using gradient descent.
For the function approximators F j given in section 5, we use exponential-kernel linear regression models.
[Table]
We use random projection to project the feature space of scene to avoid ill-conditioned behavior.
We compare against three baselines: (i) baseline-IS : A non-robust baseline that maximizes thereward calculated using importance sampling policy evaluation on the training data, (ii) baselinePOEM
[Swaminathan and Joachims, 2015a]
POEM improves the baseline-IS method by addingvariance regularization to account for ﬁnite sample error, and (iii) baseline-DRO [Si et al., 2020]as mentioned throughout the text, considers the shift over the joint (context, reward) distribution asopposed to separating them.
6.1 Evaluation Domains
We ﬁrst provide a brief overview of our evaluation settings before discussing our results. See theappendix for full descriptions of the environments.
In all settings, to construct example shifts wecluster the datapoints by one of the context covariates into what we denote as Coarse Groups (cg).
We then consider settings where the training distribution may consist of a different composition ofcgs than the deployment/testing distribution.
For each environment we give details in Table 1, suchas the distribution of cgs in the train and test data as well and the average reward on the training setof each (cg, action) pair.
Simulated Setting:
Our simulated setting has 4 cgs.
To minimic potential misspeciﬁcation that mayoccur when using real data, unknown to the designer, cg 1 and 3 ( cg 1 , 3 ) are identical, as are cgs 2 and4 ( cg 2 , 4 ).
Each context is composed of the cg and 4 other values sampled from uniform distributions.
For each of the 4 actions and cg, each (cg, action) pair has a set of true reward parameters  ⇤ .
Thelinear combination of  ⇤ and the contexts determine the bernoulli parameter of the observed reward.
Voting Setting: The voting dataset by Gerber et al. [2008] contains data collected from a randomizedcontrolled trial-style study testing various actions (1 control and 4 treatments) to increase voter turnoutin the Michigan 2006 primary election. We consider the recorded binary voting outcome as ourreward.
Because the non-control actions require additional effort of printing and mailing a lettercompared to the control (do nothing), we slightly modify the reward by lowering the non-controlaction rewards by 0.09 to simulate this cost of effort.
We use the context feature of whether theindividual voted in the 2004 primary to divide data into coarse groups.
To create a context distributionshift between train and test, we sampled different proportions of datapoints from each group to put inthe train distribution.
Additionally in this experiment we consider the case where we do not haveaccess to the coarse group at test time so it cannot be used in learning the policy.
This illustrates amore challenging and a realistic case where there is shift in the unobserved context features.
Scene Setting: We additionally test in a setting derived from the multiclass-supervised learning
Scene classiﬁcation dataset from the LibSVM [Chang and Lin, 2011] repository.
We use a slightlymodiﬁed version of the supervised to bandit conversion method from prior literature [Swaminathanand Joachims, 2015a, Agarwal et al., 2014].
In this method we assume we do not observe the correctlabel directly.
Instead each label is an action and selecting to correct action/label gives positivereward, while the wrong action/label gives negative reward.
We follow a similar procedure to create ashift as in the voting dataset so we give those details in the Appendix.
6.2 Results
We use the training dataset to optimize the policies and evaluate on the test dataset.
For the votingdataset we use importance sampling to evaluate the policies.
For Simulated and Scene, we have the true rewards for all action and use them for evaluation.
Recall that all the robust algorithms haverobustness parameters.
For FDRO we have two, the KL-radius of the reward distribution uncertaintyset  c and the radius of the context distribution uncertainty set  x .
Due to our train-test construction,our many of our settings predominantly have context shifts.
We leverage this information to set our FDRO parameter for reward shift to  c = 0 .
We show our main results in the top row of Figure 1where we consider the effect of  x .
Baseline-DRO has a single parameter for the KL-radius over thejoint context-reward distribution ambiguity set,  which we compare with.
Baseline-POEM also hasa single parameter, the variance regularization strength,  .
However, Faury et [-al.[
-]al.
[2020] prove that POEM can be equivalently viewed from a DRO formulation.
So for comparison when reporting theresults, we ﬁrst translate  to the corresponding DRO-radius  = n  2 .
We demonstrate in all threesettings, FDRO is able to outperform all the baselines.
In Figure 1a, we consider the case of testing an algorithm ﬁt with a ﬁxed  x in different populations.
We plot the performance gain of each algorithm over baseline-IS as the test population is increasinglydifferent.
We made this environment especially challenging for FDRO by ﬁxing  x for FDRO whilesolving for the best  for each different test population for baseline-DRO and baseline-POEM.
Wesee that as expected when the test context distribution is nearly identical to the training contextdistribution, the robust policy learned by FDRO does not perform as well (negative reward difference).
However as the test distribution differs more and more, we see FDRO starts performing much betterthan the non-robust policy.
As predicted, in the contextual bandit setting with binary rewards, the jointformulation of baseline-DRO suffers from the same issue proven. Similar issue of degeneration tolearning the optimal train-distribution policy has also been drawn in the supervised DRO literature [Hu et al., 2018].
We additionally see that even with the advantage of using the best parameter per testpopulation, baseline-POEM still does not perform as well.
In Figures 1b and c. We consider ﬁxed training and test distributions and plot performance at differentvalues of  /  x for the voting and scene settings.
We see that with increasing  x , FDRO can learnrobust policies that perform well on the test set.
Baseline-DRO suffers from the same issue previouslydiscussed.
The performance of baseline-POEM decreases on both distributions.
As the varianceregularization weight increases, POEM is incentivized to learn low variance policies, such as choosinglow variance but low reward action, or reducing the variance from importance sampling by choosinga policy close to the data-collection policy, which may be suboptimal.
6.3 Discussion
Continuous Rewards and Reward Shift: In the main result we focus on the binary reward setting asit is a common setting for bandits however we also show results for a continuous rewards setting. Wemodify our simulations environment to use gaussian-distributioned rewards as opposed to bernoulli.
We see in Figure 1d that baseline-DRO is able to learn polices that perform well on the test set. However FDRO with extra structure information was still able to learn a better policy with better trainperformance at the same test performance.
Additionally, in our main results we primarily consideredcontext distribution shifts.
We create a new simple reward simulation environment (described in theappendix) where we can easily create a test environment with shifted rewards.
We see in Figure 1d,both FDRO and baseline-DRO are able to handle this this case well.
Utilizing unavailable features: We note that in two of the settings - voting and scene, we assumedthere were context features in the train set that are unavailable at deployment.
This situation occurs incases where those features would be too expensive or slow to measure at deployment.
It may alsooccur in cases, such as fairness considerations, where there are sensitive features that should notinﬂuence the decision policy.
The policy learned by FDRO does not leverage this unavailable feature,however if available in the training data, the method can still better utilize this information in the ﬁrst
DRO step when calculating the worst case rewards ˆ Q X,A .
Setting the  parameters: One limitation common to all DRO methods is that setting the parameters requires may be tricky.
Choosing  too large and accounting for too large a shift mayresult in overly conservative policies at deployment.
However, we can see from the results, often ourmethod works for many values of  .
For example in Figure 1a, we see that while we may incur asmall penalty if the distribution does not shift, we can learn a policy that has much larger gains if thedistribution does shift.
A standard method is to Another approach could be to utilize a human in theloop examination approach. For any policy and  , we can examine the worst case distribution. Theworst case distribution for DRO formulation with KL distance bounds is an Softmax scaled weightingof the datapoints with a temperature determined by optimal dual variable ↵ ⇤ .
If we have discretecontext features, or a easily discretized features, examine the worst case distribution of that feature. We plot the worst case distributions of the cg context feature of FDRO in Figure 1f in the Bernoullisimulated environment (exact calculation details in the appendix). For each  value, the end user canexamine the estimated performance on the training distribution as well as the worst-case distributionof some of the features to decide if the policy is reasonable. Potential negative societal impacts: As shown in our results, FDRO is able to learn robust policiesthat perform well in domains, even under shift.
However if the shift is not as large as predicted, thepolicy learned using FDRO may be overly conservative which can potentially cause negative impactin terms of lower than necessary performance.
To mitigate this risk, experimenters can carefully set using methods proposed in the previous section.
Conclusion
In this work we consider the batch learning of contextual bandit policies when the deploymentdistribution may shift from that of data collection.
We propose a distributionally robust optimizationformulation that separates the reward and context shift.
We demonstrated that our proposed estimatorsconverge asymptotically and then developed a practical algorithm FDRO that is able to overcomesome failure cases of prior methods and learn better performing robust policies.