# Multi-User Reinforcement Learning with Low Rank Rewards

Dheeraj Nagaraj  Suhas S Kowshik  Naman Agarwal  Praneeth Netrapalli  Prateek
Jain

###### Abstract

We consider collaborative multi-user reinforcement learning, where multiple
users have the same state-action space and transition probabilities but
different rewards. Under the assumption that the reward matrix of the $N$
users has a low-rank structure – a standard and practically successful
assumption in the collaborative filtering setting – we design algorithms with
significantly lower sample complexity compared to the ones that learn the MDP
individually for each user. Our main contribution is an algorithm which
explores rewards collaboratively with $N$ user-specific MDPs and can learn
rewards efficiently in two key settings: tabular MDPs and linear MDPs. When
$N$ is large and the rank is constant, the sample complexity per MDP depends
logarithmically over the size of the state-space, which represents an
exponential reduction (in the state-space size) when compared to the standard
“non-collaborative” algorithms. Our main technical contribution is a method to
construct policies which obtain data such that low rank matrix completion is
possible (without a generative model). This goes beyond the regular RL
framework and is closely related to mean field limits of multi-agent RL.

Machine Learning, ICML

##  1 Introduction

Reinforcement learning (RL) has recently seen tremendous empirical and
theoretical success (Mnih et al., 2015; Sutton et al., 1992; Jin et al.,
2020b; Gheshlaghi Azar et al., 2013; Dann & Brunskill, 2015). Near optimal
algorithms have been proposed to explore and learn a given MDP with sample
access to trajectories. Multi-agent RL, where multiple agents interact among
themselves and the environment to collect rewards, has gained a lot of
interest due to immense practical applications. However, even simple instances
of multi-agent RL, like restless bandits, can be provably hard (Papadimitriou
& Tsitsiklis, 1999).

In this work, we consider the problem of learning optimal policies for
multiple MDPs collaboratively so that the total number of trajectories sampled
per MDP is smaller than the number of trajectories required to learn them
individually. We assume that the various users have the same transition
matrices, but different rewards and the rewards have a low rank structure.
This is closely related to mean-field limits of certain instances of multi-
agent RL as described in Section 3.

Motivation From the point of view of RL, this is an instance of multi-task
reinforcement learning (MTRL), various versions of which have been considered
in the literature (Brunskill & Li, 2013; D’Eramo et al., 2020; Teh et al.,
2017; Hessel et al., 2019; Lazaric, 2012). Here, an agent learns different
MDPs together with certain common structures. This shared structure could be a
common domain (such as moving towards different target points, but in the same
environment) or similarity in the task to be performed (such as picking up and
moving different kinds of objects). Our problem setup falls in the former
category.

Recently, collaborative filtering has been studied in the online learning
setting (Bresler & Karzand, 2021; Jain & Pal, 2022; Ariu et al., 2020;
Huleihel et al., 2021; Nguyen-Thanh et al., 2019), where multiple bandit
instances are simultaneously explored under low rank assumptions in order to
learn the preferences of multiple users simultaneously. From this point of
view, our work adds temporal dynamics, such as change in preferences over
time, based on past actions. That is, we consider non-stationary environments
via Markov Decision Processes.

To motivate our setup, we consider the example of recommendation systems in
e-commerce or video streaming. In this context, classical collaborative
filtering setup assumes that the preferences are static and not influenced by
the recommendation system itself. However, this is untrue in the real world.
Buying an item changes the preferences of the customer. In case the system
recommends a TV and the customer buys a TV, a TV stand might be the most
relevant recommendation. In case they watch a recommended video about
astrophysics, they might want to watch other astrophysics videos because the
first video piqued their interest in the topic.

Our work captures such scenarios by formulating this as an RL problem where
the state is affected by the actions of the recommendation system and the
state in turn affects the rewards (i.e, the user preferences). This idea has
gained traction recently, as shown by the survey paper (Afsar et al., 2022).
However, in the papers which have been discussed in the survey, the user
information is apriori encoded into the state embedding and the resulting
system is treated as a single agent MDP. In this work, we bring forth the
multi-agent aspect by viewing the RL approach as an enhancement of the
classical collaborative filtering allowing us discover similarity among users.
We believe the theoretical insights gained in such a setting can lead to Deep
RL algorithms which utilize this multi-agent structure effectively and
explicitly. The assumption of a common transition matrix can be relaxed in
practice by clustering users based on side information and modeling each
cluster to have a common transition matrix (see (Mate et al., 2022) and
references therein).

#### Our Contributions

We introduce the setting of multi-user collaborative reinforcement learning in
the case of tabular and linear MDPs. In our study, we isolate and overcome
several technical and conceptual challenges in order achieve sample efficient
learning. The main technical challenge we encounter is obtaining the right
distribution of state-action pairs from users such that we can successfully
run low-rank matrix estimation algorithms, without access to a generative
model (i.e, we can only deploy policies and query trajectories corresponding
to this policy). This requires clever algorithm design since some states can
be hard to even reach. In fact, this endeavor goes beyond standard RL methods
and is related to functional reward maximization and mean field limits of
multi-agent RL as explained in Section 3. To summarize our contributions:

a) Improved Sample Complexity: We provide sample efficient algorithms for both
these scenarios without access to a generative model. Under the low rank
assumption on the reward matrix, the total sample complexity required to learn
the near-optimal policies for every user scales as
$\tilde{O}(N+|\mathcal{S}||\mathcal{A}|)$ instead of
$O(N|\mathcal{S}||\mathcal{A}|)$ for tabular MDPs and $\tilde{O}(N+d)$ instead
of $O(Nd^{2})$ for linear MDPs.

b) Collaborative Exploration: In order to learn the rewards of all the users
efficiently under the low-rank assumption, we need to deploy standard low rank
matrix estimation algorithms. These require specific kinds of linear
measurements (See Section 1.1). Without access to a generative model, the main
challenge in this setting is to obtain these linear measurements by querying
trajectories of carefully designed policies. We design such algorithms in
Section 4.

c) Functional Reward Maximization: In the case of linear MDPs, matrix
completion is more challenging since we observe measurements of the form
$e_{i}^{\intercal}\Theta^{*}\psi$ where $\Theta^{*}\in\mathbb{R}^{N\times d}$,
corresponding to the reward obtained by user $i$, with respect to an embedding
$\psi$. Estimating $\Theta^{*}$ under low rank assumptions requires the
distribution of $\psi$ to have certain isotropy properties (see Section 7). We
design a procedure which can sample-efficiently estimate policies which lead
to these isotropic measurements (Section 6).

d) Matrix Completion With Row-Wise Linear Measurements: For the linear MDP
setting, the low rank matrix estimation problem lies somewhere in between the
matrix completion (Recht, 2011; Jain et al., 2013) and matrix estimation with
restricted strong convexity (Negahban et al., 2009). We give a novel active
learning based algorithm where we estimate $\Theta^{*}$ row by row without any
assumptions like incoherence. This algorithm maybe of independent interest.
This is described in Section 7.

###  1.1 Related Works

Related Settings: Multi-task Reinforcement learning has been studied
empirically and theoretically (Brunskill & Li, 2013; Taylor & Stone, 2009;
D’Eramo et al., 2020; Teh et al., 2017; Hessel et al., 2019; Sodhani et al.,
2021). (Modi et al., 2017) considers learning a sequence of MDPs with side
information, where the parameters of the MDP varies smoothly with the context.
(Shah et al., 2020) assumes the optimal Q function $Q^{*}(s,a)$, when
represented as a $\mathcal{S}\times\mathcal{A}$ matrix, has low rank. With a
generative model, they obtain algorithms which makes use of this structure to
obtain a smaller sample complexity whenever the discount factor is bounded by
a constant. (Sam et al., 2022) improves the results in this setting with
additional assumptions on the transition matrices. Our setting is different in
that we consider multiple users, and do not assume a generative model. Our
main contribution is to efficiently obtain measurements conducive to matrix
completion. (Hu et al., 2021) considers a multi-task RL problem with linear
function approximation similar to our setting, but with the assumption of low-
rank Bellman closure, where the application of the Bellman operator retains
the low rank structure. They obtain a bound depending on the quantity
$N\sqrt{d}$ instead of $(N+d)$ like in our work. (Lei & Li, 2019) RL with low
rank assumptions in an experimental context.

Low Rank Matrix Estimation:The low rank assumption is popular in the
collaborative filtering literature and has been deployed successfully in a
variety of tasks (Bell & Koren, 2007; Gleich & Lim, 2011; Hsieh et al., 2012).
Low rank matrix estimation has been extensively studied in the statistics and
ML community for decades in the context of supervised learning (Candès & Tao,
2010; Negahban & Wainwright, 2011; Fazel, 2002; Chen et al., 2019; Jain et
al., 2013, 2017; Recht, 2011; Chen et al., 2020; Chi et al., 2019) in multi-
user collaborative filtering settings. The basic question is to estimate a
$d_{1}\times d_{2}$ matrix $M$ given linear measurements
$(x_{i}^{\intercal}My_{i})_{i=1}^{n}$ when the number of samples is much
smaller than $d_{1}\times d_{2}$ using the assumption that $M$ has low rank.

a) _Matrix Completion:_ $x_{i}$ and $y_{i}$ are standard basis vectors.
Typically $x_{i}$ and $y_{i}$ are picked uniformly at random and recovery
guarantees are given whenever the matrix $M$ is incoherent (Recht, 2011).

b) _Matrix Estimation:_ $x_{i}$ and $y_{i}$ are not restricted to be standard
basis vectors. Typically, they are chosen i.i.d such that the restricted
strong convexity holds (Negahban et al., 2009).

For the case of tabular MDPs, we use the matrix completion setting and for the
case of linear MDPs, our setting lies some where in between settings a) and b)
as explained above.

###  1.2 Notation

By $\|\cdot\|$ we denote the Euclidean norm and by $e_{1},\dots,e_{m}$ the
standard basis vectors of the space $\mathbb{R}^{m}$ for some
$m\in\mathbb{N}$. Let $\mathcal{S}^{d-1}:=\\{x\in\mathbb{R}^{d}:\|x\|=1\\}$,
$\mathcal{B}_{d}(r):=\\{x\in\mathbb{R}^{d}:\|x\|\leq r\\}$. For any $m\times
n$ matrix $A$ and a set $\Omega\subseteq[n]$ by $A^{\Omega}$, we denote the
sub-matrix of $A$ where the columns corresponding to $\Omega^{\complement}$
are deleted. By $\Delta(\mathcal{A})$, we denote the set of all Borel
probability measures over the set $\mathcal{A}$. In the sequel,

##  2 Problem Setting

We consider $N$ users indexed by $[N]$, each of them associated with an MDP
with the same state-space $\mathcal{S}$, action space $\mathcal{A}$, horizon
$H$ and transition matrices $\mathcal{P}=(P_{1},\dots,P_{H-1})$. Here
$P_{h}(\cdot|s_{h},a_{h})$ is a probability measure over $\mathcal{S}$, which
gives the distribution of the state at time $h+1$ given the action
$a_{h}\in\mathcal{A}$ was taken in state $s_{h}\in\mathcal{S}$ at time $h$.
Each user has a different reward denoted by
$\mathcal{R}_{u}=(R_{1u},\dots,R_{Hu})$ where
$R_{hu}:\mathcal{S}\times\mathcal{A}\to[0,1]$. Denote the MDP associated with
the user $u$ by
$\mathcal{M}_{u}:=(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R}_{u})$. For
the sake of simplicity, we will assume that the rewards are deterministic.

Assume that all the MDPs start at a random state $S_{1}$ with the same
distribution. Consider a policy $\Pi:=(\pi_{1},\dots,\pi_{H})$ where
$\pi_{h}:\Delta(\mathcal{A})\times\mathcal{S}\to\mathbb{R}^{+}$ is a kernel -
i.e, $\pi_{h}(\cdot|s)$ gives the probability distribution over actions given
a state $s$ at time $h$. By $(S_{1:H},A_{1:H})$ we denote the trajectory
$(S_{1},A_{1}),(S_{2},A_{2}),\dots,(S_{H},A_{H})\in\mathcal{S}\times\mathcal{A}$.
By $(S_{1:H},A_{1:H})\sim\mathcal{M}(\Pi)$ we mean the random trajectory under
the policy $\Pi$ \- where $A_{h}\sim\pi_{h}(\cdot|S_{h})$ and $S_{h+1}\sim
P_{h}(\cdot|S_{h},A_{h})$. That is, it is the trajectory of the MDP under the
policy $\Pi$. Define the value function of $\mathcal{M}_{u}$ under policy
$\Pi$ as:
$V(\Pi,\mathcal{M}_{u}):=\mathbb{E}_{(S_{1:H},A_{1:H})\sim\Pi}\sum_{h=1}^{H}R_{hu}(S_{h},A_{h})$.
We will call a policy $\hat{\Pi}_{u}$ to be $\epsilon$ optimal for
$\mathcal{M}_{u}$ if
$V(\hat{\Pi}_{u},\mathcal{M}_{u})\geq\sup_{\Pi}V(\Pi,\mathcal{M}_{u})-\epsilon$.
Our goal is to find $\epsilon$ optimal policies for every $u\in[N]$ under low
rank assumptions on the rewards $R_{uh}$. We assume that we are allowed to
pick any user $u$ and query a trajectory corresponding to any policy $\Pi$.

#### Reward Free Exploration:

The objective of reward free RL is to explore an MDP (without looking at the
rewards) such that we can obtain the optimal policy for every possible reward.
After collecting $K$ trajectories from the MDP sequentially (denoted by
$\mathcal{D}_{K}$), the algorithm outputs functions $\hat{\Pi}$ and $\hat{V}$
whose input is a reward function $\mathcal{R}=(R_{h}(\cdot,\cdot))_{h=1}^{H}$
(bounded between $[0,1]$) and the output is a nearly-optimal policy
$\hat{\Pi}(\mathcal{R})$ and its estimated value
$\hat{V}(\hat{\Pi}(\mathcal{R}))$ for this reward function. Denote the MDP
with this reward function by $\mathcal{M}_{\mathcal{R}}$. Given $\epsilon>0$
and $\delta\in[0,1]$, we let $K^{\mathsf{rf}}(\epsilon,\delta)$ to be such
that whenever $K\geq K^{\mathsf{rf}}(\epsilon,\delta)$, with probability at-
least $1-\delta$ we have:

a)
$\sup_{\mathcal{R}}|V(\hat{\Pi},\mathcal{M}_{\mathcal{R}})-\hat{V}(\hat{\Pi}(\mathcal{R}))|\leq\epsilon$
and b) $\hat{\Pi}$ is an $\epsilon$ optimal policy for
$\mathcal{M}_{\mathcal{R}}$ for every $\mathcal{R}$.

This setting was introduced in (Jin et al., 2020a). In this work, we will use
the reward free exploration algorithms in (Zhang et al., 2020) for tabular
MDPs and (Wagenmaker et al., 2022) for linear MDPs.

#### Tabular MDP Setting

$\mathcal{S}$ and $\mathcal{A}$ are finite sets. Denote the reward
$R_{hu}(s,a)$ by the $N\times|\mathcal{S}||\mathcal{A}|$ matrix $R_{h}$ where
$R_{h}(u,(s,a))=R_{hu}(s,a)$. We have the following low-rank assumption:

######  Assumption (Tab) 2.1.

The matrix $R_{h}$ has rank $r$ for some
$r\leq\frac{1}{2}\min(N,|\mathcal{S}||\mathcal{A}|)$.

#### Linear MDP Setting

Our definition is slightly different from the one in (Jin et al., 2020b): a)
we use two different embeddings for rewards and transitions and b) we impose
an $\ell_{1}$ constraint on the transition embedding and an $\ell_{2}$
constraint on the reward embedding (instead of $\ell_{2}$ on both). This is a
natural choice since transition embeddings describe a mixture of probability
measures as the law of the next state. On a technical level, $\ell_{1}$ norm
is natural when controlling the statistical error due to policy search
described in Section 6, which is based on the structural result in Lemma 6.1.

We consider embeddings $\phi:\mathcal{S}\times\mathcal{A}\to\mathbb{R}^{d}$,
$\psi:\mathcal{S}\times\mathcal{A}\to\mathbb{R}^{d}$ such that
$\|\phi(s,a)\|_{1}\leq 1$, $\|\psi(s,a)\|_{2}\leq 1$. We make the following
assumptions:

  1. 1.

There exists $\theta_{hu}\in\mathbb{R}^{d}$, $\|\theta_{hu}\|_{2}\leq\sqrt{d}$
such that $R_{hu}(s,a)=\langle\theta_{hu},\psi(s,a)\rangle$ and
$R_{h}(s,a)\in[0,1]$.

  2. 2.

There exist signed measures $\mu_{1h},\dots,\mu_{dh}$ over the space
$\mathcal{S}$ such that:
$P_{h}(\cdot|s,a)=\sum_{i=1}^{d}\mu_{ih}(\cdot)\langle\phi(s,a),e_{i}\rangle$

We will assume that $\mu_{i}$ are such that
$\|\int\mu_{ih}(ds)\phi(s,a)\pi(da|s)\|_{1}\leq 1$ and
$\sup_{i,h}\int|\mu_{ih}(ds)|\leq C_{\mu}$. This is true whenever $\mu_{ih}$
are probability measures. We consider different embeddings for transition
($\phi$) and reward ($\psi$) as the transition embeddings have a natural
$\|\cdot\|_{1}$ structure since they give linear combinations of measure which
make up $P_{h}(\cdot|s,a)$. We denote the $N\times d$ matrix whose $u$-th row
is $\theta_{uh}^{\intercal}$ to be $\Theta_{h}$. The low-rank assumption in
this setting takes the following form:

######  Assumption (Lin) 2.2.

. The $N\times d$ matrix $\Theta_{h}$ has rank $r\leq\tfrac{1}{2}\min(N,d)$.

We restrict our attention to policies given by some fixed policy space
$\mathbb{Q}$. As explained below in Section 3, this is necessitated by the
fact that our techniques are required to go beyond the standard RL setup and
might necessarily require non-deterministic policies. However, the space of
all possible policies can be very large and intractable. We refer to Section A
for the construction of randomized policy class $\mathbb{Q}$ such that it
contains all $\epsilon$-optimal policies for every possible linear reward.
With some abuse of notation, we define the total variation distance between
two kernels as:
$\mathsf{TV}(\pi_{h},\pi_{h}^{\prime}):=\sup_{s\in\mathcal{S}}\mathsf{TV}(\pi_{h}(\cdot|s),\pi^{\prime}_{h}(\cdot|s))$.
We define a distance over $\mathbb{Q}$ by
$D_{\mathbb{Q}}(\Pi^{(1)},\Pi^{(2)})=\sup_{h\in[H]}\mathsf{TV}(\pi_{h}^{(1)},\pi_{h}^{(2)})$,
where $\Pi^{(i)}=(\pi_{h}^{(i)}:h\in[H])$.

##  3 Connection to Multi-Agent RL

We now connect our results to multi-agent reinforcement learning in order to
demonstrate why the problem of collaborative RL as described above can be
hard. Low rank matrix estimation requires random measurements with specific
isotropy properties. For instance, matrix completion results are derived when
we observe entries from uniformly random indices ((Recht, 2011)). In our
context, this translates to sampling from specific distribution of state and
action $(S_{h},A_{h})$ at time $h$ by depolying a policy $\Pi$ over a
uniformly random user $u\in[N]$. Rather than maximizing a scalar reward, this
requires us to sample from a distribution with certain properties, going
beyond the framework of standard RL. In fact, Section A, we show that this
requires randomized policies even for simple MDPs. In the tabular MDP case, we
sidestep these issues with clever algorithm design. However, in the linear MDP
case, this does not seem to be feasible. We use the following connection to
mean-field limits of multi-agent RL as sketched below in order to solve the
sampling question.

Suppose we pick users uniformly at random ($U_{1},\dots,U_{T}\in[N]$) and
deploy a policy $\Pi$ for each of them with corresponding trajectories
$S^{(t)}_{1:H},A^{(t)}_{1:H}\sim\mathcal{M}(\Pi)$. We observe ‘linear
measurements’ of $\Theta_{h}$ of the form
$(e_{U_{t}},\psi(S^{(t)}_{h},A^{(t)}_{h}),e_{U_{t}}^{\intercal}\Theta_{h}\psi(S^{(t)}_{h},A^{(t)}_{h}))$.
To achieve matrix estimation, we need to query $(S^{(t)}_{h},A^{(t)}_{h})$
such that the distribution of $\psi(S^{(t)}_{h},A^{(t)}_{h})$ is ‘nearly
isotropic’ ( See (4) in Section 7). Let $\Gamma_{h}(\Pi)$ denote the
distribution of $S^{(1)}_{h},A^{(1)}_{h}$. In Theorem 6.2, we show that the
conditions in (4) are satisfied whenever $J(\Gamma_{h}(\Pi))>0$ for
$J:\Delta(\mathcal{S}\times\mathcal{A})\to\mathbb{R}$ given by:

$\displaystyle J(\Gamma_{h}(\Pi)):=$
$\displaystyle\quad\inf_{x\in\mathcal{S}^{d-1}}\mathbb{E}|\langle\psi(S^{(1)}_{h},A^{(1)}_{h}),x\rangle|\sqrt{d}-\xi
d\langle\psi(S^{(1)}_{h},A^{(1)}_{h}),x\rangle^{2}$

Our objective now is to find a policy by solving the following optimization
problem

$\arg\sup_{\Pi\in\mathbb{Q}}J(\Gamma_{h}(\Pi))$ (1)

