11institutetext: School of Computer Science and Applied Mathematics,
University of The Witwatersrand, Johannesburg
11email<EMAIL_ADDRESS>22institutetext: Institute for Intelligent
Systems, University of Johannesburg, Johannesburg
22email<EMAIL_ADDRESS>33institutetext: School of Physics, University of
The Witwatersrand, Johannesburg
33email<EMAIL_ADDRESS>44institutetext: MRC/Wits Rural Public Health
and Health Transitions Research Unit, School of Public Health, University of
the Witwatersrand, Johannesburg
44email<EMAIL_ADDRESS>
# Using Machine Learning To Fuse Verbal Autopsy Narratives And Binary Features
In The Analysis Of Deaths From Hyperglycaemia

Thokozile Manaka 11 0000-0001-9910-4480 Terence van Zyl 22 0000-0003-4281-630X
Deepak Kar 33 0000-0002-4238-9822 Alisha Wade 44 0000-0002-1158-2523

###### Abstract

Lower-and-middle income countries are faced with challenges arising from a
lack of data on cause of death (COD), which can limit decisions on population
health and disease management. A verbal autopsy (VA) can provide information
about a COD in areas without robust death registration systems. A VA consists
of structured data, combining numeric and binary features, and unstructured
data as part of an open-ended narrative text. This study assesses the
performance of various machine learning approaches when analyzing both the
structured and unstructured components of the VA report. The algorithms were
trained and tested via cross-validation in the three settings of binary
features, text features and a combination of binary and text features derived
from VA reports from rural South Africa. The results obtained indicate
narrative text features contain valuable information for determining COD and
that a combination of binary and text features improves the automated COD
classification task.

###### Keywords:

Diabetes Mellitus Verbal Autopsy Cause of Death Machine Learning Natural
Language Processing

##  1 Introduction

In 2014, the International Diabetes Federation (IDF) [1] estimated 175 million
undiagnosed cases of diabetes. This lack of diagnosis is compounded by the
absence of statistics on diabetes-related deaths in many lower-and-middle-
income countries (LMIC). The mortality data that is available is frequently
derived from hospital records, and it possibly understates mortality from
fatalities that occurring outside of hospital [2]. People with diabetes may
develop chronic complications such as neuropathy, nephropathy and retinopathy,
uncontrolled hyperglycaemia, and can also develop acute potentially fatal
complications such as diabetic ketoacidosis (DKA) and hyperosmolar
hyperglycaemic syndrome (HHS) [3, 4].

A verbal autopsy (VA) is a technique, endorsed by the World Health
Organization (WHO), to determine a likely cause of death (COD) in countries
without robust death registration systems [5]. It is a record of an interview
between a non-clinician field worker and a care taker of the deceased about
events around an uncertified death. A VA report consists of two kinds of
standardized questions: the closed-ended questions where the interviewee
responds to a ’yes’ or ’no’ question and the open-ended section where the
interviewee narrates the events around the period of death [6]. Like much
electronic data, an analysis of VA reports may benefit from the application of
machine learning techniques and while some studies suggest the narrative
section of the VA report is unnecessary and of limited use due to its high
dimensionality and sparsity, we hypothesized that this narrative text would
improve the ability of verbal autopsy reports to predict a cause of death. We
compared the binary features, text features and a combination of binary and
text features in VA reports from Agincourt, South Africa, for their accuracy
in classifying COD from uncontrolled hyperglycaemia. The machine learning
techniques of logistic regression, random forest, XGBoost and neural networks
were employed to automate COD classification in the three feature settings. We
relied on data labeled by physician experts [2] as our gold standard for
determining whether a death had been caused by uncontrolled hyperglycaemia
(positive; coded as 1) or not (negative; coded as 0).

##  2 Background

Two-thirds of the 60 million annual deaths in lower-and-middle-income
countries occur outside of health facilities, often in countries with weak
death registration systems and therefore do not have a recorded cause [7].
Information about cause of death (COD) is, however, vital for public health
researchers and policy makers and verbal autopsy technology was developed to
address this need.

Verbal autopsy, a tool in which the caretaker of the decedent is interviewed
about their health preceding their death, can be formulated as a semi-
supervised learning problem where binary or continuous questionnaire responses
form features of the disease and the label of the disease is the categorical
response [8]. VAs have proved to be highly specific ($98\%$) for various types
of CODs [9].

Several machine learning approaches have been utilized in a number of
research, including Nave Bayes, neural networks, support vector machines,
random forests, InterVA4, InsilicoVA, AdaBoost, XGBoost, and Tariff [10, 11,
12, 13] and have demonstrated the ability to automate the coding of COD from
binary features of VA reports. While some studies suggest the narrative
section of the VA report is unnecessary and of limited use in COD
determination [14], others argue that the narrative section is more convenient
as it does not require a speciﬁc questionnaire format and takes less time to
collect compared to the closed-ended questions [15]. Information in free-text
narrative, such as the history of the disease symptoms and treatment, may
often be essential to making a correct diagnosis.

Boag et al. [16] evaluated common representations of clinical text in a
variety of clinical tasks and demonstrated that the best method of
representing narratives from clinical text for prediction remains debatable.
This is due to the fact that word embeddings do not apprehend medical
knowledge but describe the meaning of a word based on its context. Grnarova et
al. [17] also showed that no text representation performs better than all
others and that simple representations outscore advanced representations on
tasks like age prediction, hospital admission while complex models perform
best on diagnosis and length of stay tasks.

Clinical text representation models, ClinicalBERT [18], MeDAL [19], Publicly
Available Clinical BERT Embeddings [20], BioBERT [21], BioELMo [22] and
SciBERT [23] have been created to improve natural language processing tasks of
predicting mortality and hospital readmission from clinical notes. Datasets on
which these state-of-the-art narrative text representations are trained
include a collection of clinical and electronic health records from hospitals,
MIMIC-II (Medical Information Mart for Intensive Care) [24], MIMIC-III, PubMed
abstracts and PMC full text articles and Wikipedia. The training tasks are
disease diagnosis, hospital readmission and mortality prediction.

Other work has focused on natural language processing in VA narratives and
numerous methods have been used to classify CoD from VA narratives including
term frequency and TF-IDF (term frequency–inverse document frequency) [25],
the Tariff method [26], the neural network classifier [15] and random forest
[27]. Danso et al. [13] demonstrated improved performances of classification
algorithms when using linguistic features of part-of-speech tags, noun phrases
and word pairs for COD diagnosis from VA reports in addition to frequency-
based features. Yan et al. [28] showed that using models with key phrases as
additional features outperforms topic models which depend on features present
in the corpus when using a multi-task learning model to learn COD categories.
Addition of character information appears to improve model classiﬁcation for
smaller datasets (500 to 1000 records) calling for character-based
convolutional neural networks (CNNs) and is a promising technique for
automated VA coding [29].

While the bag of words model excels at tasks that require it to predict
categories that are directly represented by single words in their notes, such
as the frequency with which the words ’diabetes’ and ’sugar’ might appear in a
VA report. This may make it a good model to predict diabetes as the cause of
death for a VA case, however, this approach ignores the order of words inside
a document and so fails to reveal all of the underlying conditions behind a
patient’s death.

Doc2vec has successfully been used to classify documents in various fields
including health care [30, 16]. Its variants, the continuous distributed bag
of words paragraph vector (PV-DBOW) and the distributed memory paragraph
vector (DM-PV) work on the idea that the element values of a word are
influenced by the values of other words in its vicinity, and this idea is
embodied as a neural network structure as with word2vec. On sentiment analysis
and text classification tasks, the model surpass the bag of words model and
advanced models based on recursive and recurrent neural networks, which is why
it was chosen for this study. While both the binary and VA text representation
techniques for COD automation outlined above used raw VA reports, we used the
traditional approach of physician evaluation of VA forms to incorporate
clinical knowledge of diabetes and uncontrolled hyperglycemia into our
embeddings.

The outcomes of this work will aid the public health sector in areas lacking
adequate death registration systems by improving automation of COD from binary
and text features of VA reports for research, surveillance, and diagnosis of
diabetes. Improved VA text representations will provide accurate information
regarding diabetes as a cause of mortality, an early detection of which will
prevent the lethal complication hyperglycemia.

##  3 Methodology and Research Instruments

###  3.1 Algorithms

Natural Language Processing Algorithms:

Doc2vec [30] with both its variants, the continuous distributed bag-of-words
paragraph vector (PV-DBOW) and the distributed memory paragraph vector (PV-
DM).

Classification Algorithms:

Logistic regression, artificial neural network, random forest and XGBoost.

Hardware and Software:

All experiments were run with Python v3.6 using Sci-kit Learn [31] for random
forest, XGBoost and logistic regression classifiers and Keras [32] with Theano
[33] backend for neural networks. Text preprocessing was done with the Natural
Language Tool Kit NLTK [34].

###  3.2 Dataset

The MRC/Wits Rural Public Health and Health Transitions Research Unit
(Agincourt) [35], a population health and demographic surveillance system
located in rural South Africa, that supports research into the origins and
effects of diseases on social transitions and populations, provided the verbal
autopsy (VA) dataset for this study. Initial coding of the data was done by a
paediatrician with expertise in diabetes who reviewed VAs to identify features
suggestive of diabetes or uncontrolled hyperglycaemia. Cases where the
reviewing physician was uncertain were reviewed with colleagues with
experience in adult internal medicine, diabetes, and endocrinology until an
agreement was reached (Table 1). Positive instances (79) were those in which
uncontrolled hyperglycemia was a likely cause of death and negative (8619)
cases were those in which it was not.

Table 1: Features of the Agincourt verbal autopsy dataset.

Feature | Type | Description
---|---|---
Female | Binary | The deceased’s gender
Tuber | Binary |  The dead had been diagnosed with tuberculosis by a medical facility
Diabetes | Binary |  The dead had been diagnosed with diabetes by a medical facility
Men-con | Binary |  The dead had signs of mental confusion and memory loss
Cough | Binary | The dead had a cough
Ch-Cough | Binary | The dead had a chronic cough
Diarr | Binary | Diarrhoea was present in the deceased
Exc-Urine | Binary | The dead urinated excessively
Exc-Drink | Binary | The dead drank water excessively
Age | Numeric | The dead’s age
Description | Text |  Narration of signs and symptoms of the deceased by relative
Class | Binary | Cause of death classification by diabetes

We classified binary features as symptoms to which responses were ’yes’ or
’no’. These included excessive thirst, urination, mental confusion and the
rest listed in (Table 1). The narration describing symptoms and events around
the death of the deceased was classified as a text feature.

’Yes’ and ’no’ responses were code to ’1s’ and ’0s’ as well as ’Female’ and
’Male,’ were converted to ’1s’ and ’0s’. To leverage clinical knowledge and
incorporate domain knowledge of uncontrolled hyperglycaemia into VA
embeddings, we used a physician-annotated dataset. The words ’diabetes,
’sugar’ and their misspelled forms were removed from the VA text to uncover
symptoms that distinguish cases with diabetes from those without. Text
processing included removing punctuation, special characters, short words and
stop words as well as lower casing the text. Each VA case includes a paragraph
of 4 to 6 sentences of variable length of narratives of instances surrounding
death. We tokenized the final text and tagged each case to its respective
class as its unique ID and fed it as input to the doc2vec model.

Taking from the Word2vec [36] model, given a sequence of words
$w_{1},w_{2},...,w_{T}$ the doc2vec model is trained to predict one of the
word’s vector $w_{t}$ given the other $T-k$ words’ vectors.

The model maximizes the average log probability

$\frac{1}{T}\sum_{t=k}^{T-k}\log p(w_{t}|w_{t-k},...,w_{t+k})$ (1)

