# Few-Shot Graph Learning for Molecular Property Prediction

Zhichun Guo1, Chuxu Zhang2∗, Wenhao Yu1
John Herr1, Olaf Wiest1, Meng Jiang1, Nitesh V. Chawla1∗ 1University of Notre
Dame, IN, USA 2Brandeis University, MA, USA
<EMAIL_ADDRESS><EMAIL_ADDRESS>wyu1, jherr1, Olaf.G.Wiest.1, mjiang2,
<EMAIL_ADDRESS>
(2021)

###### Abstract.

The recent success of graph neural networks has significantly boosted
molecular property prediction, advancing activities such as drug discovery.
The existing deep neural network methods usually require large training
dataset for each property, impairing their performance in cases (especially
for new molecular properties) with a limited amount of experimental data,
which are common in real situations. To this end, we propose Meta-MGNN, a
novel model for few-shot molecular property prediction. Meta-MGNN applies
molecular graph neural network to learn molecular representations and builds a
meta-learning framework for model optimization. To exploit unlabeled molecular
information and address task heterogeneity of different molecular properties,
Meta-MGNN further incorporates molecular structures, attribute based self-
supervised modules and self-attentive task weights into the former framework,
strengthening the whole learning model. Extensive experiments on two public
multi-property datasets demonstrate that Meta-MGNN outperforms a variety of
state-of-the-art methods.

Molecular Property Prediction, Few-Shot Learning, Graph Learning

††copyright: acmcopyright††journalyear: 2021††conference: Proceedings of the
Web Conference 2021; April 19–23, 2021; Ljubljana, Slovenia††booktitle:
Proceedings of the Web Conference 2021 (WWW ’21), April 19–23, 2021,
Ljubljana, Slovenia††doi: 10.1145/3442381.3450112††isbn:
978-1-4503-8312-7/21/04

ACM Reference Format:
Zhichun Guo, Chuxu Zhang, Wenhao Yu, John Herr, Olaf Wiest, Meng Jiang, Nitesh
V. Chawla. 2021. Few-Shot Graph Learning for Molecular Property Prediction. In
Proceedings of The Web Conference 2021 (WWW ’21), April 19-23, 2021,
Ljubljana, Slovenia. ACM, New York, NY, USA, 9 pages.
https://doi.org/10.1145/3442381.3450112

††$*$ Corresponding authors††$\S$ Our code is avaiable at
https://github.com/zhichunguo/Meta-MGNN

##  1\. Introduction

Drug discovery significantly benefits all human beings, especially for public
health during this tough and special time caused by COVID-19 (Vamathevan et
al., 2019). Developing and discovering new drugs is a time, resource, and
money consuming process. A key step is to test a large number of molecules for
therapeutic activity through extensive biological studies (Sliwoski et al.,
2014). Unfortunately, these discovered ones often fail to become the approved
drug candidates for various reasons such as low activity or toxicity (Waring
et al., 2015). Researchers need to select a great number of similar molecules
as potential candidates. To find the molecules which have the same efficacious
property, these selected molecules need to be tested through a complex
experimental process. After that, only a few or even no molecules will be
remaining as possible drug candidates to be tested further for risk and
pharmaceutical activity. Therefore, it is crucial to improve the effectiveness
of filtering the most likely drug candidates before taking experiments via
wet-lab experimentation, thus wasting less time and resources on molecules
that are unlikely to proceed to the lead stage. This concept is generally
described as "fail early-fail cheap".

Virtual screening is a widely used approach to screen out molecules likely to
fail early, which avoids a large set of molecules to be investigated (Riniker
and Landrum, 2013; Sliwoski et al., 2014). Recent advances in deep learning
have played an important role in virtual screening. These deep learning
techniques have inspired novel approaches to a better understanding of
molecules and their properties through molecular representation learning (Hu
et al., 2020; Xu et al., 2017; Gilmer et al., 2017; Zhang et al., 2018b; Zheng
et al., 2019). Deep neural networks learn more about specific molecular
properties when they are fed with more instances during training. Thus, deep
learning models require a large amount of training data to achieve desired
capability and satisfactory performance (Devlin et al., 2018). However, it is
common that there are only a few known molecules that share the same set of
properties (Altae-Tran et al., 2017; Waring et al., 2015). We analyzed the
datasets in MoleculeNet (Wu et al., 2018), a well-known benchmark for
predicting molecular properties. We find that more than half of the properties
only are shared by fewer than 100 molecules across several datasets. This is a
case of the well-known problem of _few-shot_ available data, which seriously
impairs the performances of current approaches. Therefore, it is essential to
develop a deep neural model for predicting molecular properties effectively in
few-shot scenarios.

There are several challenges that need to be overcome to achieve this goal.
Molecules can be considered as a heterogeneous structure where each atom
connects to different neighboring atoms via different types of bonds. Previous
work (Wang et al., 2019) represents molecules as SMILES strings and leverages
sequence models (Mikolov et al., 2013; Wang et al., 2019) to learn molecular
embedding. This approach is not able to capture information in each bond well
(Wu et al., 2020). This is because bonds in molecules not only represent
connected relations between different atoms but also contain attributed
information that characterizes the bond type such as single, double, or
triple. Thus, the first challenge is to design a deep neural network that can
discover _effective_ molecular representations from _few-shot_ data. Because
only a limited amount of labeled molecular property data are available, the
second challenge is to exploit the useful _unlabeled_ information in molecule
data and further develop an _efficient_ learning procedure to transfer the
knowledge from other property prediction, so that the model can fast adapt to
the novel (new) molecular properties with limited data. Moreover, different
molecular properties could represent quite different molecular structures.
Thus, their data should be treated differently in the knowledge transfer
process. The third challenge is to distinguish the _different_ importance of
molecular properties when performing the efficient learning procedure.

To address the above challenges, we propose a novel model called Meta-MGNN for
few-shot molecular property prediction. First, we leverage graph neural
network with the pre-training process to fuse heterogeneous molecular graph
information as molecular embedding. Then, we develop a meta-learning framework
to transfer knowledge from different property prediction tasks and obtain a
well-initialized model which could be fast adapted to a new molecular property
with limited data. In order to exploit and capture unlabeled information in
molecule data, we design a self-supervised module which consists of a bond
reconstruction loss and an atom type prediction loss, accompanied by the main
property prediction loss. Moreover, considering different property prediction
tasks contribute differently to the few-shot learner, we further introduce a
self-attentive task weight to measure their importance. Both self-supervised
module and self-attentive task weight are incorporated into the meta-learning
procedure for strengthening the model.

Contributions. To summarize, the main contributions of this work are as
follows:

  * •

We formulate the molecular property prediction as a few-shot learning problem,
which exploits the rich information in various properties to address the lack
of laboratory data problem for each individual property.

  * •

To deal with the few-shot challenge, we propose a novel model called Meta-MGNN
by exploring graph neural network, self-supervised learning, and task weight
aware meta-learning.

  * •

We conduct extensive experiments on two public datasets and the evaluation
results demonstrate the superior performance of Meta-MGNN over state-of-the-
art methods. The effectiveness of each model component is also verified.

Figure 1. (a) The overall framework of Meta-MGNN: It first samples a batch of
training tasks. For each task, there are a few data examples in the support
set. These examples are fed into a GNN parameterized by $\theta$. Then the
support loss $\mathcal{L}_{support}$ is calculated and utilized to update the
GNN parameters to $\theta^{\prime}$. Next, the examples in the corresponding
query set are fed into the GNN parameterized by $\theta^{\prime}$ and
calculate the loss $\mathcal{L}_{query}^{\prime}$ for this task. The same
process repeats for other training tasks. Later, we compute the summation of
$\mathcal{L}_{query}^{\prime}$ over all sampled tasks and use it to further
update the GNN parameters for testing. (b) Self-supervised module: It includes
bond reconstruction and atom type prediction. The orange part shows that we
sample two atoms and use GNN to predict if there is a bond between them. The
green part shows that we mask several atoms randomly and use GNN to predict
their types. (c) Task-aware attention: It calculates the average of all the
molecular embedding from the query set of the same task to represent this
task. With the embedding of each task, we design a self-attentive layer to
compute the weight of each task, then incorporate it into a meta-training
process for updating model parameters $\theta$.

##  2\. Related Work

In this section, we review existing work including graph neural network, few-
shot learning, and molecular property prediction.

Graph Neural Network (GNN).  GNNs have gained increasing popularity due to its
capability of modeling graph-structured data (Hamilton et al., 2017;
Veličković et al., 2018; Zhang et al., 2019). Typically, a GNN model uses a
neighborhood aggregation function to iteratively update the representation of
a node by aggregating representations of its neighboring nodes and edges. GNNs
have showed attractive performance in various applications, such as
recommendation systems (Fan et al., 2019; Song et al., 2019), behavior
modeling (Yu et al., 2020), and anomaly detection (Zhao et al., 2020).
Molecular property prediction is also a popular application of GNNs since a
molecule could be represented as a topological graph by treating atoms as
nodes, and bonds as edges (Gilmer et al., 2017; Mansimov et al., 2019; Lu et
al., 2019; Guo et al., 2020). We will elaborate them in the next paragraph.

Molecular Property Prediction.  Methods can be categorized into two main
groups based on the input molecular type: (1) molecular graph, and (2)
simplified molecular-input line-entry (SMILES) (Weininger et al., 1989). For
the first group of methods, each molecule is represented as a graph associated
with different atom nodes interconnected by bond edges. One typical way is to
employ graph neural networks to learn molecular representations (Gilmer et
al., 2017; Hu et al., 2020; Mansimov et al., 2019; Lu et al., 2019; Guo et
al., 2020). For example, Lu et al. (Lu et al., 2019) proposed a novel
hierarchical GNN. It includes an embedding layer, a Radial Basis Function
layer, and an interaction layer to learn molecular representations from
different levels. Hu et al. (Hu et al., 2020) proposed several novel pre-
training strategies to pre-train GNNs at the level of individual nodes and the
entire graph to lean local and global molecular representations
simultaneously. For the second type of representation, SMILES is a sequence
notation for describing the structure of molecules. Researchers take molecules
as sequences and adopt language models to learn their representations (Zhang
et al., 2018b; Wang et al., 2019; Zheng et al., 2019). For example, Zhang et
al. (Zhang et al., 2018b) proposed a semi-supervised Seq2Seq fingerprint model
which contains three ends of one input, one supervised output, and one
unsupervised output. Zheng et al. (Zheng et al., 2019) presented a new model
to study structure-property relationships through a self-attentive linear
notation syntax analysis. Guo et al. (Guo et al., 2020) proposed a novel graph
and sequence fusion learning model to capture information both from the
molecular graph structure and SMILES. Here, we take each molecule as a graph
as it preserves the molecular inner structure better, and employ graph neural
networks to learn their representations.

