Analysis of Q-learning with Adaptation and Momentum Restart for Gradient Descent

1 Introduction

Q-learning [Watkins and Dayan, 1992], as one of the most important model-free reinforcement learning (RL) algorithms, has received considerable attention in recent years [Bertsekas and Tsitsiklis, 1996; Even-Dar and Mansour, 2003; Lu et al., 2018]. The vanilla Q-learning algorithm runs a step of empirical Bellman operator update of the Q-function and a step of stochastic gradient descent (SGD) in an alternating fashion. The convergence guarantee of Q-learning has been studied for the tabular case in [Bertsekas, 1995; Even-Dar and Mansour, 2003], for the case with linear function approximation in [Bertsekas and Tsitsiklis, 1996; Zou et al., 2019b; Chen et al., 2019b], and also for neural network parameterization in [Xu and Gu, 2019].

However, all the existing theoretical analyses focus on Qlearning algorithms that take simple SGD iterations. Such theory is not applicable to practical Q-learning algorithms that implement the Adaptive Moment Estimation (Adam) type of updates. In this paper, we study the Q-learning algorithm with Adam-type updates in terms of its theoretical convergence and the performance in benchmark experiments. It is known in optimization that Adam does not always converge, and instead, and a slightly modified variant AMSGrad proposed in [Reddi et al., 2018] has been widely accepted as an alternative to justify the theoretical performance of Adam-type methods. This motivates the first question that we aim to address.

∗Equal contribution

†Corresponding author

• Q1: Can we provide the convergence guarantee for Qlearning under AMSGrad updates (i.e., Q-AMSGrad)?

