T2Acmrcmr

# Is neural language acquisition similar to natural?
A chronological probing study

Ekaterina Voloshina
AIRI, HSE University
Moscow, Russia
<EMAIL_ADDRESS>
&Oleg Serikov
AIRI, DeepPavlov, HSE University
Moscow, Russia
<EMAIL_ADDRESS>
&Tatiana Shavrina
AIRI, SberDevices
Moscow, Russia
<EMAIL_ADDRESS>

###### Abstract

The probing methodology allows one to obtain a partial representation of
linguistic phenomena stored in the inner layers of the neural network, using
external classifiers and statistical analysis.

Pre-trained transformer-based language models are widely used both for natural
language understanding (NLU) and natural language generation (NLG) tasks
making them most commonly used for downstream applications. However, little
analysis was carried out, whether the models were pre-trained enough or
contained knowledge correlated with linguistic theory.

We are presenting the chronological probing study of transformer English
models such as MultiBERT and T5. We sequentially compare the information about
the language learned by the models in the process of training on corpora. The
results show that 1) linguistic information is acquired in the early stages of
training 2) both language models demonstrate capabilities to capture various
features from various levels of language, including morphology, syntax, and
even discourse, while they also can inconsistently fail on tasks that are
perceived as easy.

We also introduce the open-source framework for chronological probing
research, compatible with other transformer-based models.
https://github.com/EkaterinaVoloshina/chronological_probing

Keywords: probing, language acquisition, language modeling, transformers

DOI: 10.28995/2075-7182-2022-20-XX-XX

###### ?abstractname?

Пробинг-методология позволяет получить представление о явлениях языка,
хранящееся во внутренних слоях нейросети, с помощью внешних классификаторов и
статистического анализа.

Предобученные языковые модели на основе трансформерных архитектур широко
используются как для задач понимания естественного языка (NLU), так и для
задач генерации (NLG), что делает их наиболее часто используемыми для широкого
ряда применений. Однако, недостаточно проводился анализ, достаточно ли
предобучены модели и содержат ли знания, коррелирующие с теоретическими
представлениями о языке.

Мы представляем исследование на основе хронологического пробинга на примере
моделей MultiBERT и T5, в котором последовательно исследуем выучиваемую
информацию о языке в процессе предобучения моделей на корпусе. Результаты
показывают, что 1) лингвистическая информация усваивается уже на ранних этапах
обучения 2) обе языковые модели демонстрируют способность фиксировать
различные свойства языка на разных уровнях, включая морфологию, синтаксис и
дискурс, в то же время они могут не справляться с задачами, которые
воспринимается как простые.

Мы также предоставляем открытый фреймворк для хронологического пробинга,
совместимый с языковыми моделями на основе архитектур transformer.
https://github.com/EkaterinaVoloshina/chronological_probing

Ключевые слова: пробинг, усвоение языка, языковые модели, трансформеры

Усвоение языка у языковых моделей и человека: хронологическое пробинг-
исследование.

Екатерина Волошина | Олег Сериков | Татьяна Шаврина
---|---|---
AIRI, НИУ ВШЭ | AIRI, DeepPavlov, НИУ ВШЭ | AIRI, SberDevices
Москва, Россия | Москва, Россия | Москва, Россия
<EMAIL_ADDRESS>|<EMAIL_ADDRESS>|<EMAIL_ADDRESS>

##  1 Introduction

The role of deep learning language models has been increasing in the field of
methodology for linguistic research, providing new methods for both diachronic
and synchronic studies [Manning, 2015]. In particular, transformer-based
language modeling research has produced a variety of tools that may discover
regularities and structures in data, many of which have resulted in practical
applications.

In this study, we search for a match between the language competencies of
popular language models and compare their results with the levels of a first
language learner. As the transformer models are expected to acquire a language
during the training process, the probing methodology has shed light on model
training success. Probing tasks are usually classification tasks where classes
represent different values of a linguistic features, such as a subject number,
tree depth, and a connector type. Theoretical representation of language often
inquires about the levels of phonetics, morphology, syntax, and
discourse/pragmatics to be involved in a probing study 111However, some
researchers [Embick and Noyer, 2007, Caha, 2009] doubt that a language
functions as a level system. They suggest that morphology and syntax operate
at the same time. Other researchers argue that morphology and syntax are
different layers of a language..

The main focus of this work is to explore how language models acquire
measurable linguistic structures during training. The contributions of our
work are the following:

  * •

We propose a methodology for chronological probing, based on checkpoint-wise
result comparison during model training222
https://github.com/EkaterinaVoloshina/chronological_probing. We denote
chronological probing as any probing technique that refers to the training
history/iterations of the same model.

  * •

We test two models (MultiBERT[Sellam et al., 2021] and T5[Raffel et al.,
2019]) on existing 12 different probing tasks in morphology, syntax, and
discourse and present an analysis of the models’ gradual learning of language
phenomena, in comparison with the well-known facts about the acquisition of
the first language by a child.

  * •

We present the evaluation results for the named models and state that the
models tend to learn the linguistic phenomena in a specific order, and some
parts of grammar are “acquired” first.

The presented framework and methodology are available open-source under Apache
2.0 license.

##  2 Related work

Until recent years, the task of learning syntax, which every five-year-old
child performs effortlessly, has eluded brute language modeling force. This
makes the language models a particular object of study, considered both from
the interpretability and modeling language acquisition. As [De Villiers and
Roeper, 2011, p. 119] states, “computational models of language acquisition
must begin and end as an integral part of the empirical study of child
language.”

Following this thesis, we turn our attention to the probing methodology and
comparable case studies in the field of language acquisition, focusing on the
transformer architectures.

###  2.1 Probing and approaches to the black box of language modeling

An increasing number of works are devoted to interpreting language models from
a linguistic point of view. The quickly advancing field of probing received
lots of researchers’ attention when the hegemony of the large black-box models
was set up. Researchers question the extent of the models’ “understanding” of
the language in probing. They inspect if, and to what limits, the language
models’ behavior agrees with the insights of the theory of language. Following
the hierarchy of language levels (morphology, syntax, discourse) [Dalrymple,
2001], the probing studies often suppose the experiments related to models’
proficiency on a certain level of language.

This line of research typically comes down to analyzing how linguistic
structures are represented in a model’s knowledge. Such structures represent
syntagmatic/paradigmatic mechanisms (how language units combine and alternate,
respectively) of language. It is believed [McCoy et al., 2020], that
rediscovering these structures would help models to get closer to human
performance on a variety of tasks.

Probing, in general, considers how interpretable the behavior of the language
model wrt the linguistic properties of the data. A huge body of probing
studies rely on linear models (e.g., external classifiers [Belinkov, 2016])
that try to establish the relationship between internal representations from
the language model and the desired linguistic phenomena. Thus, the linear
correlation is measured between the model’s forward pass embeddings and the
linguistic properties of the passed data. A sample study [Tenney et al., 2019]
could measure the strength of correlation between a model’s particular layer
activations on some word and word’s part-of-speech.

Strong correlations have been recorded when comparing the models’ forward pass
activations with the passed data underlying linguistic structure [Belinkov et
al., 2017, Tenney et al., 2019, Conneau et al., 2018a, Hewitt and Liang,
2019a] using probing methods.

Such a high performance could be misleading. The properties of the model and
the properties of the used data impact the resulting score of the correlation
probing study. Thus, given only a correlation score, one does not know if it
reflects the model’s (but not the corpus itself) linguistic informativeness.
As a result, several approaches to conducting more reliable studies have been
proposed. [Hewitt and Liang, 2019a, Zhang and Bowman, 2018, Voita and Titov,
2020, Pimentel et al., 2020].

The probing methodology combining various annotated data is commonly used as
the benchmark for language model comparison and evaluation of their
generalizing ability. The SentEval toolkit [Conneau and Kiela, 2018] has led
to the popularization of the 10 tasks used to distinguish between random and
justified, brittle, and robust results of model training, including different
types of architectures. However, analogous research on the same architecture
or even the same model is in its early development stage. The first work on
probing of neural networks across time was carried by [Saphra and Lopez,
2018]. The authors showed that first, LSTM acquires syntactic and semantic
features and later information structure. [Chiang et al., 2020a] looked at the
training process of ALBERT and concluded that semantic and syntactic
information is acquired during the early steps while world knowledge
fluctuates during the training. [Liu et al., 2021] showed similar results on
RoBERTa: the model shows good results on linguistic probing tasks starting
from early stages, and later it learns factual and commonsense knowledge.

Chronological probing could enrich the interpretable documentation of model
training in time and thus explore the new aspects of model training and more
clearly expose its problems.

###  2.2 Language acquisition and language models

Language learning is one of the quintessential human traits. First language
acquisition(LA), unites both neurocognitive research, psycholinguistics, and
computational approaches, focusing on the ability to acquire the capacity to
perceive and comprehend language.

#### Statistical language acquisition

Language modeling has formed a branch in language acquisition studies named
statistical language acquisition. Various aspects of language, including
phonological, syntactic, lexical, morphological, and semantic features, were
investigated in terms of statistical patterns children receive with the
linguistic input. Recent studies postulating qualitative and quantitative
measures of LA include:

  * •

Morphology and Syntax Morphology and syntax studies across language
acquisition studies are definitely those explored the most. Starting with the
poverty of stimulus problem and the argument between innateness and learning
of grammar, it has led to typologically various sets of descriptive works and
even computational models of the acquisition process. Thus, [Lewis and Elman,
2001] train a simple RNN to discriminate between grammatical strings that
follow the inversion rule and those that do not (e.g., moving the first
auxiliary verb such as “Is the man that tall is nice?”). The training data for
the study is generated artificially and fails to prove that such a network
generalizes on a mixture of diverse syntactic constructions. [Reali and
Christiansen, 2005] use bigram models to capture the patterns of auxiliary
inversion based on lifelike data from child-directed speech. The model can
consistently assign higher probabilities to grammatical strings than
ungrammatical strings, which was interpreted as having successfully learned
the correct inversion rule. However, as [Kam et al., 2008] note, this result
is because bigrams such as “who are” are much more frequent than the
ungrammatical strings. [Prefors et al., 2006] approaches the structure
dependency problem with Bayesian learning and attempts to learn a grammar that
could generate additional sentences. The model evaluates and selects between
two grammars, a finite state grammar and a context-free grammar constructed by
the authors based on a simplified subset of child-directed English.