and the task of prediction is

$p(w_{t}|w_{t-k},...,w_{t+k})=\frac{e^{{y}_{w_{t}}}}{\sum_{i}e^{y_{i}}}$ (2)

In addition the model also takes a paragraph vector $p_{i}$ with $i$
identifying which body The authors then propose cost-sensitive classification,
which differs from random oversampling in that it provides a more continuous-
valued and consistent method of weighting samples of imbalanced training data;
random oversampling, for example, will invariably favor some training
instances over otherof text $w_{1},w_{2},...,w_{T}$ comes from.

Each paragraph and word inside a paragraph were mapped to a unique vector and
the vectors were next concatenated for prediction of the next word forming a
continuous distributed vector representations for VA texts. The contexts were
drawn from a sliding window of a defined length. Stochastic gradient descent
and back propagation were used to train both word and paragraph vectors. The
size of embedding dimension and the length of the context window (the number
of words around the target word), were the top parameters for training Doc2vec
embeddings and were both tuned experimentally. The link to the code is
available at [37].

Table 2: Random forest performance on cause of death classification with feature vectors dimensions from PV-DBOW, PV-DM and their combinations. Method | Dims | Recall | Precision | F1-Score | AUC-ROC | Accuracy
---|---|---|---|---|---|---
PV-DM | 50 | 0.3702 | 0.0601 | 0.0831 | 0.6468 | 0.9186
100 | 0.1226 | 0.0147 | 0.0261 | 0.5297 | 0.9302
200 | 0.2226 | 0.0282 | 0.0494 | 0.5853 | 0.9420
300 | 0.1619 | 0.0422 | 0.0667 | 0.5674 | 0.9663
400 | 0.1952 | 0.0528 | 0.0802 | 0.5846 | 0.9675
500 | 0.1333 | 0.0384 | 0.0582 | 0.5540 | 0.9678
PV-DBOW | 50 | 0.8464 | 0.8166 | 0.8243 | 0.9224 | 0.9971
100 | 0.7964 | 0.6731 | 0.7219 | 0.8964 | 0.9948
200 | 0.8048 | 0.7606 | 0.7794 | 0.9012 | 0.9962
300 | 0.7964 | 0.8024 | 0.7980 | 0.8974 | 0.9969
400 | 0.7774 | 0.8322 | 0.7930 | 0.8879 | 0.9969
500 | 0.8607 | 0.8196 | 0.8320 | 0.9296 | 0.9974
PV-DM + PV-DBOW | 50 | 0.8298 | 0.8571 | 0.8324 | 0.9141 | 0.9972
100 | 0.7548 | 0.5726 | 0.6192 | 0.8746 | 0.9926
200 | 0.8405 | 0.7526 | 0.7879 | 0.9189 | 0.9960
300 | 0.8298 | 0.6504 | 0.7231 | 0.9129 | 0.9946
400 | 0.8131 | 0.7738 | 0.7781 | 0.9053 | 0.9962
500 | 0.7798 | 0.7528 | 0.7578 | 08889 | 0.9963

Le and Mikolov [38] show that combining the two variations of a paragraph
vector, namely the distributed memory paragraph vector (PV-DM) and the
distributed bag of words paragraph (PV-DBOW) improves performance. We sought
to test these predictions and experimented with a number of dimensions for
feature vectors on both models and their combination. From Table 2 the optimal
vector dimension is 50 on each model. We next cross-validated the window size
using the validation set and the optimal context window size was 9 words, with
the 10th one as the one to be predicted.

####  3.2.1 Preprocessing

To address the imbalanced dataset, a technique that combines the synthetic
minority oversampling technique (SMOTE) [39] and the Tomek Links (T-Links)
[40] undersampling was used.

By first defining the minority class vector and then determining the number of
nearest numbers (k), SMOTE synthesizes new examples in the minority class from
existing examples along the boundary joining all k-nearest neighbors in the
minority class. A line is generated and a synthetic point is placed between
the minority data points and any of their neighbors. This process is performed
for each minority data point and its k neighbors until the data is balanced.

The Tomek Links (T-Links) undersampling approach identifies all the pairs of
data points that are nearest to each other but belong to different classes.
Two points $a$ and $b$ are termed Tomek Links points if for $a$ as an instance
of class $A$ and $b$ an instance of class $B$, the distance $d(a,b)$ between
$a$ and $b$ $d(a,b)<d(a,c)$ or $d(a,b)<d(b,c)$ for any instance $c$.

In the preprocessing of the VA data, we combined Tomek-Link undersampling with
SMOTE oversampling. On classification tasks, the combination of SMOTE and
undersampling approaches (ENN and Tomek Links) has been shown to be
beneficial, with better area under a receiver operating characteristic AUC-ROC
[41]. A 5-fold cross validation step was included in the training process to
evaluate the classifiers’ performance. The SMOTETomek approach was used to
sample VA data in each fold, and the classifier was trained on the training
folds and validated on the remaining folds. The ideal hyperparameters of each
classifier were found using a grid search technique.

The performance of the classifiers on the COD classification task with respect
to physician expert diagnosis was evaluated using accuracy which is defined as
the number of correct predictions from all predictions made and precision,
which is defined as the number of deaths correctly identified as deaths due to
uncontrolled hyperglycaemia out of all the of positive predictions made. We
also used recall, defined as the number of deaths correctly identified as due
to uncontrolled hyperglycaemia for all deaths due to uncontrolled
hyperglycaemia.

Both recall and precision capture pertinent cases. A model’s precision tells
us which cases to ignore, even though they were predicted by our models as
deaths due to uncontrolled hyperglycemia. Recall tells us what fraction of
deaths due to uncontrolled hyperglycaemia to focus on even through they were
predicted by our models as deaths not due to hyperglycaemia. To account for
the balance of recall and precision, we used the F1-score, which is a weighted
average of precision and recall. We were interested in how well our
classifiers predict negative cases and positive cases across different
thresholds for labeling a death as due to uncontrolled hyperglycaemia. This
was determined by a receiver operating characteristic (ROC) curve. We also
used the area under the ROC curve (AUC-ROC) which measures the overall
performance of an algorithm.

To understand the relationships captured by the vector representations in the
text setting, we trained the principal components analysis (PCA) algorithm
[42] to reduce the dimensions of the feature vectors (100-D to 2-D vector
space) by integrating feature variances and visualized the results on a plot.

We performed an error analysis where two endocrinologists independently
reviewed the false positive (188) and false negative (10) cases predicted by
the best classifier in the text features setting. Where there was disagreement
between the two reviewers, consensus was reached through discussion. To
uncover predictors of uncontrolled hyperglycaemia, we examined the positive
class and false positive predictions word embeddings to determine if they
captured the same relationships between words by visualizing their alignments.
The positive class was made of a vocabulary of 849 unique words. To allow for
better visualizations, we looked at the top 150 words which included those of
words describing symptoms of uncontrolled hypeglycaemia. The false positive
cases vocabulary was 283 words, all of which were used.

##  4 Results

Table 3: Performance comparison of logistic regression, random forest, xgboost and neural network classifiers on cause of death classification in the binary, narrative text and combined features settings. Features | Classifier | Recall | Precision | F1-Score | Accuracy
---|---|---|---|---|---
Binary | Logistic Regression | 0.8471 | 0.2110 | 0.3199 | 0.9636
Random Forest | 0.8000 | 0.2772 | 0.3885 | 0.9741
XGBoost | 0.7529 | 0.0921 | 0.1637 | 0.9384
Neural Network | 0.7882 | 0.1069 | 0.1873 | 0.9449
Text | Logistic Regression | 0.7529 | 0.2801 | 0.4067 | 0.9824
Random Forest | 0.7412 | 0.2197 | 0.3377 | 0.9770
XGBoost | 0.6000 | 0.6437 | 0.6166 | 0.9942
Neural Network | 0.5294 | 0.4976 | 0.5001 | 0.9914
Combined | Logistic Regression | 0.7294 | 0.3783 | 0.4962 | 0.9882
Random Forest | 0.8117 | 0.3281 | 0.4601 | 0.9844
XGBoost | 0.8588 | 0.7170 | 0.7807 | 0.9962
Neural Network | 0.5647 | 0.6107 | 0.5786 | 0.9935

The results show that the accuracy of the models is higher than precision,
recall and F1-score values across all settings. For the binary features
setting, Table 3 shows that the logistic regression classifier achieved the
best performance in terms of recall (0.8471), i.e out of all of the cases for
which uncontrolled hyperglycaemia was a likely cause of death, the model was
able to classify 85% as deaths due to uncontrolled hyperglycaemia. The random
forest classifier followed next in terms of performance for the metrics of
precision and F1-score. In the text features setting, the logistic regression
performed best in terms of recall (0.7529) while the XGBoost classifier
performed best in terms of precision, F1-score and accuracy. In the setting of
combined binary and text features, the XGBoost classifier performs across all
metrics, followed by the random forest classifier in terms of recall and the
neural network in terms of precision, F1-score and accuracy.

(a)  (b)

Figure 1: ROC curves and AUC-ROC of Random Forest and Logistic Regression
classifiers in the binary features, text features and combined binary and text
features setting.

(a)  (b)

Figure 2: ROC curves and AUC-ROC of XGBoost and Neural Network classifiers in
the binary features, text features and combined binary and text features
setting. Figure 3: PCA plot: The distribution of keywords indicative of death
due to uncontrolled hyperglycaemia.

Of interest were high recall scores for all classifiers, but they came at the
cost of precision across all the settings except the XGBoost in the text
features setting. The best model was then selected by taking the balance
between precision and recall (F1-Score) into consideration. In the binary
features setting, the best classifier was the random forest while the XGBoost
performed best in the text and combined binary and text features settings. The
receiver operating characteristic (ROC) curve plots of the random forest and
logistic regression classifiers across the three settings are given by Fig. 1
and those of the XGBoost and Neural Network classifiers by Fig. 2. In the
three settings, all the classifier ROC curves climbed toward the top left,
meaning the models correctly predict positive and negative cases. We aimed to
discover a large number of deaths owing to uncontrolled hyperglycemia from the
data, thus we were looking for a threshold with the best sensitivity (recall).
In the binary features setting, the logistic regression and the random forest
classifier had the best performance at a threshold that gave a true positive
rate (TPR) of around 90%; i.e where 90% of all deaths due to uncontrolled
hyperglycaemia were being identified and a false positive rate (FPR) of around
30%. Since more than 90% of deaths from the VA report were not due to
uncontrolled hyperglycaemia, operating at this threshold for both classifiers
would be ideal. In the text setting the logistic regression and the XGBoost
gave the best area under the ROC curve (AUC-ROC) with the best cut-off point
at around (TPR 98%, FPR 24%) and (TPR 96%, FPR 30%) for the two classifiers.
For the combined features setting the same classifiers had the best AUC-ROC
scores with ideal threshold values for logistic regression at (TPR 96%, FPR
18%) and the XGBoost classifier at (TPR 98%, FPR 10%).

The PCA projections to the 2-D vector space are visualized in Fig. 3 showing
that words with similar part of speech (POS) tags are next to each other (i.e
spatially correlated in the embedding space) and that the paragraph2vec model
learned words by identifying the neighbouring words.

####  4.0.1 Error Analysis

After the review of the false positive cases (i.e those identified by the best
classifier as deaths likely due to uncontrollable hyperglycaemia, but not by
the physician-coders) and the false negative cases (cases identified by the
best classifier as deaths not likely due to uncontrolled hyperglycaemia, but
by the physician-coders as deaths due to uncontrolled hyperglycaemia), a
consensus was reached that of the hundred and eighty eight false positive
cases, six were indeed likely to be due to uncontrolled hyperglycaemia and of
the ten false negative cases, three were as identified as deaths not likely
due to uncontrolled hyperglycaemia.

