Offline Multi-Agent Reinforcement Learning with Knowledge Distillation
Anonymous Author(s)
AffiliationAddress
Abstract
We introduce an offline multi-agent reinforcement learning (offline MARL) framework that utilizes previously collected data without additional online data collection.
Our method reformulates offline MARL as a sequence modeling problem and thusbuilds on top of the simplicity and scalability of the Transformer architecture.
Inthe fashion of centralized training and decentralized execution, we propose to firsttrain a teacher policy who has the privilege to access every agent’s observations,actions, and rewards.
After the teacher policy has identified and recombined the"good" behavior in the dataset, we create separate student policies and distill notonly the teacher policy’s features but also its structural relations among differentagents’ features to student policies.
We show that our framework significantlyimproves performances on a range of tasks and outperforms state-of-the-art offline
MARL baselines.
Furthermore, we demonstrate that the proposed method has abetter convergence rate, is more sample efficient, and is more robust to variousdemonstration qualities compared with baselines.
Introduction
The online learning paradigm assumed by existing multi-agent reinforcement learning (MARL)algorithms is one of the biggest obstacles to their widespread adoption.
To apply MARL, one has torepeatedly perform the following two steps: (a) collect experiences by deploying multiple agents,typically with their latest learned policies and (b) use the collected experiences to improve thepolicies.
In many scenarios, frequently performing the first step is impractical because deployingmulti-agents to the environment can be expensive and dangerous (e g., self-driving cars).
Therefore,devising offline MARL algorithms that can simply learn from previously collected datasets withoutinteraction with the environment is an important step toward solving real-world problems.
Recently, several works have extended offline RL algorithms under the single-agent setting to offline
MARL [45].
These works focus on addressing the distribution shift issue that causes the values ofunseen state-action pairs to be erroneously estimated.
However, these methods are based on temporaldifference (TD) learning and thus require bootstrapping for credit assignment, a problem that isespecially challenging under the multi-agent setting as the interactions between the agents and theenvironment can be highly complex.
To mitigate these issues, we explore the possibility to transform offline MARL into a sequencemodeling problem.
This paradigm shift, first introduced by Decision Transformer [3] under thesingle-agent setting, allows us to bypass bootstrapping and perform credit assignment directly viaself-attention.
A concurrent work, MADT [26], proposes to adapt Decision Transformer to themulti-agent setting by (a) sharing model parameters across agents and (b) attaching one-hot agent
IDs to observations.
Different from MADT, we propose a framework based on policy distillation thatachieves better performance (as demonstrated in the experiments).
[Figure]
Two agents are tasked to explore as many blocks as possible.
Explored blocks are colored g walk
Our method first trains a teacher policy (instantiated as a decision transformer) to model the entireoffline MARL dataset sequentially by accessing each agent’s observation, action, and reward.
Theprivilege to access every agent’s information helps the teacher policy understand the underlyinginteraction across agents and predict actions that encourage cooperative behavior.
However, theteacher policy cannot be deployed at test time in the fashion of decentralized execution.
To addressthis issue, we initialize a separate student policy for each agent and distill the knowledge of theteacher policy into the student policy.
In addition to allowing decentralized execution, this step ishelpful because (a) the teacher policy can identify and recombine "good" behavior in a suboptimaldataset [3], (b) the centralized teacher allows credit assignment across agents via self-attention, and(c) distillation from a privileged model (teacher) provides richer and more stable learning signal tothe student policy [2].
Furthermore, we propose a novel distillation objective that transfers structuralrelations of student policies’ features rather than actual values of individual features.
In our empiricalresults, we show that this objective is complementary to the classical policy distillation and helpsus outperform state-of-the-art baselines on a range of benchmarks.
Additionally, we provide furtheranalysis on convergence rate, sample efficiency, and robustness to different demonstration qualities.
Qualitative results are presented in https://shorturl.at/ipxGW .
In summary, our contributions are as follows:
• A framework that reformulates offline MARL as sequential modeling and policy distillationfor using self-attention to perform multi-agent credit assignment.
• A novel multi-agent policy distillation objective that focuses on preserving structural relationships among policies.
• State-of-the-art results on a broad range of tasks in terms of performance, convergence rate,sample efficiency and robustness to demonstration quality compared with baselines.
Related Works
MARL.
Applying RL to multi-agent settings has been an active research domain for years.
One lineof work adapts value-based RL algorithms to perform credit assignment among agents given the jointreward signal [40, 9, 32, 38, 13].
COPA [23] and REFIL [15] uses dynamic team composition tofurther improve the training efficiency for tasks that require cooperative behaviors.
Another line of work extends the actor-critic framework to the multi-agent setting.
For instance,
MADDPG[25] uses DDPG [21] to jointly update decentralized policies and a centralized critic.
Forbetter scalability, several works [16, 17] further introduce the attention mechanism into the critic toprevent noises from irrelevant information .
A unique challenge in multi-agent environments [14, 30] is local observation.
Previous works haveproposed various communication mechanisms to tackle this issue.
By exchanging informationbetween agents, they can better understand the tasks and other agents.
Some works [17, 8, 28, 24, 20]incorporate the attention mechanism to improve communication efficiency, while others [37] designa gating mechanism to decide when to communicate.
Offline RL.
Different from the typical RL setting, offline RL trains the policy with a static, collecteddataset and does not need any online interaction with the environment.
Then, this policy is leveraged to interact with the online environment to obtain promising results or assisit exploration [43].
However,trajectories existing in offline datasets and interaction with the online environment have differentdistributions.
One of the strategies to mitigate this issue is applying policy constraints [27, 48].
Another direction to alleviate this difficulty is to regularize the value estimation in reinforcementlearning [18] or take uncertainty into consideration [1].
Recently, the Decision Transformer [3] outperforms many state-of-the-art offline RL algorithmsby regarding the training process as a sequential modeling phase and testing on the online environment.
This approach can bypass the drawback of TD-learning and overcome sparse rewarddifficulty.
MADT [26] further extends Decision Transformer to a multi-agent domain by makingagents controlled by a shared weight transformer-based policy.
Knowledge Distillation.
Knowledge distillation (KD) transfers knowledge from one deep learningmodel (the teacher) to another (the student).
The objective originally proposed by [12] minimizes the
KL divergence between the teacher and student outputs.
This formulation makes intuitive sense whenthe output is a distribution, such as a probability mass function over classes.
Recent works focuson how to leverage the relationships between instances in the dataset to further provide meaningfulrepresentation learning [42, 41].
RKD [29] extracts the relation between instances from the featurespace by defining distance-wise relation and angle-wise relation.
Knowledge distillation is also adapted to RL scenario.
[6] proposes to distill multiple task-specificpolicies into a single policy which is more parameter-efficient.
DPD [19] utilizes two policies thatinteract with the same environment with different initialization to explore different perspectives of theenvironment and extract knowledge from each other to enhance their learning.
M&M [7] combinescurriculum learning as well as distillation to allow agents to perform well in an environment withlarge action space.
The typical goal of knowledge distillation is model compression by making the student model smaller.
Our method is different because our goal is to enable decentralized execution in the multi-agentsetting.
Method
In this section, we introduce our framework based on sequence modeling and policy distillation.
We first describe a baseline, Independent Decision Transformer (IDT), that naively adapts Decision
Transformer to the multi-agent setting in Section 3.1.
Then, we present our novel structural relation
Input: offline dataset D : { τ i ( o it , a it , r it )
Tt =1 )) } ni =
Initialize: θ as the parameters of π teacher , ϕ 1: n as the parameters of π 1: n // training centralized decision transformer for τ in D do ˆ a 1: n t = argmax π centralized ( a 1: n | ˆ a 1: n<t , θ ) θ
= θ − α ∇ θ L centralized end for // training decision transformers for agents
Freeze the weight θ for τ in D do ˆ a 1: n t = argmax π teacher ( a 1: n | ˆ a 1: n<t , θ ) for i=1:n do a it =
argmax π i ( a i | ˆ a i<t , ϕ i ) ϕ
i = ϕ
i − α ∇ ϕ
i ( L iaction + αL irel + βL iKL ) end for end for distillation for multi-agent policy distillation in Section 3.2.
The overall algorithm is presented in
Algorithm 1.
3.1 Independent Decision Transformer
[Figure]
An intuitive method to transform offline
MARL into a sequence modeling problemis to treat each agent’s trajectory as independent sequence.
For each agent i , a separate decision transformer π i is trained topredict the next action a it +1 given the pasttrajectories ( o i ( t − B ): t , a i ( t − B ): t , r i ( t − B ): t )where B stands for the maximum number of past steps to consider.
Since theseagents are trained independently withoutaccess to other agents’ information, wecall this method Independent Decision
Transformer (IDT).
Its learning objective can be formulated as a sequence modelling problemthat aims to predict the next action, L iaction =
|| a it +1 − a
it +
||
We note that although each agent canleverage self-attention to perform credit assignment along their own trajectories, this algorithm cannotperform credit assignment across agents .
This limitation prohibits agents to achieve cooperativebehavior that attains maximal returns.
3.2 Teacher-Student Policy Distillation
To address IDT’s cross-agents credit assignment issue, we propose to first train a centralized teacherpolicy π teacher that takes in all agents’ observations, actions, and rewards ( o 1: n ( t − B ): t , a 1: n ( t − B ): t , r 1: n ( t − B ): t )combined via concatenation to predict all agents’ actions ˆ a 1: n t +1 at the next step.
We represent theteacher policy π teacher with a decision transformer and train it to minimize the prediction error of allagents’ actions.
Since the teacher policy has access to all agents’ information, it can perform creditassignment across agents with self-attention to better foster cooperative behavior.
This is similarto how centralized critic helps actor critic-based methods in online MARL [16, 25].
However, thecentralized teacher policy cannot be deployed distributedly.
To fix this, we propose to distill theteacher policy’s features to n separate student policies.
We note that even though student policiesoperate independently during test time, they are very different from policies learned by IDT because(a) the supervision comes from a teacher policy that can perform credit assignment across agents, and(b) distillation provides a more stable learning signal to the student policy [2].
One intuitive policydistillation objective is to minimize the KL divergence between the actions predicted by the teacher policy ˆ a i and the student’s policy a i :
[Equation]
The other common distillation objective is to minimize the euclidean distance between the teacher’sfeatures
ˆ f i and the student’s features f i [33]:
[Equation]
However, both of these widely-used distillation strategies do not consider the structural relationbetween multi-agent students.
Structural Relation Distillation.
To help the student policies learn as much as possible from theteacher policy, we introduce a new distillation objective tailored to the multi-agent setting.
Our maininsight is that in addition to transferring the teacher policy’s actual feature values, we wish to preservethe structural relation among multi-agents’ features.
For example, if two agents belong to the sametype of units and have similar observations in SMAC [35], their outputs should be close to each other.
To this end, we define the relation between agents as the angle between feature vectors, and we intendto make the relation between student policies π 1: n mimic the relation between agents controlled byπ teacher .
In other words, we would like to minimize
[Equation]
where H denotes the Huber loss.
During the policy distillation procedure, the weight of the centralizedpolicy π teacher is frozen.
The comparision of IDT, MADT [26] and our approach is summarized in
Fig.
Mapping Networks.
Inspired by recent works [4, 11, 5] which use an MLP projection headto provide flexibility for contrastive representation learning, we propose to use a pair of mappingnetworks ( M, N ) to help policy distillation and prevent from loss of information induced by relationaldistillation loss.
By leveraging mapping networks that remove information irrelevant to structuralrelation yet potentially useful in downstream policy learning, more information can be formed andmaintained in the feature f .
Specifically, M and N are simple MLPs that transform the features beforerelational policy distillation.
Note that the weight of mapping networks are non-shared ( M ̸ = N )in our setting since the teacher and student policies are inherently asymmetric in the amount ofinformation allowed for reasoning (the teacher has access to privileged information from all agents).
Therefore, the relational policy distillation becomes
[Equation]
However, a learnable and non-fixed M effectively forms a moving target in the objective, causingunstable learning.
To stabilize the distillation process, we adopt a momentum-like update for mappingnetworks.
We make the update frequency of the teacher’s mapping network M lower than thestudent’s mapping network N .
In other words, we update M every e updates of N .
We empiricallyfound that setting e = 4 improves the convergence speed, but the converged performance is insensitiveto e .
In summary, the overall learning objective for agent i is
[Equation]
where α and β are hyperparameters that determine the importance of the proposed policy distillation.
Experiments
We execute a series of experiments to evaluate whether the proposed method is effective at solvingoffline MARL problems.
Specifically, our experiments seek to answer the following questions: first,does our method perform favorably against a wide range of existing approaches based on sequencemodeling, imitation learning, and offline reinforcement learning (Section 4.1)?
Specifically, wecompare our method with model-free offline MARL methods based on TD-learning and MADT, a
[Table] [Table]
concurrent work that also solves MARL via sequence modeling.
Next, we conduct careful ablationstudies to evaluate the contribution of each component within our framework (Section 4.2).
To test thesample efficiency of various approaches, we further show the performance of various methods whendifferent numbers of demonstrations are provided (Section 4.4).
Finally, we analyze the convergencerate (Section 4.5) and discuss the scalability (Section 5) of our method.
Baselines.
We compare our approach with three group of baselines including sequence modeling( IDT , MADT ), offline RL ( MA-CQL , MA-ICQ , MA-BCQ ) and imitation learning/inverse RL ( BC ,
MA-GAIL , MA-AIRL )
[Table]
and we briefly introduce them as follow.
BC : behavior cloning.
IDT : each agent is represented by a decision transformer
[3] independently.
MADT [26]: a weight-sharing decision transformer is used to represent each agent’s policy.
MACQL [18]: conservative Q-learning (CQL), which aims to address the extrapolation error by learninga conservative Q-function such that the expected value of a policy under this Q-function lower-boundsits true value.
MA-ICQ [45]: implicit constraint Q-learning (ICQ), which improves upon MA-CQLby only trusting the state-action pairs given in the dataset for value estimation under the multi-agentsetting.
MA-BCQ [10]: batch-constrained reinforcement learning, which restricts the action spacein order to force the agent towards behaving close to on-policy with respect to a subset of the givendata.
MA-GAIL [39]: an inverse reinforcement learning for multi-agent scenario.
MA-AIRL [47]: aframework for multi-agent inverse reinforcement learning, which is effective and scalable for Markovgames with high-dimensional state-action space and unknown dynamics.
For detail experimentalsetting about our approach and baselines, please refer to appendix A.
Environments.
We test our approach on three multiagent environments.
Fill-In : a grid-world environment that agents are required to pass all the blocksin map.
Equal Space : a particle system environmentthat agents need to keep the same distance betweeneach other.
SMAC [28]: a collaborative multi-agentreinforcement learning based on Blizzard’s StarCraft
II RTS game, and we choose four different scenarios,2s3z, 3s5z, 8m9m, and 3s5z vs 3s6z.
Grid-World :a simple grid-world environment that allow agentsto conduct discrete actions to move and collect thecorresponding objects.
Highway : a multi-agent environment that simulate highway traffic scenario.
Eachcontrollable agent drives a car and the goal is to reacha high speed while not colliding into neighboring vehicles.
For detail experimental setting, please refer toappendix B.
Offline Dataset
The offline MARL dataset for nagents is represented as a set of trajectories τ :=( o 1: n t , a 1: n t , r 1: n t )
Tt =1 .
For Grid-World and Highway,we train a centralized policy with PPO [36] and treat itas the expert to generate demonstrated trajectories.
The offline datasets for SMAC [35] are collectedby running a policy trained with MAPPO [46].
We also provide more information about datasets inappendix C.
Fill-In
Equal Space
Grid-World
Highway
Ours (conventional distillation) − 4 .
11 ± 0 .− 5 .
43 ± 0 .1 .
62 ± 0 .18 .
33 ± 1 . 27 Ours ( M = N ) − 3 .
11 ± 0 .− 2 .
61 ± 0 .1 .
42 ± 0 .
[Table]
Ablation study.
We show the average and standard deviation of return pre-agent, and all the experiments are merged with 10 random seeds.
4.1 Quantitative Results
We present the quantitative results in Table 1.
For all the experiments, we report the mean andstandard deviation of rewards based on 10 runs using different random seeds.
The results showthat our approach outperforms other baselines in all three environments.
In general, we find thatoffline RL approaches reach higher performance compared with imitation learning and inverse RLmethods.
It is because offline RL approach can use reward signals to judge the quality of the policyand try to infer optimal strategy.
Sequence modeling methods slightly outperform offline RL onessince Sequence modeling methods allow us to bypass bootstrapping and perform credit assignmentsdirectly via self-attention.
Besides, comparing ours and IDT, we can find that the proposed structuralknowledge distillation does provide performance benefits.
We also find that our approach performs better than inverse MARL methods.
Although inverse MARLdiscards the necessity of reward signal, it assumes the quality of demonstration trajectories is perfect.
Under this assumption, it is possible to infer reward signals in an adversarial manner.
However, sinceit’s hard to guarantee the quality of multi-agent trajectories due to its dependency on other agents’behaviors, the drawback of inverse MARL becomes obvious.
By observing the behavior of the agents, our approach encourage agents to behave more efficiently.
In Table 2, we present the quantitativeresults with different qualities of offlinedemonstrations and find that our methodoutperforms all baselines except for onesetting.
The results suggest that ourmethod is more robust to various demonstration qualities across tasks.
We hypothesize our method performs on parwith baselines on Highway-poor becauselearning a strong centralized decisiontransformer from poor demonstrations ischallenging.
This in turns affects the performances of student policies.
4.2 Ablation Study
To validate the effectiveness of each component in our approach.
To be more specific, we compareour approach with following variant.
Ours - Momentum Update : the variant that both mappingnetwork ( M and N ) are updated with the same frequency.
Ours - Mapping Network : the variantthat relation knowledge distillation is directly applied to feature space of decision transformer.
OursKL divergence : our approach without the KL divergence between teacher and student as distillation.
Ours (conventional distillation) : our approach but use KL divergence for knowledge distillationfrom the teacher policy.
Ours ( M = N ) : the variant of our approach that the mapping networks forteacher policy and student policies share the same weights.
The results show that all the components in our approach are required.
Specifically, Ours - Mapping
Network achieves poor performance on 7 out of 8 tasks.
It shows that Mapping Networks thattransform the students’ and teacher’s features into the same space for distillation are essential.
Additionally, according to the standard deviation of the performance, we find that approaches that donot use momentum update are less reliable.
KL divergence also provide performance benefit, but notsignificant compared with mapping networks and momentum update.
We also find that ours (conventional distillation) performs the worst, which highlight that theproposed our knowledge distillation is more effective than transitional knowledge distillation.
It isbecause structural relation among multi-agent can potentially represent how agent interact with eachother, which is the critical signal for a multi-agent task.
Fig.
5 shows the behavior of the policiestrained by our approach and ours with conventional distillation, respectively.
To make agents keepthe same distance from each other, policy obtained by our method tends to make agent move the theclosest corner and form an equilateral triangle structure.
On the other hand, policy obtained by offlinelearning and conventional distillation encourage specific agents to go to specific corners, which is asub-optimal solution.
The performance Ours ( M = N ) drops a lot, which represents that making theweights of the two mapping networks non-shared is a reasonable choice.
4.3 Finetuning
To verify whether the offline pre-trained framework can be further improved, we finetune the pretrained framework with MAPPO [46].
In Fig 6, we show the comparison of learning curves betweenfinetuning based on the pre-trained policy and policy trained from scratch.
The experimental resultsshow that the pre-trained policy can get improvement with an online multi-agent reinforcementlearning approach.
4.4 Data Efficiency
We also show the performance with a different number of trajectories in Fig.
The results showthat our relational knowledge distillation achieves better performance than conventional knowledgedistillation when the number of offline trajectories is relatively low.
Besides, we also find that offline
RL approaches perform well with limited offline trajectories compared with our approach.
Wehypothesize that the cause is that the transformer-based approach requires a large amount of data.
We stress that a common offline MARL setting allows abundant yet non-interactive demonstrationand thus our approach is still meaningful and desirable.
Finally, we find that naive imitation
learningand behavior cloning perform the worst in all the cases, and we hypothesize that imitation learningsuffers from covariate shift, which can be more severe when the training data is few and unable tocover a sufficiently large domain.
4.5 Convergence Rate
We observe better or comparable convergence with the centralized decision transformer for allenvironments compared to MADT, e.g., in SMAC 3s5z vs 3s6z with poor quality data, the centralizedtransformer is ∼ 8% better at 20k iterations (not fully converged).
For knowledge distillation, itrequires roughly similar total training iterations for the student policies to converge.
Besides, adding acentralized decision transformer doesn’t introduce a large overhead in training time.
In general, addingone more centralized DT only requires less than an additional 10 % training time in our experiment,the training time as well as learning curves for each experiment are presented in Appendix B.
Limitation and Broader Impact
The scalability of multi-agent frameworks is important.
Our approach depends on a centralizeddecision transformer, which predicts all actions of each agent based on the returns and actions of allthe agents.
Therefore, the centralized decision transformer may fall short in scalability.
However, ourexperiments find that we can still reach comparable performance compared with MADT when thenumber of agents increases.
We leave scalability in multi-agent policy distillation for future work.
Additionally, we note that offline MARL algorithms may lead to negative applications such as flyinga fleet of drones for surveillance.
Conclusion
In this work, we propose an offline MARL algorithms that first trains a teacher policy as if the datasetis generated by a single agent.
Once the teacher policy is trained, we use its predictions to supervisethe student using a novel distillation objective that considers agents’ relations.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] We provide summary of our contribution at the end of
introduction.
(b) Did you describe the limitations of your work?
[Yes] We discuss the limitation of ourapproach in 5.(c)
Did you discuss any potential negative societal impacts of your work?
[Yes] Please seeSec. 5(d)
Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A] There aren’ttheoretical results in this paper.
(b) Did you include complete proofs of all theoretical results?
[N/A]
There aren’t theoretical results in this paper.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[No] We will release
code about experiments after the paper is accepted.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] We include all the experimental setting in the supplementary B
(c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)?
[Yes] All the experiments are conducted with ten different
random seeds, and we present the mean and standard devaition of the performance.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] The computational infrastructure
are presented in the supplementary B
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] Some of the multi-
agent environments used in the experiments based on existing assets, and we clearly
cited the related works and present the url in the supplementary.
(b) Did you mention the license of the assets?
[Yes] The licenses of the assets are mentioned
in the appendix.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
This work doesn’t intend to propose a new dataset etc.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
This work doesn’t propose a new dataset, and there aren’t data
collected.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
This work doesn’t propose a new dataset, and
there aren’t data collected
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A] We don’t use crowdsourcing.
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A] We don’t use crowdsourcing.
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A] We don’t use crowdsourcing.