Bilinear Classes: A Structural Framework for Provable Generalization in RL

Introduction

Tackling large state-action spaces is a central challenge in reinforcement learning (RL). Here, function approximation and supervised learning schemes are often employed for generalization across large state-action spaces. While there have been a number of successful applications (Mnih et al., 2013; Kober et al., 2013; Silver et al., 2017; Wu et al., 2017). there is also a realization that practical RL approaches are quite sample inefficient.

Theoretically, there is a growing body of results showing how sample efficiency is possible in RL for particular model classes (often with restrictions on the model dynamics though in some cases on the class of value functions), e.g. State Aggregation (Li, 2009; Dong et al., 2020c), Linear MDPs (Yang and Wang, 2019; Jin et al., 2020), Linear Mixture MDPs (Modi et al., 2020a; Ayoub et al., 2020), Reactive POMDPs (Krishnamurthy et al., 2016), Block MDPs (Du et al., 2019a), FLAMBE (Agarwal et al., 2020b), Reactive PSRs (Littman et al., 2001), Linear Bellman Complete (Munos, 2005; Zanette et al., 2020).

More generally, there are also a few lines of work which propose more general frameworks, consisting of structural conditions which permit sample efficient RL; these include the low-rankness structure (e.g. the Bellman rank (Jiang et al., 2017) and Witness rank (Sun et al., 2019)) or under a complete condition (Munos, 2005; Zanette et al., 2020). The goal in these latter works is to develop a unified theory of generalization in RL, analogous to more classical notions of statistical complexity (e.g. VC-theory and Rademacher complexity) relevant for supervised learning. These latter frameworks are not contained in each other (see Table 1), and, furthermore, there are a number of natural RL models that cannot be incorporated into each of these frameworks (see Table 2).

Motivated by this latter line of work, we aim to understand if there are simple and natural structural conditions which capture the learnability in a general class of RL models.

This work provides a simple structural condition on the hypothesis class (which may be either model-based or value-based), where the Bellman error has a particular bilinear form, under which sample efficient learning is possible; we refer such a framework as a Bilinear Class. This structural assumption can be seen as generalizing the Bellman rank (Jiang et al., 2017); furthermore, it not only contains existing frameworks, it also covers a number of new settings that are not easily incorporated in previous frameworks (see Tables 1 and 2).

Our main result presents an optimization-based algorithm, BiLin-UCB, which provably enjoys a polynomial sample complexity guarantee for Bilinear Classes (cf. Theorem 5.2). Although our framework is more general than existing ones, our proof is substantially simpler – we give a unified analysis based on the elliptical potential lemma, developed for the theory of linear bandits (Dani et al., 2008; Srinivas et al., 2009).

Furthermore, as a point of emphasis, our results are non-parametric in nature (stated in terms of an information gain quantity (Srinivas et al., 2009)), as opposed to finite dimensional as in prior work. From a technical point of view, it is not evident how to extend prior approaches to this non-parametric setting. Notably, the non-parametric regime is particularly relevant to RL due to that, in RL, performance bounds do not degrade gracefully with approximation error or model mis-specification (e.g. see Du et al. (2020a) for discussion of these issues); the relevance of the non-parametric regime is that it may provide additional flexibility to avoid the catastrophic quality degradation due to approximation error or model mis-specification.

Definition of Bilinear Class: Our key conceptual contribution is the definition of the Bilinear Class, which isolates two key critical properties. The first property is that the Bellman error can be upper bounded by a bilinear form depending on the hypothesis. The second property is that the corresponding bilinear form for all hypothesis in the hypothesis class can be estimated with the same dataset. Analogous to supervised learning, this allows for efficient data reuse to estimate the Bellman error for all hypothesis simultaneously and eliminate those with high error.

A reduction to supervised learning: One appealing aspect of this framework is that the our main sample complexity result for RL is quantified via a reduction to the generalization error of a supervised learning problem, where we have a far better understanding of the latter. This is particularly important due to that we make no explicit assumptions on the hypothesis class $\mathcal{H}$ itself, thus allowing for neural hypothesis classes in some cases (the Bilinear Class posits an implicit relationship between $\mathcal{H}$ and the underlying MDP $\mathcal{M}$).

New models: We show our Bilinear Class framework incorporates new natural models, that are not easily incorporated into existing frameworks, e.g. linear $Q^{*}/V^{*}$, Low Occupancy Complexity, along with (infinite-dimensional) RKHS versions of linear MDPs and linear mixture MDPs. The linear $Q^{*}/V^{*}$ result is particularly notable due to a recent and remarkable lower bound which showed that if we only assume $Q^{*}$ is linear in some given set of features, then sample efficient learning is information theoretically not possible (Weisz et al., 2020). In perhaps a surprising contrast, our works shows that if we assume that both $Q^{\star}$ and $V^{\star}$ are linear in some given features then sample efficient learning is in fact possible.

Non-parametric rates: Our work is applicable to the non-parametric setting, where we develop new analysis tools to handle a number of technical challenges. This is notable as non-parametric rates for RL are few and far between. Our results are stated in terms of the critical information gain which can viewed as an analogous quantity to the critical radius, a quantity which is used to obtain sharp rates in non-parametric statistical settings (Wainwright, 2019).

Flexible Framework: The Bilinear Class framework is easily modified to include cases that do not strictly fit the definition. We show several examples of this in Section 6, where we show simple modifications of Bilinear Class framework include Witness Rank and Kernelized Nonlinear Regulator.

Section 2 provides further related work. Section 3 introduce some technical background and notation. Section 4 introduces our Bilinear Class framework, where we instantiate it on the several RL models, and Section 5 describes our algorithm and provides our main theoretical results. In Section 6, we introduce further extensions of Bilinear Classes. We conclude in Section 7. Appendix A provides additional examples of the Bilinear Class including the feature selection model Agarwal et al. (2020b), $\mathcal{Q}^{\ast}$ state aggregation, LQR, Linear MDP, and Block MDP. Appendix B provides missing proofs of Section 5. Appendix C provides a key technical theorem to attain non-parametric convergence rates in terms of the information gain, and Appendix D uses this to show concentration inequalities for all the models in a unified approach. Appendix E provides proofs for Section 6. Finally, Appendix G shows that low information gain is necessary in both Bellman Complete and Linear MDP by showing that small RKHS norm is not sufficient for sample-efficient reinforcement learning.

Related Work: Frameworks and Models

We first review existing frameworks and the relations among them. See Table 1 for a summary.

Jiang et al. (2017) defines a notion, Bellman Rank (B-Rank in Tables), in terms of the roll-in distribution and the function approximation class for $Q^{*}$, and give an algorithm with a polynomial sample complexity in terms of the Bellman Rank. They also showed a class of models, including tabular MDP, LQR, Reactive POMDP (Krishnamurthy et al., 2016), and Reactive PSR (Littman and Sutton, 2002) admit a low Bellman Rank, and thus they can be solved efficiently. Some recently proposed models, such as Block MDP (Du et al., 2019a), linear MDP (Yang and Wang, 2019; Jin et al., 2020) can also be shown to have a low Bellman rank. One caveat is that their algorithm requires a finite number of actions, so cannot be directly applied to (infinite-action) linear MDP and LQR. Subsequently, Sun et al. (2019) proposed a new framework, Witness Rank (W-Rank in tables), which generalizes Bellman Rank to model-based setting.

Bellman Complete (B-Complete in tables) is a framework of another style, which assumes that the class used for approximating the $Q$-function is closed under the Bellman operator. As shown in Table 1, neither the low-rank-style framework (Bellman Rank and Witness Rank) nor the complete-style framework (B-Complete) contains the other (See e.g., Zanette et al. (2020)).

Eluder dimension (Russo and Van Roy, 2014) is another structural condition which directly assumes the function class allows for strong extrapolation after observing dimension number of samples. With appropriate representation conditions (stronger than Bellman Complete), there is an efficient algorithm for function classes with small eluder dimension (Wang et al., 2020). However due to Eluder dimension requiring extrapolation, there are few examples of function classes with small eluder dimension beyond linear functions and monotone transformations of linear functions both of which are captured by the bilinear class.

