1 Introduction↩︎

As a general framework of reinforcement learning (RL), the sequential decision-making problem aims at maximizing a cumulative reward by making a sequence of decisions based on a history of observations and actions over time. This framework is powerful to include and generalize Markov Decision Processes (MDPs) and Partially Observable Markov Decision Processes (POMDPs), and captures a wide range of real-world applications such as recommender systems [1], [2], business management [3], economic simulation [4], robotics [5], strategic games [6], [7], and medical diagnostic systems [8].

However, tackling POMDPs alone presents significant challenges, not to mention general sequential decision-making problems. Many hardness results have been developed [9]–[13], showing that learning POMDPs is already computationally and statistically intractable in the worst case. The reason is that the non-Markovian property of these problems implies that the sufficient statistics or belief about the current environmental state encompasses all observations and actions from past interactions with the environment. This dramatically increases the computational burden and statistical complexity, since even for a finite observation space and action space, the possibilities of the beliefs of the environment state are exponentially large in terms of the number of observations and actions.

In order to tackle these challenges, recent research has introduced various structural conditions for POMDPs and general sequential decision-making problems, such as reactiveness [14], decodability [15], revealing conditions [16], hindsight observability [17], and low-rank representations with regularization conditions [18]–[20]. These conditions have opened up new possibilities for achieving polynomial sample complexities in general sequential decision-making problems. Among them, predictive state representations (PSRs) [21] has been proved to capture and generalize a rich subclass of sequential decision-making problems such as MDPs and observable POMDPs. Yet, most existing solutions for PSRs involve oracles that might not be computationally tractable. For instance, Optimistic MLE (OMLE) [19] involves a step that maximizes the optimal value function over a confidence set of models. Typically, such a confidence set does not exhibit advantageous structures, resulting in a potentially combinatorial search within the set. Another popular posterior sampling based approach [22], [23] requires to maintain a distribution over the entire set of models, which is highly memory inefficient. In addition, most existing results lack a last-iterate guarantee and produce only mixture policies, which often exhibit a very large variance in practical applications. On the other hand, the upper confidence bound (UCB) based approach has been proved to be computationally efficient and provide last-iterate guarantee in many decision-making problems such as bandits [24] and MDPs [25]. However, due to the non-Markovian property of POMDPs and general sequential decision-making problems, designing an explicit UCB is extremely challenging. To the best of our knowledge, such a design has been seldomly explored in POMDPs and beyond. Thus, we are motivated to address the following important open question:

Q1: Can we design a UCB-type algorithm for learning PSRs that (a) is both computationally tractable and statistically efficient, and (b) enjoys the last-iterate guarantee?

Another important research direction in RL is offline learning [26], where the learning agent has access to a pre-collected dataset and aims to design a favorable policy without any interaction with the environment. While offline MDPs have been extensively studied [27]–[29], there exist very limited studies of offline POMDPs from the theoretical perspective [30], [31]. To our best knowledge, offline learning for a more general model of PSRs has never been explored. Thus, we will further address the following research question:

Q2: Can we design a UCB-type algorithm for offline* learning of PSRs with guaranteed policy performance and sample efficiency?*

Main Contributions: We provide affirmative answers to both aforementioned questions by making the following contributions.

2 Related Work↩︎

Learning MDPs and POMDPs. The MDP is a basic model in RL that assumes Markovian property in the model dynamics, i.e. the distribution of the future states only depends on the current system state. Researchers show that learning tabular MDPs (with finite state and action spaces) is both computationally and statistically efficient in both online setting [32]–[37] and offline setting [27], [38]–[40]. Learning MDPs with function approximations is also well-studied by establishing favorable statistical complexity and computation efficiency [14], [27]–[29], [41]–[50]. Notably, several algorithms designed for learning MDPs with general function approximations can be extended to solve a subclass of POMDPs. In particular, OLIVE [14] and GOLF [49], which are originally designed for MDPs with low Bellman rank and low Bellman-Eluder dimension, respectively, can efficiently learn reactive POMDPs, where the optimal policy only depends on the current observation. Besides, [15], [51] study decodable RL where the observations determine the underlying states, and [52] investigate latent MDPs where there are multiple MDPs determined by some latent variables.

To directly address the partial observability in POMDPs, some works assume exploratory data or reachability property and provide polynomial sample complexity for learning these POMDPs [53]–[55]. Others tackle the challenge of exploration and exploitation tradeoff in POMDPs by considering various sub-classes of POMDPs such as low-rank POMDPs [56], observable POMDPs [57], [58], hindsight observability [17], and weakly-revealing POMDPs [16], [59]. Furthermore, [59], [60] propose computationally efficient algorithms for POMDPs with deterministic latent transitions. Notably, [57] propose a provably efficient algorithm for learning observable tabular POMDPs without computationally intractable oracles. Finally, [31] study the offline POMDPs in the presence of confounders, and [30] provide provably efficient algorithm for offline linear POMDPs. After the initial submission of this work, we notice that a UCB-type algorithm has been studied by [61] under low-rank \(L\)-step decodable POMDPs, which is a subclass of the PSRs considered here.

Learning PSRs and general sequential decision-making problems. The PSR is first introduced by [21], [62] and considered as a general representation to model dynamic systems. A line of research [63]–[66] obtain polynomial sample complexity with observability assumption and spectral techniques. Later, [18] demonstrate that learning regular PSRs is sample efficient and can avoid poly\((|\mathcal{O}|^m)\) in the sample complexity. [67] propose a PO-bilinear class that captures a rich class of tractable RL problems with partial observations, including weakly revealing POMDPs and PSRs and design an actor-critic style algorithm. For the works most closely related to ours, [19] propose a universal algorithm known as OMLE, which is capable of learning PSRs and its generalizations under certain conditions; [20] enhance the sample complexity upper bounds for three distinct algorithms, including OMLE, the model-based posterior sampling, and the estimation-to-decision type algorithm [48]; [23] address the general sequential decision-making problem by posterior sampling under a newly proposed low generalized Eluder coefficient. However, as elaborated in 1, those approaches are not efficient in terms of computational complexity or memory.

3 Preliminaries↩︎

Problem Setting. We consider a finite horizon episodic sequential decision-making problem, defined by a tuple \(\mathtt{P} = (\mathcal{O}, \mathcal{A}, H, \mathbb{P}, R)\), where \(\mathcal{O}\) represents the observation space, \(\mathcal{A}\) is a finite action space, \(H\) is the number of time steps within an episode, \(\mathbb{P}= \{\mathbb{P}_h\}\) determines the model dynamics, i.e., \(\mathbb{P}_h(o_h|o_1,\ldots,o_{h-1},a_1,\ldots,a_{h-1})\), where \(o_t\in\mathcal{O}\) is the observation at time step \(t\), and \(a_t\in\mathcal{A}\) is the action taken by the agent at time step \(t\) for all \(t\in\{1,\ldots,h\}\), and \(R: (\mathcal{O}\times\mathcal{A})^H \rightarrow [0,1]\) is the reward function defined on trajectories of one episode. We denote a historical trajectory at time step \(h\) as \(\tau_h:=(o_1,a_1,\ldots,o_h,a_h)\), and denote a future trajectory as \(\omega_h:=(o_{h+1},a_{h+1},\ldots, o_H,a_H)\). The set of all \(\tau_h\) is denoted by \(\mathcal{H}_h = (\mathcal{O}\times\mathcal{A})^{h}\) and the set of all future trajectories is denoted by \(\Omega_h = (\mathcal{O}\times\mathcal{A})^{H-h}\). In addition, let \(\omega_h^o = (o_{h+1},\ldots,o_H)\) and \(\omega_h^a = (a_{h+1},\ldots,a_H)\) be the observation sequence and the action sequence contained in \(\omega_h\), respectively. Similarly, for a history \(\tau_h\), we denote \(\tau_h^o\) and \(\tau_h^a\) as the observation and action sequences in \(\tau_h\), respectively. Notably, the general framework of sequential decision-making problem subsume not only fully observable MDPs but also POMDPs as special cases, because in MDPs, \(\mathbb{P}_h(o_h|\tau_{h-1}) = \mathbb{P}_h(o_h|o_{h-1},a_{h-1})\) and in POMDPs, \(\mathbb{P}_h(o_h|\tau_{h-1})\) can be factorized as \(\mathbb{P}_h(o_h|\tau_{h-1}) = \sum_s\mathbb{P}_h(o_h|s)\mathbb{P}_h(s|\tau_{h-1})\), where \(s\) represents unobserved states. The interaction between an agent and \(\mathtt{P}\) proceeds as follows. At the beginning of each episode, the environment initializes a fixed observation \(o_1\) at time step 1. After observing \(o_1\), the agent takes action \(a_1\), and the environment transits to \(o_2\), which is sampled according to the distribution \(\mathbb{P}_1(o_2|o_1,a_1)\). Then, at any time step \(h\geq 2\), due to the non-Markovian nature of the problem, the agent takes action \(a_h\) based on all past information \((\tau_{h-1},o_h)\), and the environment transits to \(o_{h+1}\), sampled from \(\mathbb{P}_h(o_{h+1}|\tau_{h})\). The interaction terminates after time step \(H\).

The policy \(\pi=\{\pi_h\}\) of the agent is a collection of \(H\) distributions where \(\pi_h(a_h|\tau_{h-1},o_h)\) is the probability of choosing action \(a_h\) at time step \(h\) given the history \(\tau_{h-1}\) and the current observation \(o_h\). For simplicity, we use \(\pi(\tau_h) = \pi(a_h|o_h,\tau_{h-1})\cdots\pi(a_1|o_1)\) to denote the probability of the sequence of actions \(\tau_h^a\) given the observations \(\tau_h^o\). We denote \(\mathbb{P}^{\pi}\) as the distribution of trajectories induced by policy \(\pi\) under dynamics \(\mathbb{P}\). The value of a policy \(\pi\) under \(\mathbb{P}\) and the reward \(R\) is denoted by \(V_{\mathbb{P},R}^{\pi} = \mathbb{E}_{\tau_H\sim\mathbb{P}^{\pi}}[R(\tau_H)]\).

The goal of the agent is to find an \(\epsilon\)-optimal policy \(\hat{\pi}\) that satisfies \(\max_{\pi}V_{\mathbb{P},R}^{\pi} - V_{\mathbb{P},R}^{\hat{\pi}} \leq \epsilon\). Since finding a near-optimal policy for a general decision-making problem incurs exponentially large sample complexity in the worst case, in this paper we follow the line of research in [18], [20], [23] and focus on the low-rank class of problems. To define a low-rank problem, we introduce the dynamic matrix \(\mathbb{D}_h \in \mathbb{R}^{|\mathcal{H}_h|\times |\Omega_h|}\) for each \(h\), where the entry at the \(\tau_h\)-th row and \(\omega_h\)-th column of \(\mathbb{D}_h\) is \(\mathbb{P}(\omega_h^o, \tau_h^o | \tau_h^a, \omega_h^a)\).

Predictive State Representation (PSR). To exploit the low-rank structure, we assume that for each \(h\), there exists a set of future trajectories, namely, core tests (known to the agent) \(\mathcal{Q}_h = \{\mathbf{q}_{h}^{1},\ldots, \mathbf{q}_h^{d_h}\} \subset \Omega_h\), such that the submatrix restricted to these tests \(\mathbb{D}_h[\mathcal{Q}_h]\) has rank \(r\), where \(d_h\geq r\) is a positive integer. This special set \(\mathcal{Q}_h\) allows the system dynamics to be factorized as \(\mathbb{P}(\omega_h^o,\tau_h^o | \tau_h^a,\omega_h^a) = \mathbf{m}(\omega_h)^{\top}\psi(\tau_h)\), where \(\mathbf{m}(\omega_h), \psi(\tau_h)\in\mathbb{R}^{d_h}\) and the \(\ell\)-th coordinate of \(\psi(\tau_h)\) is the joint probability of \(\tau_h\) and the \(\ell\)-th core test \(\mathbf{q}_h^{\ell}\). Mathematically, if we use \(\mathbf{o}_h^{\ell}\) and \(\mathbf{a}_h^{\ell}\) to denote the observation sequence and the action sequence of \(\mathbf{q}_h^{\ell}\), respectively, then \(\mathbb{P}(\mathbf{o}_h^{\ell}, \tau_h^o | \tau_h^a, \mathbf{a}_h^{\ell}) = [\psi(\tau_h)]_{\ell}\). By Theorem C.1 in [19], any low-rank decision-making problem admits a (self-consistent) predictive state representation \(\theta = \{\phi_h,\mathbf{M}_h\}_{h=1}^H\) given core tests \(\{\mathcal{Q}_h\}_{h=0}^{H-1}\), such that for any \(\tau_h\in\mathcal{H}_h, \omega_h\in\Omega_h\), \[\begin{align} &\psi(\tau_h) = \mathbf{M}_h(o_h,a_h)\cdots\mathbf{M}_1(o_1,a_1)\psi_0,\quad \mathbf{m}(\omega_h)^{\top} = \phi_H^{\top}\mathbf{M}_H(o_H,a_H)\cdots\mathbf{M}_{h+1}(o_{h+1},a_{h+1})\\ & \sum_{o_{h+1}}\phi_{h+1}^{\top}\mathbf{M}_{h+1}(o_{h+1},a_{h+1}) = \phi_h^{\top},\quad \mathbb{P}(o_h,\ldots,o_1 | a_1,\dots, a_h) = \phi_h^{\top} \psi(\tau_h), \end{align}\] where \(\mathbf{M}_h: \mathcal{O}\times\mathcal{A}\rightarrow\mathbb{R}^{d_h\times d_{h-1}}\), \(\phi_h\in\mathbb{R}^{d_h}\), and \(\psi_0\in\mathbb{R}^{d_0}\). For ease of presentation, we assume \(\psi_0\) is known to the agent³. Notably, the normalized version of \(\psi(\tau_h)\) with respect to \(\phi_h^{\top}\psi(\tau_h)\), denoted as \(\bar{\psi}(\tau_h) = \psi(\tau_h)/\phi_h^{\top}\psi(\tau_h)\), is known as the prediction vector [21] or prediction feature of \(\tau_h\), since \([\bar{\psi}(\tau_h)]_{\ell} = \mathbb{P}(\mathbf{o}_h^{\ell} | \tau_h, \mathbf{a}_h^{\ell})\). As illustrated in [sec:online PSR,sec:offline PSR], the prediction feature plays an important role in our algorithm design.

In the following context, we use \(\mathbb{P}_{\theta}\) to indicate the model determined by the representation \(\theta = \{\phi_h,\mathbf{M}_h\}_{h=1}^H\). For simplicity, we denote \(V_{\mathbb{P}_{\theta}, R}^{\pi} = V_{\theta,R}^{\pi}\). Moreover, let \(\mathcal{Q}_h^A = \{\mathbf{a}_h^{\ell}\}_{\ell=1}^{d_h}\) be the set of action sequences that are part of core tests, constructed by eliminating any repeated action sequence. \(\mathcal{Q}_h^A\), known as the set of core action sequences, plays a crucial role during the exploration process. Selecting from these sequences is sufficient to sample core tests, leading to accurate estimate of \(\theta\).

We further assume that the PSRs studied in this paper are well-conditioned, as specified in 1. Such an assumption and its variants are commonly adopted in the study of PSRs [19], [20], [23].

1 requires that the error of estimating \(\theta\) does not significantly blow up when the estimation error \(x\) of estimating the probability of core tests is small.

