# Representation Learning for Online and Offline RL in Low-rank MDPs

Masatoshi Uehara<EMAIL_ADDRESS>Department of Computer Science, Cornell
University  Xuezhou Zhang<EMAIL_ADDRESS>Department of Electrical and
Computer Engineering, Princeton University  Wen Sun<EMAIL_ADDRESS>Department of Computer Science, Cornell University

###### Abstract

This work studies the question of _Representation Learning in RL:_ how can we
learn a compact low-dimensional representation such that on top of the
representation we can perform RL procedures such as exploration and
exploitation, in a sample efficient manner. We focus on the low-rank Markov
Decision Processes (MDPs) where the transition dynamics correspond to a low-
rank transition matrix. Unlike prior works that assume the representation is
known (e.g., linear MDPs), here we need to learn the representation for the
low-rank MDP. We study both the online RL and offline RL settings. For the
online setting, operating with the same computational oracles used in
Flambe(Agarwal et al., 2020b)—-the state-of-art algorithm for learning
representations in low-rank MDPs, we propose an algorithm Rep-UCB—Upper
Confidence Bound driven Representation learning for RL, which _significantly
improves_ the sample complexity from
$\widetilde{O}(A^{9}d^{7}/(\epsilon^{10}(1-\gamma)^{22}))$ for Flambe to
$\widetilde{O}(A^{4}d^{4}/(\epsilon^{2}(1-\gamma)^{3}))$ with $d$ being the
rank of the transition matrix (or dimension of the ground truth
representation), $A$ being the number of actions, and $\gamma$ being the
discount factor. Notably, Rep-UCB is simpler than Flambe, as it directly
balances the interplay between representation learning, exploration, and
exploitation, while Flambe is an explore-then-commit style approach and has to
perform reward-free exploration step-by-step forward in time. For the offline
RL setting, we develop an algorithm that leverages pessimism to learn under a
partial coverage condition: our algorithm is able to compete against _any
policy_ (including non-Markovian history dependent ones) as long as it is
covered by the offline data distribution.

##  1 Introduction

When applying Reinforcement Learning (RL) to large-scale problems where data
is complex and high-dimensional, learning effective transformations of the
data, i.e., representation learning, can often significantly improve the
sample and computation efficiency of the RL procedure. Indeed, several
empirical works have shown that leveraging representation learning techniques
developed in supervised or unsupervised learning settings can accelerate the
search for good decision-making strategies (Silver et al., 2018; Stooke et
al., 2021; Srinivas et al., 2020; Yang and Nachum, 2021). However,
representation learning in RL is far more subtle than it is for non-sequential
and non-interactive learning tasks (e.g., supervised learning). Prior works
have shown that even if one is given the magic representation that exactly
linearizes the optimal policy (Du et al., 2019b) or the optimal value
functions (Wang et al., 2020; Weisz et al., 2021), RL is still challenging
(i.e., one may still need exponentially many samples to learn). This indicates
that an effective representation that permits efficient RL needs to encode
more information about the underlying Markov Decision Processes (MDPs).
Despite the recent empirical success of representation learning in RL , its
statistical guarantee and theoretical properties remain under-investigated.

In this work, we study the representation learning question under the low-rank
MDP assumption. Concretely, a low-rank MDP assumes that the MDP transition
matrix admits a low-rank factorization, i.e., there exists two _unknown_
mappings $\mu(s^{\prime}),\phi(s,a)$, such that
$P(s^{\prime}|s,a)=\mu(s^{\prime})^{\top}\phi(s,a)$ for all $s,a,s^{\prime}$,
where $P(s^{\prime}|s,a)$ is the probability of transiting to the next state
$s^{\prime}$ under the current state and action $(s,a)$. The representation
$\phi$ in a low-rank MDP not only linearizes the optimal state-action value
function of the MDP (Jin et al., 2020a), but also linearizes the transition
operator. A low-rankness assumption on large stochastic matrices is a common
assumption and has enabled successful development of algorithms for real world
applications such as movie recommendation systems (Koren et al., 2009). We
note that a low-rank MDP strictly generalizes the linear MDP model (Yang and
Wang, 2020; Jin et al., 2020a) which assumes $\phi$ is known a priori. The
unknown representation $\phi$ makes learning in low-rank MDPs much more
challenging than that in linear MDPs since one can no longer directly use
linear function approximations. On the other hand, the fact that linear MDPs
can be solved statistical and computational efficiently if $\phi$ is known a
priori implies that if one could learn the representation of the low-rank MDP,
one could then efficiently learn the optimal policy.

