Selective Uncertainty Propagation in Offline RL

Abstract

We consider the finite-horizon offline reinforcement learning (RL) setting, and are motivated by the challenge of learning the policy at any step $h$ in dynamic programming (DP) algorithms. To learn this, it is sufficient to evaluate the treatment effect of deviating from the behavioral policy at step $h$ after having optimized the policy for all future steps. Since the policy at any step can affect next-state distributions, the related distributional shift challenges can make this problem far more statistically hard than estimating such treatment effects in the stochastic contextual bandit setting. However, the hardness of many real-world RL instances lies between the two regimes. We develop a flexible and general method called selective uncertainty propagation for confidence interval construction that adapts to the hardness of the associated distribution shift challenges. We show benefits of our approach on toy environments and demonstrate the benefits of these techniques for offline policy learning.

1 Introduction

We study the finite-horizon offline reinforcement learning (RL) problem, focusing on algorithms that adapt to instance hardness. At a high-level, we study algorithms that provide better guarantees for contextual bandit (CB) like instances while being able to plan in more dynamic RL-like instances.

Our work is motivated by real-world RL problems, such as user interaction with an e-commerce search engine (recommendation system). Here, the state can be a user query, and the action is the product recommendation from the engine. When the user wants to buy a particular product, the user often only enters a single product query unrelated to the previous one; thus resembling a sequence of CB problems. On the other hand, when the user explores products, the exploration queries are related through the user’s intent, and the recommendation system may want to steer the user toward the ideal product. Hence, this resembles the RL setting. This indicates the need to develop unified solutions that integrate CB and RL techniques – adapting to instance hardness. We now introduce the CB and RL frameworks in more detail.

Stochastic contextual bandits (CBs) [Langford and Zhang, 2008, Li et al., 2010] and finite-horizon reinforcement learning (RL) [Sutton, 1988, Williams, 1992, Sutton and Barto, 1998] are two fundamental frameworks for decision-making under uncertainty. In stochastic CBs, the environment samples the context and corresponding rewards (for each action) from a fixed but unknown distribution; the agent then observes the context and learns to select the most rewarding action conditioned on the context.

Finite-horizon RL is a generalization of CBs where contexts become states and a sequence of decisions are to be made over $H$ steps. Similar to the CB problem, at each step, the agent observes the current state, selects an action conditioned on the current state, and receives a reward sampled by the environment from a corresponding conditional distribution. However, unlike the CB problem, while the initial state is sampled from a fixed but unknown distribution, the next state at any step depends on the current state and the agent’s action. Hence, the agent can plan to attain high cumulative reward by learning to reach high-value future states.

Unfortunately, the fact that actions can influence future states implies that the agent needs to learn under state-distribution shifts making the RL setting much more statistically harder than CBs in the worst case. For example, [Foster et al., 2021] show that the worst-case sample complexity to learn a non-trivial offline RL policy is either polynomial in the state space size or exponential in other parameters. On the other hand, if actions do not influence next-state distributions at any step, the RL instance would be equivalent to solving $H$ stochastic CB instances. On such instances, offline bandit algorithms [Foster and Syrgkanis, 2019] would enjoy a polynomial sample complexity for policy learning with no dependence on state space size. Hence, for such instances, state-of-the-art offline RL algorithms such as pessimistic value function optimization [Jin et al., 2021] may be unnecessarily conservative.

We formalize this dichotomy and show that the statistical hardness of offline RL instances can be captured by the size of actions’ impact on the next state’s distribution. To show this, we consider the high-level structure of dynamic programming (DP) algorithms for offline RL [e.g. Jin et al., 2021]. DP algorithms construct a policy iteratively starting from the policy for the final step and ending by constructing the policy for the first step. At any step $h$, DP algorithms can be viewed to select the policy at step $h$ that maximizes the treatment effect of deviating from the behavioral policy at step $h$ after having optimized the policy for all future steps. The goal of this paper is to estimate and construct good confidence intervals for this treatment effect at step $h$.

Our primary focus is on confidence interval (CI) construction, which is motivated by the fact that many successful offline RL algorithms learn a policy that maximizes the lower bound of constructed CIs [Jin et al., 2021]. To account for estimation errors from future steps, standard methods for CI construction at any step propagate uncertainty from future steps to the current step $h$. This paper seeks to construct better CIs that adapt to instance hardness by selectively propagating uncertainty. In cases where all actions have zero estimated impact on next-state distributions, our procedure does not propagate any uncertainty from later steps and still constructs valid CIs for the treatment effect of deviating from the behavioral policy at step $h$ after having optimized the policy for all future steps. It treats the instance like a CB problem – hence enjoying a polynomial sample complexity with no dependence on state space size for treatment effect estimation. For more dynamic instances, our procedure must unavoidably propagate more uncertainty from future steps in order to continue constructing valid CIs. In this way, we adapt to the hardness of the instance for CI construction at any step. We also show the benefits of this approach for offline policy learning by proposing an algorithm that optimizes our constructed CIs. Simple simulations further support our claim.

Related Work: Both bandits and RL have been studied extensively [Lattimore and Szepesvari, 2019, Sutton and Barto, 1998, Foster and Rakhlin, 2023]. In bandits, the focus has been on achieving higher statistical efficiency by using the reward distribution of actions [Garivier and Cappe, 2011], prior distribution of model parameters [Thompson, 1933, Agrawal and Goyal, 2012, Chapelle and Li, 2012, Russo et al., 2018], parametric structure [Dani et al., 2008, Abbasi-Yadkori et al., 2011, Agrawal and Goyal, 2013], or agnostic methods [Agarwal et al., 2014]. In RL, the focus has been on different means of learning to plan for longer horizons, such as the value function [Sutton, 1988], policy [Williams, 1992], or their combination [Sutton et al., 2000]. Just as in our work, causal inference insights have helped improve the statistical efficiency of both CB and RL algorithms [Krishnamurthy et al., 2018, Carranza et al., 2023, Syrgkanis and Zhan, 2023]. However, bridging the gap between bandits and RL is an exciting and relatively under-explored research direction. One way to define this gap is to argue that in bandit-like environments, the state never changes once initially sampled. These bandit-like environments can be viewed as a special case of the situation where actions do not impact next-state distributions. With bridging this gap as one motivation, [Zanette and Brunskill, 2019, Yin and Wang, 2021] have used variance-dependent Bernstein bounds to limit uncertainty propagation when there is a lack of next-step value function heterogeneity. Another approach is to define this gap in a binary fashion. Either there is no impact of actions on next state distributions, or we are in a dynamic MDP environment. In an online setting, [Zhang et al., 2022] develop hypothesis tests to differentiate between the two situations and then select the most appropriate exploration algorithm. While their higher-level framing is similar to ours and their approach is novel, their approach cannot outperform existing RL algorithms in MDP environments. By interpolating between the two regimes, we hope to outperform bandit and existing RL algorithms that either forgo planning or are too conservative in accounting for actions’ impact on next-state distributions.

2 Preliminaries

Setting: We consider an episodic Markov Decision Process (MDP) setting with state space $\mathcal{X}$, action space $\mathcal{A}$, horizon $H$, and transition kernel $P=({P}^{(h)})_{h=1}^{H}$. At every episode, the environment samples a starting state ${x}^{(1)}$ and a set of realized rewards $r=({r}^{(h)})_{h=1}^{H}$ from a fixed but unknown distribution $D$. Here ${r}^{(h)}$ is a map from $\mathcal{X}\times\mathcal{A}$ to $[0,1]$. For any states $x,x^{\prime}\in\mathcal{X}$ and action $a\in\mathcal{A}$, ${P}^{(h)}(x^{\prime}|x,a)$ denotes the probability density of transitioning to state $x^{\prime}$ conditional on taking action $a$ at state $x$ during step $h$. A trajectory $\tau$ is a sequence of states, actions, and rewards. That is, any trajectory $\tau$ is given by $\tau=({x}^{(h)},{a}^{(h)},{r}^{(h)}({x}^{(h)},{a}^{(h)}))_{h=1}^{H}$.