This is similar to the mean field multi-agent control problem presented in
(Cammardella et al., 2020). To demonstrate the connection to multi-agent
systems, consider $n$ agents with the same MDP $\mathcal{M}$ and embeddings
$\phi,\psi$. A trajectory here corresponds to jointly and independently
running MDP associated with each agent with the same policy. The collective
reward of the system is given by $J(\hat{\Gamma}_{h})$, where
$\hat{\Gamma}_{h}$ denotes the empirical distribution of state-actions of the
$n$ agents at time $h$. Note that, picking a policy $\Pi_{n}$ to maximize this
reward is a reward maximization problem on the joint multi-agent system. And,
for any fixed policy $\Pi$, as $n\to\infty$,
$\hat{\Gamma}_{h}(\Pi)\to\Gamma_{h}(\Pi)$ under reasonable assumptions on the
state space via the law of large numbers. Hence $J(\hat{\Gamma}_{h}(\Pi))\to
J(\Gamma_{h}(\Pi))$ under continuity (in some appropriate distance between
probability measures). Therefore the planning problem in (1) is the same as
the multi-agent RL problem described above in the limit $n\to\infty$.

##  4 The Algorithm

Our algorithm proceeds in 4 phases. In phase 1, we run reward free RL, which
selects trajectories from uniformly random users since all the users share the
same MDP. Thus, this does not incur a large per-user sample complexity. This
step allows us to find near-optimal for any reward function of our choice. At
the end of phase 1, the main unknown will the reward matrix of the users.

Phase 2 is the main technical contribution of work. In phase 2, we use the
reward free RL output from phase 1 in order to design collaborative
exploration policies. These policies obtain the right linear measurements of
the low-rank reward matrix so that we can successfully apply matrix estimation
algorithms in phase 3. Phase 4 uses the reward estimate from phase 3 and the
reward free RL output from phase 1 in order to learn the optimal policy for
every user.

###  4.1 Tabular MDP Case:

Phase 1: Reward Free Exploration We run the reward free RL algorithm in (Zhang
et al., 2020) for
$K^{\mathsf{rf}}(\tfrac{\epsilon}{8},\tfrac{\delta}{2})=C\frac{|\mathcal{S}||\mathcal{A}|H^{2}\left(|\mathcal{S}|+\log(\tfrac{1}{\delta})\right)}{\epsilon^{2}}\mathsf{polylog}(\tfrac{|\mathcal{S}||\mathcal{A}|H}{\epsilon})$
time steps by picking the MDP corresponding to a uniformly random user
whenever the reward free RL algorithm queries a trajectory. Let the output of
the reward free RL algorithm be $\hat{\Pi}$ and $\hat{V}$.

Phase 2: Querying the Reward Matrix In this phase we query a ‘uniform mask’
with the parameter $p$ for the reward matrix $R_{h}$ using Algorithm 1. For
each $(s,a)\in\mathcal{S}\times\mathcal{A}$ and $h\in[H]$, maintain a counter
$T_{h,(s,a)}$ for $(s,a)\in\mathcal{S}\times\mathcal{A}$ and $h\in[H]$,
initialized at $0$. Given the ‘active sets’
$\mathcal{G}=(\mathcal{G}_{h})_{h\in[H]}\subseteq\mathcal{S}\times\mathcal{A}$
and $h\in[H]$, we define the reward
$\mathcal{J}(;\mathcal{G})=(J_{1},\dots,J_{h})$ by

$J_{h}(s,a;\mathcal{G})=\mathbbm{1}((s,a)\in\mathcal{G}_{h})$ (2)

We will denote this reward by $\mathcal{J}(\cdot;\mathcal{G})$. Initialize
active set $\mathcal{G}=(\mathcal{G}_{h})_{h\in[H]}$ such that
$\mathcal{G}_{h}=\mathcal{S}\times\mathcal{A}$. We initialize the reward
matrix $\hat{R}_{h}(u,(s,a))=*$, where $*$ denotes unknown entry. This
algorithm terminates when it detects that sufficient number of samples have
been collected for matrix completion.

Algorithm 1 Uniform Mask Sampler for Tabular MDPs

Output:Active sets $\mathcal{G}=(\mathcal{G}_{h})_{h\in[H]}$, Partially
complete matrix $\hat{R}_{h}$

$t\leftarrow 0$
$\hat{P}^{\mathcal{G}}\leftarrow\hat{V}(\mathcal{J}(\cdot;\mathcal{G}))$

$\hat{\Pi}^{\mathcal{G}}\leftarrow\hat{\Pi}(\mathcal{J}(\cdot;\mathcal{G}))$

while $\hat{P}^{\mathcal{G}}>\tfrac{\epsilon}{2}$ do

$U_{t}\leftarrow\mathsf{Unif}([N])$ {Pick a user uniformly at random}

$S^{(t)}_{1:H},A^{(t)}_{1:H},R^{(t)}_{1:H}\sim\mathcal{M}_{U_{t}}(\hat{\Pi}^{\mathcal{G}})$
{Query trajectory}

for $h\in[H]$ do

if $(S^{(t)}_{h},A^{(t)}_{h})\in\mathcal{G}_{h}$ and
$R_{h}(U_{t},(S^{(t)}_{h},A^{(t)}_{h}))=*$ then

$T_{h,(S^{(t)}_{h},A^{(t)}_{h})}\leftarrow T_{h,(S^{(t)}_{h},A^{(t)}_{h})}+1$
{Increment count}

$\hat{R}_{h}(U_{t},(S^{(t)}_{h},A^{(t)}_{h}))\leftarrow R^{(t)}_{h}$ {Fill
Missing Entry}

end if

if $T_{h,(S^{(t)}_{h},A^{(t)}_{h})}=Np$ then

$\mathcal{G}_{h}\leftarrow\mathcal{G}_{h}\setminus\\{(S^{(t)}_{h},A^{(t)}_{h})\\}$
{Remove element from Active Set}

end if

end for

$t\leftarrow t+1$

$\hat{P}^{\mathcal{G}}\leftarrow\hat{V}(\mathcal{J}(\cdot;\mathcal{G}))$

$\hat{\Pi}^{\mathcal{G}}\leftarrow\hat{\Pi}(\mathcal{J}(\cdot;\mathcal{G}))$

end while

Phase 3: Reward Matrix Completion We receive $\mathcal{G}_{h}$ and the
partially observed matrix $\hat{R}_{h}$ for each $h$ as the output of
Algorithm 1. By $\hat{R}_{h}^{\mathcal{G}^{\complement}_{h}}$, we denote the
sub-matrix where the columns corresponding to $\mathcal{G}_{h}$ are deleted.
We use the nuclear norm minimization algorithm given in (Recht, 2011) to
recover $R_{h}^{\mathcal{G}_{h}^{\complement}}$ from
$\hat{R}_{h}^{\mathcal{G}_{h}^{\complement}}$ for every $h\in[H]$.

Phase 4: Computing the Optimal Policy Phase 3 outputs the completed sub-matrix
$R_{h}^{\mathcal{G}_{h}^{\complement}}$, where only the columns corresponding
to $|\mathcal{G}_{h}^{\complement}|$ are recovered. We construct the recovered
matrix $\bar{R}_{h}$ by setting
$\bar{R}_{h}^{\mathcal{G}_{h}^{\complement}}=R_{h}^{\mathcal{G}_{h}^{\complement}}$
and $\bar{R}_{h}^{\mathcal{G}_{h}}=0$. We compute the optimal policy for each
user using the rewards from $\bar{R}_{h}$ via the output of the reward free
RL, $\hat{\Pi}$, from Phase 1.

###  4.2 Linear MDP Case:

Phase 1 : Reward Free RL We run the reward free RL algorithm for Linear MDPs
from (Wagenmaker et al., 2022), with error $\epsilon$ and probability of
failure $\frac{\delta}{4}$. We use trajectories from random users whenever a
trajectory is queried. Here,
$K^{\mathsf{rf}}(\epsilon,\delta/4)=\frac{CdH^{5}(d+\log(\tfrac{1}{\delta}))}{\epsilon^{2}}+\frac{Cd^{9/2}H^{6}\log^{4}(\tfrac{1}{\delta})}{\epsilon}$.

Phase 2: Querying Linear Measurements of the Reward Matrix We obtain policies
whose trajectory data allows low rank matrix estimation of the reward matrix.

Step 1: For each time step $h\in[H]$, we want to query obtain samples
$(s^{(t)}_{h},a^{(t)}_{h})$ such that
$\sum_{t=1}^{T}\phi(s^{(t)}_{h},a^{(t)}_{h})\phi^{\intercal}(s^{(t)}_{h},a^{(t)}_{h})\succeq\kappa^{2}I$.
This can be done by Algorithm 2. Given a projector $Q$ to some subspace of
$\mathbb{R}^{d}$, by $\mathcal{Q}_{h,Q}$ denote the reward
$\|Q\phi(s,a)\|^{2}$ at time $h$ and $0$ otherwise. The termination condition
ensures that we see enough data in all directions $\phi$, which allows us to
find collaborative exploration policy below.

Step 2: Using the observations given in Step 1, we compute the policy
$\hat{\Pi}^{f,h}$ which approximately satisfies the property given in
Assumption 5.6. This procedure is described in Section 6.

Algorithm 2 Well Conditioned Matrix Sampler

Input:Total time $T$; Tolerance $\gamma$; lower isometry $\kappa$

Output:$\phi_{ht},S_{(h+1)t}$ for $h\in[H-1]$, $t\in[T]$

for $h=1$ to  $h=H-1$ do

$Q\leftarrow I$ $\hat{\Pi}^{Q}\leftarrow\hat{\Pi}(\mathcal{Q}_{h,Q})$

$G_{\phi,h}\leftarrow 0$ {Grammian initialized to $0$}

for $t=1$ : to $t=T$ do

$U_{t}\sim\mathsf{Unif}([N])$ {Pick a uniformly random user}

$S_{1:H},A_{1:H}\sim\mathcal{M}_{U_{t}}(\hat{\Pi}^{Q})$ {Obtain Trajectory}

$\phi_{ht}\leftarrow\phi(S_{h},A_{h})$ $S_{(h+1)t}\leftarrow S_{h+1}$

$G_{\phi,h}\leftarrow G_{\phi,h}+\phi_{ht}\phi_{ht}^{\intercal}$ {Update
Grammian}

if $\exists y\in\mathbb{S}^{d-1}:y^{\intercal}G_{\phi,h}y<\kappa^{2}$ then

$Q\leftarrow$ eigenspace of $G_{\phi,h}$ with eigenvalues $<\kappa^{2}$

$\hat{\Pi}^{Q}\leftarrow\hat{\Pi}(\mathcal{Q}_{h,Q})$

end if

end for

end for

Phase 3: Estimating Low Rank Reward Matrix For this, we use the active
learning procedure given in Section 7 via row-wise linear measurements along
with the policy $\Pi^{\mathsf{MC},h}=\hat{\Pi}^{f,h}$, which was computed in
Phase 2.

Phase 4: Computing the Optimal Policy Once the reward matrix $\Theta_{h}$ have
been reconstructed for every $h$ in Phase 3, we use the output of reward free
RL in order to compute the $\epsilon$ optimal policy for each user.

##  5 Main Results

###  5.1 Tabular MDP:

Incoherence is a standard assumption for low rank matrix completion. This
ensures that the matrix is not too sparse so that sparse measurements are
sufficient to learn it. The following definition is used in (Recht, 2011).

######  Definition 5.1.

Given a $r$ dimensional sub-space $U$ of $\mathbb{R}^{n}$ , we define the
coherence of $U$ as:

$\mu(U):=\frac{n}{r}\sup_{1\leq i\leq n}\|P_{U}e_{i}\|^{2}\,.$

A $n_{1}\times n_{2}$ matrix $M$ with singular value decomposition $U\Sigma
V^{\intercal}$ is called $(\mu_{0},\mu_{1})$ coherent if:

a) The coherence of the row and column spaces of $M$ are at-most $\mu_{0}$ b)
The absolute value of every entry of $UV^{\intercal}$ is bounded above by
$\mu_{1}\sqrt{\frac{r}{n_{1}n_{2}}}$.

Given a policy $\Pi$, and $\Omega\subseteq\mathcal{S}\times\mathcal{A}$, by
$P^{\Pi}_{h}(\Omega)$ we denote the probability that at time $h$ we have
$(S_{h},A_{h})\in\Omega$ under the policy $\Pi$.

######  Assumption (Tab) 5.2.

Given the reward matrix $R_{h}$ and
$\Omega_{h}\subset\mathcal{S}\times\mathcal{A}$, recall the notation for the
sub-matrix $R_{h}^{\Omega_{h}}$ of $R_{h}$. If
$\sup_{\Pi}P^{\Pi}(\Omega_{h}^{\complement})<\epsilon$ have:

1) $R_{h}^{\Omega_{h}}$ is $(\mu_{0},\mu_{1})$ incoherent 2)
$|\Omega^{\complement}_{h}|\leq\frac{|\mathcal{S}|}{2}$

The incoherence assumption for $R_{h}^{\Omega}$ makes sense since the set
$\Omega^{\complement}$ cannot be easily reached with _any_ policy with a
probability larger than $\epsilon$. In fact we can arrive at an $\epsilon$
optimal policy for the original reward by just setting the rewards at
$\Omega^{\complement}$ to be $0$. These can be thought of as redundant states
which do not matter for our RL model with _any_ reward.

######  Theorem 5.3.

Suppose Assumption (Tab) 2.1, 5.2 hold. Let the parameter
$p=C\frac{\max(\mu_{1}^{2},\mu_{0})r(N+|\mathcal{S}||\mathcal{A}|)\log^{2}|\mathcal{S}||\mathcal{A}|\log(\tfrac{H}{\delta})}{N|\mathcal{S}||\mathcal{A}|}\,.$
for some large enough constant $C$. Assume that $|\mathcal{S}||\mathcal{A}|$
and $N$ are large enough such that $p<1/2$. Then, with probability at-least
$1-\delta$, we can find an $\epsilon$ optimal policy $\hat{\Pi}_{u}$ for every
user $u\in[N]$ whenever the total number of trajectories queried is:

$\displaystyle
C\frac{|\mathcal{S}||\mathcal{A}|H^{2}\left(|\mathcal{S}|+\log(\tfrac{1}{\delta})\right)}{\epsilon^{2}}\mathsf{polylog}(\tfrac{|\mathcal{S}||\mathcal{A}|H}{\epsilon})$
$\displaystyle\quad+\frac{C\max(\mu_{1}^{2},\mu_{0})r(N+|\mathcal{S}||\mathcal{A}|)H\log^{2}|\mathcal{S}||\mathcal{A}|\log(\tfrac{H}{\delta})}{\epsilon}$

######  Remark 5.4.

For large $N$, the number of trajectories queried per user is
$\tilde{O}(\frac{rH\log^{2}(|\mathcal{S}||\mathcal{A}|)}{\epsilon})$, which is
an exponential improvement in the state-space size dependence when compared to
the minimax rate of $\frac{|\mathcal{S}||\mathcal{A}|H^{2}}{\epsilon^{2}}$
(Dann & Brunskill, 2015) for learning a single MDP. Every phase in the
algorithm has polynomial computational complexity in
$N,|\mathcal{S}||\mathcal{A}|$ and $\frac{1}{\epsilon}$. The probability $p$
is chosen such that
$p|\mathcal{S}||\mathcal{A}|N=\tilde{O}(r(|\mathcal{S}||\mathcal{A}|+N))$,
which is the number of free parameters required to describe a rank $r$ matrix.

###  5.2 Linear MDP

######  Assumption (Lin) 5.5.

There exists a $\gamma>0$ such that for every $x\in\mathcal{S}^{d-1}$, and
every $h\in[H]$ there exists a policy $\Pi_{x,h}$ such that whenever
$S_{1:H},A_{1:H}\sim\Pi_{x,h}$,
$\mathbb{E}\langle\phi(S_{h},A_{h}),x\rangle^{2}\geq\gamma$

The assumption above shows that we can obtain information about all
directions. If this does not hold for any $\gamma$, then $\phi(S_{h},A_{h})$
does not have any component in some direction $x_{0}$ with _any_ policy. Thus,
we can remove the sub-space spanned by $x_{0}$ and make the embedding space
$\mathbb{R}^{d-1}$ at time $h$.

######  Assumption (Lin) 5.6.

There exist $\zeta,\xi>0$ such that for every $h\in[H]$, there exists a policy
$\Pi^{h,\zeta,\xi}\in\mathbb{Q}$ such that whenever
$S_{1:H},A_{1:H}\sim\mathcal{M}(\Pi^{h,\zeta,\xi})$, we have:

$\inf_{x\in\mathcal{S}^{d-1}}\mathbb{E}\left[|\langle
x,\psi(S_{h},A_{h})\rangle|\sqrt{d}-\xi d\langle
x,\psi(S_{h},A_{h})\rangle^{2}\right]\geq\zeta$

The assumption above ensures that there exist measurements $\psi(S_{h},A_{h})$
which are conducive to low rank matrix estimation as considered in Section 7.
This means that $|\langle\psi(S_{h},A_{h}),x\rangle|=\Omega(1/\sqrt{d})$ just
like an isotropic random vector, which gives us information about all
directions. However, this condition is much looser than the assumption of
uniform distribution on the sphere.

######  Assumption (Lin) 5.7.

For any $1\geq\eta>0$, there exists an $\eta$ net for $\mathbb{Q}$, denoted by
$\hat{\mathbb{Q}}_{\eta}$ such that $\log|\hat{\mathbb{Q}}_{\eta}|\leq
D\log(\frac{1}{\eta})$.

We refer to Section A, where we justify this assumption. We first demonstrate
that deterministic policies which are sufficient for reward maximization (as
used in (Jin et al., 2020b)) cannot be used in this context, so a set of
stochastic policies is required. We then construct such policy classes with
$D=O(dH\log dH\log\log(|\mathcal{A}|))$.

######  Theorem 5.8.

Suppose Assumptions (Lin) 2.2 5.5 5.6 and 5.7 hold and suppose
$\epsilon<\frac{\gamma}{2}$. In Algorithm 2, we set
$\kappa=\frac{CC_{\mu}dH\sqrt{dH+D}(\sqrt{d}+\xi
d)}{\zeta}\sqrt{\log\left(\tfrac{C_{\mu}H(d+D)}{\zeta\gamma\delta}\right)}$
and $T=C\frac{\kappa^{2}d}{\gamma^{2}}\log\tfrac{d\kappa}{\gamma}$.

Then, with probability at least $1-\delta$, our algorithm finds $\epsilon$
optimal policy for every user $u\in[N]$ with the total number of trajectories
being bounded by: $T_{\mathsf{rf}}+T_{\mathsf{pol}}+T_{\mathsf{mat-comp}}$,
where:

$T_{\mathsf{rf}}=\frac{dH^{5}(d+\log(\tfrac{1}{\delta}))}{\epsilon^{2}}+\frac{d^{9/2}H^{6}\log^{4}(\tfrac{1}{\delta})}{\epsilon}\,,$
$T_{\mathsf{pol}}=\frac{C_{\mu}^{2}d^{5}H^{3}(dH+D)\log^{2}\left(\tfrac{C_{\mu}H(d+D)}{\zeta\gamma\delta}\right)}{\zeta^{2}\gamma^{2}}\,,$
$\displaystyle T_{\mathsf{mat-comp}}$ $\displaystyle=C\frac{Hr(N+d\log
N)\log\tfrac{d}{\zeta\xi}}{{\zeta^{2}\xi^{2}}}$
$\displaystyle\quad+\frac{H\log N\log\left(\frac{\log
N}{\delta}\right)}{\zeta^{2}\xi^{2}}\,.$ (3)

######  Remark 5.9.

When $N$ is very large, the per user sample complexity is $O(Hr)$, which is
much better than the mini-max optimal complexity of $\Omega(d^{2}H^{2})$
(Wagenmaker et al., 2022). While Phases 1 and 2 of the algorithm have a
computational complexity which is polynomial in $d$ and $\frac{1}{\epsilon}$,
the optimization problems posed in Phase 3 and 4 are not necessarily
polynomial time. We leave the computational aspects to future work. The sample
complexity $\tilde{O}(r(N+d))$ corresponds to the number of free parameters
required to describe a rank $r$ matrix.

##  6 Obtaining Policies With Given Statistics

In this section, we consider the Linear MDP setting and describe the sub-
routine described in Step 2 of Phase 2 of the algorithm where we compute a
policy $\hat{\Pi}^{f,h}$ such that the law of $\phi(S_{h},A_{h})$ under this
policy approximately satisfies the property given in Assumption 5.6. This is
required in order to use the guarantees for low matrix estimation in Phase 3,
which is described in Section 7. We first state a structural lemma which
characterizes the law of $S_{h+1},A_{h+1}$ under any policy $\Pi$.

######  Lemma 6.1.

Consider any policy $\Pi=(\pi_{1},\dots,\pi_{H-1},\pi_{H})$ to the MDP
$\mathcal{M}$. Let $S_{1:H},A_{1:H}\sim\mathcal{M}(\Pi)$. Then for any
bounded, measurable function $g:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$, we
have:

$\mathbb{E}g(S_{h+1},A_{h+1})=\sum_{i=1}^{d}\nu_{i}\int
g(s,a)d\mu_{ih}(ds)\pi_{h+1}(da|s)$

Where $\nu_{i}:=\langle\mathbb{E}\phi(S_{h},A_{h}),e_{i}\rangle$

We now want to estimate certain statistics under any policy using available
data, obtained from the output of Algorithm 2. Notice that the output of
Algorithm 2 gives a sequence of random variables
$(\phi_{h1},s_{(h+1)1}),\dots,(\phi_{hT},s_{(h+1)T})\in\mathbb{R}^{d}\times\mathcal{S}$
such that
$(s_{(h+1)l})_{l=1}^{T}|(\phi_{hl})_{l=1}^{T}\sim\prod_{l=1}^{T}\left(\sum_{i=1}^{d}\langle\phi_{hl},e_{i}\rangle\mu_{hi}(\cdot)\right)$
and $G_{\phi,h}:=\sum_{t=1}^{T}\phi_{ht}\phi_{ht}^{\intercal}$. For any
measurable function $g:\mathcal{S}\times\mathcal{A}\to\mathbb{R}^{K}$,
$\nu\in\mathbb{R}^{d}$ such that $\|\nu\|\leq 1$ and any randomized policy
$\Pi=(\pi_{1},\dots,\pi_{H})$ we define:

  1. 1.