Indeed, prior works have shown that learning in low-rank MDPs is statistically
feasible (Jiang et al., 2017; Sun et al., 2019; Du et al., 2021) via
leveraging rich function approximators. However, these algorithms are version
space algorithms and are not computationally efficient. Recent work Flambe
proposes an oracle-efficient algorithm111The oracle generally refers to
empirical risk minimization oracle employing non-bilevel optimization. We seek
to design an algorithm that runs in polynomial time, with each oracle call
counting as $O(1)$. This algorithmic framework has lead to many successful
algorithms in the literature of contextual bandit (Agarwal et al., 2014; Dudík
et al., 2017; Foster and Rakhlin, 2020). that learns in low-rank MDPs with a
polynomial sample complexity, where the computation oracle is Maximum
Likelihood Estimation (MLE). In this work, we follow the same setup from
Flambe (Agarwal et al., 2020b), and propose a new algorithm — _Upper
Confidence Bound driven Representation Learning, Exploration and Exploitation
( Rep-UCB)_, which can learn a near optimal policy for a low-rank MDP with a
polynomial sample complexity and is oracle-efficient. Comparing to Flambe, our
algorithm _significantly improves_ the sample complexity from
$O(d^{7}A^{9}/(\epsilon^{10}(1-\gamma)^{22})$ for Flambe to
$O(d^{4}A^{4}/(\epsilon^{2}(1-\gamma)^{3})$, where $d$ is the rank of the
transition matrix (or dimension of the true representation), $A$ is the number
of actions, $\epsilon$ is the suboptimality gap and $\gamma\in[0,1)$ is the
discount factor in the MDP. Our algorithm is also arguably much simpler than
Flambe: Flambe is an explore-then-commit algorithm, has to explore in a layer-
by-layer forward way, and does not permit data sharing across different time
steps. In contrast, Rep-UCB carefully trades exploration versus exploitation
by combining the reward signal and exploration bonus (constructed using the
latest learned representation), and enables data sharing across all time
steps. Our sample complexity nearly matches the ones from those
computationally inefficient algorithms (Jiang et al., 2017; Sun et al., 2019;
Du et al., 2021). We summarize the comparison with the prior works that study
representation learning in Table 1.

Methods | Setting | Sample Complexity | Computation
---|---|---|---
OLIVE (Jiang et al., 2017) | Low Bellman rank | $\frac{d^{2}A}{\epsilon^{2}(1-\gamma)^{4}}$ | Inefficient
Witness rank (Sun et al., 2019) | Low Witness rank | $\frac{d^{2}A}{\epsilon^{2}(1-\gamma)^{4}}$ | Inefficient
BLin-UCB (Du et al., 2021) | Bilinear Class | $\frac{d^{2}A}{\epsilon^{2}(1-\gamma)^{7}}$ | Inefficient
Moffle (Modi et al., 2021) | Low-nonnegative-rank MDP | $\frac{d^{6}A^{13}}{\epsilon^{2}\eta^{5}(1-\gamma)^{5}}$ | Oracle efficient
Flambe (Agarwal et al., 2020b) | Low-rank MDP | $\frac{d^{7}A^{9}}{\epsilon^{10}(1-\gamma)^{22}}$ | Oracle efficient
Rep-UCB (Ours) | Low-rank MDP | $\frac{d^{4}A^{4}}{\epsilon^{2}(1-\gamma)^{3}}$ | Oracle efficient

Table 1: Comparison among different provable representation learning
algorithms in online RL. Algorithms such as OLIVE, Witness rank, and BLin-UCB
work for settings which are more general than low-rank MDPs and have tight
sample complexity. However, these algorithms are version space algorithms and
thus are not computationally efficient. Moffle is an oracle-efficient
algorithm, but the assumptions under which Moffle operates essentially imply
that the MDP’s transition has low non-negative matrix rank (nnr) (see the
detailed discussion in the related work section). Note that a nnr is at least
as large as and could be exponentially larger than the rank (Agarwal et al.,
2020b). Finally, Flambe operates exactly under the same setting of ours, i.e.,
the same function approximation and computation oracle. Our algorithm
_significantly improves_ the sample complexity from Flambe in _all
parameters_. Note the horizon dependence is not exactly comparable as these
prior works originally considered the finite horizon setting with
nonstationary transition, and we convert their results to the discounted
setting by simply replacing the finite horizon H by $\Theta(1/(1-\gamma))$.
Also these prior works count total number of trajectories _while we count
total number of samples_.

In addition to the online exploration setting, we also show that our new
techniques can be directly used for designing offline RL algorithms for low-
rank MDPs under partial coverage. More specifically, we propose an algorithm
Rep-LCB— _Lower Confidence Bound driven Reprepresentation Learning for offline
RL_ , that given an offline dataset, can learn to compete against any policy
(including history-dependent policies) as long as it is covered by the offline
data where the coverage is measured using the relative condition number
(Agarwal et al., 2021) associated with the ground truth representation. Thus,
our offline RL result generalizes prior offline RL works on linear MDPs (Jin
et al., 2020b; Zhang et al., 2021b) which assume representation is known a
priori and use linear function approximation. Computation-wise, our approach
uses one call to the MLE computation oracle, and hence is oracle-efficient.
Rep-LCB is _the first oracle efficient offline algorithm for low-rank MDP
enjoying the aforementioned statistical guarantee._ See Section 2 for a more
detailed comparison with the existing literature on representation learning in
offline RL.

#### Our contributions.

We develop new representation learning RL algorithms that enable sample
efficient learning in low-rank MDPs under both online and offline settings:

  1. 1.

In the online episodic learning setting, our new algorithm Rep-UCB integrates
representation learning, exploration, and exploitation together, and
significantly improves the sample complexity of the prior state-of-art
algorithm Flambe;

  2. 2.

In the offline learning setting, we propose a natural concentrability
coefficient (i.e., relative condition number under the true representation)
that captures the partial coverage condition in low-rank MDP, and our
algorithm Rep-LCB learns to compete against any policy (including history-
dependent ones) under such a partial coverage condition.

The rest of the paper is organized as follows. In Section 2, we discuss
related works on representation learning in RL. In Section 3 we introduce the
definition of low-rank MDPs and our function approximation setting. We present
our online RL algorithm Rep-UCB in Section 4, followed by our offline RL
algorithm Rep-LCB in Section 5.

##  2 Related Work

#### Online Setting

We list the comparison as follows, which is summarized in Table 1.

Flambe (Agarwal et al., 2020b) was a state-of-the-art oracle-efficient
algorithm for low-rank MDPs. In all parameters, the statistical complexity is
much worse than Rep-UCB . Our algorithm and Flambe operate under the same
computation oracle. Flambe does not balance exploration and exploitation, and
uses explore-then-committee style techniques (i.e., constructions of absorbing
MDPs (Brafman and Tennenholtz, 2002)) which results in its worse sample
complexity.

With a more complex oracle, Moffle (Modi et al., 2021) is a model-free
algorithm for low-rank MDPs, with two additional assumptions: (1) the
transition has low _non-negative rank (nnr)_ , and (2) reachability in latent
states. The first assumption significantly restricts the scope of low-rank
MDPs as there are matrices whose nnr is exponentially larger than the rank
(Agarwal et al., 2020b). The sample complexity of Moffle can scale
$O(d^{6}|\mathcal{A}|^{13}/(\epsilon^{2}\eta^{5}(1-\gamma)^{5}))$, where
$\eta$ is the reachability probability, and $1/\eta$ could be as large as
$nnr^{1/2}$ (Proposition 4 in Agarwal et al. (2020b)), which essentially means
that Moffle has a polynomial dependence on the nnr. Thus, Moffle needs the nnr
of the transition matrix to be small.

OLIVE (Jiang et al., 2017), Witness rank (Sun et al., 2019) and Bilinear-UCB
(Du et al., 2021), when specialized to low-rank MDPs, have slightly tighter
dependence on $d$ (e.g., $O(d^{2}/\epsilon^{2})$). But these algorithms are
computationally inefficient as they are version space algorithms. Dann et al.
(2021) shows that with a policy class, solving a low-rank MDP can take
$\Omega(2^{d})$ samples. In this work, similar to Witness rank (Sun et al.,
2019) and Flambe, we use function approximators to model the transition. Thus
our positive result is not in contradiction to the result from Dann et al.
(2021).

VALOR (Dann et al., 2018), PCID (Du et al., 2019a), HOMER (Misra et al.,
2020), RegRL (Foster et al., 2020), and the approach from Feng et al. (2020)
are algorithms for block MDPs which is a more restricted setting than low-rank
MDPs. These works require additional assumptions such as deterministic
transitions (Dann et al., 2018), reachability (Misra et al., 2020; Du et al.,
2019a), strong Bellman completion (Foster et al., 2020), and strong
unsupervised learning oracles (Feng et al., 2020).

Finally, we mention works that tackle representation learning in quite
different settings. Zhang et al. (2021a) consider how to choose the best
representation among _correct representations_ inspired by Papini et al.
(2021) (i.e., the MDP is a linear MDP under any representation in the function
class $\Phi$). In contrast, we only assume the MDP is linear under _some_
$\phi^{\star}\in\Phi$. Hao et al. (2021b) consider feature selection in sparse
linear MDPs with a given exploratory distribution.

#### Offline Setting

We discuss related works in offline RL.

Uehara and Sun (2021) obtained similar results for offline RL on low-rank
MDPs. Though the sample complexity in their algorithm is slightly tighter, our
algorithm is oracle-efficient, while the CPPO algorithm from Uehara and Sun
(2021) is a version space algorithm.

Xie et al. (2021) propose a (general) pessimistic model-free algorithm in the
offline setting. We can also apply their algorithm to low-rank MDPs and show
some finite-sample guarantee. However, it is unclear whether the final bounds
in their results can be characterized by the relative condition number only
using the true representation, and whether they can compete with history-
dependent policies. Thus, our result is still considered superior on low-rank
MDPs. The detail is given in Section D.

In addition to the above two works, the pessimistic approach in offline RL has
been extensively investigated. Empirically, it can work on simulation control
tasks (Kidambi et al., 2020; Yu et al., 2020; Kumar et al., 2020; Liu et al.,
2020; Chang et al., 2021). On the theoretical side, pessimism allows us to
obtain the PAC guarantee on various models when a comparator policy is covered
by offline data in some forms (Jin et al., 2020b; Rashidinejad et al., 2021;
Yin et al., 2021; Zanette et al., 2021b; Zhang et al., 2021b; Chang et al.,
2021). However, these algorithms and their analysis rely on a _known
representation_. We emphasize the differences in Section 5.

Finally, we mention works that tackle representation learning from different
viewpoints and settings. Lu et al. (2021) consider multitask representation
learning under a generative model assumption. Hao et al. (2021a) study the
feature selection problem in sparse linear MDPs and Ni et al. (2021) study
dimensionality reduction in a given kernel space, under the assumption that
the offline data admits some form of full coverage condition.

##  3 Preliminaries

We consider a Markov Decision Process
$\mathcal{M}=\langle{\mathcal{S}},\mathcal{A},P,r,\gamma,d_{0}\rangle$
specified by a state space ${\mathcal{S}}$, a discrete action space
$\mathcal{A}$, a transition model
$P:{\mathcal{S}}\times\mathcal{A}\to\Delta({\mathcal{S}})$, a reward function
$r:{\mathcal{S}}\times\mathcal{A}\to\mathbb{R}$, a discount factor
$\gamma\in[0,1)$, and an initial distribution $d_{0}\in\Delta({\mathcal{S}})$.
To simplify the presentation, we assume $r(s,a)$ and $d_{0}$ are known (e.g.,
when $d_{0}$ is a probability mass only on $s_{0}$, agent always starts from a
fixed initial state $s_{0}$)222Extension to the unknown case is
straightforward. Recall the major challenging of RL is due to the unknown
transition model.. Following prior work (Jiang et al., 2017; Sun et al.,
2019), we assume _trajectory reward is normalized_ , i.e., for any trajectory
$\\{s_{h},a_{h}\\}_{h=0}^{\infty}$, we have
$\sum_{h=0}^{\infty}\gamma^{h}r(s_{h},a_{h})\in[0,1]$. Since the ground truth
$P^{\star}$ is unknown, we need to learn it by interacting with environments
in an online manner or utilizing offline data at hand. We remark that the
extension of our all results to the finite horizon nonstationary case is
straightforward.

Figure 1: A latent state model captured by low-rank MDP. Here
$\phi^{\star}(s,a)$ is a distribution over a discrete latent state space
$\mathcal{Z}$. Note that this is still a Markovian model as there is no
transition among latent states.

We use the following notation. Given a policy
$\pi:{\mathcal{S}}\to\Delta(\mathcal{A})$, which maps from state to
distribution over actions and some model $P$, we define the value function
$V^{\pi}_{P}(s)=\mathbb{E}\left[\sum_{h=0}^{\infty}\gamma^{h}r(s_{h},a_{h})|s_{0}=s,P,\pi\right]$
to represent the expected total discounted reward of $\pi$ under $P$ starting
at $s$. Similarly, we define the state-action Q function
$Q^{\pi}_{P}(s,a):=r(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim
P(\cdot|s,a)}V^{\pi}_{P}(s^{\prime})$. The expected total discounted reward of
a policy $\pi$ is denoted as $V^{\pi}_{P,r}:=\mathbb{E}_{s_{0}\sim
d_{0}}V^{\pi}_{P}(s_{0})$. We define the discounted state-action occupancy
distribution
$d^{\pi}_{P}(s,a)=(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}d^{\pi}_{P,t}(s,a)$,
where $d^{\pi}_{P,t}(s,a)$ is the probability of $\pi$ visiting $(s,a)$ at
time step $t$ under $\pi$ and $P$. We slightly abuse the notation, and denote
$d^{\pi}_{P}(s)$ as the state visitation, which is equal to
$\sum_{a\in\mathcal{A}}d^{\pi}_{P}(s,a)$. When $P$ is the ground truth
transition model $P^{\star}$, we drop the subscript and simply use $d^{\pi}$
to denote its state-action distribution. Unless otherwise noted, $\Pi$ denotes
the class of all polices $\\{S\to\Delta(\mathcal{A})\\}$. We denote total
variation distance of $P_{1}$ and $P_{2}$ by $\|P_{1}-P_{2}\|_{1}$. Finally,
given a vector $a$, we define
$\|a\|_{2}=\sqrt{a^{\top}a},\|a\|_{B}=\sqrt{a^{\top}Ba}$. $c_{0},c_{1},\cdots$
are universal constants.

We study low-rank MDPs defined as follows (Jiang et al., 2017; Agarwal et al.,
2020b). The conditions on the upper bounds of the norm of
$\phi^{\star},\mu^{\star}$ are just for normalization.

######  Definition 1 (Low-rank MDP).

Given function classes
$\Phi\subset{\mathcal{S}}\times\mathcal{A}\to\mathbb{R}^{d}$ and
$\Psi\subset{\mathcal{S}}\to\mathbb{R}^{d}$, a transition model
$P^{\star}:{\mathcal{S}}\times\mathcal{A}\to\Delta(\mathcal{A})$ admits a low
rank decomposition with dimension $d\in\mathbb{N}$ if there exists two
embedding functions $\phi^{\star}\in\Phi,\mu^{\star}\in\Psi$ such that

$\displaystyle\forall
s,s^{\prime}\in{\mathcal{S}},a\in\mathcal{A}:P^{\star}(s^{\prime}\mid
s,a)=\mu^{\star}(s^{\prime})^{\top}{\phi^{\star}(s,a)}$

where $\|\phi^{*}(s,a)\|_{2}\leq 1$ for all $(s,a)$ and for any function
$g:{\mathcal{S}}\to[0,1]$,
$\|\int\mu^{\star}(s)g(s){\textnormal{d}}(s)\|_{2}\leq\sqrt{d}$. An MDP is a
low rank MDP if $P^{\star}$ admits such a low rank decomposition.

Low-rank MDPs capture the latent variable model (Agarwal et al., 2020b) shown
in Figure 1 where $\phi^{\star}(s,a)$ is a distribution over a discrete latent
state space $\mathcal{Z}$. The block-MDP model (Du et al., 2019a) is a special
instance of the latent variable model with $\phi^{\star}(s,a)$ being a one-hot
encoding vector. Note the linear MDPs (Yang and Wang, 2020; Jin et al., 2020a)
assume $\phi^{\star}$ is known.

Next, we explain two settings: the online learning setting and the offline
learning setting. Then, we present our function approximation setup and
computational oracles.

#### Episodic Online learning

In online learning, our overall goal is to learn a stationary policy
$\hat{\pi}$ so that it maximizes $V^{\hat{\pi}}_{P^{\star},r}$, where
$P^{\star}$ is the ground truth transition. We assume that we operate under
the episodic learning setting where we can only reset to states sampled from
the initial distribution $d_{0}$ (e.g., to emphasize the challenge from
exploration, we can consider the special case where we can only reset to a
fixed $s_{0}$). In the episodic setting, given a policy $\pi$, sampling a
state $s$ from the state visitation $d_{P}^{\pi}$ is done by the following
_roll-in_ procedure: starting at $s_{0}\sim d_{0}$, at every time step $t$, we
terminate and return $s_{t}$ with probability $1-\gamma$, and otherwise we
execute $a_{t}\sim\pi(s_{t})$ and move to $t+1$, i.e., $s_{t+1}\sim
P(\cdot|s_{t},a_{t})$. Such a sampling procedure is widely used in the
episodic learning setting for discounted infinite horizon MDPs (e.g., (Kakade
and Langford, 2002; Agarwal et al., 2021, 2020a)).

#### Offline learning

In the offline RL, we are given a static dataset in the form of quadruples:

$\displaystyle\mathcal{D}=\\{s^{(i)},a^{(i)},r^{(i)},s^{\prime(i)}\\}_{i=1}^{n}\sim\rho(s,a)\delta(r=r(s,a))P^{\star}(s^{\prime}\mid
s,a).$

For simplicity, we assume $\rho=d^{\pi_{b}}_{P^{\star}}$, where
$\pi_{b}\in{\mathcal{S}}\to\Delta(\mathcal{A})$ is a fixed behavior policy. We
denote
$\mathbb{E}_{\mathcal{D}}[f(s,a,s^{\prime})]=1/n\sum_{(s,a,s^{\prime})\in\mathcal{D}}f(s,a,s^{\prime})$.
To succeed in offline RL, we in general need some coverage property of $\rho$.
One common assumption is that $\rho$ globally covers every possible policies’
state-action distribution, namely we have
$\max_{\pi,s,a}d^{\pi}_{P^{\star}}(s,a)/\rho(s,a)<\infty$ (Antos et al.,
2008). In this work, we relax such a global coverage assumption and work under
the _partial_ coverage condition where $\rho$ may not cover distributions of
all possible policies. Instead of competing against the optimal policy under
the global coverage, we aim to compete against any policies covered by the
offline data. In section 5, we will precisely define the partial coverage
condition using the concept of the relative condition number (Agarwal et al.,
2021).

#### Function approximation setup and computational oracles

Since $\mu^{\star}$ and $\phi^{\star}$ are unknown, we use function classes to
capture them. Our function approximation and computational oracles are
_exactly the same as the ones used in Flambe_. For completeness, we state the
function approximation and computational oracles below.

######  Assumption 2.

We have a model class $\mathcal{M}=\\{(\mu,\phi):\mu\in\Psi,\phi\in\Phi\\}$,
where $\mu^{\star}\in\Psi$, $\phi^{\star}\in\Phi$.

Following the norm bounds on $\mu^{\star},\phi^{\star}$ we similarly assume
that the same norm bounds hold for our function approximator, i.e., for any
$\mu\in\Phi,\phi\in\Phi$, $\|\phi(s,a)\|_{2}\leq 1$, $\forall(s,a)$ and
$\|\int\mu(s)g(s){\textnormal{d}}(s)\|_{2}\leq\sqrt{d},\forall
g:{\mathcal{S}}\to[0,1]$, and
$\int\mu^{\top}(s^{\prime})\phi(s,a){\textnormal{d}}(s^{\prime})=1,\,\forall(s,a)$.

As for computational oracles, we use a supervised learning style MLE oracle.

######  Definition 3 (Maximum Likelihood Oracle (MLE)).

Consider the model class $\mathcal{M}$ and a dataset $\mathcal{D}$ in the form
of $(s,a,s^{\prime})$, the MLE oracle returns the maixmum likelihood
estimator:

$\hat{P}:=(\hat{\mu},\hat{\phi})=\operatorname*{arg\,max}_{(\mu,\phi)\in\mathcal{M}}\mathbb{E}_{\mathcal{D}}\ln\left(\mu(s^{\prime})^{\top}\phi(s,a)\right).$

We also invoke a planning procedure for _known linear MDPs with potentially
nonlinear rewards_ , which can be done in polynomial time (after all, we know
that online learning in linear MDPs can be done statistically and
computationally efficient). Given a reward $r$ and a model $P:=(\mu,\phi)$
with $P(s^{\prime}|s,a)=\mu(s^{\prime})^{\top}\phi(s,a)$ (i.e., a _known
linear transition_ with a _known_ feature $\phi$), we can compute the optimal
policy $\arg\max_{\pi}V^{\pi}_{P}$ by standard least square value iteration
which uses linear regression. A planning procedure for a known linear MDP is
also used in Flambe, see Section 5.1 in Agarwal et al. (2020b) for details of
how to implement this procedure for a known linear MDP with polynomial
computation complexity.

##  4 Representation Learning in Online Setting

We consider the online episodic learning setting where the agent can only
reset based on the initial state distribution $d_{0}$. Thus, to find a near-
optimal policy for a low-rank MDP efficiently, we need to carefully interleave
representation learning, exploration, and exploitation.

###  4.1 Algorithm

We present our proposed algorithm in the online setting. This is described in
Algorithm 1. Every iteration, Algorithm 1 rollouts its current policy $\pi$ to
collect a triple $(s,a,s^{\prime})$ where $s\sim d_{P^{\star}}^{\pi}$, $a\sim
U(\mathcal{A}),s^{\prime}\sim P^{\star}(\cdot|s,a)$ where $U(\mathcal{A})$ is
a uniform distribution over actions. Note that to sample $s\sim
d^{\pi}_{P^{\star}}$, we start at $s_{0}\sim d_{0}$, at every time step $t$,
we terminate and return $s_{t}$ with probability $1-\gamma$, and otherwise we
execute $a_{t}\sim\pi(s_{t})$ and move to $t+1$, i.e., $s_{t+1}\sim
P^{\star}(\cdot|s_{t},a_{t})$. After collecting new data and concatenating it
with the existing data, we perform representation learning, i,e, learning a
factorization and a representation by MLE (line 6), set the bonus based on the
learned feature (Eq. 1), and update the policy via planning inside the learned
model with the bonus-enhanced reward (Line 9). Note the learned transition
$\hat{P}$ from MLE is linear with the learned feature $\hat{\phi}$, and
planning in a known linear MDP is known as computationally efficient (Jin et
al., 2020a) (see the explanation after Definition 3 as well). Our computation
MLE oracle is the same as the one used in Flambe (Agarwal et al., 2020b).

Algorithm 1 UCB-driven representation learning, exploration, and exploitation
(Rep-UCB)

1: Input: Regularizer $\lambda_{n}$, parameter $\alpha_{n}$, Models
$\mathcal{M}=\\{(\mu,\phi):\mu\in\Psi,\phi\in\Phi\\}$, Iteration $N$

2: Initialize $\pi_{0}(\cdot\mid s)$ to be uniform; set
$\mathcal{D}_{0}=\emptyset$

3: for episode $n=1,\cdots,N$ do

4: Collect a transition $(s,a,s^{\prime})$ where $s\sim
d_{P^{\star}}^{\pi_{n-1}}$, $a\sim U(\mathcal{A})$, $s^{\prime}\sim
P^{\star}(\cdot|s,a)$

5: Add $(s,a,s^{\prime})$ to the dataset:
$\mathcal{D}_{n}=\mathcal{D}_{n-1}+\\{s,a,s^{\prime}\\}$

6: Learn representation via MLE:
$\hat{P}_{n}:=(\hat{\mu}_{n},\hat{\phi}_{n})=\operatorname*{arg\,max}_{(\mu,\phi)\in\mathcal{M}}\mathbb{E}_{\mathcal{D}_{n}}\left[\ln\mu^{\top}(s^{\prime})\phi(s,a)\right]$

7: Update empirical covariance matrix
$\hat{\Sigma}_{n}=\sum_{s,a\in\mathcal{D}_{n}}\hat{\phi}_{n}(s,a)\hat{\phi}_{n}(s,a)^{\top}+\lambda_{n}I$

8: Set the exploration bonus:
$\hat{b}_{n}(s,a):=\min\left(\alpha_{n}\sqrt{\hat{\phi}_{n}(s,a)^{\top}\hat{\Sigma}^{-1}_{n}\hat{\phi}_{n}(s,a)},2\right)\vspace{-10pt}$
(1)

9: Update policy
$\pi_{n}=\operatorname*{arg\,max}_{\pi}V^{\pi}_{\hat{P}_{n},r+\hat{b}_{n}}$

10: end for

11: Return  $\pi_{1},\cdots,\pi_{N}$

###  4.2 Analysis

We show that our algorithm has the following PAC bound.

######  Theorem 4 (PAC Bound for Rep-UCB).

Fix $\delta\in(0,1),\epsilon\in(0,1)$. Let $\hat{\pi}$ be a random output of
$\pi_{1},\cdots,\pi_{N}$ and
$\pi^{\star}=\operatorname*{arg\,max}_{\pi}V^{\pi}_{P^{\star},r}$. By setting
parameters as follows:

$\displaystyle\alpha_{n}=O\left(\sqrt{\left(|\mathcal{A}|+d^{2}\right)\gamma\ln(|\mathcal{M}|n/\delta)}\right),\quad\lambda_{n}=O\left(d\ln(|\mathcal{M}|n/\delta)\right),$

with probability at least $1-\delta$, we have

$V^{\pi^{\star}}_{P^{\star},r}-V^{\hat{\pi}}_{P^{\star},r}\leq\epsilon,$

where the number of collected samples is at most

$\displaystyle N\log(N/\delta),\quad N\coloneqq
O\left(\frac{d^{4}|\mathcal{A}|^{4}\ln(|\mathcal{M}|/\delta)}{(1-\gamma)^{3}\epsilon^{2}}\cdot\ln^{2}\left(1+\frac{d^{4}|\mathcal{A}|^{4}\ln(|\mathcal{M}|/\delta)}{(1-\gamma)^{2}\epsilon^{2}}\right)\right).$

The theorem shows that Rep-UCB learns in low-rank MDPs in a statistically
efficient and oracle-efficient manner. To the best of our knowledge, this
algorithm has the best sample complexity among all oracle efficient algorithms
for low-rank MDPs.

#### Highlight of the analysis

Below we highlight our key lemmas and proof techniques.

First, why learning in a low-rank MDP is harder than learning in models with
linear structures? Unlike standard linear models such as linear MDPs (Yang and
Wang, 2020; Jin et al., 2020a), KNRs (Kakade et al., 2020; Abbasi-Yadkori and
Szepesvári, 2011; Mania et al., 2020; Song and Sun, 2021), and GP / kernel
models (Chowdhury and Gopalan, 2019; Curi et al., 2020), we cannot get
uncertainty quantification on the model in a point-wise manner. When models
are linear, one can get the following style of point-wise uncertainty
quantification for the learned model $\hat{P}$: $\forall
s,a:{\ell}(\hat{P}(\cdot|s,a),P^{\star}(\cdot|s,a))\leq\sigma(s,a)$ where
$\sigma(s,a)$ is the uncertainty measure, and ${\ell}$ is some distance metric
(e.g., $\ell_{1}$ norm). With proper scaling, the uncertainty measure
$\sigma(s,a)$ is then used for the bonus. For example, in linear MDPs, given a
dataset $\mathcal{D}=\\{s,a,s^{\prime}\\}$, we can learn a model
$\hat{P}(s^{\prime}|s,a):=\hat{\mu}(s^{\prime})^{\top}\phi^{\star}(s,a)$ , and
get point-wise uncertainty quantification:

$\displaystyle\textstyle\forall(s,a),\lvert\int
f(s^{\prime})\hat{\mu}^{\top}(s^{\prime})\phi^{\star}(s,a){\textnormal{d}}(s^{\prime})-\int
f(s^{\prime}){\mu^{\star}}^{\top}(s^{\prime})\phi^{\star}(s,a){\textnormal{d}}(s^{\prime})\rvert\leq
c\|\phi^{\star}(s,a)\|_{\Sigma^{-1}_{\phi^{\star}}}$ (2)

for some family of functions $f:{\mathcal{S}}\to\mathbb{R}$ with
$\Sigma_{\phi^{\star}}=\sum_{s,a\in\mathcal{D}}\phi^{\star}(s,a)\phi^{\star}(s,a)+\lambda
I$ (Lykouris et al., 2021; Neu and Pike-Burke, 2020). To set the scaling $c$
properly, since $\phi^{\star}$ is known a priori, the linear regression
analysis applies here, and one can apply Cauchy-Schwarz inequality to the LHS
of (2) to pull out $\phi^{\star}$ and get an upper bound in the form of

$\displaystyle\textstyle\underbrace{\|\phi^{\star}(s,a)\|_{\Sigma^{-1}_{\phi^{\star}}}}_{(a)}\underbrace{\|\int
f(s)\\{\hat{\mu}(s)-\mu^{\star}(s)\\}{\textnormal{d}}(s)\|_{\Sigma_{\phi^{\star}}}}_{(b)}$

where $c$ is set to be the linear regression training error measured in the
term $(b)$ above.

However, when we jointly learn $\mu$ and $\phi$, since nonlinear function
approximation is used333There are nonlinear models using Eluder dimension
(Russo and Van Roy, 2014) as the complexity measure. However, to the best of
our knowledge, the only known examples that admit low Eluder dimension are
(generalized) linear models., we cannot get point-wise uncertainty
quantification via linear regression-based analysis. We stress that our bonus
is not designed to capture the uncertainty quantification on the model error
between $\hat{P}(\cdot|s,a)=\hat{\mu}^{\top}\hat{\phi}(s,a)$ and
$P^{\star}(\cdot|s,a)={\mu^{\star}}^{\top}\phi^{\star}(s,a)$ in a point-wise
way, which is not tractable as $\hat{P}$ and $P^{\star}$ does not even share
the same representation. Instead, the bonus is carefully designed so that it
only provides _near-optimism at the initial state distribution_. This is
formalized as follows.

######  Lemma 5 (Almost Optimism at the Initial State Distribution).

Set the parameters as in Theorem 4. With probability $1-\delta$,

$\displaystyle\forall
n\in[1,\cdots,N],\forall\pi\in\Pi,V^{\pi}_{\hat{P}_{n},r+\hat{b}_{n}}-V^{\pi}_{P^{\star},r}\geq-
c_{1}\sqrt{\frac{|\mathcal{A}|\ln(|\mathcal{M}|n/\delta)(1-\gamma)}{n}}.$

We remark that the idea of optimism with respect to the initial state
distribution has been used in prior works (Jiang et al., 2017; Sun et al.,
2019; Du et al., 2021; Zanette et al., 2020). However, these algorithms are
not computationally efficient (i.e., they use version space instead of reward
bonus), and their version-space based analysis is different from ours.

#### Proof sketch for Lemma 5

We start by using the simulation lemma (Lemma 20) _inside the learned model_
which is important since our bonus $\hat{b}_{n}$ uses $\hat{\phi}_{n}$
associated with the learned model $\hat{P}_{n}$:

$\displaystyle V^{\pi}_{\hat{P}_{n},r+\hat{b}_{n}}-V^{\pi}_{P^{\star},r}$
$\displaystyle=\mathbb{E}_{s,a\sim
d^{\pi}_{\hat{P}_{n}}}\left[\hat{b}_{n}(s,a)+\mathbb{E}_{s^{\prime}\sim\hat{P}_{n}(\cdot|s,a)}V^{\pi}_{P^{\star}}(s^{\prime})-\mathbb{E}_{s^{\prime}\sim{P^{\star}}(\cdot|s,a)}V^{\pi}_{P^{\star}}(s^{\prime})\right]$
$\displaystyle\geq\mathbb{E}_{s,a\sim
d^{\pi}_{\hat{P}_{n}}}\left[\hat{b}_{n}(s,a)-\|\hat{P}_{n}(\cdot|s,a)-{P^{\star}}(\cdot|s,a)\|_{1}\right],$

from where we show that $\mathbb{E}_{s,a\sim
d^{\pi}_{\hat{P}_{n}}}\|\hat{P}(\cdot|s,a)-P^{\star}(\cdot|s,a)\|_{1}$ _as a
whole_ nearly lower bounds $\mathbb{E}_{s,a\sim
d^{\pi}_{\hat{P}_{n}}}[\hat{b}_{n}(s,a)]$. The detailed procedure is
illustrated in Lemma 8 in Appendix A.

Second, our bonus is using representation $\hat{\phi}_{n}$ that is being
updated every episode, and our empirical covariance matrix $\hat{\Sigma}_{n}$
is also updated whenever we update $\hat{\phi}_{n}$, which means that standard
elliptical potential based analysis (i.e., analysis used in linear
bandits/MDPs with known features) cannot work here as our feature changes
every episode. Instead, in our analysis, we have to keep tracking a potential
function that is defined using the unknown ground truth representation
$\phi^{\star}$, i.e., the elliptical potential
$\|\phi^{\star}(s,a)\|^{2}_{\Sigma^{-1}_{\rho_{n},\phi^{\star}}}$, where

$\Sigma_{\rho_{n},\phi^{\star}}=n\mathbb{E}_{(s,a)\sim\rho_{n}}\phi^{\star}(s,a)\phi^{\star}(s,a)^{\top}+\lambda_{n}I,$

and $\rho_{n}$ is the marginal state-action distribution of $\mathcal{D}_{n}$
at an iteration $n$, i.e.,
$\rho_{n}(s,a)=\sum_{i=0}^{n-1}d^{\pi_{i}}_{P^{\star}}(s)U(a)/n$. Since this
potential function uses the fixed representation $\phi^{\star}$, we can apply
the standard elliptical potential argument to track the progress that our
algorithm makes during learning. Below we illustrate the procedure of linking
the bonus under $\hat{\phi}_{n}$ to the potential function
$\|\phi^{\star}(s,a)\|^{2}_{\Sigma^{-1}_{\rho_{n},\phi^{\star}}}$ defined with
respect to the true feature $\phi^{\star}$.

#### Linking bonus under $\hat{\phi}_{n}$ to the elliptical potential function
under $\phi^{\star}$

With near optimism, using the simulation lemma (Lemma 20), we can upper bound
the per-iteration regret as follows: c where
$f_{n}(s,a):=\|\hat{P}_{n}(\cdot\mid s,a)-P^{\star}(\cdot\mid s,a)\|_{1}$. To
connect the right-hand side of the above inequality to the elliptical
potential under the fixed feature $\phi^{\star}$, we show that for any
function $g\in{\mathcal{S}}\times\mathcal{A}\to[0,B]$ for
$B\in\mathbb{R}^{+}$,

$\displaystyle\mathbb{E}_{(s,a)\sim d^{\pi_{n}}_{P^{\star}}}[g(s,a)]$
$\displaystyle\leq\mathbb{E}_{(s,a)\sim
d^{\pi_{n}}_{P^{\star}}}\left[\|\phi^{\star}(s,a)\|_{\Sigma^{-1}_{\rho_{n},\phi^{\star}}}\right]\sqrt{n\gamma|\mathcal{A}|\mathbb{E}_{\rho_{n}}[g^{2}(s,a)]+\gamma\lambda_{n}dB^{2}}$
$\displaystyle\qquad\qquad+\sqrt{(1-\gamma)|\mathcal{A}|\mathbb{E}_{\rho_{n}}g^{2}(s,a)}.$

See Lemma 13 in Appendix A. By substituting $g$ with
$\hat{b}_{n}+f_{n}/(1-\gamma)$, the first term of the RHS of the above
inequality can be upper bounded as:

$\displaystyle 2\underbrace{\mathbb{E}_{(s,a)\sim
d^{\pi_{n}}_{P^{\star}}}\left[\|\phi^{\star}(s,a)\|_{\Sigma^{-1}_{\rho_{n},\phi^{\star}}}\right]}_{(\mathcal{G}_{1})}\underbrace{\sqrt{n|\mathcal{A}|\mathbb{E}_{\rho_{n}}\left[\frac{f^{2}_{n}(s,a)}{(1-\gamma)^{2}}+\hat{b}^{2}_{n}(s,a)\right]+\lambda_{n}d}}_{(\mathcal{G}_{2})}.$

In the term ($\mathcal{G}_{2}$), we expect
$n\mathbb{E}_{\rho_{n}}[f_{n}^{2}(s,a)]$ to be $O(1)$ as
$\mathbb{E}_{s,a\sim\rho_{n}}[f^{2}_{n}(s,a)]$ is in order of $1/n$ due to the
fact that it is the generalization bound of the MLE estimator $\hat{P}_{n}$
under the training distribution $\rho_{n}$. For
$n\mathbb{E}_{\rho_{n}}[\hat{b}^{2}_{n}(s,a)]$, we expect it to be in the
order of $d$ as the (unnormalized) data covariance matrix $\hat{\Sigma}_{n}$
in the bonus $\hat{b}_{n}$ uses training data from $\rho_{n}$, i.e., we are
measuring the expected bonus under the training distribution. In other words,
the term ($\mathcal{G}_{2}$) scales in order of $\text{poly}(d)$. For the term
$(\mathcal{G}_{1})$, since it contains the potential function based on
$\phi^{\star}$, the sum of the term ($\mathcal{G}_{1}$) over all episodes can
be controlled by the standard elliptical potential argument (see Lemma 18 and
Lemma 19).

##  5 Representation Learning in Offline Setting

In this section, we study representation learning in the offline setting. We
consider the setting where the offline data does not have a full global
coverage. We present our algorithm _Lower Confidence Bound driven
Representation Learning in offline RL_ (Rep-LCB) in Algorithm 2. Our proposed
algorithm consists of three parts. The first part is MLE which learns a model
$\hat{P}$ and a representation $\hat{\phi}$. The second part is the
construction of a penalty term $\hat{b}$. Using the learned representation
$\hat{\phi}$, we use a standard bonus in linear bandits as the penalty term as
if $\hat{\phi}$ were the true feature. The third part is planning with the
learned model $\hat{P}$ and reward $r-\hat{b}$. Again this is just planning
inside a known linear MDP.

Algorithm 2 LCB-driven Representation Learning in offline RL (Rep-LCB)

1: Input: Regularizer $\lambda$, Parameter $\alpha$, Model classes
$\mathcal{M}=\\{\mu^{\top}\phi:\mu\in\Psi,\phi\in\Phi\\}$, Dataset
$\mathcal{D}$.

2: Learn a model $\hat{P}$ by MLE:
$\hat{P}=\hat{\mu}^{\top}\hat{\phi}=\operatorname*{arg\,max}_{P\in\mathcal{M}}\mathbb{E}_{\mathcal{D}}[\ln
P(s^{\prime}\mid s,a)]$.

3: Set the empirical covariance matrix
$\hat{\Sigma}=\sum_{(s,a)\in\mathcal{D}}\hat{\phi}(s,a)\hat{\phi}^{\top}(s,a)+\lambda
I$.

4: Set the reward penalty:
$\hat{b}(s,a)=\min\left(\alpha\sqrt{\hat{\phi}(s,a)^{\top}\hat{\Sigma}^{-1}\hat{\phi}(s,a)},2\right).$

5: Solve
$\hat{\pi}=\operatorname*{arg\,max}_{\pi}V^{\pi}_{\hat{P},r-\hat{b}}$.

###  5.1 Analysis

We present the PAC guarantee of Rep-LCB. Before proceeding, we define a
relative condition number as a mean to measure the deviation between a
comparator policy $\pi$ and the offline data:

$\displaystyle
C^{\star}_{\pi}=\sup_{x\in\mathbb{R}^{d}}\frac{x^{\top}\mathbb{E}_{d^{\pi}_{P^{\star}}}[\phi^{\star}(s,a){\phi^{\star}}^{\top}(s,a)]x}{x^{\top}\mathbb{E}_{\rho}[\phi^{\star}(s,a){\phi^{\star}}^{\top}(s,a)]x}.$

In the special case where the MDP is just a tabular MDP (i.e., $\phi^{\star}$
is a one-hot encoding vector), this is reduced to a density ratio
$C^{\star}_{\infty}=\max_{s,a}d^{\pi}_{P^{\star}}(s,a)/\rho(s,a)$. The
relative condition number $C^{\star}_{\pi}$ is always no larger than the
density ratio and could be much smaller for MDPs with large state spaces. Note
that we quantify the relative condition number using the unknown true
representation $\phi^{\star}$. With the above setup, now we are ready to state
the main theorem for Rep-LCB.

######  Theorem 6 (PAC Bound for Rep-LCB).

Let $\omega=\max_{a,s}(1/\pi_{b}(a\mid s))$. Denote $\hat{\pi}$ as the output
of Rep-LCB. There exists a set of parameters such that with probability at
least $1-\delta$, for _any_ policy $\pi$ (including history-dependent non-
Markovian policies),

$\displaystyle V^{\pi}_{P^{\star},r}-V^{\hat{\pi}}_{P^{\star},r}\leq
c\sqrt{\frac{d^{4}\omega^{2}C^{\star}_{\pi}\log(|\mathcal{M}|/\delta)}{(1-\gamma)^{2}n}}.$

See Theorem 15 in Appendix B for the detailed parameters. We explain several
implications. First of all, this theorem shows that we can uniformly compete
with any policy including history-dependent non-Markovian policies 444Given
$\pi=\\{\pi_{i}\\}_{i=0}^{\infty}$ where $\pi_{i}$ depends on
$s_{0},a_{0},\dots s_{i}$, $V^{\pi}_{P^{\star},r}$ and
$d^{\pi}_{P^{\star}}(s,a)$ are still well-defined.  satisfying the partial
coverage $C^{\star}_{\pi}<\infty$. Particularly, if the optimal policy
$\pi^{\star}$ is covered by the offline data, i.e.,
$C^{\star}_{\pi^{\star}}<\infty$, then our algorithm is able to compete
against it 555We also require $\omega<\infty$, which is a mild assumption
since it does not involve $P^{\star}$. Indeed, it is much weaker than the
global coverage type assumption $1/\rho(s,a)<\infty,\forall(s,a)$. . Note that
assuming offline data covers $\pi^{\star}$ is still a weaker assumption than
the global coverage such as
$\sup_{\pi}\sup_{(s,a)}d^{\pi}_{P^{\star}}(s,a)/\rho(s,a)$ in prior offline RL
works (Antos et al., 2008; Chen and Jiang, 2019). Second, our coverage
condition is measured by a relative condition number defined using the unknown
ground truth representation $\phi^{\star}$ but not depending on other
features. Prior works that use relative condition numbers as measures of
coverage are restricted to the settings where the ground truth representation
$\phi^{\star}$ is known (Jin et al., 2020b; Zhang et al., 2021b; Chang et al.,
2021; Zanette et al., 2021b).

To sum up, our algorithm is the first oracle efficient algorithm which does
not need to know $\phi^{\star}$, and requires partial coverage only in terms
of $\phi^{\star}$. Note while Uehara and Sun (2021) has a similar guarantee on
low-rank MDPs, their algorithm is not oracle-efficient as it is a version
space algorithm.

#### Highlight of the analysis

Prior offline RL works that use reward penalties (Rashidinejad et al., 2021;
Zhang et al., 2021b; Chang et al., 2021) all assume the representation
$\phi^{\star}$ is known a priori, which allows them to use linear regression
analysis to derive point-wise model uncertainty quantification which in turn
serves as a penalty. In contrast, our function class is nonlinear. As in the
online setting, what we can show is near pessimism in the initial state
distribution $d_{0}$.

######  Lemma 7 (Almost Pessimism at the Initial State Distribution).

There exists a set of parameters that with probability $1-\delta$, we have

$\forall\pi\in\Pi:V^{\pi}_{\hat{P},r-b}-V^{\pi}_{P^{\star},r}\leq
c_{1}\sqrt{\frac{\omega\log(|\mathcal{M}|/\delta)(1-\gamma)}{n}}.$

We leave the detailed proof to Lemma 14 in the Appendix B, which is similar to
the proof of Lemma 5. Though the pessimism in the initial state distribution
was recently also used in Xie et al. (2021); Zanette et al. (2021b), the
derivation is totally different.

Similar to the online setting, to obtain a result only depending on the
relative condition number using the ground truth representation $\phi^{\star}$
but not the learned feature $\hat{\phi}$, we again need to translate the
penalty defined with $\hat{\phi}$ to the potential function
$\|\phi^{\star}(s,a)\|_{\Sigma^{-1}_{\rho,\phi^{\star}}}$. The same techniques
that we used for the online setting can be leveraged here to achieve the above
goal.

##  6 Conclusion

We study online/offline RL on low-rank MDPs, where the ground truth feature is
not known a priori. For online RL, our new algorithm Rep-UCB significantly
improves the sample complexity of the piror state-of-the-art algorithm Flambe
in all parameters while using the same computational oracles. Rep-UCB has the
best sample complexity among existing oracle efficient algorithms for low-rank
MDPs by a margin. Comparing to prior representation learning works on low-rank
MDPs and block MDPs that rely on a forward step-by-step reward-free
exploration framework, our algorithm interleaves representation learning,
exploration, and exploitation together, and learns a single stationary policy.
For offline RL, our new algorithm Rep-LCB is the first oracle efficient
algorithm for low-rank MDPs that has a PAC guarantee under a partial coverage
condition measured by the relative condition number defined with the true
feature representation.

## Acknowledgement

The authors would like to thank Alekh Agarwal, Praneeth Netrapalli and Ming
Yin for valuable feedback.

## References

  * Abbasi-Yadkori and Szepesvári (2011) Yasin Abbasi-Yadkori and Csaba Szepesvári.  Regret bounds for the adaptive control of linear quadratic systems.  In _Proceedings of the 24th Annual Conference on Learning Theory_ , pages 1–26. JMLR Workshop and Conference Proceedings, 2011.
  * Agarwal et al. (2014) Alekh Agarwal, Daniel Hsu, Satyen Kale, John Langford, Lihong Li, and Robert Schapire.  Taming the monster: A fast and simple algorithm for contextual bandits.  In _International Conference on Machine Learning_ , pages 1638–1646. PMLR, 2014.
  * Agarwal et al. (2020a) Alekh Agarwal, Mikael Henaff, Sham Kakade, and Wen Sun.  Pc-pg: Policy cover directed exploration for provable policy gradient learning.  _NeurIPS_ , 2020a.
  * Agarwal et al. (2020b) Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun.  Flambe: Structural complexity and representation learning of low rank mdps.  In _Advances in Neural Information Processing Systems_ , volume 33, pages 20095–20107, 2020b.
  * Agarwal et al. (2021) Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan.  On the theory of policy gradient methods: Optimality, approximation, and distribution shift.  _Journal of Machine Learning Research_ , 22(98):1–76, 2021.
  * Antos et al. (2008) András Antos, Csaba Szepesvári, and Rémi Munos.  Learning near-optimal policies with bellman-residual minimization based fitted policy iteration and a single sample path.  _Machine Learning_ , 71:89–129, 2008.
  * Brafman and Tennenholtz (2002) Ronen I Brafman and Moshe Tennenholtz.  R-max-a general polynomial time algorithm for near-optimal reinforcement learning.  _Journal of Machine Learning Research_ , 3(Oct):213–231, 2002.
  * Chang et al. (2021) Jonathan D Chang, Masatoshi Uehara, Dhruv Sreenivas, Rahul Kidambi, and Wen Sun.  Mitigating covariate shift in imitation learning via offline data without great coverage.  2021\.
  * Chen and Jiang (2019) Jinglin Chen and Nan Jiang.  Information-theoretic considerations in batch reinforcement learning.  In _Proceedings of the 36th International Conference on Machine Learning_ , volume 97, pages 1042–1051, 2019.
  * Chowdhury and Gopalan (2019) Sayak Ray Chowdhury and Aditya Gopalan.  Online learning in kernelized markov decision processes.  In _The 22nd International Conference on Artificial Intelligence and Statistics_ , pages 3197–3205. PMLR, 2019.
  * Curi et al. (2020) Sebastian Curi, Felix Berkenkamp, and Andreas Krause.  Efficient model-based reinforcement learning through optimistic policy search and planning.  _arXiv preprint arXiv:2006.08684_ , 2020.
  * Dann et al. (2018) Christoph Dann, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire.  On oracle-efficient pac rl with rich observations.  _arXiv preprint arXiv:1803.00606_ , 2018.
  * Dann et al. (2021) Christoph Dann, Yishay Mansour, Mehryar Mohri, Ayush Sekhari, and Karthik Sridharan.  Agnostic reinforcement learning with low-rank mdps and rich observations.  _arXiv preprint arXiv:2106.11519_ , 2021.
  * Du et al. (2019a) Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford.  Provably efficient rl with rich observations via latent state decoding.  In _International Conference on Machine Learning_ , pages 1665–1674. PMLR, 2019a.
  * Du et al. (2019b) Simon S Du, Sham M Kakade, Ruosong Wang, and Lin F Yang.  Is a good representation sufficient for sample efficient reinforcement learning?  _arXiv preprint arXiv:1910.03016_ , 2019b.
  * Du et al. (2021) Simon S Du, Sham M Kakade, Jason D Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, and Ruosong Wang.  Bilinear classes: A structural framework for provable generalization in rl.  _ICML_ , 2021.
  * Dudík et al. (2017) Miroslav Dudík, Nika Haghtalab, Haipeng Luo, Robert E Schapire, Vasilis Syrgkanis, and Jennifer Wortman Vaughan.  Oracle-efficient online learning and auction design.  In _2017 ieee 58th annual symposium on foundations of computer science (focs)_ , pages 528–539. IEEE, 2017.
  * Feng et al. (2020) Fei Feng, Ruosong Wang, Wotao Yin, Simon S Du, and Lin F Yang.  Provably efficient exploration for reinforcement learning using unsupervised learning.  _arXiv preprint arXiv:2003.06898_ , 2020.
  * Foster and Rakhlin (2020) Dylan Foster and Alexander Rakhlin.  Beyond ucb: Optimal and efficient contextual bandits with regression oracles.  In _International Conference on Machine Learning_ , pages 3199–3210. PMLR, 2020.
  * Foster et al. (2020) Dylan J Foster, Alexander Rakhlin, David Simchi-Levi, and Yunzong Xu.  Instance-dependent complexity of contextual bandits and reinforcement learning: A disagreement-based perspective.  _arXiv preprint arXiv:2010.03104_ , 2020.
  * Hao et al. (2021a) Botao Hao, Yaqi Duan, Tor Lattimore, Csaba Szepesvári, and Mengdi Wang.  Sparse feature selection makes batch reinforcement learning more sample efficient.  In _International Conference on Machine Learning_ , pages 4063–4073. PMLR, 2021a.
  * Hao et al. (2021b) Botao Hao, Tor Lattimore, Csaba Szepesvári, and Mengdi Wang.  Online sparse reinforcement learning.  In _International Conference on Artificial Intelligence and Statistics_ , pages 316–324. PMLR, 2021b.
  * Jiang et al. (2017) Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire.  Contextual decision processes with low bellman rank are pac-learnable.  In _International Conference on Machine Learning_ , 2017.
  * Jin et al. (2020a) Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I Jordan.  Provably efficient reinforcement learning with linear function approximation.  In _Conference on Learning Theory_ , pages 2137–2143. PMLR, 2020a.
  * Jin et al. (2020b) Ying Jin, Zhuoran Yang, and Zhaoran Wang.  Is pessimism provably efficient for offline rl?  _arXiv preprint arXiv:2012.15085_ , 2020b.
  * Kakade and Langford (2002) Sham Kakade and John Langford.  Approximately optimal approximate reinforcement learning.  In _In Proc. 19th International Conference on Machine Learning_. Citeseer, 2002.
  * Kakade et al. (2020) Sham Kakade, Akshay Krishnamurthy, Kendall Lowrey, Motoya Ohnishi, and Wen Sun.  Information theoretic regret bounds for online nonlinear control.  In _Advances in Neural Information Processing Systems_ , volume 33, pages 15312–15325, 2020.
  * Kidambi et al. (2020) Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims.  Morel: Model-based offline reinforcement learning.  In _Advances in Neural Information Processing Systems_ , volume 33, pages 21810–21823. Curran Associates, Inc., 2020.
  * Koren et al. (2009) Yehuda Koren, Robert Bell, and Chris Volinsky.  Matrix factorization techniques for recommender systems.  _Computer_ , 42(8):30–37, 2009.
  * Kumar et al. (2020) Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.  Conservative q-learning for offline reinforcement learning.  _arXiv preprint arXiv:2006.04779_ , 2020.
  * Liu et al. (2020) Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill.  Provably good batch off-policy reinforcement learning without great exploration.  In _Advances in Neural Information Processing Systems_ , volume 33, pages 1264–1274, 2020.
  * Lu et al. (2021) Rui Lu, Gao Huang, and Simon S Du.  On the power of multitask representation learning in linear mdp.  _arXiv preprint arXiv:2106.08053_ , 2021.
  * Lykouris et al. (2021) Thodoris Lykouris, Max Simchowitz, Alex Slivkins, and Wen Sun.  Corruption-robust exploration in episodic reinforcement learning.  In _Conference on Learning Theory_ , pages 3242–3245. PMLR, 2021\.
  * Mania et al. (2020) Horia Mania, Michael I Jordan, and Benjamin Recht.  Active learning for nonlinear system identification with guarantees.  _arXiv preprint arXiv:2006.10277_ , 2020.
  * Misra et al. (2020) Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford.  Kinematic state abstraction and provably efficient rich-observation reinforcement learning.  In _International conference on machine learning_ , pages 6961–6971. PMLR, 2020.
  * Modi et al. (2021) Aditya Modi, Jinglin Chen, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal.  Model-free representation learning and exploration in low-rank mdps.  _arXiv preprint arXiv:2102.07035_ , 2021.
  * Neu and Pike-Burke (2020) Gergely Neu and Ciara Pike-Burke.  A unifying view of optimism in episodic reinforcement learning.  _arXiv preprint arXiv:2007.01891_ , 2020.
  * Ni et al. (2021) Chengzhuo Ni, Anru Zhang, Yaqi Duan, and Mengdi Wang.  Learning good state and action representations via tensor decomposition.  _arXiv preprint arXiv:2105.01136_ , 2021.
  * Papini et al. (2021) Matteo Papini, Andrea Tirinzoni, Marcello Restelli, Alessandro Lazaric, and Matteo Pirotta.  Leveraging good representations in linear contextual bandits.  _arXiv preprint arXiv:2104.03781_ , 2021.
  * Rashidinejad et al. (2021) Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell.  Bridging offline reinforcement learning and imitation learning: A tale of pessimism.  _arXiv preprint arXiv:2103.12021_ , 2021.
  * Russo and Van Roy (2014) Daniel Russo and Benjamin Van Roy.  Learning to optimize via posterior sampling.  _Mathematics of Operations Research_ , 39(4):1221–1243, 2014.
  * Silver et al. (2018) David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al.  A general reinforcement learning algorithm that masters chess, shogi, and go through self-play.  _Science_ , 362(6419):1140–1144, 2018.
  * Song and Sun (2021) Yuda Song and Wen Sun.  Pc-mlp: Model-based reinforcement learning with policy cover guided exploration.  In _International Conference on Machine Learning_ , pages 9801–9811. PMLR, 2021.
  * Srinivas et al. (2020) Aravind Srinivas, Michael Laskin, and Pieter Abbeel.  Curl: Contrastive unsupervised representations for reinforcement learning.  _arXiv preprint arXiv:2004.04136_ , 2020.
  * Stooke et al. (2021) Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin.  Decoupling representation learning from reinforcement learning.  In _International Conference on Machine Learning_ , pages 9870–9879. PMLR, 2021.
  * Sun et al. (2019) Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford.  Model-based rl in contextual decision processes: Pac bounds and exponential improvements over model-free approaches.  In _Conference on learning theory_ , pages 2898–2933. PMLR, 2019\.
  * Uehara and Sun (2021) Masatoshi Uehara and Wen Sun.  Pessimistic model-based offline rl: Pac bounds and posterior sampling under partial coverage.  _arXiv preprint arXiv:2107.06226_ , 2021.
  * Wang et al. (2020) Ruosong Wang, Dean P Foster, and Sham M Kakade.  What are the statistical limits of offline rl with linear function approximation?  _arXiv preprint arXiv:2010.11895_ , 2020.
  * Weisz et al. (2021) Gellért Weisz, Philip Amortila, and Csaba Szepesvári.  Exponential lower bounds for planning in mdps with linearly-realizable optimal action-value functions.  In _Algorithmic Learning Theory_ , pages 1237–1264. PMLR, 2021.
  * Xie et al. (2021) Tengyang Xie, Ching-An Cheng, Nan Jiang, Paul Mineiro, and Alekh Agarwal.  Bellman-consistent pessimism for offline reinforcement learning.  _arXiv preprint arXiv:2106.06926_ , 2021.
  * Yang and Wang (2020) Lin Yang and Mengdi Wang.  Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound.  In _Proceedings of the 37th International Conference on Machine Learning_ , pages 10746–10756, 2020.
  * Yang and Nachum (2021) Mengjiao Yang and Ofir Nachum.  Representation matters: Offline pretraining for sequential decision making.  _arXiv preprint arXiv:2102.05815_ , 2021.
  * Yin et al. (2021) Ming Yin, Yu Bai, and Yu-Xiang Wang.  Near-optimal offline reinforcement learning via double variance reduction.  _arXiv preprint arXiv:2102.01748_ , 2021.
  * Yu et al. (2020) Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma.  Mopo: Model-based offline policy optimization.  In _Advances in Neural Information Processing Systems_ , volume 33, pages 14129–14142, 2020.
  * Zanette et al. (2020) Andrea Zanette, Alessandro Lazaric, Mykel Kochenderfer, and Emma Brunskill.  Learning near optimal policies with low inherent bellman error.  In _International Conference on Machine Learning_ , pages 10978–10989. PMLR, 2020.
  * Zanette et al. (2021a) Andrea Zanette, Ching-An Cheng, and Alekh Agarwal.  Cautiously optimistic policy optimization and exploration with linear function approximation.  _arXiv preprint arXiv:2103.12923_ , 2021a.
  * Zanette et al. (2021b) Andrea Zanette, Martin J Wainwright, and Emma Brunskill.  Provable benefits of actor-critic methods for offline reinforcement learning.  _arXiv preprint arXiv:2108.08812_ , 2021b.
  * Zhang et al. (2021a) Weitong Zhang, Jiafan He, Dongruo Zhou, Amy Zhang, and Quanquan Gu.  Provably efficient representation learning in low-rank markov decision processes.  _arXiv preprint arXiv:2106.11935_ , 2021a.
  * Zhang et al. (2021b) Xuezhou Zhang, Yiding Chen, Jerry Zhu, and Wen Sun.  Corruption-robust offline reinforcement learning.  _arXiv preprint arXiv:2106.06630_ , 2021b.

##  Appendix A Proof of the theoretical property of Rep-UCB

Hereafter, we assume $c_{0},c_{1},\cdots,$ are some universal constants, and
the notation

$f(1/(1-\gamma),|\mathcal{A}|,\ln(1/\delta),\ln(|\mathcal{M}|),d,n)\lesssim
g(1/(1-\gamma),|\mathcal{A}|,\ln(1/\delta),\ln(|\mathcal{M}|),d,n)$

means there exists some constant $c_{1}>0$, such that

$f(1/(1-\gamma),|\mathcal{A}|,\ln(1/\delta),\ln(|\mathcal{M}|),d,n)\leq
c_{1}g(1/(1-\gamma),|\mathcal{A}|,\ln(1/\delta),\ln(|\mathcal{M}|),d,n)$

for any $0\leq\gamma<1,|\mathcal{A}|,\ln(1/\delta),\ln(|\mathcal{M}|),d,n$.

We first prove a lemma to show the optimism. Hereafter, we use the following
notation to denote the data sampling distribution

$\displaystyle d^{\bar{\pi}}(s,a)\coloneqq d^{\pi}_{P^{\star}}(s)\cdot u(a).$

where $u(a)$ denotes the uniform distribution over actions.

First, we prove the optimism at the initial distribution. This is proved by
using a simulation lemma inside the learned model which is important since
both the bonus and the learned model use $\hat{\phi}$. In high level, we will
show that the expected bonus $\mathbb{E}_{s,a\sim
d^{\pi}_{\hat{P}_{n}}}\hat{b}_{n}(s,a)$ is in the same order of the expected
model error $\mathbb{E}_{s,a\sim
d^{\pi}_{\hat{P}_{n}}}\|\hat{P}_{n}(\cdot|s,a)-P^{\star}(\cdot|s,a)\|_{1}$.
Note that the expectation is with respect to $d^{\pi}_{\hat{P}_{n}}$.

######  Lemma 8 (Almost Optimism at the Initial Distribution).

Consider an episode $n(1\leq n\leq N)$ and set

$\displaystyle\alpha_{n}=O(\sqrt{\left(|\mathcal{A}|+d^{2}\right)\gamma\ln(|\mathcal{M}|n/\delta)}),\quad\lambda_{n}=O\left(d\ln(|\mathcal{M}|n/\delta)\right),\zeta_{n}=O\left(\frac{\ln(|\mathcal{M}|n/\delta)}{n}\right).$

With probability $1-\delta$, we have

$\displaystyle\forall
n\in[1,\cdots,N],\forall\pi\in\Pi,V^{\pi}_{\hat{P}_{n},r+\hat{b}_{n}}-V^{\pi}_{P^{\star},r}\geq-\sqrt{(1-\gamma)|\mathcal{A}|\zeta_{n}}.$

###### Proof.

We define

$\displaystyle\rho_{n}=\frac{1}{n}\sum_{j=0}^{n}d^{\bar{\pi}_{j}}(s,a),\,\Sigma_{\rho_{n},\phi}=n\mathbb{E}_{(s,a)\sim\rho_{n}}[\phi\phi^{\top}]+\lambda_{n}I.$

In this proof, letting $f_{n}(s,a)=\|\hat{P}_{n}(\cdot\mid
s,a)-P^{\star}(\cdot\mid s,a)\|_{1}$, we condition on the event

$\displaystyle\forall
n,\quad\mathbb{E}_{(s,a)\sim\rho_{n}}[f^{2}_{n}(s,a)]\leq\zeta_{n},\quad\forall
n,\forall\phi,\|\phi(s,a)\|_{\hat{\Sigma}^{-1}_{n},\phi}=\Theta(\|\phi(s,a)\|_{\Sigma^{-1}_{\rho_{n},\phi}}),$

where

$\displaystyle\hat{\Sigma}_{n,\phi}=n\mathbb{E}_{(s,a)\sim\mathcal{D}_{n}}[\phi\phi^{\top}]+\lambda_{n}I.$

From Lemma 11 and Lemma 17, this event happens with probability $1-\delta$.
Then, for any policy $\pi$, from simulation lemma 20,

$\displaystyle V^{\pi}_{\hat{P}_{n},r+\hat{b}_{n}}-V^{\pi}_{P^{\star},r}$
$\displaystyle=\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left[\hat{b}_{n}(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim\hat{P}_{n}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]-\gamma\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]\right]$
$\displaystyle\gtrsim\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left[\min\left(\alpha_{n}\|\hat{\phi}_{n}(s,a)\|_{\Sigma^{-1}_{\rho_{n},\hat{\phi}_{n}}},2\right)+\gamma\mathbb{E}_{s^{\prime}\sim\hat{P}_{n}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]-\gamma\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]\right]$ (3)