Notations. We denote the complete set of model parameters as \(\Theta\) and the true model parameter as \(\theta^*\). For a vector \(x\), \(\|x\|_A\) stands for \(\sqrt{x^{\top}Ax}\), and the \(i\)-th coordinate of \(x\) is represented as \([x]_i\). For functions \(\mathbb{P}\) and \(\mathbb{Q}\) (not necessarily probability measures) over a set \(\mathcal{X}\), the total variation distance between them is \(\mathtt{D}_{\mathtt{TV}}\left(\mathbb{P}(x),\mathbb{Q}(x)\right) = \sum_x|\mathbb{P}(x)-\mathbb{Q}(x)|\), while the hellinger-squared distance is defined as \(\mathtt{D}_{\mathtt{H}}^2 (\mathbb{P}(x),\mathbb{Q}(x)) = \frac{1}{2}\sum_x (\sqrt{\mathbb{P}(x)} - \sqrt{\mathbb{Q}(x)} )^2\). Note that our definition of the total variation distance is slightly different from convention by a constant factor. We define \(d = \max_h d_h\) and \(Q_A = \max_h|\mathcal{Q}_h^A|\). We use \(\nu_h(\pi,\pi')\) to denote the policy that takes \(\pi\) at the initial \(h-1\) steps and switches to \(\pi'\) from the \(h\)-th step. Lastly, \(\mathtt{u}_{\mathcal{X}}\) represents the uniform distribution over the set \(\mathcal{X}\).

4 Online Learning for Predivtive State Representations↩︎

In this section, we propose a model-based algorithm PSR-UCB, which features three main novel designs: (a) a stable model estimation step controlling the quality of the estimated model such that its prediction features of the empirical data are useful in the design of UCB, (b) an upper confidence bound that captures the uncertainty of the estimated model, and (c) a termination condition that guarantees the last-iterate model is near-accurate and the corresponding greedy policy is near-optimal.

4.1 Algorithm↩︎

The pseudo-code for the PSR-UCB algorithm is presented in 1. We highlight the key idea of PSR-UCB: in contrast to the design of UCB-type algorithms in MDPs that leverage low-rank structures, we exploit the physical meaning of the prediction feature \(\bar{\psi}(\tau_h)\) to design bonus terms that enable efficient exploration. In particular, “physical meaning” refers to that each coordinate of a prediction feature of \(\tau_h\) represents the probability of visiting a core test conditioned on \(\tau_h\) and taking the corresponding core action sequence. Therefore, it suffices to explore a set of \(\tau_h\) whose prediction features can span the entire feature space, and use core action sequences to learn those “base” features. We next elaborate the main steps of the algorithm in greater detail as follows.

Exploration. At each iteration \(k\), PSR-UCB constructs a greedy policy \(\pi^{k-1}\) based on a previous dataset \(\mathcal{D}^{k-1} = \{\mathcal{D}_{h}^{k-1}\}_{h=0}^{H-1}\), together with an estimated model \(\hat{\theta}^{k-1}\). Intuitively, to enable efficient exploration, \(\pi^{k-1}\) is expected to sample \(\tau_h\) that “differs” the most from previous collected samples \(\tau_h\in\mathcal{D}_h^{k-1}\). How to quantify such differences forms the foundation of our algorithm design and will be elaborated later.

Then, for each \(h\in[H]\), PSR-UCB first uses \(\pi^{k-1}\) to get a sample \(\tau_{h-1}^{k,h}\), then follows the policy \(\mathtt{u}_{\mathcal{Q}_{h-1}^{\exp}}\) to get \(\omega_{h-1}^{k,h}\), where \(\mathcal{Q}_{h-1}^{\exp} = (\mathcal{A}\times \mathcal{Q}_h^A) \cup \mathcal{Q}_{h-1}^A\). Here, superscripts \(k,h\) represent the index of episodes. In other words, PSR-UCB adopts the policy \(\nu_h(\pi^{k-1}, \mathtt{u}_{\mathcal{Q}_{h-1}^{\exp}})\) to collect a sample trajectory \((\tau_{h-1}^{k,h}, \omega_{h-1}^{k,h})\), which, together with \(\mathcal{D}_{h-1}^{k-1}\), forms the dataset \(\mathcal{D}_{h-1}^k\).

The importance of uniformly selecting actions from \(\mathcal{Q}_{h-1}^{\exp}\) can be explained as follows. First, the actions in \(\mathcal{Q}_{h-1}^A\) assist us to learn the prediction feature \(\bar{\psi}^*(\tau_{h-1}^{k,h})\) as its \(\ell\)-th coordiate equals to \(\mathbb{P}_{\theta^*}(\mathbf{o}_{h-1}^{\ell} | \tau_{h-1}^{k,h}, \mathbf{a}_{h-1}^{\ell})\). Second, the action sequence \((a_h,\mathbf{a}_h^{\ell}) \in \mathcal{A}\times \mathcal{Q}_h^A\) helps us to estimate \(\mathbf{M}_h^*(o_h,a_h)\bar{\psi}^*(\tau_{h-1}^{k,h})\) because its \(\ell\)-th coordinate represents \(\mathbb{P}_{\theta^*}( \mathbf{o}_h^{\ell},o_h |\tau_{h-1}^{k,h}, a_h, \mathbf{a}_h^{\ell} )\). Therefore, by uniform exploration on \(\mathcal{Q}_{h-1}^{\exp}\) given that \(\tau_{h-1}^{k,h}\) differs from previous dataset \(\mathcal{D}_{h-1}^{k-1}\), PSR-UCB collects the most informative samples for estimating the true model \(\theta^*\).

Stable model estimation. With the updated dataset \(\mathcal{D}^k = \{\mathcal{D}_h^{k}\}\), we estimate the model by maximizing the log-likelihood functions with constraints. Specifically, PSR-UCB extracts any model \(\hat{\theta}^k\) from \(\mathcal{B}^k\) defined as: \[\begin{align} &\Theta_{\min}^k = \left\{\theta: \forall h, (\tau_h,\pi)\in\mathcal{D}_h^k,~~ \mathbb{P}_{\theta}^{\pi}(\tau_h) \geq p_{\min} \right\},\nonumber\\ &\mathcal{B}^k = \left\{\theta \in \Theta_{\min}^k: \sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\mathbb{P}_{\theta}^{\pi}(\tau_H) \geq \max_{\theta'\in\Theta_{\min}^k}\sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\mathbb{P}_{\theta'}^{\pi}(\tau_H) - \beta \right\}, \label{eqn:32confidence32set} \end{align}\tag{1}\] where the estimation margin \(\beta\) is a pre-defined constant. Here the threshold probability \(p_{\min}\) is sufficiently small to guarantee that, with high probability, \(\theta^*\in\Theta_{\min}^k\). Note that compared to existing vanilla maximum likelihood estimators [16], [19], [20], PSR-UCB has an additional constraint \(\Theta_{\min}^k\). This is a crucial condition, as \(\tau_h^{k,h+1}\) is sampled to infer the conditional probability \(\mathbb{P}_{\theta^*}(\omega_h^o|\tau_h^{k,h+1},\omega_h^a)\) or the prediction feature \(\bar{\psi}^*(\tau_h^{k,h+1})\), and learning this feature is useless or even harmful if \(\mathbb{P}_{\hat{\theta}^k}^{\pi^{k-1}}(\tau_h^{k,h+1})\) is too small, as it is proved that \(\theta^*\in\Theta_{\min}^k\) (7).

Design of UCB with prediction features. From the discussion of the previous two steps, we see that the prediction features are vital since (a) actions in \(\mathcal{Q}_h^{\exp}\) can efficiently explore the coordinates of these features, and (b) constraint \(\Theta_{\min}^k\) ensures the significance of the learned features of the collected samples. Therefore, in the next round of exploration, our objective is to sample \(\tau_h\) whose prediction feature exhibits the greatest dissimilarity compared with those of the previously collected samples \(\tau_h\in\mathcal{D}_h^{k}\). Towards that, PSR-UCB constructs an upper confidence bound \(V_{\hat{\theta}^k,\hat{b}^k}^{\pi}\) for \(\mathtt{D}_{\mathtt{TV}} (\mathbb{P}_{\hat{\theta}^k}^{\pi} (\tau_H) , \mathbb{P}_{\theta }^{\pi}(\tau_H) )\), where the bonus \(\hat{b}^k\) is defined as: \[\begin{align} \label{eqn:32bonus} &\hat{U}_{h}^k = \lambda I + \sum_{\tau_h'\in\mathcal{D}_h^{k}} \bar{\hat{\psi}}^k ({\tau_h'}) \bar{\hat{\psi}}^k ({\tau_h'})^{\top},\quad \hat{b}^k(\tau_H) = \min\left\{ \alpha \sqrt{ \sum_{h=0}^{H-1} \left\| \bar{\hat{\psi}}^k(\tau_h) \right\|^2_{(\hat{U}_{h}^k)^{-1}}}, 1 \right\}, \end{align}\tag{2}\] with pre-specified regularizer \(\lambda\) and UCB coefficient \(\alpha\). Note that large \(\hat{b}^k(\tau_H)\) indicates that \(\tau_H\) is “perceived” to be under explored since the estimated prediction feature \(\bar{\hat{\psi}}^k(\tau_h)\) is significantly different from \(\{\bar{\hat{\psi}}^k(\tau_h)\}_{ \tau_h\in\mathcal{D}_h^k}\).

Design of the greedy policy and last-iterate guaranteed termination. The construction of UCB implies that \(\hat{\theta}^k\) is highly uncertain on the trajectories \(\tau_H\) with large bonuses. Thus, PSR-UCB finds a greedy policy \(\pi^k = \arg\max_{\pi} V_{\hat{\theta}^k,\hat{b}^k}^{\pi}\) and terminates if \(V_{\hat{\theta}^k,\hat{b}^k}^{\pi}\) is sufficiently small, indicating that the estimated model \(\hat{\theta}^k\) is sufficiently accurate on any trajectory. Otherwise, \(\pi^k\) serves for the next iteration \(k+1\), as it tries to sample the most dissimilar \(\bar{\hat{\psi}}^k(\tau_h)\) compared with the previous samples \(\bar{\hat{\psi}}^k(\tau_h^{k,h+1})\) for efficient exploration. We remark that the termination condition favors the last-iterate guarantee, where a single model and a greedy policy are identified. Compared with algorithms that output a mixture policy (e.g. uniform selection in a large policy set), this guarantee may present lower variance in practical applications.

4.2 Theoretical Results↩︎

In this section, we present the theoretical results for PSR-UCB. To make a general statement, we first introduce the notion of the optimistic net of the parameter space.

Note that if \(\Theta\) is the parameter space for tabular PSRs (including finite observation and action spaces) with rank \(r\), we have \(|\bar{\Theta}_{\varepsilon}|\leq r^2|\mathcal{O}||\mathcal{A}|H^2\log\frac{H|\mathcal{O}||\mathcal{A}|}{\varepsilon}\) (see 3 or Theorem 4.7 in [19]). Now, we are ready to present the main theorem for PSR-UCB.

The following remarks highlight a few insights conveyed by 1. First, 1 explicitly states that PSR-UCB features last-iterate guarantee, i.e., the guaranteed performance is on the last output of the algorithm. This is in contrast to the previous studies [16], [19], [20] on POMDP and/or PSRs, where the performance guarantee is on a mixture of policies obtained over the entire execution of algorithms. Such policies often have a large variance. Second, in the regime with a low PSR rank such that \(r<\frac{ Q_A^2H}{\gamma^2}\), our result matches the best known sample complexity [20] in terms of rank \(r\) and \(\epsilon\). Third, thanks to the explicitly constructed bonus function \(\hat{b}\), whose computation complexity is polynomial in terms of the iteration number and the dimension \(d\), PSR-UCB is computationally tractable, provided that there exist oracles for planning in POMDPs (e.g. Line 9 in PSR-UCB) [11], [61] and maximum likelihood estimation.

Proof Sketch of 1. The proof relies on three main steps corresponding to three new technical developments, respectively. (a) A new MLE guarantee. Due to the novel stable model estimation design, we establish a new MLE guarantee (which has not been developed in the previous studies) that the total variation distance between the conditional distributions of future trajectories conditioned on \((\tau_h,\pi)\in\mathcal{D}_h^k\) is small. Mathematically, \(\sum_{ (\tau_h,\pi)\in\mathcal{D}_h^k }\mathtt{D}_{\mathtt{TV}} (\mathbb{P}_{\hat{\theta}^k}^{\pi} (\omega_h|\tau_h) , \mathbb{P}_{\theta^* }^{\pi}( \omega_h | \tau_h) )\) is upper bounded by a constant. (b) A new confidence bound. Based on (a) and the observation that the estimation error of the prediction feature \(\bar{\hat{\psi}}^k(\tau_h)\) is upper bounded by the total variation distance \(\mathtt{D}_{\mathtt{TV}} (\mathbb{P}_{\hat{\theta}^k}^{\pi} (\omega_h|\tau_h) , \mathbb{P}_{\theta^* }^{\pi}( \omega_h | \tau_h) )\), the estimation error of other prediction features is captured by the function \(\|\bar{\hat{\psi}}^k(\tau_h)\|_{(\hat{U}_h^k)^{-1}}\) up to some constant, leading to the valid UCB design. (c) A new relationship between the empirical bonus and the ground-truth bonus. To characterize the sample complexity, we need to show that \(V_{\hat{\theta}^k,\hat{b}^k}^{\pi^k}\) can be small for some \(k\). One approach is to prove that \(\sum_{k}V_{\hat{\theta}^k,\hat{b}^k}^{\pi^k}\) is sub-linear. We validate this sub-linearity by establishing \(\hat{b}^k \leq O\left(\sum_h \|\bar{\psi}^*(\tau_h)\|_{(U_h^k)^{-1}}\right)\), where \(U_h^k = \lambda I + \sum_{\tau_h\in\mathcal{D}_h^k} \bar{\psi}^*(\tau_h)\bar{\psi}^*(\tau_h)^{\top}\). This inequality bridges the empirical and the ground-truth bonuses, and yields the final result when combined with the elliptical potential lemma [69]. The complete proof can be found in 4.

When PSR is specialized to \(m\)-step decodable POMDPs [19], as elaborated in 9.5, the sample complexity of PSR-UCB does not depend on \(d\) with slight modification of the algorithm.

5 Offline Learning for Predictive State Representations↩︎

In this section, we develop a computationally tractable algorithm PSR-LCB to learn PSRs in the offline setting. In offline PSRs, a dataset \(\mathcal{D}\) contains \(K\) pre-collected trajectories that are independently sampled from a behavior policy \(\pi^b\). We slightly generalize the learning goal to be finding a policy that can compete with any target policy whose coverage coefficient is finite. We highlight that PSR-LCB is the first offline algorithm for learning PSRs. When specialized to POMDPs, our proposed algorithm enjoys better computational complexity than confidence region based algorithms [30].

5.1 Algorithm↩︎

We design a PSR-LCB algorithm for offline PSRs, whose pseudo-code is presented in 2. In the offline setting, we also leverage the prediction feature \(\bar{\psi}(\tau_h)\) to design lower confidence bound (LCB) of the true value function \(V_{\theta^*,r}^{\pi}\). We explain the main steps of the algorithm in detail as follows.

stable model estimation. Inspired by PSR-UCB, where the prediction features \(\bar{\psi}(\tau_h)\) at each step \(h\) are learned through separate datasets, PSR-LCB first randomly and evenly divides \(\mathcal{D}\) into \(H\) datasets \(\mathcal{D}_0,\ldots, \mathcal{D}_{H-1}\). The goal of this division is to separately learn \(\bar{\psi}^*(\tau_h)\) for each \(h\in\{0,1,\ldots,H-1\}\). Then, PSR-LCB extracts a model \(\hat{\theta}\) from \(\mathcal{B}_{\mathcal{D}}\) defined as: \[\begin{align} & \Theta_{\min} = \left\{\theta: \forall h, \tau_h \in \mathcal{D}_h,~~ \mathbb{P}_{\theta}^{\pi^b}(\tau_h) \geq p_{\min} \right\}, \nonumber \\ &\mathcal{B}_{\mathcal{D}} = \left\{\theta \in \Theta_{\min}: \sum_{\tau_H \in\mathcal{D}} \log\mathbb{P}_{\theta}^{\pi^b}(\tau_H) \geq \max_{\theta'\in\Theta_{\min}}\sum_{ \tau_H \in\mathcal{D}} \log\mathbb{P}_{\theta'}^{\pi^b} (\tau_H) - \hat{\beta} \right\}, \label{eqn:offline32confidence32set} \end{align}\tag{3}\] where the estimation margin \(\hat{\beta}\) is a pre-specified constant, and \(p_{\min}\) guarantees that, with high probability, \(\theta^*\in\Theta_{\min}\).

Following similar reasons of the design of PSR-UCB, the constraint \(\Theta_{\min}\) controls the quality of the estimated model such that the prediction features of the behavior samples are non-negligible and useful for learning.

Design of LCB with prediction features. Given the estimated model \(\hat{\theta}\), PSR-LCB constructs an upper confidence bound of \(\mathtt{D}_{\mathtt{TV}} \left(\mathbb{P}_{\hat{\theta} }^{\pi} (\tau_H) , \mathbb{P}_{\theta }^{\pi}(\tau_H) \right)\) in the form of \(V_{\hat{\theta} ,\hat{b} }^{\pi}\), where \(\hat{b}(\tau_H)\) is defined as: \[\begin{align} \label{eqn:32LCB32bonus} &\hat{U}_{h} = \hat{\lambda} I + \sum_{\tau_h\in\mathcal{D}_h } \bar{\hat{\psi}} (\tau_h) \bar{\hat{\psi}} (\tau_h)^{\top},\quad \hat{b} (\tau_H) = \min\left\{ \hat{\alpha} \sqrt{ \sum_h \left\| \bar{\hat{\psi}} (\tau_h) \right\|^2_{(\hat{U}_{h} )^{-1}}}, 1 \right\}, \end{align}\tag{4}\] with pre-defined regularizer \(\hat{\lambda}\) and LCB coefficient \(\hat{\alpha}\). Note that we refer to \(\hat{\alpha}\) as the LCB coefficient, as we adopt the pessimism principle in offline learning. Differently from PSR-UCB, we aim to select a policy where the estimated model exhibits the least uncertainty. Therefore, the output policy should allocate a high probability to \(\tau_H\) if \(\hat{b}(\tau_H)\) is small.

Output policy design. Building on the discussion above, PSR-LCB outputs a policy \(\bar{\pi} = \arg\max_{\pi} V_{ \hat{\theta}, R }^{\pi} - V_{ \hat{\theta}, \hat{b} }^{\pi}\), which maximizes a lower confidence bound of \(V_{\theta^*,R}^{\pi}\).

5.2 Theoretical Results↩︎

In this section, we develop the theoretical guarantee for PSR-LCB. To capture the distribution shift in offline PSRs, we require that the \(\ell_{\infty}\) norm of the ratio between the probabilities over the entire trajectory under the target policy and under the behavior policy is finite. Mathematically, if \(\pi\) is the target policy, then, \(C_{\pi^b,\infty}^{\pi} := \max_h \max_{\tau_h}\frac{ \mathbb{P}_{\theta^*}^{\pi}(\tau_h) }{\mathbb{P}_{\theta^*}^{\pi^b}(\tau_h)} < \infty.\) We note that, for general sequential decision-making problems, no Markovian property or hidden state is assumed. Therefore, defining the coverage assumption on a single observation-action pair as in [29] for offline MDP, or on hidden states as in [30] for offline POMDP, is not suitable for PSRs.

2 states that for any target policy with finite coverage coefficient \(C_{\pi^b,\infty}^{\pi}\), the performance degradation of the output policy \(\bar{\pi}\) with respect to the target policy \(\pi\) is at most \(\tilde{O}\big( \frac{C_{\pi^b,\infty}^{\pi}}{\sqrt{K}} \big)\), which is negligible if the size of the offline dataset \(K\) is sufficiently large. Thee full proof is in  5.

Proof Sketch of 2. Similar to the online setting, the proof first shows that \(V_{\hat{\theta}, \hat{b}}^{\pi}\) serves as a valid upper confidence bound of the total variation distance \(\mathtt{D}_{\mathtt{TV}}\left(\mathbb{P}_{\hat{\theta}}^{\pi}, \mathbb{P}_{\theta^*}^{\pi}\right)\) for any policy \(\pi\). We emphasize that the validity is again due to the development of the estimation guarantee of the distributions of future trajectories conditioned on empirical samples \(\tau_h\in\mathcal{D}_h\). Then, we translate the empirical bonus term \(\hat{b}\) to \(\sum_h\|\bar{\psi}^*(\tau_h)\|_{U_h^{-1}}\), where \(U_h = \lambda I + \sum_{\tau_h\in\mathcal{D}_h}\bar{\psi}^*(\tau_h)\bar{\psi}^*(\tau_h)^{\top}\). Since \(\bar{\psi}^*\) is a fixed representation function, we leverage the coverage assumption to change the distribution induced by the target policy \(\pi\) to the distribution induced by the behavior policy \(\pi^b\), and then apply Azuma-Hoeffding’s inequality to obtain the final result.

6 Conclusion↩︎

We studied learning predictive state representations (PSRs) for low-rank sequential decision-making problems. We developed a novel upper confidence bound for the total variation distance of the estimated model and the true model that enables both computationally tractable with MLE oracles and statistically efficient learning for PSRs with only supervised learning oracles. Specifically, we proposed PSR-UCB for online learning for PSRs with last-iterate guarantee, i.e. producing not only a near-optimal policy, but also a near-accurate model. The statistical efficiency was validated by a polynomial sample complexity in terms of the model parameters. In addition, we extended this UCB-type approach to the offline setting and proposed PSR-LCB. Our theoretical result offers an initial perspective on offline PSRs by demonstrating that PSR-LCB outputs a policy that can compete with any policy with finite coverage coefficient.

Author Contributions↩︎

If you’d like to, you may include a section for author contributions as is done in many journals. This is optional and at the discretion of the authors.

Acknowledgments↩︎

Use unnumbered third level headings for the acknowledgments. All acknowledgments, including those to funding agencies, go at the end of the paper.

7 Properties of PSRs↩︎

In this section, we present a few important properties of PSRs, which will be intensively used in the algorithm analysis.

First, for any model \(\theta = \{\phi_h, \mathbf{M}_h(o_h,a_h)\}\), we have the following identity \[\begin{align} &\mathbf{M}_h(o_h,a_h) \bar{\psi}(\tau_{h-1}) = \mathbb{P}_{\theta}(o_h|\tau_{h-1})\bar{\psi}(\tau_h).\label{eqn:Mpsi32is32conditional32probability} \end{align}\tag{5}\]

The following proposition is directly adapted from Lemma C.3 in [19]. Note that \(\psi_0\) is known to the agent.

The next proposition characterizes well-conditioned PSRs [23], which is obtained by noting that \(\mathbf{m}(\mathbf{q}_h^{\ell})^{\top}\) is the \(\ell\)-th row of \(\mathbf{M}_{h+1}(o_{h+1},a_{h+1})\).

The following proposition characterizes the log-cardinality of the minimal optimistic net of the rank-\(r\) PSRs. We note that this notion is closely related to the bracketing number, a typical complexity measure in the maximum likelihood estimation (MLE) anslysis [70]. The proof follows directly from Theorem C.9 in [19].

\[\begin{align} &\mathbb{P}_{\theta}^{\pi}(\tau_h) = \mathbf{v}_{h}^{\top}\mathbf{B}(o_{h},a_{h})\cdots\mathbf{B}(o_1,a_1)\psi_0 \pi(\tau_h) \end{align}\]

Now let \(\bar{\Theta}\) be a \(\epsilon_0\) cover of the the parameters with respect to the \(\ell_{\infty}\)-norm, since all parameters are bounded.

\[\begin{align} \big[\bar{\mathbf{B}}(o_h,a_h) \big]_{i,j} = \bigg\lceil \frac{\big[ \mathbf{B}(o_h,a_h) \big]_{i,j}}{\epsilon_0} \bigg\rceil \epsilon_0. \end{align}\]

Then, we have, for any \(\theta\), there exists a \(\bar{\theta}\) such that

\[\begin{align} &\forall h,\tau_h, 0\leq \mathbb{P}_{ \theta}^{\pi}(\tau_h) - \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_h) \leq \epsilon_0 (OA)^{H} \end{align}\]

Let \(\epsilon_0 = \frac{\epsilon}{(OA)^{H}}\), we have \[\begin{align} &\forall \pi, h,~~ \sum_{\tau_h}\left| \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_h) - \mathbb{P}_{\theta}^{\pi}(\tau_h)\right|\leq \epsilon \end{align}\]

The covering number is \[\begin{align} |\bar{\Theta}| = \left(\frac{(OA)^H}{\epsilon_0}\right)^{r^2OAH} = \left(\frac{OA}{\epsilon}\right)^{2r^2OAH^2} \end{align}\]

8 General MLE Analysis↩︎

In this section, we present four general propositions that characterize the performance of MLE.

We start with a proposition that states that the log-likelihood of the true model is relatively high compared to any model.

Proof:  We start with the first inequality. Suppose the data in \(\mathcal{D}_h^k\) is indexed by \(t\). Then, \[\begin{align} \mathbb{E}&\left[ \exp\left(\sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k}\log\frac{ \mathbb{P}_{\bar{\theta}}^{\pi}( \tau_h ) }{\mathbb{P}_{\theta^*}^{\pi}(\tau_h )} \right)\right]\\ & = \mathbb{E}\left[ \prod_{t\leq k}\prod_h \frac{ \mathbb{P}_{\bar{\theta}}^{\pi^t}(\tau_h^t) }{ \mathbb{P}_{ \theta^*}^{\pi^t}(\tau_h^t) } \right] \\ & = \mathbb{E}\left[ \prod_{t\leq k-1}\prod_h \frac{ \mathbb{P}_{\bar{\theta}}^{\pi^t}(\tau_h^t) }{ \mathbb{P}_{ \theta^*}^{\pi^t}(\tau_h^t) } \mathbb{E}\left[ \frac{ \mathbb{P}_{\bar{\theta}}^{\pi^{k}}(\tau_h^{k}) }{ \mathbb{P}_{ \theta^*}^{\pi^{k}}(\tau_h^{k}) } \right] \right] \\ & = \mathbb{E}\left[ \prod_{t\leq k-1}\prod_h \frac{ \mathbb{P}_{\bar{\theta}}^{\pi^t}(\tau_h^t) }{ \mathbb{P}_{ \theta^*}^{\pi^t}(\tau_h^t) } \prod_h\sum_{\tau_h} \mathbb{P}_{\bar{\theta}}^{\pi^{k}}(\tau_h ) \right] \\ &\overset{(a)}\leq \left(1 + \varepsilon\right)^H \mathbb{E}\left[ \prod_{t\leq k-1}\prod_h \frac{ \mathbb{P}_{\bar{\theta}}^{\pi^t}(\tau_h^t) }{ \mathbb{P}_{ \theta^*}^{\pi^t}(\tau_h^t) } \right] \\ &\leq \left(1+\varepsilon\right)^{KH} \\ &\overset{(b)}\leq e, \end{align}\] where \((a)\) follows because \(\sum_{\tau_H}|\mathbb{P}_{\bar{\theta}}^{\pi}(\tau_h) - \mathbb{P}_{\theta}^{\pi}(\tau_h) | \leq \varepsilon\), and \((b)\) follows because \(\varepsilon\leq \frac{1}{KH}\).

By the Chernoff bound and the union bound over \(\bar{\Theta}_{\epsilon}\) and \(k\in[K]\), with probability at least \(1-\delta\), we have, \[\begin{align} \forall k\in[K], \bar{\theta}\in\bar{\Theta}_{\epsilon},~~ \sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k}\log\frac{ \mathbb{P}_{\bar{\theta}}^{\pi}( \tau_h ) }{\mathbb{P}_{\theta^*}^{\pi}(\tau_h )} \leq 3\log\frac{K\left|\bar{\Theta}_{\epsilon}\right|}{\delta}, \end{align}\] which yields the first result of this proposition.

To show the second inequality, we follow an argument similar to that for the first inequality. We have \[\begin{align} \mathbb{E}\left[ \exp\left( \sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\frac{ \mathbb{P}_{ \theta }^{\pi}(\tau_H) }{\mathbb{P}_{ \theta^* }^{\pi}(\tau_H)} \right)\right] &\leq \mathbb{E}\left[ \exp\left( \sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log \frac{ \mathbb{P}_{ \bar{\theta} }^{\pi}(\tau_H) }{\mathbb{P}_{ \theta^* }^{\pi}(\tau_H)} \right)\right] \\ &\overset{(a)}\leq (1+\varepsilon)^{KH}\leq e, \end{align}\] where \((a)\) follows from the tower rule of the expectation and because \(\sum_{\tau_H}\mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H)\leq \varepsilon\).

Thus, with probability at least \(1-\delta\), for any \(k\in[K]\) and any \(\bar{\theta}\in\Theta\), the following inequality holds \[\begin{align} \sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\frac{ \mathbb{P}_{ \theta }^{\pi}(\tau_H) }{\mathbb{P}_{ \theta^* }^{\pi}(\tau_H)} \leq 3\log\frac{K|\bar{\Theta}_{\epsilon}|}{\delta}, \end{align}\] which completes the proof.

 \(\blacksquare\)

The following proposition upper bounds the total variation distance between the conditional distributions over the future trajectory conditioned on the empirical history trajectories. This proposition is crucial to ensure that the model estimated by PSR-UCB is accurate on those sample trajectories.

