# Extended Mind Transformers

Phoebe Klett
Normal Computing
&Thomas Ahle
Normal Computing
Correspondence to<EMAIL_ADDRESS>
###### Abstract

Pre-trained language models demonstrate general intelligence and common sense,
but long inputs quickly become a bottleneck for memorizing information at
inference time. We resurface a simple method, Memorizing Transformers [ICLR
2022], that gives the model access to a bank of pre-computed memories. We show
that it is possible to fix many of the shortcomings of the original method,
such as the need for fine-tuning, by critically assessing how positional
encodings should be updated for the keys and values retrieved. This intuitive
method uses the model’s own key/query system to select and attend to the most
relevant memories at each generation step, rather than using external
embeddings. We demonstrate the importance of external information being
retrieved in a majority of decoder layers, contrary to previous work. We open
source 111Code and data: https://github.com/normal-computing/extended-mind-
transformers/ a new counterfactual long-range retrieval benchmark, and show
that Extended Mind Transformers outperform today’s state of the art by 6% on
average.

##  1 Introduction and History of Memory Augmented Neural Networks

Today’s large language models (LLMs) are impressive general learners. During
training they memorize large vocabularies, grammars, and wide-ranging
knowledge. [brown2020language] In many application settings, however, the
model also needs to make use of specialized or topical information. The facts
seen during pre-training may have changed, or new information may need to be
memorized. Keeping track of the state of the world as the model is queried is
crucial for sound reasoning, but efficiently loading a detailed perspective of
the world remains a challenge.

This problem can largely be split into the following related but distinct sub-
problems: $(1)$ Improve performance by extending the maximum input token
sequence length $(2)$ Improve the efficiency of the attention mechanism $(3)$
Improve performance by memorizing or retrieving the relevant information.
Longer input sequences allow us to memorize more information at query time,
improving the efficiency of the attention computation allows us to attend to
more information, and retrieval methods allow us to determine which
information from the past is relevant to the current query. Many solutions to
the problem of leveraging a history of information or long input sequences
address more than one of the above sub-problems, and solutions to all three
are well-explored [Biggs2015NeuralMN]. We describe their history here to
situate our contributions.

Figure 1: Overview of attention over memories and local context, where Q is
length of queries, KV is length of key-values, HD is head dimension, and K is
a memory hyper-parameter. Arrows show queries retrieving memories, gradient
colors represent inner product scores and softmax.

###  1.1 Extending sequence length

Auto-regressive language models are trained on sequences of a fixed length.
Following the influential introduction of the transformer model in
[vaswani2023attention], Transformer-XL (Dai et al., 2019)
[dai2019transformerxl] proposed to generalize beyond that fixed context length
by conditioning each time step’s hidden states on the previous step’s hidden
states. In particular, the previous time step’s hidden states are concatenated
to the current time step’s hidden states before computing the key and value
states. Importantly, the authors in [dai2019transformerxl] also introduce the
first relative position embedding, a crucial part of generalizing beyond
training sequence length. Because transformers are position agnostic
[yun2020transformers], vectors encoding position information are added (or
otherwise incorporated) to the input embeddings. Absolute position vectors are
unique per position, whereas relative position vectors simply encode which
tokens are close (far). Compressive transformers [rae2019compressive] built on
this work by compressing those hidden states furthest away, those that would
fall out of the Transformer-XL [dai2019transformerxl] cache.

Today, strategies to extend the context window have been enabled by
improvements in position encodings. Attention with Linear Biases (ALiBi)
[press2022train] uses linear biases to damp down the relative importance of
long-range tokens rather than using explicit position information. Rotary
position embeddings [su2023roformer] generalize the intuition of using the
angle between two-dimensional vectors as a distance metric. By interpolating
between positions, we can create more unique rotary position embeddings within
the expected range [chen2023extending]. Both linear biases and rotary position
embeddings enable models to generalize to longer input sequences without fine-
tuning, but performance on very long inputs requires additional training.

###  1.2 Approximate Attention

Along a concurrent, parallel path, approximations to the usual causal
attention mechanism enabled longer training sequences by reducing
computational costs. Sparse attention factorizations of the attention
[child2019generating] reduce the quadratic cost of attention
[vaswani2023attention] to $O(n\sqrt{n})$. Important drop-in attention
approximations (they can be substituted for regular attention without further
fine-tuning) were proposed by both Longformer [beltagy2020longformer] and
Memory efficient transformers [gupta2021memoryefficient]. The former
introduces a variety of sliding window attention patterns, used now in popular
open-sourced LLMs such as Mistral [jiang2023mistral]. The latter proposes
top-k attention; attention is only computed for those key-value pairs where
the key-query inner product is the largest. This line of research has
continued to be fruitful, but it’s those hardware-aware improvements
[dao2022flashattention] that have become standard.

###  1.3 Retrieval augmentations

Retrieval methods predate the transformer model. Indeed, many transformer-
oriented retrieval methods cite and build on the original ideas proposed in
[graves2014neural] Neural turing machines (Graves et al., 2014) and Memory
Networks (Weston et al., 2015) [weston2015memory], which were influential but
computationally expensive and difficult to train. Nearest Neighbor Language
Models (kNN-LM) [khandelwal2020generalization] interpolate the model’s output
distribution with a distribution over labels of similar inputs seen during
training. KNN Information Fetching Modules[fan2020augmenting] operate by
retrieving relevant (potentially multi-modal) external sources and
concatenating with the encoded input before feeding into the decoder modules
in a classic encoder-decoder transformer. Retrieval-Augmented generation (RAG)
(Lewis et al., 2021) [lewis2021retrievalaugmented] builds directly on both the
previous ideas and proposes to retrieve external documents based on similarity
with the current input sequence, and condition on those documents as
additional context during generation. The encoding used to perform retrieval
need not be the same encoding used by the seq2seq model, since the tokens are
appended to the input sequence as text, previous to any generation.

All of the aforementioned methods represent query and long-range information
using standard embeddings. Much of the recent work on retrieval has explored
the optimal way to represent long-range information for retrieval. The Memory
Transformer (Burtsev et al., 2021) [burtsev2021memory], similar to
Transformer-XL [dai2019transformerxl], demonstrates that long-range
information can be compressed and memorized in dedicated [MEM] tokens, both
using shared and disjoint weight matrices from the query tokens. Finally,
Memorizing Transformers (Wu et al., 2022) [wu2022memorizing] proposes to cache
previous key-value representations, and perform top-k attention
[gupta2021memoryefficient] over those key-value pairs in a single decoder
layer. We believe this method has been underestimated, and it forms the basis
for our work in this paper.

Figure 2: Average perplexity on sequences of increasing input lengths. Results
shown for baselines and Extended Mind Llama-2-7b.

Figure 3: Average perplexity on sequences of increasing input lengths.
Increasing k corresponds to retrieving more key-values pairs.

###  1.4 Extended Mind Transformers

After so much experimentation, what have we learned? We’ve traded the
recurrent mechanism of caching and conditioning on previous hidden states, as
implemented by Transformer-XL [dai2019transformerxl], for the simplicity and
interpretability of similarity-based retrieval methods. We see this learning
as analogous to the Transformer advantage over LSTM models. Recent advances in
position encodings allow models to generalize to longer inputs more natively,
and some approximations to classic attention have made their way into
practice. But the two most popular methods, fine-tuning models to memorize
long input sequences and RAG [lewis2021retrievalaugmented] both leave much to
be desired.

Putting the entire history into context suffers the quadratic cost of
attention unnecessarily. More often than not, only a subset of the new
information is needed at any given generation step. Further, our experiments 5
show that attention over long inputs continues to struggle on long-range
retrieval tasks even after fine-tuning. A more optimal solution would combine
attention over long(er) inputs with a smart retrieval mechanism. This is
achieved today by combining fine-tuned models with RAG. But while RAG
[lewis2021retrievalaugmented] has become widespread and easy to implement
[llamaindex], it faces many challenges. All retrieval mechanisms external to
the model are fundamentally limited by relying on an external embedding model,
rather than the transformer’s own key/query system. This decision is made once
per query and used for all generation steps. Assuming that the entire long
inputs do not fit into context (otherwise RAG would be unnecessary), we cannot
pose queries that synthesize over all inputs. Challenges that face both fine-
tuning and RAG include recent evidence that language models struggle to ignore
irrelevant and false information [chen2023benchmarking].

Retrieval methods internal to the model avoid these issues by retrieving only
the most relevant information at each step, where relevance is measured using
the model’s representation of the data. They have yet to be adopted in
practice due to the staleness problem, or the fear that representations will
become out of distribution as the model weights evolve, and because the
additional complexity and overhead of fine-tuning is difficult to justify when
RAG is often effective. We hypothesized that just as the relative position
embeddings enable native generalization to longer inputs, they should enable
models to use past key-values retrieved within decoder layers, as proposed by
[wu2022memorizing]. If true, this would eliminate both the need to fine-tune
and thus the staleness problem as well.

###  1.5 Our Contributions

In this work we make the following contributions:

  * •

We propose Extended Mind Transformers, a variety of decoder-only transformers
closely related to Memorizing Transformers [wu2022memorizing] that retrieve
and attend to an external cache of key-value pairs (or memories) without
finetuning.

  * •

We demonstrate that both models trained with ALiBi and rotary position
embeddings can leverage retrieved information natively, and present the hyper-
parameter decisions that enable this. Namely, we show $(1)$ how to add
position information to retrieved key-value pairs and $(2)$ that the retrieved
information must be accessible to a majority of decoder layers.

  * •

We open-source a new counterfactual long-range retrieval dataset, and
demonstrate that Extended Mind Transformers outperform today’s SoTA models.

  * •

We demonstrate that Extended Mind Transformers are time-efficient. The upfront
cost of generating the cached key-values pairs is quickly amortized so our
method is a low cost performance boost for tasks that query the same documents
many times.

  * •

We demonstrate the benefit of causal citations and new active learning
generation paradigms enabled by Extended Mind Transformers.

Figure 4: Fact retrieval accuracy over various document lengths for Extended
Mind Llama-2-70b, RAG and state of the art baselines.

Figure 5: Retrieval accuracy over various document lengths for Extended Mind
Llama-2-7b, and long context baselines.

##  2 Methods