Overall, the six correctly classified false positive cases are an 8% increase
in the positive cases and the improved scores from the text features setting
over the combined features setting prove the importance of information
available in the free text of VA reports, further highlighting the need for
both parts of the VA report in identifying a COD. The results further
illustrate that the removal of the keywords, ’sugar’ and ’diabetes’ from the
corpus aided in the classifiers learning the underlying symptoms indicative of
hyperglycaemia.

##  5 Discussion

The random forest classifier achieved the best performance for the binary
features setting and the XGBoost performed best in the text and combined
features settings. While both techniques use ensemble learning, the XGBoost
algorithm uses gradient boosting to better capture patterns of high-
dimensional features in text and combined text and binary data than the random
forest classifier, which depends on making prediction by random chance. In
comparison to the binary features and text features settings, the combined
features setting had higher scores across all metric because the narrative
text provides more depth and context on events around death than the binary
features alone.

Our results are in agreement with existing literature. Sakr et al. [43],
Knuiman et al. [44] did a comparative study of the same techniques trained
with the Synthetic Minority Over-Sampling Technique (SMOTETomek technique) to
predict mortality from medical reports using binary features. The results of
their study indicated that neural networks perform worse than the other
techniques. This is further supported by Knuiman et al. [44], who compared
artificial neural networks, Naive Bayes, logistic regression, decision trees,
and Bayesian network classifiers on myocardial infarction, with logistic
regression outperforming the others (AUC-ROC = 0.82).

The additional features in the combined binary and text setting increase the
dimensionality of the feature space, which in turn increases the search space
and limits the amount of extraction of valuable information from data [45,
46]. However, Pes [46] has shown the effectiveness of ensemble techniques in
handling high-dimensional data which was the case with XGBoost and random
forest in the text features and combined features settings in our work.
Clermont et al. [47] demonstrated how the non-parametric techniques of neural
networks, random forests and XGBoost performed better than logistic regression
in predicting hospital mortality from clinical notes in ICU patients and
Pirracchio et al. [48] added that this was because the linear and additive
relationship of the logistic regression technique cannot constrain the complex
processes surrounding a cause of death and its predictors or variables.

Word embeddings from the positive class and those from the false positive
predictions did capture the same relationships between symptoms of
uncontrolled hyperglycaemia as depicted by their vector alignments. This
explains the false positive predictions by the best classifier predicts false
positives which corroborate the ’distributional hypothesis’ [49], which states
that words with comparable or related meanings exist in the same context, and
embeddings for semantically and syntactically related keywords from the VA
corpus were closer together, as expected.

####  5.0.1 Limitations of the study

This study has three main limitations, the first of which is the dataset size
and its class distribution. With 79 positive and 8619 negative cases, the VA
dataset from Agincourt has the extremely imbalanced and small minority (EISM)
problem. Oversampling and undersampling techniques for redressing this
imbalance are not as effective when the absolute number of minority examples
is this small [50], and for this reason we believe significant relationships
from the data could have been missed. Additionally, the paragraph2vec model is
benchmarked on two very large datasets, the IMDB dataset [51] (100 000 movie
reviews) and Standford Sentiment Treebank dataset [52], vector dimensions of
which were 400. Tuning this parameter on our data, we were below the optimal
dimension of 100.

The second limitation was the way in which we extrapolated medical/ disease
knowledge concepts into our embeddings. While there is prior research on
domain adaptation techniques to health and clinical text for natural language
processing, the VA dataset by nature of its collection, transcription and
language cannot be described as a clinical document. Due to lack of publicly
available VA word embeddings, our analysis was also not validated against any
VA data outside of our institution.

The last limitation of this study is interpretability of VA embeddings. It was
not straight forward to provide an explanation for our classifier’s judgments
based on paragraph and word embeddings because like most embeddings, the
paragraph2vec encodes but does not distinguish notions of relatedness and
similarity [53]. While some symptoms like ”excessive drinking” and ”excessive
urination” are related, they are not similar. For analysis of cause of death
by uncontrolled hyperglycaemia, we would be interested in all symptoms of
uncontrolled hyperglycaemia and their combinations and not similar words and
synonyms.

##  6 Conclusion

The four machine learning techniques may be useful in the automation of cause
of death (COD) classification from verbal autopsy (VA) reports. In this study,
they could accurately identify COD from uncontrolled hyperglycaemia, a
complication of diabetes. Ensemble learning techniques of random forest and
XGBoost are highly accurate, sensitive and specific to this task in all
settings of binary, text and combined binary and text features. Our results
further suggest that the narrative text of the VA report may contain vital
information for determining COD and we therefore encourage further studies to
incorporate both binary and text features, separately or in combination. In
future work, we will use the transfer learning adaptation techniques of
feature extraction, variable selection and fine tuning to make use of broader
language representations from the general health and English language domain
to improve and refine representations of text from VA reports.

Acknowledgements We would like to thank Prof. Justine Davies and the Agincourt
Health and Socio-demographic Surveillance System (HDSS) for the VA dataset. We
also acknowledge the assistance of Dr. Faheem Seedat in reviewing cases for
the error analysis. We thank the United Nations’ Organization of Women in
Science for the Developing World (OWSD) for supporting and funding this work.
AW is funded by the Fogarty International Centre of the National Institutes of
Health (grant number K43TW010698). This paper describes the views of the
authors and does not necessarily represent the official views of the National
Institutes of Health (USA).