In conventional optimization problems, restart has been incorporated into the gradient descent algorithm with momentum as a simple yet effective scheme to facilitate the acceleration performance [O'donoghue and Candes, 2015]. Hence, it is natural to incorporate such a restart technique into Q-learning and ask the following question about its performance.

• Q2: Does Q-AMSGrad with momentum restart (i.e., QAMSGradR) still converge?

As aforementioned, both Q-AMSGrad and Q-AMSGradR update alternatingly between one step of Bellman operator update of the Q-function and one step of adaptive momentum update. This is in contrast to the well-known deep Q-Network (DQN) learning [Mnih et al., 2015], which runs in a nestedloop manner with the outer loop consisting of an one-step update of the Q-function and the inner loop consisting of many iterations of supervised learning to fit a target Q-function. It is conventionally known that taking just one gradient step toward the target Q-function results in the instability of original Q-learning update. Thus the simple alternating update manner mainly aroused interest in theory instead of practice. We are interested in whether the adaptive momentum updates can improve the stability of Q-learning without using the supervised learning process for target network fitting. Therefore, the third question we want to address is to compare these two types of Q-learning algorithms from the experimental perspective.

• Q3: Do Q-AMSGrad and Q-AMSGradR perform competitively or even better than DQN in experiment?

This paper addresses the above theoretical and experimental questions with affirmative answers.

1.1 Main Contributions

Theoretically, we show that with linear function approximation (which is almost the only structure that the current tools for analysis of Q-learning can handle), both Q-AMSGrad and Q-AMSGradR converge to the global optimal solution under standard assumptions for Q-learning. To the best of our knowledge, this is the first non-asymptotic convergence guarantee on Q-learning that incorporates Adam-type update. Furthermore, a slight adaptation of our proof provides the convergence rate for the AMSGrad for conventional strongly convex optimization which has not been studied before and can be of independent interest.

Experimentally, we demonstrate that the practical versions of Q-AMSGrad and Q-AMSGradR (referred to as Q-Adam and Q-AdamR) exhibit appealing experimental performance. In a batch of 23 Atari 2600 games, our experiments show that both Q-Adam and Q-AdamR outperform DQN by 50% on average. Furthermore, Q-AdamR effectively reduces the performance variance and achieves a much more stable learning process. In our experiments for the linear quadratic regulator (LQR) problems, Q-AdamR converges even faster than the model-based value iteration (VI) solution. This is a rather surprising result given that the model-based VI has been treated as the performance upper bound for the Q-learning (including DQN) algorithms with target update [Lewis and Vrabie, 2009; Yang et al., 2019].

Detailed proofs and more experimental results will be available in the extended version of this paper on arXiv.org after the official publication of IJCAI Proceedings.

1.2 RelatedWork

We briefly review the related work as follows.

Theoretical analysis of Q-learning. Since proposed in [Watkins and Dayan, 1992], the convergence of Q-learning has been extensively studied, particularly for the case with linear function approximation such as [Bertsekas and Tsitsiklis, 1996; Zou et al., 2019b; Chen et al., 2019b; Du et al., 2019], to name a few, and more recently for the case with neural networks in [Xu and Gu, 2019], where the analysis exploits the approximate linear structure of neural networks in the overparamterized regime. All these existing analysis of Qlearning considers the vanilla SGD update, whereas our study is the first to analyze the more involved case with Adam-type updates.

Convergence analysis of Adam-type algorithms in conventional optimization. Adam was proposed in [Kingma and Ba, 2015] for speeding up the training of deep neural networks, and the regret bounds were characterized for Adam/AMSGrad in [Kingma and Ba, 2015; Reddi et al., 2018; Tran and others, 2019] for online convex optimization. Recently, convergence analysis of Adam/AMSGrad was provided for nonconvex optimization in [Zou et al., 2019a; Zhou et al., 2018; Chen et al., 2019a] and policy gradient [Xiong et al., 2020], in which such Adam-type algorithms were guaranteed to converge to a stationary point. To the best of our knowledge, our study provides the first convergence analysis of the Adam-type algorithms for Q-learning.

Empirical performance of Q-learning. DQN learning and its improved variants of dueling network structure [Wang et al., 2016], double Q-learning [Van Hasselt et al., 2016]

and variance exploration and sampling schemes [Schaul et al., 2015] have achieved significant success due to their superb performance in practice. In contrast to such nestedloop algorithms (which involves the fitting of a target Qfunction periodically), the Q-learning algorithms that strictly follow the alternating updates are much less explored in practice. [Mnih et al., 2016] proposed the asynchronous alternating Q-learning with competitive performance against DQN. However, the algorithm still relies on a slowly moving target network similar to DQN, and the multi-thread learning also complicates the computational setup. [Lu et al., 2018] studied the problem of value overestimation and proposed the non-delusional Q-learning algorithm that employs the so-called pre-conditioned Q-networks, which is also computationally complex. [Knight and Lerner, 2018] proposed a natural gradient propagation to improve the performance, where the gradient implementation is complex. Our experiments in this paper demonstrate that simple alternating Qlearning algorithms Q-AMSGrad and Q-AMSGradR without the complex designs as in [Mnih et al., 2016; Lu et al., 2018; Knight and Lerner, 2018] have competitive and sometimes better performance than DQN.

Notations We use ‖x‖ := ‖x‖2 to denote the ℓ2 norm of a vector x, and use ‖x‖∞ to denote the infinity norm. When x, y are both vectors, x/y, xy, x2 , √x are all calculated in the element-wise manner, which will be used in the update of Adam and AMSGrad. We denote [n] = 1, 2, . . . , n, and ⌊x⌋ ∈ Z as the integer such that ⌊x⌋ ≤ x < ⌊x⌋ + 1.

2 Preliminaries

We consider a Markov decision process with a considerably large or continuous state space S ⊂ RM and action space A ⊂ RN , a non-negative bounded reward function R : S × A → [0,Rmax], and a transition kernel P(s′ |s, a) that indicates the probability from a state-action pair (s, a) to a state s′ . We define U(s) ⊂ A as the admissible set of actions at state s, and π : S → A as a feasible stationary policy. We seek to solve a discrete-time sequential decision problem as follows:

where γ ∈ (0, 1) is the discount factor. Let J ⋆ (s) := Jπ⋆(s) be the optimal value function when applying the optimal policy π⋆ . The corresponding optimal Q-function can be defined as

where s′ ∼ P (·|s, a) and we use the same notation hereafter when no confusion arises. In other words, Q⋆ (s, a) represents the reward of an agent who starts from state s and takes action a at the first step and then follows the optimal policy π⋆ thereafter.

2.1 Q-learning Algorithm

This paper focuses on the Q-learning algorithm that uses a parametric function ˆQ(s, a; θ) to approximate the Q-function with a parameter θ having finite and relatively small dimensions. The update rule of Q-learning is given by

where αt is the step size at time t. It is clear that Q-learning performs the update by taking one step of temporal target update and one step of parameter learning in an alternating fashion.

2.2 Linear Function Approximation

Like most of the related work, we focus on the convergence analysis under the linear function approximation. A linear approximation of the Q-function ˆQ(s, a; θ) can be written as

where θ ∈ Rd , and φ : S ×A → Rd is a vector function of size d, and the elements of Φ represent the nonlinear kernel (feature) functions.

Remark 1. We note that recent work [Xu and Gu, 2019] established the convergence rate of Q-learning with neural network approximation, which exploits the approximate linear structure of the neural network in the overparameterized regime. Thus, our analysis under the linear function approximation can be generalized to the function class of overparameterized neural networks by applying the techniques developed in recent work [Xu and Gu, 2019].

3 Convergence Analysis of Q-AMSGrad

In this section, we characterize the convergence guarantee for Q-learning under Adam-type updates.

3.1 Q-AMSGrad Algorithm

Although Adam has obtained great success as an optimizer in deep learning, it is well known that Adam by nature is nonconvergent even for simple convex loss functions [Reddi et al., 2018]. Instead, a slightly modified version called AMSGrad [Reddi et al., 2018] is widely used to study the convergence property of Adam-type algorithms in conventional optimization. Here, we apply the update rule of AMSGrad to the Q-learning algorithm and refer to such an algorithm as Q-AMSGrad. Algorithm 1 describes Q-AMSGrad in detail.

More specifically, the iterations of Q-AMSGrad evolve by updating the exponentially decaying average of historical gradients (mt) and squared historical gradients (vt). The hyperparameters β1, β2 are used to exponentially decrease the rate of the moving averages. The difference between AMSGrad and Adam lies in the fact that AMSGrad makes the sequence ˆ vt,i increasing along the time step t for each entry i ∈ [d], whereas Adam does not guarantee such a property.

Algorithm 1 Q-AMSGrad

• 1: Input: α, λ, θ1, β1, β2,m0 = 0, ˆ v0 = 0.
• 2: for t = 1, 2, . . . , T do
• 3: αt = α √t , β1t = β1λt
• 4: Observe data (st, at, st+1) from policy π and transition probability P
• 5: bt = R(st, at) +γmax a′ φT (st+1, a′ )θt
• 6: gt = (φT (st, at)θt - bt ) φ(st, at)
• 7: mt = (1 -β1t)mt-1 +β1tgt
• 8: vt = (1 -β2)ˆ vt-1 + β2g2 t
• 9: ˆ vt = max(ˆ vt-1, vt), ˆ Vt = diag(ˆ v1, . . . , ˆ vd)
• 10: θt+1 = ΠD, ˆ V 1/4 t (θt - αt ˆ V - 1 2 t mt) where ΠD, ˆ V 1/4 t (θ′ ) = min θ∈D ∥∥∥ ˆ V 1/4 t (θ′ - θ)∥∥∥.

• 11: end for
• 12: Output: 1 T ∑Tt=1 θt

3.2 Convergence Result

Before stating the main theorem, we first introduce some technical assumptions and lemmas for our analysis.

Assumption 1. For any state-action pair (s, a) ∈ S ×A, the kernel function φ is uniformly bounded and we have

This assumption is mild since we can normalize the kernel function if the kernel function is uniformly bounded. It is widely applied in the literature to simplify the analysis of RL algorithms with linear function approximation [Bhandari et al., 2018; Chen et al., 2019b].

Assumption 2. [Chen et al., 2019b, Lemma 6.7] At each iteration t, the noisy gradient is unbiased, i.e. gt = ¯ gt + ξt with Eξt = 0 where ¯ gt = E[gt]. The equation ¯ g(θ) = 0 has a unique solution θ⋆ , and there exists a c > 0, such that for any θ ∈ Rd we have(θ - θ⋆ )T ¯ g(θ) ≥ c ‖θ - θ⋆ ‖2 . (8)

Assumption 2 has been proved as a key technical lemma in [Chen et al., 2019b] under certain assumptions, which appears to be the weakest among the existing studies for establishing the convergence guarantee for Q-learning with linear function approximation. It is the standard assumption in the related literature to analyze the convergence Qlearning with linear function approximation [Zou et al., 2019b; Chen et al., 2019b; Xu and Gu, 2019].

Assumption 3. The domain D ⊂ Rd of approximation parameters is a ball originating at θ = 0 with bounded diameter containing θ⋆ . That is, there exists D∞, such that ‖θm - θn‖ < D∞, ∀θm, θn ∈ D, and θ⋆ ∈ D.

This Assumption can be easily satisfied when we apply a projected algorithm, and is standard in the theoretical analysis of Adam-type algorithms [Chen et al., 2019a; Zhou et al., 2018].

Based on the above assumptions, we can immediately obtain the bounded property of the gradient, which is stated in the following lemma.

Lemma 1. Under Assumptions 1 and 3, at each iteration t, the gradient estimator gt in Q-AMSGrad is uniformly bounded. That is,

In addition, we denote G∞ = Rmax + (1 + γ)D∞ and let {mt, ˆ vt} for t = 1, 2, . . . be sequences generated by Algorithm 2. Then we have

We next provide the non-asymptotic convergence of QAMSGrad in the following theorem.

Theorem 1. (Convergence of Q-AMSGrad) Suppose αt = α √t , β1t = β1λt and δ = β1/β2 with δ, λ ∈ (0, 1) for t = 1, 2, . . . in Algorithm 1. Given Assumptions 1 ∼ 3, the output of Q-AMSGrad satisfies:

In Theorem 1, B1,B2,B3 in the bound in Equation (10) are constants and independent of time. Therefore, under the choice of the stepsize and hyper-parameters in Algorithm 1, Q-AMSGrad achieves a convergence rate of O ( 1 √T ) when ∑di=1 ‖g1:T,i‖ << √T [Reddi et al., 2018].

Remark 2. Our proof of convergence here has two major differences from that for AMSGrad in [Reddi et al., 2018] in conventional optimization: (a) The two algorithms are quite different. Q-AMSGrad is a Q-learning algorithm alternatively finding the best policy with a moving target, whereas AMSGrad is an optimizer for conventional optimization and does not have alternating nature. (b) Our analysis is on the convergence rate whereas [Reddi et al., 2018] provides regret bound. In fact, a slight modification of our proof also provides the convergence rate of AMSGrad for conventional strongly convex optimization, which can be of independent interest. Moreover, our proof avoids the theoretical error in the proof in [Reddi et al., 2018] as pointed out by [Tran and others, 2019].

4 Convergence Analysis of Q-AMSGradR

In this section, we propose to incorporate a momentum restart technique to Q-AMSGrad in order to improve its performance. We first introduce the algorithm and then provide the convergence analysis for such an algorithm. We demonstrate its desired experimental performance in Section 5.

4.1 Q-AMSGrad AlgorithmwithMomentum Restart

We introduce the restart technique to Q-AMSGrad and propose Q-AMSGradR as shown in Algorithm 2. Q-AMSGradR applies the same update rule as Algorithm 1, but periodically

Algorithm 2 Q-AMSGradR

• 1: Input: α, λ, θ1, β1, β2,m0 = 0, ˆ v0 = 0.
• 2: for t = 1, 2, . . . , T do
• 3: if mod(t, r) = 0 then
• 4: mt = 0, ˆ vt = 0
• 5: end if
• 6: αt = α √t , β1t = β1λt
• 7: Observe data (st, at, st+1) from policy π and transition probability P
• 8: bt = R(st, at) +γmax a′ φT (st+1, a′ )θt

• 10: mt = (1 -β1t)mt-1 +β1tgt
• 11: vt = (1 -β2)ˆ vt-1 + β2g2 t
• 12: ˆ vt = max(ˆ vt-1, vt), ˆ Vt = diag(ˆ v1, . . . , ˆ vd)

• 14: end for
• 15: Output: 1 T ∑Tt=1 θt

resets mt, ˆ vt with a period of r, i.e., mt = 0, ˆ vt = 0, ∀t = kr, k = 1, 2, · · · . We explain such an idea further as follows.

Traditional momentum-based algorithms largely depend on the historical gradient direction. When part of the historical information is incorrect, the estimation error tends to accumulate. The restart technique can be employed to deal with such an issue. One way to restart the momentum-based methods is to initialize the momentum at some restart iteration. That is, at restart iteration r, we reset mr, vr, i.e., mr = 0, vr = 0, which yields θr+1 = θr. It is an intuitive implementation technique to adjust the trajectory from time to time, and can usually help mitigate the aforementioned problem while keeping fast convergence property. For the implementation, we execute the restart periodically with a period r. It turns out that the restart technique can significantly improve the numerical performance, which can be seen in Section 5.

4.2 Convergence Result

In the following theorem, we provide the non-asymptotic convergence for Q-AMSGradR.

Theorem 2. (Convergence of Q-AMSGradR) Suppose αt = α √t , β1t = β1λt and δ = β1/β2 with δ, λ ∈ (0, 1) for t = 1, 2, . . . in Algorithm 2. Given Assumptions 1 ∼ 3, the output of Q-AMSGradR satisfies:

Theorem 2 indicates that for Q-AMSGradR to enjoy a convergence rate of O ( 1 √T ), the restart period r needs to be sufficiently large and ∑di=1 ‖g1:T,i‖ << √T . In practice as demonstrated by the experiments in Section 5, Q-AMSGradR typically performs well, not necessarily under the theoretical conditions.

5 Experimental Performance

In this section, we empirically evaluate the Q-learning algorithms studied in this paper. We first study the linear quadratic regulator (LQR) problem, which serves as a direct numerical demonstration of the convergence analysis under linear function approximation. We then use the Atari 2600 games [Brockman et al., 2016], a classic benchmark for DQN evaluations, to demonstrate the effectiveness of the Q-learning algorithms for complicated tasks. Note that our experiments test the performance of Q-Adam and Q-AdamR, which serve as practical versions of Q-AMSGrad and Q-AMSGradR by adopting Adam for Q-learning and keeping all major properties of QAMSGrad and Q-AMSGradR including the alternating update between Q-function update and parameter fitting.

Our main focus here lies in: (a) comparison between vanilla Q-Adam and that with momemtum restart (Q-AdamR), (b) comparison between Q-Adam/Q-AdamR and vanilla Qlearning (through LQR), and (c) comparison between QAdam/Q-AdamR and DQN (through Atari 2600 games). Both of our experiments show that Q-AdamR outperforms Q-Adam, vanilla Q-learning and DQN in terms of convergence speed and variance reduction. Compared with DQN in the empirical experiments of Atari games, under the same hyper-parameter settings, Q-Adam and Q-AdamR improve the performance of DQN by 50% on average.

5.1 DQNAlgorithm

As DQN is also included in this work for performance comparison. We recall the update of DQN in the following as reference. Differently from the vanilla Q-learning, DQN updates the parameters in a nested loop. Within the t-th inner loop, DQN first obtains the target Q-function as in Equation (12), and then uses a neural network to fit the target Q-function by running Y steps of a certain optimization algorithm as Equation (13). The update rule of DQN is given as follows.

where Optimizer can be SGD or Adam for example, and Equation (13) is thus a supervised learning process with T ˆQ(s, a; θ0 t )) as the "supervisor". At the t-th outer loop, DQN performs the so-called target update as