We introduce a simple kNN-based retrieval mechanism into the decoder-only
transformer architecture, just as in [wu2022memorizing]. This enables the
model to attend to a select number of external memories during each generation
step. The model decides for each token, within a particular decoder, which
memories are important. We describe the process of generating the external
memories, the addition of top-k attention [gupta2021memoryefficient] into the
attention module, and the important choices of which layers to augment with
this retrieval mechanism, as well as how to assign position information to
external memories.

###  2.1 Generating external memories

Given a set of long inputs, generating the external memories amounts to
passing all inputs through the model and caching its internal representations
of the data. This mechanism is already built into the model’s architecture to
speed up generation; we don’t need to recompute these values for previously
seen tokens during the generation process. We compute these representations
using a stride of length $s$ so that each token is conditioned on at least $s$
tokens. Figure 6 shows how we can balance precision and efficiency by tuning
the stride length. Smaller strides generate higher-quality representations,
while larger strides require fewer computations. We generate the cache once,
and documents may be processed in parallel. The upfront cost for each document
is $n$ passes through the model with $L$ number of input tokens, where $n$ is
the smallest integer such that $(sn)+L\geq$ document length.

The cost of attention at inference time (over all documents) is
$O(d(n_{q}(n_{q}+k))+n_{q}*RT(n_{D},k))$ where $n_{q}$ is the length of the
query, $n_{D}$ is the length of the long inputs, $k$ is the number of
retrieved memories (per query token), $d$ is the dimension of the
representation, and $RT$ is a choice of $topk$ retrieval function (could be an
approximate or exact method). Classic self-attention over all documents costs
$O(d(n_{d}+n_{q})^{2})$ for the same inference query. When $D$ is large, we
realize time-savings which is visualized in an inference times experiment 6.

###  2.2 Extended Mind Attention

The retrieval-augmentation is quite straightforward. In addition to self-
attention, we introduce top-k attention [gupta2021memoryefficient] over the
cached key-value pairs. I.e., we allow each query token to attend to a certain
number of external key-value pairs where the cosine similarity between the
query and key is largest. We use cosine similarity rather than inner product
to retrieve key-value pairs to mitigate distribution shift, following
[wu2022memorizing]. Explicitly, the retrieval-augmented attention mechanism
can be described by the following equations:

$\displaystyle\operatorname{scores}=\frac{K_{R}}{\mid\mid
K_{R}\mid\mid}\times\frac{Q}{\mid\mid Q\mid\mid}\hskip 14.22636pt,\hskip
14.22636pt\break\operatorname{softmax}\left(\frac{Q(K_{M}\oplus
K_{L})^{T}}{\sqrt{d}}\right)\times\left(V_{M}\oplus V_{L}\right)$ (1)

where $\oplus$ refers to tensor concatenation, $(K_{L},V_{L})$ are key-value
pairs from local context, $(K_{R},V_{R})$ are all key-value pairs from
external memories, and $(K_{M},V_{M})$ are those $k$ key-value pairs with
largest score. As usual, $Q$ are queries, and $d$ is head dimension. Figure 1
provides a visual aid for how self-attention is combined with top-k attention
over external memories.

We mask the attention weights such that each query token can only attend to
its own retrieved keys, and not those retrieved by previous or following query
tokens. This could enable even cheaper computations, but we do not implement
it in our code. We find this masking does not diminish performance as compared
to causal masking methods. To further speed up the retrieval computations we
can use a vector database designed exactly for this purpose. We support using
FAISS [douze2024faiss] in our implementation.

###  2.3 Augmenting a subset of layers

Although previous work [wu2022memorizing] suggests using retrieval-
augmentation methods on a subset or single decoder layer, we found it crucial
to use our top-k augmentation on a majority of decoder layers. As proposed by
Clark and Chalmers in the influential work “The Extended Mind” [emh], the
external information must be constantly and immediately accessible in order to
successfully contribute to the model’s memory. We name our class of models
after this pioneering work, and provide evidence for this claim in the
experiments we describe below. In the experiments that follow, we show results
in the Appendix 6.1 for Extended Mind Transformers that use memory-augmented
attention on only the last half or third of the decoders. These results,
especially for retrieval tasks, are quite poor. Although not explicitly
discussed in [wu2022memorizing], we believe it’s likely the cost of fine-
tuning was a main barrier to using more than one memory-augmented layer in
previous work.

###  2.4 Position Information

As hypothesized, we find that both models trained with ALiBi and rotary
position embeddings can use retrieved key-value pairs natively. The position
information assigned to those retrieved key-value pairs is dependent on the
position encoding used.

For models trained with ALiBi, such as MPT [MosaicML2023Introducing], we found
that placing the cache at a position just after the inputs was effective. In
this case, the model interprets the retrieved memories as being a constant
distance (one position) away from the tokens it considers local context. For
models trained using rotary position embeddings, such as Llama2
[touvron2023llama], we found that using no position information at all on
retrieved key-value pairs was most effective. Whereas all the local tokens are
rotated to encode their distance from other tokens, tokens retrieved from the
memory are not rotated at all.

###  2.5 Pruning Memories

Retrieving too many memories can cause generations to decrease in quality. We
implement two simple pruning techniques to keep generations high quality while
retrieving many memories to maximize recall. This can be seen as
regularization.

In the first method, we mask retrieved tokens that don’t meet some low
similarity threshold. We found $0.25$ to be a good choice. We implement this
for both model architectures we open source, but find it particularly
important for models trained with ALiBi. This can dampen retrieval
capabilities, however. Key-query similarity scores even for useful tokens can
be surprisingly low and even negative in models trained with rotary position
embeddings.

In contrast, we found our second pruning technique to be crucial for models
trained with rotary positions, and less so for those trained with ALiBi. This
method removes memories that correspond to special tokens; tokens that encode
the beginning or end of a sequence or an unknown token. We remove these tokens
after memories are generated, rather than removing the tokens from the
history. More detail on both regularization techniques can be found in the
Appendix 6.4, and each experiment’s corresponding Appendix section describes
which techniques were used.

##  3 Experiments

We perform three types of experiments: we run perplexity experiments to
measure how models trained with ALiBi and rotary position embeddings are able
to leverage the memory tokens during next-token prediction, we run
counterfactual retrieval experiments to measure the model’s ability to recall
and report facts in the external memory, and we run inference-time experiments
to measure the time saving of using Extended Mind Transformers over regular
attention. All code for reproducing experiment results can be found in the
attached code and will be open-sourced.

Figure 6: Average inference time for many queries over documents of length
4000 tokens.

###  3.1 Perplexity Experiments

We use perplexity as a metric for overall model performance. Perplexity is a
measure of uncertainty over each generated token, related to the cross-entropy
loss function. We batch documents into sequences of increasing input lengths,
and compute perplexity using a stride of $512$ tokens. We use the Wikitext-103
dataset [merity2016pointer]. Detail on batching can be found in the Appendix
6.1.

We provide multiple baselines to compare against Extended Mind Transformers.
The naive baseline sees all inputs in context, regardless of their length. The
truncate baseline only sees the last $N$ tokens in context, where $N$ is the
sequence length the model saw during training. Extended Mind Transformers see
the last $N$ tokens in context (the same as the truncate baseline), but also
attend to the previous tokens as external memories. We only compute perplexity
on the last $N$ tokens for all methods. We run the experiment for two
7-billion parameter models, MPT-7b [MosaicML2023Introducing] and Llama-2-7b
[touvron2023llama]. In the case of Llama-2-7b, we compute naive baseline
methods both with and without interpolated position embeddings. Results for
these experiments are shown in Figures 3, 3 and above.

Table 1: Perplexity Results for a sample of input lengths. Model | Input length: | 2560 | 4608 | 6656 | 8704
---|---|---|---|---|---
Extended Mind Llama2-7B (k=3) |  | 6.091 | 6.031 | 6.062 | 6.143
Llama2-7B (truncate) |  | 6.207 | 6.286 | 6.383 | 6.477
Llama2-7B (naive) |  | 5.100 | 10.052 | 49.657 | n/a
Llama2-7B (naive, interpolated) |  | 6.245 | 8.137 | 9.807 | n/a
Extended Mind MPT-7B (k=8) |  | 7.134 | 7.199 | 7.376 | 7.444
MPT-7B (truncate) |  | 7.252 | 7.347 | 7.560 | 7.636
MPT-7B (naive) |  | 7.233 | 23.273 | n/a | n/a

The naive baseline demonstrates the phenomenon methods to extend the input
token sequence length seek to ameliorate: after sequences exceed lengths
greater than 2-3k tokens, the performance quickly drops off. In this case,
perplexity quickly blows up. Using interpolated position embeddings staves off
the drop temporarily, but not past input sequence lengths greater than around
6k.

Both Extended Mind Transformers, MPT-7b, and Llama-2-7b are able to leverage
retrieved key-value pairs without any fine-tuning, providing strong evidence
that a variety of relative positional encodings enable this capability.
Extended Mind Transformers outperform both baselines. Even more, perplexity
continues to decrease as we increase the number of retrieved key-value pairs
per query token for models trained with ALiBi. We also provide results in the
Appendix 6.1 for Extended Mind Llama-2-7b where only a subset of decoder
layers retrieve key-value pairs from the cache. These results are
significantly worse than retrieving on all decoder layers.

###  3.2 Counterfactual Retrieval Experiments

We propose that Extended Mind Transformers address both the issue of extending
the maximum input token sequence length (the memory cache is unlimited) and
integrating retrieval into generation. Thus, we measure the quality of the
retrieval augmented attention in Extended Mind Transformers with that of pre-
trained models, fine-tuned models, and composite RAG methods. Extended Mind
Transformers are just as compatible with RAG as other models, and we report
these results as well. We found MPT models to be much less competitive on this
benchmark, and omit performance metrics.

We proxy the quality of attention with a challenging counterfactual retrieval
task. Our dataset is a modified wikiQA benchmark [abacusLongContext], which we
provide along with our code. The dataset is composed of Wikipedia articles (of
2-16 thousand tokens) and corresponding questions. We modify the dataset by
changing the labeled answers to realistic but wrong answers, to control for
facts memorized during pre-training. For example, we replace every instance of
“Lee Hazlewood” with “Terry Allen” in the Wikipedia entry for the song “These
Boots Were Made For Walking”, and then ask the model to produce the
songwriter’s name, with the correct answer now being “Terry Allen”. Details
for dataset creation can be found in the Appendix 6.2.