## References

  * Federation [2015] International Diabetes Federation.  International diabetes federation, brussels, belgium 2015.  _IDF Diabetes Atlas. 7th ed_ , 7, 2015.
  * Blackstock et al. [2019] S. Blackstock, M. D Witham, A. N Wade, A. Crampin, D. Beran, G. D Ogle, and J. I Davies.  Ability of verbal autopsy data to detect deaths due to uncontrolled hyperglycaemia: testing existing methods and development and validation of a novel weighted score.  _BMJ Open_ , 9:e026331, 2019.
  * Gerich et al. [1976] J. E Gerich, M. Lorenzi, D. M Bier, E. Tsalikian, V. Schneider, J. H Karam, and P. H Forsham.  Effects of physiologic levels of glucagon and growth hormone on human carbohydrate and lipid metabolism: studies involving administration of exogenous hormone during suppression of endogenous hormone secretion with somatostatin.  _J Clin Invest_ , 57:875–884, 1976.
  * Umpierrez et al. [2002] G. U Umpierrez, M. B Murphy, and A. E Kitabchi.  Diabetic ketoacidosis and hyperglycemic hyperosmolar syndrome.  _Diabetes Spectrum_ , 15(1):28–36, 2002.
  * WHO [2007] WHO.  Verbal autopsy standards: ascertaining and attributing cause of death, geneva, switzerland.  _WHO Press_ , 2007.
  * de Savigny et al. [2017] D. de Savigny, I. Riley, D. Chandramohan, F. Odhiambo, E. Nichols, and S. Notzo.  Classification of diabetes maculopathy images using data-adaptive neuro-fuzzy inference classifier.  _Global Health Action, Integrating community-based verbal autopsy into civil registration and vital statistics (CRVS): system-level considerations_ , pages 272–278, 2017.
  * Nations [2013] United Nations.  Department of economic and social affairs, population division, united nations.  _World Population Prospects: The 2012 revision_ , 2013.
  * Green and Flaxman [2010] S. T Green and A. D Flaxman.  Machine learning methods for verbal autopsy in developing countries.  _In AAAI Spring Symposium on Artificial Intelligence for Development (AI-D),Stanford: Stanford University_ , pages 22–24, 2010.
  * Todd et al. [1997] J. Todd, R. Balira, H. Grosskurth, P. Mayaud, and F. Mosha.  Hiv-associated adult mortality in a rural tanzania population.  _AIDS_ , 11:801–807, 1997.
  * Byass et al. [2010] P. Byass, K. Kahn, E. Fottrell, M. A Collinson, and S. M Tollman.  Moving from data on deaths to public health policy in agincourt, south africa: Approaches to analysing and understanding verbal autopsy findings.  _PLoS Medicine_ , 7, 2010.
  * Byass et al. [2006] P. Byass, E. Fottrell, D. Lan Huong, Y. Berhane, T. Corrah, K. Kahn, L. Muhe, and D. Duc Van.  Refining a probabilistic model for interpreting verbal autopsy data.  _Scandinavian Journal of Public Health_ , 34(1):26–31, 2006.
  * Fottrell and Byass [2010] E. Fottrell and P. Byass.  Verbal autopsy: methods in transition.  _Epidemiol Rev_ , 32:38–55, 2010.
  * Danso et al. [2013] S. Danso, E. Atwell, and O. Johnson.  Linguistic and statistically derived features for cause of death prediction from verbal autopsy text.  _Linguistic Processing and Knowledge in the Web_ , page 47–60, 2013\.
  * King et al. [2016] C. King, C. Zamawe, M. Banda, N. Bar-Zee, and J. Bird.  The quality and diagnostic value of open narratives in verbal autopsy: A mixed-methods analysis of partnered interviews from malawi.  _BMC Med Res Methodol_ , 16(13), 2016.
  * Jeblee et al. [2019] S. Jeblee, M. Gomes, P. Jha, F. Rudzicz, and G. Hirst.  Automatically determining cause of death from verbal autopsy narratives.  _BMC Medical Informatics and Decision Making_ , 34(19):26–31, 2019.
  * Boag et al. [2018] W. Boag, D. Doss, T. Naumann, and P. Szolovits.  What’s in a note? unpacking predictive value in clinical note representations.  _AMIA Joint Summits on Translational Science_ , 2018.
  * Grnarova et al. [2016] P. Grnarova, F. Schmidt, S. Hyland, and C. Eickhoff.  Neural document embeddings for intensive care patient mortality prediction.  _NIPS 2016 Workshop on Machine Learning for Health Workshop_ , 2016\.
  * Huang et al. [2019] K. Huang, J. Altosaar, and R. Ranganath.  Clinicalbert: Modeling clinical notes and predicting hospital readmission.  _arXiv:1904.05342_ , 2019.
  * Wen et al. [2020] Z. Wen, X. Lu, and S. Reddy.  Medal: Medical abbreviation disambiguation dataset for natural language understanding pretraining.  _Proceedings of the 3rd Clinical Natural Language Processing Workshop_ , pages 130–135, 2020.
  * Alsentzer et al. [2019] E. Alsentzer, J. R Murphy, W. Boag, W. H Weng, D. Jin, T. Naumann, and M. B. A McDermott.  Publicly available clinical bert embeddings.  _arXiv:1904.03323_ , 2019.
  * Lee et al. [2019] J. Lee, W. Yoon, S. Kim, D. Kim, S. Kim, C. H. So, , and J. Kang.  Biobert: a pre-trained biomedical language representation model for biomedical text mining.  _arXiv:1901.08746_ , 2019.
  * Jin et al. [2019] Q. Jin, B. Dhingra, W. Cohen, and X. Lu.  Probing biomedical embeddings from language models.  _ArXiv_ , abs/1904.02181:82–89, 2019.
  * Beltagy et al. [2019] Iz Beltagy, Arman Cohan, and Kyle Lo.  Scibert: Pretrained contextualized embeddings for scientific text.  _CoRR_ , abs/1903.10676, 2019.
  * Johnson et al. [2016] A. E. W. Johnson, T. J. Pollard, L. Shen, L. H. Lehman, M. Feng, M. Ghassemi, B. Moody, P. Szolovits, L. A. Celi, and R. G. Mark.  Mimic-iii, a freely accessible critical care database.  _Sci. Data_ , 2016.
  * Danso et al. [2014] S. Danso, E. Atwell, and O. Johnson.  A comparative study of machine learning methods for verbal autopsy text classification.  _International Journal of Computer Science Issues_ , 10(2):47–60, 2014.
  * James et al. [2011] S. L James, A. D Flaxman, and C. J Murray.  Performance of the tariff method: Validation of a simple additive algorithm for analysis of verbal autopsies.  _Popul Health Metrics_ , 9(1):31–47, 2011.
  * Flaxman et al. [2011] A. D Flaxman, A. Vahdatpour, S. Green, S. L James, and C. J. L Murray.  Random forests for verbal autopsy analysis: multisite validation study using clinical diagnostic gold standards.  _Population Health Metrics_ , 9:29, 2011.
  * Yan et al. [2019] Z. Yan, S. Jeblee, and G. Hirst.  Can character embeddings improve cause-of-death classification for verbal autopsy narratives?  _In Proceedings of the BioNLP 2018 Workshop and Shared Task_ , 34(19):234–239, 2019.
  * Jeblee et al. [2018] S. Jeblee, M. Gomes, and G. Hirst.  Multi-task learning for interpretable cause of death classification using key phrase predictions.  _In Proceedings of the BioNLP 2018 Workshop_ , 34(19):12–27, 2018.
  * Le and Mikolov [2014a] Q. Le and T. Mikolov.  Distributed representations of sentences and documents.  _International Conference on Machine Learning_ , page 1188–1196, 2014a.
  * Pedregosa et al. [2011] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Perrot, and É. Duchesnay.  Scikit-learn: Machine learning in python.  _Journal of Machine Learning Research_ , 370:2825–30, 2011\.
  * François [2015] C. François.  Keras.  https://github.com/fchollet/keras, 2015.
  * Bergstra et al. [2010] J. Bergstra, O. Breuleux, F. Bastien, P. Lamblin, R. Pascanu, G. Desjardins, J. Turian, D. Warde-Farley, and Y. Bengio.  Theano: A cpu and gpu math expression compiler.  _Python for Scientific Computing Conference (SciPy)_ , 2010.
  * Steven et al. [2008] B. Steven, K. Ewan, L. Edward, and B. Jason.  Multidisciplinary instruction with the natural language toolkit.  _Third Workshop on Issues in Teaching Computational Linguistics, ACL_ , 2008.
  * Kahn et al. [2012] K. Kahn, M. Collinson, F. Gómez-Olivé, O. Mokoena, R. Twine, and P. Mee.  Profile: Agincourt health and socio-demographic surveillance system.  _Int J Epidemiol_ , 41(4):988–1001, 2012.
  * Mikolov et al. [2013] T. Mikolov, I. Sutskever, K. Chen, G. Corrado, and J. Dean.  Eﬃcient estimation of word representations in vector space.  _preprint arXiv:1301.3781_ , 2013.
  * Manaka [2021] T. Manaka.  Machine-learning-text-binary.  https://github.com/Thokozile123/Machine-Learning-Text-Binary, 2021\.
  * Le and Mikolov [2014b] Q. Le and T. Mikolov.  Distributed representations of sentences and documents.  _Proceedings of the 31st International Conference on International Conference on Machine Learning_ , 32:1188–1196, 2014b.
  * Bowyer et al. [2011] K. W Bowyer, N. V Chawla, L. O Hall, and W. P Kegelmeyer.  SMOTE: synthetic minority over-sampling technique.  _CoRR_ , abs/1106.1813, 2011.
  * N. Thai-Nghe [2000] L. Schmidt-Thieme N. Thai-Nghe, T. N Do.  Learning optimal threshold on resampling data to deal with class imbalance.  _Proc. of the 8th IEEE International Conference on Computing_ , 2000\.
  * Purnajaya and Hanggara [2021] A. Purnajaya and F. Hanggara.  Perbandingan performa teknik sampling data untuk klasifikasi pasien terinfeksi covid-19 menggunakan rontgen dada.  _JAIC_ , 5(1):37–42, 2021.
  * Maćkiewicz and Ratajczak [1993] A. Maćkiewicz and W. Ratajczak.  Principal components analysis (pca).  _Computers & Geosciences_, 19(3):303–342, 1993\.
  * Sakr et al. [2017] S. Sakr, R. Elshawi, A. M Ahmed, W. T Qureshi, C. A Brawner, and S. J Keteyian.  Comparison of machine learning techniques to predict all-cause mortality using fitness data: the henry ford exercise testing (fit) project.  _BMC medical informatics and decision making_ , 174(17):1, 2017.
  * Knuiman et al. [1997] M. W Knuiman, H. T Vu, and M. R Segal.  An empirical comparison of multivariable methods for estimating risk of death from coronary heart disease.  _J Cardiovasc Risk_ , 2(4):127–34, 1997.
  * Bolón-Canedo et al. [2015] V. Bolón-Canedo, N. Sánchez-Maroño, and A. Alonso-Betanzos.  Recent advances and emerging challenges of feature selection in the context of big data.  _Knowl-Based Syst_ , (86):33–45, 2015.
  * Pes [2020] B. Pes.  Ensemble feature selection for high-dimensional data: a stability analysis across multiple domains.  _Neural Comput & Applic_, (32):5951–5973, 2020\.
  * Clermont et al. [2001] G. Clermont, D. C Angus, S. M DiRusso, M. Griffin, and W. T Linde-Zwirble.  Predicting hospital mortality for patients in the intensive care unit: a comparison of artificial neural networks with logistic regression models.  _Crit. Care Med. 29_ , page 291–296, 2001.
  * Pirracchio et al. [2015] R. Pirracchio, M. L. Petersen, M. Carone, M. R. Rigon, S. Chevret, and M. J. van der Laan.  Mortality prediction in intensive care units with the super icu learner algorithm (sicula): a population based study.  _The Lancet Respiratory Medicine_ , 1(3):42–52, 2015.
  * McDonald and Ramscar [2001] S. McDonald and M. Ramscar.  Testing the distributional hypothesis: The influence of context on judgements of semantic similarity.  _Proceedings of the Annual Meeting of the Cognitive Science Society_ , 23:303–342, 2001.
  * David [2009] C. A. David.  Finding problems in, proposing solutions to, and performing analysis on imbalanced data.  _PhD Dissertation, University of Notre Dame, Indiana_ , (1), 2009.
  * Maas et al. [2011] A. L Maas, R. E Daly, P. T Pham, D. Huang, A. Y Ng, and C. Potts.  Learning word vectors for sentiment analysis.  _In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics_ , 2011.
  * Socher et al. [2013] R. Socher, A. Perelygin, J. Y Wu, J. Chuang, C. D Manning, A. Y Ng, and C. Potts.  Recursive deep models for semantic compositionality over a sentiment treebank.  _In Conference on Empirical Methods in Natural Language Processing_ , 2013.
  * Khattak et al. [2019] F. K Khattak, S. Jeblee, C. Pou-Prom, M. Abdalla, C. Meaney, and F. Rudzicz.  A survey of word embeddings for clinical text.  _Journal of Biomedical Informatics_ , 100:100057, 2019.


# A Simple Approach for State-Action Abstraction using a Learned MDP
Homomorphism

Augustine N. Mavor-Parker
Centre for Artificial Intelligence
University College London
United Kingdom
Matthew J. Sargent
Centre for Artificial Intelligence
University College London
United Kingdom
Andrea Banino
Deepmind
London
United Kingdom
Lewis D. Griffin
Department of Computer Science
University College London
United Kingdom
Caswell Barry11footnotemark: 1
Department of Cell and Developmental Biology
University College London
United Kingdom
Equal Contribution. Send correspondence to<EMAIL_ADDRESS>
###### Abstract

Animals can rapidly infer, from limited experience, when sets of state-action
pairs have equivalent reward and transition dynamics. On the other hand,
reinforcement learning systems must painstakingly learn through trial and
error that sets of state-action pairs are value equivalent—requiring a
prohibitively large amount of samples from their environment. Previous works
use Markov decision process (MDP) homomorphisms, which are abstracted MDPs
enabling sample efficient policy learning. Consequently, impressive
improvements in sample efficiency have been achieved when a suitable MDP
homomorphism can be constructed a priori—usually by exploiting a
practitioner’s knowledge of environment symmetries. Here, we propose an
approach for constructing homomorphisms in discrete action spaces, which uses
a partial model of environment dynamics to infer which state-action pairs lead
to the same state—reducing the size of the state-action space by a factor
approximately equal to the cardinality of the action space. On MDP
homomorphism benchmarks, we demonstrate improved sample efficiency over
previous attempts to learn homomorphisms, and achieve comparable sample
efficiency to approaches that rely on prior knowledge of environment
symmetries.

##  1 Introduction

Reinforcement learning (RL) agents outperform humans on previously impregnable
benchmarks such as Go (Silver et al.,, 2016) and Starcraft (Vinyals et al.,,
2019). However, the computational expense of RL hinders its deployment in real
world applications. State of the art deep RL can agents demand hundreds of
millions of samples (or even hundreds of billions) to learn a policy—either
within an environment model (Hafner et al.,, 2020) or by direct interaction
(Badia et al.,, 2020).

In contrast, humans can abstract away details about state-action pairs that do
not effect their values. Consider the Atari game Asterix, where an agent
collects treasure while avoiding enemies in a 2D navigation task—a human gamer
understands that collecting treasure when approaching from the left has the
same value as collecting the same treasure but approaching from the right. In
this work, we identify this abstraction as a Markov decision process (MDP)
homomorphism Ravindran and Barto, (2001); van der Pol et al., 2020b . MDP
homomorphisms collapse equivalent state-actions in an observed MDP onto a
single abstract state-action pair in an abstract MDP van der Pol et al., 2020b
.

Given a mapping from experienced MDP to an abstract MDP, policies can be
learned efficiently in the smaller abstract space and then mapped back to the
experienced MDP when interacting with the environment (Ravindran and Barto,,
2001). However, obtaining mappings to the abstract state-action space is
challenging. Previous works hard code homomorphisms into a policy (van der Pol
et al., 2020b, ) but developing approaches that learn homomorphic mappings
from experience is an active area of research Rezaei-Shoshtari et al., (2022).

We developed equivalent effect abstraction, a method for constructing MDP
homomorphisms from experience via a dynamics model—leveraging the fact that
state-action pairs leading to the same next state frequently have equivalent
values. Consider again the Asterix example—moving to a given state has the
same value whether your previous action was right, left, up or down.
Therefore, if we know the value of approaching from the right then, when
equipped with a dynamics model, we demonstrate one can also infer the value of
approaching a given state from the left, from below or from above.
Consequently, equivalent effect abstraction extrapolates value judgements
between equivalent state-action pairs, reducing the amount of experience
required to learn a policy. We make the following contributions.

  1. 1.

We develop a novel approach for constructing MDP homomorphisms (equivalent
effect abstraction) that requires no prior knowledge of symmetries

  2. 2.

In the tabular setting, we show equivalent effect abstraction improves the
planning efficiency of model-based RL and the sample efficiency of model-free
RL

  3. 3.

In the deep RL setting, we show equivalent effect abstraction can be learned
from experience for Cartpole and Predator-Prey environments (which are MDP
homomorphisms benchmarks from van der Pol et al., 2020b )

  4. 4.

When a transition model is available, we show impressive gains in sample
efficiency on MinAtar’s Asterix environment