In practice, when one of the momentum-based optimizers is adopted for Equation (13), such as Adam, it is only initialized once at the beginning of the first inner loop. The historical gradient terms then accumulate throughout multiple inner loops with different targets.

Generally speaking, the difference between Q-Adam/QAdamR and DQN mainly lies in Q-Adam/Q-AdamR takes one-step Q-function update and one-step model parameter fitting alternatively, whereas DQN takes one-step Q-function update followed by a sufficient large number of steps for model parameter fitting (towards a target Q-function).

5.2 Linear Quadratic Regulator

We numerically validate the performance of Q-Adam and QAdamR through an infinite-horizon discrete-time LQR problem. A typical model-based solution (with known dynamics), known as the discrete-time algebraic Riccati equation (DARE), is adopted to derive the optimal policy u⋆t = -K⋆xt. The performance of the learning algorithm is then evaluated at each step of iterate t with the Euclidean norm ‖Kt -K⋆ ‖. Given the problem nature of LQR, we also re-scale the loss term of L(θt) := ˆQt(s, a; θt) - T ˆQ(s, a; θt) in Equation (4) as ˜L(θt) = ˜ τ 2L(θt) with some scaling factor ˜ τ ∈ (0, 1], which is beneficial for stabilizing the learning process. The performance result for each method is averaged over 10 trials with different random seeds. All algorithms share the same set of random seeds and are initialized with the same θ0. The hyperparameters of the learning settings are also consistent and further details are shown in Table 1. Note that for all the implementations, we also adopt the double Q-update [Van Hasselt et al., 2016] to help prevent over-estimations of the Q-value. The performance results are provided in Figure 1. Here we highlight the main observations from the LQR experiments.

