On the Global Convergence of Natural Actor-Critic with Two-layer Neural Network Parametrization

Motivation. The use of neural networks in Actor-critic (AC) algorithms is widespread in various machine learning applications, such as games (Vinyals et al., 2017; Bonjour et al., 2022), robotics (Morgan et al., 2021), autonomous driving (Kiran et al., 2022), ride-sharing (Zheng et al., 2022), networking (Geng et al., 2020), and recommender systems (Li et al., 2020). AC algorithms sequentially update the estimate of the actor (policy) and the critic (value function) based on the data collected at each iteration, as described in (Konda and Tsitsiklis, 1999). An empirical and theoretical improvement of the AC, known as Natural Actor-Critic (NAC), was proposed in (Peters and Schaal, 2008). NAC replaced the stochastic gradient step of the actor with the natural gradient descent step described in (Kakade, 2001b) based on the theory of (Rattray et al., 1998). Despite its widespread use in state-of-the-art reinforcement learning (RL) implementations, there are no finite-time sample complexity results available in the literature for a setup where neural networks (NNs) are used to represent the critic. Such results provide estimates of the required data to achieve a specified accuracy in a loss function, typically in terms of the difference between the average reward obtained in following the policy obtained through the algorithm being analyzed and the optimal policy.

Challenges. The primary challenge in obtaining sample complexity for a NAC is that the critic loss function with NN parametrization becomes non-convex, making finite-time sample guarantees impossible. Although recent work in (Fu et al., 2020) obtains upper bounds on the error in estimating the optimal action-value function in NN critic settings, these upper bounds are asymptotic and depend on the NN parameters. Similarly, other asymptotic sample complexity results, such as those in (Kakade, 2001b; Konda and Tsitsiklis, 1999; Bhatnagar, 2010), require i.i.d. sampling of the state action pairs at each iteration of the algorithm. While sample complexity results for linear MDP structures have been derived in (Wu et al., 2020; Xu et al., 2020b) (see Related work), the structure allows for a closed-form critic update. The relaxation of the linear MDP assumption requires a critic parametrization to learn the Q-function, limiting the ability to find globally optimal policies due to non-global optimality of critic parameters.

Approach. As the optimal critic network update is not available for general MDPs, a parametrization of the critic network is used. To address the challenge of optimizing the critic parameters, we assume a two-layer neural network structure for the critic parametrization. Recent studies have shown that the parameter optimization for two-layer ReLU neural networks can be transformed into an equivalent convex program, which is computationally tractable and solvable exactly (Pilanci and Ergen, 2020). Convex formulations for convolutions and deeper models have also been investigated (Sahiner et al., 2020a; Sahiner et al., 2020b). In this work, we use these approaches to estimate the parameterized Q-function in the critic. In the policy gradient analysis, the error caused by the lack of knowledge of the exact gradient term is assumed to be upper bounded by a constant (see related work for details), where the sample complexity of this term is not accounted due to no explicit critic estimation. This paper considers the sample complexity of the critic function to derive the overall sample complexity of the natural actor-critic algorithm. Resolving the optimality of the critic parameter update, this paper considers the following question:

Contributions. In this work, we provide an affirmation to the above-mentioned question. We summarize our contributions as follows.

• •

  We propose a novel actor-critic algorithm, employing a general parametrization for the actor and two-layer ReLU neural network parametrization for the critic.

• •

  Building upon the insights presented in (Agarwal et al., 2020b), we leverage the inherent smoothness property of the actor parametrization to derive an upper bound on the estimation error of the optimal value function. This upper bound is expressed in terms of the error incurred in attaining the compatible function approximation term, as elucidated in (Sutton et al., 1999b). Our analysis dissects this error into distinct components, each of which is individually bounded from above.

• •

  To address the estimation error arising from the critic function, we leverage the interesting findings of (Wang et al., 2021; Pilanci and Ergen, 2020), converting the problem of finding critic function parameters into a convex optimization problem. This allows us to establish a finite-time sample complexity bound, a significant improvement over previous results obtained for NAC algorithms. Notably, our approach eliminates the reliance on a linear function approximation for the critic, presenting the first finite-time sample complexity result for NAC without such an approximation.