Figure 1: We find equivalent states by moving forward and then backward
through a model of the environment. Unlike previous approaches (van der Pol et
al., 2020b, ), we are not exploiting the horizontal symmetry of the
environment. In contrast, we take advantage of the fact that state-action
pairs that lead to the same state usually have equivalent values.

We are not the first to exploit the equivalent values of state action pairs
that lead to the same next state. (Tesauro et al.,, 1995) use the concept of
afterstates (states immediately before an opponents response to an action
(Antonoglou et al.,, 2021)) in their backgammon program. This is similar to
the equivalence we exploit and has been scaled within the MuZero framework to
other board games (Antonoglou et al.,, 2021), but in general afterstates have
not found wider applicability. (Misra et al.,, 2020)’s kinematic
inseparability is a related but distinct abstraction that uses contrastive
learning to classify whether two states will lead to the same next state for
the same action. In contrast, our approach looks for state-action pairs (with
different actions) that lead to the same next state, which is a broader form
of abstraction. Another difference is that our equivalent effect abstraction
does not require states to be discrete, which makes it readily applicable to
common RL environments where it is rare to reach the same state twice. Finally
Anand et al., (2015) defines a state-action equivalence that is more
comprehensive than the equivalence defined by (Ravindran and Barto,, 2001),
which is satisfied by various equivalent state-action pairs like those found
in (Ravindran and Barto,, 2001; van der Pol et al., 2020b, ; Misra et al.,,
2020) as well as the state-action pairs satisfying the equivalent effect
abstraction we present. However, they assume a tree of discrete states can be
searched over for equivalences, which is challenging in deep RL.

Our work is organised into the following sections. In Section 2 we formally
describe MDP homomorphisms and then introduce equivalent effect abstraction in
Section 3. In Section 4 we validate our approach using benchmarks from the MDP
homomorphism literature. Related work, limitations and conclusions are in
Sections 5, 6 and 7.

##  2 MDP Homomorphisms

Using the definition from (Silver,, 2015), an MDP $\mathcal{M}$ can be
described by a tuple
$\left<\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},\gamma\right>$ where
$\mathcal{S}$ is the set of all states, $\mathcal{A}$ is the set of all
actions, $\mathcal{R}=\mathbb{E}[R_{t+1}|S_{t}=s,A_{t}=a]$ is the reward
function that determines the scalar reward received at each state,
$\mathcal{P}=\mathbb{P}[S_{t+1}=s^{\prime}|S_{t}=s,A_{t}=a]$ is the transition
function of the environment describing the probability of moving from state to
another for a given action and $\gamma\in[0,1]$ is the discount factor
describing how much an agent should favour immediate rewards over those in
future states. An agent interacts with an environment through its policy
$\pi(a|s)=\mathbb{P}[A_{t}=a|S_{t}=s]$ (Silver,, 2015) which maps the current
state to a given action. To solve an MDP, an RL agent must develop a policy
that maximises the return $G$, which is equal to the sum of discounted future
rewards $G=\sum_{t=0}^{T}\gamma^{t}R_{t+1}$ (where $t$ is the current timestep
and $T$ is the number of timesteps in a learning episode).

(Ravindran and Barto,, 2001) introduced the concept of a homomorphism which,
using the notation and definitions from (van der Pol et al., 2020b, ), is a
homomorphism $h=\left<\sigma,\\{\alpha_{s}|s\in\mathcal{S}\right>\\}$ from an
agent’s experienced MDP $\mathcal{M}$ to an abstract MDP
$\bar{\mathcal{M}}=\left<\bar{\mathcal{S}},\bar{\mathcal{A}},\bar{\mathcal{P}},\bar{\mathcal{R}},\gamma\right>$,
which satisfies the following. $\forall
s,s^{\prime}\in\mathcal{S},a\in\mathcal{A}$:

$\bar{\mathcal{P}}(\sigma(s^{\prime})|\sigma(s),\alpha_{s}(a))=\sum_{s^{\prime\prime}\in\sigma^{-1}(s^{\prime})}\mathcal{P}(s^{\prime\prime}|s,a)$
(1) $\bar{\mathcal{R}}(\sigma(s),\alpha_{s}(a))=\mathcal{R}(s,a)$ (2)

Where $\sigma$ is a mapping from experienced states to abstract states, while
$\alpha_{s}$ is a state dependent mapping between experienced and abstract
actions. Equations (1) and (2) demonstrate that a homomorphic map of states
and actions must maintain the transition dynamics and reward functions of the
experienced MDP within the constructed abstract MDP. As proved in (Ravindran
and Barto,, 2001) and leveraged more recently by (van der Pol et al., 2020b,
), if the true Q-values in the abstract MDP are known,
$\bar{Q}^{*}(\sigma({s}),\alpha_{s}(a))$, then equipped with a homomorphism
$h$, the true Q-values in the experienced MDP can be “lifted” (Ravindran and
Barto,, 2001) from their abstract counterparts.

$\bar{Q}^{*}(\sigma({s}),\alpha_{s}(a))=Q^{*}(s,a)\quad\forall
s\in\mathcal{S},a\in\mathcal{A}$ (3)

Where $Q$ is the value of state $s$ and action $a$ and ∗ signifies the value
function is optimal (van der Pol et al., 2020b, ). The size of abstract state-
action space $\bar{\mathcal{S}}\times\bar{\mathcal{A}}$ is often smaller than
the size of the experienced state-action space $\mathcal{S}\times\mathcal{A}$.
Consequently, if an agent is equipped with a homomorphic map, a policy can be
learned efficiently in the abstract state-action space and then “lifted” to
the experienced state-action space (Ravindran and Barto,, 2001). In general,
obtaining a homomorphic map is difficult. Next, we present a novel method for
inferring a homomorphic map from experience.

##  3 Equivalent Effect Abstraction

Equivalent effect abstraction is based on a simple observation that is
guaranteed to be true in most RL benchmarks if $\mathcal{R}$ and $\mathcal{P}$
are deterministic. Firstly, state-action pairs that lead to the same next
state have equivalent reward functions in the majority of RL benchmarks (e.g.
DeepMind control Tassa et al., (2018), classic control Brockman et al., (2016)
and Atari Bellemare et al., (2013)). While common MDP formulations (like the
one used in Section 2) define reward functions as depending on the current
state and previous action, in many environments it is possible to drop the
dependence on the previous action. Nevertheless, we note that we assume no
dependence on the previous action for the reward function, so our approach is
not suitable for optimal control settings where different actions may be
defined to have a different cost (Liberzon,, 2011, p. 145). Next, we develop a
method to learn the mapping $\sigma$ between state-action pairs that lead to
the same state.

Consider an agent that takes an action $a_{t}$ in a state $\mathbf{s}_{t}$ and
ends up in state $\mathbf{s}_{t+1}$, once in state $\mathbf{s}_{t+1}$, how can
we infer what other state-action pairs could also lead to state
$\mathbf{s}_{t+1}$? We could make this inference if we had both a forwards and
backwards model of the environment. To train our forwards model $\mathcal{F}$
and backwards model $\mathcal{B}$, we collect experience tuples
$\left<\mathbf{s}_{t},a_{t},\mathbf{s}_{t+1}\right>$ from the environment and
then optimise the parameters of our models $\theta$ and $\phi$ with a MSE loss
function.

$\operatorname*{arg\,min}_{\theta}\|\mathcal{F}_{\theta}(\mathbf{s}_{t},a_{t})-\mathbf{s}_{t+1}\|^{2}_{2}$
(4)
$\operatorname*{arg\,min}_{\phi}\|\mathcal{B}_{\phi}(\mathbf{s}_{t+1},a_{t})-\mathbf{s}_{t}\|^{2}_{2}$
(5)

Figure 2: PyTorch (Paszke et al.,, 2017) style pseudo-code for equivalent
effect abstraction Q-network. Our Q-network only has an output for one
action—the hypothetical action. To transform input states into the
hypothetical reference frame we move forwards and backwards through a dynamics
model, with the input action on the forwards model giving the order of the
equivalent Q-value outputs.

With a learned forwards model $\mathcal{F}_{\theta}$ and backwards model
$\mathcal{B}_{\phi}$ we are ready to reduce the size of our MDP. First, we
select an action $\bar{a}$ to be our hypothetical action. Next for a given
state-action pair $(\mathbf{s},a)$ the equivalent (hypothetical) state-action
pair can be computed by moving forwards through our model and then backwards
with the hypothetical action $\bar{a}$.

$\bar{\mathbf{s}}=\sigma(\mathbf{s},\bar{a})=\mathcal{B}_{\phi}(\mathcal{F}_{\theta}(\mathbf{s},a),\bar{a})$
(6)

In theory, one could calculate multiple equivalent state-action pairs for a
given state-action pair by querying the backwards model with multiple
different actions. However, our aim is to simplify the state-action space that
our policy learns in, so instead we select one action to be our “hypothetical”
action and all other state-action pairs are mapped into its reference frame.
In Figure 1, the hypothetical action is moving right—meaning our Q-network
only has to learn about moving right and we can map all right values to their
equivalent left values with our forwards and backwards models. The
hypothetical action is a hyperparameter. In some environments such as
Cartpole, the hypothetical action can be selected randomly. However, in other
environments backwards state predictions are not always possible for a given
state-action pair—meaning hypothetical actions need to be selected more
carefully.

By finding an equivalent hypothetical state-action pair for each experienced
state-action pair we reduce the number of state-action pairs that an off-
policy algorithm must learn to estimate—reducing the size of the space of
values from $\mathcal{S}\times\mathcal{A}$ to $\mathcal{S}\times 1+N$, where
$N\in\mathcal{S}\times\mathcal{A}$ are the states where the hypothetical
action is not computable because of edge cases that will be discussed shortly
(e.g. stochastic environments, no unique backwards action or transitions that
do not exist in the environment).

It is important to stress that as our mapping satisfies Equation 3, it also
assumes the MDP homomorphism properties developed by Ravindran and Barto,
(2001), meaning the policy we learn with our hypothetical state action pairs
can be “lifted” van der Pol et al., 2020b  and used as a policy in a given
experienced environment. Concretely, when performing value inference on a
given state-action pair we we apply Equation (6) to find the equivalent state
and then predict a single action value with a Q-network.

We have assumed that for every state-action pair, at least one equivalent
hypothetical state-action can be computed. This is often true but it is not
guaranteed. For example, near borders in a grid world there is no way to
travel left to a border state that has a border on its right. We only found
this to significantly effect results in the tabular setting, which we mitigate
by reverting to vanilla Q-learning table when necessary. This backup
Q-learning value table is learned in parallel. Furthermore, in stochastic
environments it may be difficult to predict equivalent states. In Section 6,
we suggest recurrent state space models could be used to sample from a
distribution of predicted transitions in future work (Hafner et al.,, 2019).
Lastly, when forward dynamics are deterministic the previous state could still
be unpredictable as a unique previous state for a given previous action may
not exist (Yu et al.,, 2021). Again the model based RL provides some ideas for
addressing this (Yu et al.,, 2021).

##  4 Experiments

We test our approach on a range of tasks, from tabular RL to using a
convolutional DQN Silver et al., (2016). Where possible, we try to have
overlap in our experiments with the MDP homomorphism literature—using Cartpole
to overlap with the well known van der Pol et al., 2020b  and van der Pol et
al., 2020a  as well as Predator Prey to overlap with van der Pol et al., 2020a
. Hyperparameter details can be found in the Appendix. Shaded regions indicate
the standard error.

###  4.1 Tabular Maze