where in the last step, we replaced the empirical covariance by the population
covariance. Note the notation $\lesssim$ is up to universal constants. Here,
since $\|V^{\pi}_{P^{\star},r}\|_{\infty}\leq 1$ (since we assume trajectory-
wise total reward is normalized between $[0,1]$), we have:

$\displaystyle\left|\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left\\{\mathbb{E}_{s^{\prime}\sim\hat{P}_{n}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]-\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]\right\\}\right|\leq\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left\\{f_{n}(s,a)\right\\}.$

This is further bounded by Lemma 12:

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left\\{f_{n}(s,a)\right\\}\lesssim\sqrt{\alpha^{\prime}_{n}}\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}_{n}}}\|\hat{\phi}_{n}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\hat{\phi}_{n}}^{-1}}+\sqrt{|\mathcal{A}|\zeta_{n}(1-\gamma)}.$
(4)

where

$\displaystyle\alpha^{\prime}_{n}=\gamma\\{n|\mathcal{A}|\zeta_{n}+\lambda_{n}d+n\zeta_{n}\\}\lesssim\gamma\left(|\mathcal{A}|+d^{2}\right)\ln(|\mathcal{M}|n/\delta).$

Note we here use $f_{n}(s,a)\leq 2$ and
$\mathbb{E}_{\rho_{n}}[f_{n}(s,a)^{2}]\leq\zeta_{n}$.