Natural Actor-Critic (NAC). Actor critic methods, first conceptualised in (Sutton, 1988), aim to combine the benefits of the policy gradient methods and $Q$-learning based methods. The policy gradient step in these methods is replaced by a Natural Policy Gradient proposed in (Kakade, 2001b) to obtain the so-called Natural Actor Critic in (Peters et al., 2005). Sample complexity results for Actor Critic were first obtained for MDP with finite states and actions in (Williams and Baird, 1990), and more recently in (Lan, 2023; Zhang et al., 2020b). Finite time convergence for natural actor critic using a linear MDP assumption has been obtained in (Chen and Zhao, 2022; Khodadadian et al., 2021; Xu et al., 2020c) with the best known sample complexity of $\tilde{\mathcal{O}}\left(\frac{1}{{\epsilon}^{2}(1-\gamma)^{4}}\right)$ (Xu et al., 2020b). Finite time sample complexity results are however, not available for Natural Actor Critic setups for general MDP where neural networks are used to represent the critic. (Fu et al., 2020) obtained asymptotic upper bounds on the estimation error for a natural actor critic algorithm where a neural network is used to represent the critic. The key related works here are summarized in Table 1.

Sample Complexity of PG Algorithms. Policy Gradient methods (PG) were first introduced in (Sutton et al., 1999b), with Natural Policy Gradients in (Kakade, 2001b). Sample complexity were obtained in tabular setups in (Even-Dar et al., 2009) and then improved in (Agarwal et al., 2020b). Using linear dynamics assumption convergence guarantees were obtained in (Fazel et al., 2018), (Dean et al., 2020), (Bhandari and Russo, 2019) . (Zhang et al., 2020a) obtains sample complexity for convergence to second order stationary point policies. In order to improve the performance of Policy Gradient methods, techniques such as REINFORCE (Williams, 1992), function approximation (Sutton et al., 1999a), importance weights (Papini et al., 2018), adaptive batch sizes (Papini et al., 2017), Adaptive Trust Region Policy Optimization (Shani et al., 2020) have been used. We note that the Natural Policy Gradient approach has been shown to achieve a sample complexity of $\tilde{\mathcal{O}}\left(\frac{1}{\epsilon^{3}(1-\gamma)^{6}}\right)$ (Liu et al., 2020). However, they assume that the error incurred due to the lack of knowledge of the exact gradient term is assumed to be upper bounded by a constant (Assumption 4.4, (Liu et al., 2020)). Additionally, each estimate of the critic requires on average a sample of size $\left(\frac{1}{1-\gamma}\right)$ and the obtained estimate is only known to be an unbiased estimate of the critic, with no error bounds provided. Explicit dependence of the constant of (Assumption 4.4, (Liu et al., 2020)) in terms of number of samples is not considered in their sample complexity results. For the Natural Actor Critic analysis in this paper, we provide explicit dependence of the error in estimating the gradient update and incorporate the samples required for the estimation in the sample complexity analysis. We describe this difference in detail in Appendix A.

We consider a discounted Markov Decision Process (MDP) given by the tuple ${\mathcal{M}}:=(\mathcal{S},\mathcal{A},P,R,\gamma)$, where $\mathcal{S}$ is a bounded measurable state space, $\mathcal{A}$ is the finite set of actions. ${P}:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{P}(\mathcal{S})$ is the probability transition kernel¹¹ 1 For a measurable set $\mathcal{X}$, let $\mathcal{P}(\mathcal{X})$ denote the set of all probability measures over $\mathcal{X}$. , ${R}:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{P}([0,R_{\max}])$ is the reward kernel on the state action space with $R_{\max}$ being the absolute value of the maximum reward, and $0<\gamma<1$ is the discount factor. A policy $\pi:\mathcal{S}\rightarrow\mathcal{P}(\mathcal{A})$ maps a state to a probability distribution over the action space. Here, we denote by $\mathcal{P}(\mathcal{S}),\mathcal{P}(\mathcal{A}),\mathcal{P}([a,b])$, the set of all probability distributions over the state space, the action space, and a closed interval $[a,b]$, respectively. With the above notation, we define the action value function for a given policy $\pi$ as

