# Linear Bellman Completeness Suffices for Efficient Online Reinforcement
Learning with Few Actions

Noah Golowich
MIT Email<EMAIL_ADDRESS>Supported by a Fannie & John Hertz Foundation
Fellowship and an NSF Graduate Fellowship. Ankur Moitra
MIT Email<EMAIL_ADDRESS>Supported in part by a Microsoft Trustworthy AI
Grant, an ONR grant and a David and Lucile Packard Fellowship.

(June 17, 2024)

###### Abstract

One of the most natural approaches to reinforcement learning (RL) with
function approximation is _value iteration_ , which inductively generates
approximations to the optimal value function by solving a sequence of
regression problems. To ensure the success of value iteration, it is typically
assumed that _Bellman completeness_ holds, which ensures that these regression
problems are well-specified. We study the problem of learning an optimal
policy under Bellman completeness in the _online_ model of RL with linear
function approximation. In the linear setting, while statistically efficient
algorithms are known under Bellman completeness (e.g., [JKA+17, ZLKB20a]),
these algorithms all rely on the principle of _global optimism_ which requires
solving a nonconvex optimization problem. In particular, it has remained open
as to whether _computationally efficient_ algorithms exist. In this paper we
give the first polynomial-time algorithm for RL under linear Bellman
completeness when the number of actions is any constant.

##  1 Introduction

_Reinforcement learning (RL)_ describes the problem of solving sequential
decision-making problems in a stochastically changing environment, and is
typically studied using the formalism of _Markov Decision Processes (MDPs)_.
In an MDP, a learning agent must choose a sequence of _actions_ over some
number of time steps, each of which affects the _state_ of the environment and
potentially yields some _reward_ to the agent. The agent aims to find a
_policy_ , or a mapping that describes which action to take at each state,
that maximizes its expected total reward. In order for RL to yield effective
learning strategies in its various application domains, including robotics
[GHLL17], economics [ZTS+22], and healthcare [YLNY21], it is necessary to come
up with efficient strategies for exploring complex state spaces, which may be
infinite or exponentially large.

A general approach to RL, which dates back decades [BB96, MR07, SB18] and yet
still forms the basis for many current empirical approaches [HGS16, SWD+17],
involves the use of value function approximation. Recall that the optimal
_value function_ maps a state-action pair to the agent’s expected reward under
the optimal policy starting from that state-action pair. Then this approach
posits that the optimal value function belongs to some known function class
$\mathcal{F}$, such as a class of linear functions or neural networks. A key
question is: _under what assumptions on the class $\mathcal{F}$ can we
efficiently learn a near-optimal policy, i.e., one with near-maximal expected
reward?_

##### Value iteration and Bellman completeness.

A popular and time-tested approach to finding a near-optimal policy with value
function approximation is _value iteration_. To explain this procedure, we
consider the _finite-horizon_ setting, whereby interactions with the
environment occur in _episodes_ lasting $H$ time steps. Letting $\mathcal{X}$
denote the state space and $\mathcal{A}$ denote the action space, value
iteration computes mappings
$\hat{Q}_{h}:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}$ in a
backwards-inductive manner, i.e., for $h=H,H-1,\ldots,1$. The values
$\hat{Q}_{h}(x,a)$ should be interpreted as estimates of the optimal
value111I.e., the value of the optimal policy. at step $h$ given that the
state-action pair taken at step $h$ is $(x,a)$. Since the environment’s reward
functions and transitions are unknown, $\hat{Q}_{h}$ must be estimated
empirically. Given a dataset $\mathcal{D}$ consisting of tuples
$(x_{h},a_{h},r_{h},x_{h+1})$ of states, actions, and rewards drawn from the
environment at step $h$, together with the subsequent state at step $h+1$,
$\hat{Q}_{h}$ is typically chosen to be the function in $\mathcal{F}$ which
minimizes the following square-loss objective whose labels are defined in
terms of $\hat{Q}_{h+1}$:

$\displaystyle\hat{Q}_{h}:=\operatorname*{arg\,min}_{Q_{h}\in\mathcal{F}}\sum_{(x_{h},a_{h},r_{h},x_{h+1})\in\mathcal{D}}\left(Q_{h}(x_{h},a_{h})-\left(r_{h}+\max_{a^{\prime}}\hat{Q}_{h+1}(x_{h+1},a^{\prime})\right)\right)^{2}.$
(1)

The procedure defined by Eq. 1 is often known as _least-squares value
iteration_ (LSVI) [BB96, OVRW16]. A natural condition under which LSVI might
yield good estimates of the optimal value function, and thereby a near-optimal
policy, is that the least-squares problem Eq. 1 be _well-specified_. This
requires that for any $\hat{Q}_{h+1}\in\mathcal{F}$, there is some
$Q_{h}^{\prime}\in\mathcal{F}$ so that, for all
$x_{h}\in\mathcal{X},a_{h}\in\mathcal{A}$,

$\displaystyle
Q_{h}^{\prime}(x_{h},a_{h})=\mathbb{E}\left[r_{h}+\max_{a^{\prime}}\hat{Q}_{h+1}(x_{h+1},a^{\prime})\
\mid\ (x_{h},a_{h})\right].$ (2)

Eq. 2 is known to be necessary for LSVI to succeed, in the sense that without
it, the value functions computed by LSVI may be wildly divergent [TvR96].
Moreover, classical results [Mun05, MS08] showed that if the dataset
$\mathcal{D}$ is sufficiently exploratory, then Eq. 2 is sufficient for LSVI
to succeed. The requirement that Eq. 2 holds for any
$\hat{Q}_{h+1}\in\mathcal{F}$ is often known as _Bellman completeness_ ; it is
a property of the MDP and the value function class $\mathcal{F}$.

In this paper, our focus is on finding computationally efficient algorithms.
Since regression problems such as Eq. 1 can be computationally intractable for
even relatively simple nonlinear classes $\mathcal{F}$ such as shallow neural
networks [DWX20, BJW19, GKMR20], we focus on the case where $\mathcal{F}$ is
simply the class of linear functions in $\phi_{h}(x,a)$, for some known
_feature mappings_
$\phi_{h}:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}^{d}$. In this
setting, the requirement of Bellman completeness in Eq. 2 is known as _linear
Bellman completeness_ (see Definition 2.1 for a formal definition). In the
remainder of this section we discuss our results on computationally efficient
learning of MDPs satisfying linear Bellman completeness.

###  1.1 Exploration and linear Bellman completeness

The results of [Mun05, MS08] referenced above showing sufficiency of Bellman
completeness do not address the following central problem in RL: _how can we
find a dataset $\mathcal{D}$ which is sufficiently exploratory to compute a
near-optimal policy_? More precisely, we would like that for tuples
$(x_{h},a_{h},r_{h},x_{h+1})\in\mathcal{D}$, the feature vectors
$\phi_{h}(x_{h},a_{h})$ span a “sufficient number” of distinct directions in
$\mathbb{R}^{d}$. We adopt the standard online setting in RL, allowing the
learning algorithm to form $\mathcal{D}$ by repeatedly sampling trajectories
from the MDP using adaptively chosen policies.222See Section 2.1 for a formal
definition. We first note that this problem is known to be statistically
tractable using a technique known as _global optimism_ (see Section 3.1). This
technique has various instantiations as specific algorithms in the setting of
linear Bellman completeness, including ELEANOR [ZLKB20a], GOLF [JLM21], OLIVE
[JKA+17], and BilinUCB [DKL+21]. While these algorithms require only
polynomially many samples (i.e., rounds of interaction) to output a near-
optimal policy, all are computationally inefficient, even when specialized to
the linear case, since they require solving nonconvex optimization problems.
Our goal is to find a computationally efficient algorithm which achieves the
same guarantee:

######  Question 1.1.

Is there an algorithm which learns an $\epsilon$-optimal policy in an unknown
linear Bellman complete MDP using
$\operatorname{poly}(H,d,|\mathcal{A}|,\epsilon^{-1})$ samples and time?

A sizeable portion of the work on computationally efficient RL in the last
several years has been focused on answering Question 1.1 for settings which
are strict special cases of linear Bellman completeness. The simplest such
setting is the _tabular setting_ , which describes the case that
$|\mathcal{X}|,|\mathcal{A}|$ are finite and the goal is to obtain sample and
computational complexities scaling as
$\operatorname{poly}(H,|\mathcal{X}|,|\mathcal{A}|)$. In this setting, there
are several computationally efficient algorithms which can be viewed as
variants of value iteration that are _optimistic_ in the sense that they add
bonuses to the rewards to induce exploration: these include UCBVI [AOM17] and
Q-learning-UCB [JAZBJ18, ZZJ20], which are known to obtain near-optimal rates.
Tabular MDPs are generalized by the _linear MDP_ setting, in which feature
vectors $\phi_{h}(x,a)\in\mathbb{R}^{d}$ are given, and the state-action
transition probabilities are assumed to be linear in $\phi_{h}$. Here too
there are computationally efficient algorithms, namely LSVI-UCB [JYWJ20], an
optimistic version of LSVI, as well as more recent rate-optimal variants
[AJZ22, HZZG23].

In the setting of linear Bellman completeness, which is a strict
generalization of linear MDPs [ZLKB20a, Proposition 3], the algorithm FRANCIS
of [ZLKB20b] is computationally efficient and learns a near-optimal policy in
the special case that the MDP is _reachable_ , meaning that any direction in
$\mathbb{R}^{d}$ can be reached under some policy. The setting of linear
Bellman completeness has also been studied in the special case of
_deterministic_ dynamics: [WR16] (Theorem 1) showed that there is a
polynomial-time learning algorithm when the transitions and rewards are
deterministic, and [DLMW20] established a polynomial-time algorithm for the
setting of deterministic transitions, stochastic rewards, and positive
suboptimality gap. Finally, in recent and concurrent work, [WSKS24]
established a polynomial-time algorithm for the setting of deterministic
transitions and stochastic rewards (without the assumption of positive
suboptimality gap). The algorithm of [WSKS24] has the additional advantage of
working under weaker norm assumptions (namely, the second item of Assumption
2.1 is not required). Finally, we mention that in the related _offline_
setting of reinforcement learning, the analogue of Question 1.1 was recently
resolved in concurrent work [GM24].

Despite the above line of work, Question 1.1 in its full generality has
remained open. Part of the reason for this is that the two principal
techniques to perform computationally efficient exploration in RL both break
down in the general setting of linear Bellman completeness:

  1. 1.

The first technique is _local optimism_ ,333“Local” is used to distinguish
this technique from _global optimism_ , which, as discussed above, works in
the setting of linear Bellman completeness but is computationally inefficient.
which adds exploration bonuses to the reward at each state which scale
inversely with how often the state has been visited. It includes the UCBVI and
Q-learning-UCB algorithms for the tabular setting, and the LSVI-UCB algorithm
for the linear MDP setting, among others. Local optimism requires that the
value function class be complete with respect to the exploration bonuses,
which is satisfied for linear MDPs but which fails more generally (see Section
3).

  2. 2.

The second technique is to construct a _policy cover_ , which is a small set
of policies which, roughly speaking, covers all states to the maximum extent
possible. This technique includes the FRANCIS algorithm, as well as
computationally efficient and oracle efficient algorithms for related tasks in
RL such as representation learning [DKJ+19, MHKL20, MFR23, MBFR23, GMR24] and
learning in POMDPs [GMR22]. In order for this approach to work in the absence
of reachability, it is necessary to analyze a _truncated_ version of the true
MDP. But doing so seems impossible in our setting, since truncating the MDP
breaks the property of linear Bellman completeness.

Thus, in addition to generalizing a long line of work on computationally
efficient learning of MDPs, Question 1.1 captures exactly the point where
known exploration paradigms in RL break down.

####  1.1.1 Main result for the online setting

Our main result is a positive answer to Question 1.1 in the case that
$|\mathcal{A}|$ is any constant:

######  Theorem 1.1 (Informal version of Theorem 5.10).

Suppose the ground-truth MDP satisfies linear Bellman completeness. Then for
any $\epsilon>0$, there is an algorithm (PSDP-UCB; Algorithm 1) which with
high probability learns an $\epsilon$-optimal policy using
$(Hd|\mathcal{A}|\epsilon^{-1})^{O(|\mathcal{A}|)}$ samples and time.

We remark that the exponential dependence of the sample and computational
complexities on $|\mathcal{A}|$ is somewhat unusual in RL. We are not aware of
any prior work on an RL problem for which $\operatorname{poly}(|\mathcal{A}|)$
dependence is possible but for which a “natural” algorithm for even the
$|\mathcal{A}|=2$ case does not extend to the general case. Theorem 1.1 thus
presents an intriguing challenge: to either improve the guarantee to fully
answer Question 1.1 in the affirmative, or to find a lower bound.

##### Proof idea of Theorem 1.1.

The proof of Theorem 1.1 proceeds via the local optimism approach discussed in
Item 1 above. Our main insight is that it is possible to design exploration
bonuses for which the class of linear value functions _is_ complete with
respect to them, i.e., whose Bellman-backup is linear (formally, we say that
such exploration bonuses are _Bellman-linear_ ; see Definition 3.1). Doing so
requires extensive work to “build up” to a function which satisfies several
properties required of the exploration bonuses; see Section 3 for further
details. We remark that our algorithm PSDP-UCB takes the structure of an
optimistic variant of the classical PSDP (Policy Search by Dynamic
Programming) algorithm [BKSN03], which has not previously appeared in the
literature. This choice of algorithm structure plays more nicely with the
exploration bonuses we introduce.

We find the fact that local optimism provides a solution to Question 1.1, even
in the case $|\mathcal{A}|=O(1)$, to be somewhat surprising. On the one hand,
it indicates that Bellman completeness is a stronger property than one might
initially anticipate. On the other hand, since Bellman-completeness is
satisfied in other settings for which the only known computationally efficient
algorithms are policy-cover based (Item 2 above), our results perhaps provide
some hope that local optimism can be made to work in such settings as well.

##### Organization of the paper.

We discuss preliminaries in Section 2, and overview the proof of Theorem 1.1
in Section 3, which is proved formally in Sections 4 and 5.

##  2 Preliminaries

A finite-horizon _Markov Decision Process (MDP)_ is given by a tuple
$M=(H,\mathcal{X},\mathcal{A},(P_{h}^{\scriptscriptstyle{M}})_{h=1}^{H},(r_{h}^{\scriptscriptstyle{M}})_{h=1}^{H},d_{1}^{\scriptscriptstyle{M}})$,
where $H\in\mathbb{N}$, $\mathcal{X}$ is a measure space denoting the state
set, $\mathcal{A}$ denotes the action set,
$P_{h}^{\scriptscriptstyle{M}}(\cdot|x,a)\in\Delta(\mathcal{X})$ (for
$h\in[H]$) denotes the probability transition kernels,
$r_{h}^{\scriptscriptstyle{M}}:\mathcal{X}\times\mathcal{A}\rightarrow[0,1]$
(for $h\in[H]$) denotes the reward functions, and
$d_{1}^{\scriptscriptstyle{M}}\in\Delta(\mathcal{X})$ denotes the initial
state distribution. When the MDP $M$ is clear from context, we will drop the
superscript $M$ in these notations. We let $A:=|\mathcal{A}|$ denote the
number of actions in the MDP.

A _policy_ $\pi$ consists of a tuple $\pi=(\pi_{1},\ldots,\pi_{H})$, where
each $\pi_{h}:\mathcal{X}\rightarrow\Delta(\mathcal{A})$ is a mapping from
states to distributions over actions. Note that we allow policies to be
nonstationary and randomized; $\Pi$ denotes the set of all such policies. A
policy $\pi\in\Pi$ defines a distribution over _trajectories_
$(x_{1},a_{1},r_{1},\ldots,x_{H},a_{H},r_{H})\in(\mathcal{X}\times\mathcal{A}\times[0,1]^{H}$,
as follows: first, $x_{1}\sim d_{1}$, and then for each $h\in[H]$, we draw
$a_{h}\sim\pi_{h}(x_{h})$, observe $r_{h}(x_{h},a_{h})$, and transition to
$x_{h+1}\sim P_{h}(\cdot|x_{h},a_{h})$. Let $\mathscr{H}_{H}$ denote the set
of trajectories. For a function $f:\mathscr{H}_{H}\rightarrow\mathbb{R}^{k}$,
we often write
$\mathbb{E}^{\scriptscriptstyle{M,\pi}}[f(x_{1},a_{1},r_{1},\ldots,x_{H},a_{H},r_{H})]$
to denote the expectation of $f$ over trajectories drawn from $M$ under policy
$\pi$. If $M$ is clear from context, we will simply write
$\mathbb{E}^{\pi}[f(x_{1},a_{1},r_{1},\ldots,x_{H},a_{H},r_{H})]$. Given
policies
$\pi=(\pi_{1},\ldots,\pi_{H}),\pi^{\prime}=(\pi_{1}^{\prime},\ldots,\pi_{H}^{\prime})\in\Pi$
and a step $h\in[H]$, we let $\pi\circ_{h}\pi^{\prime}\in\Pi$ denote the
policy $(\pi_{1},\ldots,\pi_{h-1},\pi_{h}^{\prime},\ldots,\pi_{H}^{\prime})$,
i.e., which acts according to $\pi$ during the first $h-1$ steps and
thereafter acts according to $\pi^{\prime}$.

For a function $f:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}$, a
(randomized) policy $\pi$ and $h\in[H]$, we write
$f(x,\pi_{h}(x)):=\mathbb{E}_{a\sim\pi_{h}(x)}[f(x,a)]$. The state-action
value function (or _$Q$ -function_) and state-value function (or _$V$
-function_) of a policy $\pi\in\Pi$ are then defined as follows: for
$h\in[H],x\in\mathcal{X},a\in\mathcal{A}$,

$\displaystyle
Q_{h}^{\pi}(x,a):=r_{h}(x,a)+\mathbb{E}^{\pi}\left[\sum_{g=h+1}^{H}r_{g}(x_{g},a_{g})\
\mid\ (x_{h},a_{h})=(x,a)\right],\qquad
V_{h}^{\pi}(x):=Q_{h}^{\pi}(x,\pi_{h}(x)).$

The _optimal policy_ is defined as
$\pi^{\star}:=\operatorname*{arg\,max}_{\pi\in\Pi}\mathbb{E}[V_{1}^{\pi}(x_{1})]$
(where expectation is over $x_{1}\sim d_{1}$). It is known that there is
always a deterministic optimal policy $\pi^{\star}$ (i.e., so that
$\pi_{h}^{\star}(x)$ is a singleton for all $x,h$). We often abbreviate
$Q_{h}^{\star}(x,a):=Q_{h}^{\pi^{\star}}(x,a)$ and
$V_{h}^{\star}(x):=V_{h}^{\pi^{\star}}(x)$.

###  2.1 The online learning problem

We assume that the transitions, rewards, and initial state distribution of the
ground-truth MDP $M$ are _unknown_ to the algorithm. To learn information
about the MDP, the algorithm interacts with $M$ via the _episodic online
learning model_ , as follows. The interaction proceeds over a series of $T$
episodes. In each step $h\in[H]$ of each episode, the algorithm observes the
current state $x_{h}$, specifies an action $a_{h}\in\mathcal{A}$ to take, and
observes a reward of $r_{h}=r_{h}(x_{h},a_{h})$. Then the environment
transitions to a new state $x_{h+1}\sim P_{h}(\cdot|x_{h},a_{h})$. We assume
that the algorithm can query $\phi_{h}(x,a)\in\mathbb{R}^{d}$ for any $x,a,h$.
We remark that our algorithm only needs to query $\phi_{h}(x,a)$ for states
$x$ which are visited at some point in some episode.

The goal is as follows: for $\epsilon,\delta\in(0,1)$, to give an algorithm
which interacts with the environment for $T=T(\epsilon,\delta)$ episodes in
the manner described above, and then to output a policy $\hat{\pi}$ so that,
with probability $1-\delta$,
$\mathbb{E}[V_{1}^{\star}(x_{1})-V_{1}^{\hat{\pi}}(x_{1})]\leq\epsilon$.
Moreover, we wish the algorithm to be computationally efficient (both in terms
of the computation required when interacting with the environment over the
course of $T$ episodes and when outputting $\hat{\pi}$ at termination).

###  2.2 Function approximation

In order for the above learning problem to be tractable, it is necessary to
make some assumptions on the ground-truth MDP $M$ being learned. While it is
well-known that boundedness of $|\mathcal{X}|,|\mathcal{A}|$ implies efficient
learning algorithms whose computational and statistical costs scale with
$\operatorname{poly}(|\mathcal{X}|,|\mathcal{A}|)$ [AOM17, JAZBJ18], realistic
RL environments typically have enormous state spaces. Thus we aim for weaker
function approximation assumptions: a common and longstanding such assumption
[BB96, MR07, SB18, YW20, JYWJ20] is _linearity_ of the value functions, with
respect to some known features. In particular, given $d\in\mathbb{N}$, we
assume that functions
$\phi_{h}:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}^{d}$ are given for
all $h\in[H]$, mapping each state-action pair $(x,a)$ to a collection of $d$
_features_ which should be interpreted as capturing all relevant information
about $(x,a)$.

