{'fraction_non_alphanumeric': 0.07427070845464406, 'fraction_numerical': 0.023919620939658617, 'mean_word_length': 3.6694655471144877, 'pattern_counts': {'":': 0, '<': 1, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Generative Adversarial Imitation Learning with Neural Networks: Global Optimality and Convergence Rate

Generative adversarial imitation learning (GAIL) demonstrates tremendous success in practice, especially when combined with neural networks. Different from reinforcement learning, GAIL learns both policy and reward function from expert (human) demonstration. Despite its empirical success, it remains unclear whether GAIL with neural networks converges to the globally optimal solution. The major difficulty comes from the nonconvex-nonconcave minimax optimization structure. To bridge the gap between practice and theory, we analyze a gradient-based algorithm with alternating updates and establish its sublinear convergence to the globally optimal solution. To the best of our knowledge, our analysis establishes the global optimality and convergence rate of GAIL with neural networks for the first time.


Introduction
The goal of imitation learning (IL) is to learn to perform a task based on expert demonstration (Ho and Ermon, 2016). In contrast to reinforcement learning (RL), the agent only has access to the expert trajectories but not the rewards. The most straightforward approach of IL is behavioral cloning (BC) (Pomerleau, 1991). BC treats IL as the supervised learning problem of predicting the actions based on the states. Despite its simplicity, BC suffers from the compounding errors caused by covariate shift (Ross et al., 2011;Ross and Bagnell, 2010). Another approach of IL is inverse reinforcement learning (IRL) (Russell, 1998;Ng and Russell, 2000;Levine and Koltun, 2012;Finn et al., 2016), which jointly learns the reward function and the corresponding optimal policy. IRL formulates IL as a bilevel optimization problem. Specifically, IRL solves an RL subproblem given a reward function at the inner level and searches for the reward function which makes the expert policy optimal at the outer level. However, IRL is computationally inefficient as it requires fully solving an RL subproblem at each iteration of the outer level. Moreover, the desired reward function may be nonunique. To address such issues of IRL, Ho and Ermon (2016) propose generative adversarial imitation learning (GAIL), which searches for the optimal policy without fully solving an RL subproblem given a reward function at each iteration. GAIL solves IL via minimax optimization with alternating updates. In particular, GAIL alternates between (i) minimizing the discrepancy in expected cumulative reward between the expert policy and the learned policy and (ii) maximizing such a discrepancy over the reward function class. Such an alternating update scheme mirrors the training of generative adversarial networks (GANs) (Goodfellow et al., 2014;Arjovsky et al., 2017), where the learned policy acts as the generator while the reward function acts as the discriminator.
Incorporated with neural networks, which parameterize the learned policy and the reward function, GAIL achieves significant empirical success in challenging applications, such as natural language processing (Yu et al., 2016), autonomous driving (Kuefler et al., 2017), human behavior modeling (Merel et al., 2017), and robotics (Tai et al., 2018). Despite its empirical success, GAIL with neural networks remains less understood in theory. The major difficulty arises from the following aspects: (i) GAIL involves minimax optimization, while the existing analysis of policy optimization with neural networks (Anthony and Bartlett, 2009;Liu et al., 2019;Bhandari and Russo, 2019;Wang et al., 2019) only focuses on a minimization or maximization problem. (ii) GAIL with neural networks is nonconvex-nonconcave, and therefore, the existing analysis of convex-concave optimization with alternating updates is not applicable (Nesterov, 2013). There is an emerging body of literature (Rafique et al., 2018;Zhang et al., 2019b) that casts nonconvex-nonconcave optimization as bilevel optimization, where the inner level is solved to a high precision as in IRL. However, such analysis is not applicable to GAIL as it involves alternating updates.
In this paper, we bridge the gap between practice and theory by establishing the global optimality and convergence of GAIL with neural networks. Specifically, we parameterize the learned policy and the reward function with two-layer neural networks and consider solving GAIL by alternatively updating the learned policy via a step of natural policy gradient (Kakade, 2002;Peters and Schaal, 2008) and the reward function via a step of gradient ascent. In particular, we parameterize the state-action value function (also known as the Q-function) with a two-layer neural network and apply a variant of the temporal difference algorithm (Sutton and Barto, 2018) to solve the policy evaluation subproblem in natural policy gradient. We prove that the learned policyπ converges to the expert policy π E at a 1/ √ T rate in the R-distance (Chen et al., 2020), which is defined as D R (π E ,π) = max r∈R J(π E ; r) − J(π; r). Here J(π; r) is the expected cumulative reward of a policy π given a reward function r(s, a) and R is the reward function class. The core of our analysis is constructing a potential function that tracks the R-distance. Such a rate of convergence implies that the learned policyπ (approximately) outperforms the expert policy π E given any reward func-tion r ∈ R within a finite number of iterations T . In other words, the learned policyπ is globally optimal. To the best of our knowledge, our analysis establishes the global optimality and convergence of GAIL with neural networks for the first time. It is worth mentioning that our analysis is straightforwardly applicable to linear and tabular settings, which, however, are not our focus.
Related works. Our work extends an emerging body of literature on RL with neural networks (Xu et al., 2019a;Zhang et al., 2019a;Bhandari and Russo, 2019;Liu et al., 2019;Wang et al., 2019;Agarwal et al., 2019) to IL. This line of research analyzes the global optimality and convergence of policy gradient for solving RL, which is a minimization or maximization problem. In contrast, we analyze GAIL, which is a minimax optimization problem.
Our work is also related to the analysis of apprenticeship learning (Syed et al., 2008) and GAIL (Cai et al., 2019a;Chen et al., 2020). Syed et al. (2008) analyze the convergence and generalization of apprenticeship learning. They assume the state space to be finite, and thus, do not require function approximation for the policy and the reward function. In contrast, we assume the state space to be infinite and employ function approximation based on neural networks. Cai et al. (2019a) study the global optimality and convergence of GAIL in the setting of linear-quadratic regulators.
In contrast, our analysis handles general MDPs without restrictive assumptions on the transition kernel and the reward function. Chen et al. (2020) study the convergence and generalization of GAIL in the setting of general MDPs. However, they only establish the convergence to a stationary point. In contrast, we establish the global optimality of GAIL.
Notations. Let [n] = {1, . . . , n} for n ∈ N + and [m : n] = {m, m + 1, . . . , n} for m < n. Also, let N (µ, Σ) be the Gaussian distribution with mean µ and covariance Σ. We denote by P(X ) the set of all probability measures over the space X . For a function f : X → R, a constant p ≥ 1, and a probability measure µ ∈ P(X ), we denote by f p,µ = ( X |f (x)| p dµ(x)) 1/p the L p (µ) norm of the function f . For any two functions f, g : X → R, we denote by f, g X = X f (x) · g(x)dx the inner product on the space X .

Background
In this section, we introduce reinforcement learning (RL) and generative adversarial imitation learning (GAIL).

Reinforcement Learning
We consider a Markov decision process (MDP) (S, A, r, P, ρ, γ). Here S ⊆ R d 1 is the state space, A ⊆ R d 2 is the action space, which is assumed to be finite throughout this paper, r : S × A → R is the reward function, P : S × A → P(S) is the transition kernel, ρ ∈ P(S) is the initial state distribution, and γ ∈ (0, 1) is the discount factor. Without loss of generality, we assume that S × A is compact and that (s, a) 2 ≤ 1 for any (s, a) ∈ S ×A ⊆ R d , where d = d 1 +d 2 . An agent following a policy π : S → P(A) interacts with the environment in the following manner. At the state s t ∈ S, the agent takes the action a t ∈ A with probability π(a t | s t ) and receives the reward r t = r(s t , a t ).
The environment then transits into the next state s t+1 with probability P (s t+1 | s t , a t ). Given a policy π and a reward function r(s, a), we define the state-action value function Q π r : S × A → R as follows, Here the expectation E π is taken with respect to a t ∼ π(· | s t ) and s t+1 ∼ P (· | s t , a t ). Correspondingly, we define the state value function V π r : S → R and the advantage function A π r : S × A → R as follows, The goal of RL is to maximize the following expected cumulative reward, The policy π induces a state visitation measure d π ∈ P(S) and a state-action visitation measure ν π ∈ P(S × A), which take the forms of It then holds that J(π; r) = E (s,a)∼νπ [r(s, a)]. Meanwhile, we assume that the policy π induces a state stationary distribution ̺ π over S, which satisfies that We denote by ρ π (s, a) = ̺(s) · π(a | s) the state-action stationary distribution over S × A.

Generative Adversarial Imitation Learning
The goal of imitation learning (IL) is to learn a policy that outperforms the expert policy π E based We denote by ν E = ν π E and d E = d π E the state-action and state visitation measures induced by the expert policy, respectively, and assume that the expert trajectory {(s i , a i )} i∈[T E ] is drawn from ν E . To this end, we parameterize the policy and the reward function by π θ for θ ∈ X Π and r β (s, a) for β ∈ X R , respectively, and solve the following minimax optimization problem known as GAIL (Ho and Ermon, 2016), (2.4) Here J(π; r) is the expected cumulative reward defined in (2.2), ψ : X R → R + is the regularizer, and λ ≥ 0 is the regularization parameter. Given a reward function class R, we define the R-distance between two policies π 1 and π 2 as follows, When R is the class of 1-Lipschitz functions, D R (π 1 , π 2 ) is the Wasserstein-1 metric between the state-action visitation measures induced by π 1 and π 2 . However, D R (π 1 , π 2 ) is not a metric in general. When D R (π 1 , π 2 ) ≤ 0, the policy π 2 outperforms the policy π 1 for any reward function r ∈ R. Such a notion of R-distance is used in Chen et al. (2020). We denote by R β = {r β (s, a) | β ∈ X R } the reward function class parameterized with β. Hence, the optimization problem in (2.4) minimizes the R β -distance D R β (π E , π θ ) (up to the regularizer λ · ψ(β)), which searches for a policyπ that (approximately) outperforms the expert policy given any reward function r β ∈ R β .

Algorithm
In this section, we introduce an algorithm with alternating updates for GAIL with neural networks, which employs natural policy gradient (NPG) to update the policy π θ and gradient ascent to update the reward function r β (s, a).

Parameterization with Neural Networks
We define a two-layer neural network with rectified linear units (ReLU) as follows, (3.2) It then holds that u W,b (s, a) = W ⊤ φ W,b (s, a). Note that the feature vector φ W,b (s, a) depends on the parameters W and b. We consider the following random initialization, Throughout the training process, we keep the parameter b fixed while updating W . For notational simplicity, we write u W,b (s, a) as u W (s, a) and φ W,b (s, a) as φ W (s, a) in the sequel. We denote by E init the expectation taken with respect to the random initialization in (3.3). For an absolute constant B > 0, we define the parameter domain as which is the ball centered at W 0 with the domain radius B.
In the sequel, we consider the following energy-based policy, where τ ≥ 0 is the inverse temperature parameter and u θ (s, a) is the energy function parameterized by the neural network defined in (3.1) with W = θ. In the sequel, we call θ the policy parameter.
Meanwhile, we parameterize the reward function r β (s, a) as follows, where u β (s, a) is the neural network defined in (3.1) with W = β and γ is the discount factor.
Here we use the scaling parameter (1 − γ) −1 to ensure that for any policy π the state-action value function Q π r β (s, a) defined in (2.1) is well approximated by the neural network defined in (3.1). In the sequel, we call β the reward parameter and define the reward function class as where S B β is the parameter domain of β defined in (3.4) with domain radius B β . To facilitate algorithm design, we establish the following proposition, which calculates the explicit expressions of the gradients ∇L(θ, β) and the Fisher information I(θ). Recall that the Fisher information is defined as (3.7) Proposition 3.1 (Gradients and Fisher Information). We call ι θ (s, a) = τ −1 · ∇ θ log π θ (a | s) the temperature-adjusted score function. It holds that It then holds that where Q π θ r β (s, a) is the state-action value function defined in (2.1) with π = π θ and r = r β , ν π θ is the state-action visitation measure defined in (2.3) with π = π θ , and I(θ) is the Fisher information defined in (3.7).
Proof. See §C.1 for a detailed proof.
Note that the expression of the policy gradient ∇ θ L(θ, β) in (3.10) of Proposition 3.1 involves the state-action value function Q π θ r β (s, a). To this end, we estimate the state-action value function Q π r (s, a) by Q ω (s, a), which is parameterized as follows, (3.12) Here u ω (s, a) is the neural network defined in (3.1) with W = ω. In the sequel, we call ω the value parameter.

GAIL with Alternating Updates
We employ an actor-critic scheme with alternating updates of the policy and the reward function, which is presented in Algorithm 1. Specifically, we update the policy parameter θ via natural policy gradient and update the reward parameter β via gradient ascent in the actor step, while we estimate the state-action value function Q π r (s, a) via neural temporal difference (TD) (Cai et al., 2019c) in the critic step.

Actor
Step. In the k-th actor step, we update the policy parameter θ and the reward parameter β as follows, where Here η > 0 is the stepsize, S B θ and S B β are the parameter domains of θ and β defined in (3.4) with domain radii B θ and B β , respectively, Proj S B β : R md → S B β is the projection operator, τ k is the inverse temperature parameter of π θ k , and I(θ k ), ∇ θ L(θ k , β k ), ∇ β L(θ k , β k ) are the estimators of , respectively, which are defined in the sequel. In (3.13), we update the policy parameter θ k along the direction δ k , which approximates the natural policy gradient I(θ) −1 · ∇ θ L(θ, β), and in (3.15) we update the inverse temperature parameter τ k to ensure that θ k+1 ∈ S B θ . Meanwhile, in (3.14), we update the reward parameter β via (projected) gradient ascent. Following from (3.9) and (3.10) of Proposition 3.1, we construct the estimators of I(θ k ) and ∇ θ L(θ k , β k ) as follows, is sampled from the state-action visitation measure ν π θ k given θ k with the batch size N , and Q ω k (s, a) is the estimator of Q π θ k r β k (s, a) computed in the critic step. Meanwhile, following from (3.11) of Proposition 3.1, we construct the estimator of ∇ β L(θ k , β k ) as follows, is the expert trajectory. For notational simplicity, we write π k = π θ k , r k (s, a) = r β k (s, a), d k = d π k and ν k = ν π k for the k-th step hereafter, where π θ is the policy, r β (s, a) is the reward function, and d π , ν π are the visitation measures defined in (2.3).

Critic
Step. Note that the estimator ∇ θ L(θ, β) in (3.17) involves the estimator Q ω k (s, a) of To this end, we parameterize Q ω (s, a) as in (3.12) and adapt neural TD (Cai et al., 2019c), which solves the following minimization problem, Here S Bω is the parameter domain with domain radius B ω , ρ k is the state-action stationary distribution induced by π k , and T π k r k is the Bellman operator. Note that the Bellman operator T π r is defined as follows, where the expectation is taken with respect to s ′ ∼ P (· | s, a) and a ′ ∼ π(· | s ′ ). In neural TD, we iteratively update the value parameter ω via where δ(j) is the Bellman residual, α > 0 is the stepsize, (s, a) is sampled from the state-action stationary distribution ρ k , and s ′ ∼ P (· | s, a), a ′ ∼ π k (· | s ′ ) are the subsequent state and action.
We defer the detailed discussion of neural TD to §B.
To approximately obtain the compatible function approximation (Sutton et al., 2000;Wang et al., 2019), we share the random initialization among the policy π θ , the reward function r β (s, a), and the state-action value function Q ω (s, a). In other words, we set , 1999). Specifically, the mixed policyπ of π 0 , . . . , π T −1 is executed by randomly selecting a policy π k for k ∈ [0 : T − 1] with equal probability before time t = 0 and exclusively following π k thereafter. It then holds for any reward function r(s, a) that Update value parameter ω k via Algorithm 2 with π k , r k , W 0 , b, T TD , and α as the input.

Main Results
In this section, we first present the assumptions for our analysis. Then, we establish the global optimality and convergence of Algorithm 1.

Assumptions
We impose the following assumptions on the stationary distributions ̺ k ∈ P(S), ρ k ∈ P(S × A) and the visitation measures d k ∈ P(S), ν k ∈ P(S × A).
Assumption 4.1. We assume that the following properties hold.
(a) Let µ be either ρ k or ν k . We assume for an absolute constant c > 0 that Here dd E /dd k , dν E /dν k , dd E /d̺ k , and dν E /dρ k are the Radon-Nikodym derivatives. For notational simplicity, we write with W = W 0 , and W 0 is the random initialization in (3.3). We impose the following assumptions on the neural network u 0 (s, a) and the transition kernel P .
Assumption 4.2. We assume that the following properties hold.
(a) LetŪ = sup (s,a)∈S×A |u 0 (s, a)|. We assume for absolute constants M 0 > 0 and v > 0 that We assume that the transition kernel P belongs to the following class, Here B P > 0 is an absolute constant, q is the probability density function of N (0, I d /d), and ϑ(s, a; w) is defined as ϑ(s, a; w) = 1{w ⊤ (s, a) > 0} · (s, a). and that S × A is compact, it is reasonable to assume that sup (s,a)∈S×A |u 0 (s, a)| is sub-Gaussian, which further implies the existence of the absolute constants M 0 and v in (4.1) of Assumption 4.2 (a). Moreover, if we assume that m is even and initialize the parameters W 0 , b as follows, we have that u 0 (s, a) = 0 for any (s, a) ∈ S × A, which allows us to set M 0 = 0 and v = +∞ in Assumption 4.2 (a). Also, it holds that 0 = u 0 (s, a) ∈ R β , which implies that D R β (π 1 , π 2 ) ≥ 0 for any π 1 and π 2 . The proof of our results with the random initialization in (4.2) is identical.
Assumption 4.3. We assume that the following properties hold.
(a) We assume for an absolute constant σ > 0 that where τ k is the inverse temperature parameter in (3.5), N ∈ N + is the batch size, and S B θ is the parameter domain of θ defined in (3.4) with the domain radius B θ . Here the expectation E k is taken with respect to the k-th batch, which is drawn from ν k given θ k .
(b) We assume that the regularizer ψ(β) in (2.4) is convex and L ψ -Lipschitz continuous over the uniformly upper bounded variances across i ∈ [m] and k, and the Markov chain that generates mixes sufficiently fast (Zhang et al., 2019a). Similar assumptions are also used in the analysis of policy optimization (Xu et al., 2019a,b). Also, Assumption 4.3 (b) holds for most commonly used regularizers (Ho and Ermon, 2016).

Global Optimality and Convergence
In this section, we establish the global optimality and convergence of Algorithm 1. The following proposition adapted from Cai et al. (2019c) characterizes the global optimality and convergence of neural TD, which is presented in Algorithm 2.
Proposition 4.4 (Global Optimality and Convergence of Neural TD). In Algorithm 2, we set ) for an absolute constant c > 0. Let π k , r k be the input and ω k be the output of Algorithm 2. Under Assumptions