where $r(s_{t},a_{t})\sim R(\cdot|s_{t},a_{t})$, $a_{t}\sim\pi(\cdot|s_{t})$ and $s_{t+1}\sim P(\cdot|s_{t},a_{t})$ for $t=\{0,\cdots,\infty\}$. For a discounted MDP, we define the optimal action value functions as $Q^{*}(s,a)=\sup_{\pi}Q^{\pi}(s,a),\hskip 14.22636pt\forall(s,a)\in\mathcal{S}\times\mathcal{A}$. A policy that achieves the optimal action value functions is known as the optimal policy and is denoted as $\pi^{*}$. It holds that $\pi^{*}$ is the greedy policy with respect to $Q^{*}$ (Bertsekas and Shreve, 2007). Hence finding $Q^{*}$ is sufficient to obtain the optimal policy. In a similar manner, we can define the value function as $V^{\pi}(s)=\mathbb{E}\left[\sum_{t=0}^{\infty}\gamma^{t}r^{\prime}(s_{t},a_{t})|s_{0}=s\right],$ and from the definition in (1), it holds that $V^{\pi}(s)=\mathbb{E}_{a\sim\pi}\left[Q^{\pi}(s,a)\right]$. Hence, we can define the optimal value function as $V^{*}(s)=\sup_{\pi}V^{\pi}(s),\hskip 14.22636pt\forall s\in\mathcal{S}$. We define $\rho_{\pi}(s)$ as the stationary state distribution induced by the policy $\pi$ starting at state $s$ and $\zeta_{\pi}(s,a)$ is the corresponding stationary state action distribution defined as $\zeta_{\pi}(s,a)=\rho_{\pi}(s).\pi(a|s)$. We overload notation to define $V^{\pi}(\rho)=\mathbb{E}_{s_{0}\sim\rho}[V^{\pi}(s_{0})]$, where $\rho$ is an initial state distribution. We can define the visitation distribution as $d^{\pi}(s_{0})=(1-\gamma)\sum_{t=0}^{\infty}{\gamma}^{t}Pr^{\pi}(s_{t}=s|s_{o})$. Here $Pr^{\pi}(s_{t}=s|s_{o})$ denotes the probability the state at time $t$ is $s$ given a starting state of $s_{o}$. Hence, we can write $d^{\pi}_{\rho}(s)=\mathbb{E}_{s_{o}\sim\rho}[d^{\pi}(s_{0})]$.

We now describe our Neural Actor-Critic (NAC) algorithm. In a natural policy gradient algorithm (Kakade, 2001a), the policy is parameterized as $\{\pi_{\lambda},\lambda\in\Lambda\}$. We have $K$ total iterations of the algorithm. At iteration $k$, the policy parameters are updated using a gradient descent step of the form

From the policy gradient theorem (Sutton et al., 1999b) we have

where $s\sim d^{\pi_{\lambda_{k}}}_{\rho},a\sim\pi_{\lambda_{k}}(.|s)$. From (Sutton et al., 1999b), the principle of compatible function approximation implies that we have

and $s\sim d^{\pi_{\lambda_{k}}}_{\rho},a\sim\pi_{\lambda_{k}}(.|s)$ Here $(A^{\pi_{\lambda_{k}}}(s,a)=Q^{\pi_{\lambda_{k}}}(s,a)-V^{\pi_{\lambda_{k}}}(s))$. For natural policy gradient algorithms such as in (Agarwal et al., 2020b) and (Liu et al., 2020) an estimate of $Q^{\pi_{\lambda_{k}}}$ (and from that an estimate of $A^{\pi_{\lambda_{k}}}(s,a)$) is obtained through a sampling procedure that requires on average $\left(\frac{1}{1-\gamma}\right)$ for each sample of $Q^{\pi_{\lambda_{k}}}$. For the natural actor critic setup we maintain a parameterised estimate of the $Q$-function is updated at each step and is used to approximate $Q^{\pi_{\lambda_{k}}}$.

(Xu et al., 2020b) and (Wu et al., 2020) assume that the $Q$ function can be represented as $Q^{\lambda_{k}}(s,a)=w^{t}(\phi(s,a))$ where $\phi(s,a)$ is a known feature vector and $w^{t}$ is a parameter that is updated in what is known as the critic step. The policy gradient step is known as the actor update step. In case a neural network is used to represent the $Q$ function, at each iteration $k$ of the algorithm, an estimate of its parameters are obtained by solving an optimization of the form

Due to the non-convex activation functions of a neural network the optimization in Equation (7) is non-convex and hence finite sample guarantees for actor critic algorithms with neural network representation of $Q$ functions is not possible. Thus in order to perform a finite time analysis, a linear MDP structure is assumed which is very restrictive and not practical for large state space problems. For a Natural Policy Gradient setup, a Monte Carlo sample is obtained, which makes the process require additional samples as causing high variance in critic estimation.

The detailed proof of Theorem 1 is given in Appendix F. The difference between our estimated value function denoted by $V^{\pi_{\lambda_{k}}}$ and the optimal value function denoted by $V^{*}$ (where $\pi_{\lambda_{k}}$ is the policy obtained at the step $k$ of algorithm 1) is first expressed as a function of the compatible function approximation error, which is then split into different components which are analysed separately. The proof is thus split into two stages. In the first stage, we demonstrate how the difference in value functions is upper bounded as a function of the errors incurred till the final step $K$. The second part is to upper bound the different error components.