Q-AdamR outperforms DARE. In ideal cases where data sampling perfectly emulates the system dynamics and the target is accurately learned in each inner loop, DARE for LQR would become equivalent to the DQN-like update if the neural network is replaced with a parameterized linear function. In practice, such ideal conditions are difficult to satisfy, and hence the actual Q-learning with target update is usually far slower (in terms of the number of steps of target updates) than DARE. Note that Q-AdamR performs significantly well and even converges faster than DARE, and thus implies it is faster than the most well-performing Q-learning with target update.

Q-AdamRoutperforms Q-Adam. Overall, under the same batch sampling scheme and restart period, Q-AdamR achieves a faster convergence and smaller variance than Q-Adam.

5.3 Atari Games

We apply the Q-Adam and Q-AdamR algorithms to the more challenging tasks of deep convolutional neural network playing a group of Atari 2600 games. The particular DQN we train to compare against adopts the dueling network structure [Wang et al., 2016], double Q-learning setup [Van Hasselt et al., 2016], ϵ-greedy exploration and experience replay [Mnih et al., 2015]. Adam is also adopted, without momentum restart, as the optimizer for the inner-loop supervised learning process. Q-Adam and Q-AdamR are implemented using the identical setup of network construction, exploration and sampling strategies.

We test all three algorithms with a batch of 23 Atari games. The choice of 10 million steps of iteration is a common setup for benchmark experiments with Atari games. Although this does not guarantee the best performance in comparison with more time-consuming training with 50 million steps or more, it is sufficient to illustrate different performances among the selected methods. The software infrastructure is based on the baseline implementation of OpenAI. Selections of the hyperparameters are listed in Table 2. We summarize the results in Figure 2. The overall performance is illustrated by first normalizing the return of each method with respect to the results obtained from DQN, and then averaging the performance of all 23 games to obtain the mean return and standard deviation. Considering we use a smaller buffer size than common practice, DQN is not consistently showing improved return over all tested games. Therefore, the self-normalized average return of DQN in Figure 2 is not strictly increasing from 0 to 100%.