It is worth noting how similar all the problem formulations are to the modern
formulations of probing classification problems described below. They are also
far from a complete description of the process of mastering grammar.

  * •

Discourse The creation of texts, not sentences, with various discourse
features, such as competence in speech acts, conversations, speech registers,
and extended speaking turns, is more often considered a later stage of speech
development. There are no computer models for the assimilation of discursive
properties comparable to models for morphology and syntax. However, research
in this direction is underway.

In [Ororbia et al., 2019], authors examine whether neural language models
acquire language better when trained in a multi-modal setting (namely,
accompanied with visual contexts) compared to traditional purely textual pre-
training. They show that indeed providing models with perceptual context is
beneficial for training language models. Authors claim this evidence to
correspond with the theory of situated cognition introduced in [Greeno and
Moore, 1993].

In this work, we propose the first methodological step for chronological
interpretation of traditional transformer language models in the framework of
LA.

##  3 Experimental setup

###  3.1 Models

We calculated the accuracy of two transformer models on 12 probing tasks. As
we want to know how universal patterns of language acquisition in models are,
we experiment with two different transformer architectures: BERT and T5. While
BERT has only encoder layers, T5 includes both encoder and decoder layers.
Therefore, embeddings from BERT come from the encoder, and T5 embeddings are
calculated after going through decoder after encoder.

For this work, we use already published models with available checkpoints. It
means that they were trained on different data and computational powers.
Moreover, they were trained with different batch sizes (256 for MultiBERT and
32 for T5). However, we follow [Chiang et al., 2020b, Liu et al., 2021] and
measure the training progress in iterations. The further comparison of the two
models is indicative only.

MultiBERT [Sellam et al., 2021] is based on BERT-base-uncased architecture,
and it is the model of the same size (12 layers and embedding size 768).
Unlike the original BERT[Devlin et al., 2019], it was trained with 25
different seeds. The authors also released checkpoints of the first five
models. We use the model with seed 0 in our experiments. MultiBERT was trained
on both literary and non-fiction texts. The former comes from BookCorpus [Zhu
et al., 2015], which includes 11,038 books of 16 different genres. The non-
fiction texts are taken from English Wikipedia collected by [Turc et al.,
2019].

T5-small model is trained within the t5-experiments
framework.333https://github.com/yurakuratov/t5-experiments and follows the
Hugging Face implementation of T5 [Raffel et al., 2019]. It consists of 6
layers with 512 embedding size.

Following the previous language, modeling works [Devlin et al., 2019,
Bojanowski et al., 2017], we use the Wikipedia data to train the model. The
raw Wikipedia data is provided by The Pile project [Gao et al., 2020] contains
$\approx 19\text{Gb}$ of expository prose texts of various domains, and is
treated as a language modeling dataset of reasonably well quality.

Baseline As a baseline, we use the method described in [Hewitt and Liang,
2019b]. We train logistic regression on top of embeddings of models mentioned
above with shuffled class labels.

Task |  Sentence examples | Labels
---|---|---
Subject number |  Her employer had escaped with his wife for several afternoons this summer. | NN
|  Your Mackenzie in-laws have sordid reputations few decent families wish to be connected with . | NNS
Person |  So I still can recomend them but prepare pay twice as much as they tell you initially . | has a person marker
|  The service was friendly and fast , but this just does nt make up for the lack - luster product . | does not have a person marker
Tree depth |  We have done everything we can for her . | 11
|  Alvin Yeung of Civic Party | 3
Top constituents |  Did it belong to the owner of the house ? | VBD_NP_VP_.
|  How long before you leave us again ? | WHNP_SQ_.
Connectors |  He ’d almost forgotten about that man . Sarah had somehow brought him back , just as she had his nightmares . | but
|  I let out a slow , careful breath . Felt tears sting my eyes . | and
Sentence position |  Quneitra Governorate ( / ALA-LC : “ Muhafzat Al-Qunaytrah “ ) is one of the fourteen governorates ( provinces ) of Syria . The governorate had a population of 87,000 at the 2010 estimate . Its area varies , according to different sources , from 685 km ² to 1,861 km ² . It is situated in southern Syria , notable for the location of the Golan Heights . The governorate borders Lebanon , Jordan and Israel . | 1
|  The bossom and the part of the xhubleta covered by the apron are made out of crocheted black wool . The bell shape is accentuated in the back part . The xhubleta is an undulating , bell-shaped folk skirt , worn by Albanian women . It usually is hung on the shoulders using two straps . Part of the Albanian traditional clothing it has 13 to 17 strips and 5 pieces of felt . | 4
Penn Discourse Treebank |  Solo woodwind players have to be creative,they want to work a lot | Pragmatic Cause
|  The U.S. , along with Britain and Singapore , left the agencyl, its anti-Western ideology , financial corruption and top leadership got out of hand | List
Discourse Coherence |  Within the fan inlet case , there are anti-icing air bosses and probes to sense the inlet pressure and temperature .’, ’High speed center of pressure shifts along with fin aeroelasticity were major factors . At the 13th ( i.e .’, ’the final ) compressor stage , air is bled out and used for anti-icing . The amount is controlled by the Pressure Ratio Bleed Control sense signal ( PRBC ) . The “ diffuser case “ at the aft end of the compressor houses the 13th stage . | a text is not coherent
|  This experience of digital circuitry and assembly language programming formed the basis of his book “ Code : The Hidden Language of Computer Hardware and Software ” . Petzold purchased a two-diskette IBM PC in 1984 for $ 5,000 . This debt encouraged him to use the PC to earn some revenue so he wrote an article about ANSI.SYS and the PROMPT command . This was submitted to PC Magazine for which they paid $ 800 . This was the beginning of Petzold ’s career as a paid writer . In 1984 , PC Magazine decided to do a review of printers . | a text is coherent

Table 1: Examples of tasks

Task |  Acceptable sentence |  Unacceptable sentence
---|---|---
Transitive |  The pedestrians question some people. |  The pedestrians wave some people.
Passive |  Tracy isn’t fired by Jodi’s daughter. |  Tracy isn’t muttered by Jodi’s daughter.
Principle A c command |  This lady who is healing Charles wasn’t hiding herself. |  This lady who is healing Charles wasn’t hiding himself.
Adjunct Island |  Who does John leave while alarming Beverly? |  Who does John leave Beverly while alarming?

Table 2: BLiMP Minimal pairs examples

###  3.2 Probing tasks

Probing tasks come from several datasets published earlier: SentEval [Conneau
et al., 2018b], Morph Call [Mikhailov et al., 2021], DisSent [Nie et al.,
2019], DiscoEval [Chen et al., 2019], and BLiMP [Warstadt et al., 2020]. The
class balance of first eight tasks is illustrated with Figure 4 in Appendix.
We choose these tasks to make our results comparable to other works on
probing.

As we want to show another perspective on language acquisition, we balance
classifier probing tasks with BLiMP tasks. As BLimP only covers morphology and
syntax, all discourse-based tasks are evaluated with a classifier.

The datasets from Benchmark of minimal linguistic pairs (BLiMP) have a
structure different from other tasks: every task includes pairs with minimal
differences to illustrate one of the grammatical features of English. One
sentence of the pair is grammatical, whereas another one is unacceptable. We
chose four BLiMP tasks for our experiments: transitive and passive verbs,
Principle A of C command, and Island effects. For the first two tasks, pairs
have different verbs, where only one verb is transitive or can be used in a
passive form. These tasks are categorized as morphological (see Table 2).

The second two tasks reflect syntactic effects on English. The Principle A
task shows the use of reflexives. According to [Chomsky, 1981], a reflexive
should have a local antecedent.

The task on Island effects tests a model’s sensibility to syntactic order.
Island is a structure from which a word can not be moved [Ross, 1967]. The
sentence is unacceptable if a word is moved out of an island.

The tasks from other datasets are summarized below:

  * •

Subject number (SentEval): this task is supposed to show how models acquire
morphology. It is a binary classification task with labels NNS and NN (plural
and singular number, respectively). The classifier should decide on a sentence
class based on the number of sentence subjects.

  * •

Person (Morph Call): this task is also morphological. It is a binary
classification with labels 0 and 1, which signifies if a subject has a person
marker or not.

  * •

Tree depth (SentEval): this task contains six classes, which stands for the
depth of the syntactic tree of a given sentence. Hence, this task is meant to
show the level of syntax acquisition.

  * •

Top constituents (SentEval): this multiclass task belongs to the group of
syntactic tasks. The aim is to choose a class that includes constituents
located right below the sentence (S) node.

  * •

Connectors (DisSent): this dataset includes pairs of sentences originally
connected with one of 5 prepositions, and the task is to choose the omitted
preposition. It is supposed to show how models catch discourse relations.

  * •

Sentence position (DiscoEval): this dataset contains sequences of 5 sentences,
and the first sentence is placed in the wrong place. Therefore, the aim is to
detect the original position of these sentences. This task is also meant to
show models’ accuracy in discourse.

  * •

Penn Discourse Treebank (DiscoEval): the task is based on Penn Discourse
Treebank annotation. The aim is to choose the right discourse relation between
two discourse items from Penn Treebank.

  * •

Discourse coherence (DiscoEval): this task is a binary classification with
classes 1 and 0. Class 1 means that the given paragraph is coherent, and class
0 should be assigned to paragraphs with shuffled sentences.

###  3.3 Probing Methods

Sentence embedding classification: Token embeddings from transformer models
are turned into sentence embedding via mean pooling. Then logistic regression
classifies embeddings’ sentences. This method is used with tasks from SentEval
and Morph Call.

Positional sentence classification: For the Sentence Position task, we used
the following method. First, we count sentence embeddings as described above.
Then the difference between the first embedding and the other is calculated
pair-wisely. The first embedding and its differences with others are
concatenated and put as input to logistic regression.

Sentence embedding concatenation & classification: We concatenated sentence
embeddings for other discourse tasks, which were calculated as the mean of
token embeddings. The concatenated sentence embeddings served as inputs for
logistic regression.