Upper Bounding Error in Separate Error Components: We use the smoothness property assumed in Assumption 1 to obtain a bound on the expectation of the difference between our estimated value function and the optimal value function.

where

and $s\sim d^{\pi^{*}}_{\nu},a\sim\pi^{*}(.|s)$ and $W$ is a constant such that $||w^{k}||_{2}\leq W$ $\forall k$, where $k$ denotes the iteration of the outer for loop of Algorithm 1. We split the term in (26) into the errors incurred due to the actor and critic step as follows

Note that $I$ is the difference between the true $A^{\pi_{\lambda_{k}}}$ function corresponding to the policy $\pi_{\lambda_{k}}$ and $A_{k,J}$ is our estimate. This estimation is carried out in the first inner for loop of Algorithm 1. Thus $I$ is the error incurred in the critic step. $II$ is the error incurred in the estimation of the actor update. This is incurred in the stochastic gradient descent steps in the second inner for loop of Algorithm 1.

Also note that the expectation is with respect to the discounted state action distribution of the state action pairs induced by the optimal policy $\pi^{*}$. The state action samples that we obtain are obtained from the policy $\pi_{\lambda_{k}}$. Thus using assumption 6 we convert the expectation in Equation (28) to an expectation with respect to the stationary state action distribution induced by the policy $\pi_{\lambda_{k}}$.

Upper Bounding Error in Critic Step: For each iteration $k$ of the Algorithm 1. We show that minimzing $I$ is equivalent to solving the following problem

We will rely upon the analysis laid out in (Farahmand et al., 2010) and instead of the iteration of the value functions, we will apply a similar analysis to the action value function to obtain an upper bound for the error incurred in solving the problem in Equation (29) using the Fitted Q-Iteration. We recreate the result for the value function from Lemmas 2 of (Munos, 2003) for the action value function $Q$ to obtain

where $\epsilon_{k,j}=T^{\pi_{\lambda_{k}}}Q_{k,j-1}-Q_{k,j}$ is the Bellman error incurred at iteration $j$ where $T^{\pi_{\lambda_{k}}}Q_{k,j-1},P^{\pi_{\lambda_{k}}}$ are defined as in Equations (58), (60) respectively. $Q_{k,J}$ denotes the $Q$ estimate at the final iteration $J$ of the first inner for loop and iteration $k$ of the outer for loop of Algorithm 1

The first term on the right hand side is called as the algorithmic error, which depends on how good our approximation of the Bellman error is. The second term on the right hand side is called as the statistical error, which is the error incurred due to the random nature of the system and depends only on the parameters of the MDP as well as the number of iterations of the FQI algorithm. Intuitively, this error depends on how much data is collected at each iteration, how efficient our solution to the optimization step is to the true solution, and how well our function class can approximate $T^{\pi_{\lambda_{k}}}Q_{k,j-1}$. Building upon this intuition, we split $\epsilon_{k,j}$ into four different components as follows.

We use the Lemmas 14, 15, 16, and 17 to bound the error terms in Equation (31).

Upper Bounding Error in Actor Step: Note that we require the minimization of the term $\mathbb{E}_{s,a}(A_{k,J}-w^{k}{\nabla}log(\pi_{\lambda_{k}}(a|s)))$. Here the expectation is with respect to stationary state action distribution corresponding to $\pi_{\lambda_{k}}$. But we do not have samples of states action pairs from the stationary distribution with respect to the policy $\pi_{\lambda_{k}}$, we only have samples from the Markov chain induced by the policy $\pi_{\lambda_{k}}$. We thus refer to the theory in (Doan, 2022) and Assumption 5 to upper bound the error incurred.

In this paper, we study a Natural Actor Critic algorithm with a neural network used to represent both the actor and the critic and find the sample complexity guarantees for the algorithm. Using the conversion of the optimization of a 2 layer ReLU Neural Network to a convex problem for estimating the critic, we show that our approach achieves a sample complexity of $\tilde{\mathcal{O}}(\epsilon^{-4}(1-\gamma)^{-4})$. This demonstrates the first approach for achieving sample complexity beyond linear MDP assumptions for the critic.

Limitations: Relaxing the different stated assumptions is an interesting direction for the future. Further, the results assume 2-layer neural network parametrization for the critic. One can likely use the framework described in (Belilovsky et al., 2019) to extend the results to a multi layer setup.