Few-shot Learning.  Successes of few-shot learning have been accomplished in
various application domains such as computer vision (Finn et al., 2017; Kim et
al., 2019) and graph learning (Yao et al., 2019; Yao et al., 2020; Garcia and
Bruna, 2018; Zhang et al., 2020a; Zhang et al., 2020b) . There are two notable
types of few-shot learning approaches: (1) metric-based learning and (2)
gradient-based learning. The former learns a generative metric to compare and
match few-examples (Vinyals et al., 2016; Sung et al., 2018; Bertinetto et
al., 2018). Vinyals et al. (Vinyals et al., 2016) proposed a novel matching
metric, named Matching Nets, to match unlabeled examples to the class of few-
shot labeled examples. Sung et al. (Sung et al., 2018) proposed relation
network which learns a deep distance metric to compute relation scores of
different images and further classify images. The latter aims to employ a
specific meta-learner to learn well-initialized parameters of the base model
for different tasks (Finn et al., 2017; Lee and Choi, 2018; Zhang et al.,
2018a). For instance, Finn et al. (Finn et al., 2017) proposed MAML which
designs this kind of meta-learner to effectively initialize a base-learner
that could be fast adapted to new tasks. In this work, our few-shot learning
strategy is gradient-based learning.

##  3\. Preliminary

In this section, we first define the few-shot molecular property prediction
problem, then present the details of using graph neural network (GNN) for
learning molecular representations.

###  3.1. Problem Definition

Let $G=(\mathcal{V},\mathcal{E})$ denote a molecular graph where $\mathcal{V}$
is the set of nodes and $\mathcal{E}\subseteq\mathcal{V}\times\mathcal{V}$ is
the set of edges. Particularly, a node in a molecular graph represents a
chemical atom and an edge represents a chemical bond between two atoms. Given
a set of molecular graphs $\mathcal{G}=\\{G_{1},\cdots,G_{N}\\}$ and their
labels $\mathcal{Y}=\\{y_{1},\cdots,y_{N}\\}$, the goal of molecular property
prediction is to learn a molecular representation vector for predicting its
label (i.e., molecular property) of each $G_{i}\in\mathcal{G}$, i.e., to learn
a mapping function $f_{\theta}:\mathcal{G}\rightarrow\mathcal{Y}$.

Unlike previous studies where there are enough examples for each new property
prediction task, this work considers a more practical scenario that only few-
shot samples are given. Specifically, we aim to develop a classifier which can
be fast adapted to predict new molecular properties that are unseen during the
training process, given only a few samples of these new properties. Formally,
the problem is defined as follows.

###### Problem 1.

Few-Shot Molecular Property Prediction Given molecular properties
$\mathcal{Y}=\\{y_{1},\cdots,y_{N}\\}$ and their corresponding few-shot
molecular graph sets $\\{{\mathcal{G}_{1}}\in
y_{1},\cdots,{\mathcal{G}_{N}}\in y_{N}\\}$ (training data), the task is to
design a machine learning model to predict molecular graphs of new properties
that only have few-shot examples (test data).

###  3.2. Molecular Graph Neural Network

By viewing molecular structure as graph data (i.e., molecular graph), recent
deep learning methods for graphs, such as graph neural networks (GNNs) (Wu et
al., 2020; Zhang et al., 2019), can be utilized to learn molecular
representations which are fed to downstream machine learning models for
molecular property prediction (Liu et al., 2019). In this section, we will
present the details of employing GNNs to obtain molecular representations.

A GNN model is able to utilize both graph structure and node/edge features
information to learn a representation vector $\textbf{h}_{v}$ for each node
$v\in\mathcal{V}$. Specifically, a GNN model uses a neighborhood aggregation
function to iteratively update the representation of a node by aggregating
representations of its neighboring nodes and edges. After $l$ iterations, a
node representation $\textbf{h}_{v}^{(l)}$ is able to capture the information
within its $l$-hop neighborhoods. In a molecular graph, each node represents
an atom and each edge represents a chemical bond between two atoms. As the
input layer of GNN, we first initialize representations of both nodes and
edges using their attributes in molecular graph. The node attributes include
atom number (AN) and chirality tag (CT), and edge attributes include bond type
(BT) and bond direction (BD). Formally, we initialize node representation as
$\textbf{h}_{v}^{(0)}=\textbf{v}_{AN}\oplus\textbf{v}_{CT}$ and edge
representation as $\textbf{h}_{e}^{(0)}=\textbf{e}_{BT}\oplus\textbf{e}_{BD}$,
where v and e denote node/edge attributes and $\oplus$ is concatenation
operator. Then, the node representation $\textbf{h}_{v}^{(l)}$ at the $l$-th
layer of GNN is formulated as:

(1)
$\displaystyle\textbf{h}^{(l)}_{\mathcal{N}(v)}=\textsc{Agg}_{l}(\\{\textbf{h}^{(l-1)}_{u}:\forall
u\in\mathcal{N}(v)\\},\\{\textbf{h}^{(l-1)}_{e}:e=(v,u)\\}),$ (2)
$\displaystyle\textbf{h}^{(l)}_{v}=\sigma(\textbf{W}^{(l)}\cdot\textsc{Concat}(\textbf{h}^{(l-1)}_{v},\textbf{h}^{(l)}_{\mathcal{N}(v)})),$

where $\mathcal{N}(v)$ is the neighbor set of $v$, $\sigma(\cdot)$ is a non-
linear activation function (e.g., LeakyReLU). $\textsc{Agg}(\cdot)$ is an
aggregating function. A number of architectures for $\textsc{Agg}(\cdot)$ have
been proposed in recent years such as graph convolutional neural network (GCN)
(Kipf and Welling, 2017) and graph attention network (GAT) (Veličković et al.,
2018). Here, we use graph isomorphism network (GIN) (Xu et al., 2019), which
has demonstrated state-of-the-art performance on a variety of benchmark tasks.
After that, we can learn the representation of each node in molecular graph:
$\textbf{h}_{v}=\textbf{h}^{(l)}_{v}/||\textbf{h}^{(l)}_{v}||_{2}$. To obtain
the graph-level representation $\textbf{h}_{G}$ for a molecular graph, we
calculate the average node embeddings at the final layer:

(3)
$\displaystyle\textbf{h}_{G}=\textsc{Mean}(\\{\textbf{h}_{v}^{(l)}:v\in\mathcal{V}\\}),$

The graph-level molecular representation $\textbf{h}_{G}$ can be further fed
into a classifier (e.g., a multi-layer perception) for molecular property
prediction, as we will present in the next section.

Pre-trained Molecular Graph Neural Network. Pre-trained models have been
widely used in natural language processing, computer vision, and graph
analysis in recent years (Devlin et al., 2018; Hu et al., 2020). In general,
pre-training allows a model to learn universal representations, provides a
better parameter initialization, and avoids overfitting on downstream tasks
with small training data. Models with pre-training have been demonstrated to
obtain superior performance than models without that. Therefore, we are
motivated to leverage the recent pre-trained graph neural network technique
(PreGNN) (Hu et al., 2020) to obtain parameter initialization of molecular
graph neural network.

##  4\. Meta-MGNN

In this section, we present the details of proposed the Meta-MGNN for few-shot
molecular property prediction. Meta-MGNN is built on MGNN and employs a meta-
learning framework for model initialization and adaption. Molecular structure
and feature based self-supervised module and self-attentive task weight are
further incorporated into the former framework for model enhancement.

###  4.1. Meta-learning Setup

We build the meta-learning framework based on MAML (Finn et al., 2017). Given
the model $f_{\theta}$ with learnable parameters $\theta$ that maps molecular
graph to specific properties such as toxicity, i.e.,
$f_{\theta}:\mathcal{G}\rightarrow\mathcal{Y}$. In meta-learning, the model is
expected to adapt to a number of different tasks, i.e., predicting different
kinds of molecular properties. Particular, in the $k$-shot meta-learning, for
each task $\mathcal{T}_{\tau}$ sampled from distribution $p(\mathcal{T})$, the
model is trained using only $k$ data samples and further tested on remaining
data samples of $\mathcal{T}_{\tau}$. In this setting, we refer to the
corresponding training and test sets of each task as support set and query
set, denoted as
$\mathcal{T}_{\tau}=\\{\mathcal{G}_{\tau},\mathcal{Y}_{\tau},{\mathcal{G}}_{\tau}^{\prime},{\mathcal{Y}}_{\tau}^{\prime}\\}$,
where $\mathcal{G}_{\tau},\mathcal{Y}_{\tau}$ are support sets of input
molecular graphs and property labels, and
$\mathcal{G}_{\tau}^{\prime},\mathcal{Y}_{\tau}^{\prime}$ are query sets of
input molecular graphs and property labels. During meta-training, the model
$f_{\theta}$ is first updated to task-specific model using support set of each
task, then further optimized to task-agnostic model using prediction loss over
the query set of all tasks in training data. After sufficient training, the
learned model can be further utilized to predict new tasks (new molecular
properties) with only $k$ data samples as support set, which is called meta-
testing. To avoid data overlapping, data of tasks used for meta-testing are
held out during meta-training. The whole framework is illustrated in Figure
1(a).

###  4.2. Meta-training

In meta-training, the goal is to obtain well initialized model parameters
$\theta$ that can be generally applicable to different tasks, and explicitly
encourage the initialized parameters to perform well after a small number of
gradient descent updates on a new task with few-shot data. When adapting to a
task $\mathcal{T}_{\tau}$, we begin with feeding the support set to the model
and calculate the loss $\mathcal{L}_{\mathcal{T}_{\tau}}$ to update parameters
$\theta$ to $\theta^{\prime}_{\tau}$ through gradient descent:

(4)
$\theta^{\prime}_{\tau}=\theta-\alpha\nabla_{\theta}\mathcal{L}_{\mathcal{T}_{\tau}}(\theta),$

where $\alpha$ is the step size. It should be noted that Eq.(4) only shows
one-step gradient update while we can take multiple-steps gradient update in
practice.

####  4.2.1. Loss Function.

Typically, the above loss $\mathcal{L}_{\mathcal{T}_{\tau}}$ is calculated by
the supervised signals from downstream tasks (Finn et al., 2017; Zhou et al.,
2019), i.e., molecular property labels in this study. However, simply using
supervised signals maybe not effective since only a few samples are given for
each task. In addition, the complexity of molecules inherently bring useful
unlabeled information in both structure and attribute. Therefore, to enhance
the above meta-training process, we propose to exploit and leverage unlabeled
information in molecular graphs. In particular, we design a self-supervised
module which consists of a bond reconstruction loss and an atom type
prediction loss, accompanying with the property prediction loss.

Molecular Property Prediction Loss. To predict molecular property, we
introduce a multi-layer perception (MLP) on top of the graph-level molecular
representation h ( Eq.(3)), i.e., $\hat{y}=\textsc{MLP}(\textbf{h})$. The loss
of prediction is defined as the cross entropy loss between the predicted
labels and ground-truth labels:

(5)
$\mathcal{L}_{label}(\theta)=-\frac{1}{k}\sum_{i=1}^{k}\textsc{CrossEntropy}(y_{i},\hat{y}_{i})$

where $k$ is the number of data samples.

Require : {$\mathcal{G}_{\tau},\mathcal{Y}_{\tau}$}: support data ;
$\\{{\mathcal{G}}_{\tau}^{\prime},{\mathcal{Y}}_{\tau}^{\prime}$}: query data;
$\alpha$, $\beta$: step sizes (i.e., learning rates)

1 $\theta$ $\leftarrow$ Pre-trained by PreGNN (Hu et al., 2020)

2 while  _not done_ do

3 Sample batch of tasks $\mathcal{T}_{\tau}\sim p(\mathcal{T})$