and 4.2, it holds for an absolute constant
Here the expectation E init is taken with respect to the random initialization in (3.3).
Proof. See §B.1 for a detailed proof.
The term B 2 ω · exp(−C v · B 2 ω ) in (4.6) of Proposition 4.4 characterizes the hardness of estimating the state-action value function Q π k r k (s, a) by the neural network defined in (3.1), which arises because Q π k r k (s, a) ∞ is not uniformly upper bounded across k. Note that if we employ the random initialization in (4.2), we have that C v = +∞. And consequently, such a term vanishes. We are now ready to establish the global optimality and convergence of Algorithm 1. . (4.7) the expectation is taken with respect to the random initialization in (3.3) and the T batches, and the error term ε k satisfies that where C h is defined in Assumption 4.1, L ψ and σ are defined in Assumption 4.3, and ) is the error induced by neural TD (Algorithm 2).
Proof. See §5 for a detailed proof.
The optimality gap in (4.7) of Theorem 4.5 is measured by the expected between the expert policy π E and the learned policyπ. of Assumption 4.3, which vanishes as the batch size N in Algorithm 1 goes to infinity. Term (iii.b) is the error of estimating Q π r (s, a) by Q ω (s, a) using neural TD (Algorithm 2). As shown in Proposition 4.4, the estimation error ǫ Q,k vanishes as m and B ω go to infinity. Term (iii.c) corresponds to the linearization error of the neural network defined in (3.1), which is characterized in Lemma A.2. Following from Theorem 4.5, it holds for B ω = Ω((C −1 v · log T ) 1/2 ), m = Ω(B 10 · T 6 ), and N = Ω(B 2 · T · σ 2 ) that E D R β (π E ,π) = O(T −1/2 + λ), which implies the 1/ √ T rate of convergence of Algorithm 1 (up to the bias induced by the regularizer).

