A VOIDING O VERFITTING TO THE I MPORTANCE W EIGHTS IN O FFLINE P OLICY O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Ofﬂine policy optimization has a critical impact on many real-world decisionmaking problems, as online learning is costly and concerning in many applications.
Importance sampling and its variants are a widely used type of estimator in ofﬂine policy evaluation, which can be helpful to remove assumptions on the chosen function approximations used to represent value functions and process models.
In this paper we identify an important overﬁtting phenomenon in optimizing the importance weighted return and propose an algorithm to avoid this overﬁtting.
We provide a theoretical justiﬁcation of the proposed algorithm through a better per-state-neighborhood normalization condition and show the limitation of previous attempts to this approach through an illustrative example.
We further test our proposed method in a healthcare-inspired simulator and a logged dataset collected from real hospitals.
These experiments show the proposed method with less overﬁtting and better test performance compared with state-of-the-art batch reinforcement learning algorithms
I NTRODUCTION
Learning decision policies from ofﬂine data has a critical impact on real-world decision-making applications such as healthcare (Lei et al., 2012; Komorowski et al., 2018), recommendation systems (Li et al., 2010; Chen et al., 2019), and intelligent tutoring systems (Mandel et al., 2014).In real-world decision-making applications, we often need to use limited amounts of data about past decisions to yield better future performance in complex, hard-to-model domains.
Such ofﬂine policy optimization encounters two key generalization challenges: distribution shift and modeling bias.
Distribution shift arises because, under a new decision policy, the distribution of observations and actions will be different than the distribution of observations and actions in the ofﬂine data.
Modeling bias can occur when popular assumptions over the domain structure or model class (such as the Markov assumption, realizability, or the value function class being closed under the Bellman), fail to hold in the problem setting.
Learning from limited data, algorithms ignoring these challenges can produce policies and estimates of policy performance that are far more optimistic than they achieve in practice, similar to the overﬁtting in supervised learning.
To address this, prior work in batch reinforcement learning (RL) and contextual bandits mostly falls into two categories: (1) changing the estimator or loss function used in computing ofﬂine value estimates, or (2) constraining the policy search space.
In category (1), self-normalized IS (or weighted
IS) (Hesterberg, 1995; Swaminathan & Joachims, 2015b), variance regularization (Swaminathan &
Joachims, 2015a; Metelli et al., 2018), doubly robust methods(Dud´ık et al., 2011; Jiang & Li, 2016;
Thomas & Brunskill, 2016), and uncertainty regularized model-based batch RL (Yu et al., 2020;
Kidambi et al., 2020) are all approaches that aim to provide more stable and variance-penalized estimates of a policy’s performance to address the uncertainty due to limited data and distribution shift.
The second category of direct constraints on the policy class includes bandit learning under action/policy constraints with behavior probability (Sachdeva et al., 2020), value-based batch RL algorithms by constraining the policy disagreement (Kumar et al., 2019; Buckman et al., 2020), thresholding the action probability (Fujimoto et al., 2019;
Futoma et al., 2020), or thresholding the marginalized state-action probability (Liu et al., 2020), under the behavior policy.
Many of the works above rely on learning MDP models or value functions based on MDP assumptions.
However, the Markov assumption and function approximation assumptions are often unrealistic in real-world problems and hard to verify, which limits their eligibility in certain application problems.
In this paper, we consider learning an policy to optimize an importance weighted return, as it is free from Markov and other types of model assumptions.
We identify an unique overﬁtting that is not solved in prior work about policy optimization and propose a new constraint on policy named eligible actions (EA) to address the overﬁtting issue.
The overﬁtting issue comes from ﬁtting the importance weights in the weighted return objective.
The importance weighted return can effectively remove a certain subgroup of samples in the dataset by decreasing all weights to zero.
This context avoidance is unrealistic when the subgroup consists of a subset of initial states, as the policy does not get to control the initial state distribution.
Consider the following motivating example: an AI decision policy may learn from historical treatment logs to avoid all the logged treatment in the history for patients with severe conditions as they lead to low rewards.
This avoidance can be achieved with ﬁnite data if the action space is large, and such avoidance will result in high off-policy estimates.
However, in the real environment, a policy still should make good decisions for patients who arrive with severe conditions.
This propensity overﬁtting problem was studied in contextual bandit problems (Swaminathan & Joachims, 2015b).
While their solution effectively constrained the upper bound on importance weights on high-reward samples, the challenge of overﬁtting to avoid the low-reward samples remains.
Our contribution To address the issue that optimizing with propensity weights may avoid certain initial states, we propose the P olicy O ptimization with EL igible A ctions (POELA) algorithm.
We constrain the potential action set to the set of observed actions with similar states to prevent improper avoidance to lower-reward initial states.
Thus within a radius of a particular state, the sum of observed importance weights are always lower bounded.
We then study the empirical performance of POELA in a medical simulator and a real-world medical dataset.
Our approach ﬁnds much better decision policies than prior state-of-the-art algorithms in both domains.
O FFLINE P OLICY O PTIMIZATION
We study the problem of ofﬂine policy optimization in sequential decision-making under uncertainty.
Let the environment be a ﬁnite-horizon Contextual Decision Process (CDP)
(Jiang et al., 2017).
A
CDP can capture more general, non-Markovian settings (also sometimes refered to as a Non-Markov
DP (Kallus & Uehara, 2019b)).
A CDP is deﬁned as a tuple (cid:104)X , A , H, P, R (cid:105) , where X is the context space, A is the action space, and H is the horizon.
P =
{ P h } Hh =1 is the unknown transition model, where P h : ( X × A ) h − 1 → ∆( X ) is the distribution over next context given the history
P 1 : ∆( X ) is the initial context distribution.
Similarly, R = { R h } Hh =1 is the reward model and
R h : ( X × A ) h → ∆([ − R max , R max ]) .
In this paper we focus on learning policies that map from the most recent context to an action distribution, π : X → ∆( A ) .
This is optimal when the domain is Markov and can often be more interpretable and more feasible to optimize given ﬁnite data in the ofﬂine setting.
In ofﬂine policy optimization settings, we have a dataset with n trajectories collected by a ﬁxed behavior policy µ : X → ∆( A ) , and we aim to ﬁnd a policy π in a policy class Π with the highest value.
Policy gradient and optimization approaches do not rely on a Markov assumption on the underlying domain, and have had some encouraging success in ofﬂine RL (Chen et al., 2019).
Often these methods leverage an importance sampling (IS) estimator in policy evaluation: ˆ v IS ( π ) = 1 n (cid:80) ni =
[Equation]
(cid:16)(cid:80)
Hh
[Equation]
( i )
[Equation]
(cid:17) (cid:81)
Ht =
[Equation]
(cid:18) π ( a ( i ) h | x ( i ) h ) µ ( a ( i ) h | x ( i ) h ) (cid:19) .
The IS estimator is an unbiased and consistent estimate of the value under the following two assumptions:
Assumption 1 (Overlap) .
For any π ∈ Π
and any x ∈ X , a ∈ A , π ( a | s )
µ
( a | s ) < ∞
Assumption 2 (No Confounding / Sequential ignorability) .
For any policy π ∈ Π and µ , conditioning on the current context x h , the sampled action a h is independent of the outcome r h : H and x h
+1: H .
IS often suffers from high variance, which has prompted work into extensions such as doubly robust methods (Jiang & Li, 2016; Thomas & Brunskill, 2016) and/or methods that balance variance and bias.
Truncating the weights and using self-normalization has been shown to be empirically beneﬁcial both in bandit and RL settings (Swaminathan & Joachims, 2015b; Futoma et al., 2020): we refer to this as Self-Normalized Truncated IS (SNTIS):
[Equation]
v
[Equation]
where W ( i ) h :
=
π ( a ( i ) h | x ( i ) h ) µ ( a ( i ) h | x ( i ) h ) .
For the ease of notation, in the rest of this paper we deﬁne W ( i ) 1: h :=
(cid:81)
ht
=1 W ( i ) t , W ( i ) :
= W ( i ) 1: H , W = (cid:80) ni
=1 W ( i ) , and r ( i ) =
(cid:80)
Hh =1 r ( i ) h .
T HE O VERFITTING P ROBLEM IN I MPORTANCE W EIGHTED P OLICY O PTIMIZATION
We now highlight an important problem with using important sampling estimators such as the above during ofﬂine policy optimization.
For concreteness we will describe this in the contextual bandit setting using the self-normalized importance sampling estimator, but a similar results hold for unnormalized estimator.
We will also shortly describe in the main text how this issue persists in existing importance weighted off-policy learning methods and about the multi-step case in the appendix.
The key issue is that given a ﬁnite dataset, during optimization the IS estimator objective can beneﬁt policies that improperly avoid low reward states, inducing a lack of overlap on the empirical behavior data.
To see this, let v π ( x ) =
E [ r | x, a ∼ π ] , ˆ p ( x ) be the empirical probability mass/density over the contexts x in the dataset, and W ( x )
= (cid:80) i : x ( i ) =
x W ( i ) W .
We now decompose the importance weighted off-policy estimator into three parts.
[Equation]
The ﬁrst term is a supervised empirical value estimate whose only error is due to the error in the empirical context distribution sampled in the dataset, vs the true context distribution.
The second term captures the error caused by the difference between context distribution introduced by weights and empirical context distribution in the dataset.
The third term computes the difference between the weighted IS estimate of the value of the policy in a speciﬁc context x versus its true value v π ( x ) , and then sums this over all contexts.
The second term is of particular interest, because it highlights how the IS estimator of a policy may effectively shift the relative weight on the context space.
Note in the bandit setting (and in the initial starting state distribution for the RL setting) such shifting should not be allowed: the policy may control what actions are taken, but cannot change the initial context distribution.
Before illustrating the problematic behavior that can result from this type of objective function for ofﬂine policy learning, we ﬁrst further analyze if it is possible for W ( x ) to differ signiﬁcantly from hatp ( x ) , given the normalization constraint of π .
Expanding the difference between the empirical distribution over a context and W ( x ) :
[Equation]
While the ﬁrst term is asymptotically zero because W → N , the second term will be non-zero whenever the policy π puts non-zero weight on actions in contexts where the empirical behavior distribution differs signiﬁcantly from the true behavior distribution.
Such a difference between the empirical and true behavior distributions can easily happen.
This is especially the case in settings where the state-action space is large and the number of observed actions per state is small, which is common in recommendation settings, healthcare and many other applications.
A policy learning algorithm maximizing the objective in Equation 2 can exploit this difference in context weights and return a policy with an over inﬂated value estimate, as we now illustrate:
Example 1.
Consider a contextual bandit problem with |X| contexts and |A| actions in each context.
For half the contexts S p , the reward is 1 for one action and zero for others.
For the other half of the contexts, S n , the reward is -1 for half the actions, and -5 for the rest.
The true distribution over contexts is uniform.
Note in this setting the optimal policy would have a reward of 0.
The behavior dataset is draw from a uniform distribution over contexts and actions.
When the sample size |D| < |A| , we can assume there is only one observed positive reward in the dataset.
Then computing the policy to maximize the important sampling weighted objective will yield weights that are 0 on all contexts except for the context with the positive reward, through a policy π that selects actions that are not present in the dataset for all contexts whose observed actions lead to only zero or negative rewards: let A s i
= { a i : r ( s i , a i ) ≤ 0 } then π ( s i ) =
a j where a j / ∈
A s i .
The resulting IS/SNIS estimator of the value of π is 1 , which is much higher than the optimal policy value.
In addition, note that the returned policy π is likely to be worse than the optimal policy for any context where r ( s i , a i ) = − 1 , since that is the optimal reward possible for such states s i , and by π selecting an unobserved action a j in that state ( π ( s i ) = a j ), the policy π may select an action with worse true reward, r ( s i , a j )
= − 5 .
Note that in this example, there is no error in the importance weighted return for the contexts with positive W ( x ) (which is the third term in the decomposition shown in Equation 2).
The issue arises because to optimize the objective, it is possible to select a policy whose W ( x ) distribution over the contexts has zero weight on at least the half of the contexts with negative rewards S n , by assigning them actions not selected in the dataset.
This causes the second term in the decomposition (Equation 2) to be signiﬁcantly overestimated.
This is a unique overﬁtting phenomena in the counterfactual (or called off-policy) learning settings since in supervised learning it is not possible for a hypothesis π to manipulate the weights over the input contexts.
While this issue can arise in contextual bandits with large state and action space in small datasets, we show that it is even easier for this to occur in sequential reinforcement learning settings, even when only 2 actions are available in Example 2 and 3 in the Appendix.
Intuitively, the issue arises because in estimating the value of a new decision policy, it is acceptable to choose a policy that re-distributes the weights of actions within an initial context but not re-distribute the weights across initial contexts , since it is not a function of the actions selected.
It is well known that in importance sampling, the expected ratio of the weights should be 1: E y ∼ µ [ π ( y ) /µ
( y )]
= 1 .
In contextual policies, we expect that the expected weights should also be 1 for all initial contexts x 0 : E a ∼ µ ( a | x 0 )
[ π ( a | x 0 ) /µ ( a | x 0 ) | x 0 ] = 1 .
However, optimizing for a standard importance sampling objective (such as Equation 2) does not involve constraints that the empirical expectation of weights given an initial context ˆ E (cid:104) W ( i ) h | x ( i ) h )
| x ( i ) h (cid:105) (or the weights of n -step given initial context ˆ E (cid:104) W ( i ) | x ( i ) 1 ) | x ( i ) 1 (cid:105) ) is still close to one.
In order to optimize the future return, the optimization algorithm can optimize the weights of initial contexts in the dataset, instead of optimizing the weights on different actions/ action sequences in a given initial context.
This result may seem surprizing, given that under mild assumptions which are satisﬁed here (Assumptions 1 and 2) importance sampling is well known to provide an unbiased estimate of the value of a policy.
Our observations do not contradict this fact: while importance sampling will still provide an unbiased estimate given a policy, the policy optimization can exploit the ﬁnite sample error and its heterogeneousness across the policy class.
P RIOR WORK ON O FFLINE P OLICY O PTIMIZATION
Before we describe a method to alleviate this problem, we ﬁrst review prior work.
There is increasing interest in multi-armed bandits and ofﬂine RL to avoiding overly optimistic estimates of policies computed from ﬁnite datasets that can cause suboptimal policy learning, and prior work has discussed extrapolation error(Fujimoto et al., 2019), propensity overﬁtting (Swaminathan & Joachims, 2015b), and bandit error (Brandfonbrener et al., 2020).
The problem we identiﬁed arises speciﬁcally from deﬁcient support in the observed dataset.
To address potential overﬁtting, one line of work focuses on changing the objective used for policy optimization.
For example, Swaminathan & Joachims (2015b) has shown that policy optimization using vanilla importance sampling in contextual bandits will over-maximize the weights if the reward is positive and over-minimize the weights for negative reward.
Shifting the rewards will help with the latter (and avoid part of the context avoidance we describe) but can worsen the former.
The self-normalized estimator proposed to address the issue of over-maximizing reward (Swaminathan &
Joachims, 2015b) is equivariant to any constant shift in rewards, and will still suffer from the context avoidance issues we describe above.
Counterfactual risk minimization (Swaminathan & Joachims, 2015a) uses variance regularization based on the empirical Bernstein’s inequality.
However, this penalization is at the policy level and does not directly address the problem with avoiding contexts with low reward.
In Appendix Figure 2c we show the counterfactual risk minimization regularization with or without self-normalization requires a large dataset to perform well.
Later work (Joachims et al., 2018) extended norm-POEM to stochastic gradient descent settings for large-scale training.
Recent work (Brandfonbrener et al., 2020) discussed a similar overﬁtting issue as we describe and compared the performance of ofﬂine policy optimization and model/value-based method on this overﬁtting issue.
Those authors focus primarily on the negative result of the policy optimization approach and the advantage of the model/value-based method.
Doubly robust estimators (Dud´ık et al., 2011; Jiang & Li, 2016; Thomas & Brunskill, 2016; Kallus & Uehara, 2019a;b) have multiple beneﬁts but as long as the learned Q function is imperfect, policy learning can still overﬁt to the high/positive residual r − Q .
Pessimism under uncertainty approaches are promising (Kidambi et al., 2020; Yu et al., 2020) but have so far only been developed for Markov settings and are not robust to model class misspeciﬁcation.
Another line of ofﬂine batch policy optimization constrains the policy search space, for example to constrain target policies to be close in some distance to the behavior policy (Kumar et al., 2019;
Buckman et al., 2020).
A common approach is to constrain target policies to have a minimum conditional action probability under the behavior policy µ ( a | s ) (e.g. Sachdeva et al.
(2020); Fujimoto et al.
; Futoma et al.
; Liu et al.
This work has focused on algorithms and analysis for the Markov setting with additional model realizability assumptions, except Sachdeva et al.
(2020) in the bandit case and Futoma et al.
(2020) in the RL setting.
For example, the MDP value based approaches often require the value function class is closed respect to the Bellman operator and the value function is realizable.
Such assumptions are hard to verify in real-world domains.
This prior work does not address the issues we outlined, that occur due to deﬁcient support in the observed ﬁnite data rather than in the expected behaviors.
Our work can be broadly viewed as following in the recent line of work on pessimism under uncertainty, but adapted to provide policy search based ofﬂine learning method that does not require the Markov assumption or model realizability assumptions, and can achieve strong performance given a ﬁnite dataset.
L OWER B OUND OF ˆ E [ W | x ] AND E LIGIBLE A CTIONS
We have observed that the issue rises in policy learning because weight can be placed on unobserved actions for certain contexts, causing the empirical conditional expectation of weights given a context to be zero for such contexts.
To address this, one possibility is to constrain the conditional expectation of weights given a context to be 1 or lower bounded.
However doing so in inﬁnite/continuous context spaces is subtle: each context likely only appears once in the dataset, and requiring ˆ E
[ W | x ] =would be equivalent to only allowing a policy that exactly matches the observed logged actions.
We now propose a slight relaxation of the above proposal.
Intuitively the issue arises because without further constraints, policy learning can place large weight on unobserved actions in the dataset, for which our reward uncertainty is high.
The recent line of pessimism under uncertainty for model and value based MDP ofﬂine RL explicitly accounts for such statistical uncertainty through constraining or penalizing actions and/or states and actions for which there are limited observed data.
Similar to this, we introduce a pessimistic constraint on the policy class to be considered, and then demonstrate this allows us to constrain the empirical conditional expectation of weights given a context.
For policy learning, we create local constraints on the eligible policy class by deﬁning for each context x , dataset D and a given threshold δ , the eligible action set A ( x ; D
[Equation]
[Equation]
Intuitively this allows any action that was taken for a given context, or actions taken in contexts within a given distance of the observed contexts.
We now show that any policy that only selects actions in eligible action sets will ensure that the empirical sum of weights in any hypersphere of any context can be lower bounded, as desired.
To do so we ﬁrst introduce an assumption about the policy’s smoothness in the context space.
Assumption
[Equation] [Equation]
Under this and the former assumptions, we can now show the following.
All proofs, when omitted, are provided in the appendix:
Theorem
[Equation]
Given the likelihood ratio is lower bounded, we can further show that the self-normalized truncated weights are also lower bounded in the one-step settings.
Corollary 1.
For H = 1 , (cid:80)
[Equation]
for M > 1 .
In n -step sequential settings, it is necessary to have the 1-step weights be greater than zero in order to have n -step weights greater than zero.
Proposition 1.
For any x , δ , E [
[Equation]
[Equation]
If the policy overﬁts the weights on an h-step context such that ˆ E [ W ( i ) h | x ( i ) h ∈ B ( x, δ )] = 0 , the n -step weights will also be zero even the roll-in probability under π is non-zero.
The above results highlight that introducing a constraint on the policy class to have local overlap with actions taken in the dataset, is sufﬁcient to ensure that the weights on the contexts are lower bounded.
This will help address the overﬁtting issue highlighted in the prior sections.
Policy learning can now be done by ﬁnding the policy which satisfy the eligible actions constraints given the input dataset:
[Equation]
J ( π ; D ) can be any objective function such as ˆ v IS , ˆ v SNTIS or with the counterfactual risk minimization regularization (Swaminathan & Joachims, 2015a) in prior work:
[Equation]
A natural question is whether these action eligibility local constraints limits the expressivity of the policy class Π .
We show that asymptotically the expressivity is the same, and the maximizer in
Equation 6 will converge to the optimal policy in the policy class if J ( π, D ) is consistent.
Theorem 2.
i)
Fixed δ , for any x , A h ( x ; D , δ ) → { a : µ ( a | x ) > 0 } as n → ∞ with probability 1.
Thus the solution to Equation 6 is the same as ˆ π D ,J = arg max π J ( π, D ) .
ii)
If J ( π, D ) is the objective in Equation 7, the truncation threshold M as a function satisﬁes M → ∞ and M/n →as n → ∞ , and | Π | < ∞ , then v ˆ π D ,J → max π ∈ Π v π in probability.
A LGORITHM : P OLICY O PTIMIZATION WITH EL IGIBLE A CTIONS
Now in Algorithm 1 we introduce our POELA (Policy Optimization with Eligible Actions) algorithm that implements the learning objective in Equation 6 (see Algorithm1—. We use the counterfactual risk minimization objective function Equation 7 as the J ( π ; D ) , where the estimator (cid:99) Var (ˆ v SNTIS ) is constructed using the Normal approximation in [ (Owen, 2013, Equation 9.9) (cid:99)
Var (ˆ v SNTIS ) =
(cid:80) ni
=1 ( r ( i ) − ˆ v SNTIS ) 2 (min { W ( i ) ,M } )
( (cid:80) ni =1 min { W ( i ) ,M } ).
After each gradient step, we enforce the policy to satisfy the eligible action constraints by re-normalizing the output probability on A h ( x ; D , δ ) for x ∈ D .
The eligible action set for each training sample is static and can be stored to reduce computational cost.
In the experiments, we use Euclidean distance in over nearby states at any time index.
E XPERIMENT
We now compare POELA with several prior methods for ofﬂine RL.
Perhaps the most relevant work in avoiding overﬁtting when using importance sampling is norm-POEM (Swaminathan & Joachims, 2015b).
Here we extend it here to be suitable for sequential decision settings, use a neural network policy class and refer to the resulting algorithm as PO-CRM.
A second baseline PO- µ constrains the policy class to only include policies which take actions with a sufﬁcient probability under the behavior policy µ ( a | s ) (see e.g.Futoma et al.
We also compared with recent deep value-based MDP methods in batch RL: BCQ (Fujimoto et al., 2019) and PQL (Liu et al., 2020).
For all algorithms we use a feedforward neural network for the relevant policy and/or value, function approximators.
We use the same procedures for each algorithm to select its hyper-parameters.
An algorithm is trained on the training set multiple times, using different hyper-parameters and several restarts.
Intermittent policies generated during the training process are saved at checkpoints.
Across this set of potential policies, we identify the policy with the highest self-normalized truncated IS (SNTIS) estimates on a held-out validation set.
Finally, we report the test performance of selected policy either through online Monte-Carlo estimation if a simulator is available, or using SNTIS estimates on a held out test set.
Full experimental details are provided are in the Appendix.
7. E XPERIMENT IN LGG T UMOR G ROWTH I NHIBITION S IMULATOR
The Tumor Growth Inhibition (TGI) simulator (Ribba et al., 2012) describes low-grade gliomas (LGG) growth kinetics in response to chemotherapy in a horizon of 30 steps (months), with an non-Markov context and a binary action of drug dosage (Yauney & Shah, 2018).
The reward consist of an immediate penalty proportional to the drug concentration, and a delayed reward of the decrease in mean tumor diameter.
The behavior policy selects from a ﬁxed dosing schedule of 9 months (the median duration from Peyre et al.
with 70% probability and else selects actions at random.
In this experiment, the behavior policy can only take values in { 0 . 15 , 0 . 85 } .
This means constraining the policy class to have a minimal probability under µ ( a | s ) , as in baseline PO- µ , is only a non-trivial constraint for thresholds greater than 0 .
15 : this produces a single potential target policy, which is the deterministic ﬁxed-dosage part of behavior policy.
We include this as 9 -mon (short for 9 month dosing) in Table 1.
The training and validation set both have 1000 episodes, and we repeat the experiment 5 times with 5 different train and validation sets.
Policy values are normalized between(uniform random) and 100 (best policy from online RL).
As shown in Table 1 (Non-MDP rows), our POELA achieves the highest test value as well as smaller variability compared with the baselines.
Does POELA reduce overﬁtting?
We calculate the difference between ˆ v SNTIS on the validation set and the online test value.
Although for a given algorithm, we select the ﬁnal policy using off-policy evaluation on the validation set, most algorithms still result in a policy whose value is a signiﬁcant overestimate of its true performance (see Table1, rows for ˆ v SNTIS − v π ).
In contrast, our approach yields a policy whose value is much more accurately estimated and performs better.
Performance comparison in MDP environment We also repeat the experiment with an MDP modiﬁcation of the simulator, including an immediate Markovian reward and additional features for a Markovian state space.
Note that we expect BCQ and PQL to do very well: both are designed to avoid overﬁtting in ofﬂine MDP learning and in particular PQL uses a pessimism under uncertainty approach to penalize policies that put weight on state–action pairs with little support.
Though our POELA algorithm makes no Markov assumptions, it performs only slightly worse than the conservative MDP methods.
POELA also substantially outperforms other policy-based approaches.
[Table] [Table]
7. E XPERIMENT IN ICU DATA - MIMIC III
We next apply our method in a real-world example of learning policies for sepsis treatment in medical intensive care units (ICU).
We use an extracted cohort (Komorowski et al., 2018) of patients fulﬁlling the sepsis-3 criteria from the MIMIC III data set (Johnson et al., 2016) and obtain a dataset ofpatients, 44 context features, 25 actions and a 20 step maximum horizon.
Full details are in the
Appendix.
We hold out 20% of data for validation and 20% of data for the ﬁnal test.
The treatment logs do not include the probabilities of clinicians’ actions.
Instead, as suggested by prior work(Raghu et al., 2018), we estimate the probabilities of the behavior clinicians’ policy by k -NN with k = 100 .
To ensure overlap, for all policy optimization algorithms we allow π ( a | s ) >
only if ˆ µ ( a | s ) > 0 .
Using self-normalized truncated importance sampling to evaluate the performance on a test set is appealing because it makes little assumptions on the underlying domain.
However, it may be that very few test behavior policy trajectories match a candidate test policy, which can make the resulting value estimate unreliable.
The amount of overlap between the test set and a desired policy to evaluate can be measured quantitatively by the effective sample size (ESS) (Owen, 2013).
To help increase the chance that the test effective sample size is sufﬁcient to yield reliable estimates, during the policy selection process for a given algorithm, only policies with an effective sample size of at leaston the validation set are considered.
1 Similar to prior work (Thomas et al., 2015), in addition to the
SNTIS estimator on the test set, we also report a 95% upper and lower bound from bias-corrected and accelerated (BCa) bootstrap.
The clinician’s column is the test dataset rewards and sample size.
Table 2 shows our POELA is the best on all metrics, achieving the highest evaluation on the test set, the highest upper and lower bounds, and the highest effective sample size.
Is the variance regularization helpful?
To demonstrate the effect of variance penalty on the
POELA algorithm, we demonstrate the test performance POELA without the variance penalty.
The performance is worse than using the variance penalty but still higher than the baseline algorithms.
Is there a trade-off between effective sample size (ESS) and performance estimates?
A tension in conservative ofﬂine optimization is that the most reliable and conservative policy estimates come from effectively imitating the behavior policy (which will maximize effective sample size).
Policies that differ substantially from the behavior policy may yield higher performance, but have less overlap with the existing logged data (and lower ESS).
In Figure 1a we plot each hyper-parameter and each re-start from different algorithms.
The plots show that POELA achieve a better Pareto frontier between performance estimates and effective sample size than other algorithms.
[Figure]
Do learned policies take medically reasonable actions on sicker patients?
As we discuss, prior
IS methods can overﬁt to achieve high estimated IS/SNIS score by avoiding initial contexts with low expected rewards.
We now explore the learned policies for patients with high logged SOFA score (measuring organ failure) initially in the test dataset.
Figure 1b depicts the number of actions that would be taken by different policies, as well as by the clinicians on the high-SOFA contexts.
POELA mainly takes high-vasopressors high-IV-ﬂuid treatments but also some IV-ﬂuid-only treatments, which is similar to the clinician’s policy but more concentrated on high-vasopressors treatments.
In contrast, PO-CRM and value-based methods, take treatments that are different from the logged clinician decisions.
This suggests these policies may be overﬁtting to avoid the contexts with high
SOFA in the distribution.
However, some patients arrive with high SOFA scores, so a policy must have a suitable policy to support such individuals, which our method helps to ensure.
While these results are promising, before deploying any ofﬂine reinforcement learning algorithm in a high stakes clinical setting, further investigation and collaborations with clinicians would be essential.
C ONCLUSION & F UTURE W ORK
To conclude, we identify a novel overﬁtting challenge that arises when using importance sampling as part of an ofﬂine policy learning objective.
In particular, the objective can result in a policy that under-weigh certain (lower performing) initial contexts, to achieve higher average value.
To address this, we constrain the policy class to only consider logged actions taken by nearby samples.
This can be viewed as a similar pessimism constraint that has been used in MDP ofﬂine policy learning, but now developed for a non-Markov, direct policy search setting.
Our approach yielded strong performance relative to state-of-the-art prior approaches in a tumor growth simulator and a real world dataset on ICU sepsis treatment.
Our method may be particularly useful for many applied settings such as healthcare, education and customer interactions, which have a short or medium length decision horizon, but are unlikely to be Markov in the observed per-step variables.
A reader might wonder if a similar beneﬁt might be possible if the algorithm ﬁrst learned an empirical behavior policy from the logged dataset, and then constraining the policy class to be close to the empirical behavior policy.
Our current results suggest that such an approach is likely to require additional innovation: indeed the PO- µ baseline for the sepsis dataset uses an empirical behavior policy, and does not perform as well as our approach.
This is interesting because prior work has shown that using such an empirical behavior policy can yield better off-policy estimates than using true behavior policy (Xie et al., 2018; Hanna et al., 2021) settings, and lead to beneﬁts in off policy learning for contextual bandits
Xie et al.
However, this prior work did not consider constraining the policy class nor performing ofﬂine RL, and there are some signiﬁcant subtleties in learning an empirical behavior policy for large continuous spaces.
For example, Hanna et al.
(2021) demonstrated that learning a behavior policy that maximized the likelihood of the logged data in a
IS-based estimator did not yield the most accurate estimate.
An interesting direction for future work is whether different ways of learning an empirical behavior policy might yield similar or additional beneﬁts to our locally constrained approach proposed.
A C OUNTER E XAMPLES IN RL S ETTINGS
In the main text, we gave an example about the overﬁtting issue in contextual bandits with large state and action space in small datasets.
Here we show that it is even easier for this to occur in sequential reinforcement learning settings, even when only 2 actions are available in the next two examples with or without state aliasing.
Example 2.
Consider a sequential treatment problem as shown in Figure 2a.
There are two actions available in each state.
From the ﬁrst state, action a 1 has a 50% chance of leading to an immediate terminal positive reward r = 1 and a 50% chance of leading to an immediate terminal negative reward r = − 1 .
From the ﬁrst state, action a 2 also has 50% chance of leading to an immediate terminal positive reward r = 1 .
For the other 50% of states, action a 2 results in transitions to additional states, which are followed by additional actions, for another H − 1 steps; however, all transitions eventually end in a large negative outcome (e.g. r = − 5 ).
For example, one could consider a risky surgical procedure that results in many subsequent additional operations and but is ultimately typically unsuccessful.
Assume the behavior policy is uniform over each action, yielding µ ( a = 0 | x 1 )
=
µ ( a = 1 | x 1 ) = 0.5 and a probability of each action sequence following a 2 of 1 | A | H − 1 .
With even minimal data the value of π ( x 0 ) =
a 1 will be accurately estimated as 0.
However, when
H is large relative to a function of the dataset size, there always exists a action sequence after an initial selection of a 2 that is not observed in the dataset.
This means that a policy π 2 that starts with π ( x 0 ) =
a 2 and then selects an unobserved action sequence will essentially put 0 weight on the resulting contexts that incur r = − 5 outcomes, even though such outcomes will occur 50% of the time after taking action a 2 .
In this case, the value of π 2 will be overestimated signiﬁcantly by IS or self-normalized IS.
Thus the ofﬂine policy optimization will prefer taking action 2 at the ﬁrst step as a result of overﬁtting even though the true value of ﬁrst taking a 2 is − 1 .
5 and the optimal policy value is 0 , obtained by taking action a 1 .
Now we add a slight change in the transitions shown in Figure 2a.
We can see that model/value-based approach will also fail.
Example 3.
In this example, we add another action in the ﬁrst step.
The action 3 and action 1 will lead to the same next state.
However in the next state, no matter which action taken, the reward will depends on the action taken in the last step: If a 1 = 1 , then we have the same reward for a = 1 in the example in Figure 2a.
If a 1 = 3 then we have a reward − 5 .
Thus model and value based method will mix the reward for a 1 = 1 and a 1 = 3 so fail in this example.
Other method is not affected by the additional structure as it only add an action with minimum reward.
[Figure] [Figure]
[Figure] [Figure]
B P ROOFS FOR S ECTION
Proof of Theorem
Proof.
[Equation]
Proof of Corollary
Proof.
[Equation]
Proof of
Proposition
Proof.
This is due to π ( a | x ( i ) h ) and µ ( a | x ( i ) h ) are independent from history given x ( i ) h .
So W ( i ) 1: h and
W ( i ) h are conditionally independent given x ( i ) h .
Proof of Theorem
Proof.
We ﬁrst prove the ﬁrst part.
Let P h ( x ; µ ) to be the distribution of context at h -th step with roll-in policy µ .
For any ﬁxed a , we can deﬁne the distribution P h ( x | a ; µ ) =
µ ( a | x )P h ( x ; µ ) / (cid:80) a µ ( a | x )P h ( x ; µ ) .
For a such that µ ( a | x ) > 0 , P h ( x | a ; µ ) is also greater than zero.
All x ( i ) h with a ( i ) h
= a are i.i.d.
samples draw from the distribution P h ( x ; µ ) .
By the property of nearest neighbor (Cover & Hart, 1967), with probability min x
[Equation]
i
[Equation]
That means with probability
a ∈
A h ( x ; D , δ ) for all a such that µ ( a | x ) > 0 .
Given the overlap assumption, we have a ∈ A h ( x ; D , δ ) for for all a such that π ( a | x ) > 0 with probability 1.
Thus the solution to Equation 6 is the same as arg max π J ( π, D )
=
ˆ π J, D .
Now we prove the second part.
By the condition that M → ∞ and Mn → 0 as n → ∞ , we have that the truncated IS estimator is mean square consistent (Ionides, 2008):
[Equation]
as n → ∞ .
Similarly, we have that the mean of weights converge to 1 in quadratic mean:
[Equation]
By continuous mapping theorem, we have that the self-normalized truncated IS converge to v π in probability ˆ v SNTIS p −→ n .
The empirical variance penalty, also converge to 0 almostly surely, since
M/n converge to 0 :
[Equation]
Thus the objective function J ( π ; D ) converge to v π in probability:
[Equation]
Since we assume | Π | < ∞ , we have
[Equation]
So with probability | Π | δ n , for any (cid:15) :
[Equation] [Equation]
where π (cid:63) is arg max π ∈ Π v π .
As | Π | δ n → 0 , we proved the true value of empirical maximizer v ˆ π J, D converge to the maximum of value max π ∈ Π v π in probability.
C E XPERIMENT D ETAILS
For all experiment, we follows the 3-phases pipelines to decide the test score we ﬁnally report in the paper.
The pipeline is described in the main text and summarized in Figure ??
here as well.
C. E
[Equation]
ETAILS IN TGI S IMULATOR
The TGI simulator describes low-grade gliomas (LGG) growth kinetics in response to chemotherapy in a horizon of 30 months using an ordinary differential equation model.
The parameter in ODEs are estimated using data from adult diffuse LGG during and after chemotherapy was used, in a horizon of 30 months.
The goal in this environment is to achieve a reduction in mean tumor diameters (MTD) while reducing the drug dosage (Yauney & Shah, 2018).
We includes the MTD, the drug concentration, and the number of month (time-step) in the context space.
Notice that this context space is non-Markov as it does not include all parameters in the ODEs.
Actions are binary representing taking the full dose or no dose which is same as prior work (Yauney & Shah, 2018).
The reward at each step consist of an immediate penalty proportional to the drug concentration, and a delayed reward at the end measures the decrease of MTD compared with the beginning.
Each episodes, the parameters including the initial MTD are sampled from a log-Normal distribution as (Ribba et al., 2012) representing the difference in individuals.
The behavior policy is a ﬁxed dosing schedule ofmonths (the median duration from Peyre et al.
plus 30% of a uniformly random choice of actions.
We run all algorithms on a training set with 1000 episodes with different hyperparameters (listed below), and 5 restarts, saving checkpoints along the training.
Then we select the best policy for each algorithm by ˆ v SNTIS on the validation set with 1000 episodes as well.
Hyperparameters In the ﬁrst part of Table 3 we show the searched hyperparameters of each algorithm, except that the parameter b in PQL is set adaptively as the 2 -percentile of the score on the training set as in the original paper Liu et al.
As we know the behavior policy, we use the true behavior policy in BCQ and PQL algorithm.
So BCQ threshold takes only two values as the behavior policy is (cid:15) -deterministic
so there are only two distinct values.
In the second part of Table 3 we specify some ﬁxed hyperparameters/settings for all algorithm.
All policy/Q functions are approximated by fully-connected neural networks with two hidden layers with 32 units.
Hyperparameters used in algorithms values
[Table]
CRM Var coefﬁcient
POELA, PO
[Table]
BCQ
[Table] [Table]
M in
[Table]
All
Max
[Table]
steps
[Table] [Table]
[Table]
The difference in the max update steps and checkpoints frequency is caused by the fact that BCQ and
PQL is updated by stochastic gradient descent and all policy optimization based on SNTIS is using gradient descent.
C. E XPERIMENT D ETAILS IN MIMIC III D ATASET
The
MIMIC
III sepsis dataset is available upon application and training: https://mimic.mit.edu/iii/gettingstarted/.
The code to extract the cohort is available on: https://gitlab.doc.ic.ac.uk/AIClinician/AIClinician.
This cohort consists of data forpatients.
The contexts for each patient consist of 44 features, summarized in 4-hour intervals, for at most 20 steps.
The actions we consider are the prescription of IV ﬂuids and vasopressors.
Each of the two treatments is binned into 5 discrete actions according to the dosage amounts, resulting inpossible actions.
The rewards are deﬁned from the 90-day mortality in the logs, 100 if the patient survives and 0 otherwise.
We now provide details of the experiment on MIMIC III sepsis dataset here.
We run all algorithms on a training set with 8982 trajectories with different hyperparameters (listed below), and 3 restarts, saving checkpoints along the training.
Then we select the best policy for each algorithm by ˆ v SNTIS on the validation set with 2994 trajectories.
Finally we get the ˆ v SNTIS evaluation on the test set with 2995 trajectories.
In the ﬁrst part of Table 4 we list the hyperparameters that we searched on the validation set for each algorithm, except that the parameter b in PQL is set adaptively as the 2 -percentile of the score on the training set as in the original paper Liu et al.
In the second part of Table 3 we specify some ﬁxed hyperparameters/settings for all algorithm.
All policy/Q functions are approximated by fully-connected neural networks with two hidden layers with 256 units.
As we explained, the difference in the max update steps and checkpoints frequency is caused by the fact that BCQ and PQL is updated by stochastic gradient descent and all policy optimization based on SNTIS is using gradient descent.
C.2. A CTION VISUALIZATION FOR MID / LOW -SOFA PATIENTS
In the main text we show the action visualization for high-SOFA ( > 15 ) as a diagnose of how algorithm behave for the high-risk patients.
For the completeness of results here we show the visualization of mid-SOFA ( 5 − 15 ) and low-SOFA ( < 5 ) patient contexts.
[Table]
algorithms
[Table]
C.2. A DDITIONAL ABLATION STUDY : WITHOUT EFFECTIVE SAMPLE SIZE CONSTRAINTS FOR HYPER - PARAMETER SELECTION ON VALIDATION SET
In the main text
[Table]
we set an effective sample size threshold of 200 for a policy/hyper-parameter to be selected on validation set.
This is to make sure we have large enough effective sample size on the test set to provide reliable off-policy test estimates.
Here we show the result if we do not threshold the effective sample size on validation set.
Generally, all algorithms will prefer a high off-policy estimates without enough effective sample size.
On the test set, all algorithms yields a small effective sample size
[Table]
thus unreliable off-policy estimates and large bootstrap conﬁdence interval.
The proposed methods is better than baselines but also has much smaller 95% bootstrap lower bound than with the effective sample size constraint
[Table]
[Figure]
C. I MPORTANCE WEIGHTS IN LOW - REWARD TRAJECTORIES
To examine if the proposed overﬁtting phenomenon exists in real experimental datasets, we compute the importance weights of the learned policy on the low-reward trajectories in the training data for our MIMIC III dataset and our tumor simulator.
Our hypothesis is that overﬁtting of the importance weights in policy gradient methods may result in the algorithm avoiding initial states with low rewards, which motivated our proposed algorithm.
In MIMIC III dataset the reward for a trajectory is either 0 or 100 . We deﬁne the low-reward trajectories as those with 0 reward.
Low-reward trajectories are over 60% of all trajectories in the dataset. In the Tumor simulation experiment we deﬁne a low-reward trajectory when reward is less than − 2 . Over 95% of trajectories in the Tumor simulation dataset are low-reward trajectories. The table below shows, for each algorithm and setting, the sum of the SNTIS weights of the learned policy on the training set, for low-reward trajectory states.
Our primary interest is to illustrate that alternate policy gradient methods that are also suitable for non-Markov domains, can exhibit the importance sampling overﬁtting of avoiding low reward trajectories. We indeed see that our POELA method has a much larger weight on low-reward trajectories than these alternate ofﬂine policy search methods: [Table]
The Q-learning baselines we consider (BCQ and PQL) do not directly use the importance weights, but they do try to avoid actions and/or states and actions with little support. Our POELA method can be viewed as being similarly inspired, but for non-Markovian settings where policy gradient is beneﬁcial.
We also compute the SNTIS weights of the BCQ/PQL policy on the training set in the Markov domain that satisﬁes the Markov assumptions of BCQ/PQL. Here we can see that our
POELA, BCQ and PQL all still give signiﬁcantly more weight to low reward trajectories than the alternate policy gradient methods:
[Table]
These results help illustrate that the over avoidance of low-reward trajectories can be observed by past policy gradient methods in our datasets. Of course, one challenge is that in real settings, an excellent policy may have low importance weights in avoidable low-reward states and trajectories, but should have higher importance weights in non-avoidable low reward starting states and trajectories. To get a fuller picture of performance, it is helpful to look both at the weights on trajectories with low rewards and the test evaluation results. Compared with strong policy gradient baselines, our proposed regularization method have larger importance weights on low-reward trajectories, and the gap between training/validation evaluation and online test performance is also smaller, suggesting that we are less likely to learn policies that erroneously believe they can avoid unavoidable low reward settings.
C. T HE EFFECT OF ELIGIBLE ACTION CONSTRAINTS δ
In this section we explore how the choice of δ , which constrains the policy class through impacting the eligible actions, impacts empirical performance.
Larger δ corresponds to a less constrained policy class.
Other hyper-parameters are selected by the same procedure as described in previous sections.
As δ increases, the policy search operates with less constraints. The results show that in this case, our policy gradient method produces a policy with a higher value in the training set, but that policy may not perform as well in the test evaluation, and may have a smaller effective sample size than when a smaller δ is used.
The best hyper-parameter value δ lies in the middle of the explored range. δ can be selected based on performance and effective sample size.
[Table]