4 for  _all $\mathcal{T}_{\tau}$_ do

5 Sample $k$ examples $\\{G_{\tau 1},G_{\tau 2},\cdots,G_{\tau
k}\\}\in\mathcal{G}_{\tau}$

6 for  _i=1 to k_ do

7 $y_{\tau i},\textbf{h}_{\tau i}=\mathrm{GNN}(G_{\tau i},\ \theta)$

8

9 end for

10 $\textbf{H}_{\tau}=$ Mean ($\textbf{h}_{\tau 1},\textbf{h}_{\tau
2},\cdots,\textbf{h}_{\tau k}$)

11 $\mathcal{L}_{\tau}\leftarrow$ Eq. (9) with $\\{y_{\tau 1},y_{\tau
2},\cdots,y_{\tau k}\\}$

12 $\theta_{\tau}^{\prime}=\theta\ -\ \alpha\nabla\mathcal{L}_{\tau}$

13 Sample n examples $\\{G_{\tau 1}^{\prime},G_{\tau 2}^{\prime},...G_{\tau
n}^{\prime}\\}\in\mathcal{G}_{\tau}^{\prime}$

14 for  _j = 1 to n_ do

15 $y_{\tau j}^{\prime},\textbf{h}_{\tau j}^{\prime}=\mathrm{GNN}(G_{\tau
j}^{\prime},\ \theta_{\tau}^{\prime})$

16

17 end for

18 $\mathcal{L}_{\tau}^{\prime}\leftarrow$ Eq. (9) with $\\{y_{\tau
1}^{\prime},y_{\tau 2}^{\prime},\cdots,y_{\tau n}^{\prime}\\}$

19

20 end for

21 $\\{\eta(\mathcal{T}_{1}),\cdots,\eta(\mathcal{T}_{t})\\}\leftarrow$ Eq.
(11) with $\\{\textbf{H}_{1},\cdots,\textbf{H}_{t}\\}$

22 $\theta\leftarrow\theta-\beta\nabla_{\theta}\sum_{\mathcal{T}_{\tau}\sim
p(\mathcal{T})}\eta(\mathcal{T}_{i})\cdot\mathcal{L}_{i}^{\prime}$

23

24 end while

Algorithm 1 Meta-MGNN

Bond Reconstruction Loss. To perform bond reconstruction in molecular graphs,
we first sample a set of positive edges (existing bonds) in the molecular
graph, then sample a set of negative edges (non-existing bonds) by choosing
node pairs that do not have an edge in the original molecular graph. We denote
$\mathcal{E}_{s}$ as the union set of sampled positive edges and negative
edges. In practice, we set $|\mathcal{E}_{s}|$ = 10 including 5 positive
samples and 5 negative samples. The bond reconstruction score is computed by
the inner product of embeddings between the sampled pair of nodes, i.e.,
$\hat{e}_{uv}=\textbf{h}_{v}^{\top}\cdot\textbf{h}_{u}$. The bond
reconstruction loss is defined as the binary cross entropy loss between the
predicted bonds and ground-truth bonds:

(6)
$\mathcal{L}_{edge}(\theta)=-\frac{1}{|\mathcal{E}_{s}|}\sum_{e_{uv}\in\mathcal{E}_{s}}\textsc{BinaryCrossEntropy}(e_{uv},\hat{e}_{uv})$

Atom Type Prediction Loss. In a molecule, different atoms are connected in a
certain way (e.g., carbon–carbon bond, carbon–oxygen bond), leading to
different molecular structure. The atom type determines how a node in the
molecular graph connects with neighboring nodes. Thus, we utilize the
contextual sub-graph of a node (atom) to predict its type. Specifically, we
first sample a set of nodes in a molecular graph, denoted as
$\mathcal{V}_{ct}\subseteq\mathcal{V}$. For each node $v$ in
$\mathcal{V}_{ct}$, the contextual sub-graph is defined as its neighbors
within $l$-hops, i.e., $G_{sub}=(\mathcal{U}_{sub},\mathcal{E}_{sub})$ where
$\mathcal{U}_{sub}=\\{v\\}\cup\mathcal{N}_{l}(v)$,
$\mathcal{E}_{sub}\subseteq\mathcal{U}_{sub}\times\mathcal{U}_{sub}$, and
$\mathcal{N}_{l}(v)$ represents the the set of neighboring nodes of node $v$.
In practice, we choose $\mathcal{V}_{ct}=15\%\ $nodes in the graph and $l=1$.
Later, we use a a multi-layer perception (MLP) on the top of mean pooling of
all nodes in the contextual sub-graph excluding the central node and the atom
type prediction loss is formulated as the cross entropy loss between predicted
node type and ground-truth node type:

(7)
$\hat{v}_{i}=\textsc{Mlp}~{}(\textsc{Mean}(\\{\textbf{h}_{u}:u\in\mathcal{N}_{l}(v)\\})),$
(8)
$\mathcal{L}_{node}(\theta)=-\frac{1}{|\mathcal{V}_{c}|}\sum_{i=1}^{|\mathcal{V}_{c}|}\textsc{CrossEntropy}(v_{i},\hat{v}_{i}),$

The self-supervised module (of both bond reconstruction and atom type
prediction) is illustrated in Figure 1(b).

Joint Loss. The loss for task $\mathcal{T}_{\tau}$ in the meta-training
process is formulated as the summation over the above three losses,

(9)
$\mathcal{L}_{\mathcal{T}_{\tau}}(\theta)=\mathcal{L}_{node}(\theta)+\lambda_{1}\mathcal{L}_{edge}(\theta)+\lambda_{2}\mathcal{L}_{label}(\theta)$

where $\lambda_{1}$ and $\lambda_{2}$ are trade-off parameters that control
the importance of different losses. In practice, we set
$\lambda_{1}=\lambda_{2}=0.1$.

####  4.2.2. Task-aware Attention.

With the new model parameter $\theta^{\prime}_{\tau}$ obtained from the
support set data of task $\mathcal{T}_{\tau}$, i.e.,
$\theta^{\prime}_{\tau}=\theta-\alpha\nabla_{\theta}\mathcal{L}_{\mathcal{T}_{\tau}}(\theta)$,
the model is further updated as follows:

(10)
$\theta\leftarrow\theta-\beta\nabla_{\theta}\sum_{{\mathcal{T}_{\tau}}\sim
p(\mathcal{T})}\eta(\mathcal{T}_{\tau})\cdot\mathcal{L}_{\mathcal{T}_{\tau}}^{\prime}(\theta^{\prime}_{\tau}),$

where $\beta$ is meta-learning rate,
$\mathcal{L}_{\mathcal{T}_{\tau}}^{\prime}$ is the joint loss over query set
of $\mathcal{T}_{\tau}$. In other words, the model parameters $\theta$ are
further updated over losses of all sampled tasks through gradient descent. The
traditional meta-learning methods (e.g., MAML (Finn et al., 2017)) treat each
task with the same weight when optimizing the meta-leaner (i.e.,
$\eta(\mathcal{T}_{\tau})$ are same for all tasks), which cannot reflect how
important of different property prediction tasks are. Therefore, considering
different property prediction tasks contribute differently to the meta-learner
optimization, we further introduce a self-attentive weight to measure task
importance. Particular, we use self-attentive mechanism (Lin et al., 2017) to
calculate importance of each task:

(11)
$\eta({\mathcal{T}_{\tau}})=\frac{\exp(\textsc{Mlp}(\textbf{H}_{\mathcal{T}_{\tau}}))}{\sum_{\mathcal{T}_{\tau^{\prime}}\in\mathcal{T}}\exp(\textsc{Mlp}(\textbf{H}_{\mathcal{T}_{\tau^{\prime}}}))},~{}\textbf{H}_{\mathcal{T}_{\tau}}=\textsc{Mean}(\\{\textbf{h}_{\mathcal{T}_{\tau},i}\\}_{i=1}^{k}).$

where $\mathcal{T}$ is the set of all tasks, $\textbf{H}_{\mathcal{T}_{\tau}}$
denotes the task embedding which is computed by averaging all molecular
embeddings of $\mathcal{T}_{\tau}$. Figure 1(c) illustrates the details of
self-attentive task weight computation. The meta-training process is described
in Algorithm 1.

####  4.2.3. Meta-testing.

During meta-testing, we first utilize the few-shot support set of new tasks to
update parameters $\theta$ of Meta-MGNN via one or a small number of gradient
descent steps using Eq. (4), then evaluate performance in query set.

##  5\. Experiments

Table 1. Details of atom and bond features.

# Atom Type | 118
---|---
Atom Chirality Tag | Unspecified, Tetrahedral cw,
Tetrahedral ccw, Other
Bond Type | Single, Double, Triple, Aromatic
Bond Direction | -, Endupright, Enddownright

In this section, we conduct extensive experiments on two public datasets
(Tox21 and Sider) to compare performances of different models and show related
analysis.

Table 2.  The performances of all methods on both datasets. Our proposed
method Meta-MGNN can outperform all baseline methods. The last column reports
the average improvements (in percentage) of Meta-MGNN over the best baseline
method in different tasks. Bond indicates the best performance. Underline
represents the best baseline performance.

Dataset | Task | GraphSAGE (Hamilton et al., 2017) | GCN (Kipf and Welling, 2017) | MAML (Finn et al., 2017) | Seq3seq (Xu et al., 2019) | EGNN (Kim et al., 2019) | PreGNN (Hu et al., 2020) | Meta-MGNN | $\Delta$AUC
---|---|---|---|---|---|---|---|---|---
(2017) | (2017) | (2017) | (2018) | (2019) | (2020)
|  | 1-shot
Tox21 | SR-HS | 65.97 | 65.00 | 68.56 | 73.18 | 72.51 | 73.09 | 73.81 | +0.63
SR-MMP | 71.23 | 71.20 | 76.34 | 79.08 | 76.90 | 76.20 | 79.09 | +0.01
SR-p53 | 58.05 | 66.60 | 71.28 | 75.23 | 78.03 | 76.87 | 77.71 | -0.32
Average | 65.10 | 67.60 | 72.06 | 75.83 | 75.81 | 75.39 | 76.87 | +1.04
Sider | Si-T1 | 65.23 | 63.60 | 66.82 | 66.50 | 71.39 | 73.04 | 75.41 | +2.37
Si-T2 | 60.47 | 62.01 | 63.62 | 57.03 | 67.87 | 66.06 | 69.39 | +1.52
Si-T3 | 61.45 | 64.52 | 67.50 | 61.38 | 68.23 | 70.36 | 70.65 | +0.29
Si-T4 | 64.41 | 65.28 | 69.02 | 63.45 | 72.67 | 72.34 | 72.69 | +0.02
Si-T5 | 77.85 | 74.95 | 77.07 | 74.83 | 78.88 | 77.99 | 79.95 | +1.07
Si-T6 | 61.19 | 63.20 | 67.01 | 63.70 | 66.31 | 69.45 | 71.97 | +2.52
Average | 65.10 | 65.60 | 68.51 | 64.48 | 70.89 | 71.54 | 73.34 | +1.80
|  | 5-shots
Tox21 | SR-HS | 69.09 | 68.13 | 69.02 | 74.07 | 73.23 | 73.39 | 74.80 | +0.73
SR-MMP | 72.22 | 69.06 | 76.43 | 80.40 | 79.07 | 78.25 | 80.26 | -0.14
SR-p53 | 61.45 | 72.01 | 73.95 | 77.07 | 78.12 | 78.01 | 79.00 | +0.88
Average | 67.59 | 69.73 | 73.13 | 77.18 | 76.81 | 76.55 | 78.02 | +0.84
Sider | Si-T1 | 67.61 | 65.66 | 70.12 | 68.99 | 72.76 | 74.77 | 76.32 | +1.55
Si-T2 | 59.86 | 64.62 | 64.46 | 56.53 | 68.13 | 65.69 | 69.34 | +1.21
Si-T3 | 60.61 | 64.90 | 68.20 | 64.20 | 70.11 | 71.07 | 72.29 | +1.22
Si-T4 | 64.82 | 64.85 | 67.75 | 67.15 | 72.73 | 73.42 | 74.46 | +1.04
Si-T5 | 78.33 | 76.93 | 78.61 | 78.55 | 79.61 | 80.67 | 81.79 | +1.12
Si-T6 | 61.91 | 62.06 | 67.74 | 66.30 | 67.17 | 71.48 | 74.12 | +2.64
Average | 65.52 | 66.50 | 69.48 | 66.95 | 71.75 | 72.85 | 74.72 | +1.87