The weakest assumption on value function linearity is simply that
$Q_{h}^{\star}$ is linear, i.e., for some $w_{h}^{\star}\in\mathbb{R}^{d}$, we
have $Q_{h}^{\star}(x,a)=\langle w_{h}^{\star},\phi_{h}(x,a)\rangle$ for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$. Unfortunately, unless
$\mathsf{NP}=\mathsf{RP}$, it is not possible to computationally efficiently
learn a near-optimal policy under this assumption [KLLM22, LMK+23].
Accordingly, we make a stronger assumption known as _linear Bellman
completeness_ , which states, roughly speaking, that the Bellman backup of all
linear functions is linear. To formally state this assumption (in Definition
2.1 below), we need to introduce the following notation: for $h\in[H]$, define

$\displaystyle\mathcal{B}_{h}:=\\{\theta\in\mathbb{R}^{d}\ :\
|\langle\phi_{h}(x,a),\theta\rangle|\leq 1\ \
\forall(x,a)\in\mathcal{X}\times\mathcal{A}\\}.$

In words, $\mathcal{B}_{h}$ denotes the set of coefficient vectors inducing
bounded (linear) functions on $\mathcal{X}\times\mathcal{A}$.

######  Definition 2.1 (Linear Bellman Completeness).

The MDP $M$ is defined to be _linear Bellman complete_ with respect to the
feature mappings $(\phi_{h})_{h\in[H]}$ if, for each $h\in[H]$, there is a
mapping $\mathcal{T}_{h}:\mathcal{B}_{h+1}\rightarrow\mathcal{B}_{h}$ so that,
for all $\theta\in\mathcal{B}_{h+1}$ and all
$(x,a)\in\mathcal{X}\times\mathcal{A}$,

$\displaystyle\langle\phi_{h}(x,a),\mathcal{T}_{h}\theta\rangle=\mathbb{E}_{x^{\prime}\sim
P_{h}(x,a)}\left[\max_{a^{\prime}\in\mathcal{A}}\langle\phi_{h+1}(x^{\prime},a^{\prime}),\theta\rangle\right].$

Moreover, we require that for all $h\in[H]$,
$(x,a)\in\mathcal{X}\times\mathcal{A}$, the reward at $(x,a)$ at step $h$ is
given by:

$\displaystyle
r_{h}(x,a):=\langle\phi_{h}(x,a),\theta_{h}^{\mathsf{r}}\rangle,$

for some vectors $\theta_{h}^{\mathsf{r}}\in\mathcal{B}_{h}$.

It is immediate from Definition 2.1 that if $M$ is linear Bellman complete,
then there are vectors $w_{h}^{\star}\in H\cdot\mathcal{B}_{h}$ so that
$Q_{h}^{\star}(x,a)=\langle w_{h}^{\star},\phi_{h}(x,a)\rangle$ for all $x,a$.
In addition to linear Bellman completeness, we make the following standard
boundedness assumptions on the coefficient and feature vectors:

######  Assumption 2.1 (Boundedness).

We assume the following:

  1. 1.

For all $h\in[H],x\in\mathcal{X},a\in\mathcal{A}$, we have
$\|\phi_{h}(x,a)\|_{2}\leq 1$.

  2. 2.

For some parameter $B\in\mathbb{R}_{+}$: for all $w_{h}\in\mathcal{B}_{h}$, it
holds that $\|w_{h}\|_{2}\leq B$.

  3. 3.

For all $h\in[H]$, $\|\theta_{h}^{\mathsf{r}}\|_{2}\leq 1$ (and hence
$\sup_{x,a,h}|r_{h}(x,a)|\leq 1$).

##### Polytope of actions.

We assume that $\mathcal{A}$ is finite and write $A:=|\mathcal{A}|$. Let
$\mathscr{P}^{d}$ denote the space of $d$-dimensional (bounded) polyhedra. For
$A\in\mathbb{N}$, let $\mathscr{P}_{A}^{d}\subset\mathscr{P}^{d}$ denote the
space of $d$-dimensional polyhedra $\Phi$ so that $\dim({\rm Span}\Phi)\leq
A$. For each $x\in\mathcal{X}$, we define $\Phi_{h}(x):=\\{\phi_{h}(x,a):\
a\in\mathcal{A}\\}\subset\mathbb{R}^{d}$ and
$\bar{\Phi}_{h}(x):=\mathrm{co}(\Phi_{h}(x))\in\mathscr{P}_{A}^{d}\subset\mathscr{P}^{d}$.

##### Gaussian smoothing.

For $\theta\in\mathbb{R}^{d}$ and $\sigma>0$, we write

$\displaystyle\mathcal{N}_{\sigma}(\theta):=\mathcal{N}(0,\sigma^{2}\cdot
I_{d})(\theta)=\frac{1}{(2\pi)^{d/2}\sigma^{d}}\cdot\exp\left(-\frac{1}{2\sigma^{2}}\|\theta\|_{2}^{2}\right)$

to denote the probability density function of the standard normal distribution
with covariance $\sigma^{2}I_{d}$. Furthermore, for
$f:\mathbb{R}^{d}\rightarrow\mathbb{R}$, we write
$\mathsf{S}_{\sigma}f(\theta)$ to denote the convolution of $f$ with
$\mathcal{N}_{\sigma}$, namely

$\displaystyle\mathsf{S}_{\sigma}f(\theta):=\int_{\mathbb{R}^{d}}f(z)\mathcal{N}_{\sigma}(\theta-z)dz=\int_{\mathbb{R}^{d}}f(\theta-z)\mathcal{N}_{\sigma}(z)dz=\mathbb{E}_{z\sim\mathcal{N}(0,\sigma^{2}\cdot
I_{d})}[f(\theta-z)].$

##### Miscellaneous notation.

Given $d\in\mathbb{N}$, we let $\mathbb{S}^{d}$ denote the space of symmetric
$d\times d$ matrices, $\mathbb{S}_{+}^{d}\subset\mathbb{S}^{d}$ denote the
space of positive semidefinite (PSD) matrices, and
$\mathbb{S}_{++}^{d}\subset\mathbb{S}_{+}^{d}$ denote the space of positive
definite matrices. For a PSD matrix $T\in\mathbb{S}_{+}^{d}$, we let $T^{1/2}$
denote the unique PSD matrix whose square is $T$. Given a subset
$\mathcal{S}\subset\mathbb{R}^{d}$ and a matrix $T\in\mathbb{R}^{d\times d}$,
we let $T\cdot\mathcal{S}:=\\{Tv:\ v\in\mathcal{S}\\}$. For vectors
$v,v^{\prime}\in\mathbb{R}^{d}$, we let $[v,v^{\prime}]$ denote the segment
$[v,v^{\prime}]:=\\{\lambda v+(1-\lambda)v^{\prime}:\ \lambda\in[0,1]\\}$. Let
$S^{d-1}=\\{u\in\mathbb{R}^{d}:\ \|u\|_{2}=1\\}$ denote the $d$-dimensional
unit sphere. For a square matrix $T$, we let $\|T\|$ denote its spectral norm.

###  2.3 Linear policies

Under the assumption of linear Bellman completeness, there is always an
optimal policy with the additional structure of being a _linear policy_.
Roughly speaking, linear policies take an action defined by the argmax with
respect to some fixed coefficient vector: in particular, given
$w\in\mathbb{R}^{d}$ and $h\in[H]$, the associated linear policy at step $h$
is $\pi_{h,w}(x):=\operatorname*{arg\,max}_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle$ _if the argmax is unique_. It requires some care to
appropriately break ties for states $x$ at which the
$\operatorname*{arg\,max}$ is not unique.444In particular, an appropriate tie-
breaking procedure is necessary for Lemma 4.3 to hold. To do so, given
$w\in\mathbb{R}^{d},h\in[H],x\in\mathcal{X}$, define
$\mathcal{A}_{h,w}(x):=\operatorname*{arg\,max}_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\subset\mathcal{A}$, where $\operatorname*{arg\,max}$ is
interpreted as the set of all actions maximizing $\langle
w,\phi_{h}(x,a)\rangle$. Then we set, for each $a\in\mathcal{A}_{h,w}(x)$,

$\displaystyle\mathcal{G}_{h,w}(x,a):=\left\\{\theta\in\mathbb{R}^{d}:\
\|\theta\|_{2}=1,\
\langle\theta,\phi_{h}(x,a)>\max_{a^{\prime}\in\mathcal{A}_{h,w}(x)\backslash\\{a\\}}\langle\theta,\phi_{h}(x,a^{\prime})\right\\}.$

Let $\nu_{d}$ denote the spherical measure on $S^{d-1}$. It is straightforward
to see that, for all $w$ and $x$,
$S^{d-1}\backslash\bigcup_{a\in\mathcal{A}_{h,w}(x)}\mathcal{G}_{h,w}(x,a)$
has measure 0 with respect to $\nu_{d}$. We now define,
$\pi_{h,w}:x\rightarrow\Delta(\mathcal{A})$ to be the following randomized
policy: for all $w\in\mathbb{R}^{d}$, $h\in[H]$, $x\in\mathcal{X}$,
$a\in\mathcal{A}$,

$\displaystyle\pi_{h,w}(a|x):=\mathbbm{1}\\{{a\in\mathcal{A}_{h,w}(x)}\\}\cdot\nu_{d}(\mathcal{G}_{h,w}(x,a)).$
(3)

It is straightfroward to see that $\pi_{h,w}(\cdot|x)\in\Delta(\mathcal{A})$.
We say that a policy $\pi$ is a _linear policy_ if it is of the form
$\pi=(\pi_{1,w_{1}},\ldots,\pi_{H,w_{H}})$, for
$w_{1},\ldots,w_{H}\in\mathbb{R}^{d}$. We let
$\Pi^{\mathsf{lin}}_{h}:=\\{\pi_{h,w}:\ w\in\mathbb{R}^{d}\\}$ denote the
space of linear policies at step $h$, and
$\Pi^{\mathsf{lin}}:=\\{(\pi_{1},\ldots,\pi_{H}):\
\pi_{h}\in\Pi^{\mathsf{lin}}_{h}\\}$ denote the space of linear policies.

##  3 Technical overview

In this section, we overview the proof of Theorem 1.1 (stated formally in
Theorem 5.10), which shows how, in the online setting (Section 2.1), Algorithm
1 can efficiently learn a near-optimal policy for an unknown MDP which is
linear Bellman complete. To simplify our notation, we assume in this section
that the parameter $B$ in Assumption 2.1 is bounded by $B\leq O(1)$. The bulk
of the challenge is to perform the task of _exploration_ : how can we interact
with the environment so as to reach state-action pairs $(x,a)$ for which
$\phi_{h}(x,a)$ points in a new direction?

###  3.1 Prior work: exploration via optimism

A popular approach to exploration in RL involves the use of _optimism_ , which
describes, loosely speaking, the technique of perturbing the algorithm’s
estimates of the MDP’s value function, typically to _increase_ the estimated
values, so as to induce the algorithm to visit new directions in feature
space. Two distinct flavors of optimism have emerged in the literature: the
first, _global optimism_ ,555See, e.g., [ZLKB20a], which used global optimism
to computationally inefficiently learn a near-optimal policy under linear
Bellman completeness; many other papers, including [JLM21, JKA+17, DKL+21],
use global optimism in settings with more general nonlinear function
approximation. constructs a confidence set consisting of all possible vectors
$w=(w_{1},\ldots,w_{H})$ which could be consistent with the optimal value
function given the data observed so far. The global optimism procedure then
chooses some $\bar{w}$ in this confidence set which maximizes
$\mathbb{E}[\max_{a}\langle\bar{w}_{1},\phi_{1}(x_{1},a)\rangle]$. It then
executes the linear policy defined by $\bar{w}$, namely the policy
$(\pi_{1,\bar{w}_{1}},\ldots,\pi_{H,\bar{w}_{H}})$, uses the resulting data to
update the confidence sets, and repeats. Unfortunately, the optimization
problem of finding a maximizing $\bar{w}$ is nonconvex, and seems unlikely to
have an efficient algorithm.

The second type of optimism-based exploration technique in RL is a more
_local_ approach: at each episode $t$, for some function
$B_{h}^{(t)}:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}_{\geq 0}$, an
_exploration bonus_ of $B_{h}^{(t)}(x,a)$ is added to the reward $r_{h}(x,a)$
for the pair $(x,a)$ at step $h$. Then, the algorithm uses the data gathered
prior to episode $t$ to estimate a policy $\pi^{(t)}$ which, roughly speaking,
maximizes the expected sum of rewards $r_{h}(x_{h},a_{h})$ and bonuses
$B_{h}^{(t)}(x_{h},a_{h})$ over a trajectory. This technique has proved
successful for computationally efficient learning of tabular MDPs (i.e., the
setting where $|\mathcal{X}|,|\mathcal{A}|$ are finite) [AOM17, JAZBJ18], as
well as _linear MDPs_ [JYWJ20], which constitute a generalization of tabular
MDPs and a strict subclass of linear Bellman complete MDPs.

Unfortunately, the local optimism approach fails in the more general setting
of linear Bellman completeness. At a high level, this failure of local
optimism results from the fact that the exploration bonuses $B_{h}^{(t)}(x,a)$
may not be linear. To illustrate, we consider the LSVI-UCB approach of
[JYWJ20]. This approach computes optimistic $Q$-function estimates,
$\hat{Q}_{h}^{(t)}$, in a backwards-inductive manner. In particular, it
defines:

$\displaystyle\hat{Q}_{h}^{(t)}(x,a):=\min\\{\langle\hat{w}_{h}^{(t)},\phi_{h}(x,a)\rangle+B_{h}^{(t)}(x,a),H\\},$
(4)

where $\hat{w}_{h}^{(t)}$ is the solution to a least-squares objective
function whose labels are given by
$\max_{a\in\mathcal{A}}\hat{Q}_{h+1}^{(t)}(x^{i},a)$, for various states
$x^{i}$.666The choice of the scalar $H$ in the minimum in Eq. 4 results from
the fact that the $Q$-function of any policy is bounded above in absolute
value by $H$, as the reward at each step has absolute value at most 1. The
bonus $B_{h}^{(t)}(x,a)$ is defined as follows: let $\Sigma_{h}^{(t)}$ be the
covariance matrix of features at step $h$ observed prior to episode $t$. Then
$B_{h}^{(t)}$ is given by a scaling of the _quadratic bonus_ , i.e., for some
scalar $\beta_{h}$, $B_{h}^{(t)}(x,a):=\beta_{h}\cdot
B^{\mathsf{quad}}_{h}(x,a;(\Sigma_{h}^{(t)})^{-1})$, where the quadratic
bonus, $B^{\mathsf{quad}}_{h}(x,a;\Sigma)$, is defined for a general PSD
matrix $\Sigma$, by

$\displaystyle
B^{\mathsf{quad}}_{h}(x,a;\Sigma):=(\phi_{h}(x,a)^{\top}\cdot\Sigma\cdot\phi_{h}(x,a))^{1/2}.$
(5)

The intuition behind the quadratic bonus is as follows:
$B^{\mathsf{quad}}_{h}(x,a;(\Sigma_{h}^{(t)})^{-1})$ will be particularly
large if $\phi_{h}(x,a)$ points in the direction of eigenvectors of
$\Sigma_{h}^{(t)}$ with small eigenvalues, i.e., directions which have not
been explored in prior episodes. Thus such _unexplored_ states receive larger
bonuses, and should be explored more during later episodes. More formally,
$B^{\mathsf{quad}}_{h}(\cdot)$ takes the same form as the standard error bound
from least-squares regression (see Lemma 5.7), meaning that by adding it to
the reward function, one can “cancel out” regression errors and thus show that
$\hat{Q}_{h}^{(t)}$ in Eq. 4 is optimistic (see [JYWJ20, AJZ22, ZSU+22,
HZZG23]).

To ensure that $\hat{w}_{h}^{(t)}$, used in the definition of
$\hat{Q}_{h}^{(t)}$, is a good estimator of the future rewards and bonuses, we
certainly need the regression problem to be _well-specified_ , i.e., the
expectation of the regression labels is a linear function in the features
$\phi_{h}(x,a)$. In particular, this approach crucially relies on the fact
that, in a linear MDP, $\mathbb{E}_{x^{\prime}\sim
P_{h}(x,a)}[\max_{a^{\prime}}\hat{Q}_{h+1}^{(t)}(x^{\prime},a^{\prime})]$, is
a linear function of $\phi_{h}(x,a)$. In fact, an even stronger statement
holds for linear MDPs: for _any_ function
$F:\mathcal{X}\rightarrow\mathbb{R}$, its _Bellman backup_ , namely
$\mathbb{E}_{x^{\prime}\sim P_{h}(x,a)}[F(x^{\prime})]$, is linear in
$\phi_{h}(x,a)$. Unfortunately, this fact fails to hold under the weaker
assumption of linear Bellman completeness. We provide an example in
Proposition A.1 for which $\mathbb{E}_{x^{\prime}\sim
P_{h}(x,a)}[\max_{a^{\prime}}\hat{Q}_{h+1}^{(t)}(x^{\prime},a^{\prime})]$ is
_not linear_ , for some $\hat{Q}_{h+1}^{(t)}$ as in Eq. 4.

Our main innovation is to show that if the bonus $B_{h}^{(t)}(x,a)$ is
carefully defined to be somewhat different from a quadratic bonus, then a
variant of Eq. 4 _does_ have a linear Bellman backup. In the remainder of this
section, we discuss in detail how to execute this strategy: in Section 3.2, we
first discuss the overall structure of our algorithm, PSDP-UCB, which is a
variant of the LSVI-UCB algorithm discussed above but which lends itself to a
simpler analysis for the setting of linear Bellman completeness. Then, in
Sections 3.3 and 3.4, we discuss how to construct exploration bonuses for use
in PSDP-UCB which do have a linear Bellman backup.

###  3.2 Overview of PSDP-UCB

Algorithm 1 PSDP-UCB$(T,H,\lambda,\beta,\lambda_{1})$: PSDP with upper
confidence bounds

1:Number of episodes $T$, horizon $H$, non-negative parameters
$\lambda,\beta,\lambda_{1}$.

2:for Round $t=1,\ldots,T$ do

3: for Step $h=H,\ldots 1$ do

4: For $i\in[n]$, draw samples
$\\{(x_{k}^{(t,i,h)},a_{k}^{(t,i,h)},r_{k}^{(t,i,h)})\\}_{k=1}^{H}$ i.i.d.
from
$\mathsf{unif}(\\{\hat{\pi}^{(s)}\circ_{h}\tilde{\pi}^{(s)}\circ_{h+1}\hat{\pi}^{(t)}\\}_{s=1}^{t-1}\\})$.

5: Set $\Sigma_{h}^{(t)}\leftarrow\lambda
I+\sum_{i=1}^{n}\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})^{\top}$.

6: For each $g>h$ and $i\in[n]$, define $\triangleright$ _$F_{g}^{(t)}$
defined in Eq. 33._
$\displaystyle\hat{r}_{g}^{(t,i,h)}:=r_{g}^{(t,i,h)}+F_{g}^{(t)}(x_{g}^{(t,i,h)}).$
(6)

7: Set
$\hat{w}_{h}^{(t)}\leftarrow(\Sigma_{h}^{(t)})^{-1}\cdot\sum_{i=1}^{n}\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})\cdot\left(r_{h}^{(t,i,h)}+\sum_{g=h+1}^{H}\hat{r}_{g}^{(t,i,h)}\right)$.

8: Define $\hat{\pi}_{h}^{(t)}$ by
$\hat{\pi}_{h}^{(t)}(x):=\operatorname*{arg\,max}_{a\in\mathcal{A}}\langle\phi_{h}(x,a),\hat{w}_{h}^{(t)}\rangle$.
_(If $t=1$, let $\hat{\pi}^{(t)}$ be arbitrary.)_

9: Define
$(\Sigma^{\prime},\Lambda^{\prime}):=\mathsf{trunc}_{\sigma_{\mathsf{tr}}}({\frac{\beta}{\lambda_{1}}\cdot(\Sigma_{h}^{(t)})^{-1/2}})$
(per Definition 5.2).

10: Define $\tilde{\pi}_{h}^{(t)}:\mathcal{X}\rightarrow\Delta(\mathcal{A})$
by, for all $(x,a)\in\mathcal{X}\times\mathcal{A}$, .
$\displaystyle\tilde{\pi}_{h}^{(t)}(x)(a):=\operatorname{Pr}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left(a=\operatorname*{arg\,max}_{a^{\prime}\in\mathcal{A}}\langle
w,\phi_{h}(x,a^{\prime})\right).$

Our algorithm, PSDP-UCB, is presented in Algorithm 1. For some
$T\in\mathbb{N}$, the algorithm proceeds over $T$ _rounds_. In each round
$t\in[T]$, PSDP-UCB constructs a policy
$\hat{\pi}^{(t)}=(\hat{\pi}_{1}^{(t)},\ldots,\hat{\pi}_{H}^{(t)})$, which may
be interpreted as an estimate of the optimal policy. We will show that, for
$T$ sufficiently large, with high probability there is some $t\in[T]$ so that
$\mathbb{E}[V_{1}^{\star}(x_{1})-V_{1}^{\hat{\pi}^{(t)}}(x_{1})]$ is small.