We run two sizes of Extended Mind Transformers on our new benchmark, based on
the 7-billion and 70-billion parameter Chat Llama-2 [touvron2023llama] models.
No Extended Mind Transformers have been further fine-tuned. For comparators to
the 7-billion parameter model we run Chat Llama-2-7b with interpolated
position embeddings (with and without RAG), as well as two base Llama-2-7b
models that have been fine-tuned on longer input sequences. “Yarn-
Llama-2-7b-64k” [peng2023yarn] (also instruct fine-tuned) and
“LLaMA-2-7B-32K”[together] were developed by Nous Research and Together
Computer, respectively. As a comparator to the 70-billion parameter model, we
test the popular closed-source model GPT-4 [openai2024gpt4]. Details on the
RAG implementation can be found in the Appendix 6.2.

All comparator models see the entire prompt and document in context. RAG
methods select the best $\approx 2500$ tokens from the document to include in
context. Extended Mind Transformers see only the prompt and question in
context, and attend to the document only through the retrieval-augmented
attention. When combined with RAG, they see the best $\approx 2500$ tokens in
context, and attend to the document in memory. Details on prompting can be
found in the Appendix 6.2, as well as results for Extended Mind Chat
Llama-2-7b where only some decoders retrieve key-value pairs from the cache.

Table 2: Retrieval results for all models. Model |  | 2k | 4k | 8k | 16k
---|---|---|---|---|---
Extended Mind Chat Llama2-7b |  | 38.59% | 37.95% | 36.00 % | 33.54%
Chat Llama2 |  | 42.74% | 41.96% | 33.50% | 17.68%
Nous |  | 34.85% | 35.71% | 37.50% | 33.54%
Together AI |  | 37.34% | 40.63% | 37.00% | 36.59%
Extended Mind Chat Llama2-70b + RAG |  | 59.34% | 53.13% | 53.00% | 50.00%
Chat Llama2-70b + RAG |  | 50.21% | 47.32% | 44.50% | 45.12%
GPT-4 |  | 52.28% | 50.00% | 46.00% | 43.29%

Figure 5 shows the overall retrieval accuracies of Extended Chat Llama-2-7b as
compared to other methods for extending input length. Extended Mind Chat Llama
outperforms both fine-tuned models on short inputs, and the baseline model on
long inputs, competitive with the fine-tuned models. Figure 5 shows the
overall retrieval accuracies of Extended Chat Llama-2-70b as compared to other
retrieval methods. While neither RAG nor our top-k retrieval alone are enough
to beat GPT-4 across the board, Extended Mind Chat Llama combined with RAG
improves upon GPT-4 by a large margin, $6\%$ when averaged across input
lengths. We also include more detailed heatmap visuals in the appendix 6.2
that introduce the added dimension of how many times a retrieved fact appeared
in the document.

###  3.3 Inference Times

We demonstrate the time-efficiency of Extended Mind Transformers by plotting
average inference time for queries over documents of 4k tokens. Extended Mind
Transformers incur an upfront cost to generate the memories, and then need
only compute self-attention on the query and retrieved memories. We compare
this to both a naive and more sophisticated benchmark. The naive method (this
is the only option for closed-source models) puts the document into context
alongside the query, for all queries. Hence the constant inference time, as
shown in Figure 6. As a more sophisticated comparison, we cache the key-value
pairs generated for the document after the first query and avoid re-computing
them each successive query.

We show performance for Extended Mind Transformers that use a stride length
equal to both 512 (default) and 1024 (recommended for efficiency) when
generating memories. Further detail on hyper-parameters can be found in the
Appendix 6.3. While Extended-Mind models do incur the additional upfront cost,
this is quickly amortized over additional queries, making our performance
improvement a low cost improvement.

##  4 Interpretability and Impact to Reasoning

In this final section, we discuss a new kind of citation and active learning
generation technique enabled by Extended Mind Transformers.

Figure 7: Highlighted tokens are those most retrieved memories when generating
the token “1971”.

###  4.1 Citations

Beyond performance improvements, Extended Mind Transformers also enable a new
kind of causal citations. Local attention weights are difficult to interpret
[jain2019attention] due to their scale. Simply reporting similar in-context
information or which information from external documents was included in
context (for instance, as retrieved by RAG) has no guaranteed causal relation
to what information was used during generation. Instead, we can now easily
report which memories were retrieved and attended to at each step of the
generation. This gives us a much better idea of which information the model
used to generate each next token. We enable this in our open-source models,
returning the indices of retrieved memories along with model outputs. Figure 7
demonstrates an example of this. Since we retrieve key-value pairs within each
layer, we make some standard aggregation choices that are detailed further in
the Appendix 6.5.

###  4.2 Uncertainty and Active Learning

Extended Mind Transformers also enable new insight into when a model is
uncertain about its generation. When a model is uncertain about its
generation, or equivalently it didn’t memorize the required information during
training, its generations evolve as we allow the model to retrieve more
information from the cache. Conversely, if the model was sure of its answer
with a small amount of retrieved information, adding new memories will not
change its output. We provide examples in the Appendix 6.6.

Figure 8: Example of using active learning generation with Extended Mind Chat
Llama-2.

More than this, Extended Mind Transformers suggest a new way to reduce
hallucinations. Using simple uncertainty metrics, i.e. token level entropy, we
can detect when the model is uncertain about its generated token, and allow
the model to regenerate the step using more information from the cache. We
leave full-scale experiments for this technique to future papers, but provide
a demonstrative example in Figure 4.2, and further details in 6.6. We suspect
this active learning aware generation will also improve given more robust
uncertainty metrics. We are particularly interested in using second-order
uncertainty from language models that have been fine-tuned to be Bayesian.

##  5 Conclusion

We provide an overview of the history of Memory augmented neural networks, and
situate our new method, Extended Mind Transformers in this rich history. We
propose a new variety of transformer models that retrieve and attend to
external documents without finetuning, based on the work in
[wu2022memorizing]. Extended Mind Transformers retrieve cached key-value pairs
within each decoder layer, for each generated token. We demonstrate that
relative positional embeddings allow today’s pretrained LLMs to use retrieved
key-value pairs without fine-tuning, eliminating the previous barriers to
using retrieval mechanisms internal to language models. We also demonstrate
the importance of the external information being constantly and immediately
accessible by comparing with models which only use retrieval on a subset of
decoder layers.

We open-source a new counterfactual long-range retrieval benchmark, and show
that Extended Mind Transformers are competitive with models that have been
fine-tuned on long inputs, despite not being fine-tuned, and provide new SoTA
results when combined with RAG. We measure and report the time savings of
using Extended Mind Transformers when making multiple queries over long
documents. Finally, we discuss the new kind of citations enabled by Extended
Mind Transformers, as well as the potential for new active learning inspired
generation techniques.

##  6 Appendix

All code and instructions for reproducing experiment results can be found in
the accompanying code.

###  6.1 Experiment Details: Perplexity

Resources used: (2x) A100-80GB GPUs, $3-24$ hours, depending on method and
parameters.

Additional Experiment Settings:

  1. 1.

All models use $N=2048$. I.e. we cut off (or cache) inputs longer than 2048
tokens for truncate (memorizing) runs.

  2. 2.

Llama baseline with interpolation uses dynamic position interpolation with an
$\alpha=8$, which is standard.

  3. 3.

We do not use similarity masking 6.4 in any perplexity experiments.

  4. 4.

We do remove memories associated to unknown and special tokens 6.4 for Llama
models.

  5. 5.

The below graphic illustrates how different methods handle input sequences,
batched documents of subsequences of increasing lengths. For simplicity, the
graphic shows a stride equal to the input length.

Additional Results Perplexity results for models with retrieval on only a
subset (the last third and half) of decoder layers:

Model | Input length: | 2560 | 4608 | 6144
---|---|---|---|---
Extended Mind Llama2-7B (retrieval on last third) |  | 6.181 | 6.212 | 6.271
Extended Mind Llama2-7B (retrieval on last third) |  | 6.162 | 6.167 | 6.212

###  6.2 Experiment Details: Counterfactual Retrieval

Resources used: (2-4x) A100-80GB GPUs, $3-24$ hours, depending on method and
parameters.

Heatmap visuals The below figures illustrate the percent of correctly
retrieved facts by bucketed number of fact appearances. Each cell represents
samples where the number of fact appearances falls within the range specified
on the x-axis, and is colored to represent the overall success of the
retrievals. The results are further stratified by sequence length, on the
y-axis. Since we did not create the dataset from scratch, this isn’t
consistent from document to document. Intuitively, if the fact appears many
times, it should be easier to retrieve.

Figure 9: Retrieval results Base Model, Fine-tuned Models, and Extended Mind
Transformer. Cells represent a bucket (x-axis), and sequence length (y-axis).
Colors to represent the overall success of the retrievals.

Figure 10: Retrieval results Base Model, Fine-tuned Models, and Extended Mind
Transformer. Cells represent a bucket (x-axis), and sequence length (y-axis).
Colors to represent the overall success of the retrievals.

Dataset Creation:

Substitutions were made in a few ways.

  1. 1.

Substitution of an entity with a semantically similar entity. E.g. England vs.
France, or Hank Aaron vs Joe Adcock.

  2. 2.

Substitution of an entity with a syntactically similar entity. E.g. Beatrix vs
Beatrice.

  3. 3.

Number / Date substitution. E.g. 1990 vs 1999, or 30 ml vs. 100 ml.

Additional Experiment Settings:

  1. 1.

Chat models are prompted using the appropriate assistant syntax, and models
that haven’t been instruction fine-tuned are given a one-shot example.
Extended Mind Transformers are prompted using one-shot example, no matter the
base model. This is all included in the code.

  2. 2.

Extended Mind Llama-2-7b uses a $k$ value of 12; each query token is allowed
to retrieve 12 keys from memory. Extended Mind Llama-2-70b uses a $k$ value of
4.

  3. 3.

We use a stride length of 512 when generating the memory cache, and allow 3072
input tokens. While the models saw sequences of 2048 tokens during training,
perplexity experiments confirm that input lengths up to 4096 are
comprehensible.

  4. 4.