$\mathcal{T}_{1}(g,\pi_{1})=\mathbb{E}\int g(S_{1},a)\pi_{1}(da|S_{1})$

  2. 2.

$\mathcal{T}_{h+1}(g;\nu,\pi_{h+1})=\sum_{i=1}^{d}\langle\nu,e_{i}\rangle\int\mu_{ih}(ds)\pi_{h+1}(da|s)g(s,a)\,\text{
when }h\leq H-1$

  3. 3.

$E^{\nu}_{1}(\Pi):=\|\mathcal{T}_{1}(\phi,\pi_{1})-\nu\|_{1}$

  4. 4.

$E^{\nu}_{h}(\Pi)=\inf_{\nu_{1},\dots,\nu_{h-1}\in\mathcal{B}_{d}(1)}F(\Pi,\nu_{1},\dots,\nu_{h-1},\nu)\,$
whenever $h>1$

Define $\alpha_{ht,\nu}=\phi_{ht}^{\intercal}G_{\phi,h}^{-1}\nu$. We estimate
these operators from data as follows:

  1. 1.

$\hat{\mathcal{T}}_{1}(g,\pi_{1})=\frac{1}{T}\sum_{t=1}^{T}\int
g(s_{1t},a)\pi_{1}(da|s_{1t})$

  2. 2.

$\hat{\mathcal{T}}_{h+1}(g,\nu,\pi_{h+1})=\sum_{t=1}^{T}\alpha_{ht,\nu}\int
g(s_{(h+1)t},a)\pi_{h+1}(da|s_{(h+1)t})$

  3. 3.

$\hat{E}^{\nu}_{1}(\Pi):=\|\hat{\mathcal{T}}_{1}(\phi,\pi_{1})-\nu\|_{1}$

  4. 4.

$\hat{E}^{\nu}_{h}(\Pi)=\inf_{\nu_{1},\dots,\nu_{h-1}\in\mathcal{B}_{d}(1)}\hat{F}(\Pi,\nu_{1},\dots,\nu_{h-1},\nu)\,$
whenever $h>1$

Where, for $h>1$ and $\nu_{1},\dots,\nu_{h-1}\in\mathcal{B}_{d}(1)$, we have
defined:

  1. 1.

$F(\Pi,\nu_{1},\dots,\nu_{h-1},\nu_{h}):=E_{1}^{\nu_{1}}(\Pi)+\sum_{j=2}^{h}\|\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\nu_{j}\|_{1}$

  2. 2.

$\hat{F}(\Pi,\nu_{1},\dots,\nu_{h-1},\nu_{h}):=\hat{E}_{1}^{\nu_{1}}(\Pi)+\sum_{j=2}^{h}\|\hat{\mathcal{T}}_{j}(\phi,\nu_{j-1},\pi_{j})-\nu_{j}\|_{1}$

Define $f(s,a;x):=|\langle x,\psi(s,a)\rangle|\sqrt{d}-\xi d\langle
x,\psi(s,a)\rangle^{2}$. The output of our method is:

  1. 1.

$\displaystyle\hat{\Pi}^{f,1}=$
$\displaystyle\arg\sup_{\Pi=(\pi_{1},\dots,\pi_{H})\in\mathbb{Q}}\inf_{x\in\mathcal{S}^{d-1}}\hat{\mathcal{T}}_{1}(f(\cdot;x),\pi_{1})$

  2. 2.

$\displaystyle(\hat{\Pi}^{f,h},\hat{\nu})=$
$\displaystyle\quad\arg\sup_{\begin{subarray}{c}\nu\in\mathcal{B}(1)\\\
\Pi=(\pi_{1},\dots,\pi_{H})\in\mathbb{Q}\end{subarray}}\inf_{x\in\mathcal{S}^{d-1}}\hat{\mathcal{T}}_{h}(f(\cdot;x);\nu,\pi_{h})$

whenever $h>1$, subject to
$\hat{E}^{\hat{\nu},h-1}(\hat{\Pi}^{f,h})\leq\eta_{0}$

  3. 3.

Assign output: $\hat{\Pi}^{\zeta,\xi,h}=\hat{\Pi}^{f,h}$

The idea behind this method is as follows. First, using the output of
algorithm 2, we construct $\hat{\mathcal{T}}_{h}(g,\nu,\pi_{h})$, which
approximates the functional $\mathcal{T}_{h}(g,\nu,\pi_{h})$ uniformly for
every $\nu,\pi_{h}$. This is shown in Lemma E.1 in the appendix. We will show
in Theorem 6.2 that obtaining policies which can be used with the matrix
completion routine reduces to picking a policy
$\Pi^{f,h}=(\pi_{1},\dots,\pi_{H})$ such that whenever
$S_{1:H},A_{1:H}\sim\mathcal{M}(\Pi^{f,h})$, we must have:
$\inf_{x\in\mathcal{S}^{d-1}}\mathbb{E}f(S_{h},A_{h};x)\geq\zeta$. Now, we use
Lemma 6.1 to conclude that if such a policy exists, then there exist
$\nu_{1},\dots,\nu_{h-1}$ such that $\mathbb{E}\phi(S_{j},A_{j})=\nu_{j}$ and
$\inf_{x\in\mathcal{S}^{d-1}}\mathcal{T}_{h}(f(;x);\nu_{h},\pi_{h})\geq\zeta$.
Since we only have sample access, we find such a policy approximately by
optimizing using the estimates $\hat{\mathcal{T}}$ instead of the exact
functional $\mathcal{T}$ as described above.

######  Theorem 6.2.

We condition on the event $G_{\phi,h}\geq\kappa^{2}I$ for every $h\in[H]$. Let
$\kappa,\eta,\eta_{0}$ be such that for some small enough constants
$c_{0},c>0$ and a large enough constant $C>0$:

  1. 1.

$\eta\leq c\frac{\zeta}{C_{\mu}dH(\sqrt{d}+\xi
d)H}\sqrt{\tfrac{\kappa^{2}}{T}}\,;\quad\eta_{0}=c_{0}\frac{\zeta}{C_{\mu}(\sqrt{d}+d\xi)}$

  2. 2.

$\kappa\geq C\frac{C_{\mu}(\sqrt{d}+\xi
d)dH}{\zeta}\sqrt{\log\left(\tfrac{dH|\hat{\mathbb{Q}}_{\eta}|}{\delta}\right)+dH\log\left(\tfrac{d}{\eta}\right)}$

Recall the policy $\hat{\Pi}^{f,h}$. Suppose the Assumption 5.6 holds. Then,
with probability at-least $1-\frac{\delta}{4}$ we obtain the policy
$\hat{\Pi}^{f,h}$ is such that whenever
$S_{1:H},A_{1:H}\sim\mathcal{M}(\hat{\Pi}^{f,h})$, we have:

$\inf_{x\in\mathcal{S}^{d-1}}\mathbb{E}f(S_{h},A_{h};x)\geq\frac{\zeta}{2}$

This implies that $\psi(S_{h},A_{h})$ satisfies
$\mathbb{E}|\langle\psi(S_{h},A_{h}),x\rangle|\geq\frac{\zeta}{2\sqrt{d}};\quad\mathbb{E}\psi(S_{h},A_{h})\psi_{ik}^{\intercal}\leq\frac{1}{d\xi^{2}}$

##  7 Matrix Estimation with Row-wise Linear Measurements

We now describe the active learning based low rank matrix estimation
procedure. For an unknown rank $r$ matrix $\Theta^{*}$ (corresponding to
$\Theta^{*}_{h}$ in the definition of Linear MDPs) of dimensions $N\times d$,
we are allowed to query samples of the form
$(e_{i},\psi,e_{i}^{\intercal}\Theta^{*}\psi)$ for any $i\in[N]$ of our choice
and $\psi=\psi(S_{h},A_{h})$ where
$S_{1:H},A_{1:H}\sim\mathcal{M}(\Pi^{\mathsf{MC},h})$, for some input policy
$\Pi^{\mathsf{MC},h}$. This corresponds to running the MDP of user $i$, with
the policy $\Pi^{\mathsf{MC},h}$ and observing the reward at time $h$, given
by $\langle e_{i},\Theta_{h}^{*}\Psi(S_{h},A_{h})\rangle$. We want to estimate
the matrix $\Theta^{*}$ from these samples with high-probability.

###  7.1 The Estimator

Given any $N\times d$ matrix $\Delta$, by $\Delta_{i}^{\intercal}$, we denote
its $i$-th row. Given $K\in\mathbb{N}$, and a sequence of vectors
$\Psi=(\psi_{ik}\in\mathbb{R}^{d})_{i\in[N],k\in[K]}$.

$L(\Delta,\Psi):=\frac{1}{NK}\sum_{i=1}^{N}\sum_{k=1}^{K}|\langle\Delta_{i},\psi_{ik}\rangle|^{2}$

We estimate $\Theta^{*}$ row-wise using the following iterative procedure,
where recover some rows of $\Theta^{*}$ into $\hat{\Theta}$ in each iteration
and obtain the corresponding linear measurements of $\Theta^{*}$. Letting the
set of unknown rows at iteration $t$ to be $\bar{I}_{t-1}$ (with
$\bar{I}_{0}=[N]$). We draw a fresh sequence of vectors $\Psi^{(t)}$ from some
distribution, we then recover some rows
$\bar{I}_{t}^{\complement}\subseteq\bar{I}_{t-1}$ of $\Theta^{*}$ and store
them in $\hat{\Theta}$.

  1. 1.

Draw $\Psi^{(t)}=(\psi_{ik}^{(t)})_{k\in[K_{t}],i\in\bar{I}_{t-1}}$, we obtain
$\theta^{*}_{ik}=e_{i}^{\intercal}\Theta^{*}\psi_{ik}^{(t)}$.

  2. 2.

Consider the loss function

$\displaystyle L(\Theta-\Theta^{*},\Psi^{(t)})$
$\displaystyle:=\frac{1}{K_{t}|\bar{I}_{t-1}|}\sum_{i\in\bar{I}_{t-1}}\sum_{k=1}^{K}|\langle\Theta_{i},\psi^{(t)}_{ik}\rangle-\theta_{ik}^{*}|^{2}\,.$

  3. 3.

Find a matrix $\Theta$ with rank $\leq r$ such that
$L(\Theta-\Theta^{*},\Psi^{(t)})=0$.

  4. 4.

Initialize $\bar{I}_{t}\leftarrow\emptyset$.

  5. 5.

For every $i\in\bar{I}_{t-1}$, draw $K$ fresh samples using
$\tilde{\psi}_{i1}^{(t)},\cdots,\tilde{\psi}_{iK}^{(t)}$ and compute
$\sum_{k=1}^{K}|\langle\Theta_{i},\tilde{\psi}^{(t)}_{ik}\rangle-\theta_{ik}^{*}|^{2}$.
If
$\sum_{k=1}^{K}|\langle\Theta_{i},\tilde{\psi}^{(t)}_{ik}\rangle-\theta_{ik}^{*}|^{2}>0$
then add $i$ to $\bar{I}_{t}$ i.e.,
$\bar{I}_{t}\leftarrow\bar{I}_{t}\cup\\{i\\}$.

  6. 6.

End routine when $\bar{I}_{t}=\emptyset$.

Suppose $\psi_{ik}$ are i.i.d random vectors such that there exist
$\zeta,\xi>0$ such that for any $x\in\mathbb{R}^{d}$, $\|x\|=1$ we have:

$\displaystyle\|\psi_{ik}\|\leq 1\text{ almost
surely};\quad\mathbb{E}|\langle\psi_{ik},x\rangle|\geq\frac{\zeta}{\sqrt{d}};$
$\displaystyle\quad\mathbb{E}\psi_{ik}\psi_{ik}^{\intercal}\leq\frac{1}{d\xi^{2}}$
(4)

To give some intuition, the second condition above means that given any vector
$x$, there is some overlap between the random vector $\psi$ and $x$, ensuring
that every measurement gives us some information helping us to complete the
matrix. The third assumption is a standard bound on the covariance matrix.
Then we have the following theorem whose proof is presented in Section F.

######  Theorem 7.1.

Assume that $\sup_{i}\|\Theta^{*}_{i}\|\leq C_{\theta}$ and that the
distribution of $\psi^{(t)}_{ik}$ satisfies (4). Suppose
$K_{t}|\bar{I}_{t-1}|=C\frac{r|\bar{I}_{t-1}|+dr}{\zeta^{2}\xi^{2}}\log\tfrac{d}{\zeta\xi}+C\frac{\log\left(\frac{\log
N}{\delta}\right)}{\zeta^{2}\xi^{2}}$. With probability at-least $1-\delta$,
the algorithm terminates after $\log N$ iterations and the output
$\hat{\Theta}$ satisfies $\hat{\Theta}=\Theta^{*}$. Therefore, with
probability at-least $1-\delta$, the sample complexity for estimation of
$\Theta^{*}$ is:

$C\frac{r(N+d\log N)}{\zeta^{2}\xi^{2}}\log\tfrac{d}{\zeta\xi}+C\frac{\log
N\log\left(\frac{\log N}{\delta}\right)}{\zeta^{2}\xi^{2}}$

##  8 Discussion

In this work, we designed methods to perform collaborative exploration of a
number of MDPs with near optimal sample complexity. In particular, we
encountered and solved the important problem of exploring such that the data
can be used down-stream to learn the optimal policy for every one of the MDPs.
We also established connections to mean-field limits of multi-agent
reinforcement learning problems. In future work, we hope to use the
observations in the current work in order to design collaborative RL
algorithms based on practically deployed RL algorithms like PPO, DQN and TD3.

## References

  * Afsar et al. (2022) Afsar, M. M., Crump, T., and Far, B.  Reinforcement learning based recommender systems: A survey.  _ACM Computing Surveys_ , 55(7):1–38, 2022.
  * Agarwal et al. (2019) Agarwal, A., Jiang, N., Kakade, S. M., and Sun, W.  Reinforcement learning: Theory and algorithms.  _CS Dept., UW Seattle, Seattle, WA, USA, Tech. Rep_ , pp. 10–4, 2019.
  * Ariu et al. (2020) Ariu, K., Ryu, N., Yun, S.-Y., and Proutière, A.  Regret in online recommendation systems.  _Advances in Neural Information Processing Systems_ , 33:21141–21150, 2020.
  * Bell & Koren (2007) Bell, R. M. and Koren, Y.  Scalable collaborative filtering with jointly derived neighborhood interpolation weights.  In _Seventh IEEE international conference on data mining (ICDM 2007)_ , pp. 43–52. IEEE, 2007.
  * Boucheron et al. (2013) Boucheron, S., Lugosi, G., and Massart, P.  _Concentration inequalities: A nonasymptotic theory of independence_.  Oxford university press, 2013.
  * Bresler & Karzand (2021) Bresler, G. and Karzand, M.  Regret bounds and regimes of optimality for user-user and item-item collaborative filtering.  _IEEE Transactions on Information Theory_ , 67(6):4197–4222, 2021.
  * Brunskill & Li (2013) Brunskill, E. and Li, L.  Sample complexity of multi-task reinforcement learning.  _arXiv preprint arXiv:1309.6821_ , 2013.
  * Cammardella et al. (2020) Cammardella, N., Bušić, A., and Meyn, S.  Kullback-leibler-quadratic optimal control.  _arXiv preprint arXiv:2004.01798_ , 2020.
  * Candès & Tao (2010) Candès, E. J. and Tao, T.  The power of convex relaxation: Near-optimal matrix completion.  _IEEE Transactions on Information Theory_ , 56(5):2053–2080, 2010.
  * Chen et al. (2019) Chen, Y., Fan, J., Ma, C., and Yan, Y.  Inference and uncertainty quantification for noisy matrix completion.  _Proceedings of the National Academy of Sciences_ , 116(46):22931–22937, 2019.
  * Chen et al. (2020) Chen, Y., Chi, Y., Fan, J., Ma, C., and Yan, Y.  Noisy matrix completion: Understanding statistical guarantees for convex relaxation via nonconvex optimization.  _SIAM journal on optimization_ , 30(4):3098–3121, 2020.
  * Chi et al. (2019) Chi, Y., Lu, Y. M., and Chen, Y.  Nonconvex optimization meets low-rank matrix factorization: An overview.  _IEEE Transactions on Signal Processing_ , 67(20):5239–5269, 2019.
  * Dann & Brunskill (2015) Dann, C. and Brunskill, E.  Sample complexity of episodic fixed-horizon reinforcement learning.  _Advances in Neural Information Processing Systems_ , 28, 2015.
  * D’Eramo et al. (2020) D’Eramo, C., Tateo, D., Bonarini, A., Restelli, M., Peters, J., et al.  Sharing knowledge in multi-task deep reinforcement learning.  In _8th International Conference on Learning Representations_ , pp. 1–11. OpenReview. net, 2020.
  * Dubhashi & Ranjan (1996) Dubhashi, D. P. and Ranjan, D.  Balls and bins: A study in negative dependence.  _BRICS Report Series_ , 3(25), 1996.
  * Fazel (2002) Fazel, M.  _Matrix rank minimization with applications_.  PhD thesis, PhD thesis, Stanford University, 2002.
  * Gheshlaghi Azar et al. (2013) Gheshlaghi Azar, M., Munos, R., and Kappen, H. J.  Minimax pac bounds on the sample complexity of reinforcement learning with a generative model.  _Machine learning_ , 91(3):325–349, 2013.
  * Gleich & Lim (2011) Gleich, D. F. and Lim, L.-h.  Rank aggregation via nuclear norm minimization.  In _Proceedings of the 17th ACM SIGKDD international conference on Knowledge discovery and data mining_ , pp. 60–68, 2011.
  * Hessel et al. (2019) Hessel, M., Soyer, H., Espeholt, L., Czarnecki, W., Schmitt, S., and van Hasselt, H.  Multi-task deep reinforcement learning with popart.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pp. 3796–3803, 2019.
  * Hsieh et al. (2012) Hsieh, C.-J., Chiang, K.-Y., and Dhillon, I. S.  Low rank modeling of signed networks.  In _Proceedings of the 18th ACM SIGKDD international conference on Knowledge discovery and data mining_ , pp. 507–515, 2012.
  * Hu et al. (2021) Hu, J., Chen, X., Jin, C., Li, L., and Wang, L.  Near-optimal representation learning for linear bandits and linear rl.  In _International Conference on Machine Learning_ , pp. 4349–4358. PMLR, 2021.
  * Huleihel et al. (2021) Huleihel, W., Pal, S., and Shayevitz, O.  Learning user preferences in non-stationary environments.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 1432–1440. PMLR, 2021.
  * Jain & Pal (2022) Jain, P. and Pal, S.  Online low rank matrix completion.  _arXiv preprint arXiv:2209.03997_ , 2022.
  * Jain et al. (2013) Jain, P., Netrapalli, P., and Sanghavi, S.  Low-rank matrix completion using alternating minimization.  In _Proceedings of the forty-fifth annual ACM symposium on Theory of computing_ , pp. 665–674, 2013.
  * Jain et al. (2017) Jain, P., Kar, P., et al.  Non-convex optimization for machine learning.  _Foundations and Trends® in Machine Learning_ , 10(3-4):142–363, 2017.
  * Jin et al. (2020a) Jin, C., Krishnamurthy, A., Simchowitz, M., and Yu, T.  Reward-free exploration for reinforcement learning.  In _International Conference on Machine Learning_ , pp. 4870–4879. PMLR, 2020a.
  * Jin et al. (2020b) Jin, C., Yang, Z., Wang, Z., and Jordan, M. I.  Provably efficient reinforcement learning with linear function approximation.  In _Conference on Learning Theory_ , pp. 2137–2143. PMLR, 2020b.
  * Lazaric (2012) Lazaric, A.  Transfer in reinforcement learning: a framework and a survey.  In _Reinforcement Learning_ , pp. 143–173. Springer, 2012.
  * Lei & Li (2019) Lei, Y. and Li, W.  When collaborative filtering meets reinforcement learning.  _arXiv preprint arXiv:1902.00715_ , 2019.
  * Mate et al. (2022) Mate, A., Madaan, L., Taneja, A., Madhiwalla, N., Verma, S., Singh, G., Hegde, A., Varakantham, P., and Tambe, M.  Field study in deploying restless multi-armed bandits: Assisting non-profits in improving maternal and child health.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 36, pp. 12017–12025, 2022.
  * Mnih et al. (2015) Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al.  Human-level control through deep reinforcement learning.  _nature_ , 518(7540):529–533, 2015.
  * Modi et al. (2017) Modi, A., Jiang, N., Singh, S., and Tewari, A.  Markov decision processes with continuous side information.  _arXiv preprint arXiv:1711.05726_ , 2017.
  * Negahban & Wainwright (2011) Negahban, S. and Wainwright, M. J.  Estimation of (near) low-rank matrices with noise and high-dimensional scaling.  _The Annals of Statistics_ , 39(2):1069–1097, 2011.
  * Negahban et al. (2009) Negahban, S., Yu, B., Wainwright, M. J., and Ravikumar, P.  A unified framework for high-dimensional analysis of $m$-estimators with decomposable regularizers.  _Advances in neural information processing systems_ , 22, 2009.
  * Nguyen-Thanh et al. (2019) Nguyen-Thanh, N., Marinca, D., Khawam, K., Rohde, D., Vasile, F., Lohan, E. S., Martin, S., and Quadri, D.  Recommendation system-based upper confidence bound for online advertising.  _arXiv preprint arXiv:1909.04190_ , 2019.
  * Papadimitriou & Tsitsiklis (1999) Papadimitriou, C. H. and Tsitsiklis, J. N.  The complexity of optimal queuing network control.  _Mathematics of Operations Research_ , 24(2), 1999.
  * Recht (2011) Recht, B.  A simpler approach to matrix completion.  _Journal of Machine Learning Research_ , 12(12), 2011.
  * Sam et al. (2022) Sam, T., Chen, Y., and Yu, C. L.  Overcoming the long horizon barrier for sample-efficient reinforcement learning with latent low-rank structure.  _arXiv preprint arXiv:2206.03569_ , 2022.
  * Shah et al. (2020) Shah, D., Song, D., Xu, Z., and Yang, Y.  Sample efficient reinforcement learning via low-rank matrix estimation.  _Advances in Neural Information Processing Systems_ , 33:12092–12103, 2020.
  * Sodhani et al. (2021) Sodhani, S., Zhang, A., and Pineau, J.  Multi-task reinforcement learning with context-based representations.  In _International Conference on Machine Learning_ , pp. 9767–9779. PMLR, 2021.
  * Sutton et al. (1992) Sutton, R. S., Barto, A. G., and Williams, R. J.  Reinforcement learning is direct adaptive optimal control.  _IEEE control systems magazine_ , 12(2):19–22, 1992.
  * Taylor & Stone (2009) Taylor, M. E. and Stone, P.  Transfer learning for reinforcement learning domains: A survey.  _Journal of Machine Learning Research_ , 10(7), 2009.
  * Teh et al. (2017) Teh, Y., Bapst, V., Czarnecki, W. M., Quan, J., Kirkpatrick, J., Hadsell, R., Heess, N., and Pascanu, R.  Distral: Robust multitask reinforcement learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Vershynin (2018) Vershynin, R.  _High-dimensional probability: An introduction with applications in data science_ , volume 47.  Cambridge university press, 2018.
  * Wagenmaker et al. (2022) Wagenmaker, A. J., Chen, Y., Simchowitz, M., Du, S., and Jamieson, K.  Reward-free rl is no harder than reward-aware rl in linear markov decision processes.  In _International Conference on Machine Learning_ , pp. 22430–22456. PMLR, 2022.
  * Zhang et al. (2020) Zhang, Z., Du, S. S., and Ji, X.  Nearly minimax optimal reward-free reinforcement learning.  _arXiv preprint arXiv:2010.05901_ , 2020.

