SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section, we provide regret analysis of learning in MDPs with stochastic delays. For the sake of simplicity, we assume the reward is known, however, extension to unknown reward causes no real difficulty. Motivated by the augmented MDP reformulation, we introduce our learning algorithm in Algorithm 2. In Line 5, unobserved states and rewards are returned to the agent as described in Protocol 1. Using the data set, we construct bonus functions compensating the uncertainty in one-step transitions of the original MDP. This largely sharpens the confidence region, yet still ensures a valid optimism. We emphasize that in Line 9, we are planning on  $\widehat{\mathbf{MDP}}_{\mathrm{aug}}$  involving the augmented transitions and expanded states of  $\tau \in \widetilde{S}_{\mathrm{aug}}$ . Only in this way, we can obtain an executable policy in delayed MDPs. The planning complexity is  $SA^H$  though.
4.1 Regret bound
We define regret in delayed MDP as
$$
\operatorname {R e g r e t} (K) = \sum_ {k = 1} ^ {K} \max _ {\pi \in \Pi_ {\mathrm {e x e c}}} V _ {1} ^ {\pi} (s _ {1} ^ {k}) - \sum_ {k = 1} ^ {K} V _ {1} ^ {\pi_ {k}} (s _ {1} ^ {k}),
$$
where  $V_{1}^{\pi}$  is the value function of the original MDP. Although the regret here is defined on the original MDP, it is equivalent to the regret of the same policy on  $\mathrm{MDP}_{\mathrm{aug}}$  and further  $\widetilde{\mathrm{MDP}}_{\mathrm{aug}}$  by Proposition 3.1. Note that we are comparing with the best executable policy. The performance degradation caused by observation delay is discussed in Section 4.2. The following theorem bounds the regret.
Theorem 4.1 (Regret bound for Delayed MDP). Suppose Assumption 2.1 holds. Let  $\gamma \in (0,1)$  be any failure probability. With probability  $1 - \gamma$ , the regret of Algorithm 2 satisfies
$$
\operatorname {R e g r e t} (K) \leq c \left(H ^ {4} \sqrt {S A K \iota} + H ^ {4} S ^ {2} A \iota^ {2}\right),
$$
where  $\iota = \log \frac{SAHK}{\gamma}$  and  $c$  is a constant.
The proof is provided in Appendix B.1. We discuss several implications.
Algorithm 2 Policy learning for delayed MDPs using  $\widetilde{\mathsf{MDP}}_{\mathrm{aug}}$
1: Input: Original horizon  $H$ , extended horizon  $\tilde{H}$ , policy class  $\Pi_{\mathrm{exec}}$ , failure probability  $\gamma$ .
2: Init:  $V_{\widetilde{H} +1}(\tau) = 0$  and  $Q_{\widetilde{H}}(\tau ,a) = H$  for any  $\tau$  and  $a$ , data set  $\mathcal{D}^0 = \emptyset$ , initial policy  $\pi^0$ .
3: for episode  $k = 1, \dots, K$  do
4: Execute policy  $\pi^{k - 1}$  for  $\tilde{H}$  steps.
5: After the episode ends, collect data  $\mathcal{D}^k = \mathcal{D}^{k - 1}\cup \{(s_h^k,a_h^k,r_h^k,\Delta_h^k)\}_{h = 1}^H$
6: On data set  $\mathcal{D}^k$ , compute counting numbers  $N_h^k(s_h, a_h)$ ,  $N_h^k(s_h, a_h, s_{h+1})$  and  $N_h^k(s_h, a_h, \delta_h) = \sum_{j=1}^{k} \mathbb{1}\{s_h^j = s_h, a_h^j = a_h, \Delta_h^j = \delta_h\}$ .
7: Estimate transition probabilities and delay distributions via
$$
\widehat {p} _ {h} ^ {k} (s _ {h + 1} | s _ {h}, a _ {h}) = \frac {N _ {h} ^ {k} (s _ {h} , a _ {h} , s _ {h + 1})}{N _ {h} ^ {k} (s _ {h} , a _ {h})}, \quad \text {a n d} \quad \widehat {\theta} _ {\text {d e l a y}} ^ {k} (s _ {h}, a _ {h}, \delta_ {h}) = \frac {N _ {h} ^ {k} (s _ {h} , a _ {h} , \delta_ {h})}{\sum_ {\delta \geq \delta_ {h}} N _ {h} ^ {k} (s _ {h} , a _ {h} , \delta)}.
$$
Then estimators of  $\widetilde{p}_{h,\mathrm{aug}}$  in  $\widetilde{\mathrm{MDP}}_{\mathrm{aug}}$  is computed using  $\widehat{p}_h^k$  and  $\widehat{\theta}_{\mathrm{delay}}^k$
8: Set bonus function as
$$
b _ {h} ^ {k} (\tau_ {h}, a _ {h}) = c H \left(\sqrt {\frac {(H \wedge D) \iota}{N _ {t _ {h}} ^ {k} (s _ {t _ {h}} , a _ {t _ {h}} , \delta_ {t _ {h}})}} + \sqrt {\frac {(H \wedge D) \iota}{N _ {t _ {h}} ^ {k} (s _ {t _ {h}} , a _ {t _ {h}})}}\right)
$$
for  $\iota = \log \frac{SAKH}{\gamma}$  and  $c$  sufficiently large.
9: Run optimistic value iteration in  $\widetilde{\mathrm{MDP}}_{\mathrm{aug}}$  for  $\tilde{H}$  steps and obtain  $\pi^k \in \Pi_{\mathrm{exec}}$ .
10: end for
11: Return: Learned policy  $\pi_{1:H}^{k}$  for  $k = 1, \ldots, K$ .
Sharp dependence on  $S$  and  $A$  Theorem 4.1 has a sharp dependence on  $S$  and  $A$ , although the expanded state space  $\widetilde{S}_{\mathrm{aug}}$  has a cardinality bounded by  $SA^H$ . Naively learning and planning in  $\widetilde{\mathrm{MDP}}_{\mathrm{aug}}$  would suffer from the exponential enlargement of  $A^H$ . However, we identify the sparse structures in the transition probabilities. As can be seen,  $\widetilde{p}_{h,\mathrm{aug}}$  only involves one-step transitions in the original MDP and some conditionally independent delay distributions. Such structures lead to a rather easy estimation of  $\widetilde{p}_{h,\mathrm{aug}}$ , which can be constructed from the estimators of one-step transitions in the original MDP. Meanwhile, the sparse structures make exploration in  $\widetilde{\mathrm{MDP}}_{\mathrm{aug}}$  efficient, due to many unreachable states.
Effect of the delay distribution and delay length Theorem 4.1 holds for arbitrary conditionally independent delay distributions, even include heavy-tailed distributions. In the worst case of unbounded delays, Theorem 4.1 gives rise to a  $\mathcal{O}(H^4\sqrt{SAK\iota})$  regret. The reason to this is that if the delay is larger than  $H$ , then the corresponding state will only be observed after an episode ends and won't be used in planning. Therefore, we can truncate the delay at  $H$ , regardless of its tail distributions.
When the maximal length of delay is bounded by  $D < H$ , e.g., CDMDPs with  $d_h = D$  for any  $h$ , Theorem 4.1 implies that the regret is bounded by
$$
\operatorname {R e g r e t} (K) \leq c \left((D + 1) ^ {5 / 2} \sqrt {H ^ {3} S A K \iota} + H ^ {4} S ^ {2} A \iota^ {2}\right)
$$
for a constant  $c$ . A proof is provided in Appendix B.2. As can be seen, as the length of delay increases, the regret bound enlarges, reflecting the increased difficulty of long delays. Moreover, when  $D = 0$ , that is, no observation delays, the regret bound recovers that in standard MDPs.
4.2 Performance degradation of policy class  $\Pi_{\mathrm{exec}}$
This section devotes to quantify the performance degradation caused by delayed observations. In particular, we bound the value difference between the best executable policy and the best Markov policy in a no delay environment. Recall that  $V_{1}$  is the value function of the original MDP. We denote
$$
\pi_ {\text {n o d e l a y}} ^ {*} = \arg \max  _ {\pi} V _ {1} ^ {\pi} (s _ {1}) \quad \text {a n d} \quad \pi_ {\text {d e l a y}} ^ {*} = \arg \max  _ {\pi \in \Pi_ {\text {e x e c}}} V _ {1} ^ {\pi} (s _ {1})
$$
as the best vanilla optimal policy and executable policy, respectively. The values achieved by  $\pi_{\mathrm{nodelay}}^*$  and  $\pi_{\mathrm{delay}}^*$  are denoted as  $V_{1,\mathrm{nodelay}}^{*}(s_1)$  and  $V_{1,\mathrm{delay}}^{*}(s_1)$ , respectively. The gap between  $V_{1,\mathrm{nodelay}}^{*}$  and  $V_{1,\mathrm{delay}}^{*}$  quantifies the performance degradation, which is denoted as  $\mathrm{gap}(s_1) = V_{1,\mathrm{nodelay}}^{*}(s_1) - V_{1,\mathrm{delay}}^{*}(s_1)$ . We bound gap in Proposition B.2 in Appendix due to space limit.
In a nutshell, we show that the performance degradation gap is highly relevant to the belief distribution  $\mathfrak{b}_h(\cdot |\tau)$ . When  $\mathfrak{b}_h(\cdot |\tau)$  is evenly spread, meaning that the entropy of  $\mathfrak{b}_h$  is high and inferring the current unseen state is difficult, we potentially suffer from a large gap. On the contrary, when  $\mathfrak{b}_h(\cdot |\tau)$  is nearly deterministic, the performance degradation is small. In the special case of deterministic transitions, we have  $\mathrm{gap} = 0$ .
4.3 The (mysterious) effect of delay on the optimal value
To further understand the effect of the delay on the optimal value, we provide the following dichotomy. On the one hand, we show that there exists an MDP instance, such that a constant delay of  $d$  steps does not hurt the performance. On the other hand, in the same MDP instance, a constant delay of  $d + 1$  steps suffers from a constant performance drop.
Proposition 4.2. Consider constant delayed MDPs. Fix a positive integer  $d < H$ . Then there exists an MDP instance such that the following two items hold simultaneously.
- When delay is  $d$ , it holds that  $\frac{1}{K}\sum_{k=1}^{K}\mathsf{gap}(s_1^k) = 0$ .
- When delay is  $d + 1$ , it holds that  $\frac{1}{K}\sum_{k = 1}^{K}\mathsf{gap}(s_1^k)\geq \frac{1}{2} -\sqrt{\frac{1}{2K}\log\frac{1}{\gamma}}$ , with probability  $1 - \gamma$
The proof is provided in Appendix B.4. We remark that Proposition 4.2 says that observation delay can be dangerous, even with the slightest possible number of steps. The idea behind Proposition 4.2 is consistent with the analysis on gap. In particular, we construct an MDP instance demonstrated in Figure 2. The reward vanishes at all times but  $d + 1$ . When delay is  $d$ , the initial state  $s_1$  is revealed and the policy can choose the best action to receive a reward. When delay is  $d + 1$ , however, there is always a  $1/2$  probability of missing the best action for any policy, which leads to a constant performance degradation.
Figure 2: MDP instance on two states with two actions. The transition is lazy until time  $d$ . Then the transition is uniform regardless of actions for time  $d + 1$ . Reward is nonzero only at time  $d + 1$ . This is an example with a delay of length  $d$  causes no degradation and a delay of  $d + 1$  causes a constant performance degradation.
5 RL with missing observations and regret analysis
We now switch our study to MDPs with missing observations. In such an environment, executable policies share the same structures as delayed MDPs, where an action is taken based on available history information. Compared to delayed observations, learning with missing observations is more challenging. Since unobserved states and rewards are never revealed, we are suffering from information loss. Besides, we will frequently deal with multi-step transitions, due to missing observations between two consecutive visible states.
5.1 Optimistic planning with missing observations
Despite the difficulty, we present here algorithms that are efficient in learning and planning for MDPs with missing observations. We begin with an optimistic planning algorithm in Algorithm 3. To unify the notation, we denote  $s_h^k = \emptyset$  and  $r_h^k = \emptyset$  as missing the observation.
Algorithm 3 Optimistic planning for MDPs with missing observations
1: Input: Horizon  $H$ , observable rate  $\lambda_h$
2:  $\mathbf{Init}\colon \mathcal{B}^0 = \Theta$  to be all possible tabular MDPs, data set  $\mathcal{D}^0 = \emptyset$ .
3: for episode  $k = 1, \dots, K$  do
4: Set policy  $\pi^k = \mathrm{argmax}_{\pi \in \Pi_{\mathrm{exec}}}\max_{\theta \in \mathcal{B}^k}V_{1,\theta}^\pi (s_1^k)$
5: Play policy  $\pi^k$  and collect data  $\mathcal{D}^{k - 1}\cup \{(s_h^k,a_h^k,r_h^k)\}_{h = 1}^H$
6: Compute counting number  $N_h^k (s,a) = \sum_{j = 1}^k\mathbb{1}\{s_h^j = s,a_h^j = a,s_{h + 1}^j\neq \emptyset \}$
7: Update confidence set
$$
\mathcal {B} ^ {k} = \left\{\theta : \| \widehat {p _ {h} ^ {k}} (\cdot | s, a) - p _ {h} ^ {\theta} (\cdot | s, a) \| _ {\mathrm {T V}} \leq c \sqrt {\frac {S \iota}{N _ {h} ^ {k} (s , a)}} \text {f o r a l l} (h, s, a) \right\} \cap \mathcal {B} ^ {k - 1},
$$
where  $\widehat{p}_h^k (s'|s,a) = \frac{N_h^k(s,a,s')}{N_h^k(s,a)}$  and  $c$  is some constant.
The majority of the algorithm resembles the typical optimistic planning Jaksch et al. [2010] but with some notable differences. In Line 4, the value function  $V_{1,\theta}$  is for the original MDP with transition probabilities parameterized by  $\theta$ . Different from the typical optimistic planning, the underlying MDP here obeys the stochastic observable model in Assumption 2.2. Therefore, the value  $V_{1,\theta}$  is the sum of all possible values under missing observations. When counting  $N_h^k (s,a)$  in Line 6, we exclude data tuples missing the next state, which inevitably slows down the learning curve. Nonetheless, the effect of missing only contributes as a scaling factor in the regret.
Proposition 5.1. Suppose Assumption 2.2 holds with  $\lambda_h$  known. Given a failure probability  $\gamma$ , with probability  $1 - \gamma$ , the regret of Algorithm 4 satisfies
$$
\operatorname {R e g r e t} (K) \leq c \left(\left\lceil \frac {1}{- \log \left(1 - \lambda_ {0} ^ {2}\right)} \right\rceil \sqrt {H ^ {3} S ^ {2} A K \iota^ {3}} + \sqrt {H ^ {4} K \iota}\right),
$$
The proof is provided in Appendix C.1. Proposition 5.1 is optimal in the  $K$  dependence and achieves an  $S^2 A$  dependence on the complexity of the underlying MDP. In the extreme case of  $\lambda_0\approx 0$ , which implies that every state and reward are hardly observable, we have Regret  $(K) = \tilde{\mathcal{O}}\left(\frac{1}{\lambda_0^2}\sqrt{H^3S^2AK}\right)$ . Here  $\lambda_0^2$  is the probability of observing two consecutive states for estimating the transition probabilities. Proposition 5.1 requires knowledge of observable rate  $\lambda_{h}$ . This is not a restrictive condition, as estimating  $\lambda_{h}$  from Bernoulli random variables is much easier than estimating transition probabilities.
5.2 Model-based planning using augmented MDPs
Proposition 5.1 has a lenient dependence on the missing rate  $1 - \lambda_0^2$ , nonetheless, is not sharp on the dependence of  $S$ . We next show that the augmented MDP approach is effective to tackle missing observations, when the observable rate satisfies additional conditions. Specifically, we assume that the observable rate  $\lambda_h$  is independent of  $(s,a)$ . We utilize the  $\mathrm{MDP}_{\mathrm{aug}}$  reformulation, except that we redefine the transition probabilities as
$$
p _ {h, \mathrm {a u g}} (\tau_ {h + 1} | \tau_ {h}, a _ {h}) = \left\{ \begin{array}{l l} \lambda_ {h} p _ {h} (s _ {h + 1} | s _ {t _ {h}}, \mathbf {a} _ {t _ {h}: h}) & \text {i f} t _ {h + 1} = h + 1 \\ \mathbb {M} _ {a} (\tau_ {h + 1}, \tau_ {h}) (1 - \lambda_ {h}) & \text {i f} t _ {h + 1} = t _ {h} \\ 0 & \text {o t h e r w i s e} \end{array} \right..
$$
The first case in  $p_{h,\mathrm{aug}}$  corresponds to receiving the state observation at time  $h + 1$ . In contrast to the delayed MDPs, the transition probabilities here potentially rely on multi-step transitions in the original MDP. The second case of the transition corresponds to missing the observation. We summarize the policy learning procedure in Algorithm 4 in Appendix C.2, which is similar to Algorithm 2, but with a new bonus function. The following theorem shows that Algorithm 4 is asymptotically efficient when the observable rate is relatively high.
Theorem 5.2. Suppose Assumption 2.2 holds with  $\lambda_0 \geq 1 - A^{-(1 + v)}$  for some positive constant  $v$ . Given a failure probability  $\gamma$ , with probability  $1 - \gamma$ , the regret of Algorithm 4 satisfies
$$
\operatorname {R e g r e t} (K) \leq c \left(H ^ {4} \sqrt {S A K \iota^ {3}} + S ^ {2} \sqrt {H ^ {9} K ^ {\frac {1}{(1 + v)}} \iota^ {6}}\right),
$$
The proof is provided in Appendix C.2. Some remarks are in order.
$SA$  rate when  $K$  is large When the number of episodes  $K \geq S^{3(1 + v) / v}$ , the first term  $H^4\sqrt{SAK\iota^3}$  in the regret bound dominates and attains a sharp dependence on  $S$  and  $A$ . However, when the number of episodes are limited, the regret bound has a worse dependence on the state space size  $S$ . We also observe that as the missing rate  $\lambda$  becomes small (equivalently,  $v$  becomes large), the regret is close to  $\widetilde{O}(H^4\sqrt{SAK\iota^3})$ .
Observable rate smaller than  $1 - 1 / A$  Theorem 5.2 holds for an observable rate  $\lambda_0 > 1 - 1 / A$ . The intuition behind is that to fully explore all the actions when a state observation is missing takes  $A$  trials. Therefore, in expectation, we will encounter a missing observation at least every  $A$  episodes as long as  $\lambda_0 > 1 - 1 / A$ . Nonetheless, when  $\lambda_0 \leq 1 - 1 / A$ , the regret bound remains curiously underexplored. We conjecture that  $\lambda_0 = 1 - 1 / A$  is a critical point distinguishes unique strategies for learning and planning in MDPs with missing observations. A detailed analysis goes beyond the scope of the current paper.
Proof sketch The proof of Theorem 5.2 adapts the analysis of model-based UCBVI algorithms Azar et al. [2017]. Let  $m$  denote the maximal length of consecutive missing observations. We denote  $\mathcal{E}_m$  as the event when the maximal length of consecutive missing is less than  $m$ . On event  $\mathcal{E}_m$ , a naive analysis leads to a  $\widetilde{\mathcal{O}}\left(\sqrt{\mathrm{poly}(H)SA^{m+1}K}\right)$  regret, in observation to the size of the expanded state space  $\mathcal{S}_{\mathrm{aug}}$ . However, our analysis circumvents the  $A^m$  dependence by exploiting the occurrence of consecutive missing observations is rare (Lemma C.3). On the complement of event, the regret is bounded by  $KH(1 - \mathbb{P}(\mathcal{E}_m))$ . Summing up the two parts and choosing a proper  $m$  yield our result.
In this paper, we have studied learning and planning in impaired observability MDPs. We focus on MDPs with delayed and missing observations. Specifically, for delayed observations, we have shown an efficient  $\widetilde{O}(H^4\sqrt{SAK})$  regret. For missing observations, we have provided an optimistic planning algorithm achieving an  $\widetilde{O}(\sqrt{H^3S^2AK})$  regret. If the missing rate is relatively small, we have shown an efficient  $\widetilde{O}(H^4\sqrt{SAK})$  regret bound. Further, we have characterized the performance drop caused by impaired observability compared to full observability.
The work of H. V. Poor was supported in part by a grant from the C3.ai Digital Transformation Institute. Mengdi Wang acknowledges the support by NSF grants DMS-1953686, IIS-2107304, CMMI-1653435, CPS-2312093, ONR grant 1006977, and C3.AI.
Alekh Agarwal and John C Duchi. Distributed delayed stochastic optimization. Advances in Neural Information Processing Systems, 24, 2011.
Mridul Agarwal and Vaneet Aggarwal. Blind decision making: Reinforcement learning with delayed observations. Pattern Recognition Letters, 150:176-182, 2021.
Shipra Agrawal and Randy Jia. Optimistic posterior sampling for reinforcement learning: worst-case regret bounds. Advances in Neural Information Processing Systems, 30, 2017.
Claire Vernade, Alexandra Carpentier, Tor Lattimore, Giovanni Zappella, Beyza Ermis, and Michael Brueckner. Linear bandits with stochastic delayed feedback. In Proceedings of the International Conference on Machine Learning, pages 9712-9721. PMLR, 2020.
Martin J Wainwright. High-dimensional Statistics: A Non-asymptotic Viewpoint, volume 48. Cambridge University Press, 2019.
Thomas J Walsh, Ali Nouri, Lihong Li, and Michael L Littman. Planning and learning in environments with delayed feedback. In Proceedings of Machine Learning: ECML 2007: 18th European Conference on Machine Learning, Warsaw, Poland, September 17-21, 2007, pages 442-453. Springer, 2007.
Tsachy Weissman, Erik Ordentlich, Gadiel Seroussi, Sergio Verdu, and Marcelo J Weinberger. Inequalities for the 11 deviation of the empirical distribution. Hewlett-Packard Labs, Technical Report, 2003.
CB Winsten. Geometric distributions in the theory of queues. Journal of the Royal Statistical Society: Series B (Methodological), 21(1):1-22, 1959.
Yunchang Yang, Han Zhong, Tianhao Wu, Bin Liu, Liwei Wang, and Simon S Du. A reduction-based framework for sequential decision making with delayed feedback. arXiv preprint arXiv:2302.01477, 2023.
Andrea Zanette and Emma Brunskill. Tighter problem-dependent regret bounds in reinforcement learning without domain knowledge using value function bounds. In Proceedings of the International Conference on Machine Learning, pages 7304-7312. PMLR, 2019.
Wenhao Zhan, Masatoshi Uehara, Wen Sun, and Jason D Lee. PAC reinforcement learning for predictive state representations. arXiv preprint arXiv:2207.05738, 2022.
Zihan Zhang, Yuan Zhou, and Xiangyang Ji. Almost optimal model-free reinforcement learning via reference-advantage decomposition. Advances in Neural Information Processing Systems, 33: 15198-15207, 2020.
Han Zhong, Wei Xiong, Sirui Zheng, Liwei Wang, Zhaoran Wang, Zhuoran Yang, and Tong Zhang. A posterior sampling framework for interactive decision making. arXiv preprint arXiv:2211.01962, 2022.
A Omitted proof in Section 3
A.1 Proof of Proposition 3.1
Proof. Consider an arbitrary fixed inter-arrival pattern  $\Delta_0, \Delta_1, \ldots, \Delta_{H - 1}$ . We show that the expected accumulated rewards under this inter-arrival pattern are identical for  $\mathrm{MDP}_{\mathrm{aug}}$  and  $\widetilde{\mathrm{MDP}}_{\mathrm{aug}}$ . In  $\widetilde{\mathrm{MDP}}_{\mathrm{aug}}$ , we have
$$
\begin{array}{l} \mathbb {E} ^ {\pi} \left[ \sum_ {h = 1} ^ {\widetilde {H}} \widetilde {r} _ {h, \text {a u g}} \left(\tau_ {h}, a _ {h}\right) \mid \tau_ {1}, \Delta_ {0}, \dots , \Delta_ {H - 1} \right] \\ \stackrel {(i)} {=} \mathbb {E} ^ {\pi} \left[ \sum_ {h = 1} ^ {\tilde {H}} \widetilde {r} _ {t _ {h}, \operatorname {a u g}} \left(s _ {t _ {h}}, a _ {t _ {h}}\right) \mathbb {1} \left\{\delta_ {t _ {h}} = 0 \right\} \mathbb {1} \left\{t _ {h} \in \{1, \dots , H \} \right\} \mid \tau_ {1}, \Delta_ {0}, \dots , \Delta_ {H - 1} \right] \\ \stackrel {(i i)} {=} \mathbb {E} ^ {\pi} \left[ \sum_ {h = 1} ^ {H} r (s _ {h}, a _ {h}) \Big | \tau_ {1}, \Delta_ {0}, \ldots , \Delta_ {H - 1} \right] \\ = \mathbb {E} ^ {\pi} \left[ \sum_ {h = 1} ^ {H} r _ {h, \text {a u g}} \left(\tau_ {h}, a _ {h}\right) \mid \tau_ {1}, \Delta_ {0}, \dots , \Delta_ {H - 1} \right], \\ \end{array}
$$
where equality  $(i)$  invokes the definition of  $\widetilde{r}_{h,\mathrm{aug}}$  and equality  $(ii)$  eliminates zero reward terms. Now taking expectation over all possible inter-arrival patterns, we deduce
$$
\mathbb {E} ^ {\pi} \left[ \sum_ {h = 1} ^ {\widetilde {H}} \widetilde {r} _ {\mathrm {a u g}} (\tau_ {h}, a _ {h}) \Big | \tau_ {1} \right] = \mathbb {E} ^ {\pi} \left[ \sum_ {h = 1} ^ {H} r _ {h, \mathrm {a u g}} (s _ {h}, a _ {h}) \Big | \tau_ {1} \right].
$$
B Omitted proofs in Section 4
B.1 Proof of Theorem 4.1
Proof. We adapt the main steps from Azar et al. [2017] for proving the theorem. The proof consists of verifying a valid optimism and developing a regret analysis. We denote  $\widetilde{Q}_{h,\mathrm{aug}}^{*}$  as the optimal  $Q$ -function for  $\widetilde{\mathrm{MDP}}_{\mathrm{aug}}$ . When analyzing the regret, we also denote  $\widetilde{Q}_{h,\mathrm{aug}}^{k}$  as the optimal  $Q$ -function in the  $k$ -th episode.
Valid optimism To begin with, we verify that the choice of the bonus functions leads to a valid optimism in the following lemma.
Lemma B.1. Given any failure probability  $\gamma < 1$ , we set a bonus as
$$
b _ {h} ^ {k} (\tau_ {h}, a _ {h}) = c _ {A} H \left(\sqrt {\frac {H \iota}{N _ {t _ {h}} (s _ {t _ {h}} , a _ {t _ {h}} , \delta_ {t _ {h}})}} + \sqrt {\frac {H \iota}{N _ {t _ {h}} (s _ {t _ {h}} , a _ {t _ {h}})}}\right),
$$
where  $\iota = \log \left(\frac{SAHK}{\gamma}\right)$  and  $c_{A}$  is a constant. Then with probability  $1 - \gamma$ , it holds
$$
\widetilde {Q} _ {h, \mathrm {a u g}} ^ {k} (\tau_ {h}, a _ {h}) \geq \widetilde {Q} _ {h, \mathrm {a u g}} ^ {*} (\tau_ {h}, a _ {h}), \quad \widetilde {V} _ {h, \mathrm {a u g}} ^ {k} (\tau_ {h}) \geq \widetilde {V} _ {h, \mathrm {a u g}} ^ {*} (\tau_ {h}) \quad \text {f o r a n y} \quad (k, h, \tau_ {h}, a _ {h}).
$$
Proof of Lemma B.1. We compute the cardinality of the expanded state space  $\widetilde{S}_{\mathrm{aug}}$  as
$$
\left| \widetilde {S} _ {\mathrm {a u g}} \right| \stackrel {(i)} {=} \sum_ {i = 0} ^ {H} H S A ^ {i} = H S \frac {A ^ {H + 1} - 1}{A - 1} \leq 2 H S A ^ {H}.
$$
For a fixed episode  $k$ , we show by backward induction that the assertion in Lemma B.1 holds. To ease the presentation, we omit all superscripts  $k$ , all subscripts "aug", as well as the tilde notation.