We implement a simple RAG mechanism. We chunk the documents into token
sequences of length 500 with no overlap, and embed them using OpenAI’s Ada
model. We then compare similarity to the embedded query, and retrieve the
top-5 chunks (about 2500 tokens).

  5. 5.

There are two natural ways to combine RAG with Extended Mind Transformers.
$(1)$ Put whatever cannot fit in context into memory, or $(2)$ Put everything
into memory. While we observe benefit from both, the results reported reflect
the second strategy.

  6. 6.

We do not use similarity masking 6.4 in any retrieval experiments.

  7. 7.

Baseline models use dynamic position interpolation with $\alpha=8$.

  8. 8.

We do remove memories associated to unknown and special tokens 6.4 for
Extended Mind Llama models.

Additional Results Perplexity results for models with retrieval on only a
subset (the last third and half) of decoder layers. Results are total number
correct, rather than percentage, since recall is so low. Datapoint counts for
each split are listed alongside.

Model |  | 2k | 4k | 8k | 16k
---|---|---|---|---|---
Extended Mind Llama2-7B (retrieval on last third) |  | 5 | 9 | 5 | 6
Extended Mind Llama2-7B (retrieval on last half) |  | 5 | 7 | 5 | 4

Datapoint counts: 241 (2k), 224 (4k), 200 (8k), 164 (16k)

###  6.3 Experiment Details: Inference Times

Resources used: (2x) A100-80GB GPUs, $1-3$ hours, depending on method and
parameters.

Additional Experiment Settings:

  1. 1.

We use base Llama-2-7b for these experiments.

  2. 2.

We allow inputs up to 4096 tokens when generating the memories. While the
model saw sequence of length 2048 during training, the model can comprehend
sequences up to twice that length, as shown in our perplexity experiments.

###  6.4 Further Discussion: Pruning Memories

When to use similarity masking: We found similarity masking useful for both
models trained with ALiBi and rotary position embeddings when doing long form
text generation. However, we did not find it necessary on short retrieval
tasks for the Llama models.

When to remove special tokens: We found it crucial to remove special tokens
when retrieving memories from raw web documents, especially those with HTML
remnants. This prevents unknown tokens from being retrieved and confusing the
model. While removing these tokens does not appear to have much downside, we
did not find it as necessary for models trained with ALiBi.

###  6.5 Further Discussion: Citations

Since we retrieve key-value pairs within each layer, we simply report those
tokens which correspond to the key value pairs which were most retrieved
throughout the entire model.

It would be interesting to do a further analysis on which tokens are retrieved
by layers earlier / later in the decoder stack, we leave this to future work.

###  6.6 Further Discussion: Active Learning

Figure 4.2 uses Chat Llama-2-7b. The original generation retrieves using a
constant $k=1$, and the Active learning generation uses a baseline constant
$k=1$, and allows tokens with an entropy above $0.60$ to regenerate using an
additional $2$ memories, $k=3$ in total.

Evolution of uncertain and certain generations. This example uses MPT-7b with
default parameters, and $k$ specified as below. All settings can be found in
the code, and are enumerated below as well.

Experiment Settings

  1. 1.

Model used is MPT-7b [MosaicML2023Introducing].

  2. 2.

We do use similarity masking, with a threshold of $0.25$.

  3. 3.

We use $k$ as specified below.

Uncertain Example
Prompt: “When did Alexander Grothendieck get his French citizenship?”
Baseline Generation: “When did Alexander Grothendieck get his French
citizenship? I am trying to find out when Alexander Grothendieck got his
French citizenship. I know that he was born in Germany and that he got his
French citizenship in the late 1950s. I am trying to find out when he”
Generation for k=5: “A: I think he got it in the early 1960s”
Generation for k=6: “A: I think he got it in the early 1970s.”
Generation for k=7: “A: He was born in France, and he was naturalized in 1971.
<|endoftext|>”
Generation for k=8: “A: I think he got it in 1971. <|endoftext|>Q”
Generation for k=30: “A: He was born in Germany, and became a French citizen
in 1971.”

Certain Example
Prompt: “What was did Alexander Grothendieck’s profession?”
Baseline Generation: “What was did Alexander Grothendieck’s profession?
Alexander Grothendieck was a French mathematician”
Generation for k=2: “Alexander Grothendieck was a mathematician.”
Generation for k=8: “A: He was a mathematician. <|endoftext|>Q: What”


* [149] C. Zhou, H. Zhao, J. Han, C. Xu, C. Xu, T. Huang, and B. Shi, “Unmodnet: Learning to unwrap a modulo image for high dynamic range imaging,” _NeurIPS_ , vol. 33, 2020.
  * [150] Y. Huang, S. Qiu, C. Wang, and C. Li, “Learning representations for high-dynamic-range image color transfer in a self-supervised way,” _TM_ , vol. 23, pp. 176–188, 2020.
  * [151] H. Xu, J. Ma, J. Jiang, X. Guo, and H. Ling, “U2fusion: A unified unsupervised image fusion network,” _TPAMI_ , 2020.
  * [152] H. Jung, Y. Kim, H. Jang, N. Ha, and K. Sohn, “Unsupervised deep image fusion with structure tensor representations,” _TIP_ , vol. 29, pp. 3845–3858, 2020\.
  * [153] R. Li, C. Wang, S. Liu, J. Wang, G. Liu, and B. Zeng, “Uphdr-gan: Generative adversarial network for high dynamic range imaging with unpaired data,” _arXiv preprint_ , 2021.
  * [154] A. Jaiswal, A. R. Babu, M. Z. Zadeh, D. Banerjee, and F. Makedon, “A survey on contrastive self-supervised learning,” _Technologies_ , vol. 9, no. 1, p. 2, 2021.
  * [155] M. Long, H. Zhu, J. Wang, and M. I. Jordan, “Unsupervised domain adaptation with residual transfer networks,” _arXiv preprint_ , 2016.
  * [156] L. Wang and K.-J. Yoon, “Psat-gan: Efficient adversarial attacks against holistic scene understanding,” _IEEE Transactions on Image Processing_ , vol. 30, pp. 7541–7553, 2021.
  * [157] V. A. Patel, P. Shah, and S. Raman, “A generative adversarial network for tone mapping hdr images,” in _NCVPRIPG_. Springer, 2017, pp. 220–231.
  * [158] A. Rana, P. Singh, G. Valenzise, F. Dufaux, N. Komodakis, and A. Smolic, “Deep tone mapping operator for high dynamic range images,” _IEEE TIP_ , vol. 29, pp. 1285–1298, 2019.
  * [159] M. Mirza and S. Osindero, “Conditional generative adversarial nets,” _arXiv preprint arXiv:1411.1784_ , 2014.
  * [160] S. Lee, S. Y. Jo, G. H. An, and S.-J. Kang, “Learning to generate multi-exposure stacks with cycle consistency for high dynamic range imaging,” _TM_ , 2020.
  * [161] A. Jolicoeur-Martineau, “The relativistic discriminator: a key element missing from standard gan,” _ICLR_ , 2018.
  * [162] S. W. Park and J. Kwon, “Sphere generative adversarial network based on geometric moment matching,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2019, pp. 4292–4301.
  * [163] P. Isola, J.-Y. Zhu, T. Zhou, and A. A. Efros, “Image-to-image translation with conditional adversarial networks,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2017, pp. 1125–1134.
  * [164] J. Wang, W. Zhou, G.-J. Qi, Z. Fu, Q. Tian, and H. Li, “Transformation gan for unsupervised image synthesis and representation learning,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2020, pp. 472–481.
  * [165] M. Alghamdi, Q. Fu, A. Thabet, and W. Heidrich, “Transfer deep learning for reconfigurable snapshot hdr imaging using coded masks,” in _CGF_. Wiley Online Library, 2021.
  * [166] A. Krizhevsky, I. Sutskever, and G. E. Hinton, “Imagenet classification with deep convolutional neural networks,” _NIPS_ , vol. 25, pp. 1097–1105, 2012\.
  * [167] Y. Zhang, T. Xiang, T. M. Hospedales, and H. Lu, “Deep mutual learning,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2018, pp. 4320–4328.
  * [168] C. Finn, P. Abbeel, and S. Levine, “Model-agnostic meta-learning for fast adaptation of deep networks,” in _ICML_. PMLR, 2017, pp. 1126–1135.
  * [169] J. W. Soh, S. Cho, and N. I. Cho, “Meta-transfer learning for zero-shot super-resolution,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2020, pp. 3516–3525.
  * [170] Q. Sun, Y. Liu, T. Chua, and B. Schiele, “Meta-transfer learning for few-shot learning,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , June 2019.
  * [171] J. Li and P. Fang, “Hdrnet: Single-image-based hdr reconstruction using channel attention cnn,” ser. ICMSSP 2019. ACM, 2019, p. 119–124.
  * [172] Y. Deng, Q. Liu, and T. Ikenaga, “Multi-scale contextual attention based hdr reconstruction of dynamic scenes,” in _ICDIP_ , vol. 11519, 2020, p. 115191F.
  * [173] Z. Pan, M. Yu, G. Jiang, H. Xu, Z. Peng, and F. Chen, “Multi-exposure high dynamic range imaging with informative content enhanced network,” _Neurocomputing_ , vol. 386, pp. 147–164, 2020.
  * [174] R. Yu, W. Liu, Y. Zhang, Z. Qu, D. Zhao, and B. Zhang, “Deepexposure: Learning to expose photos with asynchronously reinforced adversarial learning,” in _NIPS_ , 2018, pp. 2153–2163.
  * [175] J.-G. Wang and L.-B. Zhou, “Traffic light recognition with high dynamic range imaging and deep learning,” _IEEE Transactions on Intelligent Transportation Systems_ , vol. 20, no. 4, pp. 1341–1352, 2018.
  * [176] X. Wang, Z. Sun, Q. Zhang, Y. Fang, L. Ma, S. Wang, and S. Kwong, “Multi-exposure decomposition-fusion model for high dynamic range image saliency detection,” _TCSVT_ , vol. 30, no. 12, pp. 4409–4420, 2020.
  * [177] E. Onzon, F. Mannan, and F. Heide, “Neural auto-exposure for high-dynamic range object detection,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2021, pp. 7710–7720.
  * [178] P. Satilmis, T. Bashford-Rogers, A. Chalmers, and K. Debattista, “Per-pixel classification of clouds from whole sky hdr images,” _SPIC_ , vol. 88, p. 115950, 2020.
  * [179] S. Im, H.-G. Jeon, and I. S. Kweon, “Robust depth estimation using auto-exposure bracketing,” _TIP_ , vol. 28, no. 5, pp. 2451–2464, 2018.
  * [180] R. Wang, Q. Zhang, C.-W. Fu, X. Shen, W.-S. Zheng, and J. Jia, “Underexposed photo enhancement using deep illumination estimation,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , June 2019\.
  * [181] J. Zhang, K. Sunkavalli, Y. Hold-Geoffroy, S. Hadap, J. Eisenman, and J.-F. Lalonde, “All-weather deep outdoor lighting estimation,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , June 2019.
  * [182] Y. Hold-Geoffroy, K. Sunkavalli, S. Hadap, E. Gambaretto, and J.-F. Lalonde, “Deep outdoor illumination estimation,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , July 2017.
  * [183] Y. Hold-Geoffroy, A. Athawale, and J.-F. Lalonde, “Deep sky modeling for single image outdoor lighting estimation,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2019, pp. 6927–6935.
  * [184] J. Hu, O. Gallo, K. Pulli, and X. Sun, “Hdr deghosting: How to deal with saturation?” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2013, pp. 1163–1170.
  * [185] T.-H. Oh, J.-Y. Lee, Y.-W. Tai, and I. S. Kweon, “Robust high dynamic range imaging by rank minimization,” _TPAMI_ , 2014.
  * [186] A. O. Akyüz, R. Fleming, B. E. Riecke, E. Reinhard, and H. H. Bülthoff, “Do hdr displays support ldr content? a psychophysical evaluation,” _TOG_.
  * [187] Y. Huo, F. Yang, L. Dong, and V. Brost, “Physiological inverse tone mapping based on retina response,” _The Visual Computer_.
  * [188] R. P. Kovaleski and M. M. Oliveira, “High-quality reverse tone mapping for a wide range of exposures,” in _SIBGRAPI_. IEEE.
  * [189] B. Masia, A. Serrano, and D. Gutierrez, “Dynamic range expansion based on image statistics,” _MTA_ , 2017.
  * [190] Y. Chen, X. Dai, M. Liu, D. Chen, L. Yuan, and Z. Liu, “Dynamic convolution: Attention over convolution kernels,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2020, pp. 11 030–11 039.
  * [191] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton, “A simple framework for contrastive learning of visual representations,” in _ICML_. PMLR, 2020, pp. 1597–1607.
  * [192] J. Yang, R. Xu, R. Li, X. Qi, X. Shen, G. Li, and L. Lin, “An adversarial perturbation oriented domain adaptation approach for semantic segmentation,” in _AAAI_ , vol. 34, no. 07, 2020, pp. 12 613–12 620.
  * [193] H. Chen, Y. Wang, C. Xu, Z. Yang, C. Liu, B. Shi, C. Xu, C. Xu, and Q. Tian, “Data-free learning of student networks,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2019, pp. 3514–3522.
  * [194] Q. Sun, Y. Liu, T.-S. Chua, and B. Schiele, “Meta-transfer learning for few-shot learning,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , June 2019.
  * [195] Z. Liu, M. Sun, T. Zhou, G. Huang, and T. Darrell, “Rethinking the value of network pruning,” _arXiv preprint_ , 2018.
  * [196] M. Sandler, A. Howard, M. Zhu, A. Zhmoginov, and L.-C. Chen, “Mobilenetv2: Inverted residuals and linear bottlenecks,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2018, pp. 4510–4520.
  * [197] A. R. Zamir, A. Sax, N. Cheerla, R. Suri, Z. Cao, J. Malik, and L. J. Guibas, “Robust learning through cross-task consistency,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2020, pp. 11 197–11 206.
  * [198] H. R. V. Joze, A. Shaban, M. L. Iuzzolino, and K. Koishida, “Mmtm: multimodal transfer module for cnn fusion,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2020, pp. 13 289–13 299.
  * [199] S. Zhou, J. Zhang, W. Zuo, and C. C. Loy, “Cross-scale internal graph neural network for image super-resolution,” _Neurips_ , 2020.