Combining all things together,

$\displaystyle\left|\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left\\{\mathbb{E}_{s^{\prime}\sim\hat{P}_{n}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]-\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]\right\\}\right|\leq
2\mathbb{E}_{(s,a)\sim d^{\pi}_{\hat{P}_{n}}}\left\\{f_{n}(s,a)\right\\}$
$\displaystyle\lesssim\sqrt{\alpha^{\prime}_{n}}\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}_{n}}}\|\hat{\phi}_{n}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\hat{\phi}_{n}}^{-1}}+\sqrt{(1-\gamma)|\mathcal{A}|\zeta_{n}}$
$\displaystyle\leq\alpha_{n}\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}_{n}}}\|\hat{\phi}_{n}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\hat{\phi}_{n}}^{-1}}+\sqrt{(1-\gamma)|\mathcal{A}|\zeta_{n}},\quad\text{where
}\alpha_{n}:=\sqrt{\alpha^{\prime}_{n}}.$ (5)

Going back to (3), we have

$\displaystyle V^{\pi}_{\hat{P}_{n},r+\hat{b}_{n}}-V^{\pi}_{P^{\star},r}$
$\displaystyle\gtrsim\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left[\min\left(\alpha_{n}\|\hat{\phi}_{n}(s,a)\|_{\Sigma_{\rho_{n},\hat{\phi}_{n}}^{-1}},2\right)+\gamma\mathbb{E}_{s^{\prime}\sim\hat{P}_{n}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]-\gamma\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]\right]$
$\displaystyle\geq\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left[\min\left(\alpha_{n}\|\hat{\phi}_{n}(s,a)\|_{\Sigma_{\rho_{n},\hat{\phi}_{n}}^{-1}},2\right)-\min\left(\alpha_{n}\|\hat{\phi}_{n}(s,a)\|_{\Sigma_{\rho_{n},\hat{\phi}_{n}}^{-1}}+\sqrt{(1-\gamma)|\mathcal{A}|\zeta_{n}},2\right)\right]$
$\displaystyle\geq-\sqrt{(1-\gamma)|\mathcal{A}|\zeta_{n}}.$