Concurrently, Jin et al. (2021) proposes a new structural model called Bellman Eluder dimension (BE dimension) which takes both the MDP structure and the function class into consideration. We note that neither BE nor Bilinear Class capture each other. Notably, Bilinear Classes, via use of flexible Bellman error estimators, naturally captures model-based settings including linear mixture MDPs, KNRs, and factored MDPs, which are hard for model-free algorithms and frameworks to capture since the value functions of these models could be arbitrarily complicated. Specifically, Sun et al. (2019) shows that for factored MDPs, model-free algorithms such as OLIVE Jiang et al. (2017) suffer exponential sample complexity in worst case which implies that both BE dimension and Bellman rank are large for factored MDPs. However, Bilinear Class and Witness rank Sun et al. (2019) properly capture the complexity of factored MDPs. Similar situation may also apply to KNRs. For instance, Dong et al. (2020a) showed that for a simple piecewise linear dynamics (thus captured by KNRs) and piecewise reward functions, the optimal policy could contain exponentially many linear pieces and the optimal Q and V functions are fractals which are not differentiable anywhere and cannot be approximated by any neural networks with a polynomial width. It is unclear if such models have low BE dimension.

The primary difference is that the two complexity measures are applied to different structural aspects of the MDP: Bellman eluder framework is applied to the Bellman error and the bilinear class is applied to any loss estimator of the Bellman error. The actual complexity measures of eluder dimension and information gain are very similar and in fact equivalent for RKHS (Huang et al., ; Jin et al., 2021). As these two complexity measures are different in general, an interesting direction for further work is to understand how eluder dimension can address new settings of practical interest beyond (generalized) linear models and whether Bellman eluder dimension can be broadened to capture model-based approaches (like the linear mixture model). Finally, we comment that there are models (e.g., deterministic linear $Q^{\star}$ and $Q^{\star}$ state-action aggregation) that are captured by neither frameworks; we leave to future work to propose a framework that can capture these models that do not have error amplification.

With an additional Bellman completeness assumption on the function class, Jin et al. (2021) gives an algorithm which extends Eleanor from Zanette et al. (2020) to nonlinear function approximation that achieves a regret guarantee with faster rates than our algorithm. We note that our algorithm and OLIVE (as shown by Jin et al. (2021)) does not require Bellman completeness which is a much stronger assumption than realizability. As examples, the low occupancy complexity, feature selection model, linear mixture model, and many other model-based models are not Bellman complete. While our work focuses on PAC bounds, we conjecture that the techniques from Dong et al. (2020b) can be used for deriving regret bounds without completeness.

Now we discuss existing RL models. A summary on whether a model can be incorporated into a framework is provided in Table 2.

Tabular MDP is the most basic model, which has a finite number of states and actions, and all frameworks incorporate this model. When the state-action space is large, different RL models have been proposed to study when one can generalize across the state-action pairs.

Reactive POMDP (Krishnamurthy et al., 2016) assumes there is a small number of hidden states and the $Q^{*}$-function belongs to a pre-specified function class. Block MDP (Du et al., 2019a) also assumes there is a small number of hidden states and further assumes the hidden states are decodable. Reactive PSR (Littman et al., 2001) considers partial observable systems whose parameters are grounded in observable quantities. FLAMBE (Agarwal et al., 2020b) considers the feature selection and removes the assumption of known feature in linear MDP. These models all admit a low-rank structure, and thus can be incorporated into the Bellman Rank or Witness Rank and our Bilinear Classes.

The Linear Bellman Complete model (Munos, 2005) uses linear functions to approximate the $Q$-function, and assumes the linear function class is closed under the Bellman operator. Zanette et al. (2020) presented a statistically efficient algorithm for this model. This model does not have a low Bellman Rank or Witness Rank but can be incorporated into the Bellman Complete framework and ours.

Linear MDP (Yang and Wang, 2019; Jin et al., 2020) assumes the transition probability and the reward are linear in given features. This model not only admits a low-rank structure, but also satisfies the complete condition. Therefore, this model belongs in all frameworks. However, when the number of action is infinite, the algorithms for Bellman Rank and Witness Rank are not applicable because their sample complexity scales with the number of actions. Linear mixture MDP (Modi et al., 2020a; Ayoub et al., 2020) assumes the transition probability is a linear mixture of some base models. This model cannot be included in Bellman Rank, Witness Rank, or Bellman Complete, but our Bilinear Classes includes this model.

LQR is a fundamental model for continuous control that can be efficiently solvable (Dean et al., 2019). While LQR has a low Bellman Rank and low Witness Rank, since the algorithms for Bellman Rank and Witness Rank scale with the number of actions and LQR’s action set is uncountable, these two frameworks cannot incorporate LQR.

There is a line of work on state-action aggregation. $Q^{*}$ “irrelevance” state aggregation assumes one can aggregate states to a meta-state if these states share the same $Q^{*}$ value, and the number of meta-states is small (Li, 2009; Jiang et al., 2015). $Q^{*}$ state-action aggregation aggregates state-action pairs to a meta-state-action pair if these pairs have the same $Q^{*}$-value (Dong et al., 2020c; Li, 2009).

Lastly, when only assuming $Q^{*}$ is linear, there exists an exponential lower bound (Weisz et al., 2020), but with the additional assumption that the MDP is (nearly) deterministic and has large sub-optimality gap, there exists sample efficient algorithms (Wen and Van Roy, 2013; Du et al., 2019b, 2020b).

Setting

A deterministic, stationary policy $\pi:\mathcal{S}\mapsto\mathcal{A}$ specifies a decision-making strategy in which the agent chooses actions adaptively based on the current state, i.e. $a_{h}\sim\pi(s_{h})$. We denote a non-stationary policy $\pi=\{\pi_{0},\dots,\pi_{H-1}\}$ as a sequence of stationary policies where $\pi_{h}:\mathcal{S}\mapsto\mathcal{A}$.

Given a policy $\pi$ and a state-action pair $(s,a)\in\mathcal{S}\times\mathcal{A}$, the $Q$-function at time step $h$ is defined as

and, similarly, a value function time step $h$ of a given state $s$ under a policy $\pi$ is defined as

where both expectations are with respect to $s_{0},a_{0},\ldots s_{H-1},a_{H-1}\sim d^{\pi}$. We use $Q_{h}^{\star}$ and $V_{h}^{\star}$ to denote the $Q$ and $V$-functions of the optimal policy.

Throughout the paper, we will consider an algorithm as sample-efficient, if it uses number of trajectories polynomial in the problem horizon $H$, inherent dimension $d$, accuracy parameter $1/\epsilon$ and poly-logarithmic in the number of candidate value-functions.

For any two vectors $x,y$, we denote $[x,y]$ as the vector that concatenates $x,y$, i.e., $[x,y]:=[x^{\top},y^{\top}]^{\top}$. For any set $S$, we write $\triangle(S)$ to denote the probability simplex. We often use $U(S)$ as the uniform distribution over set $S$. We will let $\mathcal{V}$ denote a Hilbert space (which we assume is either finite dimensional or separable).

We let $[H]$ denote the set $\{0,\ldots H-1\}$. We slightly abuse notation (overloading $d^{\pi}$ with its marginal distributions), where $s_{h}\sim d^{\pi},(s_{h},a_{h})\sim d^{\pi},(r_{h},s_{h},a_{h},s_{h+1})\sim d^{\pi}$ and most frequently $o_{h}\sim d^{\pi}$ denotes the marginal distributions at timestep $h$. We also use the shorthand notation $s_{0},a_{0},\ldots s_{H-1},a_{H-1}\sim\pi$, $s_{h},a_{h}\sim\pi$ for $s_{0},a_{0},\ldots s_{H-1},a_{H-1}\sim d^{\pi}$, $s_{h},a_{h}\sim d^{\pi}$.

Bilinear Classes

Before, we define our structural framework – Bilinear Class, we first define our hypothesis class.