Figure 3: $\mathbf{(a),(b)}$ In a gridworld, removing redundancies in the
Q-table improves sample efficiency. For the vanilla baselines, an agent must
learn the value of each state-action pair for all actions in the action space.
With equivalent effect abstraction, we reduce the number of Q-values that need
to be learned. This yields improvements in both model-free Q-learning and
model-based approaches Q-planning. 50 seeds are used for Q-learning while 10
seeds are used for Q-planning.

We begin with the maze environment from (Sutton and Barto,, 2018, p. 165). The
maze consists of $6\times 9$ cells. The agent starts on the far west of the
maze and must navigate east and find its way around walls in the middle of the
environment. After passing the walls, the agent must then travel further east
to reach a corridor, before moving north to reach a goal location. In this
particular experiment, we assume a model of the environment is known
beforehand—which we move through forwards and then backwards to create a
homomorphic map.

We use an open-source Q-learning (Watkins and Dayan,, 1992) implementation as
a baseline111github.com/thehawkgriffith/dyna-maze (along with the maze code
from the same authors). We then adapt the Q-learning implementation to make
use of a homomorphic map as shown in Algorithm 1. We show in Figure 3(a) that
our approach compares favorably to vanilla Q-learning, converging much faster
to the optimal policy. In the gridworld environment, we reduce the size of the
state action space and hence the number of Q-values (when going left) from
$216$ to $67$. In this case, the number of states in the reduced action space
is not reduced by the cardinality of the action space exactly. This is because
of the 13 states where it is not possible to reach them by traveling left
because there is a border to the right.

To test planning efficiency, we performed a further experiment on the
gridworld with model-based Q-planning (Sutton and Barto,, 2018, p. 161). Our
Q-planning method is equivalent to the Q-learning implementation but altered
to learn by randomly sampling model transitions. At each Q-planning update we
pause training and evaluate performance in a 100 step episode. Planning
performance is shown in Figure 3(b). Similarly to the Q-learning results,
leveraging knowledge of equivalent actions is able to significantly improve
learning efficiency.

###  4.2 Cartpole

Figure 4: $\mathbf{(a)}$ Cartpole: we improve upon the DQN baseline
significantly (the model is learned in the initial 3 episodes of experience—a
different 3 episodes per repeat— _that is shown on the x-axis of the plot_).
For vanilla DQN and equivalent effect abstraction we use 50 seeds, 5 seeds are
used for each PRAE result and 30 seeds are used for MDP homomorphic networks.
For clarity, PRAE 10 refers to using PRAE with 10 episodes of data to
construct an environment model, which it then plans in. $\mathbf{(b)}$
Predator Prey: we learn a mapping to improve sample efficiency over vanilla
DQN. Unlike in Cartpole _the experience used to train the homomorphic map is
not included in the plot_ for equivalent effect abstraction DQN. For
equivalent effect abstraction we reuse one set of pre-trained backwards and
forwards models for each RL run.

Next, we apply our approach to the Cartpole benchmark Brockman et al., (2016).
In Cartpole, an agent controls a cart with a pendulum attached to it. The goal
of the task is to learn to balance the pendulum upright by moving the cart
horizontally left or right. The states are a four-dimensional vector
$(\text{position},\text{velocity},\text{angle},\text{angular velocity})$ and
Cartpole has a discrete action space (move the cart left or right). To learn a
model, we use 3 episodes of learning experience at the beginning of training
to train the homomorphic map (two simple linear models, one for each action
optimised with Adam (Kingma and Ba,, 2015)) to generate equivalent state-
action pairs for the opposite action to the one experienced. These initial
training steps _are_ included for equivalent effect abstraction in Figure
4(a). We integrate the learned homomorphism into a DQN implementation Silver,
(2015) 222https://gist.github.com/Pocuston/13f1a7786648e1e2ff95bfad02a51521,
with Q-values only being learned for an arbitrarily chosen hypothetical
action.

We compare our approach to three baselines: vanilla DQN, MDP homomorphic
Networks van der Pol et al., 2020b  and PRAE van der Pol et al., 2020a . For
all experiments we use the baselines’ open source implementations and the
libraries they build upon Stooke and Abbeel, (2019). MDP homomorphic networks
use specially constructed weights that are equivariant to environment
symmetries (requiring prior knowledge of symmetries). PRAE trains a
contrastive model to learn a mapping to a latent “plannable” MDP, that
satisfies the definitions of an MDP homomorphism. PRAE then performs planning
on the learned abstract MDP. For Cartpole, we learned a mapping for PRAE with
datasets of 10, 100, 10000 episodes of random experience. We plot the average
convergence performance of PRAE’s planning algorithm. In the Cartpole task, we
did not need to avoid any hypothetical state action pairs that are not
computable—meaning a reduction of the size of the state action space from
$\mathcal{S}\times\mathcal{A}$ to $\mathcal{S}\times 1$.

As shown in Figure 4(a), in the low sample regime we improve upon both PRAE
and vanilla DQN—with our approach converging at around episode 12 while
vanilla Q-learning takes around 16 episodes to converge—note that this
improvement includes the number of episodes required to learn our mapping. We
found that in many cases MDP homomorphic networks were able to achieve good
performance but worst case runs brought down mean performance significantly.
To demonstrate this, we plot the mean and median for MDP homomorphic networks
(only the mean is plotted for other methods). The median MDP homomorphic
network is also able to similarly improve upon the vanilla DQN baseline, by
leveraging a practitioner’s prior knowledge of environment symmetries. MDP
homomorphic networks are not necessarily a competing approach to equivalent
effect abstraction and future work could conceivably combine these two
approaches for even greater sample efficiency.

###  4.3 Stochastic Predator Prey

Following the seminal work of van der Pol et al., 2020b , we also benchmark
our method on the Predator Prey environment, where an agent must chase a
stochastically moving prey in a 2D world van der Pol et al., 2020a . The
observations are a $7\times 7\times 3$ tensor encoding agent position and prey
position. The objective of the agent is to catch the prey in fewest steps
possible. Rewards are set to -0.1 unless the predator catches the prey, in
which case the episode ends and a reward of 1 is provided.

For our method, we train action-dependent forwards and backwards models on a
dataset of transitions created by taking random actions in the environment for
$10^{4}$ environment steps (equivalent to around $170$ episodes of random
experience). Note that the environment is stochastic, so learning a perfect
environment model is impossible. We compare to the baselines introduced in
Section 4.2. For PRAE, we benchmark with both 10000 and 100 episodes of
experience data and perform planning to convergence. Similar to the Cartpole
results, PRAE can be effective, but only when a relatively large amount of
model training experience is available.

In Predator Prey, it is easy to learn a policy with around 5 episodes with a
vanilla DQN. As a result, in this particular environment it is not practical
to use equivalent effect abstraction as by the time a model can be learned the
policy has already converged. However, what Figure 4(b) does show is that if a
model can be obtained before online policy learning, then equivalent effect
abstraction can deliver an improvement over vanilla DQN—even though the
environment is stochastic. This is consistent with previous literature on MDP
homomorphisms van der Pol et al., 2020b , which designed the Predator Prey
environment to test whether their MDP homomorphism was robust to
stochasticity.

###  4.4 MinAtar Asterix

MinAtar Young and Tian, (2019) emulates a subset of games from the Atari
benchmark Bellemare et al., (2013) but at a lower dimensionality—allowing
policies similar to those found in the full Atari learning environment to be
trained at lower computational cost. Consequently, we can do similar
benchmarking to previous works on MDP homomorphisms (i.e. van der Pol et al.,
2020b ’s Atari experiments) but using tasks that are easy to reproduce. We use
_Asterix_ to benchmark equivalent effect abstraction. In Asterix, an agent
must discover moving treasure pots, while avoiding enemies. The effective
actions are left, right, up, down and “do nothing”. Asterix has no obvious
global symmetries, making it difficult for approaches that rely on symmetries
van der Pol et al., 2020b .

Further, the MinAtar environments follow the “sticky action” protocol Machado
et al., (2018), meaning actions are only executed with a probability of 90%
while the other 10% of the time the previous action is repeated (Figure 5(a)).
To further validate our algorithm in stochastic environments we increase the
stickiness of actions to 25% in Figure 5(b).

In this section, we assume a model of the environment is known, which we
construct programmatically to perform equivalent effect abstraction as if the
environment were deterministic—meaning every time a sticky action is chosen
the model’s prediction is incorrect (see supplementary code). Thus the
question addressed here is _not_ about learning backwards and forwards models
as we did for Cartpole (section 4.2) and Predator Prey (section 4.3). Instead
we seek to understand, given a model of the environment (which may be
imperfect due to stochasticity), what approach can we take to perform value-
based deep RL in the smallest number of learning steps.

We use the convolutional DQN Silver et al., (2016) implementation from the
MinAtar repository Young and Tian, (2019) as the backbone for equivalent
effect abstraction, leaving their hyperparameters unchanged. As is the case in
many games (and RL environments more generally), MinAtar allows a “do nothing”
action. Interestingly, it is always possible to construct our homomorphism
with the “do nothing” as our hypothetical action, meaning we never have to
fall back to vanilla Q-learning (which we did previously in the tabular maze
when the homomorphism was impossible to construct). In Asterix it is always
possible to reach every state by doing the hypothetical action—meaning we
reduce the size of the state-action space by a factor of 6 (the number of
actions in the action space). We compare equivalent effect abstraction to the
vanilla DQN from the MinAtar Young and Tian, (2019) repository. As a further
baseline, we construct a DQN with a equivariant feature extractor, similar to
the network used to learn Pong in van der Pol et al., 2020b . We implement
feature extractors that are equivariant to both horizontal and vertical
reflections. The hyperparameter selection process for this baseline is
described in the Appendix.

Asterix is a relatively complex game requiring an agent to explore for moving
treasures while simultaneously avoiding enemies. During this exploration, we
show empirically that traversing the state space is much more efficient when
the equivalences between different state-action pairs are known. Equivalent
effect abstraction converges $~{}3\times$ faster than vanilla Q-learning,
consistent with out results in previous sections. As there is no obvious
symmetry—in contrast to Pong for example—homomorphic DQN does not yield an
advantage over the vanilla DQN.

Figure 5: $\mathbf{(a)}$ In Asterix, Equivalent effect abstraction allows DQN
to reach a return of 20 in around $3.5\times 10^{5}$ frames, while vanilla DQN
requires approximately $(1\times 10^{6})$ frames to reach this return. We run
each method for five seeds. $\mathbf{(b)}$ Equivalent effect abstraction
maintains its advantage in the presence of even more stochastic transitions
(25% stochasticity as opposed to 10% previously). Both plots have been
smoothed with a Gaussian kernel where $\sigma=1$.

##  5 Related Work

(Ravindran and Barto,, 2001) developed a framework for abstraction under
symmetry—deriving theoretical guarantees when lifting abstract policies.
(Givan et al.,, 2003) proposed MDP model minimisation using bisimulation
(Larsen and Skou,, 1991), where states could be mapped to an abstract state if
their expected transition dynamics and reward sequences were
indistinguishable. These early applications of MDP homomorphisms required
prior knowledge from a practitioner, with no obvious method for integrating
homomorphisms into function approximation.

(Biza et al.,, 2021) use bisimulation to inform structured hidden Markov model
priors to infer a reduced discrete state space from learnt embeddings.
Bisimulation can be relaxed to a similarity metric (Ferns et al.,, 2011; Ferns
and Precup,, 2014)—these metrics have been used to condition embeddings with
contrastive losses to generate invariant representations of states. (Zhang et
al.,, 2020) augment their embedding encoder with an additional contrastive
loss term that results in encoder distances forming an approximate
bisimulation metric—allowing agents to perform control tasks in the presence
of distractors. (Rezaei-Shoshtari et al.,, 2022) formalise continuous
homomorphisms—building a homomorphism through a bisimulation metric. Zhu et
al., (2022) use a contrastive loss and define state-action pairs with the same
actions as equivalent in an abstract space. Agarwal et al., (2021) outline a
new policy similarity metric that is used in conjunction with a contrastive
loss—pulling together embeddings that result in similar trajectories. In
general, contrastive learning achieves impressive results but can be sample
hungry van der Pol et al., 2020a ).

