summaries: Many Information Retrieval systems rely on two components: a retriever that identifies a small set of "support" documents from a large corpus, followed by a reader that re-scores these support documents more finely. For retrievers, metrics like BM25 were once common but they are increasingly replaced by machine learned components. However, most datasets do not provide direct supervision information for the retriever.

summaries: The paper targets an important problem in open-domain QA - the training of the retriever for the purpose of determining a segment that may contain the answer. In the most traditional setting, the retriever is just a traditional IR system such as BM25. In some existing work, the retriever has been trained to locate the documents containing the answer (e.g. inverse cloze task, or DPR). This paper goes in the same direction. The difference is that it uses the attention weights as relevance signals to train the retriever, instead of the inclusion of the answer in the passage.

summaries: The authors propose a training technique for information retrieval models in the context of (open domain) question answering. Assuming the existence of some reader model, the idea is to use internal information of that model as a training signal for a retriever. Specifically, they use the attention activations over the input documents as synthetic labels for the retriever.

abstractText: The task of information retrieval is an important component of many natural language processing systems, such as open domain question answering. While traditional methods were based on hand-crafted features, continuous representations based on neural networks recently obtained competitive results. A challenge of using such methods is to obtain supervised data to train the retriever model, corresponding to pairs of query and support documents. In this paper, we propose a technique to learn retriever models for downstream tasks, inspired by knowledge distillation, and which does not require annotated pairs of query and documents. Our approach leverages attention scores of a reader model, used to solve the task based on retrieved documents, to obtain synthetic labels for the retriever. We evaluate our method on question answering, obtaining state-of-the-art results.

authors:
Wei-Cheng Chang
Felix X Yu
Yin-Wen Chang
Yiming Yang
Sanjiv Kumar

title: Pre-training tasks for embedding-based large-scale retrieval

venue: arXiv preprint arXiv:2002.03932,

authors:
Danqi Chen
Adam Fisch
Jason Weston
Antoine Bordes

title: Reading Wikipedia to answer opendomain questions

venue: In Proc. ACL,

authors:
Scott Deerwester
Susan T Dumais
George W Furnas
Thomas K Landauer
Richard Harshman

title: Indexing by latent semantic analysis

venue: Journal of the American society for information science,

authors:
Mostafa Dehghani
Hamed Zamani
Aliaksei Severyn
Jaap Kamps
W. Bruce Croft

title: Neural ranking models with weak supervision

venue: In Proceedings of the 40th International ACM SIGIR Conference on Research and Development in Information Retrieval,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: BERT: Pre-training of deep bidirectional transformers for language understanding

venue: In Proc. NAACL,

authors:
Daniel Gillick
Alessandro Presta
Gaurav Singh Tomar

title: End-to-end retrieval in continuous space

venue: arXiv preprint arXiv:1811.08008,

authors:
Kelvin Guu
Kenton Lee
Zora Tung
Panupong Pasupat
Ming-Wei Chang

title: Realm: Retrievalaugmented language model pre-training

venue: arXiv preprint arXiv:2002.08909,

authors:
Geoffrey Hinton
Oriol Vinyals
Jeff Dean

title: Distilling the knowledge in a neural network

authors:
Po-Sen Huang
Xiaodong He
Jianfeng Gao
Li Deng
Alex Acero
Larry Heck

title: Learning deep structured semantic models for web search using clickthrough data

venue: In Proceedings of the 22nd ACM international conference on Information & Knowledge Management,

authors:
Samuel Humeau
Kurt Shuster
Marie-Anne Lachaux
Jason Weston

title: Poly-encoders: Transformer architectures and pre-training strategies for fast and accurate multi-sentence scoring

authors:
Gautier Izacard
Edouard Grave

title: Leveraging passage retrieval with generative models for open domain question answering

venue: arXiv preprint arXiv:2007.01282,

authors:
Jeff Johnson
Matthijs Douze
Hervé Jégou

title: Billion-scale similarity search with gpus

venue: IEEE Transactions on Big Data,

title: A statistical interpretation of term specificity and its application in retrieval

venue: Journal of documentation,

authors:
Mandar Joshi
Eunsol Choi
Daniel S. Weld
Luke Zettlemoyer

title: Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension

authors:
Vladimir Karpukhin
Barlas Oğuz
Sewon Min
Ledell Wu
Sergey Edunov
Danqi Chen
Wen-tau Yih

title: Dense passage retrieval for open-domain question answering

venue: arXiv preprint arXiv:2004.04906,

authors:
Omar Khattab
Christopher Potts
Matei Zaharia

title: Relevance-guided supervision for openqa with colbert

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
Tomáš Kočiskỳ
Jonathan Schwarz
Phil Blunsom
Chris Dyer
Karl Moritz Hermann
Gábor Melis
Edward Grefenstette

title: The NarrativeQA reading comprehension challenge

authors:
Tom Kwiatkowski
Jennimaria Palomaki
Olivia Redfield
Michael Collins
Ankur Parikh
Chris Alberti
Danielle Epstein
Illia Polosukhin
Matthew Kelcey
Jacob Devlin
Kenton Lee
Kristina N. Toutanova
Llion Jones
Ming-Wei Chang
Andrew Dai
Jakob Uszkoreit
Quoc Le
Slav Petrov

title: Natural Questions: a benchmark for question answering research

venue: TACL,

authors:
Kenton Lee
Ming-Wei Chang
Kristina Toutanova

title: Latent retrieval for weakly supervised open domain question answering

authors:
Mike Lewis
Yinhan Liu
Naman Goyal
Marjan Ghazvininejad
Abdelrahman Mohamed
Omer Levy
Ves Stoyanov
Luke Zettlemoyer

title: BART: Denoising sequence-to-sequence pretraining for natural language generation, translation, and comprehension

authors:
Mike Lewis
Marjan Ghazvininejad
Gargi Ghosh
Armen Aghajanyan
Sida Wang
Luke Zettlemoyer

title: Pre-training via paraphrasing

venue: arXiv preprint arXiv:2006.15020,

authors:
Patrick Lewis
Ethan Perez
Aleksandara Piktus
Fabio Petroni
Vladimir Karpukhin
Naman Goyal
Heinrich Küttler
Mike Lewis
Wen-tau Yih
Tim Rocktäschel

title: Retrieval-augmented generation for knowledge-intensive nlp tasks

venue: arXiv preprint arXiv:2005.11401,

authors:
Ilya Loshchilov
Frank Hutter

title: Decoupled weight decay regularization

venue: In International Conference on Learning Representations,

authors:
Yi Luan
Jacob Eisenstein
Kristina Toutanova
Michael Collins

title: Sparse, dense, and attentional representations for text retrieval

venue: arXiv preprint arXiv:2005.00181,

authors:
Christopher D Manning
Hinrich Schütze
Prabhakar Raghavan

title: Introduction to information retrieval

venue: Cambridge university press,

authors:
Rodrigo Nogueira
Kyunghyun Cho

title: Passage re-ranking with bert

venue: arXiv preprint arXiv:1901.04085,

authors:
Hamid Palangi
Li Deng
Yelong Shen
Jianfeng Gao
Xiaodong He
Jianshu Chen
Xinying Song
Rabab Ward

title: Deep sentence embedding using long short-term memory networks: Analysis and application to information retrieval

venue: IEEE/ACM Transactions on Audio, Speech, and Language Processing,

authors:
Colin Raffel
Noam Shazeer
Adam Roberts
Katherine Lee
Sharan Narang
Michael Matena
Yanqi Zhou
Wei Li
Peter J Liu

title: Exploring the limits of transfer learning with a unified text-to-text transformer

venue: arXiv preprint arXiv:1910.10683,

authors:
Pranav Rajpurkar
Jian Zhang
Konstantin Lopyrev
Percy Liang

title: SQuAD: 100,000+ questions for machine comprehension of text

venue: In Proc. EMNLP,

authors:
Stephen E Robertson
Steve Walker
Susan Jones
Micheline M Hancock-Beaulieu
Mike Gatford

title: Okapi at TREC-3

venue: NIST Special Publication Sp,

authors:
Yelong Shen
Xiaodong He
Jianfeng Gao
Li Deng
Grégoire Mesnil

title: Learning semantic representations using convolutional neural networks for web search

venue: In Proceedings of the 23rd international conference on world wide web,

authors:
James Thorne
Andreas Vlachos
Christos Christodoulopoulos
Arpit Mittal

title: Fever: a large-scale dataset for fact extraction and verification

venue: arXiv preprint arXiv:1803.05355,

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Ł ukasz Kaiser
Illia Polosukhin

