Going Beyond Linear RL: Sample Efficient Neural Function Approximation
Year: 2022

Abstract:
Deep Reinforcement Learning (RL) powered by neural net approximation of the Q function has had enormous empirical success. While the theory of RL has traditionally focused on linear function approximation (or eluder dimension) approaches, little is known about nonlinear RL with neural net approximations of the Q functions. This is the focus of this work, where we study function approximation with two-layer neural networks (considering both ReLU and polynomial activation functions). Our first result is a computationally and statistically efficient algorithm in the generative model setting under completeness for two-layer neural networks. Our second result considers this setting but under only realizability of the neural net function class. Here, assuming deterministic dynamics, the sample complexity scales linearly in the algebraic dimension. In all cases, our results significantly improve upon what can be attained with linear (or eluder dimension) methods.

heading: 1 Introduction
text: In reinforcement learning (RL), an agent aims to learn the optimal decision-making rule by interacting with an unknown environment [71]. Deep Reinforcement Learning, empowered by deep neural networks [48, 32], has achieved tremendous success in various real-world applications, such as Go [68], Atari [54], Dota2 [7], Texas Holdém poker [56], and autonomous driving [66]. Those modern RL applications are characterized by large state-action spaces, and the empirical success of deep RL corroborates the observation that deep neural networks can extrapolate well across state-action spaces [35, 55, 50].
Although in practice non-linear function approximation scheme is prevalent, theoretical understandings of the sample complexity of RL mainly focus on tabular or linear function approximation settings [69, 38, 5, 41, 63, 88, 1, 43, 44, 77]. These results rely on finite state space or exact linear approximations. Recently, sample efficient algorithms under non-linear function approximation settings are proposed [82, 11, 20, 12, 51, 72, 13, 91, 86]. Those algorithms are based on Bellman rank [40], eluder dimension [64], neural tangent kernel [37, 4, 14, 92], or sequential Rademacher complexity [60, 61]. Yet, the understanding of how deep RL learns and generalizes in large state spaces is far from complete. While the aforementioned works study function approximation structures that possess the nice properties of linear models, such as low information gain and low eluder dimensions, the highly non-linear nature of neural networks renders challenges on their applicability to deep RL. For one thing, recent wisdoms in deep learning theory cast doubt on the ability of neural tangent kernel and random features to model the actual neural networks. Indeed, the neural tangent kernel
∗Alphabetical order. Correspondence to: Baihe Huang, baihehuang@pku.edu.cn, Jason D. Lee, Jasondl@ princeton.edu.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
approximately reduces neural networks to linear models, but the RKHS norm of neural networks is exponential [87]. Moreover, it remains unclear what neural networks models have low eluder dimensions. For example, recent work [13] shows that two layer neural networks have exponential eluder dimension in the dimension of features. Thus, the mismatch between the empirical success of deep RL and its theoretical understanding remains significant, which yields the following important question:
What are the structural properties that allow sample-efficient algorithms for RL with neural network function approximation?
Recent work in RL suggests that learning RL with neural networks function approximation is exponentially hard [13, 49, 52]. In this paper, however, we advance the understanding of the above question by displaying several structural properties that allow efficient RL algorithms with neural function approximations. We consider several value function approximation models that possess high information gain and high eluder dimension. Specifically, we study two structures, namely twolayer neural networks and structured polynomials (i.e. two-layer neural networks with polynomial activation functions), under two RL settings, namely RL with simulator model and online RL. In the simulator (generative model) setting [45, 67], the agent can simulate the MDP at any state-action pair. In online RL, the agent can only start at an initial state and interact with the MDP step by step. The goal in both settings is to find a near-optimal policy while minimizing the number of samples used.
We obtain the following results. For the simulator setting, we propose sample-efficient algorithms for RL with two-layer neural network function approximation. Under either policy completeness, Bellman completeness, or gap conditions, our method provably learns near-optimal policy with polynomial sample complexities. For online RL, we provide sample-efficient algorithms for RL with structured polynomial function approximation. When the transition is deterministic, we also present sample-efficient algorithms under only the realizability assumption [18, 78]. Our main techniques are based on neural network recovery [90, 39, 28], and algebraic geometry [53, 65, 8, 76].