Proof of Main Results
In this section, we present the proof of Theorem 4.5, which establishes the global optimality and convergence of Algorithm 1. For notational simplicity, we write π s (a) = π(a | s) for any policy π, state s ∈ S, and action a ∈ A. For any policies π 1 , π 2 and distribution µ over S, we denote the expected Kullback-Leibler (KL) divergence by KL µ , which is defined as KL µ (π 1 π 2 ) = E s∼µ [KL(π s 1 π s 2 )]. For any visitation measures d π ∈ P(S) and ν π ∈ P(S × A), we denote by E dπ and E νπ the expectations taken with respect to s ∼ d π and (s, a) ∼ ν π , respectively. Following from the property of the mixed policyπ in (3.21), we have that (5.1) We now upper bound the optimality gap in (5.1) by upper bounding the following difference of expected cumulative rewards, where β ′ ∈ S B β is chosen arbitrarily and L(θ, β) is the objective function defined in (2.4). Following from Assumption 4.3 and the fact that β k , β ′ ∈ S B β , we have that which upper bounds term (iii) of (5.2). It remains to upper bound terms (i) and (ii) of (5.2), which hinges on the one-point convexity of J(π; r) with respect to π and the (approximate) convexity of L(θ, β) with respect to β.
Upper bound of term (i) in (5.2). In what follows, we upper bound term (i) of (5.2). We first introduce the following cost difference lemma (Kakade and Langford, 2002), which corresponds to the one-point convexity of J(π; r) with respect to π. Recall that d E ∈ P(S) is the state visitation measure induced by the expert policy π E .
Lemma 5.1 (Cost Difference Lemma, Lemma 6.1 in Kakade and Langford (2002)). For any policy π and reward function r(s, a), it holds that where γ is the discount factor.
Furthermore, we establish the following lemma, which upper bounds the right-hand side of (5.4) in Lemma 5.1.
Lemma 5.2. Under Assumptions 4.1-4.3, we have that Here M 0 is defined in Assumption 4.2, σ is defined in Assumption 4.3, N is the batch size in (3.16)-(3.18), and ǫ Q, )) for an absolute constant C v > 0, which depends on the absolute constant v in Assumption 4.2.
Proof. See §C.2 for a detailed proof.
Combining Lemmas 5.1 and 5.2, we have that In what follows, we upper bound term (ii) of (5.2). We first establish the following lemma, which characterizes the (approximate) convexity of L(θ, β) with respect to β.
Proof. See §C.3 for a detailed proof.
The term O(B 3/2 β ·m −1/4 ) in Lemma 5.3 arises from the linearization error of the neural network, which is characterized in Lemma A.2. Based on Lemma 5.3, we establish the following lemma, which upper bounds term (ii) of (5.2).
Lemma 5.4. Under Assumptions 4.1 and 4.3, we have that Proof. See §C.4 for a detailed proof.
By Lemma 5.4, we have that which upper bounds term (ii) of (5.2). Here ∆ k is upper bounded in (5.7) of Lemma 5.4. Plugging (5.3), (5.6), and (5.8) into (5.2), we obtain that are upper bounded in (5.5) and (5.7) of Lemmas 5.2 and 5.4, respectively. Note that the upper bound of ∆ k does not depend on θ and β. Upon telescoping (5.9) with respect to k, we obtain that Following from the fact that τ 0 = 0 and the parameterization of π θ in (3.5), it holds that π s 0 is the uniform distribution over A for any s ∈ S. Thus, we have KL d E (π E π 0 ) ≤ log |A|. Meanwhile, following from the fact that β ′ ∈ S B β , it holds that β ′ − β 0 2 ≤ B β . Finally, by setting η = T −1/2 , τ k = k · η, andB = max{B θ , B β } in (5.10), we have that Here ε k is upper bounded as follows, )) for an absolute constant C v > 0. Thus, we complete the proof of Theorem 4.5.

A Neural Networks
In what follows, we present the properties of the neural network defined in (3.1). First, we define the following function class.
Definition A.1 (Function Class). For B > 0 and m ∈ N + , we define where φ 0 (s, a) is the feature vector defined in (3.2) with W = W 0 .
As shown in Rahimi and Recht (2008) 2019c)). Under Assumption 4.1, it holds for any W, W 1 , W 2 ∈ S B that, where φ W (s, a) is the feature vector defined in (3.2) and µ ∈ P(S ×A) is a distribution that satisfies Assumption 4.1.
Proof. See §A.1 for a detailed proof.
Following from Lemma A.2, the function class F B,m defined in Definition A.1 is a first-order approximation of the class of the neural networks defined in (3.1). Meanwhile, we establish the following lemma to characterize the sub-Gaussian property of the neural network defined in (3.1).
Proof. See §A.2 for a detailed proof.

A.1 Proof of Lemma A.2
Proof. We consider any W, W ′ ∈ S B . By the definition of φ W (s, a) in (3.2) and the triangle inequality, we have that We now upper bound the right-hand side of (A.1). For the term |[W ] ⊤ l (s, a)| in (A.1), we have that where the first inequality follows from the triangle inequality and the second inequality follows from the Cauchy-Schwartz inequality and the fact that (s, a) 2 ≤ 1. To upper bound the Thus, we have that

Plugging (A.2) and (A.3) into (A.1), we have that
By the fact that W, W ′ ∈ S B , we obtain that Taking the expectation with respect to the random initialization in (3.3) and using the Cauchy-Schwartz inequality, we have that where the second inequality follows from the fact that W ′ − W 0 2 ≤ B, the third inequality follows from Jensen's inequality, and the last inequality follows from Assumption 4.1 and Lemma A.2. Thus, for any W, W 1 , W 2 ∈ S B , we have that Moreover, following from the Cauchy-Schwartz inequality, we have that · 1,µ ≤ · 2,µ . Thus, by Jensen's inequality, we have that which completes the proof of Lemma A.2.

A.2 Proof of Lemma A.3
In what follows, we present the proof of Lemma A.3.
Proof. Recall that we write u W (s, a) = W ⊤ φ W (s, a) and u 0 (s, a) = u W 0 (s, a). Then, we have where the last inequality follows from the Cauchy-Schwartz inequality. It suffices to upper bound the three terms on the right-hand side of (A.4). Note that we have W, W ′ ∈ S B and φ W ′ (s, a) 2 ≤ 1.
We have that It remains to upper bound the term |u W ′ (s, a) − u 0 (s, a)| in (A.4). Note that u W (s, a) is almost everywhere differentiable with respect to W . Also, it holds that ∇ W u W (s, a) = φ W (s, a). Thus, following from the mean-value theorem and the Cauchy-Schwartz inequality, we have that where the second inequality follows from the fact that φ W (s, a) 2 ≤ 1 and W ′ ∈ S B . Plugging Following from Assumption 4.2, we have that sup (s,a)∈S×A |W ⊤ φ W ′ (s, a)| is sub-Gaussian. Furthermore, it holds that for t > 2M 0 + 6B. Thus, we complete the proof of Lemma A.3.

B Neural Temporal Difference
In this section, we introduce neural TD (Cai et al., 2019c), which computes ω k in Algorithm 1.
Specifically, neural TD solves the optimization problem in (3.19) via the update in (3.20), which is presented in Algorithm 2.
Algorithm 2 Neural TD Require: Policy π, reward function r, initialization W 0 , b, number of iterations T TD of neural TD, and stepsize α of neural TD.
where F Bω,m is defined in Definition A.1.
Recall that we denote by φ 0 (s, a) the feature vector corresponding to the random initialization in (3.3). We establish the following lemma to upper bound the bias E init [ Proj F Bω ,m Q π r (s, a) − Q π r (s, a) 2 2,ρπ ] in (B.1) of Proposition B.1 when the reward function r(s, a) belongs to the reward function class R β .
Combining Proposition B.1 and Lemma B.2, for we have for any π that Finally, by setting T TD = Ω(m), we have that Thus, since r β = (1 − γ) −1 · u β (s, a), we have that where the second equality follows from (B.2) and the last equality follows from Fubini's theorem.
To prove Lemma B.2, we first approximate Q π r β (s, a) bȳ Q(s, a) = u β (s, a) + ϑ(s, a; w) ⊤ᾱ (w)dq(w), (B.4) whereᾱ(w) = α(w) · 1{ α(w) 2 ≤ K} for an absolute constant K > 0 specified later. Then, it holds for any (s, a) ∈ S × A that where the second inequality follows from the Cauchy-Schwartz inequality and the last inequality follows from the fact that ϑ(s, a; w) 2 ≤ 1. Thus, we have that We now upper bound the right-hand side of (B.5). To this end, we show that sup w α(w) 2 is sub-Gaussian in the sequel. By the definition of α(w) in (B.3), we have that where the second inequality follows from Assumption 4.2 and the third inequality follows from the fact that V π r β (s) = E (s ′ ,a ′ )∼νπ(s) [r β (s ′ , a ′ )]. Here we denote by ν π (s) the state-action visitation measure starting from the state s and following the policy π. Following from Lemma A.3, we have that sup w α(w) 2 is sub-Gaussian. By Lemma A.3 and (B.6), it holds for t > (1 − γ) −1 · γ · B P · Following from (B.5) and (B.7), we have that We now construct Q(s, a) ∈ F K,m , which approximatesQ(s, a) defined in (B.4). We define f (s, a) = ϑ(s, a; w) ⊤ᾱ (w)dq(ω).
Then, we have thatQ(s, a) = u β (s, a) + f (s, a). Note that f (s, a) belongs to the following function class, We now show that f (s, a) is well approximated by the following function class, where φ 0 (s, a) is the feature vector corresponding to the random initialization. We obtain the following lemma from Rahimi and Recht (2009), which characterizes the approximation error of F K,∞ by F K,m . (2009)). For any f (s, a) ∈ F K,∞ , it holds with probability at least 1 − δ that
Lemma B.3 implies that there exists f (s, a) ∈ F K,m such that By the fact that f (s, a) ∈ F K,m and the definition of F K,m in Definition A.1, we have that f (s, a) ∈ We then have that Q(s, a) ∈ F B β +K,m and that where the last inequality follows from Assumption 4.1, Lemma A.2, and (B.9).