In each round $t\in[T]$, the algorithm iterates through steps
$h=H,H-1,\ldots,1$: for each such $h$, it will define the mapping
$\hat{\pi}_{h}^{(t)}:\mathcal{X}\rightarrow\Delta(\mathcal{A})$. Thus, at step
$h$, $\hat{\pi}_{h+1}^{(t)},\ldots,\hat{\pi}_{H}^{(t)}$ have already been
defined. At each step $h$, in 4, the algorithm samples some number $n$ of
trajectories from the MDP according to a uniformly random policy from
(roughly) the set
$\\{\hat{\pi}^{(s)}\circ_{h+1}\hat{\pi}^{(t)}\\}_{s\in[t-1]}$.777 Technically,
this statement is slightly inaccurate: Algorithm 1 actually samples a uniform
policy from the set
$\\{\hat{\pi}^{(s)}\circ_{h}\tilde{\pi}_{h}^{(s)}\circ_{h+1}\hat{\pi}^{(t)}\\}_{s\in[t-1]}$,
where $\tilde{\pi}_{h}^{(s)}$ is defined for each $s$ to be a policy which
performs a sort of “uniform” exploration at each step $h$ (10). This choice is
made to simplify the analysis and is immaterial for our present discussion.
Next, in 5, the algorithm sets $\Sigma_{h}^{(t)}$ to be the empirical
covariance matrix of features at step $h$ under the $n$ trajectories just
sampled.

PSDP-UCB then modifies the rewards in each of these $n$ trajectories by adding
to the true rewards an exploration bonus, denoted by $F_{g}^{(t)}(\cdot)$ in
6. We call these modified rewards, denoted by $\hat{r}_{h}^{(t,i,h)}$ in
Algorithm 1, _optimistic rewards_. The remainder of this section will be
focused on defining and explaining the intuition behind the bonus function
$F_{g}^{(t)}$ (the formal definition is given in Eq. 33). Then, in 7, PSDP-UCB
sets $\hat{w}_{h}^{(t)}$ to be the solution to a least-squares regression
problem where the covariance matrix is $\Sigma_{h}^{(t)}$ and the labels are
the cumulative optimistic rewards from step $h$ to $H$. The algorithm then
repeats this procedure for step $h-1$, and so on.

We have omitted a description of 9 and 10, in which a policy at step $h$,
$\tilde{\pi}_{h}^{(t)}:\mathcal{X}\rightarrow\Delta(\mathcal{A})$, is defined.
This policy plays a relatively minor role (see Footnote 7) and may be ignored
for the purpose of the present discussion.

##### Overview of the analysis of PSDP-UCB.

To analyze PSDP-UCB, we define

$\displaystyle Q_{h}^{(t)}(x,a):=$ $\displaystyle
r_{h}(x,a)+\mathbb{E}^{\hat{\pi}^{(t)}}\left[\sum_{g=h+1}^{H}r_{g}(x_{g},a_{g})+F_{g}^{(t)}(x_{g})\
\mid\ (x_{h},a_{h})=(x,a)\right],$ (7)

and $V_{h}^{(t)}(x):=Q_{h}^{(t)}(x,\hat{\pi}_{h}^{(t)}(x))$. $Q_{h}^{(t)}$
represents the $Q$-value function corresponding to the policy
$\hat{\pi}^{(t)}$ and the optimistic rewards discussed above. Our main
technical lemma is the following statement showing that adding the exploration
bonuses $F_{h}^{(t)}(\cdot)$ to the rewards ensures that $Q_{h}^{(t)}$ is an
upper bound on $Q_{h}^{\star}$:

######  Lemma 3.1 (Informal version of Lemma 5.9).

With high probability over the execution of PSDP-UCB (Algorithm 1), for all
$x,a,h,t$, we have $Q_{h}^{(t)}(x,a)\geq Q_{h}^{\star}(x,a)-o(1)$ and
$V_{h}^{(t)}(x)+F_{h}^{(t)}(x)\geq V_{h}^{\star}(x)-o(1)$.

In light of Lemma 3.1, we can bound the average suboptimality of the policies
$\hat{\pi}^{(t)}$ as follows:

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\left(\mathbb{E}[V_{1}^{\star}(x_{1})-V_{1}^{\hat{\pi}^{(t)}}(x_{1})]\right)\leq$
$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\left(\mathbb{E}[V_{1}^{(t)}(x_{1})+F_{1}^{(t)}(x_{1})-V_{1}^{\hat{\pi}^{(t)}}(x_{1})]\right)-o(1)$
$\displaystyle=$
$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}^{\hat{\pi}^{(t)}}\left[\sum_{h=1}^{H}F_{h}^{(t)}(x_{h})\right]-o(1),$
(8)

where the inequality uses Lemma 3.1, and the equality uses the definition of
$V_{h}^{(t)}$. To upper bound Eq. 8, we need to show, roughly speaking, the
following: for states $x$ which are likely to be visited by many of the
policies $\hat{\pi}^{(t)}$, the sum of exploration bonuses at $x$ is sublinear
in $T$, i.e., $\frac{1}{T}\sum_{t=1}^{T}F_{h}^{(t)}(x)=o(T)$. We show a
version of this statement “in feature space”: in particular, we first upper
bound $|F_{h}^{(t)}(x)|$, for all states $x$, as follows.

######  Lemma 3.2 (Informal version of Lemmas 5.5, 5.3 and 5.2).

For all $x,h,t$ in the execution of PSDP-UCB, it holds that

$\displaystyle|F_{h}^{(t)}(x)|\leq\operatorname{poly}(d,H)^{A}\cdot\max_{a\in\mathcal{A}}\left(\phi_{h}(x,a)^{\top}\cdot(\Sigma_{h}^{(t)})^{-1}\cdot\phi_{h}(x,a)\right)^{1/2}.$

Lemma 3.2 shows that for states $x$ for which all actions $a\in\mathcal{A}$
point in directions of $\Sigma_{h}^{(t)}$ corresponding to eigenvectors with
large eigenvalues (i.e., “well-explored” directions), $|F_{h}^{(t)}(x)|$ must
be small. Using Lemma 3.2, the quantity in Eq. 8 may be bounded by $o(1)$
using a variant of the elliptic potential lemma (Lemma 6.4). The elliptic
potential lemma formalizes the notion that there are finitely many “different
directions” in $\mathbb{R}^{d}$, so if $T$ is sufficiently large, all
directions $\mathbb{R}^{d}$ frequently explored by policies $\hat{\pi}^{(t)}$
must correspond to large-eigenvalue directions for $\Sigma_{h}^{(t)}$.

Thus, it remains to discuss the proofs of Lemmas 3.1 and 3.2. Notice that the
statements of these two lemmas are in tension: Lemma 3.1 requires the bonuses
$F_{h}^{(t)}$ be sufficiently large so that $Q_{h}^{(t)}$ is an upper bound on
$Q_{h}^{\star}$, while Lemma 3.2 requires that $F_{h}^{(t)}$ be sufficiently
small, as controlled by $\Sigma_{h}^{(t)}$. As discussed in the following
subsections, by carefully defining $F_{h}^{(t)}$, we can ensure that both of
these criteria are met.

###  3.3 Implementing local optimism via Bellman-linearity

Essentially any approach to proving Lemma 3.1 requires that $Q_{h}^{(t)}(x,a)$
be a linear function of $\phi_{h}(x,a)$.888In particular, linearity of
$Q_{h}^{(t)}$ establishes that the regression problem solved in 7 of Algorithm
1 is well-specified, which allows us to show upper bounds on the prediction
error $\left|\langle\hat{w}_{h}^{(t)},\phi_{h}(x,a)\rangle-
Q_{h}^{(t)}(x,a)\right|$ and is in turn used to establish optimism. In turn,
as discussed in Section 3.1, this requires that the bonuses
$F_{h+1}^{(t)}(\cdot)$ have a Bellman backup which is linear in
$\phi_{h}(x,a)$. It will be useful to formalize the following notion of
_Bellman-linearity_ , which describes those functions whose Bellman backup is
linear in the features:

######  Definition 3.1.

Fix $h\in[H]$ with $h>1$ and $k\in\mathbb{N}$. We say that a function
$G:\mathcal{X}\rightarrow\mathbb{R}^{k}$ is _Bellman-linear at step $h$_ if
there exists $W_{G}\in\mathbb{R}^{k\times d}$ so that, for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$, it holds that
$W_{G}\cdot\phi_{h-1}(x,a)=\mathbb{E}_{x^{\prime}\sim
P_{h-1}(x,a)}[G(x^{\prime})]$. In particular, in the case that $k=1$,
$G:\mathcal{X}\rightarrow\mathbb{R}$ is Bellman-linear at step $h$ if there is
$w_{G}\in\mathbb{R}^{d}$ so that $\langle
w_{G},\phi_{h-1}(x,a)\rangle=\mathbb{E}_{x^{\prime}\sim
P_{h-1}(x,a)}[G(x^{\prime})]$.

If the value of $h$ is clear, we will simply say the function $G$ is Bellman-
linear. The canonical example of a Bellman-linear function is the mapping
$x\mapsto\max_{a}\langle w,\phi_{h}(x,a)\rangle$, for any
$w\in\mathbb{R}^{d}$. Given that such functions are Bellman-linear, we can
construct additional Bellman-linear functions by averaging, taking limits
(Lemma 4.1), and differentiating (Lemma 4.2). The following lemma shows that
Bellman-linearity of the bonuses $F_{h}^{(t)}$ is indeed a sufficient
condition for linearity of $Q_{h}^{(t)}$:

######  Lemma 3.3 (Variant of Lemma 5.1).

Suppose that, in PSDP-UCB, for each $h\in[H],t\in[T]$,
$F_{h}^{(t)}:\mathcal{X}\rightarrow\mathbb{R}$ is Bellman-linear at step $h$.
Then for all $h,t$, there is some $w_{h}^{(t)}\in\mathbb{R}^{d}$ so that
$Q_{h}^{(t)}(x,a)=\langle\phi_{h}(x,a),w_{h}^{(t)}\rangle$ for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$.

The proof of Lemma 3.3 requires some work to show that the expectation of the
bonuses $F_{g}^{(t)}$ for $g>h+1$ is linear at step $h$: this fact crucially
relies on the fact that the policies $\hat{\pi}^{(t)}$ are linear policies
(see Lemma 4.3).

##### A Bellman-linear approximation of $B^{\mathsf{quad}}_{h}$.

In light of our discussion pertaining to the quadratic bonus
$B^{\mathsf{quad}}_{h}$ around Eq. 5 above, a natural strategy is to show that
$\max_{a}B^{\mathsf{quad}}_{h}(\cdot,a;\Sigma)$ is Bellman-linear at step $h$
for any $\Sigma$. Though this is not the case (Proposition A.2), Lemma 3.4
below suggests a workaround: the quadratic bonus
$B^{\mathsf{quad}}_{h}(x,a;\Sigma)$ can be _approximated_ by a Bellman-linear
function.

######  Lemma 3.4 (Simplified version of Lemma 5.3).

Fix $h\in[H]$ and a PSD matrix $\Sigma$, and define

$\displaystyle
F^{\mathsf{normal}}_{h}(x;\Sigma):=\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma)}\left[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\right].$ (9)

Then $F^{\mathsf{normal}}_{h}(\cdot;\Sigma)$ is Bellman-linear at step $h$,
and for any $x\in\mathcal{X}$,

$\displaystyle\frac{1}{\sqrt{2\pi}}\max_{a,a^{\prime}\in\mathcal{A}}\left((\phi_{h}(x,a)-\phi_{h}(x,a^{\prime}))^{\top}\Sigma(\phi_{h}(x,a)-\phi_{h}(x,a^{\prime}))\right)^{1/2}\leq
F^{\mathsf{normal}}_{h}(x;\Sigma)\leq\sqrt{d}\max_{a\in\mathcal{A}}\left(\phi_{h}(x,a)^{\top}\Sigma\phi_{h}(x,a)\right)^{1/2}.$

The proof of Lemma 3.4 is straightforward: that $x\mapsto
F^{\mathsf{normal}}_{h}(x;\Sigma)$ is Bellman-linear at step $h$ follows from
the fact that $F^{\mathsf{normal}}_{h}(x;\Sigma)$ is defined as an average of
Bellman-linear functions. The upper and lower bounds on
$F^{\mathsf{normal}}_{h}(x;\Sigma)$ follow from a direct computation.

Lemma 3.4 suggests the following approach: for some scalar $\beta_{h}$,
suppose we define

$\displaystyle F_{h}^{(t)}(x):=\beta_{h}\cdot
F^{\mathsf{normal}}_{h}(x;(\Sigma_{h}^{(t)})^{-1}).$ (10)

Since $F^{\mathsf{normal}}_{h}(\cdot;\Sigma)$ is Bellman-linear and
approximates the quadratic bonus (Lemma 3.4), we might hope that it acts
similarly to the quadratic bonus in LSVI-UCB to induce an optimistic property
of the value functions $Q_{h}^{(t)}(x,a)$ in PSDP-UCB, thereby allowing us to
establish Lemma 3.1. Unfortunately, this argument encounters a few snags:

  1. 1.

The lower bound on $F^{\mathsf{normal}}_{h}(x;\Sigma)$ in Lemma 3.4 is not
exactly a constant times $B^{\mathsf{quad}}_{h}(x;\Sigma)$, but instead
involves a _difference_ between feature vectors at the state $x$. Thus,
strictly speaking, $F^{\mathsf{normal}}_{h}(x;\Sigma)$ does not give a
multiplicative approximation of the quadratic bonus
$B^{\mathsf{quad}}_{h}(x;\Sigma)$.

  2. 2.

Even if $F^{\mathsf{normal}}_{h}(x;\Sigma)$ did give a multiplicative
approximation of $B^{\mathsf{quad}}_{h}(x;\Sigma)$, we would need to ensure
that this multiplicative approximation is good enough for all arguments to go
through.

  3. 3.

Ignoring the previous two items, we still suffer from the following
exponential blowup. Roughly speaking, the proof of Lemma 3.1 proceeds by
showing that the loss in value from any sub-optimal action taken by
$\hat{\pi}^{(t)}$ at step $h$ is canceled out by the bonus
$F_{h}^{(t)}(x_{h})$. To show this, we need to show that the exploration bonus
$\beta_{h}\cdot F^{\mathsf{normal}}_{h}(x_{h};(\Sigma_{h}^{(t)})^{-1})$ is at
least as large as the regression error in our estimate of $\hat{w}_{h}^{(t)}$
from 7. But the labels in this linear regression involve the sum of bonuses
$F_{g}^{(t)}(\cdot)$ for $h+1\leq g\leq H$, and thus scale with
$\beta_{h+1}+\cdots+\beta_{H}$. Working through standard high-confidence
bounds for linear regression (see Lemma 5.7 for details) establishes that we
would need
$\beta_{h}\geq\operatorname{poly}(d)\cdot(\beta_{h+1}+\cdots+\beta_{H})$. This
recursion leads to $\beta_{1}\geq d^{O(H)}$, which would require exponentially
many samples to cancel out exponentially large error terms.

Items 1 and 2 above turn out not to be major issues. In the case of Item 1, a
lower bound on $F_{h}(x;\Sigma)$ growing only with the difference between
feature vectors still induces a sufficient exploration property on the induced
value functions: see the proof of Lemma 5.9. Item 2 may be handled by
increasing the scaling factors $\beta_{h}$ by a factor proportional to the gap
between the upper and lower bounds of Lemma 3.4.

The main technical challenge is dealing with Item 3 above, namely the
exponential blowup of $\beta_{h}$. Typically, the solution to the above
dilemma is to _truncate_ the bonus, i.e., define
$F_{h}^{(t)}(x):=\min\\{\beta\cdot
F^{\mathsf{normal}}_{h}(x;(\Sigma_{h}^{(t)})^{-1}),\bar{F}_{h}(x)\\}$, for
some $\bar{F}_{h}(x)\ll\beta$ which is a uniform upper bound on
$V_{h}^{\star}(x)$. (Note that, when using such an approach, the parameter
$\beta$ does not need any dependence on $h$.) Indeed, this truncation is used
in LSVI-UCB, where the $\hat{Q}_{h}^{(t)}$ function in Eq. 4 is truncated by
the constant $\bar{F}_{h}(x)\equiv H$. Unfortunately, in the setting of linear
Bellman completeness, truncation, in general, breaks Bellman-linearity of
$F_{h}^{(t)}(x)$. In the following section, we discuss how to implement this
truncation in a way that preserves Bellman-linearity.

###  3.4 Constructing truncated bonuses

##### Approximating $(\Sigma_{h}^{(t)})^{-1}$.

For each $h\in[H]$ and $t\in[T]$, the inverse covariance matrix
$(\Sigma_{h}^{(t)})^{-1}$ captures which directions of feature space are
explored by policies $\hat{\pi}^{(s)}\circ_{h}\tilde{\pi}^{(s)}$, for $s\leq
t-1$. We split $(\Sigma_{h}^{(t)})^{-1}$ into matrices representing “well-
explored directions” and “unexplored directions”, as follows: for an
appropriate threshold value $\sigma$, we let $\Sigma^{\prime}$ denote the
projection matrix onto the subspace spanned by eigenvectors of
$(\Sigma_{h}^{(t)})^{-1/2}$ with eigenvalues greater than $\sigma$, and
$\Lambda^{\prime}$ denote the projection matrix onto the subspace spanned by
eigenvectors of $(\Sigma_{h}^{(t)})^{-1/2}$ with eigenvalues less than
$\sigma$. Formally, as defined in Definition 5.2, we write
$(\Sigma^{\prime},\Lambda^{\prime})=\mathsf{trunc}_{\sigma}({(\Sigma_{h}^{(t)})^{-1/2}})$.
We denote subspaces onto which $\Sigma^{\prime},\Lambda^{\prime}$ project by
$\mathcal{S}_{\Sigma^{\prime}},\mathcal{S}_{\Lambda^{\prime}}$, respectively.
Note that $\Sigma^{\prime}\Lambda^{\prime}=\Lambda^{\prime}\Sigma^{\prime}=0$
and $\Sigma^{\prime}+\Lambda^{\prime}=I_{d}$; we call the pair
$(\Sigma^{\prime},\Lambda^{\prime})$ an _orthogonal pair_ (Definition 4.1).

If $\sigma$ is chosen sufficiently small to be less than our desired error
threshold, then the estimations $\hat{w}_{h}^{(t)}$ produced in 7 of Algorithm
1 are sufficiently accurate in directions spanned by
$\mathcal{S}_{\Lambda^{\prime}}$. Thus, we only have to add an exploration
bonus at a state $x$ which scales in proportion to the width of the features
$\phi_{h}(x,a),\ a\in\mathcal{A}$ in directions spanned by the orthogonal
complement of $\mathcal{S}_{\Lambda^{\prime}}$, namely
$\mathcal{S}_{\Sigma^{\prime}}$. This leads us to aim to construct an
exploration bonus which scales roughly as $x\mapsto\min\\{\beta\cdot
F^{\mathsf{normal}}_{h}(x;\Sigma^{\prime}),\bar{F}_{h}(x)\\}$, for some
parameter $\beta>1$ and an appropriate choice of $\bar{F}_{h}$ which should
(a) be uniformly upper bounded by a parameter which does not scale with
$\beta$, and (b) be an upper bound on $V_{h}^{\star}$. It turns out that
condition (b) can be relaxed to require only that $\bar{F}_{h}(x)$ be an upper
bound on
$\max_{a\in\mathcal{A}}Q_{h}^{\star}(x,a)-\min_{a\in\mathcal{A}}Q_{h}^{\star}(x,a)=\max_{a,a^{\prime}\in\mathcal{A}}\langle
w_{h}^{\star},\phi_{h}(x,a)-\phi_{h}(x,a^{\prime})\rangle.$ Thus, to
summarize, we aim to find a bonus function which satisfies the following
properties:

######  Problem 3.2 (Bonus function).

Fix $h,t$ (which determine
$(\Sigma^{\prime},\Lambda^{\prime})=\mathsf{trunc}_{\sigma}({(\Sigma_{h}^{(t)})^{-1/2}})$
as above), and some $\beta>1$. Can we find a Bellman-linear function
$F_{h}^{(t)}:\mathcal{X}\to\mathbb{R}_{\geq 0}$ which satisfies:

  1. 1.

For all $x\in\mathcal{X}$, $F_{h}^{(t)}(x)$ is lower bounded as

$\displaystyle F_{h}^{(t)}(x)\geq$ $\displaystyle
BH\cdot\max_{a,a^{\prime}\in\mathcal{A}}\min\left\\{\beta\cdot\|\Sigma^{\prime}\cdot(\phi_{h}(x,a)-\phi_{h}(x,a^{\prime}))\|_{2},\|\Lambda^{\prime}\cdot(\phi_{h}(x,a)-\phi_{h}(x,a^{\prime}))\|_{2}\right\\}.$
(11)

  2. 2.

For all $x\in\mathcal{X}$, $F_{h}^{(t)}(x)$ is upper bounded as

$\displaystyle F_{h}^{(t)}(x)\leq$
$\displaystyle\operatorname{poly}(d,H,B)\cdot\max_{a\in\mathcal{A}}\min\left\\{\beta\|\Sigma^{\prime}\cdot\phi_{h}(x,a)\|_{2},\|\phi_{h}(x,a)\|_{2}\right\\}.$

##### The case $d=A=2$.

To proceed, we consider the special case in which $d=2$. Let us suppose that
$\mathcal{S}_{\Sigma^{\prime}},\mathcal{S}_{\Lambda^{\prime}}$ are each one
dimensional, and are spanned by unit vectors $u,v\in\mathbb{R}^{2}$,
respectively, which are orthogonal. Let us consider the function