More recently, homomorphisms in MDPs have been used in tasks with symmetries.
PRAE (van der Pol et al., 2020a, ) trained a contrastive world model and
accompanying loss function that satisfies the transition dynamics definition
(Equation (1)) by design. (van der Pol et al., 2020b, ) derived a method for
learning network weights that are equivariant to environment symmetries but
their algorithm requires a practitioner to hardcode symmetric groups
beforehand. Similar approaches have been adopted for environments with
continuous symmetries Wang et al., (2022). Another recent approach to using
MDP homomorphisms is detailed in (Biza and Platt,, 2019). Here the problem of
finding MDP homomorphisms is approached using Online Partition Iteration—the
mappings from state to abstract states are learnt by predicting which
partition a state should fall into given an action, and refining the
partitions through splitting. Lastly, (Asadi et al.,, 2019) and (Lyu et al.,,
2023) explore equivalences where state-actions pairs lead to separate states
but with the same probabilities.

Perhaps the most relevant related approach is the use of afterstates (Sutton
and Barto,, 2018, p. 136). Afterstates slightly shift an MDP out of phase with
conventional state transitions, creating environment states that are in-
between the initial effect of a policy’s action and the reaction of the
environment to said action. Applications of the afterstates framework have
generally been constrained to board games Tesauro et al., (1995) and focus on
dealing with the stochasticity of an opponent rather than improved sample
efficiency Antonoglou et al., (2021). Also closely related, (Misra et al.,,
2020) elegantly deploy a similar concept—using contrastive learning to learn
an abstraction that groups together state-action pairs that will pass through
or have passed through same state with the same previous action. However, this
approach is not necessarily scalable to common RL environments such as Atari
and classic control tasks where experiencing the exact same state twice is
unlikely.

More broadly, model-based RL has enabled sample-efficient superhuman
performance in Atari (Hafner et al.,, 2020), but the number of planning steps
required to learn a policy is still large. Backwards models have been proposed
to improve sample efficiency of world model representation learning (Yu et
al.,, 2021), which would be interesting if integrated with equivalent effect
abstraction. Our work is related to approaches that learn data augmentation
policies for RL (Raileanu et al.,, 2021), however instead of augmenting
experience with a more diverse distribution of states, we are trying to learn
values for a narrower distribution of hypothetical states.

##  6 Limitations and Future Work

Equivalent effect abstraction requires a dynamics model. For our approach to
be practical, the number of transitions required to learn the environment
model must be smaller than those required to learn a policy—this is true in
many environments (see Cartpole results) but it is not guaranteed (see
Predator Prey results). What is more, despite recent progress Hafner et al.,
(2019, 2020); Saxena et al., (2021); Ye et al., (2021), learning a dynamics
model can be difficult to learn for two reasons.

Stochastic dynamics can make environment transitions multimodal, meaning point
predictions of future or previous states are not adequate. In theory, this can
be dealt with using stochastic models Hafner et al., (2019) of forwards and
backwards dynamics, that could be sampled when abstracting state-action pairs.
Even worse, dynamics can be completely unpredictable Kendall and Gal, (2017);
Mavor-Parker et al., (2021), in which case uncertainty predictions could be
used to signal that equivalent effect abstraction should temporarily revert to
vanilla Q-learning. Further, backwards transitions may be unpredictable
because an equivalent state does not exist—see Section 1(Yu et al.,, 2021).
Uncertainty predictions could also be used to avoid out-of-distribution
mappings Kendall and Gal, (2017). We note that there is a silver lining for
equivalent effect abstraction, which is that only one step transition models
are required, avoiding the challenge of learning multi-step models Saxena et
al., (2021).

Equivalent effect abstraction is formulated within the framework of value
based RL using discrete actions. While many of the recent breakthroughs in
deep RL are value based (e.g. Silver et al., (2017), Badia et al., (2020)),
actor-critic approaches Schulman et al., (2017); Mnih et al., (2016) are often
the natural choice for control tasks. Embedding equivalent effect abstraction
into actor-critic architectures is a fruitful avenue for future research. It
is also conceivable to formulate equivalent effect abstraction within a
continuous action space by simply discretising the action space, which has
allowed discrete action space methods to obtain state of the art performance
on RL tasks Banino et al., (2021). Additionally, equivalent effect abstraction
could also be integrated into existing homomorphic MDP methods that rely on
symmetries to further reduce the size of the abstract state-action space van
der Pol et al., 2020b . In practice, this would mean using a homomorphic
Q-network van der Pol et al., 2020b  but modified to only have one
hypothetical action.

##  7 Conclusion

Equivalent effect abstraction is a simple method that reduces the size of a
state-action space. It is easy to implement and, provided with dynamics
models, it requires no prior knowledge of environment symmetries. We show that
equivalent effect abstraction improves the sample efficiency of policy
learning in tabular environments, control tasks with continuous state spaces,
stochastic deep RL environments, and in MinAtar’s Asterix. An exciting next
step would be to integrate equivalent effect abstraction into model-based RL
(Hafner et al.,, 2020; Yu et al.,, 2021) to improve planning efficiency.

## References

  * Agarwal et al., (2021) Agarwal, R., Machado, M. C., Castro, P. S., and Bellemare, M. G. (2021).  Contrastive behavioral similarity embeddings for generalization in reinforcement learning.  arXiv preprint arXiv:2101.05265.
  * Anand et al., (2015) Anand, A., Grover, A., Mausam, M., and Singla, P. (2015).  Asap-uct: Abstraction of state-action pairs in uct.  In Proceedings of the 24th International Conference on Artificial Intelligence, pages 1509–1515.
  * Antonoglou et al., (2021) Antonoglou, I., Schrittwieser, J., Ozair, S., Hubert, T. K., and Silver, D. (2021).  Planning in stochastic environments with a learned model.  In International Conference on Learning Representations.
  * Asadi et al., (2019) Asadi, M., Talebi, M. S., Bourel, H., and Maillard, O.-A. (2019).  Model-based reinforcement learning exploiting state-action equivalence.  In Asian Conference on Machine Learning, pages 204–219. PMLR.
  * Badia et al., (2020) Badia, A. P., Piot, B., Kapturowski, S., Sprechmann, P., Vitvitskyi, A., Guo, Z. D., and Blundell, C. (2020).  Agent57: Outperforming the atari human benchmark.  In International Conference on Machine Learning, pages 507–517. PMLR.
  * Banino et al., (2021) Banino, A., Badia, A. P., Walker, J., Scholtes, T., Mitrovic, J., and Blundell, C. (2021).  Coberl: Contrastive bert for reinforcement learning.  arXiv preprint arXiv:2107.05431.
  * Bellemare et al., (2013) Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. (2013).  The arcade learning environment: An evaluation platform for general agents.  Journal of Artificial Intelligence Research, 47:253–279.
  * Biewald, (2020) Biewald, L. (2020).  Experiment tracking with weights and biases.  Software available from wandb.com.
  * Biza and Platt, (2019) Biza, O. and Platt, R. (2019).  Online abstraction with mdp homomorphisms for deep learning.  In Proceedings of the 18th International Conference on Autonomous Agents and MultiAgent Systems.
  * Biza et al., (2021) Biza, O., Platt, R., van de Meent, J., and Wong, L. (2021).  Learning discrete state abstractions with deep variational inference.  Advances in Approximate Bayesian Inference.
  * Brockman et al., (2016) Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang, J., and Zaremba, W. (2016).  Openai gym.  arXiv preprint arXiv:1606.01540.
  * Ferns et al., (2011) Ferns, N., Panangaden, P., and Precup, D. (2011).  Bisimulation metrics for continuous markov decision processes.  SIAM Journal on Computing, 40(6):1662–1714.
  * Ferns and Precup, (2014) Ferns, N. and Precup, D. (2014).  Bisimulation metrics are optimal value functions.  In Proceedings of the Thirtieth Conference on Uncertainty in Artificial Intelligence, UAI’14, page 210–219, Arlington, Virginia, USA. AUAI Press.
  * Givan et al., (2003) Givan, R., Dean, T., and Greig, M. (2003).  Equivalence notions and model minimization in markov decision processes.  Artificial Intelligence, 147(1-2):163–223.
  * Hafner et al., (2019) Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J. (2019).  Learning latent dynamics for planning from pixels.  In International conference on machine learning, pages 2555–2565. PMLR.
  * Hafner et al., (2020) Hafner, D., Lillicrap, T. P., Norouzi, M., and Ba, J. (2020).  Mastering atari with discrete world models.  In International Conference on Learning Representations.
  * Kendall and Gal, (2017) Kendall, A. and Gal, Y. (2017).  What uncertainties do we need in bayesian deep learning for computer vision?  Advances in neural information processing systems, 30.
  * Kingma and Ba, (2015) Kingma, D. P. and Ba, J. (2015).  Adam: A method for stochastic optimization.  In ICLR (Poster).
  * Larsen and Skou, (1991) Larsen, K. G. and Skou, A. (1991).  Bisimulation through probabilistic testing.  Information and computation, 94(1):1–28.
  * Liberzon, (2011) Liberzon, D. (2011).  Calculus of variations and optimal control theory: a concise introduction.  Princeton university press.
  * Lyu et al., (2023) Lyu, Y., Côme, A., Zhang, Y., and Talebi, M. S. (2023).  Scaling up q-learning via exploiting state–action equivalence.  Entropy, 25(4):584.
  * Machado et al., (2018) Machado, M. C., Bellemare, M. G., Talvitie, E., Veness, J., Hausknecht, M., and Bowling, M. (2018).  Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents.  Journal of Artificial Intelligence Research, 61:523–562.
  * Mavor-Parker et al., (2021) Mavor-Parker, A. N., Young, K. A., Barry, C., and Griffin, L. D. (2021).  How to stay curious while avoiding noisy TVs using aleatoric uncertainty estimation.  Proceedings of the 39th International Conference on Machine Learning, 162:15220–15240.
  * Misra et al., (2020) Misra, D., Henaff, M., Krishnamurthy, A., and Langford, J. (2020).  Kinematic state abstraction and provably efficient rich-observation reinforcement learning.  In International conference on machine learning, pages 6961–6971. PMLR.
  * 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, pages 1928–1937. PMLR.
  * pandas development team, (2020) pandas development team, T. (2020).  pandas-dev/pandas: Pandas.
  * Paszke et al., (2017) Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga, L., and Lerer, A. (2017).  Automatic differentiation in pytorch.
  * Raileanu et al., (2021) Raileanu, R., Goldstein, M., Yarats, D., Kostrikov, I., and Fergus, R. (2021).  Automatic data augmentation for generalization in reinforcement learning.  Advances in Neural Information Processing Systems, 34.
  * Ravindran and Barto, (2001) Ravindran, B. and Barto, A. G. (2001).  Symmetries and model minimization in markov decision processes.
  * Rezaei-Shoshtari et al., (2022) Rezaei-Shoshtari, S., Zhao, R., Panangaden, P., Meger, D., and Precup, D. (2022).  Continuous mdp homomorphisms and homomorphic policy gradient.  arXiv preprint arXiv:2209.07364.
  * Saxena et al., (2021) Saxena, V., Ba, J., and Hafner, D. (2021).  Clockwork variational autoencoders.  Advances in Neural Information Processing Systems, 34:29246–29257.
  * 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, (2015) Silver, D. (2015).  Lecture notes in reinforcement learning.  Computer Science Department, University College London.
  * Silver et al., (2016) Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., et al. (2016).  Mastering the game of go with deep neural networks and tree search.  nature, 529(7587):484–489.
  * 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–359.
  * Stooke and Abbeel, (2019) Stooke, A. and Abbeel, P. (2019).  rlpyt: A research code base for deep reinforcement learning in pytorch.  arXiv preprint arXiv:1909.01500.
  * Sutton and Barto, (2018) Sutton, R. S. and Barto, A. G. (2018).  Reinforcement learning: An introduction.  MIT press.
  * Tassa et al., (2018) Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. L., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A., et al. (2018).  Deepmind control suite.  arXiv preprint arXiv:1801.00690.
  * Tesauro et al., (1995) Tesauro, G. et al. (1995).  Temporal difference learning and td-gammon.  Communications of the ACM, 38(3):58–68.
  * (40) van der Pol, E., Kipf, T., Oliehoek, F. A., and Welling, M. (2020a).  Plannable approximations to mdp homomorphisms: Equivariance under actions.  In AAMAS.
  * (41) van der Pol, E., Worrall, D., van Hoof, H., Oliehoek, F., and Welling, M. (2020b).  Mdp homomorphic networks: Group symmetries in reinforcement learning.  Advances in Neural Information Processing Systems, 33:4199–4210.
  * Vinyals et al., (2019) Vinyals, O., Babuschkin, I., Czarnecki, W. M., Mathieu, M., Dudzik, A., Chung, J., Choi, D. H., Powell, R., Ewalds, T., Georgiev, P., et al. (2019).  Grandmaster level in starcraft ii using multi-agent reinforcement learning.  Nature, 575(7782):350–354.
  * Wang et al., (2022) Wang, D., Walters, R., and Platt, R. (2022).  So (2) equivariant reinforcement learning.  In International conference on learning representations (ICLR).
  * Watkins and Dayan, (1992) Watkins, C. J. and Dayan, P. (1992).  Q-learning.  Machine learning, 8(3):279–292.
  * Ye et al., (2021) Ye, W., Liu, S., Kurutach, T., Abbeel, P., and Gao, Y. (2021).  Mastering atari games with limited data.  Advances in Neural Information Processing Systems, 34.
  * Young and Tian, (2019) Young, K. and Tian, T. (2019).  Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments.  arXiv preprint arXiv:1903.03176.
  * Yu et al., (2021) Yu, C., Li, D., Hao, J., Wang, J., and Burgess, N. (2021).  Learning state representations via retracing in reinforcement learning.  arXiv preprint arXiv:2111.12600.
  * Zhang et al., (2020) Zhang, A., McAllister, R. T., Calandra, R., Gal, Y., and Levine, S. (2020).  Learning invariant representations for reinforcement learning without reconstruction.  In International Conference on Learning Representations.
  * Zhu et al., (2022) Zhu, Z.-M., Jiang, S., Liu, Y.-R., Yu, Y., and Zhang, K. (2022).  Invariant action effect model for reinforcement learning.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 9260–9268.

