Actor-critic is implicitly biased towards high entropy optimal policies

1 OVERVIEW . Reinforcement learning methods navigate an environment and seek to maximize their reward ( Sutton & Barto , 2018 ) . A key tension is the tradeoff between exploration and exploitation : does a learner ( also called an agent or policy ) explore for new high-reward states , or does it exploit the best states it has already found ? This is a sensitive part of RL algorithm design , as it is easy for methods to become blind to parts of the state space ; to combat this , many methods have an explicit exploration component , for instance the ϵ-greedy method , which forces exploration in all states with probability ϵ ( Sutton & Barto , 2018 ; Tokic , 2010 ) . Similarly , many methods must use projections and regularization to smooth their estimates ( Williams & Peng , 1991 ; Mnih et al. , 2016 ; Cen et al. , 2020 ) . This work considers actor-critic methods , where a policy ( or actor ) is updated via the suggestions of a critic . In this setting , prior work invokes a combination of explicit regularization and exploration to avoid getting stuck , and makes various fast mixing assumptions to help accurate exploration . For example , recent work with a single trajectory in the tabular case used both an explicit ϵ-greedy component and uniform mixing assumptions ( Khodadadian et al. , 2021 ) , neural actor-critic methods use a combination of projections and regularization together with various assumptions on mixing and on the path followed through policy space ( Cai et al. , 2019 ; Wang et al. , 2019 ) , and even direct analyses of the TD subroutine in our linear MDP setting make use of both projection steps and an assumption of starting from the stationary distribution ( Bhandari et al. , 2018 ) . Contribution . This work shows that a simple linear actor-critic ( cf . Algorithm 1 ) in a linear MDP ( cf . Assumption 1.3 ) with a finite but non-tabular state space ( cf . Assumption 1.1 ) finds an ϵ-optimal policy in poly ( 1/ϵ ) samples , without any explicit exploration or projections in the algorithm and without any uniform mixing assumptions on the policy space ( cf . Theorem 1.4 ) . The algorithm and analysis avoid both via an implicit bias towards high entropy policies : the actor-critic policy path never leaves a Kullback-Leibler ( KL ) divergence ball of the maximum entropy optimal policy , and this firstly ensures implicit exploration , and secondly ensures fast mixing . In more detail : 1 . Actor analysis via mirror descent . We write the actor update as an explicit mirror descent . While on the surface this does not change the method ( e.g. , in the tabular case , the method is identical to natural policy gradient ( Agarwal et al. , 2021b ) ) , it gives a clean optimization guarantee which carries a KL-based implicit bias consequence for free , and decouples concerns between the actor and critic . 2 . Critic analysis via projection-free sampling tools within KL balls . The preceding mirror descent component guarantees that we stay within a small KL ball , if the statistical error of the critic is controlled . Concordantly , our sampling tools guarantee this statistical error is small , if we stay within a small KL ball . Concretely , we provide useful lemmas that every policy in a KL ball around the high entropy policy has uniformly upper bounded mixing times , and separately give a projection-free ( implicitly regularized ! ) analysis of the standard temporal-difference ( TD ) update from any starting state ( Sutton , 1988 ) , whereas the closest TD analysis in the literature uses projections and requires the sampling process to be started from the stationary distribution ( Bhandari et al. , 2018 ) . The mixing assumptions here contrast in general with prior work , which either makes explicit use of stationary distributions ( Cai et al. , 2019 ; Wang et al. , 2019 ; Bhandari et al. , 2018 ) , or makes uniform mixing assumptions on all policies ( Xu et al. , 2020 ; Khodadadian et al. , 2021 ) . Our final proof has the above actor and critic components feeding off each other : we use an induction to simultaneously establish that since the actor stayed in a small KL ball in previous iterations , then the new critic update is accurate , and similarly the accuracy of the previous critic updates ensures the actor continues to stay in a small KL ball . We feel these tools will be useful in other work . 1.1 SETTING AND MAIN RESULTS . We will now give the setting , main result , and algorithm in full . Further details on MDPs can be found in Section 1.3 , but the actor-critic method appears in Algorithm 1 . To start , the environment and policies are as follows . Assumption 1.1 . The Markov Decision Process ( MDP ) has states s ∈ Rd and finitely many actions a ∈ A : = { e1 , . . . , ek } , and finite rewards r ∈ [ 0 , 1 ] . States are observed in some feature encoding s ∈ Rd , but the state space S ⊆ { s ∈ Rd : 1/2 ≤ ∥s∥ ≤ 1 } is assumed finite : |S| < ∞ . Policies are linear softmax policies : a policy is parameterized by a weight matrix W ∈ Rd×k , and given a state s ∈ Rd , uses a per-state softmax to sample a new action a : a ∼ ϕ ( sTW · ) , where ϕ ( sTWa ) = exp ( s TWa ) ∑ b∈A exp ( s TWb ) . ( 1.1 ) Let As denote the set of optimal actions for a given state s. It is assumed that As is nonempty for every s ∈ S ; equivalently , for every state , there exists an optimal policy whose stationary distribution exists and places positive mass on that state . ♢ The choice of linear policies simplifies presentation and analysis , but the tools here should be applicable to other settings . This choice also allows direct comparison to the widely-studied implicit bias Algorithm 1 Single-trajectory linear actor-critic . Inputs : actor iterations t and step size θ ; critic iterations N and step size η. Initialize : actor weights W0 = 0 ∈ Rd×k , pre-softmax mapping p0 ( s , a ) : = sTW0a , policy π0 : = ϕ ( p0 ) ( cf . eq . ( 1.1 ) ) ; sample initial state/action/reward triple ( s0,0 , a0,0 , r0,0 ) . for i = 0 , 1 , 2 , . . . , t− 1 do Critic update : use πi to interact with the MDP , obtaining state/action/reward triples ( si , j , ai , j , ri , j ) j≤N by continuing the existing trajectory , and form TD estimates Ui , j+1 : = Ui , j − ηsi , j ( sTi , jUi , jai , j − γsTi , j+1Ui , jai , j+1 − ri , j ) aTi , j , with initial condition Ui,0 = 0 . Set Ûi : = 1N ∑ j < N Ui , j and Q̂i ( s , a ) : = sTÛia , and also ( si+1,0 , ai+1,0 , ri+1,0 ) : = ( si , N , ai , N , ri , N ) to continue the existing trajectory in next iteration . Actor update : set Wi+1 : =Wi + θÛi and pi+1 : = pi + θQ̂i and πi+1 : = ϕ ( pi+1 ) . end for of gradient descent in linear classification settings ( Soudry et al. , 2017 ; Ji & Telgarsky , 2018 ) , as will be discussed further in Section 1.2 . The choice of finite state space is to remove measure-theoretic concerns and to allow a simple characterization of the maximum entropy optimal policy . Lemma 1.2 ( simplification of Lemma A.1 ) . Under Assumption 1.1 , there exists a unique maximum entropy policy π , which satisfies π ( s , · ) = Uniform ( As ) for every state s. To round out this introductory presentation of the actor , the last component is the update : pi+1 : = pi + Q̂i and πi+1 : = ϕ ( pi+1 ) . This is explicitly a mirror descent or dual averaging representation of the policy , where we use a mirror mapping ϕ to obtain the policy πi+1 from pre-softmax values pi+1 . As mentioned before , this update appears in prior work in the tabular setting with natural policy gradient and actor-critic ( Agarwal et al. , 2021b ; Khodadadian et al. , 2021 ) , and will be related to other methods in Section 1.2 . We will further motivate this update in Section 2 . The final assumption and description of the critic are as follows . As will be discussed in Section 2 , the policy becomes optimal if Q̂i is an accurate estimate of the true Q function . We employ a standard TD update with no projections or constraints . To guarantee that this linear model of Qi is accurate , we make a standard linear MDP assumption ( Bradtke & Barto , 1996 ; Melo & Ribeiro , 2007 ; Jin et al. , 2020 ) . Assumption 1.3 . In words , the linear MDP assumption is that the MDP rewards and transitions are modeled by linear functions . In more detail , for convenience first fix a canonical vector form for state/action pairs ( s , a ) ∈ Rd×k : let xsa ∈ Rdk denote the vector obtained via unrolling the matrix saT row-wise ( whereby vector inner products with xsa match matrix inner products with saT ) . The linear MDP assumption is then that there exists a fixed vector y ∈ Rdk and a fixed matrix M ∈ Rd×dk so that for any state/action pair xsa and any subsequent state s′ ∈ Rd , E [ r | ( s , a ) ] = xTsay , and E [ s′ | ( s , a ) ] =Mxsa . Lastly , suppose 1/2 ≤ ∥s∥ ≤ 1 for all s ∈ S. ♢ Though a strong assumption , it is not only common , but note also that since TD must continually interact with the MDP , then it would have little hope of accuracy if it can not model short-term MDP dynamics . Indeed , as is shown in Lemma C.3 ( but appears in various forms throughout the literature ) , Assumption 1.3 implies that the fixed point of the TD update is the true Q function . This assumption and the closely related compatible linear function approximation assumption will be discussed in Section 1.2 . We now state our main result , which bounds not just the value function V ( cf . Section 1.3 ) but also the KL divergence Kvsπ ( π , πi ) = Es′∼vsπ ∑ a π ( s ′ , a ) ln π ( s ′ , a ) πi ( s′ , a ) , where vsπ is the visitation distribution of the maximum entropy optimal policy π when run from state s ( cf . Section 1.3 ) . Theorem 1.4 . Suppose Assumptions 1.1 and 1.3 ( which imply the ( unique ) maximum entropy optimal policy π is well-defined , and also irreducible and aperiodic ) . Given iteration count t , choose θ = Θ ( 1 t13/16 ln ( t ) 1/4 ) , N = Θ ( t2 ln t ) , η = Θ ( 1√ N lnN ) , where the constants hidden inside each Θ ( · ) depend only on π and the MDP , but not on t. With these parameters in place , invoke Algorithm 1 , and let ( πi ) i < t be the resulting sequence of policies . Then with probability at least 1− 1/t1/8 , simultaneously for every state s ∈ Rd and every i ≤ t , Kvsπ ( π , πi ) + θ ( 1− γ ) ∑ j < i ( Vπ ( s ) − Vj ( s ) ) ≤ ln k + 1 ( 1− γ ) 2 . Before outlining the proof structure and organization of the rest of the paper , a few comments on the interpretion of Theorem 1.4 are as follows . Remark 1.5 ( Discussion of Theorem 1.4 ) . 1 . Implicit bias . Since π is optimal , the second term can be deleted , and the bound implies max i≤t s∈S Kvsπ ( π , πi ) ≤ ln k + 1 ( 1− γ ) 2 ; since this holds for all i ≤ t , it controls the optimization path . This term is a direct consequence of our mirror descent setup , and is used to control the TD errors at every iteration . This implicit bias of the policy path stands therefore in stark contrast to the worstcase KL divergence between arbitrary softmax policies and π , which is infinite : to show sup { Kvsπ ( π , ϕ ( ⟨W , ·⟩ ) ) : W ∈ R d×k , s ∈ S } = ∞ , it suffices to consider a sequence ( Wr ) r≥1 so that π ( s′ , a ) /ϕ ( ( s′ ) TWra ) → ∞ as r → ∞ for some ( s′ , a ) with vsπ ( s′ ) > 0 and π ( s′ , a ) > 0 . 2 . Mixing time constants . The critic iterations N and step size η hide mixing time constants ; these mixing time constants depend only on the KL bound ln k+ 1/ ( 1− γ ) 2 , and in particular there is no hidden growth in these terms with t. That is to say , mixing times are uniformly controlled over a fixed KL ball that does not depend on t ; prior work by contrast makes strong mixing assumptions ( Wang et al. , 2019 ; Xu et al. , 2020 ; Khodadadian et al. , 2021 ) . 3 . High probability guarantee . Though prior work focuses on bounds in expectation , we chose a high probability guarantee to emphasize that the bound does not blow up , despite an arguably more strenuous setting . 4 . Single trajectory . A single trajectory through the MDP is used to remove the option of the algorithm escaping from poor choices with resets ; only the implicit bias can save it . 5 . Rate . To reach a policy which whose value function is ϵ-close to optimal , a trajectory length ( number of samples ) of 1/ϵ16 is sufficient , ignoring log factors ( it suffices to divide both sides of the bound by tθ ( 1−γ ) , set t = 1/ϵ16/3 , and note the trajectory length is tN ) . This is slower than the 1/ϵ6 given in the only other single-trajectory analysis in the literature Khodadadian et al . ( 2021 ) , but by contrast that work makes uniform mixing assumptions ( cf . Khodadadian et al . ( 2021 , Lemma C.1 ) ) , requires the tabular setting , and uses ϵ-greedy for explicit exploration in each iteration . ♢ The proof of Theorem 1.4 and organization of the paper are as follows . After overviews of related work and notation in Sections 1.2 and 1.3 , the first proof component , discussed in Section 2 , is the outer loop of Algorithm 1 , namely the update to the policy πi . As discussed before , this part of the analysis writes the policy update as a mirror descent , and conveniently decouples the suboptimality error into an actor error , which is handled by standard mirror descent tools and provides the implicit bias towards high entropy policies , and a critic error , namely the error of the estimated Q function Q̂i . The second component , discussed in Section 3 , is therefore the TD analysis establishing that the estimate Q̂i is accurate , which not only requires an abstract TD guarantee ( which , as mentioned , is projection-free and run from an arbitrary starting state , unlike prior work ) , but also requires tools to explicitly bound mixing times , rather than assuming mixing times are bounded . This culminates in the proof of Theorem 1.4 , which is sketched at the end of Section 3 , and uses an induction combining the guarantees from the preceding two proof components at all times : it is established inductively that the next policy πi+1 has high entropy and low suboptimality because the previous Q function estimates ( Q̂j ) j≤i were accurate , and simultaneously that the next estimate Q̂i+1 is accurate because πi+1 has high entropy , which guarantees fast mixing . Section 4 concludes with some discussion and open problems , and the appendices contain the full proofs .

Summary: The paper proposes and analyzes the convergence rate and optimality of a natural-actor-critic-like actor critic with linear function approximation in linear MDPs. The most important discovery is that once the policy is within the ball of the maximum entropy optimal policy, it remains there forever provided that the critic is sufficiently accurate. Consequently, one now needs to make ergodicity assumption for only the optimal policy, instead of all policies along the optimization path.