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