STAS: Spatial-Temporal Return Decomposition for Multi-agent Reinforcement Learning

In recent years, significant progress has been made in the field of cooperative multi-agent reinforcement learning, demonstrating its potential to solve diverse problems across several domains (Du et al. 2023), such as Real-Time Strategy games (Vinyals et al. 2019), distributed control (Busoniu, Babuska, and De Schutter 2008), and robotics (Perrusquía, Yu, and Li 2021). However, cooperative games present a considerable challenge known as the credit assignment problem, which involves determining each agent’s contribution using the centralized controller. Furthermore, real-world environments are complex and unpredictable, leading to long-term delayed and discontinuous rewards, and causing inefficiencies in the learning process. For instance, in the game of Go, the agent may not receive rewards until the outcome is determined. These nontrivial features pose significant challenges to existing methods that mainly rely on dense rewards to learn a policy that maximizes long-term returns. Additionally, the lack of dense rewards can make finding optimal solutions difficult, resulting in high variance during training and potential failure to complete the task.

Substantial progress has been made in solving the credit assignment problem, with notable approaches including QMIX  (Rashid et al. 2020), QTRAN  (Son et al. 2019), and SQDDPG  (Wang et al. 2020a). However, these methods do not address delayed reward environments. Recently, several studies have focused on temporal return decomposition methods that aim to decompose the return on a temporal scale  (Liu et al. 2019a; Ren et al. 2022; Arjona-Medina et al. 2019). These methods use a Markovian proxy reward function to replace environmental rewards, converting the delayed reward problem into a dense reward problem. Inspired by this progress, we consider addressing the challenge of sparse or delayed rewards in multi-agent tasks.

In this paper, we consider a new task called the spatial-temporal return decomposition problem, in which the algorithm must not only decompose the delayed episodic return to each time step but also assign credits to each agent according to their contribution. To solve this problem, we propose a novel approach called STAS, which establishes a connection between credit assignment and return decomposition. The Shapley value (Shapley 1997; Bilbao and Edelman 2000; Sundararajan and Najmi 2020) is a well-known concept for computing relationships in cooperative games, as it considers all possible orderings in which players could join a coalition and determines the average marginal contribution of each player across all possible orderings. Inspired by that, STAS first uses a temporal attention to analyze state correlations within episodes, providing essential information for calculating Shapley values. Then it employs spatial attention module that uses Shapley values to allocate global rewards based on each agent’s contribution. We also introduce a random masked mechanism to overcome the computational complexity of the Shapley value by simulating different coalitions. In the spatial-temporal return decomposition problem, pinpointing key frames and the most contributory agent is vital for clear rewards that highlight actions positively impacting the final goal. By using a dual transformer structure, STAS can identify key steps with the temporal transformer and use the spatial transformer to find the key agents. This enables STAS to obtain a Markovian reward, resulting in a better simulation of dense reward situations. Finally, with successfully decomposed delayed return, single-agent can be effectively trained using policies such as PPO (Schulman et al. 2017) or SAC (Haarnoja et al. 2018).

In summary, our contributions are three-fold. Firstly, we define the Spatial-Temporal Return Decomposition Problem and propose STAS, which connects credit assignment with return decomposition to solve the problem. Secondly, we introduce a dual transformer structure that utilizes a temporal transformer to identify key steps in long-time intervals and a spatial transformer to find the most contributory agents. We also leverage the Shapley value to compute the contribution of each agent and design a random masked attention mechanism to approximate it. Lastly, we validate our approach in a challenging environment, Alice & Bob, showing superior performance and effective cooperative pattern learning. Additionally, our method excels in general tasks in the Multi-agent Particle Environment (MPE), surpassing all baselines.