Proof:  We start with a general upper bound on the total variation distance between two conditional distributions. Note that for any \(\theta,\theta'\in\Theta\cup\bar{\Theta}_{\epsilon}\) and fixed \((\tau_h,\pi)\), we have \[\begin{align} \mathtt{D}_{\mathtt{TV}}&\left( \mathbb{P}_{\theta}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\theta'}^{\pi}(\omega_h|\tau_h )\right) \\ & = \sum_{\omega_h} \bigg| \frac{\mathbb{P}_{\theta'}^{ \pi }(\omega_h,\tau_h ) \mathbb{P}_{\theta}^{\pi }(\tau_h ) - \mathbb{P}_{\theta}^{ \pi }(\omega_h,\tau_h )\mathbb{P}_{\theta'}^{\pi }(\tau_h ) }{\mathbb{P}_{\theta}^{\pi }(\tau_h )\mathbb{P}_{\theta'}^{\pi }(\tau_h )} \bigg|\\ & = \sum_{\omega_h} \left| \frac{ \left(\mathbb{P}_{\theta'}^{ \pi }(\omega_h,\tau_h ) - \mathbb{P}_{\theta}^{ \pi }(\omega_h,\tau_h ) \right) \mathbb{P}_{\theta}^{\pi }(\tau_h ) + \mathbb{P}_{\theta}^{ \pi }(\omega_h,\tau_h ) \left(\mathbb{P}_{\theta}^{\pi }(\tau_h ) -\mathbb{P}_{\theta'}^{\pi }(\tau_h ) \right) }{\mathbb{P}_{\theta}^{\pi }(\tau_h ) \mathbb{P}_{\theta'}^{\pi }(\tau_h )} \right|\\ &\leq \frac{|\mathbb{P}_{\theta}^{\pi }(\tau_h ) - \mathbb{P}_{\theta'}^{\pi }(\tau_h )|}{\mathbb{P}_{\theta'}^{\pi }(\tau_h ) } + \frac{1}{\mathbb{P}_{\theta'}^{\pi }(\tau_h )}\sum_{\omega_h} \left|\left(\mathbb{P}_{\theta'}^{ \pi }(\omega_h,\tau_h ) - \mathbb{P}_{\theta}^{ \pi }(\omega_h,\tau_h ) \right)\right|\\ &\leq \frac{2}{\mathbb{P}_{\theta'}^{\pi }(\tau_h )} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta}^{ \pi }(\tau_H), \mathbb{P}_{\theta'}^{ \pi }(\tau_H) \right). \end{align}\]

By symmetry, we also have \[\begin{align} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\theta'}^{\pi}(\omega_h|\tau_h )\right) \leq \frac{2}{ \max\left\{\mathbb{P}_{\theta}^{\pi }(\tau_h ), \mathbb{P}_{\theta'}^{\pi }(\tau_h ) \right\} } \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta}^{ \pi }(\tau_H), \mathbb{P}_{\theta'}^{ \pi }(\tau_H) \right). \end{align}\]

We replace \(\theta'\) by a \(\bar{\theta} \in\bar{\Theta}_{\varepsilon}\) that is \(\varepsilon\)-optimistic over \(\theta\) (recall 2), i.e. \(\mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta}^{\pi} , \mathbb{P}_{\bar{\theta}}^{\pi}\right) \leq \varepsilon,\) and \(\mathbb{P}_{\bar{\theta}}^{\pi}(\tau_h)\geq\mathbb{P}_{\theta}^{\pi}(\tau_h)\) holds for any \(\pi\) and \(\tau_h\). Then, due the construction of \(\Theta_{\min}^k\), we have \[\forall (\tau_h,\pi)\in\mathcal{D}_h^k,~~ \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta}^{\pi}(\omega_h|\tau_h), \mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h) \right) \leq \frac{2\epsilon}{p_{\min}}\leq \frac{2}{KH},\] which implies \[\begin{align} \sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k} & \mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\theta}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )\right) \\ &\overset{(a)}\leq \sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k} 2\mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\theta}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h )\right) + 2\mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )\right)\\ &\leq \frac{4 }{KH} + 2 \sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k} \mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )\right). \end{align}\] Here \((a)\) follows because the total variation distance satisfies the triangle inequality and \((a+b)^2\leq 2a^2 + 2b^2\).

Moreover, note that \[\begin{align} \mathtt{D}_{\mathtt{TV}}^2& \left( \mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )\right)\\ &\overset{(a)}\leq 4(2 + 2/(KH))\mathtt{D}^2_{\mathtt{H}} \left( \mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h ) \right) \\ & \leq 6\left( 1+\frac{1}{KH} - \mathop{\mathbb{E}}_{\omega_h\sim \mathbb{P}_{\theta^*}^{\pi} } \sqrt{\frac{\mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h )}{\mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )}} \right)\\ & \overset{(b)}\leq - 6 \log \mathop{\mathbb{E}}_{\omega_h\sim \mathbb{P}_{\theta^*}^{\pi}(\cdot|\tau_h ) } \sqrt{\frac{\mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h )}{\mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )}} + \frac{6}{KH}, \end{align}\] where \((a)\) is due to 12 and \((b)\) follows because \(1-x\leq -\log x\) for any \(x>0\).

Thus, the summation of the total variation distance between conditional distributions conditioned on \((\tau_h,\pi)\in\mathcal{D}_h^k\) can be upper bounded by \[\begin{align} \sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k} & \mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\theta}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )\right)\\ &\leq \frac{18}{KH} - 12\sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k} \log \mathop{\mathbb{E}}_{\omega_h\sim \mathbb{P}_{\theta^*}^{\pi}(\cdot|\tau_h ) } \sqrt{\frac{\mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h )}{\mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )}}. \end{align}\]

In addition, we have \[\begin{align} &\mathop{\mathbb{E}}_{ \substack{\forall h,(\tau_h,\pi)\in\mathcal{D}_h^k,\\ \omega_h \sim \mathbb{P}_{\theta^*}^{\pi}(\cdot|\tau_h) }} \left[\exp\left( \frac{1}{2}\sum_h\sum_{(\omega_h,\tau_h,\pi)\in\mathcal{D}_h^k} \log\frac{\mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h |\tau_h )}{\mathbb{P}_{\theta^*}^{\pi}(\omega_h |\tau_h )} - \sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k} \log \mathop{\mathbb{E}}_{\omega_h\sim \mathbb{P}_{\theta^*}^{\pi}(\cdot|\tau_h ) } \sqrt{\frac{\mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h )}{\mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )}} \right)\right]\\ & \quad\quad = \frac{ \mathop{\mathbb{E}}_{ \substack{\forall h,(\tau_h,\pi)\in\mathcal{D}_h^k,\\ \omega_h \sim \mathbb{P}_{\theta^*}^{\pi}(\cdot|\tau_h) } } \left[\prod_h\prod_{(\omega_h,\tau_h)\in\mathcal{D}_h^k} \sqrt{\frac{ \mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h |\tau_h ) }{\mathbb{P}_{\theta^*}^{\pi}(\omega_h |\tau_h )} }\right] }{\prod_h\prod_{(\tau_h,\pi)\in\mathcal{D}_h^k} \mathop{\mathbb{E}}_{\omega_h\sim \mathbb{P}_{\theta^*}^{\pi}(\cdot|\tau_h ) } \left[\sqrt{\frac{\mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h )}{\mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )}} \right] } = 1, \end{align}\] where the last equality is due to the conditional independence of \(\omega_h \in \mathcal{D}_h^k\) given \((\tau_h,\pi)\in\mathcal{D}_h^k\).

Therefore, by the Chernoff bound, with probability \(1-\delta\), we have

\[\begin{align} -\sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k} \log \mathop{\mathbb{E}}_{\omega_h\sim \mathbb{P}_{\theta^*}^{\pi}(\cdot|\tau_h ) } \sqrt{\frac{\mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h|\tau_h )}{\mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h )}} \leq \frac{1}{2} \sum_h\sum_{(\omega_h,\tau_h,\pi)\in\mathcal{D}_h^k}\log\frac{\mathbb{P}_{\theta^*}^{\pi}(\omega_h |\tau_h )}{\mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h |\tau_h )} + \log\frac{1}{\delta}. \end{align}\]

Taking the union bound over \(\bar{\Theta}_{\epsilon}\), \(k\in[K]\), and rescaling \(\delta\), we have, with probability at least \(1-\delta\), \(\forall k\in[K]\), the following inequality holds: \[\begin{align} \sum_h &\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k} \mathtt{D}_{\mathtt{TV}}^2\left(\mathbb{P}_{\theta}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h ) \right) \\ &\leq \frac{18 }{KH} + 6\sum_h\sum_{(\omega_h,\tau_h,\pi)\in\mathcal{D}_h^k}\log\frac{\mathbb{P}_{\theta^*}^{\pi}(\omega_h |\tau_h )}{\mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h |\tau_h )} + 12\log\frac{K\left|\bar{\Theta}_{\varepsilon}\right|}{\delta}\\ &\leq 6\sum_h\sum_{(\omega_h, \tau_h,\pi)\in\mathcal{D}_h^k}\log\frac{\mathbb{P}_{\theta^*}^{\pi}( \omega_h ,\tau_h )}{\mathbb{P}_{\bar{\theta}}^{\pi}(\omega_h ,\tau_h )} + 6\sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k}\log\frac{ \mathbb{P}_{\bar{\theta}}^{\pi}( \tau_h ) }{\mathbb{P}_{\theta^*}^{\pi}(\tau_h )} + 13\log\frac{K\left|\bar{\Theta}_{\varepsilon}\right|}{\delta}. \end{align}\]

Note that, following from 4, with probability at least \(1-\delta\), we have for any \(k\in[K]\), \[\begin{align} \sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k}\log\frac{ \mathbb{P}_{\bar{\theta}}^{\pi}( \tau_h ) }{\mathbb{P}_{\theta^*}^{\pi}(\tau_h )} \leq 3\log\frac{K\left|\bar{\Theta}_{\varepsilon}\right|}{\delta}. \end{align}\]

Hence, combining with the optimistic property of \(\bar{\theta}\) and rescaling \(\delta\), we have that the following inequality holds with probability at least \(1-\delta\): \[\begin{align} \sum_h\sum_{(\tau_h,\pi)\in\mathcal{D}_h^k} \mathtt{D}_{\mathtt{TV}}^2\left(\mathbb{P}_{\theta}^{\pi}(\omega_h|\tau_h ), \mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h ) \right) \leq 6\sum_h\sum_{(\tau_H,\pi)\in\mathcal{D}_h^k}\log\frac{\mathbb{P}_{\theta^*}^{ \pi }(\tau_H )}{\mathbb{P}_{ \theta}^{ \pi }(\tau_H )} + 31\log\frac{K\left|\bar{\Theta}_{\varepsilon}\right|}{\delta}, \end{align}\] which yields the final result.

 \(\blacksquare\)

The following proposition is standard in the MLE analysis, and we provide the full analysis here for completeness.

Proof:  First, by the construction of \(\bar{\Theta}_{\varepsilon}\), for any \(\theta\), let \(\bar{\theta}\) be optimistic over \(\theta\), i.e., \(\sum_{\tau_H}\left|\mathbb{P}_{ \theta}^{\pi}(\tau_H) - \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H) \right|\leq\varepsilon\). We translate the distance between \(\theta\) and \(\theta^*\) to the distance between \(\bar{\theta}\) and \(\theta^*\) as follows. \[\begin{align} \mathtt{D}_{\mathtt{H}}^2 & ( \mathbb{P}_{\theta}^{\pi}(\tau_H) , \mathbb{P}_{\theta^*}^{\pi} (\tau_H) ) \\ & = 1 - \sum_{\tau_H} \sqrt{ \mathbb{P}_{\theta}^{\pi}(\tau_H) \mathbb{P}_{\theta^*}^{\pi} (\tau_H) }\\ & = 1 - \sum_{\tau_H} \sqrt{ \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H) \mathbb{P}_{\theta^*}^{\pi} (\tau_H) + \left(\mathbb{P}_{ \theta}^{\pi}(\tau_H) - \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H) \right)\mathbb{P}_{\theta^*}^{\pi} (\tau_H)} \\ &\overset{(a)} \leq 1 - \sum_{\tau_H} \sqrt{ \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H) \mathbb{P}_{\theta^*}^{\pi} (\tau_H) } + \sum_{\tau_H}\sqrt{ \left|\mathbb{P}_{ \theta}^{\pi}(\tau_H) - \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H) \right| \mathbb{P}_{\theta^*}^{\pi} (\tau_H)} \\ &\overset{(b)} \leq - \log \mathop{\mathbb{E}}_{\tau_H\sim\mathbb{P}_{\theta^*}^{\pi}(\cdot)} \sqrt{ \frac{ \mathbb{P}_{ \bar{\theta} }^{\pi}(\tau_H) }{ \mathbb{P}_{\theta^*}^{\pi}(\tau_H) } } + \sqrt{\sum_{\tau_H}\left|\mathbb{P}_{ \theta}^{\pi}(\tau_H) - \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H) \right| } \\ & \leq - \log \mathop{\mathbb{E}}_{\tau_H\sim\mathbb{P}_{\theta^*}^{\pi}(\cdot)} \sqrt{ \frac{ \mathbb{P}_{ \bar{\theta} }^{\pi}(\tau_H) }{ \mathbb{P}_{\theta^*}^{\pi}(\tau_H) } } + \sqrt{\varepsilon}, \end{align}\] where \((a)\) follows because \(\sqrt{a+b} \geq\sqrt{a} - \sqrt{|b|}\) if \(a>0\) and \(a+b>0\), and \((b)\) follows from the Cauchy’s inequality and the fact that \(1-x\leq - \log x\).

Hence, in order to upper bound \(\sum_{\pi\in\mathcal{D}_k} \mathtt{D}_{\mathtt{H}}^2 ( \mathbb{P}_{\theta}^{\pi}(\tau_H) , \mathbb{P}_{\theta^*}^{\pi} (\tau_H) )\), it suffices to upper bound
\(\sum_{\pi\in\mathcal{D}_k} - \log \mathop{\mathbb{E}}_{\tau_H\sim\mathbb{P}_{\theta^*}^{\pi}(\cdot)} \sqrt{ \frac{ \mathbb{P}_{ \bar{\theta} }^{\pi}(\tau_H) }{ \mathbb{P}_{\theta^*}^{\pi}(\tau_H) } }.\) To this end, we observe that, \[\begin{align} \mathbb{E}&\left[\exp\left( \frac{1}{2} \sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\frac{ \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H) }{\mathbb{P}_{\theta^*}^{\pi}(\tau_H)} - \sum_{ \pi \in\mathcal{D}^k}\log \mathop{\mathbb{E}}_{\tau_H\sim\mathbb{P}_{\theta^*}^{\pi}(\cdot)} \sqrt{ \frac{ \mathbb{P}_{\theta}^{\pi}(\tau_H) }{ \mathbb{P}_{\theta^*}^{\pi}(\tau_H) } } \right)\right] \\ &\quad \overset{(a)}= \frac{ \mathbb{E}\left[\prod_{(\tau_H,\pi)\in\mathcal{D}^k} \sqrt{ \frac{ \mathbb{P}_{\theta}^{\pi}(\tau_H) }{ \mathbb{P}_{\theta^*}^{\pi}(\tau_H) } } \right] }{ \mathbb{E}\left[\prod_{(\tau_H,\pi)\in\mathcal{D}^k} \sqrt{ \frac{ \mathbb{P}_{\theta}^{\pi}(\tau_H) }{ \mathbb{P}_{\theta^*}^{\pi}(\tau_H) } } \right] } = 1, \end{align}\] where \((a)\) follows because \((\tau_H,\pi)\in\mathcal{D}^k\) form a filtration.

Then, by the Chernoff bound, we have \[\begin{align} \mathbb{P}&\left( \frac{1}{2} \sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\frac{ \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H) }{\mathbb{P}_{\theta^*}^{\pi}(\tau_H)} - \sum_{ \pi \in\mathcal{D}^k}\log \mathop{\mathbb{E}}_{\tau_H\sim\mathbb{P}_{\theta^*}^{\pi}(\cdot)} \sqrt{ \frac{ \mathbb{P}_{\theta}^{\pi}(\tau_H) }{ \mathbb{P}_{\theta^*}^{\pi}(\tau_H) } } \geq \log\frac{1}{\delta} \right) \\ & = \mathbb{P}\left( \exp\left(\frac{1}{2} \sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\frac{ \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H) }{\mathbb{P}_{\theta^*}^{\pi}(\tau_H)} - \sum_{ \pi \in\mathcal{D}^k}\log \mathop{\mathbb{E}}_{\tau_H\sim\mathbb{P}_{\theta^*}^{\pi}(\cdot)} \sqrt{ \frac{ \mathbb{P}_{\theta}^{\pi}(\tau_H) }{ \mathbb{P}_{\theta^*}^{\pi}(\tau_H) } } \right) \geq \frac{1}{\delta} \right)\\ &\overset{(a)}\leq \delta \mathbb{E}\left[\exp\left( \frac{1}{2} \sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\frac{ \mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H) }{\mathbb{P}_{\theta^*}^{\pi}(\tau_H)} - \sum_{ \pi \in\mathcal{D}^k}\log \mathop{\mathbb{E}}_{\tau_H\sim\mathbb{P}_{\theta^*}^{\pi}(\cdot)} \sqrt{ \frac{ \mathbb{P}_{\theta}^{\pi}(\tau_H) }{ \mathbb{P}_{\theta^*}^{\pi}(\tau_H) } } \right)\right]\\ &\leq \delta, \end{align}\] where \((a)\) is due to the Markov inequality.

Finally, rescaling \(\delta\) to \(\delta/(K|\bar{\Theta}_{\varepsilon}|)\) and taking the union bound over \(\bar{\Theta}_{\epsilon}\) and \(k\in[K]\), we conclude that, with probability at least \(1-\delta\), \(\forall \theta\in\Theta, k\in[K]\), \[\begin{align} ~~\sum_{\pi\in\mathcal{D}^k} & \mathtt{D}_{\mathtt{H}}^2 ( \mathbb{P}_{\theta}^{\pi}(\tau_H) , \mathbb{P}_{\theta^*}^{\pi} (\tau_H) ) \\ &\leq KH\sqrt{\varepsilon} + \frac{1}{2} \sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\frac{ \mathbb{P}_{ \theta^*}^{\pi}(\tau_H) }{\mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H)} + \log\frac{K|\bar{\Theta}_{\epsilon}|}{\delta} \\ &\overset{(a)}\leq \frac{1}{2}\sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\frac{ \mathbb{P}_{ \theta^*}^{\pi}(\tau_H) }{\mathbb{P}_{ \theta }^{\pi}(\tau_H)} + 2\log\frac{K|\bar{\Theta}_{\epsilon}|}{\delta}, \end{align}\] where \((a)\) follows because \(\varepsilon\leq \frac{1}{K^2H^2}\).

 \(\blacksquare\)

Thus, applying the Chernoff bound and taking the union bound over \(\bar{\Theta}_{\epsilon}\) and \(k\in[K]\), with probability at least \(1-\delta\), we have (JY: can you simply start with \(\mathbb{P}[\mathcal{E}_\pi]...\)?) \[\begin{align} ~~\sum_{\pi\in\mathcal{D}^k} & \mathtt{D}_{\mathtt{H}}^2 ( \mathbb{P}_{\theta}^{\pi}(\tau_H) , \mathbb{P}_{\theta^*}^{\pi} (\tau_H) ) \\ &\leq KH\sqrt{\varepsilon} + \frac{1}{2} \sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\frac{ \mathbb{P}_{ \theta^*}^{\pi}(\tau_H) }{\mathbb{P}_{\bar{\theta}}^{\pi}(\tau_H)} + \log\frac{K|\bar{\Theta}_{\epsilon}|}{\delta} \\ &\overset{(a)}\leq \frac{1}{2}\sum_{(\tau_H,\pi)\in\mathcal{D}^k} \log\frac{ \mathbb{P}_{ \theta^*}^{\pi}(\tau_H) }{\mathbb{P}_{ \theta }^{\pi}(\tau_H)} + 2\log\frac{K|\bar{\Theta}_{\epsilon}|}{\delta}, \quad \forall k\in[K], \end{align}\] where \((a)\) follows because \(\varepsilon\leq \frac{1}{K^2H^2}\).

The following proposition states that the constraint \(\Theta_{\min}^k\) does not rule out the true model.

Proof:  For any \(k\in[K]\), \(h\in\{0,1,\ldots,H-1\}\) and \((\tau_h,\pi)\in \mathcal{D}_h^k\), we have \[\begin{align} \mathbb{P}&\left( \mathbb{P}_{\theta^*}^{\pi }(\tau_h ) < p_{\min} \right) \\ & = \mathbb{E}\left[ \mathbb{P}\left( \mathbb{P}_{\theta^*}^{\pi }(\tau_h ) < p_{\min} \big| \pi \right)\right] \\ & = \mathbb{E}\left[ \sum_{\tau_h } \mathbb{P}_{\theta^*}^{\pi }(\tau_h^t)\mathbb{1}\left\{\mathbb{P}_{\theta^*}^{\pi^t}(\tau_h^t) < p_{\min} \right\} \right] \\ &\leq \mathbb{E}\left[\sum_{\tau_h } p_{\min}\right] \\ &\leq \frac{\delta}{KH} . \end{align}\]

Thus, taking the union bound over \(k\in[K]\), \(h\in\{0,1,\ldots,H-1\}\) and \((\tau_h,\pi)\in \mathcal{D}_h\), we conclude that \[\begin{align} \mathbb{P}(\mathcal{E}_{\min}) \geq 1 - \delta. \end{align}\]

 \(\blacksquare\)

9 Proof of 1 (for Online PSR-UCB)↩︎

In this section, we present the full analysis for the online algorithm PSR-UCB to show 1. In particular, the proof of 1 consists of three main steps. Step 1. We prove that the estimated model \(\hat{\theta}^k\) is not only accurate over the exploration policies, but also accurate conditioned on emipircal samples \(\tau_h\in\mathcal{D}_h^k\). Step 2. Building up on the first step, we are able to show that \(V_{\hat{\theta}^k,\hat{b}^k}^{\pi}\) is a valid upper bound on the total variation distance between \(\hat{\theta}^k\) and \(\theta^*\). Step 3. Based on a newly developed inequality that translates \(V_{\hat{\theta}^k,\hat{b}^k}^{\pi}\) to the ground-truth prediction feature \(\bar{\psi}^*(\tau_h)\), we show that the summation of \(V_{\hat{\theta}^k,\hat{b}^k}^{\pi}\) over the iteration \(k\) grows sublinear in time. This finally characterizes the optimality of the output policy, the accuracy of the output model and the sample complexity of PSR-UCB, and finishes the proof.