Masking tasks: The probing task is based on the idea of masking language
modeling. In a sentence, each word is masked, and then its probability is
summed with other words’ probabilities. The probability of an acceptable
sentence should be higher than the probability of an unacceptable sentence.
This method is for use for all tasks from BLiMP.

##  4 Results

###  4.1 Results of MultiBERT

The results of the experiments with the MultiBERT-base model are summarized in
Figure 1. The model shows the best results on Subject Number and Person tasks.
The classification of PDTB relations, Tree depth, and Principle A
acceptability are performed with the worst accuracy.

Figure 1: Comparison of MultiBERT results on different tasks. How to read this
figure: from left to right, on the X axis, we see results of intermediate
evaluation on the task during model pre-training. Each iteration is equal to
25,600,000 sentences for MultiBERT and 3,200,000 sentences for T5. The Y-axis
shows the accuracy metric on the tasks. Tasks are shown in the legend in
different colors. As we can see, in the process of model pre-training, there
already is a gradual increase in accuracy in tasks related to morphology
(shown in orange) in the early stages. The information in the model embeddings
stabilizes fairly quickly and remains stable from the 20,000th training step.
The same can not be said for tasks related to syntax (shown in purple): their
quality remains unstable and fluctuates quite a lot during pre-training.
Discourse tasks (green) remain stable at a low-quality level from the start
and tend to improve the metrics very slowly.

As seen from Figure 1, accuracy of models stop changing after 600,000
iterations. However, there is a significant difference between tasks from
BLiMP and other datasets. For example, the performance on the Adjunct Island
task remains unstable during the whole period of iterations. Another
difference between these tasks lies in the quality of the models. It is
illustrated with tasks grouped as “morphological”: Subject Number and Person
tasks, which use logistic regression on MultiBERT embeddings, are solved much
better than Transitive and Passive verbs. However, as follows from the plot,
it is hard to group tasks based on the absolute value of accuracy.

The changing dynamics provide another perspective. From this point of view,
all tasks grouped as “discourse” show a similar feature: unlike others, their
quality does not fluctuate but rather slightly grows across the training time.
On other tasks, models increase the quality during the first iterations.
“Syntactic” tasks tend to change even during later iterations. However, it is
not a strict rule, and some tasks show similar behavior to “morphological”
ones.

###  4.2 Results of T5

Figure 2: Comparison of T5 results on a different task. How to read this
figure: from left to right, on the X axis, we see results of intermediate
evaluation on the task during model pre-training. Each iteration is equal to
25,600,000 sentences for MultiBERT and 3,200,000 sentences for T5. The Y-axis
shows the accuracy metric on the tasks. Tasks are shown in the legend in
different colors. As we can see, in the case of the T5 model, the task quality
seems to be more stable from the beginning, with a few exceptions like subject
number classification. Most of the tasks show the slow yet gradual growth of
the metrics, but somehow not the verb transitivity classification.

Due to the architecture, the significant difference in T5 results is the zero-
close quality on BLiMP datasets. Except for these tasks, the quality of T5 is
similar to MultiBERT. The best performance is on the Person task, and the
worst quality is shown on PDTB relation classification and Tree depth.

Unlike MultiBERT, we first used the available checkpoints of T5 with a step of
100,000 iterations. Then we trained a new model on the same resources and
texts, but it might have a better initialization, which affected the final
results.

Similar to MultiBERT, discourse tasks show almost no significant change and
slow growth, whereas the model increases its results on syntactic and
morphological tasks during the first 100,000 iterations.

###  4.3 Comparison of models

We described the surface results of models’ performance and now can deep into
more detailed results. The results described above should be considered
relative. To illustrate how much information models acquire during these
iterations, we compare them to final models. As the process of training T5 was
not finished, we compared this model with the original T5. As seen from Figure
3, MultiBERT scores are close to the results of the final checkpoint. Hence,
there is no need to look at later iterations. The comparison with the original
T5 shows that the model we use performs worse due to the smaller resources it
was trained on. Therefore, the difference in quality should not be explained
by the difference in architecture.

However, we should consider that some tasks are performed with the same
quality as embeddings with shuffled labels (Discourse coherence and Person).
Moreover, T5 does not perform much better than the Penn Discourse Treebank
relations baseline. Consequently, models encounter difficulty with discourse
tasks.

Furthermore, MultiBERT and T5 show similar learning trajectories on several
tasks, such as Connectors and Sentence Position tasks. Another key feature
shared by the two models is the termination of increases between 500,000 and
600,000 iterations. Despite the fact that models vary in size and training
process, they show some similarities in probing tasks. Hence, the acquisition
generally does not depend on the model architecture.

Figure 3: Performance on models on different tasks. The detailed task-wise
comparison shows the difference in T5 and MultiBERT training results. The
models are compared to the final available checkpoints of the models (Original
T5, Final BERT) and with the random baseline.

##  5 Discussion

Our results show that linguistic information is acquired fast, before 600,000
training iterations. It corresponds to results of other researchers [Chiang et
al., 2020a, Liu et al., 2021] that independently showed similar results on a
fast acquisition of linguistic features. However, discourse is not fully
acquired by the end of the observed training period compared to the baseline
results. The difference between syntactic and morphological tasks is
insignificant. It correlates with ideas in morphosyntax. Although we can not
prove that morphology and syntax are regarded as the same layer in models, we
can make a less strict statement that models acquire all grammatical units
simultaneously.

BLiMP gives another perspective on the process of acquisition. MultiBERT
results remain unstable for a longer period than similar tasks with
classifiers. It might indicate the difference between two different approaches
to probing. However, from the linguistic point of view, BLiMP includes more
difficult linguistic feature cases, while SentEval tasks test more basic
knowledge. Hence, it could explain worse results.

T5 architecture does not allow to use of this dataset in the same way as for
MultiBERT since Masking Language Modeling and T5 generation are different
tasks. We leave for further research an adaptation of this dataset for T5.

###  5.1 Human language acquisition results

Many of the linguistic features used in probing tasks have been well studied
in terms of their promptness and ease of acquisition by English speakers.
First of all, it concerns morphology and syntax. Markers such as a person,
number, and gender are acquired very early by children: before age two [Clark,
2017]. Of course, in languages besides English, the acquisition of these
features varies: if the feature is marked consistently with one affix and no
morphonological alternation, children seem to find it easier to acquire. It is
shown that the earlier mastery of case marking is present in languages like
Hungarian and Turkish but not in German or Serbo-Croatian [Slobin, 1985].

Discursive features are acquired by children much later. Studies like
[Pearson, 2003] show that child’s texts become more complex and
decontextualized with age. Also, texts produced by children gradually progress
in achieving more cohesion through “referential and semantic links that bridge
across sentences; they achieve coherence through a global hierarchical
structure”. The discourse in these conversations between toddlers is tentative
when neither side can be reliably significant. A longitudinal study of
dialogues between two little girls aged 4 to 6, [McTear, 1985] traced the
emergence of more and greater thematic continuity in their conversation as
utterances began to play the dual role of responding to the preceding
utterance as well as enabling further conversation. However, [Dorval et al.,
1984] showed that second-graders (eight-year-olds) were almost as likely to
give unconditioned responses as conditionals, with no significant improvement
seen until fifth grade.

Regarding the requirements that the sphere of language acquisition imposes on
children, one can very carefully assess the limit in which the language models
under consideration lie in terms of their abilities: their embeddings
correspond to the level of language proficiency in a child under 11 years old.

##  6 Conclusion

Encoder and encoder-decoder language models have increased importance in tasks
requiring understanding the natural language. The probing methodology we
presented allows analyzing the changes within the language model during
training, from epoch to epoch. The overall results of the work show that:

  * •

T5 does not give any results on BLiMP due to the generation algorithm. Most
tasks show that T5 acquires basic morphological and syntactic features and
some discourse features.

  * •

MultiBERT shows results close to the trained model starting from 100,000
iterations. MultiBERT does not improve its quality on some discourse tasks
compared to randomly labeled embeddings. However, it could be said that
MultiBERT acquires each level to some extent.

  * •

Both T5 and MultiBERT demonstrate comparable results regarding the quality of
the language level acquisition. As we can not distinguish the factors between
these results (whether this is the model’s architecture, the training corpora,
or both), we present these results ’as is’ for the researchers that use them
in the downstream tasks.

  * •

Recording such results during training could save a lot of computational
resources and time for interpreting the results, including downstream tasks.
There are understandable context length limitations that prevent, for example,
learning the discourse tasks. However, the results of the T5 model compared to
the random embeddings on some tasks were unexpected or lower than expected.

  * •

As the results show, the easiest tasks for the models tend to be morphology
and syntax-related. These language level results are correlated and show a
similar learning trajectory. Unlike morphology and syntax tasks, results on
discourse-based tasks tend to be low, therefore, there is not enough evidence
to claim that discourse has been learned.

  * •

Using language acquisition as a trace can benefit in comparing general human
language ability and modern language modeling methods. Drawing a border from
above on the results on discursive tasks, we can say that in the current
realities, the models do not pass the bar that 8-year-old children pass.

We welcome both NLP and language acquisition research communities to reproduce
the experimental setup and use the presented approach while training other
architectures and contribute to formulating the better and more complex tasks
correlated with language learning.