##  Appendix A More Discussion Regarding Policy Space

###  A.1 Necessity of Randomized Policies

We will first show that randomized policies might be necessary in such
contexts with a simple example and show that obtaining states which satisfy
conditions like (4) goes beyond simple reward maximization. Suppose $H=1$,
$\mathcal{S}=\\{1\\}$ and $\mathcal{A}=\\{1,\dots,d\\}$. We consider the
embedding $\psi(s,a)=e_{a}$. Suppose we want to obtain a policy $\pi$ such
that whenever $S_{1},A_{1}\sim\pi$,
$\lambda_{\min}(\mathbb{E}\psi(S_{1},A_{1})\psi(S_{1},A_{1})^{\intercal})$ is
maximized (where $\lambda_{\min}$ denotes the minimum eigenvalue). This is
maximized when $\pi(da|s)$ is chosen to be the uniform distribution over
$\mathcal{A}$ and the corresponding value is $1/d$. Note that whenever $\pi$
is a deterministic policy we will have $\lambda_{\min}=0$ whenever $d>1$. This
is in contrast to reward maximization problems where, under general
conditions, a deterministic optimal policy exists (See Theorem 1.7 in (Agarwal
et al., 2019)).

If fact, we can also show that the policy which minimizes
$\|\mathbb{E}\psi(S_{1},A_{1})-\frac{1}{d}\sum_{a=1}^{d}e_{a}\|$ must also
necessarily be random.

In the case of linear MDPs, we can find such a deterministic optimal policy
$\Pi=(\pi_{1},\dots,\pi_{H})$ as
$\pi_{h}(s)=\arg\sup_{a}\langle\psi(s,a),u^{*}_{h}\rangle+\langle\phi(s,a),v^{*}_{h}\rangle$
(Jin et al., 2020b). This reduces the problem to estimating the parameters
$u_{h}^{*},v_{h}^{*}$ even when the state-action space is an infinite set.
However, when such policies are not guaranteed to exist, as in case of
functional maximization required in Section 7, the set of all policies can be
intractably large. This is the justification for picking a nice enough policy
space denoted by $\mathbb{Q}$.

###  A.2 Constructing Policy Spaces

We consider any linear MDP satisfying the definition given in Section 2 and
suppose $\mathcal{A}$ is finite. We consider the set of all probability
distributions
$\pi_{h}(a|s;u,v)\propto\exp(\langle\phi(s,a),u\rangle+\langle\psi(s,a),v\rangle)$.
We consider $\mathbb{Q}_{h}=\\{\pi_{h}(a|s,u,v):u,v\in\mathcal{B}_{d}(R)\\}$,
We let our policy space be
$\mathbb{Q}=\\{\Pi=(\pi_{1},\dots,\pi_{H}):\pi_{h}\in\mathbb{Q}_{h}\\}$.

######  Lemma A.1.

Consider the probability distribution over a finite set $[|\mathcal{A}|]$ give
by $p_{\beta}(a)\propto\exp(\beta x_{a})$ for every $a\in[|\mathcal{A}|]$ some
$x_{a}\in\mathbb{R}^{+}$ and $\beta\in\mathbb{R}^{+}$. For any $\epsilon>0$
and random variable $A\sim p_{\beta}$, we must have:

$\mathbb{P}(x_{A}<\sup_{a}x_{a}-\epsilon)\leq|\mathcal{A}|\exp(-\beta\epsilon)$

And

$\mathbb{E}x_{A}\geq(\sup_{a}x_{a}-\epsilon)(1-|\mathcal{A}|\exp(-\beta\epsilon))$

######  Lemma A.2.

Let $Q^{*}_{h}(s,a)$ be the optimal action-value function for the MDP. Then
the policy $\Pi=(\pi_{1},\dots,\pi_{h})$ given by
$\pi_{h}(a|s)\propto\exp(\beta Q^{*}_{h}(s,a))$ is $\epsilon
H+H^{2}|\mathcal{A}|\exp(-\beta\epsilon)$ sub-optimal for any $\epsilon>0$

###### Proof.

Consider the optimal value function defined by
$V^{*}_{h}(s)=\sup_{a}Q^{*}_{h}(s,a)$. Let $\bar{Q}_{h}(s,a)$ denote the
optimal action value function under the policy $\Pi$ and let
$\bar{V}(s)=\int\bar{Q}_{h}(s,a)\pi_{h}(da|s)$ denote the value at state $s$
with the policy $\Pi$. Clearly, we have:
$\bar{Q}_{H}(s,a)=Q^{*}_{H}(s,a)=R(s,a)$. $\bar{Q}_{h}(s,a)\geq
Q^{*}_{h}(s,a)-\eta$ uniformly. Then we have

$\displaystyle\bar{V}_{h}(s)$ $\displaystyle\geq\int
Q^{*}_{h}(s,a)\pi_{h}(da|s)-\eta$
$\displaystyle\geq(\sup_{a}Q_{h}^{*}(s,a)-\epsilon)(1-|\mathcal{A}|\exp(-\beta\epsilon))-\eta$
$\displaystyle\geq\sup_{a}Q_{h}^{*}(s,a)-\epsilon-H|\mathcal{A}|\exp(-\beta\epsilon)-\eta=V_{h}^{*}(s)-\epsilon-H|\mathcal{A}|\exp(-\beta\epsilon)-\eta$
(5)

In the second step, we have invoked Lemma A.1. In the last step, we have used
the fact that $Q^{*}_{h}\in[0,H]$ uniformly. Now, by the Bellman iteration, we
have:

$\displaystyle Q^{*}_{h-1}(s,a)$
$\displaystyle=R_{h-1}(s,a)+\mathbb{E}_{s^{\prime}\sim
P_{h-1}(|s,a)}V^{*}_{h}(s^{\prime})$ $\displaystyle\geq
R_{h-1}(s,a)+\mathbb{E}_{s^{\prime}\sim
P_{h-1}(|s,a)}\bar{V}_{h}(s^{\prime})-\epsilon-H|\mathcal{A}|\exp(-\beta\epsilon)-\eta$
$\displaystyle=\bar{Q}_{h-1}(s,a)-\epsilon-H|\mathcal{A}|\exp(-\beta\epsilon)-\eta$
(6)

Therefore, by induction, we conclude that $\bar{V}_{1}(s)\geq
V^{*}_{1}(s)-\epsilon H-H^{2}|\mathcal{A}|\exp(-\beta\epsilon)$

Therefore, by the definition of the value function, we conclude the claim. ∎

Now, by a simple extension of Proposition 2.3 in (Jin et al., 2020b), we
conclude that the optimal $Q_{h}^{*}$ function for any linear MDP can be
written as:

$Q_{h}^{*}(s,a)=\langle\psi(s,a),u_{h}^{*}\rangle+\langle\phi(s,a),v_{h}^{*}\rangle\,.$

Where $\|u_{h}^{*}\|_{2}\leq\sqrt{d}$ and $\|v_{h}^{*}\|_{\infty}\leq
HC_{\mu}$. Observe that choosing $\epsilon=\frac{\eta}{2H}$ and
$\beta=2\frac{\log(2H|\mathcal{A}|/\eta)}{\eta}$ will ensure that the
randomized policy $\Pi$ in the statement of Lemma A.2 is $\eta$ optimal.
Therefore, we can take $R=2dHC_{\mu}\frac{\log(2H|\mathcal{A}|/\eta)}{\eta}$
in the definition of $\mathbb{Q}_{h}$ above and conclude that this includes
every $\eta$ optimal policy for every MDP with embedding functions
$\phi,\psi$. We will now bound the covering number. Recall the definition of
the distance
$D_{\mathbb{Q}}(\Pi_{1},\Pi_{2})=\sup_{h\in[H]}\mathsf{TV}(\pi_{h}^{(1)},\pi_{h}^{(2)})$.
Therefore it is sufficient to obtain an $\eta$ cover for $\mathbb{Q}_{h}$
(denoted by $\hat{\mathbb{Q}}_{h,\eta}$) and then construct
$\hat{\mathbb{Q}}_{\eta}=\\{\Pi=(\pi_{1},\dots,\pi_{H}):\pi_{h}\in\hat{\mathbb{Q}}_{h,\eta}\forall
h\in[H]\\}=\prod_{h=1}^{H}\hat{\mathbb{Q}}_{h,\eta}$.

######  Lemma A.3.

$\pi(|s;u,v)$ be as defined in the beginning of this Subsection.

$\mathsf{TV}(\pi(\cdot|s;u,v),\pi(\cdot|s;u^{\prime},v^{\prime}))\leq\frac{1}{2}\left(\exp(2\|u-u^{\prime}\|_{2}+2\|v-v^{\prime}\|_{\infty})-1\right)$

###### Proof.

Denote $\pi(a|s;u,v)$ by $\pi(a)$ and $\pi(a|s;u^{\prime},v^{\prime})$ by
$\pi^{\prime}(a)$. Consider the corresponding partition functions denoted by
$Z:=\sum_{a\in\mathcal{A}}\exp(\langle\psi(s,a),u\rangle+\langle\phi(s,a),v\rangle)$
and
$Z^{\prime}:=\sum_{a\in\mathcal{A}}\exp(\langle\psi(s,a),u^{\prime}\rangle+\langle\phi(s,a),v^{\prime}\rangle)$.
We conclude that using Hölder’s inequality for $\langle
u-u^{\prime},\psi\rangle$ and $\langle v-v^{\prime},\phi\rangle$ that:

$\displaystyle\exp(-\|u-u^{\prime}\|_{2}-\|v-v^{\prime}\|_{\infty})\leq\frac{Z^{\prime}}{Z}\leq\exp(\|u-u^{\prime}\|_{2}+\|v-v^{\prime}\|_{\infty})$
$\displaystyle\exp(-2\|u-u^{\prime}\|_{2}-2\|v-v^{\prime}\|_{\infty})\leq\frac{\pi^{\prime}(a)}{\pi(a)}\leq\exp(2\|u-u^{\prime}\|_{2}+2\|v-v^{\prime}\|_{\infty})$
(7) $\displaystyle\mathsf{TV}(\pi,\pi^{\prime})$
$\displaystyle=\tfrac{1}{2}\sum_{a\in\mathcal{A}}|\pi(a)-\pi^{\prime}(a)|$
$\displaystyle=\tfrac{1}{2}\sum_{a\in\mathcal{A}}\pi(a)\bigr{|}1-\tfrac{\pi^{\prime}(a)}{\pi(a)}\bigr{|}$
$\displaystyle\leq\frac{1}{2}\left(\exp(2\|u-u^{\prime}\|_{2}+2\|v-v^{\prime}\|_{\infty})-1\right)$
(8)

∎

Using the lemma above, we conclude that
$\hat{\mathbb{Q}}_{h,\eta}=\\{\pi_{h}(|s;u,v):u,v\in\hat{\mathcal{B}}_{d,\eta/4}(R)\\}$
whenever $\eta\leq 1$. Here $\hat{\mathcal{B}}_{d,\eta/4}(R)$ an $\eta/4$ net
over $\mathcal{B}_{d}(R)$ with respect to the norm $\|\cdot\|_{2}$. From the
results in (Vershynin, 2018), we can therefore take:

$|\hat{\mathbb{Q}}_{h,\eta}|\leq|\hat{\mathcal{B}}_{d,\eta/4}(R)|^{2}\leq\exp(Cd\log(\tfrac{C\eta}{R}))$
(9)

Since we had $\hat{\mathbb{Q}}_{\eta}=\prod_{h=1}^{H}$, we conclude that:

$\log(|\hat{\mathbb{Q}}_{\eta}|)\leq
cdH\left(\log\tfrac{dH}{\eta}+\log\log(2H|\mathcal{A}|/\eta)\right)$

##  Appendix B Analysis - Tabular MDPs

We will call the reward free RL procedure in Phase 1 to be successful if it
outputs the $\epsilon$ optimal policy. This has probability atleast
$1-\frac{\delta}{2}$.

###  B.1 Analysis of Algorithm 1

######  Lemma B.1.

Suppose $p\leq\frac{1}{2}$, conditioned on the success of Phase 1, with
probability at-least $1-\exp(-cNp|\mathcal{S}||\mathcal{A}|H)$, Algorithm 1
terminates after querying $\frac{C|\mathcal{S}||\mathcal{A}|NHp}{\epsilon}$
trajectories. $(\mathcal{G}_{h})_{h\in[H]}$, the active sets at the
termination of the algorithm. They satisfy:

$\sup_{\pi}\sum_{h=1}^{H}P^{\pi}_{h}(\mathcal{G}_{h})\leq\frac{5\epsilon}{8}$
(10)

For any $a\times b$ matrix $R$, let $\hat{R}$ be its partially observed
version (that is, there exists a set of indices $I\subseteq[a]\times[b]$ such
that $\hat{R}_{ij}=R_{ij}$ if $(i,j)\in I$ and $\hat{R}_{ij}=*$ otherwise). We
call a random set of indices $J$ to have the distribution
$\mathsf{Unif}(m,[a],[b])$ if $J$ is drawn uniformly at random such that
$|J|=m$.

######  Lemma B.2 (Modification: Mod1).

Suppose we run, independently, a modification of algorithm 1 where on the
“Query trajectory” step the trajectories are sampled from a fixed MDP
$\mathcal{M}_{1}$ (but rewards are from the reward function corresponding to
$U_{t}$). Consider all the random variables that determine the trajectory of
this algorithm:
$\left(\hat{V},\hat{\Pi},(S^{(t)}_{1:H},A^{(t)}_{1:H},R^{(t)}_{1:H})_{t},(U_{t})_{t}\right)$.
Then the joint distribution of this collection of random variables is
unchanged under the modification.

###### Proof.

The proof follows from an induction argument on the time index $t$. We
describe the key steps here. For the ease of notation, let
$\mathcal{T}_{t}=\left((S^{(t)}_{1:H},A^{(t)}_{1:H},R^{(t)}_{1:H}),U_{t}\right)$
Let $X_{T}=\left(\hat{V},\hat{\Pi},(\mathcal{T}_{t})_{t\leq T}\right)$. Let
$\tilde{X}_{T}$ and $\tilde{\mathcal{T}}_{t}$ denote the corresponding
quantities under the modification. It is enough to show that finite
dimensional marginals have the same joint distribution under the modification.
In particular, we will show:

  1. 1.

$\mathcal{T}_{0}\overset{d}{=}\tilde{\mathcal{T}}_{0}$

  2. 2.

Suppose $X_{T}\overset{d}{=}\tilde{X}_{T}$. Then the Markov kernel
$k_{\mathcal{T}_{T+1}|X_{T}}$ is almost surely (under the common distribution
of $X_{T}$, $\tilde{X}_{T}$) equal to
$k_{\tilde{\mathcal{T}}_{T+1}|\tilde{X}_{T}}$. Thus
$X_{T+1}\overset{d}{=}\tilde{X}_{T+1}$

The first statement is straightforward since, in the zeroth step, the
distribution of $\hat{\Pi}^{\mathcal{G}}$ not affected by the modification,
and thus due to identical MDP transitions across users, the distribution of
$\mathcal{T}_{0}$ is preserved under modification. A similar argument proves
the second statement. Roughly, given a realization of $X_{T}$ the distribution
of $\mathcal{T}_{T+1}$ is same as the distribution of
$\tilde{\mathcal{T}}_{T+1}$ given the same realization of $\tilde{X}_{T}$, due
to the exact same reason presented for the first statement. A fully formal
proof requires setting up appropriate proability spaces, so we omit it here.
Furthermore, since the random variables considered are all discrete, one can
argue via PMFs as well. ∎

######  Lemma B.3.

Suppose $p\leq\frac{1}{2}$. conditioned on the success of Phase 1 and
termination of Algorithm 1, for every $h\in[H]$, the Algorithm 1 returns
partially filled reward matrices $\hat{R}_{h}$. Consider the sub-matrix
$\hat{R}_{h}^{\mathcal{G}^{\complement}_{h}}$. Let
$I_{h}\subseteq[N]\times\mathcal{G}^{\complement}_{h}$ be the sub-set of
observed indices for $\hat{R}_{h}$. Let
$J_{h}|\mathcal{G}_{h}\sim\mathsf{Unif}(\frac{Np|\mathcal{G}^{\complement}_{h}|}{2},[N],\mathcal{G}^{\complement}_{h})$.
There exists a coupling between $J_{h}$ and $I_{h}$ such that

$\mathbb{P}\left(J_{h}\subseteq I_{h}\bigr{|}\mathcal{G}_{h}\right)\geq
1-|\mathcal{S}||\mathcal{A}|\exp(-cNp)$

###### Proof.

Let us fix $\mathcal{G}_{h}$ and construct a coupling between $I_{h}$ and
$J_{h}$.

Consider any fixed, arbitrary permutations $\sigma_{g}$ over $[N]$, for
$g\in\mathcal{G}_{h}^{\complement}$. By $\sigma(I_{h})$, we denote
$\\{(\sigma_{g}(i),g):(i,g)\in I_{h}\\}$.

######  Claim B.4.

Conditioned on $\mathcal{G}_{h}$, $\sigma(I_{h})$ has the same distribution as
$I_{h}$.

###### Proof.

Let $\\{\sigma_{(s,a):[N]\to[N]}|(s,a)\in\mathcal{S}\times\mathcal{A}\\}$ be a
set of arbitrary permutations on $[N]$. From lemma B.2 it is enough to prove
the statement for the random variables under the modification described in
that lemma (call this Mod1). Now consider a further modification (call it
Mod2) where in every iteration $t$, we sample $U_{t}\sim\mathsf{Unif}([N])$,
for each horizon $h$ we set
$\tilde{U}^{(t)}_{h}=\sigma_{(S_{h}^{(t)},A_{h}^{(t)})}(U_{t})$, and then
update the entries of $R_{h}(\tilde{U}^{(t)}_{h},(S_{h}^{(t)},A_{h}^{(t)}))$
(instead of $R_{h}(U_{t},(S_{h}^{(t)},A_{h}^{(t)}))$). Next, we couple these
two modifications by using same $(\hat{V},\hat{\Pi})$ and the same set of
$U_{t}$’s for both the modifications. Further, we couple the MDP used in these
modifications to be the same, single MDP.

Now an induction argument shows that the sequence of active sets $\mathcal{G}$
obtained in these modifications are also identical for every time $t$; only
the rows of $R_{h}$ where entries are filled change according to the set of
permutations chosen.Thus, under the described coupling, Mod1 and Mod2 produce
identical trajectories (i.e., $(S_{1:H}^{(t)},A_{1:H}^{(t)})$), the columns of
reward matrices are just permutations of each other described by the chosen
set of permutations, and algorithm 1 terminate at the same time in both these
cases. However, the same induction argument also shows that for each $t$ and
$h$, conditioned on $\mathcal{G}_{h}$, trajectories (which is same in Mod1 and
Mod2) until at beginning of iteration $t$, we have
$(U_{t},(S^{(t)}_{h},A^{(t)}_{h}))\overset{d}{=}(\tilde{U}^{(t)}_{h},(S^{(t)}_{h},A^{(t)}_{h}))$.

Therefore if
$I_{h},\tilde{I}_{h}\subset[N]\times\mathcal{G}_{h}^{\complement}$ denotes the
subset of observed indices at termination (outside the active set), then
$\tilde{I}_{h}=\sigma(I_{h})\equiv\\{(\sigma_{(s,a)}(i),(s,a)):(i,(s,a)\in
I_{h})\\}$ and, conditioned on $\mathcal{G}_{h}$,
$\tilde{I}_{h}\overset{d}{=}I_{h}$

∎

######  Claim B.5.

At termination, conditioned on $\mathcal{G}_{h}$, random sets
$I_{h}^{g}=\\{(i,g):(i,g)\in I_{h}\\}$ are jointly independent.

###### Proof.

Again we work with the modification Mod1 described in lemma B.2. For each
$(s,a)$ consider the collection of $U_{t}$’s that are used populate the column
$(s,a)$ of matrix $R_{h}$ in algorithm 1. Call this collection
$\mathcal{U}_{(s,a)}$. ∎

######  Remark B.6.