heading: 1.1 Summary of our results
text: Our main results in different settings are summarized in Table 1. We consider two-layer neural networks f(x) = ⟨v, σ(Wx)⟩ (where σ is ReLU activation) and rank k polynomials (see Example 4.3).
We make the following elaborations on Table 1.
• For simplicity, we display only the dependence on the feature dimension d, network width or polynomial rank k, precision ϵ, and degree p (of polynomials).
• In the table Sim. denotes simulator model, Onl. denotes online RL, Det. denotes deterministic transitions, Stoch. denotes stochastic transitions, Gap. denotes gap condition, (R) denotes realizability assumption only, and (C) denotes completeness assumption (either policy complete or Bellman complete) together with realizability assumption.
• We apply [21] for the deterministic transition baseline, and apply [19] for the stochastic transition baseline. We are unaware of any methods that can directly learn MDP with neural network value function approximation2.
• In polynomial case, the baseline first vectorizes the tensor ( 1 x )⊗p into a vector in R(d+1)p
and then performs on this vector. In the neural network case, the baseline uses a polynomial of degree 1/ϵ to approximate the neural network with precision ϵ and then vectorizes the polynomial into a vector in Rdpoly(1/ϵ) . The baseline method for realizable model (denoted
by (*)) needs a further gap assumption of gap ≥ dpoly(1/ϵ)ϵ to avoid the approximation error from escalating [21]; note for small ϵ this condition never holds but we include it in the table for the sake of comparison.
• In rank k polynomial case, our result O(dk) in simulator model can be found in Theorem 4.7 and our result O(dk) in online RL model can be found in Theorem 4.8. These results only require a realizability assumption. Efficient explorations are guaranteed by algebraic-geometric arguments. In neural network model, our result Õ(poly(d) · exp(k)) in simulator model can be found in Theorem 3.4. This result also only relies on the realizability assumption. For stochastic transitions, our result Õ(poly(d, k)/ϵ2) works for either policy complete or Bellman complete settings, as in Theorem 3.5 and Theorem 3.6 respectively. The Õ(poly(d, k)) result for gap condition can be found in Theorem 3.8.

heading: 1.2 Related Work
text: Linear Function Approximation. RL with linear function approximation has been widely studied under various settings, including linear MDP and linear mixture MDP [44, 89, 85]. While these papers have proved efficient regret and sample complexity bounds, their analyses relied heavily on two techniques: they used the confidence ellipsoid to quantify the uncertainty, and they used the elliptical potential lemma to bound the total uncertainty [2]. These two techniques were integral to their analyses but are so restrictive that they generally do not extend to nonlinear cases.
Eluder Dimension. [62, 59] proposed eluder dimension, a complexity measure of the function space, and proved regret and sample complexity bounds that scaled with the eluder dimension, for bandits and reinforcement learning [73, 42]. They also showed that the eluder dimension is small in several settings, including generalized linear models and LQR. However, as shown in [13], the eluder dimension could be exponentially large even with a single ReLU neuron, which suggested the eluder dimension would face difficulty in dealing with neural network cases. The eluder dimension is only known to give non-trivial bounds for linear function classes and monotone functions of linear function classes. For structured polynomial classes, the eluder dimension simply embeds into an ambient linear space of dimension dp, where d is the dimension, and p is the degree. This parallels the lower bounds in linearization / neural tangent kernel (NTK) works [79, 30, 3], which show that linearization also incurs a similarly large penalty of dp sample complexity, and more advanced algorithm design is need to circumvent linearization[6, 9, 22, 83, 26, 58, 29, 57, 34, 75, 10].
Bellman Rank and Completeness. [40, 70] studied RL with general function approximation. They used Bellman rank to measure the error of the function class under the Bellman operator and gave proved bounds in the term of it. Recently, [16] propose bilinear rank and encompass more function approximation models. However, it is hard to bound either the Bellman rank or the bilinear rank for neural nets. Therefore, their results are not known to include the neural network approximation setting. Another line of work shows that exponential sample complexity is unavoidable even with good representations [19, 80, 78], which implies the realizability assumption alone might be insufficient for function approximations.

