MAP Propagation Algorithm: Faster Learning with a Team of Reinforcement Learning Agents

1 Introduction . Error backpropagation algorithm ( backprop ) [ 1 ] efficiently computes the gradient of an objective function with respect to parameters , by iterating backward from the last layer of a multi-layer artificial neural network ( ANN ) . However , backprop is generally regarded as being biologically implausible [ 2 , 3 , 4 , 5 , 6 , 7 ] . First , the learning rule given by backprop is non-local , as it relies on information other than input and output of a neuron-like unit computed in the feedforward phase ; while biologically-observed synaptic plasticity depends mostly on local information ( e.g . spike-timingdependent plasticity ( STDP ) [ 8 ] ) and possibly some global signals ( e.g . reward-modulated spiketiming-dependent plasticity ( R-STDP ) [ 8 , 9 , 10 ] ) . Second , backprop requires precise coordination between feedforward and feedback connections , because the feedforward value has to be retained until error signals arrive ; while it is unclear how a biological system can coordinate an entire network to alternate between feedforward and feedback phases precisely . Third , backprop requires synaptic symmetry in the forward and backward paths , rendering it biologically implausible . Nonetheless , recent work has demonstrated that this symmetry may not be necessary for backprop due to the ‘ feedback alignment ’ phenomenon [ 11 , 12 , 13 ] . Alternatively , REINFORCE [ 14 ] could be applied to all units in the network to train an ANN as a more biologically plausible way of learning . It is shown that the learning rule gives an unbiased estimate of the gradient of return [ 14 ] . Another interpretation of this relates to viewing each unit as a reinforcement learning ( RL ) agent , with each agent trying to maximize the global reward . Such a team of agents is also known as a coagent network [ 15 ] . However , coagent networks can only solve 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . simple tasks due to the high variance associated with the learning rule and thus the low speed of learning . The high variance stems from the lack of structural credit assignment , i.e . a single scalar reward is used to evaluate the action of all agents in the network . To address this high variance associated with REINFORCE , we propose a novel algorithm that significantly reduces the variance while retaining the local property of the learning rule . We call this newly proposed algorithm maximum a posteriori ( MAP ) propagation . Essentially , MAP propagation replaces the hidden units ’ values with their MAP estimates conditioned on the action chosen , or equivalently , minimizes the energy function of the network , before applying REINFORCE . We prove that for a network with normally distributed hidden units , by minimizing the energy function of the network , the parameter update given by REINFORCE and backprop ( with the reparametrization trick ) becomes the same , thus establishing a connection between REINFORCE and backprop . Our experiments show that a team of agents trained with MAP propagation can learn much faster than REINFORCE , such that the team of agents can solve common RL tasks at a similar ( or higher ) speed compared to an ANN trained by backprop , as well as exhibiting sophisticated exploration that differs from an ANN trained by backprop . The novel MAP propagation algorithm represents a new class of algorithm to train an ANN that is more biologically plausible than backprop and maintains a comparable learning speed to backprop at the same time . Our work also opens the prospect of the broader application of teams of agents , called coagent networks [ 15 ] , in deep RL . 2 Background and Notation . We consider a Markov Decision Process ( MDP ) defined by a tuple ( S , A , P , R , γ , d0 ) , where S is a finite set of states of an agent ’ s environment ( although this work can be extended to the infinite state case ) , A is a finite set of actions , P : S × A × S → [ 0 , 1 ] is a transition function giving the dynamics of the environment , R : S × A → R is a reward function , γ ∈ [ 0 , 1 ] is a discount factor , and d0 : S → [ 0 , 1 ] is an initial state distribution . Denoting the state , action , and reward signal at time t by St , At , and Rt respectively , P ( s , a , s′ ) = Pr ( St+1 = s′|St = s , At = a ) , R ( s , a ) = E [ Rt|St = s , At = a ] , and d0 ( s ) = Pr ( S0 = s ) , where P and d0 are valid probability mass functions . An episode is a sequence of states , actions , and rewards , starting from t = 0 and continuing until reaching the terminal state . For any learning methods , we can measure its performance as it improves with experience over multiple episodes , which makes up a run . Letting Gt = ∑∞ k=t γ k−tRk denote the infinite-horizon discounted return accrued after acting at time t , we are interested in finding , or approximating , a policy π : S ×A → [ 0 , 1 ] such that for any time t ≥ 0 , selecting actions according to π ( s , a ) = Pr ( At = a|St = s ) maximizes the expected return E [ Gt|π ] . The value function for policy π is V π where for all s ∈ S , V π ( s ) = E [ Gt|St = s , π ] , which can be shown to be independent of t for the infinite-horizon case we are considering . Here we restrict attention to policies computed by multi-layer networks consisting of L layers of stochastic units . Let H lt ∈ Rn ( l ) denote the activation values of the units in layer l at time t and n ( l ) denote the number of units in layer l. For any t ≥ 0 , we also let H0t = St , HLt = At , and Ht = { H1t , H2t , ... , HL−1t } . We call any elements in Ht a hidden layer and HLt the output layer . The distribution of H lt conditional on H l−1 t is given by πl : Rn ( l−1 ) × Rn ( l ) → [ 0 , 1 ] , such that for any t ≥ 0 , πl ( hl−1 , hl ; W l ) = Pr ( H lt = hl|H l−1t = hl−1 ; W l ) , where W l is the parameter of layer l. We also denote all parameters of the network as W = { W 1 , W 2 , ... , WL } . To sample an action At from the network , we iteratively sample H lt ∼ πl ( H l−1t , · ; W l ) from l = 1 to L. We say a layer l is normally distributed if πl ( H l−1t , · ; W l ) = N ( gl ( H l−1t ; W l ) , σ2l ) , the normal distribution with mean gl ( H l−1t ; W l ) , where gl : Rn ( l−1 ) → Rn ( l ) is a differentiable function , and a fixed standard deviation σl . A common choice of g is a linear transformation followed by an activation function ; that is , g ( H l−1t ; W l ) = f ( W lH l−1t ) where f is a non-linear activation function such as softplus or rectified linear unit ( ReLU ) and W l ∈ Rn ( l ) ×n ( l−1 ) . We also define the energy function E : Rn ( 1 ) × Rn ( 2 ) × ... × Rn ( L−1 ) → [ 0 , ∞ ) as E ( h ; s , a ) = − log Pr ( Ht = h|St = s , At = a ) , which can be shown to be independent of t. The case we consider here is one in which all the units of the network implement an RL algorithm and share the same reward signal . These networks can therefore be considered as teams of agents ( agent here refers to an RL agent [ 16 ] ) , which have also been called coagent networks [ 15 ] . We denote ∇xf as the gradient of f w.r.t . x , AT as the transpose of matrix A , and ∇Af ( Pr ( A ) ) as the shorthand for ∇af ( Pr ( A = a ) ) . For a random variable X with a distribution that depends on parameter W and a random variable Y , we call h ( Z ; W , Y ) the re-parameterization of X by Z conditioned on Y if h ( Z ; W , Y ) and X have the same conditional distribution ; that is , Pr ( h ( Z ; W , Y ) = x|Y = y ) = Pr ( X = x|Y = y ; W ) for all values of x , y and W , where Z is an independent random variable with a distribution that does not depend on parameter W and h is an invertible and differentiable function . In case X has a multi-layer structure , we denote hl ( Z ; W , Y ) as the lth layer in h ( Z ; W , Y ) and h−1 as the inverse of h. In general , we use the superscript l to denote the lth layer in a variable if the variable has a multi-layer structure . Also , for all distributions discussed in this paper , the probability mass function is replaced by probability density function if the random variable is continuous . 3 Algorithm . 3.1 MAP Propagation . MAP propagation is based on REINFORCE applied to each hidden unit with the same global reinforcement signal . To reduce the variance associated with the learning rule , we note that this variance can be reduced by using the expected parameter update conditioned on the state and the selected action instead . However , this expected parameter update is generally intractable to compute analytically . Therefore , we propose to use the maximum a posteriori ( MAP ) estimate to approximate the expected parameter update . This makes the resulting learning rule biased but reduces the variance significantly . The details of MAP propagation are as below . The gradient of return with respect to W l ( where l ∈ { 1 , 2 , ... , L } in all discussion below unless stated otherwise ) can be estimated by REINFORCE , also known as likelihood ratio estimators : ∇W l E [ Gt ] = ∞∑ k=t γ ( k−t ) E [ Gk∇W l log Pr ( Ak|Sk ) ] . ( 1 ) We can show that the terms in the summation of ( 1 ) also equal E [ Gk∇W l log πl ( H l−1k , H lk ; W l ) ] , which is the REINFORCE learning rule applied to each hidden unit with the same global reinforcement signal Gk : Theorem 1 . Let the policy be a multi-layer network of stochastic units as defined in Section 2 . For any t ≥ 0 and l ∈ { 1 , 2 , ... , L } , E [ Gt∇W l logPr ( At|St ; W ) ] = E [ Gt∇W l log πl ( H l−1t , H lt ; W l ) ] . ( 2 ) The proof is in Appendix B.1 . Note that this theorem is also proved by Williams [ 14 ] . This shows that we can apply REINFORCE to each unit of the network , and the learning rule still gives an unbiased estimate of the gradient of the return . Therefore , denoting α as the step size , we can update parameters by the following stochastic gradient ascent rule : W l ←W l + αGt∇W l log πl ( H l−1t , H lt ; W l ) . ( 3 ) However , this learning rule suffers from high variance since a single reward , which results from the exploration of all units , is used to evaluate actions of all units , making the learning rule scale poorly with the number of units in the network . To reduce the variance , we notice that we can replace ∇W l log πl ( H l−1t , H lt ; W l ) in learning rule ( 3 ) by E [ ∇W l log πl ( H l−1t , H lt ; W l ) |St , At ] , noting that ( see Appendix B.1 for the details ; at the R.H.S. , the outer expectation is taken over St , At and Gt , while the inner expectaion is taken over H l−1t and H l t ) : E [ Gt∇W l log πl ( H l−1t , H lt ; W l ) ] = E [ Gt E [ ∇W l log πl ( H l−1t , H lt ; W l ) |St , At ] ] . ( 4 ) This can reduce variance since the variance associated with the stochastic noise of hidden units is removed in the learning rule ( see Appendix B.6 for the proof ) . The learning rule now becomes : W l ←W l + αGt E [ ∇W l log πl ( H l−1t , H lt ; W l ) |St , At ] . ( 5 ) Since ( 3 ) is following gradient of return in expectation , and the expected update value of ( 3 ) and ( 5 ) is the same , we conclude that ( 5 ) is also a valid learning rule as it follows the gradient of return in expectation . However , we note that E [ ∇W l log πl ( H l−1t , H lt ; W l ) |St , At ] in ( 5 ) is generally intractable to compute analytically . Instead , we propose to use the MAP estimate to approximate this term:1 E [ ∇W l log πl ( H l−1t , H lt ; W l ) |St , At ] ≈ ∇W l log πl ( Ĥ l−1t , Ĥ lt ; W l ) , ( 6 ) where Ĥt = argmaxht Pr ( Ht = ht|St , At ) . There are many methods to approximate Ĥt , such as hill-climbing methods . In case of hidden units being continuous , we can approximate Ĥt by running gradient ascent on log Pr ( Ht|St , At ) as a function of Ht for fixed St and At , such that Ht approaches Ĥt . Ht can be initialized as the value sampled from the network when sampling action At . To be concrete , before applying learning rule ( 3 ) , we first run gradient ascent on Ht for N steps : Ht ← Ht + α∇Ht log Pr ( Ht|St , At ) . ( 7 ) For l ∈ { 1 , 2 , ... , L− 1 } , this is equivalent to ( see Appendix B.4 for the details ) : H lt ← H lt + α ( ∇Hlt log πl+1 ( H l t , H l+1 t ; W l+1 ) +∇Hlt log πl ( H l−1 t , H l t ; W l ) ) . ( 8 ) This update rule maximizes the probability of hidden units ’ value given the state and the action with respect to hidden units ’ value . Based on the definition of energy functions in Section 2 , the update rule of hidden units can also be seen as minimizing the energy function E ( Ht ; St , At ) [ 17 ] . After updating Ht for N steps by ( 8 ) , we obtain an estimate of Ĥt , denoted as H̃t , and apply the following learning rule to learn the parameters of network : W l ←W l + αGt∇W l log πl ( H̃ l−1t , H̃ lt ; W l ) . ( 9 ) We call the algorithm that applies REINFORCE after replacing the value of hidden units by their approximated MAP estimate MAP propagation . The pseudo-code of MAP propagation with gradient ascent to approximate Ĥt can be found in Algorithm 1 in Appendix A . Note that N = 0 recovers the special case of pure REINFORCE . Similar to actor-critic networks [ 16 ] , we can also train a critic network to estimate the state-value , V π ( St ) , so Gt in ( 9 ) can be replaced by TD error δt = Rt + γV π ( St+1 ) − V π ( St ) and the whole algorithm can be implemented online . To better facilitate temporal credit assignment , we can also use eligibility traces to replace the gradient in ( 9 ) , using the same idea of actor-critic networks with eligibility trace [ 16 ] . The pseudo-code of it can be found in Algorithm 2 in Appendix A . A team of agents can be trained by MAP propagation to estimate the state-value , such that a separate team of agents can fulfill the role of a critic network , and the whole actor-critic network can be trained without backprop . A simple way to achieve this is to convert the estimation of state-value to an RL task , but this conversion is inefficient since the information of optimal actions is lost ( the agent only knows a scalar reward but not the target output ) . Appendix C proposes a new learning rule to train a team of agents to estimate the state-value by MAP propagation efficiently based on the information of optimal actions . Essentially , MAP propagation is equivalent to applying REINFORCE after minimizing the energy function . As there are many studies on the biological plausibility of REINFORCE , we refer readers to the book of Sutton and Barto [ 16 ] for a review and discussion of the connection between REINFORCE and neuroscience . The main difference between MAP propagation and REINFORCE is the minimization of the energy function given by the update rule ( 8 ) . This update rule is local as it only depends on the units one layer above and below based on feedforward and feedback connections . There is much evidence that feedback signals in brains alter neural activity [ 7 , 18 ] , supporting the use of feedback connections in MAP propagation . The update rule can also be performed in parallel for all layers , removing the need for precise coordination between feedforward and feedback connections as in backprop . However , the update rule requires the feedback weight to be symmetric of the feedforward weight and different values to be propagated through feedforward and feedback connections . MAP propagation fits well into the recently proposed NGRAD hypothesis [ 7 ] , which hypothesizes that the cortex uses differences in activity states to drive learning . The main idea of NGRAD is that “ higherlevel activities can nudge lower-level activities towards values that are more consistent with the 1We let Ĥ0t = H0t = St , Ĥlt denote the lth layer in Ĥt for l ∈ { 1 , 2 , ... , L− 1 } , and ĤLt = HLt = At . higher-level activity ” , which also describes the process of energy minimization in MAP propagation . A detailed discussion of the biological plausibility of MAP propagation and its relationship with the NGRAD hypothesis can be found in Appendix F .

Summary: The paper introduces a new synapse-local method for reinforcement learning, that is similar to the REINFORCE algorithm, but reduces the high variance caused by REINFORCE's need to assess many stochastic samples, at the cost of added computation.   IIUC, it does this by 1- making REINFORCE layer-local (Eq 2), 2- Replacing the gradient of probabilities of observed responses  by its expectation given the state and action (which is uncomputable) (Eq 4-5), and 3- Approximating this uncomputable expectation by replacing actual activations with their maximum a posteriori estimate, given the chosen action and observed state. The latter turns out to be computable, with some iterative process.  The algorithm is also shown to be equivalent, update-wise, to plain backpropagation with the reparametrization trick (this is how I understood Theorem 2 and Eq. 11, but I may be wrong, see below).   Various experiments show that this method learns faster than REINFORCE and about as fast as backprop with the reparametrization trick, in terms of number of steps (though each step is computationally more intensive).