###  5.1. Datasets

We evaluate different methods on the Tox21 and Sider datasets. They are
collected from MoleculeNet (Wu et al., 2018), which is a large scale benchmark
dataset for molecular machine learning. Tox21 has 7,831 instances with 12
different tasks; Sider has 1,427 instances with 27 different tasks. In each
task, molecules are divided into positive instances and negative instances
(i.e., binary labels). A positive instance means that a molecule has a
specific property, and a negative instance means that a molecule does not have
the property. We manually split 3 tasks from Tox21 and 6 tasks from Sider for
meta-testing. The details of these two datasets are as follows:

  * •

Tox21111https://tripod.nih.gov/tox21/ challenge/: Toxicity on 12 biological
targets, including nuclear receptors and stress response pathways.

  * •

Sider: 27 system organ classes where molecules are marketed drugs and adverse
drug reactions (Kuhn et al., 2016).

Dataset Processing. The raw data of molecules are given as SMILES strings. We
transfer SMILES strings to molecular graphs by using Rdkit.Chem (Landrum,
2013). Then we extracted a set of node and bond features which can preserve
the molecular structure best to use in the experiments. The details about
features are listed in Table 1.

###  5.2. Baselines

We compare our model with multiple baseline models.

  * •

GraphSAGE (Hamilton et al., 2017). It generates the nodes’ embedding by
sampling and aggregating their neighbors’ embeddings, which can effectively
capture the graph information.

  * •

GCN (Kipf and Welling, 2017). It is a widely used graph-based model, which
contains an effective convolutional neural network component. GCN outperforms
various models by learning both local graph structure and features of nodes.

  * •

MAML (Finn et al., 2017). It builds a task-agnostic algorithm for few-shot
learning, where training a model’s parameters using a small number of gradient
updates will lead to fast learning on new tasks.

  * •

Seq3seq (Zhang et al., 2018b). It is a Seq2Seq model for molecular property
prediction. The loss function contains both self-recovery loss and inference
task loss.

  * •

EGNN (Kim et al., 2019).  It is an edge-labeling graph neural network for few-
shot learning which is proved a well-generalizable model for low-data problem.

  * •

PreGNN (Hu et al., 2020). This model develops self-supervised learning to
pretrain GNN for molecular property prediction. It captures both useful local
and global information.

(a) SR-HS

(b) SR-MMP

(c) SR-p53

(d) Si-T1

(e) Si-T2

(f) Si-T3

(g) Si-T4

(h) Si-T5

(i) Si-T6

Figure 2. Performances of different model variants in Tox21 data (the first
row) and Sider (the remaining figures) data. Different model components (i.e.,
graph neural network pre-training, self-supervised module, and task-aware
attention) indeed make effect and improve the model performance. Our proposed
model (M8) has better result than all other model variants.

###  5.3. Evaluation Metrics

We evaluate the performance of each model using ROC-AUC. We consider each
molecular property as an independent task for few-shot learning. We use 3 and
6 tasks as test tasks of Tox21 and Sider data, respectively. Each task is a
binary label classification task. Table 2 and Figure 3 report the result for
each test task. For both datasets, we consider 2-way classification with 1 and
5 shots.

###  5.4. Reproducibility Settings

We take graph isomorphism network (GIN) (Xu et al., 2019) as base graph neural
network. In our experiment, we utilize the supervised-contextpred pre-trained
GIN of PreGNN (Hu et al., 2020). The GIN layer number is set as 5. We set all
embedding dimensions to 300. The same feature will share the same initial
embedding. We set the update step in training tasks as 5 and the update step
in testing tasks as 10. We set the trade-off weight of self-supervised module
as 0.1. We use Pytorch to implement the model and run it on a GPU.

###  5.5. Comparisons with Baselines

Overall Performance. The overall performances of all methods are reported in
Table 2. According to this table, we can find that Meta-MGNN outperforms all
baseline models on both Tox21 and Sider datasets. Specifically, for 1-shot
learning, the average improvements are +1.04% and +1.80% on Tox21 and on
Sider, respectively. The values equal +0.84% and +1.87% for 5-shot learning.
In addition, we observe that PreGNN (Hu et al., 2020) and EGNN (Kim et al.,
2019) perform the best among all baseline methods on average. However, the
baseline methods do not have stable performance on different tasks. In other
words, they may perform well on one task, but perform poorly on another task.
In comparison, the performance of Meta-MGNN is stable. It has the best
performance for all tasks in both datasest.

Analyzing Meta-MGNN Structure.  MAML demonstrates superior performance than
the other two GNN models (GraphSage and GCN). It makes sense since MAML trains
the model through meta-learning, making it better adapt to new tasks with few
data samples. Besides taking advantage of meta-learning, Meta-MGNN also
utilizes pre-trained graph neural network model (PreGNN) (Hu et al., 2020) to
initialize model parameters. PreGNN uses a large amount of molecules data to
pre-train the graph neural network, which leads to better parameter
initialization. Therefore, by taking advantages of both meta-learning and pre-
training, Meta-MGNN demonstrates superior performance than the other baseline
methods.

Table 3. We implemented 8 model variants for ablation study. The abbreviations used in the table: pre-trained model (PTM), meta-learning (ML), bond reconstruction (BR), atom-type prediction (AP) and task-aware attention (T-At).  | PTM | ML | BR | AP | T-At
---|---|---|---|---|---
M1 | ✔ |  |  |  |
M2 |  | ✔ |  |  |
M3 | ✔ | ✔ |  |  |
M4 | ✔ | ✔ | ✔ |  |
M5 | ✔ | ✔ |  | ✔ |
M6 | ✔ | ✔ | ✔ | ✔ |
M7 | ✔ | ✔ |  |  | ✔
M8 | ✔ | ✔ | ✔ | ✔ | ✔

(a) Meta-MGNN

(b) PreGNN (Hu et al., 2020)

(c) MAML (Finn et al., 2017)

Figure 3. Visualizations of molecular embeddings generated by our model (Meta-
MGNN), PreGNN (Hu et al., 2020), and MAML (Finn et al., 2017). The blue dots
denote negative labels in SR-MMP (a molecular property). The orange dots
represent positive labels in SR-MMP. Our model can better discriminate
embeddings of these two kinds of labels than the other methods.

Performance on Different Datasets.  From Table 2, we can observe that our
proposed Meta-MGNN outperforms the best baseline method by +1.80% for 1-shot
learning and +1.87% for 5-shots learning on Sider dataset. However, the
average improvements on Tox21 are +1.04% for 1-shot learning and +0.84% for
5-shots learning, which are smaller than those in Sider. As we know, the major
advantage of few-shot learning model is to make model predict new tasks better
by using a small number of data samples. Obviously, training with more tasks
allows the model to learn more knowledge. The advantages of the few-shot
learning model can be better reflected on the dataset which contains more
tasks. Since Sider has more tasks than Tox21, Meta-MGNN can deliver greater
improvements on Sider than on Tox21. Additionally, we also find that the
overall performance on Tox21 is better than that on Sider for all models. This
is due to the larger size of Tox21, which improves the generalization
capabilities of these deep learning models, as reflected by the evaluation
scores.

###  5.6. Ablation Study

Settings. Besides comparing with baseline methods, we also implement model
variants (ablation studies) to show the effectiveness of different model
components. The details of different model variants are illustrated as follows
(also shown in the Table 3):

  * •

M1. Pre-trained graph neural network. We take GIN (Xu et al., 2019) as our
base graph neural network and pre-train it by both supervised and unsupervised
(Context Prediction) pre-training strategies.

  * •

M2. Graph neural network model (without pre-training) trained with the meta
learning process.

  * •

M3. Our base model, which is based on GIN (Xu et al., 2019) and learned with
meta-learning algorithms. This model is also pre-trained by supervised and
unsupervised (Context Prediction) pre-training.

  * •

M4 & M5 & M6. These models are based on M3 and augmented with the self-
supervised module. M4, M5, and M6 are augmented with the bond reconstruction,
the atom-type prediction, and both of them, respectively. This is to analyze
the effectiveness of the self-supervised module.

  * •

M7. It is based on M3 and enhanced with task-aware attention to incorporate
the importance of different tasks. This is to analyze the effectiveness of
task weight in meta-learning.

  * •

M8. It is based on M3 and augmented with both both self-supervised module and
self-attentive task weight.

Performance Comparison and Analysis.  The performances of all model variants
are shown in Figure 2. The three sub-figures in the first row are model
performance on Tox21. The sub-figures in the second row and third row are
model performance on Sider. There are several findings from these figures.
First, M2 has the worst results in all cases, illustrating the significant
impact of the pre-training step for graph neural networks. Second, the
performance of M3 is better than M1 and M2, which indicates the effectiveness
of combining both the pre-training and few-shot learning strategies. Third,
adding different self-supervised components (bond reconstruction and atom type
prediction) can further improve model performance, as reflected by the better
performances of M4, M5, and M6 over M3. Among these three variants, M6 has the
best performance as it adds both self-supervised tasks. Additionally, M7
outperforms M3, demonstrating the benefit of incorporating task-attention
weight into the meta-learning process. At last, M8 (the proposed model) has
the best performances in most cases, which shows the best capability graph
neural network model trained by meta-learning process and augmented with both
self-supervised module and task-aware attention. According to these findings,
we can conclude that different model components indeed bring benefits to model
design and improve performance.

###  5.7. Case Study of Embedding Visualization

To better show the effectiveness of our model, we visualize the molecular
embeddings generated by our proposed Meta-MGNN, PreGNN (Hu et al., 2020) and
MAML (Finn et al., 2017) using t-SNE (Maaten and Hinton, 2008), which are
shown in Figure 3. Specifically, it shows the embedding result of testing
datasets from SR-MMP (a molecular property). The blue plots and orange plots
represent molecules without SR-MMP property and with SR-MMP property,
respectively. It can be observed that our model achieves better performance in
discriminating two kinds of molecules than the other two models. In Figure
3(a), the bottom left corner of the figure is mostly occupied by orange dots
and the blues ones are mostly in the upper right corner. However, most orange
plots are mixed with blue plots in Figure 3(b) and 3(c).