heading: 2 Preliminaries
text: An episodic Markov Decision Process (MDP) is defined by the tupleM = (S,A, H,P, r) where S is the state space, A is the action set, H is the number of time steps in each episode, P is the transition kernel and r is the reward function. In each episode the agent starts at a fixed initial state s1 and at each time step h ∈ [H] it takes action ah, receives reward rh(sh, ah) and transits to sh+1 ∼ P(·|sh, ah). A deterministic policy π is a length-H sequence of functions π = {πh : S 7→ A}Hh=1. Given a policy π, we define the value function V πh (s) as the expected sum of reward under policy π starting from
2Prior work on neural function approximation has focused on neural tangent kernels, which would require dpoly(1/ϵ) to approximate a two-layer network [31].
sh = s:
V πh (s) := E [ H∑ t=h rt(st, at)|sh = s ] and we define the Q function Qπh(s, a) as the the expected sum of reward taking action a in state sh = s and then following π:
Qπh(s, a) := E [ H∑ t=h rt(st, at)|sh = s, ah = a ] .
The Bellman operator Th applied to Q-function Qh+1 is defined as follow Th(Qh+1)(s, a) := rh(s, a) + Es′∼P(·|s,a)[max
a′ Qh+1(s
′, a′)].
There exists an optimal policy π∗ that gives the optimal value function for all states, i.e. V π ∗
h (s) =
supπ V π h (s) for all h ∈ [H] and s ∈ S. For notational simplicity we abbreviate V π
∗ as V ∗ and
correspondingly Qπ ∗ as Q∗. Therefore Q∗ satisfies the following Bellman optimality equations for all s ∈ S, a ∈ A and h ∈ [H]:
Q∗h(s, a) = Th(Q∗h+1)(s, a).
The goal is to find a policy π that is ϵ-optimal in the sense that V ∗1 (s1)− V π1 (s1) ≤ ϵ, within a small number of samples. We consider two query models of interacting with MDP:
• In the simulator model ([45], [67]), the agent interacts with a black-box that simulates the MDP. At each time step h ∈ [H], the agent can start at a state-action pair (s, a) and interact with the black box by executing some policy π chosen by the agent.
• In online RL, the agent can only start at the initial state and interact with the MDP by using a policy and observing the rewards and the next states. In each episode k, the agent proposes a policy πk based on all history up to episode k − 1 and executes πk to generate a single trajectory {skh, akh}Hh=1 with akh = πkh(skh) and skh+1 ∼ P(·|skh, akh).

heading: 2.1 Function approximation
text: In reinforcement learning with value function approximation, the learner is given a function class F = F1×· · ·×FH , whereFh ⊂ {f : S×A 7→ [0, 1]} is a set of candidate functions to approximate Q∗. The following assumption is commonly adopted in the literature [43, 74, 42, 17]. Assumption 2.1 (Realizability). Q∗h ∈ Fh for all h ∈ [H].
The function approximation is equipped with feature mapping ϕ : S ×A 7→ {u ∈ Rd : ∥u∥2 ≤ Bϕ} that is known to the agent. We focus the continuous action setting (e.g. in control and robotics problems) and make the following regularity assumption on the feature function ϕ. Assumption 2.2 (Bounded features). Assume ϕ(s, a) ≤ Bϕ,∀(s, a) ∈ S ×A.
Notation For any vector x ∈ Rd, let xmax := maxi∈[d] xi and xmin := mini∈[d] xi. Let si(·) denote the i-th singular value, smin(·) denotes the minimum eigenvalue and smax(·) denotes the maximum eigenvalue. The conditional number is defined by κ(·) := smax(·)/smin(·). We use ⊗ to denote Kronecker product and ◦ to denote Hadamard product. For a given integer H , we use [H] to denote the set {1, 2, . . . ,H}. For a function f : X 7→ Y, we use f−1(y) := {x ∈ X : f(x) = y} to denote the preimage of y ∈ Y. We use the shorthand x ≲ y (x ≳ y) to indicate x ≤ O(y) (x ≥ Ω(y)).

heading: 3 Neural Network Function Approximation
text: In this section we show sample-efficient algorithms with neural network function approximations. The function class of interest is given in the following definition. More general neural network class is discussed in Appendix B.5.
Definition 3.1 (Neural Network Function Class). We use FNN to denote the function class of f(ϕ(s, a)) : S ×A 7→ R where f(x) = ⟨v, σ(Wx)⟩ : Rd 7→ R is a two-layer neural network where σ is ReLU, ∥W∥F ≤ BW , v ∈ {±1}k, ∏k i=1 si(W )/smin(W ) ≤ λ, smax(W )/smin(W ) ≤ κ and k ≤ d. Here ϕ : A × S 7→ Rd is a known feature map whose image contains a ball {u ∈ Rd : ∥u∥2 ≤ δϕ} with δϕ ≥ d · polylog(d).3
We introduce the following completeness properties in the setting of value function approximations. Along with Assumption 2.1, they are commonly adopted in the literature . Definition 3.2 (Policy complete). Given MDPM = (S,A,P, r,H), function class Fh : S ×A 7→ R, h ∈ [H] is called policy complete if for all π and h ∈ [H], Qπh ∈ Fh. Definition 3.3 (Bellman complete). Given MDPM = (S,A,P, r,H), function class Fh : S ×A 7→ R, h ∈ [H] is called Bellman complete if for all h ∈ [H] and Qh+1 ∈ Fh+1, Th(Qh+1) ∈ Fh.

heading: 3.1 Warmup: Realizable Q∗ with deterministic transition
text: We start by considering the case when the transition kernel is deterministic. In this case only Assumption 2.1 is required for the expressiveness of neural network function approximations. Algorithm 1 learns optimal policy from time step H to 1. Suppose we have learned policies πh+1, . . . , πH at level h and they are exactly the optimal policies. We first explore features ϕ(sih, a i h) over a standard Gaussian distribution, and if ∥ϕ(sih, aih)∥2 ≥ δϕ then we simply skip this trial. Recall that δϕ ≥ d ·poly log(d), so with high probability (w.r.t d) almost all feature samples will be explored. We next construct an estimate Q̂ih of Q ∗(sih, a i h) by collecting cumulative rewards using πh+1, . . . , πH as the roll-out. Since the transition is deterministic, Q̂ih = Q ∗(sih, a i h) for all explored samples (s i h, a i h). Recall that Q∗h is a two-layer neural network, we can now recover its parameters in Line 12 exactly by invoking techniques in neural network optimization (see, e.g. [39, 27, 90]). Details of this step can be found in Appendix B.5, where the method is mainly based on [90]. This means the reconstructed Q̂h in Line 13 is precisely Q∗, and the algorithm can thus find optimal policy π∗h in the h-th level.
Algorithm 1 Learning realizable Q∗ with deterministic transition 1: for h = H, . . . 1 do 2: Sample xih, i ∈ [n] from standard Gaussian N (0, Id) 3: for i ∈ [n] do 4: if ∥xih∥ ≤ δϕ then 5: Find (sih, a i h) ∈ ϕ−1(xih) and locate the state sih in the generative model
6: Pull action aih and use πh+1, . . . , πH as the roll-out to collect rewards r (i) h , . . . , r (i) H 7: Construct estimation
Q̂ih ← r (i) h + · · ·+ r (i) H
8: else 9: Let Q̂ih ← 0.
10: end if 11: end for 12: Compute (vh,Wh)← NEURALNETRECOVERY({(xih, Q̂ih) : i ∈ [n]}) 13: Set Q̂h(s, a)← v⊤h σ(Whϕ(s, a)) 14: Let πh(s)← argmaxa∈S Q̂h(s, a) 15: end for 16: Return π1, . . . , πH
Theorem 3.4. (Informal) If n ≥ d · poly(κ, k, λ, log d,BW , Bϕ, H), then with high probability Algorithm 1 learns the optimal policy.
The formal statement and complete proof are deferred to the Appendix B.1. The main idea of exact neural network recovery can be summarized in the following. We first use method of moments to
3Here the δϕ is chosen only for simplicity. In general this assumption can be relaxed to that the image of ϕ contains an arbitrary dense ball near the origin, since one can always rescale the feature mapping in the neural function approximation.
find a ‘rough’ parameter recovery. If this ‘rough’ recovery is sufficiently close to the true parameter, the empirical ℓ2 loss function is locally strongly convex and there is unique global minimum. Then we can apply gradient descent to find this global minimum which is exactly the true parameter.