## References

  * [Belinkov et al., 2017]  Yonatan Belinkov, Lluís Màrquez, Hassan Sajjad, Nadir Durrani, Fahim Dalvi, and James Glass.  2017\.  Evaluating layers of representation in neural machine translation on part-of-speech and semantic tagging tasks.  // Proceedings of the Eighth International Joint Conference on Natural Language Processing (Volume 1: Long Papers), P 1–10, Taipei, Taiwan, November. Asian Federation of Natural Language Processing.
  * [Belinkov, 2016]  Yonatan Belinkov.  2016\.  Probing classifiers: Promises, shortcomings, and advances.  Computational Linguistics, P 1–12.
  * [Bojanowski et al., 2017]  Piotr Bojanowski, Edouard Grave, Armand Joulin, and Tomas Mikolov.  2017\.  Enriching word vectors with subword information.  Transactions of the Association for Computational Linguistics, 5:135–146.
  * [Caha, 2009]  Pavel Caha.  2009\.  The nanosyntax of case.
  * [Chen et al., 2019]  Mingda Chen, Zewei Chu, and Kevin Gimpel.  2019\.  Evaluation benchmarks and learning criteria for discourse-aware sentence representations.  // Proc. of EMNLP.
  * [Chiang et al., 2020a]  Cheng-Han Chiang, Sung-Feng Huang, and Hung-yi Lee.  2020a.  Pretrained language model embryology: The birth of albert.  arXiv preprint arXiv:2010.02480.
  * [Chiang et al., 2020b]  Cheng-Han Chiang, Sung-Feng Huang, and Hung-yi Lee.  2020b.  Pretrained language model embryology: The birth of ALBERT.  // Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), P 6813–6828, Online, November. Association for Computational Linguistics.
  * [Chomsky, 1981]  Noam Chomsky.  1981\.  Lectures on government and binding (dordrecht: Foris).  Studies in generative grammar, 9.
  * [Clark, 2017]  Eve V Clark.  2017\.  Morphology in language acquisition.  The handbook of morphology, P 374–389.
  * [Conneau and Kiela, 2018]  Alexis Conneau and Douwe Kiela.  2018\.  Senteval: An evaluation toolkit for universal sentence representations.  arXiv preprint arXiv:1803.05449.
  * [Conneau et al., 2018a]  Alexis Conneau, German Kruszewski, Guillaume Lample, Loïc Barrault, and Marco Baroni.  2018a.  What you can cram into a single $&!#* vector: Probing sentence embeddings for linguistic properties.  // Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), P 2126–2136, Melbourne, Australia, July. Association for Computational Linguistics.
  * [Conneau et al., 2018b]  Alexis Conneau, German Kruszewski, Guillaume Lample, Loic Barrault, and Marco Baroni.  2018b.  What you can cram into a single vector: Probing sentence embeddings for linguistic properties.  arXiv preprint arXiv:1805.01070.
  * [Dalrymple, 2001]  Mary Dalrymple.  2001\.  Lexical functional grammar.  Brill.
  * [De Villiers and Roeper, 2011]  Jill G De Villiers and Thomas Roeper.  2011\.  Handbook of generative approaches to language acquisition, volume 41.  Springer.
  * [Devlin et al., 2019]  Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  2019\.  Bert: Pre-training of deep bidirectional transformers for language understanding.
  * [Dorval et al., 1984]  Bruce Dorval, Carol O Eckerman, and Susan Ervin-Tripp.  1984\.  Developmental trends in the quality of conversation achieved by small groups of acquainted peers.  Monographs of the Society for Research in Child Development, P 1–91.
  * [Embick and Noyer, 2007]  David Embick and Rolf Noyer.  2007\.  Distributed morphology and the syntax/morphology interface.  The Oxford handbook of linguistic interfaces, 289324.
  * [Gao et al., 2020]  Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al.  2020\.  The pile: An 800gb dataset of diverse text for language modeling.  arXiv preprint arXiv:2101.00027.
  * [Greeno and Moore, 1993]  James G Greeno and Joyce L Moore.  1993\.  Situativity and symbols: Response to vera and simon.
  * [Hewitt and Liang, 2019a]  John Hewitt and Percy Liang.  2019a.  Designing and interpreting probes with control tasks.  // Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), P 2733–2743, Hong Kong, China, November. Association for Computational Linguistics.
  * [Hewitt and Liang, 2019b]  John Hewitt and Percy Liang.  2019b.  Designing and interpreting probes with control tasks.  arXiv preprint arXiv:1909.03368.
  * [Kam et al., 2008]  Xuân-Nga Cao Kam, Iglika Stoyneshka, Lidiya Tornyova, Janet D Fodor, and William G Sakas.  2008\.  Bigrams and the richness of the stimulus.  Cognitive science, 32(4):771–787.
  * [Lewis and Elman, 2001]  John D Lewis and Jeffrey L Elman.  2001\.  Learnability and the statistical structure of language: Poverty of stimulus arguments revisited.  // Proceedings of the 26th annual Boston University conference on language development, volume 1, P 359–370. Citeseer.
  * [Liu et al., 2021]  Leo Z Liu, Yizhong Wang, Jungo Kasai, Hannaneh Hajishirzi, and Noah A Smith.  2021\.  Probing across time: What does roberta know and when?  arXiv preprint arXiv:2104.07885.
  * [Manning, 2015]  Christopher D Manning.  2015\.  Last words: Computational linguistics and deep learning.
  * [McCoy et al., 2020]  R. Thomas McCoy, Junghyun Min, and Tal Linzen.  2020\.  BERTs of a feather do not generalize together: Large variability in generalization across models with similar test set performance.  // Proceedings of the Third BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP, P 217–227, Online, November. Association for Computational Linguistics.
  * [McTear, 1985]  Michael McTear.  1985\.  Children’s conversation.  B. Blackwell.
  * [Mikhailov et al., 2021]  Vladislav Mikhailov, Oleg Serikov, and Ekaterina Artemova.  2021\.  Morph call: Probing morphosyntactic content of multilingual transformers.  // Proceedings of the Third Workshop on Computational Typology and Multilingual NLP, P 97–121, Online, June. Association for Computational Linguistics.
  * [Nie et al., 2019]  Allen Nie, Erin Bennett, and Noah Goodman.  2019\.  Dissent: Learning sentence representations from explicit discourse relations.  // Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, P 4497–4510.
  * [Ororbia et al., 2019]  Alexander Ororbia, Ankur Mali, Matthew Kelly, and David Reitter.  2019\.  Like a baby: Visually situated neural language acquisition.  // Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, P 5127–5136, Florence, Italy, July. Association for Computational Linguistics.
  * [Pearson, 2003]  Barbara Zurer Pearson.  2003\.  Language acquisition: Discourse, narrative and pragmatics.  Disertasi. USA: Department of Communication Disorders.
  * [Pimentel et al., 2020]  Tiago Pimentel, Josef Valvoda, Rowan Hall Maudslay, Ran Zmigrod, Adina Williams, and Ryan Cotterell.  2020\.  Information-theoretic probing for linguistic structure.  CoRR, abs/2004.03061.
  * [Prefors et al., 2006]  Amy Prefors, Terry Regier, and Joshua B Tenenbaum.  2006\.  Poverty of the stimulus? a rational approach.  // Proceedings of the Annual Meeting of the Cognitive Science Society, volume 28.
  * [Raffel et al., 2019]  Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  2019\.  Exploring the limits of transfer learning with a unified text-to-text transformer.  arXiv preprint arXiv:1910.10683.
  * [Reali and Christiansen, 2005]  Florencia Reali and Morten H Christiansen.  2005\.  Uncovering the richness of the stimulus: Structure dependence and indirect statistical evidence.  Cognitive Science, 29(6):1007–1028.
  * [Ross, 1967]  John Robert Ross.  1967\.  Constraints on variables in syntax.
  * [Saphra and Lopez, 2018]  Naomi Saphra and Adam Lopez.  2018\.  Understanding learning dynamics of language models with svcca.  arXiv preprint arXiv:1811.00225.
  * [Sellam et al., 2021]  Thibault Sellam, Steve Yadlowsky, Jason Wei, Naomi Saphra, Alexander D’Amour, Tal Linzen, Jasmijn Bastings, Iulia Turc, Jacob Eisenstein, Dipanjan Das, et al.  2021\.  The multiberts: Bert reproductions for robustness analysis.  arXiv preprint arXiv:2106.16163.
  * [Slobin, 1985]  Dan I Slobin.  1985\.  Crosslinguistic evidence for the language-making capacity.  The crosslinguistic study of language acquisition, 2:1157–249.
  * [Tenney et al., 2019]  Ian Tenney, Patrick Xia, Berlin Chen, Alex Wang, Adam Poliak, R Thomas McCoy, Najoung Kim, Benjamin Van Durme, Samuel R. Bowman, Dipanjan Das, and Ellie Pavlick.  2019\.  What do you learn from context? Probing for sentence structure in contextualized word representations.  arXiv e-prints, P arXiv:1905.06316, May.
  * [Turc et al., 2019]  Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  2019\.  Well-read students learn better: On the importance of pre-training compact models.  arXiv preprint arXiv:1908.08962.
  * [Voita and Titov, 2020]  Elena Voita and Ivan Titov.  2020\.  Information-theoretic probing with minimum description length.  arXiv preprint arXiv:2003.12298.
  * [Warstadt et al., 2020]  Alex Warstadt, Alicia Parrish, Haokun Liu, Anhad Mohananey, Wei Peng, Sheng-Fu Wang, and Samuel R Bowman.  2020\.  Blimp: The benchmark of linguistic minimal pairs for english.  Transactions of the Association for Computational Linguistics, 8:377–392.
  * [Zhang and Bowman, 2018]  Kelly Zhang and Samuel Bowman.  2018\.  Language modeling teaches you more than translation does: Lessons learned through auxiliary syntactic task analysis.  // Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, P 359–361, Brussels, Belgium, November. Association for Computational Linguistics.
  * [Zhu et al., 2015]  Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler.  2015\.  Aligning books and movies: Towards story-like visual explanations by watching movies and reading books.  // Proceedings of the IEEE international conference on computer vision, P 19–27.

## Appendix

Figure 4: The class balance of datasets


# Let Offline RL Flow: Training Conservative Agents in the Latent Space of
Normalizing Flows

Dmitriy Akimov, Alexander Nikulin, Vladislav Kurenkov, Denis Tarasov & Sergey
Kolesnikov
Tinkoff AI
<EMAIL_ADDRESS>Dmitry
Akimov
Tinkoff
<EMAIL_ADDRESS>
&Vladislav Kurenkov
Tinkoff
<EMAIL_ADDRESS>
&Alexander Nikulin
Tinkoff
<EMAIL_ADDRESS>
&Denis Tarasov
Tinkoff
<EMAIL_ADDRESS>
&Sergey Kolesnikov
Tinkoff
<EMAIL_ADDRESS>

###### Abstract