Return decomposition. Return decomposition, or temporal credit assignment, seeks to partition delayed returns into individual step contributions. Earlier methods relied on naive principles, such as count-based (Gangwani, Zhou, and Peng 2020; Liu et al. 2019a) or return-equivalent contribution analysis (Arjona-Medina et al. 2019; Patil et al. 2022; Zhang et al. 2023). Han et al. (2022) remodeled the problem and divided the value network into a history return predictor and a current reward one as HC-decompostion. Ren et al. (2022) proposed a concise method to utilize Monte Carlo Sampling as a time step sampler to decrease the computing complexity as well as the high variance. Although She, Gupta, and Kochenderfer (2022) initially addressed sparse rewards in credit assignment, their method falls short in extremely delayed reward scenarios, with only terminal rewards at the trajectory’s end. More recently, Xiao, Ramasubramanian, and Poovendran (2022) utilized attention structure to expand the return decomposition to multi-agent tasks as well as assign the reward to each agent in an implicit way by an end-to-end model.

Credit assignment. Credit assignment methods can be divided into implicit and explicit methods. Implicit methods treat the global value and the local value as a whole for end-to-end learning. Earlier methods ((Sunehag et al. 2018; Rashid et al. 2020)) developed a network to aggregate the Q-values of each agent into a global Q-value based on additivity or monotonicity. Subsequent methods attempt to delve deeper into the characteristics of IGM to achieve more accurate estimation (Son et al. 2019; Wang et al. 2021) or employ other approaches such as Taylor expansion to estimate global value (Yang et al. 2020; Wang et al. 2020b).

On the other hand, explicit methods attempt to establish an explicit connection between the global value and local value, enhancing the interpretability of the model by introducing algebraic relationships. Foerster et al. (2018) compute each agent’s advantage with a counterfactual baseline, while Wang et al. (2020a) introduced Shapley value to calculate the contribution of each agent. Li et al. (2021) combined counterfactual baseline and Shapley value to calculate each agent’s contribution in different coalitions. Du et al. (2019) learns per-agent intrinsic reward functions to evade credit assignment issue. Explicit methods typically employ the Actor-Critic architecture, which allows us to observe the contribution of each agent to the global reward, but this may result in high computational complexity.

Cooperative Multi-Agent Reinforcement Learning   In a fully cooperative multi-agent task  (Panait and Luke 2005), agents need to work independently to achieve a common goal. Usually it can be described as a tuple: $\langle\mathcal{N},\mathcal{S},\mathcal{U},P,r,\gamma,\rho_{0}\rangle$. Let $\mathcal{N}=\{1,2,\cdots,n\}$ denote the set of $n$ agents. Denote the joint state space of the agents as $\mathcal{S}=\prod_{i=1}^{n}S_{i}$ and the joint action space of the agents as $\mathcal{U}=\prod_{i=1}^{n}U_{i}$ respectively. $P:\mathcal{S}\times\mathcal{U}\times\mathcal{S}\to[0,1]$ is the state transition function. At time step $t$, let $\boldsymbol{s}_{t}=\{s^{i}_{t}\}^{n}_{i=1}$ with each $s^{i}_{t}\in S_{i}$ being the state from agent $i$. Accordingly, let $\boldsymbol{u}_{t}=\{u^{i}_{t}\}^{n}_{i=1}$ with each $u^{i}_{t}\in U_{i}$ indicating the action taken by the agent $i$. All agents will share a global reward from environment $r(\boldsymbol{s}_{t},\boldsymbol{u}_{t}):\mathcal{S}\times\mathcal{U}\to\mathbb{R}$. $\gamma\in[0,1)$ is a discount factor and $\rho_{0}:\mathcal{S}\to\mathbb{R}$ is the distribution of the initial state $\boldsymbol{s}_{0}$. The goal of the agents is to find optimal policies that achieve maximum global return, $J=\mathbb{E}_{\boldsymbol{s}_{0},\boldsymbol{u}_{0},\cdots}[\sum_{t=0}^{T}\gamma^{t}r_{t}]$ where $\boldsymbol{s}_{0}\sim\rho_{0}(\boldsymbol{s}_{0})$ and $T$ is the length of horizon.