##  6\. Conclusions

In this work, we proposed a few-shot learning approach for the molecular
property prediction problem, which is important and has not been well studied.
We proposed a novel model called Meta-MGNN. Meta-MGNN utilized a graph neural
network (with pre-training) to learn molecular embeddings and further employed
a meta-learning process to learn well-initialized model parameters that could
be fast adapted to new molecular properties with few-shot data samples. A
self-supervised module and self-attentive task weight were further proposed
and incorporate into the meta-learning framework, which benefited the whole
model. We evaluated our model on two public multi-task datasets and the
comparison of the experimental results showed that our model can outperform
state-of-art methods. The effectiveness of each model component was also
verified. The initial success of this study suggests following studies. The
future work might consider better task embedding formulation when computing
task weight in meta-learning. It is also possible to fuse both graph and
sequence model to learn molecular embeddings for the meta-learning process.

## Acknowledgements

This work was supported in part by National Science Foundation grant
CCI-1925607. We thank the anonymous referees for their valuable comments and
helpful suggestions.

## References

  * (1)
  * Altae-Tran et al. (2017) Han Altae-Tran, Bharath Ramsundar, Aneesh S Pappu, and Vijay Pande. 2017.  Low data drug discovery with one-shot learning. In _ACS Central Science_.
  * Bertinetto et al. (2018) Luca Bertinetto, Joao F Henriques, Philip HS Torr, and Andrea Vedaldi. 2018.  Meta-learning with differentiable closed-form solvers. In _International Conference for Learning Representation (ICLR)_.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018.  Bert: Pre-training of deep bidirectional transformers for language understanding. In _Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT)_.
  * Fan et al. (2019) Wenqi Fan, Yao Ma, Qing Li, Yuan He, Eric Zhao, Jiliang Tang, and Dawei Yin. 2019\.  Graph neural networks for social recommendation. In _The Web Conference_.
  * Finn et al. (2017) Chelsea Finn, Pieter Abbeel, and Sergey Levine. 2017\.  Model-agnostic meta-learning for fast adaptation of deep networks. In _International Conference on Machine Learning (ICML)_.
  * Garcia and Bruna (2018) Victor Garcia and Joan Bruna. 2018.  Few-shot learning with graph neural networks. In _International Conference for Learning Representation (ICLR)_.
  * Gilmer et al. (2017) Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. 2017\.  Neural message passing for quantum chemistry. In _International Conference on Machine Learning (ICML)_.
  * Guo et al. (2020) Zhichun Guo, Wenhao Yu, Chuxu Zhang, Meng Jiang, and Nitesh Chawla. 2020.  GraSeq: graph and sequence fusion learning for molecular property prediction. In _International Conference on Information and Knowledge Management (CIKM)_.
  * Hamilton et al. (2017) Will Hamilton, Zhitao Ying, and Jure Leskovec. 2017\.  Inductive representation learning on large graphs. In _Advances in Neural Information Processing Systems_.
  * Hu et al. (2020) Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. 2020.  Strategies for pre-training graph neural networks. In _International Conference for Learning Representation (ICLR)_.
  * Kim et al. (2019) Jongmin Kim, Taesup Kim, Sungwoong Kim, and Chang D Yoo. 2019\.  Edge-labeling graph neural network for few-shot learning. In _IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_.
  * Kipf and Welling (2017) Thomas N Kipf and Max Welling. 2017.  Semi-supervised classification with graph convolutional networks. In _International Conference for Learning Representation_.
  * Kuhn et al. (2016) Michael Kuhn, Ivica Letunic, Lars Juhl Jensen, and Peer Bork. 2016.  The SIDER database of drugs and side effects. In _Nucleic Acids Research_.
  * Landrum (2013) Greg Landrum. 2013\.  Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling.
  * Lee and Choi (2018) Yoonho Lee and Seungjin Choi. 2018.  Gradient-based meta-learning with learned layerwise metric and subspace. In _International Conference on Machine Learning_.
  * Lin et al. (2017) Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. 2017.  A structured self-attentive sentence embedding. In _International Conference for Learning Representation (ICLR)_.
  * Liu et al. (2019) Ke Liu, Xiangyan Sun, Lei Jia, Jun Ma, Haoming Xing, Junqiu Wu, Hua Gao, Yax Sun, Florian Boulnois, and Jie Fan. 2019.  Chemi-Net: a molecular graph convolutional network for accurate drug property prediction. In _International Journal of Molecular Sciences_.
  * Lu et al. (2019) Chengqiang Lu, Qi Liu, Chao Wang, Zhenya Huang, Peize Lin, and Lixin He. 2019\.  Molecular property prediction: A multilevel quantum interactions modeling perspective. In _AAAI Conference on Artificial Intelligence (AAAI)_.
  * Maaten and Hinton (2008) Laurens van der Maaten and Geoffrey Hinton. 2008.  Visualizing data using t-SNE. In _Journal of Machine Learning Research_.
  * Mansimov et al. (2019) Elman Mansimov, Omar Mahmood, Seokho Kang, and Kyunghyun Cho. 2019.  Molecular geometry prediction using a deep generative graph neural network. In _Scientific Reports_.
  * Mikolov et al. (2013) Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. 2013.  Distributed representations of words and phrases and their compositionality. In _Advances in Neural Information Processing Systems (NeurIPS)_.
  * Riniker and Landrum (2013) Sereina Riniker and Gregory A Landrum. 2013.  Similarity maps-a visualization strategy for molecular fingerprints and machine-learning methods. In _Journal of Cheminformatics_.
  * Sliwoski et al. (2014) Gregory Sliwoski, Sandeepkumar Kothiwale, Jens Meiler, and Edward W Lowe. 2014.  Computational methods in drug discovery. In _Pharmacological Reviews_.
  * Song et al. (2019) Weiping Song, Zhiping Xiao, Yifan Wang, Laurent Charlin, Ming Zhang, and Jian Tang. 2019\.  Session-based social recommendation via dynamic graph attention networks. In _The ACM International Conference on Web Search and Data Mining_.
  * Sung et al. (2018) Flood Sung, Yongxin Yang, Li Zhang, Tao Xiang, Philip HS Torr, and Timothy M Hospedales. 2018.  Learning to compare: relation network for few-shot learning. In _IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_.
  * Vamathevan et al. (2019) Jessica Vamathevan, Dominic Clark, Paul Czodrowski, Ian Dunham, Edgardo Ferran, George Lee, Bin Li, Anant Madabhushi, Parantu Shah, Michaela Spitzer, et al. 2019\.  Applications of machine learning in drug discovery and development. In _Nature Reviews Drug Discovery_.
  * Veličković et al. (2018) Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. 2018.  Graph attention networks. In _International Conference for Learning Representation (ICLR)_.
  * Vinyals et al. (2016) Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. 2016\.  Matching networks for one shot learning. In _Advances in Neural Information Processing Systems (NeurIPS)_.
  * Wang et al. (2019) Sheng Wang, Yuzhi Guo, Yuhong Wang, Hongmao Sun, and Junzhou Huang. 2019.  SMILES-BERT: Large scale unsupervised pre-training for molecular property prediction. In _ACM International Conference on Bioinformatics, Computational Biology, and Health Informatics (BCB)_.
  * Waring et al. (2015) Michael J Waring, John Arrowsmith, Andrew R Leach, Paul D Leeson, Sam Mandrell, Robert M Owen, Garry Pairaudeau, William D Pennie, Stephen D Pickett, Jibo Wang, et al. 2015\.  An analysis of the attrition of drug candidates from four major pharmaceutical companies. In _Nature Reviews Drug Discovery_.
  * Weininger et al. (1989) David Weininger, Arthur Weininger, and Joseph L Weininger. 1989\.  SMILES. 2. Algorithm for generation of unique SMILES notation. In _Journal of Chemical Information and Computer Sciences_.
  * Wu et al. (2020) Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. 2020\.  A comprehensive survey on graph neural networks. In _IEEE Transactions on Neural Networks and Learning Systems_.
  * Wu et al. (2018) Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu, Karl Leswing, and Vijay Pande. 2018\.  MoleculeNet: a benchmark for molecular machine learning. In _Chemical Science_.
  * Xu et al. (2019) Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. 2019\.  How powerful are graph neural networks?. In _International Conference for Learning Representation (ICLR)_.
  * Xu et al. (2017) Zheng Xu, Sheng Wang, Feiyun Zhu, and Junzhou Huang. 2017\.  Seq2seq fingerprint: An unsupervised deep molecular embedding for drug discovery. In _ACM International Conference on Bioinformatics, Computational Biology, and Health Informatics (BCB)_.
  * Yao et al. (2019) Huaxiu Yao, Yiding Liu, Ying Wei, Xianfeng Tang, and Zhenhui Li. 2019.  Learning from multiple cities: A meta-learning approach for spatial-temporal prediction. In _The Web Conference (WWW)_.
  * Yao et al. (2020) Huaxiu Yao, Chuxu Zhang, Ying Wei, Meng Jiang, Suhang Wang, Junzhou Huang, Nitesh Chawla, and Zhenhui Li. 2020\.  Graph few-shot learning via knowledge transfer. In _AAAI Conference on Artificial Intelligence (AAAI)_.
  * Yu et al. (2020) Wenhao Yu, Mengxia Yu, Tong Zhao, and Meng Jiang. 2020\.  Identifying referential intention with heterogeneous contexts. In _The Web Conference (WWW)_.
  * Zhang et al. (2019) Chuxu Zhang, Dongjin Song, Chao Huang, Ananthram Swami, and Nitesh V Chawla. 2019.  Heterogeneous graph neural network. In _ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD)_.
  * Zhang et al. (2020a) Chuxu Zhang, Huaxiu Yao, Chao Huang, Meng Jiang, Zhenhui Li, and Nitesh V Chawla. 2020a.  Few-shot knowledge graph completion. In _AAAI Conference on Artificial Intelligence (AAAI)_.
  * Zhang et al. (2020b) Chuxu Zhang, Lu Yu, Mandana Saebi, Meng Jiang, and Nitesh Chawla. 2020b.  Few-shot multi-hop relation reasoning over knowledge bases. In _Conference on Empirical Methods in Natural Language Processing: Findings (EMNLP: Findings)_.
  * Zhang et al. (2018a) Ruixiang Zhang, Tong Che, Zoubin Ghahramani, Yoshua Bengio, and Yangqiu Song. 2018a.  Metagan: An adversarial approach to few-shot learning. In _Advances in Neural Information Processing Systems (NeurIPS)_.
  * Zhang et al. (2018b) Xiaoyu Zhang, Sheng Wang, Feiyun Zhu, Zheng Xu, Yuhong Wang, and Junzhou Huang. 2018b.  Seq3seq fingerprint: towards end-to-end semi-supervised deep drug discovery. In _ACM International Conference on Bioinformatics, Computational Biology, and Health Informatics (BCB)_.
  * Zhao et al. (2020) Tong Zhao, Bo Ni, Wenhao Yu, and Meng Jiang. 2020\.  Early anomaly detection by learning and forecasting behavior. In _arXiv preprint arXiv:2010.10016_.
  * Zheng et al. (2019) Shuangjia Zheng, Xin Yan, Yuedong Yang, and Jun Xu. 2019\.  Identifying structure–property relationships through SMILES syntax analysis with self-attention mechanism. In _Journal of Chemical Information and Modeling_.
  * Zhou et al. (2019) Fan Zhou, Chengtai Cao, Kunpeng Zhang, Goce Trajcevski, Ting Zhong, and Ji Geng. 2019\.  Meta-GNN: On few-shot node classification in graph meta-learning. In _International Conference on Information and Knowledge Management (CIKM)_.


