SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\| \widehat {P} - P \| _ {1 \to \infty} \lesssim \sqrt {\frac {n}{T}} \log (n T), \| \widehat {P} - P \| _ {\infty} \lesssim \frac {1}{\sqrt {n T}} \log (n T).
$$
As for the generative model, for a homogeneous frequency matrix, our estimates achieve a relative error scaling at most as  $\sqrt{n / T}$  up to the logarithmic factor for all metrics. Note that up to a logarithmic factor, the upper bound for  $\| \widehat{P} - P \|_{1 \to \infty}$  (and similarly for  $\widehat{M}$ ) matches the minimax lower bound derived in [63].
3.4 Elements of the proofs
The proofs of the three above theorems share similar arguments. We only describe elements of the proof of Theorem 5, corresponding to the most challenging model. The most difficult result concerns the singular subspace recovery (the upper bounds (i) in our theorems), and it can be decomposed into the following three steps. The first two steps are meant to deal with the Markovian nature of the data. The third step consists in applying a leave-one-out analysis to recover the singular subspaces.
Step 1: Multinomial approximation of Markovian data. We treat the matrix  $\widetilde{M}^{(k)}$  arising from one subset of data, and for simplicity, we remove the superscript  $(k)$ , i.e.,  $\widetilde{M} = \widetilde{M}^{(k)}$ . Note that  $T_{\tau}\widetilde{M}$  is a matrix recording the numbers of transitions observed in the data for any pair of states: denote by  $N_{i,j}$  this number for  $(i,j)$ . We approximate the joint distribution of  $N = (N_{i,j})_{(i,j)}$  by a multinomial distribution with  $n^2$  components and parameter  $T_{\tau}M_{i,j}$  for component  $(i,j)$ . Denote by  $Z = (Z_{i,j})_{(i,j)}$  the corresponding multinomial random variable. Using the mixing property of the Markov chain and the choice of  $\tau$ , we establish (see Lemma 21 in App. C) that for any subset  $\mathcal{Z}$  of  $\{z\in \mathbb{N}^{n^2}:\sum_{(i,j)}z_{i,j} = T_\tau \}$ , we have  $\mathbb{P}[N\in \mathcal{Z}]\leq 3\mathbb{P}[Z\in \mathcal{Z}]$ .
Step 2: Towards Poisson random matrices with independent entries. The random matrix  $Z$  does not have independent entries. Independence is however a requirement if we wish to apply the leave-one-out argument. Consider the random matrix  $Y$  whose entries are independent Poisson random variables with mean  $T_{\tau}M_{i,j}$  for the  $(i,j)$ -th entry. We establish the following connection between the distribution of  $Z$  and that of  $Y$ : for any  $\mathcal{Z} \subset \mathbb{N}^{n^2}$ , we have  $\mathbb{P}[Z \in \mathcal{Z}] \leq e\sqrt{T_{\tau}}\mathbb{P}[Y \in \mathcal{Z}]$ . Refer to Lemma 22 in App. C for details.
Step 3: The leave-one-out argument for Poisson matrices. Combining the two first steps provides a connection between the observation matrix  $\widetilde{M}$  and a Poisson matrix  $Y$  with independent entries. This allows us to apply a leave-one-out analysis to  $\widetilde{M}$  as if it had independent entries (replacing  $\widetilde{M}$  by  $Y$ ). The analysis starts by applying the standard dilation trick (see Section 4.10 in [14]) so as to make  $\widetilde{M}$  symmetric. Then, we can decompose the error  $\|U - \widehat{U}(\widehat{U}^\top U)\|_{2\to \infty}$  (see Lemma 32 in App. E) into several terms. The most challenging of these terms is  $\|(M - \widetilde{M})(U - \widehat{U}(\widehat{U}^\top U))\|_{2\to \infty} = \max_{l\in [n]} \|(M_{l,:} - \widetilde{M}_{l,:})(U - \widehat{U}(\widehat{U}^\top U))\|_2$  because of inherent dependence between  $M - \widetilde{M}$  and  $U - \widehat{U}(\widehat{U}^\top U)$ . The leave-one-out analysis allows us to decouple this statistical dependency. It consists in exploiting the row and column independence of matrix  $\widetilde{M}$  to approximate  $\|(M_{l,:} - \widetilde{M}_{l,:})(U - \widehat{U}(\widehat{U}^\top U))\|_2$  by  $\|(M_{l,:} - \widetilde{M}_{l,:})(U - \widehat{U}^{(l)}((\widehat{U}^{(l)})^\top U)\|_2$  where  $\widehat{U}^{(l)}$  is the matrix of eigenvectors of matrix  $\widetilde{M}^{(l)}$  obtained by zeroing the  $l$ -th row and column of  $\widetilde{M}$ . By construction,  $(M_{l,:} - \widetilde{M}_{l,:})$  and  $U - \widehat{U}^{(l)}((\widehat{U}^{(l)})^\top U)$  are independent, which simplifies the analysis. The proof is completed by a
further appropriate decomposition of this term, combined with concentration inequalities for random Poisson matrices (see App. D).
4 Regret Minimization in Low-Rank Bandits
Consider a low-rank bandit problem with a homogeneous rank- $r$  reward matrix  $M$ . We wish to devise an algorithm  $\pi$  with low regret.  $\pi$  selects in round  $t$  an entry  $(i_t^\pi, j_t^\pi)$  based on previous observations, and receives as a feedback the noisy reward  $M_{i_t^\pi, j_t^\pi} + \xi_t$ . The regret up to round  $T$  is defined by  $R^\pi(T) = TM_{i^\star, j^\star} - \mathbb{E}[\sum_{t=1}^T M_{i_t^\pi, j_t^\pi}]$ , where  $(i^\star, j^\star)$  is an optimal entry. One could think of a simple Explore-Then-Commit (ETC) algorithm, where in the first phase entries are sampled uniformly at random, and where in a second phase, the algorithm always selects the highest entry of  $\widehat{M}$  built using the samples gathered in the first phase and obtained by spectral decomposition. When the length of the first phase is  $T^{2/3}(n+m)^{1/3}$ , the ETC algorithm would yield a regret upper bounded by  $O(T^{2/3}(n+m)^{1/3})$  for  $T = \Omega((n+m) \log^3(n+m))$ .
To get better regret guarantees, we present SME-AE (Successive Matrix Estimation and Arm Elimination), an algorithm meant to identify the best entry as quickly as possible with a prescribed level of certainty. After the SME-AE has returned the estimated best entry, we commit and play this entry for the remaining rounds. The pseudo-code of SME-AE is presented in Algorithm 1. The algorithm runs in epochs: in epoch  $\ell$ , it samples  $T_{\ell}$  entries uniformly at random among all entries (in  $T_{\ell}$ , the constant  $C$  just depends on upper bounds of the parameters  $\mu, \kappa$ , and  $\|M\|_{\infty}$ , refer to App. G); from these samples, a matrix  $\widehat{M}^{(\ell)}$  is estimated and  $\mathcal{A}_{\ell}$ , the set of candidate arms, is pruned. The pruning procedure is based on the estimated gaps:  $\widehat{\Delta}_{i,j}^{(\ell)} = \widehat{M}_{\star}^{(\ell)} - \widehat{M}_{i,j}^{(\ell)}$  where  $\widehat{M}_{\star}^{(\ell)} = \max_{i,j} \widehat{M}_{i,j}^{(\ell)}$ .
Algorithm 1: Successive Matrix Estimation and Arm Elimination (SME-AE)
Input: Arms  $[m]\times [n]$  , confidence level  $\delta$ $\ell = 1$  .   
 $\mathcal{A}_1 = [m]\times [n]$  .   
while  $|\mathcal{A}_{\ell}| > 1$  do   
 $\begin{array}{l}\delta_{\ell} = \delta /\ell^{2};\\ T_{\ell} = \left[C\left(2^{\ell +2}\right)^{2}(m + n)\log^{3}\left(2^{2\ell +4}(m + n) / \delta_{\ell}\right)\right];\\ \mathrm{Sample~uniformly~at~random~}T_{\ell}\mathrm{~entries~from~}\mathcal{A}_{1}:(M_{i_{t},j_{t}} + \xi_{t})_{t = 1,\dots ,T_{\ell}};\\ \mathrm{Estimate~}\widehat{M}^{(\ell)}\mathrm{~via~spectral~decomposition~as~described~in~Section~3.1~};\\ \mathcal{A}_{\ell +1} = \big\{(i,j)\in \mathcal{A}_{\ell}:\widehat{\Delta}_{i,j}^{(\ell)}\leq 2^{-(\ell +2)}\big\} ;\ell = \ell +1; \end{array}$    
end   
Output: Recommend the remaining pair  $(\hat{i}_{\tau},\hat{j}_{\tau})$  in  $\mathcal{A}_{\ell}$
The following theorem characterizes the performance of SME-AE and the resulting regret. To simplify the notation, we introduce the gaps: for any entry  $(i,j)$ ,  $\Delta_{i,j} = (M_{i^{\star},j^{\star}} - M_{i,j})$ ,  $\Delta_{\mathrm{min}} = \min_{(i,j):\Delta_{i,j} > 0}\Delta_{i,j}$ ,  $\Delta_{\mathrm{max}} = \max_{(i,j)}\Delta_{i,j}$ , and  $\bar{\Delta} = \sum_{(i,j)}\Delta_{i,j} / (mn)$ . We define the function  $\psi(n,m,\delta) = \frac{c(m + n)\log(e / \Delta_{\mathrm{min}})}{\Delta_{\mathrm{min}}^2}\log^3\left(\frac{e(m + n)\log(e / \Delta_{\mathrm{min}})}{\Delta_{\mathrm{min}}\delta}\right)$  for some universal constant  $c > 0$ .
Theorem 7. (Best entry identification) For any  $\delta \in (0,1)$ , SME-AE( $\delta$ ) stops at time  $\tau$  and recommends arm  $(\hat{i}_{\tau},\hat{j}_{\tau})$  with the guarantee  $\mathbb{P}\big((\hat{i}_{\tau},\hat{j}_{\tau}) = (i^{\star},j^{\star}),\tau \leq \psi (n,m,\delta)\big) \geq 1 - \delta$ . Moreover, for any  $T \geq 1$  and  $\alpha > 0$ , the sample complexity  $\tau$  of SME-AE(1/T $^{\alpha}$ ) satisfies  $\mathbb{E}[\tau \wedge T] \leq \psi (n,m,T^{-\alpha}) + T^{1 - \alpha}$ . (Regret) Let  $T \geq 1$ . Consider the algorithm  $\pi$  that first runs SME-AE(1/T $^2$ ) and then commits to its output  $(\hat{i}_{\tau},\hat{j}_{\tau})$  after  $\tau$ . We have:  $R^{\pi}(T) \leq \bar{\Delta}\left(\psi (n,m,T^{-2}) + 1\right) + \frac{\Delta_{\max}}{T}$ .
The proof of Theorem 7 is given in App. G. Note that the regret upper bounds hold for any time horizon  $T \geq 1$ , and that it scales as  $O((m + n) \log^3 (T) \bar{\Delta} / \Delta_{\mathrm{min}}^2)$  (up to logarithmic factors in  $m, n$  and  $1 / \Delta_{\mathrm{min}}$ ). The cubic dependence in  $\log^3 (T)$  is an artifact of our proof techniques. More precisely, it is due to the Poisson approximation used to obtain entry-wise guarantees. Importantly, for any time horizon, the regret upper bound only depends on  $(m + n)$  rather than  $mn$  (the number
of arms / entries), and hence, the low-rank structure is efficiently exploited. If we further restrict our attention to problems with gap ratio  $\Delta_{\mathrm{max}} / \Delta_{\mathrm{min}}$  upper bounded by  $\zeta$ , our regret upper bound becomes  $O(\zeta(m + n) \log^3(T) / \Delta_{\mathrm{min}})$ , and can be transformed into the minimax gap-independent upper bound  $O(\zeta((m + n)T)^{1/2} \log^2(T))$ , see App. G. Finally note that  $\Omega((m + n)T)^{1/2})$  is an obvious minimax regret lower bound for our low-rank bandit problem.
A very similar low-rank bandit problem has been investigated in [7]. There, under similar assumptions (see Assumption 1 and Definition 1), the authors devise an algorithm with both gap-dependent and gap-independent regret guarantees. The latter are difficult to compare with ours. Their guarantees exhibit a better dependence on  $T$  and  $\Delta_{\mathrm{min}}$ , but worse in the matrix dimensions  $n$  and  $m$ . Indeed in our model,  $b^{\star}$  in [7] corresponds to  $\| M\|_{2\to \infty}$  and scales as  $\sqrt{n}$ . As a consequence, the upper bounds in [7] have a dependence in  $n$  and  $m$  scaling as  $\sqrt{n} (n + m)$  in the worst case for gap-dependent guarantees and even  $nm$  (through the constant  $C_2$  in [7]) for gap-independent guarantees.
5 Representation Learning in Low-Rank MDPs
The results derived for Models II(a) and II(b) are instrumental towards representation learning and hence towards model-based or reward-free RL in low-rank MDPs. In this section, we provide an example of application of these results, and mention other examples in Section 7. A low-rank MDP is defined by  $(\mathcal{S},\mathcal{A},\{P^a\}_{a\in \mathcal{A}},R,\gamma)$  where  $\mathcal{S}$ ,  $\mathcal{A}$  denote state and action spaces of cardinalities  $n$  and  $A$ , respectively,  $P^a$  denotes the rank- $r$  transition matrix when taking action  $a$ ,  $R$  is the reward function, and  $\gamma$  is the discount factor. We assume that all rewards are in  $[0,1]$ . The value function of a policy  $\pi : \mathcal{S} \to \mathcal{A}$  is defined as  $V_R^\pi(x) = \mathbb{E}[\sum_{t=1}^\infty \gamma^{t-1} R(x_t^\pi, \pi_t(x_t^\pi)) | x_1^\pi = x]$  where  $x_t^\pi$  is the state visited under  $\pi$  in round  $t$ . We denote by  $\pi^\star(R)$  an optimal policy (i.e., with the highest value function).
Reward-free RL. In the reward-free RL setting (see e.g. [36, 30, 66]), the learner does not receive any reward signal during the exploration process. The latter is only used to construct estimates  $\{\widehat{P}^a\}_{a\in \mathcal{A}}$  of  $\{P^a\}_{a\in \mathcal{A}}$ . The reward function  $R$  is revealed at the end, and the learner may compute  $\hat{\pi}(R)$  an optimal policy for the MDP  $(S, \mathcal{A}, \{\widehat{P}^a\}_{a\in \mathcal{A}}, R, \gamma)$ . The performance of this model-based approach is often assessed through  $\Gamma = \sup_{R} \| V_R^{\pi^*(R)} - V_R^{\widehat{\pi}(R)} \|_\infty$ . In tabular MDP, to identify an  $\epsilon$ -optimal policy for all reward functions, i.e., to ensure that  $\Gamma \leq \epsilon$ , we believe that the number of samples that have to be collected should be  $\Omega(\text{poly}(\frac{1}{1-\gamma})\frac{n^2A}{\epsilon^2})$  (the exact degree of the polynomial in  $1/(1-\gamma)$  has to be determined). This conjecture is based on the sample complexity lower bounds derived for reward-free RL in episodic tabular MDP [30, 43]. Now for low-rank MDPs, the equivalent lower bound would be  $\Omega(\text{poly}(\frac{1}{1-\gamma})\frac{nA}{\epsilon^2})$  [28] (this minimax lower bound is valid for Block MDPs, a particular case of low-rank MDPs).
Leveraging our low-rank matrix estimation guarantees, we propose an algorithm matching the aforementioned sample complexity lower bound (up to logarithmic factors) at least when the frequency matrices  $\{M^a\}_{a\in \mathcal{A}}$  are homogeneous. The algorithm consists of two phases: (1) in the model estimation phase, it collects  $A$  trajectories, each of length  $T / A$ , corresponding to the Markov chains with transition matrices  $\{P^a\}_{a\in \mathcal{A}}$ . From this data, it uses the spectral decomposition method described in §3 to build estimates  $\{\widehat{P}^a\}_{a\in \mathcal{A}}$ . (2) In the planning phase, based on the reward function  $R$ , it computes the best policy  $\hat{\pi}(R)$  for the MDP  $(S, \mathcal{A}, \{\widehat{P}^a\}_{a\in \mathcal{A}}, R, \gamma)$ . The following theorem summarizes the performance of this algorithm. To simplify the presentation, we only provide the performance guarantees of the algorithm for homogeneous transition matrices (guarantees for more general matrices can be derived plugging in the results from Theorem 5).
Theorem 8. Assume that for any  $a \in \mathcal{A}$ ,  $M^a$  is homogeneous (as defined in Corollary 4). If  $T \geq cnA\log^2(nAT)$  for some universal constant  $c > 0$ , then we have with probability at least  $1 - \min\{n^{-2}, T^{-1}\}$ :  $\Gamma = \sup_R \|V_R^{\pi^*(R)} - V_R^{\hat{\pi}(R)}\|_{\infty} \lesssim \frac{1}{(1-\gamma)^2}\sqrt{\frac{nA}{T}}\log(nAT)$ .
Theorem 8 is a direct consequence of Corollary 6 and of the fact that for any reward function  $R$ :  $\| V_R^{\pi^*(R)} - V_R^{\widehat{\pi} (R)}\|_{\infty}\leq \frac{2\gamma}{(1 - \gamma)^2}\max_{a\in \mathcal{A}}\| P^a -\widehat{P}^a\|_{1\to \infty}$ , see App. A. The theorem implies that if we wish to guarantee  $\Gamma \leq \epsilon$ , we just need to collect  $O(\frac{nA}{\epsilon^2(1 - \gamma)^4})$  samples up to a logarithmic factor.
This sample complexity is minimax optimal in  $n$ ,  $A$ , and  $\epsilon$  in view of the lower bound presented in [28].
Low-rank matrix estimation. Until recently, the main efforts on low-rank matrix recovery were focused on guarantees w.r.t. the spectral or Frobenius norms, see e.g. [19] and references therein. The first matrix estimation and subspace recovery guarantees in  $\ell_{2\to \infty}$  and  $\ell_{\infty}$  were established in [22], [21] via a more involved perturbation analysis than the classical Davis-Kahan bound. An alternative approach based on a leave-one-out analysis was proposed in [2], and further refined in [10, 12, 17], see [14] for a survey. Some work have also adapted the techniques beyond the independent noise assumption [39, 1, 5], but for very specific structural dependence. We deal with a stronger dependence, and in particular with Markovian data (an important scenario in RL).
The estimation of low-rank transition matrices of Markov chains has been studied in [63, 9] using spectral methods and in [40, 67] using maximum-likelihood approaches. [63] does not conduct any fine-grained subspace recovery analysis (such as the leave-one-out), and hence the results pertaining to the  $\| \cdot \|_{1\to \infty}$ -guarantees are questionable; refer to App. H for a detailed justification. All these papers do not present entry-wise guarantees.
It is worth mentioning that there exist other methods for matrix estimation that do not rely on spectral decompositions like ours, yet enjoy entry-wise matrix estimation guarantees [51, 3, 50]. However, these methods require different assumptions than ours that may be too strong for our purposes, notably having access to the so-called anchor rows and columns. Moreover, we do not know if these methods also lead to guarantees for subspace recovery in the norm  $\| \cdot \|_{2\to \infty}$ , nor how to extend those results to settings with dependent noise.
Low-rank bandits. Low-rank structure in bandits has received a lot of attention recently [35, 37, 32, 57, 41, 7, 34, 27]. Different set-ups have been proposed (refer to App. H for a detailed exposition, in particular, we discuss how the settings proposed in [32, 7] are equivalent), and regret guarantees in an instance dependent and minimax sense have been both established.
Typically minimax regret guarantees in bandits scale as  $\sqrt{T}$ , but the scaling in dimension may defer when dealing with a low rank structure [32, 34, 7]. In [32], the authors also leverage spectral methods. They reduce the problem to a linear bandit of dimension  $nm$  but where only roughly  $n + m$  dimensions are relevant. This entails that a regret lower bound of order  $(n + m)\sqrt{T}$  is inevitable. Actually, in their reduction to linear bandits, they only use a subspace recovery in Frobenius norm, which perhaps explains the scaling  $(n + m)^{3/2}$  in their regret guarantees. It is worth noting that in [34], the authors manage to improve upon the work [32] and obtain a scaling order  $(m + n)$  in the regret. Our algorithm leverages entry-wise guarantees which rely on a stronger subspace recovery guarantee. This allows us to obtain a scaling  $\sqrt{n + m}$  in the regret. The work of [27] is yet another closely related work to ours. There, the authors propose an algorithm achieving a regret of order  $\text{polylog}(n + m)\sqrt{T}$  for a contextual bandit problem with low rank structure. However, their result only holds for rank 1 and their observation setup is different than ours because in their setting, the learner observes  $m$  entries per round while in ours the learner only observes one entry per round. In [7], the authors use matrix estimation with nuclear norm penalization to estimate the matrix  $M$ . Their regret guarantees are already discussed in §4.
Some instance-dependent guarantees with logarithmic regret for low rank bandits have been established in [35, 37, 57]. However, these results suffer what may be qualified as serious limitations. Indeed, [35, 57] provide instance dependent regret guarantees but only consider low-rank bandits with rank 1, and the regret bounds of [35] are expressed in terms of the so-called column and row gaps (see their Theorem 1) which are distinct from the standard gap notions. [37] extend the results in [35] to rank  $r$  with the limitation that they require stronger assumptions than ours. Moreover, the computational complexity of their algorithm depends exponentially on the rank  $r$ ; they require a search over spaces of size  $\binom{m}{r}$  and  $\binom{n}{r}$ . Our proposed algorithm does not suffer from such limitations.
We wish to highlight that our entry-wise guarantees for matrix estimation are the key enabling tool that led us to the design and analysis of our proposed algorithm. In fact, the need for such guarantees arises naturally in the analysis of gap-dependent regret bounds (see Appendix G.1). Therefore, we
believe that such guarantees can pave the way towards better, faster, and efficient algorithms for bandits with low-rank structure.
Low-rank Reinforcement Learning. RL with low rank structure has been recently extensively studied but always in the function approximation framework [29, 18, 20, 44, 24, 65, 56, 4, 46, 59, 60, 49]. There, the transition probabilities can be written as  $\phi(x,a)^\top \mu(x')$  where the unknown feature functions  $\phi(x,a), \mu(x') \in \mathbb{R}^r$  belong to some specific class  $\mathcal{F}$  of functions. The major issue with algorithms proposed in this literature is that they rely on strong computational oracles (e.g., ERM, MLE), see [33, 25, 64] for detailed discussions. In contrast, we do not assume that the transition matrices are constructed based on a given restricted class of functions, and our algorithms do not rely on any oracle and are computationally efficient. In [51, 50], the authors also depart from the function approximation framework. There, they consider a low rank structure different than ours. Their matrix estimation method enjoys an entry-wise guarantee, but requires to identify a subset of rows and columns spanning the range of the full matrix. Moreover, their results are only limited the generative models, which allows to actually rely on independent data samples.
7 Conclusion and Perspectives
In this paper, we have established that spectral methods efficiently recover low-rank matrices even in correlated noise. We have investigated noise correlations that naturally arise in RL, and have managed to prove that spectral methods yield nearly-minimal entry-wise error. Our results for low-rank matrix estimation have been applied to design efficient algorithms in low-rank RL problems and to analyze their performance. We believe that these results may find many more applications in low-rank RL. They can be applied (i) to reward-free RL in episodic MDPs (this setting is easier than that presented in §5 since successive episodes are independent); (ii) to scenarios corresponding to offline RL [62] where the data consists of a single trajectory generated under a given behavior policy (from this data, we can extract the transitions  $(x, a, x')$  where a given action  $a$  is involved and apply the spectral method to learn  $\widehat{P}^a$ ); (iii) to traditional RL where the reward function  $R$  has to be learnt (learning  $R$  is a problem that lies in some sense between the inference problems in our Models I and II); (iv) to model-free RL where we would directly learn the  $Q$  function as done in [52] under a generative model; (v) to low-rank RL problems with continuous state spaces (this can be done if the transition probabilities are smooth in the states, and by combining our methods to an appropriate discretization of the state space).
This research was supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation.