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 provides a novel analysis to the linear MDP setting through an actor critic setup where the policy is softmax-parameterized. The claim is that we can avoid mixing time and exploration based assumptions by using this analysis. This is attractive since its much simpler/cleaner, but also comes with a slower convergence rate. The main steps include bounding the policy within a KL ball of the maximum entropy optimal policy through a mirror descent style analysis, followed by controlling the critic approximation error given the policy is within the KL ball.