A policy $\pi$ is a sequence of $H$ action sampling kernels $\{{\pi}^{(h)}\}_{h=1}^{H}$, where ${\pi}^{(h)}(a|x)$ denotes the probability of sampling action $a$ at state $x$ during step $h$ under the policy $\pi$. We let $D(\pi)$ denote the induced distribution over trajectories under the policy $\pi$. For any policy $\pi$, we define the (state-) value function ${V}^{(h)}_{\pi}:\mathcal{X}\rightarrow[0,H-h+1]$ at each step $h\in[H]$ such that,

$${V}^{(h)}_{\pi}(x)=\mathbb{E}_{D(\pi)}\bigg{[}\sum_{i=h}^{H}{r}^{(i)}({x}^{(i)},{a}^{(i)})\bigg{|}{x}^{(h)}=x\bigg{]}.$$ (1)

The value of policy $\pi$ is given by $\mathbb{E}_{D}[{V}^{(1)}_{\pi}({x}^{(1)})]$. We can take expectation over $D$ instead of $D(\pi)$ here since the only random variable in ${V}^{(1)}_{\pi}({x}^{(1)})$ is the initial state ${x}^{(1)}$ which does not depend on the choice of the policy $\pi$.

For any step $h\in[H]$, we let ${R}^{(h)}$ be a function from $\mathcal{X}\times\mathcal{A}$ to $[0,1]$ denoting the expected reward function for step $h$. That is, ${R}^{(h)}(x,a)=\mathbb{E}_{D}[{r}^{(h)}(x,a)]$. With some abuse of notation, for any $x,x^{\prime}\in\mathcal{X}$, we let ${R}^{(h)}(x,\pi)=\sum_{a}{\pi}^{(h)}(a|x){R}^{(h)}(x,a)$ and ${P}^{(h)}(x^{\prime}|x,\pi)=\sum_{a}{\pi}^{(h)}(a|x){P}^{(h)}(x^{\prime}|x,a)$. That is, ${R}^{(h)}(x,\pi)$ is the expected reward at state $x$ and step $h$ under the policy $\pi$. Similarly, ${P}^{(h)}(x^{\prime}|x,\pi)$ is the expected transition probability from $x$ to $x^{\prime}$ at step $h$ under the policy $\pi$. For any step $h\in[H]$, we also let ${V_{\text{max}}}^{(h)}$ denote a bound on the maximum value ${V}^{(h)}_{\pi}(x)$ can take for any state $x$ and policy $\pi$.

It is also equivalent to define the value functions (${V}^{(h)}_{\pi}$) using the iterative definition in (2), where ${V_{\pi}}^{(H+1)}\equiv 0$.

$$\forall h\in[H],x\in\mathcal{X},$$ (2)

$${V_{\pi}}^{(h)}(x)={R}^{(h)}(x,\pi)+\int_{x^{\prime}}{V_{\pi}}^{(h+1)}(x^{\prime}){P}^{(h)}(x^{\prime}|x,\pi).$$

Data Collection Process: In this paper, we focus on the offline setting [Levine et al., 2020] with training data collected under a behavioral policy $\pi_{b}$. Apart from the policy $\pi_{b}$, the learner only has access to a dataset $S$ consisting of $T$ trajectories sampled from the distribution $D(\pi_{b})$, where $D(\pi_{b})$ is the data sampling distribution induced by $\pi_{b}$. That is, $S=\{\tau_{t}\}_{t=1}^{T}$, where $\tau_{t}=({x}^{(h)}_{t},{a}^{(h)}_{t},{r}^{(h)}_{t}({x}^{(h)},{a}^{(h)}))_{h=1}^{H}\sim D(\pi_{b})$. Since the transitions in these trajectories are induced by the behavioral policy, for notational convenience, we let $P_{b}=({P_{b}}^{(h)})_{h=1}^{H}$ denote the transition kernel under the policy $\pi_{b}$. That is, for any $x,x^{\prime}\in\mathcal{X}$, ${P_{b}}^{(h)}(x^{\prime}|x)={P}^{(h)}(x^{\prime}|x,\pi_{b})$.

2.1 Estimand of Interest

We now turn our attention to defining our target estimand, which refers to the specific quantity we aim to estimate. Consider a fixed policy $\pi$ and suppose we would like to estimate its value. Since estimating the value of the behavioral policy $\pi_{b}$ is easy (empirical average of total observed reward in each trajectory), we argue that that it is sufficient to estimate $\mathbb{E}_{D}[{V}^{(1)}_{\pi}({x}^{(1)})-{V}^{(1)}_{\pi_{b}}({x}^{(1)})]$ – the difference in values between evaluation and behavioral policy. This difference can be further decomposed. For each step $h$, let $\tilde{\pi}_{h}=({\pi_{b}}^{(1)},\dots,{\pi_{b}}^{(h-1)},{\pi}^{(h)},\dots,{\pi}^{(H)})$ be the policy that follows the behavioral policy upto step $h-1$ and then follows the evaluation policy. In (3), we decompose the difference in policy value between the evaluation and behavioral policy into the sum of differences in policy value between $\tilde{\pi}_{h}$ and $\tilde{\pi}_{h+1}$ for each step $h$.

$$\mathbb{E}_{D}[{V}^{(1)}_{\pi}({x}^{(1)})-{V}^{(1)}_{\pi_{b}}({x}^{(1)})]$$ (3)

$$\stackrel{{\scriptstyle(i)}}{{=}}\sum_{h=1}^{H}\mathbb{E}_{D}[{V}^{(1)}_{\tilde{\pi}_{h}}({x}^{(1)})-{V}^{(1)}_{\tilde{\pi}_{h+1}}({x}^{(1)})]$$

$$\stackrel{{\scriptstyle(ii)}}{{=}}\sum_{h=1}^{H}\mathbb{E}_{D(\pi_{b})}[{V}^{(h)}_{\tilde{\pi}_{h}}({x}^{(h)})-{V}^{(h)}_{\tilde{\pi}_{h+1}}({x}^{(h)})].$$

Here (i) follows from telescoping and (ii) follows from the fact that the policies $\tilde{\pi}_{h}$ and $\tilde{\pi}_{h+1}$ agree with the behavioral policy for the first $h-1$ steps. We let ${\alpha}^{(h)}_{\pi}$, the term corresponding to step $h$ in the above decomposition, be our estimand of interest. That is, our estimand (${\alpha}^{(h)}_{\pi}$) is the difference in value of policies $\tilde{\pi}_{h}$ and $\tilde{\pi}_{h+1}$ – these policies only differ in the current step $h$, which may cause difference in immediate rewards and may also cause a difference in in next-state distributions (affecting future rewards even if the policies at future steps are the same).

$${\alpha}^{(h)}_{\pi}=\mathbb{E}_{D(\pi_{b})}[{V}^{(h)}_{\tilde{\pi}_{h}}({x}^{(h)})-{V}^{(h)}_{\tilde{\pi}_{h+1}}({x}^{(h)})]$$ (4)

We now seek to justify ${\alpha}^{(h)}_{\pi}$ as an important estimand, and start by arguing that it is a reasonable estimand to care about. Note that, given the decomposition in (3), estimating and constructing CIs for $\{{\alpha}^{(h)}_{\pi}\}_{h=1}^{H}$ allows us to estimate and construct CIs for $\mathbb{E}_{D}[{V}^{(1)}_{\pi}({x}^{(1)})-{V}^{(1)}_{\pi_{b}}({x}^{(1)})]$ (the difference in policy value between evaluation and behavioral policies) – and thus allows us to estimate and construct CIs for $\mathbb{E}_{D}[{V}^{(1)}_{\pi}({x}^{(1)})]$ (evaluation policy value).