title: Attention is all you need

venue: In Advances in Neural Information Processing Systems

title: The TREC-8 question answering track report

venue: In TREC,

authors:
Sohee Yang
Minjoon Seo

title: Is retriever merely an approximator of reader

venue: arXiv preprint arXiv:2010.10999,

authors:
Wei Yang
Yuqing Xie
Aileen Lin
Xingyu Li
Luchen Tan
Kun Xiong
Ming Li
Jimmy Lin

title: End-to-end open-domain question answering with BERTserini

venue: In Proc. NAACL (Demonstrations),

heading: 1 INTRODUCTION

text:
Information retrieval is an important component for many natural language processing tasks, such as question answering (Voorhees et al., 1999) or fact checking (Thorne et al., 2018). For example, many real world question answering systems start by retrieving a set of support documents from a large source of knowledge such as Wikipedia. Then, a finer-grained model processes these documents to extract the answer. Traditionally, information retrieval systems were based on hand-crafted sparse representations of text documents, such as TF-IDF or BM25 (Jones, 1972; Robertson et al., 1995). Recently, methods based on dense vectors and machine learning have shown promising results (Karpukhin et al., 2020; Khattab et al., 2020). Deep neural networks based on pre-training, such as BERT (Devlin et al., 2019), have been used to encode documents into fixed-size representations. These representations are then queried using approximate nearest neighbors (Johnson et al., 2019). These techniques have lead to improved performance on various question answering tasks.
A challenge of applying machine learning to information retrieval is to obtain training data for the retriever. To train such models, one needs pairs of queries and the corresponding list of documents that contains the information corresponding to the queries. Unfortunately, hand-labeling data to that end is time consuming, and many datasets and applications lack such annotations. An alternative approach is to resort to heuristics, or weakly supervised learning, for example by considering that all documents containing the answer are positive examples. However, these approaches suffer from the following limitations. First, frequent answers or entities might lead to false positive examples. As an example, consider the question “where was Ada Lovelace born?”. The sentence “Ada Lovelace died in 1852 in London” would be considered as a positive example, because it contains the answer “London”. A second limitation is that for some tasks, such as fact checking or long form question answering, such heuristics might not be applicable directly.
In this paper, we propose a procedure to learn retriever systems without strong supervision in the form of pairs of queries and documents. Following previous work (Chen et al., 2017), our approach uses two models: the first one retrieves documents from a large source of knowledge (the retriever), the second one processes the support documents to solve the task (the reader). Our method is inspired by knowledge distillation (Hinton et al., 2015), and uses the reader model to obtain synthetic labels to train the retriever model. More precisely, we use a sequence-to-sequence model as the reader, and use the attention activations over the input documents as synthetic labels to train the retriever. Said otherwise, we assume that attention activations are a good proxy for the relevance of
documents. We then train the retriever to reproduce the ranking of documents corresponding to that metric.
We make the following contributions:
• First, we show that attention scores from a sequence-to-sequence reader model are a good measure of document relevance (Sec. 3.2) ;
• Second, inspired by knowledge distillation, we propose to iteratively train the retriever from these activations, and compare different loss functions (Sec. 3.4) ;
• Finally, we evaluate our method on three question-answering benchmarks, obtaining stateof-the-art results (Sec. 4).
Our code is available at: github.com/facebookresearch/FiD.

heading: 2 RELATED WORK

text:
We briefly review information retrieval based on machine learning. We refer the reader to Manning et al. (2008) and Mitra et al. (2018) for a more exhaustive introduction to the subject.
Vector space models. In traditional information retrieval systems, documents and queries are represented as sparse vectors, each dimension corresponding to a different term. Different schemes have been considered to weigh the different term, the most well known being based on inverse document frequency, or term specificity (Jones, 1972). This technique was later extended, leading to the BM25 weighting scheme which is still widely used today (Robertson et al., 1995). A limitation of sparse representations is that the terms of the query need to match the terms of the returned documents. To overcome this, Deerwester et al. (1990) proposed to use latent semantic analysis for indexing, leading to low-dimension dense representations of documents.
Neural information retrieval. Following the success of deep learning for other natural processing tasks, neural networks were applied to the task of information retrieval. Huang et al. (2013) proposed a deep bag-of-words model, where queries and documents were embedded independently, a technique known as bi-encoder. Documents were then ranked by using the cosine similarity with the query, and the model was trained on clickthrough data from a search engine. This technique was later extended by using convolutional neural networks (Shen et al., 2014) and recurrent neural networks (Palangi et al., 2016). A limitation of independently embedding documents and query is that it does not capture fine-grained interactions between the query and documents. This lead Nogueira & Cho (2019) and Yang et al. (2019) to use a BERT model to jointly embed documents and query, a technique known as cross-encoder.
End-to-end retrieval. Most of the methods described in the previous paragraph were used to rerank a small number of documents, usually returned by a traditional IR systems. In the context of ad-hoc document retrieval, Gillick et al. (2018) showed that bi-encoder models could be competitive with traditional IR systems. For open domain question-answering, Karpukhin et al. (2020) introduced dense passage retrieval (DPR), which uses dense embeddings and nearest neighbors search. More precisely, question and passage embeddings are obtained using a BERT-based biencoder model, which is trained on a small dataset of question and passage pairs. Then, the full knowledge source (Wikipedia) is encoded using this model, and passages are queried by computing the k-nearest neighbors of the embedding of the question. Jointly embedding the query and documents makes the application of cross-encoder models intractable to large database. To address this limitation, Humeau et al. (2019) introduced the poly-encoder architecture, in which each documents is represented by multiple vectors instead of one. Similarly, Khattab et al. (2020) proposed a scoring function where each term of the query and documents is represented by a single vector. To make the method tractable, their system retrieves documents with an approximate score, which are then re-ranked with the exact one. Finally, Luan et al. (2020) conducts a theoretical and empirical study of sparse, dense and cross-attention information retrieval systems.
Unsupervised learning. Closest to our work, there is growing body of work trying to learn information retrieval systems from unsupervised data. Lee et al. (2019) introduced the inverse cloze task
for pre-training retrievers, which can then be fine-tuned end-to-end on question-answering tasks. This pre-training scheme was later evaluated for ad-hoc document retrieval by Chang et al. (2020). Guu et al. (2020) proposed to augment language model pre-training with a retriever module, which is trained using the masked language modeling objective. Similarly, Lewis et al. (2020a) introduced a sequence-to-sequence model that is pre-trained by generating a target text, after retrieving a set of related texts. Lewis et al. (2020b) further train the retriever obtained in Karpukhin et al. (2020) by backpropagating to the retriever the error between the generated output and the gold answer.
Simultaneously to our work, Yang & Seo (2020) proposes to train a retriever with knowledge distillation. The main difference with our method is the nature of the synthetic labels that are used to train the retriever. Yang & Seo (2020) uses the DPR reader, which includes a classifier that predicts which passage contains the answer, and can be seen as a cross-encoder reranker. This technique thus performs the distillation of a cross-encoder retriever to a bi-encoder retriever. In contrast, our method uses the internal attention scores of the reader, which does not require additional supervision besides pairs of question and answer.

heading: 3 METHODOLOGY

text:
Our system is composed of two modules, the retriever and the reader, following the standard pipeline for open-domain question answering. Given an input question these modules are used in a two-step process to generate an answer. First the retriever selects support passages in a large knowledge source. Then these passages are processed by the reader, along with the question, to generate an answer. For the reader module we use the Fusion-in-Decoder model (Izacard & Grave, 2020), which achieves state-of-the-art performance when combined with BM25 or DPR (Karpukhin et al., 2020). It is based on a sequence-to-sequence architecture, and is initialized from pre-trained models such as T5 or BART (Raffel et al., 2019; Lewis et al., 2019).
The focus of this work is to train the retriever without strong supervision or weakly supervised learning based on heuristics. For this we propose to train the retriever by learning to approximate the attention score of the reader. The training scheme outlined here can be seen as a student-teacher pipeline, where the teacher, the reader module, produces targets which are used to train a student network, the reader. By doing so, we hope to leverage the signal extracted from the question-answer pairs by the reader. Since the goal of the retriever is to retrieve the most relevant passages, by training the retriever to estimate the reader attention scores, we implicitly make the assumption that these scores are a good proxy for the usefulness of a passage to answer the question.
In this section we will first describe the Fusion-in-Decoder architecture, before elaborating on the signal which is used to train the retriever, the design of the retriever, and how this module is trained.

heading: 3.1 CROSS-ATTENTION MECHANISM