# Comparing Observation and Action Representations for
Deep Reinforcement Learning in $\mu$RTS

Shengyi Huang and Santiago Ontañón
Drexel University, Philadelphia, Pennsylvania, 19104
<EMAIL_ADDRESS>
###### Abstract

This paper presents a preliminary study comparing different observation and
action space representations for Deep Reinforcement Learning (DRL) in the
context of Real-time Strategy (RTS) games. Specifically, we compare two
representations: (1) a global representation where the observation represents
the whole game state, and the RL agent needs to choose which unit to issue
actions to, and which actions to execute; and (2) a local representation where
the observation is represented from the point of view of an individual unit,
and the RL agent picks actions for each unit independently. We evaluate these
representations in $\mu$RTS showing that the local representation seems to
outperform the global representation when training agents with the task of
harvesting resources.

## Introduction

Real-time strategy (RTS) games pose a significant challenge for artificial
intelligence (AI) (?; ?). They are complex due to a variety of reasons: (1)
players need to issue actions in real-time, which means agents have a very
limited time to produce what is the next action to execute, (2) most RTS games
are partially observable, i.e., a player might not always able to observe the
opponents’ strategies and actions, and (3) RTS games have very large action
spaces.

Recent application of Deep Reinforcement Learning (DRL) to RTS games
introduces additional challenges such as (1) dealing with extremely sparse
rewards and (2) designing efficient observation and action space
representations. In this paper, we attack the latter challenge by comparing
two intuitive representations. The first one is a global representation that
enables the DRL agent to observe the entire game state as a series of
features, and issue global commands to choose both which unit to issue actions
to and which actions to execute. The second one is a local representation
where the DRL agent sequentially selects actions for individual units
independently. We used $\mu$RTS as our testbed to evaluate the performance of
these representations and our experiments show that: (1) the local
representation seems to outperform global representation, probably because the
agent using local representation does not have to learn how to select a unit;
(2) both local and global representation significantly outperform a random
baseline agent.

The remaining of this paper is organized as follows. In the next section, we
will present some background on RTS games and reinforcement learning. Then, we
present our training methods, experiments setup, and discussion of training
results. Lastly, we make our conclusion and consider the future works.

## Background

This section provides the necessary background on RTS games, and reinforcement
learning in RTS games.

### Real Time Strategy Games

Real-time Strategy (RTS) games are complex adversarial domains, typically
simulating battles between a large number of military units, that pose a
significant challenge to both human and artificial intelligence (?). Designing
AI techniques for RTS games is challenging because (1) they have huge decision
spaces: the branching factor of a typical RTS game, StarCraft, has been
estimated to be on the order of $10^{50}$ or higher (?); (2) they are real-
time, which means that these games typically execute at 10 to 50 decision
cycles per second, leaving players with just a fraction of a second to decide
the next action, players can issue actions simultaneously, and actions are
durative, and (3) most of them are partially observable and players must sent
units to scout the map.

Figure 1: A screenshot of $\mu$RTS. Square units are “bases” (light grey, that
can produce workers), “barracks” (dark grey, that can produce military units),
and “resources mines” (green, from where workers can extract resources to
produce more units), the circular units are “workers” (small, dark grey) and
military units (large, yellow or light blue).

In the experiments reported in this paper, we employed
$\mu$RTS111https://github.com/santiontanon/microrts, a simple RTS game
maintaining the essential features that make RTS games challenging from an AI
point of view: simultaneous and durative actions, large branching factors and
real-time decision making. Although the game can be configured to be partially
observably and non-deterministic, those settings are turned off for all the
experiments presented in this paper.

A screenshot of the game is shown in Figure 1. The squared units in green are
Minerals with numbers on them indicating the remaining resources. The units
with blue outline belong to player 1 and those with red outline belong to
player 2. The light grey squared units are Bases with numbers indicating the
amount of resources owned by the player, while the darker grey squared units
are the Barracks. Movable units have round shapes with grey units being
Workers, orange units being Lights, yellow being Heavy units (not shown in the
figure) and blue units being Ranged.

Moreover, in order to ease reinforcement learning research in $\mu$RTS, we
have prepared an OpenAI gym wrapper for the game, which we have made available
to the research community222https://github.com/vwxyzjn/gym-microrts.

### Reinforcement Learning in RTS Games

One of the earliest applications of reinforcement learning to RTS games is the
work of Concurrent Hierarchical Reinforcement Learning (CHRL) with Q-learning
(?). Marthi et al. experimented with the Wargus game and demonstrated the use
of Alisp language to specify a list of desired tasks and use Q-learning to
tune the parameters. They showed that the CHRL could significantly improve the
agents’ performance compared to flat concurrent hierarchical Q-learning.

In addition, Jaidee and Muñoz-Avila have leveraged Q-learning to learn a
policy for each class of units (peasants, knights, barracks, etc), which
drastically reduced the memory requirement of Q-learning (?). The trade-off is
the lack of coordination between agents since each agent only learns to
control a certain type of units or buildings. That being said, they were able
to demonstrate good performance by defeating the built-in AI 80% of the time.

Another idea explored in RTS games is that of options, which are temporally
extended actions (?). The agent first needs to evaluate which option to
choose, and then needs to determine when to terminate the options chosen. The
options framework also exhibits a sense of Hierarchical Reinforcement Learning
(HRL) since raw actions are primitives compared to options. Researchers have
tried to combine options and heuristic algorithms to simplify the game space
(?).

In recent years, Deep Reinforcement Learning (DRL) approaches have received
significant attention in a number of games such as classic Atari games (?) and
Go (?). Over the years, researchers have introduced different type of DRL
algorithms. Most notably, Asynchronous Advantage Actor-Critic (A3C), Deep
Q-Network (DQN), and Proximal Policy Gradient (PPO) have gained the state-of-
the-art results in a variety of games (?; ?; ?).

In 2017, researchers at Deepmind started tackling the challenge of training
agents for StarCraft II, one of the most complex and popular RTS games in
history (?). They conducted a series of experiments using A3C (?) on the full
game as well as a collection of mini-games such as training the “marines”
units to defeat the “roaches” units in StarCraft II. They were able to show
good converging performance for the mini-games but had no material success in
the full game setting. However, they recently demonstrated a bot “AlphaStar”
that defeats a professional StarCraft II Protoss player in the full game
setting.

Researchers at Facebook also joined this line of research by publishing an
open-sourced RTS game research platform ELF (?). For simplicity, they
conducted experiments utilizing hard-coded actions such as BUILD_BARRACKS,
which “automatically picks a worker to build barracks at an empty location, if
the player can afford”. Through these high-level actions, they trained agents
using A3C, curriculum training, MCTS, and show the trained agent is able to
defeat the built-in scripted bots in the full-game setting. Notably,
researchers at Tencent used similar approach to hard-code high-level actions
in StarCraft II and defeated the cheating AI in the full-game setting by
commencing early attacks (?). Compared to this work, we are interested in RL
settings that do not require hand-crafted macro actions.

Moreover, Liu et al. studied the application of HRL to incorporate the
generation of macro actions through expert replays (?). Lee et al. also hard-
coded a collection of macro actions for zergs and trained agents combining DQN
and LSTM, obtaining a 83% winning rate in the AIIDE 2017 StarCraft AI
competition (?).

To address the action coordination problem, Samvelyan, Rashid et al. built the
StarCraft Multi-Agent Challenge (SMAC) for Multi-agent Reinforcement Learning
(MARL), where each unit is controlled by a separate agent (?). They compared
popular MARL algorithms such as QMIX, COMA, VDN, and IQL in mini maps that
usually features battle between different agents (?; ?; ?; ?).

Though these previous research show important results, there has not been any
published study comparing the relative merits of different types of
observation and action representations in RTS games, which is the key issue
that our paper starts to address.

Global Representation
---
Features | Possible values
Hit Points | 0, 1, 2, 3, 4, 5, $\geq 6$
Resources | 0, 1, 2, 3, 4, 5, $\geq 6$
Owner | player 1, -, player 2
Unit Types | -, resource, base, barrack, worker, light, heavy
Action | -, move, harvest, return, produce, attack

Local Representation
---
Features | Possible values
Hit Points | wall, 0, 1, 2, 3, 4, 5, $\geq 6$
Resources | wall, 0, 1, 2, 3, 4, 5, $\geq 6$
Owners | wall, player 1, -, player 2
Unit Types | wall, -, resource, base, barrack, worker, light, heavy
Action | wall, -, move, harvest, return, produce, attack

Table 1: The list of feature maps and their descriptions.

## Our Approach

We designed two pairs of observation and action representation for comparing
performance. The first one is a global representation that enables the RL
agent to observe the entire game and issue global commands to choose both
which unit to issue actions to and which actions to execute. The second one is
a local representation where the agent is externally given a specific unit to
control and just has to pick actions for that unit. Both representations are
designed to train agents to harvest resources as fast as possible.

### Global Representation

We utilized _flattened one-hot_ feature maps as our observation representation
similar to PySC2 (?) and a couple pre-existing works on observation evaluation
functions on $\mu$RTS (?; ?). Assuming the game map is represented as a grid,
a feature map is a matrix of the same dimensions as the game map, where in
each cell of the matrix we have a value that represents some aspect of the
corresponding game map grid cell. For simplicity, we assume that each feature
map takes discrete values. The exhaustive list of feature maps and their
description is presented at Table 1.

Let $h$ and $w$ be the height and width of the map respectively, we construct
the observation as a set of $n_{f}=5$ feature maps and that all feature maps
can take the same number of values $n_{c}$. Thus, each feature map is a vector
of length $h*w$, and the observation can be represented as a
$n_{f}\times(h*w)$ matrix. Finally, in order to facilitate learning, we use a
one-hot-encoding representation, and thus the observation is represented as a
$n_{f}\times(h*w)\times n_{c}$ tensor ($n_{c}$ is the number of values of the
feature plane that can take the larger number of values). The dimensions of
these tensors are usually referred to as their shape in common deep learning
libraries, and thus we will use this term from now on. Moreover, in the
experiments reported in this paper $n_{f}=5$ and $n_{c}=7$.

Then, to execute an action in $\mu$RTS, the RL agent needs to specify the
unit, the action types, and the parameters of the action to execute. For our
experiments that only focus on the task of harvesting resources, the set of
target action types are NOOP (no operation, represented simply as a - in Table
1), Move, Harvest, Return. Some of these action types requires an action
parameter that specifies the direction at which the action is issued. The list
of available action parameter is Up, Right, Down, Left. So if the agent wants
to command a worker to harvest the resources located left to it, the agent has
to specify the worker, set action type to Harvest, and action parameter to
Left.

Since unit IDs change from game to game, in order to learn generalizable
policies, the global representation does not require the agent to specify the
unit ID. Instead, the RL agent predicts the coordinates of the unit it wants
to control.