We assume access to a hypothesis class $\mathcal{H}=\mathcal{H}_{0}\times\ldots\times\mathcal{H}_{H-1}$, which can be abstract sets that permit for both model-based and value-based hypotheses. The only restriction we make is that for all $f\in\mathcal{H}$, we have an associated state-action value function $Q_{h,f}$ and a value function $V_{h,f}$. We next provide some examples:

An example of value-based hypothesis class $\mathcal{H}$ is an explicit set of state-action value $Q$ and value functions $V$ i.e.

Note that in this case, for any hypothesis $f:=((Q_{0},V_{0}),(Q_{1},V_{1}),\ldots,(Q_{H-1},V_{H-1}))\in\mathcal{H}$, we can take the associated $Q_{h,f}=Q_{h}$ and associated $V_{h,f}=V_{h}$.

Another example of value-based hypothesis class $\mathcal{H}$ is when $\mathcal{H}$ is just a set of state-action value $Q$ functions i.e.

In this case, for any hypothesis $f:=(Q_{0},Q_{1},\ldots,Q_{H-1})\in\mathcal{H}$, we can take the associated $Q_{h,f}=Q_{h}$ and the associated $V_{h,f}$ function to be greedy with respect to the $Q_{h,f}$ function i.e. $V_{h,f}(\cdot)=\max_{a\in\mathcal{A}}Q_{h,f}(\cdot,a)$.

An example of model-based hypothesis class is when $\mathcal{H}_{h}$ is a set of models/transition kernels $P_{h}$ and reward functions $R_{h}$ i.e.

In this case, for any hypothesis $f:=((P_{0},R_{0}),(P_{1},R_{1}),\ldots,(P_{H-1},R_{H-1}))\in\mathcal{H}$, we can take the associated $Q_{h,f}$ and $V_{h,f}$ functions to be the optimal value functions corresponding to the transition kernels $\{P_{h}\}_{h=0}^{H-1}$ and reward functions $\{R_{h}\}_{h=0}^{H-1}$.

Furthermore, we assume the hypothesis class is constrained so that $V_{h,f}(s)=\max_{a}Q_{h,f}(s,a)$ for all $f\in\mathcal{H}$, $h\in[H]$, and $s\in\mathcal{S}$, which is always possible as we can remove hypothesis for which this is not true. We let $\pi_{h,f}$ be the greedy policy with respect to $Q_{h,f}$, i.e., $\pi_{h,f}(s)=\mathop{{}\textrm{argmax}}_{a\in\mathcal{A}}Q_{h,f}(s,a)$, and $\pi_{f}$ as the sequence of time-dependent policies $\{\pi_{h,f}\}_{h=0}^{H-1}$.

1 Warmup: Bellman rank, the Q𝑄Q and V𝑉V versions.

As a motivation for our structural framework, we next discuss Bellman rank framework considered in Jiang et al. (2017). In this case, the hypothesis class $\mathcal{H}_{h}$ contains Q value functions, i.e.,

In this case, for any hypothesis $f:=(Q_{0},Q_{1},\ldots,Q_{H-1})\in\mathcal{H}$, we take the associated state-action value function $Q_{h,f}=Q_{h}$ and the associated state value $V_{h,f}$ function to be greedy with respect to the $Q_{h,f}$ function i.e. $V_{h,f}(\cdot)=\max_{a\in\mathcal{A}}Q_{h,f}(\cdot,a)$.

Even though Jiang et al. (2017) only considered $V$-Bellman Rank, as a natural extension of this definition, we can also consider the $Q$-Bellman Rank.

Let us interpret how the two definitions differ in the usage of functions $V_{h,f}$ vs $Q_{h,f}$ (along with the usage of the “estimation” policies $a_{0:h}\sim\pi_{f}$ vs $a_{0:h-1}\sim\pi_{f}$ and $a_{h}\sim\pi_{g}$). Recall that the Bellman equations can be written in terms of the value functions or the state-action values; here, the intuition is that the former definition corresponds to enforcing Bellman consistency of the value functions while the latter definition corresponds to enforcing Bellman consistency of the state-action value functions. Our more general structural framework, Bilinear Classes, will cover both these definitions for infinite dimensional hypothesis class (note that Jiang et al. (2017) only considered finite dimensional hypothesis class).

2 Bilinear Classes

We now introduce a new structural framework – the Bilinear Class.

We say that $\mathcal{H}$ is realizable for an MDP $\mathcal{M}$ if, for all $h\in[H]$, there exists a hypothesis $f^{\star}\in\mathcal{H}$ such that $Q_{h}^{\star}(s,a)=Q_{h,f^{\star}}(s,a)$, where $Q_{h}^{\star}$ is the optimal state-action value at time step $h$ in the ground truth MDP $\mathcal{M}$. For instance, for the model-based perspective, the realizability assumption is implied if the ground truth transition $P$ belongs to our hypothesis class $\mathcal{H}$.

Now we are ready to introduce the Bilinear Class.

Typically, $\pi_{\textrm{est}}(f)$ will be either the uniform distribution on $\mathcal{A}$ or $\pi_{f}$ itself; in the latter case, we refer to the estimation strategy as being on-policy.

We also define $\mathcal{X}_{h}:=\{X_{h}(f)\colon f\in\mathcal{H}\}$ and $\mathcal{X}:=\{\mathcal{X}_{h}:h\in[H]\}$.

We now provide some intuition for definition of Bilinear Class. The first part of the definition (Equation 1) basically relates the Bellman error for hypothesis $f$ (and hence sub-optimality) to the sum of bilinear forms $\left\lvert\langle W_{h}(f)-W_{h}(f^{\star}),X_{h}(f)\rangle\right\rvert$ (see for example proof of Lemma 5.5). Crucially, the second part of the definition (Equation 2), allows us to “reuse” data from hypothesis $f$ to estimate the bilinear form $\left\lvert\langle W_{h}(g)-W_{h}(f^{\star}),X_{h}(f)\rangle\right\rvert$ for all hypothesis $g$ in our hypothesis class! This is reminiscent of uniform convergence guarantees in supervised learning, where data can be reused to simultaneously estimate the loss for all hypothesis and eliminate those with high loss.

2.1 Finite Bellman rank ⟹\implies Bilinear Class

Its straightforward to see that in this case, both Equation 1 and Equation 2 are satisfied. ∎

Note that for $g=f$, we have that for observed transition info $o_{h}=(r_{h},s_{h},a_{h},s_{h+1})$

Therefore, to prove that this is a Bilinear Class, we will show that a stronger “equality” version of Equation 2 holds (which will also prove Equation 1 holds). Observe that for any $h$,

3 Examples

We now provide examples of Bilinear Classes: two known models (Linear Bellman Complete and Linear Mixture Models) and two new models that we propose (Linear $Q^{\star}/V^{\star}$ and Low Occupancy Complexity). We return to these examples to give non-parametric sample complexities in Section 5.3. See Appendix A for additional examples of Bilinear Classes.

First, we show our definition naturally captures model-based hypothesis class.

We say that a MDP $\mathcal{M}$ is a Linear Mixture Model if there exists (known) features $\phi:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\mapsto\mathcal{V}$ and $\psi:\mathcal{S}\times\mathcal{A}\mapsto\mathcal{V}$; and (unknown) $\theta^{\star}\in\mathcal{V}$ for some Hilbert space $\mathcal{V}$ such that for all $h\in[H]$ and $(s,a,s^{\prime})\in\mathcal{S}\times\mathcal{A}\times\mathcal{S}$

We denote hypothesis in our hypothesis class $\mathcal{H}$ as tuples $(\theta_{0},\ldots\theta_{H-1})$, where $\theta_{h}\in\mathcal{V}$. Recall that given a model $f\in\mathcal{H}$ (i.e. $f$ is the time-dependent transitions, i.e., $f_{h}:\mathcal{S}\times\mathcal{A}\mapsto\Delta(\mathcal{S})$), we denote $V_{h,f}$ as the optimal value function under model $f$ and corresponding reward function (in this case defined by $\psi$). Specifically, for any hypothesis $g=\{\theta_{0},\dots,\theta_{H-1}\}\in\mathcal{H}$, $V_{h,g}$ and $Q_{h,g}$ satisfy the following Bellman optimality equation:

Note that in this example, discrepancy function will explicitly depend on $f$. For hypothesis $g=\{\theta_{0},\dots,\theta_{H-1}\}\in\mathcal{H}$ and observed transition info $o_{h}=(r_{h},s_{h},a_{h},s_{h+1})$, we define

Observe that for $g=f$, using Equation 3, for observed transition info $o_{h}=(r_{h},s_{h},a_{h},s_{h+1})$,

We consider on-policy estimation $\pi_{est}=\pi_{f}$. To prove that linear mixture MDP is a Bilinear Class, we only need to show that an “equality” version of Equation 2 holds (which implies Equation 1 holds by the frame above). For $g=\{\theta_{0},\dots,\theta_{H-1}\}\in\mathcal{H}$, observe:

where we defined the $W_{h},X_{h}$ functions as follows:

This concludes that Linear Mixture Model also forms a Bilinear Class. ∎

We introduce a new model: linear $Q^{\star}/V^{\star}$ where we assume both the optimal $Q^{\star}$ and $V^{\star}$ are linear functions in features that lie in (possibly infinite dimensional) Hilbert space.

We say that a MDP $\mathcal{M}$ is a linear $Q^{\star}/V^{\star}$ model if there exist (known) features $\phi:\mathcal{S}\times\mathcal{A}\mapsto\mathcal{V}_{1}$, $\psi:\mathcal{S}\mapsto\mathcal{V}_{2}$ and (unknown) $(w^{\star},\theta^{\star})\in\mathcal{V}_{1}\times\mathcal{V}_{2}$ for some Hilbert spaces $\mathcal{V}_{1},\mathcal{V}_{2}$ such that for all $h\in[H]$ and for all $(s,a,s^{\prime})\in\mathcal{S}\times\mathcal{A}\times\mathcal{S}$,

Here, our hypothesis class $\mathcal{H}=\mathcal{H}_{0}\times\ldots,\mathcal{H}_{H-1}$ is a set of linear functions i.e. for all $h\in[H]$, the set $\mathcal{H}_{h}$ is defined as:

Note that we will show that a stronger “equality” version of Equation 2 holds, which will also prove Equation 1 holds since for observed transition info $o_{h}=(r_{h},s_{h},a_{h},s_{h+1})$,

3.3 Bellman Complete and Linear MDPs

We now consider Bellman Complete which captures the linear MDP model (see Section A.4 for more detail on linear MDP model). Here, our hypothesis class $\mathcal{H}$ is set of linear functions with respect to some (known) feature $\phi:\mathcal{S}\times\mathcal{A}\mapsto\mathcal{V}$, where $\mathcal{V}$ is a Hilbert space. We denote hypothesis in our hypothesis class $\mathcal{H}$ as tuples $(\theta_{0},\ldots\theta_{H-1})$, where $\theta_{h}\in\mathcal{V}$.

We say our hypothesis class $\mathcal{H}$ is Linear Bellman Complete with respect to $\mathcal{M}$ if $\mathcal{H}$ is realizable and there exists $\mathcal{T}_{h}:\mathcal{V}\rightarrow\mathcal{V}$ such that for all $(\theta_{0},\ldots\theta_{H-1})\in\mathcal{H}$ and $h\in[H]$,

for all $(s,a)\in\mathcal{S}\times\mathcal{A}$.

Note that in this case, we will show that a stronger version of Equation 2 holds i.e with equality instead of $\leq$ inequality, which will also prove Equation 1 holds since for observed transition info $o_{h}=(r_{h},s_{h},a_{h},s_{h+1})$,

Observe that $W_{h}(f^{\star})=0$ for all $h$. ∎

3.4 Low Occupancy Complexity (new model).

We introduce another new model: Low Occupancy Complexity.

We say that a MDP $\mathcal{M}$ and hypothesis class $\mathcal{H}$ has low occupancy complexity with respect to a (possibly unknown) feature mapping $\phi_{h}:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{V}$ (where $\mathcal{V}$ is a Hilbert space) if $\mathcal{H}$ is realizable and there exists a (possibly unknown) $\beta_{h}:\mathcal{H}\mapsto\mathcal{V}$ for $h\in[H]$ such that for all $f\in\mathcal{H}$ and $(s_{h},a_{h})\in\mathcal{S}\times\mathcal{A}$ we have that:

To see why this is a Bilinear Class, as in previous proofs, we will show that an “equality” version of Equation 2 holds, which will also prove Equation 1 holds since

Observe that for any $h$ (here observed transition info $o_{h}=(r_{h},s_{h},a_{h},s_{h+1})$):

Note that $W_{h}(f^{\star})=0$. This completes the proof. ∎

Note that as such the hypothesis class $\mathcal{H}$ could be arbitrary and unlike other models where we assume linearity, here it could be a neural state-action value class. Our model can also capture the setting where the state-only occupancy has low complexity, i.e., $d^{\pi_{f}}(s_{h})=\beta_{h}(f)\mu_{h}(s_{h})$, for some $\mu_{h}:\mathcal{S}\to\mathcal{V}$. In this case, we will use $\pi_{est}=U(\mathcal{A})$.

The Algorithm and Theory

There are two ways to collect batch samples. For the case where $\pi_{est}=\pi_{f_{t}}$, then for data collection in Line 4, we can generate $m$ length-H trajectories by executing $\pi_{f_{t}}$ starting from $s_{0}$. For the general case (e.g. consider setting $\pi_{est}$ to be a uniform distribution over $\mathcal{A}$), we gather the data for each $h\in[H]$ independently. For $h\in[H]$, we first roll-in with $\pi_{f_{t}}$ to generate $s_{h}$; then execute $a_{h}\sim\pi_{est}$; and then continue to generate $s_{h+1}\sim P_{h}(\cdot|s_{h},a_{h})$ and $r_{h}\sim R(\cdot|s_{h},a_{h})$. Repeating this process for all $h$, we need $Hm$ trajectories to form the batch datasets $\{\mathcal{D}_{t;h}\}_{h=0}^{H-1}$.

For a set $\mathcal{D}\subset\mathcal{S}\times\mathcal{A}\times\mathcal{S}$, we will also use $\mathcal{D}$ to represent the uniform distribution over this set.

It is helpful to separate the dependence of generalization error on failure probability $\delta$ and number of samples $m$ in order to state Theorem 5.2 concisely. $\varepsilon_{\textrm{gen}}(m,\mathcal{H})$ is related to uniform convergence and measures the generalization error of hypothesis class $\mathcal{H}$ and for the hypothesis classes discussed in this paper, $\varepsilon_{\textrm{gen}}(m,\mathcal{H})\to 0$ as $m\to\infty$. One example is when $\pi_{est}=\pi_{f}$, and $\mathcal{H}$ is a discrete function class, then we have $\varepsilon_{\textrm{gen}}(m,\mathcal{H})=O\left(\sqrt{(1+\ln(|\mathcal{H}|))/m}.\right)$. In Appendix D, we also discuss uniform convergence via a novel covering argument for infinite dimensional RKHS.

Set the parameters as: number of iterations $T=\widetilde{d}_{m}$ and confidence radius $R=\sqrt{T}\varepsilon_{\textrm{gen}}(m,\mathcal{H})\cdot\textrm{conf}(\delta/(TH))$. With probability at least $1-\delta$, Algorithm 1 uses at most $mHT$ trajectories and returns a hypothesis $f$ such that:

number of iterations T=c_{2}dH\ln\Big{(}B_{X}B_{W}m\Big{)} and confidence radius $R=c_{3}\sqrt{T}\cdot H\sqrt{\ln(|\mathcal{H}|)/m}\cdot\ln(TH/\delta)$, with probability at least $1-\delta$, Algorithm 1 returns a hypothesis $f$ such that $V^{\star}(s_{0})-V^{\pi_{f}}(s_{0})\leq\epsilon$ using at most

The proof for this corollary follows from bounds on $\varepsilon_{\textrm{gen}}(m,\mathcal{H})$ and $\textrm{conf}(\delta)$ using Hoeffding’s inequality (Lemma F.1). We present the complete proof in Appendix B.