##  Appendix A Appendix

###  A.1 Impact Statement

We do not believe that EEA introduces any novel dangers or concerns in terms
of broader impact. Like all reinforcement learning algorithms, our agents
likely do not generalise very well and could behave unexpectedly when deployed
in real world environments. We suggest extreme caution when using any RL
algorithm for consequential applications.

###  A.2 Hyperparameter Search

Below we show the hyperparameters swept through for homomorphic MDP, DQN and
Equivalent Effect Abstraction agents, broken down by environment. The PRAE
architectures were generally kept the same as the hyperparameters provided in
(van der Pol et al., 2020a, ), with the exception of the learning rate which
we evaluate at $0.0001,0.001$ and $0.1$

####  A.2.1 Sutton and Barto Tabular Gridworld

We use the hyperparameters specified in (Sutton and Barto,, 2018, p. 165):
namely, Learning Rate$=0.1$, $\gamma=0.95$ and $\epsilon=0.1$

####  A.2.2 Cartpole

Table 1: Hyperparameters swept through for the Cartpole environment. Learning rate decay refers to decaying the learning rate by a factor of ten at after a specified number of episodes have elapsed. Hyperparameter | Values
---|---
Learning Rate | $0.00001,0.0001,0.001,0.01$
$\epsilon$ decay schedule | No decay, exponential $\tau=\frac{-1}{200}$
$\gamma$ | $0.8,0.99$
Activation | ReLU, tanh
Learning Rate decay | No decay, 5, 10, 15, 20

Homomorphic MDP best hyperparameters: Learning Rate = $0.001$, $\epsilon$
decay schedule = No decay, $\gamma=0.8$, activation = tanh, Learning Rate
decay = No decay

Equivalent Effect Abstraction best hyperparameters: Learning Rate = $0.001$,
$\epsilon$ decay schedule = No decay, $\gamma=0.8$, activation = tanh,
Learning Rate decay = 10

Vanilla DQN best hyperparameters: Learning Rate = $0.001$, $\epsilon$ decay
schedule = No decay, $\gamma=0.8$, activation = tanh, Learning Rate decay = 15

####  A.2.3 Stochastic Predator Prey

Table 2: Hyperparameters swept through for the Predator Prey environment. Hyperparameter | Values
---|---
Learning Rate | 0.0001, 0.001, 0.01
$\gamma$ | 0.8, 0.99

Homomorphic MDP best hyperparameters: Learning Rate = $0.001$, $\gamma=0.99$

Equivalent Effect Abstraction best hyperparameters: Learning Rate = $0.01$,
$\gamma=0.8$

Vanilla DQN best hyperparameters: Learning Rate = $0.001$, $\gamma=0.99$

####  A.2.4 Minatar

All hyperparameters were kept at the tuned DQN values provided in Young and
Tian, (2019). For the Homomorphic MDP network, we use a symmetric group for
reflection about the vertical axis when training on Breakout; on Asterix we
show results for both reflection about the vertical and horizontal axes.

###  A.3 Model Architectures

####  A.3.1 Cartpole

Listing 1: Homomorphic MDP Network (van der Pol et al., 2020b, )

⬇

BasisLinear*(repr_in=4, channels_in=1, repr_out=2, channels_out=64)

ReLU() / tanh()

BasisLinear(repr_in=2, channels_in=64, repr_out=2, channels_out=64)

ReLU() / tanh()

BasisLinear(repr_in=2, channels_in=64, repr_out=2, channels_out=1)

*BasisLinear refers to the symmeterised layers used in (van der Pol et al., 2020b, ) to create homomorphic networks. This network is identical to the Cartpole network presented in that paper, but with only one output head that outputs state-action values.

Listing 2: Value Network architecture for DQN and Equivalent Effect
Abstraction

⬇

Linear(input_size=4, output_size=1024)

tanh()

Linear(input_size=1024, output_size=1024)

tanh()

Linear(input_size=8, output_size=1024)

tanh()

Linear(input_size=1024, output_size=2)

Listing 3: Transition Model Architecture for Equivalent Effect Abstraction

⬇

Linear(input_size=2, output_size=2)

Listing 4: PRAE Architectures (van der Pol et al., 2020a, )

⬇

# state encoder

Linear(input_size=4 ,output_size=64)

ReLU()

Linear(input_size=64, output_size=32)

ReLU()

Linear(input_size=32, output_size=50)

#action encoder

Linear(input_size=54 ,output_size=100)

ReLU()

Linear(input_size=100, output_size=2)

# reward prediction network

Linear(input_size=50 ,output_size=64)

ReLU()

Linear(input_size=64, output_size=1)

####  A.3.2 Predator Prey

Listing 5: Homomorphic MDP Network (van der Pol et al., 2020b, )

⬇

BasisConv2d(repr_in=1, channels_in=1, repr_out=4, channels_out=4,

filter_size=(7,7), stride=2, padding=0)

ReLU()

BasisConv2d(repr_in=4, channels_in=4, repr_out=4, channels_out=8,

filter_size=(5,5), stride=1, padding=0)

ReLU()

GlobalMaxPool()

BasisLinear(repr_in=4, channels_in=8, repr_out=4, channels_out=128)

ReLU()

BasisLinear(repr_in=4, channels_in=8, repr_out=4, channels_out=128)

ReLU()

BasisLinear(repr_in=4, channels_in=128, repr_out=5, channels_out=1)

This is again the same network used in van der Pol et al., 2020b , albeit with
a different output head.

Listing 6: Value Network Architecture for DQN and Equivalent Effect
Abstraction

⬇

Linear(input_size=441, output_size=1024)

ReLU()

Linear(input_size=1024, output_size=8)

ReLU()

Linear(input_size=8, output_size=1024)

ReLU()

Linear(input_size=1024, output_size=5)

Listing 7: Transition Model Architecture for Equivalent Effect Abstraction

⬇

Linear(input_size=882, output_size=512)

ReLU()

Linear(input_size=512, output_size=8)

ReLU()

Linear(input_size=8, output_size=512)

ReLU()

Linear(input_size=512, output_size=441)

Listing 8: PRAE Architectures (van der Pol et al., 2020a, )

⬇

# state encoder

Linear(input_size=441, output_size=64)

ReLU()

Linear(input_size=64, output_size=32)

ReLU()

Linear(input_size=32, output_size=50)

#action encoder

Linear(input_size=54, output_size=100)

ReLU()

Linear(input_size=100, output_size=2)

# reward prediction network

Linear(input_size=50, output_size=64)

ReLU()

Linear(input_size=64, output_size=1)

####  A.3.3 MinAtar

Listing 9: Homomorphic MDP Network (van der Pol et al., 2020b, )

⬇

BasisConv2d(repr_in=1, channels_in=4, repr_out=2, channels_out=16,

filter_size=(3,3), stride=1, padding=0)

ReLU()

Linear(input_size=1408, output_size=256)

ReLU()

Linear(input_size=256, output_size=6)

Listing 10: Value Network Architecture for DQN and Equivilent Effect
Abstraction

⬇

Conv2d(channels_in=4, channels_out=16, filter_size=(3,3), stride=1, padding=0)

ReLU()

Linear(input_size=1024, output_size=128)

ReLU()

Linear(input_size=128, output_size=6)

###  A.4 Additional Minatar Results with an Incorrect Model

Here we present some initial results where we tried to naively reuse the
Asterix dynamics model for Breakout (ignoring upwards and downwards movements
as they are not possible in Breakout). However, this model is imperfect as it
always assumes the ball is reflected from the paddle back towards the
direction it came in, which is not always true in Breakout depending on the
side of the paddle that the ball hits. What is more, the vanilla DQN agent
cannot solve the game and all value based methods almost instantly converge to
a suboptimal policy, suggesting that alternative learning algorithms should be
used in the future on this environment.

Figure 6: Here we get similar results to van der Pol et al., 2020b  using our
imperfect model. On this particular environment all versions of DQN converge
almost instantly to a suboptimal policy. To obtain more meaningful results,
future work should integrate equivalent effect abstraction into actor-critic
methods that are able to solve Breakout Young and Tian, (2019).

##  Appendix B Syncing Log arrays

For the sticky MinAtar results we logged frames and return values with wandb
while performing our experiments Biewald, (2020). We performed this logging
suboptimally as we logged steps for frames and returns separately—meaning we
ended up with two arrays where one was NaN when the other was logging a float
and vice versa. We used Panda’s (pandas development team,, 2020) linear
interpolator to estimate the NaN values in our final plots. This process does
not effect our results significantly as the points were logged at a very high
frequency.