Therefore, we construct the _action_ at time step $t$ as a vector with 4
elements: $\left[a_{t}^{\text{x}},a_{t}^{\text{y}},a_{t}^{\text{action
type}},a_{t}^{\text{action param}}\right]$, where
$a_{t}^{\text{x}},a_{t}^{\text{y}},a_{t}^{\text{action
type}},a_{t}^{\text{action param}}$ are integers signifying the selected
x-coordinate, y-coordinate, action type, and action parameters, respectively,
predicted using one-hot encoding representations. Moreover:

  1. 1.

If the action produced is invalid in $\mu$RTS, the action will be replaced by
a NOOP action. Thus, it is up to the RL agent to learn which are the valid
actions.

  2. 2.

This representation varies from PYSC2’s action representation that uses the
actions themselves to select units and perform actions on them.

  3. 3.

The agent can only issue one action to one unit at each time step.

In summary, the global representation encodes the entire game state for the RL
agent, who needs to learn to choose which unit to control (by outputting the
unit’s coordinates) and which action type and parameter to be issued at each
time step.

Figure 2: The local feature maps of shape $(2w+1)\times(2w+1)=3\times 3$
outlined in red of the an unit (the red circle) when $w=1$. The blue crosses
indicate the cells are the walls of the game map.

### Local Representation

In this representation, the agent perceives the game from the point of view of
a particular unit and only needs to issue action to that unit. In other words,
the agent does _not_ have to learn to pick which unit to control.

We utilized flattened one-hot _local_ feature maps as our observation
representation. Consider a parameter of window size $w$ that specifies “how
far can an unit see”. Specifically, a window size $w=1$ means each unit is
able to observe the cell it is located and all the cells that are 1 distance
away from it in either axis, as shown in 2. Note that if the unit sees a cell
beyond the boundary of the map, we consider that cell a “Wall” (Note that the
global representation does not need such “Wall” unit type because the agent
always sees the entire map and does _not_ see beyond the boundary of the map).
A local feature map is thus a matrix of dimension $(2w+1)\times(2w+1)$. A
comprehensive list of local feature maps and their descriptions can be found
in Table 1.

At each time step, the agent is given a target unit (we rotate which unit is
the focus at consecutive game frames). For instance, assuming there are 3
units $u_{1},u_{2},u_{3}$, then $u_{1},u_{2},u_{3},u_{1},u_{2},...$ will be
selected at time step $1,2,3,4,5,...$, respectively. Given the selected unit
at each time step, we construct _observation_ to be a set of $n_{f}$ local
feature maps from the point of view of the said unit according to $w$, with
all of them can take the same number of discrete values $n_{c}$. Similarly, we
use one-hot-encoding representation and the observation is represented as a
$n_{f}\times(2w+1)^{2}\times n_{c}$ tensor, where $n_{f}=5,n_{c}=8$.

For action execution, since the unit is already selected, the agent only needs
to predict the action type and parameter. Therefore, we construct the _action_
at time step $t$ as a vector with 2 element: $\left[a_{t}^{\text{action
type}},a_{t}^{\text{action param}}\right]$, where $a_{t}^{\text{action
type}},a_{t}^{\text{action param}}$ are integers that signifies the selected
action type, and action parameters, respectively (predicted as one-hot
vectors). As before, if the action issued is not valid in $\mu$RTS, it will be
replaced by a NOOP action.

In summary, the local representation contains the local game state features
from the point of view of a selected unit for the RL agent, and it needs to
learn which action type and parameter to be issued at each time step. Note
that it does _not_ have to learn how to select a unit to control, but only how
to interact with its surrounding observations.

Figure 3: The neural network architecture that demonstrates the flow from the
observation vector to action probabilities. The number at each box suggests
the input or output shapes

### Reward Function

The game environment will give the agent a reward of 10 when a worker has
successfully harvested resources and another 10 when it returns the harvested
resources back to the base. Otherwise the game environment gives a reward of
0.

It’s worth pointing out that, like in many other reinforcement learning
problems, rewards can be very sparse as maps grow and workers get further away
from the resources and the base, requiring thousands of exploration steps
before getting a non-zero reward.

## Training of Agents

Figure 4: The mini-games that focuses on harvesting resources with different
map sizes of $4\times 4$, $6\times 6$, and $8\times 8$.

For the global representation, we selected Advantage Actor-Critic (A2C), one
of the most popular RL algorithms, to evaluate the agents’ performance (?). We
created the Policy Network and Value Networks. As shown in Figure 3 that
demonstrates the forward pass, the observation tensor is fed into the policy
network that produces a dense vector of shape $(h\cdot w\cdot n_{a}\cdot
n_{p})$, which is split into 4 different subvectors of logits. Each of those
logits is passed into a Categorical Distribution to gain the probability
distribution for sampling values to produce the a combined action vector
$a=\begin{bmatrix}a_{t}^{\text{x}}&a_{t}^{\text{y}}&a_{t}^{\text{action
type}}&a_{t}^{\text{action param}}\end{bmatrix}$. In addition, the observation
vector is also used by the value network to create a valuation of the current
observation.

Parameter Names | Parameter Values
---|---
Maps |  $4\times 4$, $6\times 6$, $8\times 8$
Learning Rate of Adam Optimizer | 0.0007
Random Seed | 1, 2, 3
Episode Length | 2,000 time steps
Total Time steps | 2,000,000 time steps
$w$ (Local Representation Window Size) | 1, 2
$\gamma$ (Discount Factor) | 0.99
$\beta$ (Value Function Coefficient) | 0.25
$\eta$ (Entropy Regularization Coefficient) | 0.01
$\omega$ (Gradient Norm Threshold ) | 0.5
Table 2: The list of experiment parameters and their values.

Given the forward pass, we run the game for a full episode, record the reward
$r_{t}$ at each time step $t$ and compute

$\displaystyle\log\pi_{\theta}(a_{t}|s_{t})$
$\displaystyle=\log\pi_{\theta}(a_{t}^{\text{x}}|s_{t})$
$\displaystyle+\log\pi_{\theta}(a_{t}^{\text{y}}|s_{t})$
$\displaystyle+\log\pi_{\theta}(a_{t}^{\text{action type}}|s_{t})$
$\displaystyle+\log\pi_{\theta}(a_{t}^{\text{action param}}|s_{t})$

where $s_{t}$ is the observation tensor at time step $t$. In addition,
analogously compute the entropy
$\sum_{a}\pi_{\theta}(a|s_{t})\log\pi_{\theta}(a|s_{t})$ as the sum of entropy
for $a_{t}^{\text{x}},a_{t}^{\text{y}},a_{t}^{\text{action type}}$ and
$a_{t}^{\text{action param}}$. After each episode finishes, we could train the
agent based on the following gradient

$\displaystyle\underbrace{(G_{t}-v_{\theta^{\prime}}(s_{t}))\nabla_{\theta}\log\pi_{\theta}(a_{t}|s_{t})}_{\text{policy
gradient}}+$
$\displaystyle\underbrace{\beta(G_{t}-v_{\theta^{\prime}}(s_{t}))\nabla_{\theta^{\prime}}v_{\theta^{\prime}}(s_{t})}_{\text{value
estimation gradient}}+$
$\displaystyle\eta\underbrace{\sum_{a}\pi_{\theta}(a|s_{t})\log\pi_{\theta}(a|s_{t})}_{\text{entropy
regularization}}$

where $\theta^{\prime}$ and $\theta$ are weights for the value network and the
policy network, respectively, $\beta$ and $\eta$ are hyperparameters for
controlling the contribution of value gradients and entropy regularization
gradients, and $G_{t}$ is the discounted rewards from step $t$. A
comprehensive list of training parameters and their values can be found at
Table 2. For more details on the algorithm, please refer to the original paper
(?).

Regarding the local representation, we used a similar Policy Network
architecture that produces a dense vector of shape $(n_{a}\cdot n_{p})$. We
similarly calculate

$\log\pi_{\theta}(a_{t}|s_{t})=\log\pi_{\theta}(a_{t}^{\text{action
type}}|s_{t})+\log\pi_{\theta}(a_{t}^{\text{action param}}|s_{t})$

and the entropy $\sum_{a}\pi_{\theta}(a|s_{t})\log\pi_{\theta}(a|s_{t})$ as
the sum of entropy for $a_{t}^{\text{action type}}$ and $a_{t}^{\text{action
param}}$; then use A2C to train the agent.

## Experimental Evaluation

In this section, we compare the aforementioned approaches on maps of different
sizes.

### Experimental Setup

We created three maps of size $4\times 4$, $6\times 6$, and $8\times 8$, all
of which are shown in the Figure 4. They all have two workers, a block
containing 230 resources, a base for the workers to return the harvested
resources, and a dummy enemy base that is useless. Note that it takes a worker
10 time steps to harvest a resource and another 10 time steps to return the
resources to the base, excluding the time steps to move around the map. So in
the most efficient harvesting setup as shown in the $4\times 4$ map, where
each worker is right next to the base and the resources, the most resources
that the two workers can gather is 2 per 20 frames. In fact, since both
representations only accept one action at each time step, the two workers have
to wait until 21 frames for the resources to be returned. Simple math shows
that $2\cdot 2000/20=200$ is the most that the two workers can harvest within
the 2,000 time steps per episode. A comprehensive list of experimental
parameters and their values are presented in Table 2. For values of
hyperparameters, we simply used the default values from the OpenAI’s A2C
implementation (?).

### Experimental Results

| map | $t_{\text{first harvest}}$ | $t_{\text{first return}}$ | $r$
---|---|---|---|---
RandomAI | $4\times 4$ | 51.33 | 142.67 | 11.87
Global | $4\times 4$ | 99.00 | 167.73 | 13.13
Local ($w=1$) | $4\times 4$ | 29.87 | 172.47 | 67.20
Local ($w=2$) | $4\times 4$ | 45.00 | 73.73 | 33.40
RandomAI | $6\times 6$ | 421.33 | 797.33 | 2.00
Global | $6\times 6$ | 533.33 | 1931.20 | 0.07
Local ($w=1$) | $6\times 6$ | 59.20 | 567.40 | 3.53
Local ($w=2$) | $6\times 6$ | 62.33 | 408.73 | 3.93
RandomAI | $8\times 8$ | 878.67 | 1480.67 | 0.87
Global | $8\times 8$ | 1464.53 | - | 0.00
Local ($w=1$) | $8\times 8$ | 167.20 | 1844.20 | 0.20
Local ($w=2$) | $8\times 8$ | 89.87 | - | 0.00
Table 3: The list of representations and their performance according to our
metrics. The “-” in $t_{\text{first return}}$ indicates the agent never
returned any resources.

Figure 5 shows the average reward of runs with different random seeds per
episode as a function of training steps. As shown in the plots, local
representation with $w=1$ yields the best episode rewards in all three maps.
As mentioned previously, we suspect local representation is performing better
just because the agent doesn’t have to learn to pick a unit as the agent does
under the global representation, i.e., the units are pre-selected under the
local representation.

Figure 5: Episode rewards (y axis) as a function of training time steps
(x-axis) for the 3 map sizes.