From the second line to the third line, we again use
$\|V^{\pi}_{P^{\star},r}\|_{\infty}\leq O(1)$ and (4). This concludes the
proof. ∎

Next, we obtain the upper bound of
$\sum_{n=0}^{N}V^{\pi^{\star}}_{P^{\star},r}-V^{\pi_{n}}_{P^{\star},r}$.
Recall $\pi^{\star}$ is the optimal policy. Though this form is the same as a
standard regret form, since we are not exactly deploying $\pi_{n}$ in episode
$n$ (recall that we play a uniform action at the end of the episode), we
cannot get the regret guarantee. However, it suffices for the PAC guarantee.

######  Lemma 9 (Regret).

With probability $1-\delta$, we have

$\displaystyle\sum_{n=1}^{N}V^{\pi^{\star}}_{P^{\star},r}-V^{\pi_{n}}_{P^{\star},r}$
$\displaystyle\lesssim\sqrt{N\ln\left(1+\frac{N}{d^{2}\ln(|\mathcal{M}|/\delta)}\right)\ln(N|\mathcal{M}|/\delta)}\frac{|\mathcal{A}|^{2}d^{2}}{(1-\gamma)}.$

###### Proof.

Similar to Lemma 8, letting $f_{n}(s,a)=\|\hat{P}_{n}(\cdot\mid
s,a)-P^{\star}(\cdot\mid s,a)\|_{1}$, we condition on the event

$\displaystyle\forall
n,\quad\mathbb{E}_{(s,a)\sim\rho_{n}}[f^{2}_{n}(s,a)]\leq\zeta_{n},\quad\forall
n,\forall\phi,\|\phi(s,a)\|_{\hat{\Sigma}^{-1}_{n},\phi}=\Theta(\|\phi(s,a)\|_{\Sigma^{-1}_{\rho_{n},\phi}}).$
(6)

From Lemma 11 and Lemma 17, this event happens with probability $1-\delta$.

For any fixed episode $n$ and any policy $\pi$, we have

$\displaystyle V^{\pi^{\star}}_{P^{\star},r}-V^{\pi_{n}}_{P^{\star},r}$
$\displaystyle\leq
V^{\pi^{\star}}_{\hat{P}_{n},r+\hat{b}_{n}}-V^{\pi_{n}}_{P^{\star},r}+\sqrt{|\mathcal{A}|\zeta_{n}(1-\gamma)}$
(Lemma 8) $\displaystyle\leq
V^{\pi_{n}}_{\hat{P}_{n},r+\hat{b}_{n}}-V^{\pi_{n}}_{P^{\star},r}+\sqrt{|\mathcal{A}|\zeta_{n}(1-\gamma)}$
($\pi_{n}=\operatorname*{arg\,max}_{\pi}V^{\pi}_{\hat{P}_{n},r+\hat{b}_{n}}$)
$\displaystyle=\mathbb{E}_{(s,a)\sim
d^{\pi_{n}}_{P^{\star}}}[\hat{b}_{n}(s,a)+\gamma\mathbb{E}_{\hat{P}_{n}(s^{\prime}\mid
s,a)}[V^{\pi_{n}}_{\hat{P}_{n},r+\hat{b}_{n}}(s^{\prime})]-\gamma\mathbb{E}_{P^{\star}(s^{\prime}\mid
s,a)}[V^{\pi_{n}}_{\hat{P}_{n},r+\hat{b}_{n}}(s^{\prime})]]+\sqrt{|\mathcal{A}|\zeta_{n}(1-\gamma)}.$

We use the 2nd form of simulation Lemma 20 in the last display.

Then, noting $\|\hat{b}_{n}\|_{\infty}\leq 2$, we have
$\|V^{\pi_{n}}_{\hat{P}_{n},r+\hat{b}_{n}}\|_{\infty}\leq 2/(1-\gamma)$.
Combining this fact with the above expansion, we have

$\displaystyle V^{\pi}_{P^{\star},r}-V^{\pi_{n}}_{P^{\star},r}$
$\displaystyle\leq\mathbb{E}_{(s,a)\sim
d^{\pi_{n}}_{P^{\star}}}[\hat{b}_{n}(s,a)]+\left(\frac{2}{1-\gamma}\right)\mathbb{E}_{(s,a)\sim
d^{\pi_{n}}_{P^{\star}}}[f_{n}(s,a)]+\sqrt{|\mathcal{A}|\zeta_{n}(1-\gamma)}$
$\displaystyle\leq\underbrace{|\mathcal{A}|\mathbb{E}_{(s,a)\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\left[\hat{b}_{n}(s,a)\right]}_{\text{(a)}}+\underbrace{\left(\frac{2}{1-\gamma}\right)|\mathcal{A}|\mathbb{E}_{(s,a)\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}[f_{n}(s,a)]}_{\text{(b)}}+\sqrt{{|\mathcal{A}|\zeta_{n}}{(1-\gamma)}}$
(7)

where in the last inequality we use the notation
$d^{\bar{\pi}}_{P^{\star}}(s,a)\coloneqq d^{\pi}_{P^{\star}}(s)\cdot U(a)$ and
an important weighting step.

First, we calculate the first term (a) in Inequality 7. Following Lemma 13 and
noting the bonus $\hat{b}_{n}$ is $O(1)$, we have

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\left[\hat{b}_{n}(s,a)\right]$
$\displaystyle\lesssim\mathbb{E}_{(s,a)\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\left[\min\left(\alpha_{n}\|\hat{\phi}_{n}(s,a)\|_{\Sigma^{-1}_{\rho_{n},\hat{\phi}_{n}}},2\right)\right]$
( From (6) ) $\displaystyle\lesssim\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi^{\star}}^{-1}}\sqrt{{n\gamma|\mathcal{A}|\alpha^{2}_{n}}\mathbb{E}_{(s,a)\sim\rho_{n}}\left[\|\hat{\phi}_{n}(s,a)\|^{2}_{\Sigma^{-1}_{\rho_{n},\hat{\phi}_{n}}}\right]+d\gamma\lambda_{n}}$
$\displaystyle+\sqrt{|\mathcal{A}|\alpha^{2}_{n}\mathbb{E}_{(s,a)\sim\rho_{n}}\left[\|\hat{\phi}_{n}(s,a)\|^{2}_{\Sigma^{-1}_{\rho_{n},\hat{\phi}_{n}}}\right](1-\gamma)}.$

Note that we use the fact that $B=2$ when applying Lemma 13. In addition, we
have

$\displaystyle
n\mathbb{E}_{(s,a)\sim\rho_{n}}\left[\|\hat{\phi}_{n}(s,a)\|^{2}_{\Sigma^{-1}_{\rho_{n},\hat{\phi}_{n}}}\right]=n\operatorname{Tr}(\mathbb{E}_{\rho_{n}}[\hat{\phi}_{n}\hat{\phi}^{\top}_{n}]\\{n\mathbb{E}_{\rho_{n}}[\hat{\phi}_{n}\hat{\phi}^{\top}_{n}]+\lambda_{n}I\\}^{-1})\leq
d.$

Then,

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\left[\hat{b}_{n}(s,a)\right]\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi^{\star}}^{-1}}\sqrt{\gamma
d|\mathcal{A}|\alpha^{2}_{n}+\gamma
d\lambda_{n}}+\sqrt{{d|\mathcal{A}|\alpha^{2}_{n}}(1-\gamma)/n}.$

Second, we calculate the term (b) in inequality 7. Following Lemma 13 and
noting $f^{2}_{n}(s,a)$ is upper-bounded by $4$ (i.e., $B=4$ in Lemma 13), we
have

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}[f_{n}(s,a)]$
$\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi^{\star}}^{-1}}\sqrt{\left\\{n|\mathcal{A}|\gamma\mathbb{E}_{(s,a)\sim\rho_{n}}\left[f^{2}_{n}(s,a)\right]\right\\}+4\gamma\lambda_{n}d}+\sqrt{|\mathcal{A}|\mathbb{E}_{(s,a)\sim\rho_{n}}\left[{f^{2}_{n}(s,a)}{(1-\gamma)}\right]}$
$\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi^{\star}}^{-1}}\sqrt{n|\mathcal{A}|\gamma\zeta_{n}+4\gamma\lambda_{n}d}+\sqrt{|\mathcal{A}|\zeta_{n}(1-\gamma)}$
$\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi^{\star}}^{-1}}\alpha_{n}+\sqrt{|\mathcal{A}|\zeta_{n}(1-\gamma)},$

where in the second inequality, we use
$\mathbb{E}_{s,a\sim\rho}f_{n}^{2}(s,a)\leq\zeta_{n}$, and in the last line,
recall
$\sqrt{\gamma}\sqrt{{{n|\mathcal{A}|}\zeta_{n}}+\lambda_{n}d+n\zeta_{n}}\lesssim\alpha_{n}$.

Then, by combining the above calculation of the term (a) and term (b) in
inequality 7, we have:

$\displaystyle V^{\pi}_{P^{\star},r}-V^{\pi_{n}}_{P^{\star},r}$
$\displaystyle\lesssim|\mathcal{A}|\left(\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi^{\star}}^{-1}}\sqrt{d|\mathcal{A}|\alpha^{2}_{n}+{d\lambda_{n}}{}}+\sqrt{\frac{d|\mathcal{A}|\alpha^{2}_{n}(1-\gamma)}{n}}\right)$
$\displaystyle+\frac{|\mathcal{A}|}{(1-\gamma)}\left(\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi^{\star}}^{-1}}\alpha_{n}+\sqrt{|\mathcal{A}|\zeta_{n}(1-\gamma)}\right).$

Hereafter, we take the dominating term out. First, recall

$\alpha_{n}\lesssim\sqrt{\left\\{|\mathcal{A}|+d^{2}\right\\}\ln(N|\mathcal{M}|/\delta))}\lesssim\sqrt{|\mathcal{A}|d^{2}\ln(N|\mathcal{M}|/\delta)}.$

Second, we also use

$\displaystyle\sum_{n=1}^{N}\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi^{\star}}^{-1}}\leq\sqrt{N\sum_{n=1}^{N}\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}[\phi^{\star}(\tilde{s},\tilde{a})^{\top}\Sigma^{-1}_{\rho_{n},\phi^{\star}}\phi^{\star}(\tilde{s},\tilde{a})]}$
(CS inequality)
$\displaystyle\lesssim\sqrt{N\left(\ln\det(\sum_{n=1}^{N}\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\bar{\pi}_{n}}_{P^{\star}}}[\phi^{\star}(\tilde{s},\tilde{a})\phi^{\star}(\tilde{s},\tilde{a})^{\top}])-\ln\det(\lambda_{1}I)\right)}$
(Lemma 18 and $\lambda_{1}\leq\cdots\leq\lambda_{N}$)
$\displaystyle\leq\sqrt{dN\ln\left(1+\frac{N}{d\lambda_{1}}\right)}.$
(Potential function bound, Lemma 19 noting $\|\phi^{\star}(s,a)\|_{2}\leq 1$
for any $(s,a)$.)

Finally,