C Proofs of Auxiliary Results
In what follows, we present the proofs of the lemmas in §3-5.

C.1 Proof of Proposition 3.1
Proof. By the definition of the neural network in (3.1), we have for any (s, a) ∈ S × A that ∇ W u W (s, a) = φ W (s, a) almost everywhere. We first calculate ∇ θ L(θ, β). Following from the policy gradient theorem (Sutton and Barto, 2018) and the definition of L(θ, β) in (2.4), we have that Following from the parameterization of π θ in (3.5) and the definition of ι θ (s, a) in (3.8) of Proposition 3.1, we have that It remains to calculate I(θ) and ∇ β L(θ, β). By (C.2) and the definition of I(θ) in (3.7), it holds that By the definition of the objective function L(θ, β) in (2.4), it holds that Thus, we complete the proof of Proposition 3.1.

C.2 Proof of Lemma 5.2
Proof. The proof of Lemma 5.2 is similar to that of Lemmas 5.4 and 5.5 in Wang et al. (2019). By direct calculation, we have that The following lemmas upper bound ∆ where C h is defined in Assumption 4.1 and σ is defined in Assumption 4.3.
Proof. See §D.1 for a detailed proof.
Lemma C.2 (Upper Bound of Term (i.b) in (C.3)). Under Assumption 4.1, we have that where ǫ Q,k takes the form of Proof. See §D.2 for a detailed proof.
where M 0 is defined in Assumption 4.2.
Proof. See §D.3 for a detailed proof.
Finally, by Lemmas C.1-C.3, under Assumptions 4.2 and 4.3, we obtain from (C.3) that Here M 0 is defined in Assumption 4.2, τ k+1 is the inverse temperature parameter of π k+1 defined in (3.5), σ is defined in Assumption 4.3, and ǫ Q,k is defined in (C.4) of Lemma C.2. Following from Proposition 4.4, we have that Thus, we complete the proof of Lemma 5.2.