$\displaystyle F^{\mathsf{tl}}_{h}(x;u,v):=\max_{a\in\mathcal{A}}\langle
u,\phi_{h}(x,a)\rangle+\max_{a\in\mathcal{A}}\langle
v,\phi_{h}(x,a)\rangle-\max_{a\in\mathcal{A}}\langle
u+v,\phi_{h}(x,a)\rangle,$

for some parameter $\beta>1$. If we suppose that there are only 2 actions
(i.e., $|\mathcal{A}|=2$), then
$x\mapsto\sum_{\epsilon_{u},\epsilon_{v}\in\\{\pm
1\\}}F^{\mathsf{tl}}_{h}(x;\beta\epsilon_{u}\cdot u,\epsilon_{v}\cdot v)$
(scaled by a $\operatorname{poly}(B,H)$ factor) satisfies the conditions of
Problem 3.2. To see this, let us write $\mathcal{A}=\\{0,1\\}$, and set
$\phi_{h}^{\Delta}(x):=\phi_{h}(x,1)-\phi_{h}(x,0)$. Fix $x\in\mathcal{X}$,
and choose $\epsilon_{u},\epsilon_{v}\in\\{\pm 1\\}$ so that
$\langle\epsilon_{u}\cdot u,\phi_{h}^{\Delta}(x)\rangle$ and
$\langle\epsilon_{v}\cdot v,\phi_{h}^{\Delta}(x)\rangle$ have opposite signs.
Then

$\displaystyle F^{\mathsf{tl}}_{h}(x;\beta\epsilon_{u}u,\epsilon_{v}v)=$
$\displaystyle|\langle\beta\epsilon_{u}u,\phi_{h}^{\Delta}(x)\rangle|+|\langle\epsilon_{v}v,\phi_{h}^{\Delta}(x)\rangle|-|\langle\beta\epsilon_{u}u+\epsilon_{v}v,\phi_{h}^{\Delta}(x)\rangle|$
$\displaystyle\geq$ $\displaystyle\min\\{|\langle\beta
u,\phi_{h}^{\Delta}(x)\rangle|,|\langle v,\phi_{h}^{\Delta}(x)\rangle|\\},$
(12)

where the inequality uses the property that for real numbers $a,b$ with
opposite signs, we have $|a|+|b|-|a+b|\geq\min\\{|a|,|b|\\}$. By combining Eq.
12 with non-negativity of $F^{\mathsf{tl}}_{h}(\cdot)$, we obtain the first
property of Problem 3.2. The second property may be verified in a similar
manner (see Lemma 4.8).

##### The case $d=2$ and general $A$.

Let us try to extend the above argument for the case $d=2$ and $A=2$ to the
case of general $A$ (still keeping $d=2$ fixed). Unfortunately, the function
$x\mapsto\sum_{\epsilon_{u},\epsilon_{v}\in\\{\pm
1\\}}F^{\mathsf{tl}}_{h}(x;\beta\epsilon_{u}\cdot u,\epsilon_{v}\cdot v)$ we
used above no longer satisfies the first condition of Problem 3.2. To see
this, suppose that $\mathcal{A}=\\{0,1,2,3\\}$ and the feature vectors at some
state $x$ form a rectangle, i.e.,
$\phi_{h}(x,0)=(0,0),\phi_{h}(x,1)=(u,0),\phi_{h}(x,2)=(u+v,u+v),\phi_{h}(x,3)=(0,v)$.
It is straightforward to see that $F^{\mathsf{tl}}_{h}(x;u,v)=0$, whereas
$\max_{a,a^{\prime}\in\mathcal{A}}\min\\{|\langle\beta
u,\phi_{h}(x,a)-\phi_{h}(x,a^{\prime})\rangle|,|\langle
v,\phi_{h}(x,a)-\phi_{h}(x,a^{\prime})\rangle|\\}\geq 1$ for any $\beta\geq 1$
(e.g., take $a=0,a^{\prime}=2$). Thus, Eq. 12 cannot hold, even if we only
wish for it to hold up to a $\operatorname{poly}(H,B,d)$ factor.

The above issue runs fairly deep: it can be shown that, under some mild
conditions on $F^{\mathsf{tl}}_{h}$ which ensure that it be Bellman-
linear,999Namely, that
$F^{\mathsf{tl}}_{h}(x;u,v)=\int_{w\in\mathbb{R}^{d}}f(w)\cdot\max_{a\in\mathcal{A}}\langle\phi_{h}(x,a),w\rangle
d\mu(w)$, for some Borel measure $\mu$ and function
$f:\mathbb{R}^{d}\to\mathbb{R}$. there is _no function_ $F_{h}^{(t)}$
satisfying both conditions of Problem 3.2. To rectify this issue, we will
relax the first condition of Problem 3.2, as follows:

######  Problem 3.3.

In the setting of Problem 3.2, can we find a Bellman-linear function
$F_{h}^{(t)}:\mathcal{X}\to\mathbb{R}_{\geq 0}$ which satisfies the
constraints of Problem 3.2 where Item 1 is replaced by the following
constraint:

  1. 1.

For all $x\in\mathcal{X}$, $F_{h}^{(t)}(x)$ is lower bounded as

$\displaystyle F_{h}^{(t)}(x)\geq$ $\displaystyle
BH\cdot\max_{a,a^{\prime}\in\mathcal{A}}\min_{\psi\in\bar{\Phi}(x)}\left\\{\|\beta\Sigma^{\prime}\cdot(\phi_{h}(x,a)-\psi)\|_{2}+\|\Lambda^{\prime}\cdot(\psi-\phi_{h}(x,a^{\prime})\|_{2}\right\\},$
(13)

where $\bar{\Phi}_{h}(x):={\rm co}\\{\phi_{h}(x,a):\ a\in\mathcal{A}\\}$
denotes the convex hull of the feature vectors $\phi_{h}(x,a),\
a\in\mathcal{A}$.

Note that Eq. 13 is weaker than Eq. 12, as can be verified by considering the
choices $\psi=\phi_{h}(x,a)$ and $\psi=\phi_{h}(x,a^{\prime})$ for given
$a,a^{\prime}$ in the outer maximum. As we discuss in Section 5.3, this weaker
lower bound on $F_{h}^{(t)}(x)$ is still sufficient to ensure optimism.
Moreover, for the example discussed above in which the feature vectors at a
state $x$ form a rectangle, the right-hand side of Eq. 13 is $0$: if
$a,a^{\prime}$ in the outer maximum are chosen to represent opposing corners
of the rectangle, then taking $\psi$ to be one of the other two corners yields
$|\langle\beta u,\phi_{h}(x,a)-\psi\rangle|=|\langle
v,\psi-\phi_{h}(x,a^{\prime})\rangle|=0$. For this example, we therefore have
that $x\mapsto\sum_{\epsilon_{u},\epsilon_{v}\in\\{\pm
1\\}}F^{\mathsf{tl}}_{h}(x;\beta\epsilon_{u}\cdot u,\epsilon_{v}\cdot v)$
satisfies the properties of Problem 3.3.

In fact, this holds more generally in the $d=2$ case, as can be seen by the
following argument: fix any $a,a^{\prime}\in\mathcal{A}$, and choose

$\displaystyle\psi^{\star}:=\operatorname*{arg\,min}_{\psi\in\bar{\Phi}_{h}(x)}\left\\{|\langle\beta
u,\phi_{h}(x,a)-\psi\rangle|+|\langle
v,\psi-\phi_{h}(x,a^{\prime})\rangle|\right\\}.$ (14)

Moreover choose $\epsilon_{u},\epsilon_{v}\in\\{\pm 1\\}$ so that
$\langle\phi_{h}(x,a)-\psi^{\star},\epsilon_{u}u\rangle\geq 0$ and
$\langle\phi_{h}(x,a^{\prime})-\psi^{\star},\epsilon_{v}v\rangle\geq 0$. Note
that Eq. 14 is a convex optimization problem; its KKT conditions tell us that
$\langle\psi^{\star}-\psi,\beta\epsilon_{u}u+\epsilon_{v}v\rangle\geq 0$ for
all $\psi\in\bar{\Phi}_{h}(x)$. Then, recalling that
$a,a^{\prime}\in\mathcal{A}$ were fixed, we have

$\displaystyle F^{\mathsf{tl}}_{h}(x;\beta\epsilon_{u}u,\epsilon_{v}v)\geq$
$\displaystyle\langle\beta\epsilon_{u}u,\phi_{h}(x,a)\rangle+\langle\epsilon_{v}v,\phi_{h}(x,a^{\prime})\rangle-\max_{\bar{a}\in\mathcal{A}}\langle\beta\epsilon_{u}u+\epsilon_{v}v,\phi_{h}(x,\bar{a})\rangle$
$\displaystyle\geq$
$\displaystyle\langle\beta\epsilon_{u}u,\phi_{h}(x,a)\rangle+\langle\epsilon_{v}v,\phi_{h}(x,a^{\prime})\rangle-\langle\beta\epsilon_{u}u+\epsilon_{v}v,\psi^{\star}\rangle$
$\displaystyle=$
$\displaystyle\langle\beta\epsilon_{u}u,\phi_{h}(x,a)-\psi^{\star}\rangle+\langle\epsilon_{v}v,\phi_{h}(x,a^{\prime})-\psi^{\star}\rangle$
$\displaystyle=$ $\displaystyle|\langle\beta
u,\phi_{h}(x,a)-\psi^{\star}\rangle|+|\langle
v,\phi_{h}(x,a^{\prime})-\psi^{\star}\rangle|,$ (15)

where the second inequality uses the KKT conditions of Eq. 14 and the final
equality uses the choices of $\epsilon_{u},\epsilon_{v}$. Eq. 15, combined
with non-negativity of $F^{\mathsf{tl}}_{h}(\cdot)$, verifies the first
property of Problem 3.3, as desired; the second property (namely, Item 2 of
Problem 3.2) can be verified by a direct computation (Lemma 4.8).

##### The case of general $d$ and general $A$.

The argument presented above for $d=2$ generalizes to the case of arbitrary
$d\in\mathbb{N}$. In particular, we will aim to construct a bonus of the form
$G_{h}^{(t)}(x):=O(BH)\cdot\mathbb{E}_{u\sim\mathcal{D}_{\Sigma^{\prime}},v\sim\mathcal{D}_{\Lambda^{\prime}}}[F^{\mathsf{tl}}_{h}(x;u,v)]$,
where $\mathcal{D}_{\Sigma^{\prime}},\mathcal{D}_{\Lambda^{\prime}}$ are
distributions supported on
$\mathcal{S}_{\Sigma^{\prime}},\mathcal{S}_{\Lambda^{\prime}}$, respectively.
A key consideration is: _how do we choose the distributions
$\mathcal{D}_{\Sigma^{\prime}},\mathcal{D}_{\Lambda^{\prime}}$?_ To answer
this question, it is instructive to consider how the above argument verifying
Item 1 of Problem 3.3 generalizes to higher dimensions. In particular, given
arbitrary $a,a^{\prime}\in\mathcal{A}$, in lieu of Eq. 14, we define

$\displaystyle\psi^{\star}:=\operatorname*{arg\,min}_{\psi\in\bar{\Phi}_{h}(x)}\left\\{\|\beta\Sigma^{\prime}\cdot(\phi_{h}(x,a)-\psi)\|_{2}+\|\Lambda^{\prime}\cdot(\psi-\phi_{h}(x,a^{\prime}))\|_{2}\right\\}.$
(16)

Using the KKT conditions of Eq. 16, it can be verified that, for
$u=\frac{\Sigma^{\prime}\cdot(\phi_{h}(x,a)-\psi^{\star})}{\|\Sigma^{\prime}\cdot(\phi_{h}(x,a)-\psi^{\star})\|_{2}}$
and
$v=\frac{\Lambda^{\prime}\cdot(\phi_{h}(x,a^{\prime})-\psi^{\star})}{\|\Lambda^{\prime}\cdot(\phi_{h}(x,a^{\prime})-\psi^{\star})\|_{2}}$,
we have that

$\displaystyle F^{\mathsf{tl}}_{h}(x;\beta
u,v)\geq\|\beta\Sigma^{\prime}\cdot(\phi_{h}(x,a)-\psi^{\star})\|_{2}+\|\Lambda^{\prime}\cdot(\phi_{h}(x,a^{\prime})-\psi^{\star})\|_{2}.$
(17)

Unfortunately these choices of $u,v$ depend on the state $x$, and
$\mathcal{D}_{\Sigma^{\prime}},\mathcal{D}_{\Lambda^{\prime}}$ must be
independent of $x$ in order to ensure Bellman-linearity of the bonus function.

Fortunately, as shown in Lemma 4.9, under some additional technical
conditions, Eq. 17 holds up to an additive term of $O(\varepsilon)$ when the
left-hand side is replaced by $F^{\mathsf{tl}}_{h}(x;\beta
u^{\prime},v^{\prime})$, for any vectors $u^{\prime},v^{\prime}$ which are
$\varepsilon$-close in $\ell_{2}$ norm to $u,v$, respectively. If we take
$\mathcal{D}_{\Sigma^{\prime}},\mathcal{D}_{\Lambda^{\prime}}$ to be the
uniform distributions on the unit sphere in the subspaces
$\mathcal{S}_{\Sigma^{\prime}},\mathcal{S}_{\Lambda^{\prime}}$, respectively,
then
$u^{\prime}\sim\mathcal{D}_{\Sigma^{\prime}},v^{\prime}\sim\mathcal{D}_{\Lambda^{\prime}}$
satisfy $\|u^{\prime}-u\|_{2}\leq\varepsilon,\|v^{\prime}-v\|_{2}\leq\epsilon$
with probability $\varepsilon^{O(d)}$. With these choices of
$\mathcal{D}_{\Sigma^{\prime}},\mathcal{D}_{\Lambda^{\prime}}$, the bonus
$G_{h}^{(t)}$ satisfies the constraint Eq. 13 up to a factor of
$\varepsilon^{-O(d)}$. In fact, this can easily be improved to a factor of
$\varepsilon^{-O(A)}$ by using the fact that closeness of $u,u^{\prime}$ and
of $v,v^{\prime}$ only needs to hold in the subspace spanned by
$\phi_{h}(x,a)$ (for $a\in\mathcal{A}$), which has at most $A$ dimensions –
see Lemma 4.10.

It turns out to simplify the analysis to take
$\mathcal{D}_{\Sigma^{\prime}},\mathcal{D}_{\Lambda^{\prime}}$ to be the
distributions
$\mathcal{N}(0,\Sigma^{\prime}),\mathcal{N}(0,\Lambda^{\prime})$,
respectively, which permit essentially the same analysis as described above.
Moreover, to satisfy certain technical conditions in Lemma 4.9, we add a term
to the bonus function which is proportional to
$F^{\mathsf{normal}}_{h}(x;\Sigma^{\prime})=\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle]$. To sum up, we define the bonus function as

$\displaystyle F_{h}^{(t)}(x):=$
$\displaystyle\beta_{1}\cdot\mathbb{E}_{u^{\prime}\sim\mathcal{N}(0,\Sigma^{\prime})}\mathbb{E}_{v^{\prime}\sim\mathcal{N}(0,\Lambda^{\prime})}\left[F^{\mathsf{tl}}_{h}(x;\beta
u^{\prime},v^{\prime})\right]+\beta_{2}\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\right],$ (18)

where $\beta_{1},\beta_{2},\beta$ are carefully chosen parameters bounded
above by $(BHD/\epsilon)^{O(A)}$, where $\epsilon$ is the desired error level
of the algorithm’s output policy. The precise values of these parameters
appear in the formal definition of $F_{h}^{(t)}$ in Eq. 33. We remark that,
crucially, $\beta_{2}\ll\beta_{1}\cdot\beta$ – the entire point of the
truncation procedure we have just described is to ensure that the size of
$F_{h}^{(t)}(x)$ (which may be as large as $O(\beta_{2})$) does _not_ scale
with $\beta$.

In Sections 4.2 and 4.3, we formally verify, per the arguments outlined above,
that the bonus $F_{h}^{(t)}$ in Eq. 18 satisfies the properties laid out in
Problem 3.3. Then, in Sections 5.1, 5.2 and 5.3, we show how this implies the
statements of Lemmas 3.1 and 3.2, which suffices to prove our main theorem
(Theorem 1.1).

##  4 Bellman linear functions

Recall the definition of Bellman linearity in Definition 3.1. In this section,
we discuss several properties of Bellman linear functions.

######  Lemma 4.1 (Limiting preserves Bellman-linearity).

Fix $h\in[H]$. Suppose that for each $i\in\mathbb{N}$, we are given a function
$G_{i}:\mathcal{X}\rightarrow\mathbb{R}$ which is Bellman-linear at step $h$
and so that $\sup_{x\in\mathcal{X},i\in\mathbb{N}}|G_{i}(x)|\leq C$, for some
$C>0$. Furthermore suppose that, for some
$G:\mathcal{X}\rightarrow\mathbb{R}$, $\lim_{i\rightarrow\infty}G_{i}(x)=G(x)$
for all $x\in\mathcal{X}$. Then $G$ is Bellman-linear at step $h$.

###### Proof.

For each $i\in\mathbb{N}$, we know that there is some $w_{i}\in\mathbb{R}^{d}$
so that, for all $(x,a)\in\mathcal{X}\times\mathcal{A}$,
$\langle\phi_{h-1}(x,a),w_{i}\rangle=\mathbb{E}_{x^{\prime}\sim
P_{h-1}(x,a)}[G_{i}(x^{\prime})]$. For each $i,a,x$, we know that
$|\mathbb{E}_{x^{\prime}\sim P_{h-1}(x,a)}[G_{i}(x^{\prime})]|\leq C$ by
uniform boundedness of $G_{i}$, so $|\langle\phi_{h-1}(x,a),w_{i}\rangle|\leq
C$ for all $i,a,x$. By definition of $\mathcal{B}_{h-1}$, it follows that
$w_{i}\in C\cdot\mathcal{B}_{h-1}$, which gives $\|w_{i}\|_{2}\leq CB$ by
Assumption 2.1. Hence we may find a subsequence $(w_{i_{j}})_{j\in\mathbb{N}}$
so that $\lim_{j\rightarrow\infty}w_{i_{j}}=w^{\star}$, for some $w^{\star}\in
C\cdot\mathcal{B}_{h-1}$. Thus, for any
$(x,a)\in\mathcal{X}\times\mathcal{A}$, we have

$\displaystyle\langle\phi_{h-1}(x,a),w^{\star}\rangle=\lim_{j\rightarrow\infty}\langle\phi_{h-1}(x,a),w_{i_{j}}\rangle=\lim_{j\rightarrow\infty}\mathbb{E}_{x^{\prime}\sim
P_{h-1}(x,a)}[G_{i_{j}}(x^{\prime})]$ $\displaystyle=$
$\displaystyle\mathbb{E}_{x^{\prime}\sim
P_{h-1}(x,a)}\left[\lim_{j\rightarrow\infty}G_{i_{j}}(x^{\prime})\right]=\mathbb{E}_{x^{\prime}\sim
P_{h-1}(x,a)}[G(x^{\prime})],$

where the second-to-last equality uses the dominated convergence theorem and
uniform boundedness of the functions $G_{i}$. ∎

######  Lemma 4.2 (Differentiation preserves Bellman-linearity).

Fix any $h,k\in\mathbb{N}$ and consider an open set $U\in\mathbb{R}^{k}$ as
well as a function $F:U\times\mathcal{X}\rightarrow\mathbb{R}$ so that
$F(u,x)$ is continuously differentiable in $u$ for all $x\in\mathcal{X}$, and
that $\max_{x\in\mathcal{X},u\in
U}\max\\{F(u,x),\|\nabla_{u}F(u,x)\|_{2}\\}\leq C$. Suppose that, for all
$u\in U$, the function $F(u,\cdot)$ is Bellman-linear at step $h$. Then the
function $\nabla_{u}F(u,\cdot)$ is Bellman-linear at step $h$, for all $u\in
U$.

###### Proof.

By Bellman-linearity, we know that for all $u\in U$, there is some
$w_{u}\in\mathbb{R}^{d}$ so that, for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$, it holds that

$\displaystyle\langle\phi_{h}(x,a),w_{u}\rangle=\mathbb{E}_{x^{\prime}\sim
P_{h-1}(x,a)}\left[F(u,x^{\prime})\right].$

By the dominated convergence theorem (which may be applied as a consequence of
uniform boundedness of $\|\nabla_{u}F(u,x)\|_{2}$ over $u\in U$), it holds
that $u\mapsto\langle\phi_{h}(x,a),w_{u}\rangle$ is continuously
differentiable for all $x,a$, and moreover,

$\displaystyle\nabla_{u}\langle\phi_{h}(x,a),w_{u}\rangle=\mathbb{E}_{x^{\prime}\sim
P_{h-1}(x,a)}[\nabla_{u}F(u,x^{\prime})].$ (19)

Choose
$(x_{1},a_{1}),\ldots,(x_{\ell},a_{\ell})\in\mathcal{X}\times\mathcal{A}$ so
that $\\{\phi_{h}(x_{i},a_{i})\\}_{i=1}^{\ell}$ is linearly independent and
spans $\\{\phi_{h}(x,a)\\}_{(x,a)\in\mathcal{X}\times\mathcal{A}}$; then
$\ell\leq d$. Let $W_{u}\in\mathbb{R}^{d\times d}$ be any matrix satisfying