$\displaystyle\sum_{n=1}^{N}V^{\pi}_{P^{\star},r}-V^{\pi_{n}}_{P^{\star},r}$
$\displaystyle\lesssim|\mathcal{A}|\left(\sqrt{dN\ln\left(1+\frac{N}{d\lambda_{1}}\right)}\sqrt{{d|\mathcal{A}|\alpha^{2}_{N}}+d\lambda_{N}}+\sum_{n=1}^{N}\sqrt{\frac{d|\mathcal{A}|\alpha^{2}_{n}(1-\gamma)}{n}}\right)$
$\displaystyle+\frac{|\mathcal{A}|}{(1-\gamma)}\left(\sqrt{dN\ln\left(1+\frac{N}{d\lambda_{1}}\right)}\alpha_{N}+\sum_{n=1}^{N}\sqrt{|\mathcal{A}|\zeta_{n}(1-\gamma)}\right)$
$\displaystyle\lesssim|\mathcal{A}|\sqrt{dN\ln\left(1+\frac{N}{d\lambda_{1}}\right)}\sqrt{d|\mathcal{A}|\alpha^{2}_{N}}+\frac{|\mathcal{A}|}{(1-\gamma)}\sqrt{dN\ln\left(1+\frac{N}{d\lambda_{1}}\right)}\alpha_{N}$
(Some algebra. We take the dominating term out. )
$\displaystyle\lesssim\sqrt{dN\ln\left(1+\frac{N}{d\lambda_{1}}\right)}\frac{|\mathcal{A}|^{2}d^{3/2}\ln^{1/2}(N|\mathcal{M}|/\delta)}{(1-\gamma)}.$

This concludes the proof. ∎

Using Lemma 9, we can immediately obtain the PAC guarantee.

######  Theorem 10 (PAC guarantee of Rep-UCB).

By interacting with the environment for a number of steps at most

$\displaystyle N\log(N/\delta),\quad N\coloneqq
O\left(\frac{d^{4}|\mathcal{A}|^{4}\ln(|\mathcal{M}|/\delta)}{(1-\gamma)^{3}\epsilon^{2}}\ln^{2}\left(1+\frac{d^{4}|\mathcal{A}|^{4}\ln(|\mathcal{M}|/\delta)}{(1-\gamma)^{2}\epsilon^{2}}\right)\right).$

with probability $1-\delta$, we can ensure
$V^{\pi^{\star}}_{P^{\star},r}-V^{\hat{\pi}}_{P^{\star},r}\leq\epsilon.$

###### Proof.

From Lemma 9 and Lemma 22, when $N$ is

$\displaystyle
O\left(\frac{d^{4}|\mathcal{A}|^{4}\ln(|\mathcal{M}|/\delta)}{(1-\gamma)^{2}\epsilon^{2}}\ln^{2}\left(1+\frac{d^{4}|\mathcal{A}|^{4}\ln(|\mathcal{M}|/\delta)}{(1-\gamma)^{2}\epsilon^{2}}\right)\right),$

with probability $1-\delta$, we can ensure

$\displaystyle\frac{1}{N}\sum_{n=1}^{N}V^{\pi^{\star}}_{P^{\star},r}-V^{\pi_{n}}_{P^{\star},r}\leq\epsilon.$

With probability $1-\delta$, we need $\ln(1/\delta)$ interactions with the
environment to get one triple $(s,a,s^{\prime})$ from
$d^{\pi}_{P^{\star}}(s,a)$. Thus, the total sample complexity is
$O(N(1-\gamma)^{-1}\ln(N/\delta))$.

∎

Next, we provide an important lemma to ensure the concentration of the bonus
term. The version for fixed $\phi$ is proved in Zanette et al. (2021a, Lemma
39). Here, we take a union bound over the whole feature $\phi\in\Phi$. Recall

$\displaystyle\rho_{n}=\frac{1}{n}\sum_{i=0}^{n-1}d^{\bar{\pi}_{i}}_{P^{\star}}(s,a),\quad(s^{(i)},a^{(i)})\sim
d^{\bar{\pi}_{i}}_{P^{\star}}(s,a).$

The extension to the offline case is again straightforward. Recall the
definition of $d^{\bar{\pi}}(s,a)$ in the beginning of Appendix.

######  Lemma 11 (Concentration of the bonus term).

Set $\lambda_{n}=\Theta(d\ln(n|\Phi|/\delta))$ for any $n$. Define

$\displaystyle\Sigma_{\rho_{n},\phi}=n\mathbb{E}_{(s,a)\sim\rho_{n}}[\phi(s,a)\phi^{\top}(s,a)]+\lambda_{n}I,\quad\hat{\Sigma}_{n,\phi}=\sum_{i=0}^{n-1}\phi(s^{(i)},a^{(i)})\phi^{\top}(s^{(i)},a^{(i)})+\lambda_{n}I.$

With probability $1-\delta$, we have

$\displaystyle\forall
n\in\mathbb{N}^{+},\forall\phi\in\Phi,c_{1}\|\phi(s,a)\|_{\Sigma^{-1}_{\rho_{n},\phi}}\leq\|\phi(s,a)\|_{\hat{\Sigma}^{-1}_{n,\phi}}\leq
c_{2}\|\phi(s,a)\|_{\Sigma^{-1}_{\rho_{n},\phi}}.$

Next, we show our key lemma leveraging low-rank MDP structures. For any
$g\in{\mathcal{S}}\times\mathcal{A}\to\mathbb{R}$, This lemma shows that
$\mathbb{E}_{(s,a)\sim d^{\pi}_{\hat{P}_{n}}}\left\\{g(s,a)\right\\}$ can be
upper-bounded using $\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}_{n}}}\|\hat{\phi}_{n}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\hat{\phi}_{n}}^{-1}}$
as long as we have the convergence guarantee for
$\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]$.

######  Lemma 12 (One-step back inequality for the learned model).

Take any $g\in{\mathcal{S}}\times\mathcal{A}\to\mathbb{R}$ s.t.
$\|g\|_{\infty}\leq B$. We condition on the event where the MLE guarantee
(15):

$\displaystyle\mathbb{E}_{(s,a)\sim\rho_{n}}[\|\hat{P}_{n}(\cdot\mid
s,a)-P^{\star}(\cdot\mid s,a)\|^{2}_{1}]\lesssim\zeta_{n},$

holds. Then, for any policy $\pi$, we have

$\displaystyle|\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left\\{g(s,a)\right\\}|$
$\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}_{n}}}\|\hat{\phi}_{n}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\hat{\phi}_{n}}^{-1}}\sqrt{\gamma}\sqrt{\left\\{n|\mathcal{A}|\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]\right\\}+\lambda_{n}dB^{2}+nB^{2}\zeta_{n}}$
$\displaystyle+\sqrt{(1-\gamma)|\mathcal{A}|\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]}.$

Recall
$\Sigma_{\rho_{n},\hat{\phi}_{n}}=n\mathbb{E}_{(s,a)\sim\rho_{n}}[\hat{\phi}_{n}(s,a)\hat{\phi}^{\top}_{n}(s,a)]+\lambda_{n}I$.

###### Proof.

First, we have an equality:

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left\\{g(s,a)\right\\}=\gamma\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}_{n}},s\sim\hat{P}_{n}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left\\{g(s,a)\right\\}+(1-\gamma)\mathbb{E}_{s\sim
d_{0},a\sim\pi(s_{0})}\left\\{g(s,a)\right\\},$ (8)

The second term in (8) is upper-bounded by

$\displaystyle(1-\gamma)\sqrt{\max_{(s,a)}\frac{d_{0}(s)\pi(a\mid
s)}{\rho_{n}(s,a)}\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]}$
$\displaystyle\leq(1-\gamma)\sqrt{\max_{(s,a)}\frac{d_{0}(s)\pi(a\mid
s)}{(1-\gamma)d_{0}(s)u(a)}\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]}\leq\sqrt{(1-\gamma)|\mathcal{A}|\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]}.$

Next we consider the first term in (8). By CS inequality, we have

$\displaystyle\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}_{n}},s\sim\hat{P}_{n}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left\\{g(s,a)\right\\}=\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}_{n}}}\hat{\phi}_{n}(\tilde{s},\tilde{a})^{\top}\int\sum_{a}\hat{\mu}_{n}(s)\pi(a\mid
s)g(s,a)d(s)$ $\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}_{n}}}\|\hat{\phi}_{n}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\hat{\phi}_{n}}^{-1}}\left\|\int\sum_{a}\hat{\mu}_{n}(s)\pi(a\mid
s)g(s,a)d(s)\right\|_{\Sigma_{\rho_{n},\hat{\phi}_{n}}}.$

Then,

$\displaystyle\|\int\sum_{a}\hat{\mu}_{n}(s)\pi(a\mid
s)g(s,a)d(s)\|^{2}_{\Sigma_{\rho_{n},\hat{\phi}_{n}}}$
$\displaystyle\leq\left\\{\int\sum_{a}\hat{\mu}_{n}(s)\pi(a\mid
s)g(s,a)d(s)\right\\}^{\top}\left\\{n\mathbb{E}_{(s,a)\sim\rho_{n}}[\hat{\phi}_{n}\hat{\phi}^{\top}_{n}]+\lambda_{n}I\right\\}\left\\{\int\sum_{a}\hat{\mu}_{n}(s)\pi(a\mid
s)g(s,a)d(s)\right\\}$ $\displaystyle\leq
n\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho_{n}}\left\\{\left[\int\sum_{a}\hat{\mu}_{n}(s)^{\top}\hat{\phi}_{n}(\tilde{s},\tilde{a})\pi(a\mid
s)g(s,a)d(s)\right]^{2}\right\\}+B^{2}\lambda_{n}d$ (Use the assumption
$\|\sum_{a}\pi(a\mid s)g(s,a)\|_{\infty}\leq B$ and
$\int\|\hat{\mu}_{n}(s)h(s){\textnormal{d}}(s)\|_{2}\leq\sqrt{d}$ for any
$h:{\mathcal{S}}\to[0,1]$. )
$\displaystyle=n\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho_{n}}\left[\left\\{\mathbb{E}_{s\sim\hat{P}_{n}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left[g(s,a)\right]\right\\}^{2}\right]+B^{2}\lambda_{n}d$
$\displaystyle=n\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho_{n}}\left[\left\\{\mathbb{E}_{s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left[g(s,a)\right]\right\\}^{2}\right]+B^{2}\lambda_{n}d+n\zeta_{n}B^{2}$
(MLE guarantee) $\displaystyle\leq
n\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho_{n},s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left[g^{2}(s,a)\right]+B^{2}\lambda_{n}d+n\zeta_{n}B^{2}.$
(Jensen)
$\displaystyle\leq\frac{n}{\min\\{u(a)\\}}\left\\{\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho_{n},s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim
U(\mathcal{A})}\left[g^{2}(s,a)\right]\right\\}+\lambda_{n}dB^{2}+n\zeta_{n}B^{2}$
(Importance sampling) $\displaystyle\leq
n|\mathcal{A}|\left\\{\frac{1}{(\gamma)}\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]\right\\}+\lambda_{n}dB^{2}+n\zeta_{n}B^{2},$

Then, the final statement is immediately concluded. ∎

Here, we show a similar lemma as Lemma 12. The difference is we aim for
calculating $\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\star}}}\left\\{g(s,a)\right\\}$ instead of $\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}_{n}}}\left\\{g(s,a)\right\\}$ . For any
$g\in{\mathcal{S}}\times\mathcal{A}\to\mathbb{R}$, this lemma shows that
$\mathbb{E}_{(s,a)\sim d^{\pi}_{P^{\star}}}\left\\{g(s,a)\right\\}$ can be
upper-bounded using $\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\hat{\phi}^{\star}}^{-1}}$
as long as we have the convergence guarantee for
$\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]$. Note comparing to
Lemma 12, this is not a probabilistic statement. Note that
$\|\phi^{\star}(s,a)\|_{\Sigma^{-1}_{\rho_{n},\phi^{\star}}}$ is the usual
elliptical potential function under the fixed representation $\phi^{\star}$.

######  Lemma 13 (One-step back inequality for the true model ).

Take any $g\in{\mathcal{S}}\times\mathcal{A}\to\mathbb{R}$ such that
$\|g\|_{\infty}\leq B$. Then,

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\star}}}\left\\{g(s,a)\right\\}$
$\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi{\star}}^{-1}}\sqrt{\gamma}\sqrt{{n|\mathcal{A}|}\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]+\lambda_{n}dB^{2}}$
$\displaystyle+\sqrt{(1-\gamma)|\mathcal{A}|\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]}.$

Recall
$\Sigma_{\rho_{n},\phi^{\star}}=n\mathbb{E}_{(s,a)\sim\rho_{n}}[\phi^{\star}\\{\phi^{\star}\\}^{\top}]+\lambda_{n}I$.

###### Proof.

First, we have

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\star}}}\left\\{g(s,a)\right\\}=\gamma\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}},s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left\\{g(s,a)\right\\}+(1-\gamma)\mathbb{E}_{s\sim
d_{0},a\sim\pi(s_{0})}\left\\{g(s,a)\right\\}.$ (9)

The second term in (9) is upper-bounded by

$\displaystyle(1-\gamma)\sqrt{\max_{(s,a)}\frac{d_{0}(s)\pi(a\mid
s)}{\rho_{n}(s,a)}\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]}\leq\sqrt{|\mathcal{A}|\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right](1-\gamma)}.$

By CS inequality, the first term in (9) is further bounded as follows:

$\displaystyle\mathbb{E}_{(\tilde{s},\tilde{a})\sim d^{\pi}_{P^{\star}},s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left\\{g(s,a)\right\\}=\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\phi^{\star}(\tilde{s},\tilde{a})^{\top}\int\sum_{a}\mu^{\star}(s)\pi(a\mid
s)g(s,a)d(s)$ $\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi^{\star}}^{-1}}\left\|\int\sum_{a}\mu^{\star}(s)\pi(a\mid
s)g(s,a)d(s)\right\|_{\Sigma_{\rho_{n},\phi^{\star}}}.$

Here, we have

$\displaystyle\|\int\sum_{a}\mu^{\star}(s)\pi(a\mid
s)g(s,a)d(s)\|^{2}_{\Sigma_{\rho_{n},\phi^{\star}}}$
$\displaystyle\leq\left\\{\int\sum_{a}\mu^{\star}(s)\pi(a\mid
s)g(s,a)d(s)\right\\}^{\top}\left\\{n\mathbb{E}_{(s,a)\sim\rho_{n}}[\phi^{\star}(s,a)\\{\phi^{\star}(s,a)\\}^{\top}]+\lambda_{n}I\right\\}\left\\{\int\sum_{a}\mu^{\star}(s)\pi(a\mid
s)g(s,a)d(s)\right\\}$ $\displaystyle\leq
n\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho_{n}}\left\\{\left[\int\sum_{a}\mu^{\star}(s)^{\top}\phi^{\star}(\tilde{s},\tilde{a})\pi(a\mid
s)g(s,a)d(s)\right]^{2}\right\\}+\lambda_{n}dB^{2}$ $\displaystyle\leq
n\left\\{\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho_{n},s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left[g^{2}(s,a)\right]\right\\}+\lambda_{n}dB^{2}.$
(Jensen)

Therefore,

$\displaystyle n\left\\{\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho_{n},s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left[g^{2}(s,a)\right]\right\\}+\lambda_{n}dB$
$\displaystyle\leq\frac{n}{\min\\{u(a)\\}}\left\\{\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho_{n},s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim
U(\mathcal{A})}\left[g^{2}(s,a)\right]\right\\}+\lambda_{n}dB^{2}$
$\displaystyle\leq
n|\mathcal{A}|\left\\{\frac{1}{\gamma}\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]\right\\}+\lambda_{n}dB^{2}.$

Finally, we have

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\star}}}\left\\{g(s,a)\right\\}$
$\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho_{n},\phi{\star}}^{-1}}\sqrt{\left\\{\frac{n|\mathcal{A}|}{\gamma}\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right]\right\\}+\lambda_{n}dB^{2}}$
$\displaystyle+\sqrt{|\mathcal{A}|\mathbb{E}_{(s,a)\sim\rho_{n}}\left[g^{2}(s,a)\right](1-\gamma)}.$

This concludes the proof. ∎

##  Appendix B Proof of the theoretical property of Rep-LCB

This section provides the detailed proofs for our results in the offline
setting.

Below we first prove that $V^{\pi}_{\hat{P},r-\hat{b}}$ is an almost
pessimistic estimator of $V^{\pi}_{P^{\star},r}$.

######  Lemma 14 (Almost Pessimism at the Initial Distribution).

Let $\omega=\max_{a,s}1/\pi_{b}(a\mid s)$. Set

$\displaystyle\alpha={c_{1}}\sqrt{\left(\omega+d^{2}\right)\gamma\ln(|\mathcal{M}|/\delta)},\quad\lambda=O(d\ln(|\mathcal{M}|/\delta)),\zeta=O\left(\frac{\ln(|\mathcal{M}|/\delta)}{n}\right).$

With probability $1-\delta$, for any policy $\pi$, we have

$\displaystyle
V^{\pi}_{\hat{P},r-\hat{b}}-V^{\pi}_{P^{\star},r}\leq\sqrt{\frac{\omega\ln(|\mathcal{M}|/\delta)(1-\gamma)}{n}}.$

###### Proof.

We define