Beyond being an effective surrogate for policy evaluation, ${\alpha}^{(h)}_{\pi}$ is an important quantity to consider in dynamic programming (DP) algorithms. DP algorithms construct the policy for the final step (${\pi}^{(H)}$) and iteratively construct policies for earlier steps. At step $h$, the policy at steps $h+1$ to $H$ are already fixed/computed. Hence at this step, one can interpret DP algorithms as attempting to select ${\pi}^{(h)}$ in order to maximize ${\alpha}^{(h)}_{\pi}$ – that is, maximize the treatment effect of deviating from the behavioral policy at step h after having optimized the policy for all future steps. Hence, for any step $h$, ${\alpha}^{(h)}_{\pi}$ is a helpful estimand to consider for decision-making at step $h$.

Importantly for us, when actions at step $h$ do not affect next state distributions, the problem of choosing a policy at step $h$ can be viewed as a CB problem. Helpfully in this case, unlike policy value, ${\alpha}^{(h)}_{\pi}$ only depends on immediate rewards and can be estimated via offline stochastic CB techniques. However, when actions at step $h$ do influence next state distributions, RL techniques are necessary for estimating ${\alpha}^{(h)}_{\pi}$. Hence, beyond being a critical quantity for decision-making at step $h$, it is also a quantity that is amenable to interpolating between CB and RL techniques. Thus, our paper focuses on estimating and constructing tight confidence intervals (CIs) for this estimand (${\alpha}^{(h)}_{\pi}$).

3 Shift Model

Offline RL is more challenging than offline policy learning in the stochastic CB setting [Foster et al., 2021]. The primary reason for the difference between the two settings is due to state distribution shift induced due to deviating from the behavioral policy. Distribution shift makes any statistical learning theory problem challenging [Vergara et al., 2012, Bobu et al., 2018, Farshchian et al., 2018]. Hence methods that adapt to instance hardness must rely on some implicit or explicit approach to measure this state-distribution shift. To this end, we model the “heterogeneous treatment effect" [Künzel et al., 2019, Nie and Wager, 2021] of actions on the next-state distribution and refer to this effect as the “shift model". More precisely, we define the shift model $\Delta=({\Delta}^{(h)})_{h=1}^{H}$ in (5).

$$\forall(x,a),{\Delta}^{(h)}(\cdot|x,a)={P}^{(h)}(\cdot|x,a)-{P_{b}}^{(h)}(\cdot|x).$$ (5)

Here $\Delta^{(h)}(x^{\prime}|x,a)$ captures the shift in the probability of transitioning from $x$ to $x^{\prime}$ due to selecting action $a$ at state $x$ instead of following the behavioral policy. With some abuse of notation, for any $x,x^{\prime}\in\mathcal{X}$, we let ${\Delta}^{(h)}(x^{\prime}|x,\pi)=\sum_{a}{\pi}^{(h)}(a|x){\Delta}^{(h)}(x^{\prime}|x,a)$. That is, ${\Delta}^{(h)}(x^{\prime}|x,\pi)$ is the expected shift (w.r.t to $P_{b}$) in probability of transitioning from $x$ to $x^{\prime}$ at step $h$ under the policy $\pi$. It is worth noting that shifts are bounded. For all $(x,a)$, since the $\Delta^{(h)}(\cdot|x,a)$ is a difference of two state-distributions, we have $\|\Delta^{(h)}(\cdot|x,a)\|_{1}\leq 2$ from triangle inequality.

We argue that shift helps capture instance hardness for estimating ${\alpha}^{(h)}_{\pi}$. To see this, we provide a shift-dependent expression for ${\alpha}^{(h)}_{\pi}$.

$${\alpha}^{(h)}_{\pi} \stackrel{{\scriptstyle(i)}}{{=}}\mathbb{E}_{D(\pi_{b})}[{V}^{(h)}_{\tilde{\pi}_{h}}({x}^{(h)})-{V}^{(h)}_{\tilde{\pi}_{h+1}}({x}^{(h)})]$$ (6)

$$\stackrel{{\scriptstyle(ii)}}{{=}}\mathbb{E}_{D(\pi_{b})}[{R}^{(h)}({x}^{(h)},\pi)-{R}^{(h)}({x}^{(h)},\pi_{b})]$$

$$+\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}{V_{\pi}}^{(h+1)}(x^{\prime}){\Delta}^{(h)}(x^{\prime}|{x}^{(h)},\pi)\bigg{]}.$$

Here (i) follows from (4) (definition of ${\alpha}^{(h)}_{\pi}$); and (ii) follows from (2) and (5). Note that, in the final expression of (6), the first term can be estimated using stochastic CB techniques and the dependence on next-step value function is scaled by the size of this shift. This hints at the possibility of developing methods that interpolate between CB and RL techniques. More formally, in Section 4, we show shift estimates enable us to adapt to the hardness of our setting – when estimating and constructing CIs for ${\alpha}^{(h)}_{\pi}$.

4 Theory: Selective Propagation

In Section 2, we motivated and defined our estimand ${\alpha}^{(h)}_{\pi}$ (see (4)) – which is the treatment effect for deviating from the behavioral policy at step $h$ after having already deviated from the behavioral policy for all future steps. We now present an approach to estimate and construct tight valid CIs for ${\alpha}^{(h)}_{\pi}$ – with interval size adapting to instance hardness. Here harder instances have a larger next-state distribution shifts when deviating from the behavioral policy. When shifts are smaller, we can rely more on statistically efficient CB methods. However when shifts are larger (instance is more dynamic), we unavoidably must rely more on RL methods that account for worst-case distribution shifts.

Our approach to estimate and construct tight valid CIs for ${\alpha}^{(h)}_{\pi}$ requires several inputs. These inputs, described in Section 4.1, allow us to abstract away existing approaches to tackle well-studied estimation problems in CB and RL settings. In Section 4.2, we describe how to combine these existing tools to achieve guarantees that adapt to instance hardness.

4.1 Inputs

Our method interpolates between existing tools for CB and RL settings, by leveraging shift estimates. To simplify our analysis and generalize our results, we assume access to these estimates as inputs to our interpolation method. In particular, we take as input: (1) offline CB treatment effect estimate and corresponding CI, (2) optimistic and pessimistic offline RL value function estimates, and (3) shift estimates with average error bounds. As the quality of our inputs improve (potentially as better estimators get developed), the quality of our outputs will correspondingly improve.

We now formally describe these inputs – requiring all the associated high-probability bounds to hold simultaneously with probability at least $1-\delta_{\text{in}}$. We start by describing the first input, which is based on CB methods.

Input 1 (CB estimates): This input provides an estimate and CI for ${\theta}^{(h)}$ (formally defined in (7)) – which is the average treatment effect on the immediate reward for deviating from the behavioral policy at step $h$.

$${\theta}^{(h)}_{\pi}=\mathbb{E}_{D(\pi_{b})}\Big{[}{R}^{(h)}({x}^{(h)},\tilde{\pi}_{h})-{R}^{(h)}({x}^{(h)},\tilde{\pi}_{h+1})\Big{]}$$ (7)

Since ${\theta}^{(h)}_{\pi}$ only depends on the immediate reward, well-established offline CB techniques [e.g., Dudik et al., 2014] can be used to estimate and construct CIs for the difference (in terms of immediate rewards) between these policies. We let ${\hat{\theta}}^{(h)}_{\pi}$ be our input estimate and let ${\kappa}^{(h)}_{\pi,\theta}$ be the input CI radius. That is, the confidence interval is given by (8).