Offline reinforcement learning aims to train a policy on a pre-recorded and
fixed dataset without any additional environment interactions. There are two
major challenges in this setting: (1) extrapolation error caused by
approximating the value of state-action pairs not well-covered by the training
data and (2) distributional shift between behavior and inference policies. One
way to tackle these problems is to induce conservatism - i.e., keeping the
learned policies closer to the behavioral ones. To achieve this, we build upon
recent works on learning policies in latent action spaces and use a special
form of Normalizing Flows for constructing a generative model, which we use as
a conservative action encoder. This Normalizing Flows action encoder is pre-
trained in a supervised manner on the offline dataset, and then an additional
policy model - controller in the latent space - is trained via reinforcement
learning. This approach avoids querying actions outside of the training
dataset and therefore does not require additional regularization for out-of-
dataset actions. We evaluate our method on various locomotion and navigation
tasks, demonstrating that our approach outperforms recently proposed
algorithms with generative action models on a large portion of datasets.111Our
implementation is available at https://github.com/tinkoff-ai/cnf

##  1 Introduction

Offline Reinforcement Learning (ORL) addresses the problem of training new
decision-making policy from a static and pre-recorded dataset collected by
some other policies without any additional data collection (Lange et al.,
2012; Levine et al., 2020). One of the main challenges in this setting is the
extrapolation error (Fujimoto et al., 2019) – i.e. inability to properly
estimate values of state-action pairs not well-supported by the training data,
which in turn leads to overestimation bias. This problem is typically resolved
with various forms of conservatism, for example, Implicit Q-Learning
(Kostrikov et al., 2021) completely avoids estimates of out-of-sample actions,
Conservative Q-Learning (Kumar et al., 2020) penalizes q-values for out-of-
distribution actions and others (Fujimoto and Gu, 2021; Kumar et al., 2019)
put explicit constraints to stay closer to the behavioral policies.

An alternative approach to constraint-trained policies was introduced in PLAS
(Zhou et al., 2020), where authors proposed to construct a latent space that
maps to the actions well-supported by the training data. To achieve this, Zhou
et al. (2020) use Variational Autoencoder (VAE) (Kingma et al., 2019) to learn
a latent action space and then train a controller within it. However, as was
demonstrated in Chen et al. (2022), their specific use of VAE leads to a
necessity for clipping the latent space. Otherwise, the training process
becomes unstable, and the optimized controller can exploit the newly
constructed action space, arriving at the regions resulting in out-of-
distribution actions in the original space. While the described clipping
procedure was found to be effective, this solution is rather ad-hoc and
discards some of the in-dataset actions which could potentially limit the
performance of the trained policies.

In this work, inspired by the recent success of Normalizing Flows (NFs) (Singh
et al., 2020) in the online reinforcement learning setup, we propose a new
method called Conservative Normalizing Flows (CNF) for constructing a latent
action space useful for offline RL. First, we describe why a naive approach
for constructing latent action spaces with NFs is also prone to extrapolation
error, and then outline a straightforward architectural modification that
allows avoiding this without a need for manual post-hoc clipping. Our method
is schematically presented in Figure 1, where we highlight key differences
between our method and the previous approach. We benchmark our method against
other competitors based on generative models and show that it performs
favorably on a large portion of the D4RL (Fu et al., 2020) locomotion and
maze2d datasets.

##  2 Preliminaries

Offline RL The goal of offline RL is to find a policy that maximizes the
expected discounted return given a static and pre-recorded dataset
$\mathcal{D}$ consisting of state-action-reward tuples. Normally, the
underlying decision-making problem is formulated via Markov Decision Process
(MDP) that is defined as a 4-elements tuple, consisting of state space
$\mathbf{S}$, action space $\mathbf{A}$, state transition probability
$p:\mathbf{S}\times\mathbf{S}\times\mathbf{A}\rightarrow[0,\infty]$, which
represents probability density of the next state $s^{\prime}\in\mathbf{S}$
given the current state $s\in\mathbf{S}$ and action $a\in\mathbf{A}$; bounded
reward function
$r:\mathbf{S}\times\mathbf{A}\times\mathbf{S}\rightarrow[r_{min},r_{max}]$ and
a scalar discount factor $\gamma\in(0,1)$. We denote the reward
$r(\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}_{t+1})$ as $r_{t}$. The discounted
return is defined as $R_{t}=\sum_{k=0}^{\infty}\gamma^{k}r_{t+k}$. Also, the
notion of the advantage function $A(\mathbf{s},\mathbf{a})$ is introduced – a
difference between state-action value $Q(\mathbf{s},\mathbf{a})$ and state
value $V(\mathbf{s})$ functions:

$\begin{split}Q^{\pi}(\mathbf{s}_{t},\mathbf{a}_{t})&=r_{t}+\mathbb{E}_{\pi}[\sum_{k=0}^{\infty}\gamma^{k}r_{t+k}]\\\
V^{\pi}(\mathbf{s})&=\mathbb{E}_{\mathbf{a}\sim\pi}[Q^{\pi}(\mathbf{s},\mathbf{a})]\\\
A^{\pi}(\mathbf{s},\mathbf{a})&=Q^{\pi}(\mathbf{s},\mathbf{a})-V^{\pi}(\mathbf{s})\end{split}$
(1)

Advantage Weighted Actor Critic One way to learn a policy in an offline RL
setting is by following the gradient of the expected discounted return
estimated via importance sampling (Levine et al., 2020), however, methods
employing estimation of the Q-function were found to be more empirically
successful (Kumar et al., 2020; Nair et al., 2020; Wang et al., 2020). Here,
we describe Advantage Weighted Actor Critic (Nair et al., 2020) – where the
policy is trained by optimization of log-probabilities of the actions from the
data buffer re-weighted by the exponentiated advantage.

In practice, there are two trained models: policy $\pi_{\theta}$ with
parameters $\theta$ and critic $Q_{\psi}$ with parameters $\psi$. The training
process consists of two alternating phases: policy evaluation and policy
improvement. During the policy evaluation phase, the critic $Q^{\pi}(s,a)$
estimates the action-value function for the current policy, and during the
policy improvement phase, the actor $\pi$ is updated based on the current
estimation of advantage. Combining all together, two following losses are
minimized using the gradient descent:

$\begin{split}L_{\pi}(\theta)&=\mathbb{E}_{(\mathbf{s},\mathbf{a})\sim\mathcal{D}}[-\log\pi_{\theta}(\mathbf{a}|\mathbf{s})\cdot\exp(A_{\psi}(\mathbf{s},\mathbf{a})/\lambda)]\\\
L_{TD}(\psi)&=\mathbb{E}_{(\mathbf{s},\mathbf{a},r,\mathbf{s^{\prime}})\sim\mathcal{D}}[(r+\gamma
Q_{\psi}(\mathbf{s^{\prime}},\mathbf{a^{\prime}}\sim\pi_{\theta}(\cdot|\mathbf{s^{\prime}}))-Q_{\psi}(\mathbf{s},\mathbf{a}))^{2}]\end{split}$
(2)

Where $A_{\phi}(\mathbf{s},\mathbf{a})$ is computed according to Equation 1
using critic $Q_{\phi}$ and $\lambda$ is a temperature hyperparameter.

Normalizing Flows Given a dataset $\mathcal{D}=\\{x^{(i)}\\}_{i=1}^{N}$, with
points $x^{(i)}$ from unknown distribution with density $p_{\mathbf{X}}$ the
goal of a Normalizing Flow model (Dinh et al., 2016; Kingma and Dhariwal,
2018) is to train an invertible mapping
$\mathbf{z}=\mathbf{f}_{\phi}(\mathbf{x})$ with parameters $\phi$ to a simpler
base distribution with density $p_{\mathbf{Z}}$, typically spherical Gaussian:
$\mathbf{z}\sim N(0,I)$. This mapping is required to be invertible by design
to sample new points from data distribution by applying the inverse mapping to
samples from the base distribution:
$\mathbf{x}=\mathbf{f}_{\phi}^{-1}(\mathbf{z})$. A full flow model is a
composition of $K$ invertible functions $\mathbf{f}_{i}$ and the relationship
between $\mathbf{z}$ and $\mathbf{x}$ can be written as:

$\mathbf{x}\xleftrightarrow[\text{}]{\mathbf{f}_{1}}\mathbf{h}_{1}\xleftrightarrow[\text{}]{\mathbf{f}_{2}}\mathbf{h}_{2}\cdots\xleftrightarrow[\text{}]{\mathbf{f}_{K}}\mathbf{z}$
(3)

Log-likelihood of a data point $x$ is obtained by using the change of variable
formula and can be written as:

$\begin{split}\log p_{\phi}(\mathbf{x})&=\log
p_{\mathbf{Z}}(\mathbf{z})+\log|\det(d\mathbf{z}/d\mathbf{x})|\\\ &=\log
p_{\mathbf{Z}}(\mathbf{z})+\sum_{i=1}^{K}\log|\det(d\mathbf{h}_{i}/d\mathbf{h}_{i-1})|\end{split}$
(4)

Normalizing Flows models are optimized to directly maximize log-likelihood of
data points using Equation 4.

##  3 Method

Figure 1: Schematic visualization and comparison of PLAS and CNF (ours)
approaches. Both methods use an action encoder-decoder model trained in a
supervised manner on an offline dataset and a controller model to select
actions from the latent space of the encoder. PLAS algorithm uses VAE with
normal latent distribution with unbounded support (represented as the blue
circle) and restricts latent policy outputs to only a part of the latent space
(represented as the black borders inside of latent space). Our algorithm uses
Normalizing Flow instead of VAE and bound base distribution itself, allowing
the latent policy to use the whole latent space.

In this section, we describe our approach named Conservative Normalizing Flows
(CNF) in detail (Figure 1). We start by delineating the proposed architecture
and then describe both the pre-training and policy optimization phases.

###  3.1 Conservative Normalizing Flows

Normalizing Flows models are trainable bijective mappings, composed of fully
invertible layers, that transform the original data space to the base
distribution space. The latter is often referred to as the latent space.
Typically, the base distribution is modeled as normal since it provides a
tractable density and is easy to sample from (Kingma and Dhariwal, 2018).
However, similar to the PLAS (Zhou et al., 2020), this distribution would
result in unbounded support. Therefore, when learning a policy in the latent
space to maximize q-values it may exploit the regions that lead to the out-of-
distribution actions in the original space exaggerating the extrapolation
error.