# Distributed Multi-robot Online Sampling with Budget Constraints

Azin Shamshirgaran1 Sandeep Manjanna2 Stefano Carpin1 1A. Shamshirgaran and S.
Carpin are with the Department of Computer Science and Engineering, University
of California, Merced, CA, USA. A. Shamshirgaran is supported by USDA-NIFA
under award # 2021-67022-33452 (National Robotics Initiative). S. Carpin is
partially supported by the IoT4Ag Engineering Research Center funded by the
National Science Foundation (NSF) under NSF Cooperative Agreement Number
EEC-1941529. Any opinions, findings, conclusions, or recommendations expressed
in this publication are those of the author(s) and do not necessarily reflect
the view of the U.S. Department of Agriculture or the National Science
Foundation.2S. Manjanna is at Plaksha University, Mohali, India.

###### Abstract

In multi-robot informative path planning the problem is to find a route for
each robot in a team to visit a set of locations that can provide the most
useful data to reconstruct an unknown scalar field. In the budgeted version,
each robot is subject to a travel budget limiting the distance it can travel.
Our interest in this problem is motivated by applications in precision
agriculture, where robots are used to collect measurements to estimate domain-
relevant scalar parameters such as soil moisture or nitrates concentrations.
In this paper, we propose an online, distributed multi-robot sampling
algorithm based on Monte Carlo Tree Search (MCTS) where each robot iteratively
selects the next sampling location through communication with other robots and
considering its remaining budget.

We evaluate our proposed method for varying team sizes and in different
environments, and we compare our solution with four different baseline
methods. Our experiments show that our solution outperforms the baselines when
the budget is tight by collecting measurements leading to smaller
reconstruction errors.

##  I Introduction

There is an increasing number of applications for autonomous robots in
agriculture [22, 7], and while the most obvious interest may be in fruit
harvesting [4], there is also sustained demand for robots supporting data
collection at scale, especially for measuring and/or estimating scalar
parameters such as soil moisture and nitrate concentration that cannot be
easily determined through remote sensing with satellites and drones (see Fig.
1).

Figure 1: Our robot has been retrofit with a soil moisture sensor that can be
inserted in the soil to collect data at preassigned locations, or at locations
decided by the robot on-the-fly (the soil moisture probe is visible on the
left).

Robots are expected to play a vital role in the implementation of farm
monitoring systems in support of _precision agriculture_. Precision
agriculture is defined as “the matching of agronomic inputs and practices to
localized conditions within a field and the improvement of the accuracy of
their application” [6]. Key to this vision is the ability to perform scalable
data collection on demand. In this context, the capability of deploying a
coordinated team of robots to collect data is instrumental, as a team of
robots can collect more data per time unit, and also offers increased
robustness to individual failures.

For this approach to be efficient, it is necessary for robots to coordinate
their efforts to avoid unnecessary duplicate work or negative interferences.
As part of the measurement collection task, sample locations can be provided
in advance chosen by experts, or randomly, or selected along the way in
response to collected data. Central to our work is the necessity to perform
task allocation being aware of the distance a robot can travel before its
battery is depleted. From a practical standpoint, a robot running out of
energy in the middle of its data collection mission is a major problem, as it
will be necessary to manually recover it. Our task allocation strategy,
therefore, focuses both on avoiding duplicate work and on managing energy
constraints.

Communication among robots is another key dimension to be considered in multi-
robot scenarios [1] where in some works all robots can share limited amounts
of data with one another irrespective of the distance [3], whereas in other
approaches more data is exchanged, but only when robots are sufficiently close
to each other [11].

In this work we present a distributed algorithm based on communication between
robots to avoid having multiple robots collecting measurements at the same
locations, because according to our model multiple measurements at the same
site do not increase the quality of the estimate. Additionally, the proposed
algorithm manages exploration and exploitation for each agent by considering
the amount of remaining energy to prevent the robots from running out of power
before a preassigned final recharging location is reached. The main
contributions of this paper are 1) a budget aware, online distributed method
for multi-robot coordination based on Monte Carlo Tree Search (MCTS); 2) a
strategy to update sampling locations on the fly, to leverage date acquired
during the mission; 3) a thorough comparison with other methods on different
benchmark datasets to evaluate the quality of the proposed method.

The remainder of this paper is organized as follows. Selected related work is
presented in Section II. The problem formulation is given in Section III,
while our proposed method is presented in Section IV. Extensive simulation
results are discussed in Section V, and conclusions are offered in Section VI.

##  II Related Work

In this section, we provide pointers to selected works relevant to the problem
we present in this manuscript. In multi-robot sampling and monitoring,
sectioning and Voronoi partitioning are common approaches. In [9], the authors
propose a dynamic Voronoi approach, where robots repeatedly compute Voronoi
partitions and each robot performs sampling within its partition. Although
this method shares the exploration task efficiently between robots, the
iterative recomputation of Voronoi regions may lead to many unnecessary
motions that are problematic in our application where robots are subject to a
limited travel budget.

In [21], we studied multirobot routing in a vineyard using a formulation based
on the team orienteering problem. In the orienteering problem (OP) an agent is
required to traverse a graph in which each vertex has a predetermined reward
and each edge has a fixed cost. A path should be computed to maximize the sum
of collected rewards while ensuring that the sum of the costs of traversed
edges does not exceed a preassigned budget. Team Orienteering Problem (TOP) is
a version of OP in which a team of robots works together to collect rewards.
The main limitation of this line of research is that it assumes that rewards
are predetermined in advance. However, in many situations this is not
realistic, as data collected on the fly may lead to revised estimates about
the value of reaching a certain location.

Recently, MCTS has gained a great deal of attention for multi-robot
informative path planning. In [8, 2], the authors proposed a method combining
Gaussian processes (GP) and MCTS to solve the problem of environmental
monitoring. In [11, 12], a decentralized approach is proposed using a policy
gradient method for multirobot environmental monitoring and sampling. To
encourage robots to spread away from other robots, each robot has a reward
function dependent on its distance from the others. Furthermore, the method
considers a communication range between robots to exchange locations and the
history of previous locations with co-working robots. In [13], a distributed
adaptive sampling method for multi-agent scenarios was proposed which is
robust to the failure of robots and communication. To estimate the optimal
policy, deep neural networks and policy gradient methods are used. These
works, however, do not explicitly incorporate limits on the distance traveled
by robots.