$$|{\theta}^{(h)}_{\pi}-{\hat{\theta}}^{(h)}_{\pi}|\leq{\kappa}^{(h)}_{\pi,\theta}$$ (8)

When deviating from the behavioral policy at step $h$ has no impact on next-state distributions, the estimate and CI for ${\theta}^{(h)}_{\pi}$ can be used as the estimate and CI for ${\alpha}^{(h)}_{\pi}$. However, when there is an impact on next-state distributions, valid estimation and CI construction for ${\alpha}^{(h)}_{\pi}$ requires us to propagate estimates and uncertainty from future steps to the current step. To enable this propagation, we take estimates for ${V_{\pi}}^{(h+1)}$ as our second input.

Input 2 (RL estimates): This input provides pessimistic, standard, and optimistic estimates for ${V_{\pi}}^{(h+1)}$ – denoted by ${\hat{V}_{\pi,p}}^{(h+1)}$, ${\hat{V}_{\pi}}^{(h+1)}$, and ${\hat{V}_{\pi,o}}^{(h+1)}$ respectively – such that the ordering in (9) holds. Further, with high probability, we require (10) holds – that is, the true value function is bounded by the pessimistic and optimistic value function estimates.

$$\forall x,\;0\leq{\hat{V}_{\pi,p}}^{(h+1)}(x)\leq{\hat{V}_{\pi}}^{(h+1)}(x)\leq{\hat{V}_{\pi,o}}^{(h+1)}(x)\leq{V}^{(h+1)}_{\text{max}}$$ (9)

$$\forall x,\;{V_{\pi}}^{(h+1)}(x)\in[{\hat{V}_{\pi,p}}^{(h+1)}(x),{\hat{V}_{\pi,o}}^{(h+1)}(x)]$$ (10)

There is a large and growing literature on value function estimation in RL, including optimistic and pessimistic value function estimation that are designed to satisfy (10) [e.g., Martin et al., 2017, Wang et al., 2019, Jin et al., 2021]. Thus, we can employ the most cutting-edge methods to construct these next-step value function estimates.

Input 2 gave us estimates for ${V_{\pi}}^{(h+1)}$ (next-step value), which we may need to propagate to the current step $h$ – when constructing an estimate and CI for ${\alpha}^{(h)}_{\pi}$. Since our goal is to interpolate between tight CB guarantees and always valid RL guarantees, unlike traditional RL algorithms, we want to be selective in propagating next-step estimates/uncertainties. Our final input, shift estimates, allows us to only propagate these estimates when required – enabling our adaptation to instance hardness.

Input 3 (Shift estimates): This input provides a estimate for ${{\Delta}}^{(h)}$ (see (5)) and an associated error bound – denoted by ${\hat{\Delta}}^{(h)}$ and ${\kappa}^{(h)}_{\pi,\Delta}$ respectively – such that (11) holds (recall from Section 3 that ${{\Delta}}^{(h)}$ satisfies the same bound). We also require (12) holds with high-probability.

$$\forall(x,a),\;\|{\hat{\Delta}}^{(h)}(\cdot|x,a)\|_{1}\leq 2$$ (11)

$$\mathbb{E}_{D(\pi_{b})}\big{\|}{\hat{\Delta}}^{(h)}(\cdot|{x}^{(h)},\pi)-{\Delta}^{(h)}(\cdot|{x}^{(h)},\pi)\big{\|}_{1}\leq{\kappa}^{(h)}_{\pi,\Delta}$$ (12)

Since the true shift model is a function of the true transition model (see Section 3), shift can be estimated via first estimating transition model and then calculating the treatment effect (due to deviating from the behavioral policy) of transitioning between any pair of states [see Moerland et al., 2023, for a survey on model-based RL and transition model estimation.].

4.2 Combining Inputs

We now have all our required input estimates, and can state our main result (Theorem 4.1) on constructing an estimate and CI for ${\alpha}^{(h)}_{\pi}$ – in a way that adapts to instance hardness.

Theorem 4.1.

Suppose we have: (1) CB inputs $({\hat{\theta}}^{(h)}_{\pi},{\kappa}^{(h)}_{\pi,\theta})$; (2) RL inputs $({\hat{V}_{\pi,p}}^{(h+1)},{\hat{V}_{\pi}}^{(h+1)},{\hat{V}_{\pi,o}}^{(h+1)})$ satisfying (9); and (3) shift inputs $({\hat{\Delta}}^{(h)},{\kappa}^{(h)}_{\pi,\Delta})$ satisfying (11) – such that (8), (10), and (12) hold with probability at least $1-\delta_{\text{in}}$.
Moreover, suppose we have a (holdout) dataset of $T$ trajectories $S=\{\tau_{t}\}_{t=1}^{T}$ – sampled from the distribution $D(\pi_{b})$ – that were not used for constructing input estimates. Our estimate for ${\alpha}^{(h)}_{\pi}$ is then denoted by ${\hat{\alpha}}^{(h)}_{\pi}$ and given by (13).

$${\hat{\alpha}}^{(h)}_{\pi}={\hat{\theta}}^{(h)}_{\pi}+\frac{1}{T}\sum_{t=1}^{T}\int_{x^{\prime}}{\hat{V}_{\pi}}^{(h+1)}(x^{\prime}){\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)}_{t},\pi)$$ (13)

Now for some fixed $\delta>0$, with probability at least $1-\delta-\delta_{\text{in}}$, we have the confidence interval in (14) holds.

$$|{{\alpha}}^{(h)}_{\pi}-{\hat{\alpha}}^{(h)}_{\pi}|\leq{L}^{(h)}_{\pi}.$$ (14)

Here ${L}^{(h)}_{\pi}$ is given by (15).

$${L}^{(h)}_{\pi}={\kappa}^{(h)}_{\pi,\theta}+{V_{\max}}^{(h+1)}{\kappa}^{(h)}_{\pi,\Delta}+6{V}^{(h+1)}_{\max}\sqrt{\frac{\ln(4/\delta)}{2T}}$$ (15)

$$+\frac{1}{T}\sum_{t=1}^{T}\int_{x^{\prime}}|\mathbb{E}_{D(\pi_{b})}[{\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)}_{t},\pi)]|{\Gamma_{\pi}}^{(h+1)}(x^{\prime})$$

Here ${\Gamma_{\pi}}^{(h+1)}$ is the difference between the optimistic and pessimistic estimates – it captures the uncertainty in the next-step value function estimates. That is, for all $x\in\mathcal{X}$, ${\Gamma_{\pi}}^{(h+1)}(x)={\hat{V}_{\pi,o}}^{(h+1)}(x)-{\hat{V}_{\pi,p}}^{(h+1)}(x)$.

One of the advantages of in-distribution supervised learning is that excess risk bounds only depend on complexity of hypothesis class (and number of training samples), with no dependence on size of feature space [see Shalev-Shwartz and Ben-David, 2014]. As discussed in Section 1, the statistical challenges of RL stem from the fact that learning under (state) distribution shifts is hard. For example, without additional assumptions, optimistic/pessimistic value function estimation have an unavoidable polynomial dependency on state-space size [Foster and Rakhlin, 2023]. Our goal is to avoid/minimize such dependencies when possible. The key benefit of Theorem 4.1 is that both our estimate $({\hat{\alpha}}^{(h)}_{\pi})$ and our CI width $({L}^{(h)}_{\pi})$ are “selective" in propagating/utilizing the RL estimates from input 2 – allowing us to only suffer from the slower worst-case RL estimation rates on harder instances. To better understand this, let us dig deeper into the terms in our CI width (${L}^{(h)}_{\pi}$).