(a) NF-Normal, Density

(b) NF-Normal, Samples

(c) NF-Uniform, Density

(d) NF-Uniform, Samples

Figure 2: A toy example to demonstrate that both NFs either with Normal or
Uniform latent distributions can recover the training data. However, as we
demonstrated above, a potential controller trained in the latent space of
normal-based NFs is still able to sample actions outside of the training
dataset.

To illustrate this problem better, consider a toy 2-d task with modeling moons
dataset which consists of points of two interleaving half circles (Buitinck et
al., 2013) using normalizing flows. We train two normalizing flow models,
first with the normal latent distribution and second with the uniform latent
distribution. We additionally squeeze data points to lie in $(-1,+1)^{2}$ 2-d
region via linear transformation in this experiment, and add the inverse of
$\tanh$ function as the first layer of NFs models, so that $\tanh$ function
applied as the last transformation during sampling from the model. Both
trained NFs model the underlying training distribution well (Figure 2). We
then select the NF model with normal latent space and gradually increase the
amplitude $a$ of the latent samples drawn from $a\cdot N(0,I)$ during sampling
new data points from the model, this process is presented in Figure 3. And, as
expected, higher amplitude values result in more out-of-distribution data
points.

One solution to avoid exploiting regions leading to out-of-distribution
actions is to restrict policy output amplitude by some value. This approach
was proposed in Zhou et al. (2020); Chen et al. (2022), where the output of
the deterministic policy in the latent space was modeled as
$z=z_{max}\cdot\tanh(\pi_{\theta}(\mathbf{s}))$ and $z_{max}$ was set to 2.

However, it was shown that the optimal clipping value $z_{max}$ is different
for every locomotion dataset in the D4RL benchmark (Zhou et al. (2020),
Appendix C). Since $z_{max}$ value is essentially part of the latent space
policy model, it is necessary to evaluate multiple values online after
training to select the best one, which may not be feasible in some tasks.

To tackle this problem we would like to avoid this post-hoc clipping and
construct a latent action space model that would prohibit the exploitation of
out-of-distribution actions by design. To do so, we make use of NFs
versatility, and add invertible $\tanh$ activation after the last layer of the
Normalizing Flows model - it makes NFs outputs lie in a bounded
$n-$dimensional222$n$ is the action space dimensionality. interval
$(-1,+1)^{n}$, which in turn allows us to substitute normal base distribution
with $n-$dimensional uniform $\mathbf{U}(-1,+1)^{n}$. With these changes, the
potential actor model should not be capable of generating actions outside of
the training distribution.

(a) Amplitude = 2

(b) Amplitude = 4

(c) Amplitude = 10

(d) Amplitude = 30

Figure 3: A toy example demonstrates that the normal-based NF can be
manipulated by a controller to produce out-of-distribution training points. We
model this by increasing the amplitude of the latent space samples, as we
found this to happen in the preliminary experiments when the controller tries
to maximize q-values. Note that the model with uniform latent space does not
suffer from this problem because it already uses the whole latent distribution
support during the training and sampling processes.

###  3.2 Latent Action Space Pre-Training

First, we start by pre-training conditional Normalizing Flow model
$\mathbf{f}(\mathbf{a}|\mathbf{s})$ with parameters $\phi$ on the actions and
states from the offline dataset, this is essentially a supervised learning
problem. We use the same Normalizing Flow conditioning scheme as in the Singh
et al. (2020) (for more details, see Figure 8 in the original paper): for each
NF layer, we add additional input for the state vector. Given a conditional
NFs $\mathbf{f}$ with parameters $\phi$, we use it to compute log-likelihood
$\log p_{\phi}(\mathbf{a}|\mathbf{s})$ of actions $\mathbf{a}$ conditioned on
states $\mathbf{s}$ from the offline dataset $\mathcal{D}$, minimizing the
following loss:

$L_{\mathbf{f}}(\phi)=\mathbb{E}_{(\mathbf{s},\mathbf{a})\sim\mathcal{D}}[-\log
p_{\phi}(\mathbf{a}|\mathbf{s})]$ (5)

Algorithm 1 CNF, Pre-training

1:Input:  offline data replay buffer $\mathcal{D}$

2:Initialize flow action encoder parameters $\phi$

3:repeat until convergence

4: Sample a mini-batch $B=\left\\{\left(\mathbf{s},\mathbf{a}\right)\right\\}$
from $\mathcal{D}$

5: Compute flow loss using Equation 5, compute loss gradient and update flow
model:
$\nabla_{\theta_{1}}\frac{1}{|B|}\sum_{\left(\mathbf{s},\mathbf{a}\right)\in
B}[-\log p_{\phi}(\mathbf{a|s})]$

The pre-training algorithm is summarized in Algorithm 1 and the final result
of this optimization process is a mapping $\mathbf{f}:A\times S\rightarrow Z$,
where $Z$ is the latent action space bounded by the $(-1,1)$ interval. As this
mapping is invertible, we can further transform latent vectors into the
original action space for policy optimization.

###  3.3 Policy Optimization

The whole procedure is outlined in Algorithm 2. Here, we describe this phase
in detail as follows. Given a pre-trained latent action space model, we freeze
its parameters and add it as an action encoder during RL training. We modify
policy optimization loss from Advantage Weighted Actor Critic (Nair et al.,
2020) to use it with our flow model. Specifically, we train a stochastic
latent policy model $\pi(z|s)$ with parameters $\theta$, which predicts $\mu$
and $\sigma^{2}$ for $\tanh(N(\mu,\sigma^{2}))$ distribution, and two critic
models $Q_{1}(\mathbf{s},\mathbf{a})$, $Q_{2}(\mathbf{s},\mathbf{a})$ with
parameters $\psi_{1}$ and $\psi_{2}$ to mitigate positive bias in the policy
improvement step that is known to degrade the performance of value-based
methods (Hasselt, 2010; Fujimoto et al., 2018). Note that our policy model
operates in the latent space, not in the original space. By combining latent
policy and a pre-trained normalizing flow model, we obtain the following loss
function for policy optimization:

$L_{\pi}(\theta)=\mathbb{E}_{(\mathbf{s},\mathbf{a})\sim\mathcal{D}}[\omega(\mathbf{s},\mathbf{a})\cdot|\mathbf{a}-\mathbf{f}^{-1}(z\sim\pi_{\theta}(\cdot|\mathbf{s})|\mathbf{s})|],$
(6)

where weights $\omega$ comes from exponentiation of the Advantage function:

$\begin{split}\omega(\mathbf{a},\mathbf{s})&=\exp\left(\left(Q(\mathbf{s},\mathbf{a})-Q(\mathbf{s},\mathbf{f}^{-1}(z\sim\pi(\cdot|\mathbf{s})|\mathbf{s}))\right)/\lambda\right)\\\
&=\exp\left(A(\mathbf{s},\mathbf{a})/\lambda\right)\end{split}$ (7)

and Q-function is set to the minimum between two trained models:
$Q(\mathbf{s},\mathbf{a})=\min_{i=1,2}Q_{i}(\mathbf{s},\mathbf{a})$, as
proposed by Fujimoto et al. (2018). Here, $\lambda\in(0,\infty)$ is a
temperature hyperparameter: for higher values training objective behaves
similarly to behavioral cloning and for lower values it aims to maximize
advantage. Overall, this loss function train a policy that maximizes the
Q-values subject to a distribution constraint (Nair et al., 2020).

Together with the latent policy model, we optimize two critics with the
standard Q-learning loss:

$\begin{split}L_{Q}(\psi_{1},\psi_{2})&=\mathbb{E}_{(\mathbf{s},\mathbf{a},r,\mathbf{s^{\prime}})\sim\mathcal{D}}[(Q_{1}(\mathbf{s},\mathbf{a})-y)^{2}+(Q_{2}(\mathbf{s},\mathbf{a})-y)^{2}]\\\
y&=r+\gamma\mathbb{E}_{\mathbf{a^{\prime}}\sim\mathbf{f}^{-1}\left(\pi(\cdot|\mathbf{s^{\prime}})\right)}[\min_{i=1,2}Q_{i}(\mathbf{s^{\prime}},\mathbf{a^{\prime}})]\end{split}$
(8)

Since NFs are differentiable, we can compute the gradient of the given loss
functions for the policy model weights using the chain rule and
reparametrization trick (Kingma et al., 2019). Note, that one could bypass the
differentiation through the NFs model, and optimize policy and critics in the
latent space directly. However, we show in the Appendix (Section A.4), that
this results in worse policies.

Algorithm 2 Offline RL with CNF training

1:Input:  Pre-trained flow action encoder model
$\mathbf{f(\mathbf{a}|\mathbf{s})}$ with parameters $\phi$, offline data
replay buffer $\mathcal{D}$

2:Initialize actor $\pi$ with parameters $\theta$ and two critics $Q_{1}$ and
$Q_{2}$ with parameters $\psi_{1}$ and $\psi_{2}$

3:repeat for a given number of train-ops

4: Sample a mini-batch
$B=\left\\{\left(\mathbf{s},\mathbf{a},r,\mathbf{s^{\prime}}\right)\right\\}$
from $\mathcal{D}$

5: Sample next-state actions $\mathbf{a^{\prime}}$ using policy
$\pi(\cdot|\mathbf{s^{\prime}})$ and flow $\mathbf{f}(\cdot|s^{\prime})$
models, compute Q-target:
$\mathbf{a^{\prime}}=\mathbf{f}^{-1}(z\sim\pi(\cdot|\mathbf{s^{\prime}}))\text{
for all }\mathbf{s^{\prime}}\in B$
$y=r+\gamma(\min_{i=1,2}Q_{i}(\mathbf{s^{\prime}},\mathbf{a^{\prime}}))$

6: Compute critics loss using Equation 8, compute loss gradient and update
models:
$\nabla_{\psi_{1},\psi_{2}}\frac{1}{|B|}\sum_{\left(\mathbf{s},\mathbf{a},r,\mathbf{a^{\prime}}\right)\in
B}[(Q_{1}(\mathbf{s},\mathbf{a})-y)^{2}+(Q_{2}(\mathbf{s},\mathbf{a})-y)^{2}]$