Our next results will be non-parametric in nature and therefore it is helpful to introduce the maximum information gain (Srinivas et al., 2009), which captures an important notion of the effective dimension of a set. Let $\mathcal{X}\subset\mathcal{V}$ , where $\mathcal{V}$ is a Hilbert space. For $\lambda>0$ and integer $n>0$, the maximum information gain $\gamma_{n}(\lambda;\mathcal{X})$ is defined as:

If $\mathcal{X}$ is of the form $\mathcal{X}=\{\mathcal{X}_{h}:h\in[H]\}$, we use the notation

Define critical information gain, denoted by $\widetilde{\gamma}(\lambda;\mathcal{X})$, as the smallest integer $k>0$ s.t. $k\geq\gamma_{k}(\lambda;\mathcal{X})$, i.e.

(where $k$ is an integer). Note that such a $\widetilde{\gamma}(\lambda;\mathcal{X})$ exists provided that the information gain $\gamma_{n}(\lambda;\mathcal{X})$ has a sufficiently mild growth condition in both $n$ and $1/\lambda$. The critical information gain can viewed as an analogous quantity to the critical radius, a quantity which arises in non-parametric statistics (Wainwright, 2019).

We now present our main theorem. Recall the definitions $\mathcal{X}_{h}:=\{X_{h}(f)\colon f\in\mathcal{H}\}$ and $\mathcal{X}:=\{\mathcal{X}_{h}:h\in[H]\}$.

Set the parameters as: number of iterations $T=\widetilde{d}_{m}$ and confidence radius $R=\sqrt{\widetilde{d}_{m}}\varepsilon_{\textrm{gen}}(m,\mathcal{H})\cdot\textrm{conf}(\delta/(\widetilde{d}_{m}H))$. With probability at least $1-\delta$, Algorithm 1 uses at most $mH\widetilde{d}_{m}$ trajectories and returns a hypothesis $f$ such that:

Next, we provide an elementary and detailed proof for our main theorem using an elliptical potential argument.

2 Proof of Theorem 5.1 and Theorem 5.2

In this subsection, we prove our main theorems – Theorem 5.1 and Theorem 5.2.

For all $t\in[T]$ and $g\in\mathcal{H}$ and $h\in[H]$, with probability at least $1-\delta$, we have:

This follows from the uniform convergence (5.1) and then union bounding over all $t\in[T]$ and $h\in[H]$. ∎

We now complete the proof of Theorem 5.1 and Theorem 5.2 using Lemma 5.1, Lemma 5.2 and setting the parameters using the definition of critical information gain.

Fix $\lambda=\varepsilon_{\textrm{gen}}^{2}(m,\mathcal{H})/B_{W}^{2}$. From definition of critical information gain (Equation 6), it follows that for $T=\widetilde{\gamma}(\lambda,\mathcal{X})$,

Observing that for our choice of $T$, $\gamma_{T}(\lambda;\mathcal{X})/T\leq 1$ and $e-1<2$ , we get

where the second last equality uses the definition of $\lambda$.

Moreover, each iteration of the algorithm, takes only $mH$ trajectories, this gives the total trajectories as $mHT=mH\widetilde{\gamma}(\lambda,\mathcal{X})$. This proves Theorem 5.2. Theorem 5.1 follows from the upper bound on $\widetilde{\gamma}(\lambda,\mathcal{X})$ for finite dimensional $\mathcal{X}_{h}$ using Lemma F.3. ∎

In the rest of the section, we will prove our main lemma – Lemma 5.2. The first step shows that under 5.1, our $R$ is set properly so that $f^{\star}$ is always a feasible solution of the constrained optimization program in Algorithm 1.

Assume the event in Lemma 5.1 holds. Then for all $t\in[T]$, we have that $f^{\star}$ is always a feasible solution.

Note that $\mathcal{L}_{\mu_{i;h},f_{i}}(f^{\ast})=0$ (Equation 2). Thus using Lemma 5.1, we have:

Noting that $t\leq T$ and in our parameter setup $R=\sqrt{T}\varepsilon_{\textrm{gen}}$ completes the proof. ∎

The feasibility result immediately leads to optimism.

Assume the event in Lemma 5.1 holds. Then for all $t\in[T]$, we have $V^{\star}\leq V_{f_{t};0}(s_{0})$.

Lemma 5.3 implies $f^{\star}$ is a feasible solution for the optimization program for all $t\in[T]$. This proves the claim. ∎

The following lemma relates the sub-optimality to a sum of bilinear forms. Using the performance difference lemma, we first show that sub-optimality is upper bounded by the Bellman errors of $Q_{h,f_{t}}$, which are further upper bounded by sum of bilinear forms via our assumption (Equation 1).

Assume the event in Lemma 5.1 holds. Then, the following holds for all $t\in[T]$:

where the last step follows Equation 1 in the Bilinear Class definition. ∎

The following is a variant of the Elliptical Potential Lemma, central in the analysis of linear bandits (Dani et al., 2008; Srinivas et al., 2009; Abbasi-Yadkori et al., 2011).

By definition of $\Sigma_{t}$ and matrix determinant lemma, we have:

Now, we will finish the proof of Lemma 5.2 by showing that the sum of bilinear forms in Lemma 5.5 is small for at least for one $t\in[T]$. More precisely, using Equation 2 together with elliptical potential argument (Lemma 5.6), we can show that after $\widetilde{d}_{m}$ many iterations, we must have found a policy $\pi_{f_{t}}$ such that $\left\lvert\langle W_{h}(f_{t})-W_{h}(f^{\star}),X_{h}(f_{t})\rangle\right\rvert$ is small for all $h$.

Our goal (as per Lemma 5.5 and Equation 1) is to find $t\in[T]$ such that

for appropriately chosen $A$. We will show existence of such $X_{h}(f_{t})$ and $A$ (Equation 7) using the potential argument (Lemma 5.6) and conditions on $W_{h}(f_{t})-W_{h}(f^{\star})$ follow from our optimization program. We now show this in more detail.

where we have used definition of maximum information gain $\gamma_{T}(\lambda;\mathcal{X}_{h})$ (Equation 4) and

where the last equality follows from Equation 5. Since, each of these terms is $\geq 0$, we get that there exists $t\in[T]$ such that

Again, since each of these terms is $\geq 0$, we get that for all $h\in[H]$

and simplifying, we get that for all $h\in[H]$,

Also, by construction of our program, for all iterations and in particular for $t$, it holds that for all $h\in[H]$

where the first inequality follows from $(a+b)^{2}\leq 2a^{2}+2b^{2}$ and the last step follows from the frame above and $t\in[T]$. Using the definition of Bilinear Class (Equation 2), for all $h\in[H]$

where the first inequality follows from the frame above and definition of $\Sigma_{t;h}$. Using Equation 7 and the frame above, this immediately shows that for all $h\in[H]$

Using Lemma 5.5, this gives the desired result. ∎

3 Corollaries for Particular Models

In this section, we apply our main theorem to special models: linear $Q^{\star}/V^{\star}$, RKHS bellman complete, RKHS linear mixture model, and low occupancy complexity model. While linear bellman complete and linear mixture model have been studied, our results extends to infinite dimensional RKHS setting.

In this subsection, we provide the sample complexity result for the linear $Q^{\star}/V^{\star}$ model (Definition 4.5). To state our results for linear $Q^{\star}/V^{\star}$, we define the following sets:

and define the concatenation setFor infinite dimensional $\Phi$ and $\Psi$, we consider the natural inner product space where $\langle[x_{1},y_{1}],[x_{2},y_{2}]\rangle=\langle x_{1},x_{2}\rangle+\langle y_{1},y_{2}\rangle$.

trajectories for some absolute constant $c_{1},c_{2}$.

To prove this, we will prove a more general sample complexity result for the infinite dimensional RKHS case.