Before we proceed, we use \(\mathcal{E}\) to denote the event \(\mathcal{E}_{\omega}\cap\mathcal{E}_{\pi}\cap \mathcal{E}_{\min}\), where these three events are defined in 8. Due to [proposiiton: empirical distance less than log likelihood difference,proposition: hellinger distance less than log likelihood distance,proposition: propobability of empirical data is high] and union bound, we immediately have \(\mathbb{P}(\mathcal{E})\geq 1-3\delta.\)

9.1 Step 1: Estimation Guarantee↩︎

We show that the estimated model is accurate with the past exploration policies and dataset.

Proof:  To show the first inequality, note that by the selection of \(\hat{\theta}^k\), we have \(\hat{\theta}^k\in\mathcal{B}^k\) (defined in 1 ).

Following from 5, we have \[\begin{align} \sum_h & \sum_{(\tau_h,\pi)\in\mathcal{D}_h^k} \mathtt{D}_{\mathtt{TV}}^2 \left( \mathbb{P}_{\hat{\theta}^k}^{\pi}(\omega_h|\tau_h), \mathbb{P}_{\theta^*}^{\pi}(\omega_h|\tau_h) \right) \\ &\leq 6 \sum_h\sum_{(\tau_H,\pi)\in\mathcal{D}_h^k} \log \mathbb{P}_{\theta^*}^{\pi}(\tau_H) - 6\sum_h\sum_{(\tau_H,\pi)\in\mathcal{D}_h^k} \log \mathbb{P}_{\bar{\theta}^k}^{\pi}(\tau_H) + 31\log\frac{K|\bar{\Theta}_{\varepsilon}|}{\delta}\\ &\overset{(a)}\leq 6\max_{\theta'\in\Theta_{\min}^k}\sum_h\sum_{(\tau_H,\pi)\in\mathcal{D}_h^k} \log \mathbb{P}_{\theta'}^{\pi}(\tau_H) - 6\sum_h\sum_{(\tau_H,\pi)\in\mathcal{D}_h^k} \log \mathbb{P}_{\bar{\theta}^k}^{\pi}(\tau_H) + 31\log\frac{K|\bar{\Theta}_{\varepsilon}|}{\delta}\\ &\leq 7\beta, \end{align}\] where \((a)\) follows from \(\theta^*\in\Theta_{\min}^k\) (7).

To show the second inequality, 6 implies that \[\begin{align} \sum_{\pi\in\mathcal{D}^k}& \mathtt{D}_{\mathtt{H}}^2\left( \mathbb{P}_{\hat{\theta}^k}^{\pi}(\tau_H), \mathbb{P}_{\theta^*}^{\pi} (\tau_H) \right) \\ &\leq \sum_{\pi\in\mathcal{D}^k} \log\mathbb{P}_{\theta^*}^{\pi}(\tau_H) - \sum_{\pi\in\mathcal{D}^k}\log\mathbb{P}_{\hat{\theta}^k}^{\pi}(\tau_H) + 2\log\frac{K|\bar{\Theta}_{\varepsilon}|}{\delta}\\ &\overset{(a)}\leq \max_{ \theta'\in\Theta_{\min}^k } \sum_{\pi\in\mathcal{D}^k} \log\mathbb{P}_{\theta'}^{\pi}(\tau_H) - \sum_{\pi\in\mathcal{D}^k}\log\mathbb{P}_{\hat{\theta}^k}^{\pi}(\tau_H) + 2\log\frac{K|\bar{\Theta}_{\varepsilon}|}{\delta}\\ &\leq 7\beta, \end{align}\] where \((a)\) follows from \(\theta^*\in\Theta_{\min}^k\) (7).

 \(\blacksquare\)

9.2 Step 2: UCB for Total Variation Distance↩︎

Following from 1, the total variation distance between two PSRs is controlled by the estimation error. Hence, we first characterize the estimation error of \(\mathbf{M}_h^*(o_h,a_h)\) in the following lemma.

Proof:  We index future trajectory \(\omega_{h-1} = (o_{h},a_h,\ldots,o_H,a_H)\) by \(i\), and history trajectory \(\tau_{h-1}\) by \(j\). For simplicity, we denote \(\mathbf{m}^*(\omega_{h})^{\top}\left(\hat{\mathbf{M}}_h^{(k)}(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right)\) as \(w_{i}^{\top}\), and \(\bar{\hat{\psi}}_{h-1}^{(k)}(\tau_{h-1}) = \frac{ \hat{\psi}_{h-1}^{k}(\tau_{h-1})}{\hat{\phi}_{h-1}^{k\top}\hat{\psi}_{h-1}^{k}(\tau_{h-1})}\) as \(x_{j}\). We also denote \(\pi(\omega_{h-1}|\tau_{h-1})\) as \(\pi_{i|j}\).

Then, the LHS of ?? can be written as \[\begin{align} \text{LHS} &= \sum_{\tau_H} \left| \mathbf{m}^* (\omega_h)^{\top} \left(\hat{\mathbf{M}}_h^{k}(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right) \hat{\psi}_{h-1}^{k}(\tau_{h-1}) \right| \pi(\tau_H) \\ &\overset{(a)} = \sum_{i }\sum_{j }|w_{i}^{\top}x_{j}|\pi_{i|j}\mathbb{P}_{ \hat{\theta}^k }^{\pi}(j)\\ & = \sum_{j } \sum_{i } (\pi_{i|j}\cdot\mathtt{sgn}(w_{i}^{\top}x_{j})\cdot w_{i})^{\top} x_{j}\cdot\mathbb{P}_{ \hat{\theta}^k }^{\pi}(j)\\ & = \sum_{j} \left( \sum_{i} \pi_{i|j} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j} )\cdot w_{i} \right)^{\top} x_{j}\cdot\mathbb{P}_{ \hat{\theta}^k }^{\pi}(j)\\ &\overset{(b)}\leq \mathop{\mathbb{E}}_{j\sim\mathbb{P}_{ \hat{\theta}^k }^{\pi}}\left[ \big\| x_j \big\|_{\left(\hat{U}_{h-1}^{k}\right)^{-1}} \sqrt{ \left\|\sum_{i} \pi_{i|j} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j} )\cdot w_{i} \right\|_{\hat{U}_{h-1}^{k}}^2 } \right], \end{align}\] where \((a)\) follows because \(\mathbb{P}_{\theta}^{\pi}(\tau_h) = \phi_h^{\top}\psi(\tau_h)\), and \((b)\) follows from the Cauchy’s inequality.

Fix \(\tau_{h-1} = j_0\). We aim to upper bound the term \(I_1:= \left\|\sum_{i} \pi_{i|j_0} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j_0} )\cdot w_{i} \right\|_{\hat{U}_{h-1}^{k}}^2\). Then, we have \[\begin{align} &I_1 = \underbrace{\lambda \left\| \sum_{i} \pi_{i|j_0} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j_0} )\cdot w_{i} \right\|_2^2}_{I_2} + \underbrace{\sum_{j\in\mathcal{D}_{h-1}^{\tau}} \left[ \left(\sum_{i} \pi_{i|j_0} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j_0} )\cdot w_{i}\right)^{\top} x_{j} \right]^2 }_{I_3}. \end{align}\]

We first upper bound the first term \(I_2\) as follows: \[\begin{align} \sqrt{I_2} &= \sqrt{\lambda} \max_{x\in\mathbb{R}^{d_{h-1}}: \|x\|_2=1} \left|\sum_{i} \pi_{i|j_0}\mathtt{sgn}(w_i^{\top}x_{j_0})w_i^{\top} x\right|\\ & \leq \sqrt{\lambda} \max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{\omega_{h-1}} \left| \mathbf{m}^*(\omega_h)^{\top} \left( \hat{\mathbf{M}}_h^{k}(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right) x \right| \pi(\omega_{h-1}|j_0)\\ & \overset{(a)}\leq \sqrt{\lambda} \max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{\omega_{h-1}}\left|\mathbf{m}^{*}(\omega_h)^{\top} \hat{\mathbf{M}}_h^{k}(o_h,a_h) x \right| \pi(\omega_{h-1}|j_0) \\ & \quad + \sqrt{\lambda} \max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{\omega_{h-1}}\left|\mathbf{m}^*(\omega_h)^{\top} \mathbf{M}_h^*(o_h,a_h) x \right| \pi(\omega_{h-1}|j_0) \\ & \overset{(b)}\leq \frac{\sqrt{\lambda}}{\gamma}\max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{o_h,a_h}\left\| \hat{\mathbf{M}}_h^{k}(o_h,a_h) x \right\|_1 \pi(a_h|o_h,j_0) + \frac{\sqrt{\lambda}}{\gamma} \max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \| x\|_1 \\ & \overset{(c)}\leq \frac{2Q_A\sqrt{d\lambda}}{\gamma^2}, \end{align}\] where \((a)\) follows because \(|a+b|\leq |a|+|b|\), \((b)\) follows from 1, and \((c)\) follows from 2 and because \(\max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1}\| x\|_1 \leq \sqrt{d}\).

We next upper bound the second term \(I_3\) as follows. \[\begin{align} I_3 &\leq \sum_{\tau_{h-1}\in \mathcal{D}_{h-1}^{k}} \left( \sum_{\omega_{h-1}} \left| \mathbf{m}^* (\omega_{h})^{\top}\left(\hat{\mathbf{M}}_h^{k}(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right)\bar{\hat{\psi}}^{k}(\tau_{h-1}) \right| \pi(\omega_{h-1}|j_0) \right)^2 \\ & \overset{(a)}\leq \sum_{\tau_{h-1}\in \mathcal{D}_{h-1}^{k}} \left( \sum_{\omega_{h-1}} \left|\mathbf{m}^*(\omega_h)^{\top} \left( \hat{\mathbf{M}}_h^{k}(o_h,a_h) \bar{\hat{\psi}}^{k}(\tau_{h-1}) - \mathbf{M}_h^*(o_h,a_h) \bar{\psi}^*(\tau_{h-1}) \right) \right| \pi(\omega_{h-1}|j_0) \right. \\ & + \left. \sum_{\omega_{h-1}}\left| \mathbf{m}^*(\omega_h)^{\top} \mathbf{M}_h^*(o_h,a_h) \left(\bar{\hat{\psi}}^{(k)}(\tau_{h-1}) - \bar{\psi}^*(\tau_{h-1}) \right) \right| \pi(\omega_{h-1}|j_0) \right)^2 \\ & \overset{(b)}\leq \sum_{\tau_{h-1}\in \mathcal{D}_{h-1}^{k} } \left( \frac{1}{\gamma} \sum_{o_h,a_h} \left\| \mathbb{P}_{\hat{\theta}^{k}}(o_h|\tau_{h-1})\bar{\hat{\psi}}_{h}^{k}(\tau_{h}) - \mathbb{P}_{\theta}(o_h|\tau_{h-1}) \bar{\psi}_{h}^*(\tau_{h}) \right\|_1 \pi(a_h|o_h,j_0) \right. \\ &+ \left. \frac{1}{\gamma} \sum_{\tau_{h-1}\in \mathcal{D}_{h-1}^{\tau}} \left\| \bar{\hat{\psi}}_{h-1}^{k}(\tau_{h-1}) - \bar{\psi}_{h-1}^*(\tau_{h-1}) \right\|_1 \right)^2 \\ & \overset{(c)}= \frac{1}{\gamma^2} \sum_{\tau_{h-1}\in \mathcal{D}_{h-1}^{k} } \left( \sum_{o_h,a_h} \sum_{\ell=1}^{|\mathcal{Q}_h|} \left| \mathbb{P}_{\hat{\theta}^k}(\mathbf{o}_h^{\ell},o_h|\tau_{h-1},a_h,\mathbf{a}_h^{\ell}) - \mathbb{P}_{ \theta^* }(\mathbf{o}_h^{\ell},o_h|\tau_{h-1},a_h,\mathbf{a}_h^{\ell}) \right|\pi(a_h|o_h,j_0) \right. \\ &+ \left. \sum_{\ell=1}^{|\mathcal{Q}_{h-1}|} \left| \mathbb{P}_{\hat{\theta}^k}(\mathbf{o}_{h-1}^{\ell} | \tau_{h-1}, \mathbf{a}_{h-1}^{\ell} ) - \mathbb{P}_{ \theta^* } (\mathbf{o}_{h-1}^{\ell} | \tau_{h-1}, \mathbf{a}_{h-1}^{\ell} ) \right| \right)^2 \\ &\leq \frac{1}{\gamma^2} \sum_{\tau_{h-1}\in \mathcal{D}_{h-1}^{k} } \left( \;\sum_{\mathbf{a}_{h-1}\in \mathcal{Q}_h^{\exp} } \sum_{\omega_{h-1}^o }\left|\mathbb{P}_{\hat{\theta}^k}(\omega^o_{h-1}|\tau_{h-1}, \mathbf{a}_{h-1} ) - \mathbb{P}_{\theta^*}(\omega^o_{h-1}|\tau_{h-1}, \mathbf{a}_{h-1} ) \right| \right)^2 \\ &\leq \frac{4|\mathcal{A}|^2Q_A^2}{\gamma^2} \sum_{\tau_{h-1}\in\mathcal{D}_{h-1}^{k}} \mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\hat{\theta}^k }^{\mathtt{u}_{\mathcal{Q}_{h-1}^{\exp} }}(\omega_{h-1} | \tau_{h-1} ) , \mathbb{P}_{ \theta^* }^{\mathtt{u}_{\mathcal{Q}_{h-1}^{\exp} } } (\omega_{h-1} | \tau_{h-1} ) \right), \end{align}\] where \((a)\) follows because \(|a+b|\leq |a|+|b|\), \((b)\) follows from 5 and 1, and \((c)\) follows from the physical meaning of the prediction feature, i.e. \([\bar{\psi}(\tau_h)]_{\ell} = \mathbb{P}_{\theta}(\mathbf{o}_h^{\ell}|\tau_h,\mathbf{a}_h^{\ell})\).

By combining the upper bounds for \(I_2\) and \(I_3\), we conclude that \[\begin{align} I_1\leq \frac{4\lambda Q_A^2 d}{\gamma^4} + \frac{ 4|\mathcal{A}|^2Q_A^2}{\gamma^2}\sum_{\tau_{h-1}\in\mathcal{D}_{h-1}^{k}} \mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\hat{\theta}^k }^{\mathtt{u}_{\mathcal{Q}_{h-1}^{\exp} }}(\omega_{h-1} | \tau_{h-1} ) , \mathbb{P}_{ \theta^* }^{\mathtt{u}_{\mathcal{Q}_{h-1}^{\exp} } } (\omega_{h-1} | \tau_{h-1} ) \right) = \alpha_{h-1}^k, \end{align}\] which completes the proof.

 \(\blacksquare\)

The following lemma validates that \(V_{\hat{\theta}^k,\hat{b}^k}^{\pi}\) is an upper bound on the total variation distance between the estimated model \(\hat{\theta}^k\) and the true model \(\theta^*\).

Proof:  We proceed the proof as follows: \[\begin{align} \sum_{h}(\alpha_{h-1}^k)^2 & \leq \frac{4\lambda HQ_A^2d }{\gamma^4} + \frac{ 4|\mathcal{A}|^2Q_A^2}{\gamma^2} \sum_h\sum_{(\tau_{h-1},\pi)\in\mathcal{D}_{h-1}^{k}} \mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\hat{\theta}^k }^{ \pi }(\omega_{h-1} | \tau_{h-1} ) , \mathbb{P}_{ \theta^* }^{ \pi } (\omega_{h-1} | \tau_{h-1} ) \right) \\ &\overset{(a)}\leq \frac{ 4 \lambda H Q_A^2 d }{\gamma^4} + \frac{ 28\beta |\mathcal{A}|^2 Q_A^2 }{\gamma^2}, \end{align}\] where \((a)\) follows from 1. The proof then follows directly from 1, 2 and the Cauchy’s inequality.

 \(\blacksquare\)

Note that the reward function \(R\) is within \([0,1]\). We hence obtain the following corollary.

9.3 Step 3: Sublinear Summation↩︎

To prove that \(\sum_k V_{\hat{\theta}^k,\hat{b}^k}^{\pi^k}\) is sublinear, i.e., scales as \(O(\sqrt{K})\), we first prove the following lemma that relates the estimated feature and the ground-truth feature via the total variation distance between the estimated model and the true model.

Proof:  Recall that \[\begin{align} \hat{U}_h^k = \lambda I + \sum_{\tau\in\mathcal{D}_h^k} \bar{\hat{\psi}}^{k}(\tau_h)\bar{\hat{\psi}}^{k}(\tau_h)^{\top}. \end{align}\]

We define the ground-truth counterpart of \(\hat{U}_h^k\) as follows: \[\begin{align} U_h^k = \lambda I + \sum_{\tau\in\mathcal{D}_h^k} \bar{\psi}^*(\tau_h) \bar{ \psi}^*(\tau_h)^{\top}. \end{align}\]

Then, following from 13, we have \[\begin{align} &\sqrt{ \sum_{h=0}^{H-1} \left\| \bar{\hat{\psi}}^{k}(\tau_h) \right\|^2_{\left(\hat{U}_h^{k}\right)^{-1}} }\\ &\quad \leq \sum_{h=0}^{H-1} \left\|\bar{\hat{\psi}}^{k}(\tau_h) \right\|_{\left(\hat{U}_h^{k}\right)^{-1}} \\ &\quad \leq \frac{1}{\sqrt{\lambda}} \sum_{h=0}^{H-1} \left\| \bar{\hat{\psi}}^{k}(\tau_h) - \bar{\psi}^* (\tau_h) \right\|_2 + \sum_{h=0}^{H-1} \left(1 + \frac{ \sqrt{r} \sqrt{ \sum_{\tau_h\in\mathcal{D}_h^k } \left\| \bar{\hat{\psi}}^{k}(\tau_h) - \bar{\psi}^* (\tau_h) \right\|_2^2 } }{ \sqrt{\lambda} } \right) \left\|\bar{\psi}^* (\tau_h) \right\|_{\left(U_h^{k}\right)^{-1}}. \end{align}\]

Furthermore, note that \[\begin{align} \left\| \bar{\hat{\psi}}^{k}(\tau_h) - \bar{\psi}^* (\tau_h) \right\|_2 & \leq \left\| \bar{\hat{\psi}}^{k}(\tau_h) - \bar{\psi}^* (\tau_h) \right\|_1 \\ &\overset{(a)} \leq 2|\mathcal{A}|Q_A \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{ \hat{\theta}^k }^{ \mathtt{u}_{\mathcal{Q}_h^{\exp}} }(\omega_h|\tau_h), \mathbb{P}_{ \theta^* }^{ \mathtt{u}_{\mathcal{Q}_h^{\exp}} }(\omega_h|\tau_h) \right), \end{align}\] where \((a)\) follows from the physical meaning of the prediction feature.

Following from 1, we conclude that \[\begin{align} &\sqrt{ \sum_{h=0}^{H-1} \left\|\bar{\hat{\psi}}^{k}(\tau_h) \right\|^2_{\left(\hat{U}_h^{k}\right)^{-1}} }\\ &\quad \leq \frac{1}{\sqrt{\lambda}} \sum_{h=0}^{H-1} \left\| \bar{\hat{\psi}}^{k}(\tau_h) - \bar{\psi}^* (\tau_h) \right\|_2 + \left(1 + \frac{2 |\mathcal{A}|Q_A\sqrt{7r\beta} }{ \sqrt{\lambda} } \right) \sum_{h=0}^{H-1} \left\|\bar{\psi}^* (\tau_h) \right\|_{\left(U_h^{k}\right)^{-1}}. \end{align}\]

For the first term, taking expectation, we have \[\begin{align} \sum_{h=0}^{H-1} & \mathbb{E}_{\tau_h\sim \mathbb{P}_{\theta^*}^{\pi}} \left[ \left\| \bar{\hat{\psi}}^{k}(\tau_h) - \bar{\psi}^*(\tau_h) \right\|_1 \right] \\ &\leq \sum_{h=0}^{H-1} \sum_{\tau_h} \left( \left\| \bar{\hat{\psi}}^{k}(\tau_h)\left( \mathbb{P}_{\theta}^{\pi}(\tau_h) - \mathbb{P}_{\hat{\theta}^k}^{\pi}(\tau_h) \right) + \bar{\hat{\psi}}^{k}(\tau_h) \mathbb{P}_{\hat{\theta}^k}^{\pi}(\tau_h) - \bar{\psi}^*(\tau_h) \mathbb{P}_{\theta^*}^{\pi}(\tau_h) \right\|_1 \right)\\ &\leq \sum_{h=0}^{H-1} \sum_{\tau_h} \left( \left\| \bar{\hat{\psi}}^{k}(\tau_h) \right\|_1 \left| \mathbb{P}_{\theta}^{\pi}(\tau_h) - \mathbb{P}_{\hat{\theta}^k}^{\pi}(\tau_h) \right| + \left\| \hat{\psi}^{k}(\tau_h) - \psi^*(\tau_h) \right\|_1 \pi(\tau_h)\right)\\ &\overset{(a)}\leq 2\sum_{h=0}^{H-1} Q_A \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi}(\tau_h) , \mathbb{P}_{\hat{\theta}^k}^{\pi}(\tau_h)\right) \\ & \leq 2 H Q_A \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi}(\tau_H) , \mathbb{P}_{\hat{\theta}^k}^{\pi}(\tau_H)\right), \end{align}\] where \((a)\) follows from \(\|\bar{\hat{\psi}}^k(\tau_h)\|_1 \leq |\mathcal{Q}_h^A|\leq Q_A\), and the physical meaning of \(\psi(\tau_h)\).

Thus, \[\begin{align} \mathop{\mathbb{E}}_{\tau_H\sim\mathbb{P}_{\theta^*}^{\pi}}&\left[\sqrt{ \sum_{h=0}^{H-1} \left\|\bar{\hat{\psi}}^{k}(\tau_h) \right\|^2_{\left(\hat{U}_h^{k}\right)^{-1}} } \right] \\ &\leq \left(1 + \frac{2|\mathcal{A}|Q_A\sqrt{7r\beta}}{\sqrt{\lambda}} \right) \sum_{h=0}^{H-1} \mathop{\mathbb{E}}_{\tau_h\sim\mathbb{P}_{\theta^*}^{\pi}} \left[ \left\|\bar{\psi}^* (\tau_h) \right\|_{\left(U_h^{k}\right)^{-1}} \right] + \frac{2HQ_A}{\sqrt{\lambda}} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi}(\tau_H) , \mathbb{P}_{\hat{\theta}^k}^{\pi}(\tau_H)\right). \end{align}\]

 \(\blacksquare\)