Note that ${\kappa}^{(h)}_{\pi,\theta}$ and ${\kappa}^{(h)}_{\pi,\Delta}$ (see Inputs 1 and 3) bound errors averaged under the behavioral policy state-distribution – that is, they bound in-distribution average errors. Hence, with appropriate inputs, the first two terms in ${L}^{(h)}_{\pi}$ shrink quickly with no dependency on state-space size [Dudik et al., 2014, Shalev-Shwartz and Ben-David, 2014]. The third term in ${L}^{(h)}_{\pi}$, which enjoys a $\mathcal{O}(\sqrt{\log(1/\delta)/T})$ rate, also shrinks quickly to zero and has no dependency on state-space size.

We now only need to discuss the fourth and final term in ${L}^{(h)}_{\pi}$. Unlike the previous terms which bound in-distribution average errors, this term does depend on per-state (point-wise) errors (${\Gamma_{\pi}}^{(h+1)}(x^{\prime})$). The reason RL algorithms seek to bound per-state errors is because these guarantees do not depend on the state-distribution and are valid under any shift. We now illustrate how this robustness to state-distribution shift comes at a cost of larger error bounds, and argue that it is advantageous to scale these terms down with the estimated shift. First, as a sanity check, we show that this term is finite.

$$\int_{x^{\prime}}|{\hat{\Delta}}^{(h)}(x^{\prime}|{x_{t}}^{(h)},\pi)|\cdot{\Gamma_{\pi}}^{(h+1)}(x^{\prime})$$ (16)

$$\stackrel{{\scriptstyle(i)}}{{\leq}}\|{\hat{\Delta}}^{(h)}(\cdot|{x_{t}}^{(h)},\pi)\|_{1}\|{\Gamma_{\pi}}^{(h+1)}\|_{\infty}\stackrel{{\scriptstyle(ii)}}{{\leq}}2\|{\Gamma_{\pi}}^{(h+1)}\|_{\infty}.$$

where (i) follows from Hölder’s inequality, and (ii) follows from (11). Now that we know this term is finite, we can argue that it shrinks to zero. Since ${\Gamma_{\pi}}^{(h+1)}(x^{\prime})$ captures the size of per-state errors for pessimistic/optimistic value function estimates, we can expect this term to converges to zero in the limit with infinite data. The size of ${\Gamma_{\pi}}^{(h+1)}(x^{\prime})$ must depend on how often states similar to $x^{\prime}$ were visited at step $h+1$ in the training data for the RL input. For simplicity, let us consider the scenario when all states are visited uniformly at step $h+1$ under the distribution $D(\pi_{b})$. In such a scenario, the frequency at which states similar to $x^{\prime}$ were visited at step $h+1$ would depend on some measure of the size of the state space $\mathcal{X}$. This would imply that ${\Gamma_{\pi}}^{(h+1)}(x^{\prime})$ shrinks at a rate that depends on some measure of the size of the state space $\mathcal{X}$. That is, while these terms shrink, they shrink slowly. Hence per-state bounds, while independent of state-distribution, come at a cost of slower statistical rates. As shown in [Foster et al., 2021], such a dependence of confidence interval width on state-space size is unavoidable in the worst-case.

The key message of Theorem 4.1 is that we can move beyond this worst-case scenario by scaling these point-wise errors with the estimated shifts $({\hat{\Delta}}^{(h)})$. For example, when ${\hat{\Delta}}^{(h)}\equiv 0$, the fourth term in ${L_{\pi}}^{(h)}$ is zero, allowing us to recover contextual bandit-style guarantees that are independent of state-space size. It is worth noting that, even when state-space size is not a concern, being selective about error/estimate propagation can improve the resulting interval widths.

5 Selectively Pessimistic Value Iteration

Algorithm 1 Selectively Pessimistic Value Iteration

Inputs: Estimated reward model $\hat{R}$, estimated transition model $\hat{P}$, induced shift model $\hat{\Delta}$, and estimation uncertainty bonuses $b$.

Initialize ${\hat{V}}^{(H+1)},{\hat{V}_{p}}^{(H+1)},{\hat{V}_{o}}^{(H+1)}\equiv 0$.

for $h=H\mbox{ to }1$ do

     For all $(x,a)\in\mathcal{X}\times\mathcal{A}$, $\triangleright$ Estimate Q value.

$${\hat{Q}}^{(h)}(x,a)\leftarrow{\hat{R}}^{(h)}(x,a)+\int_{x^{\prime}}{\hat{V}}^{(h+1)}(x^{\prime}){\hat{P}}^{(h)}(x^{\prime}|x,a)$$ (17)

     For all $(x,a)\in\mathcal{X}\times\mathcal{A}$, $\triangleright$ Selective pessimism.

$${\hat{Q}_{\text{sp}}}^{(h)}(x,a)\leftarrow{\hat{Q}}^{(h)}(x,a)-{b}^{(h)}(x,a)$$ (18)

$$-\int_{x^{\prime}}|{\hat{\Delta}}^{(h)}(x^{\prime}|x,a)|\cdot({\hat{V}_{o}}^{(h+1)}(x^{\prime})-{\hat{V}_{p}}^{(h+1)}(x^{\prime})),$$

$${{\pi}}^{(h)}(x)\leftarrow\operatorname*{arg\,max}_{a}{\hat{Q}_{\text{sp}}}^{(h)}(x,a)$$

     For all $x\in\mathcal{X}$, $\triangleright$ Estimate pessimistic value.

$${\hat{V}_{\text{p}}}^{(h)}(x)\leftarrow\max\Big{(}0,\;{\hat{R}}^{(h)}(x,{\pi}^{(h)}(x))$$ (19)

$$+\int_{x^{\prime}}{\hat{V}_{p}}^{(h+1)}(x^{\prime}){\hat{P}}^{(h)}(x^{\prime}|x,a)-{b}^{(h)}(x,a)\Big{)}$$

     For all $x\in\mathcal{X}$, $\triangleright$ Estimate optimistic value.

$${\hat{V}_{o}}^{(h)}(x)\leftarrow\min\Big{(}{V_{\max}}^{(h)},\;{\hat{R}}^{(h)}(x,{\pi}^{(h)}(x))$$ (20)

$$+\int_{x^{\prime}}{\hat{V}_{o}}^{(h+1)}(x^{\prime}){\hat{P}}^{(h)}(x^{\prime}|x,a)+{b}^{(h)}(x,a)\Big{)}$$

     For all $x\in\mathcal{X}$, $\triangleright$ Estimate value.

$${\hat{V}}^{(h)}(x)\leftarrow\min\Big{(}{\hat{V}_{o}}^{(h)}(x),$$ (21)

$$\max\Big{(}{\hat{V}_{p}}^{(h)}(x),{\hat{Q}}^{(h)}(x,{\pi}^{(h)}(x))\Big{)}\Big{)}$$

end for

To demonstrate the benefits of our insights from Section 4, we now present an algorithm for offline policy learning. At a high-level, we modify pessimistic value iteration (PVI) [Jin et al., 2021], a dynamic programming (DP) algorithm that chooses the policy at each step $h$ by maximizing pessimistic value function estimates. Pessimistic value maximization helps PVI avoids model exploitation – that is, PVI avoids picking policies with inaccurately high estimated values at step $h$ by penalizing value function estimate uncertainty. We argue that, depending on the instance, such penalties due to estimate uncertainty may be larger than necessary for selecting the policy at step $h$.

In particular, based on the results in Section 4, uncertainty from later steps does not always need to be fully propagated to estimate lower bounds for the effect of deviating from the behavioral policy at step $h$ (after fixing the policy for all future steps). Since we can view maximizing this treatment effect as the goal of DP algorithms at any step $h$, maximizing the tighter lower bounds from Section 4 should allow us to do better on easier CB-like instances (while always avoiding model exploitation) – motivating the design of Algorithm 1.