7: Sample actions $\mathbf{a}$ using policy $\pi(\cdot|\mathbf{s})$ and flow
$\mathbf{f}(\cdot|s)$ models, compute advantage weights 7:
$\hat{\mathbf{a}}=\mathbf{f}^{-1}(z\sim\pi(\cdot|\mathbf{s})|\mathbf{s})\text{
for all }\mathbf{s}\in B$
$A(\mathbf{s},\mathbf{a})=Q(\mathbf{s},\mathbf{a})-Q(\mathbf{s},\hat{\mathbf{a}})$
$\omega=\exp(A(\mathbf{s},\mathbf{a})/\lambda)$

8: Compute policy loss using Equation 6, compute loss gradient and update
policy model:
$\nabla_{\theta_{2}}\frac{1}{|B|}\sum_{\left(\mathbf{s},\mathbf{a}\right)\in
B}\exp(A(\mathbf{s},\mathbf{a})/\lambda)\cdot|\mathbf{a}-\hat{\mathbf{a}}|$

##  4 Experiments

To show how the proposed method works, we benchmark it on various locomotion
and navigation tasks from the popular D4RL benchmark (Fu et al., 2020)
comparing it to the other methods based on generative models - PLAS (Zhou et
al., 2020) and LAPO (Chen et al., 2022). We also include AWAC (Nair et al.,
2020) algorithm in our comparisons because we build our policy optimization
method on top of it, and IQL (Kostrikov et al., 2021) algorithm because of its
competitive performance across non-ensemble methods.

###  4.1 D4RL benchmark

Locomotion We focus on three locomotion environments from the D4RL dataset:
Walker2d-v2, Hopper-v2, and HalfCheetah-v2. For the Normalizing Flows model’s
pre-training phase, we divide the training dataset into two parts by
separating $10\%$ portion of randomly selected data for validation. We run 50
experiments with the random search of hyperparameters from Table 3, and then
select the best model according to the log-likelihood on the validation
dataset for RL training. We train all models, including Normalizing Flows,
latent policies, and critics, using Adam optimizer Kingma and Ba (2014).
During the RL phase, we run 1 million training steps for actor and critic
models on all environments except HalfCheetah-v2, where we use only $200.000$
training steps as it is enough for convergence. We evaluate the agent by
running 10 episodes and averaging the scores over them once per 5000 (100 for
HalfCheetah-v2) training steps. Hyperparameters for RL training are listed in
Table 2.

For comparison, we implemented AWAC and IQL algorithms in our code base and
used the publicly available PLAS implementation. We run AWAC and IQL
algorithms with hyperparameters found in the papers Kostrikov et al. (2021);
Nair et al. (2020) and PLAS with hyperparameters recommended in Zhou et al.
(2020). For LAPO, we use scores reported in the paper. The final scores of
benchmarked algorithms are summarized in Table 1. Missing environments are
labeled as ’-’ in the table.

To highlight the performance of our method, we include training curves in
Figure 4. It can be seen that the proposed algorithm exhibits preferable
performance on all 9 locomotion datasets, especially on the HalfCheetah-v2
environment.

Figure 4: Average normalized performance on D4RL locomotion tasks. The x-axis
denotes the training steps. Each curve is averaged over 3 random seeds. Shaded
area represents one standard deviation.

Maze2d For our next experiment, we choose three maze2d datasets with
increasing layout complexity: umaze-v1, medium-v1, and large-v1. We use the
same flow pre-training scheme as before: we run 20 experiments with random
hyperparameters from Table 3 and select the model with the highest log-
likelihood on the validation dataset. We list the hyperparameters for RL
training used in this experiment in Table 2. For comparison, we use our
implementation of IQL, rely on a publicly available PLAS implementation, and
report scores from the LAPO paper. As can be seen in Table 1, our method
outperforms baselines in 2 out of 3 environments.

Table 1: Normalized performance on D4RL benchmark. The scores are averaged over 10 final evaluations and 3 random seeds. The results are reproduced for all of the algorithms except LAPO, for which we take the values stated in the original paper. CNF outperforms other methods on 8 out of 9 locomotion datasets, and on two out of three maze2d datasets. Environment | AWAC | IQL | PLAS | LAPO | CNF
---|---|---|---|---|---
walker2d-medium-v2 | 78.20 | 78.80 | 71.2 | 80.75 |  83.60 $\pm$ 3.01
walker2d-medium-replay-v2 | 76.76 | 74.49 | 2.74 | - |  81.96 $\pm$ 1.98
walker2d-medium-expert-v2 | 112.97 | 111.96 | 108.13 | - |  112.32 $\pm$ 0.21
hopper-medium-v2 | 62.59 | 63.7 | 52.93 | 51.63 |  69.32 $\pm$ 1.04
hopper-medium-replay-v2 | 73.12 | 87.72 | 3.17 | - |  89.04 $\pm$ 10.39
hopper-medium-expert-v2 | 109.64 | 109.05 | 106.5 | - | 108.6 $\pm$ 5.45
halfcheetah-medium-v2 | 43.15 | 47.4 | 43.78 | 45.97 |  50.55 $\pm$ 0.53
halfcheetah-medium-replay-v2 | 42.00 | 43.2 | 44.8 | - |  45.84 $\pm$ 0.31
halfcheetah-medium-expert-v2 | 87.40 | 78.95 | 86.63 | - |  96.23 $\pm$ 0.20
locomotion-v2 average | 76.20 | 77.25 | 57.76 | - | 81.94
maze2d-umaze-v1 | - | 37.69 | 53.9 | 118.9 | 62.9 $\pm$ 10.36
maze2d-medium-v1 | - | 35.45 | 66.4 | 142.8 |  155.89 $\pm$ 10.49
maze2d-large-v1 | - | 49.64 | 107.2 | 200.6 |  212.8 $\pm$ 2.23

###  4.2 Ablations

To study the importance of each major component in the proposed CNF method, we
conduct additional ablation experiments. We begin by comparing training
policies in uniform and normal latent spaces. Then, we examine additional
clipping for the latent policy to see if it can operate in normal latent
space. And finally, we integrate the VAE model into our approach and compare
performance between Normalizing Flows and VAEs for training policies in latent
spaces.

Normalizing Flows with normal latent distribution, no clipping In the first
experiment, we compare the performance of the proposed method but with normal
latent distribution in Normalizing Flow and latent policy models. We use a
conventional Normalizing Flows encoder without $\tanh$ activation after the
last layer and with the normal latent distribution. We pre-train action
encoders in the same way as before, using only the best hyperparameters from
previous experiments (note that these hyperparameters result in very similar
models in terms of the training and validation losses, as depicted in Figure
8). We also remove $\tanh$ activation after the last layer from the latent
policy model, letting it operate over the whole latent space of the pre-
trained Normalizing Flow action encoder, which we make to predict $\mu$ and
$\sigma^{2}$ for $N(\mu,\sigma^{2})$ distribution. We again select the
HalfCheetah-v2 environment and run 3 experiments with different random seeds
per dataset. We plot the results in Figure 5. One can see that the performance
degraded substantially, indicating that latent policy without clipping could
not be trained to produce a competitive performance.

Figure 5: Comparison of the proposed method with uniform (CNF) and normal (NF)
latent spaces. Policy performance is significantly worse when the latent space
is normal.

Normalizing Flows with normal latent distribution, manual clipping As it was
shown in the previous ablation, latent policies without clipping in normal
space perform poorly. To test how different clipping values affect the agent’s
performance, we did the following experiment, which is similar to the previous
one, except we manually clip latent policy output by some value. To do so, we
return $\tanh$ activation at the end of the policy model and multiply it by
$a$, treated as a hyperparameter. Latent policy output is modeled as $z\sim
a\cdot\tanh(N(\pi_{\theta}(.|s)))$, where $N$ is a normal distribution with
parameters predicted by the policy model. This experiment is similar to Zhou
et al. (2020) Ablation C, but instead of VAEs, we use Normalizing Flows to
extract latent policies using Algorithm 2. Also, we use datasets from the
Walker2d-v2 environment. We examine several values for parameter $a$ and
compare them with the proposed method. The results are averaged over three
random seeds and are shown in Figure 6. One can see that optimal clipping
values are different for each dataset: for the medium dataset, there is no
disparity in performance, but the clipping value of 3 produces slightly better
results and almost matches the performance of the uniform latent space; for
the medium-replay, it equals to 2 and for the medium-expert, it equals to 1.
On the other hand, CNF, parameterized with the uniform latent distribution,
does not add extra clipping value as a hyperparameter and performs better on
each dataset.

Figure 6: Normalized average performance of the proposed method (uniform on
the X-axis) and latent policies with clipped normal latent distribution.
Number on the X-axis is the clipping value. Optimal clipping value is
different for each dataset and final performance is better for the proposed
method.

Action encoder model: Normalizing Flow and VAE In this experiment, we compare
the performance of latent policies obtained by the CNF with different action
encoders, namely, Normalizing Flow and VAE. We adopt the best VAE architecture
and training parameters from Zhou et al. (2020), and then integrate pre-
trained models into the latent policy training as in Algorithm 2. For this
comparison, we rely on the HalfCheetah-v2 environment and plot training curves
for PLAS, CNF, and CNF-VAE. Results are presented in Figure 7. One can see
that the CNF performs best with the Normalizing Flow action encoder. On the
other hand, the use of the VAE encoder shows performance similar to PLAS. On
the halfcheetah-medium-replay-v2 dataset, it starts from a commensurate
performance to CNF, but exhibit marginal improvement during the training
process. This experiment indicates the importance of the whole pipeline with
the use of Normalizing Flows as opposed to the VAEs.

Figure 7: Comparison of CNF, CNF with VAE (CNF-VAE) and PLAS. CNF with
Normalizing Flow and uniform distribution performs above all. For normal
distribution optimal clipping value is different for each dataset.

##  5 Conclusion

In this work, we presented a new deep offline RL method called Conservative
Normalizing Flow (CNF). It constructs a latent action space with the use of
the NFs model and then runs a policy optimization within it. This approach
makes trained policies capable of fully utilizing the latent space as opposed
to the post-hoc manual clipping procedures in PLAS Zhou et al. (2020) and LAPO
Chen et al. (2022). We benchmarked our method against other competitors based
on generative models and showed that it performs favorably on a large portion
of D4RL Fu et al. (2020) locomotion and maze2d datasets.