The following lemma can be proved via the \(\ell_2\) Eluder argument [20], [23]. Since we have a slightly different estimation oracle and guarantee, we provide the full proof here for completeness.

Proof:  First, by the first inequality in 1, we have \[\begin{align} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{ \theta^* }^{\pi^k}(\tau_h) , \mathbb{P}_{\hat{\theta}^k}^{\pi^k}(\tau_h)\right) \leq \sum_h\sum_{\tau_H} \left| \hat{\mathbf{m}}^k(\omega_h)^{\top}\left(\hat{\mathbf{M}}_h^k(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h)\right) \psi^* (\tau_{h-1}) \right| \pi^k(\tau_H). \end{align}\]

It suffices to upper bound \(\sum_{\tau_H} \left| \hat{\mathbf{m}}^k(\omega_h)^{\top}\left(\hat{\mathbf{M}}_h^k(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h)\right) \psi^* (\tau_{h-1}) \right| \pi(\tau_H)\) for any policy \(\pi\). For simplicity, we use similar notations as in 2. We index the future trajectory \(\omega_{h-1} = (o_{h},a_h,\ldots,o_H,a_H)\) by \(i\), and the history trajectory \(\tau_{h-1}\) by \(j\). We represent \(\bar{\psi}^*(\tau_{h-1})\) as \(x_j\), and \(\hat{\mathbf{m}}^k(\omega_h)^{\top}\left(\hat{\mathbf{M}}_h^k(o_h,a_h) - \mathbf{M}_h^* (o_h,a_h)\right)\) as \(w_{i}\). We also denote \(\pi(\omega_{h-1}|\tau_{h-1})\) by \(\pi_{i|j}\).

Let \(\lambda_0\) be a constant determined later and define the matrix \[\begin{align} \Lambda_h^k = \lambda_0 I + \sum_{t<k} \mathbb{E}_{j\sim\mathbb{P}_{\theta^*}^{\pi^t}} \left[ x_jx_j^{\top} \right]. \end{align}\]

Then, for any \(\pi\), we have \[\begin{align} \sum_{\tau_H} & \left| \hat{\mathbf{m}}^k(\omega_h)^{\top}\left(\hat{\mathbf{M}}_h^k(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h)\right) \psi^* (\tau_{h-1}) \right| \pi(\tau_H)\\ &= \mathbb{E}_{j\sim \mathbb{P}_{\theta^*}^{\pi^k}} \left[ \sum_{i} | \pi_{i|j} w_{i }^{\top} x_j|\right] \\ & = \mathbb{E}_{j\sim \mathbb{P}_{\theta^*}^{\pi^k}} \left[ \left(\sum_{i} \pi_{i|j} \mathtt{sgn}(w_{i }^{\top} x_j) w_{i } \right)^{\top} x_j\right] \\ &\overset{(a)} \leq \mathbb{E}_{j\sim \mathbb{P}_{\theta^*}^{\pi^k}} \left[ \|x_j\|_{\Lambda_h^{\dagger}} \left\|\sum_{i} \pi_{i|j} \mathtt{sgn}(w_{i }^{\top} x_j) w_{i } \right\|_{\Lambda_h} \right], \end{align}\] where \((a)\) follows from the Cauchy’s inequality.

Now we fix \(j=j_0\) and consider the following term. \[\begin{align} &\left\|\sum_{i} \pi_{i|j_0}\mathtt{sgn}(w_{i }^{\top}x_{j_0}) w_{i} \right\|^2_{\Lambda_h} \\ & \quad = \underbrace{\lambda_0 \left\|\sum_{i} \pi_{i|j_0}\mathtt{sgn}(w_{i }^{\top}x_{j_0}) w_{i} \right\|^2_2}_{I_1} + \underbrace{\sum_{t<k} \mathbb{E}_{j\sim \mathbb{P}_{\theta^*}^{\pi^t}} \left[ \left( \sum_{i} \pi_{i|j_0}\mathtt{sgn}(w_{i }^{\top}x_{j_0}) w_{i }^{\top} x_j \right)^2 \right] }_{I_2} \end{align}\] For the first term \(I_1\), we have \[\begin{align} \sqrt{I_1} &= \sqrt{\lambda_0} \max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \left|\sum_{i} \pi_{i|j_0}\mathtt{sgn}(w_{i }^{\top}x_{j_0}) w_{i}^{\top} x \right| \\ & \leq \sqrt{\lambda_0} \max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{\omega_{h-1}} \pi(\omega_{h-1} | j_0 ) \left| \hat{\mathbf{m}}^k(\omega_{h-1})^{\top} x \right| \\ &\quad + \sqrt{\lambda_0} \max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{\omega_{h-1}} \pi (\omega_{h-1} | j_0 ) \left| \hat{\mathbf{m}}^k(\omega_h)^{\top} \mathbf{M}_h^*(o_h,a_h) x \right| \\ &\overset{(a)}\leq \frac{\sqrt{d\lambda_0}}{\gamma} + \frac{ Q_A\sqrt{d \lambda_0 } }{\gamma^2}\\ &\leq \frac{ 2Q_A\sqrt{d \lambda_0 } }{\gamma^2}, \end{align}\] where \((a)\) follows from 1, 2, and the fact that \(\max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \|x\|_1\leq \sqrt{d_{h-1}}\leq \sqrt{d}\).

For the second term \(I_2\), we have \[\begin{align} I_2 &\leq \sum_{t<k}\mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^t}} \left[ \left( \sum_{\omega_{h-1}} \pi(\omega_{h-1} | j_0) \left| \hat{\mathbf{m}}^k(\omega_h)^{\top} \left( \hat{\mathbf{M}}_h^k(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right) \bar{\psi}^*(\tau_{h-1}) \right| \right)^2 \right] \\ & \leq \sum_{t<k}\mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^t}} \left[ \left( \sum_{\omega_{h-1}} \pi(\omega_{h-1} | j_0) \left| \hat{\mathbf{m}}^k(\omega_{h-1})^{\top} \left( \bar{\psi}^*(\tau_{h-1}) - \bar{\hat{\psi}}^k (\tau_{h-1}) \right) \right| \right. \right.\\ & + \left. \left. \sum_{\omega_{h-1}} \pi(\omega_{h-1}|j_0)\left| \hat{\mathbf{m}}^k(\omega_h)^{\top} \left( \hat{\mathbf{M}}_h^k(o_h,a_h)\bar{\hat{\psi}}^k (\tau_{h-1}) - \mathbf{M}_h^*(o_h,a_h) \bar{\psi}^*(\tau_{h-1}) \right) \right| \right)^2 \right] \\ &\overset{(a)}\leq \sum_{t<k} \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^t}} \left[ \left( \frac{1}{\gamma} \left\| \bar{\psi}^* (\tau_{h-1}) - \bar{\hat{\psi}}^k (\tau_{h-1}) \right\|_1 \right. \right.\\ & + \left. \left. \frac{1}{\gamma} \sum_{o_h,a_h} \pi^k(a_h|o_h,j_0)\left\| \hat{\mathbf{M}}^k(o_h,a_h)\bar{\hat{\psi}}^k (\tau_{h-1}) - \mathbf{M}_h^*(o_h,a_h) \bar{\psi}^*(\tau_{h-1}) \right\|_1 \right)^2 \right] \\ & = \frac{1}{\gamma^2} \sum_{t<k} \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^t}} \left[ \left( \sum_{\ell=1}^{|\mathcal{Q}_{h-1}|}\left| \mathbb{P}_{\hat{\theta}^k}(\mathbf{o}_{h-1}^{\ell} | \tau_{h-1}, \mathbf{a}_{h-1}^{\ell} ) - \mathbb{P}_{\theta^*}(\mathbf{o}_{h-1}^{\ell} | \tau_{h-1}, \mathbf{a}_{h-1}^{\ell} )\right| \right. \right.\\ & + \left.\left.\sum_{o_h,a_h} \sum_{\ell=1}^{|\mathcal{Q}_h|} \pi^k(a_h|o_h,j_0)\left| \mathbb{P}_{\hat{\theta}^k}(\mathbf{o}_h^{\ell}, o_h | \tau_{h-1} ,a_h, \mathbf{a}_h^{\ell} ) - \mathbb{P}_{\theta^*} (\mathbf{o}_h^{\ell}, o_h | \tau_{h-1} ,a_h, \mathbf{a}_h^{\ell} ) \right| \right)^2\right] \\ & \leq \frac{|\mathcal{Q}_{h-1}^{\exp}|^2 }{\gamma^2} \sum_{t<k} \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^t}} \left[ \mathtt{D}^2_{\mathtt{TV}} \left( \mathbb{P}_{\hat{\theta}^k }^{ \mathtt{u}_{\mathcal{Q}_{h-1}^{\exp}} } (\omega_{h-1} | \tau_{h-1} ) , \mathbb{P}_{\theta^* }^{ \mathtt{u}_{\mathcal{Q}_{h-1}^{\exp}} } (\omega_{h-1} | \tau_{h-1} )\right) \right] \\ & \overset{(b)}\leq \frac{4|\mathcal{A}|^2Q_A^2}{\gamma^2} \sum_{t<k} \mathtt{D}^2_{\mathtt{H}} \left( \mathbb{P}_{\hat{\theta}^k }^{ \nu_{h}(\pi^t, \mathtt{u}_{\mathcal{Q}_{h-1}^{\exp}}) } ( \tau_{H} ) , \mathbb{P}_{\theta^* }^{ \nu_h(\pi^t, \mathtt{u}_{\mathcal{Q}_{h-1}^{\exp}}) } ( \tau_{H} ) \right), \end{align}\] where \((a)\) follows from 1, and \((b)\) follows because \(|\mathcal{Q}_{h-1}^{\exp}|=|\mathcal{A}||\mathcal{Q}_h^A| + |\mathcal{Q}_{h-1}^A| \leq 2|\mathcal{A}|Q_A\).

Thus, we have \[\begin{align} \sum_{\tau_H}& \left| \hat{\mathbf{m}}^k(\omega_h)^{\top}\left(\hat{\mathbf{M}}_h^k(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h)\right)\psi^*(\tau_{h-1}) \right| \pi(\tau_H)\\ &\leq \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi} } \left[ \tilde{\alpha}_{h-1}^k \left\| \bar{\psi}^* (\tau_{h-1}) \right\|_{\Lambda_{h-1}^{\dagger} } \right], \end{align}\] where \[\begin{align} (\tilde{\alpha}_{h-1}^k)^2 = \frac{4\lambda_0Q_A^2d}{\gamma^4} + \frac{4|\mathcal{A}|^2Q_A^2}{\gamma^2} \sum_{t<k} \mathtt{D}^2_{\mathtt{H}} \left( \mathbb{P}_{\hat{\theta}^k }^{ \nu_h(\pi^t, \mathtt{u}_{\mathcal{Q}_{h-1}^{\exp}}) } ( \tau_{H} ) , \mathbb{P}_{\theta^* }^{ \nu_h(\pi^t, \mathtt{u}_{\mathcal{Q}_{h-1}^{\exp}}) } ( \tau_{H} ) \right) . \end{align}\]

By choosing \(\lambda_0 = \frac{\gamma^4}{4Q_A^2d}\), and recalling 1, we further have \[\begin{align} \sum_{h}( \tilde{\alpha}_{h-1}^k)^2\leq \frac{28 |\mathcal{A}|^2Q_A^2 \beta }{\gamma^2} \triangleq \tilde{\alpha}. \end{align}\]

Hence, with the Cauchy’s inequality, we have \[\begin{align} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\hat{\theta}^k}^{\pi^k}(\tau_H), \mathbb{P}_{\theta}^{\pi^k} (\tau_H) \right) \leq \min\left\{\tilde{\alpha} \sqrt{ \sum_{h=1}^H \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^k}}\left[ \left\| \bar{\psi}^*(\tau_{h-1}) \right\|^2_{\Lambda_{h-1}^{-1}} \right] } , 2\right\}. \end{align}\]

Taking the summation and applying the elliptical potential lemma (i.e., 14), we have \[\begin{align} \sum_{k=1}^K& \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\hat{\theta}^k}^{\pi^k}(\tau_H), \mathbb{P}_{\theta^*}^{\pi^k} (\tau_H) \right) \\ &\leq \sqrt{K} \sqrt{ \sum_{k=1}^K \sum_{h=1}^H \min\left\{ \tilde{\alpha} \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^k}}\left[ \left\| \bar{\psi}^*(\tau_{h-1}) \right\|^2_{\Lambda_{h-1}^{\dagger}} \right] , 4 \right\}} \\ & \lesssim \frac{|\mathcal{A}|Q_A}{\gamma} \sqrt{rHK\beta\log(1+dQ_AK/\gamma)}, \end{align}\] which yields the final result.

 \(\blacksquare\)

The next lemma shows that the summation \(\sum_{k=1}^K V_{\hat{\theta}^{k}, \hat{b}^k}^{\pi^k}\) grows sublinearly in \(K\).

Proof:  First, since \(\hat{b}^k(\tau_H)\in[0,1]\), we have \[\begin{align} V_{\hat{\theta}^{k}, \hat{b}^k}^{\pi^k} &\leq V_{ \theta^* , \hat{b}^k}^{\pi^k} + \frac{1}{2}\mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\hat{\theta}^{k}}^{\pi^k}, \mathbb{P}_{\theta^*}^{\pi^k}\right). \end{align}\]

Then, following from 4, we have \[\begin{align} \sum_{k=1}^K& V_{ \theta^* , \hat{b}^k}^{\pi^k} \\ & \leq \sum_{k=1}^K\min\left\{\alpha\left(1 + \frac{2|\mathcal{A}|Q_A\sqrt{7r\beta}}{\sqrt{\lambda}} \right)\sum_{h=0}^{H-1} \mathop{\mathbb{E}}_{\tau_H\sim \mathbb{P}_{\theta^*}^{\pi^k}} \left[ \left\| \bar{\psi}^*(\tau_h) \right\|_{(U_h^{k})^{-1}} \right] + \sum_{k=1}^K\frac{\alpha HQ_A}{\sqrt{\lambda}} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi^k}(\tau_H) , \mathbb{P}_{\hat{\theta}^k}^{\pi^k}(\tau_H)\right) ,1 \right\} \\ &\leq \underbrace{ \sum_{k=1}^K \min\left\{ \alpha\left(1 + \frac{2|\mathcal{A}|Q_A\sqrt{7r\beta}}{\sqrt{\lambda}} \right) \sum_{h=0}^{H-1} \mathop{\mathbb{E}}_{\tau_h\sim \mathbb{P}_{\theta^*}^{\pi^k}} \left[ \left\| \bar{\psi}^*(\tau_h) \right\|_{(U_h^{k})^{-1}} \right] , 1 \right\} }_{I_1} + \sum_{k=1}^K\frac{\alpha HQ_A}{\sqrt{\lambda}} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi^k}(\tau_H) , \mathbb{P}_{\hat{\theta}^k}^{\pi^k}(\tau_H)\right). \end{align}\]

We next analyze the term \(I_1\). Recall that \(U_h^k = \lambda I + \sum_{\tau_h\in\mathcal{D}_h^k} \bar{\psi}^*(\tau_h)\bar{\psi}^*(\tau_h)^{\top}\). Further note that \[\left\{ \mathbb{E}_{\tau_h\sim\mathbb{P}_{\theta^*}^{\pi^k}}\left[\|\bar{\psi}^*(\tau_h)\|_{(U_h^k)^{-1}} \right] - \|\bar{\psi}^*(\tau_h^{k+1,h+1})\|_{(U_h^k)^{-1}} \right\}_{k=1}^K\] forms a martingale. Applying the Azuma-Hoeffding inequality (see 11), we have, with probability at least \(1-\delta\), \[\begin{align} I_1 & \leq \sqrt{2K\log(2/\delta)} + \sum_{k=1}^K\min\left\{ \alpha\left(1 + \frac{2|\mathcal{A}|Q_A\sqrt{7r\beta}}{\sqrt{\lambda}} \right) \sum_{h=0}^{H-1} \left\| \bar{\psi}^*(\tau_h^{k+1,h+1}) \right\|_{(U_h^{k})^{-1}} , 1 \right\}\\ &\overset{(a)}\lesssim \sqrt{2K\log(2/\delta)} + \alpha\left(1 + \frac{2|\mathcal{A}|Q_A\sqrt{7r\beta}}{\sqrt{\lambda}} \right) H \sqrt{rK\log(1+K/\lambda)}\\ &\lesssim \alpha\left(1 + \frac{|\mathcal{A}|Q_A\sqrt{r\beta}}{\sqrt{\lambda}} \right) H \sqrt{rK\log(1+K/\lambda)}, \end{align}\] where \((a)\) follows from 14.

Let \(\beta:=\max\{\log(1+K/\lambda), \log(1+dQ_AK/\gamma) \}\) for simplicity. By 5, we have, \[\begin{align} \sum_k & V_{\hat{\theta}^k,\hat{b}^k }^{\pi^k} \\ & \lesssim \alpha\left(1 + \frac{|\mathcal{A}|Q_A\sqrt{r\beta}}{\sqrt{\lambda}} \right) H\sqrt{rK\beta_0} + \frac{\alpha H }{\sqrt{\lambda} } \frac{|\mathcal{A}|Q_A^2\sqrt{\beta} }{\gamma}\sqrt{rHK\beta_0} \\ &\leq \alpha \left(1 + \frac{|\mathcal{A}|Q_A\sqrt{r\beta}}{\sqrt{\lambda}} + \frac{|\mathcal{A}|Q_A^2\sqrt{\beta H}}{\gamma\sqrt{\lambda}}\right) H\sqrt{rK\beta_0} \\ &\lesssim \left( \frac{|\mathcal{A}|Q_A\sqrt{\beta}}{\gamma} + \frac{Q_A\sqrt{dH}}{\gamma^2} \sqrt{\lambda} \right)\left( 1 + \frac{|\mathcal{A}|Q_A\sqrt{\beta} \max\{ \sqrt{r}, Q_A\sqrt{ H}/\gamma \} }{ \sqrt{\lambda}}\right) H\sqrt{rK\beta_0}\\ & = \left( 1 + \frac{ \sqrt{dH}}{ |\mathcal{A}|\sqrt{\beta}\gamma} \sqrt{\lambda} \right)\left(1 + \frac{ |\mathcal{A}| Q_A\sqrt{\beta} \max\{ \sqrt{r}, Q_A\sqrt{ H}/\gamma \} }{ \sqrt{\lambda}}\right) \frac{|\mathcal{A}| Q_A H \sqrt{r\beta K\beta_0}}{\gamma} \\ &\overset{(a)}\lesssim \left( 1 + \frac{ Q_A \sqrt{dH} \max\{ \sqrt{r}, Q_A\sqrt{ H}/\gamma \} }{\gamma} \right) \frac{|\mathcal{A}| Q_A H\sqrt{r\beta K \beta_0}}{\gamma} \\ &\leq \left( \sqrt{r} + \frac{Q_A\sqrt{H}}{\gamma} \right) \frac{ |\mathcal{A}|Q_A^2H\sqrt{drH\beta K \beta_0} }{\gamma^2}, \end{align}\] where \((a)\) follows by choosing \[\begin{align} \lambda = \frac{\gamma |\mathcal{A}|^2Q_A \beta \max\{ \sqrt{r}, Q_A\sqrt{ H}/\gamma \} }{ \sqrt{d H} }. \end{align}\] This completes the proof.

 \(\blacksquare\)

9.4 Proof of 1↩︎

Proof:  The proof is under event \(\mathcal{E}\), which occurs with probability at least \(1-3\delta\).

Following from 2, if PSR-UCB terminates, we have \[\begin{align} \forall \pi,~~ \mathtt{D}_{\mathtt{TV}}\left(\mathbb{P}_{ \theta^{\epsilon} }^{\pi}(\tau_H), \mathbb{P}_{\theta^*}^{\pi}(\tau_H) \right) = 2\max_{R} \left|V_{\theta^\epsilon,R}^{\pi} - V_{\theta^*,R}^{\pi}\right| \leq 2 V_{\theta^\epsilon,\hat{b}^{\epsilon}}^{\pi} \leq \epsilon, \end{align}\] where the last inequality follows from the termination condition of PSR-UCB.

In addition, \[\begin{align} V_{\theta^*,R}^{\pi^*} &- V_{\theta^*,R}^{\bar{\pi}}\\ & = V_{\theta^*,R}^{\pi^*} - V_{\theta^\epsilon,R}^{\pi^*} + V_{\theta^\epsilon,R}^{\pi^*} - V_{\theta^\epsilon,R}^{\bar{\pi}} + V_{\theta^\epsilon,R}^{\bar{\pi}} - V_{\theta^*,R}^{\bar{\pi}} \\ &\overset{(a)}\leq 2\max_{\pi} V_{\theta^\epsilon,\hat{b}^{\epsilon}}^{\pi}\leq \epsilon, \end{align}\] where \((a)\) is due to the design of \(\bar{\pi}\) and 2.

Finally, recall that 6 states that \[\begin{align} \sum_{k=1}^K V_{\hat{\theta}^k,\hat{b}^k} & \lesssim \left( \sqrt{r} + \frac{Q_A\sqrt{H}}{\gamma} \right) \frac{ |\mathcal{A}|Q_A^2H\sqrt{drH\beta K \beta_0} }{\gamma^2}. \end{align}\]

By the pigeon-hole principle and the termination condition of PSR-UCB, if \[K = \tilde{O}\left( \left( r + \frac{Q_A^2 H}{\gamma^2}\right) \frac{r d H^2 |\mathcal{A}|^2 Q_A^4 \beta}{\gamma^4\epsilon^2} \right),\] PSR-UCB must terminate within \(K\) episodes, implying that the sample complexity of PSR-UCB is at most

\[\tilde{O}\left( \left( r + \frac{Q_A^2 H}{\gamma^2}\right) \frac{r d H^3 |\mathcal{A}|^2 Q_A^4 \beta}{\gamma^4\epsilon^2} \right).\]

 \(\blacksquare\)

9.5 Proof of 1 (For POMDPs)↩︎

When PSR is specialized to an \(m\)-step decodable POMDP, we follow the setup described in Section D in [19]. In POMDPs, there exists a state space \(\mathcal{S}\) such that the system dynamics proceeds as follows. When \(h=1\), the system is at the state \(s_1\) and the agent observes \(o_1\) that is determined by \(s_1\). At any step \(h\geq 1\), if the system is at a state \(s_h\in\mathcal{S}\) and the agent takes an action \(a_h\in\mathcal{A}\), the system transits to a state \(s_{h+1}\) sampled from a transition distribution \(\mathbf{T}_{h,a_h}(\cdot|s_h)\), and then the agent observes \(o_{h+1}\) sampled from an emission distribution \(\mathbf{O}_h(\cdot|s_{h+1})\). Notably, all states are unobservable to the agent.