heading: 3.2 Policy complete neural function approximation
text: Now we consider general stochastic transitions. Difficulties arise in this scenario due to noises in the estimation of Q-functions. In the presence of model misspecification, these noises cause estimation errors to amplify through levels and require samples to be exponential in H . In this section, we show that neural network function approximation is still learnable, assuming the function class FNN is policy complete with regard to MDP M. Thus for all π ∈ Π, we can denote Qπh(s, a) = ⟨vπ, σ(Wπϕ(s, a))⟩.
Algorithm 2 Learn policy complete NN with simulator. 1: for h = H, . . . 1 do 2: Sample xih, i ∈ [n] from standard Gaussian N (0, Id) 3: for i ∈ [n] do 4: if ∥xih∥ ≤ δϕ then 5: Find (sih, a i h) ∈ ϕ−1(xih) and locate the state sih in the generative model
6: Pull action aih and use πh+1, . . . , πH as the roll-out to collect rewards r (i) h , . . . , r (i) H 7: Construct unbiased estimation of Qπh+1,...,πHh (s i h, a i h)
Q̂ih ← r (i) h + · · ·+ r (i) H
8: else 9: Let Q̂ih ← 0.
10: end if 11: end for 12: Compute (vh,Wh)← NEURALNETNOISYRECOVERY({(xih, Q̂ih) : i ∈ [n]}) 13: Set Q̂h(s, a)← v⊤h σ(Whϕ(s, a)) 14: Let πh(s)← argmaxa∈S Q̂h(s, a) 15: end for 16: Return π1, . . . , πH
Algorithm 2 learns policy from level H,H − 1, . . . , 1. In level h, the algorithm has learned policy πh+1, . . . , πH that is only sub-optimal by (H − h)ϵ/H . Then it explores features ϕ(s, a) from N (0, Id). The algorithm then queries (s, a) and uses learned policy πh+1, . . . , πH as roll out to collect an unbiased estimate of the Q-function Qπh+1,...,πHh (s, a). Since Q πh+1,...,πH h (s, a) ∈ FNN is a two-layer neural network, it can then be recovered from samples. Details of this step can be found in Appendix B.5, where the methods are mainly based on [90]. The algorithm then reconstructs this Q-function and finds its optimal policy πh.
Theorem 3.5. (Informal) Fix ϵ, t, if n ≥ ϵ−2 · d · poly(κ, k,BW , Bϕ, H, log(d/t)), then with probability at least 1− t Algorithm 2 returns an ϵ-optimal policy π.
The formal statement and complete proof are deferred to Appendix B.2. Notice that unlike the case of Theorem 3.4, the sample complexity does not depend on λ, thus avoiding the potential exponential dependence in k.
The main idea of the proof is that at each time step a neural network surrogate of Q∗ can be constructed by the policy already learned. Suppose we have learned πh+1, . . . , πH in level h, then from policy completeness Qπh+1,...,πHh belongs to FNN and we can interact with the simulator to obtain its estimate Q̂h. If ∥Q̂h −Q πh+1,...,πH h ∥∞ is small, the optimistic planning based on Q̂h is not far from the optimal policy of Qπh+1,...,πHh . Therefore the errors can be decoupled into the errors in recovering Q
πh+1,...,πH h and the suboptimality of Q πh+1,...,πH h , which depends on level h + 1. This reasoning
can then be recursively performed to level H , and thus we can bound the suboptimality of πh.

heading: 3.3 Bellman complete neural function approximation
text: In addition to policy completeness, we show that neural network function approximation can also learn efficiently under the setting where the function class FNN is Bellman complete with regard to MDPM. Specifically, for Qh+1 ∈ Fh+1, there are vQh+1 and WQh+1 such that Th(Qh+1)(s, a) = ⟨vQh+1 , σ(WQh+1ϕ(s, a))⟩. Algorithm 3 is similar to the algorithm in previous section. Suppose in level h, the algorithm has constructed the Q-function Q̂h+1(s, a) = v⊤h+1σ(Wh+1ϕ(s, a)) that is (H − h)ϵ/H-close to the optimal Q∗h+1. It then recovers weights vh,Wh from Th(Q̂h+1)(s, a) = ⟨vQ̂h+1 , σ(W Q̂h+1ϕ(s, a))⟩, using unbiased estimates rh(sih, a i h) + V̂h+1(s i h+1). The Q-function Q̂h(s, a) = v ⊤ h σ(Whϕ(s, a)) reconstructed from weights vh,Wh is thus (H − h+ 1)ϵ/H-close to the Q∗h.
Algorithm 3 Learn Bellman complete NN with simulator. 1: for h = H, . . . 1 do 2: Sample xih, i ∈ [n] from standard Gaussian N (0, Id) 3: for i ∈ [n] do 4: if ∥xih∥ ≤ δϕ then 5: Find (sih, a i h) ∈ ϕ−1(xih) and locate the state sih in the generative model
6: Pull action aih and and observe rh(s i h, a i h), s i h+1 7: Construct unbiased estimation of Th(Q̂h+1)(sih, aih)
Q̂ih ← rh(sih, aih) + V̂h+1(sih+1)
8: else 9: Let Q̂ih ← 0.
10: end if 11: end for 12: Compute (vh,Wh)← NEURALNETNOISYRECOVERY({(xih, Q̂ih) : i ∈ [n]}) 13: Set Q̂h(s, a)← v⊤h σ(Whϕ(s, a)) and V̂h ← maxa∈A Q̂h(s, a) 14: Let πh(s)← argmaxa∈A Q̂h(s, a) 15: end for 16: Return π1, . . . , πH
Theorem 3.6. (Informal) Fix ϵ, t, if n ≥ ϵ−2 · d · poly(κ, k,BW , Bϕ, H, log(d/t)), then with probability at least 1− t Algorithm 3 returns an ϵ-optimal policy π.
Due to Bellman completeness, the error of estimation Q̂h can be controlled recursively. In fact, we can show ∥Q̂h − Q∗(s, a)∥∞ is small by induction. The formal statement and detailed proof are deferred to Appendix B.3. Similar to Theorem 3.5, the sample complexity does not explicitly depend on λ, thus avoiding potentially exponential dependence in k.