Suppose MDP $\mathcal{M}$ is a linear $Q^{\star}/V^{\star}$ model. Assume $\sup_{(w,\theta)\in\mathcal{H}_{h},h\in[H]}\lVert[w,\theta]\rVert_{2}\leq B_{W}$ and $\sup_{x\in\Phi\circ\Psi}\lVert x\rVert_{2}\leq B_{X}$. Fix $\delta\in(0,1/3)$, batch sample size $m$, and define:

where \nu:=\ln\left(1+3B_{X}B_{W}\sqrt{m\widetilde{\gamma}\Big{(}\frac{1}{8B_{W}^{2}m};\Phi\circ\Psi\Big{)}}\right).

Set the parameters as: R=(12H/\sqrt{m})\sqrt{\widetilde{d}_{m}(\mathcal{X})\cdot\widetilde{d}_{m}(\Phi\circ\Psi)}\cdot\sqrt{\ln\big{(}(\widetilde{d}_{m}(\mathcal{X})H)/\delta\big{)}} and $T=\widetilde{d}_{m}(\mathcal{X})$. With probability greater than $1-\delta$, Algorithm 1 uses at most $mH\widetilde{d}_{m}(\mathcal{X})$ trajectories and returns a hypothesis $f$:

where $v:=\sqrt{\ln\left((\widetilde{d}_{m}(\mathcal{X})H)/\delta\right)}$.

First, using Corollary D.3, we get that for any distribution $\mu$ over $\mathcal{S}\times\mathcal{A}\times\mathcal{S}$ and for any $\delta\in(0,1)$, with probability of at least $1-\delta$ over choice of an i.i.d. sample $\mathcal{D}\sim\mu^{m}$ of size $m$, for all $g=([w_{0},\theta_{0}],\ldots,[w_{H-1},\theta_{H-1}])\in\mathcal{H}$ (note that $\mathcal{L}_{\mu}(g)$ only depends on $[w_{h},\theta_{h}]$ for distribution $\mu$ over observed transitions $o_{h}=(r_{h},s_{h},a_{h},s_{h+1})$ at timestep $h$.)

where we have used that $\ln(1/\delta)>1$ and $\widetilde{\gamma}_{m}=\widetilde{\gamma}(1/(8B_{W}^{2}m);\Phi\circ\Psi)$ (as defined in Equation 6). Define

Substituting this in Theorem 5.2 gives the result

Similarly, as $\sup_{z\in\mathcal{X}}\lVert z\rVert\leq\sup_{x\in\Phi\circ\Psi}\lVert x\rVert$, using Lemma F.3 and similar analysis as above (and $144H^{2}\widetilde{d}_{m}(\Phi\circ\Psi)\geq 1$), we get

To get $\epsilon$-optimal policy (from Equation 11), we have to set

Further upper bounding the right hand side of the above inequality by substituting in upper bounds for $\widetilde{d}_{m}(\mathcal{X})$ and $\widetilde{d}_{m}(\Phi\circ\Psi)$ from frames above, we can set $m$ to be as large as:

Using Lemma F.2 for $\alpha=4$, $a=32\cdot(72)^{2}d^{2}H^{5}\ln(1/\delta)/\epsilon^{2}$, $b=25B_{X}^{2}B^{2}_{W}dH^{2}$ and $c=5^{4}$, we get that

Substituting this in the expression above for $\widetilde{d}_{m}(\mathcal{X})$ and setting this upper bound to $T$, we get

Since, we use on policy estimation, i.e., $\pi_{est}=\pi_{f_{t}}$ for all $t$, the trajectory complexity is $mT$ which completes the proof. ∎

3.2 RKHS Bellman Complete.

In this subsection, we provide the sample complexity result for the Linear Bellman Complete model (Definition 4.6). To state our results, we define

trajectories for some absolute constant $c_{1},c_{2}$.

In comparison, Jin et al. (2020) has sample complexity $\widetilde{O}(d^{3}H^{3}/\epsilon^{2}\log(1/\delta))$ and Zanette et al. (2020) has $\widetilde{O}(d^{2}H^{3}/\epsilon^{2}\log(1/\delta))$. To prove this, we will prove a more general sample complexity result for the infinite dimensional RKHS case. Note that RKHS Linear MDP is a special instance of RKHS Bellman Complete. Prior works that studied RKHS Linear MDP either achieves worse rate (Agarwal et al., 2020a) or further assumes finite covering dimension of the space of all possible upper confidence bound Q functions which are algorithm dependent quantities (Yang et al., 2020).

Suppose $\mathcal{H}$ is Bellman Complete with respect to MDP $\mathcal{M}$ for some Hilbert space $\mathcal{V}$. Assume $\sup_{h\in[H],\theta\in\mathcal{H}_{h}}\lVert\theta\rVert_{2}\leq B_{W}$ and $\sup_{x\in\Phi}\lVert x\rVert_{2}\leq B_{X}$. Fix $\delta\in(0,1/3)$, batch sample size $m$, and define:

where \nu=\ln\left(1+3B_{X}B_{W}\sqrt{m\widetilde{\gamma}\Big{(}\frac{1}{8B_{W}^{2}m};\Phi\Big{)}}\right).

where v=\sqrt{\ln\big{(}(\widetilde{d}_{m}(\mathcal{X})H)/\delta\big{)}}.

First, using Corollary D.2, we get that for any distribution $\mu$ over $\mathcal{S}\times\mathcal{A}\times\mathcal{S}$ and for any $\delta\in(0,1)$, with probability of at least $1-\delta$ over choice of an i.i.d. sample $\mathcal{D}\sim\mu^{m}$ of size $m$, for all $g=(\theta_{0},\ldots,\theta_{H-1})\in\mathcal{H}$ (note that $\mathcal{L}_{\mu}(g)$ only depends on $\theta_{h}$ for distribution $\mu$ over observed transitions $o_{h}=(r_{h},s_{h},a_{h},s_{h+1})$ at timestep $h$.)

where we have used that $\ln(1/\delta)>1$ and $\widetilde{\gamma}_{m}=\widetilde{\gamma}(1/(8B_{W}^{2}m);\Phi)$ (as defined in Equation 6). Define

Substituting this in Theorem 5.2 gives the result

Since the proof follows similar to proof of Corollary 5.2, we will only provide a proof sketch here. First, from Lemma F.3, we have that

Similarly, as $\sup_{z\in\mathcal{X}}\lVert z\rVert\leq\sup_{x\in\Phi}\lVert x\rVert$, using Lemma F.3 (and since $400H^{2}\widetilde{d}_{m}(\Phi)\geq 1$), we get

To get $\epsilon$-optimal policy, we have to set

The rest of the proof follows similarly to proof of Corollary 5.2. ∎

3.3 RKHS linear mixture model

In this subsection, we provide the sample complexity result for the Linear Mixture model (Definition 4.4). To present our sample complexity results, we define:

trajectories for some absolute constant $c_{1},c_{2}$.

In comparison, Modi et al. (2020a) has sample complexity $\widetilde{O}(d^{2}H^{2}/\epsilon^{2}\log(1/\delta))$. To prove this, we will prove a more general sample complexity result for the infinite dimensional RKHS case. We omit proof of Corollary 5.6 since it follows same as proof of Corollary 5.2.

Suppose MDP $\mathcal{M}$ is a linear Mixture Model. Assume $\sup_{\theta\in\mathcal{H}_{h},h\in[H]}\lVert\theta\rVert_{2}\leq B_{W}$ and $\sup_{x\in\Phi_{h},h\in[H]}\lVert x\rVert_{2}\leq B_{X}$. Fix $\delta\in(0,1/3)$, batch sample size $m$, and define:

where \nu_{h}=\ln\left(1+3B_{X}B_{W}\sqrt{m\widetilde{\gamma}\Big{(}\frac{1}{8B_{W}^{2}m};\Phi_{h}\Big{)}}\right).

Set parameters as: R=(12H/\sqrt{m})\sqrt{\widetilde{d}_{m}(\mathcal{X})\cdot\widetilde{d}_{m}(\Phi)}\cdot\sqrt{\ln\big{(}(\widetilde{d}_{m}(\mathcal{X})H)/\delta\big{)}} and $T=\widetilde{d}_{m}(\mathcal{X})$. With probability greater than $1-\delta$, Algorithm 1 uses at most $mH\widetilde{d}_{m}(\mathcal{X})$ trajectories and returns a hypothesis $f$