Overall, both Q-Adam and Q-AdamR achieve significant improvement in comparison with the DQN results. In particular, AltQ-Adam increases the performance by over 100% in some of the tasks including Asterix, BeamRider, Enduro, Gopher, etc. However, it also illustrates certain instability with complete failure on Amidar and Assault. This is also capture by the higher variance illustrated on Figure 2. Periodic restart (Q-AdamR) resolves this issue efficiently with an on-par performance on average and far smaller variance. In terms of the maximum average return, Q-Adam and Q-AdamR perform no worse then DQN on 17 and 20 games respectively out of the 23 games being evaluated. Furthermore, if we consider having a final score that is smaller or equal to the start score as a learning failure, DQN fails the learning in 5 out of 23 games (Asteroids, DoubleDunk, Gravitar, Pitfall, Tennis), Q-Adam fails in 3 out of 23 games (Amidar, Assault, Asteroids) and Q-AdamR does not fail in any of the tasks. That is, Q-AdamR not only reduces the variance, but also provides a more consistent performance across the task domain. This implies that momentum restart effectively corrects the accumulated error and stabilizes the training process.

6 Conclusion

We study two Q-learning algorithms with Adam-type updates, and demonstrate their superior performance over the vanilla Q-learning and DQN algorithms through a linear quadratic regulator problem and a batch of 23 Atari 2600 games.

It is of considerable future interest to further investigate the potential of the restart scheme. One possible direction is to develop an adaptive restart mechanism with changing period determined by an appropriately defined signal of restart. This will potentially relieve the effort in hyper-parameter tuning of finding a good fixed period.

Acknowledgements

The work was supported in part by the U.S. National Science Foundation under Grants CCF-1761506, ECCS-1818904, CCF-1909291 and CCF-1900145, and the startup funding of the Southern University of Science and Technology (SUSTech), China.