heading: 3.4 Realizable Q∗ with optimality gap
text: In this section we consider MDPs where there is a non-zero gap between the optimal policy and any other ones. This concept, known as optimality gap, is widely used in reinforcement learning and bandit literature [20, 19, 21].
Definition 3.7. The optimality gap is defined as
ρ = inf a:Q∗(s,a) ̸=V ∗(s)
V ∗(s)−Q∗(s, a).
We show that in the presence positive optimality gap, there exists an algorithm that can learn the optimal policy with polynomial samples even without the completeness assumptions. Intuitively, this is because one only needs to recover the neural network up to precision ρ/4 in order to make sure the greedy policy is identical to the optimal one. The formal statement and proof are deferred to Appendix B.4.
Theorem 3.8. (Informal) Fix t ∈ (0, 1), if n = dρ2 · poly(κ, k,BW , Bϕ, H, log(d/t)), then with probability at least 1− t there exists an algorithm that returns the optimal policy π∗. Remark 3.9. In all aforementioned methods, there are two key components that allow efficient learning. First, the exploration is conducted in a way that guarantees an ℓ∞ recovery of candidate functions. By ℓ∞ recovery we mean the algorithm recovers a candidate Q-function in this class deviating from the target function Q∗ by at most ϵ uniformly for all state-action pairs in the domain of interest. This notion of learning guarantee has received study in active learning [33, 46] and recently gain interest in contextual bandits [25]. Second, the agent constructs unbiased estimators of certain approximations to Q∗ that lie in the neural function approximation class. This allows the recovery error to decouple linearly across time steps, which is made possible in several well-posed MDP instances, such as deterministic MDPs, MDPs with completeness assumptions, and MDPs with gap conditions. In principle, we note that provably efficient RL algorithms with general function approximation is possible as long as the above two components are present. We will see in the next section another example of learning RL with highly non-convex function approximations, where the function class of interest, admissible polynomial families, also allows for exploration schemes to achieve ℓ∞ recovery.