In [15], we presented an offline path planning method based on Q-learning to
solve the sampling problem for a single robot in a stochastic environment
subject to a preassigned constraint on the distance it can travel. In [3], we
instead considered the problem of reconstructing a spatial field using
multiple robots, Gaussian processes, and MCTS. As part of this work, robots
communicate with one another and send their current location and observations
to other robots and explore the spatial properties of GPs to attempt to spread
the robots in the environment. In this paper, we extend this line of research
by adding the ability to select new sampling locations as well as considering
other robots’ actions during planning.

##  III Background and Problem Formulation

###  III-A Informative Path Planning

We start defining the multi-robot informative path planning (IPP) problem we
study in the following. Let $\mathcal{U}\subset\Re^{2}$ be the environment of
interest. By visiting a set of locations and collecting samples, our goal is
to estimate a scalar function $h:\mathcal{U}\rightarrow\mathbb{R}$ which
represents a parameter of interest (e.g., soil moisture). We assume there are
$\kappa$ robots in the team, each indicated as $R_{i}\mbox{ with
}i\in[1,2,3,...,\kappa]$. The location of each robot ($x,y$-position) is
defined by $s_{s}^{R_{i}}=(s_{x}^{R_{i}},s_{y}^{R_{i}})$. All robots start
from pre-assigned start positions $s_{init}^{R_{i}}$ (e.g., the points where
the robots are deployed), and must terminate their mission at assigned final
locations $s_{f}^{R_{i}}$ (e.g., where their batteries will either be swapped
or recharged). Each robot has a predetermined travel budget $B^{R_{i}}$
limiting the distance it can travel. This constraint models, for example, the
limited energy provided by the battery onboard the robot. There are $n$ sample
locations of interest111In agricultural applications these preassigned
locations are often identified a-priori by domain experts based on past
experience and/or data such as yield, plant stress, and the like. in
$\mathcal{U}$ denoted by the set $\mathcal{V}=\\{s_{1},s_{2},\dots,s_{n}\\}$.
The number of locations and their placement in $\mathcal{U}$ is such that no
robot has a sufficient budget to visit all of them. The goal of multi-robot
informative path planning (IPP) is to select a path for each robot ${R_{i}}$
to visit a subset of locations in $\mathcal{V}$ such that none of the robots
exceeds the travel budget and the quality of the collected information is
maximized. Although $\mathcal{V}$ is given in advance, to allow for greater
flexibility the set $\mathcal{V}$ can also be updated on the fly, i.e., new
sample locations can be determined by the robots as the mission unfolds. When
a robot reaches a location in $\mathcal{V}$ it uses its onboard sensor to
measure the parameter of interest, and this value is then used to estimate the
unknown function $h$. More formally, assuming $\rho^{R_{i}}$ is a path taken
by robot ${R_{i}}$ visiting a subset of sample locations of $\mathcal{V}$,
$f(\rho^{R_{i}})$ is a generic function measuring the quality of the estimate
of $h$ obtained from measurements collected at the locations visited along the
path $\rho^{R_{i}}$. Let $C(\rho^{R_{i}})$ be the travel cost associated with
traversing $\rho^{R_{i}}$ for robot ${R_{i}}$. The multi-robot informative
path planning problem (IPP) can then be expressed as the problem of solving
the following constrained optimization problem for each robot

$\rho_{\ast}^{R_{i}}=\operatorname*{argmax}_{\rho^{R_{i}}\in\psi^{R_{i}}}f(\rho^{R_{i}})\,\,\textrm{s.t.}\,C(\rho^{R_{i}})\leq
B^{R_{i}}$

where $\psi^{R_{i}}$ is the set of all paths from $s_{init}^{R_{i}}$ (start
location) to $s_{f}^{R_{i}}$ (final location) of robot ${R_{i}}$. In our
problem setting, for given path $\rho^{R_{i}}$, the cost $C(\rho^{R_{i}})$ is
not deterministic, but rather a random variable whose realization is obtained
only at run time. This models the fact that the energy or time spent to move
between two locations is in general stochastic, and we assume that the robot
has a probability density function describing such random variable. Hence,
while robot $R_{i}$ moves along $\rho^{R_{i}}$ from location to location, it
is necessary to monitor the energy spent to ensure it does not exceed
$B^{R_{i}}.$

###  III-B Gaussian Process Regression

We model the spatial distribution of the scalar field $h$ being estimated
using Gaussian Processes (GP). GPs are extensively used in geostatistics [16,
18] to model environmental parameters (in geostatistics this approach is known
as kriging.) We denote with $x_{g}^{R_{i}}$ the scalar reading collected by
robot ${R_{i}}$ when sampling location is $s_{g}\in\mathcal{V}$, and will
denote with $\chi_{g}^{R_{i}}$ the random variable modeling $x_{g}^{R_{i}}$
that follows a Gaussian distribution with mean $\mu_{g}$ and variance
$\sigma_{g}^{2}$. Using the data collected at the sample locations, a
posterior of $h$ can be estimated using standard GP regression algorithms (the
reader is referred to [14] for a comprehensive introduction to this topic,
including GP regression algorithms.)

###  III-C Monte Carlo tree search (MCTS)

MCTS is an online method for solving sequential, stochastic decision making
problems. MCTS builds a tree with a root node representing the current state
and edges connecting states that can be reached by executing a single action.
Nodes subsequently added to the tree represent states that can be reached
through a sequence of actions originating at the root. Each action is assigned
a $Q$-value representing _how good_ the action is, which is an estimate of the
value that will be obtained through a complete execution starting with that
action. Once the tree has been constructed, an action is selected from those
available at the root. Upon execution of the selected action, the tree is
discarded and rebuilt with the next state as its root. A basic version of MCTS
consists of the following four steps [19, 5]:

  * •

Selection: Using the so-called _tree policy_ , a path from the root to a leaf
node is selected.

  * •

Expansion: From the selected leaf node, one or more child nodes are added to
the tree.

  * •

Rollout: A complete episode is simulated from the selected leaf node, or from
one of its newly added child nodes (if any). During this simulation, a simple,
suboptimal policy is used to decide the actions.

  * •

Backup: Based on the return generated by the simulated episode, the action
values attached to the tree edges traversed by the tree policy are updated, or
initialized.

A critical component is the tree policy for action selection (“selection” step
in the list above). One popular criterion for action selection is the $UCB$
rule defined in Eq. (1) and first introduced [10] Each candidate action $a$ is
assigned a $UCB(a)$ value defined as

$UCB(a)=Q_{t}(a)+c\sqrt{\frac{\ln t}{N_{t}(a)}}$ (1)

and eventually the action with the highest $UCB$ value is selected for
execution. In Eq. (1), $Q_{t}(a)$ denotes the action value estimate,
$N_{t}(a)$ is the number of times that action $a$ has been selected prior to
time $t$, and $c$ is a constant controlling the exploration. Initially,
$N_{t}(a)$ is zero for all actions222$UCB(a)$ is assumed to be $\infty$ when
$N_{t}(a)=0$, thus forcing exploration.. Every time $a$ is selected, $t$ and
$N_{t}(a)$ increase, and every time $a$ is not selected, $t$ increases but not
$N_{t}(a)$, ensuring that all actions will eventually be selected, but actions
with lower value estimates or those that have already been selected frequently
will be selected less frequently. This criterion balances exploration and
exploitation, with the balance determined by the parameter $c.$

##  IV Proposed Algorithm

Starting from the problem formulation, we here describe our proposed method.
For each robot $R_{i}$, let $\mathcal{A}^{R_{i}}\subseteq\mathcal{V}$ be the
set containing the locations visited by robot $R_{i}$. Initially
$\mathcal{A}^{R_{i}}=\\{s_{init}^{R_{i}}\\}$ and, by definition of
$\mathcal{A}^{R_{i}},$ throughout the execution of the algorithm the current
location of the robot $s_{s}^{R_{i}}$ is one of the elements of
$\mathcal{A}^{R_{i}}$. Each set $\mathcal{A}^{R_{i}}$ is iteratively expanded
by the execution of an action $a$ representing a possible next sampling
location $s_{g}$ in $\mathcal{V}$ for robot ${R_{i}}$. The execution of $a$
implies that the robot will move to $s_{g}$ and collect a sample at that
location. Using MCTS, the goal is to select a good sequence of sample
locations, $\mathcal{A}^{R_{i}}$, for robot ${R_{i}}$ while considering the
travel budget, $B^{R_{i}}$, and other robots’ decisions. The meaning of good
depends on the choice of objective function(s) and will be discussed shortly.

In our proposed method each robot $R_{i}$ shares its visited locations with
other robots with the objective of avoiding having multiple robots visiting
the same locations. This will lead to more distinct samples being collected
and ultimately to a better estimate for $h$. Our communication model assumes
that robots can exchange limited information (such as locations) at long
range. This is in line with the current technology (e.g., LoRa [17]) used by
robots in agricultural applications and previous works [3, 11].

To select the next location visit, each robot uses a reward function defined
as follows. For each unvisited location $s_{g}$, robot $R_{i}$ defines a value
$r_{g}^{R_{i}}$. A possible candidate could be $r_{g}^{R_{i}}=\sigma_{g}^{2}$
where $\sigma_{g}^{2}$ is the variance of the posterior estimate of $h$
provided by GP regression based on the samples collected up to that moment.
With this choice, high reward values would be assigned to locations with high
uncertainty [3]. In our case, we scale the predicted variance by the distance
between the robot’s current location and $s_{g}$ location. The reward
$r_{g}^{R_{i}}$ associated with a potential sampling location $s_{g}$
considered by robot $R_{i}$ is therefore defined as

$r_{g}^{R_{i}}=\frac{\sigma^{2}_{g}}{d(s_{g},s_{s}^{R_{i}})}$ (2)