Shapley Value Shapley value is a popular method for computing individual contributions in cooperative games (Shapley 1997; Bilbao and Edelman 2000; Sundararajan and Najmi 2020). Given a coalitional game $(v,N)$ where $N$ is the number of players cooperating together and $v$ refers to the value function describing the payoff that a coalition can obtain. For a particular player $i$, let $C$ denote an arbitrary subset that does not contain player $i$ and $C\cup\{i\}$ denote the attendance of player $i$ in subset $C$. Then the marginal contribution of player $i$ in $C$ can be defined as:

The Shapley Value of player $i$ can be computed from the marginal contribution of player $i$ in all subsets of $N$:

One crucial characteristic of the Shapley value is its efficiency property, which guarantees that the sum of the Shapley values of all agents equals the value of the grand coalition, i.e., $v(N)=\sum_{i=1}^{N}\Phi_{v}(i)$.

Return Decomposition In Episodic Reinforcement Learning. Commonly, agents receive a reward $r_{t}$ immediately after execution of action $\boldsymbol{u}_{t}$ at state $\boldsymbol{s}_{t}$. However, in the setting of episodic reinforcement learning (Liu et al. 2019b; Xiao, Ramasubramanian, and Poovendran 2022), agents can only obtain one global reward feedback at the end of the trajectory. Let $\tau=(\boldsymbol{s}_{0},\boldsymbol{u}_{0},\boldsymbol{s}_{1},\boldsymbol{u}_{1},\cdots,\boldsymbol{s}_{T})$ denote a trajectory of length $T$. We assume all trajectories terminate in finite steps. Then the episodic return function $r_{\text{ep}}(\tau)$ is defined on the trajectory space. The goal of episodic reinforcement learning in the multi-agent setting is to maximize the trajectory return, $J=\mathbb{E}_{\tau=(\boldsymbol{s}_{0},\boldsymbol{u}_{0},\cdots)}[r_{\text{ep}}(\tau)]$. Therefore, the extremely sparse rewards will introduce large bias and variance (Arjona-Medina et al. 2019; Ng, Harada, and Russell 1999) into the process of training, let alone the lower sample efficiency when agents share such global rewards. In practice, we assume that the episodic return has some structure in nature, e.g., a sum-decomposable form (Ren et al. 2022):

This assumption is quite feasible and widely adopted in empirical studies (Liu et al. 2019a), where the task can be quantified by an additive metric, e.g., number of defeated enemies or team score in a sports match.

In this paper, we address the challenge of fully cooperative multi-agent systems with episodic global return. We propose the STAS framework in Fig 1, which decomposes the return in both temporal and spatial dimensions.

Following Eq. (3), we assume that the sum of dense global rewards along the episode can reconstruct the final return given at the last time step. Similarly, the global reward equates to the sum of all distributed individual rewards.

where $\Phi_{v}(i,t)$ is the distributed reward for each agent. As mentioned previously, the use of Shapley Value in cooperative games has proven beneficial, distributing the payoff of each agent and leading to faster convergence rates. Consequently, we calculate the Shapley value as the actual distributed reward for each agent.

In this section, we conduct experiments to evaluate the performance and effectiveness of our proposed method on the challenging problem of spatial-temporal return decomposition. We evaluate our approach using two different implementations and test it on three different environments: a newly designed environment “Alice & Bob”, as well as the cooperative navigation and predator-prey scenarios in the general multi-agent benchmark MPE (Lowe et al. 2017). The source code is available at https://github.com/zowiezhang/STAS.

In this paper, we focus on the credit assignment problem in multi-agent learning with delayed episodic return. We introduce return decomposition methods to this problem, and we propose an novel approach called STAS. Our approach utilizes the Shapley value to evaluate each agent’s contribution and combines these contributions on both spatial and temporal scales to fit with the real episodic return from the environment. Our experimental results demonstrate that STAS outperforms all baselines, showing lower variance and faster convergence speed. In future work, we plan to include more baselines while extending it to more challenging tasks.

This work is sponsored by National Natural Science Foundation of China (62376013), Collective Intelligence & Collaboration Laboratory (Open Fund Project No. QXZ23014101)