Going Beyond Linear RL: Sample Efficient Neural Function Approximation

1 Introduction . 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 ] . 1.1 Summary of our results . 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 . 1.2 Related Work . 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 . 2 Preliminaries . 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 ) .

Summary: 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.