Since the columns of $I_{h}$ have exactly $Np$ entries, permutation invariance
proved in the above claim implies that

For any set $\bar{J}\subseteq[N]\times\mathcal{G}_{h}^{\complement}$, define
the count function $(N^{\bar{J}}_{g})_{g\in\mathcal{G}_{h}^{\complement}}$
such that $N^{\bar{J}}_{g}=|\\{i\in[N]:(i,g)\in\bar{J}\\}|$.

We are now ready to give the coupling: given $\mathcal{G}_{h}$, draw uniformly
random, independent permutations $\sigma_{g}$ for
$g\in\mathcal{G}_{h}^{\complement}$. Draw
$(N_{g})_{g\in\mathcal{G}^{\complement}_{h}}$ independent of $\sigma_{g}$ and
to have the joint law of
$(N^{J_{h}}_{g})_{g\in\mathcal{G}_{h}^{\complement}}$. Define:

$\tilde{J}_{h}=\\{(\sigma_{g}(i),g):i\leq
N_{g},g\in\mathcal{G}_{h}^{\complement}\\}$
$\tilde{I}_{h}=\\{(\sigma_{g}(i),g):i\leq
Np,g\in\mathcal{G}_{h}^{\complement}\\}$

######  Claim B.7.

The marginal distributions of $\tilde{J}_{h}$ and $\tilde{I}_{h}$ are
respectively the distributions of $J_{h}$ and $I_{h}$.

###### Proof.

First we will prove a general statement about $J\sim\mathsf{Unif}(r,[N],[M])$.
Let $X\in\\{0,1\\}^{N\times M}$ with $X_{i,m}=1$ iff $(i,m)\in J$. Let
$(N_{m})_{m\in[M]}$ be the count functions corresponding to $J$ i.e.,
$N_{m}=\sum_{i}X_{i,m}$. Let $Y_{m}=(X_{1,m},\cdots,X_{N,m})$. We will argue
that conditional on $\\{N_{m}:m\in[M]\\}$, the random vectors $Y_{m}$ are
jointly independent. Indeed, pick any $x\in\\{0,1\\}^{N\times M}$ and
$(n_{1},\cdots,n_{m})$. Let $y_{m}$ be the $m$’th column of $x$. Then

$\mathbb{P}\left[X=x,\cap_{m}\\{N_{m}=n_{m}\\}\right]=\left(\prod_{m}1\left[\sum_{i}x_{i,m}=n_{m}\right]\right)1\left[\sum_{i,m}x_{i,m}=r\right]\frac{1}{\binom{MN}{r}}$

The above can also be written as

$\mathbb{P}\left[\cap_{m}\\{Y_{m}=y_{m}\\},\cap_{m}\\{N_{m}=n_{m}\\}\right]=\left(\prod_{m}1\left[\sum_{i}x_{i,m}=n_{m}\right]\right)1\left[\sum_{m}n_{m}=r\right]\frac{1}{\binom{MN}{r}}$

Let $\mathbf{1}$ denote the all $1$ vector in $\mathbb{R}^{N}$. Note that
$y_{m}=(x_{1,m},\cdots,x_{N,m})^{\top}$. Marginalizing the above, we see that

$\mathbb{P}\left[\cap_{m}\\{N_{m}=n_{m}\\}\right]=\left(\prod_{m}\frac{1}{\binom{N}{n_{m}}}\right)1\left[\sum_{m}n_{m}=r\right]\frac{1}{\binom{MN}{r}}$

Thus the conditional distribution can be expressed as

$\mathbb{P}\left[\cap_{m}\\{Y_{m}=y_{m}\\}\biggr{|}\cap_{m}\\{N_{m}=n_{m}\\}\right]=\left(\prod_{m}\frac{1\left[\mathbf{1}^{\top}y_{m}=n_{m}\right]}{\binom{N}{n_{m}}}\right)\frac{1\left[\sum_{m}n_{m}=r\right]}{\binom{MN}{r}}$

Since the (conditional) joint PMF factors, it is an easy calculation to show
the conditional independence i.e.,

$\mathbb{P}\left[\cap_{m}\\{Y_{m}=y_{m}\\}\biggr{|}\cap_{m}\\{N_{m}=n_{m}\\}\right]=\prod_{m}\mathbb{P}\left[Y_{m}=y_{m}\biggr{|}\cap_{m}\\{N_{m}=n_{m}\\}\right]$

Furthermore, for any $n_{1},\cdots n_{m}$ such that $\sum_{m}n_{m}=r$,
marginalization shows

$\mathbb{P}\left[Y_{m}=y_{m}|\cap_{m}\\{N_{m}=n_{m}\\}\right]=\frac{1\left[\mathbf{1}^{\top}y_{m}=n_{m}\right]}{\binom{N}{n_{m}}}$

Let $N_{-m}=(N_{1},\cdots,N_{m-1},N_{m+1},\cdots,N_{M})$, and similarly for
$n_{-m}$. Then

$\displaystyle\mathbb{P}\left[Y_{m}=y_{m},N_{-m}=n_{-m}|N_{m}=n_{m}\right]$
$\displaystyle=\mathbb{P}\left[Y_{m}=y_{m}|\cap_{m}\\{N_{m}=n_{m}\\}\right]\mathbb{P}\left[N_{-m}=n_{-m}|N_{m}=n_{m}\right]$
$\displaystyle=\begin{cases}0,&\sum_{m}n_{m}\neq r\\\
\frac{1\left[\mathbf{1}^{\top}y_{m}=n_{m}\right]}{\binom{N}{n_{m}}}\mathbb{P}\left[N_{-m}=n_{-m}|N_{m}=n_{m}\right],&\mathrm{otherwise}\end{cases}$

The above factorization directly implies that $Y_{m}$, conditioned on $N_{m}$
is uniformly distributed on its support $\\{y:\mathbf{1}^{\top}y=N_{m}\\}$ and
is independent of $N_{-m}$. Thus

$\displaystyle\mathbb{P}\left[\cap_{m}\\{Y_{m}=y_{m}\\}\biggr{|}\cap_{m}\\{N_{m}=n_{m}\\}\right]$
$\displaystyle=\prod_{m}\mathbb{P}\left[Y_{m}=y_{m}\biggr{|}N_{m}=n_{m}\right]$
$\displaystyle=\prod_{m}\frac{1\left[\mathbf{1}^{\top}y_{m}=n_{m}\right]}{\binom{N}{n_{m}}}$

Observation: The above calculations give another way to generate $Y$: first
generate $N_{1},\cdots,N_{M}$ from the right distribution, and then
conditioned on $N_{m}$ generate each $Y_{m}$ uniformly such that
$\mathbf{1}^{\top}Y_{m}=N_{m}$.

Next we apply the above calculations and observation to
$J=J_{h}|\mathcal{G}_{h}\sim\mathsf{Unif}(\frac{Np|\mathcal{G}^{\complement}_{h}|}{2},[N],\mathcal{G}^{\complement}_{h})$.
For a uniformly random permutation $\sigma$ on $[N]$, the set
$\\{\sigma(i):1\leq i\leq k\\}$ is uniformly distributed on all $k$-sized
subsets of $[N]$. In the statement of the claim the permutations are chosen
independently for each $g\in\mathcal{G}_{c}^{\complement}$. Thus from the
above observation, we have $J_{h}\overset{d}{=}\tilde{J}_{h}$ conditioned on
$\mathcal{G}_{h}$.

The claim about $\tilde{I}_{h}$ follows directly from permutation invariance
proved by claim B.4.

∎

######  Claim B.8.

$\mathbb{P}(N_{g}>Np|\mathcal{G}_{h})\leq\exp(-c_{0}Np)$ for every
$g\in\mathcal{G}_{h}^{\complement}$

###### Proof.

Throughout this proof, we will condition on the terminal active set
$\mathcal{G}_{h}$. We will show this using the results on concentration with
negative regression property as established in Proposition 29 in (Dubhashi &
Ranjan, 1996). $N_{g}=\sum_{i=1}^{N}\mathbbm{1}((i,g)\in\tilde{J}_{h})$. Now
we will show that the collection $X_{ig}:=\mathbbm{1}((i,g)\in\tilde{J}_{h})$
for $i\in[N],g\in\mathcal{G}_{h}^{\complement}$ satisfy the negative
regression property. By the definition of negative regression, we can conclude
that the sub-collection $(X_{ig})_{i\in[N]}$ also satisfies this property for
every $g\in\mathcal{G}_{h}^{\complement}$.

Consider the partial order over binary vectors $X\succeq Y$ iff $X_{l}\geq
Y_{l}$ for every $l$. The negative regression property is satisfied iff for
every $K_{1},K_{2}\subseteq[N]\times\mathcal{G}_{h}^{\complement}$ such that
$K_{1}\cap K_{2}=\emptyset$, and a real valued function $f(X_{m}:m\in K_{1})$
which is non-decreasing with respect to the partial order, we must have:

$g(t_{l}:l\in K_{2}):=\mathbb{E}\left[f(X_{m}:m\in
K_{1})\bigr{|}X_{l}=t_{l},\forall l\in K_{2}\right]$

be such that $g$ is a non-increasing function in $t_{l}$ with respect to the
partial order. Note that in the case of uniform distribution as in
$\tilde{J}_{h}$, the distribution $(X_{m})_{m\in K_{1}}$ is the uniform,
permutation invariant distribution with constant sum almost surely. The sum
being $\frac{Np|\mathcal{G}_{h}^{\complement}|}{2}-\sum_{l\in K_{2}}t_{l}$.
Therefore, whenever $t_{l}^{\prime}\geq t_{l}$ for every $l\in K_{2}$, we have
the following stochastic dominance:

$\left[(X_{m})_{m\in K_{1}}\biggr{|}X_{l}=t^{\prime}_{l}\forall l\in
K_{2}\right]\preceq\left[(X_{m})_{m\in K_{1}}\biggr{|}X_{l}=t_{l}\forall l\in
K_{2}\right]$

Therefore, this coupling leads us to conclude that:

$\displaystyle g(t_{l}:l\in K_{2})$
$\displaystyle:=\mathbb{E}\left[f(X_{m}:m\in K_{1})\bigr{|}X_{l}=t_{l}\forall
l\in K_{2}\right]$ $\displaystyle\geq\mathbb{E}\left[f(X_{m}:m\in
K_{1})\bigr{|}X_{l}=t_{l}^{\prime}\forall l\in K_{2}\right]$
$\displaystyle=g(t^{\prime}_{l}:l\in K_{2})$ (11)

The second step follows from stochastic dominance. This implies that the
function $g$ is non-increasing which establishes the negative regression
property. Now, we consult Proposition 29 in (Dubhashi & Ranjan, 1996) to show
that we can take Chernoff bounds on $N_{g}=\sum_{i\in[N]}X_{ig}$ as though
$X_{ig}$ were i.i.d $\mathsf{Ber}(p)$. Therefore, from an application of
Bernstein’s inequality (Boucheron et al., 2013), we conclude the statement of
the claim. ∎

Now, $J_{h}\subseteq I_{h}$ if and only if $N_{g}\leq Np$ for every
$g\in\mathcal{G}^{\complement}_{h}$. Therefore, from the claim above, we have
$\mathbb{P}(J_{h}\subseteq I_{h}|\mathcal{G}_{h}^{\complement})\geq
1-|\mathcal{S}||\mathcal{A}|\exp\left(-c_{0}Np\right)$. ∎

We are now ready to prove Theorem 5.3.

###### Proof of Theorem 5.3.

In order to establish the result, we need to show that with $p$ as set in the
statement, the algorithm returns $\epsilon$ optimal policies $\hat{\Pi}_{u}$
for every user $u\in[N]$ with probability at-least $1-\delta$.

The total sample complexity is the number of trajectories queried in Phase 1
plus the number of trajectories queried in Phase 2. Phase 1 queries
$K^{\mathsf{rf}}(\tfrac{\epsilon}{8},\tfrac{\delta}{2})$ trajectories, which
is
$C\frac{|\mathcal{S}||\mathcal{A}|H^{2}\left(|\mathcal{S}|+\log(\tfrac{1}{\delta})\right)}{\epsilon^{2}}\mathsf{polylog}(\tfrac{|\mathcal{S}||\mathcal{A}|H}{\epsilon})$
by the results of (Zhang et al., 2020). By Lemma B.1, we conclude that the
sample complexity of phase 2 is
$\frac{C|\mathcal{S}||\mathcal{A}|NHp}{\epsilon}$ and with the value of $p$
given in the statement of the theorem, this succeeds with probability at-least
$1-\frac{\delta}{4}$ when conditioned on the success of Phase 1.

We will show that conditioned on the success of Phase 2, with probability at-
least $1-\frac{\delta}{4}$, the nuclear norm minimization algorithm of (Recht,
2011) successfully obtains $R_{h}^{\mathcal{G}_{h}^{\complement}}$. Indeed by
Theorem 1 in (Recht, 2011), we see that whenever co-ordinates of
$R_{h}^{\mathcal{G}_{h}^{\complement}}$ corresponding to random indices drawn
from $\mathsf{Unif}(m,[N],\mathcal{G}_{h}^{\complement})$ are observed with
$m=C_{1}\max(\mu_{1}^{2},\mu_{0})r(N+|\mathcal{G}_{h}^{\complement}|)\log^{2}|\mathcal{G}_{h}^{\complement}|\log(\tfrac{H}{\delta})$,
the algorithm succeeds at recovering $R_{h}^{\mathcal{G}_{h}^{\complement}}$
with probability at-least $1-\frac{\delta}{8H}$. The number of co-ordinates we
observe is

$Np|\mathcal{G}_{h}^{\complement}|\geq\frac{Np|\mathcal{S}\mathcal{A}|}{2}\geq
2C_{1}\max(\mu_{1}^{2},\mu_{0})r(N+|\mathcal{G}_{h}^{\complement}|)\log^{2}|\mathcal{G}_{h}^{\complement}|\log(\tfrac{H}{\delta})$

In the last step, we have used Assumption 5.2 to conclude that
$|\mathcal{G}_{h}^{\complement}|\geq\frac{|\mathcal{S}||\mathcal{A}|}{2}$. For
the constant $C$ in the definition of $p$ large enough, we must have:

$m\leq\frac{Np|\mathcal{G}_{h}^{\complement}|}{2}$

Note that the results of (Recht, 2011) requires at-least $m$ observations to
be chosen uniformly at random co-ordinates, but we do not obtain observations
which are uniformly at uniformly random co-ordinates. Here, we will use the
results of Lemma B.3. Let $J_{h}$ be a fictitious subset of co-ordinates
distributed as $\mathsf{Unif}(m,[N],\mathcal{G}_{h}^{\complement})$ when
conditioned on $\mathcal{G}_{h}^{\complement}$. If the observed co-ordinates
are $J_{h}$, then we can successfully estimate the reward matrix $R_{h}$ with
proability at-least $1-\frac{\delta}{8H}$ in this case. Now, suppose that the
actually observed co-ordinates are $I_{h}$, which is a strict super-set of
$J_{h}$. Then we check that the matrix completion algorithm, which is based on
constrained nuclear-norm minimization, still succeeds with observed co-
ordinates corresponding to $I_{h}$ whenever it succeeds with the observed co-
ordinates correspond to $J_{h}$.

We now refer to the coupling in Lemma B.3, which shows that we can couple
$J_{h}$ to the real distribution $I_{h}$ such that $J_{h}\subseteq I_{h}$ with
probability at-least $1-\frac{\delta}{8H}$ When the constant $C_{1}$ in the
definition of $p$ is large enough, we conclude by invoking Lemma B.3 that:
$J_{h}\subseteq I_{h}$ with probability at-least $1-\frac{\delta}{8H}$.
Applying union bound for $h\in[H]$, we conclude that Phase 3 succeeds with
probability at-least $1-\frac{\delta}{4}$ when conditioned on the success of
Phases 1 and 2.

Therefore, from the arguments above, we conclude that Phases 1,2 and 3 succeed
with probability at-least $1-\delta$ and give us the reward matrices
$R_{h}^{\mathcal{G}_{h}^{\complement}}$ where the sets satisfy the following
equation from Lemma B.1.

$\sup_{\pi}\sum_{h=1}^{H}P^{\pi}_{h}(\mathcal{G}_{h})\leq\frac{5\epsilon}{8}$
(12)

It now remains to show that we obtain $\epsilon$ optimal policies for each
user after Phase 4. Note that whenever Phase 1 succeeds, we can compute
$\epsilon/4$ optimal policies for every possible reward function bounded in
$[0,1]$. Since we do not know the rewards over the set $\mathcal{G}_{h}$, we
set it to zero as described in the algorithm to obtain $\bar{R}_{h}$. It
remains to show that planning with $\bar{R}_{h}$ and using it with the reward
free RL algorithm gives us an $\epsilon$ optimal policy. Suppose $\Pi^{*}_{u}$
is the optimal policy for user $u$ and suppose $\bar{\Pi}_{u}$ be the optimal
policy for user $u$ under rewards $\bar{R}_{h}(u,(s,a))$. Note that combined
with the guarantees for the reward free RL, in order to complete the proof of
the theorem, it is sufficient to show that the policy $\bar{\Pi}_{u}$ is
$3\epsilon/4$ optimal with respect to the actual rewards $R_{h}(u,(s,a))$. Let
$S^{*}_{1:H},A^{*}_{1:H}\sim\mathcal{M}(\Pi_{u}^{*})$ and
$\bar{S}_{1:H},\bar{A}_{1:H}\sim\mathcal{M}(\bar{\Pi}_{u})$.