heading: 4 Polynomial Realizability
text: In this section, we study the sample complexity to learn deterministic MDPs under polynomial realizability. We identify sufficient and necessary conditions for efficiently learning the MDPs for two different settings — the generative model setting and the online RL setting. Specifically, we show that if the image of the feature map ϕh(sh, ah) satisfies some positive measure conditions, then by random exploring, we can identify the optimal policy with samples linear in the algebraic dimension of the underlying polynomial class. We also provide a lower bound example showing the separation between the two settings.
Next, we introduce the notion of Admissible Polynomial Families, which are the families of structured polynomials that enable efficient learning. Definition 4.1 (Admissible Polynomial Families). For x ∈ Rd, denote x̃ = [1, x⊤]⊤. Let X :={ x̃⊗p : x ∈ Rd } . For any algebraic variety V , we define FV := {fΘ(x) = ⟨Θ, x̃⊗p⟩ : Θ ∈ V} as the polynomial family parameterized by Θ ∈ V . We say FV is admissible4 w.r.t. X , if for any Θ ∈ V , dim(X ∩ {X ∈ X : ⟨X,Θ⟩ = 0⟩}) < dim(X ) = d. We define the dimension D of the family to be the dimension of V .
The following theorem shows that to learn an admissible polynomial family, the sample complexity only scales with the algebraic dimension of the polynomial family.
Theorem 4.2 ([36]). Consider the polynomial family FV of dimension D. For n ≥ 2D, there exists a Lebesgue-measure zero set N ∈ Rd × . . .Rd, such that if (x1, · · · , xn) /∈ N , for any yi, there is a unique f (or no such f ) to the system of equations yi = f(xi) for f ∈ FV .
We give two important examples of admissible polynomial families with low dimension.
Example 4.3. (Low-rank Polynomial of rank k) The function f ∈ FV is a polynomial with k terms, that is
F (x) = k∑ i=1 λi⟨vi, x⟩pi ,
where p = max{pi}. The dimension of this family is upper bounded by D ≤ dk. Neural network with monomial/polynomial activation functions are low-rank polynomials.
Example 4.4. The function f ∈ FV is of the form f(x) = q(Ux), where U ∈ Rk×d and q is a degree p polynomial. The polynomial q and matrix U are unknown. The dimension of this family is upper bounded by D ≤ d(k + 1)p.
Next, we introduce the notion of positive measure.
4Admissible means the dimension of X decreases by one when there is an additional linear constraint ⟨Θ, X⟩ = 0
Definition 4.5. We say a measurable set E ∈ Rd is of positive measure if µ(E) > 0, where µ is the standard Lebesgue measure on Rd.
If a measurable set E satisfies µ(E) > 0, then there exists a procedure to draw samples from E, such that for any N ⊂ Rd of Lebesgue-measure zero, the probability that the sample falls in N is zero. In fact, the sampling probability can be given by Px∈N (0,Id)(·|x ∈ E). The intuition behind its definition is that for all admissible polynomial families, the set of (x1, · · · , xn) with "redundant information" about learning the parameter Θ is of Lebesgue-measure zero. Therefore, a positive measure set allows you to query randomly and avoids getting coherent measurements.
Next two theorems identify the sufficent conditions for efficiently learning deterministic MDPs under polynomial realizability. Specifically, under online RL setting, we require the strong assumption that the set {ϕh(s, a)|a ∈ A} is of positive measure for all h ∈ [H] and all s ∈ S , while under generative model setting, we only require the union set ⋃ s∈S{ϕh(s, a)|a ∈ A} to be of positive measure for all h ∈ [H]. The algorithms for solving the both cases are summarized in Algorithms 4 and 5. Assumption 4.6 (Polynomial Realizability). For all h ∈ [H], Q∗h(sh, ah), viewed as the function of ϕh(sh, ah), lies in some admissible polynomial family FVh with dimension bounded by D. Theorem 4.7. For the generative model setting, assume that the set {ϕh(s, a) | s ∈ S, a ∈ A} is of positive measure at any level h. Under the polynomial realizability, Algorithm 4 almost surely learns the optimal policy π⋆ with at most N = 2DH samples. Theorem 4.8. For the online RL setting, assume that {ϕh(s, a) | a ∈ A} is of positive measure for every state s at every level h. Under polynomial realizability, within T = 2DH episodes, Algorithm 5 learns the optimal policy π⋆ almost surely.
Algorithm 4 Dynamic programming under generative model settings 1: for h = H, · · · , 1 do 2: Sample 2D points {ϕh(s(i)h , a (i) h )}2Di=1 according to Px∈N (0,Id)(·|x ∈ Eh) where Eh =
{ϕh(s, a) | s ∈ S, a ∈ A}. 3: Query the generative model with state-action pair (s(i)h , a (i) h ) at level h for i = 1, . . . , 2D,
and observe the next state s̃(i)h and reward r (i) h .
4: Solve for Q∗h with the 2D equations Q ∗ h(s (i) h , a (i) h ) = r (i) h + V ∗ h+1(s̃ (i) h ). 5: Set π∗h(s) = argmaxa Q ∗ h(s, a) and V ∗ h (s) = maxa Q ∗ h(s, a). 6: end for 7: Output π∗
Algorithm 5 Dynamic programming under online RL settings 1: for h = H, · · · , 1 do 2: Fix any action sequence a1, · · · , ah−1. 3: Play a1, · · · , ah−1 for the first h − 1 levels and reach a state sh. Sample 2D points {ϕh(sh, a(i)h )}2Di=1 according to Px∈N (0,Id)(·|x ∈ Eh) where Eh = {ϕh(sh, a) | a ∈ A}.
4: Play a(i)h at sh for i = 1, . . . , 2D, and observe the next state s̃ (i) h and reward r (i) h . 5: Solve for Q∗h with the 2D equations Q ∗ h(s (i) h , a (i) h ) = r (i) h + V ∗ h+1(s̃ (i) h ). 6: Set π∗h(s) = argmaxa Q ∗ h(s, a) and V ∗ h (s) = maxa Q ∗ h(s, a). 7: end for 8: Output π∗
We remark that our Theorem 4.8 for learning MDPs under the online RL setting relies on a very strong assumption that allows the learner to explore randomly for any state. However, this assumption is necessary in some sense, as is suggested by our lower bound example in the next subsection.