$\displaystyle\Sigma_{\rho,\phi}=n\mathbb{E}_{(s,a)\sim\rho}[\phi\phi^{\top}]+\lambda
I,\quad\hat{\Sigma}_{\phi}=n\mathbb{E}_{\mathcal{D}}[\phi\phi^{\top}]+\lambda
I.$

where $\lambda=O(d\ln(|\mathcal{M}|/\delta))$. In this proof, letting
$f(s,a)=\|\hat{P}(\cdot\mid s,a)-P^{\star}(\cdot\mid s,a)\|_{1}$, we condition
on the events:

$\displaystyle\mathbb{E}_{(s,a)\sim\rho}[f^{2}(s,a)]\leq\zeta,\quad\forall\phi\in\Phi:\|\phi(s,a)\|_{\hat{\Sigma}^{-1}_{\phi}}=\Theta(\|\phi(s,a)\|_{\Sigma^{-1}_{\rho,\phi}}).$
(10)

where $\zeta=O(\ln(|\mathcal{M}|/\delta)/n)$. From the offline version of
Lemma 17 and Lemma 11 666We can remove $\ln n$ since $n$ is fixed in the
offline setting. , this event happens with probability $1-\delta$.

Then, from simulation lemma (Lemma 20),

$\displaystyle V^{\pi}_{\hat{P},r-\hat{b}}-V^{\pi}_{P^{\star},r}$
$\displaystyle=\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left[-\hat{b}(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim\hat{P}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]-\gamma\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]\right]$
$\displaystyle\lesssim\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left[-\min\left(\alpha\|\hat{\phi}(s,a)\|_{\Sigma^{-1}_{\rho,\hat{\phi}}},2\right)+\gamma\mathbb{E}_{s^{\prime}\sim\hat{P}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]-\gamma\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]\right].$ (From
(10))

Here, we have

$\displaystyle\left|\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left\\{\mathbb{E}_{s^{\prime}\sim\hat{P}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]-\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]\right\\}\right|\leq\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left\\{f(s,a)\right\\},$

noting $\|V^{\pi}_{P^{\star},r}\|_{\infty}\leq 1$. This is further bounded by
Lemma 16:

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left\\{f(s,a)\right\\}\lesssim\sqrt{\alpha^{\prime}}\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}}}\|\hat{\phi}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\hat{\phi}}^{-1}}+\sqrt{\omega\zeta(1-\gamma)}.$
(11)

where

$\displaystyle\alpha^{\prime}=n\gamma\omega\zeta+\gamma^{2}\lambda
d+\gamma^{2}n\zeta\lesssim\left(\omega+d^{2}\right)\gamma\ln(|\mathcal{M}|/\delta).$

Here, we use $f(s,a)\leq 2$ in Lemma 16 and
$\mathbb{E}_{(s,a)\sim\rho}[f^{2}(s,a)]\leq\zeta$.

Thus,

$\displaystyle\left|\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left\\{\mathbb{E}_{s^{\prime}\sim\hat{P}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]-\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]\right\\}\right|\leq\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left\\{f(s,a)\right\\}$
$\displaystyle\leq\sqrt{\alpha^{\prime}}\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}}}\|\hat{\phi}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\hat{\phi}}^{-1}}+\sqrt{\omega\zeta(1-\gamma)}$
$\displaystyle=\alpha\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}}}\|\hat{\phi}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\hat{\phi}}^{-1}}+\sqrt{\omega\zeta(1-\gamma)},\quad\alpha=\sqrt{\alpha^{\prime}}.$

Going back to the simulation lemma 20, we have

$\displaystyle V^{\pi}_{\hat{P},r-\hat{b}}-V^{\pi}_{P^{\star},r}$
$\displaystyle\lesssim\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left[-\min\left(\alpha\|\hat{\phi}(s,a)\|_{\Sigma_{\rho,\hat{\phi}}^{-1}},2\right)+\mathbb{E}_{s^{\prime}\sim\hat{P}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]-\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}\left[V^{\pi}_{P^{\star},r}(s^{\prime})\right]\right]$
$\displaystyle\leq\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left[-\min\left(\alpha\|\hat{\phi}(s,a)\|_{\Sigma_{\rho,\hat{\phi}}^{-1}},2\right)+\min\left(\alpha\|\hat{\phi}(s,a)\|_{\Sigma_{\rho,\hat{\phi}}^{-1}}+\sqrt{\omega\zeta(1-\gamma)},2\right)\right]$
$\displaystyle\leq\sqrt{\omega\zeta(1-\gamma)}.$

This concludes the proof. ∎

With the above lemma, now we can proceed to prove the main theorem.

######  Theorem 15 (PAC guarantee of Rep-LCB).

Set the parameters as in Lemma 14. With probability $1-\delta$, for any
comparator policy $\pi$ including history-dependent non-Markovian policies, we
have

$\displaystyle
V^{\pi}_{P^{\star},r}-V^{\hat{\pi}}_{P^{\star},r}\lesssim\frac{\omega
d^{2}}{(1-\gamma)}\sqrt{\frac{C^{\star}_{\pi}\ln(|\mathcal{M}|/\delta)}{n}},$

where $C^{\star}_{\pi}$ is the relative condition number under $\phi^{\star}$:

$\displaystyle
C^{\star}_{\pi}\coloneqq\sup_{x\in\mathbb{R^{d}}}\frac{x^{\top}\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\star}}}[\phi^{\star}(s,a)\\{\phi^{\star}(s,a)\\}^{\top}]x}{x^{\top}\mathbb{E}_{(s,a)\sim\rho}[\phi^{\star}(s,a)\\{\phi^{\star}(s,a)\\}^{\top}]x}.$

###### Proof.

In this proof, letting $f(s,a)=\|\hat{P}(\cdot\mid s,a)-P^{\star}(\cdot\mid
s,a)\|_{1}$ we condition on the events:

$\displaystyle\mathbb{E}_{(s,a)\sim\rho}[f^{2}(s,a)]\leq\zeta,\quad\forall\phi\in\Phi:\|\phi(s,a)\|_{\hat{\Sigma}^{-1}_{\phi}}=\Theta(\|\phi(s,a)\|_{\Sigma^{-1}_{\rho,\phi}}).$
(12)

From Lemma 11 and Lemma 17, this event happens with probability $1-\delta$.

For any policy $\pi$, we have

$\displaystyle V^{\pi}_{P^{\star},r}-V^{\hat{\pi}}_{P^{\star},r}$
$\displaystyle\leq
V^{\pi}_{P^{\star},r}-V^{\hat{\pi}}_{\hat{P},r-\hat{b}}+\sqrt{\omega\zeta(1-\gamma)}$
(Lemma 14) $\displaystyle\leq
V^{\pi}_{P^{\star},r}-V^{\pi}_{\hat{P},r-\hat{b}}+\sqrt{\omega\zeta(1-\gamma)}$
$\displaystyle\lesssim\underbrace{\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\star}}}[\hat{b}(s,a)]}_{\text{(a)}}+\underbrace{\left(\frac{1}{1-\gamma}\right)\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\star}}}[f(s,a)]}_{\text{(b)}}+\sqrt{\omega\zeta(1-\gamma)}.$

Recall $f(s,a)=\|\hat{P}(\cdot\mid s,a)-P^{\star}(\cdot\mid s,a)\|_{1}$.

From the second line to the third line, note though $\hat{\pi}$ is the argmax
over Markoovian polices, $\hat{\pi}$ is also the argmax over all history-
dependent polices. In the last line, we use a simulation lemma 20, which is
tailored to a time-inhomogeneous policy. We here use
$\|V^{\pi}_{\hat{P},r-\hat{b}}\|_{\infty}\leq 2/((1-\gamma))$. noting
$\|\hat{b}\|_{\infty}=O(1)$.

We further calculate the first term (a). Considering the offline version of
Lemma 13 and noting $\|\hat{b}\|_{\infty}\leq 2$, we have

$\displaystyle\mathbb{E}_{(s,a)\sim d^{\pi}_{P^{\star}}}[\hat{b}(s,a)]$
$\displaystyle\lesssim\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\phi^{\star}}^{-1}}\sqrt{n\omega\left\\{\gamma\mathbb{E}_{(s,a)\sim\rho}\left[\hat{b}^{2}(s,a)\right]\right\\}+\gamma\lambda
d}$
$\displaystyle+\sqrt{\omega(1-\gamma)}\\{\mathbb{E}_{\rho}[\hat{b}^{2}(s,a)]\\}^{1/2}.$

From (12), we have

$\displaystyle n\mathbb{E}_{(s,a)\sim\rho}\left[\hat{b}^{2}(s,a)\right]\leq
n\mathbb{E}_{(s,a)\sim\rho}\left[\min\left(\alpha^{2}\|\hat{\phi}(s,a)\|^{2}_{\Sigma^{-1}_{\rho,\hat{\phi}}},4\right)\right]\leq
n\mathbb{E}_{(s,a)\sim\rho}\left[\alpha^{2}\|\hat{\phi}(s,a)\|^{2}_{\Sigma^{-1}_{\rho,\hat{\phi}}}\right]$
(13)
$\displaystyle\leq\operatorname{Tr}[n\mathbb{E}[\hat{\phi}\hat{\phi}^{\top}]\\{n\mathbb{E}[\hat{\phi}\hat{\phi}^{\top}]+\lambda
I\\}^{-1}]\leq\operatorname{Tr}[n(\mathbb{E}[\hat{\phi}\hat{\phi}^{\top}]+\lambda
I)\\{n\mathbb{E}[\hat{\phi}\hat{\phi}^{\top}]+\lambda I\\}^{-1}]\leq d.$

Thus,

$\displaystyle\mathbb{E}_{(s,a)\sim d^{\pi}_{P^{\star}}}[\hat{b}(s,a)]$
$\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\phi^{\star}}^{-1}}\sqrt{\omega
d\alpha^{2}\gamma+\gamma\lambda d}+\sqrt{\frac{\omega
d\alpha^{2}(1-\gamma)}{n}}.$

Second, we further calculate the second term (b). Considering the offline
version of Lemma 13 and noting $f^{2}(s,a)$ is upper-bounded by $4$,

$\displaystyle\mathbb{E}_{(s,a)\sim d^{\pi}_{P^{\star}}}[f(s,a)]$
$\displaystyle=\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\phi^{\star}}^{-1}}\sqrt{n\omega\left\\{\gamma\mathbb{E}_{(s,a)\sim\rho}\left[f^{2}(s,a)\right]\right\\}+4\gamma\lambda
d}+\sqrt{\omega\mathbb{E}_{(s,a)\sim\rho}\left[f^{2}(s,a)\right](1-\gamma)}$
$\displaystyle\lesssim\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\phi^{\star}}^{-1}}\sqrt{\omega\left\\{n\gamma\zeta\right\\}+\gamma\lambda
d}+\sqrt{\omega\zeta(1-\gamma)}$
$\displaystyle\lesssim\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\phi^{\star}}^{-1}}\alpha+\sqrt{\omega\zeta(1-\gamma)}.$

In the final line, recall
$\sqrt{\omega\left\\{n\gamma\zeta\right\\}+\gamma\lambda d+\gamma
n\zeta}\leq\alpha.$

Finally, by combining the calculation of the first term (a) and the second
term (b), we have

$\displaystyle V^{\pi}_{P^{\star},r}-V^{\hat{\pi}}_{P^{\star},r}$
$\displaystyle\lesssim\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\phi^{\star}}^{-1}}\sqrt{d\alpha^{2}\omega\gamma+\gamma\lambda
d}+\sqrt{\frac{\omega\alpha^{2}d(1-\gamma)}{n}}$
$\displaystyle+\frac{\alpha}{(1-\gamma)}\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\phi^{\star}}^{-1}}+\sqrt{\frac{\omega\zeta}{(1-\gamma)}}$

Now, we use the fact $\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\phi^{\star}}^{-1}}$
is upper-bounded as

$\displaystyle\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\phi^{\star}}^{-1}}$
$\displaystyle\leq\sqrt{\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{P^{\star}}}\|\phi^{\star}(\tilde{s},\tilde{a})\|^{2}_{\Sigma_{\rho,\phi^{\star}}^{-1}}}\leq\sqrt{C^{\star}\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho}\|\phi^{\star}(\tilde{s},\tilde{a})\|^{2}_{\Sigma_{\rho,\phi^{\star}}^{-1}}}$
(Refer to Lemma 21 ) $\displaystyle\leq\sqrt{C^{\star}d/n}.$ (From (13))

Finally, we have

$\displaystyle V^{\pi}_{P^{\star},r}-V^{\hat{\pi}}_{P^{\star},r}$
$\displaystyle\lesssim\sqrt{\frac{C^{\star}d}{n}}\sqrt{d\alpha^{2}\omega\gamma+\gamma\lambda
d}+\frac{\alpha}{(1-\gamma)}\sqrt{\frac{C^{\star}d}{n}}+\sqrt{\frac{\omega\alpha^{2}d(1-\gamma)}{n}}+\sqrt{\frac{\omega\zeta}{(1-\gamma)}}$
$\displaystyle\lesssim\sqrt{\frac{C^{\star}d}{n}}\sqrt{{d\alpha^{2}\omega}{\gamma}}+\frac{\alpha}{(1-\gamma)}\sqrt{\frac{C^{\star}d}{n}}$
(Take out two dominating terms) $\displaystyle\lesssim\frac{\omega
d^{2}}{(1-\gamma)}\sqrt{\frac{C^{\star}\ln(|\mathcal{M}|/\delta)}{n}}.$

∎

The lemma below is a key technical lemma for our proof. It shows that one can
relate the expected value of any function $f(s,a)$ with respect to
$d^{\pi}_{\hat{P}}$ (i.e., inside the learned model $\hat{P}$) to the
potential function with respect to $d^{\pi}_{\hat{P}}$, i.e.,
$\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}}}\|\hat{\phi}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\hat{\phi}}^{-1}}$.
Pairing $\hat{\phi}$ and $\hat{P}$ is important since $\hat{P}$ is the low-
rank transition model defined using $\hat{\phi}$. As we have seen in the above
analysis, when using the lemma below, we instantiate
$f(s,a):=\|\hat{P}(\cdot|s,a)-P^{\star}(\cdot|s,a)\|_{1}$.

######  Lemma 16 (One-step back inequality for the learned model in offline
setting).

Take any $f\subset{\mathcal{S}}\times\mathcal{A}\to\mathbb{R}$ s.t.
$\|f\|_{\infty}\leq B$. We condition on the event where the MLE guarantee
holds:

$\displaystyle\mathbb{E}_{(s,a)\sim\rho}\|\hat{P}(\cdot\mid
s,a)-P^{\star}(\cdot\mid s,a)\|^{2}_{1}\lesssim\zeta.$

Then, letting $\omega=\max_{s,a}(1/\pi_{b}(a\mid s))$, for any policy $\pi$,
we have

$\displaystyle|\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left\\{f(s,a)\right\\}|$
$\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}}}\|\hat{\phi}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\hat{\phi}}^{-1}}\sqrt{\left\\{n\omega\mathbb{E}_{(s,a)\sim\rho}\left[f^{2}(s,a)\right]\right\\}+\gamma^{2}\lambda
dB^{2}+n\gamma^{2}\zeta B^{2}}$
$\displaystyle+\sqrt{\omega\mathbb{E}_{(s,a)\sim\rho}\left[f^{2}(s,a)\right](1-\gamma)}.$

###### Proof.

First, we have an equality:

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\pi}_{\hat{P}}}\left\\{f(s,a)\right\\}=\gamma\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}},s\sim\hat{P}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left\\{f(s,a)\right\\}+(1-\gamma)\mathbb{E}_{s\sim
d_{0},a\sim\pi(s_{0})}\left\\{f(s,a)\right\\}.$ (14)

The second term in (14) is upper-bounded by

$\displaystyle\mathbb{E}_{s\sim
d_{0},a\sim\pi(s_{0})}\left\\{f(s,a)\right\\}\leq\mathbb{E}_{s\sim
d_{0},a\sim\pi(s_{0})}\left\\{f^{2}(s,a)\right\\}\\}^{1/2}=\sqrt{\omega\mathbb{E}_{(s,a)\sim\rho}\left[f^{2}(s,a)\right]/(1-\gamma)}.$

Next we consider the first term in (14). By CS inequality, we have

$\displaystyle\left|\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}},s\sim\hat{P}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left\\{f(s,a)\right\\}\right|=\left|\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}}}\hat{\phi}(\tilde{s},\tilde{a})^{\top}\int\sum_{a}\hat{\mu}(s)\pi(a\mid
s)f(s,a)d(s)\right|$ $\displaystyle\leq\mathbb{E}_{(\tilde{s},\tilde{a})\sim
d^{\pi}_{\hat{P}}}\|\hat{\phi}(\tilde{s},\tilde{a})\|_{\Sigma_{\rho,\hat{\phi}}^{-1}}\|\int\sum_{a}\hat{\mu}(s)\pi(a\mid
s)f(s,a)d(s)\|_{\Sigma_{\rho,\hat{\phi}}}.$