$\displaystyle\mathbb{E}\sum_{h=1}^{H}R_{h}(u,(S^{*}_{h},A^{*}_{h}))\leq\mathbb{E}\sum_{h=1}^{H}R_{h}(u,(S^{*}_{h},A^{*}_{h}))\mathbbm{1}((S^{*}_{h},A^{*}_{h})\in\mathcal{G}_{h}^{\complement})+\mathbbm{1}((S^{*}_{h},A^{*}_{h})\in\mathcal{G}_{h})$
$\displaystyle=\mathbb{E}\sum_{h=1}^{H}\bar{R}_{h}(u,(S^{*}_{h},A^{*}_{h}))+\mathbbm{1}((S^{*}_{h},A^{*}_{h})\in\mathcal{G}_{h})=\mathbb{E}\sum_{h=1}^{H}\bar{R}_{h}(u,(S^{*}_{h},A^{*}_{h}))+\sum_{h=1}^{H}P^{\Pi^{*}}_{h}(\mathcal{G}_{h})$
$\displaystyle\leq\mathbb{E}\sum_{h=1}^{H}\bar{R}_{h}(u,(S^{*}_{h},A^{*}_{h})+\frac{5\epsilon}{8}$
$\displaystyle\leq\mathbb{E}\sum_{h=1}^{H}\bar{R}_{h}(u,(\bar{S}_{h},\bar{A}_{h})+\frac{5\epsilon}{8}$
$\displaystyle\leq\mathbb{E}\sum_{h=1}^{H}R_{h}(u,(\bar{S}_{h},\bar{A}_{h})+\frac{5\epsilon}{8}$

In the first step we have used the fact that the rewards are uniformly bounded
in $[0,1]$. In the second step, we have used the definition of
$\bar{R}_{h}(u,(s,a)):=R_{h}(u,(s,a))\mathbbm{1}((s,a)\in\mathcal{G}_{h}^{\complement})$.
In the third step, we have used the guarantee in (12). In the fourth step, we
have used the fact that $\bar{Pi}$ maximizes the reward $\bar{R}_{h}$. In the
fifth step, we have used the fact that
$R_{h}(u,(s,a))\geq\bar{R}_{h}(u,(s,a))$ uniformly. From the discussion above,
this concludes the proof of the theorem. ∎

##  Appendix C Analysis - Linear MDPs

######  Lemma C.1.

Suppose Assumption 5.5 holds. Let $\kappa>1$ and $T\geq
C\frac{d\kappa^{2}}{(\gamma-\epsilon)^{2}}\log\tfrac{d\kappa}{\gamma-\epsilon}$.
With probability at-least $1-H\exp(-c(\gamma-\epsilon)T)$, the output of
Algorithm 2 returns $\phi_{th}$ such that
$\sum_{t=1}^{T}\phi_{th}\phi_{th}^{\intercal}\succeq\kappa^{2}I$ for every
$h\in[H]$

###### Proof of Theorem 5.8.

By Theorem 1 in (Wagenmaker et al., 2022), we take
$K^{\mathsf{rf}}(\epsilon,\delta/4)=\frac{CdH^{5}(d+\log(\tfrac{1}{\delta}))}{\epsilon^{2}}+\frac{Cd^{9/2}H^{6}\log^{4}(\tfrac{1}{\delta})}{\epsilon}$.
Phase 1 succeeds with probability $1-\tfrac{\delta}{4}$.

Note that this is the quantity $T_{\mathsf{rf}}$ in the statement of the
theorem. We now condition on the success of Phase 1. The number of
trajectories queried by Algorithm 2 which is given by $HT=T_{\mathsf{pol}}$.
By Lemma C.1, we conclude that for the given values of $T$ and $\kappa$, this
algorithm successfully outputs $\phi_{ht}$ such that
$G_{\phi,h}\succeq\kappa^{2}I$ for every $h\in[H]$, with probability at-least
$1-\frac{\delta}{4}$.

Now, condition on the success of Algorithm 2. By theorem 6.2, we conclude that
with these conclude that with proabability at-least $1-\tfrac{\delta}{4}$,
with the values of the given parameters, for every $h\in[H]$, the procedure in
Step 2 of Phase 2 outputs a policy $\hat{\Pi}^{f,h}$ such that whenever
$S_{1:H},A_{1:H}\sim\mathcal{M}(\hat{\Pi}^{f,h})$, the conditions in (4) is
satisfied for the random vector $\psi(S_{h},A_{h})$ with $\zeta$ replaced by
$\zeta/2$. We then use the active learning based matrix completion procedure
given in Section 7, where the vectors $\psi_{jk}$ are sample using the policy
$\hat{\Pi}^{f,h}$ on the given user. By theorem 7.1, we conclude that
conditioned on the success of all the steps above, with probability
$1-\frac{\delta}{4}$, we can exactly estimate each of the matrices
$\Theta_{h}^{*}$ for $h\in[H]$ with $T_{\mathsf{mat-comp}}$ number of samples.

Upon the success of Phases 1, 2, 3 (which occurs with probability at-least
$1-\delta$ by union bound), we conclude that Phase 4 gives the $\epsilon$
optimal policy for each user $u\in[N]$ because of the guarantees of reward
free RL.

∎

##  Appendix D Deferred Proofs

###  D.1 Proof of Lemma B.1

###### Proof.

We suppose that the reward free RL in Phase 1 succeeds and returns the
$\frac{\epsilon}{8}$ optimal policy for every choice of rewards bounded in
$[0,1]$. The algorithm terminates whenever the active sets are such that

$\hat{V}(\mathcal{J}(;\mathcal{G}))\leq\frac{\epsilon}{2}$ (14)

Note that by the definition of $\mathcal{J}(;\mathcal{G})$, the maximum value
for the MDP with reward $\mathcal{J}(;\mathcal{G})$ is
$\sup_{\Pi}\sum_{h=1}^{H}P^{\Pi}(\mathcal{G}_{h})$. Since $\hat{V}$ is the
output of the reward free RL algorithm, we conclude that we have:

$|\hat{V}(\mathcal{J}(;\mathcal{G}))-\sup_{\Pi}\sum_{h=1}^{H}P^{\Pi}_{h}(\mathcal{G}_{h})|\leq\frac{\epsilon}{8}$
(15)

We conclude via (14) and (15) that (10) holds, which establishes the second
part of the theorem. We now consider the termination time.

Suppose $\mathcal{G}^{(t)}$ is the sequence of active sets before termination
at step $t$ (i.e, it satisfies
$\hat{V}(\mathcal{J}(;\mathcal{G}^{(t)}))>\frac{\epsilon}{2}$). Recall
$\hat{\Pi}$, the output of the reward free RL algorithm. It follows from the
guarantees for reward free RL that:

$|\sum_{h=1}^{H}P_{h}^{\hat{\Pi}^{\mathcal{G}}}(\mathcal{G}^{(t)}_{h})-\sup_{\Pi}\sum_{h=1}^{H}P^{\Pi}_{h}(\mathcal{G}^{(t)}_{h})|\leq\frac{\epsilon}{8}$

Combining this with Equation (15) and the fact that
$\hat{V}(\mathcal{J}(;\mathcal{G}^{(t)}))>\frac{\epsilon}{2}$, we conclude:

$\sum_{h=1}^{H}P_{h}^{\hat{\Pi}^{\mathcal{G}}}(\mathcal{G}^{(t)}_{h})\geq\frac{\epsilon}{4}$
(16)

We consider the potential function with $\varphi(0)=0$ and
$\varphi(t)=\sum_{h\in[H]}\sum_{(s,a)\in\mathcal{S}\times\mathcal{A}}T^{(t)}_{h,(s,a)}$,
where $T^{(t)}_{h,(s,a)}$ is the counter $T_{h,(s,a)}$ inside Algorithm 1 at
the beginning of the step $t$.

Whenever $\mathcal{G}^{(t)}$ is such that
$\hat{V}(\mathcal{J}(;\mathcal{G}^{(t)}))>\tfrac{\epsilon}{2}$, we define
$N_{t}:=\varphi(t+1)-\varphi(t)$ (i.e, before termination). Just for the sake
of theoretical arguments, we define the fictious random variables
$N_{t}=\mathsf{Ber}(\frac{\epsilon}{8})$ i.i.d after termination. Let
$\mathcal{F}_{t}=\sigma(\mathcal{G}^{(s)},S_{1:H}^{(s)},A_{1:H}^{(s)},R_{1:H}^{(s)},U^{(s)}:s\leq
t)$

######  Claim D.1.

The following relations hold:

  1. 1.

$\mathbb{E}\left[N_{t}|\mathcal{F}_{t}\right]\geq\frac{\epsilon}{8}$

  2. 2.

$\mathbb{E}\left[N_{t}^{2}|\mathcal{F}_{t}\right]\leq\frac{\mathbb{E}\left[N_{t}|\mathcal{F}_{t}\right]H}{4}$

  3. 3.

$|N_{t}|\leq H$ almost surely.

###### Proof.

The inequalities are clear when $\mathcal{G}^{(t)}$ is such that
$\hat{V}(\mathcal{J}(;\mathcal{G}^{(t)}))\leq\tfrac{\epsilon}{2}$. Now
consider the case
$\hat{V}(\mathcal{J}(;\mathcal{G}^{(t)}))>\tfrac{\epsilon}{2}$. By definition,
conditioned on this event, we have almost surely:

$N_{t}=\sum_{h=1}^{H}\mathbbm{1}((S^{(t)}_{h},A^{(t)}_{h})\in\mathcal{G}^{(t)}_{h}).\mathbbm{1}(\hat{R}_{h}^{(t)}(U_{t},(S_{h}^{(t)},A_{h}^{(t)}))=*)$

That is, we increment the $T_{h,(s,a)}$ only when we encounter an element of
the active set such that the entry for this user has not been observed before.
Observe that for any arbitrary $(s,a)\in\mathcal{S}\times\mathcal{A}$

$\displaystyle\mathbb{P}\left(\hat{R}_{h}^{(t)}(U_{t},(s,a))=*\bigr{|}\mathcal{F}_{t},(S_{h}^{(t)},A_{h}^{(t)})=(s,a)\right)=\frac{|\\{u:\hat{R}^{(t)}_{h}(s,a)=*\\}|}{N}\,.$
(17)

This is true since the law of $S_{h}^{(t)},A_{h}^{(t)}$ is independent of
$U_{t}$ (since all users share the same MDP), when conditioned on
$\mathcal{F}_{t}$. Now, the algorithm only fills the column corresponding to
$(s,a)$ until the number of entries is smaller than $Np\leq\frac{N}{2}$. We
conclude that:

$|\\{u:\hat{R}^{(t)}(h,(s,a))=*\\}|\geq N-Np\geq\frac{N}{2}\,.$

This allows us to conclude
$\mathbb{P}\left(\hat{R}_{h}^{(t)}(U_{t},(s,a))=*\bigr{|}\mathcal{F}_{t},(S_{h}^{(t)},A_{h}^{(t)})=(s,a)\right)\geq\frac{1}{2}$
and hence:

$\displaystyle\mathbb{E}N_{t}$
$\displaystyle=\sum_{h=1}^{H}\mathbb{E}\mathbbm{1}((S^{(t)}_{h},A^{(t)}_{h})\in\mathcal{G}^{(t)}_{h}).\mathbbm{1}(R_{h}^{(t)}(U_{t},(S_{h}^{(t)},A_{h}^{(t)}))=*)$
$\displaystyle\geq\frac{1}{2}\sum_{h=1}^{H}\mathbb{E}\mathbbm{1}((S^{(t)}_{h},A^{(t)}_{h})\in\mathcal{G}^{(t)}_{h})$
$\displaystyle=\frac{1}{2}\sum_{h=1}^{H}P^{\hat{\Pi}^{\mathcal{G}}}_{h}(\mathcal{G}_{h}^{(t)})\geq\frac{\epsilon}{8}$
(18)

In the last step we have used (16). The bound $|N_{t}|\leq H$ almost surely
follows from definition. Now note that
$\mathbb{E}\left[N_{t}^{2}|\mathcal{F}_{t}\right]\leq
H\mathbb{E}\left[N_{t}|\mathcal{F}_{t}\right]$.

∎

######  Claim D.2.

For any $\tau\in\mathbb{N}$ and some $c_{0}>0$ small enough, we have:

$\mathbb{P}\left(\sum_{t=0}^{\tau-1}N_{t}<\frac{\epsilon\tau}{16}\right)\leq\exp(-c_{0}\tfrac{\epsilon\tau}{H})$

###### Proof.

For $\frac{3}{4H}>\lambda>0$, consider
$M_{t}=-\tfrac{\lambda^{2}\mathbb{E}\left[N_{t}^{2}|\mathcal{F}_{t}\right]}{1-\tfrac{\lambda
H}{3}}+\lambda\left(\mathbb{E}\left[N_{t}|\mathcal{F}_{t}\right]-N_{t}\right)$.
Now consider:

$\displaystyle\mathbb{E}\exp(\sum_{t=0}^{\tau-1}M_{t})$
$\displaystyle=\mathbb{E}\left[\mathbb{E}\left[\exp(M_{\tau-1})|\mathcal{F}_{\tau-1}\right]\exp\left(\sum_{t=0}^{\tau-1}M_{t}\right)\right]$
$\displaystyle=\mathbb{E}\left[\exp(\lambda\mathbb{E}\left[N_{\tau-1}|\mathcal{F}_{\tau-1}\right]-\lambda
N_{t})|\mathcal{F}_{\tau-1}\right]\exp\left(\sum_{t=0}^{\tau-2}M_{t}\right)\exp\left(-\tfrac{\lambda^{2}\mathbb{E}\left[N_{\tau-1}^{2}|\mathcal{F}_{\tau-1}\right]}{1-\tfrac{\lambda
H}{3}}\right)$
$\displaystyle\leq\mathbb{E}\exp\left(-\tfrac{\lambda^{2}\mathbb{E}\left[N_{\tau-1}^{2}|\mathcal{F}_{\tau-1}\right]}{1-\tfrac{\lambda
H}{3}}\right)\exp\left(\sum_{t=0}^{\tau-2}M_{t}\right)\exp\left(-\tfrac{\lambda^{2}\mathbb{E}\left[N_{\tau-1}^{2}|\mathcal{F}_{\tau-1}\right]}{1-\tfrac{\lambda
H}{3}}\right)$ $\displaystyle=\mathbb{E}\exp(\sum_{t=0}^{\tau-2}M_{t})$ (19)

In the first step we have used the fact that $\sum_{t=0}^{\tau-2}M_{t}$ is
$\mathcal{F}_{\tau-1}$ measurable and the towering property of conditional
expectation. In the third step, we have used the exponential moment bound
given in Exercise 2.8.5 in (Vershynin, 2018), as applied to
$N_{\tau}-\mathbb{E}\left[N_{\tau-1}|\mathcal{F}_{\tau-1}\right]$ along with
the fact that $N_{t}\in[0,H]$ almost surely. From (19), we conclude that
$\mathbb{E}\exp(\sum_{t=0}^{\tau}M_{t})\leq 1$ and thus applying the Chernoff
bound, we conclude that for any $\beta>0$

$\mathbb{P}\left(\sum_{t=0}^{\tau-1}-\tfrac{\lambda\mathbb{E}\left[N_{t}^{2}|\mathcal{F}_{t}\right]}{1-\tfrac{\lambda
H}{3}}+\left(\mathbb{E}\left[N_{t}|\mathcal{F}_{t}\right]-N_{t}\right)>\beta\right)\leq\exp(-\lambda\beta)$

Now, using item 2 from Claim D.1, we conclude that

$\mathbb{P}\left(\sum_{t=0}^{\tau-1}N_{t}<-\beta+\tfrac{3-4\lambda
H}{3-\lambda
H}\sum_{t=0}^{\tau-1}\mathbb{E}\left[N_{t}|\mathcal{F}_{t}\right]\right)\leq\exp(-\lambda\beta)$

Now, using item 1 from Claim D.1, we note that
$\sum_{t=0}^{\tau-1}\mathbb{E}\left[N_{t}|\mathcal{F}_{t}\right]\geq\frac{\epsilon\tau}{8}$
almost surely. Setting $\lambda=\frac{1}{4H}$ and $\beta=c_{0}\epsilon\tau$
for some small enough constant $\epsilon$, we conclude:

$\mathbb{P}\left(\sum_{t=0}^{\tau-1}N_{t}<\frac{\epsilon\tau}{16}\right)\leq\exp(-c_{0}\tfrac{\epsilon\tau}{H})$

∎

Let $\tau^{\mathsf{term}}$ denote the termination time for the algorithm. This
is true since $\varphi(t)$ is increasing in $t$, $\varphi(t)\leq
NpH|\mathcal{S}||\mathcal{A}|$, and strict inequality holds when
$t<\tau^{\mathsf{term}}$. For every $\tau<\tau^{\mathsf{term}}$ we have
$\varphi(\tau)=\sum_{t=0}^{\tau-1}N_{\tau}<NpH|\mathcal{S}||\mathcal{A}|$.
Therefore, we have the following relationship between the events:

$\\{\tau^{\mathsf{term}}>\tau\\}\subseteq\bigr{\\{}\sum_{t=1}^{\tau}N_{\tau}<Np|\mathcal{S}||\mathcal{A}|H\bigr{\\}}$

Setting $\tau=\frac{16Np|\mathcal{S}||\mathcal{A}|H}{\epsilon}$, we have:

$\mathbb{P}(\tau^{\mathsf{term}}>\tau)\leq\mathbb{P}\left(\sum_{t=1}^{\tau}N_{\tau}<Np|\mathcal{S}||\mathcal{A}|H\right)\leq\exp(-cNp|\mathcal{S}||\mathcal{A}|)$

∎

###  D.2 Proof of Lemma C.1

###### Proof.

By Remark 4.3 in (Wagenmaker et al., 2022), we show that non-linear rewards
can be handled by the reward free RL algorithm in Phase 1 as long all the
reward are uniformly bounded in $[0,1]$. Let $B_{th}$ be the matrix
$I+A_{\phi}$ in Algorithm 2 at step $t$ for horizon $h$. Similarly, let the
corresponding projection $Q$ be $Q_{th}$. Recall that $Q_{th}$ is the
projection onto an eigenspace of $B_{th}$. Now, suppose
$S_{1:H},A_{1:H}\sim\mathcal{M}_{U_{t}}(\hat{\Pi}^{Q_{t,h}})$ as in the
algorithm. Let $\phi_{th}:=\phi(S_{h},A_{h})$. Now, if $Q_{th}\neq 0$, then:

$\displaystyle\phi_{th}^{\intercal}B_{th}^{-1}\phi_{th}$
$\displaystyle\geq\phi_{th}^{\intercal}Q_{th}B_{th}^{-1}Q_{th}\phi_{th}$
$\displaystyle\geq\phi_{th}^{\intercal}Q_{th}\frac{I}{1+\kappa^{2}}Q_{th}\phi_{th}$
$\displaystyle=\frac{\|Q_{th}\phi_{th}\|^{2}}{1+\kappa^{2}}$ (20)

In the first step, we have used the fact that $Q_{th}$ is the projector to the
eigenspace of $B_{th}^{-1}$. In the second step, we have used the fact that
over the eigenspace corresponding to $Q_{th}$, the eigenvalues of
$B_{th}^{-1}$ are at-least $\frac{1}{1+\kappa^{2}}$. We now invoke Assumption
5.5 in order to show that, along with the guarantees of reward free RL in
phase 1, we conclude that:

$\mathbb{E}\left[\|Q_{th}\phi_{th}\|^{2}|Q_{th}\neq
0,B_{th}\right]\geq\gamma-\epsilon$ (21)

Now, note by the fact that $Q_{th}$ is a projector and that $\|\phi_{th}\|\leq
1$, we have:

$\mathbb{E}\left[\|Q_{th}\phi_{th}\|^{4}\bigr{|}Q_{th}\neq
0,B_{th}\right]\leq\mathbb{E}\left[\|Q_{th}\phi_{th}\|^{2}\bigr{|}Q_{th}\neq
0,B_{th}\right]$ (22)

Recall the Paley-Zygmund inequality which states that for any positive random
variable $Z$, we must have:
$\mathbb{P}(Z>\frac{\mathbb{E}Z}{2})\geq\frac{1}{4}\frac{(\mathbb{E}Z)^{2}}{\mathbb{E}Z^{2}}$.
Therefore,

$\displaystyle\mathbb{P}\left[\phi_{th}^{\intercal}B_{th}^{-1}\phi_{th}>\frac{\gamma-\epsilon}{2(1+\kappa^{2})}\biggr{|}Q_{th}\neq
0,B_{th}\right]\geq\mathbb{P}\left[\|Q_{th}\phi_{th}\|^{2}>\frac{\gamma-\epsilon}{2}\biggr{|}Q_{th}\neq
0,B_{th}\right]$
$\displaystyle\geq\mathbb{P}\left[\|Q_{th}\phi_{th}\|^{2}>\frac{1}{2}\mathbb{E}\left[\|Q_{th}\phi_{th}\|^{2}\bigr{|}Q_{th}\neq
0,B_{th}\right]\biggr{|}Q_{th}\neq 0,B_{th}\right]$
$\displaystyle\geq\frac{1}{4}\frac{\mathbb{E}\left[\|Q_{th}\phi_{th}\|^{2}\bigr{|}Q_{th}\neq
0,B_{th}\right]^{2}}{\mathbb{E}\left[\|Q_{th}\phi_{th}\|^{4}\bigr{|}Q_{th}\neq
0,B_{th}\right]}\geq\frac{1}{4}\mathbb{E}\left[\|Q_{th}\phi_{th}\|^{2}\bigr{|}Q_{th}\neq
0,B_{th}\right]$ $\displaystyle\geq\frac{\gamma-\epsilon}{4}$ (23)

In the first step, we have used (20). In the second step, we have used (21).
In the third step, we have used the Paley-Zygmund inequality and the moment
bound in (22).

Define the stopping time $\tau=\inf\\{t\leq T:Q_{th}=0\\}$ and $\tau=\infty$
if the set in the RHS is empty. Let $\Xi^{0}_{t}$ for
$t\in\\{0\\}\cup\mathbb{N}$ be a sequence of i.i.d random variables with the
law
$\frac{\gamma-\epsilon}{2(1+\kappa^{2})}\mathsf{Ber}(\frac{\gamma-\epsilon}{4})$.
We consider the sequence of random variables
$\Xi_{t}=\phi_{th}^{\intercal}B_{th}^{-1}\phi_{th}$ for $t<\tau$ and
$\Xi_{t}=\Xi^{0}_{t}$ for $t\geq\tau$

Now, we apply the matrix determinant lemma which states that
$\det(B+uu^{\intercal})=\det(B)(1+u^{\intercal}B^{-1}u)$. We note that
$B_{(t+1)h}=B_{th}+\phi_{th}\phi_{th}^{\intercal}$. Therefore, whenever
$t<\tau$, we must have:

$\det(B_{(t+1)h})=\det(B_{th})(1+\Xi_{t})$ (24)

Since $\|\phi_{th}\|\leq 1$ almost surely, we must have

$\mathsf{Tr}(B_{th})=\sum_{i=1}^{d}\langle e_{i},B_{th}e_{i}\rangle\leq d+t$

It is easy to show that for any PSD matrix, $A$, if
$\mathsf{Tr}(A)\leq\alpha$, then $\det(A)\leq(\frac{\alpha}{d})^{d}$ (since
trace is the sum of the eigenvalues and the determinant is the product).
Combining the equations above, we conclude that whenever $t<\tau$, we must
have:

$\left(\frac{t+1+d}{d}\right)^{d}\geq\prod_{s=0}^{t}(1+\Xi_{t})$

Therefore, the event

$\\{\tau>T\\}\subseteq\\{\left(\tfrac{T+1+d}{d}\right)^{d}\geq\prod_{s=0}^{T}(1+\Xi_{t})\\}$
(25)

######  Claim D.3.

$\mathbb{P}\left[\prod_{s=0}^{T}(1+\Xi_{t})\geq\left(1+\frac{\gamma-\epsilon}{2(1+\kappa^{2})}\right)^{\frac{(\gamma-\epsilon)T}{8}}\right]\geq
1-\exp\left(-c_{0}T(\gamma-\epsilon)\right)$

Let $\kappa>1$ and $T\geq
C\frac{d\kappa^{2}}{(\gamma-\epsilon)^{2}}\log\tfrac{d\kappa}{\gamma-\epsilon}$,
we have:

$\mathbb{P}\left[\prod_{s=0}^{T}(1+\Xi_{t})\geq\left(\tfrac{T+1+d}{d}\right)^{d}\right]\geq
1-\exp\left(-c_{0}T(\gamma-\epsilon)\right)$

###### Proof.

Let $N_{T}$ be the number of variables $(\Xi_{t})_{t=0}^{T}$ such that
$\Xi_{t}\geq\frac{\gamma-\epsilon}{2\kappa^{2}}$. Then, it is clear that
$\prod_{s=0}^{T}(1+\Xi_{t})\geq(1+\frac{\gamma-\epsilon}{2\kappa^{2}})^{N_{T}}$.

Therefore,

$\displaystyle\mathbb{P}\left[\prod_{s=0}^{T}(1+\Xi_{t})\geq\left(1+\frac{\gamma-\epsilon}{2(1+\kappa^{2})}\right)^{\frac{(\gamma-\epsilon)T}{8}}\right]$
$\displaystyle\geq\mathbb{P}\left(N_{T}\geq\frac{(\gamma-\epsilon)T}{8}\right)$
$\displaystyle\geq\mathbb{P}\left(\mathsf{Bin}(T,\tfrac{\gamma-\epsilon}{4})\geq\frac{(\gamma-\epsilon)T}{8}\right)$
$\displaystyle\geq 1-\exp\left(-c_{0}T(\gamma-\epsilon)\right)$ (26)

Here $\mathsf{Bin}$ refers to the law of a binomial random variable. The first
step follows from the fact that
$\prod_{s=0}^{T}(1+\Xi_{t})\geq(1+\frac{\gamma-\epsilon}{2(1+\kappa^{2})})^{N_{T}}$
almost surely. The second step follows from (23), which shows that conditioned
on $Q_{th},B_{th}$, the random variable
$\mathbbm{1}(\Xi_{t}\geq\frac{\gamma-\epsilon}{2(1+\kappa^{2})})$
stochastically dominates $\mathsf{Ber}(\tfrac{\gamma-\epsilon}{4})$. The last
step follows from an application of Bernstein’s inequality for binomial random
variables. ∎

Now, using (25) along with Claim D.3, we conclude:

$\mathbb{P}(\tau>T)\leq\mathbb{P}\left(\left(\tfrac{T+1+d}{d}\right)^{d}\geq\prod_{s=0}^{T}(1+\Xi_{t})\right)\leq\exp(-c_{0}T(\gamma-\epsilon))$

∎

###  D.3 Proof of Lemma 6.1

###### Proof.

By the definition of Linear MDP, we must have
$S_{h+1}|S_{h},A_{h}\sim\sum_{i=1}^{d}\langle\phi(S_{h},A_{h}),e_{i}\rangle\mu_{ih}(\cdot)$
and $A_{h+1}|S_{h+1}\sim\pi_{h+1}(\cdot|S_{h})$. Therefore, for any bounded,
measurable function $g:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$, we must
have:

$\displaystyle\mathbb{E}g(S_{h+1},A_{h+1})$
$\displaystyle=\mathbb{E}\left[\mathbb{E}\left[g(S_{h+1},A_{h+1})\bigr{|}S_{h},A_{h}\right]\right]$
$\displaystyle=\mathbb{E}\sum_{i=1}^{d}\langle\phi(S_{h},A_{h}),e_{i}\rangle\int\mu_{i(h-1)}(ds)\pi_{h+1}(da|s)g(s,a)$
$\displaystyle=\sum_{i=1}^{d}\nu_{ih}\int\mu_{ih}(ds)\pi_{h+1}(da|s)g(s,a)$
(27)

∎

###  D.4 Proof of Lemma E.1

###### Proof.

It is clear from the assumption that $\mathbb{E}\left[\int
g(s_{(h+1)t},a)\pi_{h+1}(da|s_{t})|(\phi_{ht})_{t\leq
T}\right]=\sum_{i=1}^{d}\langle\phi_{ht},e_{i}\rangle\int\mu_{ih}(ds)\pi_{h+1}(da|s)g(s,a)$.

Note that

$\displaystyle\sum_{t=1}^{T}\alpha_{ht,\nu}\phi_{ht}$
$\displaystyle=\sum_{t=1}^{T}\phi_{ht}^{\intercal}G_{\phi,h}^{-1}\nu\phi_{ht}$
$\displaystyle=\sum_{t=1}^{T}\phi_{ht}\phi_{ht}^{\intercal}G_{\phi,h}^{-1}\nu=(\sum_{t=1}^{T}\phi_{ht}\phi_{ht}^{\intercal})G_{\phi,h}^{-1}\nu$
$\displaystyle=G_{\phi,h}G_{\phi,h}^{-1}\nu=\nu$ (28)

Therefore,

$\displaystyle\mathbb{E}\left[\hat{\mathcal{T}}(g;\nu,\pi_{h})|(\phi_{l})_{t\in[T]}\right]$
$\displaystyle=\sum_{i=1}^{d}\langle\sum_{t=1}^{T}\alpha_{t,\nu}\phi_{t},e_{i}\rangle\int\mu_{i(h-1)}(ds)\pi_{h}(da|s)g(s,a)$
$\displaystyle=\sum_{i=1}^{d}\langle\nu,e_{i}\rangle\int\mu_{i(h-1)}(ds)\pi_{h}(da|s)g(s,a)=\mathcal{T}(g;\nu,\pi_{h})$
(29)

Note that, conditioned on $(\phi_{t})_{t\in[T]}$, $\alpha_{ht,\nu}\int
g(s_{(h+1)t},a)\pi_{h+1}(da|s_{(h+1)t})$ are independent random variables
bounded above by $\alpha_{ht,\nu}B$. Therefore, applying the Azuma-Hoeffding
inequality, we conclude:

$\mathbb{P}\left(|\hat{\mathcal{T}}(g;\nu,\pi_{h})-\mathcal{T}(g;\nu,\pi_{h})|>\beta\biggr{|}(\phi_{t})_{t\in[T]}\right)\leq
2\exp\left(-\tfrac{t^{2}}{2B^{2}\sum_{t}\alpha_{t,\nu}^{2}}\right)$

Now, observe that
$\sum_{t}\alpha_{ht,\nu}^{2}=\sum_{t}\nu^{\intercal}G_{\phi,h}^{-1}\nu\leq\frac{1}{\kappa^{2}}$
whenever $G_{\phi,h}\succeq\kappa^{2}I$ This concludes the proof. ∎

###  D.5 Proof of Lemma E.2

###### Proof.

Notice that:

$\displaystyle\bigr{|}E_{1}^{\nu_{1}}(\Pi)-E_{1}^{\nu_{1}^{\prime}}(\Pi^{\prime})\bigr{|}\leq\bigr{|}E_{1}^{\nu_{1}}(\Pi)-E_{1}^{\nu_{1}}(\Pi^{\prime})\bigr{|}+\bigr{|}E_{1}^{\nu_{1}}(\Pi^{\prime})-E_{1}^{\nu_{1}^{\prime}}(\Pi^{\prime})\bigr{|}$
$\displaystyle\leq\bigr{|}E_{1}^{\nu_{1}}(\Pi)-E_{1}^{\nu_{1}}(\Pi^{\prime})\bigr{|}+\|\nu_{1}-\nu_{1}^{\prime}\|_{1}$
$\displaystyle\leq\bigr{\|}\mathbb{E}\int\phi(S_{1},a)\pi_{1}(da|S_{1})-\mathbb{E}\int\phi(S_{1},a)\pi^{\prime}_{1}(da|S_{1})\bigr{\|}_{1}+\|\nu_{1}-\nu_{1}^{\prime}\|_{1}$
$\displaystyle\leq\sup_{(s,a)}\|\phi(s,a)\|_{1}\mathsf{TV}(\pi_{1},\pi_{1}^{\prime})+\|\nu_{1}-\nu_{1}^{\prime}\|_{1}\leq\mathsf{TV}(\pi_{1},\pi_{1}^{\prime})+\|\nu_{1}-\nu_{1}^{\prime}\|_{1}$
(30)

In the first, second and third steps we have used the triangle inequality. In
the last step, we have used the fact that for any bounded function, and any
probability measures $\mu,\nu$, we have $|\int f(x)\mu(dx)-\int
f(x)\nu(dx)|\leq\sup_{x}|f(x)|\mathsf{TV}(\nu,\mu)$.

Now consider:

$\displaystyle\bigr{|}\|\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\nu_{j}\|_{1}-\|\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi^{\prime}_{j})-\nu^{\prime}_{j}\|_{1}\bigr{|}$
$\displaystyle\leq\|\nu_{j}-\nu_{j}^{\prime}\|_{1}+\bigr{\|}\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi^{\prime}_{j})\bigr{\|}_{1}$
$\displaystyle\leq\|\nu_{j}-\nu_{j}^{\prime}\|_{1}+\bigr{\|}\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi_{j})\bigr{\|}_{1}+\bigr{\|}\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi_{j})-\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi^{\prime}_{j})\bigr{\|}_{1}$
(31)