C.3 Proof of Lemma 5.3
Proof. We consider a fixed β ′ ∈ S B β . For notational simplicity, we write r ′ = r β ′ (s, a), r k = r k (s, a) and φ β = φ β (s, a). By the parameterization of r β (s, a) in (3.6), we have that where the last inequality follows from (3.10) of Proposition 3.1. Then, we have that where the last inequality follows from Assumption 4.1, Lemma A.2, and the fact that β ′ , β k ∈ S B β .
Thus, we complete the proof of Lemma 5.3.

C.4 Proof of Lemma 5.4
Proof. By the update of β k in (3.14), it holds for any β ′ ∈ S B β that Combining (C.5) and (C.6), we have that We now upper bound terms (ii.a), (ii.b), and (ii.c) on the right-hand side of (C.7). Following from Assumption 4.1 and Lemma A.2, we have that which upper bounds term (ii.c) of (C.7). For term (ii.b) of (C.7), we have that Here the first inequality follows from the Cauchy-Schwartz inequality, the second inequality follows from the fact that β k , β ′ ∈ S B β , and the last inequality follows from Assumption 4.3. To upper bound term (ii.a) in (C.7), we have that where the first inequality follows from the Cauchy-Schwartz inequality and the second inequality follows from the update of β in (3.14). Furthermore, we have where the first inequality follows from Jensen's inequality and the second inequality follows from the fact that φ W (s, a) 2 ≤ 1 and the Lipschitz continuity of ψ(β) in Assumption 4.3. By plugging (C.11) into (C.10), we have that where the last inequality follows from Assumption 4.3. Finally, by plugging (C.8), (C.9), and (C.12) into (C.7), we have that k | = η · (2 + λ · L ψ ) 2 + σ 2 · N −1 + 2B β · σ · N −1/2 + O(B 3/2 β · m −1/4 ).
Thus, we complete the proof of Lemma 5.4.

D Proofs of Supporting Lemmas
In what follows, we present the proofs of the lemmas in §C.
In what follows, we upper bound terms (i) and (ii) on the right-hand side of (D.3).