where v=\sqrt{\ln\big{(}(\widetilde{d}_{m}(\mathcal{X})H)/\delta\big{)}}.

First, using Corollary D.3 and Lemma F.1, we get that for any distribution $\mu$ over $\mathcal{S}\times\mathcal{A}\times\mathcal{S}$ and for any $\delta\in(0,1)$, with probability of at least $1-\delta$ over choice of an i.i.d. sample $\mathcal{D}\sim\mu^{m}$ of size $m$, for all $g=(\theta_{0},\ldots,\theta_{H-1})\in\mathcal{H}$ (note that $\mathcal{L}_{\mu}(g)$ only depends on $\theta_{h}$ for distribution $\mu$ over observed transitions $o_{h}=(r_{h},s_{h},a_{h},s_{h+1})$ at timestep $h$.)

where we have used that $\ln(1/\delta)>1$ and $\widetilde{\gamma}_{m}=\max_{h\in[H]}\widetilde{\gamma}(1/(8B_{W}^{2}m);\Phi_{h})$ (as defined in Equation 6). Define

Substituting this in Theorem 5.2 gives the result

3.4 Low Occupancy Complexity

Recall the low occupancy complexity model in Definition 4.7.

Suppose $\mathcal{H}$ has low occupancy complexity. Assume $\sup_{f\in\mathcal{H}_{h},h\in[H]}\lVert W_{h}(f)\rVert_{2}\leq B_{W}$. Fix $\delta\in(0,1/3)$, batch sample size $m$, and define:

Set $T=\widetilde{d}_{m}(\mathcal{X})$ and R=(2\sqrt{2}H/\sqrt{m})\cdot\sqrt{\widetilde{d}_{m}(\mathcal{X})}\cdot\sqrt{1+\ln\big{(}|\mathcal{H}|\big{)}}\cdot\sqrt{\ln\big{(}\widetilde{d}_{m}(\mathcal{X})H\big{)}+\ln\big{(}1/\delta\big{)}}. With probability greater than $1-\delta$, Algorithm 1 uses at most $mH\widetilde{d}_{m}(\mathcal{X})$ trajectories and returns a hypothesis $f$ such that:

where v=\sqrt{\ln\big{(}\widetilde{d}_{m}(\mathcal{X})H\big{)}+\ln\big{(}1/\delta\big{)}}.

First, using Lemma F.1, we get that for any distribution $\mu$ over $\mathcal{S}\times\mathcal{A}\times\mathcal{S}$ and for any $\delta\in(0,1)$, with probability of at least $1-\delta$ over choice of an i.i.d. sample $\mathcal{D}\sim\mu^{m}$ of size $m$, for all $g\in\mathcal{H}$

Substituting this in Theorem 5.2 gives the result

3.5 Finite Bellman Rank

In this section, we will prove sample complexity bounds for MDPs with finite Bellman Rank introduced in Jiang et al. (2016) (also defined as $V$-Bellman rank in Section 4.1).

For a given MDP $\mathcal{M}$, suppose a hypothesis class $\mathcal{H}$ has Bellman rank $d$. Assume $\sup_{f\in\mathcal{H}_{h},h\in[H]}\lVert W_{h}(f)\rVert_{2}\leq B_{W}$ and $\sup_{f\in\mathcal{H},h\in[H]}\lVert X_{h}(f)\rVert\leq B_{X}$ for some $B_{W},B_{X}\geq 1$. Fix $\delta\in(0,1/3)$ and $\epsilon\in(0,H)$. There exists an appropriate setting of batch sample size $m$, number of iteration $T$ and confidence radius $R$ such that with probability at least $1-\delta$, Algorithm 1 returns a hypothesis $f$ such that $V^{\star}(s_{0})-V^{\pi_{f}}(s_{0})\leq\epsilon$ using at most

trajectories for some absolute constant $c_{1},c_{2}$.

Note that in comparison, Jiang et al. (2016) has sample complexity $\widetilde{O}(d^{2}H^{5}|\mathcal{A}|/\epsilon^{2}\log(1/\delta))$. We now present the proof.

First, as observed in Jiang et al. (2016)[Lemma 14], we get that for any distribution $\mu$ over $\mathcal{S}\times\mathcal{A}\times\mathcal{S}$ and for any $\delta\in(0,1)$, with probability of at least $1-\delta$ over choice of an i.i.d. sample $\mathcal{D}\sim\mu^{m}$ of size $m$, for all $g\in\mathcal{H}$

where the second inequality holds as long as $m>2H|\mathcal{A}|\ln(|\mathcal{H}|/\delta)$. This satisfies our 5.1 with

Substituting this in Theorem 5.2 gives the result

where the second last step follows from Lemma F.3. Substituting $\varepsilon_{\textrm{gen}}$ and conf in Theorem 5.2 also gives

To get $\epsilon$-optimal policy, we have to set

Further simplifying the RHS, we can write it as

Using Lemma F.2 for $\alpha=2$, $a=4608dH^{5}|\mathcal{A}|(1+\ln(|\mathcal{H}|))/\epsilon^{2}$, $b=16dH^{2}B_{W}^{2}B_{X}^{2}/\delta$ and $c=9$, we get that

Substituting this in the expression above for $\widetilde{d}_{m}(\mathcal{X})$ and setting this upper bound to $T$, we get

Since, we use on policy estimation, i.e., $\pi_{est}=U(\mathcal{A})$ for all $t$, the trajectory complexity is $mTH$ which completes the proof. ∎

Extended Bilinear Classes

While Bilinear Classes captures most existing models, in this section, we discuss several straightforward extensions of it to incorporate additional models such as Kernelized Nonlinear Regulator (KNR), generalized linear Bellman complete model, and Witness Rank.

Typically, $\pi_{\textrm{est}}(f)$ will be either the uniform distribution on $\mathcal{A}$ or $\pi_{f}$ itself; in the latter case, we refer to the estimation strategy as being on-policy.

We also define $\mathcal{X}_{h}:=\{X_{h}(f)\colon f\in\mathcal{H}\}$ and $\mathcal{X}:=\{\mathcal{X}_{h}:h\in[H]\}$.

We make the following assumptions on the two nonlinear transformations. We assume the slope of $\zeta$ is lower bounded, and $\xi$ is non-decreasing and concave. Similar assumption has been used in generalized linear bandit model (e.g, Russo and Van Roy (2014)).

For $\zeta$, we assume $\zeta(0)=0$ and $\zeta$ is continuously differentiable, and

For $\xi$, we assume $\xi(0)=0$, and $\xi$ is concave and non-decreasing.

We again rely on a reduction to supervised learning style generalization error by extending 5.1 to the following new assumption such that it now includes the additional function class $\mathcal{F}_{h}$.

One simple example of the $\varepsilon_{\textrm{gen}}(m,\mathcal{H},\mathcal{F})$ is when $\mathcal{H}$ and $\mathcal{F}$ are both discrete, $\varepsilon_{\textrm{gen}}(m,\mathcal{H},\mathcal{F})$ will scale in the order of $\widetilde{O}\left(\sqrt{\ln(|\mathcal{H}||\mathcal{F}|)/m}\right)$ via standard uniform convergence analysis.

With the above assumptions, we can show that our algorithm achieves the following regret.

For Generalized Bilinear Class under 6.1, setting parameters properly, we have that with probability at least $1-\delta$:

The proof of the above theorem largely follows the proof of Theorem 5.2, and is deferred to Appendix E.

In this section, we show how the above definition captures KNR (Kakade et al., 2020) which we define next. We note that neither Bellman rank nor Witness rank could capture KNR directly. Specifically, since $\phi(s,a)$ could be nonlinear transformation and reward could be arbitrary (except being bounded in $$), it is not possible to leverage model-free approaches to solve KNR as the value functions and Q functions of a KNR could be too complicated to be captured by function classes with bounded complexity.