PVI often relies on point-wise estimation uncertainty bonus $b=({b}^{(h)})_{h=1}^{H}$ to account for reward and transition estimation errors. For any $(x,a,h)$, ${b}^{(h)}(x,a)$ bounds the total model estimation errors at $(x,a,h)$. At any $h$, PVI uses bonuses to account for estimation errors at step $h$ and propagate estimation errors from future steps (by propagating pessimistic values from step $h+1$ that were computed using these bonuses). We also similarly rely on these bonuses to account for model estimation errors.

Our algorithm (Algorithm 1), selectively pessimistic value iteration (SPVI), takes as input: an estimated reward model $\hat{R}$, an estimated transition model $\hat{P}$, an induced shift model $\hat{\Delta}$ such that ${\hat{\Delta}}^{(h)}(x^{\prime}|x,a)\equiv{\hat{P}}^{(h)}(x^{\prime}|x,a)-{\hat{P}}^{(h)}(x^{\prime}|x,\pi_{b})$ for all $(x,a,x^{\prime},h)$, and a point-wise estimation uncertainty bonus $b=({b}^{(h)})_{h=1}^{H}$. No additional data is required. The algorithm uses these inputs to construct a policy $\pi=({\pi}^{(h)})_{h=1}^{H}$ iteratively. That is, SPVI constructs ${\pi}^{(h)}$ while iterating over steps $h=H$ to $h=1$.

At step $h$, SPVI computes: (1) the Q-value estimate for step $h$ (which is well defined since $({\pi}^{(h+1)},..,{\pi}^{(H)})$ is fixed) denoted by ${\hat{Q}}^{(h)}$ (see (17)), (2) the policy at step $h$ denoted by ${\pi}^{(h)}$ (see (18)), and (3) pessimistic (${\hat{V}_{p}}^{(h)}$), optimistic (${\hat{V}_{o}}^{(h)}$), and standard (${\hat{V}}^{(h)}$) value functions for step $h$ (see (19), (20), and (21)). Steps one and three are fairly standard, as several RL algorithms construct Q-values and pessimistic/standard/optimistic value functions. Hence, for brevity, we only explain step 2 (constructing ${\pi}^{(h)}$) – the crux of our modification.

As we argued earlier, one can view the goal at step $h$ as selecting ${\pi}^{(h)}$ in order to maximize ${\alpha}^{(h)}_{\pi}$. We construct a tight lower bound for ${\alpha}^{(h)}_{\pi}$ and argue that step 2 maximizes this bound. Similar to standard pessimistic value estimation, we let the bonus ${b}^{(h)}(x,a)$ bound the total model estimation errors at $(x,a,h)$. Now from the analysis in Theorem 4.1, we have (22) is a valid lowed bound on ${\alpha}^{(h)}_{\pi}$.

$${\alpha}^{(h)}_{\pi}=\mathbb{E}_{D(\pi_{b})}[{V}^{(h)}_{\tilde{\pi}_{h}}({x}^{(h)})-{V}^{(h)}_{\tilde{\pi}_{h+1}}({x}^{(h)})]$$ (22)

$$\geq\mathbb{E}_{D(\pi_{b})}\bigg{[}{\hat{Q}}^{(h)}({x}^{(h)},{\pi}^{(h)})-{\hat{Q}}^{(h)}({x}^{(h)},{\pi_{b}}^{(h)})$$

$$-{b}^{(h)}({x}^{(h)},{\pi}^{(h)})-\int_{x^{\prime}}|{\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)},\pi)|{\Gamma}^{(h+1)}(x^{\prime})\bigg{]}$$

$$=\mathbb{E}_{D(\pi_{b})}[{\hat{Q}_{\text{sp}}}^{(h)}({x}^{(h)},{\pi}^{(h)})-{\hat{Q}}^{(h)}({x}^{(h)},{\pi_{b}}^{(h)})]$$

Here ${\Gamma}^{(h+1)}(x)$ is a shorthand for ${\hat{V}_{o}}^{(h+1)}(x)-{\hat{V}_{p}}^{(h+1)}(x)$; and for any $g\in\{{\hat{Q}}^{(h)},{\hat{Q}_{\text{sp}}}^{(h)},{b}^{(h)}\}$ we let $g(x,{\pi}^{(h)})$ be a shorthand for $g(x,{\pi}^{(h)}(x))$. Now from (18), since ${{\pi}}^{(h)}(x)\equiv\operatorname*{arg\,max}_{a}{\hat{Q}_{\text{sp}}}^{(h)}(x,a)$ for all $x$, we have that step 2 (constructing ${\pi}^{(h)}$) is equivalent to maximizing the lower bound in (22). Importantly, by maximizing this tight lower bound we can do better on easier CB-like instances (while always avoiding model exploitation by penalizing uncertainty in estimating ${\alpha}^{(h)}_{\pi}$). This completes our justification of Algorithm 1 (SPVI).

6 Simulation

Figure 1: ChainBandit MDP with horizon/length of 4 (this is an adjustable environment parameter). The environment has two chains, a top chain, and a bottom chain. The environment also has three actions given by $a_{1},a_{2},a_{3}$. The top chain states are the most rewarding. The agent starts at (1,0). At any state in the bottom chain, all the actions lead to the same transition (which is to move to the next state in the bottom chain) and are essentially bandit states. In the top chain, both $a_{1}$ and $a_{2}$ lead to the same transitions (which is to move to the next state in the top chain), and $a_{3}$ makes the agent move to the next state in the bottom chain. In the top chain, the highest cumulative reward comes from never taking action $a_{3}$; however, the highest immediate reward comes from selecting the action $a_{3}$ (which makes planning beneficial in this environment). Note that at every state, action $a_{3}$ is a sub-optimal action.

To illustrate our insights, we consider a simple toy environment called “ChainBandit". As described in Figure 1, this environment is constructed to have both dynamic (some actions result in a different next-state distribution) and bandit-like (non-dynamic) elements. Ensuring that while planning and some estimate/uncertainty propagation is necessary, complete uncertainty propagation can be unnecessary to evaluate policies of interest. Throughout this section, we consider ChainBandit with a chain length of 3 and consider the following behavioral policy ($\pi_{b}$) – at every state and step, $\pi_{b}$ selects action $a_{3}$ with probability $0.8$ and selects the other two actions with probability $0.1$ respectively. From the data collected, we evaluate standard and selective uncertainty propagation for tasks of: (1) estimating upper/lower bounds for ${\alpha}^{(h)}_{\pi}$; and (2) offline policy learning.

Since uncertainty propagation is the focus of this simulation, in order to make a fair comparison, both standard and selective propagation: use the same tabular approach to estimate a (step independent) reward/transition model; and use the same (step independent) count-based bonuses to account for model estimation errors. Here bonus is given by $b(x,a)=\sqrt{\ln(|\mathcal{X}|*|\mathcal{A}|*H/\delta)/n(x,a)}$ – where $n(x,a)$ is the number of times action $a$ was taken at state $x$, and confidence parameter $\delta=0.05$.

For the step $h$ and policy $\pi$ of interest, standard CIs for ${\alpha}^{(h)}_{\pi}$ are constructed using pessimistic/optimistic value estimates at step $h$ for policies $\tilde{\pi}_{h}$ and $\tilde{\pi}_{h+1}$ – i.e, utilizing (23).

$$\mathbb{E}_{D(\pi_{b})}[{\hat{V}_{\tilde{\pi}_{h},o}}^{(h)}({x}^{(h)})-{\hat{V}_{\tilde{\pi}_{h+1},p}}^{(h)}({x}^{(h)})]\geq{\alpha}^{(h)}_{\pi}$$ (23)

$$\geq\mathbb{E}_{D(\pi_{b})}[{\hat{V}_{\tilde{\pi}_{h},p}}^{(h)}({x}^{(h)})-{\hat{V}_{\tilde{\pi}_{h+1},o}}^{(h)}({x}^{(h)})]$$