To define \(m\)-step decodability, we introduce the \(m\)-step emission-action matrix \(\{\mathbf{G}_h\}\) where \[\begin{align} [\mathbf{G}_h]_{(\mathbf{a},\mathbf{o}), s} = \mathbb{P}(o_h,\ldots, o_{\min\{h+m-1, H\} }= \mathbf{o} | s_h = s, a_h\ldots,a_{\min\{h+m-2,H\} } = \mathbf{a} ). \end{align}\]

Therefore, for \(h\leq H-h\), we have \(\mathbf{G}_h\in\mathbb{R}^{(|\mathcal{A}|^{m-1}|\mathcal{O}|^m)\times |\mathcal{S}|}\). To make a unified argument, we define \(m\)-step decodability as follows.

We remark that the key difference between our definition and the definition introduced in [16] is that we allow \(h>H-m\). This indicates that after step \(h>H-m\), the process is \((H-h)\)-step decodable, which is reasonable in most real-world examples, since the closer we approach to the end, the clearer our understanding of the current state becomes. Moreover, the new definition admits more convenient notations with almost not impact on the logic of the proof.

In this section, a PSR \(\theta = \{\phi_h, \mathbf{M}_h(o_h,a_h)\}\) is reparametrized by \(\{\mathbf{T}_{h,a}, \mathbf{O}_h\}_{h\in[H],a\in\mathcal{A}}\). More importantly, we have the following equations and properties.

\[\begin{align} \left\{ \begin{aligned} &\mathbf{Y}_h \in \arg\min_{ \mathbf{Y}_h: \mathbf{Y}_h\mathbf{G}_h = 0 } \|\mathbf{G}_h^{\dagger} + \mathbf{Y}_h\|_1\\ &\mathbf{M}_h(o_h,a_h) = \mathbf{G}_{h+1} \mathbf{T}_{h,a_h} \text{diag}(\mathbf{O}_h(o_h|\cdot)) (\mathbf{G}_h^{\dagger} + \mathbf{Y}_h) \\ &\phi_H^{\top} M_H(o_H,a_H) = \mathbf{e}_{(o_H,a_H)} \in \mathbb{R}^{\mathcal{O}\times\mathcal{A}}\\ &\mathcal{Q}_h = (\mathcal{O}\times \mathcal{A})^{\min\{m-1,H-h\} }\\ &\mathcal{Q}_h^A = \mathcal{A}^{ \min\{m-1,H-h\} }\\ & d\leq |\mathcal{O}|^m|\mathcal{A}|^m\\ &r = |\mathcal{S}|\\ &Q_A = |\mathcal{A}|^{m-1}\\ &\gamma = \frac{|\mathcal{S}| + |\mathcal{A}|^{m-1}}{\alpha} = \frac{r + Q_A}{\alpha}. \end{aligned} \right.\label{eqn:POMDP32notation} \end{align}\tag{6}\]

Proof:  We use the notations defined in this section and particularly those in 6 .

The proof differentiates from the proof of 1 in two lemmas.

First, in 2, the meaning of \(x_j\) changes from \(\bar{\hat{\psi}}^{k}(\tau_{h-1})\) to \((\hat{\mathbf{G}}_h^k)^{\dagger} \bar{\hat{\psi}}^{k}(\tau_{h-1})\), and \(w_i^{\top}\) changes from \(\mathbf{m}^*(\omega_h)(\hat{\mathbf{M}}_h^k(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h))\) to \(\mathbf{m}^*(\omega_h)(\hat{\mathbf{M}}_h^k(o_h,a_h) - \mathbf{M}_h^*(o_h,a_h)) \mathbf{G}_h\).

Then, the proof of the upper bound of the term \(I_2\) will become \(Q_A\sqrt{r\lambda}/\gamma^2\) (without dependency on \(d\)), simply because the dimension of bonus term \(\hat{b}^k\) is now \(r = |\mathcal{S}|\) instead of \(d\).

Second, in 4, we need to show that the new bonus term \(\|(\hat{\mathbf{G}}_{h+1}^k)^{\dagger}\bar{\hat{\psi}}(\tau_h)\|_{(\hat{U}_h^k)^{-1}}\) can still be bounded in a similar way, namely, \(\|(\hat{\mathbf{G}}_{h+1}^k)^{\dagger}\bar{\hat{\psi}}(\tau_h)\|_{(\hat{U}_h^k)^{-1}}\leq \tilde{O} \left(\|(\mathbf{G}_{h+1}^*)^{\dagger}\bar{\psi}^*(\tau_h)\|_{(U_h^k)^{-1}} \right)\).

We note that the matrices \(\hat{U}_h^{k}\) and \(U_h^k\) now have new definitions: \[\begin{align} &\hat{U}_h^{k} = \lambda I + (\hat{\mathbf{G}}_{h+1}^k)^{\dagger}\sum_{\tau_h\in\mathcal{D}_h^k} \left[\bar{\hat{\psi}}^{k}(\tau_h) \bar{\hat{\psi}}^k(\tau_h)^{\top} \right ] (\hat{\mathbf{G}}_{h+1}^k)^{\dagger\top}, \\ & U_h^{k} = \lambda I + (\mathbf{G}_{h+1}^*)^{\dagger}\sum_{\tau_h\in\mathcal{D}_h^k} \left[\bar{\psi}^{k}(\tau_h) \bar{\psi}^*(\tau_h)^{\top} \right ] (\mathbf{G}_{h+1}^*)^{\dagger\top}. \end{align}\]

To show such an upper bound, we next provide the key steps.

First, we apply 13 and obtain \[\begin{align} & \left\| (\hat{\mathbf{G}}_{h+1}^k)^{\dagger}\bar{\hat{\psi}}^{k}(\tau_h) \right\|_{\left(\hat{U}_h^{k}\right)^{-1}} \\ &\quad \leq \frac{1}{\sqrt{\lambda}} \left\| (\hat{\mathbf{G}}_{h+1}^k)^{\dagger}\bar{\hat{\psi}}^{k}(\tau_h) - (\mathbf{G}_{h-1}^*)^{\dagger}\bar{\psi}^* (\tau_h) \right\|_2 \\ &\quad \quad + \left(1 + \frac{ \sqrt{r} \sqrt{ \sum_{\tau_h\in\mathcal{D}_h^k } \left\| (\hat{\mathbf{G}}_{h+1}^k)^{\dagger}\bar{\hat{\psi}}^{k}(\tau_h) - (\mathbf{G}_{h+1}^*)^{\dagger}\bar{\psi}^* (\tau_h) \right\|_2^2 } }{ \sqrt{\lambda} } \right) \left\| (\mathbf{G}_{h+1}^*)^{\dagger}\bar{\psi}^* (\tau_h) \right\|_{\left(U_h^{k}\right)^{-1}}. \end{align}\]

For any \(\tau_h\), we can derive \[\begin{align} &\left\| (\hat{\mathbf{G}}_{h+1}^k)^{\dagger} \bar{\hat{\psi}}^{k}(\tau_h) - (\mathbf{G}_{h+1}^*)^{\dagger}\bar{\psi}^* (\tau_h) \right\|_2\\ & = \left\| \left( (\hat{\mathbf{G}}_{h+1}^k)^{\dagger} + \hat{\mathbf{Y}}_{h+1}^k \right) \bar{\hat{\psi}}^{k}(\tau_h) - (\mathbf{G}_{h+1}^*)^{\dagger}\bar{\psi}^* (\tau_h) \right\|_2\\ &\leq \left\| \left( (\hat{\mathbf{G}}_{h+1}^k)^{\dagger} + \hat{\mathbf{Y}}_{h+1}^k \right) (\bar{\hat{\psi}}^{k}(\tau_h) - \bar{\psi}^* (\tau_h) )\right\|_2 + \left\| \left( (\hat{\mathbf{G}}_{h+1}^k)^{\dagger} + \hat{\mathbf{Y}}_{h+1}^k - (\mathbf{G}_{h+1}^{* })^{\dagger} \right) \bar{\psi}^* (\tau_h) \right\|_2\\ &\leq \left\| (\hat{\mathbf{G}}_{h+1}^k)^{\dagger} + \hat{\mathbf{Y}}_{h+1}^k \right\|_1 \left\|\bar{\hat{\psi}}^{k}(\tau_h) - \bar{\psi}^* (\tau_h) \right\|_1 + \left\| \left((\hat{\mathbf{G}}_{h+1}^k)^{\dagger} + \hat{\mathbf{Y}}_{h+1}^k \right) \left( I - \hat{\mathbf{G}}_{h+1}^k (\mathbf{G}_{h+1}^{* })^{\dagger} \right) \bar{\psi}^* (\tau_h) \right\|_2 \\ &\leq \frac{1}{\gamma}\left\|\bar{\hat{\psi}}^{k}(\tau_h) - \bar{\psi}^* (\tau_h) \right\|_1 + \left\| \left((\hat{\mathbf{G}}_{h+1}^k)^{\dagger} + \hat{\mathbf{Y}}_{h+1}^k \right) \left( \mathbf{G}_{h+1}^* - \hat{\mathbf{G}}_{h+1}^k \right) (\mathbf{G}_{h+1}^{* })^{\dagger} \bar{\psi}^* (\tau_h) \right\|_2, \end{align}\]

where the last inequality is due to the \(m\)-step decodability.

The first term above has been bounded in 4. Then, we bound the second term by the Cauchy’s inequality as follows: \[\begin{align} &\left\| \left((\hat{\mathbf{G}}_{h+1}^k)^{\dagger} + \hat{\mathbf{Y}}_{h+1}^k \right) \left( \mathbf{G}_{h+1}^* - \hat{\mathbf{G}}_{h+1}^k \right) (\mathbf{G}_{h+1}^{* })^{\dagger} \bar{\psi}^* (\tau_h) \right\|_2\\ &\quad \leq \left\| (\mathbf{G}_{h+1}^{* })^{\dagger} \bar{\psi}^* (\tau_h) \right\|_{\left( \bar{U}_h^{k}\right)^{-1}} \\ &\quad \quad \times \sqrt{ rQ_A\lambda/\gamma + \sum_{t<k}\mathbb{E}_{\tau_h\sim \pi^t} \left\| \left((\hat{\mathbf{G}}_{h+1}^k)^{\dagger} + \hat{\mathbf{Y}}_{h+1}^k \right) \left( \mathbf{G}_{h+1}^* - \hat{\mathbf{G}}_{h+1}^k \right) (\mathbf{G}_{h+1}^{* })^{\dagger} \bar{\psi}^* (\tau_h) \right\|_2^2 }\\ &\quad \leq 5\left\| (\mathbf{G}_{h+1}^{* })^{\dagger} \bar{\psi}^* (\tau_h) \right\|_{\left(U_h^{k}\right)^{-1}} \sqrt{ rQ_A\lambda/\gamma + |\mathcal{A}|Q_A^2\beta/\gamma^2 }, \end{align}\]

where \(\bar{U}_h^k\) denotes the expected matrix of \(U_h^k\), and has the following form \[\begin{align} \bar{U}_h^k = \mathbb{E}[U_h^k] = \lambda I + (\mathbf{G}_{h+1}^*)^{\dagger} \sum_{t<k}\mathbb{E}_{\tau_h\sim\pi^t}\left[ \bar{\psi}^{k}(\tau_h) \bar{\psi}^*(\tau_h)^{\top}\right] (\mathbf{G}_{h+1}^*)^{\dagger\top}. \end{align}\]

In particular, the last inequality above follows from two facts. First, we use Lemma 39 in [71] to transform the expected matrix to the empirical matrix \(U_h^k\). Second, we use Lemma 12 to upper-bound \(\mathbb{E}_{\tau_h\sim \pi^t} \left\| \left((\hat{\mathbf{G}}_{h+1}^k)^{\dagger} + \hat{\mathbf{Y}}_{h+1}^k \right) \left( \mathbf{G}_{h+1}^* - \hat{\mathbf{G}}_{h+1}^k \right) (\mathbf{G}_{h+1}^{* })^{\dagger} \bar{\psi}^* (\tau_h) \right\|_2^2\) by the Hellinger squared distance between the estimated model \(\hat{\theta}^k\) and the true model \(\theta^*\). Then, due to the estimation guarantee (1), we have \[\begin{align} \sum_{t<k} &\mathbb{E}_{\tau_h\sim \mathbb{P}_{\theta^*}^{\pi^t} } \left\| \left((\hat{\mathbf{G}}_{h+1}^k)^{\dagger} + \hat{\mathbf{Y}}_{h+1}^k \right) \left( \mathbf{G}_{h+1}^* - \hat{\mathbf{G}}_{h+1}^k \right) (\mathbf{G}_{h+1}^{* })^{\dagger} \bar{\psi}^* (\tau_h) \right\|_2^2 \\ & \leq \frac{1}{\gamma^2} \sum_{t<k} \mathbb{E}_{\tau_h\sim \mathbb{P}_{\theta^*}^{\pi^t} } \left\| \left( \mathbf{G}_{h+1}^* - \hat{\mathbf{G}}_{h+1}^k \right) (\mathbf{G}_{h+1}^{* })^{\dagger} \bar{\psi}^* (\tau_h) \right\|_1 ^2 \\ & = \frac{1}{\gamma^2} \sum_{t<k} \mathbb{E}_{\tau_h\sim \mathbb{P}_{\theta^*}^{\pi^t} } \left( \sum_{\mathbf{o}\in\mathcal{Q}_{h+1} ,\mathbf{a}\in\mathcal{Q}_{h+1}^A } \sum_{s\in\mathcal{S} } \left( \mathbb{P}_{\hat{\theta}^k}(\mathbf{o}| \mathbf{a}, s_{h+1} =s ) - \mathbb{P}_{\theta^* }(\mathbf{o}| \mathbf{a}, s_{h+1} =s ) \right) \mathbb{P}_{\theta^*}(s|\tau_h) \right)^2 \\ &\leq \frac{Q_A^2}{\gamma^2} \sum_{t<k} \mathbb{E}_{\tau_h\sim \mathbb{P}_{\theta^*}^{\pi^t} } \left( \mathbb{E}_{s\sim \mathbb{P}_{\theta^*}(\cdot|\tau_h) } \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\hat{\theta}^k}^{\mathtt{u}_{\mathcal{Q}_{h+1}^A}}(\cdot | s_{h+1} =s ) , \mathbb{P}_{\theta^* }^{\mathtt{u}_{\mathcal{Q}_{h+1}^A}}(\cdot| s_{h+1} =s ) \right) \right)^2 \\ &\leq \frac{Q_A^2}{\gamma^2} \sum_{t<k} \mathbb{E}_{\tau_h\sim \mathbb{P}_{\theta^*}^{\pi^t} } \mathbb{E}_{s\sim \mathbb{P}_{\theta^*}(\cdot|\tau_h) } \mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\hat{\theta}^k}^{\mathtt{u}_{\mathcal{Q}_{h+1}^A}}(\cdot | s_{h+1} =s ) , \mathbb{P}_{\theta^* }^{\mathtt{u}_{\mathcal{Q}_{h+1}^A}}(\cdot| s_{h+1} =s ) \right)\\ &\lesssim \frac{|\mathcal{A}|Q_A^2}{\gamma^2} \sum_{\pi\in\mathcal{D^k} } \mathtt{D}_{\mathtt{H}}^2 \left( \mathbb{P}_{\hat{\theta}^k}^{\pi} , \mathbb{P}_{\theta^* }^{\pi} \right) = \tilde{O}(\beta). \end{align}\]

Combining above key steps, we obtain that \[\begin{align} &\mathbb{E}_{\tau_h\sim\pi} \left\| (\hat{\mathbf{G}}_{h+1}^k)^{\dagger}\bar{\hat{\psi}}^{k}(\tau_h) \right\|_{\left(\hat{U}_h^{k}\right)^{-1}} \\ &\leq \text{poly}(r, |\mathcal{A}|,Q_A, 1/\gamma)\left( \mathbb{E}_{\tau_h\sim\pi} \|(\mathbf{G}_{h+1}^*)^{\dagger}\bar{\psi}^*(\tau_h)\|_{(U_h^k)^{-1}} + \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi}(\tau_H) , \mathbb{P}_{\hat{\theta}^k}^{\pi}(\tau_H)\right) \right). \end{align}\]

The rest of the proof follow the same argument as that for 1.

 \(\blacksquare\)

10 Proof of 2 (for Offline PSR-LCB)↩︎

In this section, we present the full analysis for the offline algorithm PSR-LCB to show 2. In particular, the proof of 2 consists of three main steps. Step 1: We provide the offline estimation guarantee for the estimated model \(\hat{\theta}\). Step 2: Building up on the first step, we are able to show that \(V_{\hat{\theta},\hat{b}}^{\pi}\) is a valid upper bound of the total variation distance between \(\hat{\theta}\) and \(\theta^*\). Hence, \(V_{\hat{\theta}, R}^{\pi} - V_{\hat{\theta},\hat{b}}^{\pi}\) is a valid LCB for the true value \(V_{\theta^*,R}^{\pi}\). Step 3: We translate \(V_{\hat{\theta},\hat{b}}^{\pi}\) to the ground-truth prediction feature \(\bar{\psi}^*(\tau_h)\). Finally, we show that the \(V_{\hat{\theta},\hat{b}}^{\pi}\) scales in the order of \(1/\sqrt{K}\), which characterizes the performance of PSR-LCB and completes the proof.

We first introduce the following definitions for good events.

Good Events. Recall that \(\Theta_{\min} = \left\{\theta: \forall h, \tau_h \in \mathcal{D}_h,~~ \mathbb{P}_{\theta}^{\pi^b}(\tau_h) \geq p_{\min} \right\}\), where \(p_{\min}\leq \frac{\delta}{KH(|\mathcal{O}||\mathcal{A}|)^H}\). Let \(\varepsilon\leq\frac{p_{\min}}{KH}\). Analogous to the proof for online learning, we introduce three events defined as follows. \[\begin{align} \mathcal{E}_{\omega}^{o} &= \left\{ \forall \theta\in\Theta_{\min}, ~~ \sum_h \sum_{ \tau_h \in \mathcal{D}_h } \mathtt{D}_{\mathtt{TV}}^2 \left( \mathbb{P}_{\theta}^{\pi^b}(\omega_h|\tau_h ), \mathbb{P}_{\theta^*}^{\pi^b}(\omega_h|\tau_h ) \right) \right. \\ &\quad \left. \leq 6\sum_h\sum_{ \tau_H \in\mathcal{D}_h }\log\frac{\mathbb{P}_{\theta^*}^{ \pi^b }(\tau_H)}{\mathbb{P}_{ \theta}^{ \pi^b }(\tau_H)} + 31\log\frac{3K\left|\bar{\Theta}_{\varepsilon}\right|}{\delta} \right\}, \\ \mathcal{E}_{\pi}^{o} &= \left\{ \forall \theta\in\Theta, K,~~ \mathtt{D}_{\mathtt{H}}^2 ( \mathbb{P}_{\theta}^{\pi^b}(\tau_H) , \mathbb{P}_{\theta^*}^{\pi} (\tau_H) ) \leq \sum_{ \tau_H \in\mathcal{D}} \log\frac{ \mathbb{P}_{ \theta^*}^{\pi^b}(\tau_H) }{\mathbb{P}_{ \theta }^{\pi^b}(\tau_H)} + 2\log\frac{3K|\bar{\Theta}_{\varepsilon}|}{\delta} \right\}, \\ \mathcal{E}_{\min}^{o} &= \bigg\{ \forall h, \tau_h \in\mathcal{D}_h ,~~ \mathbb{P}_{\theta^*}^{\pi^b}(\tau_h) \geq p_{\min}\bigg\}. \end{align}\]

Following the proof steps similar to those in 8, we can conclude that \(\mathcal{E}^o:= \mathcal{E}_{\omega}^{o}\cap \mathcal{E}_{\pi}^{o} \cap \mathcal{E}_{\min}^{o}\) occurs with probability at least \(1-\delta\).

10.1 Step 1: Estimation Guarantee↩︎

Proof:  The proof follows the steps similar to those in 1, except that we replace the exploration policies \(\nu_h(\pi^k,\mathtt{u}_{\mathcal{Q}_{h-1}^{\exp}})\) by the behavior policy.

 \(\blacksquare\)

10.2 Step 2: UCB for Total Variation Distance and LCB for Value Function↩︎

The following lemma provides an explicit upper bound on the total variation distance between the estimated model and the true model.