Now, observe that:

$\displaystyle\bigr{\|}\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi_{j})\bigr{\|}_{1}\leq\sum_{i=1}^{d}|\langle\nu_{j-1}-\nu^{\prime}_{j-1},e_{i}\rangle|\biggr{\|}\int\phi(s,a)\mu_{i(j-1)}(ds)\pi_{j}(da|s)\biggr{\|}_{1}$
$\displaystyle\leq\sum_{i=1}^{d}|\langle\nu_{j-1}-\nu^{\prime}_{j-1},e_{i}\rangle|=\|\nu_{j-1}-\nu_{j-1}^{\prime}\|_{1}$
(32)

Where we recall $\sup_{i,h,\pi}\|\int\phi(s,a)\mu_{ih}(ds)\pi(da|s)\|_{1}\leq
1$ as given in the definition of Linear MDP. Using the Hahn-Jordan
decomposition of a signed measure, we conclude:

$\displaystyle\bigr{\|}\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi_{j})-\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi_{j})\bigr{\|}_{1}$
$\displaystyle\leq\sum_{i=1}^{d}|\langle\nu^{\prime}_{j-1},e_{i}\rangle|\biggr{\|}\int\phi(s,a)\mu_{i(j-1)}(ds)(\pi_{j}(da|s)-\pi_{j}^{\prime}(da|s))\biggr{\|}_{1}$
$\displaystyle\leq\sum_{i=1}^{d}C_{\mu}|\langle\nu^{\prime}_{j-1},e_{i}\rangle|\mathsf{TV}(\pi_{j},\pi_{j}^{\prime})\leq
C_{\mu}\|\nu_{j-1}^{\prime}\|_{1}\mathsf{TV}(\pi_{j},\pi_{j}^{\prime})$ (33)

Combining (31), (32) and (33) we conclude:

$\displaystyle\bigr{|}\|\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\nu_{j}\|_{1}-\|\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi^{\prime}_{j})-\nu^{\prime}_{j}\|_{1}\bigr{|}$
$\displaystyle\leq\|\nu_{j}-\nu_{j}^{\prime}\|_{1}+C_{\mu}\|\nu_{j-1}^{\prime}\|_{1}\mathsf{TV}(\pi_{j},\pi_{j}^{\prime})+\|\nu_{j-1}-\nu_{j-1}^{\prime}\|_{1}$
(34)

Combining (30) and (34), we conclude the first inequality in the statement of
the lemma.

With a reasoning very similar to that in (30), we have:

$\displaystyle\bigr{|}\hat{E}_{1}^{\nu_{1}}(\Pi)-\hat{E}_{1}^{\nu_{1}^{\prime}}(\Pi^{\prime})\bigr{|}\leq\mathsf{TV}(\pi_{1},\pi_{1}^{\prime})+\|\nu_{1}-\nu_{1}^{\prime}\|_{1}$
(35)

Using similar reasoning as in (34):

$\displaystyle\bigr{|}\|\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\nu_{j}\|_{1}-\|\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi^{\prime}_{j})-\nu^{\prime}_{j}\|_{1}\bigr{|}$
$\displaystyle\leq\|\nu_{j}-\nu_{j}^{\prime}\|_{1}+\left(\sum_{t=1}^{T}|(\nu_{j-1}-\nu_{j-1}^{\prime})^{\intercal}G_{\phi,j-1}^{-1}\phi_{(j-1)t}|\right)+\left(\sum_{t=1}^{T}|(\nu^{\prime}_{j-1})^{\intercal}G_{\phi,j-1}^{-1}\phi_{(j-1)t}|\right)\mathsf{TV}(\pi_{j},\pi_{j}^{\prime})$
(36)

Now note that for any $\nu\in\mathbb{R}^{d}$, we have:

$\displaystyle\sum_{t=1}^{T}\bigr{|}\nu^{\intercal}G_{\phi,j-1}^{-1}\phi_{(j-1)t}\bigr{|}$
$\displaystyle\leq\sqrt{T\sum_{t}\bigr{|}\nu^{\intercal}G_{\phi,j-1}^{-1}\phi_{(j-1)t}\bigr{|}^{2}}$
$\displaystyle=\sqrt{T\sum_{t=1}^{T}\nu^{\intercal}G_{\phi,j-1}^{-1}\phi_{(j-1)t}\phi_{(j-1)t}^{\intercal}G_{\phi,j-1}^{-1}\nu}$
$\displaystyle=\sqrt{T\nu^{\intercal}G_{\phi,j-1}^{-1}\nu}$
$\displaystyle\leq\sqrt{\tfrac{T}{\kappa^{2}}}\|\nu\|_{2}$ (37)

Here, in the first step we have used the fact that whenever
$x\in\mathbb{R}^{K}$, we must have $\|x\|_{1}\leq\sqrt{K}\|x\|_{2}$. In the
third step, we have used the fact that
$\sum_{t=1}^{T}\phi_{(j-1)t}\phi_{(j-1)t}^{\intercal}=G_{\phi,j-1}$ by
definition. In the last step, we have used the fact that
$\hat{G}_{\phi,j-1}\succeq\kappa^{2}I$. Plugging this into (36), we conclude:

$\displaystyle\bigr{|}\|\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\nu_{j}\|_{1}-\|\mathcal{T}_{j}(\phi,\nu^{\prime}_{j-1},\pi^{\prime}_{j})-\nu^{\prime}_{j}\|_{1}\bigr{|}$
$\displaystyle\leq\|\nu_{j}-\nu_{j}^{\prime}\|_{1}+\sqrt{\tfrac{T}{\kappa^{2}}}\|\nu_{j-1}-\nu_{j-1}^{\prime}\|_{2}+\sqrt{\tfrac{T}{\kappa^{2}}}\|\nu_{j-1}^{\prime}\|_{2}\mathsf{TV}(\pi_{j},\pi_{j}^{\prime})$
(38)

Using this and the definition of $\hat{F}$ we conclude the second inequality
in the statement of the lemma. (45) and (E.2) follow from a similar reasoning.
∎

###  D.6 Proof of Lemma E.3

###### Proof.

First consider the case $h=1$. Let $g(s,a):=\phi(s,a)$. In this case,
$\sup_{\nu\in\mathcal{B}_{d}(1)}|\hat{E}^{\nu,1}(\Pi)-E^{\nu,1}(\Pi)|\leq\|\mathcal{T}_{0}(g;\pi_{1})-\hat{\mathcal{T}}_{0}(g;\pi_{1})\|_{1}$.
By (30) and (35), we conclude that $\pi_{1}\to\mathcal{T}_{0}(\phi;\pi_{1})$
and $\pi_{1}\to\hat{\mathcal{T}}_{0}(\phi;\pi_{1})$ are $1$-Lipschitz with
respect to $\mathsf{TV}()$ and $\|\cdot\|_{1}$.

$\sup_{\Pi={\pi_{1},\dots,\pi_{H}}\in\mathbb{Q}}\|\mathcal{T}_{0}(g;\pi_{1})-\hat{\mathcal{T}}_{0}(g;\pi_{1})\|_{1}\leq\sup_{\Pi={\pi_{1},\dots,\pi_{H}}\in\hat{\mathbb{Q}}_{\eta}}\|\mathcal{T}_{0}(g;\pi_{1})-\hat{\mathcal{T}}_{0}(g;\pi_{1})\|_{1}+2\eta$

We apply Lemma E.1 co-ordinate wise to the co-ordinates of $\phi$ and union
bound over $\hat{\mathbb{Q}}_{\eta}$. We have:

$\mathbb{P}\left(\sup_{\Pi={\pi_{1},\dots,\pi_{H}}\in\mathbb{Q}}\|\mathcal{T}_{0}(g;\pi_{1})-\hat{\mathcal{T}}_{0}(g;\pi_{1})\|_{1}>2\eta+d\beta\right)\leq
d|\hat{\mathbb{Q}}_{\eta}|\exp(-\tfrac{\beta^{2}\kappa^{2}}{2})$ (39)

Now, consider $h>1$. Consider any $\eta$ net over $\mathcal{B}_{d}(1)$,
denoted by $\hat{\mathcal{B}}_{d,\eta}$ with respect to the norm
$\|\cdot\|_{1}$. We can take
$|\hat{\mathcal{B}}_{d,\eta}|\leq\exp(Cd\log(d/\eta))$ (Vershynin, 2018).
Invoking Lemma E.2, we conclude:

$\displaystyle\sup_{\Pi\in\mathbb{Q}}\sup_{\nu\in\mathcal{B}_{d}(1)}|E^{\nu,h}(\Pi)-\hat{E}^{\nu,h}(\Pi)|\leq\sup_{\begin{subarray}{c}\nu_{1},\dots,\nu_{h}\in\mathcal{B}_{d}(1)\\\
\Pi\in\mathbb{Q}\end{subarray}}|\hat{F}(\Pi,\nu_{1},\dots,\nu_{h})-F(\Pi,\nu_{1},\dots,\nu_{h})|$
$\displaystyle\leq\sup_{\begin{subarray}{c}\nu_{1},\dots,\nu_{h}\in\hat{\mathcal{B}}_{d,\eta}\\\
\Pi\in\hat{\mathbb{Q}}_{\eta}\end{subarray}}|\hat{F}(\Pi,\nu_{1},\dots,\nu_{h})-F(\Pi,\nu_{1},\dots,\nu_{h})|+2\left(1+C_{\mu}+\sqrt{\tfrac{T}{\kappa^{2}}}\right)\eta
h$ (40)

Now, by the triangle inequality, we have:

$\displaystyle|\hat{F}(\Pi,\nu_{1},\dots,\nu_{h})-F(\Pi,\nu_{1},\dots,\nu_{h})|$
$\displaystyle\leq\|\mathcal{T}_{0}(\phi,\pi_{1})-\hat{\mathcal{T}}_{0}(\phi,\pi_{1})\|_{1}+\sum_{j=1}^{h-1}\|\mathcal{T}_{j}(\phi,\nu_{j},\pi_{j+1})-\hat{\mathcal{T}}_{j}(\phi,\nu_{j},\pi_{j+1})\|_{1}$
(41)

Therefore, by invoking Lemma E.1, along with union bound over every component
in the sum in (D.6) and over the net in (40) we conclude that:

$\displaystyle\mathbb{P}\left[\sup_{\begin{subarray}{c}\nu_{1},\dots,\nu_{h}\in\hat{\mathcal{B}}_{d,\eta}\\\
\Pi\in\hat{\mathbb{Q}}_{\eta}\end{subarray}}|\hat{F}(\Pi,\nu_{1},\dots,\nu_{h})-F(\Pi,\nu_{1},\dots,\nu_{h})|>\beta
dh\right]$ $\displaystyle\leq
2dh|\hat{\mathbb{Q}}_{\eta}||\hat{\mathcal{B}}_{d,\eta}|^{h}\exp(-\tfrac{\beta^{2}\kappa^{2}}{2})$
(42)

Combining (40) and (42), we conclude the second item in the statement of the
lemma.

The concentration of $X_{1}$ and $X_{h}$ follow in a similar fashion, but here
we consider an $\eta$ net even over $x$ and use the Lipschitzness results
given in Lemma E.2 and the fact that $x\to f(\phi;x)$ is $1$ Lipschitz.

∎

##  Appendix E Proof of Theorem 6.2

######  Lemma E.1.

Suppose $h\in[H-1]$, and $g:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$ be such
that $|g(s,a)|\leq B$ for every $(s,a)$. For any policy $\pi_{h}$ and any
$\nu$ such that $\|\nu\|_{2}\leq 1$, we must have:

$\mathbb{P}\left(|\hat{\mathcal{T}}_{h}(g;\nu,\pi_{h})-\mathcal{T}_{h}(g;\nu,\pi_{h})|>\beta\biggr{|}(\phi_{ht})_{t\in[T]},G_{\phi,h}\succeq\kappa^{2}I\right)\leq
2\exp\left(-\tfrac{\beta^{2}\kappa^{2}}{2B^{2}}\right)$

######  Lemma E.2.

Let $\Pi=(\pi_{1},\dots,\pi_{H})$,
$\Pi^{\prime}=(\pi^{\prime}_{1},\dots,\pi^{\prime}_{H})$ be policies in
$\mathbb{Q}$. Conditioned on the event $G_{\phi,h}\succeq\kappa^{2}I$, the
following hold:

$\displaystyle|F(\Pi,\nu_{1},\dots,\nu_{h-1},\nu_{h})-F(\Pi^{\prime},\nu^{\prime}_{1},\dots,\nu^{\prime}_{h-1},\nu^{\prime}_{h})|$
$\displaystyle\leq\left(\sum_{j=2}^{h}C_{\mu}\mathsf{TV}(\pi_{j},\pi_{j}^{\prime})\|\nu_{j}\|_{1}+2\|\nu_{j}-\nu_{j}^{\prime}\|_{1}\right)+\mathsf{TV}(\pi_{1},\pi_{1}^{\prime})+2\|\nu_{1}-\nu_{1}^{\prime}\|_{1}$
(43)
$\displaystyle|\hat{F}(\Pi,\nu_{1},\dots,\nu_{h-1},\nu_{h})-\hat{F}(\Pi^{\prime},\nu^{\prime}_{1},\dots,\nu^{\prime}_{h-1},\nu^{\prime}_{h})|$
$\displaystyle\leq\sqrt{\tfrac{T}{\kappa^{2}}}\left(\sum_{j=2}^{h}\mathsf{TV}(\pi_{j},\pi_{j}^{\prime})\|\nu_{j}\|_{2}+\|\nu_{j}-\nu_{j}^{\prime}\|_{2}\right)+\mathsf{TV}(\pi_{1},\pi_{1}^{\prime})+\sum_{j=1}^{h}\|\nu_{j}-\nu_{j}^{\prime}\|_{1}$
(44)

Suppose $x\in\mathcal{S}^{d-1}$

$\displaystyle|\mathcal{T}_{h}(f(\cdot;x),\nu,\pi_{h})-\mathcal{T}_{h}(f(\cdot;x^{\prime}),\nu^{\prime},\pi^{\prime}_{h})|$
$\displaystyle\leq 2C_{\mu}(\sqrt{d}+\xi
d)\left(\|\nu-\nu^{\prime}\|_{1}+\mathsf{TV}(\pi_{h},\pi_{h}^{\prime})+\|x-x^{\prime}\|_{2}\|\nu\|_{1}\right)$
(45)

$\displaystyle|\hat{\mathcal{T}}_{h}(f(\cdot;x),\nu,\pi_{h})-\hat{\mathcal{T}}_{h}(f(\cdot;x^{\prime}),\nu^{\prime},\pi^{\prime}_{h})|$
$\displaystyle\leq 2\sqrt{\tfrac{T}{\kappa^{2}}}(\sqrt{d}+\xi
d)\left(\|\nu-\nu^{\prime}\|_{2}+\mathsf{TV}(\pi_{h},\pi_{h}^{\prime})+\|x-x^{\prime}\|_{2}\|\nu\|_{2}\right)$
(46)

######  Lemma E.3.