For selective CIs, we use (22) to construct the lower bound and similarly construct the upper bound. Note that converting inequalities like (22) and (23) into empirical bounds is straightforward since our training is from $D(\pi_{b})$. Figure 2 plots CIs for both selective and standard uncertainty propagation, when varying the evaluation policy. As expected, benefits over standard pessimism are larger when next-state distribution shift is smaller – that is, when evaluation policy is closer to the behavioral policy.

In Figure 3, we plot the value of learnt policy from various algorithms as we vary the number of training episodes. In particular, we compare SPVI, PVI [Jin et al., 2021], and pessimistic supervised learning (PSL). Here PSL refers to a pessimistic bandit policy optimization applied to each step without planning. The ChainBandit environment benefits from planning, so PSL performs poorly as expected.

On all Chain Bandit simulations we tried, SPVI was by far the best-performing algorithm. The reason we considered the behavioral policy described earlier was that it was more disadvantageous for SPVI. In particular, since selective pessimism has an initial bias against policies that lead to significant shifts, we chose a highly sub-optimal behavioral policy (selecting $a_{3}$ with probability $0.8$). While this leads to a worse start for SPVI than PVI, eventually, SPVI outperforms the other algorithms. We also run simulations for CI construction and policy learning on the standard GridWorld (see Appendix B) – since this is a very dynamic environment, both standard and selective propagation have similar performance.

Figure 2: We plot CIs for ${\alpha}^{(2)}_{\pi}$ while varying the evaluation policy. These evaluation policies are parameterized by $\lambda\in[0,1]$. For all states/steps, the probability of selecting $a_{1},a_{2}$ and $a_{3}$ are $(1-\lambda)/2,(1-\lambda)/2,$ and $\lambda$ respectively. Note that the evaluation policy is the same as the behavioral policy for $\lambda=0.8$. The number of training episodes is 10000, and the plots are averaged over 10 runs.

Figure 3: Policy learning with a bad behavioral policy

Conclusion: We introduce selective propagation, a general approach to interpolate between CB and RL techniques – achieving guarantees that adapt to instance hardness. Further developing this could impact real world problems (e.g., recommendation systems, mHealth, EdTech) that lie in between the two settings.

Selective Uncertainty Propagation in Offline RL
(Supplementary Material)

Appendix A Proof for Section 4

In this section we prove our main theoretical result, Theorem 4.1. For convenience, we restate it below. See 4.1

We now begin the proof of Theorem 4.1. We start by proving Lemma A.1, which shows ${\tilde{\alpha}}^{(h)}_{\pi}$ is close to our target estimand. Where ${\tilde{\alpha}}^{(h)}_{\pi}$ is defined in (24) and can be viewed as a less empirical version of ${\hat{\alpha}}^{(h)}_{\pi}$ (our treatment effect estimator at step $h$).

$${\tilde{\alpha}}^{(h)}_{\pi}={\hat{\theta}}^{(h)}_{\pi}+\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}{\hat{V}_{\pi}}^{(h+1)}(x^{\prime}){\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)},\pi)\bigg{]}$$ (24)

Lemma A.1.

Under the conditions of Theorem 4.1, we show the following bound holds with probability $1-\delta_{\text{in}}$:

$$|{\tilde{\alpha}}^{(h)}_{\pi}-{\alpha}^{(h)}_{\pi}|\leq{\tilde{L}_{\pi}}^{(h)}.$$ (25)

Where ${\tilde{L}_{\pi}}^{(h)}$ is given by (26).

$${\tilde{L}_{\pi}}^{(h)}={\kappa}^{(h)}_{\pi,\theta}+{V_{\max}}^{(h+1)}{\kappa}^{(h)}_{\pi,\Delta}+\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}|\mathbb{E}_{D(\pi_{b})}[{\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)},\pi)]|\cdot{{\Gamma}_{\pi}}^{(h+1)}(x^{\prime})\bigg{]}$$ (26)

Proof.

We start by simplifying ${\alpha}^{(h)}_{\pi}$, defined in (4).

$${\alpha}^{(h)}_{\pi}=\mathbb{E}_{D(\pi_{b})}[{V}^{(h)}_{\tilde{\pi}_{h}}({x}^{(h)})-{V}^{(h)}_{\tilde{\pi}_{h+1}}({x}^{(h)})]$$ (27)

$$\stackrel{{\scriptstyle(i)}}{{=}}\mathbb{E}_{D(\pi_{b})}\bigg{[}\bigg{(}{R}^{(h)}(x,\tilde{\pi}_{h})+\int_{x^{\prime}}{V_{\pi}}^{(h+1)}(x^{\prime}){P}^{(h)}(x^{\prime}|x,\tilde{\pi}_{h})\bigg{)}-\bigg{(}{R}^{(h)}(x,\tilde{\pi}_{h+1})+\int_{x^{\prime}}{V_{\pi}}^{(h+1)}(x^{\prime}){P}^{(h)}(x^{\prime}|x,\tilde{\pi}_{h+1})\bigg{)}\bigg{]}$$

$$\stackrel{{\scriptstyle(ii)}}{{=}}\mathbb{E}_{D(\pi_{b})}\bigg{[}\bigg{(}{R}^{(h)}(x,\tilde{\pi}_{h})-{R}^{(h)}(x,\tilde{\pi}_{h+1})\bigg{)}+\bigg{(}\int_{x^{\prime}}{V_{\pi}}^{(h+1)}(x^{\prime})\Big{(}{P}^{(h)}(x^{\prime}|x,\tilde{\pi}_{h})-{P}^{(h)}(x^{\prime}|x,\tilde{\pi}_{h+1})\Big{)}\bigg{)}\bigg{]}$$

$$\stackrel{{\scriptstyle(iii)}}{{=}}{\theta}^{(h)}_{\pi}+\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}{V_{\pi}}^{(h+1)}(x^{\prime}){\Delta}^{(h)}(x^{\prime}|x,\pi)\bigg{]}$$

Here (i) follows from (2), (ii) follows from re-arranging terms, and (iii) follows from (7) and (5).

With probability $1-\delta_{\text{in}}$, we have the input guarantees of Section 4 hold. Hence, utilizing these guarantees, we can bound the distance between ${\tilde{\alpha}}^{(h)}_{\pi}$ and the target estimand ${\alpha}^{(h)}_{\pi}$.

$$|{\tilde{\alpha}}^{(h)}_{\pi}-{\alpha}^{(h)}_{\pi}|$$ (28)

$$\stackrel{{\scriptstyle(i)}}{{=}}\bigg{|}\bigg{(}{\hat{\theta}}^{(h)}_{\pi}+\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}{\hat{V}_{\pi}}^{(h+1)}(x^{\prime}){\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)},\pi)\bigg{]}\bigg{)}-\bigg{(}{\theta}^{(h)}_{\pi}+\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}{V_{\pi}}^{(h+1)}(x^{\prime}){\Delta}^{(h)}(x^{\prime}|x,\pi)\bigg{]}\bigg{)}\bigg{|}$$

$$\stackrel{{\scriptstyle(ii)}}{{\leq}}{\kappa}^{(h)}_{\pi,\theta}+\bigg{|}\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}\bigg{(}{\hat{V}_{\pi}}^{(h+1)}(x^{\prime}){\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)},\pi)-{V_{\pi}}^{(h+1)}(x^{\prime}){{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)},\pi)\bigg{)}\bigg{]}\bigg{|}$$

$$\stackrel{{\scriptstyle(iii)}}{{\leq}}{\kappa}^{(h)}_{\pi,\theta}+\bigg{|}\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}\Big{(}{\hat{V}_{\pi}}^{(h+1)}(x^{\prime})-{V_{\pi}}^{(h+1)}(x^{\prime})\Big{)}{\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)},\pi)\bigg{]}\bigg{|}$$