Proof:  Similarly to the analysis in that for 9, we index \(\omega_{h-1} = (o_{h},a_h,\ldots,o_H,a_H)\) by \(i\), and \(\tau_{h-1}\) by \(j\). In addition, we denote \(\hat{\mathbf{m}}(\omega_{h})^{\top}\left(\hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right)\) by \(w_{i}^{\top}\), denote \(\bar{ \psi}^* (\tau_{h-1})\) by \(x_{j}\), and denote \(\pi(\omega_{h-1}|\tau_{h-1})\) by \(\pi_{i|j}\). Then, we have \[\begin{align} \sum_{\tau_H}& \left| \hat{\mathbf{m}} (\omega_h)^{\top} \left(\hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right) \psi^* (\tau_{h-1}) \right| \pi(\tau_H) \\ & = \sum_{i }\sum_{j }|w_{i}^{\top}x_{j}|\pi_{i|j}\mathbb{P}_{ \theta^* }^{\pi}(j)\\ & = \sum_{j } \sum_{i } (\pi_{i|j}\cdot\mathtt{sgn}(w_{i}^{\top}x_{j})\cdot w_{i})^{\top} x_{j}\cdot\mathbb{P}_{\theta^*}^{\pi}(j)\\ & = \sum_{j} \left( \sum_{i} \pi_{i|j} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j} )\cdot w_{i} \right)^{\top} x_{j}\cdot\mathbb{P}_{ \theta^* }^{\pi}(j)\\ &\leq \mathop{\mathbb{E}}_{j\sim\mathbb{P}_{\theta^*}^{\pi}}\left[ \big\| x_j \big\|_{\Lambda_{h-1}^{-1}} \sqrt{ \left\|\sum_{i} \pi_{i|j} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j} )\cdot w_{i} \right\|_{ \Lambda_{h-1} }^2 } \right], \end{align}\] where \(\Lambda_{h-1} = \lambda_0 I + \frac{K}{H} \mathbb{E}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^b}} \left[ \bar{\psi}^*(\tau_{h-1}) \bar{\psi}^*(\tau_{h-1})^{\top} \right]\) and \(\lambda_0\) will be determined later. We fix \(\tau_{h-1} = j_0\) and aim to analyze the coefficient of \(\|x_{j_0}\|_{\Lambda_{h-1}^{-1}}\). We have \[\begin{align} &\left\| \sum_{i} \pi_{i|j_0} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j_0} )\cdot w_{i} \right\|_{\Lambda_{h-1} }^2\\ &\quad = \underbrace{\lambda_0 \left\| \sum_{i} \pi_{i|j_0} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j_0} )\cdot w_{i} \right\|_2^2}_{I_1} + \underbrace{\frac{K}{H} \mathop{\mathbb{E}}_{j\sim\mathbb{P}_{\theta^*}^{\pi^b}} \left[ \left(\sum_{i} \pi_{i|j_0} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j_0} )\cdot w_{i}\right)^{\top} x_{j} \right]^2 }_{I_2}. \end{align}\] For the first term \(I_1\), we have \[\begin{align} \sqrt{I_1} & = \sqrt{\lambda_0} \max_{x\in\mathbb{R}^{d_{h-1}}: \|x\|_2=1} \left|\sum_{i} \pi_{i|j_0}\mathtt{sgn}(w_i^{\top}x_{j_0})w_i^{\top} x\right|\\ &\quad \leq \sqrt{\lambda_0}\max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{\omega_{h-1}} \left| \mathbf{m}^*(\omega_h)^{\top} \left( \hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right) x \right| \pi(\omega_{h-1}|j_0)\\ &\quad \leq \sqrt{\lambda_0}\max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{\omega_{h-1}}\left| \hat{\mathbf{m}} (\omega_h)^{\top} \hat{\mathbf{M}}_h (o_h,a_h) x \right| \pi(\omega_{h-1}|j_0) \\ &\quad\quad + \sqrt{\lambda_0}\max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{\omega_{h-1}}\left| \hat{\mathbf{m}}(\omega_h)^{\top} \mathbf{M}_h^*(o_h,a_h) x \right| \pi(\omega_{h-1}|j_0) \\ & \overset{(a)}\leq \frac{\sqrt{d\lambda_0}}{\gamma} + \frac{\sqrt{\lambda_0}}{\gamma}\max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{o_h,a_h}\left\| \hat{\mathbf{M}}_h (o_h,a_h) x \right\|_1 \pi(a_h|o_h,j_0) \\ &\overset{(b)}\leq \frac{2Q_A\sqrt{d\lambda_0}}{\gamma^2}, \end{align}\] where \((a)\) follows from 1, and \((b)\) follows from 2.

For the second term \(I_2\), we have \[\begin{align} I_2 &\leq \frac{K}{H} \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^b}} \left[\left( \sum_{\omega_{h-1}} \left| \hat{\mathbf{m}} (\omega_{h})^{\top}\left(\hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right)\bar{\psi}^* (\tau_{h-1}) \right| \pi(\omega_{h-1}|j_0) \right)^2 \right] \\ & \leq \frac{K}{H} \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^b}} \left[ \left( \sum_{\omega_{h-1}} \left| \hat{\mathbf{m}}(\omega_h)^{\top} \left( \hat{\mathbf{M}}_h (o_h,a_h) \bar{ \hat{\psi}} (\tau_{h-1}) - \mathbf{M}_h^*(o_h,a_h) \bar{\psi}^*(\tau_{h-1}) \right) \right| \pi(\omega_{h-1}|j_0) \right. \right. \\ & + \left. \left. \sum_{\omega_{h-1}}\left| \hat{\mathbf{m}} (\omega_h)^{\top} \hat{\mathbf{M}}_h(o_h,a_h) \left(\bar{\hat{\psi}} (\tau_{h-1}) - \bar{\psi}^*(\tau_{h-1}) \right) \right| \pi(\omega_{h-1}|j_0) \right)^2 \right] \\ & \overset{(a)}\leq \frac{K}{H} \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^b} } \left[ \left( \frac{1}{\gamma} \sum_{o_h,a_h} \left\| \mathbb{P}_{ \hat{\theta} } (o_h|\tau_{h-1})\bar{\hat{\psi}}_{h} (\tau_{h}) - \mathbb{P}_{\theta}(o_h|\tau_{h-1}) \bar{\psi}_{h}^*(\tau_{h}) \right\|_1 \pi(a_h|o_h,j_0) \right. \right.\\ &+ \left. \left. \frac{1}{\gamma} \left\| \bar{\hat{\psi}} (\tau_{h-1}) - \bar{\psi}^*(\tau_{h-1}) \right\|_1 \right)^2 \right] \\ & = \frac{K}{H\gamma^2} \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^b}} \left[\left( \sum_{o_h,a_h} \sum_{\ell=1}^{|\mathcal{Q}_h|} \left| \mathbb{P}_{ \hat{\theta} }(\mathbf{o}_h^{\ell},o_h|\tau_{h-1},a_h,\mathbf{a}_h^{\ell}) - \mathbb{P}_{ \theta^* }(\mathbf{o}_h^{\ell},o_h|\tau_{h-1},a_h,\mathbf{a}_h^{\ell}) \right|\pi(a_h|o_h,j_0) \right. \right.\\ &+ \left. \left. \sum_{\ell=1}^{|\mathcal{Q}_{h-1}|} \left| \mathbb{P}_{ \hat{\theta} }(\mathbf{o}_{h-1}^{\ell} | \tau_{h-1}, \mathbf{a}_{h-1}^{\ell} ) - \mathbb{P}_{ \theta^* } (\mathbf{o}_{h-1}^{\ell} | \tau_{h-1}, \mathbf{a}_{h-1}^{\ell} ) \right| \right)^2 \right] \\ &\leq \frac{K}{H\gamma^2} \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^b}} \left[ \left( \;\sum_{\mathbf{a}_{h-1}\in \mathcal{Q}_h^{\exp} } \sum_{\omega_{h-1}^o }\left|\mathbb{P}_{ \hat{\theta} }(\omega^o_{h-1}|\tau_{h-1}, \mathbf{a}_{h-1} ) - \mathbb{P}_{\theta^*}(\omega^o_{h-1}|\tau_{h-1}, \mathbf{a}_{h-1} ) \right| \right)^2 \right] \\ &\overset{(b)}\leq \frac{ K }{ H \iota^2 \gamma^2} \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^b}} \left[\mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\hat{\theta} }^{ \pi^b }(\omega_{h-1} | \tau_{h-1} ) , \mathbb{P}_{ \theta^* }^{ \pi^b } (\omega_{h-1} | \tau_{h-1} ) \right) \right]\\ &\overset{(c)}\leq \frac{64K}{H\iota^2\gamma^2}\mathtt{D}_{\mathtt{H}}^2 \left(\mathbb{P}_{\hat{\theta}}^{\pi^b}(\tau_H), \mathbb{P}_{\theta^*}^{\pi^b}(\tau_H) \right)\\ &\overset{(d)}\lesssim \frac{ \hat{\beta}}{H\iota^2\gamma^2}, \end{align}\] where \((a)\) follows from 1 and 5 , \((b)\) follows because \(\pi^b(\mathbf{a}_{h-1})\geq\iota\) for all \(\mathbf{a}_{h-1}\in\mathcal{Q}_{h-1}^{\exp}\), \((c)\) follows from 12, and \((d)\) follows from 7.

Thus, by choosing \(\lambda_0 = \frac{\gamma^4}{4Q_A^2d}\), we have \[\begin{align} \sum_{\tau_H}& \left| \hat{\mathbf{m}}(\omega_h)^{\top} \left(\hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right) \psi^* (\tau_{h-1}) \right| \pi(\tau_H) \\ &\lesssim \sqrt{\frac{\hat{\beta}}{H\iota^2\gamma^2}} \mathop{\mathbb{E}}_{\tau_{h-1}\sim \mathbb{P}_{ \theta^* }^{\pi}}\left[ \left\|\bar{\psi}^* (\tau_{h-1})\right\|_{ \Lambda_{h-1}^{-1}} \right]. \end{align}\] Following from 1, we conclude that \[\begin{align} \mathtt{D}_{\mathtt{TV}}&\left( \mathbb{P}_{\theta^*}^{\pi}(\tau_H) , \mathbb{P}_{\hat{\theta}}^{\pi}(\tau_H)\right) \lesssim \min \left\{\sqrt{\frac{\hat{\beta}}{H\iota^2\gamma^2}} \sum_{h=1}^H \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi}} \left[ \; \left\|\bar{\psi}^* (\tau_{h-1})\right\|_{ \Lambda_{h-1}^{-1}} \right] ,2 \right\}. \end{align}\]

 \(\blacksquare\)

Following from 1, we provide the following lemma that upper bounds the estimation error by the bonus function.

Proof:  We index \(\omega_{h-1} = (o_{h},a_h,\ldots,o_H,a_H)\) by \(i\), and \(\tau_{h-1}\) by \(j\). In addition, we denote
\(\mathbf{m}^*(\omega_{h})^{\top}\left(\hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right)\) by \(w_{i}^{\top}\), \(\bar{\hat{\psi}}_{h-1} (\tau_{h-1}) = \frac{ \hat{\psi}_{h-1}(\tau_{h-1})}{\hat{\phi}_{h-1}^{\top}\hat{\psi}_{h-1} (\tau_{h-1})}\) by \(x_{j}\), and \(\pi(\omega_{h-1}|\tau_{h-1})\) by \(\pi_{i|j}\). Then, we have \[\begin{align} \sum_{\tau_H} &\left| \mathbf{m}^* (\omega_h)^{\top} \left(\hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right) \hat{\psi}_{h-1} (\tau_{h-1}) \right| \pi(\tau_H) \\ & = \sum_{i }\sum_{j }|w_{i}^{\top}x_{j}|\pi_{i|j}\mathbb{P}_{ \hat{\theta} }^{\pi}(j)\\ & = \sum_{j } \sum_{i } (\pi_{i|j}\cdot\mathtt{sgn}(w_{i}^{\top}x_{j})\cdot w_{i})^{\top} x_{j}\cdot\mathbb{P}_{\theta}^{\pi}(j)\\ & = \sum_{j} \left( \sum_{i} \pi_{i|j} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j} )\cdot w_{i} \right)^{\top} x_{j}\cdot\mathbb{P}_{ \hat{\theta} }^{\pi}(j)\\ &\leq \mathop{\mathbb{E}}_{j\sim\mathbb{P}_{ \hat{\theta} }^{\pi}}\left[ \big\| x_j \big\|_{\left(\hat{U}_{h-1}^{k}\right)^{-1}} \sqrt{ \left\|\sum_{i} \pi_{i|j} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j} )\cdot w_{i} \right\|_{ \hat{U}_{h-1} }^2 } \right]. \end{align}\]

We fix an index \(j = j_0\), and aim to analyze \(\left\| \sum_{i} \pi_{i|j_0} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j_0} )\cdot w_{i} \right\|_{\hat{U}_{h-1}^{k}}^2\), which can be written as \[\begin{align} &\left\| \sum_{i} \pi_{i|j_0} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j_0} )\cdot w_{i} \right\|_{\hat{U}_{h-1}^{k}}^2\\ &\quad = \underbrace{\hat{\lambda} \left\| \sum_{i} \pi_{i|j_0} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j_0} )\cdot w_{i} \right\|_2^2}_{I_1} + \underbrace{\sum_{ j\in\mathcal{D}_{h-1} } \left[ \left(\sum_{i} \pi_{i|j_0} \cdot \mathtt{sgn}(w_{i}^{\top}x_{j_0} )\cdot w_{i}\right)^{\top} x_{j} \right]^2}_{I_2} . \end{align}\]

For the first term \(I_1\), we have \[\begin{align} \sqrt{I_1} & = \sqrt{\hat{\lambda } } \max_{x\in\mathbb{R}^{d_{h-1}}: \|x\|_2=1} \left|\sum_{i} \pi_{i|j_0}\mathtt{sgn}(w_i^{\top}x_{j_0})w_i^{\top} x\right|\\ & \leq \sqrt{\hat{\lambda } } \max_{x\in\mathbb{R}^{d_{h-1}}:\|x\|_2=1} \sum_{\omega_{h-1}} \left| \mathbf{m}^*(\omega_h)^{\top} \left( \hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right) x \right| \pi(\omega_{h-1}|j_0)\\ & \leq \sqrt{\hat{\lambda } } \max_{x:\|x\|_2=1} \sum_{\omega_{h-1}}\left|\mathbf{m}^{*}(\omega_h)^{\top} \hat{\mathbf{M}}_h (o_h,a_h) x \right| \pi(\omega_{h-1}|j_0) \\ &\quad + \sqrt{\hat{\lambda } } \max_{x:\|x\|_2=1} \sum_{\omega_{h-1}}\left|\mathbf{m}^*(\omega_h)^{\top} \mathbf{M}_h^*(o_h,a_h) x \right| \pi(\omega_{h-1}|j_0) \\ &\leq \frac{\sqrt{\hat{\lambda } }}{\gamma}\max_{x:\|x\|_2=1} \sum_{o_h,a_h}\left\| \hat{\mathbf{M}}_h (o_h,a_h) x \right\|_1 \pi(a_h|o_h,j_0) + \frac{\sqrt{\hat{\lambda } }}{\gamma} \| x\|_1 \\ &\leq \frac{ 2Q_A\sqrt{d\hat{\lambda}} }{\gamma^2}. \end{align}\]

For the second term \(I_2\), we have \[\begin{align} I_2 & \leq \sum_{ \tau_{h-1}\in \mathcal{D}_{h-1} } \left( \sum_{\omega_{h-1}} \left| \mathbf{m}^* (\omega_{h})^{\top}\left(\hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right)\bar{\hat{\psi}} (\tau_{h-1}) \right| \pi(\omega_{h-1}|j_0) \right)^2 \\ & \leq \sum_{ \tau_{h-1}\in \mathcal{D}_{h-1} } \left( \sum_{\omega_{h-1}} \left|\mathbf{m}^*(\omega_h)^{\top} \left( \hat{\mathbf{M}}_h (o_h,a_h) \bar{\hat{\psi}} (\tau_{h-1}) - \mathbf{M}_h^*(o_h,a_h) \bar{\psi}^*(\tau_{h-1}) \right) \right| \pi(\omega_{h-1}|j_0) \right. \\ & + \left. \sum_{\omega_{h-1}}\left| \mathbf{m}^*(\omega_h)^{\top} \mathbf{M}_h^*(o_h,a_h) \left( \bar{\hat{\psi}} (\tau_{h-1}) - \bar{\psi}^*(\tau_{h-1}) \right) \right| \pi(\omega_{h-1}|j_0) \right)^2 \\ & \overset{(a)}\leq \sum_{\tau_{h-1}\in \mathcal{D}_{h-1} } \left( \frac{1}{\gamma} \sum_{o_h,a_h} \left\| \mathbb{P}_{ \hat{\theta} } (o_h|\tau_{h-1})\bar{\hat{\psi}} (\tau_{h}) - \mathbb{P}_{\theta}(o_h|\tau_{h-1}) \bar{\psi}^*(\tau_{h}) \right\|_1 \pi(a_h|o_h,j_0) \right. \\ &+ \left. \frac{1}{\gamma} \left\| \bar{\hat{\psi}} (\tau_{h-1}) - \bar{\psi}^*(\tau_{h-1}) \right\|_1 \right)^2 \\ & = \frac{1}{\gamma^2} \sum_{\tau_{h-1}\in \mathcal{D}_{h-1} } \left( \sum_{o_h,a_h} \sum_{\ell=1}^{|\mathcal{Q}_h|} \left| \mathbb{P}_{ \hat{\theta} }(\mathbf{o}_h^{\ell},o_h|\tau_{h-1},a_h,\mathbf{a}_h^{\ell}) - \mathbb{P}_{ \theta^* }(\mathbf{o}_h^{\ell},o_h|\tau_{h-1},a_h,\mathbf{a}_h^{\ell}) \right|\pi(a_h|o_h,j_0) \right. \\ &+ \left. \sum_{\ell=1}^{|\mathcal{Q}_{h-1}|} \left| \mathbb{P}_{ \hat{\theta} }(\mathbf{o}_{h-1}^{\ell} | \tau_{h-1}, \mathbf{a}_{h-1}^{\ell} ) - \mathbb{P}_{ \theta^* } (\mathbf{o}_{h-1}^{\ell} | \tau_{h-1}, \mathbf{a}_{h-1}^{\ell} ) \right| \right)^2 \\ &\leq \frac{1}{\gamma^2} \sum_{\tau_{h-1}\in \mathcal{D}_{h-1} } \left( \;\sum_{\mathbf{a}_{h-1}\in \mathcal{Q}_{h-1}^{\exp} } \sum_{\omega_{h-1}^o }\left|\mathbb{P}_{ \hat{\theta} }(\omega^o_{h-1}|\tau_{h-1}, \mathbf{a}_{h-1} ) - \mathbb{P}_{\theta^*}(\omega^o_{h-1}|\tau_{h-1}, \mathbf{a}_{h-1} ) \right| \right)^2 \\ &\overset{(b)}\leq \frac{ 1 }{ \iota^2 \gamma^2} \sum_{\tau_{h-1}\in\mathcal{D}_{h-1} } \mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\hat{\theta} }^{ \pi^b }(\omega_{h-1} | \tau_{h-1} ) , \mathbb{P}_{ \theta^* }^{ \pi^b } (\omega_{h-1} | \tau_{h-1} ) \right), \end{align}\] where \((a)\) follows from 1, and \((b)\) follows from the condition \(\pi^b(\mathbf{a}_h)\geq \iota\) for any \(\mathbf{a}_h\in\mathcal{Q}_{h-1}^{\exp}\).

Thus, we conclude that \[\begin{align} \sum_{\tau_H}& \left| \mathbf{m}^*(\omega_h)^{\top} \left(\hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right) \hat{\psi} (\tau_{h-1}) \right| \pi(\tau_H) \\ &\leq \mathbb{E}_{\tau_{h-1}\sim \mathbb{P}_{ \hat{\theta} }^{\pi}}\left[ \hat{\alpha}_{h-1} \left\|\bar{\hat{\psi}} (\tau_{h-1})\right\|_{ \left( \hat{U}_{h-1} \right)^{-1}} \right], \end{align}\] where \[\begin{align} ( \hat{\alpha}_{h-1} )^2 = \frac{ 4\hat{\lambda}Q_A^2d}{\gamma^4} + \frac{1 }{ \iota^2 \gamma^2} \sum_{\tau_{h-1} \in\mathcal{D}_{h-1} } \mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{ \hat{\theta} }^{\pi^b}(\omega_{h-1} | \tau_{h-1} ) , \mathbb{P}_{ \theta^* }^{ \pi^b } ( \omega_{h-1} | \tau_{h-1} ) \right). \end{align}\]

 \(\blacksquare\)

Proof:  By the definition of the total variation distance, we have \[\begin{align} \left| V_{\hat{\theta} , R }^{\pi} - V_{\theta^* , R}^{\pi} \right| &\leq \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\hat{\theta}}^{\pi}, \mathbb{P}_{\theta^*}^{\pi} \right)\\ &\overset{(a)} \leq \sum_{h=1}^H \sum_{\tau_H} \left| \mathbf{m}^*(\omega_{h})^{\top} \left(\hat{\mathbf{M}}_h (o_h,a_h) - \mathbf{M}_h^*(o_h,a_h) \right) \hat{\psi} (\tau_{h-1}) \right| \pi(\tau_H)\\ &\overset{(b)} \leq \min\left\{ \sum_{h=1}^H \mathbb{E}_{\tau_{h-1}\sim \mathbb{P}_{ \hat{\theta} }^{\pi}}\left[ \hat{\alpha}_{h-1} \left\|\bar{\hat{\psi}} (\tau_{h-1})\right\|_{ \left( \hat{U}_{h-1} \right)^{-1}} \right] , 1 \right\}\\ &\overset{(c)} \leq \min\left\{ \sqrt{\sum_{h=1}^H \hat{\alpha}_{h-1}^2}\sqrt{ \sum_{h = 0}^{H-1} \left\| \bar{\hat{\psi}}_h (\tau_h) \right\|^2_{ (\hat{U}_h )^{-1}} } , 1 \right \}, \end{align}\] where \((a)\) follows from 1, \((b)\) follows from 9 and because \(R(\tau_H)\in[0,1]\), and \((c)\) follows from the Cauchy’s inequality.

By 7, we further have \[\begin{align} \sum_{h=1}^H \hat{\alpha}_{h-1}^2 & \leq \frac{ 4\hat{\lambda} H Q_A^2d }{ \gamma^4 } + \frac{1}{ \iota^2 \gamma^2} \sum_h\sum_{ \tau_{h-1}\in\mathcal{D}_{h-1} } \mathtt{D}_{\mathtt{TV}}^2\left( \mathbb{P}_{\hat{\theta}^k }^{ \pi^b }(\omega_{h-1} | \tau_{h-1} ) , \mathbb{P}_{ \theta^* }^{ \pi^b } (\omega_{h-1} | \tau_{h-1} ) \right) \\ &\leq \frac{ 4\lambda H Q_A^2d }{\gamma^4} + \frac{ 7 \hat{\beta} }{ \iota^2 \gamma^2}, \end{align}\] which concludes the proof.

 \(\blacksquare\)

10.3 Step 3: Relationship between Empirical Bonus and Ground-Truth Bonus↩︎

Proof:  Recall that \[\begin{align} \hat{U}_h = \hat{\lambda} I + \sum_{\tau_h\in\mathcal{D}_h} \bar{\hat{\psi}}(\tau_h)\bar{\hat{\psi}}(\tau_h)^{\top}. \end{align}\] We define the ground-truth counterpart of \(\hat{U}_h\) as follows: \[\begin{align} U_h = \hat{\lambda} I + \sum_{\tau_h\in\mathcal{D}_h} \bar{ \psi }(\tau_h)\bar{ \psi }(\tau_h)^{\top}. \end{align}\]

Then, by 13, we have \[\begin{align} &\sqrt{ \sum_h \left\|\bar{\hat{\psi}} (\tau_h) \right\|^2_{\left( \hat{U}_h \right)^{-1}} }\\ &\quad \leq \sum_h \left\|\bar{\hat{\psi}}_h (\tau_h) \right\|_{\left( \hat{U}_h \right)^{-1}} \\ &\quad \leq \frac{1}{\sqrt{\hat{\lambda}}} \sum_h \left\| \bar{\hat{\psi}} (\tau_h) - \bar{\psi}_h^* (\tau_h) \right\|_2 + \sum_h \left(1 + \frac{ \sqrt{r} \sqrt{ \sum_{\tau_h\in\mathcal{D}_h } \left\| \bar{\hat{\psi}}(\tau_h) - \bar{\psi}^* (\tau_h) \right\|_2^2 } }{ \sqrt{\hat{\lambda}} } \right) \left\|\bar{\psi}^* (\tau_h) \right\|_{\left(U_h\right)^{-1}} . \end{align}\]

