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: This is the first paper to pre-train a multi-agent model with offline datasets and online fine-tune it to a particular task. The authors propose MADT an offline pre-training architecture for multi-agent RL integrating transformers with imitation learning on an a collected data set of offline data. Towards this end, the authors contribute a data set of offline multi-agent interactions collected from running multi-agent PPO on 5 tasks from SMAC. On this data set, the proposed MADT method is shown to outperform popular single agent offline RL algorithms from the literature and excel at fine-tuning a pre-trained offline model.