$$+\bigg{|}\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}{V_{\pi}}^{(h+1)}(x^{\prime})\Big{(}{\hat{\Delta}}^{(h)}-{{\Delta}}^{(h)}\Big{)}(x^{\prime}|{x}^{(h)},\pi)\bigg{)}\bigg{]}\bigg{|}$$

$$\stackrel{{\scriptstyle(iv)}}{{\leq}}{\kappa}^{(h)}_{\pi,\theta}+\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}|\mathbb{E}_{D(\pi_{b})}[{\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)},\pi)]|\cdot{{\Gamma}_{\pi}}^{(h+1)}(x^{\prime})\bigg{]}+{V_{\max}}^{(h+1)}{\kappa}^{(h)}_{\pi,\Delta}=:{\tilde{L}_{\pi}}^{(h)}$$

Here (i) follows from (24) and (27), (ii) follows from triangle inequality and the input guarantee, (iii) follows from triangle inequality, and finally (iv) follows from Cauchy-Schwarz inequality and the input guarantees. ∎

Having shown in Lemma A.1 that ${\tilde{\alpha}}^{(h)}_{\pi}$ is close to our target estimand, we only need to show two more facts to complete the proof of Theorem 4.1: (i) we need to show ${\tilde{\alpha}}^{(h)}_{\pi}$ is close to our treatment effect estimator ${\hat{\alpha}}^{(h)}_{\pi}$ (defined in (13)), and (ii) we need to show ${\tilde{L}_{\pi}}^{(h)}$ is sufficiently smaller that ${{L}_{\pi}}^{(h)}$. We show both these statements hold with high-probability.

We start with showing ${\tilde{\alpha}}^{(h)}_{\pi}$ is close to ${\hat{\alpha}}^{(h)}_{\pi}$ with high-probability. In particular, we have (29) holds with probability at least $1-\delta/2$.

$$|{\tilde{\alpha}}^{(h)}_{\pi}-{\hat{\alpha}}^{(h)}_{\pi}|$$ (29)

$$=\bigg{|}\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}{\hat{V}_{\pi}}^{(h+1)}(x^{\prime}){\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)},\pi)\bigg{]}-\frac{1}{T}\sum_{t=1}^{T}\int_{x^{\prime}}{\hat{V}_{\pi}}^{(h+1)}(x^{\prime}){\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)}_{t},\pi)\bigg{|}$$

$$\leq 2{V_{\max}}^{(h+1)}\sqrt{\frac{\ln(4/\delta)}{2T}}.$$

In (29), the first equality follows from (13) and (24). The first inequality follows from Hoeffding’s inequality the fact that $\int_{x^{\prime}}{\hat{V}_{\pi}}^{(h+1)}(x^{\prime}){\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)}_{t},\pi)\in[-{V_{\max}}^{(h+1)},{V_{\max}}^{(h+1)}]$.

From Hoeffding’s inequality and the fact that (16) holds, we also have (30) holds with probability at least $1-\delta/2$.

$$\mathbb{E}_{D(\pi_{b})}\bigg{[}\int_{x^{\prime}}|\mathbb{E}_{D(\pi_{b})}[{\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)},\pi)]|\cdot{{\Gamma}_{\pi}}^{(h+1)}(x^{\prime})\bigg{]}-\frac{1}{T}\sum_{t=1}^{T}\int_{x^{\prime}}|\mathbb{E}_{D(\pi_{b})}[{\hat{\Delta}}^{(h)}(x^{\prime}|{x}^{(h)}_{t},\pi)]|\cdot{{\Gamma}_{\pi}}^{(h+1)}(x^{\prime})$$ (30)

$$\leq 4{V_{\max}}^{(h+1)}\sqrt{\frac{\ln(4/\delta)}{2T}}.$$

Note that (30) implies ${\tilde{L}}^{(h)}_{\pi}+2{V_{\max}}^{(h+1)}\sqrt{\frac{\ln(4/\delta)}{2T}}$ is no larger than ${L}^{(h)}_{\pi}$.

Hence, with probability at least $1-\delta_{\text{in}}-\delta$, we have (25), (29), and (30) hold. We now use these equation to show that (31) holds.

$$|{\alpha}^{(h)}_{\pi}-{\hat{\alpha}}^{(h)}_{\pi}|$$ (31)

$$\stackrel{{\scriptstyle(i)}}{{\leq}}|{\alpha}^{(h)}_{\pi}-{\tilde{\alpha}}^{(h)}_{\pi}|+|{\tilde{\alpha}}^{(h)}_{\pi}-{\hat{\alpha}}^{(h)}_{\pi}|$$

$$\stackrel{{\scriptstyle(ii)}}{{\leq}}{\tilde{L}_{\pi}}^{(h)}+2{V_{\max}}^{(h+1)}\sqrt{\frac{\ln(4/\delta)}{2T}}$$

$$\stackrel{{\scriptstyle(iii)}}{{\leq}}{{L}_{\pi}}^{(h)}.$$

Here (i) follows from triangle inequality; (ii) follows from (25) and (29); and (iii) follows from (30). Therefore, we have shown that (31) holds with probability at least $1-\delta_{\text{in}}-\delta$. This completes the proof of Theorem 4.1.

Appendix B Additional Simulation

Similar to Section 6, we now run simulations for the standard GridWorld environment (with width $8$ and height $3$). Here states are discrete points on a bounded two-dimensional grid. The agent’s starting state is sampled uniformly at random from the grid, and the agent should learn to reach a specified goal state (which is an absorbing state). Upon transitioning to the goal state, the agent receives a reward of one and receives a reward of zero otherwise. The agent has 4 actions (left, right, up, and down). These actions make the agent move one step in that direction if possible. If the agent is at the boundary of the grid and can’t move in the direction selected, the agent continues to stay in the same state. Since we make the goal state an absorbing state, all actions at this state lead to the agent continuing to stay in this state. We set the starting state for the GridWorld as (1,1), the terminal state as (2,2), and the horizon as 3. For all states/steps, our behavioral policy samples actions (left, right, up, and down) with the probabilities $(0.20,0.10,0.50,0.20)$. Other choices for GridWorld environment parameters and behavioral policy appear to generate similar plots. All our plots in this section are averaged over five simulation runs.

Figure 4: Plotting ${\alpha}^{(2)}_{\pi}$ with varying the evaluation policy $\pi$. Here $\lambda$ (evaluation policy probability of taking the down action) corresponds to the X-axis.

We constructed intervals for ${\alpha}^{(2)}_{\pi}$, the treatment effect at step $2$, using both selective and standard uncertainty propagation. In Figure 4, we plot the CIs for ${\alpha}^{(2)}_{\pi}$ as we vary the evaluation policy. With $\lambda$ parameterizing the evaluation policies. For all states/steps, the evaluation policy corresponding to $\lambda$ samples actions (left, right, up, and down) with the probabilities $(0.25,0.20,0.55-\lambda,\lambda)$. Our CIs are constructed from sampling a dataset of $2000$ episodes. In Figure 4, we see that both selective and standard uncertainty propagation have a similar performance – this is understandable because GridWorld is a dynamic environment (each action has a different next-state distribution), hence estimate/uncertainty propagation is less avoidable here for valid CI construction.

Figure 5: Learning Experiment on GridWorld

We also compare the policy learning algorithms on the same GridWorld environment, with the same behavioral policy. In Figure 5, we plot the value of policy learnt by SPVI, PVI, and PSL – as we vary the number of training episodes. We see PSL still performs terribly since planning is necessary. However, since GridWorld is a dynamic environment, PVI and SPVI have a similar performance.