where $\sigma_{g}$ is variance of the candidate location
$s_{g}\in\mathcal{V}$, and $d(s_{g},s_{s}^{R_{i}})$ is the distance between
the current location of robot $R_{i}$ and $s_{g}$. By introducing the distance
into the reward function we bias the algorithm to favor closer locations if
the predicted variance of two candidates is the same. The reason for this
preference is that we have a limited budget. The selection of the next
location is performed online, i.e., the reward associated with each location
is not predetermined, but re-estimated iteratively based on the locations
already visited and data previously collected. The reward is then used to
calculate the expected return, or the action value estimate $Q_{t}(a)$. In our
work we deal with an _episodic_ task, i.e., the task always ends after a
finite amount of time, either because the robot reaches the final location, or
because it runs out of energy. In this case, it is typical to define
$Q_{t}(s_{g})$ as a function of reward sequence

$Q_{t}(s_{g})=r_{g}^{t}+\lambda
r_{g^{\prime}}^{t+1}+...+\lambda^{T-t}r_{g^{\prime\prime}}^{T};\,\,0\leq\lambda\leq
1$ (3)

where $\lambda$ is a factor discounting future rewards and $T$ is the time of
the last action. $g,g^{\prime},...,g^{\prime\prime}$ are the selected sample
locations and $r_{g}$ is the reward associated with the sampling location $g$.
Algorithm 1 shows the planning algorithm independently executed by each robot.

Algorithm 1 Online MCTS based planner for robot $R_{i}$ with resampling

1:Input: $\mathcal{V}$, $s_{init}^{R_{i}}$, $s_{f}^{R_{i}}$, $B^{R_{i}}$

2:$\mathcal{A}^{R_{i}}\leftarrow\\{s_{init}^{R_{i}}\\}$

3:$s_{s}^{R_{i}}\leftarrow s_{init}^{R_{i}}$

4:while $B^{R_{i}}>0$ and  $s_{s}^{R_{i}}\neq s_{f}^{R_{i}}$ do

5: $cand^{R_{i}}\leftarrow\Psi[s_{s}^{R_{i}}]\setminus\mathcal{A}^{R_{i}}$

6: $cand^{R_{i}}\leftarrow cand^{R_{i}}\setminus\mathcal{A}^{R_{j}}$ for all
$j\neq i$

7: $\mathcal{T},s_{g}\leftarrow$ MCTS($s_{s}^{R_{i}}$, $cand^{R_{i}}$)

8: Move to $s_{g}$, collect reading $x_{g}^{R_{i}}$, and measure consumed
energy $c_{s}^{g}$

9: $\sigma_{g}^{2}\leftarrow$ update GP with new observation $x_{g}^{R_{i}}$

10: $\mathcal{V}\leftarrow$ Resampling based on new $\sigma_{g}^{2}$

11: $B^{R_{i}}\leftarrow B^{R_{i}}-c_{s}^{g}$

12: $\mathcal{A}^{R_{i}}\leftarrow\mathcal{A}^{R_{i}}\cup\\{s_{g}\\}$

13: $s_{s}^{R_{i}}\leftarrow s_{g}$

14: Brodacast($s_{s}^{R_{i}}$)

15:end while

16:return $B^{R_{i}}$, $\mathcal{A}^{R_{i}}$

The algorithm takes as input the set of $\mathcal{V}$, the initial and final
locations $s_{init}^{R_{i}}$ and $s_{f}^{R_{i}}$, and the budget $B^{R_{i}}$
for the each robot.

A children map $\Psi[s_{g}]$ contains the locations that can be reached from
each location $s_{g}$. As we consider problem instances with tens or hundreds
of possible locations, considering all of them would lead to search trees with
extremely high branching factors. Therefore, to minimize planning time, we
limit the set of locations that are considered from each location, and this
set is returned by the function $\Psi$. More specifically, we limit the
branching factor to $M$ (an even number). For each location $s_{g}$, $M/2$
elements in $\Psi$ are the nearest elements in $\mathcal{V}$ and $M/2-1$ are
randomly chosen from the remaining locations. Moreover the final location,
$s_{f}^{R_{i}}$ is always added to $\Psi$. This selection balances global
exploration and local exploitation. The addition of $s_{f}^{R_{i}}$ to $\Psi$
ensures that from any location robot $R_{i}$ can always consider moving to the
final goal location. This is useful when the travel budget is about to expire.
$cand^{R_{i}}$ is the set of candidate locations for the current location and
is obtained by removing already visited locations $\mathcal{A}^{R_{i}}$ from
the set of reachable locations returned by $\Psi$.

The current location of the robot, $s_{s}^{R_{i}}$, is considered as a root
node of the MCTS tree (line 7 in Alg. 1). The MCTS is expanded for a fixed
number of iterations. Each time, the path and leaf are chosen using UCB, as
per Eq. (1). When a leaf is reached, a rollout is executed. During rollout,
the planner continues to select additional random locations until it either
reaches the final destination or runs out of energy. During the MCTS expansion
and rollout, every time a candidate location is included in the tree, a
generative model is used to estimate how much energy would be consumed. This
estimate is given by the formula

$c_{s}^{g}=\alpha d(s_{s}^{R_{i}},s_{g})+U(\Lambda)$ (4)

where $d(s_{s}^{R_{i}},s_{g})$ is the distance between the current location,
$s_{s}^{R_{i}}$ and candidate location, $s_{g}$, and $U(\Lambda)$ is a random
sample from a uniform distribution over the interval $[0,\Lambda]$.

After the tree $\mathcal{T}$ has been built and the next location,
$s_{g}^{R_{i}}$, is selected, the budget for robot $R_{i}$ is updated. Based
on the new sample reading, the GP is updated, and based on the updated GP, it
will generate a new set of random sample locations with the normalized
variance as probabilities associated with each location. In the event that one
of the robots reaches its final destination or runs out of energy, other
robots will continue their sampling task.

##  V Results and discussion

Methods: To assess the effectiveness of the proposed method (dubbed RMCTS in
the following), we compared it with four different alternatives. The first is
a non-coordinated MCTS method (NCMCTS), the second is MCTS without resampling
(MCST), the third is the Orienteering method (Or) [20] and the fourth is
Multirobot Planning for Informed Spatial Sampling (MRS) [11]. NCMCTS is the
same as our proposed method but without any data shared between robots. MCTS
is also the same as our proposed method, but without the resampling step (line
10 in the algorithm). The Or method builds a graph with all sample locations
and determines the path that collects the maximum reward without exceeding the
preassigned budget. In this case, it is necessary to assign a value to each
element of $\mathcal{V}$ in advance, consistent with the fact that in
orienteering one must know the rewards of the vertices beforehand. To make the
method comparable with ours, we assigned equal rewards to all vertices, as our
method does not require prior knowledge of the value of collecting a sample at
a given location. This way, Or will attempt to visit as many locations as
possible. MRS is a decentralized sampling approach where each robot in a team
performs an informed survey using a policy-gradient-based sampling strategy.
This method uses policy gradient search to directly optimize the policy
parameters $\theta$ based on simulated experiences. In its original
formulation, this method does not consider budget constraints during training.
To account for it, during testing the next location proposed by the algorithm
is rejected if there is not enough budget left to reach the designated
location and then the robot moves to the final location $s_{f}^{R_{i}}$. In
this case, the algorithm selects $s_{f}^{R_{i}}$ and tries to reach the final
location. The MRS algorithm and its implementation are described in greater
detail in [11].

Performance metrics: To assess the performance of the various algorithms, we
consider two metrics. The first is the mean square error between $\hat{h}$
(the estimate of $h$) and $h$ itself. In our implementation, GP regression is
computed using the _scikit-learn_ Python library and its GP regression module
using Mattérn Kernel with length scale of 1 and smoothness parameter of
$\nu=1.5$. The choice of the kernel and of the parameters was made after
having experimentally evaluated different alternatives and having assessed
that these are the best choices. The second metric is the remaining budget,
which is the amount of budget that has not been used when the mission
terminates.

###  V-A Synthetic data-set

We start evaluating our proposed method on a bidimensional grid of size
$30\times 30$ where the scalar field $h$ is defined as a mixture of Gaussian
distributions. The set $\mathcal{V}$ consists of 100 locations. Each robot
starts from (0, 0) (upper left corner), and the final location is located at
(30, 30) (lower right corner). We consider different numbers of robots (3 and
5) and different budgets (100 and 200). For each case, displayed data are
averages over 100 independent runs. In all simulations, the MCTS, NCMCTS, and
RMCTS algorithms add 1000 nodes to the tree. For the function $\Psi$ we set
$M=30$. In Eq. (1) we set $c=3$, in Eq. (3) we set $\lambda=1$, and in Eq. (4)
we set $\alpha=0.5$ and $\Lambda=1$ and we use Manhattan distance. In RMCTS,
the resampling process is applied every other iteration and after each
resampling, $|\mathcal{V}|=30$. The MRS method seeks to collect samples at
locations with high values for the function $h$ at earlier stages of the
exploration. During training, MRS runs 20 simulated trajectories to update and
learn the parameter of policy $\pi$, and in the test phase it uses that
learned policy to generate an explicit action plan.

Table I summarizes the metrics for MCTS, RMCTS, NCMCTS, and MRS (summary of
100 runs). The table displays the initial budget $B$, the number of robots in
a team $N_{R_{i}}$, the average MSE error and average remaining budget for
each robot $B_{re}$. The numerical comparison confirms the superiority of
RMCTS across the board. As the budget and number of robots increase, the
performance of NCMCTS becomes similar to MCTS. Indeed, with a greater number
of robots and a large budget, it is possible to visit numerous locations even
without coordination. However, this is not the case when the number of robots
is smaller or the budget is tight, and in such cases coordination is key. As a
result of resampling, RMCTS achieves a lower MSE because based on the updated
GP it generates a better set of sample locations to explore. MCTS and RMCTS do
not need any prior knowledge about the prior model of the environment whereas
MRS needs to know the prior model as a reward map. Also, MRS requires pre-
training, while MCTS and RMCTS are online and can choose the next locations on
the fly. To give an order of magnitude, training time for one robot in MRS is
about 17 minutes, while RMCTS and MCTS do not need training and planning time
for both is below 15 seconds (cumulative time for all planning stages
alternating with execution).