$\displaystyle
W_{u}\cdot\phi_{h}(x_{i},a_{i})=\nabla_{u}\langle\phi_{h}(x_{i},a_{i}),w_{u}\rangle\quad\forall
i\in[\ell],$

which is possible by linear independence of $\phi_{h}(x_{i},a_{i}),\
i\in[\ell]$. Now consider any $(x,a)\in\mathcal{X}\times\mathcal{A}$; we can
find $\alpha_{1},\ldots,\alpha_{\ell}\in\mathbb{R}$ so that
$\phi_{h}(x,a)=\sum_{i=1}^{\ell}\alpha_{i}\cdot\phi_{h}(x_{i},a_{i})$. Then

$\displaystyle
W_{u}\cdot\phi_{h}(x,a)=\sum_{i=1}^{\ell}\alpha_{i}\cdot\nabla_{u}\langle\phi_{h}(x_{i},a_{i}),w_{u}\rangle=\nabla_{u}\sum_{i=1}^{\ell}\alpha_{i}\cdot\langle\phi_{h}(x_{i},a_{i}),w_{u}\rangle=\nabla_{u}\langle\phi_{h}(x,a),w_{u}\rangle,$

which completes the proof by Eq. 19. ∎

Lemma 4.3 below states that the feature vectors induced by any linear policy
are Bellman-linear. It is a special case of Corollary 3.2 of [GM24] (by taking
the inherent Bellman error to be 0).

######  Lemma 4.3.

Suppose that $M$ is linear Bellman complete. Then for each $h\in[H]$ and each
$w\in\mathbb{R}^{d}$, the mapping
$x^{\prime}\mapsto\phi_{h+1}(x^{\prime},\pi_{h+1,w}(x^{\prime}))$ is Bellman-
linear at step $h+1$, i.e., there is a linear map
$L_{h}(w):\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ so that, for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$,

$\displaystyle L_{h}(w)^{\top}\cdot\phi_{h}(x,a)=\mathbb{E}_{x^{\prime}\sim
P_{h}(x,a)}[\phi_{h+1}(x^{\prime},\pi_{h+1,w}(x^{\prime}))].$

###### Proof.

Fix $h\in[H]$. Define
$V(x^{\prime},\theta):=\max_{a^{\prime}\in\mathcal{A}}\langle\phi_{h+1}(x^{\prime},a^{\prime}),\theta\rangle$
and $Q(x,a,\theta)=\langle\phi_{h}(x,a),\mathcal{T}_{h}\theta\rangle$, so that
$Q(x,a,\theta)=\langle\phi_{h}(x,a),\mathcal{T}_{h}\theta\rangle=\mathbb{E}_{x^{\prime}\sim
P_{h}(x,a)}[V(x^{\prime},\theta)]$. By Fubini’s theorem, for any $\sigma>0$,
we have

$\displaystyle\mathsf{S}_{\sigma}Q(x,a,\theta)=\mathbb{E}_{z\sim\mathcal{N}(0,\sigma^{2}I_{d})}[Q(x,a,\theta-z)]=\mathbb{E}_{x^{\prime}\sim
P_{h}(x,a)}\mathbb{E}_{z\sim\mathcal{N}(0,\sigma^{2}I_{d})}[V(x^{\prime},\theta-x)]=\mathbb{E}_{x^{\prime}\sim
P_{h}(x,a)}[\mathsf{S}_{\sigma}V(x^{\prime},\theta)].$

Linear Bellman completeness of $M$ gives that for each
$\theta\in\mathbb{R}^{d}$, the function $x^{\prime}\mapsto
V(x^{\prime},\theta)$, and thus
$x^{\prime}\mapsto\mathsf{S}_{\sigma}V(x^{\prime},\theta)$ is Bellman-linear
at step $h+1$. The function $\mathsf{S}_{\sigma}V(x^{\prime},\theta)$ is
infinitely differentiable as a function of $\theta$, for all $x^{\prime}$;
moreover, since

$\displaystyle\nabla\mathsf{S}_{\sigma}V(x^{\prime},\theta)=\mathbb{E}_{z\sim\mathcal{N}(0,\sigma^{2}I_{d})}[\nabla
V(x^{\prime},\theta-z)]=\mathbb{E}_{z\sim\mathcal{N}(0,\sigma^{2}I_{d})}[\phi_{h+1}(x^{\prime},\pi_{h+1,\theta-z}(x^{\prime}))],$
(20)

we have that
$\max\\{\mathsf{S}_{\sigma}V(x^{\prime},\theta),\|\nabla\mathsf{S}_{\sigma}V(x^{\prime},\theta)\|_{2}\\}$
is bounded uniformly over $x^{\prime}$ in the neighborhood of any $\theta$. We
remark that the first equality in Eq. 20 uses the dominated convergence
theorem, and the second equality uses that the $V(x^{\prime},\theta-z)$ is
differentiable almost surely over $z\sim\mathcal{N}(0,\sigma^{2}I_{d})$. Thus,
by Lemma 4.2, the function
$x^{\prime}\mapsto\nabla\mathsf{S}_{\sigma}V(x^{\prime},\theta)$ is Bellman-
linear at step $h$, for all $\theta\in\mathbb{R}^{d}$. But for
$x^{\prime}\in\mathcal{X}$ and $w\in\mathbb{R}^{d}$,

$\displaystyle\nabla\mathsf{S}_{\sigma}V(x^{\prime},w)=\phi_{h+1}(x^{\prime},\pi_{h+1,w,\sigma}(x^{\prime})),$

where the equality above uses Eq. 20 together with Definition A.1. Next, Eq.
72 ensures that for each $x^{\prime}\in\mathcal{X}$, $\lim_{\sigma\rightarrow
0^{+}}\phi_{h+1}(x^{\prime},\pi_{h+1,w,\sigma}(x^{\prime})=\phi_{h+1}(x^{\prime},\pi_{h+1,w}(x^{\prime}))$.
Then Lemma 4.1 (applied to each coordinate of the function
$x^{\prime}\mapsto\phi_{h+1}(x^{\prime},\pi_{h+1,w,\sigma}(x^{\prime})$, for
$\sigma>0$) yields that
$x^{\prime}\mapsto\phi_{h+1}(x^{\prime},\pi_{h+1,w}(x^{\prime}))$ is Bellman-
linear at step $h+1$, for any $w\in\mathbb{R}^{d}$. ∎

Given any $\pi_{h+1}\in\Pi^{\mathsf{lin}}_{h+1}$ (so that
$\pi_{h+1}=\pi_{h+1,v}$ for some $v\in\mathbb{R}^{d}$) and
$w\in\mathbb{R}^{d}$, we define $\mathcal{T}_{h}^{\pi}w:=L_{h}(v)\cdot w$.
Then by Lemma 4.3, for all $(x,a)\in\mathcal{X}\times\mathcal{A}$,

$\displaystyle\langle\phi_{h}(x,a),\mathcal{T}_{h}^{\pi}w\rangle=\mathbb{E}_{x^{\prime}\sim
P_{h}(x,a)}[\langle\phi_{h+1}(x^{\prime},\pi_{h+1}(x^{\prime})),w\rangle].$
(21)

Note that $\mathcal{T}_{h}^{\pi}$ depends only on $\pi_{h+1}$; thus with
slight abuse of notation, we will sometimes write
$\mathcal{T}_{h}^{\pi_{h+1}}w$ in place of $\mathcal{T}_{h}^{\pi}w$. We also
remark that by linearity of expectation, for any $\pi\in\Pi^{\mathsf{Plin}}$
and $w\in\mathbb{R}^{d}$, there is a vector
$\mathcal{T}_{h}^{\pi}w\in\mathbb{R}^{d}$ so that Eq. 21 holds for the
perturbed linear policy $\pi$.

As an immediate consequence of Lemma 4.3, we can show that the $Q$-function
for a linear policy is linear.

######  Corollary 4.4.

Suppose that $M$ is linear Bellman complete, and that
$\pi\in\Pi^{\mathsf{Plin}}$. Then for each $h\in[H]$, there is a vector
$w_{h}^{\pi}\in H\cdot\mathcal{B}_{h}\subset\mathbb{R}^{d}$ so that for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$, $Q_{h}^{\pi}(x,a)=\langle
w_{h}^{\pi},\phi_{h}(x,a)\rangle$. Moreover, $\|w_{h}^{\pi}\|_{2}\leq HB$.

###### Proof.

Since $Q_{h}^{\pi}(x,a)=r_{h}(x,a)+\mathbb{E}_{x^{\prime}\sim
P_{h}(x,a)}[Q_{h+1}^{\pi}(x,\pi_{h+1}(x))]$, by induction and the fact that
$r_{h}(x,a)$ is linear, it suffices to show that for all $v\in\mathbb{R}^{d}$,
there is some $w\in\mathbb{R}^{d}$ so that
$\langle\phi_{h}(x,a),w\rangle=\mathbb{E}_{x^{\prime}\sim
P_{h}(x,a)}[\langle\phi_{h}(x,\pi_{h+1}(x)),v\rangle]$. We may write
$\pi_{h+1}=\pi_{h+1,y,\sigma}$ for some $y\in\mathbb{R}^{d}$ and $\sigma\geq
0$; then by Lemma 4.3, we may take
$w=\mathbb{E}_{z\sim\mathcal{N}(y,\sigma^{2})}[L_{h}(z)]\cdot v$.

To see the upper bound on $\|w_{h}^{\pi}\|_{2}$, note that, by definition of
$Q_{h}^{\pi}$, we have $|\langle w_{h}^{\pi},\phi_{h}(x,a)\rangle|\leq H$ for
all $x,a,h$. Then it follows that $\|w_{h}^{\pi}\|_{2}\leq HB$ by Assumption
2.1. ∎

###  4.1 Construction of Bellman-linear functions for bonuses

A central component of the proof consists of the construction of functions, to
be used as exploration bonuses, which (a) are Bellman-linear and (b)
approximate a truncated version of arbitrary norms induced by PSD matrices. In
this section, we establish Bellman linearity of certain functions which form
the building blocks of these bonuses.

Given orthogonal vectors $u,v\in\mathbb{R}^{d}$, we let
$\Pi_{u,v}:=\frac{uu^{\top}}{\|u\|_{2}}+\frac{vv^{\top}}{\|v\|_{2}}$. Note
that if $u,v$ are unit vectors, then $\Pi_{u,v}$ is the projection matrix onto
${\rm Span}\\{u,v\\}$.

######  Definition 4.1 (Orthogonal pair).

We define a tuple of $d\times d$ PSD matrices $(\Sigma,\Lambda)$ to be an
_orthogonal pair_ if the following equalities hold:

$\displaystyle\Sigma^{2}=\Sigma,\
\Lambda^{2}=\Lambda,\qquad\Sigma\Lambda=\Lambda\Sigma=0,\qquad\Sigma+\Lambda=I_{d}.$

For an orthogonal pair $(\Sigma,\Lambda)$, $\Sigma,\Lambda$ are projections
onto subspaces of $\mathbb{R}^{d}$, which we denote by
$\mathcal{S}_{\Sigma}:=\\{\Sigma v:\
v\in\mathbb{R}^{d}\\},\mathcal{S}_{\Lambda}:=\\{\Lambda v:\
v\in\mathbb{R}^{d}\\}$, respectively. Since $\Sigma\Lambda=\Lambda\Sigma=0$,
the subspaces $\mathcal{S}_{\Sigma},\mathcal{S}_{\Lambda}$ are orthogonal.
Since $\Sigma+\Lambda=I_{d}$, we have that
$\dim(\mathcal{S}_{\Sigma})+\dim(\mathcal{S}_{\Lambda})=d$.

######  Definition 4.2 (Truncated linear bonus).

Consider orthogonal vectors $u,v\in\mathbb{R}^{d}$. For
$\Phi\in\mathscr{P}^{d}$, we define

$F^{\mathsf{tl}}(\Phi;u,v)=\max_{\phi\in\Phi}\langle
u,\phi\rangle+\max_{\phi\in\Phi}\langle
v,\phi\rangle-\max_{\phi\in\Phi}\langle u+v,\phi\rangle.$

For $x\in\mathcal{X}$ and $h\in[H]$, we then define
$F^{\mathsf{tl}}_{h}(x;u,v):=F^{\mathsf{tl}}(\bar{\Phi}_{h}(x);u,v)$.

The following lemma is immediate from the definition of
$F^{\mathsf{tl}}(\cdot)$.

######  Lemma 4.5 (Bellman-linearity of truncated linear bonus).

Given any $h>1$ and vectors $u,v\in\mathbb{R}^{d}$, the mapping $x\mapsto
F^{\mathsf{tl}}_{h}(x;u,v)$ is Bellman-linear at step $h$.

We next record the following lemmas for future use, which follow as an
immediate consequence of the definition of $F^{\mathsf{tl}}$.

######  Lemma 4.6.

Let $\Phi\in\mathscr{P}^{d}$. Suppose that
$u,u^{\prime},v,v^{\prime}\in\mathbb{R}^{d}$ satisfy $\langle
u,\phi\rangle=\langle u^{\prime},\phi\rangle$ and $\langle
v,\phi\rangle=\langle v^{\prime},\phi\rangle$ for all $\phi\in\Phi$. Then

$\displaystyle
F^{\mathsf{tl}}(\Phi;u,v)=F^{\mathsf{tl}}(\Phi;u^{\prime},v^{\prime}).$

###### Proof.

We have that $\max_{\phi\in\Phi}\langle
w,\phi\rangle=\max_{\phi\in\Phi}\langle w^{\prime},\phi\rangle$ for each pair
$(w,w^{\prime})\in\\{(u,u^{\prime}),(v,v^{\prime}),(u+v,u^{\prime}+v^{\prime})\\}$.
∎

######  Lemma 4.7.

Let $\Phi\in\mathscr{P}^{d}$. Suppose that $u,v\in\mathbb{R}^{d}$ are given
and that $\alpha_{u},\alpha_{v}\geq 0$ are real numbers. Then

$\displaystyle F^{\mathsf{tl}}(\Phi;\alpha_{u}\cdot u,\alpha_{v}\cdot
v)\geq\min\\{\alpha_{u},\alpha_{v}\\}\cdot F^{\mathsf{tl}}(\Phi;u,v).$

###### Proof.

Since $F^{\mathsf{tl}}(\Phi;\alpha u,\alpha v)=\alpha\cdot
F^{\mathsf{tl}}(\Phi;u,v)$ for any $\alpha\geq 0$, it suffices to consider the
case that $\min\\{\alpha_{u},\alpha_{v}\\}=1$. Without loss of generality let
us suppose that $\alpha_{u}=1\leq\alpha_{v}$. Then

$\displaystyle F^{\mathsf{tl}}(\Phi;u,\alpha_{v}\cdot
v)-F^{\mathsf{tl}}(\Phi;u,v)=$
$\displaystyle(\alpha_{v}-1)\cdot\max_{\phi\in\Phi}\langle
v,\phi\rangle-\max_{\phi\in\Phi}\langle u+\alpha_{v}\cdot
v,\phi\rangle+\max_{\phi\in\Phi}\langle u+v,\phi\rangle\geq 0,$

where the final inequality follows since $u+\alpha_{v}\cdot
v=(u+v)+(\alpha_{v}-1)\cdot v$. ∎

###  4.2 An upper bound on the truncated linear bonus

Lemma 4.8 below shows that the truncated linear bonus
$F^{\mathsf{tl}}(\Phi;u,v)$ can be bounded above by the length of the
projection of $\Phi$ onto each of $u$ and $v$.

######  Lemma 4.8.

Suppose $u,v\in\mathbb{R}^{d}$ and $\Phi\in\mathscr{P}^{d}$. Then

$\displaystyle 0\leq F^{\mathsf{tl}}(\Phi;u,v)\leq
2\cdot\min\left\\{\max_{\phi,\phi^{\prime}\in\Phi}\langle
u,\phi-\phi^{\prime}\rangle,\max_{\phi,\phi^{\prime}\in\Phi}\langle
v,\phi-\phi^{\prime}\rangle\right\\}.$

###### Proof.

For any fixed $\phi$, we have

$\displaystyle\langle u+v,\phi\rangle\leq\max_{\phi^{\prime}\in\Phi}\langle
u,\phi^{\prime}\rangle+\max_{\phi^{\prime}\in\Phi}\langle
v,\phi^{\prime}\rangle,$

which implies $F^{\mathsf{tl}}(\Phi;u,v)\geq 0$.

To upper bound $F^{\mathsf{tl}}(\Phi;u,v)$, we first note that

$\displaystyle F^{\mathsf{tl}}(\Phi;u,v)\leq$ $\displaystyle
F^{\mathsf{tl}}(\Phi;u,v)+F^{\mathsf{tl}}(\Phi;-u,-v)$ $\displaystyle=$
$\displaystyle\max_{\phi,\phi^{\prime}\in\Phi}\langle
u,\phi-\phi^{\prime}\rangle+\max_{\phi,\phi^{\prime}\in\Phi}\langle
v,\phi-\phi^{\prime}\rangle-\max_{\phi,\phi^{\prime}\in\Phi}\langle
u+v,\phi-\phi^{\prime}\rangle.$

We will now upper bound
$F^{\mathsf{tl}}(\Phi;u,v)+F^{\mathsf{tl}}(\Phi;-u,-v)$ as follows. By
symmetry, it is without loss of generality to assume that
$\max_{\phi,\phi^{\prime}\in\Phi}\langle
u,\phi-\phi^{\prime}\rangle\leq\max_{\phi,\phi^{\prime}\in\Phi}\langle
v,\phi-\phi^{\prime}\rangle$. Now we have

$\displaystyle\max_{\phi,\phi^{\prime}\in\Phi}\langle
u+v,\phi-\phi^{\prime}\rangle\geq$
$\displaystyle\max_{\phi,\phi^{\prime}\in\Phi}\langle
v,\phi-\phi^{\prime}\rangle-\max_{\phi,\phi^{\prime}\in\Phi}\langle-u,\phi-\phi^{\prime}\rangle$
$\displaystyle=$ $\displaystyle\max_{\phi,\phi^{\prime}\in\Phi}\langle
v,\phi-\phi^{\prime}\rangle-\max_{\phi,\phi^{\prime}\in\Phi}\langle
u,\phi-\phi^{\prime}\rangle.$

Rearranging, we see that
$F^{\mathsf{tl}}(\Phi;u,v)+F^{\mathsf{tl}}(\Phi;-u,-v)\leq
2\max_{\phi,\phi^{\prime}\in\Phi}\langle u,\phi-\phi^{\prime}\rangle$, as
desired. ∎

###  4.3 A lower bound on the modified truncated linear bonus

Given $\Phi\in\mathscr{P}^{d}$ and a matrix $\Gamma\in\mathbb{R}^{d\times d}$,
let $\mathcal{S}_{\Phi}\subset\mathbb{R}^{d}$ denote the subspace
$\mathcal{S}_{\Phi}:={\rm Span}\\{\phi:\phi\in\Phi\\}$, and
$\mathcal{S}_{\Phi,\Gamma}\subset\mathbb{R}^{d}$ denote the subspace
$\mathcal{S}_{\Phi,\Gamma}:={\rm Span}\\{\Gamma\cdot\phi:\
\phi\in\Phi\\}=\Gamma\cdot\mathcal{S}_{\Phi}$. Let
$\Pi_{\Phi,\Gamma}\in\mathbb{R}^{d\times d}$ denote the matrix which projects
onto $\mathcal{S}_{\Phi,\Gamma}$. Recall that for an orthogonal pair
$(\Sigma,\Lambda)$ (Definition 4.1), we have
$\mathcal{S}_{\Sigma}=\Sigma\cdot\mathbb{R}^{d},\mathcal{S}_{\Lambda}=\Lambda\cdot\mathbb{R}^{d}$.
Then for any $\Phi\in\mathscr{P}^{d}$,
$\mathcal{S}_{\Phi,\Sigma}\subset\mathcal{S}_{\Sigma}$ and
$\mathcal{S}_{\Phi,\Lambda}\subset\mathcal{S}_{\Lambda}$.

Moreover, given $\Phi\in\mathscr{P}^{d}$ and $\phi_{1},\phi_{2}\in\Phi$, we
define

$\displaystyle\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}:=\operatorname*{arg\,min}_{\xi\in\Phi}\left\\{\|\beta\Sigma\cdot(\phi_{1}-\xi)\|_{2}+\|\Lambda\cdot(\xi-\phi_{2})\|_{2}\right\\}.$
(22)

The point $\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}$ can be thought of as a sort
of “midpoint” between $\phi_{1}$ and $\phi_{2}$ in $\Phi$, where the distance
between $\phi_{1}$ and $\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}$ is measured
with respect to $\beta\Sigma$ and the distance between
$\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}$ and $\phi_{2}$ is measured with
respect to $\Lambda$. As it turns out, this particular notion of midpoint will
be useful in proving that our exploration bonuses induce optimistic value
functions.

Finally, we introduce some notation regarding the normal cone of a convex
body. For a convex subset $\mathcal{C}\subset\mathbb{R}^{d}$ and
$z\in\mathcal{C}$, the normal cone of $\mathcal{C}$ at $z$, denoted
$N_{\mathcal{C}}(z)$ is defined as
$N_{\mathcal{C}}(z):=\\{y\in\mathbb{R}^{d}:\ \langle y,z-z^{\prime}\rangle\geq
0\ \forall z^{\prime}\in\mathcal{C}\\}$.

######  Lemma 4.9.

