SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zikang Tian $^{1,2,3}$ , Ruizhi Chen $^{4}$ , Xing Hu $^{1,5}$ , Ling Li $^{2,4}$ , Rui Zhang $^{1}$ , Fan Wu $^{2,3,4}$ , Shaohui Peng $^{4}$ , Jiaming Guo $^{1}$ , Zidong Du $^{1,5}$ , Qi Guo $^{1}$ , Yunji Chen $^{1,2*}$
$^{1}$ SKL of Processors, Institute of Computing Technology, CAS, Beijing, China  $^{2}$ University of Chinese Academy of Sciences, Beijing, China  $^{3}$ Cambricon Technologies, Beijing, China
$^{4}$ Intelligent Software Research Center, Institute of Software, CAS, Beijing, China  
 $^{5}$ Shanghai Innovation Center for Processor Technologies, SHIC, Shanghai, China  
{tianzikang21s, cyj}@ict.ac.cn,
In recent years, Multi-Agent Reinforcement Learning (MARL) techniques have made significant strides in achieving high asymptotic performance in single task. However, there has been limited exploration of model transferability across tasks. Training a model from scratch for each task can be time-consuming and expensive, especially for large-scale Multi-Agent Systems. Therefore, it is crucial to develop methods for generalizing the model across tasks. Considering that there exist task-independent subtasks across MARL tasks, a model that can decompose such subtasks from the source task could generalize to target tasks. However, ensuring true task-independence of subtasks poses a challenge. In this paper, we propose to decompose a task into a series of generalizable subtasks (DT2GS), a novel framework that addresses this challenge by utilizing a scalable subtask encoder and an adaptive subtask semantic module. We show that these components endow subtasks with two properties critical for task-independence: avoiding overfitting to the source task and maintaining consistent yet scalable semantics across tasks. Empirical results demonstrate that DT2GS possesses sound zero-shot generalization capability across tasks, exhibits sufficient transferability, and outperforms existing methods in both multi-task and single-task problems.
In the last few years, many works in MARL field prospers, including value-based algorithms [26, 25, 20, 27, 30, 39, 10, 17], and policy-based algorithms [6, 36, 35, 11, 24, 7]. However, these works mainly focused on the model's asymptotic performance in a single task, neglecting its transferability across tasks. Training an optimal model on a single task requires millions of interactions with the environment [21, 16], particularly when dealing with large-scale Multi-Agent Systems, whereas transferring the model across tasks can reduce training cost by dozens of times [2, 5]. As the number of tasks increases, the reduction in training costs due to model transfer will become even more significant.
Knowledge reuse is a common approach for model generalization across tasks in the MARL field [4]. Recent knowledge reuse methods for online MARL can be roughly classified into two categories: network-design-based methods and task-embedding-based methods. Network-design-based methods [33, 9, 1] implicitly reuse knowledge extracted from the source task to the target task by constructing
universal model structure across tasks by utilizing Population Invariant Structure such as Transformer [9] or GNN [1]. However, it's unclear whether the knowledge extracted from the source task is suitable for the target task. Task-embedding-based methods [19, 23, 14] reuse knowledge by calculating task similarity using learned task embeddings that capture task dynamics. However, accurately mapping task dynamics to task embeddings requires numerous tasks as samples. In a word, current knowledge reuse methods for online MARL still have limitations, that is, inefficient knowledge reuse and reliance on a large number of tasks samples. Additionally, we provide a section for related work in Appendix A to introduce these methods more specifically.
Compared with the common knowledge reused in these methods above, there exists an alternative knowledge named task-independent subtasks that can relieve these two limitations. These task-independent subtasks, such as "hit and run", "focus fire", "disperse and line up", etc, are certainly applicable across tasks [41], improving the efficiency of knowledge reuse. Besides, these task-independent subtasks can be decomposed from few tasks, so as to remove the reliance on a large number of tasks samples. However, ensuring the task-independence of decomposed subtasks is a challenge. Task-independence of subtasks refers to their effectiveness across tasks, which requires two essential properties: (1) avoiding overfitting to the source task, (2) maintaining consistent yet scalable semantics across tasks. In this paper, we propose DT2GS, a novel framework devoted to generalize model across tasks by decomposing subtasks from the source tasks and endowing them with these two properties necessary for task-independence.
The proposed framework, DT2GS, aims to decompose a task into a series of generalizable subtasks, as shown in Figure 1. DT2GS is primarily composed of two parts: the scalable subtask encoder and the adaptive action decoder. The scalable subtask encoder assigns a subtask to each agent based on its {subtask, entity-observation} history rather than the {action, observation} history. This approach helps prevent the process of assigning subtasks from overfitting to the source task. The adaptive action decoder leverages the assigned subtasks and current entity-observations to calculate the specific actions for interacting with environment. Within this decoder, the adaptive subtask semantic module ensures that the assigned subtasks have consistent yet scalable semantics across tasks based on their effects on entities. Based on the experimental results, our framework exhibits several desirable properties as follows:
- zero-shot generalization capability: the trained model can be effectively depolyed to multiple target tasks without any fine-tuning;
- robust transferability: significantly accelerated model convergence on complex target tasks and achieving an average speedup of  $100 \times$ :
- better asymptotic performance: achieved state-of-the-art (SOTA) results on multi-task and single-task problems;
- better subtask interpretability: decomposed task-independent subtasks with practical behavioral semantics that is consistent yet scalable across tasks.
In this paper, a fully cooperative multi-agent system (MAS) is described as a decentralized partially observable Markov decision processes (Dec-POMDP) [18], which is defined by a tuple  $G = < S, \mathcal{A}, O, \gamma, n, \mathbb{A}, \Omega, R, P >$ .  $\mathcal{S}$  is the global state space.  $\mathcal{A}$  is the action space shared for all agents.  $O$  is the shared individual observation space for all agents.  $\gamma$  is the discount factor.  $n$  is the number of agents.  $\mathbb{A} = \{1, \dots, n\}$  is the set of agents our algorithm controlled. At each timestep  $t$ , agent  $i$  obtains an individual observation  $o_i^t \in O$  from dynamic environment according to the observation function  $\Omega(s^t, i)$ . And if we suppose agent  $i$  are controlled by policy  $\pi$ , which takes  $o_i^t$  or history individual observations  $\tau_i^t$  as input and parameterized by  $\theta_i$ , agent  $i$  will select an action  $a_i^t$  according to  $\pi(a_i^t | o_i^t)$ . Therefore, a joint action  $A^t = (a_1, a_2, \dots, a_n) \in \mathcal{A}$  will be formed where  $a_i$  corresponds to agent  $i$ . After passing the joint action  $A^t$  into the environment, a global reward signal  $r_t = R(s^t, A^t)$  shared by all agents will be received and the environment will transmit into next state  $s^{t+1}$  according to the transition function  $P(s^{t+1} | s^t, A^t): S \times \mathcal{A} \times S \to [0, 1]$ . The goal of agents is defined as finding the
policy  $\pi$  to maximize the objective function
$$
J (\pi) = \mathbb {E} \left[ \sum_ {t ^ {\prime} = 0} ^ {\infty} \gamma^ {t ^ {\prime}} R \left(s ^ {t ^ {\prime}}, A ^ {t ^ {\prime}}\right) \right] \tag {1}
$$
where  $\gamma^{t'}R(s^{t'},A^{t'})$  is the discounted return of all agents.
2.2 Generalizable model structure in MARL
For model generalization in MAS, a generalizable model structure is necessary for addressing the problem of varying state/observation/action space  $(S/O/A)$  across tasks. Here we define the agents controlled by MARL policy and agents built-in tasks as entities. In this paper, we use  $n, m, n_{ally}, n_{enemy}$  denotes the number of agents, entities, allies and enemies, respectively. And the following equation holds:  $n = n_{ally} + 1$ ,  $m = n + n_{enemy}$ . As demonstrated in ASN [33], an agent's observation  $o_i$  can be constructed as a concatenation of  $m$  entity-observations:  $o_i = [o_{i,1}, o_{i,2}, \dots, o_{i,m}]$ , where  $o_{i,1}$  is the observation of agent  $i$  on itself and environment, and the rest are the observations of agent  $i$  on other  $m - 1$  entities. Additionally, action space  $\mathcal{A}$  can be decomposed into two components:  $\mathcal{A}^{self}$ , which consists of actions affecting the agent itself or the environment, and  $\mathcal{A}^{interactive}$ , which contains actions that directly interact with other entities. This alignment between entity-observations and actions, which is referred as action semantics, forms the foundation for computing the value or probability of an action based on its aligning entity-observation, leading in a generalizable model structure across tasks.
Figure 1: DT2GS Framework. The DT2GS Framework comprises two modules: the Scalable Subtask Encoder and the Adaptive Action Decoder, where the Adaptive Subtask Semantic module serves as the core of the Adaptive Action Decoder. The Scalable Subtask Encoder effectively assigns subtasks to agents without overfitting to the source task. With the Adaptive Subtask Semantic module endowing assigned subtasks with consistent yet scalable semantics across tasks, the action decoder takes the entity-observations and subtasks as inputs to generate actions for interacting with the environment.
3 DT2GS Framework
DT2GS primarily consists of two components: the scalable subtask encoder and the adaptive action decoder, as shown in Figure 1. DT2GS begins with the scalable subtask encoder, which assigns a subtask to each agent based on its {subtask, entity-observation} history instead of the {action, observation} history, so as to avoid overfitting to the source tasks. Then the adaptive action decoder calculates the specific actions for interacting with environment by leveraging the assigned subtasks
and current entity-observations. Within this decoder, the adaptive subtask semantic module, which plays a core role, ensures that the assigned subtasks have consistent yet scalable semantics across tasks based on their effects on entities. With the proposed scalable subtask encoder and the adaptive subtask semantic module, DT2GS endows the decomposed subtasks with task-independence, leading model generalizable across tasks. A pseudocode is provided in Appendix H to illustrate the DT2GS framework.
Figure 2: Scalable Subtask Encoder. The Scalable Subtask Encoder assigns subtasks to each agent based on its behavioral history representation, which is depicted by its {subtask, entity-observation} history instead of commonly used {action, observation} history. The red dashed-line box denotes the {action, observation} history, which is task-dependent. Specifically, at each timestep  $t$ , agent  $i$  obtains its subtask  $k_{i}^{t}$  by utilizing its observation  $o_{i}^{t}$ , last action  $a_{i}^{t-1}$ , history latent embedding  $h_{i}^{t-1}$  produced by RNN. On the contrary, the red solid-line box denotes the {subtask, entity-observation} history, which is task-independent. Specifically, at each timestep  $t+1$ , agent  $i$  obtains its subtask  $k_{i}^{t+1}$  by utilizing its entity-observation  $\{o_{i,j}^{t+1}\} = [o_{i,1}^{t+1}, o_{i,2}^{t+1}, \dots, o_{i,m}^{t+1}]$ , last subtask  $k_{i}^{t}$ , history latent embedding  $h_{i}^{t}$  produced by RNN.
3.1 Scalable Subtask Encoder
An agent's history, which comprises its observations and behavioral history, contains valuable information about its characteristics, such as velocity, attack distance, etc. Therefore, analyzing an agent's history can help us identify the most suitable subtask for the agent to perform. Typically [38], a recurrent neural network, such as LSTM [8] or GRU [3], is employed to construct the agent's history representation from its {action, observation} history. The history representation is then passed to a softmax operator, which samples a one-hot subtask for the agent. However, since actions and observations are task-dependent, this approach may lead to subtask encoder that overfits to the source task, making the source model non-generalizable. Considering that the subtasks history of an agent not only reflects its behavior history, but also is task-independent, and the entity-observations history is more general across tasks than observations history, our scalable subtask encoder assigns a subtask to each agent based on its {subtask, entity-observation} history rather than {action, observation} history, as illustrated in the red solid-line box in Figure 2.
The detail of the scalable subtask encoder is demonstrated in Figure 1 right. By replacing observations with entity-observations, we first design a structure that can obtain fixed-dimensional observation embeddings, regardless of the number of entity-observations contained in the observation. Additionally, since permuting the order of entity-observations in the observation does not change the information[13], we use a Permutation Invariant operator, namely, Gaussian Product, to obtain the observation embedding. Specifically, we first use a MLP parameterized by  $\theta_{e}$  to embed the entity-observation as an entity-embedding:
$$
\left(\mu_ {i, j} ^ {t}, \sigma_ {i, j} ^ {t}\right) = f _ {\theta_ {e}} \left(o _ {i, j} ^ {t}\right), j = 1, \dots , m \tag {2}
$$
and the observation embedding  $e_i^t$ , which is also referred to as Env Cognition of agent  $i$  in Figure 1, is constructed as:
$$
e _ {i} ^ {t} \sim \mathcal {N} \left(\mu_ {i} ^ {t}, \sigma_ {i} ^ {t}\right), w h e r e \mathcal {N} \left(\mu_ {i} ^ {t}, \sigma_ {i} ^ {t}\right) \propto \prod_ {j = 1} ^ {m} \mathcal {N} \left(\mu_ {i, j} ^ {t}, \sigma_ {i, j} ^ {t}\right) \tag {3}
$$
Then we utilize a trajectory encoder based on GRU, which is parameterized by  $\theta_h$ , to obtain an agent's history representation. The GRU takes an agent's observation embedding  $e_i^t$ , last subtask  $k_i^{t-1}$  and previous hidden state  $h_i^{t-1}$  as inputs, and generates a new hidden state  $h_i^t$  as its history representation:
$$
h _ {i} ^ {t} = f _ {\theta_ {h}} \left(e _ {i} ^ {t}, k _ {i} ^ {t - 1}, h _ {i} ^ {t - 1}\right) \tag {4}
$$
Afterwards, we use the Gumbel-Softmax trick with a categorical reparameterization [12] to assign a subtask  $k_{i}^{t}$  to agent  $i$  based on its current history representation  $h_{i}^{t}$ :
$$
k _ {i} ^ {t} \sim \text {G u m b e l - S o f t m a x} \left(h _ {i} ^ {t}\right) \tag {5}
$$
where  $k_{i}^{t} \in \mathbb{R}^{n_{k}}$  is a  $n_{k}$ -dimensional one-hot vector and  $n_{k}$  is a hyperparameter denoting the total number of subtasks. The Gumbel-Softmax operator allows our process of subtask assignment trainable by the way of gradient backpropagation.
Figure 3: Adaptive Subtask Semantic module. The Adaptive Subtask Semantic Module endows subtasks and actions with adaptive semantics.  $z_{k}$  denotes the embedding of subtask  $k$ . Entity-observations are denoted by  $o = [o^{self}, o^{ally_{1}}, \dots, o^{ally_{n_{ally}}}, o^{enemy_{1}}, \dots, o^{enemy_{enemy}}]$  (we replace  $n_{enemy}$  with  $n_{e}$  in Figure 3 for convenience). With subtask embedding and entity-observations, we obtain the semantics  $\phi_{k}$  of subtask  $k$  and the semantics of actions  $\varphi = [\varphi^{self}, \varphi^{ally_{1}}, \dots, \varphi^{ally_{n_{ally}}}, \varphi^{enemy_{1}}, \dots, \varphi^{enemy_{enemy}}]$  by utilizing the Attention mechanism [29].
3.2 Adaptive Subtask Semantics
The generalizable subtasks should maintain consistent yet scalable semantics across tasks. The traditional method of constructing subtask semantics is embedding one-hot subtasks into vector representations based on MLP. However, this approach endows target subtasks with completely identical semantics as source subtasks when deploying the source model to the target tasks, neglecting the difference between source and target tasks so as to restrict the zero-shot generalization of subtask semantics across tasks. To overcome this limitation, our adaptive subtask semantic module leverages the inductive bias that subtask semantics refers to the effects of an agent on entities when it performs a given subtask. This approach endows subtasks with semantics that can adaptively adjust based on the agent's effects on entities, thereby enabling greater zero-shot generalization of subtask semantics across tasks.
Since a subtask's semantics actually refers to the effects of an agent on entities when it performs this subtask, we construct subtask semantics as the weighted sum of entities, where the weight represents the degree of effects. And considering that we have obtained one-hot subtask and entity-observations, we utilize the Attention mechanism [29] without position embedding, which is scalable across tasks regardless of the number of entities and permutation invariant about the order of entity-observaitons in the observation, to obtain the subtask semantics, as illustrated in Figure 3. Specifically, we take one-hot subtask embedding
$$
z _ {i} ^ {t} = \operatorname {E m b e d d i n g} \left(k _ {i} ^ {t}\right) \tag {6}
$$
as the query and entity-observations' embedding as the keys as well as values:  $\hat{Q}_i^t = W_Qz_i^t, K_i^t = W_Ko_i^t, V_i^t = W_Vo_i^t$ , where  $W_{Q}, W_{K}, W_{V}$  are learnable parameters and  $o_i^t = [o_{i,1}^t, o_{i,2}^t, \dots, o_{i,m}^t]$ . Then we construct adaptive subtask semantics  $\phi_i^t$  as follows:
$$
\phi_ {i} ^ {t} = \operatorname {s o f t m a x} \left(\frac {\hat {Q} _ {i} ^ {t} K _ {i} ^ {t T}}{\sqrt {d _ {K}}}\right) V _ {i} ^ {t}, \quad \hat {Q} _ {i} ^ {t} = W _ {Q} z _ {i} ^ {t} \tag {7}
$$
where  $d_{K}$  is the feature dimension of  $K_{i}^{t}$ . Additionally, to make the process of obtaining actions scalable across tasks, we extended the model structure of ASN [33]. Since the alignment between actions and entity-observations may discrepant from the prior, we construct the adaptive action semantics  $\varphi_{i}^{t} = [\varphi_{i,1}^{t},\varphi_{i,2}^{t},\dots,\varphi_{i,m}^{t}]$  by utilizing the Self-Attention mechanism taking the entity-observations as inputs:
$$
\varphi_ {i} ^ {t} = \operatorname {s o f t m a x} \left(\frac {Q _ {i} ^ {t} K _ {i} ^ {t T}}{\sqrt {d _ {K}}}\right) V _ {i} ^ {t}, \quad Q _ {i} ^ {t} = W _ {Q} o _ {i} ^ {t} \tag {8}
$$
Subsequently, we calculate the value or probability of each action in current observation by comparing the similarity of subtask semantics with action semantics:
$$
Q _ {\text {v a l u e}} \left(a _ {j} \mid o _ {i} ^ {t}\right) \text {o r} P r \left(a _ {j} \mid o _ {i} ^ {t}\right) = \text {s i m i l a r i t y} \left(\phi_ {i} ^ {t}, \varphi_ {i, j} ^ {t}\right) \tag {9}
$$
where similarity is a trainable MLP taking the concatenate of  $\phi_i^t$  and  $\varphi_{i,j}^{t}$  as inputs.
We evaluated the performance of DT2GS on the StarCraft Multi-Agent Challenge (SMAC) [22] and the multi-agent particle world environments (MPE) [16] (shown in Appendix D). SMAC contains several tasks that are similar but different, such as the marine-series tasks ( $\{3\mathrm{m}, 8\mathrm{m}, 8\mathrm{m\_vs\_9m}, 10\mathrm{m\_vs\_11m}, 25\mathrm{m}, 27\mathrm{m\_vs\_30m}\}$ ) and the stalker_zealot-series tasks ( $2\mathrm{s}3\mathrm{z}, 3\mathrm{s}5\mathrm{z}, 3\mathrm{s}5\mathrm{z\_vs\_3s6z}$ ). Besides, changing the number of agents and landmarks in MPE can also form a series of tasks. These tasks met our requirements of cross-task generalization, allowing us to evaluate the ability of DT2GS to generalize across different tasks.
In the generalization capability part of experiments, we selected ASN [33] and UPDeT [9] as baselines. ASN was chosen because it promotes the development of universal models across tasks in MARL. For the sake of fairness in comparison, we make ASN generalizable across tasks by utilizing the attention mechanism and use "ASN_G" to denote this generalizable ASN. UPDeT was selected because it constructs a universal model across tasks via policy decoupling with self-attention, which is also used in the adaptive subtask semantic module of DT2GS in Sec 3.2. In addition, we implemented DT2GS, UPDeT, and ASN_G based on MAPPO [40], which is considered SOTA in on-policy MARL.
In the asymptotic performance part of experiments, we added another five baselines for comparison: MAPPO[40], LDSA [38], ROMA [31], RODE [32], and HSD [37]. LDSA, ROMA, RODE, and HSD focus on concepts like skills/options/roles in MARL, which are similar to subtasks studied in our method.
Additionally, We set the number of subtasks  $n_k$  to 4 and averaged all results over 4 random seeds. And the experiments are arranged as follows: Firstly, we evaluated the zero-shot generalization capability of DT2GS across tasks. Secondly, we analyzed the practical semantics of subtasks on the marine-series tasks. Thirdly, we conducted several experiments to exhibit the transferability of DT2GS across tasks. Finally, we demonstrated that DT2GS achieves SOTA performance on multi-task and most single-task problems in terms of asymptotic performance.
4.2 Zero-Shot Generalization across Tasks
In this section, we designed 8 different zero-shot generalization scenarios, where each scenario includes a target task that is more difficult than the source task. That is to say, the relationship from the source task to the target task is extrapolated. For example, the target task may be larger in scale (like  $8\mathrm{m}\rightarrow 25\mathrm{m}$ ) or contain a greater disparity in military strength than the source task (3s_vs_4z  $\rightarrow$  3s_vs_5z). We deployed the model trained on the source tasks to the target tasks without any finetune. As shown in Figure 4, DT2GS significantly outperforms UPDeT and ASN_G in terms of zero-shot generalization capability, achieving an average test winning rate surpass of about  $22\%$  and  $34\%$ , respectively, over all 8 zero-shot generalization scenarios.
Figure 4: The figure shows a comparison of zero-shot generalization capability between DT2GS, UPDeT, and ASN_G across various source and target tasks. The horizontal axis represents the source task  $\rightarrow$  the target task, where (a) 3s_vs_4z  $\rightarrow$  3s_vs_5z, (b) 2s3z  $\rightarrow$  3s5z, (c) 3s5z  $\rightarrow$  3s5z_vs_3s6z, (d)  $8\mathrm{m}\rightarrow 8\mathrm{m\_vs\_9m}$ , (e)  $8\mathrm{m}\rightarrow 10\mathrm{m\_vs\_11m}$ , (f)  $8\mathrm{m}\rightarrow 25\mathrm{m}$ , (g)  $8\mathrm{m\_vs\_9m}\rightarrow 25\mathrm{m}$ , and (h)  $8\mathrm{m\_vs\_9m}\rightarrow 5\mathrm{m\_vs\_6m}$ . The vertical axis represents the winning rate when deploying the source model to target task without any finetue. The red, green, and purple histograms correspond to DT2GS, UPDeT, and ASN_G, respectively. The missing histograms indicate that the winning rate of deploying the source model to target task is  $0\%$ .
4.3 Analysis of Subtask Semantics