text:
First, let us briefly review the Fusion-in-Decoder model (FiD, Izacard & Grave, 2020). The underlying architecture is a sequence-to-sequence model, composed of an encoder and a decoder. The encoder independently processes np different text inputs (sk)1≤k≤np . In the case of open-domain question answering based on Wikipedia, each input sk is the concatenation of the question q and a support passage, with special tokens question:, title: and context: added before the question, the title of the Wikipedia article and the text of each passage. The output representations of the encoder are then concatenated to form a global representation X of dimension ( ∑ k `k) × d, where `k is the length of the k-th segment and d is the dimension of the embeddings and hidden representations of the model. Then, the decoder processes this representation as a regular autoregressive model, alternating self-attention, cross-attention and feed-forward modules.
Only the cross-attention module explicitly takes as input the global output representation X of the encoder. If H ∈ Rd denotes the output of the previous self-attention layer of the decoder, the crossattention operation consists in the following operations. First, queries Q, keys K and values V are computed by applying linear transformations:
Q = WQH, K = WKX, V = WVX.
Then a similarity score between the query at position i, Qi, and the key at position j, Kj , is obtained by computing the dot-product between these two elements, and normalized over the dimension:
αi,j = Q T i Kj , α̃i,j = exp(αi,j)∑ m exp(αi,m) .
A new representation is obtained as a sum of the values, weighted by the attention probabilities, before going through a final linear transformation Wo:
Oi = WO ∑ j α̃i,jVi,j
The operations described above are performed in parallel with different linear transformations in the case of multi-head attention. Finally a normalization layer is applied, and this pipeline is wrapped by a skip connection. See Vaswani et al. (2017) for more details on the structure of Transformers.

heading: 3.2 CROSS-ATTENTION SCORE AS A RELEVANCE MEASURE FOR PASSAGE RETRIEVAL

text:
In some sense, the attention scores α:,j involving the j-th key measures the importance of this key, and corresponding value, to compute the next representation. We hypothesize that it is good proxy to estimate the relevance of a passage — the more the tokens in a text segment are attended to, the more relevant the text segment is to answer the question.
Given the reader model, an input question q and a corresponding set of support passages Dq = (pk)1≤k≤n, we obtain relevance scores (Gq,pk)1≤k≤n for each passage by aggregating attention scores. In particular, the score Gq,pk is obtained by averaging the pre-attention scores α0,: over all the tokens in the input sk corresponding to the passage pk, all the layers and all the heads of the decoder. Note that the FiD decoder jointly processes the passages, and thus the score Gq,pk depends on the other support passages. We consider other pooling operators, such as max, to aggregate attention scores over layers, heads and tokens and empirically compare them in Sec. 5.2.
Before we proceed, let us consider the following simple experiment, which is a first indication that reader attention scores are indeed a strong relevance signal. Given a question and 100 passages retrieved with DPR, our goal is to select the 10 best passages. When using the top 10 passages from DPR instead of the top 100, the performance of our reader drops from 48.2 EM to 42.9 EM. On the other hand, if we select the top 10 documents according to the attention scores, the performance only drops to 46.8 EM.

heading: 3.3 DENSE BI-ENCODER FOR PASSAGE RETRIEVAL

text:
Ideally, we would like to rank passages according to the reader cross-attention scores. In practice however, since the passages and the question need to be processed simultaneously by the reader module it is impractical to query a large knowledge source this way. Thus, we use a retriever model composed of an embedder functionE that maps any text passage to a d-dimensional vector, such that the similarity score between a question q and a passage p is defined as Sθ(q, p) = E(q)TE(p). This similarity metric enables us to index all passages in the knowledge source as a preprocessing step. Then at runtime, passages with the highest similarity score with the input question are retrieved, by using an efficient similarity search library such as FAISS (Johnson et al., 2019).
For the embedder we use BERT and follow DPR by considering that the encodings E(q) and E(p) are obtained by extracting the representation of the initial [CLS] token. This leads to a representation of dimension d = 768 in the case of a base model. Differently from DPR, we use the same encoding function E for the questions and passages by sharing parameters.

heading: 3.4 DISTILLING THE CROSS-ATTENTION SCORE TO A BI-ENCODER

text:
In this section, we describe how to train the retriever model, based on the relevance scores obtained in Sec. 3.2. For the training objective of the retriever, we propose to minimize the KL-divergence between the output Sθ(q, p) and the score Gq,p after normalization:
LKL(θ,Q) = ∑
q∈Q,p∈Dq
G̃q,p(log G̃q,p − log S̃θ(q, p)),
where
G̃q,p = exp(Gq,p)∑
p′∈Dq exp(Gq,p′) , S̃θ(q, p) = exp(Sθ(q, p))∑ p′∈Dq exp(Sθ(q, p ′)) .
In Sec. 5.1 we present results obtained when using alternatives to this training objective. We consider two other objectives which have been used in Dehghani et al. (2017), where BM25 is used as a teacher model to train a neural ranker. A first option consists in training the retriever with a regression approach by minimizing the mean squared error:
LMSE(θ,Q) = ∑
q∈Q,p∈Dq
(Sθ(q, p)−Gq,p)2.
The second option we consider is to use a max-margin loss that explicitly penalizes inversions in the ranking estimated by the retriever:
Lranking(θ,Q) = ∑
q∈Q,p1,p2∈Dq
max (0, γ − sign(Gq,p1 −Gq,p2)(Sθ(q, p1)− Sθ(q, p2))) .
In words, if p1 is more relevant to answer the question q than p2 , i.e. Gq,p1 > Gq,p2 , the loss pushes the retriever score of p1 to be larger than the score of p2 by at least a margin of γ.

heading: 3.5 ITERATIVE TRAINING

text:
In this section, we explain how iterative training can be used with the student-teacher scheme described in the previous section, similarly to Khattab et al. (2020). This iterative procedure can be interpreted as using the current retriever to sample negative examples, in order to train a new retriever. When learning a retriever with discriminative training, negative samples play an important role, and various strategies have been considered in previous work. Karpukhin et al. (2020) compared random sampling with using the top-k passages from BM25 which do not contain the answer and with using the positive passages from other queries. Consider that for each question, we have an initial set of support documents D0q . We propose to use an iterative pipeline where each iteration can be described as the following 4-step process:
1. Train the reader R using the set of support documents for each question D0q . 2. Compute aggregated attention scores (Gq,p)q∈Q,p∈D0q with the reader R.
3. Train the retriever E using the scores (Gq,p)q∈Q,p∈D0q .
4. Retrieve top-passages with the new trained retriever E.
This multi-step procedure can be repeated multiple times. A critical point of the training procedure is the initial set of documents corresponding to each question. In Sec. 4, we compare retrievers obtained by starting from documents obtained using BM25 or cosine similarity from a BERT model. In particular, we show that while the initial performance with BERT is low, the iterative procedure allows to greatly improve the performance of the model.

heading: 4 EXPERIMENTS

text: In this section we evaluate the student-teacher training procedure from the previous section. We show that we obtain competitive performance without strong supervision for support documents.

heading: 4.1 EXPERIMENTAL SETTING

text:
Datasets. We perform experiments on TriviaQA (Joshi et al., 2017) and NaturalQuestions (Kwiatkowski et al., 2019), two standard benchmarks for open-domain question answering. TriviaQA is made of questions from trivia and quiz league websites, and does not contain gold support documents. NaturalQuestions contains questions corresponding to web search queries, and gold support documents from Wikipedia. Following the setting from Lee et al. (2019); Karpukhin
et al. (2020), we use the original evaluation set as test set, and keep 10% of the training data for validation. We use the Wikipedia dump from Dec. 20, 2018 for support documents, splitting articles into non-overlapping passages of 100 tokens, and applying the same preprocessing as Chen et al. (2017).
We also evaluate on NarrativeQuestions (Kočiskỳ et al., 2018), using a publicly available preprocessed version.1 This is a reading comprehension dataset built on a corpus of books and movie scripts. For each story, questions are generated by human annotators based on a summary of the given document. We consider the full story setting, where the task is to answer questions given the entire story and not the summary used to generate question-answer pairs. Here the knowledge source is not the same for all questions: given a question the retrieval operation is performed on all passages of the associated story. These passages are obtained by dividing the story in chunks of 100 words. These stories are long documents, with an average of 60k words. While part of the documents could be processed entirely by the Fusion-in-Decoder module, it is interesting to limit the number of support passages to reduce the computational cost of the reading step.
While answers in TriviaQA and NaturalQuestions are short, NarrativeQA answers are about five words long on average, with medium length answers such as ”He dismantles it and attaches it to his mother’s jeep” which answers the question ”What does Mark do with his radio station?”. Notably a significant number of answers do not correspond to spans in the story. It is thus not straightforward to train the retriever with heuristics using question-answer pairs. In our case we use the same pipeline as for TriviaQA and NaturalQuestions, demonstrating the flexibility of our approach.
Evaluation. The model performance is assessed in two ways. First, following previous work such as DPR and ColbertQA, we report the top-k retrieval accuracy (P@k), which is the percentage of questions for which at least one passage of the top-k retrieved passages contains the gold answer. It is unclear how well this metric evaluates the retriever performance, since the answer can be contained in a passage without being related to the question. This is notably true for common words or entities.
We also report the final end-to-end performance of the question answering system composed of the retriever and reader modules. This is the metric we are fundamentally interested in. For TriviaQA and NaturalQuestions, predicted answers are evaluated with the standard exact match metric (EM), as introduced by Rajpurkar et al. (2016). For NarrativeQA we report the metrics proposed in the original paper: ROUGE-L, BLEU-1, BLEU-4 and METEOR.

heading: 4.2 TECHNICAL DETAILS

text:
Initialization. Similarly to DPR, we initialize the retriever with the BERT base model, pretrained with uncased text. The Fusion-in-Decoder reader is initialized with the T5 base model. A critical component of the iterative training procedure is the initialization of the support passages D0q asso-
1https://cs.nyu.edu/˜kcho/NarrativeQA
ciated with each question q. For this we consider different options. The first one is to use passages retrieved using BM25. We use the implementation from Apache Lucene2 with default parameters, and tokenize questions and passages with SpaCy3. We also use passages obtained with BERT as a retriever without fine-tuning, this leads to poor initial performance. Finally in Table 2 we show that initializing D0q with passages obtained with DPR (Karpukhin et al., 2020) outperforms the two previous initializations. We train all retrievers using 100 passages. For the reader, we use 100 passages for NaturalQuestions and TriviaQA and 20 passages for NarrativeQA.
Iterative training. We apply the iterative training procedure on each dataset independently. Both the reader and the retriever are fine-tuned using the ADAM algorithm (Kingma & Ba, 2014), with a batch of size 64. The reader is trained for 10k gradient steps with a constant learning rate of 10−4, and the best model is selected based on the validation performance. The retriever is trained with a constant learning rate of 5 ·10−5 until the performance saturates. To monitor the performance of the retriever during training, we measure the similarity between the reader and the retriever rankings. At each new training iteration the reader is reinitialized from T5 base, while we pursue the training of the retriever. We found that restarting from T5 base is important for the first iterations when starting with BERT documents. We have not tried to reinitialize the retriever between each iteration. More details on the hyperparameters and the training procedure are reported in Appendix A.2.

heading: 4.3 RESULTS

text:
In Table 1, we report the performance of our approach for different number of self-training iterations. Generally, we observe that the accuracy of our system increases with the number of iterations, obtaining strong performance after a few iterations. Interestingly, while the initial performance with documents retrieved with BERT is very poor, our method still reach competitive scores on TriviaQA, and to a lesser extent, NaturalQuestions. However, a second observation is that the quality of the initial document sets plays an important role on the performance of the end system. Indeed, we observe that starting the procedure from BM25 documents, which are higher quality as indicated by the performance of the system at iteration 0, leads to stronger results than using BERT documents. An interesting research question would be to explore pre-training of the initial BERT model for retrieval, for example by using the inverse cloze task.
In Table 2, we report the performance of our approach, as well as existing state-of-the-art systems on TriviaQA and NaturalQuestions. In addition to initializing our method with documents retrieved with BM25 and BERT, we also train a system by starting from DPR documents. First, we observe that our method improve the performance over the state-of-the-art, even when starting from BM25 documents. This validates our assumption that it is possible to obtain strong retrievers without the need of supervision for the documents. Second, when starting from DPR passages, our method leads to a +4.5 EM improvement on TriviaQA and +2.3 EM improvement on NaturalQuestions when the final evaluation is carried out with a large reader.
2lucene.apache.org 3spacy.io
In Table 3, we report the performance of our method on the NarrativeQA dataset. We use the setting where the knowledge source corresponds to the whole document, and in particular, we do not use the summary. We compare our results to the best ones reported in the original paper for this setting. Similar to results obtained on NaturalQuestions and TriviaQA, we observe that training the retriever by using the attention scores of the reader leads to improvements, compared to the BM25 baseline.

heading: 5 ABLATIONS

text: In this section, we investigate design choices regarding two key elements of our approach: the training objective and the aggregation of cross-attention scores. For all experiments, we consider a simplified experimental setting: a single training iteration is performed on NaturalQuestions, starting from BM25 passages.

heading: 5.1 TRAINING OBJECTIVES

text: In Table 4 we report the performance of our model trained with the different training objectives described in Sec. 3.3. We observe that using the KL-divergence between the aggregated scores of the reader and the scores of the retriever outperforms the other objective functions.

heading: 5.2 HOW TO AGGREGATE CROSS-ATTENTION SCORES?

text:
In Section 4 the cross-attention scores α are aggregated in a specific way, in order to obtain a single scalar used to train the retriever. Formally let us denote by αi,j,k,h the cross-attention scores between token i of the output and token j of the input, for the k-th layer and h-th head. Then, the scores Gq,p for p ∈ Dq used in Section 4 are computed as follows:
Gq,p = mean j,k,h α0,j,k,h,
where j describes the input tokens corresponding to p. In Table 5 we explore alternatives to this choice by considering different aggregation schemes. In particular, we consider (1) taking the max over the input tokens corresponding to passage p instead of the average, (2) taking the average over the output tokens instead of taking the score of the first token, (3) taking the mean over the last six layers instead of all the layers, (4) taking the max over the layers instead of the average, (5) taking the max over the heads instead of the average. We observe that the performance of our approach is relatively stable to the choice of aggregation, and that the best result is obtained by averaging, except over the output tokens where it is best to only consider the first token.

heading: 6 CONCLUSION

text: In this paper, we introduce a method to train an information retrieval module for downstream tasks, without using pairs of queries and documents as annotations. Our approach is inspired by knowledge distillation, where the retriever module corresponds to the student model and the reader module corresponds to the teacher model. In particular, we use the cross-attention scores, from a sequenceto-sequence reader, to obtain synthetic targets for the retriever. We compare different ways to aggregate the scores, as well as different training objectives to learn the retriever. We show that iteratively training the reader and the retriever leads to better performance, and obtain state-of-the-art performance on competitive question answering benchmarks. In the future, we would like to explore better pre-training strategies for the retriever module, as well as better scoring functions for the retriever.

heading: A EXPERIMENTAL DETAILS

text:
A.1 SETTING
For NaturalQuestions and TriviaQA we follow the standard open-domain question answering setting used in Lee et al. (2019); Karpukhin et al. (2020). In this setting the original development set is used as test set, and 10% of the training set is used for development purpose. Moreover, for NaturalQuestions, all questions with answers longer than five tokens are discarded.
For TriviaQA we use the unique human-generated answer to train the reader. In this dataset part of the answers are in uppercase. We normalize uppercase answers by converting the first letter in each word to uppercase and remaining characters to lowercase using the title Python string method.
For NarrativeQA, questions and answers in uppercase are converted to lowercase.
A.2 TRAINING
For every datasets, both the reader and the retriever are fine-tuned with a dropout rate of 10%. All models at the exception of the large reader are trained using the ADAM algorithm (Kingma & Ba, 2014) with a constant learning rate of 10−4 for the base reader and 5 · 10−5 for the retriever. The base reader is trained for 10k gradient steps with a batch size of 64. We train the large reader with the ADAMW algorithm (Loshchilov & Hutter, 2019) with a peak learning rate of 5 · 10−5 and a linear warmup for 600 gradient steps followed by a linear decrease of the learning rate for 14.4k gradient steps.
We perform model selection on the validation performance. The retriever is trained until its performance saturates with a batch size of 64. To monitor the performance of the retriever during training, we measure the similarity between the ranking obtained with the reader score, and the ranking of the retriever. We use different metrics for this: the number of inversions between the two rankings, the proportion of passages in the retriever top-k that are also in the reader top-k and the number of passages to obtain all top-k passage of the reader.
During training and at test time, each text input of the encoder is restricted to be at most 250 token long. For NaturalQuestions and TriviaQA, we use wikipedia as a knowledge source, thus for each passage there is an associated article title. Each input is composed of the concatenation of a question, title and support passage with special tokens question:, title: and context: added before the question, the title and the text of each passage. In the case of NarrativeQA, the question and each passage are concatenated to form the different inputs.
A.3 INFERENCE
At test time, for TriviaQA and NaturalQuestions we use greedy decoding, and Beam Search with 3 beams for NarrativeQA.


summaries: The paper proposes a constrained reinforcement learning (RL) formulation relying on constraints written in a formal language. The proposed formulation is based on constrained Markov decision processes where the constraint is represented as a deterministic finite automaton that rejects any trajectory violating the constraint. The proposed solution relies on transforming the automaton's sparse binary cost into an approximate dense cost and augmenting that with the reward objective. The paper presents a series of results from simulations in Safety Gym, MuJoCo, and Atari environments.

summaries: The paper builds on the constrained MDP framework (Altman, 1999), by considering the special-case where the cost functions are defined in terms of states from a parser of a formal language. In the experiments the work uses deterministic finite automata (DFA) but in principle other more expressive classes could be used. Using a formal language to specify constraints may simplify model checking (although this is left to future work).

summaries: The authors propose to use formal languages, specifically DFAs, as a mechanism to specify constraints in a constrained MDP setting. This has the benefit of being able to rely on a large body of existing work on identification, safety verification, etc. The strategy relies on decomposing the constraint into a translation, recogniser & cost assignment function that connect the MDP to the DFA. The mentioned cost can then be combined with existing solution for solving cMDPs, such as reward shaping and Lagrangian methods. The key observation is that adding the recogniser state to the observations of the policy can result in significant gains in both performance and constraint satisfaction. A range of results are presented across different environment suites and hyper parameters.

abstractText: In order to satisfy safety conditions, an agent may be constrained from acting freely. A safe controller can be designed a priori if an environment is well understood, but not when learning is employed. In particular, reinforcement learned (RL) controllers require exploration, which can be hazardous in safety critical situations. We study the benefits of giving structure to the constraints of a constrained Markov decision process by specifying them in formal languages as a step towards using safety methods from software engineering and controller synthesis. We instantiate these constraints as finite automata to efficiently recognise constraint violations. Constraint states are then used to augment the underlying MDP state and to learn a dense cost function, easing the problem of quickly learning joint MDP/constraint dynamics. We empirically evaluate the effect of these methods on training a variety of RL algorithms over several constraints specified in Safety Gym, MuJoCo, and Atari environments.

authors:
Mohammed Alshiekh
Roderick Bloem
Rüdiger Ehlers
Bettina Könighofer
Scott Niekum
Ufuk Topcu

title: Safe reinforcement learning via shielding

venue: In Proceedings of AAAI 18,

title: Constrained Markov decision processes, volume 7

venue: CRC Press,

authors:
Dario Amodei
Chris Olah
Jacob Steinhardt
Paul Christiano
John Schulman
Dan Mané

title: Concrete problems in AI safety

venue: arXiv preprint arXiv:1606.06565,

authors:
N. Aréchiga
B. Krogh

title: Using verified control envelopes for safe controller design

venue: In 2014 American Control Conference,

authors:
Christel Baier
Boudewijn Haverkort
Holger Hermanns
Joost-Pieter Katoen

title: Model-checking algorithms for continuous-time Markov chains

venue: IEEE Transactions on software engineering,

authors:
M.G. Bellemare
Y. Naddaf
J. Veness
M. Bowling

title: The arcade learning environment: An evaluation platform for general agents

venue: Journal of Artificial Intelligence Research,

authors:
Ahmed Bouajjani
Javier Esparza
Oded Maler

title: Reachability analysis of pushdown automata: Application to model-checking

venue: In International Conference on Concurrency Theory,

authors:
Alberto Camacho
Oscar Chen
Scott Sanner
Sheila A McIlraith

title: Decision-making with non-markovian rewards: From LTL to automata-based reward shaping

venue: In Proceedings of the Multi-disciplinary Conference on Reinforcement Learning and Decision Making (RLDM),

authors:
Yinlam Chow
Mohammad Ghavamzadeh
Lucas Janson
Marco Pavone

title: Risk-constrained reinforcement learning with percentile risk criteria

venue: The Journal of Machine Learning Research,

authors:
Gal Dalal
Krishnamurthy Dvijotham
Matej Vecerik
Todd Hester
Cosmin Paduraru
Yuval Tassa

title: Safe exploration in continuous action spaces

venue: arXiv preprint arXiv:1801.08757,

authors:
Giuseppe De Giacomo
Marco Favorito
Luca Iocchi
Fabio Patrizi
Alessandro Ronca

title: Temporal logic monitoring rewards via transducers

authors:
Nathan Fulton
André Platzer

title: Verifiably safe off-model reinforcement learning

venue: In International Conference on Tools and Algorithms for the Construction and Analysis of Systems,

authors:
Ernst Moritz Hahn
Mateo Perez
Sven Schewe
Fabio Somenzi
Ashutosh Trivedi
Dominik Wojtczak

title: Omega-regular objectives in model-free reinforcement learning

authors:
Mohammadhosein Hasanbeig
Alessandro Abate
Daniel Kroening

title: Cautious reinforcement learning with logical constraints

venue: arXiv preprint arXiv:2002.12156,

authors:
Richard Helm
Ian M Holland
Dipayan Gangopadhyay

title: Contracts: specifying behavioral compositions in object-oriented systems

venue: ACM Sigplan Notices,

authors:
Rodrigo Toro Icarte
Toryn Q. Klassen
Richard Anthony Valenzano
Sheila A. McIlraith

title: Using reward machines for high-level task specification and decomposition in reinforcement learning

venue: In Proceedings of the 35th International Conference on Machine Learning,

authors:
Nils Jansen
Bettina Könighofer
Sebastian Junges
Roderick Bloem

title: Shielded decision-making in mdps

venue: arXiv preprint arXiv:1807.06096,

authors:
Orna Kupferman
Moshe Y Vardi
Pierre Wolper

title: An automata-theoretic approach to branchingtime model checking

venue: Journal of the ACM (JACM),

authors:
Marta Kwiatkowska
Gethin Norman
David Parker

title: Prism: Probabilistic symbolic model checker

venue: In International Conference on Modelling Techniques and Tools for Computer Performance Evaluation,

authors:
Xiao Li
Yao Ma
Calin Belta

title: Automata guided reinforcement learning with demonstrations

authors:
Michael L. Littman
Ufuk Topcu
Jie Fu
Charles Lee Isbell Jr.
Min Wen
James MacGlashan

title: Environment-independent task specifications via GLTL

title: Applying ‘design by contract

authors:
Seyed Sajad Mousavi
Behzad Ghazanfari
Nasser Mozayani
Mohammad Reza Jahed-Motlagh

title: Automatic abstraction controller in reinforcement learning agent via automata

venue: Appl. Soft Comput.,

authors:
Andrew Y. Ng
Daishi Harada
Stuart J. Russell

title: Policy invariance under reward transformations: Theory and application to reward shaping

venue: In Proceedings of the Sixteenth International Conference on Machine Learning (ICML

authors:
Tu-Hoa Pham
Giovanni De Magistris
Ryuki Tachibana

title: Optlayer-practical constrained optimization for deep reinforcement learning in the real world

venue: IEEE International Conference on Robotics and Automation (ICRA),

authors:
Alex Ray
Joshua Achiam
Dario Amodei

title: Benchmarking safe exploration in deep reinforcement learning

venue: arXiv preprint arXiv:1910.01708,

authors:
William Saunders
Girish Sastry
Andreas Stuhlmueller
Owain Evans

title: Trial without error: Towards safe reinforcement learning via human intervention

venue: arXiv preprint arXiv:1707.05173,

authors:
John Schulman
Filip Wolski
Prafulla Dhariwal
Alec Radford
Oleg Klimov

title: Proximal policy optimization algorithms

venue: arXiv preprint arXiv:1707.06347,

authors:
Richard S Sutton
Andrew G Barto

title: Reinforcement learning: An introduction

venue: MIT press,

authors:
Min Wen
Ivan Papusha
Ufuk Topcu

title: Learning from demonstrations with high-level side information

venue: In Proceedings of the Twenty-Sixth International Joint Conference on Artificial Intelligence,

authors:
Eric Wiewiora
Garrison W Cottrell
Charles Elkan

title: Principled methods for advising reinforcement learning agents

venue: In Proceedings of the 20th International Conference on Machine Learning

authors:
He Zhu
Zikang Xiong
Stephen Magill
Suresh Jagannathan

title: An inductive synthesis framework for verifiable reinforcement learning

venue: In Proceedings of the 40th ACM SIGPLAN Conference on Programming Language Design and Implementation,

heading: 1 INTRODUCTION

text:
The ability to impose safety constraints on an agent is key to the deployment of reinforcement learning (RL) systems in real-world environments (Amodei et al., 2016). Controllers that are derived mathematically typically rely on a full a priori analysis of agent behavior remaining within a predefined envelope of safety in order to guarantee safe operation (Aréchiga & Krogh, 2014). This approach restricts controllers to pre-defined, analytical operational limits, but allows for verification of safety properties (Huth & Kwiatkowska, 1997) and satisfaction of software contracts (Helm et al., 1990), which enables their use as a component in larger systems. By contrast, RL controllers are free to learn control trajectories that better suit their tasks and goals; however, understanding and verifying their safety properties is challenging. A particular hazard of learning an RL controller is the requirement of exploration in an unknown environment. It is desirable not only to obey constraints in the final policy, but also throughout the exploration and learning process (Ray et al., 2019).
The goal of safe operation as an optimization objective is formalized by the constrained Markov decision process (CMDP) (Altman, 1999), which adds to a Markov decision process (MDP) a cost signal similar to the reward signal, and poses a constrained optimization problem in which discounted reward is maximized while the total cost must remain below a pre-specified limit per constraint. We use this framework and propose specifying CMDP constraints in formal languages to add useful structure based on expert knowledge, e.g., building sensitivity to proximity into constraints on object collision or converting a non-Markovian constraint into a Markovian one (De Giacomo et al., 2020).
A significant advantage of specifying constraints with formal languages is that they already form a well-developed basis for components of safety-critical systems (Huth & Kwiatkowska, 1997; Clarke et al., 2001; Kwiatkowska et al., 2002; Baier et al., 2003) and safety properties specified in formal languages can be verified a priori (Kupferman et al., 2000; Bouajjani et al., 1997). Moreover, the recognition problem for many classes of formal languages imposes modest computational requirements, making them suitable for efficient runtime verification (Chen & Roşu, 2007). This allows for low-overhead incorporation of potentially complex constraints into RL training and deployment.
We propose (1) a method for posing formal language constraints defined over MDP trajectories as CMDP cost functions; (2) augmenting MDP state with constraint automaton state to more explicitly encourage learning of joint MDP/constraint dynamics; (3) a method for learning a dense cost function given a sparse cost function from joint MDP/constraint dynamics; and (4) a method based on
constraint structure to dynamically modify the set of available actions to guarantee the prevention of constraint violations. We validate our methods over a variety of RL algorithms with standard constraints in Safety Gym and hand-built constraints in MuJoCo and Atari environments.
The remainder of this work is organized as follows. Section 2 presents related work in CMDPs, using expert advice in RL and safety, as well as formal languages in similar settings. Section 3 describes our definition of a formal language-based cost function, as well as how it’s employed in state augmentation, cost shaping, and action shaping. Section 4 details our experimental setup and results and finally, discussion of limitations and future work are located in Section 5.

heading: 2 RELATED WORK

text:
Safety and CMDP Framework The CMDP framework doesn’t prescribe the exact form of constraints or how to satisfy the constrained optimization problem. Chow et al. (2017) propose conditional value-at-risk of accumulated cost and chance constraints as the values to be constrained and use a Lagrangian formulation to derive a Bellman optimality condition. Dalal et al. (2018) use a different constraint for each MDP state and a safety layer that analytically solves a linearized action correction formulation per state. Similarly, Pham et al. (2018) introduce a layer that corrects the output of a policy to respect constraints on the dynamic of a robotic arm.
Teacher Advice A subset of work in safe exploration uses expert advice with potential-based reward shaping mechanisms (Ng et al., 1999). Wiewiora et al. (2003) introduce a general method for incorporating arbitrary advice into the reward structure. Saunders et al. (2017) use a human in the loop to learn an effective RL agent while minimizing cost accumulated over training. Camacho et al. (2017a;b) use DFAs with static reward shaping attached to states to express non-Markovian rewards. We generalise this work with a learned shaping function in the case of dense soft constraints, and by generalising from reward shaping to other CMDP learning mechanisms. Similar to teacher advice is shielding (Jansen et al., 2018; Alshiekh et al., 2018), in which an agent’s actions are filtered through a shield that blocks actions that would introduce an unsafe state (similar to hard constraints; section 3), but typically requires MDP states to be enumerable and few enough that a shield can be constructed efficiently.
Formal Languages Formal languages and automata have been used before in RL for task specification or as task abstractions (options) in hierarchical reinforcement learning (Icarte et al., 2018b; Li et al., 2017; Wen et al., 2017; Mousavi et al., 2014). In some cases, these automata were derived from Linear Temporal Logic (LTL) formulae, in others LTL or other formal language formulae have been directly used to specify tasks (Icarte et al., 2018a). Littman et al. (2017) defines a modified LTL designed for use in reinforcement learning. In robotics, LTL is used for task learning (Li et al., 2017), sometimes in conjunction with teacher demonstrations (Li et al., 2018). Zhu et al. (2019) and Fulton & Platzer (2019) both study the use of formal languages for safe RL, though each makes assumptions about a prior knowledge of the environment dynamics. Hasanbeig et al. (2018) and Hasanbeig et al. (2020) learn a product MDP with a safety constraint specified in a formal language,
but require knowledge of MDP structure and don’t scale to large MDPs. Hahn et al. (2019) uses an omega-regular language to specify an objective and constructs a product MDP to ensure the objective is reached almost surely, but requires large reward shaping which prevents learning when used for a constraint rather than objective.

heading: 3 FORMAL LANGUAGE CONSTRAINED MDPS

text:
The constrained Markov decision process (CMDP) (Altman, 1999) extends the Markov decision process (Sutton & Barto, 2018) to incorporate constraints. The difference is an additional set of cost functions ci : S ⇥ A ⇥ S ! R and set of cost limits di 2 R. Then, the constrained optimization problem is argmax⇡ Jr(⇡) s.t. Jci(⇡)  di, i = 1, . . . , k where Jr(⇡) is a return-based objective, e.g., finite horizon discounted return defined Jr(⇡) = E⌧⇠⇡[ P t2⌧
trt] and Jci is a cost-based constraint function defined similarly, replacing rt with ci,t.
We propose formal language constrained MDPs (FLCMDPs) as a subset of CMDPs in which each constraint Ci ⇢ (S⇥A⇥S)⇤ is defined by a set of prohibited trajectories. (Subscript i is suppressed from this point without loss of generality). Because C is defined by a formal language, it can be recognized efficiently by an automaton, which we use to construct the cost function. We define three functions for interacting with the constraint automaton: a translation function TC : (S⇥A⇥S) ! ⌃C that converts MDP transitions into a symbol in the recognizer’s input language, a recognizer function DC : Qc⇥⌃C ! QC that steps the recognizer automaton using its current state and the input symbol and returns its next state, and finally a cost assignment GC : QC ! R that assigns a real-valued cost to each recognizer state. The composition of these three functions forms a CMDP cost function defined c = GC DC TC : (S ⇥ A ⇥ S) ⇥ Qc ! R, which requires constraint state Qc to be tracked along with the MDP state, S. The interaction of these functions with the underlying MDP framework is illustrated in Figure 1(a), where the constraint uses the MDP state and action at time t to calculate the cost signal at time t and, if action shaping is being employed as discussed below, influence the action at time t+ 1. We note that this construction does not require knowledge of the MDP transition probabilities and further, its complexity is independent of the size of the MDP, which allows it to scale to large, challenging environments.
Translation Function TC The translation function accepts as input the MDP state and action at each time step, and outputs a token in the discrete, finite language of the associated recognizer. This allows the recognizer automaton to be defined in a small, discrete language, rather than over unwieldy and potentially non-discrete MDP transitions. Further, freedom in choice of input language allows for flexible design of the constraint automaton to encode the desired inductive bias, and thus more meaningful structured states.
Recognizer Function DC Each constraint is instantiated with a finite automaton recognizer that decides whether a trajectory is in the constraint set. The only necessary assumption about the recognizer is that it defines some meaningful state that may be used for learning the constraint dynamics. Our implementation uses a deterministic finite automaton (DFA) as the recognizer for each constraint, defined as (Q,⌃, , q0, F ), where Q is the set of the DFA’s states, ⌃ is the alphabet over which the constraint is defined, : Q⇥ ⌃ ! Q is the transition function, q0 2 Q is the start state, and F ⇢ Q is the set of accepting states that represent constraint violations. The DFA is set to its initial state at the start of each episode and is advanced at each time step with the token output by the translation layer. Although our experiments use DFAs as a relatively simple recognizer, the framework can be easily modified to work with automata that encode richer formal languages like pushdown automata or hybrid automata.
Constraint State Augmentation In order to more efficiently learn constraint dynamics, the MDP state st is augmented with a one-hot representation of the recognizer state qt. To preserve the Markov property of the underlying MDP, state augmentation should contain sufficient information about the recognizer state and, if it is stateful, the translation function. To enhance performance, the one-hot state is embedded to dlog2(|Q|)e dimensions before being input into any network and the embedding is learned with gradients backpropagated through the full network.
Cost Assignment Function The cost assignment function GC assigns a real-valued cost to each state of the recognizer. This cost can be used in optimization to enforce the constraint with a Lagrangian-derived objective penalty, or via reward shaping, which updates the reward function to rt ct, where is a scaling coeffficient. Cost assignments are frequently sparse, where GC is only non-zero at accepting states that recognize a constraint violation. This poses a learning problem for optimization-based methods that use reward shaping or an objective penalty to solve the CMDP. A goal of constrained RL is to minimize accumulated constraint violations over training but, to ensure that the frequency of violations is small, the optimization penalty can be large relative to the reward signal. This can lead to a situation in which an unnecessarily conservative policy is adopted early in training, slowing exploration. We next propose a method for learning a dense cost function that takes advantage of the structure of the constraint automaton to more quickly learn constraint dynamics and avoid unnecessarily conservative behavior.
Learned Dense Cost The goal of learning a dense cost is not to change the optimality or nearoptimality of a policy with respect to the constrained learning problem. Thus, we use the form of potential-based shaping: F (st 1, at, st) = (st) (st 1), where is a potential function (see Ng et al. (1999) for details). This is added as a shaping term to the sparse cost to get the dense cost G0C(qt 1, qt) = GC(qt) + ( (qt) (qt 1)) , where scales the added dense cost relative to the sparse cost, and is a function of the recognizer state rather than the MDP state, which requires st 2 and at 2 as additional inputs to calculate qt 1. Generally, if the value of increases as the automaton state is nearer to a violation, then the added shaping terms add cost for moving nearer to a constraint violation and refund cost for backing away from a potential violation.
In our experiments, the potential ⇡(qt) is defined using tv(qt), which is a random variable defined as the number of steps between visiting recognizer state qt and an accepting recognizer state. This variable’s distribution is based on ⇡ and the MDP’s transition function. Its value is small if a violation is expected to occur soon after reaching qt and vice-versa. We then define the potential function as ⇡(qt) = 1 2
(E⇡ [Tv(qt)]/tbaselinev ) , which ensures that its value is always in [0, 1] and rises exponentially as the expected time to a violation becomes smaller. If the expected time to next violation is much larger than the provided baseline, tbaselinev , then the potential value becomes small, as shaping is unnecessary in safe states. The expected value of Tv(qt) over trajectories may be calculated as a Monte Carlo estimate from rollouts by counting the number of steps between each visit to qt and the next constraint violation. Each Tv is tracked with an exponential moving average and is updated between episodes to ensure that it’s stationary in each rollout. We set tbaselinev to be the ratio of estimated or exact length of an episode and the constraint limit di, but find empirically that the the method is resilient to the exact choice.
Hard Constraints and Action Shaping When safety constraints are strict, i.e., when the limit on the number of constraint violations d is zero, the set of available actions is reduced to ensure a violation cannot occur. If a constraint isn’t fully state-dependent (i.e., there is always an action choice that avoids violation), then action shaping can guarantee that a constraint is never violated. Otherwise, knowledge of which actions lead to future violating trajectories requires knowledge of the underlying MDP dynamics, which is possible by learning a model that converts state constraints into state-conditional action constraints as in Dalal et al. (2018).
Our implementation of hard constraints initially allows the agent to freely choose its action, but before finalizing that choice, simulates stepping the DFA with the resulting token from the translation function and, if that lookahead step would move it into a violating state, it switches to the next best choice until a non-violating action is found. For the constraints in our experiments, it is always possible to choose a non-violating action. A known safe fallback policy can be employed in the case when an episode cannot be terminated. Action shaping can be applied during training or deployment, as opposed to reward shaping, which is only applied during training. We experiment with applying action shaping only during training, only during evaluation, or in both training and evaluation.

heading: 4 EXPERIMENTAL EVALUATION

heading: 4.1 CONSTRAINTS

text:
We evaluated FLCMDPs on four families of constraints, which we define with regular expressions. The cost function defined by each constraint is binary, unless we use dense cost shaping.
No-dithering: A no-dithering constraint prohibits movements in small, tight patterns that cover very small areas. In one dimension, we define dithering as actions are taken to move left, right, left, and right in order or the opposite, i.e., .⇤ (` r)2|(r `)2. The automaton encoding this constraint is depicted in Figure 1(b). In environments with two-dimensional action spaces, such as Atari Seaquest, we generalize this to vertical and diagonal moves and constrains actions that take the agent back to where it started in at most four steps1. In MuJoCo, constraints are applied per joint and the translation function maps negative and positive-valued actions to ‘`’ and ‘r’, respectively.
No-overactuating: A no-overactuating constraint prohibits repeated movements in the same direction over a long period of time. In Atari environments, this forbids moving in the same direction four times in a row, i.e., .⇤ (`4 [ r4). In two dimensions, this is extended to include moving vertically: .⇤ (L4 [R4 [ U4 [D4). Each of the left (L), right (R), up (U ) and down (D) tokens is produced by the translation function from the primary direction it’s named after or diagonal moves that contain the primary direction, e.g., L = ` [ `+u [ `+d, where “`+u” is the atomic left-up diagonal action. In MuJoCo environments, overactuation is modelled as occurring when the sum of the magnitudes of joint actuations exceeds a threshold. This requires the translation function to discretize the magnitude in order for a DFA to calculate an approximate sum. The MDP state-based version is “dynamic actuation”, which sets the threshold dynamically based on a discretized distance from the goal.
Proximity: The proximity constraint, used in Safety Gym, encodes the distance to a collision with any of the variety of hazards found in its environments. The translation function uses the max value over all the hazard lidars, which have higher value as a hazard comes closer, and discretizes it into one of ten values. The constraint is defined as being violated if the agent contacts the hazard, which is identical to the constraint defined in the Safety Gym environments and described in Ray et al. (2019).
Domain-specific: In addition to the previously described simple constraints, we define hand-built constraints for the Breakout and Space Invaders Atari environments. These constraints are designed to mimic specific human strategies in each environment for avoiding taking actions that end the episode. In Atari Breakout, we define the “paddle-ball” constraint, which limits the allowed horizontal distance between the ball and the center of the paddle. In Atari Space Invaders, we define the “danger zone” constraint, which puts a floor on the the allowed distance between the player’s ship and the bullets fired by enemies. We provide more details of each constraint in Appendix B.

heading: 4.2 ENVIRONMENTS

text:
In Safety Gym, the Spinning Up implementation of PPO with Lagrangian optimization penalization was employed, with hyperparameters as chosen identically to Ray et al. (2019). We modified each network to concatenate the constraint state augmentation with the input and used d = 25 for the expected cost limit. All safety requirements are accounted for in a single constraint and we report the constraint violations as accounted for in the Safety Gym environments rather than as reported by the finite automaton (though these are identical when not using cost shaping). Each environment, which is randomly re-arranged each episode, is made up of a pairing of a robot and a task. The robots are Point, which turns and moves, and Car, which is wheeled with differential drive control. The tasks are Goal, which requires moving into a goal area, Button, which requires pressing a series of buttons, and Push, which requires moving a box into a goal area. More details can be found in Ray et al. (2019).
In Atari environments (Bellemare et al., 2013), we modified the DQN implemented in OpenAI Baselines (Dhariwal et al., 2017) by appending the state augmentation to the output of its final convolutional layer. Reward shaping was used for soft constraint enforcement with the penalty fixed at one of {0, 0.001, 0.0025, 0.005, 0.01}, and each agent was trained for 10M steps before collecting data from an evaluation phase of 100K steps for 15 or more train/eval seed pairs for each hyperparameter combination. For MuJoCo environments (Brockman et al., 2016), we trained the
1The regex describing this constraint is included in Appendix D.
Baselines PPO agent (Schulman et al., 2017) with dense constraints and the state augmentation concatenated to the input MDP observation. Reward shaping similar to Atari was employed for constraint enforcement. Atari and MuJoCo environments do not have constraints built in, so we report the number of constraint violations per episode from the custom constraints and minimize them without a specific goal value.

heading: 4.3 RESULTS

text:
We experiment with our methods to evaluate the usefulness of formal language constraints in optimizing three objectives. In the final policy output by the training process, it is desirable to simultaneously maximize return per episode and minimize constraint violations per episode, or keep them below the specified limit. The third objective is to minimize accumulated cost regret
Environment Constraint Dense Shaping Episode Reward Step Reward Viols/100 steps Episode Reward Step Reward Viols/100 steps
FLCMDP State Augmented Baseline Reward Mean Mean Mean Mean Mean Mean
Environment Constraint Dense Shaping Episode Reward Step Reward Viols/100 Steps Episode Reward Step Reward Viols/100 Steps
Breakout actuation False 0.001 297.12 ± 8.07 0.15 ± 0.0039 0.45 ± 0.00067 272.19 ± 43.12 0.14 ± 0.01 13.59 ± 0.025 dithering True 0.005 302.24 ± 43.81 0.14 ± 0.02 0.11 ± 0.001 272.19 ± 43.12 0.14 ± 0.01 0.12 ± 0.001 paddle ball True 0.0025 314.79 ± 15.09 0.17 ± 0.0056 6.15 ± 0.0031 272.19 ± 43.12 0.14 ± 0.01 13.40 ± 0.022
Seaquest actuation True 0.0025 2339.54 ± 442.02 0.93 ± 0.11 4.43 ± 0.016 2250.13 ± 647.92 0.96 ± 0.21 9.74 ± 0.017dithering False 0.001 1997.91 ± 539.75 0.86 ± 0.23 1.58 ± 0.011 2250.13 ± 647.92 0.96 ± 0.21 1.61 ± 0.007
SpaceInvaders actuation False 0.005 646.99 ± 50.55 0.64 ± 0.04 29 ± 0.0053 604.86 ± 44.86 0.62 ± 0.04 10.88 ± 0.011 dangerzone False 0.001 687.37 ± 16.75 0.63 ± 0.01 0.00 ± 0.00 604.86 ± 44.86 0.62 ± 0.04 0.00 ± 0.00 dithering False 0.001 640.35 ± 25.94 0.67 ± 0.09 0.17 ± 0.00027 604.86 ± 44.86 0.62 ± 0.04 0.53 ± 0.0064
over the course of training. To examine the proposed methods, we investigate two questions. First, what effect do the proposed methods have on accumulated regret? In Section 4.3.1, we compare the proposed methods against a baseline when combined with PPO using a Lagrangian approach in Safety Gym (Ray et al., 2019). Second, how should hyperparameters be chosen to minimize or maximize each objective respectively? In Section 4.3.2 we examine which hyperparameter choices worked well in various Atari and MuJoCo environments using reward shaping. Finally, in Section 4.3.3 we see what effect enforcing zero constraint violations with action shaping has on Atari environments.

heading: 4.3.1 LAGRANGIAN RESULTS AND ACCUMULATED COST

text:
Table 1 compares PPO with Lagrangian constraint enforcement with and without constraint state augmentation. The clearest trend is in the reduction of the cost rate, which measures accumulated cost regret, often by between almost one half and an order of magnitude. This results from the inclusion of the helpful inductive bias provided by the constraint structure. This result is not surprising, but does quantify the magnitude of the benefit that a low-overhead method like formal language constraints can have. Qualitatively, we noted that the earliest steps of training had decreased performance generally as the embedding of the constraint state was being learned, but quickly surpassed baseline performance once the updates of the embedded representation became small.
Ray et al. (2019) says that algorithm A1 dominates A2 when they are evaluated under the same conditions, the cost and return of A1 are at least as good as that of A2, and at least one of cost and return is strictly better for A1. By this definition, the state augmented approach strictly dominates the baseline in 6 of 12 environments, while coming close in most of the rest. Specifically, we also note that state augmentation allowed a significant step to be taken in closing the gap between unconstrained PPO return and PPO-Lagrangian in the Point-Goal2 and Car-Push2 environments, with increases of roughly an order of magnitude in each.

heading: 4.3.2 REWARD SHAPING RESULTS AND SENSITIVITY TO REWARD SHAPING

text:
The most basic function of the proposed framework is to reduce constraint violations. Figure 2 shows performance/conformance curves in Safety Gym where reward shaping was varied in {0, 0.025, 0.005, 0.1, 1}. The curves show a meaningful trade-off between the objectives episodic return and cost, which is improved by the use of proposed methods. Table 2 presents the mean and standard deviation of violations per 100 evaluation steps, episode length, and episode reward for reward shaping-enforced constraints with the choice of hyperparameters that produced
Environment Constraint Episode Reward Viols/100 Steps Episode Reward Viols/100 Steps Episode Reward Viols/100 Steps
Breakout actuation 302.00 ± 20.75 0.0 ± 0.0 320.31 ± 6.09 0.092 ± 0.00026 314.91 ± 13.80 0.0 ± 0.0 dithering 295.31 ± 29.07 0.0 ± 0.0 276.72 ± 15.50 0.0073 ± 3.1e-05 275.25 ± 12.67 0.0 ± 0.0 paddle ball 218.14 ± 22.85 0.0 ± 0.0 281.77 ± 11.03 0.11 ± 0.00029 229.00 ± 11.65 0.0 ± 0.0
Seaquest actuation 1926.97 ± 430.24 0.0 ± 0.0 1899.78 ± 502.27 8.30 ± 0.043 1895.77 ± 366.79 0.0 ± 0.0dithering 2284.78 ± 15.45 0.0 ± 0.0 2256.06 ± 30.53 0.01 ± 2.8e-05 2267.53 ± 23.94 0.0 ± 0.0
SpaceInvaders actuation 586.66 ± 58.69 0.0 ± 0.0 582.79 ± 51.47 14.62 ± 0.012 583.13 ± 60.60 0.0 ± 0.0 dangerzone 613.61 ± 24.05 0.0 ± 0.0 733.52 ± 16.95 0.0 ± 0.0 613.82 ± 27.52 0.0 ± 0.0 dithering 627.40 ± 31.43 0.0 ± 0.0 624.33 ± 25.27 0.008 ± 3e-05 626.59 ± 31.04 0.0 ± 0.0
the minimum violations for each environment/constraint pair in evaluation. We note that the highest value of reward shaping available is generally the best choice for minimizing constraint violations, which were often reduced by an order of magnitude or more from the baseline. Minimizing constraint violations has a small deleterious effect on mean episode reward, but because mean reward per step didn’t decrease, episodes were shorter as a result of constraint enforcement.
In addition to minimizing constraint violations, we found that the application of soft constraints can also increase reward per episode. Table 3 presents results for soft constraints with the choice of hyperparameters that produced the maximum reward in each environment/constraint pair. In this case, lower reward shaping values perform best. The hyperparameter values that minimized constraint violations with the Breakout actuation and paddle ball constraints also maximized reward, implying that the objectives were correlated under those constraints. Table 4 presents results for soft constraints with constraint state augmentation in three MuJoCo environments. We find, similar to Atari, that there is one value of reward shaping that is most effective in each environment/constraint pair and that reward degrades smoothly as is shifted from the optimal value.

heading: 4.3.3 HARD ACTION SHAPING RESULTS

text: Table 5 presents results for hard constraints with the hyperparameters that produced the maximum return for each environment/constraint pair. Results for cases where hard action shaping was only applied during training or only applied during evaluation are presented as well. There is a slight trend indicating that using action shaping at train time in addition to evaluation increases performance. For those constraints that are qualitatively observed to constrain adaptive behavior, performance rises when using hard shaping only in training, at the cost of allowing constraint violations.

heading: 5 DISCUSSION

text:
The ability to specify MDP constraints in formal languages opens the possibility for using model checking (Kupferman et al., 2000; Bouajjani et al., 1997), agnostic to the choice of learning algorithm, to verify properties of a safety constraint. Formal language constraints might be learned from exploration, given a pre-specified safety objective, and, because of their explicitness, used without complication for downstream applications or verification. This makes formal language constraints particularly useful in multi-component, contract-based software systems (Meyer, 1992), where one or more components is learned using the MDP formalism.
Experiments with more complex constraints are necessary to explore yet unaddressed challenges, the primary challenge being that the constraints used with action shaping in this work ensured the allowed set of actions was never empty. If this is not the case, lookahead might be required to guarantee zero constraint violations. Further, the tested hard constraints were only used with DQN, which provides a ranked choice over discrete actions. Future work might investigate how to choose optimal actions which are not the first choice in the absence of ranked choice.