Learning to communicate through imagination with model-based deep multi-agent reinforcement learning

1 INTRODUCTION . “ We use imagination in our ordinary perception of the world . This perception can not be separated from interpretation. ” ( Warnock , 1976 ) . The human brain , and the mind that emerges from its working , is currently our best example of a general purpose intelligent learning system . And our ability to imagine , is an integral part of it ( Abraham , 2020 ) . The imagination is furthermore intimately connected to other parts of our cognition such as our use of language ( Shulman , 2012 ) . In fact , Dor ( 2015 ) argues that : “ The functional specificity of language lies in the very particular functional strategy it employs . It is dedicated to the systematic instruction of imagination : we use it to communicate directly with our interlocutors ’ imaginations. ” However , the origin of language resides not only in individual cognition , but in society ( Von Humboldt , 1999 ) , grounded in part through interpersonal experience ( Bisk et al. , 2020 ) . The complexity of the world necessitates our use of individual mental models ( Forrester , 1971 ) , to store abstract representations of the information we perceive through the direct experiences of our senses ( Chang and Tsao , 2017 ) . As society expanded , the sharing of direct experiences within groups reached its limit . Growing societies could only continue to function through the invention of language , a unique and effective communication protocol where a sender ’ s coded message of abstract mental representations delivered through speech , could serve as a direct instruction to the receiver ’ s imagination ( Dor , 2015 ) . Therefore , the combination of language and imagination gave us the ability to solve complex tasks by performing abstract reasoning ( Perkins , 1985 ) and joint spatiotemporal planning ( Reuland , 2010 ) . In this work , we explore a plausible learning system architecture for the development of an artificial multi-agent communication protocol of the imagination . Based on the above discussion , the minimum set of required features of such a system include : ( 1 ) that it be constructed from multiple individual agents where , ( 2 ) each agent possesses an abstract model of the world that can serve as an imagination , ( 3 ) has access to a communication medium , or channel , and ( 4 ) jointly learns and interacts in a collective society . Consequently , these features map most directly onto the learning framework of model-based deep multi-agent reinforcement learning . Reinforcement learning ( RL ) has demonstrated close connections with neuroscientific models of learning ( Barto , 1995 ; Schultz et al. , 1997 ) . However , beside this connection , RL has proven to be an extremely useful computational framework for building effective artificial learning systems ( Sutton and Barto , 2018 ) . This is true , not only in simulated environments and games ( Mnih et al. , 2015 ; Silver et al. , 2017 ) , but also in real-world applications ( Gregurić et al. , 2020 ) . Futhermore , RL approaches are being considered for some of humanities most pressing problems , such as the need to build sustainable food supply ( Binas et al. , 2019 ) and energy forecasting systems ( Jeong and Kim , 2020 ) , brought about through global climate change ( Manabe and Wetherald , 1967 ; Hays et al. , 1976 ; Hansen et al. , 2012 ; Rolnick et al. , 2019 ) . Our system . We develop our system specifically in the context of cooperative mutli-agent RL ( OroojlooyJadid and Hajinezhad , 2019 ) , where multiple agents jointly attempt to learn how to act in a partially observable environment by maximising a shared global reward . Our agents make use of model-based reinforcement learning ( Langlois et al. , 2019 ; Moerland et al. , 2020 ) . To learn an artificial language of the imagination , each individual agent in our system is given access to a recurrent world model capable of learning rich abstract representations of real and imagined future states . We combine this world model with an encoder function to encode world model rollouts as messages and use a recurrent differentiable message passing channel for communication . To show the benefits of our system , we develop a set of ablation tests and specialised experiments using novel as well as well-known multi-agent environments and compare the performance of our system to a set of strong model-free deep MARL baselines . Our findings and contributions . We find that joint planning using learned communication through imagination can significantly improve MARL system performance when compared to a set of stateof-the-art baselines . We demonstrate this advantage of planning in a set of specialised environments specifically designed to test for the use of communication combined with imagined future prediction . Our present work is not at scale and we only consider situations containing two agents . However , to the best of our knowledge , this is the first demonstration of a model-based deep MARL system that combines world models with differentiable communication for joint planning , able to solve tasks successfully , where state-of-the-art model-free deep MARL methods fail . We see this work as a preliminary step towards building larger-scale joint planning systems using model-based deep multi-agent RL . 2 BACKGROUND AND RELATED WORK . Reinforcement learning is concerned with optimal sequential decision making within a particular environment . In single agent RL , the problem is modeled as a Markov decision process ( MDP ) defined by the following tuple ( S , A , r , p , ρ0 , γ ) ( Andreae , 1969 ; Watkins , 1989 ) . At time step t , in a state st , which is a member of the state space S , the agent can select an action at from a set of actions A . The environment state transition function p ( st+1|st , at ) provides a distribution over next states st+1 and a reward function r ( st , at , st+1 ) returns a scalar reward , given the current state , action and next state . The initial state distribution is given by ρ0 , with s0 ∼ ρ0 , and γ ∈ ( 0 , 1 ] is a discount factor controlling the influence of future reward . The goal of RL is to find an optimal policy π∗ , where the policy is a mapping from states to a distribution over actions , that maximises long-term discounted future reward such that π∗ = argmaxπE [ ∑∞ t=0 γ tr ( st , at , st+1 ) ] . If the environment state is partially observed by the agent , an observation function o ( st ) is assumed and the agent has access only to the observation ot = o ( st ) at each time step , with the full observation space defined as O = { o ( s ) |s ∈ S } . In this work , we focus only on the case of partial observability . Deep RL . Popular algorithms for solving the RL problem include value-based methods such as Qlearning ( Watkins and Dayan , 1992 ) and policy gradient methods such as the REINFORCE algorithm ( Williams , 1992 ) . Q-learning learns a value function Q ( s , a ) for state-action pairs and obtains a policy by selecting actions according to these learned values using a specific action selector , e.g . -greedy ( Watkins , 1989 ) or UCB ( Auer et al. , 2002 ) . In contrast , policy gradient methods learn a parameterised policy πθ , with parameters θ , directly by following a performance gradient signal with respect to θ . The above approaches are combined in actor-critic methods ( Sutton et al. , 2000 ) , where the actor refers to the policy being learned and the critic to the value function . In deep RL , the policy and value functions use deep neural networks as high-capacity function approximators capable of learning distributed abstract representations from raw input signals that are useful for downstream decision making . Recent state-of-the-art deep RL methods include Deep Q-Networks ( DQN ) ( Mnih et al. , 2013 ) and related variants ( Hessel et al. , 2017 ) , as well as advanced actor-critic methods such as PPO ( Schulman et al. , 2017 ) and SAC ( Haarnoja et al. , 2018 ) . See ( Arulkumaran et al. , 2017 ; Li , 2017 ) for an in-depth review of deep RL . Model-based RL . In RL , the environment transition function p is typically unknown . As a result , so-called model-free RL methods , such as DQN and PPO , rely solely on data gathered from the environment , i.e . real experience , to learn an optimal policy . However , if given access to a transition function , an agent can generate useful simulated , or imagined experience , and use it to plan . Therefore , in model-based RL , a model p̂φ ( ot+1|ot , at ) with parameters φ is learned using stored transitions gathered from either a random , heuristic or learned policy to simulate transitions from the true ( unknown ) transition function p. The model can then be used for model-based planning , which can either happen in the background , or at decisiontime . We briefly highlight the differences between these two types of planning and discuss work related to each and how this relates to our own work . – Background planning . In background planning , the model is primarily used to generate additional experience and assist learning , i.e . for updating the parameters of the policy and/or value functions . An early version of this approach is DYNA-Q ( Sutton , 1990 ) which uses the additional experience to help learn a value function . However , the usefulness of a model degrades over long time horizons as model rollout error starts to compound ( Gu et al. , 2016 ) . This has lead to different approaches that either use fixed depth rollouts based on model uncertainty ( Feinberg et al. , 2018 ) , dynamic rollout schedules ( Buckman et al. , 2018 ) , or short rollouts starting from intermediate states sampled from a buffer ( Janner et al. , 2019 ) . A promising alternative approach is to update gradients directly via imagined rollouts in a lower-dimensional latent space ( Hafner et al. , 2019 ; 2020 ; Byravan et al. , 2020 ) . – Decision-time planning . In decision-time planning , the model is used to generate imagined rollouts from a given state for the purpose of selecting the optimal action or sequence of actions . Decisiontime planning methods for discrete action spaces often rely on search methods such as Monte Carlo tree search ( MCTS ) ( Coulom , 2006 ) and have been used successfully in several works ( Silver et al. , 2017 ; Anthony et al. , 2017 ; Schrittwieser et al. , 2019 ) . In continuous action spaces , methods include trajectory optimisation approaches using trajectory sampling ( Todorov and Li , 2005 ; Theodorou et al. , 2010 ; Nagabandi et al. , 2018 ; Chua et al. , 2018 ) or collocation ( Posa et al. , 2014 ) ( optimising reward while forcing the model ’ s predictions to be close to already visited states ) . The model in our system is utilised for decision-time planning and follows the approach of Ha and Schmidhuber ( 2018 ) , who used recurrent neural world models as a way to give agent ’ s the ability to learn how to think ( Schmidhuber , 2015 ) . Specifically , we make use of a recurrent world model that takes the form of a mixture density network LSTM ( MDN-LSTM ) , as used in ( Ha and Eck , 2017 ) . The model is therefore a form of a recurrent Gaussian mixture model and allows us to sample probabilistic predictions of imagined next states . An illustration of the core features of model-based RL and the different types of planning is given in Figure 1 . Also see ( Janner , 2019 ) and ( Mordatch and Hamrick , 2020 ) for useful overviews . Multi-agent RL ( MARL ) . In the multi-agent case with N agents , we use the formalism of partially observable Markov games ( Littman , 1994 ) , defined as the tuple given above for the single agent case , but with observation and action spaces given by the following cartesian products : O = ∏N i=1Oi ⊆ S and A = ∏N i=1Ai , for agents i = 1 , ... , N . The goal in this setting is to find an optimal joint policy π∗ ( a1 , ... , aN |o1 , ... , oN ) that maximises a shared long-term discounted future reward for all agent as π∗ = argmaxπE [ ∑N i=1 ∑∞ t=0 γ tr ( oti , a t i , o t+1 i ) ] . Early work in MARL simply trained multiple independent Q-learning algorithms ( Tan , 1993 ) , which has since been extended to include deep neural networks , or more specifically independent DQNs ( Tampuu et al. , 2017 ) . However , from the perspective of an individual agent , these approaches treat all other learning agents as part of the environment , resulting in the optimal policy distribution to become non-stationary . Furthermore , if the environment is only partially observable , the learning task can become even more difficult , where agents may struggle with credit assignment due to spurious rewards received from unobserved actions of other agents ( Claus and Boutilier , 1998 ) . To mitigate the issue of non-stationarity , MARL systems are often designed within the paradigm of centralised training with decentralised execution ( CTDE ) ( Oliehoek et al. , 2008 ; Lowe et al. , 2017 ; Foerster et al. , 2017 ) . In CTDE , a centralised value function , or critic , is used during training , which conditions on the global state and joint actions from all the agents to make the learning problem stationary , but is then later removed once the individual agent ’ s policies have been learned , making it possible to use each policy independently during system execution . However , individual agent policies extracted in this way may still perform poorly because training is not specifically aligned with the goal of performing well under decentralised execution . Therefore , state-of-the-art value-based MARL approaches such as Q-mix ( Rashid et al. , 2018 ) and QTRAN ( Son et al. , 2019 ) make use of value function decomposition strategies ( Sunehag et al. , 2017 ) to more closely resemble decentralised training , where each agent is a recurrent DQN ( Hausknecht and Stone , 2015 ) that has memory to also deal with partial observability . Another clear way to help with the issue of partial observability is for agents to be able to communicate . Learned multi-agent communication has been a key innovation in helping MARL systems to scale to more complex environments and solve more challenging tasks ( Foerster et al. , 2016 ; Sukhbaatar et al. , 2016 ; Singh et al. , 2018 ; Chu et al. , 2020 ) . To facilitate communication in our work , we formally extend the Markov gameM by having agents connected to each other via communication channels according to a pre-defined neighbourhood graph G ( V , E ) . The graph G is defined by a set of nodes ( vertices ) V along with a set of edge connections E = { ( i , j ) |i , j ∈ V , i 6= j } , where each agent is a node in the graph , locally connected to other agent nodes . We define the connected neighbourhood surrounding agent i as Ni = { j ∈ V| ( i , j ) ∈ E } . This networked Markov game MG is then defined by the following tuple ( G , S , A , r , p , ρ0 , γ ) . Our communication channels are recurrent and end-to-end differentiable allowing for agent-to-agent communication protocols to be learned during training . Unlike work studying the emergence of language through communication in MARL , e.g . ( Lazaridou et al. , 2016 ; Mordatch and Abbeel , 2017 ; Kajić et al. , 2020 ) our work is more focused on communication through imagination as a useful system design for task solving , as apposed to uncovering new insights into emergent phenomena related to the human imagination . Model-based MARL . To the best of our knowledge , the literature on model-based deep MARL is quite sparse and very little work has been done in this area . A notable exception is the recent work by Krupnik et al . ( 2020 ) on multi-agent model-based latent space trajectory optimisation . Here a multi-step generative model , specifically a temporal segment model ( Mishra et al. , 2017 ) , is used to generate rollouts in a disentangled latent space and optimisation is performed directly over agent latent variables . Our work is the first we are aware of in the area of model-based deep MARL that combines communication with decision-time planning using recurrent neural world models .

Summary: This paper proposes to combine model-based and multi-agent reinforcement learning. The authors follow the typical recurrent neural world models setting to generate imagined rollouts for decision-time planning. To tackle the non-stationarity of a multi-agent environment, they build end-to-end differentiable communication channels between agents within a pre-defined neighborhood. The communication message is defined as abstract information encoded from the imagined rollout. Agents then make decisions based on the message they received and the output of recurrent neural world models. Empirical studies are performed to show the superiority of proposed methods over SOTA model-free MARL approaches. Results are shown in two simple environments, which are designed to require communication between agents to solve the task.