## References

  * Buitinck et al. (2013) Lars Buitinck, Gilles Louppe, Mathieu Blondel, Fabian Pedregosa, Andreas Mueller, Olivier Grisel, Vlad Niculae, Peter Prettenhofer, Alexandre Gramfort, Jaques Grobler, Robert Layton, Jake VanderPlas, Arnaud Joly, Brian Holt, and Gaël Varoquaux.  API design for machine learning software: experiences from the scikit-learn project.  In _ECML PKDD Workshop: Languages for Data Mining and Machine Learning_ , pages 108–122, 2013.
  * Chen et al. (2022) Xi Chen, Ali Ghadirzadeh, Tianhe Yu, Yuan Gao, Jianhao Wang, Wenzhe Li, Bin Liang, Chelsea Finn, and Chongjie Zhang.  Latent-variable advantage-weighted policy optimization for offline rl.  _arXiv preprint arXiv:2203.08949_ , 2022.
  * Dinh et al. (2016) Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real nvp.  _arXiv preprint arXiv:1605.08803_ , 2016.
  * Fu et al. (2020) Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.  D4rl: Datasets for deep data-driven reinforcement learning.  _arXiv preprint arXiv:2004.07219_ , 2020.
  * Fujimoto and Gu (2021) Scott Fujimoto and Shixiang Shane Gu.  A minimalist approach to offline reinforcement learning.  _Advances in neural information processing systems_ , 34:20132–20145, 2021.
  * Fujimoto et al. (2018) Scott Fujimoto, Herke van Hoof, and David Meger.  Addressing function approximation error in actor-critic methods, 2018\.  URL https://arxiv.org/abs/1802.09477.
  * Fujimoto et al. (2019) Scott Fujimoto, David Meger, and Doina Precup.  Off-policy deep reinforcement learning without exploration.  In _International conference on machine learning_ , pages 2052–2062. PMLR, 2019.
  * Haarnoja et al. (2018) Tuomas Haarnoja, Kristian Hartikainen, Pieter Abbeel, and Sergey Levine.  Latent space policies for hierarchical reinforcement learning.  In _International Conference on Machine Learning_ , pages 1851–1860. PMLR, 2018.
  * Hasselt (2010) Hado Hasselt.  Double q-learning.  _Advances in neural information processing systems_ , 23, 2010.
  * Jaques et al. (2020) Natasha Jaques, Judy Hanwen Shen, Asma Ghandeharioun, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Shane Gu, and Rosalind Picard.  Human-centric dialog training via offline reinforcement learning.  _arXiv preprint arXiv:2010.05848_ , 2020.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kingma et al. (2019) Diederik P Kingma, Max Welling, et al.  An introduction to variational autoencoders.  _Foundations and Trends® in Machine Learning_ , 12(4):307–392, 2019.
  * Kingma and Dhariwal (2018) Durk P Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  _Advances in neural information processing systems_ , 31, 2018.
  * Kostrikov et al. (2021) Ilya Kostrikov, Ashvin Nair, and Sergey Levine.  Offline reinforcement learning with implicit q-learning.  _arXiv preprint arXiv:2110.06169_ , 2021.
  * Kumar et al. (2019) Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine.  Stabilizing off-policy q-learning via bootstrapping error reduction.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Kumar et al. (2020) Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.  Conservative q-learning for offline reinforcement learning.  _Advances in Neural Information Processing Systems_ , 33:1179–1191, 2020.
  * Lange et al. (2012) Sascha Lange, Thomas Gabel, and Martin Riedmiller.  Batch reinforcement learning.  In _Reinforcement learning_ , pages 45–73. Springer, 2012.
  * Levine et al. (2020) Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  _arXiv preprint arXiv:2005.01643_ , 2020.
  * Nair et al. (2020) Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine.  Awac: Accelerating online reinforcement learning with offline datasets.  _arXiv preprint arXiv:2006.09359_ , 2020.
  * Singh et al. (2020) Avi Singh, Huihan Liu, Gaoyue Zhou, Albert Yu, Nicholas Rhinehart, and Sergey Levine.  Parrot: Data-driven behavioral priors for reinforcement learning.  _arXiv preprint arXiv:2011.10024_ , 2020.
  * Wang et al. (2020) Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al.  Critic regularized regression.  _Advances in Neural Information Processing Systems_ , 33:7768–7778, 2020.
  * Zhou et al. (2020) Wenxuan Zhou, Sujay Bajracharya, and David Held.  Plas: Latent action space for offline reinforcement learning.  _arXiv preprint arXiv:2011.07213_ , 2020.

##  Appendix A Appendix

###  A.1 Related Work

Normalizing Flows in RL Several prior works applied Normalizing Flows models
for reinforcement learning tasks. In the work Haarnoja et al. [2018]
Normalizing Flows were used for training hierarchical policies using the Soft
Actor-Critic RL algorithm. They choose Normalizing Flows because it provides
an expression for exact likelihood computation and has an intuitive way to
stack a sequence of models to construct one hierarchical model. In contrast to
our work, they did not modify NF’s latent space, they did not use offline pre-
training of flow models and train models in the online RL framework.

The PARROT work [Singh et al., 2020] proposes to use the Normalizing Flows
action encoder during the behavioral cloning pre-training phase before running
reinforcement learning on the target task. After pre-training, they freeze
Normalizing Flow and run reinforcement learning on a new and unseen task. They
aimed to develop a method leveraging near-optimal demonstrations during
Normalizing Flows pre-training and speed up the convergence of RL by better
and meaningful exploration, while we use Normalizing Flows in purely offline
RL setting and to make agent conservative.

Offline RL A large portion of recently proposed deep offline RL algorithms
focuses on addressing the extrapolation issue, trying to impose a certain
degree of conservatism limiting the deviation of the final policy from the
behavioral one. Researchers approached this problem from multiple angles. For
example, Kumar et al. [2020] proposed to directly penalize out-of-distribution
actions, while Kostrikov et al. [2021] avoids estimating values for out-of-
sample actions completely. Others Fujimoto and Gu [2021], Kumar et al. [2019],
Jaques et al. [2020] put explicit constraints to stay closer to the behavioral
policy. Here, we took a different approach by constructing a latent action
space that allows us to bypass the need for explicit regularizations.

Offline RL with generative models In the method named PLAS [Zhou et al.,
2020], the authors proposed to pre-train conditional variational autoencoder
on actions from an offline dataset. This idea resembles ours, but to make
agent conservative authors restrict policy outputs in the latent space to a
fixed range. In our work, we aimed to make a better action encoder model by
switching from VAEs to Normalizing Flows. This allows us to utilize the whole
latent space of the action encoder, avoid manual clipping, and we
experimentally demonstrate that our approach leads to better performance on
the popular offline RL benchmark.

One recent approach called LAPO [Chen et al., 2022] proposes to train an
action encoder together with a reinforcement learning agent. They motivated
this by observing that the action distribution does not match the return
distribution in the training data set, and therefore actions that lead to
higher returns are more important for action encoder training. In our work, we
examine an orthogonal direction, studying a different generative model for
action encoding.

###  A.2 Hyperparameters

Table 2: Hyperparameters for latent policy and critics models for RL training phase on locomotion and maze2 tasks. We ran a grid-search for the values written in square brackets. Rest of the parameter were fixed for all of the datasets. Hyperparameter | value
---|---
number of training steps | 1000000
number of training steps (HalfCheetah-v2) | 200000
number of layers (locomotion) | [3, 4]
number of layers (maze2d) | 4
hidden size | 256
learning rate | 3e-4
batch size (locomotion) | [256, 512, 1024]
batch size (maze2d) | 10240
$\lambda$-temperature (locomotion) | 1/3
$\lambda$-temperature (maze2d) | 1/10

Table 3: NFs hyperparameters for the supervised pre-training phase on locomotion and maze2 tasks. Medium, medium-replay, and medium-expert datasets are marked as m, m-r, and m-e correspondingly. We sample learning rate and weight decay from the continuous uniform distribution. Hyperparameter | value
---|---
number of training steps | 100000
number of layers (locomotion m and m-e datasets) | 12
number of layers (locomotion m-r and maze2d datasets) | 4
hidden size (locomotion m and m-e datasets) | 256
hidden size (locomotion m-r and maze2d datasets) | 64
learning rate |  $\min$ = 1e-5, $\max$ = 3e-3
weight decay |  $\min$ = 0.0, $\max$ = 1e-2
batch size | [512, 1024, 2048]

###  A.3 Normalizing Flows training curves

Figure 8: Training and validation loss for Normalizing Flows pre-training with
Uniform and Normal latent spaces. Both training and validation curves are
almost identical between models with different latent spaces, which means that
models have similar encoding and reconstruction quality. For each latent
policy training, we select NF model with the lowest validation metric.

###  A.4 Training controller in latent space

Ideologically, policy optimization can be carried out directly in the latent
space. This can be done by simply substituting the original actions with their
latent counterparts. After this substitution, a myriad of offline RL
algorithms can be used to extract a new policy. To test if this is a viable
approach, we train actor and critic models in the latent space of the
Normalizing Flows without utilizing gradients from the action encoder during
the policy optimization phase. First, we convert actions from the original
environment’s action space to the latent space by encoding them with the pre-
trained action encoder. After that, we run the AWAC algorithm with no
additional changes to the dataset with converted actions. We conduct this
experiment on HalfCheetah-v2 datasets: medium, medium-replay, and medium-
expert. Results are presented in Figure 9. It can be seen that this approach
shows promising results, but convergence speed and the final score are
slightly lower on all datasets. We conjecture that the performance is lower
because a metric in the latent space is induced from the original action space
by the flow model, and optimization of the distance between latent vectors
corresponds to optimization of the distance between actions only implicitly.
When training policy using metrics from the original action space, as
suggested by Equation 6, the Normalizing Flows model provides gradients to the
policy model that guide it to minimize the actual distance between actions.

Figure 9: Comparison of CNF trained in original (blue) and latent (red) action
spaces. Final performance is slightly lower for the training directly in
latent space.