heading: 4.1 Necessity of Generic Feature Maps in Online RL
text: In this section, we consider lower bounds for learning deterministic MDPs with polynomial realizable Q∗ under online RL setting. Our goal is to show that in the online setting the generic assumption on
the feature maps ϕh(s, ·) is necessary. On the contrary, under the generative model setting one can efficiently learn the MDPs without such a strong assumption, since at every level h the we can set the state arbitrarily.
MDP construction We briefly introduce the intuition of our construction. Consider a family of MDPs with only two states S = {Sgood, Sbad}. we set the feature map ϕh such that, for the good state Sgood, it allows the learner to explore randomly, i.e., {ϕh(Sgood, a) | a ∈ A} is of postive measure. However, for the bad state Sbad, all actions are mapped to some restricted set, which forbids random exploration, i.e., {ϕh(Sbad, a) | a ∈ A} is measure zero. This is illustrated in Figure 1. Specifically, at least Ω(dp) actions are needed to identify the groud-truth polynomial of Q∗h for Sbad, while O(d) actions suffice for Sgood.
The transition Ph is constructed as Ph(sbad|s, a) = 1 for all s ∈ S, a ∈ A, which means it is impossible for the online scenarios to reach the good state for h > 1.
Theorem 4.9. There exists a family of MDPs satisfying Assumption 4.6, such that the set {ϕh(s, a) | s ∈ S, a ∈ A} is of positive measure at any level h, but for all h there is some sbad ∈ S such that {ϕh(sbad, a) | a ∈ A} is measure zero. Under the online RL setting, any algorithm needs to play at least Ω(dp) episodes to identify the optimal policy. On the contrary, under the generative model setting, only O(d)H samples are needed.

heading: 5 Conclusions
text: In this paper, we consider neural network and polynomial function approximations in the simulator and online settings. To our knowledge, this is the first paper that shows sample-efficient reinforcement learning is possible with neural net function approximation. Our results substantially improve upon what can be achieved with existing results that primarily rely on embedding neural networks into linear function classes. The analysis reveals that for function approximations that allows for efficient ℓ∞ recovery, such as two layer neural networks and admissible polynomial families, reinforcement learning can be reduced to parameter recovery problems, as well-studied in theories for deep learning, phase retrieval, and etc. Our method can also be potentially extended to handle three-layer and deeper neural networks, with advanced tools in [23, 24].
Our results for polynomial activation require deterministic transitions, since we cannot handle how noise propagates in solving polynomial equations. We leave to future work an in-depth study of the stability of roots of polynomial systems with noise, which is a fundamental mathematical problem and even unsolved for homogeneous polynomials. In particular, noisy tensor decomposition approaches combined with zeroth-order optimization may allow for stochastic transitions [36].
In the online RL setting, we can only show efficient learning under a very strong yet necessary assumption on the feature mapping. We leave to future work identifying more realistic and natural conditions which permit efficient learning in the online RL setting.
Finally, in future work, we hope to consider deep neural networks where parameter recovery or ℓ∞ error is unattainable, and deep reinforcement learning with representation learning [84, 15].

heading: Acknowledgements
text: JDL acknowledges support of the ARO under MURI Award W911NF-11-1-0303, the Sloan Research Fellowship, NSF CCF 2002272, and an ONR Young Investigator Award. QL is supported by NSF 2030859 and the Computing Research Association for the CIFellows Project. SK acknowledges funding from the NSF Award CCF-1703574 and the ONR award N00014-18-1-2247.

Review summaries:
This paper studies sample efficient RL with two-layer neural network function approximation with a generative model. In particular, they prove that they can learn an optimal policy with poly samples with only realizability in deterministic MDP, and that the sample complexity for learning a near-optimal policy in stochastic MDPs is in an order of $\epsilon^{-2}$ in two cases: policy completeness assumption and Bellman completeness assumption. The algorithms they use are based on simple (non-optimistic) value regression. 

This paper studies reinforcement learning with a 1-hidden layer neural network and low-rank polynomial function approximation schemes.   Under mostly standard assumptions about realizability, completness etc, they provide algorithms mainly focusing on sample complexity in cases when there is access to a perfect generative model of the transitions and in the fully online setting, where planning is not possible in episodic reinforcement learning.  They improve upon the sample complexity with respect to existing baselines for the respective function approximation schemes.

The paper considers RL with neural function approximation, extending significantly previous works that only focuses on linear feature approximation. In particular, the authors propose (1) an efficient algorithm using two-layer NNs in a generative model assuming completeness, and (2) an efficient algorithm with two-layer NNs but only assuming realizability.  The results significantly extends prior research on linear methods, or methods with bounded eluder dimension.

This paper focuses on analyzing sampling efficiency of nonlinear reinforcement learning with neural network approximations of the $Q$ function. It reduces the sampling complexity from $O(d^{\text{poly}(1/\epsilon)})$ to $O(\text{poly}(d) \exp(k))$ where $k$ is the width of the hidden layer of a two-layer fully connected neural network. The idea is described accurately and the paper is overall well written.