Meta Attention For Off-Policy Actor-Critic

1 INTRODUCTION . Reinforcement Learning ( RL ) algorithms based on the Actor-Critic framework have achieved considerable success in many areas such as games , robot control , and planning . Compared with onpolicy methods , off-policy methods possess more efficient sampling since they do not require new samples to be collected for each gradient step and make better use of experience ( Haarnoja et al. , 2018a ) . However , even for off-policy methods , traditional reinforcement learning algorithms still have extremely low sample efficiency ( Yu , 2018 ) . Recently , meta-learning ( Hospedales et al. , 2020 ) has become topical as a paradigm to accelerate RL by learning an inductive bias from past experience . By learning aspects of the learning strategy , such as fast adaptation strategies ( Finn et al. , 2017 ) ( Rakelly et al. , 2019 ) , losses ( Zhou et al. , 2020 ) ( Bechtle et al. , 2020 ) , optimization strategies ( Duan et al. , 2016 ) , exploration strategies ( Gupta et al. , 2018 ) , hyperparameters ( Bechtle et al. , 2020 ) , and intrinsic rewards ( Zheng et al. , 2018 ) , meta-learning has significantly improved sample efficiency over standard RL algorithms . Improving sample efficiency through attention mechanism has also been proved to be very effective in image-based reinforcement learning ( Barati & Chen , 2019 ; Chen et al. , 2019 ) . The application of attention mechanism in multi-agent system ( Parnika et al. , 2021 ; Iqbal & Sha , 2019 ) and multi-object task ( Team et al. , 2021 ) also shows its powerful capabilities in information processing . However , in the existing works , attention mechanisms often have clear application scenarios , such as imagebased control or multiple sources of information ( multi-agent or multi-target ) . The effective combination of attention mechanisms with current algorithms in a state-based single-agent environment is still a problem to be investigated In this paper , we propose a meta attention method based on the attention mechanism . In the human decision-making process , people often modify their concepts based on the feedback and results to obtain a better decision . Inspired by this decision-making process , we use meta attention to adapt the features generated by the policy network based on the evaluation of value network . Our work differs from current attention-based work in that our meta attention approach works only within the Actor-Critic framework and does not depend on a specific scenario . We formalize the metalearning process as a bi-level optimization problem . Our approach can be flexibly combined with various algorithms by using meta attention as a meta-learner and optimizing meta attention in the outer layer . Unlike the existing meta-learning methods , our meta attention approach can improve the performance of agent through gradients in the training stage and obtain better actions by adjusting the features in the execution stage . We evaluated the proposed meta attention method in a series of continuous control tasks in Gym and Roboschool , including three 3D robot control tasks , two 2D control tasks , and one classic control task based on DDPG ( Lillicrap et al. , 2016 ) , TD3 ( Fujimoto et al. , 2018 ) and SAC ( Haarnoja et al. , 2018b ) . Besides , we also discussed the changes and impact caused by modifying the actor features through meta attention . Experimental results how that our meta-attention approach is not only effective in accelerating the learning progress of the agent in the training phase , but also improves the actions in the execution phase , further enhancing the performance of the agent . 2 RELATED WORK . Attention Mechanism Attention is a behavioral and cognitive process of selectively attending to a discrete aspect of information , whether subjective or objective , while ignoring other perceptible information ( de Santana Correia & Colombini , 2021 ) . Typically , attention mechanism is mainly applied in the fields of computational vision and natural language processing . Accordingly , although the implementation methods are different , the application of attention mechanism in reinforcement learning is mainly focused on video games ( Wu et al. , 2021 ; Chen et al. , 2019 ; Barati & Chen , 2019 ; Mott et al. , 2019 ) ( Manchin et al. , 2019 ) . Other works such as Peng et al . ( 2020 ) proposed a dynamic attention model with a dynamic encoder-decoder architecture , which dynamically explores node features to efficiently exploit hidden structural information at different construction steps . Li et al . ( 2021 ) applied the attention mechanism to generate feature vectors and input them into value and policy head during the feature extraction phase of PPO and PPG.Jiang & Lu ( 2018 ) ; Iqbal & Sha ( 2019 ) ; Mao et al . ( 2019 ) employed a multi-head attention mechanism to make one agent selectively pay attention to information from other agents . Team et al . ( 2021 ) used the attention mechanism to match multiple goals with the hidden state of the current state to obtain the goal-attention hidden state under different goals . This approach allows the agent to predict the expected return obtained by attending to a goal at the end of a scene . Meta Reinforcement Learning Meta-learning is most often understood as learning to learn , which refers to learning from historical information or multiple learning episodes to improve learning algorithms . Since works which fed historical trajectories into Recurrent Neural Network to flesh out task-level information ( Wang et al. , 2016 ; Duan et al. , 2016 ) , various meta-learning methods have been proposed to strengthen agent performance . Houthooft et al . ( 2018 ) ; Kirsch et al . ( 2020 ) ; Zhou et al . ( 2020 ) used meta-learning methods to learn a loss function rather than artificial design to improve performance of agent in single or multiple tasks.Gupta et al . ( 2018 ) ; Stadie et al . ( 2018 ) ; Xu et al . ( 2018a ) employed meta-learning methods to learn exploration instead of traditional exploration methods.Finn et al . ( 2017 ) meta learned a good model initialization of the model that can be quickly adapted to different tasks . Xu et al . ( 2018b ) improved the performance of the agent by meta-learning discount factors.Rakelly et al . ( 2019 ) Fakoor et al . ( 2020 ) treats meta information as an unobservable state of Partially Observable Markov Decision Process , further improves the agent ’ s performance in multi task learning . Although the attention mechanism may not have an explicit meta-learning object , it can also be considered as the king of the meta-learning method . Bi-level Optimization Generally , meta-learning can be formalized as a bi-level optimization ( BLO ) problem . However , the solution of BLO problems is often challenging . Franceschi et al . ( 2018 ) proposed a framework for approximating the solution of BLO problems using the gradient method . Since then , many works based on the gradient methods to optimize meta leaner have successfully proved the feasibility of gradient method , such asLi et al . ( 2019 ) ; Finn & Levine ( 2018 ) ; Lian et al . ( 2020 ) ; Flennerhag et al . ( 2020 ) . For the RL problem , Zhou et al . ( 2020 ) optimize the meta critic , the upper-level of a BLO , as an intrinsic motivation by gradient method . Kirsch et al . ( 2020 ) enables a population of agents to use and improve a single parameterized objective function through gradient learning on different tasks . Rajeswaran et al . ( 2019 ) proposed an implicit MAML algorithm by drawing upon implicit differentiation that effectively decouples the meta-gradient computation from the selection of an inner loop optimizer . Liu et al . ( 2019 ) proposed a surrogate objective function TMAML , which incorporates control variables into gradient estimation through automatic differentiation and improves the quality of the gradient estimation by reducing the variance without introducing bias . 3 METHODOLOGY . 3.1 OFF-POLICY ACTOR-CRITIC . In general , the reinforcement learning task can be considered as finding the optimal policies in Markov Decision Processes ( MDPs ) . The MDP is defined by a tuple ( S , A , P , R ) , where S is a set of states , A is a set of actions , P is a set of probabilities to switch from a state s to s′ for a given set of action a , and R : S × A → R is a scalar reward function . In the Actor-Critic framework , the policy network ( Actor ) πφ ( s ) and the value network ( Critic ) Qθ ( s , a ) are parameterized by a neural network respectively . At each time t , the agent receives an observation st and takes a action at based on its policy π : S → A , then receives a reward rt and a new state st+1 . The tuple ( st , at , rt , st+1 ) describes a state transition and will be stored in a reply buffer D for off-policy learning . The objective of RL is to find the optimal policy πφ to maximizes the expected cumulative return J : J ( φ ) = E [ ∞∑ t=0 γtR ( st , at ) | at ∼ πφ ( · | st ) ] ( 1 ) Where γ is the discount factor , and J ( φ ) also can be written as the expected value for the Q-function : J ( φ ) = Es∼pπQθ ( s , a ) |at∼πφ ( ·|st ) , ( 2 ) Where Q-function is the expected discounted sum of rewards following visitation at state s and execution of action a , and pπ is the state distribution induced by policy π . For off-policy ActorCritic architectures such as DDPG , TD3 and SAC , the loss for the actor provided by the critic may be different , but the Q-function will be learned by minimizing the loss of the below equation as same : LCriticθ = Est∼pπ , at∼πφ [ ( Qθ ( st , at ) − yt ) 2 ] yt = rt + γQθ′ ( st+1 , πφ′ ( st+1 ) ) , ( 3 ) where φ′ and θ′ represent the target network for the critic and the actor respectively . The actor loss usually differs in details according to different algorithms , however , they all follow a form of the following formula : LActorφ = −J ( φ ) = − Es∼pπQθ ( s , a ) |a=πφ ( s ) ( 4 ) 3.2 META ATTENTION METHOD . The attention mechanism has been widely used in image recognition and natural language processing since it can model the human pattern recognition process . It allocates attention to the important part of information while automatically ignores low-value features . According to ( Vaswani et al. , 2017 ) , attention function can be described as mapping a query and a set of key-value pairs to an output , where the output is computed as a weighted sum of the values , with the weight assigned to each value being computed by a compatibility function of the query with the corresponding key : Attention ( Query , Key , Value ) = ∑ i Similarity ( Queryi , Keyi ) ∗ Value i In order to introduce the attention mechanism into the Actor-Critic framework , we first split the actor and critic network into two parts . We take the last layer of actor network π̂ as an action net to produce actions and the rest as actor feature net π̄ ( s ) to extract features , with the entire policy network being denoted as πφ ( s ) = π̂ ( π̄ ( s ) ) . Similarly , the value network can also be divided into a critic net Q̂ and a critic feature net Q̄ ( s ) , and the whole value network can be expressed as Qθ ( s ) = Q̂ ( Q̄ ( s , a ) ) . We use the feature net as encoder to obtain Query , Key and Value by the following formula : Querry = π̄ ( s ) Key = Q̄ ( s , πφ ( s ) ) Value = π̄ ( s ) We input Query ( actor features ) and Key ( critic features ) into meta attention network fψ ( Query , Key ) ( a three-layers MLP parameterized by ψ ) to calculate the similarity of each feature dimension . To enhance or reduce the features in specific corresponding dimensions , we multiply the output after the sigmoid function by 2 to obtain the feature scale . In this paper , we use ◦ to denote the Hadamard product , and by calculating the Hadamard product ◦ of Value ( actor features ) and scale , we obtain the attention features π̄′ ( s ) after modifying some dimensions : Attention Features = π̄′ψ ( s ) = fψ ( π̄ ( s ) , Q̄ ( s , πφ ( s ) ) ) ◦ 2 ◦ π̄ ( s ) ( 5 ) This calculation process corresponds to the part in the green dotted line box in Figure 1 ( b ) . There are also two critical problems in the optimization process . 1 ) how to affect the agent ’ s decision-making and training process through attention features ; and 2 ) how to optimize the meta attention network to achieve the correct matching Query ( actor features ) and Key ( critic features ) to generate the proper feature scales . To address these issues , we formalize the entire optimization process as a bi-level optimization problem , referring to the meta attention method as the outer level , and to the task performed by the agent as the inner level : ψ∗ = arg min ψ LAttentionψ ( D ; φ∗ ; ψ ) ( 6 ) s.t . φ∗ = arg min φ [ LActorφ ( D ; φ|a ∼ πφ ( s ) ) + LActorφ ( D ; φ ; ψ|a ∼ π̂ ( π̄′ψ ( s ) ) ] , ( 7 ) where LAttentionψ is the meta optimization objective , and L Actor φ is the actor loss in the form of Eq. ( 4 ) . For the first problem , we first implement a traditional back-propagate dominated by the actor loss LActorφ |a=πφ ( s ) on training data dtrn : φold = φ− η ∂LActorφ ( dtrn|a ∼ πφ ( s ) ) ∂φ , ( 8 ) Where η is the learning rate , this process corresponds to the upper part of Figure 1 ( b ) and the first step of meta-training in Algorithm1 . Then , we generate a new attention action a = π̂ ( π̄′ψ ( s ) ) from the attention features through action net and feed into the critic to get the loss LActorφ ( D ; φ ; ψ ) |a=π̂ ( π̄′ψ ( s ) ) for back-propagation : φnew = φold − η ∂LActorφold ( dtrn ; ψ|a ∼ π̂ ( π̄′ψ ( s ) ) ) ∂φ ( 9 ) This allows the agent to obtain better actions by simply modifying the features compared to the original action without increasing the batch size , strengthening the agent ’ s learning of good actions and reducing the probability of producing poor actions after the gradient step . This process corresponds to the lower part of ( b ) in Figure 1 and the second step of meta-training in Algorithm1 . For the second problem , we have two basic assumptions : 1 ) a good meta attention network will inevitably generate feature scales that can generate actions with higher value since it better associates the relationship between actor features and critic features ; 2 ) in the process of back-propagation , high-value actions enhance actors more than low-value actions because good actions will strengthen the actor ’ s tendency to make such choices , while bad actions cause the actor to try other actions . Furthermore , since the meta attention network is involved in the back propagation process , we use the utility of meta attention in this process as the attention loss LAttentionψ and the attention loss on validation data dval is defined as : LAttentionψ = tanh ( LActorφnew ( dval|a ∼ πφnew ( s ) ) − L Actor φold ( dval|a ∼ πφold ( s ) ) ) ( 10 ) Under this definition , when performing gradient descent updates , it is ensured that the meta attention is always updated along the direction that improves the performance of the agent . This process corresponds to the blue line part of ( b ) in Figure 1 and the meta-test in Algorithm1 .

Summary: This paper proposes an attention-based actor-critic agent. The authors propose to parameterise the actor and the critic with two separate neural networks. Their algorithm operates in a two-stage fashion. In the first stage, the algorithm is a standard actor critic that produces an action probability distribution and a prediction of the action-value function. In the second state, the algorithm attends over the policy's and value function's hidden features. This output of the attention mechanism is a new feature representation that is fed into the policy's and value function's final layer, to produce an alternative action probability distribution and action value-function. The agent then takes the action that has the highest predicted return. The authors evaluate their proposed algorithm under DDPG and SAC on standard continuous control tasks and demonstrate that it can improve performance.