Furthermore, note that \[\begin{align} \left\| \bar{\hat{\psi}} (\tau_h) - \bar{\psi}^* (\tau_h) \right\|_2 & \leq \left\| \bar{\hat{\psi}}_h (\tau_h) - \bar{\psi}_h^* (\tau_h) \right\|_1 \\ & \leq \frac{1}{\iota} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{ \hat{\theta} }^{ \pi^b }(\omega_h|\tau_h), \mathbb{P}_{ \theta^* }^{ \pi^b }(\omega_h|\tau_h) \right), \end{align}\] where the second inequality follows because \(\pi^b(\mathbf{a}_h) \geq \iota\) for all \(\mathbf{a}_h\in\mathcal{Q}_h^{\exp}\). By 7, we conclude that \[\begin{align} &\sqrt{ \sum_h \left\|\bar{\hat{\psi}}_h (\tau_h) \right\|^2_{\left(\hat{U}_h \right)^{-1}} }\\ &\quad \leq \frac{1}{\sqrt{\hat{\lambda}}} \sum_h \left\| \bar{\hat{\psi}}_h (\tau_h) - \bar{\psi}_h^* (\tau_h) \right\|_2 + \left(1 + \frac{ \sqrt{7r\hat{\beta}} }{ \iota \sqrt{\hat{\lambda}} } \right) \sum_h \left\|\bar{\psi}_h^* (\tau_h) \right\|_{\left(U_h^{k}\right)^{-1}}. \end{align}\]

For the first term, following an argument similar to those in 4 and taking the expectation, we have \[\begin{align} \sum_h & \mathbb{E}_{\tau_h\sim \mathbb{P}_{\theta^*}^{\pi}} \left[ \left\| \bar{\hat{\psi}}(\tau_h) - \bar{\psi}(\tau_h) \right\|_1 \right] \\ &\leq \sum_h \sum_{\tau_h} \left( \left\| \bar{\hat{\psi}}(\tau_h)\left( \mathbb{P}_{\theta}^{\pi}(\tau_h) - \mathbb{P}_{\hat{\theta}}^{\pi}(\tau_h) \right) + \bar{\hat{\psi}}(\tau_h) \mathbb{P}_{\hat{\theta}}^{\pi}(\tau_h) - \bar{\psi}(\tau_h) \mathbb{P}_{\theta^*}^{\pi}(\tau_h) \right\|_1 \right)\\ &\leq \sum_h \sum_{\tau_h} \left( \left\| \bar{\hat{\psi}}(\tau_h) \right\|_1 \left| \mathbb{P}_{\theta}^{\pi}(\tau_h) - \mathbb{P}_{\hat{\theta}}^{\pi}(\tau_h) \right| + \left\| \hat{\psi}(\tau_h) - \bar{\psi}^*(\tau_h) \right\|_1 \pi(\tau_h)\right)\\ &\leq 2\sum_h |\mathcal{Q}_h^A| \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi}(\tau_h) , \mathbb{P}_{\hat{\theta}}^{\pi}(\tau_h)\right) \\ &\leq 2HQ_A \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi}(\tau_H) , \mathbb{P}_{\hat{\theta}}^{\pi}(\tau_H)\right), \end{align}\] which completes the proof.

 \(\blacksquare\)

10.4 Proof of 2↩︎

Now, we are ready to prove 2.

Proof:  Under event \(\mathcal{E}^o\), the performance difference can be upper bounded as follows. \[\begin{align} V_{\theta^*,r}^{\pi} &- V_{\theta^*,r}^{\hat{\pi}} \\ & = V_{\theta^*,r}^{\pi} - (V_{\hat{\theta}, r}^{\pi} - V_{\hat{\theta}, \hat{b}}^{\pi}) + \underbrace{( V_{\hat{\theta}, r}^{\pi} - V_{\hat{\theta}, \hat{b}}^{\pi} ) - ( V_{\hat{\theta}, r}^{\hat{\pi}} - V_{\hat{\theta}, \hat{b}}^{\hat{\pi}} )}_{I_1} + \underbrace{V_{\hat{\theta}, r}^{\hat{\pi}} - V_{\hat{\theta}, \hat{b}}^{\hat{\pi}} - V_{\theta^*,r}^{\hat{\pi}}}_{I_2}\\ &\overset{(a)}\leq \frac{1}{2}\mathtt{D}_{\mathtt{TV}} \left(\mathbb{P}_{\hat{\theta}}^{\pi} , \mathbb{P}_{\theta^*}^{\pi} \right) + V_{\hat{\theta}, \hat{b}}^{\pi} \\ &\overset{(b)}\leq \mathtt{D}_{\mathtt{TV}} \left(\mathbb{P}_{\hat{\theta}}^{\pi} , \mathbb{P}_{\theta^*}^{\pi} \right) + V_{\theta^*, \hat{b}}^{\pi}, \end{align}\] where \((a)\) follows from the design of \(\bar{\pi}\) that results in \(I_1\leq 0\), and from 4 which implies \(I_2\leq 0\), and \((b)\) follows from 4.

By 8, we have \[\begin{align} \mathtt{D}_{\mathtt{TV}} \left(\mathbb{P}_{\hat{\theta}}^{\pi} , \mathbb{P}_{\theta^*}^{\pi} \right) & \lesssim\sqrt{\frac{\hat{\beta}}{H\iota^2\gamma^2}} \sum_{h=1}^H \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi}} \left[ \; \left\|\bar{\psi}^* (\tau_{h-1})\right\|_{ \Lambda_{h-1}^{-1}} \right] \\ &\overset{(a)}\leq C_{\pi^b,\infty}^{\pi}\sqrt{\frac{\hat{\beta}}{H\iota^2\gamma^2}} \sum_{h=1}^H \mathop{\mathbb{E}}_{\tau_{h-1}\sim\mathbb{P}_{\theta^*}^{\pi^b}} \left[ \; \left\|\bar{\psi}^* (\tau_{h-1})\right\|_{ \Lambda_{h-1}^{-1}} \right] \\ & \leq C_{\pi^b,\infty}^{\pi}H\sqrt{\frac{\hat{\beta}}{H\iota^2\gamma^2}} \sqrt{\frac{rH}{K}} = C_{\pi^b,\infty}^{\pi}H\sqrt{\frac{r\hat{\beta}}{K\iota^2\gamma^2}}, \end{align}\] where \((a)\) follows from the definition of the coverage coefficient.

By 10, we have \[\begin{align} &V_{\theta^*, \hat{b}}^{\pi} \\ & \leq \min\left\{ \hat{\alpha} \left(1 + \frac{2 \sqrt{7 r\hat{\beta}}}{ \iota \sqrt{\hat{\lambda}}} \right) \sum_{h=0}^{H-1} \mathop{\mathbb{E}}_{\tau_h\sim\mathbb{P}_{\theta^*}^{\pi}} \left[\left\|\bar{\psi}^* (\tau_h) \right\|_{\left(U_h\right)^{-1}} \right] + \frac{2\hat{\alpha}HQ_A}{ \sqrt{\hat{\lambda}}} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi}(\tau_H) , \mathbb{P}_{\hat{\theta}}^{\pi}(\tau_H)\right), 1 \right\} \\ &\leq \min\left\{ \hat{\alpha} \left(1 + \frac{2 \sqrt{7 r\hat{\beta}}}{ \iota \sqrt{\hat{\lambda}}} \right) \sum_{h=0}^{H-1} \mathop{\mathbb{E}}_{\tau_h\sim\mathbb{P}_{\theta^*}^{\pi}} \left[\left\|\bar{\psi}^* (\tau_h) \right\|_{\left(U_h\right)^{-1}} \right], 1 \right\} + \frac{2\hat{\alpha}HQ_A}{ \sqrt{\hat{\lambda}}} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi}(\tau_H) , \mathbb{P}_{\hat{\theta}}^{\pi}(\tau_H)\right)\\ &\leq \underbrace{ \min\left\{ \hat{\alpha}C_{\pi^b,\infty}^{\pi} \left(1 + \frac{2 \sqrt{7 r\hat{\beta}}}{ \iota \sqrt{\hat{\lambda}}} \right) \sum_{h=0}^{H-1} \mathop{\mathbb{E}}_{\tau_h\sim\mathbb{P}_{\theta^*}^{\pi^b}} \left[\left\|\bar{\psi}^* (\tau_h) \right\|_{\left(U_h\right)^{-1}} \right], 1 \right\} }_{I_3} + \frac{2\hat{\alpha}HQ_A}{ \sqrt{\hat{\lambda}}} \mathtt{D}_{\mathtt{TV}}\left( \mathbb{P}_{\theta^*}^{\pi}(\tau_H) , \mathbb{P}_{\hat{\theta}}^{\pi}(\tau_H)\right). \end{align}\]

Recall that \(U_h = \hat{\lambda} I + \sum_{\tau_h\in\mathcal{D}_h} \bar{\psi}^*(\tau_h)\bar{\psi}^*(\tau_h)^{\top}\) and the distribution of \(\tau_h\in\mathcal{D}_h\) follows \(\mathbb{P}_{\theta^*}^{\pi^b}\). Therefore, by the Azuma-Hoeffding’s inequality (11), with probability at least \(1-\delta\), we have \[\begin{align} KI_3 &\leq \sqrt{2K\log(2/\delta)} + \min\left\{ \hat{\alpha}C_{\pi^b,\infty}^{\pi} \left(1 + \frac{2 \sqrt{7 r\hat{\beta}}}{ \iota \sqrt{\hat{\lambda}}} \right) \sum_{h=0}^{H-1} \sum_{\tau_h\in\mathcal{D}_h} \left\|\bar{\psi}^* (\tau_h) \right\|_{\left(U_h\right)^{-1}} , 1 \right\} \\ &\overset{(a)}\lesssim \sqrt{K\log(2/\delta)} + \hat{\alpha}C_{\pi^b,\infty}^{\pi}\left(1+\frac{ \sqrt{r\hat{\beta}} }{ \iota\sqrt{\hat{\lambda}} } \right) H \sqrt{\frac{rK}{H}} \\ &\lesssim \hat{\alpha}C_{\pi^b,\infty}^{\pi}\left(1+\frac{ \sqrt{r\hat{\beta}} }{ \iota\sqrt{\hat{\lambda}} } \right) \sqrt{rHK\log(2/\delta)}, \end{align}\] where \((a)\) follows from the Cauchy’s inequality.

Combining the above results, we have \[\begin{align} V_{\theta^*,r}^{\pi} &- V_{\theta^*,r}^{\hat{\pi}}\\ &\lesssim \hat{\alpha}C_{\pi^b,\infty}^{\pi}\left(1+\frac{ \sqrt{r\hat{\beta}} }{ \iota\sqrt{\hat{\lambda}} } \right) \sqrt{\frac{rH\log(2/\delta)}{K}} + \frac{\hat{\alpha}HQ_A}{\sqrt{\hat{\lambda}}} \frac{C_{\pi^b,\infty}^{\pi} H}{\iota\gamma}\sqrt{\frac{r\hat{\beta}}{K}} \\ &\lesssim C_{\pi^b,\infty}^{\pi}\left(\frac{Q_A\sqrt{dH} }{\gamma^2}\sqrt{\hat{\lambda}} + \frac{\sqrt{\hat{\beta}}}{\iota\gamma}\right) \left( 1 + \frac{ \max\{ \sqrt{r} , \sqrt{H}Q_A/\gamma\} }{\iota\sqrt{\hat{\lambda}}} \right) H\sqrt{\frac{rH\hat{\beta}\log(2/\delta)}{K}}\\ & = C_{\pi^b,\infty}^{\pi}\left(\frac{\iota Q_A\sqrt{dH} }{\gamma\sqrt{\hat{\beta}}}\sqrt{\hat{\lambda}} + 1 \right) \left( 1 + \frac{ \max\{ \sqrt{r} , \sqrt{H}Q_A/\gamma\} }{\iota\sqrt{\hat{\lambda}}} \right) \frac{H\sqrt{\hat{\beta}}}{\iota\gamma}\sqrt{\frac{rH\hat{\beta}\log(2/\delta)}{K}}\\ &\overset{(a)}\lesssim \max\left\{ \sqrt{r} , \frac{\sqrt{H}Q_A}{\gamma} \right\} C_{\pi^b,\infty}^{\pi} \frac{Q_AH^2\sqrt{d}}{\iota\gamma^2} \sqrt{\frac{r\hat{\beta}\log(2/\delta)}{K}}, \end{align}\] where \((a)\) follows because \[\begin{align} \lambda = \frac{\gamma \sqrt{\hat{\beta}} \max\{\sqrt{r}, Q_A\sqrt{H}/\gamma\}}{\iota^2Q_A\sqrt{dH}}. \end{align}\] This completes the proof.

 \(\blacksquare\)

11 Auxiliary Lemmas↩︎

The following lemma characterize the relationship between the total variation distance and the Hellinger-squared distance. Note that the result for probability measures has been proved in Lemma H.1 in [23]. Since we consider more general bounded measures, we provide the full proof for completeness.

Proof:  We first prove the first inequality. By the definition of total variation distance, we have \[\begin{align} \mathtt{D}_{\mathtt{TV}}^2(P,Q) & = \left(\sum_x |P(x) - Q(x)| \right)^2\\ & = \left(\sum_x \left(\sqrt{P(x)} - \sqrt{Q(x)}\right)\left(\sqrt{P(x)} + \sqrt{Q(x)} \right) \right)^2\\ &\overset{(a)}\leq \left(\sum_x\left(\sqrt{P(x)} - \sqrt{Q(x)}\right)^2\right) \left( 2\sum_{x}\left( P(x) + Q(x)\right) \right)\\ &\leq 4(|P| + |Q|) \mathtt{D}_{\mathtt{H}}^2(P,Q), \end{align}\] where \((a)\) follows from the Cauchy’s inequality and because \((a+b)^2\leq 2a^2+2b^2\).

For the second inequality, we have, \[\begin{align} \mathop{\mathbb{E}}_{X\sim P } &\left[ \mathtt{D}_{\mathtt{H}}^2(P_{Y|X},Q_{Y|X})\right]\\ &= \sum_{x}P(x)\left( \sum_y \left(\sqrt{P_{Y|X}(y)} - \sqrt{Q_{Y|X}(y)} \right)^2 \right) \\ & = \sum_{x,y} \left( \sqrt{P_{X,Y}(x,y)} - \sqrt{Q_{X,Y}(x,y)} + \sqrt{Q_{Y|X}(y) Q(x)} - \sqrt{Q_{Y|X}(y) P(x)} \right)^2 \\ &\leq 2\sum_{x,y} \left( \sqrt{P_{X,Y}(x,y)} - \sqrt{Q_{X,Y}(x,y)} \right)^2 + 2\sum_{x,y} Q_{Y|X}(y)\left( \sqrt{ Q(x)} - \sqrt{ P(x)} \right)^2 \\ & = 4\mathtt{D}_{\mathtt{H}}^2(P_{X,Y},Q_{X,Y}) + 2(|P|+|Q| - 2\sum_x\sqrt{P(x)Q(x)}) \\ &\overset{(a)}\leq 4\mathtt{D}_{\mathtt{H}}^2(P_{X,Y},Q_{X,Y}) + 2(|P|+|Q| - 2\sum_x\sum_y\sqrt{P_{Y|X}(y)P(x)Q_{Y|X}(y)Q(x)}) \\ & = 8 \mathtt{D}_{\mathtt{H}}^2(P_{X,Y},Q_{X,Y}), \end{align}\] where \((a)\) follows from the Cauchy’s inequality that applies on \(\sum_y\sqrt{P_{Y|X}(y)Q_{Y|X}(y)}\).

 \(\blacksquare\)

Proof:  We first write \[\begin{align} & \|x_{i}\|_{A^{-1}} - \|y_{i}\|_{B^{-1}} \\ & \quad = \|x_{i}\|_{A^{-1}} - \|y_{i}\|_{A^{-1}} + \|y_{i}\|_{A^{-1}} - \|y_{i}\|_{B^{-1}}\\ &\quad = \frac{ \overbrace{ \|x_{i}\|^2_{A^{-1}} - \|y_{i}\|^2_{A^{-1}} }^{I_1} }{ \|x_{i}\|_{A^{-1}} + |y_{i}\|_{A^{-1}} } + \frac{ \overbrace{ \|y_{i}\|^2_{A^{-1}} - \|y_{i}\|^2_{B^{-1}} }^{I_2} }{ \|y_{i}\|_{A^{-1}} + \|y_{i}\|_{B^{-1}} }. \end{align}\]

For the first term \(I_1\), we repeatedly apply the Cauchy’s inequality, and have \[\begin{align} \|x_{i} &\|_{A^{-1}}^2 - \|y_{i}\|_{A^{-1}}^2\\ & = x_{i}^{\top} A^{-1} (x_{i} - y_{i}) + y_{i}^{\top} A^{-1} (x_{i} - y_{i})\\ & \leq \|x_{i}\|_{A^{-1}}\|x_{i} - y_{i}\|_{A^{-1}} + \|y_{i}\|_{A^{-1}}\|x_{i}-y_{i}\|_{A^{-1}}\\ & \leq \frac{1}{\sqrt{\lambda}} \left(\|x_{i}\|_{A^{-1}} + \|y_{i}\|_{A^{-1}} \right) \left\| x_{i} - y_{i} \right\|_{2}. \end{align}\]

For the second term \(I_2\), we repeatedly apply the Cauchy’s inequality, and have \[\begin{align} \|y_{i} &\|_{A^{-1}}^2 - \|y_{i}\|_{B^{-1}}^2\\ & = y_{i}^{\top}A^{-1}(B-A)B^{-1}y_{i}\\ & = \sum_{j\in\mathcal{J}} \bigg(y_{i}^{\top}A^{-1}x_{j} (y_{j} - x_{j})^{\top}B^{-1}y_{i} + y_{i}^{\top}A^{-1}(y_{j} -x_{j})y_{j}^{\top}B^{-1}y_{i} \bigg)\\ &\leq \sum_{j\in\mathcal{J} } \|x_{j}\|_{A^{-1}} \|y_{j}-x_{j}\|_{B^{-1}} \|y_{i}\|_{A^{-1}}\|y_{i}\|_{B^{-1}} \\ & \quad + \sum_{j\in\mathcal{J}} \|y_{j}-x_{j}\|_{A^{-1}} \|y_{j}\|_{B^{-1}} \|y_{i}\|_{A^{-1}} \|y_{i}\|_{B^{-1}}\\ & \leq \frac{1}{\sqrt{\lambda}}\|y_{i}\|_{A^{-1}}\|y_{i}\|_{B^{-1}} \sqrt{ \sum_{j\in\mathcal{J}} \|x_{j}-y_{j}\|_{2}^2 } \sqrt{ \sum_{j\in\mathcal{J}} \|x_{j}\|^2_{A^{-1}} }\\ &\quad + \frac{1}{\sqrt{\lambda}} \|y_{i}\|_{A^{-1}} \|y_{i}\|_{B^{-1}} \sqrt{ \sum_{j\in\mathcal{J} } \|x_{j}-y_{j}\|_{2}^2 }\sqrt{ \sum_{j\in\mathcal{J}} \|y_{j}\|_{B^{-1}}^2 } \\ &\leq \frac{2\sqrt{r}}{\sqrt{\lambda}} \|y_{i}\|_{A^{-1}} \|y_{i}\|_{B^{-1}} \sqrt{ \sum_{j\in\mathcal{J} } \|x_{j_h}-y_{j_h}\|_{2}^2 }. \end{align}\] Therefore, the lemma follows from the fact that \(\|y_i\|_{A^{-1}} \leq \|y_i\|_{A^{-1}} + \|y_i\|_{B^{-1}}.\)

 \(\blacksquare\)

The following lemma and its variants has been developed in [69], [72], [73]. We slightly generalize the padding term from \(1\) to an arbitrary positive number \(B\) and provide the full proof for completeness.

Proof:  Note that the second inequality is an immediate result from the first inequality by the Cauchy’s inequality. Hence, it suffices to prove the first inequality. To this end, we have \[\begin{align} \sum_{k=1}^K \min\left\{\|x_k\|_{U_k^{-1}}^2 , B \right\} & \overset{(a)}\leq (1+B)\sum_{k=1}^K \log\left(1 + \|x_k \|_{U_k^{-1}}^2 \right)\\ & = (1+B)\sum_{k=1}^K \log \left( 1+ \mathtt{trace}\left( \left(U_{k+1} - U_k \right) U_k^{-1} \right) \right) \\ &= (1+B)\sum_{k=1}^K \log\left( 1+ \mathtt{trace}\left( U_k^{-1/2}\left(U_{k+1} - U_k \right) U_k^{-1/2} \right) \right)\\ &\leq (1+B)\sum_{k=1}^K \log\mathtt{det}\left(I_d + U_k^{-1/2}\left(U_{k+1} - U_k \right) U_k^{-1/2} \right)\\ &= (1+B)\sum_{k=1}^K \log \frac{\mathtt{det}\left(U_{k+1}\right)}{\mathtt{det}(U_k)}\\ & = (1+B) \log \frac{\mathtt{det}(U_{K+1})}{\mathtt{det}(U_1)} \\ & = (1+B)\log\mathtt{det}\left( I + \frac{1}{\lambda}\sum_{k=1}^{K}x_kx_k^{\top} \right) \\ &\overset{(b)}\leq (1+B)r\log(1+K/\lambda), \end{align}\] where \((a)\) follows because \(x\leq (1+B)\log(1+x)\) if \(0<x\leq B\), and \((b)\) follows because \(\mathtt{rank}(\mathcal{X}) \leq r.\)

 \(\blacksquare\)

References↩︎

---

1. The Pennsylvania State University, State College, PA 16801↩︎

2. The Ohio State University, Columbus, OH 43210↩︎

3. The sample complexity of learning \(\psi_0\) if it is unknown is relatively small compared with learning other parameters.↩︎

4. Computation complexity of calculating total variation distance is still an open question [68] and only some specific forms can be approximated in polynomial time.↩︎

5. \(|\mathcal{O}|\) can be the cardinality of \(\mathcal{O}\) if it is finite, or the measure of \(\mathcal{O}\) if it is a measurable set with positive and bounded measure.↩︎