Consider an orthogonal pair $(\Sigma,\Lambda)$ and $\beta\geq 1,\zeta>0$.
Suppose that $\Phi\in\mathscr{P}^{d}$ satisfies

$\displaystyle\max\left\\{\max_{\phi_{1},\phi_{2}\in\Phi}\|\beta\Sigma\cdot(\phi_{1}-\phi_{2})\|_{2},\
\max_{\phi_{1},\phi_{2}\in\Phi}\|\Lambda\cdot(\phi_{1}-\phi_{2})\|_{2}\right\\}\leq\zeta.$
(23)

Fix any $\phi_{1},\phi_{2}\in\Phi$. Then there are unit vectors
$u\in\mathcal{S}_{\Phi,\Sigma}$ and $v\in\mathcal{S}_{\Phi,\Lambda}$
(depending on $\Phi,\phi_{1},\phi_{2}$) so that for all unit vectors
$u^{\prime}\in\mathcal{S}_{\Sigma},v^{\prime}\in\mathcal{S}_{\Lambda}$
satisfying

$\displaystyle\|u-u^{\prime}\|_{2}\leq\varepsilon,\qquad\|v-v^{\prime}\|_{2}\leq\varepsilon,\qquad\langle
u^{\prime},u\rangle\geq\eta,\qquad\langle v^{\prime},v\rangle\geq\eta,$ (24)

it holds that

$\displaystyle F^{\mathsf{tl}}(\Phi;\beta u^{\prime},v^{\prime})$
$\displaystyle\geq\eta\|\beta\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}+\eta\|\Lambda\cdot(\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}-\phi_{2})\|_{2}-2\varepsilon\zeta.$

###### Proof.

Fix $\Phi\in\mathscr{P}^{d}$ satisfying Eq. 23 and $\phi_{1},\phi_{2}\in\Phi$.

##### Step 1: Choosing $u,v$.

The KKT conditions for optimality of $\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}$
(as defined in Eq. 22) give that, for all $z\in\bar{\Phi}$,

$\displaystyle\frac{\beta^{2}\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})}{\|\beta\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}}+\frac{\Lambda\cdot(\phi_{2}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})}{\|\Lambda\cdot(\phi_{2}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}}\in
N_{\Phi}(\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}),$ (25)

where $N_{\Phi}(\xi)$ denotes the normal cone of $\Phi$ at $\xi$.We now set

$\displaystyle
u:=\frac{\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})}{\|\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}},\qquad
v:=\frac{\Lambda\cdot(\phi_{2}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})}{\|\Lambda\cdot(\phi_{2}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}}.$

Then Eq. 25 states that for all $z\in\Phi$, $\langle\beta
u+v,\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}-z\rangle\geq 0$. Moreover, since
$\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}\in\Phi$, we certainly have that
$u\in\mathcal{S}_{\Phi,\Sigma}\subset\mathcal{S}_{\Sigma},\
v\in\mathcal{S}_{\Phi,\Lambda}\subset\mathcal{S}_{\Lambda}$.

##### Step 2: Relating $u,v$ to $u^{\prime},v^{\prime}$ satisfying Eq. 24.

Now consider any unit vectors
$u^{\prime}\in\mathcal{S}_{\Sigma},v^{\prime}\in\mathcal{S}_{\Lambda}$
satisfying Eq. 24. Choose

$\displaystyle\xi^{\prime}:=\operatorname*{arg\,max}_{\xi\in\Phi}\langle\beta
u^{\prime}+v^{\prime},\xi\rangle.$ (26)

Next, we compute

$\displaystyle\langle\beta
u^{\prime}+v^{\prime},\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}\rangle=$
$\displaystyle\langle(\beta u^{\prime}+v^{\prime})-(\beta
u+v),\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}\rangle+\langle\beta
u+v,\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}\rangle$
$\displaystyle\leq$ $\displaystyle\beta|\langle
u-u^{\prime},\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}\rangle|+|\langle
v-v^{\prime},\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}\rangle|$
$\displaystyle=$ $\displaystyle\beta|\langle
u-u^{\prime},\Sigma\cdot(\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\rangle|+|\langle
v-v^{\prime},\Lambda\cdot(\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\rangle|$
$\displaystyle=$
$\displaystyle\beta|\langle\Pi_{\Phi,\Sigma}\cdot(u-u^{\prime}),\Sigma\cdot(\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\rangle|+|\langle\Pi_{\Phi,\Lambda}\cdot(v-v^{\prime}),\Lambda\cdot(\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\rangle|$
$\displaystyle\leq$
$\displaystyle\beta\|\Pi_{\Phi,\Sigma}\cdot(u-u^{\prime})\|_{2}\cdot\|\Sigma\cdot(\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}+\|\Pi_{\Phi,\Lambda}\cdot(v-v^{\prime})\|_{2}\cdot\|\Lambda\cdot(\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}$
$\displaystyle\leq$ $\displaystyle 2\varepsilon\zeta,$ (27)

where the first inequality uses the triangle inequality and the fact that
$\langle\beta
u+v,\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}\rangle\leq 0$ since
$\xi^{\prime}\in\Phi$; the second equality uses the fact that
$u-u^{\prime}\in\mathcal{S}_{\Sigma}$ and
$v-v^{\prime}\in\mathcal{S}_{\Lambda}$; the third equality uses the fact that
$\Sigma\cdot(\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\in\mathcal{S}_{\Phi,\Sigma}$
and
$\Lambda\cdot(\xi^{\prime}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\in\mathcal{S}_{\Phi,\Lambda}$;
the second inequality uses Cauchy-Schwarz; and the third inequality uses Eq.
24 and Eq. 23.

We then have

$\displaystyle\langle\phi_{1}-\xi^{\prime},\beta
u^{\prime}\rangle+\langle\phi_{2}-\xi^{\prime},v^{\prime}\rangle$
$\displaystyle=$ $\displaystyle\langle\Sigma\cdot(\phi_{1}-\xi^{\prime}),\beta
u^{\prime}+v^{\prime}\rangle+\langle\Lambda\cdot(\phi_{2}-\xi^{\prime}),v^{\prime}+\beta
u^{\prime}\rangle$ $\displaystyle=$
$\displaystyle\langle\Sigma\cdot\phi_{1}+\Lambda\cdot\phi_{2}-\xi^{\prime},\beta
u^{\prime}+v^{\prime}\rangle$ $\displaystyle\geq$
$\displaystyle\langle\Sigma\cdot\phi_{1}+\Lambda\cdot\phi_{2}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}},\beta
u^{\prime}+v^{\prime}\rangle-2\varepsilon\zeta$ $\displaystyle=$
$\displaystyle\langle\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}),\beta
u^{\prime}\rangle+\langle\Lambda\cdot(\phi_{2}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}),v^{\prime}\rangle-2\varepsilon\zeta$
$\displaystyle\geq$
$\displaystyle\eta\beta\|\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}+\eta\|\Lambda\cdot(\phi_{2}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}-2\varepsilon\zeta,$
(28)

where the first equality uses that $\Sigma u^{\prime}=u^{\prime},\Sigma
v^{\prime}=v^{\prime}$, and $\Sigma v^{\prime}=\Lambda u^{\prime}=0$ since
$(\Sigma,\Lambda)$ is an orthogonal pair, the second equality uses that
$\Sigma+\Lambda=I_{d}$, the first inequality uses Eq. 27, the third equality
again uses $\Sigma+\Lambda=I_{d}$ and the fact that $\Sigma v^{\prime}=\Lambda
u^{\prime}=0$, and the second inequality uses Eq. 24 and the definition of
$u,v$.

##### Step 3: wrapping up.

Finally, we may write

$\displaystyle F^{\mathsf{tl}}(\Phi;\beta u^{\prime},v^{\prime})=$
$\displaystyle\max_{\phi\in\Phi}\langle\phi,\beta
u^{\prime}\rangle+\max_{\phi\in\Phi}\langle\phi,v^{\prime}\rangle-\max_{\phi\in\Phi}\langle\phi,\beta
u^{\prime}+v^{\prime}\rangle$ $\displaystyle=$
$\displaystyle\max_{\phi\in\Phi}\langle\phi,\beta
u^{\prime}\rangle+\max_{\phi\in\Phi}\langle\phi,v^{\prime}\rangle-\langle\xi,\beta
u^{\prime}+v^{\prime}\rangle$ $\displaystyle\geq$
$\displaystyle\langle\phi_{1},\beta
u^{\prime}\rangle+\langle\phi_{2},v^{\prime}\rangle-\langle\xi,\beta
u^{\prime}+v^{\prime}\rangle$ $\displaystyle\geq$
$\displaystyle\eta\beta\|\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}+\eta\|\Lambda\cdot(\phi_{2}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}-2\varepsilon\zeta,$

where the second equality uses the definition of $\xi$ in Eq. 26, and the
second inequality uses Eq. 28. ∎

By averaging over $u^{\prime},v^{\prime}$ drawn randomly from appropriate
subspaces, we have the following consequence of Lemma 4.9.

######  Lemma 4.10.

There is a constant $C_{\ref{cor:optimal-perimeter}}$ so that the following
holds. Fix $A\in\mathbb{N}$, an orthogonal pair $(\Sigma,\Lambda)$, $\beta\geq
1$, $\zeta>0$, $\Phi\in\mathscr{P}_{A}^{d}$, $\phi_{1},\phi_{2}\in\Phi$, and
suppose that $\Phi$ satisfies Eq. 23. Then for any $\varepsilon>0$,

$\displaystyle\left(\frac{C_{\ref{cor:optimal-
perimeter}}}{\varepsilon}\right)^{2A}\cdot\mathbb{E}_{u^{\prime}\sim\mathcal{N}(0,\Sigma)}\mathbb{E}_{v^{\prime}\sim\mathcal{N}(0,\Lambda)}\left[F^{\mathsf{tl}}(\Phi;\beta
u^{\prime},v^{\prime})\right]$ $\displaystyle\geq$
$\displaystyle\|\beta\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}+\|\Lambda\cdot(\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}-\phi_{2})\|_{2}-4\varepsilon\zeta.$

###### Proof.

Note that if $\dim(\mathcal{S}_{\Phi,\Sigma})=0$, then we could choose
$\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}=\phi_{2}$ in Eq. 22, and thus the
lemma statement becomes immediate. A symmetric argument applies to the case
that $\dim(\mathcal{S}_{\Phi,\Lambda})=0$. Thus we may assume for the
remainder of the proof that
$\dim(\mathcal{S}_{\Phi,\Sigma}),\dim(\mathcal{S}_{\Phi,\Lambda})\geq 1$.

Note that the subspace
$\mathcal{S}_{\Phi,\Sigma}\subseteq\mathcal{S}_{\Sigma}$ satisfies
$\dim(\mathcal{S}_{\Phi,\Sigma})\leq A$, and similarly,
$\dim(\mathcal{S}_{\Phi,\Lambda})\leq A$. Note that
$\mathbb{E}_{u^{\prime}\sim\mathcal{N}(0,\Sigma)}[(\Pi_{\Phi,\Sigma}u^{\prime})\cdot(\Pi_{\Phi,\Sigma}u^{\prime})^{\top}]=\Pi_{\Phi,\Sigma}\cdot\Sigma\cdot\Pi_{\Phi,\Sigma}=\Pi_{\Phi,\Sigma}$,
meaning that
$\Pi_{\Phi,\Sigma}u^{\prime}\sim\mathcal{N}(0,\Pi_{\Phi,\Sigma})$. Similarly,
we have that
$\Pi_{\Phi,\Lambda}v^{\prime}\sim\mathcal{N}(0,\Pi_{\Phi,\Lambda})$.

Therefore, we have that
$\frac{\Pi_{\Phi,\Sigma}u^{\prime}}{\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}}\sim\mathsf{unif}(S^{d-1}\cap\mathcal{S}_{\Phi,\Sigma})$.
Since $S^{d-1}\cap\mathcal{S}_{\Phi,\Sigma}$ may be covered with
$(3/\varepsilon)^{A}$ balls of radius $\varepsilon$, it follows that, for any
$u\in S^{d-1}\cap\mathcal{S}_{\Phi,\Sigma}$, we have

$\displaystyle\operatorname{Pr}_{u^{\prime}\sim\mathcal{N}(0,\Sigma)}\left(\left\|\frac{\Pi_{\Phi,\Sigma}u^{\prime}}{\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}}-u\right\|_{2}\leq\varepsilon\right)\geq(\varepsilon/3)^{A}.$

Moreover, we have that
$\operatorname{Pr}_{u^{\prime}\sim\mathcal{N}(0,\Sigma)}(\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}\geq
1/2)\geq 1/2$: this holds because $\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}^{2}$ is
a chi-squared random variable with $\dim(\mathcal{S}_{\Phi,\Sigma})\geq 1$
degrees of freedom. Using the fact that $\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}$
and $\frac{\Pi_{\Phi,\Sigma}u^{\prime}}{\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}}$
are independent random variables, it follows that

$\displaystyle\frac{1}{2}\cdot(\varepsilon/3)^{A}\leq$
$\displaystyle\operatorname{Pr}_{u^{\prime}\sim\mathcal{N}(0,\Sigma)}\left(\left\|\frac{\Pi_{\Phi,\Sigma}u^{\prime}}{\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}}-u\right\|_{2}\leq\varepsilon\mbox{
and }\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}\geq 1/2\right).$

In a symmetric manner, we obtain that

$\displaystyle\frac{1}{2}\cdot(\varepsilon/3)^{A}\leq$
$\displaystyle\operatorname{Pr}_{v^{\prime}\sim\mathcal{N}(0,\Lambda)}\left(\left\|\frac{\Pi_{\Phi,\Lambda}v^{\prime}}{\|\Pi_{\Phi,\Lambda}v^{\prime}\|_{2}}-v\right\|_{2}\leq\varepsilon\mbox{
and }\|\Pi_{\Phi,\Lambda}v^{\prime}\|_{2}\geq 1/2\right).$

Let us define the random variables
$u^{\prime\prime}:=\frac{\Pi_{\Phi,\Sigma}u^{\prime}}{\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}}$
and
$v^{\prime\prime}:=\frac{\Pi_{\Phi,\Lambda}v^{\prime}}{\|\Pi_{\Phi,\Lambda}v^{\prime}\|_{2}}$.
Let $u\in\mathcal{S}_{\Phi,\Sigma}$ and $v\in\mathcal{S}_{\Phi,\Lambda}$ be
chosen according to the statement of Lemma 4.9 given $\Phi,\phi_{1},\phi_{2}$.
Then with probability at least $(1/4)\cdot(\varepsilon/3)^{2A}$ over the
independent draws of $u^{\prime}$ and $v^{\prime}$ (which induce values of
$u^{\prime\prime},v^{\prime\prime}$), we have that:

$\displaystyle\|u^{\prime\prime}-u\|_{2}\leq\varepsilon,\quad\|v^{\prime\prime}-v\|_{2}\leq\varepsilon,\quad\langle
u^{\prime\prime},u\rangle\geq 1/2,\quad\langle v^{\prime\prime},v\rangle\geq
1/2,\quad\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}\geq
1/2,\quad\|\Pi_{\Phi,\Lambda}v^{\prime}\|_{2}\geq 1/2,$ (29)

where we have used that $\varepsilon\leq 1$ and that for unit vectors
$y,y^{\prime}$, we have $\langle
y,y^{\prime}\rangle=\frac{2-\|y-y^{\prime}\|_{2}^{2}}{2}$.

For any $\phi\in\Phi\subset\mathcal{S}_{\Phi}$, we have

$\displaystyle\langle u^{\prime},\phi\rangle=\langle
u^{\prime},\Sigma\phi\rangle=\langle\Pi_{\Phi,\Sigma}u^{\prime},\Sigma\phi\rangle=\langle\Pi_{\Phi,\Sigma}u^{\prime},\phi\rangle,$

where the first equality uses that $u^{\prime}\in\mathcal{S}_{\Sigma}$, the
second equality uses that $\Sigma\phi\in\mathcal{S}_{\Phi,\Sigma}$, and the
third equality uses that
$\Pi_{\Phi,\Sigma}u^{\prime}\in\mathcal{S}_{\Phi,\Sigma}\subseteq\mathcal{S}_{\Sigma}$.
In a similar manner, we have that, for all $\phi\in\Phi$, $\langle
v^{\prime},\phi\rangle=\langle\Pi_{\Phi,\Lambda}v^{\prime},\phi\rangle$. It
then follows from Lemma 4.6 that

$\displaystyle F^{\mathsf{tl}}(\Phi;\beta
u^{\prime},v^{\prime})=F^{\mathsf{tl}}(\Phi;\beta\cdot\Pi_{\Phi,\Sigma}u^{\prime},\Pi_{\Phi,\Lambda}v^{\prime}).$
(30)

Under the event that $\|\Pi_{\Phi,\Sigma}u^{\prime}\|_{2}\geq 1/2$ and
$\|\Pi_{\Phi,\Lambda}v^{\prime}\|_{2}\geq 1/2$, we have by Lemma 4.7 that

$\displaystyle
F^{\mathsf{tl}}(\Phi;\beta\Pi_{\Phi,\Sigma}u^{\prime},\Pi_{\Phi,\Lambda}v^{\prime})\geq\frac{1}{2}\cdot
F^{\mathsf{tl}}(\Phi;\beta u^{\prime\prime},v^{\prime\prime}).$ (31)

Combining Eqs. 30 and 31, we see that

$\displaystyle\mathbb{E}_{u^{\prime}\sim\mathcal{N}(0,\Sigma)}\mathbb{E}_{v^{\prime}\sim\mathcal{N}(0,\Lambda)}\left[F^{\mathsf{tl}}(\Phi;\beta
u^{\prime},v^{\prime})\right]$ $\displaystyle\geq$
$\displaystyle\frac{1}{2}\mathbb{E}_{u^{\prime}\sim\mathcal{N}(0,\Sigma)}\mathbb{E}_{v^{\prime}\sim\mathcal{N}(0,\Lambda)}\left[\widetilde{F}^{\mathsf{tl}}(\Phi;\beta
u^{\prime\prime},v^{\prime\prime})\right]$ $\displaystyle\geq$
$\displaystyle\frac{1}{4}\cdot\frac{1}{4}\cdot\left(\frac{\varepsilon}{3}\right)^{2A}\cdot\left(\|\beta\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}+\|\Lambda\cdot(\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}-\phi_{2})\|_{2}-4\varepsilon\zeta\right),$

where the second inequality uses the fact that Eq. 29 holds with probability
at least $(1/4)\cdot(\varepsilon/3)^{2A}$ together with Lemma 4.9 with
$\eta=1/2$, $u^{\prime}$ set to $u^{\prime\prime}$, and $v^{\prime}$ set to
$v^{\prime\prime}$. Rearranging and using that $A\geq 2$, we obtain that

$\displaystyle\left(\frac{6}{\varepsilon}\right)^{2A}\cdot\mathbb{E}_{u^{\prime}\sim\mathcal{N}(0,\Sigma)}\mathbb{E}_{v^{\prime}\sim\mathcal{N}(0,\Lambda)}\left[\widetilde{F}^{\mathsf{tl}}(\Phi;\beta
u^{\prime},v^{\prime})\right]$ $\displaystyle\geq$
$\displaystyle\|\beta\Sigma\cdot(\phi_{1}-\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}})\|_{2}+\|\Lambda\cdot(\mu^{\star}_{{\Phi,\phi_{1},\phi_{2}}}-\phi_{2})\|_{2}-4\varepsilon\zeta,$

as desired. (In particular, we may take $C_{\ref{cor:optimal-perimeter}}=6$.)
∎

##  5 The algorithm

##### The algorithm.

Algorithm 1 presents our algorithm for learning linear Bellman complete MDPs.
For some $T\in\mathbb{N}$, the algorithm proceeds for $T$ _rounds_. In each
round $t\in[T]$, the algorithm computes two policies
$\hat{\pi}_{h}^{(t)},\tilde{\pi}_{h}^{(t)}:\mathcal{X}\rightarrow\mathcal{A}$
at step $h$, in order of decreasing $h$. For each step $h$, trajectories are
gathered from the uniform mixture over
$\hat{\pi}^{(s)}\circ_{h}\tilde{\pi}^{(s)}\circ_{h+1}\hat{\pi}^{(t)}$, for
$s\in[t-1]$ (4). These trajectories are then used, together with a procedure
to construct an optimism-based upper confidence bonus, to define a policy
$\hat{\pi}_{h}^{(t)}$ (7 and 8), which represents the algorithm’s current
optimistic estimate of the optimal policy. Finally, $\tilde{\pi}_{h}^{(t)}$ is
defined to be a sort of uniform policy (10), with respect to the current
covariance matrix of features $\Sigma_{h}^{(t)}$.

##### Parameter definitions.

Below, we specify the parameters $\lambda_{1},\beta$ used in the definition of
the bonuses $F_{h}^{(t)}$ (in Eq. 33 below), the parameters $\lambda,n,T$ used
in Algorithm 1, as well as several other parameters used in the analysis of
this section:

######  Definition 5.1 (Parameter definitions).

Fix $\delta\in(0,1),\varepsilon_{\mathsf{final}}\in(0,1)$, as well as
parameters $A,H,B,d$ of the MDP. We define the following parameters, which are
used in Algorithm 1 as well in the remainder of this section:

  * •

$\lambda=C_{\ref{lem:psd-concentration}}d\log(2THn/\delta)$, where
$C_{\ref{lem:psd-concentration}}$ is a constant chosen sufficiently large in
Lemma 6.3.

  * •