Fig. 2(a) and 2(b) show sample paths for 3 robots with $B=100$. With MRS,
robots focus on areas with high values for the underlying function being
reconstructed (warmer colors), while in RMCTS, robots visit all areas. As the
goal is to reconstruct the underlying function $h$ with low RMSE error, to do
that robots must visit both areas with high and low values.

$B$ | $N_{R_{i}}$ | method | MSE | $B_{re}$
---|---|---|---|---
100 | 3 | MCTS | 0.52 | 13.72
100 | 3 | RMCTS | 0.47 | 11.23
100 | 3 | NCMCTS | 1.27 | 14.1
100 | 3 | MRS | 1.39 | 12
100 | 5 | MCTS | 0.38 | 13.21
100 | 5 | RMCTS | 0.35 | 11.29
100 | 5 | NCMCTS | 0.81 | 15.04
100 | 5 | MRS | 1.19 | 6
200 | 3 | MCTS | 0.48 | 23.72
200 | 3 | RMCTS | 0.41 | 23.13
200 | 3 | NCMCTS | 0.54 | 27.11
200 | 3 | MRS | 1.13 | 34
200 | 5 | MCTS | 0.35 | 21.32
200 | 5 | RMCTS | 0.32 | 18.71
200 | 5 | NCMCTS | 0.49 | 20.41
200 | 5 | MRS | 0.65 | 20
TABLE I: Avg. Results for 100 runs for the synthetic data-set.

(a) RMCTS

(b) MRS

(c) RMCTS

(d) MRS

(e) RMCTS

(f) MRS

(g) RMCTS

(h) MRS

Figure 2:  Figures (a)-(b) show three-robots sampling paths with budget
$B=100$ in synthetic environment using RMCTS and MRS. Figures (c)-(d) show
three-robots sampling paths with budget $B=100$ in vineyard environment using
RMCTS and MRS. Figures (e)-(f) show three-robots sampling paths with budget
$B=200$ in vineyard environment using RMCTS and MRS. Figures (g)-(h) show
five-robots sampling paths with budget $B=100$ in vineyard environment using
RMCTS and MRS.

###  V-B Experimental data-set

Next, we test all methods using a data-set for soil moisture experimentally
collected in a commercial vineyard located in the California Central Valley.
In this case, we use 100 sample locations that are distributed throughout the
environment. The parameters are the same as Sec. V-A. Table II summarizes the
results for 100 runs of MCTS, RMCTS, Or, and MRS methods for one robot. It can
be seen that RMCTS outperforms other methods with a tight budget, while MRS
achieves better MSE with a higher budget. The Or method consumes almost all of
the budget, but it should be noted that for example in our implementation, the
planning time for Or (40.21) is more than five times greater than RMCTS (7.52
s) and MCTS (6.18 s). The fact that the MSE is not significantly better in the
Or method even with a higher budget is due to the fact that in orienteering,
one aims at collecting the maximum additive reward, and this can be achieved
by visiting many nearby locations that will lead to limited additional
information to better estimate the scalar field $h$.

Finally, table III summarizes results for 100 runs of MCTS, RMCTS, and MRS
methods for teams of three, five, and ten robots. Similar to other cases, we
also find that RMCTS outperforms other methods when the budget is tight, while
MRS outperforms other methods when the budget is higher. MRS performs better
with higher budgets since it considers the samples’ locations at each step
rather than our method, which considers samples at specific locations. Figure
2 (c) and (d) show the sampling path for 3 robots with $B=100$ while figure
(e) and (f) show the sampling path for 3 robots with $B=200$ and figure (g)
and (h) show the sampling path for 5 robots with $B=100$. With RMCTS the
robots visit the most informative locations which leads to a more accurate
reconstruction of the spatial domain and lower MSE.

$B$ | method | MSE | $B_{re}$
---|---|---|---
100 | MCTS | 3.16 | 18.45
100 | RMCTS | 2.99 | 15.06
100 | Or | 4.3 | 5.78
100 | MRS | 6.67 | 16
200 | MCTS | 2.99 | 16.83
200 | RMCTS | 2.17 | 12.03
200 | Or | 2.14 | 10.17
200 | MRS | 1.42 | 14
TABLE II: Avg. Results for 100 runs for the vineyard data-set.  $B$ | $N_{R_{i}}$ | method | MSE | $B_{re}$
---|---|---|---|---
100 | 3 | MCTS | 2.83 | 12.87
100 | 3 | RMCTS | 2.50 | 10.62
100 | 3 | MRS | 6.67 | 17
100 | 5 | MCTS | 2.53 | 17.38
100 | 5 | RMCTS | 2.41 | 15.64
100 | 5 | MRS | 6.23 | 28
200 | 3 | MCTS | 2.46 | 20.66
200 | 3 | RMCTS | 2.37 | 12.27
200 | 3 | MRS | 1.08 | 13
200 | 5 | MCTS | 2.38 | 16.38
200 | 5 | RMCTS | 2.14 | 14.37
200 | 5 | MRS | 0.10 | 6
100 | 10 | MCTS | 2.17 | 27.35
100 | 10 | RMCTS | 1.84 | 25.11
100 | 10 | MRS | 0.23 | 12
TABLE III: Avg. Results for 100 runs for the vineyard data-set.

##  VI Conclusions and Future Work

In this paper, we proposed an online distributed multi-robot sampling
algorithm based on the MCTS algorithm which is scalable to the size of the
team. To minimize revisiting locations, robots share their past experiences
(visited sampling locations). A GP model of the scalar field being estimated
is updated every time a sample location is measured and is used in the process
of generation of a new set of random sample locations. In comparison to
baseline methods, our proposed approach is more accurate and makes better use
of the limited budget. Also, our proposed method does not require prior
knowledge of the environment distribution. In the future, we intend to
incorporate the estimation of other robots’ plans into our method and to test
the proposed method on the field.

## References

  * [1] D. Bertsekas.  Rollout, policy iteration, and distributed reinforcement learning.  Athena Scientific, 2021.
  * [2] G. Best, O. M Cliff, T. Patten, R. R. Mettu, and R. Fitch.  Dec-mcts: Decentralized planning for multi-robot active perception.  The International Journal of Robotics Research, 38(2-3):316–337, 2019.
  * [3] L. Booth and S. Carpin.  Distributed estimation of scalar fields with implicit coordination.  In Distributed Autonomous Robotic Systems, 2022 (to appear).
  * [4] M. Campbell, A. Dechemi, and K. Karydis.  An integrated actuation-perception framework for robotic leaf retrieval: Detection, localization, and cutting.  In 2022 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 9210–9216. IEEE, 2022.
  * [5] S. Carpin and T. C. Thayer.  Solving stochastic orienteering problems with chance constraints using monte carlo tree search.  In 2022 IEEE 18th International Conference on Automation Science and Engineering (CASE), pages 1170–1177. IEEE, 2022.
  * [6] H.J.S. Finch, A.M. Samuel, and G.P.F. Lane.  Precision farming.  In Lockhart & Wiseman’s Crop Husbandry Including Grassland, pages 235 – 244. Woodhead Publishing, ninth edition edition, 2014.
  * [7] D. V. Gealy, S. McKinley, M. Gou, L. Miller, S. Vougioukas, J. Viers, S. Carpin, and K. Goldberg.  Co-robotic device for automated tuning of emitters to enable precision irrigation.  In Proceedings of the IEEE Conference on Automation Science and Engineering, pages 922–927, 2016.
  * [8] D. Jang, J. Yoo, C. Y. Son, and H. J. Kim.  Fully distributed informative planning for environmental learning with multi-robot systems.  arXiv preprint arXiv:2112.14433, 2021.
  * [9] S. Kemna, J. G Rogers, C. Nieto-Granda, S. Young, and G. S. Sukhatme.  Multi-robot coordination through dynamic voronoi partitioning for informative adaptive sampling in communication-constrained environments.  In 2017 IEEE International Conference on Robotics and Automation (ICRA), pages 2124–2130. IEEE, 2017.
  * [10] L. Kocsis and C. Szepesvári.  Bandit based monte-carlo planning.  In European conference on machine learning, pages 282–293. Springer, 2006.
  * [11] S. Manjanna, M. A. Hsieh, and G. Dudek.  Scalable multirobot planning for informed spatial sampling.  Autonomous Robots, 46(7):817–829, 2022.
  * [12] S. Manjanna, H. van Hoof, and G. Dudek.  Reinforcement learning with non-uniform state representations for adaptive search.  In 2018 IEEE International Symposium on Safety, Security, and Rescue Robotics (SSRR), pages 1–7. IEEE, 2018.
  * [13] L. Pan, S. Manjanna, and M. A. Hsieh.  Marlas: Multi agent reinforcement learning for cooperated adaptive sampling.  arXiv preprint arXiv:2207.07751, 2022.
  * [14] C. E. Rasmussen.  Gaussian processes in machine learning.  In Summer school on machine learning, pages 63–71. Springer, 2003\.
  * [15] A. Shamshirgaran and S. Carpin.  Reconstructing a spatial field with an autonomous robot under a budget constraint.  In 2022 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 8963–8970. IEEE, 2022.
  * [16] M. L. Stein.  Interpolation of spatial data: some theory for kriging.  Springer Science & Business Media, 1999.
  * [17] J. S. P. Sundaram, W. Du, and Zh. Zhiwei.  A survey on lora networking: Research problems, current solutions, and open issues.  IEEE Communications Surveys & Tutorials, 22(1):371–388, 2019.
  * [18] V. Suryan and P. Tokekar.  Learning a spatial field in minimum time with a team of robots.  IEEE Transactions on Robotics, 36(5):1562–1576, 2020.
  * [19] R. S. Sutton and A. G. Barto.  Reinforcement learning: An introduction.  MIT press, 2018.
  * [20] T. Thayer, S. Vougioukas, K. Goldberg, and S. Carpin.  Routing algorithms for robot assisted precision irrigation.  In Proceedings of the IEEE International Conference on Robotics and Automation, pages 2221–2228, 2018.
  * [21] T. C Thayer, S. Vougioukas, K. Goldberg, and S. Carpin.  Multi-robot routing algorithms for robots operating in vineyards.  In 2018 IEEE 14th International Conference on Automation Science and Engineering (CASE), pages 14–21. IEEE, 2018.
  * [22] S. Vougioukas.  Agricultural robotics.  Annual review of control, robotics, and autonomous systems, 2:339–364, 2019.