Given features $\phi:\mathcal{S}\times\mathcal{A}\to\mathcal{V}$ with $\mathcal{V}$ being some Hilbert space, we say a MDP $\mathcal{M}$ is a Kernelized Nonlinear Regulator (KNR) if it admits the following transition function:

We follow on-policy strategy and set discriminator classes to be empty, i.e., we set $\pi_{est}=\pi_{f}$, and $\mathcal{F}_{h}=\emptyset$ for all $h\in[H]$. Thus, we have for observed transition info $o_{h}=(r_{h},s_{h},a_{h},s_{h+1})$:

2 Generalized Linear Bellman Complete

We first introduce the generalized linear Bellman complete model, and then we show how our framework captures it.

and $\sigma(\mathcal{T}_{h}(\theta_{h+1})^{\top}\phi(s,a))\in\mathcal{H}_{h}$.

Let us define discriminators $\mathcal{F}_{h}:=\left\{f-f^{\prime}:f\in\mathcal{H}_{h},f^{\prime}\in\mathcal{H}_{h}\right\}$. Note that the Bellman complete assumption indicates the following. For any $f:=\{\theta_{0},\dots,\theta_{H-1}\}$, we have $\sigma(\mathcal{T}_{h}(\theta_{h+1})^{\top}\phi(\cdot,\cdot))-\sigma(\theta_{h}^{\top}\phi(\cdot,\cdot))\in\mathcal{F}_{h}$ and $\sigma(\theta_{h}^{\top}\phi(\cdot,\cdot))-\sigma(\mathcal{T}_{h}(\theta_{h+1})^{\top}\phi(\cdot,\cdot))\in\mathcal{F}_{h}$.

Under this assumption (also used in Wang et al. (2019)), we can show that Definition 6.1 captures the generalized linear Bellman complete model.

Setting $\pi_{est}=\pi_{f}$, adding expectation with respect to $s_{h},a_{h},r_{h},s_{h+1}$ under the roll-in policy $\pi_{f}$, we get:

where the third equality uses the generalized linear Bellman complete assumption, and the first inequality uses the fact that $\sigma(\theta_{h}^{\top}\phi(s_{h},a_{h}))-\sigma(\mathcal{T}_{h}(\theta_{h+1})^{\top}\phi(s_{h},a_{h}))\in\mathcal{F}$. Now we continue with the property of the inverse link function as follows.

where the first inequality above uses mean value theorem and $\sigma^{\prime}(x)\geq a,\forall x$ (6.3). Thus, we can conclude that:

The above is captured by Equation 13 with $\zeta(x)=ax$ being a linear function.

Now we consider upper bounding the Bellman error. Denote $f:=\{\theta_{0},\dots,\theta_{H-1}\}$. We have

Thus, we have shown that generalized linear MDP is captured by Definition 6.1 with $\zeta(x)=ax$ and $\xi(x)=b\sqrt{x}$. Note that $\zeta$ and $\xi$ satisfies 6.1. ∎

3 Witness Rank

Similar to Bellman rank, the algorithm and analysis from Sun et al. (2019) rely on $d$ being finite. Below we show how definition 6.1 naturally captures witness rank.

Recall that we denote $f^{\star}$ as the ground truth which in this case means the ground truth transition $P$. This implies that $\langle W_{h}(f^{\star}),X_{h}(f)\rangle=0$ for any $f\in\mathcal{H}$. This allows us to write the above formulation as:

Therefore, it is a Bilinear Class with Discrepancy Family with $\zeta(x)=x$ and $\xi(x)=\frac{1}{\kappa}x$.

Here we also give an example for $\epsilon_{gen}$. For $\mathcal{F}$ and $\mathcal{H}$ with bounded complexity (e.g., discrete $\mathcal{F}$ and discrete $\mathcal{H}$), we still achieve the generalization error, i.e., for all $f$, for all $g\in\mathcal{H}$, with probability at least $1-\delta$:

where $s_{i}\sim d^{\pi_{f}}_{h},a_{i}\sim U(\mathcal{A}),s^{\prime}_{i}\sim P_{h}(\cdot|s,a)$, and the inequality assumes that $\ln(1/\delta)\geq 1$ (see Lemma 12 from Sun et al. (2019) for derivation).

For completeness, we consider factored MDP as a special example here. We refer readers to Sun et al. (2019) for a detailed treatment of how witness rank capturing factored MDP.

We consider state space $\mathcal{S}\subset\mathcal{O}^{d}$ where $\mathcal{O}$ is a discrete set and we denote $s[i]$ as the i-th entry of the state $s$. For each dimension $i$, we denote $\text{pa}_{i}\subset[d]$ as the set of state dimensions that directly influences state dimension $i$ (we call them the parent set of the i-th dimension). In factored MDP, the transition is governed by the following factorized transition:

where $P^{(i)}$ is the condition distribution that governs the transition from $s[\text{pa}_{i}],a$ to $s^{\prime}[i]$. Here, we do not assume any structure on reward function.

Note that the complexity of the problem is captured by the number of parameters in the transition operator, which in this case is equal to $\sum_{i=1}^{d}HA|\mathcal{O}|^{1+|\text{pa}_{i}|}$. Note that when the parent set $\text{pa}_{i}$ is not too big (e.g., a constant that is independent of $d$), this complexity could be exponentially smaller than $|\mathcal{O}|^{d}$ for a MDP that does not have factorized structure.

Thus factored MDP is captured by Definition 6.1 where $\zeta(x)=x$, and $\xi(s)=AHx$. Sun et al. (2019) shows that value function based approaches including Olive Jiang et al. (2017) in worst case requires $2^{H}$ many samples to solve factored MDPs, which in turn indicates that the prior structural complexity such as Bellman rank and Bellman Eluder (Jin et al., 2021) must be exponential in H.

Conclusion

We presented a new framework, Bilinear Classes, together with a new sample efficient algorithm, BiLin-UCB. A key emphasis of the new class and algorithm is that many learnable RL models can be analyzed with the same algorithm and proof.

Our framework is more general than existing ones, and incorporates a large number of RL models with function approximation. Along with the general framework, our work also introduces several important new models including linear $Q^{\star}/V^{\star}$, RKHS Bellman complete, RKHS linear mixture models and low occupancy complexity. Our rates are non-parametric and depend on a new information theoretic quantity—critical information gain, which is an analog to the critical radius from non-parametric statistics. With this new quantity, our results extend prior finite-dimension results to infinite dimensional RKHS setting.

The Bilinear Classes can also be flexibly extended to cover many other examples including Witness Rank and Kernelized Nonlinear Regulator. We believe many other models (potentially even those proposed in the future) can be analyzed via extensions of the Bilinear Classes.

Acknowledgements

We thank Chi Jin and Qinghua Liu for discussions on Section 6 including the generalized linear bellman complete model. We thank Akshay Krishnamurthy for a discussion regarding $Q/V$-Bellman rank.

References

Appendix A Additional Examples of Bilinear Classes

We now include some other examples of Bilinear Classes in addition to ones discussed in Section 4.3.

We consider the feature selection setting introduced by Agarwal et al. [2020b].

We say a MDP $\mathcal{M}$ is low rank feature selection model if there exists (unknown) functions $\mu_{h}^{\star}:\mathcal{S}\mapsto\mathcal{V}$ and (unknown) features $\phi^{\star}:\mathcal{S}\times\mathcal{A}\mapsto\mathcal{V}$, $\psi^{\star}:\mathcal{S}\times\mathcal{A}$ for some Hilbert space $\mathcal{V}$ such that for all $h\in[H]$ and $(s,a,s^{\prime})\in\mathcal{S}\times\mathcal{A}\times\mathcal{S}$

Note that unlike linear MDP model where $\phi^{\star}$ is assumed to be known, here $\phi^{\star}$ is unknown to the learner. We use a function class $\Phi\subset\mathcal{S}\times\mathcal{A}\mapsto\mathcal{V}$ to capture $\phi^{\star}$, i.e., we assume realizability $\phi^{\star}\in\Phi$.

We can define our function class $\mathcal{H}=\mathcal{H}_{0}\times\ldots,\mathcal{H}_{H-1}$ as follows