$T=d\cdot\left(\frac{C_{\ref{thm:policy-
learning}}H^{4}B^{3}dA^{1/2}\log^{1/2}(HABd/(\varepsilon_{\mathsf{final}}\delta))}{\varepsilon_{\mathsf{final}}}\right)^{6A+2}$,
where $C_{\ref{thm:policy-learning}}$ is a constant chosen sufficiently large
in the proof of Theorem 5.10.

  * •

$n=3T$.

  * •

$\iota={\log\left(\frac{TH(\lambda d+n)}{\delta}\right)}$.

  * •

$\lambda_{1}=BH$.

  * •

$\varepsilon_{\mathsf{bkup}}=\frac{\varepsilon_{\mathsf{final}}}{2H}$.

  * •

$\sigma_{\mathsf{tr}}=\frac{\varepsilon_{\mathsf{bkup}}}{4\lambda_{1}}$.

  * •

$\varepsilon_{\mathsf{apx}}=\frac{\varepsilon_{\mathsf{bkup}}}{128\sqrt{2\pi}C_{\ref{lem:phi-
what-wt-diff}}\lambda_{1}^{2}HBd\sqrt{\iota}}$, where $C_{\ref{lem:phi-what-
wt-diff}}$ is chosen sufficiently large in the proof of Lemma 5.7 holds.

  * •

$\beta=4C_{\ref{lem:phi-what-wt-
diff}}HB\sqrt{d\iota}\cdot{5\lambda_{1}\sqrt{d}}\cdot\left(\frac{C_{\ref{cor:optimal-
perimeter}}}{\varepsilon_{\mathsf{apx}}}\right)^{2A}$.

  * •

$\xi=\frac{\beta}{4C_{\ref{lem:phi-what-wt-diff}}HB\sqrt{d\iota}\cdot
2\sqrt{2\pi}\lambda_{1}\sqrt{d}}$. (Note that $\xi\geq 1$.)

##### Definition of bonuses.

Next we define the exploration bonus function
$F_{g}^{(t)}:\mathcal{X}\rightarrow\mathbb{R}$ used in the construction of the
rewards $\hat{r}_{h}^{(t,i,h)}$ in 6. The key component in $F_{g}^{(t)}$ is a
linear combination of truncated linear bonuses $F^{\mathsf{tl}}_{h}(x;u,v)$,
over various values of $u,v$. The vectors $u,v$ are chosen to belong to the
subspaces $\mathcal{S}_{\Sigma},\mathcal{S}_{\Lambda}$, for matrices
$\Sigma,\Lambda$ forming an orthogonal pair which are obtained from the
covariance matrix $\Sigma_{h}^{(t)}$ via a certain truncation procedure,
defined below:

######  Definition 5.2.

Let $\Gamma\in\mathbb{R}^{d\times d}$ be a PSD matrix, and suppose that
$\Gamma$ may be diagonalized as $\Gamma=UDU^{\top}$, for a diagonal matrix $D$
and an orthogonal matrix $U$. Given $\sigma>0$, we define the _$\sigma$
-truncated pair of $\Gamma$_, denoted
$(\Sigma^{\prime},\Lambda^{\prime}):=\mathsf{trunc}_{\sigma}({\Gamma})$, as
$\Sigma^{\prime}=UD_{+}U^{\top},\Lambda^{\prime}=UD_{-}U^{\top}$, where
$D_{+},D_{-}$ are diagonal matrices whose entries are given by:

$\displaystyle(D_{+})_{ii}=\begin{cases}1&:D_{ii}\geq\sigma\\\
0&:D_{ii}<\sigma\end{cases},\quad(D_{-})_{ii}=\begin{cases}0&:D_{ii}\geq\sigma\\\
1&:D_{ii}<\sigma\end{cases}.$ (32)

In words, for
$(\Sigma^{\prime},\Lambda^{\prime})=\mathsf{trunc}_{\sigma}({\Gamma})$,
$\Sigma^{\prime}$ contains the large components of $\Gamma$ (as parametrized
by $\sigma$), and $\Lambda^{\prime}$ contains the small components of
$\Gamma$. Note that the definition of $\mathsf{trunc}_{\sigma}({\Gamma})$ does
not depend on the choice of diagonalization. Moreover, note that
$(\Sigma^{\prime},\Lambda^{\prime})=\mathsf{trunc}_{\sigma}({\Gamma})$ is an
orthogonal pair.

The bonus function $F_{h}^{(t)}$ used in Algorithm 1 is defined as follows:
for each $t\in[T]$, $h\in[H]$,

$\displaystyle F_{h}^{(t)}(x):=$
$\displaystyle\lambda_{1}\cdot\left(\frac{C_{\ref{cor:optimal-
perimeter}}}{\varepsilon_{\mathsf{apx}}}\right)^{2A}\cdot\mathbb{E}_{u^{\prime}\sim\mathcal{N}(0,\Sigma^{\prime})}\mathbb{E}_{v^{\prime}\sim\mathcal{N}(0,\Lambda^{\prime})}\left[F^{\mathsf{tl}}_{h}(x;\beta
u^{\prime},v^{\prime})\right]$
$\displaystyle+2\sqrt{2\pi}\lambda_{1}\xi\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\right]$ (33) for
$\displaystyle\quad\quad(\Sigma^{\prime},\Lambda^{\prime}):=\mathsf{trunc}_{\sigma_{\mathsf{tr}}}({(\beta/\lambda_{1})\cdot(\Sigma_{h}^{(t)})^{-1/2}}).$

#####  $Q$\- and $V$-function definitions.

For each $h\in[H],t\in[T]$, we make the following definitions:

$\displaystyle Q_{h}^{(t)}(x,a):=$ $\displaystyle
r_{h}(x,a)+\mathbb{E}^{\hat{\pi}^{(t)}}\left[\sum_{g=h+1}^{H}r_{g}(x_{g},a_{g})+F_{g}^{(t)}(x_{g})\
\mid\ (x_{h},a_{h})=(x,a)\right]$ (34) $\displaystyle V_{h}^{(t)}(x):=$
$\displaystyle Q_{h}^{(t)}(x,\hat{\pi}_{h}^{(t)}(x))$ (35)
$\displaystyle\hat{Q}_{h}^{(t)}(x,a):=$
$\displaystyle\langle\hat{w}_{h}^{(t)},\phi_{h}(x,a)\rangle$
$\displaystyle\hat{V}_{h}^{(t)}(x):=$
$\displaystyle\hat{Q}_{h}^{(t)}(x,\hat{\pi}_{h}^{(t)}(x))=\max_{a\in\mathcal{A}}\hat{Q}_{h}^{(t)}(x,a).$

The functions $Q_{h}^{(t)}(\cdot)$ and $V_{h}^{(t)}(\cdot)$ represent the
$Q$\- and $V$-value functions for the policy $\hat{\pi}^{(t)}$, with respect
to the rewards $\hat{r}_{g}^{(t)}$ defined in Eq. 6. In Lemma 5.1 below, we
show that $Q_{h}^{(t)}$ is a linear function, i.e., for some
$w_{h}^{(t)}\in\mathbb{R}^{d}$,
$Q_{h}^{(t)}(x,a)=\langle\phi_{h}(x,a),w_{h}^{(t)}\rangle$ for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$. The vector $\hat{w}_{h}^{(t)}$
computed in Algorithm 1 may be viewed as an empirical approximation to
$w_{h}^{(t)}$, so that $\hat{Q}_{h}^{(t)}(\cdot)$ and
$\hat{V}_{h}^{(t)}(\cdot)$ represent empirical approximations to
$Q_{h}^{(t)},V_{h}^{(t)}$, respectively.

######  Lemma 5.1.

For all $t\in[T],h\in[h]$, there is some $w_{h}^{(t)}\in\mathbb{R}^{d}$ so
that $Q_{h}^{(t)}(x,a)=\langle\phi_{h}(x,a),w_{h}^{(t)}\rangle$ for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$.

###### Proof.

For all $x,a,h,t$, we have

$\displaystyle
Q_{h}^{(t)}(x,a)=r_{h}(x,a)+\mathbb{E}_{x^{\prime}\sim\mathbb{P}_{h}(x,a)}\left[F_{h+1}^{(t)}(x^{\prime})+Q_{h+1}^{(t)}(x^{\prime},\hat{\pi}_{h+1}^{(t)}(x^{\prime})))\right].$

For fixed $t\in[T]$, we will use reverse induction on $h$ to establish the
stated claim. The base case $h=H$ follows since
$Q_{H}^{(t)}(x,a)=r_{H}(x,a)=\langle w_{H}^{(t)},\phi_{H}(x,a)\rangle$ for
some $w_{H}\in\mathbb{R}^{d}$, by assumption.

Now, assume for the inductive hypothesis that for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$, $Q_{h+1}^{(t)}(x,a)=\langle
w_{h+1}^{(t)},\phi_{h+1}(x,a)\rangle$ for some
$w_{h+1}^{(t)}\in\mathbb{R}^{d}$. Then Lemma 4.3 gives that there is some
$\theta^{1}\in\mathbb{R}^{d}$ so that for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$,

$\displaystyle\langle\theta^{1},\phi_{h}(x,a)\rangle=\mathbb{E}_{x^{\prime}\sim\mathbb{P}_{h}(x,a)}\left[\langle\phi_{h+1}(x^{\prime},\hat{\pi}_{h+1}^{(t)}(x^{\prime})),w_{h+1}^{(t)}\rangle\right]=\mathbb{E}_{x^{\prime}\sim\mathbb{P}_{h}(x,a)}\left[Q_{h+1}^{(t)}(x^{\prime},\hat{\pi}_{h+1}^{(t)}(x^{\prime}))\right].$

Next, note that it is immediate from the definition Eq. 33 that, with
probability 1 over $u^{\prime}\sim\mathcal{N}(0,\Sigma^{\prime})$ and
$v^{\prime}\sim\mathcal{N}(0,\Lambda^{\prime})$, the function
$F^{\mathsf{tl}}_{h+1}(\cdot;u^{\prime},v^{\prime},\beta\Sigma^{\prime},\Lambda)$
is Bellman-linear at step $h+1$. Note that the function
$x\mapsto\max_{a\in\mathcal{A}}\langle w,\phi_{h+1}(x,a)\rangle$ is also
Bellman-linear at step $h+1$, for all $w\in\mathbb{R}^{d}$. Since the sum of
Bellman-linear functions is Bellman-linear, the definition of $F_{h+1}^{(t)}$
in Eq. 33 gives that $F_{h+1}^{(t)}$ is Bellman-linear at step $h+1$. Thus,
there is some $\theta^{2}\in\mathbb{R}^{2}$ so that

$\displaystyle\langle\phi_{h}(x,a),\theta^{2}\rangle=\mathbb{E}_{x^{\prime}\sim\mathbb{P}_{h}(x,a)}\left[F_{h+1}^{(t)}(x^{\prime})\right]$

for all $(x,a)\in\mathcal{X}\times\mathcal{A}$.

Finally, by assumption, there is some $\theta^{3}\in\mathbb{R}^{d}$ so that
$r_{h}(x,a)=\langle\theta^{3},\phi_{h}(x,a)\rangle$ for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$. Hence, writing
$w_{h}^{(t)}=\theta^{1}+\theta^{2}+\theta^{3}$, we have that
$Q_{h}^{(t)}(x,a)=\langle w_{h}^{(t)},\phi_{h}(x,a)\rangle$ for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$. ∎

###  5.1 Boundedness of the bonuses

Next, we prove some lemmas which establish bounds on the magnitude of the
bonus functions $F_{h}^{(t)}(\cdot)$, and thereby on the magnitude of the
optimistic $Q$-functions $Q_{h}^{(t)}(\cdot)$.

######  Lemma 5.2.

For a PSD matrix $\Gamma$ and $\sigma>0$, writing
$(\Sigma^{\prime},\Lambda^{\prime}):=\mathsf{trunc}_{\sigma}({\Gamma})$, we
have that $\Sigma^{\prime}\preceq\frac{1}{\sigma}\cdot\Gamma$.

###### Proof.

We write $\Gamma=UDU^{\top}$ and $\Sigma^{\prime}=UD_{+}U^{\top}$, where
$D_{+}$ is as defined in Eq. 32. Certainly $0=(D_{+})_{ii}\leq D_{ii}$ if
$D_{ii}<\sigma$; otherwise, we have $1=(D_{+})_{ii}\leq\frac{1}{\sigma}\cdot
D_{ii}$, which implies that $D_{+}\preceq\frac{D}{\sigma}$, as desired. ∎

######  Lemma 5.3.

Let $\Phi\in\mathscr{P}^{d}$ be a given polyhedron. Consider any PSD matrix
$\Sigma\in\mathbb{R}^{d\times d}$. Then

$\displaystyle\frac{1}{\sqrt{2\pi}}\cdot\max_{\phi,\phi^{\prime}\in\Phi}((\phi-\phi^{\prime})^{\top}\Sigma(\phi-\phi^{\prime}))^{1/2}\leq\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma)}\left[\max_{\phi\in\Phi}\langle
w,\phi\rangle\right]\leq\sqrt{d}\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma)}\left[\phi_{w}^{\top}\Sigma\phi_{w}\right]^{1/2},$

where $\phi_{w}:=\operatorname*{arg\,max}_{\phi\in\Phi}\langle w,\phi\rangle$.

###### Proof.

Define $\Psi=\\{\Sigma^{1/2}\cdot\phi\ :\ \phi\in\Phi\\}$, and, for
$v\in\mathbb{R}^{d}$, write
$\psi_{v}:=\operatorname*{arg\,max}_{\psi\in\Psi}\langle v,\psi\rangle$. Then

$\displaystyle\Sigma^{1/2}\cdot\phi_{\Sigma^{1/2}v}=\Sigma^{1/2}\cdot\operatorname*{arg\,max}_{\phi\in\Phi}\langle
v,\Sigma^{1/2}\phi\rangle=\operatorname*{arg\,max}_{\psi\in\Psi}\langle
v,\psi\rangle=\psi_{v}.$

We then have the following equalities:

$\displaystyle\max_{\phi,\phi^{\prime}\in\Phi}((\phi-\phi^{\prime})^{\top}\Sigma(\phi-\phi^{\prime}))^{1/2}=$
$\displaystyle\max_{\psi,\psi^{\prime}\in\Psi}\|\psi-\psi^{\prime}\|_{2}$ (36)
$\displaystyle\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma)}\left[\phi_{w}^{\top}\Sigma\phi_{w}\right]^{1/2}=$
$\displaystyle\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}\left[\phi_{\Sigma^{1/2}v}^{\top}\Sigma\phi_{\Sigma^{1/2}v}\right]^{1/2}=\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}\left[\|\psi_{v}\|_{2}^{2}\right]^{1/2}$
(37)
$\displaystyle\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma)}\left[\max_{\phi\in\Phi}\langle
w,\phi\rangle\right]=$
$\displaystyle\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}\left[\max_{\phi\in\Phi}\langle\Sigma^{1/2}v,\phi\rangle\right]=\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}\left[\max_{\psi\in\Psi}\langle
v,\psi\rangle\right].$ (38)

It therefore suffices to show that

$\displaystyle\frac{1}{\sqrt{2\pi}}\cdot\max_{\psi,\psi^{\prime}\in\Psi}\|\psi-\psi^{\prime}\|_{2}\leq\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}\left[\max_{\psi\in\Psi}\langle
v,\psi\rangle\right]\leq\sqrt{d}\cdot\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}[\|\psi_{v}\|_{2}^{2}]^{1/2}.$
(39)

To show the first inequality in Eq. 39, pick $\psi_{0},\psi_{1}\in\Psi$
maximizing $\|\psi_{0}-\psi_{1}\|_{2}$. Then

$\displaystyle\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}\left[\max_{\psi\in\Psi}\langle
v,\psi\rangle\right]\geq$
$\displaystyle\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}\left[\max\\{\langle
v,\psi_{0}\rangle,\langle v,\psi_{1}\rangle\\}\right]$ $\displaystyle=$
$\displaystyle\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}\left[\langle
v,(\psi_{0}+\psi_{1})/2\rangle+\max\\{\langle
v,(\psi_{0}-\psi_{1})/2\rangle,\langle
v,(\psi_{1}-\psi_{0})/2\rangle\\}\right]$ $\displaystyle=$
$\displaystyle\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}[|\langle
v,(\psi_{0}-\psi_{1})/2\rangle|]$ $\displaystyle=$
$\displaystyle\frac{\sqrt{2}}{2\sqrt{\pi}}\|\psi_{0}-\psi_{1}\|_{2},$

where the second equality uses that
$\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}[\langle
v,(\psi_{0}+\psi_{1})/2\rangle]=0$, and the final inequality uses that
$\langle
v,(\psi_{0}-\psi_{1})/2\rangle\sim\mathcal{N}(0,\|(\psi_{0}-\psi_{1})/2\|_{2}^{2})$.

To show the second inequality in Eq. 39, note that

$\displaystyle\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}\left[\max_{\psi\in\Psi}\langle
v,\psi\rangle\right]=$
$\displaystyle\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}\left[\langle
v,\psi_{v}\rangle\right]$ $\displaystyle\leq$
$\displaystyle{\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}[\|v\|_{2}^{2}]}^{1/2}\cdot\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}[\|\psi_{v}\|_{2}^{2}]^{1/2}$
$\displaystyle=$
$\displaystyle\sqrt{d}\cdot\mathbb{E}_{v\sim\mathcal{N}(0,I_{d})}[\|\psi_{v}\|_{2}^{2}]^{1/2},$

where the inequality uses Cauchy-Schwarz. ∎

######  Lemma 5.4.

Suppose that $\Gamma\in\mathbb{R}^{d}$ is a PSD matrix, and let $\sigma>0$ be
given. Let
$(\Sigma^{\prime},\Lambda^{\prime}):=\mathsf{trunc}_{\sigma}({\Gamma})$. Then
for all $h\in[H],x\in\mathcal{X}$, $a,a^{\prime}\in\mathcal{A}$, and
$v\in\bar{\Phi}_{h}(x)$, it holds that

$\displaystyle\|\Gamma\cdot(\phi_{h}(x,a)-v)\|_{2}+\|v-\phi_{h}(x,a^{\prime})\|_{2}$
$\displaystyle\leq$
$\displaystyle\|\Gamma\|\cdot\|\Sigma^{\prime}\cdot(\phi_{h}(x,a)-v)\|_{2}+\|\Lambda^{\prime}\cdot(v-\phi_{h}(x,a^{\prime}))\|_{2}+\sqrt{2\pi}\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{\bar{a}\in\mathcal{A}}\langle
w,\phi_{h}(x,\bar{a})\right]+2\sigma.$

###### Proof.

Note that $\Gamma^{2}\preceq\|\Gamma\|^{2}\cdot\Sigma^{\prime}+\sigma^{2}\cdot
I_{d}$, so that

$\displaystyle\|\Gamma\cdot(\phi_{h}(x,a)-v)\|_{2}\leq$
$\displaystyle\|\Gamma\|\cdot\|\Sigma^{\prime}\cdot(\phi_{h}(x,a)-v)\|_{2}+\sigma\cdot\|\phi_{h}(x,a)-v\|_{2}\leq\|\Sigma^{\prime}\cdot(\phi_{h}(x,a)-v)\|_{2}+2\sigma\,$
(40)

where the final inequality uses that
$\max_{\phi\in\bar{\Phi}_{h}(x)}\|\phi\|_{2}\leq 1$ (Assumption 2.1). Next,
the fact that $\Lambda^{\prime}+\Sigma^{\prime}=I_{d}$ gives that

$\displaystyle\|v-\phi_{h}(x,a^{\prime})\|_{2}\leq$
$\displaystyle\|\Lambda^{\prime}\cdot(v-\phi_{h}(x,a^{\prime}))\|_{2}+\max_{a_{1},a_{2}\in\mathcal{A}}\|\Sigma^{\prime}\cdot(\phi_{h}(x,a_{1})-\phi_{h}(x,a_{2}))\|_{2}$
$\displaystyle\leq$
$\displaystyle\|\Lambda^{\prime}\cdot(v-\phi_{h}(x,a^{\prime}))\|_{2}+\sqrt{2\pi}\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{\bar{a}\in\mathcal{A}}\langle
w,\phi_{h}(x,\bar{a})\right],$ (41)

where the second inequality uses Lemma 5.3. The conclusion of the lemma
statement follows by adding Eqs. 40 and 41. ∎

######  Lemma 5.5 (Bound on bonuses in terms of $\Sigma^{\prime}$).

There is a sufficiently large constant $C_{\ref{lem:sigmap-bound}}$ so that
the following holds. Fix $t\in[T],h\in[H]$, and let
$(\Sigma^{\prime},\Lambda^{\prime}):=\mathsf{trunc}_{\sigma_{\mathsf{tr}}}({(\beta/\lambda_{1})\cdot(\Sigma_{h}^{(t)})^{-1/2}})$.
Then for all $x\in\mathcal{X}$,

$\displaystyle|F_{h}^{(t)}(x)|\leq$
$\displaystyle\sqrt{d}\beta\lambda_{1}\cdot\left(\frac{C_{\ref{lem:sigmap-
bound}}}{\varepsilon_{\mathsf{apx}}}\right)^{2A}\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\right].$

###### Proof.

