Offline Pre-trained Multi-Agent Decision Transformer

1 INRTODUCTION . Multi-Agent reinforcement learning ( MARL ) plays an essential role for solving complex decisionmaking tasks by learning from the interaction data between machine or autonomous agents and simulated physical environments . It typically applies to self-driving ( Shalev-Shwartz et al. , 2016 ) , traffic control ( Bazzan , 2009 ) , and recommendation system ( Xian et al. , 2019 ) , and , to specifically mention , surpasses human beings on some video games ( Bellemare et al. , 2013 ; Wang et al. , 2020 ) . However , the experience-based policy learning scheme requires the algorithms with high sample efficiency because of the limited computing resources and high cost resulting from the data collection ( Haarnoja et al. , 2018 ; Munos et al. , 2016 ; Espeholt et al. , 2019 ; 2018 ) . Furthermore , even in domains where the online environment is feasible , we might still prefer to utilize previously collected data instead , for example , if the domain ’ s complex requires large datasets for effective generalization . In addition , a policy trained on one scenario usually can not perform well on another even though under the same task . Therefore , a universal policy is critical for saving the training time of general reinforcement learning . Offline RL ( Levine et al. , 2020 ) attracts more researchers motivated by improving online reinforcement learning with offline datasets ( Fu et al. , 2020 ) . Those methods like CQL ( Kumar et al. , 2020 ) , BEAR ( Kumar et al. , 2019 ) , and BCQ ( Fujimoto et al. , 2019 ) conventionally leverage the current off-policy RL algorithms , e.g. , DQN ( Mnih et al. , 2015 ) , with provable constraints ( Rashidinejad et al. , 2021 ) . They mainly focus on solving the overestimation of Q values learned in offline datasets and the distribution shift between offline and online . However , those constraints limit the performance with insufficiently offline dataset exploitation ( Kumar et al. , 2019 ) . Recently , the transformer shows its advantage of representing the sequential data for the classification or prediction tasks . Most related to our work , Chen et al . ( 2021 ) replace the constrained Q networks in conventional offline RL algorithms with causal transformer and demonstrate the superiority comparing with baselines such as Behavior Cloning ( BC ) and state-of-the-art offline algorithms in single-agent offline RL . Nevertheless , the overestimation problems emerge in the offline multi-agent RL when one directly applies the sequential modeling method . We propose leveraging the representation capability of the transformer and then load the pre-trained model as part of the policy in the fine-tuning phase . This paper proposes Multi-Agent Decision Transformers ( MADT ) for pre-training the general policy on offline datasets , capable of generalizing onto other seen/unseen environments . Firstly , we collect the offline datasets from the well-known , challenging MARL environment , StarCraft Multi-Agent Challenge ( SMAC ) ( Samvelyan et al. , 2019 ) , interacting with the state-of-the-art algorithm , multiagent PPO ( Yu et al. , 2021 ) . Then we construct transformer variants for multiple agents with parameter sharing or acing as a part . We give the paradigm of pre-training a model with the transformer and fine-tuning it with MARL . We conclude the main contribution as follows : • We propose a series of transformer variants for offline MARL via leveraging the sequential modeling of the attention mechanism . • We derive a framework with offline pretraining and online fine-tuning to improve sample efficiency and generality across different scenarios . • We are the first to apply the pre-trained model across different multi-agent scenarios to validate the transferring ability . Experimental results show fast adaptation and superior performance . • We build a dataset with different skill-level covering all SMAC scenarios for benchmarking progress in MARL pre-training . 2 PRELIMINARY . Multi-Agent Reinforcement Learning . For the Markov Game , which is a multi-agent extension from the Markov Decision Process ( MDP ) , there is a tuple representing the essential elements < S , A , R , P , n , γ > , where S denotes the state space of n agents : S1 × S2 . . . Sn → S. Ai is the action space of each agent i , P : Si × Ai → PD ( Si ) denotes the transition function emiting the distribution over the state space and A is the joint action space , Ri : S × Ai → R is the reward function of each agent and takes action following their policies π ( a|s ) ∈ Πi : S→ PD ( A ) from the policy space Πi , where Πi denotes the policy space of agent i , a ∼ Ai , and s ∼ Si . Each agent aims to maximize its long-term reward ∑ t γ trt , where rti ∈ Ri denotes the reward of agent i in time t and γ denotes the discount factor . Attention-based Model . Attention-based model has shown its stable and strong representation capability . The scale dot-production attention uses the self-attention mechanism demonstrated in ( Vaswani et al. , 2017 ) . Let Q ∈ Rtq×dq be the quries , K ∈ Rtk×dk be the keys , and V ∈ Rtv×dv be the values , where t∗ are the element numbers of different inputs and d∗ are the corresponding element dimensions . Normally , tk = tv and dq = dk . The outputs of self-attention are computed as : Attention ( Q , K , V ) = softmax ( QKT√ dk ) V ( 1 ) where the scalar 1/ √ dk is used to prevent the softmax function into regions that has very small gradients . Then we introduce the multi-head attention process as follows : MultiHead ( Q , K , V ) = Concat ( head1 , . . . , headh ) WO ( 2 ) headi = Attention ( QW Q i , KW K i , VW V i ) ( 3 ) The position-wise feed-forward network is another core module of the transformer . It consists of two linear transformations with a ReLU activation in between . The dimensionality of inputs and outputs is dmodel , and that of the feef forward layer is dff . Specially , FFN ( x ) = max ( 0 , xW1 + b1 ) W2 + b2 ( 4 ) where W1 ∈ Rdmodel×dff and W2 ∈ Rdff×dmodelare the weights , and b1 ∈ Rdff and b2 ∈ Rdmodel are the biases . Across different positions are the same linear transformations . Note that the position encoding for leveraging the order of the sequence as follows : PE ( pos , 2i ) = sin pos/100002i/dmodel PE ( pos , 2i+ 1 ) = cos pos/100002i/dmodel ( 5 ) 3 METHODOLOGY . In this section , we demonstrate how the transformer is applied to our offline pre-training MARL framework . Firstly , we introduce an offline MARL method , in which the transformer maps between the local observations and actions of each agent in the offline dataset via parameter sharing sequentially . Then we leverage the hidden representation as the input of the multi-agent decision transformer ( MADT ) to minimize the cross-entropy loss . Furthermore , we introduce how to integrate the online MARL with MADT in constructing our whole framework to train a universal MARL policy . To accelerate the online learning , we load the pre-trained model as a part of MARL algorithms and learn the policy based on experience in the latest buffer stored from the online environment . To train a universal MARL policy quickly adapting to other tasks , we bridge the gap between different scenarios from observations , actions , and available actions , respectively . Figure 1 overviews our method from the perspective of offline pre-training with supervised learning and online fine-tuning with MARL algorithms . 3.1 MULTI-AGENT DECISION TRANSFORMER . Algorithm 1 shows the training process of Multi-Agent Decision Transformer ( MADT ) , in which we autoregressively encode the trajectories from the offline datasets in offline pre-trained MARL and train the transformer-based network with supervised learning . We carefully reformulate the trajectories as the inputs of the causal transformer different from those in the Decision Transformer ( Chen et al. , 2021 ) . Denote that we deprecate the reward-to-go and actions that are encoded with states together in the single-agent DT . We will interpret the reason for this in the next section . Similar to the seq2seq models , MADT is based on the autoregressive architecture with the reformulated sequential inputs across timescale . The left part of Figure 2 shows the architecture . The causal transformer encodes the agent i ’ s trajectory sequence τ ti at the time step t to a hidden representation h t i = ( h1 , h2 , . . . , hl ) with a dynamic mask . Given ht , the output at the time step t only based on the previous data then consumes the previously emitted actions as additional inputs when predicting a new action . Algorithm 1 : MADT-Offline : Multi-Agent Decision Transformer Input : Offline dataset D : { τi : 〈sti , oti , ati , vti , dti , rti〉Tt=1 } ni=1 , vti denotes the available action Initialize θ for the Causal Transformer ; Initialize α as the learning rate , C as the context length , and n as the maximize agent number for τ = { τ1 , . . . , τi , . . . , τn } in D do Chunk the trajectory into τi = { rti , sti , ati } t∈1 : C as the ground truth samples , where C is the context length , and mask the trajectory when dti is true for τ ti = { τ1i . . . τ ti . . . τTi } in τi do Mask illegal actions via P ( aj |τi , â < j ; θ ′ ) = 0 if vti is True Predict the action ât = arg maxa P ( a|τi , â < t ; θ ′ ) Update θ with θ = arg max θ′ 1 C ∑C j=1 P ( aj ) logP ( âj |τi , â < j ; θ ′ ) end end return θ Trajectories Reformulation as Input . We model the lowest granularity at each time step as a modeling unit xt from the static offline dataset for the concise representation . MARL has many elements such as 〈global_state , local_observation〉 , different from the single agent . It is reasonable for sequential modeling methods to model them in Markov Decision Process ( MDP ) . Therefore , we formulate the trajectory as follows : τ i = ( x1 , . . . xt . . . , xT ) where xt = ( st , oit , a i t ) where sit denotes the global shared state , o i t denotes the individual observation for agent i at time step t , and ait denotes the action . We take xt as a token similar to the input of natural language processing . Output Sequence Constructing . To bridge the gap between training with the whole context trajectory and testing with only previous data , we mask the context data to autoregressively output in the time step t with previous data in 〈1 . . . t− 1〉 . Therefore , MADT predicts the sequential actions at each time step using the decoder as follows : y = at = arg max a pθ ( a|τ , a1 , . . . , at−1 ) ( 6 ) where θ denotes the parameter of MADT and τ denotes the trajectory including the global state s , local observation o before the time step t , pθ is the distribution over the legal action space under the available action v. Core Module Description . MADT differs from the transformers in conventional sequence modeling tasks that take inputs with position encoding and decode the encoded hidden representation autoregressively . We use the masking mechanism with a lower triangular matrix to compute the attention : Attention ( Q , K , V ) = softmax ( QKT√ dk +M ) V ( 7 ) where M is the mask matrix which ensures the input at the time step t can only correlate with the input from 〈1 , . . . , t− 1〉 . We employ the Cross-entropy ( CE ) as the total sequential prediction loss and utilize the available action v to ensure agents taking those illegal actions with probability zero . The CE loss can be represented as follows : LCE ( θ ) = 1 C C∑ t=1 P ( at ) logP ( ât|τt , â < t ; θ ) , where C is the context length ( 8 ) where at is the ground truth action , τt includes { s1 : t , o1 : t } . â denotes the output of MADT . The cross-entropy loss shown above aims to minimize the distribution distance between the prediction and the ground truth .

Summary: The paper studies the offline training and online fine-tuning setting for MARL, and show that it is promising to learn policy representatrions to improve the performance on downstream MARL tasks. The paper proposes to integrate the transformer architecture to improve efficiency and generalization ability. Authors conduct extensive experiments to validate its effectiveness.