The reason that local representation with $w=2$ is performing worse than local
representation with $w=1$ seems purely computational since the observation
tensor when $w=2$ is almost twice as the size when $w=1$. The other notable
characteristic is that local representation with $w=2$ seems to exhibit more
stable performance. Given more computational resources, we think local
representation with $w=2$ will eventually become more efficient than local
representation with $w=1$ because larger window size means the units are less
likely to move to a cell where the base or resources are outside of its local
feature maps (or in simpler terms, the unit moved to a cell that could not
observe the base or resources). In our experiments, the map sizes are
generally small and the position of the base and resources are stationary, so
the local representation with $w=1$ is probably enough for the unit to observe
sufficient information (and for learning to reach the resources, they just
need to learn to move to the top-left coordinate of the map).

Notably, the performance of both local and global representation suffer when
we scale the game with larger maps. As shown in the Figure 4, the larger maps
presents even challenging moving trajectories for the agents since the
distance between resources and the base is increased, which means sparser
rewards and larger exploration needed for the agents. Global representation is
especially worse off because the agent still has to learn to select units,
navigate, harvest, and return to base in the larger maps.

As a baseline, we created three metrics to evaluate the performance of agents
against the built-in RandomAI $\mu$RTS agent that uniformly selects a valid
action, either NOOP, Move, Harvest, or Return, and a valid parameter depending
on the game state. Notice that all actions produced by the RandomAI agent are
valid since it has the game rules hard-coded, whereas our reinforcement
learning agents have to learn which actions are valid. Thus, the RandomAI
agent has an unfair advantage. A way to level the playing field would be to
mask out invalid agents in the output of the reinforcement learning, and
consider the action with highest likelihood only among the valid ones.
However, we decided against giving our agent any domain knowledge beyond the
observation and action representation in this first set of experiments, just
to set a baseline.

Let $t_{\text{first harvest}}$ and $t_{\text{first return}}$ be the time steps
that the agent successfully harvested and return resources for the first time,
respectively, which will give us an idea of how much ”wandering” did the agent
did before harvesting and returning. Then, let $r$ be the number of resources
gathered in total, which will evaluate the agents’ performance as a whole. We
ran the trained agents with different seeds for 5 episodes (10,000 time steps)
and calculate the average $t_{\text{first harvest}}$, $t_{\text{first
return}}$, and $r$ across runs with different seeds. The evaluation result is
listed at Table 3. As shown, the trained agents perform significantly better
than the RandomAI agent except in the $8\times 8$ map where all agents perform
equally bad.

### Visual Behavior of Agents

When using RL to train agents to play games, it is easy to be distracted by
the numerical rewards and various metrics. However, watching the agent play in
the game remains the ultimate measure for agents’ ability. The following
paragraphs highlight some interesting behaviors of agents.

In $4\times 4$ maps, the agent under the local representation almost exhibits
optimal behavior where the two Workers harvest the resources and return them
restlessly. The agent under the global representation, however, seems to only
learn to control _one_ Worker. It seems the agent gets stuck in a local
minimum where it believes the maneuver of only one Worker is best for
harvesting resources.

In $6\times 6$ and $8\times 8$ maps, the agent under the local representation
struggles with returning the resources to the base. In fact, the agent tries
to harvest another unit of resources even though it already carries one. The
unit moves closer to the resources than to the base, but unfortunately it has
to return the resources before harvesting new ones. It almost seems that the
agent learned only to harvest resources, but has not learned how to return
them.

## Conclusions and Future Work

In this paper, we compared two different observation and action
representations in the context of RTS games: (1) a global representation that
feeds the agent an observation of feature maps that are as large as the game
map itself, and require the agent to learn to locate unit it is wants to issue
actions to and predict the intended action type and parameter; and (2) a local
representation that feeds the agent an observation of local feature maps that
are the feature maps of some distance away from the point of view of a
selected unit, and require the agent to learn which action type and parameter
to issue for the said unit. We train agents on maps that focus on harvesting
resources, establish some objective metrics to evaluate the agents’
performance and show that the local representation generally outperforms the
global representation. This advantage, however, does not necessarily hold in
larger maps, where the exploration and sparse rewards become a huge problem.

For future works, we plan to consider full-game setting, including attack and
produce actions, which have certain conditioning logic involved, i.e., the
base can only produce Workers while the barracks can only produce Light,
Heavy, and Ranged units. In addition, we hope to find approaches that scale
better with larger map sizes. Techniques such as imitation learning or using
curiosity as intrinsic rewards have been shown to help with exploring more of
the map (?). Moreover, we would like to assess whether our representation
comparison results generalize to other reinforcement learning algorithms.

## References

  * [Buro 2003] Buro, M.  2003\.  Real-time strategy games: A new ai research challenge.  In IJCAI, volume 2003, 1534–1535.
  * [Dhariwal et al. 2017] Dhariwal, P.; Hesse, C.; Klimov, O.; Nichol, A.; Plappert, M.; Radford, A.; Schulman, J.; Sidor, S.; Wu, Y.; and Zhokhov, P.  2017\.  Openai baselines.  https://github.com/openai/baselines.
  * [Foerster et al. 2018] Foerster, J. N.; Farquhar, G.; Afouras, T.; Nardelli, N.; and Whiteson, S.  2018\.  Counterfactual multi-agent policy gradients.  In Thirty-Second AAAI Conference on Artificial Intelligence.
  * [Jaidee and Muñoz-Avila 2013] Jaidee, U., and Muñoz-Avila, H.  2013\.  Modeling unit classes as agents in real-time strategy games.  Proceedings of the 9th AAAI Conference on Artificial Intelligence and Interactive Digital Entertainment, AIIDE 2013 149–155.
  * [Lee et al. 2018] Lee, D.; Tang, H.; Zhang, J. O.; Xu, H.; Darrell, T.; and Abbeel, P.  2018\.  Modular architecture for starcraft ii with deep reinforcement learning.  In Fourteenth Artificial Intelligence and Interactive Digital Entertainment Conference.
  * [Liu et al. 2019] Liu, R.-Z.; Guo, H.; Ji, X.; Yu, Y.; Xiao, Z.; Wu, Y.; Pang, Z.-J.; and Lu, T.  2019\.  Efficient reinforcement learning with a mind-game for full-length starcraft ii.  arXiv preprint arXiv:1903.00715.
  * [Marthi et al. 2005] Marthi, B.; Russell, S. J.; Latham, D.; and Guestrin, C.  2005\.  Concurrent hierarchical reinforcement learning.  In IJCAI, 779–785.
  * [Mnih et al. 2013] Mnih, V.; Kavukcuoglu, K.; Silver, D.; Graves, A.; Antonoglou, I.; Wierstra, D.; and Riedmiller, M.  2013\.  Playing atari with deep reinforcement learning.  arXiv preprint arXiv:1312.5602.
  * [Mnih et al. 2016] Mnih, V.; Badia, A. P.; Mirza, M.; Graves, A.; Lillicrap, T.; Harley, T.; Silver, D.; and Kavukcuoglu, K.  2016\.  Asynchronous methods for deep reinforcement learning.  In International conference on machine learning, 1928–1937.
  * [Ontanón et al. 2013] Ontanón, S.; Synnaeve, G.; Uriarte, A.; Richoux, F.; Churchill, D.; and Preuss, M.  2013\.  A survey of real-time strategy game ai research and competition in starcraft.  IEEE Transactions on Computational Intelligence and AI in games 5(4):293–311.
  * [Pathak et al. 2017] Pathak, D.; Agrawal, P.; Efros, A. A.; and Darrell, T.  2017\.  Curiosity-driven exploration by self-supervised prediction.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, 16–17.
  * [Rashid et al. 2018] Rashid, T.; Samvelyan, M.; de Witt, C. S.; Farquhar, G.; Foerster, J.; and Whiteson, S.  2018\.  Qmix: Monotonic value function factorisation for deep multi-agent reinforcement learning.  arXiv preprint arXiv:1803.11485.
  * [Samvelyan et al. 2019] Samvelyan, M.; Rashid, T.; de Witt, C. S.; Farquhar, G.; Nardelli, N.; Rudner, T. G.; Hung, C.-M.; Torr, P. H.; Foerster, J.; and Whiteson, S.  2019\.  The starcraft multi-agent challenge.  arXiv preprint arXiv:1902.04043.
  * [Schulman et al. 2017] Schulman, J.; Wolski, F.; Dhariwal, P.; Radford, A.; and Klimov, O.  2017\.  Proximal policy optimization algorithms.  arXiv preprint arXiv:1707.06347.
  * [Silver et al. 2017] Silver, D.; Schrittwieser, J.; Simonyan, K.; Antonoglou, I.; Huang, A.; Guez, A.; Hubert, T.; Baker, L.; Lai, M.; Bolton, A.; et al.  2017\.  Mastering the game of go without human knowledge.  Nature 550(7676):354.
  * [Stanescu et al. 2016] Stanescu, M.; Barriga, N. A.; Hess, A.; and Buro, M.  2016\.  Evaluating real-time strategy game states using convolutional neural networks.
  * [Stolle and Precup 2002] Stolle, M., and Precup, D.  2002\.  Learning options in reinforcement learning.  In International Symposium on abstraction, reformulation, and approximation, 212–223.  Springer.
  * [Sun et al. 2018] Sun, P.; Sun, X.; Han, L.; Xiong, J.; Wang, Q.; Li, B.; Zheng, Y.; Liu, J.; Liu, Y.; Liu, H.; et al.  2018\.  Tstarbots: Defeating the cheating level builtin ai in starcraft ii in the full game.  arXiv preprint arXiv:1809.07193.
  * [Sunehag et al. 2017] Sunehag, P.; Lever, G.; Gruslys, A.; Czarnecki, W. M.; Zambaldi, V.; Jaderberg, M.; Lanctot, M.; Sonnerat, N.; Leibo, J. Z.; Tuyls, K.; et al.  2017\.  Value-decomposition networks for cooperative multi-agent learning.  arXiv preprint arXiv:1706.05296.
  * [Tampuu et al. 2017] Tampuu, A.; Matiisen, T.; Kodelja, D.; Kuzovkin, I.; Korjus, K.; Aru, J.; Aru, J.; and Vicente, R.  2017\.  Multiagent cooperation and competition with deep reinforcement learning.  PloS one 12(4):e0172395.
  * [Tavares and Chaimowicz 2018] Tavares, A. R., and Chaimowicz, L.  2018\.  Tabular reinforcement learning in real-time strategy games via options.  2018 IEEE Conference on Computational Intelligence and Games (CIG) 1–8.
  * [Tian et al. 2017] Tian, Y.; Gong, Q.; Shang, W.; Wu, Y.; and Zitnick, C. L.  2017\.  Elf: An extensive, lightweight and flexible research platform for real-time strategy games.  In Advances in Neural Information Processing Systems, 2659–2669.
  * [Vinyals et al. 2017] Vinyals, O.; Ewalds, T.; Bartunov, S.; Georgiev, P.; Vezhnevets, A. S.; Yeo, M.; Makhzani, A.; Küttler, H.; Agapiou, J.; Schrittwieser, J.; et al.  2017\.  Starcraft ii: A new challenge for reinforcement learning.  arXiv preprint arXiv:1708.04782.
  * [Yang and Ontañón 2018] Yang, Z., and Ontañón, S.  2018\.  Learning map-independent evaluation functions for real-time strategy games.  2018 IEEE Conference on Computational Intelligence and Games (CIG) 1–7.