Fix any
$u^{\prime}\in\mathcal{S}_{\Sigma^{\prime}},v^{\prime}\in\mathcal{S}_{\Lambda^{\prime}}$.
Then

$\displaystyle F^{\mathsf{tl}}_{h}(x;\beta u^{\prime},v^{\prime})\leq$
$\displaystyle
2\|u^{\prime}\|_{2}\cdot\max_{a,a^{\prime}\in\mathcal{A}}\|\beta\Sigma^{\prime}\cdot(\phi_{h}(x,a)-\phi_{h}(x,a^{\prime}))\|_{2}$
$\displaystyle\leq$ $\displaystyle
2\sqrt{2\pi}\beta\|u^{\prime}\|_{2}\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\right],$

where the first inequality uses Lemma 4.8 with $u=\beta u^{\prime}$,
$v=v^{\prime}$ as well as the fact that $\Sigma u^{\prime}=u^{\prime}$, and
the second inequality uses Lemma 5.3. Using the definition of
$F_{h}^{(t)}(\cdot)$ in Eq. 33, it follows that, for a sufficiently large
constant $C>0$,

$\displaystyle|F_{h}^{(t)}(x)|\leq$
$\displaystyle{\lambda_{1}}\cdot\left(\frac{C_{\ref{cor:optimal-
perimeter}}}{\varepsilon_{\mathsf{apx}}}\right)^{2A}\cdot\mathbb{E}_{u^{\prime}\sim\mathcal{N}(0,\Sigma^{\prime})}\left[4\sqrt{2\pi}\beta\|u^{\prime}\|_{2}\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\right]\right]$
$\displaystyle+2\sqrt{2\pi}\lambda_{1}\xi\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\right]$ $\displaystyle\leq$
$\displaystyle\left({4\sqrt{2\pi}\sqrt{d}\beta\lambda_{1}}\cdot\left(\frac{C_{\ref{cor:optimal-
perimeter}}}{\varepsilon_{\mathsf{apx}}}\right)^{2A}+2\sqrt{2\pi}\lambda_{1}\xi\right)\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\right]$ $\displaystyle\leq$
$\displaystyle\sqrt{d}\beta\lambda_{1}\cdot\left(\frac{C}{\varepsilon_{\mathsf{apx}}}\right)^{2A}\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\right],$

where the final inequality uses that $\xi\leq\beta$ (Definition 5.1). ∎

######  Lemma 5.6 (Absolute bound on bonuses).

For all $h,t,x$, it holds that
$|F_{h}^{(t)}(x)|\leq\frac{\beta}{2C_{\ref{lem:phi-what-wt-
diff}}HB\sqrt{d\iota}}$, and that
$\|w_{h}^{(t)}\|_{2}\leq\frac{\beta}{C_{\ref{lem:phi-what-wt-
diff}}\sqrt{d\iota}}$.

###### Proof.

Fix $t\in[T],h\in[H]$, and set
$(\Sigma^{\prime},\Lambda^{\prime}):=\mathsf{trunc}_{\sigma_{\mathsf{tr}}}({(\beta/\lambda_{1})\cdot(\Sigma_{h}^{(t)})^{-1/2}})$,
as per Eq. 33. For any fixed $u^{\prime},v^{\prime}\in\mathbb{R}^{d}$ with
$u^{\prime}\in\mathcal{S}_{\Sigma^{\prime}},v^{\prime}\in\mathcal{S}_{\Lambda^{\prime}}$,
we have

$\displaystyle F^{\mathsf{tl}}_{h}(x;\beta u^{\prime},v^{\prime})\leq$
$\displaystyle
2\|v^{\prime}\|_{2}\cdot\max_{\phi,\phi^{\prime}\in\bar{\Phi}_{h}(x)}\|\Lambda^{\prime}\cdot(\phi-\phi^{\prime})\|_{2}\leq
4\|v^{\prime}\|_{2},$

where the first inequality uses Lemma 4.8 and the fact that $\Lambda
v^{\prime}=v^{\prime}$, and the second inequality uses that
$\max_{h,x,a}\|\phi_{h}(x,a)\|_{2}\leq 1$. Then, by the definition of
$F_{h}^{(t)}$ in Eq. 33, we have

$\displaystyle|F_{h}^{(t)}(x)|\leq$
$\displaystyle{\lambda_{1}}\cdot\left(\frac{C_{\ref{cor:optimal-
perimeter}}}{\varepsilon_{\mathsf{apx}}}\right)^{2A}\cdot\mathbb{E}_{v^{\prime}\sim\mathcal{N}(0,\Lambda)}[4\|v^{\prime}\|_{2}]+(2\sqrt{2\pi}\lambda_{1}\xi)\cdot\mathbb{E}_{w\sim\mathcal{N}(0,\Sigma^{\prime})}\left[\max_{a\in\mathcal{A}}\langle
w,\phi_{h}(x,a)\rangle\right]$ $\displaystyle\leq$
$\displaystyle{4\lambda_{1}\sqrt{d}}\cdot\left(\frac{C_{\ref{cor:optimal-
perimeter}}}{\varepsilon_{\mathsf{apx}}}\right)^{2A}+2\sqrt{2\pi}\lambda_{1}\xi\cdot\sqrt{d}\leq\frac{\beta}{2C_{\ref{lem:phi-
what-wt-diff}}HB\sqrt{d\iota}},$

where the second inequality uses Lemma 5.3, and the final inequality is
derived as follows:

$\displaystyle{4\lambda_{1}\sqrt{d}}\cdot\left(\frac{C_{\ref{cor:optimal-
perimeter}}}{\varepsilon_{\mathsf{apx}}}\right)^{2A}+(2\sqrt{2\pi}\lambda_{1}\xi)\cdot\sqrt{d}\leq$
$\displaystyle{4\lambda_{1}\sqrt{d}}\cdot\left(\frac{C_{\ref{cor:optimal-
perimeter}}}{\varepsilon_{\mathsf{apx}}}\right)^{2A}+\frac{\beta}{4C_{\ref{lem:phi-
what-wt-diff}}HB\sqrt{d\iota}}\leq\frac{\beta}{2C_{\ref{lem:phi-what-wt-
diff}}HB\sqrt{d\iota}},$

where we have used the fact that
$2\sqrt{2\pi}\lambda_{1}\sqrt{d}\cdot\xi\leq\frac{\beta}{4C_{\ref{lem:phi-
what-wt-diff}}HB\sqrt{d\iota}}$ and that $\frac{\beta}{4C_{\ref{lem:phi-what-
wt-
diff}}HB\sqrt{d\iota}}\geq{4\lambda_{1}\sqrt{d}}\cdot\left(\frac{C_{\ref{cor:optimal-
perimeter}}}{\varepsilon_{\mathsf{apx}}}\right)^{2A}$ (by Definition 5.1).

The definition of $w_{h}^{(t)}$ (per Eq. 34 and Lemma 5.1), combined with the
fact that $|r_{h}(x,a)|\leq 1$ for all $h,x,a$ (Assumption 2.1), gives that
$|\langle w_{h}^{(t)},\phi_{h}(x,a)\rangle|\leq
H+H\cdot\max_{h,x}\\{|F_{h}^{(t)}(x)|\\}$, from which it follows that
$w_{h}^{(t)}\in(H+H\cdot\max_{h,x}\\{|F_{h}^{(t)}(x)|\\})\cdot\mathcal{B}_{h}$.
Since $\max_{w\in\mathcal{B}_{h}}\|w\|_{2}\leq B$ (again by Assumption 2.1),
we have as a consequence that
$\|w_{h}^{(t)}\|_{2}\leq\frac{\beta}{C_{\ref{lem:phi-what-wt-
diff}}\sqrt{d\iota}}$. ∎

###  5.2 Linear regression lemma

We will need the following lemma which establishes a standard concentration
statement regarding closeness of $\hat{w}_{h}^{(t)}$ to $w_{h}^{(t)}$.

######  Lemma 5.7.

There is a sufficiently large constant $C_{\ref{lem:phi-what-wt-diff}}$ (which
is used in the definitions of $\varepsilon_{\mathsf{apx}},\beta,\xi$ in
Definition 5.1) so that the following holds. There is an event $\mathcal{E}$
which occurs with probability at least $1-\delta/2$ so that, for all
$h\in[H],t\in[T],\phi\in\mathbb{R}^{d}$, it holds that

$\displaystyle|\langle\phi,\hat{w}_{h}^{(t)}-w_{h}^{(t)}\rangle|\leq\beta\cdot\left(\phi^{\top}(\Sigma_{h}^{(t)})^{-1}\phi\right)^{1/2}.$
(42)

###### Proof of Lemma 5.7.

Fix any $h\in[H],t\in[T]$. Note that, by definition of $\Sigma_{h}^{(t)}$ in 5
of Algorithm 1,

$\displaystyle
w_{h}^{(t)}=(\Sigma_{h}^{(t)})^{-1}\Sigma_{h}^{(t)}w_{h}^{(t)}=(\Sigma_{h}^{(t)})^{-1}\cdot\left(\lambda
w_{h}^{(t)}+\sum_{i=1}^{n}\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})\cdot\langle\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)}),w_{h}^{(t)}\rangle\right).$

Also recall that the definition of $\hat{w}_{h}^{(t)}$ in 7 gives

$\displaystyle\hat{w}_{h}^{(t)}=(\Sigma_{h}^{(t)})^{-1}\cdot\sum_{i=1}^{n}\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})\cdot\left(r_{h}^{(t,i,h)}+\sum_{g=h+1}^{H}\hat{r}_{g}^{(t,i,h)}\right).$

Thus,

$\displaystyle w_{h}^{(t)}-\hat{w}_{h}^{(t)}=$
$\displaystyle(\Sigma_{h}^{(t)})^{-1}\cdot\left(\lambda
w_{h}^{(t)}+\sum_{i=1}^{n}\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})\right.$
$\displaystyle\quad\cdot\left(\langle\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)}),w_{h}^{(t)}\rangle\left.-r_{h}^{(t,i,h)}-\sum_{g=h+1}^{H}\hat{r}_{g}^{(t,i,h)}\right)\right).$
(43)

For each $i\in[n]$, we have, by Eq. 34, that

$\displaystyle\mathbb{E}^{\hat{\pi}^{(t)}}\left[r_{h}^{(t,i,h)}+\sum_{g=h+1}^{H}\hat{r}_{g}^{(t,i,h)}\
\mid\
(x_{h},a_{h})=(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})\right]=Q_{h}^{(t)}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})=\langle\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)}),w_{h}^{(t)}\rangle.$

(In particular, note that by definition of the sampling procedure in 4, the
trajectory $(x_{h},a_{h},\ldots,x_{H},a_{H})$ is indeed drawn from the MDP
with policy $\hat{\pi}^{(t)}$.)

Let us write $Y:=\frac{\beta}{C_{\ref{lem:phi-what-wt-diff}}\sqrt{d\iota}}$.
The definition of $\beta$ in Definition 5.1 together with Lemma 5.6 ensures
that $Y\geq
2H\cdot\max\left\\{1,\max_{h,x}|F_{h}^{(t)}(x)|\right\\}\geq\left|r_{h}^{(t,i,h)}+\sum_{g=h+1}^{H}\hat{r}_{g}^{(t,i,h)}\right|$
with probability 1. We also have from Lemma 5.6 that
$Y\geq\|w_{h}^{(t)}\|_{2}$ with probability 1.

Thus, by Corollary 6.2 with $\Gamma_{0}=\lambda I$,
$\phi_{i}=\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})$ and
$\varepsilon_{i}=\langle\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)}),w_{h}^{(t)}\rangle-
r_{h}^{(t,i,h)}-\sum_{g=h+1}^{H}\hat{r}_{g}^{(t,i,h)}$ (so that we may take
$\sigma=2Y$), under some event $\mathcal{E}^{(t,h)}$ which holds with
probability at least $1-\delta/(2TH)$, we have

$\displaystyle\left\|\sum_{i=1}^{n}\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})\cdot\left(\langle\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)}),w_{h}^{(t)}\rangle-
r_{h}^{(t,i,h)}-\sum_{g=h+1}^{H}\hat{r}_{g}^{(t,i,h)}\right)\right\|_{(\Sigma_{h}^{(t)})^{-1}}$
$\displaystyle\leq$ $\displaystyle
8Y\cdot\sqrt{\log\left(\frac{TH\cdot\det(\Sigma_{h}^{(t)})^{1/2}\det(\lambda
I)^{-1/2}}{\delta/2}\right)}$ $\displaystyle\leq$ $\displaystyle
8Y\sqrt{d}\cdot\sqrt{\log\left(TH\cdot\frac{\lambda d+n}{\lambda
d\cdot\delta/2}\right)},$ (44)

where the final inequality follows since

$\displaystyle\det(\Sigma_{h}^{(t)})\leq\left(\frac{1}{d}\operatorname{Tr}\Sigma_{h}^{(t)}\right)^{d}\leq\left(\frac{1}{d}\cdot\left(\lambda
d+\sum_{i=1}^{n}\|\phi_{h}(x_{h}^{(t,i,h)},a_{h}^{(t,i,h)})\|_{2}^{2}\right)\right)^{d}\leq\left(\frac{\lambda
d+n}{d}\right)^{d}.$

Let us define
$\omega_{h}^{(t)}:=w_{h}^{(t)}-\hat{w}_{h}^{(t)}-(\Sigma_{h}^{(t)})^{-1}\cdot\lambda
w_{h}^{(t)}$, so that Eq. 43 and Eq. 44 give that, under
$\mathcal{E}^{(t,h)}$,

$\displaystyle\|(\Sigma_{h}^{(t)})^{1/2}\omega_{h}^{(t)}\|_{2}\leq
8Y\sqrt{d}\cdot\sqrt{\log\left(TH\cdot\frac{\lambda d+n}{\lambda
d\delta/2}\right)}.$ (45)

It follows that, under $\mathcal{E}^{(t,h)}$, for all $\phi\in\mathbb{R}^{d}$,

$\displaystyle|\langle\hat{w}_{h}^{(t)}-w_{h}^{(t)},\phi\rangle|=$
$\displaystyle|\lambda\langle(\Sigma_{h}^{(t)})^{-1}w_{h}^{(t)},\phi\rangle|+|\langle\omega_{h}^{(t)},\phi\rangle|$
$\displaystyle=$
$\displaystyle|\lambda\langle(\Sigma_{h}^{(t)})^{-1}w_{h}^{(t)},\phi\rangle|+|\langle(\Sigma_{h}^{(t)})^{1/2}\cdot\omega_{h}^{(t)},(\Sigma_{h}^{(t)})^{-1/2}\cdot\phi\rangle|$
$\displaystyle\leq$ $\displaystyle\left(\sqrt{\lambda}\cdot
Y+8Y\sqrt{d}\cdot\sqrt{\log\left(TH\cdot\frac{\lambda d+n}{\lambda
d\delta/2}\right)}\right)\cdot\sqrt{\phi^{\top}\cdot(\Sigma_{h}^{(t)})^{-1}\cdot\phi},$
(46)

where the final inequality uses Eq. 45, Cauchy-Schwarz, and

$\displaystyle|\lambda\langle(\Sigma_{h}^{(t)})^{-1}w_{h}^{(t)},\phi\rangle|=$
$\displaystyle|\langle w_{h}^{(t)},\lambda(\Sigma_{h}^{(t)})^{-1}\phi\rangle|$
$\displaystyle\leq$
$\displaystyle\|w_{h}^{(t)}\|_{2}\cdot\sqrt{\phi^{\top}\cdot\lambda^{2}(\Sigma_{h}^{(t)})^{-2}\cdot\phi}$
$\displaystyle\leq$ $\displaystyle\sqrt{\lambda}\cdot
Y\cdot\sqrt{\phi^{\top}\cdot(\Sigma_{h}^{(t)})^{-1}\cdot\phi},$

where the final inequality above uses the fact that
$\lambda\cdot(\Sigma_{h}^{(t)})^{-1}\preceq I_{d}$. Finally, we note that Eq.
46 gives that Eq. 42 holds under $\mathcal{E}^{(h,t)}$ since we have

$\displaystyle\beta=Y\cdot C_{\ref{lem:phi-what-wt-diff}}\sqrt{d\iota}\geq$
$\displaystyle Y\cdot\left(\sqrt{C_{\ref{lem:psd-
concentration}}d\log(2THn/\delta)}+8\sqrt{d\log\left(\frac{TH(\lambda
d+n)}{\lambda d\delta/2}\right)}\right)$ $\displaystyle=$ $\displaystyle
Y\cdot\left(\sqrt{\lambda}+8\sqrt{d}\sqrt{\log\left(\frac{TH(\lambda
d+n)}{\lambda d\delta/2}\right)}\right),$ (47)

where the above follows by the choices of $\lambda=C_{\ref{lem:psd-
concentration}}d\log(2THn/\delta)$, $\iota=\log\left(\frac{TH(\lambda
d+n)}{\delta/2}\right)\geq\log(2THn/\delta)$, and as long as $C_{\ref{lem:phi-
what-wt-diff}}$ is chosen sufficiently large. The conclusion of the lemma
follows by taking $\mathcal{E}=\bigcap_{h\in[H],t\in[T]}\mathcal{E}^{(h,t)}$
and noting that by a union bound, $\mathcal{E}$ occurs with probability at
least $1-\delta/2$. ∎

###  5.3 Establishing optimism

In this section, we establish that that $Q_{h}^{(t)}(\cdot)$ is
(approximately) an upper bound on $Q_{h}^{\star}(\cdot)$, in Lemma 5.9 below.

First, we prove Lemma 5.8 below, which shows that for states $x$ which satisfy
a certain inequality (namely, Eq. 48), $V_{h}^{\pi}(x)$ is approximately upper
bounded by $V_{h}^{(t)}(x)$ for any linear policy $\pi$. It is used to deal
with one particular case in the proof of Lemma 5.9. Recall that for any
$\pi\in\Pi^{\mathsf{lin}}$, Corollary 4.4 guarantees the existence of a vector
$w_{h}^{\pi}\in\mathbb{R}^{d}$ so that for all
$(x,a)\in\mathcal{X}\times\mathcal{A}$, $Q_{h}^{\pi}(x,a)=\langle
w_{h}^{\pi},\phi_{h}(x,a)\rangle$.

######  Lemma 5.8.

Consider an orthogonal pair of PSD matrices $(\Sigma,\Lambda)$, and fix
$x\in\mathcal{X}$. For some $\xi\geq 1$, suppose that

$\displaystyle\max_{a_{1},a_{2}\in\mathcal{A}}\|\Lambda\cdot(\phi_{h}(x,a_{1})-\phi_{h}(x,a_{2}))\|_{2}\leq\xi\cdot\max_{a_{1},a_{2}\in\mathcal{A}}\|\Sigma\cdot(\phi_{h}(x,a_{1})-\phi_{h}(x,a_{2}))\|_{2}.$
(48)

Then for any $\pi\in\Pi^{\mathsf{lin}}$,

$\displaystyle
V_{h}^{(t)}(x)+2\xi\|w_{h}^{\pi}\|_{2}\cdot\max_{a_{1},a_{2}\in\mathcal{A}}\|\Sigma\cdot(\phi_{h}(x,a_{1})-\phi_{h}(x,a_{2}))\|_{2}\geq
V_{h}^{\pi}(x)+\left(Q_{h}^{(t)}(x,\hat{\pi}_{h}^{(t)}(x))-Q_{h}^{\pi}(x,\hat{\pi}_{h}^{(t)}(x))\right).$

###### Proof.

Note that, for any actions $a,a^{\prime}\in\mathcal{A}$, we have

$\displaystyle\|\phi_{h}(x,a)-\phi_{h}(x,a^{\prime})\|_{2}\leq$
$\displaystyle\|\Lambda\cdot(\phi_{h}(x,a)-\phi_{h}(x,a^{\prime}))\|_{2}+\|\Sigma\cdot(\phi_{h}(x,a)-\phi_{h}(x,a^{\prime}))\|_{2}$
$\displaystyle\leq$ $\displaystyle
2\xi\cdot\max_{a,a^{\prime}\in\mathcal{A}}\|\Sigma\cdot(\phi_{h}(x,a)-\phi_{h}(x,a^{\prime}))\|_{2},$
(49)

where the second inequality uses Eq. 48 and the fact that $\xi\geq 1$.

We compute

$\displaystyle V_{h}^{(t)}(x)-V_{h}^{\pi}(x)$ $\displaystyle=$ $\displaystyle
Q_{h}^{(t)}(x,\hat{\pi}_{h}^{(t)}(x))-Q_{h}^{\pi}(x,\pi_{h}(x))$
$\displaystyle=$ $\displaystyle\langle
w_{h}^{\pi},\phi_{h}(x,\hat{\pi}_{h}^{(t)}(x))-\phi_{h}(x,\pi_{h}(x))\rangle+\left(Q_{h}^{(t)}(x,\hat{\pi}_{h}^{(t)}(x))-Q_{h}^{\pi}(x,\hat{\pi}_{h}^{(t)}(x))\right)$
$\displaystyle\geq$
$\displaystyle-\|w_{h}^{\pi}\|_{2}\cdot\|\phi_{h}(x,\hat{\pi}_{h}^{(t)}(x))-\phi_{h}(x,\pi_{h}(x))\|_{2}+\left(Q_{h}^{(t)}(x,\hat{\pi}_{h}^{(t)}(x))-Q_{h}^{\pi}(x,\hat{\pi}_{h}^{(t)}(x))\right)$
$\displaystyle\geq$