Then,

$\displaystyle\|\int\hat{\mu}(s)\pi(a\mid
s)f(s,a)d(s,a)\|^{2}_{\Sigma_{\rho,\hat{\phi}}}$
$\displaystyle\leq\left\\{\int\sum_{a}\hat{\mu}(s)\pi(a\mid
s)f(s,a)d(s)\right\\}^{\top}\left\\{n\mathbb{E}_{(s,a)\sim\rho}[\hat{\phi}\hat{\phi}^{\top}]+\lambda
I\right\\}\left\\{\int\sum_{a}\hat{\mu}(s)\pi(a\mid s)f(s,a)d(s)\right\\}$
$\displaystyle\leq
n\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho}\left\\{\left[\int\sum_{a}\hat{\mu}(s)^{\top}\hat{\phi}(\tilde{s},\tilde{a})\pi(a\mid
s)f(s,a)d(s)\right]^{2}\right\\}+B^{2}\lambda d$ (Use the assumption
$\|\sum_{a}f(s,a)\|_{\infty}\leq B$ and
$\|\int\hat{\mu}(s)h(s){\textnormal{d}}(s)\|_{2}\leq\sqrt{d}$ for
$h:{\mathcal{S}}\to[0,1]$. )
$\displaystyle=n\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho}\\{\mathbb{E}_{s\sim\hat{P}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left[f(s,a)\right]^{2}\\}+B^{2}\lambda
d$
$\displaystyle=n\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho}\\{\mathbb{E}_{s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left[f(s,a)\right]^{2}\\}+B^{2}\lambda
d+nB^{2}\zeta$ (MLE guarantee and
$\|\mathbb{E}_{a\sim\pi(\cdot)}[f^{2}(\cdot,a)]\|_{\infty}\leq B^{2}.$)
$\displaystyle\leq n\left\\{\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho,s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left[f^{2}(s,a)\right]\right\\}+B^{2}\lambda
d+nB^{2}\zeta.$ (Jensen)

Finally, the first term in (14) is upper-bounded by

$\displaystyle n\left\\{\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho,s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim\pi(s)}\left[f^{2}(s,a)\right]\right\\}+\lambda
dB^{2}+nB^{2}\zeta$ $\displaystyle\leq
n\omega\left\\{\mathbb{E}_{(\tilde{s},\tilde{a})\sim\rho,s\sim
P^{\star}(\tilde{s},\tilde{a}),a\sim\pi_{b}(s)}\left[f^{2}(s,a)\right]\right\\}+\lambda
dB^{2}+nB^{2}\zeta$ (Importance sampling) $\displaystyle\leq
n\omega\left\\{\frac{1}{\gamma}\mathbb{E}_{(s,a)\sim\rho}\left[f^{2}(s,a)\right]\right\\}+\lambda
dB^{2}+nB^{2}\zeta.$ (Definition of $\rho$)

The final statement is immediately concluded.

∎

##  Appendix C Auxiliary lemmas

First, we present the MLE guarantee. Regarding the proof, refer to Agarwal et
al. (2020b, Theorem 21). Note $\hat{P}_{n}$ and $\bar{\pi}_{n}$ are the
quantities appearing in the proposed online algorithm. We can also immediately
obtain the statement to the offline case.

######  Lemma 17 (MLE guarantee).

For a fixed episode $n$, with probability $1-\delta$,

$\displaystyle\sum_{j=0}^{n-1}\mathbb{E}_{(s,a)\sim
d^{\bar{\pi}_{j}}}[\|\hat{P}_{n}(\cdot\mid s,a)-P^{\star}(\cdot\mid
s,a)\|^{2}_{1}]\lesssim
n\zeta,\quad\zeta\coloneqq\frac{\ln(|\mathcal{M}|/\delta)}{n}.$

As a straightforward corollary, with probability $1-\delta$,

$\displaystyle\forall n\in\mathbb{N}^{+},\sum_{j=0}^{n-1}\mathbb{E}_{(s,a)\sim
d^{\bar{\pi}_{j}}}[\|\hat{P}_{n}(\cdot\mid s,a)-P^{\star}(\cdot\mid
s,a)\|^{2}_{1}]\lesssim
n\zeta_{n},\quad\zeta_{n}\coloneqq\frac{\ln(|\mathcal{M}|n/\delta)}{n}.$ (15)

The following is a standard inequality to prove regret bounds for linear
models. Refer to Agarwal et al. (2020a, Lemma G.2.)

######  Lemma 18.

Consider the following process. For $n=1,\cdots,N$, $M_{n}=M_{n-1}+G_{n}$ with
$M_{0}=\lambda_{0}I$ and $G_{n}$ being a positive semidefinite matrix with
eigenvalues upper-bounded by $1$. We have that:

$\displaystyle
2\ln\det(M_{N})-2\ln\det(\lambda_{0}I)\geq\sum_{n=1}^{N}\operatorname{Tr}(G_{n}M^{-1}_{n-1}).$

######  Lemma 19 (Potential function lemma).

Suppose $\operatorname{Tr}(G_{n})\leq B^{2}$.

$\displaystyle 2\ln\det(M_{N})-2\ln\det(\lambda_{0}I)\leq
d\ln\left(1+\frac{NB^{2}}{d\lambda_{0}}\right).$

###### Proof.

Let $\sigma_{1},\cdots,\sigma_{d}$ be the set of singular values of $M_{N}$
recalling $M_{N}$ is a positive semidefinite matrix. Then, by the AM-GM
inequality,

$\displaystyle\ln\det(M_{N})/\det(\lambda_{0}I)=\ln\prod_{i=1}^{d}(\sigma_{i}/\lambda_{0})\leq\ln
d\left(\frac{1}{d}\sum_{i=1}^{d}(\sigma_{i}/\lambda_{0}))\right)$

Since we have $\sum_{i}\sigma_{i}=\operatorname{Tr}(M_{N})\leq
d\lambda_{0}+NB^{2}$, the statement is concluded. ∎

######  Lemma 20 (Simulation lemma).

Given two MDPs $(P^{\prime},r+b)$ and $(P,r)$, for any policy $\pi$, we have:

$\displaystyle
V^{\pi}_{P^{\prime},r+b}-V^{\pi}_{P,r}=\frac{1}{1-\gamma}\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\prime}}}[b(s,a)+\gamma\mathbb{E}_{P^{\prime}(s^{\prime}\mid
s,a)}[Q^{\pi}_{P,r}(s^{\prime},\pi)]-\gamma\mathbb{E}_{P(s^{\prime}\mid
s,a)}[Q^{\pi}_{P,r}(s^{\prime},\pi)]]$

and

$\displaystyle
V^{\pi}_{P^{\prime},r+b}-V^{\pi}_{P,r}=\frac{1}{1-\gamma}\mathbb{E}_{(s,a)\sim
d^{\pi}_{P}}[b(s,a)+\gamma\mathbb{E}_{P^{\prime}(s^{\prime}\mid
s,a)}[Q^{\pi}_{P,r+b}(s^{\prime},\pi)]-\gamma\mathbb{E}_{P(s^{\prime}\mid
s,a)}[Q^{\pi}_{P^{\prime},r+b}(s^{\prime},\pi)]].$

###### Proof.

We use

$\displaystyle
V^{\pi}_{P}-f(s_{0},\pi)=\frac{1}{1-\gamma}\mathbb{E}_{d^{\pi}_{P}}[r(s,a)+\gamma\mathbb{E}_{P(s^{\prime}\mid
s,a)}[f(s^{\prime},\pi)]-f(s,a)]]$

Then,

$\displaystyle V^{\pi}_{P^{\prime},r+b}-V^{\pi}_{P,r}$
$\displaystyle=\frac{1}{1-\gamma}\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\prime}}}[r(s,a)+b(s,a)+\gamma\mathbb{E}_{P^{\prime}(s^{\prime}\mid
s,a)}[Q^{\pi}_{P,r}(s^{\prime},\pi)]-Q^{\pi}_{P,r}(s,a)]]$
$\displaystyle=\frac{1}{1-\gamma}\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\prime}}}[b(s,a)+\gamma\mathbb{E}_{P^{\prime}(s^{\prime}\mid
s,a)}[Q^{\pi}_{P,r}(s^{\prime},\pi)]-\gamma\mathbb{E}_{P(s^{\prime}\mid
s,a)}[Q^{\pi}_{P,r}(s^{\prime},\pi)]].$

Similarly,

$\displaystyle V^{\pi}_{P,r}-V^{\pi}_{P^{\prime},r+b}$
$\displaystyle=\frac{1}{1-\gamma}\mathbb{E}_{(s,a)\sim
d^{\pi}_{P}}[r(s,a)+\gamma\mathbb{E}_{P(s^{\prime}\mid
s,a)}[Q^{\pi}_{P^{\prime},r+b}(s^{\prime},\pi)]-Q^{\pi}_{P^{\prime},r+b}(s,a)]]$
$\displaystyle=\frac{1}{1-\gamma}\mathbb{E}_{(s,a)\sim
d^{\pi}_{P}}[-b(s,a)+\gamma\mathbb{E}_{P(s^{\prime}\mid
s,a)}[Q^{\pi}_{P^{\prime},r+b}(s^{\prime},\pi)]-\gamma\mathbb{E}_{P^{\prime}(s^{\prime}\mid
s,a)}[Q^{\pi}_{P,r}(s^{\prime},\pi)]].$

∎

The following lemma is used to deal with the distribution shift in the offline
setting. For the proof, refer to Chang et al. (2021).

######  Lemma 21 (Distribution shift lemma).

Consider any policy $\pi$ and state-action distribution $\rho$, and any
representation $\phi^{\star}$, we have:

$\displaystyle\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\star}}}[\phi^{\star}(s,a)\\{\phi^{\star}(s,a)\\}^{\top}]\leq
C^{\star}\mathbb{E}_{\rho}[\phi^{\star}(s,a)\\{\phi^{\star}(s,a)\\}^{\top}],\quad
C^{\star}\coloneqq\sup_{x\in\mathbb{R^{d}}}\frac{x^{\top}\mathbb{E}_{(s,a)\sim
d^{\pi}_{P^{\star}}}[\phi^{\star}\\{\phi^{\star}\\}^{\top}]x}{x^{\top}\mathbb{E}_{(s,a)\sim\rho}[\phi^{\star}\\{\phi^{\star}\\}^{\top}]x}.$

This is some auxiliary lemma to convert the finite sample error bound into the
sample complexity.

######  Lemma 22 (Conversion of finite sample error bounds into sample
complexities).

By taking

$N=1/\epsilon^{\prime 2}\times\ln^{2}(1+1/\epsilon^{\prime
2}),\epsilon^{\prime}=\frac{\epsilon}{a_{1}\ln^{1/2}(e+a_{2})\ln^{1/2}(e+a_{3})}.$

It satisfies

$\displaystyle
a_{1}\sqrt{1/N}\ln^{1/2}(1+a_{2}N)\ln^{1/2}(1+a_{3}N)<c\epsilon.$

where $c$ is a constant independent of $a_{1},a_{2},a_{3}$.

###### Proof.

We first have

$\displaystyle a_{1}\sqrt{1/N}\ln^{1/2}(1+a_{2}N)\ln^{1/2}(1+a_{3}N)\leq
a_{1}\max(\ln^{1/2}(1+a_{2})\ln^{1/2}(1+a_{3}),1)\sqrt{1/N}\ln(1+N).$

Here, we use

$\displaystyle\ln^{1/2}(1+a_{2}N)\leq\\{\ln(1+a_{2})+\ln(1+N)\\}^{1/2}\leq\sqrt{\max(1,\ln(1+a_{2}))\ln(1+N)}.$

Then, we prove when $N=1/\epsilon^{2}\times\ln^{2}(1+1/\epsilon^{2})$.

$\displaystyle\sqrt{1/N}\ln(1+N)<\epsilon.$

This is proved by

$\displaystyle\sqrt{1/N}\ln(1+N)$
$\displaystyle\leq\epsilon\times\frac{\ln(1+1/\epsilon^{2}\times\ln^{2}(1+1/\epsilon^{2}))}{\ln(1+1/\epsilon^{2})}$
$\displaystyle\leq\epsilon\times\frac{\ln(1+1/\epsilon^{2})+\ln(1+\ln^{2}(1+1/\epsilon^{2}))}{\ln(1+1/\epsilon^{2})}$
$\displaystyle\leq\epsilon+\epsilon\times\frac{\ln(1+\ln^{2}(1+1/\epsilon^{2}))}{\ln(1+1/\epsilon^{2})}$
$\displaystyle\leq\epsilon+\epsilon\times\frac{0.5\\{1+\ln^{2}(1+1/\epsilon^{2}))\\}^{1/2}-1}{\ln(1+1/\epsilon^{2})}$
$\displaystyle\lesssim\epsilon.$

From the third line to the fourth line, we use $\ln(x)\leq 0.5(x^{1/2}-1)$ for
$x>0$. Then, the final statement is concluded. ∎

##  Appendix D More comparison to Xie et al. (2021)

We briefly explain the guarantee when we use Algorithm 1 (Xie et al., 2021).
For a given reward $r$, we first define a new feature class $\Phi^{+}_{r}$.

######  Definition 23 (Augmented feature).

Let $\phi=[\phi_{1},\cdots,\phi_{d}]$.

$\displaystyle\Phi^{+}_{r}=\\{\phi^{+}_{r};\phi\in\Phi\\},\quad\phi^{+}_{r}=[\phi_{1},\cdots,\phi_{d},r].$

Then, we set

$\displaystyle\mathcal{F}=\\{a^{\top}\phi^{+}_{r}\mid\|a\|_{2}\leq
c\sqrt{d}+1,\phi^{+}_{r}\in\Phi^{+}_{r}\\}.$

where $c$ is some suitable constant. Given the hypothesis class $\mathcal{F}$
for the Q-function, we can run Algorithm 1 in (Xie et al., 2021). We denote
the output policy as $\hat{\pi}$.

We check two assumptions to ensure the algorithm works. The first assumption
is realizability. This is satisfied since for any policy $\pi\in\Pi$ ($\Pi$ is
the class of all Markovian polices), we have
$Q^{\pi}_{P^{\star},r}\in\mathcal{F}$. The second assumption is completeness.
This is also satisfied since
${\mathcal{T}}^{\pi}_{P^{\star},r}\mathcal{F}\subset\mathcal{F}$ for any
policy $\pi\in\Pi$ where ${\mathcal{T}}^{\pi}_{P^{\star},r}$ is a Bellman-
operator s.t.

${\mathcal{T}}^{\pi}_{P^{\star},r}:\\{{\mathcal{S}}\times\mathcal{A}\to\mathbb{R}\\}\ni
f\mapsto r(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim
P^{\star}(s,a)}[f(s^{\prime},\pi)]\in\\{{\mathcal{S}}\times\mathcal{A}\to\mathbb{R}\\},$

where we denote $f(s,\pi)=\mathbb{E}_{a\sim\pi(s)}f(s,a)$. Then, by invoking
their Corollary 5, we have

######  Theorem 24 (PAC bound based on Xie et al. (2021)).

With probability $1-\delta$,

$\displaystyle\forall\pi\in\Pi:V^{\pi}_{P^{\star},r}-V^{\hat{\pi}}_{P^{\star},r}\leq
c\frac{\sqrt{C^{\dagger}_{\pi,r}}}{(1-\gamma)^{2}}\left(\frac{(d+1)\log(1/\delta)\log|\mathcal{A}|}{n}\right)^{1/5}.$

where

$\displaystyle
C^{\dagger}_{\pi,r}=\sup_{\phi^{+}_{r}\in\Phi^{+}_{r}}\sup_{a\in\mathbb{R}^{d+1}}\frac{a^{\top}\mathbb{E}_{d^{\pi}_{P^{\star}}}[\phi^{+}_{r}\\{\phi^{+}_{r}\\}^{\top}]a}{a^{\top}\mathbb{E}_{\rho}[\phi^{+}_{r}\\{\phi^{+}_{r}\\}^{\top}]a}.$

We compare the above result with our result in Theorem 6. First, since
$C^{\dagger}_{r}$ includes $r$ and all possible features in $\Phi$, this
partial coverage condition is stronger than ours (recall our partial coverage
condition is only related to the true representation $\phi^{\star}$), and we
always have $C^{\star}_{\pi}\leq C^{\dagger}_{\pi,r}$. Secondly, the
dependence on $n$ is much worse. Third, it is unclear whether the learned
policy can compete against any history-dependent policy. Recall in Theorem 6,
we show that our algorithm can compete with any history-dependent policies.