Multi-Agent MDP Homomorphic Networks

1 INTRODUCTION Equivariant and geometric deep learning have gained traction in recent years , showing promising results in supervised learning ( Cohen & Welling , 2016 ; Winkels & Cohen , 2018 ; Weiler et al. , 2018 ; Weiler & Cesa , 2019 ; Worrall et al. , 2017 ; Fuchs et al. , 2020 ; Thomas et al. , 2018 ) , unsupervised learning ( Dey et al. , 2021 ) and reinforcement learning ( van der Pol et al. , 2020 ; Simm et al. , 2020 ) . In single agent reinforcement learning , enforcing equivariance to group symmetries has been shown to improve data efficiency , for example with MDP homomorphic networks ( van der Pol et al. , 2020 ) , trajectory augmentation ( Lin et al. , 2020 ; Mavalankar , 2020 ) , or symmetric locomotion policies ( Abdolhosseini et al. , 2019 ) . Equivariant approaches enable a single agent to learn policies more efficiently within its environment by sharing weights between state-action pairs that are equivalent under a transformation . As a result of this weight sharing , the agent implicitly learns a policy in a reduced version of the MDP . We are interested in cooperative multi-agent reinforcement learning , where symmetries exist not only in the global environment , but also between individual agents in the larger multi-agent system . Existing work on symmetries in single agent reinforcement learning can only be generalized to the fully centralized multi-agent setting , because such approaches rely on the global symmetry in the full state-action spaces and these can result in correspondences across agents , as shown in Figure 1 . This means such approaches can not be used in distributed multi-agent systems with communication constraints . Here , we seek to be equivariant to global symmetries of cooperative multi-agent systems while still being able to execute policies in a distributed manner . Within deep multi-agent reinforcement learning , existing work has shown the potential of using permutation symmetries and invariance between agents ( Liu et al. , 2019 ; Jiang et al. , 2020 ; Sunehag et al. , 2017 ; Robbel et al. , 2016 ; Böhmer et al. , 2020 ; Sukhbaatar et al. , 2016 ; van der Pol & Oliehoek , 2016 ) . Such work takes an anonymity view of homogeneous agents , where only the agent ’ s observations , and not its identity , matter for the policy . Using permutation symmetries ensures extensive weight sharing between agents , resulting in improved data efficiency . In this work , we go beyond such permutation symmetries , and consider more general symmetries of global multi-agent systems , such as rotational symmetries . In this paper , we propose Multi-Agent MDP Homomorphic Networks , a class of distributed policy networks which are equivariant to global symmetries of the multi-agent system , as well as to standard permutation symmetries . Our contributions are as follows . ( i ) We propose a factorization of global symmetries in the joint state-action space of cooperative multi-agent systems . ( ii ) We introduce a multi-agent equivariant policy network based on this factorization . ( iii ) Our main contribution is an approach to cooperative multi-agent reinforcement learning that is globally equivariant while requiring only local agent computation and local communication between agents at execution time . We evaluate Multi-Agent MDP Homomorphic Networks on symmetric multi-agent problems and show improved data efficiency compared to non-equivariant baselines . 2 RELATED WORK . Symmetries in single agent reinforcement learning Symmetries in Markov Decision Processes have been formalized by Zinkevich & Balch ( 2001 ) ; Ravindran & Barto ( 2001 ) . Recent work on symmetries in single agent deep reinforcement learning has shown improvements in terms of data effiency . Such work revolves around symmetries in policy networks ( van der Pol et al. , 2020 ; Simm et al. , 2020 ) , symmetric filters ( Clark & Storkey , 2014 ) , invariant data augmentation ( Laskin et al. , 2020 ; Kostrikov et al. , 2020 ) or equivariant trajectory augmentation ( Lin et al. , 2020 ; Mavalankar , 2020 ; Mishra et al. , 2019 ) These approaches are only suitable for single agent problems or centralized multi-agent controllers . Here , we solve the problem of enforcing global equivariance while still allowing distributed execution . Graphs and permutation symmetries in multi agent reinforcement learning Graph-based methods in cooperative multiagent reinforcement learning are well-explored . Much work is based around coordination graphs ( Guestrin et al. , 2002b ; a ; Kok & Vlassis , 2006 ) , including approaches that approximate local Q-functions with neural networks and use max-plus to find a joint policy ( van der Pol & Oliehoek , 2016 ; Böhmer et al. , 2020 ) , and approaches that use graph-structured networks to find joint policies or value functions ( Jiang et al. , 2020 ; Sukhbaatar et al. , 2016 ) . In deep learning for multi-agent systems , the use of permutation symmetries is common , either through explicit formulations ( Sunehag et al. , 2017 ; Böhmer et al. , 2020 ) or through the use of graph or message passing networks ( Liu et al. , 2019 ; Jiang et al. , 2020 ; Sukhbaatar et al. , 2016 ) . Policies in multi-agent systems with permutation symmetries between agents are also known as functionally homogeneous policies ( Zinkevich & Balch , 2001 ) or policies with agent anonymity ( Robbel et al. , 2016 ; Varakantham et al. , 2014 ) . Here , we move beyond permutation symmetries to a broader group of symmetries in multiagent reinforcement learning . Symmetries in multi agent reinforcement learning Recently , Hu et al . ( 2020 ) used knowledge of symmetries to improve zero-shot coordination in games which require symmetry-breaking . Here , we instead use symmetries in cooperative multi-agent systems to improve data efficiency by parameter sharing between different transformations of the global system . 3 BACKGROUND . In this section we introduce the necessary definitions and notation used in the rest of the paper . 3.1 MULTI-AGENT MDPS . We will start from the definition of Multi-Agent MDPs , a class of fully observable cooperative multi-agent systems . Full observability implies that each agent can execute the same centralized policy . Later on we will define a distributed variant of this type of decision making problem . Definition 1 A Multi-Agent Markov Decision Process ( MMDP ) ( Boutilier , 1996 ) is a tuple ( N , S , A , T , R ) where N is a set of m agents , S is the state space , A = A1 × · · · × Am is the joint action space of the MMDP , T : S ×A1 × · · · ×Am × S → [ 0 , 1 ] is the transition function , and R : S ×A1 × · · · ×Am → R is the immediate reward function . The goal of an MMDP , as in the single agent case , is to find a joint policy that maps states to probability distributions over joint actions , π : S → ∆ ( A ) , ( with ∆ ( A ) the space of such distributions ) to maximize the expected discounted return of the system , Rt = E [ ∑T k=0 γ krt+k+1 ] with γ ∈ [ 0 , 1 ] a discount factor . An MMDP can be viewed as a single-agent MDP where the agent takes joint actions . 3.2 GROUPS AND TRANSFORMATIONS . In this paper we will refer extensively to group symmetries . Here , we will briefly introduce these concepts and explain their significance to discussing equivalences of decision making problems . A group G is a set with a binary operator · that obeys the group axioms : identity , inverse , closure , and associativity . Consider as a running example the set of 90 degree rotations { 0◦ , 90◦ , 180◦ , 270◦ } , which we can write as rotation matrices : R ( θ ) = [ cos θ − sin θ sin θ cos θ ] ( 1 ) with θ ∈ { 0 , π2 , π , 3π 2 } . Composing any two matrices in this set results in another matrix in the set , meaning the set is closed under composition . For example , composing R ( π2 ) and R ( π ) results in another member of the set , in this case R ( 3π2 ) . Similarly , each member of the set has an inverse that is also in the set , and R ( 0 ) is an identity element . Since matrix multiplication is associative , the group axioms are satisfied and the set is a group under composition . A group action is a function G×X → X that satisfies ex = x ( where e is the identity element ) and ( g · h ) x = g · ( hx ) . For example , the group of 90 degree rotation matrices acts on vectors to rotate them . Similar to the action of this group on vectors , we can define an action of the same group on image space : e.g. , the NumPy ( Harris et al. , 2020 ) function np.rot90 acts on the set of images . We will consider group actions on the set of states represented as image observations . We match these with group actions on policies . Since we consider discrete action spaces , a group element g acting on a policy π will be represented as a matrix multiplication of the policy with a permutation matrix . When discussing symmetries in decision making problems , we identify sets of state-action pairs that are equivalent : if the state is transformed , the policy should be transformed as well , but potentially with a different representation of the transformation . See Figure 1 . We are interested in the case where the reward and transition functions are invariant in the orbit of state-action pairs under a symmetry group . The orbit of a point v ∈ V , with V a vector space , is the set of all its transformations ( e.g . all rotations of the point ) , defined as O ( v ) = { gv|∀g ∈ G } . The orbit of a point under a group forms an equivalence class . See Figure 2 for an example of an orbit of a traffic light state . 4 DISTRIBUTING SYMMETRIES OVER MULTIPLE AGENTS . Consider the cooperative traffic light control system in Figure 1 that contains transformationequivalent global state-action pairs . We first formalize global symmetries of the system similarly to symmetries in a single agent MDP . Then , we will discuss how we can formulate distributed symmetries in a distributed MMDP . Finally , we introduce Multi-Agent MDP Homomorphic Networks . 4.1 SYMMETRIES IN MMDPS . As a starting point , we define the notion of symmetries in an MMDP , similar to an MDP with symmetries ( van der Pol et al. , 2020 ) . Definition 2 An MMDP is an MMDP with symmetries if reward and transition functions are invariant under a transformation groupG . That is , the MMDP has symmetries if there is at least one non-trivial group G of transformations Lg : S → S and for every s , Ksg : A→ A such that R ( s , a ) = R ( Lg [ s ] , K s g [ a ] ) ∀g ∈ G , s ∈ S , a ∈ A , ( 2 ) T ( s , a , s′ ) = T ( Lg [ s ] , K s g [ a ] , Lg [ s ′ ] ) ∀g ∈ G , s , s′ ∈ S , a ∈ A . ( 3 ) If two state-action pairs s , a and Lg [ s ] , Ksg [ a ] obey Eq . 2 and 3 , then they are equivalent ( van der Pol et al. , 2020 ) . Consider as an example the symmetries in Figure 1 . These symmetries can result in correspondences across agents , for example when the observation of agent i is mapped by the symmetry to another agent j that is arbitrarily far away and with which there is no communication channel . In the next section , we will resolve this problem by defining distributed symmetries in terms of local observations and the communication graph defined by the state . If we have an MMDP with symmetries , that means that there are symmetric optimal policies , i.e . if the state of the MMDP transforms , the policy transforms accordingly . The above definition of an MMDP with symmetries is only applicable to the centralized setting . If we want to be able to execute policies in a distributed manner , we will need to enforce equivariance in a distributed manner .

Summary: The authors give a generalization of MDP Homomorphic Networks (van der Pol, 2020) to multi-agent distributed settings.  This is not a special case of previous work as the agents are not assumed to be able to communicate freely.  The symmetry constraints of the single-agent case generalize to the multi-agent case with the twist that the group action also induces a permutation of the agents.  To create a neural network which matches this symmetry, the authors use a message-passing NN with equivariance constraints.  The method shows improved sample efficiency relative to networks which are only permutation equivariant trained with and without augmentation.  Only the discrete 90-degree rotation group is considered.