Condition on the event $G_{\phi,h}\succeq\kappa^{2}I$ for every $h\in[H]$. Fix
some $\eta>0$ and let $\hat{\mathbb{Q}}_{\eta}$ denote any $\eta$-net over
$\mathbb{Q}$. With probability at-least $1-\delta/4$, the following hold
simultaneously:

  1. 1.

$\sup_{\nu}\sup_{\Pi\in\mathbb{Q}}|E^{\nu}_{1}(\Pi)-\hat{E}^{\nu}_{1}(\Pi)|\leq
C\frac{d}{\kappa}\sqrt{\log\left(\tfrac{d|\hat{\mathbb{Q}}_{\eta}|}{\delta}\right)}+C\eta\,.$

  2. 2.

For $h>1$:

$\sup_{\Pi\in\mathbb{Q}}\sup_{\nu\in\mathcal{B}_{d}(1)}|E^{\nu}_{h}(\Pi)-\hat{E}^{\nu}_{h}(\Pi)|\leq\frac{CdH}{\kappa}\sqrt{\log\left(\tfrac{dH|\hat{\mathbb{Q}}_{\eta}|}{\delta}\right)+Hd\log\left(\tfrac{d}{\eta}\right)}+C\left(\sqrt{\tfrac{T}{\kappa^{2}}}\right)\eta
H$

  3. 3.

$X_{1}:=\sup_{\Pi=(\pi_{1},\dots,\pi_{H})\in\mathbb{Q}}\bigr{|}\inf_{x\in\mathcal{S}^{d-1}}\hat{\mathcal{T}}_{1}(f(\cdot;x),\pi_{1})-\inf_{x\in\mathcal{S}^{d-1}}\mathcal{T}_{1}(f(\cdot;x),\pi_{1})\bigr{|}$

$X_{1}\leq\frac{C(\sqrt{d}+\xi
d)}{\kappa}\sqrt{\log\left(\tfrac{|\hat{\mathbb{Q}}_{\eta}|}{\delta}\right)+d\log\left(\tfrac{d}{\eta}\right)}+C\eta(\sqrt{d}+\xi
d)$

  4. 4.

$X_{h}:=\sup_{\begin{subarray}{c}\nu\in\mathcal{B}(1)\\\
\Pi=(\pi_{1},\dots,\pi_{H})\in\mathbb{Q}\end{subarray}}\bigr{|}\inf_{x\in\mathcal{S}^{d-1}}\hat{\mathcal{T}}_{h}(f(\cdot;x);\nu,\pi_{h})-\inf_{x\in\mathcal{S}^{d-1}}\mathcal{T}_{h}(f(\cdot;x);\nu,\pi_{h})\bigr{|}$

$X_{h}\leq\frac{C(\sqrt{d}+\xi
d)}{\kappa}\sqrt{\log\left(\tfrac{|\hat{\mathbb{Q}}_{\eta}|H}{\delta}\right)+d\log\left(\tfrac{d}{\eta}\right)}+C\eta(\sqrt{d}+\xi
d)\left(C_{\mu}+\sqrt{\tfrac{T}{\kappa^{2}}}\right)$

######  Lemma E.4.

$\Pi=(\pi_{1},\dots,\pi_{H})$. For any $\eta\geq 0$, and $h\in[H]$, suppose
$E^{\nu}_{h}(\Pi)\leq\eta$. Then, we have:

$\|\mathbb{E}\phi(S_{h},A_{h})-\nu\|_{1}\leq\eta\,.$

###### Proof.

Let $S_{1:H},A_{1:H}\sim\mathcal{M}(\Pi)$. By Lemma 6.1, we conclude that:
$\mathbb{E}\phi(S_{1},A_{1})=\mathcal{T}_{1}(\phi,\pi_{1})$. Therefore we
conclude the lemma for the case $h=1$. Now let $h>1$.

Now, note that for $j>1$, we have:
$\mathbb{E}\phi(S_{j},A_{j})=\mathcal{T}_{j}(\phi,\mathbb{E}\phi(S_{j-1},A_{j-1}),\pi_{j})$.
There exists a sequence $\nu_{1},\dots,\nu_{h-1}$ such that

$E_{1}^{\nu_{1}}(\Pi)+\sum_{j=2}^{h}\|\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\nu_{j}\|_{1}\leq\eta_{0}$

Letting $E_{1}^{\nu_{1}}(\Pi)=:\eta_{1}$,
$\|\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\nu_{j}\|_{1}=:\eta_{j}$, we have
from the case $h=1$ :
$\|\mathbb{E}\phi(S_{1},A_{1})-\nu_{1}\|_{1}\leq\eta_{1}$.

$\displaystyle\|\mathbb{E}\phi(S_{j},A_{j})-\nu_{j}\|_{1}=\|\mathcal{T}_{j}(\phi,\mathbb{E}\phi(S_{j-1},A_{j-1}),\pi_{j})-\nu_{j}\|_{1}$
$\displaystyle\leq\|\mathcal{T}_{j}(\phi,\mathbb{E}\phi(S_{j-1},A_{j-1}),\pi_{j})-\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})\|_{1}+\|\mathcal{T}_{j}(\phi,\nu_{j-1},\pi_{j})-\nu_{j}\|_{1}$
$\displaystyle\leq\|\nu_{j-1}-\mathbb{E}\phi(S_{j-1},A_{j-1})\|_{1}+\eta_{j}$
(47)

We have used (32) in the last step. Continuing recursively, we conclude the
result ∎

###### Proof of Theorem 6.2.

We condition on the event described in Lemma E.3. We suppose that $\kappa$,
$\eta$ and $\eta_{0}$ are related as in the statement of the theorem. We will
apply these values whenever we invoke the concentration bounds obtained from
Lemma E.3 in all the inequalities below. First consider $h=1$. Let
$\hat{\Pi}^{f,1}=(\pi^{f,1}_{H},\dots,\pi^{f,H}_{H})$. By item 3 in Lemma E.3,
we have (with $X_{1}$ as defined in the lemma):

$\inf_{x\in\mathcal{S}^{d-1}}\hat{\mathcal{T}_{1}}(f(;x),\pi^{f,1}_{1})\geq\sup_{\Pi\in\mathbb{Q}}\inf_{x\in\mathcal{S}^{d-1}}\mathcal{T}_{1}(f(;x),\pi_{1})-X_{1}\geq\zeta-
X_{1}\geq\tfrac{3\zeta}{4}$

Similarly, we have:

$\inf_{x\in\mathcal{S}^{d-1}}\mathcal{T}_{1}(f(;x),\pi^{f,1}_{1})\geq\inf_{x\in\mathcal{S}^{d-1}}\hat{\mathcal{T}}_{1}(f(;x),\pi^{f,1}_{1})-\tfrac{\zeta}{4}$

Combining the two displays above, we conclude the theorem for $h=1$. Now
consider $h>1$. We will first show that the constraint
$\hat{E}^{\nu,h-1}(\Pi)\leq\eta_{0}$ is feasible for some $\Pi\in\mathbb{Q}$
and some $\nu$. Note that, for any policy $\Pi$ there exists a
$\nu_{1},\dots,\nu_{h-1}\in\mathbb{R}^{d}$ such that
$\mathbb{E}\phi(S_{j},A_{j})=\nu_{j}$ whenever
$S_{1:H},A_{1:H}\sim\mathcal{M}(\Pi)$. For the choice $\nu=\nu_{h-1}$, we must
have $E^{\nu,h-1}(\Pi)=0$. Now, by item 1 and 2 of Lemma E.3, we conclude that
$\hat{E}^{\nu,h-1}\leq\eta_{0}$. Therefore this optimization is feasible.

Consider the solutions to the optimization problem given by $\hat{\nu}$ and
$\hat{\Pi}^{f,h}$. Note again from Lemma E.3 that
$E^{\hat{\nu},h-1}(\hat{\Pi}^{f,h})\leq\hat{E}^{\hat{\nu},h-1}(\hat{\Pi}^{f,h})+\eta_{0}\leq
2\eta_{0}$. Now, applying Lemma E.4, we conclude that whenever
$S_{1:H},A_{1:H}\sim\mathcal{M}(\hat{\Pi}^{f,h})$

$\|\mathbb{E}\phi(S_{h-1},A_{h-1})-\hat{\nu}\|_{1}\leq 2\eta_{0}$

By a similar reasoning as the case $h=1$, we conclude:

$\inf_{x\in\mathcal{S}^{d-1}}\mathcal{T}_{h}(f(;x),\hat{\nu},\pi^{f,h}_{h})\geq
3\tfrac{\zeta}{4}$

Now, applying (45), we conclude:

$\displaystyle\inf_{x\in\mathcal{S}^{d-1}}\mathbb{E}f(S_{h},A_{h};x)=\inf_{x\in\mathcal{S}^{d-1}}\mathcal{T}_{h}(f(;x),\mathbb{E}\phi(S_{h-1},A_{h-1}),\pi^{f,h}_{h})$
$\displaystyle\geq\inf_{x\in\mathcal{S}^{d-1}}\mathcal{T}_{h}(f(;x),\hat{\nu},\pi^{f,h}_{h})-\sup_{x\in\mathcal{S}^{d-1}}|\mathcal{T}_{h}(f(;x),\mathbb{E}\phi(S_{h-1},A_{h-1}),\pi^{f,h}_{h})-\mathcal{T}_{h}(f(;x),\hat{\nu},\pi^{f,h}_{h})|$
$\displaystyle\geq\frac{3\zeta}{4}-2C_{\mu}(\sqrt{d}+\xi
d)\left(\|\hat{\nu}-\mathbb{E}\phi(S_{h-1},A_{h-1})\|_{1}\right)\geq\frac{\zeta}{2}$
(48)

In the last step, we have used the lipschitzness bound for $\mathcal{T}_{h}$
given in Lemma E.2. We will show that the conditions given in (4) are
satisfied for $\psi(S_{h},A_{h})$ with parameters $\zeta/2$ instead of
$\zeta$.

$\|\psi(S_{h},A_{h})\|_{2}\leq 1$ almost surely follows from the definition of
$\psi$. Now, $\mathbb{E}f(S_{h},A_{h},x)\geq\frac{\zeta}{2}$ for every
$x\in\mathcal{S}^{d-1}$ implies $\mathbb{E}|\langle
x,\psi(S_{h},A_{h})\rangle|\geq\frac{\zeta}{2\sqrt{d}}$. Using the definition
of $f(S_{h},A_{h},x)$ (see Section 6) and the fact that
$\mathbb{E}f(S_{h},A_{h},x)\geq\frac{\zeta}{2}$ as established above, we
conclude that for every $x\in\mathcal{S}^{d-1}$, we also have:

$\displaystyle d\xi\mathbb{E}\langle x,\psi(S_{h},A_{h})\rangle^{2}$
$\displaystyle\leq\sqrt{d}\mathbb{E}|\langle
x,\psi(S_{h},A_{h})\rangle|-\frac{\zeta}{2}$
$\displaystyle\leq\sqrt{d}\mathbb{E}|\langle
x,\psi(S_{h},A_{h})\rangle|\leq\sqrt{d}\sqrt{\mathbb{E}\langle
x,\psi(S_{h},A_{h})\rangle^{2}}$ (49)

In the second step, we have used Jensen’s inequality. From this, we conclude
$\mathbb{E}\langle x,\psi(S_{h},A_{h})\rangle^{2}\leq\frac{1}{d\xi^{2}}$ for
every $x\in\mathcal{S}^{d-1}$ and thence
$\mathbb{E}\psi(S_{h},A_{h})\psi(S_{h},A_{h})^{\intercal}\preceq\frac{1}{d\xi^{2}}$.
∎

##  Appendix F Proof of Theorem 7.1

Let the unknown row set in the iteration $t$ in the matrix estimation
procedure of Section 7.1 be denoted by $\bar{I}_{t-1}$. For the analysis, we
will use the convention that $\bar{I}_{t}=\emptyset$ if the procedure
terminates before the $t$-th iteration. Suppose $K_{t}$ is such that for every
$t\leq\log N$, we have: $K_{t}|\bar{I}_{t-1}|\geq
C\frac{r|\bar{I}_{t-1}|+dr}{\zeta^{2}\xi^{2}}\log\tfrac{d}{\zeta\xi}+C\frac{\log\left(\frac{\log
N}{\delta}\right)}{\zeta^{2}\xi^{2}}$. We will then show that the event
$\\{|\bar{I}_{t}|\leq\frac{1}{10}|\bar{I}_{t-1}|\forall t\leq\log
N\\}\cap\\{\hat{\Theta}_{i}=\Theta^{*}_{i},\forall i\in\bar{I}_{\log
N}^{\complement}\\}$ has probability at-least $1-\delta$. To show this, it is
sufficient to consider the step $t=1$ with $\bar{I}_{0}=[N]$, $K_{1}=K$,
$\Psi^{(1)}=\Psi$ and show that with probability $1-\frac{\delta}{\log N}$,
$\bar{I}_{1}\leq\frac{9N}{10}$ and $\hat{\Theta}_{i}=\Theta^{*}_{i}$ for every
$i\in\bar{I}_{1}^{\complement}$. The result then follows from a union bound.
We will therefore establish the following structural lemma and prove the
Theorem 7.1. The rest of the section is then dedicated to proving Lemma F.1.

######  Lemma F.1.

Suppose the distribution of $\psi_{ik}$ satisfies (4). Let
$K\geq\frac{C(r+\tfrac{dr}{N})}{\zeta^{2}\xi^{2}}\log\frac{d}{\zeta\xi}$. Let
$\mathcal{Y}(\Psi)$ denote the set of all matrices $\Delta$ with rank at most
$2r$ such that $L(\Delta,\Psi)=0$. Let
$I_{\mathcal{Z}}(\Delta)=\\{i\in[N]:\Delta_{i}\neq 0\\}$. With probability
$1-\exp(-c\zeta^{2}\xi^{2}NK)$ we must have:

$\mathcal{Y}(\Psi)\cap\bigr{\\{}\Delta:|I_{\mathcal{Z}}(\Delta)|>\tfrac{N}{10}\bigr{\\}}=\emptyset$

###### Proof of Theorem 7.1.

Let $\bar{\Theta}$ be the rank $\leq r$ matrix found satisfying
$L(\bar{\Theta}-\Theta^{*},\Psi^{(t)})=0$. By Lemma F.1, we have that
$|I_{\mathcal{Z}}(\bar{\Theta}-\Theta^{*})|\leq\frac{N}{10}$ with probability
at-least $1-\frac{\delta}{\log N}$ (by setting $K=K_{1}$ as in the statement
of Theorem 7.1). By Lemma F.2, the probability that there exists an $i\in[N]$
such that $\bar{\Theta}_{i}\neq\Theta^{*}_{i}$, and
$\sum_{k=1}^{K}\left|\langle\bar{\Theta}_{i},\tilde{\psi}_{ik}\rangle-\theta^{*}_{ik}\right|^{2}=0$
is at most $\left|I\right|\cdot\exp(-c\zeta^{2}\xi^{2}K)\leq\delta\cdot
N^{-c}$ for some large constant $c$. From this we conclude that
$\bar{\Theta}_{i}=\Theta^{*}_{i}$ for every $i\in\bar{I}_{1}^{\complement}$. ∎

######  Lemma F.2.

Fix any $\hat{\Theta}$. Suppose the distribution of
$(\psi_{ik})_{i\in[N],k\in[K]}$ satisfies (4). Then, there exists a small
enough constant $c$ such that:

$\mathbb{P}\left(\exists i\mbox{ s.t.
}\sum_{k=1}^{K}\left|\langle\hat{\Theta}_{i},\psi_{ik}\rangle-\theta^{*}_{ik}\right|^{2}<\tfrac{K\zeta^{4}\xi^{2}\left\|\hat{\Theta}_{i}-\Theta^{*}_{i}\right\|^{2}}{32{d}}\right)\leq\left|I\right|\cdot\exp(-c\zeta^{2}\xi^{2}K).$

###### Proof.

Consider the Paley-Zygmund inequality, which states that for any positive
random variable $Z$,

$\mathbb{P}\left(Z\geq\tfrac{\mathbb{E}Z}{2}\right)\geq\frac{(\mathbb{E}Z)^{2}}{4\mathbb{E}Z^{2}}\,.$

Suppose $i\in I$ and denote $\Gamma_{i}:=\hat{\Theta}_{i}-\Theta^{*}_{i}$. By
the properties of $\psi_{ik}$, we have that
$\mathbb{E}|\langle\Gamma_{i},\psi_{ik}\rangle|\geq\frac{\zeta\left\|\Gamma_{i}\right\|}{\sqrt{d}}$
and
$\mathbb{E}|\langle\Gamma_{i},\psi_{ik}\rangle|^{2}\leq\frac{\left\|\Gamma_{i}\right\|^{2}}{\xi^{2}d}$.

Applying the Paley-Zygmund inequality to the random variable
$|\langle\Gamma_{i},\psi_{ik}\rangle|$, we conclude the result in (53):

$\mathbb{P}\left(|\langle\psi_{ik},\Gamma_{i}\rangle|\geq\frac{\zeta\left\|\Gamma_{i}\right\|}{2\sqrt{d}}\right)\geq\frac{\zeta^{2}\xi^{2}}{4}$
(50)

Let $p_{0}:=\frac{\zeta^{2}\xi^{2}}{4}$. Let
$N(\Gamma_{i},\Psi_{i}):=\sum_{k=1}^{K}\mathbbm{1}\left(|\langle\psi_{ik},\Gamma_{i}\rangle|>\frac{\zeta}{2\sqrt{d}}\right)$.
Clearly,
$\sum_{k=1}^{K}\left|\langle\hat{\Theta}_{i},\psi_{ik}\rangle-\theta^{*}_{ik}\right|^{2}\geq\frac{\zeta^{2}\left\|\Gamma_{i}\right\|^{2}}{4{d}}\cdot
N(\Gamma_{i},\Psi_{i})$. Therefore, we have:

$\displaystyle\mathbb{P}\left(\sum_{k=1}^{K}\left|\langle\hat{\Theta}_{i},\psi_{ik}\rangle-\theta^{*}_{ik}\right|^{2}<\tfrac{K\zeta^{4}\xi^{2}\left\|\Gamma_{i}\right\|^{2}}{32{d}}\right)$
$\displaystyle\leq\mathbb{P}\left(N(\Gamma_{i},\Psi_{i})<\tfrac{K\zeta^{2}\xi^{2}}{8}\right)$
$\displaystyle\leq\mathbb{P}\left(\mathsf{Bin}(K,p_{0})\leq\tfrac{Kp_{0}}{2}\right)$
$\displaystyle\leq\exp(-cp_{0}K)$ (51)

Here $\mathsf{Bin}(K,p_{0})$ denotes the binomial random variable. In the
second step we have used the fact that $N(\Gamma_{i},\Psi_{i})$ is a sum of
$K$ independent Bernoulli random variables with probability of being $1$ for
each of them being at-least $p_{0}=\frac{\zeta^{2}\xi^{2}}{4}$. In the last
step, we have used Sanov’s theorem for large deviations. In the last step we
have used Bernstein’s inequality for concentration of sums of Bernoulli random
variables (see (Boucheron et al., 2013)). The statement of the result then
follows from a union bound argument over $i\in I$. ∎

###  F.1 Proof of Lemma F.1

Suppose $I\neq\emptyset$, $I\subseteq[N]$ be any fixed subset. By
$\mathcal{M}(N,d,I,2r)$, we denote the set of all $N\times d$ matrices
$\Delta$ with rank at-most $2r$ such that $\|\Delta_{i}\|>0$ for all $i\in I$.
By $\mathcal{B}(N,d,I,2r)$ we denote the set of all $N\times d$ matrices with
rank at-most $2r$ such that $\|\Gamma_{i}\|=1$ whenever $i\in I$.

######  Lemma F.3.

Suppose $\inf_{\Gamma\in\mathcal{B}(N,d,I,2r)}L(\Gamma,\Psi)>0$. Then,
$L(\Delta,\Psi)>0$ for every $\Delta\in\mathcal{M}(N,d,I,2r)$

###### Proof.

For every $\Delta\in\mathcal{M}(N,d,I,2r)$, we construct $\Gamma$ such that:

$\Gamma_{i}=\begin{cases}\frac{\Delta_{i}}{\|\Delta_{i}\|}&\text{ whenever
}i\in I\\\ 0&\text{ otherwise}\end{cases}$ (52)

Now, by hypothesis, $L(\Gamma,\Psi)>0$. This implies, there exists an $i\in I$
and $k\in K$ such that $|\langle\psi_{ik},\Gamma_{i}\rangle|>0$. This implies
$|\langle\psi_{ik},\Delta_{i}\rangle|>0$ and thence we conclude that
$L(\Delta,\Psi)>0$. ∎

######  Lemma F.4.

Suppose $\Gamma$ is such that $\|\Gamma_{i}\|=1$ for every $i\in I$. Suppose
the distribution of $(\psi_{ik})_{i\in[N],k\in[K]}$ satisfy (4). Then, there
exists a small enough constant $c$ such that:

$\mathbb{P}\left(L(\Gamma,\Psi)<\tfrac{|I|\zeta^{4}\xi^{2}}{32Nd}\right)\leq|I|^{2}K^{2}\exp(-c\zeta^{2}\xi^{2}|I|K)$

###### Proof.

Consider the Paley-Zygmund inequality, which states that for any positive
random variable $Z$,

$\mathbb{P}\left(Z\geq\tfrac{\mathbb{E}Z}{2}\right)\geq\frac{(\mathbb{E}Z)^{2}}{4\mathbb{E}Z^{2}}\,.$

Suppose $i\in I$. By the properties of $\psi_{ik}$, we have that
$\mathbb{E}|\langle\Gamma_{i},\psi_{ik}\rangle|\geq\frac{\zeta}{\sqrt{d}}$ and
$\mathbb{E}|\langle\Gamma_{i},\psi_{ik}\rangle|^{2}\leq\frac{1}{\xi^{2}d}$

Applying the Paley-Zygmund inequality to the random variable
$|\langle\Gamma_{i},\psi_{ik}\rangle|$, we conclude the result in (53):