Damped Anderson Mixing for Deep Reinforcement Learning: Acceleration, Convergence, and Stabilization

1 Introduction . In reinforcement learning ( RL ) [ 28 ] , an agent seeks an optimal policy in a sequential decisionmaking process . Deep RL has recently achieved significant improvements in a variety of challenging tasks , including game playing [ 20 , 26 , 18 ] and robust navigation [ 19 ] . A flurry of state-of-theart algorithms have been proposed , including Deep Q-Learning ( DQN ) [ 20 ] and variants such as Double-DQN [ 13 ] , Dueling-DQN [ 31 ] , Deep Deterministic Policy Gradient ( DDPG ) [ 16 ] , Soft Actor-Critic [ 12 ] and distributional RL algorithms [ 5 , 17 , 33 ] , all of which have successfully solved end-to-end decision-making problems such as playing Atari games . However , the slow convergence and sample inefficiency of RL algorithms still hinders the progress of RL research , particularly in high-dimensional state spaces where deep neural network are used as function approximators , making learning in real physical worlds impractical . To address these issues , various acceleration strategies have been proposed , including the classical Gauss-Seidel Value Iteration [ 23 ] and Jacobi Value Iteration [ 24 ] . Another popular branch of techniques accelerates RL by leveraging historical data . Interpolation methods such as Average-DQN [ 2 ] ∗Equal contributions in alphabetical order †Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . have been widely used in first-order optimization problems [ 6 ] and have been proven to converge faster than vanilla gradient methods . As an effective multi-step interpolation method , Anderson mixing [ 30 , 9 ] , also known as Anderson acceleration , has attracted great attention from RL researchers . The insight underpinning of Anderson acceleration is that RL [ 28 ] is intrinsically linked to fixed-point iterations : the optimal value function is the fixed point of the Bellman optimality operator . These fixed-points are computed recursively by repeatedly applying an operator of interest [ 11 ] . Anderson mixing is a general method to accelerate fixed-point iterations [ 30 ] and has been successfully applied to fields , such as the computational chemistry [ 22 ] or electronic structure computation [ 1 ] . In particular , Anderson acceleration leverages the m previous estimates in order to find a better estimate in a fixed-point iteration . To compute the mixing coefficients in Anderson iteration , it searches for a point with a minimal residual within the subspace spanned by these estimates . It is thus natural to explore the efficacy of Anderson acceleration in RL settings . Several works [ 15 , 25 ] have attempted to apply Anderson acceleration to reinforcement learning . Anderson mixing was first applied to value iteration in [ 11 , 15 ] and resulted in significant convergence improvements . Regularized Anderson acceleration [ 25 ] was recently proposed to further accelerate convergence and enhance the final performance of state-of-the-art RL algorithms in various experiments . However , previous applications of Anderson acceleration were typically heuristic : consequently , these empirical improvements in convergence have so far lacked a rigorous mathematical justification . In this paper , we provide deeper insights into Anderson acceleration in reinforcement learning by establishing its connection with quasi-Newton methods for policy iteration and improved convergence guarantees under the assumptions that the Bellman operator is differential and non-expansive . MellowMax operator is adopted to replace the max operator in policy iteration to simultaneously guarantee faster convergence of value function and reduce the estimated gradient variance to yield stabilization . In addition , we analyze the stability properties of Anderson acceleration in policy iteration and propose a stable regularization to further enhance the stability . These key two factors , i.e. , the stable regularization and the theoretically-inspired MellowMax operator , are the basis for our Stable Anderson Acceleration ( Stable AA ) method . Finally , our experimental results on various Atari games demonstrate that our Stable AA method enjoys faster convergence and achieves better performance relative to existing Anderson acceleration baselines . Our work provides a unified analytic framework that illuminates Anderson acceleration for reinforcement learning algorithms from the perspectives of acceleration , convergence , and stabilization . 2 Acceleration and Convergence Analysis of Anderson Acceleration on RL . We first present the notion of Anderson acceleration in the reinforcement learning and then provide deeper insights into the acceleration if affords by establishing a connection with quasi-Newton methods . Finally , a theoretical convergence analysis is provided to demonstrate that Anderson acceleration can increase the convergence radius of policy iteration by an extra contraction factor . Background Consider a Markov decision process ( MDP ) specified by the tuple ( S , A , R , p , γ ) , where S is a set of the states and A is a set of actions . The functions R : S × A → R and p : S × A × S → [ 0 , 1 ] are the reward function , with Rt = R ( s , a ) , and transition dynamics function , respectively for the MDP . The discount rate is denoted by γ ∈ [ 0 , 1 ) and determines the relative importance of immediate rewards relative to rewards received in the future . The Qvalue function evaluates the expected return starting from a given state-action pair ( s , a ) , that is , Qπ ( s , a ) = E [ ∑∞ t=0 γ tRt+1 | s0 = s , a0 = a ] . A policy π ( a|s ) is a distribution mapping the state space S to the action space A . 2.1 Anderson Acceleration in Policy Iteration . We focus on the tabular case to enable the theoretical analysis of Anderson acceleration in value ( policy ) iteration , which can be naturally applied to function approximation . Both the value iteration ( V -notation ) and the policy iteration ( Q-notation ) can have Anderson acceleration applied to them to improve convergence . However , theoretical analysis has shown that value iteration enjoys a γ-linear convergence rate , i.e. , ‖V ( t ) − V ∗‖∞ ≤ γ‖V ( t−1 ) − V ∗‖∞ , where V ( t ) is the value function in the iteration step t and V ∗ is the optimal value function , while policy iteration converges faster . This is due to the fact that policy iteration more fully evaluates the current policy than does value iteration . Additionally , policy iteration is more fundamental and scales more readily to deep reinforcement learning . For this reason , we analyze our method in the policy iteration setting under Q-notation as value iteration is a special case of policy iteration . Thus , our analysis also applies under value iteration . We first focus on the control setting where the optimal value of state-action pair Q∗ ( s , a ) is defined recursively as a function of the optimal value of the other state-action pair : Q∗ ( s , a ) = R ( s , a ) + γ ∑ s′∈S p ( s′ | s , a ) ·max a′ Q∗ ( s′ , a′ ) . ( 1 ) Combining policy evaluation and policy improvement , the resulting policy iteration algorithm is equivalent to solving for the fixed point of the Bellman optimality operator T : R|S×A| → R|S×A| with T Q ( s , a ) = R ( s , a ) + γ ∑ s′∈S p ( s ′|s , a ) ·maxa′ Q ( s′ , a′ ) . As a general technique to speed up fixed-point iteration [ 30 ] , Anderson acceleration has been successfully yet heuristically applied to reinforcement learning algorithms [ 25 , 11 , 15 ] . Specifically , Anderson acceleration linearly combines the m previous estimates to yield a better iteration target in the fixed point iteration . Geometrically , Anderson acceleration applies the operator to a point that has a minimal residual within the subspace spanned by these estimates . In policy iteration , Anderson acceleration maintains a memory of the previous Q function values and updates the iterate as a linear combination of these values with dynamic weights α ( k ) in the kth iteration step . Specifically , Q ( k+1 ) ( s , a ) = ( 1− βk ) m∑ i=0 α ( k ) i Q ( k−m+i ) ( s , a ) + βk m∑ i=0 α ( k ) i T Q ( k−m+i ) ( s , a ) , ( 2 ) where 0 ≤ βk ≤ 1 is the damping parameter . All of the coefficients α ( k ) i in the coefficient vector α ( k ) are computed following α ( k ) = argmin α∈Rm+1 ∥∥∥∥∥ m∑ i=0 αi ( T Q ( k−m+i ) −Q ( k−m+i ) ) ∥∥∥∥∥ 2 = argmin α∈Rm+1 ∥∥∥∆′kT · α∥∥∥ 2 , s.t . m∑ i=0 αi = 1 , ( 3 ) where ∆′k T = [ ek−m , · · · , ek ] ∈ R|S×A|× ( m+1 ) and ek = T Q ( k ) −Q ( k ) ∈ R|S×A| is the Bellman residuals matrix . By the Karush-Kuhn-Tucker conditions , the analytic solution of optimal coefficient vector αk is α ( k ) = ( ∆′k∆ ′ k T ) −1 1 1T ( ∆′k∆ ′ k T ) −1 1 , ( 4 ) where 1 denotes the vector with all components equal to one . 2.2 Connection Between Damped Anderson Acceleration and Quasi-Newton Methods . We know that the optimization problem is closely linked with solving a fixed-point iteration problem by directly solving its first-order condition . Inspired by [ 8 ] , we show that Anderson acceleration in policy iteration attempts to perform a special form of quasi-Newton iteration from its optimization problem behind . To illuminate this connection , we firstly show that the original constrained optimization to obtain the optimal α ( k ) in Eq . ( 3 ) can be equivalent to the unconstrained one τ ( k ) = argminτ∈Rm ‖ek −Hkτ‖ 2 , ( 5 ) where we let n = |S| × |A| , and then Hk = [ ek − ek−1 , · · · , ek−m+1 − ek−m ] ∈ Rn×m . τ ( k ) = [ τ ( k ) 0 , τ ( k ) 1 , · · · , τ ( k ) m−1 ] T ∈ Rm with τ ( k ) i = ∑m−i−1 j=0 α ( k ) j . Let δk = Q ( k ) − Q ( k−1 ) , ∆k = [ δk , δk−1 , · · · , δk−m+1 ] ∈ Rn×m . We show that the updating rule of Q ( s , a ) in Anderson acceleration can be expressed as a quasi-Newton form in Proposition 1 . Proposition 1 . By conducting the damped Anderson acceleration ( Eq . ( 2 ) and ( 3 ) ) on the policy iteration , the updating of Q ( k+1 ) can be reformulated as Q ( k+1 ) : = Q ( k ) −Gkek ( 6 ) whereGk = ( ∆k + βkHk ) ( HTk Hk ) −1 HTk −βkI can serve as an approximation of inverse Jacobian matrix of ek = TQ ( k ) −Q ( k ) , and I is an identity matrix . Proposition 1 points out that Anderson acceleration on policy iteration additionally leverages more information about the fixed-point residual ek = TQ ( k ) −Q ( k ) to update the Q function . Particularly , the first part ( ∆k + βkHk ) ( HTk Hk ) −1HTk in Gk contains partial structure matrix information about the real inverse of Jacobian matrix , which has the huge potential to speed up the convergence of the fixed-point iteration . More importantly , the results established in [ 30 ] and [ 9 ] can been seen as special cases of Proposition 1 with βk = 1 . If we directly get rid of the first part in Gk and set βk = 1 , the updating rule exactly degenerates to the Q-value function iteration without Anderson acceleration .

Summary: This paper looks to characterise the use of Anderson acceleration in Deep Q learning.  Theoretical convergence results are given in tabular settings. A novel algorithm is developed to satisfy the assumptions made in the theory. The algorithm outperforms reasonable baselines and ablation studies.