# Uncertainty-Aware Deployment of Pre-trained Language-Conditioned Imitation
Learning Policies

Bo Wu1, Bruce D. Lee1, Kostas Daniilidis1, Bernadette Bucher1,2, Nikolai
Matni1 1GRASP Lab, University of Pennsylvania: {bobwu, brucele, kostas,
<EMAIL_ADDRESS>Department, University of Michigan:
<EMAIL_ADDRESS>work is supported by NSF Award SLES-2331880 and
CAREER-2045834

###### Abstract

Large-scale robotic policies trained on data from diverse tasks and robotic
platforms hold great promise for enabling general-purpose robots; however,
reliable generalization to new environment conditions remains a major
challenge. Toward addressing this challenge, we propose a novel approach for
uncertainty-aware deployment of pre-trained language-conditioned imitation
learning agents. Specifically, we use temperature scaling to calibrate these
models and exploit the calibrated model to make uncertainty-aware decisions by
aggregating the local information of candidate actions. We implement our
approach in simulation using three such pre-trained models, and showcase its
potential to significantly enhance task completion rates. The accompanying
code is accessible at the link:
https://github.com/BobWu1998/uncertainty_quant_all.git

##  I Introduction

Inspired by advancements in general-purpose natural language processing and
computer vision models, the robotics community has invested considerable
effort in developing generalist decision making agents that operate across
various robotic platforms and perform a broad spectrum of tasks [1]. These so-
called “foundation models” are typically trained once on a large and diverse
dataset, then deployed in a specific setting with minimal task-specific fine-
tuning. Such pre-trained models have been shown to generalize in many ways
such as to new tasks, new robot platforms, and new environments [2, 3, 4, 5].

However, the extent of these agents’ generalization abilities remains poorly
understood. A significant barrier hindering the generalization abilities and
our understanding of them is the lack of well-calibrated, task-specific
uncertainty estimates. This gap precludes the possibility of making
uncertainty-aware decisions. Our work aims to close this gap by proposing a
calibration enabled uncertainty-aware decision making protocol to boost the
success rate for language-conditioned imitation learning (IL) policies.

###  I-A Contributions

We propose a straightforward yet effective modification to enhance the
generalization capabilities of pre-trained language-conditioned imitation
learning agents for robotic manipulation. The modification is composed of two
components, outlined in Figure 1:

  * •

A calibration step, tailored to the task of interest, that refines the model’s
outputs to generate confidence scores that approximate the correctness
likelihood for imitating the expert.

  * •

An uncertainty-aware action selection technique using the probability
distribution output by the calibrated model (See Algorithm 1).

Each of these steps is novel, and may be of independent interest. However, we
view our primary contribution as the holistic integration of these two
components into the decision making pipeline for language-conditioned IL
policies.

We instantiate our approach on several robotic manipulation models, including
the Perceiver-Actor model [6], the RVT model [7], and the CLIPort model [8]
_without additional training_ , and show that using a small calibration
dataset of expert demonstrations from the target task can lead to a
substantial improvement in task completion rate. Through comprehensive
experimentation, we identify scenarios where integrating uncertainty
quantification into the decision making process is most beneficial. Our
experiments also find that existing pre-trained models provide poorly
calibrated likelihood estimates when instantiated on any particular task of
interest. Therefore, appropriately calibrating the model for the task of
interest is essential to reap the rewards of uncertainty-aware decision
making.

Figure 1: Overview of uncertainty-aware action selection on top of pre-trained
models. The blue cross indicates the maximum score in the heatmap. The upper
path illustrate the standard deployment of pre-trained IL policies at test
time in which the output of the pre-trained model directly predicts the
action. The bottom path indicates our proposed approach, in which the model is
calibrated using a small dataset of expert demonstrations from the task of
interest before selecting the action in an uncertainty-aware manner.

###  I-B Related Work

##### Generalist Robotic Policies

Early efforts [2] introduce a diverse and extensive robotics dataset, and use
this dataset to train visual foresight and supervised inverse models. By
conditioning IL agents with a description of the task to complete, subsequent
work uses such diverse datasets to enable sample-efficient IL [3, 9, 10, 11].
Recent developments have embraced transformer-based architecture to train
multi-modal decision making agents [12, 13] and language-conditioned IL
policies for robotic manipulation [4, 5, 8, 6, 7]. Efforts have also been made
to extend to more general prompt conditioning using images and videos [14]. In
Shah et al. [15], the authors explore the development of a generalist visual
navigation model. For a review of generalist decision making agents, see Yang
et al. [1]. Our work improves the performance of these generalist robotic
policies by incorporating calibrated uncertainty estimates to enable more
informed decisions at inference time.

##### Uncertainty Quantification

For robotic decision making, confidence scores quantifying the model output
uncertainty enable downstream formulations of safe decisions in robotic
planning problems [16, 17, 18]. However, it has been shown that the confidence
scores output by modern classification models are poorly calibrated in that
they are not representative of the true correctness likelihood [19]. While
numerous approaches to improving calibration post-hoc exist, such as
temperature scaling [19] and conformal prediction [20], these techniques are
seldom applied in the realm of pre-trained language-conditioned imitation
learning policies to enhance generalization performance. A notable exception
is the recent work by Ren et al. [21] which applies conformal prediction to
language-based planners to request operator clarification in the face of
ambiguity. In contrast, we consider a setting in which there is no ambiguity
in the task-conditioning, but rather the pre-trained model is imperfect. We
calibrate the uncertainty of this pre-trained model, and use it to select
actions in an uncertainty-aware manner. For further discussion of uncertainty
quantification of foundation models for robotic decision making, see Firoozi
et al. [22].

##  II Approach

###  II-A Background: Language-conditioned IL for Robotic Manipulation

We consider robotic manipulation problems in a discrete-time decision making
context. Let $o_{t}\in\mathcal{O}$ be the observation of the world at time $t$
and $a_{t}\in\mathcal{A}$ be an action applied to the robot at time $t$. Here,
$\mathcal{A}$ is a finite, discretized action space.

Language-conditioned IL assumes access to $n$ expert demonstrations
$\left\\{\zeta_{1},\dots,\zeta_{n}\right\\}$ accompanied by language-
conditioning $\left\\{c_{1},\dots,c_{n}\right\\}$. Each demonstration consists
of a trajectory of observations, and the corresponding actions taken by an
expert. The language-conditioning $c_{i}$ provides information about the task
which the corresponding expert demonstration is completing. Using these
demonstrations, the learner searches for a policy which imitates the expert
behavior for the given language conditioning.

As the action space is discrete, the policy operates as a classifier for the
observation-language pair. In particular, given some observation
$o\in\mathcal{O}$ and language-conditioning $c$, the policy selects an action
as

$\displaystyle\hat{\pi}(o,c)=\operatorname*{argmax}_{a\in\mathcal{A}}\hat{f}_{j(a)}(o,c).$
(1)

Here, $j:\mathcal{A}\to[{\left|\mathcal{A}\right|}]$ is a function enumerating
$\mathcal{A}$, and $\hat{f}$ is a function which takes in an observation $o$
and a language conditioning $c$ to output a vector of logits $\hat{f}(o,c)$
with $\hat{f}_{j(a)}(o,c)$ representing its $j(a)^{\mathrm{th}}$ element.

A policy of the form mentioned above may be determined from the expert
demonstrations and language-conditioning using behavior cloning. In
particular, the learner may choose $\hat{f}$ belonging to some function class
to minimize the cross entropy loss, as in Brohan et al. [4] and Shridhar et
al. [6]. The resultant function, $\hat{f}$, serves as a pre-trained model,
which can be deployed on numerous different tasks. In particular, given some
new language conditioning $c_{n+1}$, it is deployed to select the robot’s
actions as $a_{t}=\hat{\pi}(o_{t},c_{n+1}).$ In general, the target task need
not be one of the tasks used for training; we may have
$c_{n+1}\notin\left\\{c_{1},\dots,c_{n}\right\\}$. However, many of our
experimental results are focused on the setting where
$c_{n+1}\in\left\\{c_{1},\dots,c_{n}\right\\}$.

Pre-trained models using the aforementioned strategy to select an action
neglect pertinent information output by the model about the distribution of
potential actions. This is in part because such pre-trained models tend to be
poorly calibrated and fail to reflect the actual correctness likelihood of
prediction matching the expert demonstration. Accordingly, we address these
problems by calibrating the model and then following an uncertainty-aware
action selection strategy, which is discussed in Section II-B and Section
II-C.

###  II-B Target Task Calibration

The logits can be viewed as a representation of a probability distribution $p$
over possible actions by applying the softmax function $\sigma_{SM}$ to the
output of the model $\hat{f}$: $p=\sigma_{SM}(\hat{f}(o,c)).$ While $p$ is a
viable probability distribution, it may not be well-aligned with the true
correctness likelihood for imitating the expert, especially for the specific
task $c_{n+1}$.

To improve the alignment between the distribution $p$ and the correctness
likelihood for task $c_{n+1}$, we collect a calibration set from task
$c_{n+1}$. In particular, we assume that the target language-conditioning
$c_{n+1}$ is drawn randomly from some distribution
$c_{n+1}\sim\mathcal{D}_{\mathsf{test}}$, and that we have access to a
calibration set of expert demonstrations
$\left\\{\zeta_{1}^{\mathsf{cal}},\dots,\zeta_{k}^{\mathsf{cal}}\right\\}$
along with their associated conditioning
$\left\\{c_{1}^{\mathsf{cal}},\dots,c_{k}^{\mathsf{cal}}\right\\}$, where
$c_{i}^{\mathsf{cal}}$ are drawn from the same distribution as $c_{n+1}$:
$c_{i}^{\mathsf{cal}}\sim\mathcal{D}_{\mathsf{test}}$. Calibrated models may
be obtained from this set with various methods, but we propose using
temperature scaling [19]. Specifically, we determine the optimal temperature
parameter $\hat{T}$ as the value of $T$ which minimizes the cross entropy loss
of the model $\frac{1}{T}\hat{f}$ over the calibration set. Given this
temperature parameter, we define the calibrated model outputting the
probability scores as
$\tilde{p}=\tilde{f}(o,c)=\sigma_{SM}\left(\frac{\hat{f}(o,c)}{\hat{T}}\right),$
with components $\tilde{p}_{k}=\tilde{f}_{k}(o,c)$.

###  II-C Uncertainty-Aware Action Selection

The greedy action selection approach of equation (1) does not make use of the
calibrated confidence scores from Section II-B, and overlooks important
structure of the action space. Specifically, distinct elements in
$\mathcal{A}$ may represent similar actions. This structure should be
leveraged in the manipulation problem, as many tasks can be completed with
distinct, but similar actions, e.g. a robotic gripper can open a drawer by
pulling at different positions of the handle.

We propose Algorithm 1 to remedy the above shortcoming. Rather than selecting
the action with the highest confidence score, Algorithm 1 selects the action
for which the neighboring region of potential actions has the highest sum of
confidence scores. The neighboring region of an action $a\in\mathcal{A}$ is
defined as $\left\\{a^{\prime}:d(a,a^{\prime})\leq\tau\right\\},$ where the
metric $d$ is selected based on the problem at hand to capture similarity
between two actions in $\mathcal{A}$, and $\tau$ is a similarity threshold.
While both the metric $d$ and the threshold $\tau$ are hand-designed in this
work, they could be chosen in a data-driven manner in future work. Note that
while the greedy action selection approach in equation (1) is only sensitive
to the ordering of highest confidence score output by the model, Algorithm 1
needs a well-calibrated model to adequately determine the confidence of the
neighboring region.

Algorithm 1 Uncertainty-Aware Action Selection

1:Input:  Calibrated model $\tilde{f}$, conditioning $c$, observation $o$,
robustness threshold $\tau$

2:Compute
$\tilde{p}_{1},\dots,\tilde{p}_{{\left|\mathcal{A}\right|}}=\tilde{f}(o,c).$

3:Set
$\hat{a}=\operatorname*{argmax}_{a\in\mathcal{A}}\sum_{a^{\prime}\in\mathcal{A}}\tilde{p}_{j(a^{\prime})}\mathbf{1}(d(a,a^{\prime})<\tau)$

4:Output:  Uncertainty-aware action $\hat{a}$

Model | Baseline | Uncalibrated UA | Calibrated UA
---|---|---|---
PerAct | $0.382\pm 0.012$ | $0.385\pm 0.012$ | $\bm{0.414\pm 0.012}$
RVT | $0.602\pm 0.012$ | $0.607\pm 0.012$ | $\bm{0.623\pm 0.011}$
CLIPort | $0.803\pm 0.005$ | $0.833\pm 0.005$ | $\bm{0.833\pm 0.005}$
TABLE I: Summary of results for calibrated and uncalibrated uncertainty-aware
approach and the uncertainty oblivious (baseline) approach on PerAct, RVT, and
CLIPort. For PerAct and RVT, the values in the table represent the portion of
successfully completed tasks. For CLIPort, the values represent the mean
reward, where the maximum reward for an episode is one if the desired task is
completed successfully. Standard errors are shown.

In Algorithm 1, our method assumes to take a calibrated model $\tilde{f}$, a
language conditioning of the task $c$, the observation $o$, and the robustness
threshold $\tau$ as input. At each time step, the calibrated model $\tilde{f}$
produces a vector of logits $\tilde{p}$, based on task conditioning and
observation. Then, the approach calculates the uncertainty-aware score for the
actions by aggregating the neighboring confidence scores of each candidate,
where two actions $a$ and $a^{\prime}$ are considered “neighbors” if
$d(a,a^{\prime})<\tau$. Finally, the model makes the decision about the action
based on the highest confidence score produced in the previous step.

Note that the approach for selecting the action from the calibrated model in 3
of Algorithm 1 is just one possibility for aggregating the confidence of
neighboring actions. An alternative is to convolve the output of the
calibrated model with a Gaussian kernel. This approach is discussed in Section
-B2. Another alternative is to construct prediction sets using conformal
prediction, and then select a representative from the set.

A key feature of the proposed approach is that the weights of the pre-trained
model are not updated, as it can be both computationally and statistically
prohibitive to fine-tune the model. As a result, the efficacy of an agent
using the proposed approach is still bottle-necked by the quality of the pre-
trained agent. We therefore expect improvements in the performance from
leveraging uncertainty-aware decision making to be limited to some degree by
the quality of the pre-trained model. In our experiments, we highlight several
cases where the performance boost is most notable, and offer explanations for
why this may be the case.

##  III Experimental Results

We evaluate the method discussed in Section II by deploying it on three pre-
trained models: the PerAct model [6], the RVT model [7], and the CLIPort model
[8]. PerAct and RVT are evaluated on the RLBench dataset [23], while CLIPort
is evaluated in the Ravens benchmark [24]. Each of these datasets breaks down
the desired behavior of the robot into tasks. A task consists of a small
collection of related language conditionings. For example, the RLBench task
“stack blocks” consists of language-conditionings:111See [23] for all
language-conditionings of all tasks.

  * •

“place three of the blue cubes on top of each other”

  * •

“build a tall tower out of four green cubes”

  * •

“place two of the red cubes on top of each other”

Within each task for both the RLBench dataset and the Ravens dataset, the
variations of language-conditionings are sampled uniformly at random from the
set of task-specific variations built into the environments. We therefore
calibrate the models for particular tasks by obtaining expert demonstrations
corresponding to language-conditionings sampled uniformly at random from a
particular task.

PerAct and RVT were trained on a dataset consisting of 100 expert
demonstrations for each of 18 distinct RLBench tasks. We evaluate our approach
on each model on all 18 of the tasks they were trained on. For each task, we
calibrate the model using temperature scaling with a calibration set of $25$
expert demonstrations, and evaluate on $100$ episodes.

CLIPort was trained on 1000 demonstrations from each of 9 distinct tasks from
the Ravens dataset. We evaluate our approach on four of the nine tasks that
the model was trained on (discussed in Section III-A). On each task, we
calibrate the model on 100 expert demonstrations using temperature scaling.
Evaluations are performed on 1000 episodes.222For further details about the
implementation of uncertainty-aware action selection on these models, see
Section -A and Section -B.

Task | Baseline | UA
---|---|---
| stack-block-pyramid-seq
---
seen-colors
$0.9475\pm 0.0047$ | $\bm{0.9815\pm 0.0027}$
| packing-seen-google
---
objects-seq
$0.8090\pm 0.0104$ | $\bm{0.8275\pm 0.0100}$
| packing-seen-google
---
objects-group
$0.8846\pm 0.0080$ | $\bm{0.9108\pm 0.0072}$
| assembling-kits-seq
---
seen-colors
$0.5706\pm 0.0116$ | $\bm{0.6134\pm 0.0115}$
Task with Distractors | Baseline | UA
| stack-block-pyramid-seq
---
seen-colors-diff-sizes
$0.7963\pm 0.0073$ | $\bm{0.9711\pm 0.0031}$
| assembling-kits-seq
---
seen-colors-diff-sizes
$0.4832\pm 0.0028$ | $\bm{0.5544\pm 0.0030}$
TABLE II: Task-specific mean rewards for CLIPort on four tasks from the
original Ravens dataset and two tasks with small distractors added to the
original tasks.

(a) Put Gray Block on the Lightest Brown Block

(b) Pack All the Spatula with Purple Head Objects in the Brown Box

(c) Put the Brown Ring In the Brown Box

Figure 2: Qualitative results for CLIPort with and without uncertainty-aware
action selection. The language commands are in the subfigure captions. From
top to bottom, each row represents the model’s observation, the heatmap in
pick phase, the heatmap in place phase, and the observation after executing
the actions. In 2a and 2b, our method forces the model to choose a central
location on the target object instead of the edges. In 2c, our method corrects
the model by smoothing out the action with highest raw score.

The results from running uncertainty-aware action selection on the three
models are summarized in Table I. We see that the uncertainty-aware approach
offers a benefit to the overall task completion rate for all models. To
investigate exactly when the approach provides a benefit, the subsequent
sections study three pertinent questions, summarized below.

  1. 1.

Section III-A: Why does uncertainty-aware action selection help? Uncertainty-
aware action selection prevents the end effector from moving to positions
described by isolated, high-confidence locations output by the pre-trained
model. This can prevent the gripper from completely missing the object, and
encourages a better grip when it is successful in gripping the object.

  2. 2.

Section III-B: Is calibration necessary? It depends on the model. For simple
models trained on very large datasets (CLIPort), the pre-trained model appears
to be fairly well calibrated, and calibration is not necessary to achieve the
benefit of uncertainty-aware decisions. Both PerAct and RVT are poorly
calibrated on any given task, and calibrating the model is critical to achieve
the benefit of uncertainy-aware decisions.

  3. 3.

Section III-C: Can uncertainty-awareness enable generalzation to new settings?
We demonstrate that small distractor objects to the scene can have a
devastating performance of the pre-trained model. However, uncertainty-aware
action selection achieves essentially the same performance it does in the
absence of distractors. This provides a strong indication that accounting for
uncertainty-aware action selection can improve generalization to certain task
distribution shifts.

###  III-A Why Does Uncertainty-Aware Action Selection Help?

Figure 3: Success rate score for Perceiver-Actor, uncertainty-aware strategy
with calibrated/uncalibrated confidence for each task Figure 4: Success rate
score for RVT, uncertainty-aware strategy with calibrated/uncalibrated
confidence for each task

To investigate why uncertainty-aware action selection helps, we first look at
the per-task success rates of the CLIPort model. In the first four rows of
Table II, we compare uncertainty-aware action selection with the baseline on
four tasks from the Ravens dataset. From the results, we find that
incorporating uncertainty is particularly useful when the CLIPort task
requires precise placement of target objects. This is important for stacking
blocks, packing objects, and assembling kits. The reason that uncertainty
awareness is beneficial in these settings is that it enables conservative
decisions in the picking phase by encouraging the end effector to pick up the
object at the center instead of at the edges, as illustrated in Figure 1,
Figure 2a, and Figure 2b.

In these examples, our method takes the calibrated score map of the baseline
(left column of the second row), and generates uncertainty-aware score
calculation (right column of the second row). Conditioned on the more
conservative picking location, the agent picks up the object more easily and
finds a better placing location, as shown in the third row. The bottom row of
each step shows the result of the manipulation at the current time step.

Additionally, by smoothing confidence scores across possible actions, the
uncertainty-aware approach prevents the end effector from selecting isolated
high-confidence actions. The benefit of this is illustrated in the second row
of Figure 2c.

The per task success rates of PerAct are shown in Figure 3 while the per task
success rates of RVT are shown in Figure 4. Due to the higher complexity of
the tasks from RLBench and the higher standard error in the results, it is
difficult to draw concrete conclusions about when the benefit of accounting
for uncertainty is most important. However, we do see the most pronounced
improvement on tasks with a lower success rate. This benefit appears to be due
to the larger discrepancy between the action with the highest confidence score
and the action with the largest sum of neighboring confidence scores on tasks
for which the model is overall less confident. In particular, we plot the
maximum entropy of the translation distribution output by PerAct for each task
in Figure 5. We see that for the push buttons task, it has the highest success
rate and the lowest maximum entropy. Meanwhile, the highest maximum entropy
comes from the task put groceries in cupboard, upon which we see a significant
boost coming from our method.

Figure 5: Maximum entropy over all episodes in each task of PerAct

In light of the above experimental results, we conclude that by encouraging
the model to select actions with higher neighboring confidence scores, our
method prevents the model from choosing aggresive action candidates, thereby
resulting in better grasping performance.

###  III-B Do we need calibration?

We see in Table I that for both PerAct and RVT, the uncertainty-aware action
provides only a very small improvement in the absence of calibration (0.3% and
0.5% boost in task success rate, respectively); however, they yield a
significant task success rate if a calibration step is included (3.2% and 2.1%
boosts, respectively).

(a) uncalibrated

(b) Calibrated

Figure 6: Reliability diagrams are shown for PerAct on two of the RLBench
tasks. Figure 6a illustrates the reliability diagrams before calibration,
while Figure 6b shows the reliability diagrams after running temperature
scaling on a calibration set of 25 expert demonstrations from the each task.
Overlayed on the plot is the value for the expected calibration error (ECE,
see Guo et al. [19] for further details). We see that temperature scaling
provides a marked improvement in the calibration.

To understand why this is the case, we construct reliability diagrams in
Figure 6. The figure depicts this reliability diagram for PerAct on two tasks,
with both the uncailbrated model and the model calibrated using temperature
scaling. The $y$-axis is the model accuracy (the correctness likelihood) at
each confidence level, which is the number of predicted actions that are the
same as the expert demonstration, divided by the total number of predictions
made with confidence within the thresholding interval. The $x$-axis depicts
the model confidence, which is maximum value in the probabilty distribution
representing the calibrated model: $\max_{a\in\mathcal{A}}\tilde{f}_{a}(o,c)$.
The confidence is perfectly calibrated when the accuracy and the confidence
are linearly related with slope of one. For more details on reliability
diagrams, see Guo et al. [19]. These diagrams illustrate the value of
temperature scaling for improving model calibration using a calibration set of
25 expert demonstrations from the task of interest. From these results, we
conclude that PerAct has very poorly calibrated uncertainty measurements, but
applying temperature scaling using a small number of demonstrations from a
specific task improves calibration substantially.

The importance of this calibration is further reflected by the improvement
offered by the calibrated uncertainty-aware approach over the uncalibrated
uncertainty-aware approach in Table I. As in Table I, RVT benefits from
temperature scaling. However, CLIPort does not see a benefit of temperature
scaling in improving the success rate of uncertainty-aware action selection.
We find that the calibrating cross-entropy loss barely changes during from
calibration, indicating that the pre-trained model is already fairly well-
calibrated for each of the individual tasks.

###  III-C Can Uncertainty-Awareness Improve Generalzation?

To determine whether uncertainty-aware decisions can help pre-trained models
generalize to out-of-distribution settings, we designed two new tasks. One of
these is a modification of stack blocks, and the other one is a modification
of the assemble kits task. Both of the new tasks introduce a collection of
smaller-sized target objects of various colors and shapes, which serve as
distractors to the model. In the last two rows of Table II, we find that the
pre-trained method is very fragile to such distractors, and they cause a
massive drop in performance from 94.75% completion to 79.63% for stacking
blocks, and from 57.06% to 48.32% for assembling kits. However, the
uncertainty-aware approach is robust to such distractors: performance degrades
only from 98.15% and 61.34% completion to 97.11% and 55.44%, respectively.

In Figure 7, we demonstrate the impact of the distractors on the pre-trained
model. We see that the precense of distractors causes the pre-trained model to
pick up one of the distractor blocks which has a single pixel of high
confidence. The uncertainty-aware approach smooths the heatmap, thereby
determining that the isolated high confidence event is due to a distractor,
and that the best course of action is to select a region where the model
output a collection of neighboring actions with high confidence.

Figure 7: Qualitative results of CLIPort on instruction ”Put the brown block
on the blue and red blocks” when distractors are introduced.

More generally, we believe these results indicate the potential value of
incorporating uncertainty estimates for decision making in out-of-distribution
scenarios.

##  IV Conclusion

In this paper, we proposed the use of simple classification model calibration
techniques in conjunction with an uncertainty-aware action selection approach
to improve the success rates for language-conditioned imitation learning in
robotic manipulation. Our results suggest that by appropriately calibrating
the model and incorporating an uncertainty-aware decision making strategy, we
can benefit the pre-trained “foundation models” in both seen and unseen tasks.
This approach does not require any additional training or fine-tuning of the
original model and is scalable to versatile robot manipulation models and
platforms.

The modest, though notable performance boost that occurs by incorporating
uncertainty-aware decisions for pre-trained models highlights the potential
for such approaches for generalist robotic policies in general. An exciting
avenue for future work is to obtain calibrated uncertainty models for
generalist dynamics forecasting which can then be used as part of an
uncertainty-aware trajectory planning or model predictive control scheme.

###  -A Model details

For PerAct and RVT, the action space consists of three coordinates for the
discretized desired end effector translation, three coordinates for the
discretized desired end effector rotation, and two binary values indicating
gripper status and collision avoidance module activation. The desired action
is then achieved using a sampling-based motion planner. The world state is
represented by a voxel observation of the environment. PerAct directly takes
this voxel obseration as an input (see [6]), while RVT takes five RGB-D
virtual camera projections of this voxel observations as input (see Goyal et
al. [7]). Due to differences between the output of each model, and the
structure assumed by our approach in (1) and due to computational challenges
of Algorithm 1 in high dimensions, we make several model-specific adjustments
to the uncertainty-aware action selection. See Section -B1 and Section -B2 for
PerAct and RVT, respectively.

CLIPort instead breaks tasks into a sequence of pick-and-place actions. Each
pick action is described by the two dimensional end effector location, while
each place action is described by the two dimensional end effector location
and the one dimensional end effector orientation. For each pick and place
step, the CLIPort model takes as input the language description of the task,
and a top-down virtual RGB-D image of the scene. See Shridhar et al. [8], Zeng
et al. [24] for further details. Note that both picking and placing outputs a
two dimensional heatmap of confidence scores for the appropriate end effector
location. Applying algorithm 1, we calculate the uncertainty-aware heatmap and
determine the appropriate picking location. For the orientation during the
place step, we directly trust the output of pre-trained model.

All hyperparameters are determined by performing a parameter sweep using cross
validation.

###  -B Additional Implementation Details

####  -B1 PerAct

The action space $\mathcal{A}$ is very large, and it is computationally
expensive to search over the entire action space to calculate the neighboring
confidence for each action. Meanwhile, while the probability output from a
well-trained model can span the entire action space, there are relatively few
entries with high confidence scores. For these reasons, we restrict our
uncertainty-aware action selection search to translation actions that have
relatively high confidence scores. For rotation, gripper status, and collision
avoidance module activation, we directly trust the pre-trained model.

To restrict the action search space, we make several modifications to
Algorithm 1. These are summarized in Algorithm 2. Our method assumes that we
have a calibrated model $\tilde{f}$. First we apply a probability threshold to
filter out the actions with probabilities higher than this threshold. Then, if
the resulting action space has more than k elements, we keep only the top k of
them. After this, we select the center of the remaining actions, define a
search space $\mathcal{H}$ around the center with radius $D/2$, and go loop
through the actions in the search space. For each action $a$, we extract the
remaining actions $a^{\prime}\in\mathcal{A}^{\prime}$ around it with
$d(a,a^{\prime})<\tau$, and calculate the accumulative scores for $a$.
Finally, we take the action with the greatest accumulative score to be our
desired action. This mitigates the computational problem described above by
overlooking non-informative parts in the model’s output. We choose
$\alpha=1/{\left|\mathcal{A}\right|}$ and $k=4000$.

Algorithm 2 Detailed UA Action Selection Strategy

1:Input:  Calibrated model $\tilde{f}$, conditioning $c$, observation $o$,
threshold $\tau$, probability threshold $\alpha$, search size $D$.

2:Compute action probabilities:
$\tilde{p}_{1},\dots,\tilde{p}_{{\left|\mathcal{A}\right|}}=\tilde{f}(o;c).$

3:$\mathcal{A}^{\prime}=\\{a\in\mathcal{A},\ \tilde{p}_{a}>\alpha\\}$,

4:if ${\left|\mathcal{A}^{\prime}\right|}>k$: keep $k$ actions in
$\mathcal{A}^{\prime}$ with the highest scores

5:c = mean($a^{\prime}\in\mathcal{A}^{\prime}$)

6:$\mathcal{H}$= { $a\in[c-D/2,c+D/2]^{3}\cap\mathcal{A}$ }

7:Set
$\hat{a}=\operatorname*{argmax}_{a\in\mathcal{H}}\sum_{a^{\prime}\in\mathcal{A}^{\prime}}\tilde{p}_{j(a^{\prime})}\mathbf{1}(d(a,a^{\prime})<\tau)$

8:Output:  Uncertainty-aware action $\hat{a}$

####  -B2 RVT

RVT is modified from PerAct to predict the heatmap for the translation in
virtual image planes instead of the raw 3D voxel space. As such, we adjust
Algorithm 1 to work with the 2D output accordingly. For each heatmap, we
convolve it with a Gaussian blurring kernel to get an uncertainty-aware
heatmap and reconstruct a 3D uncertainty-aware point cloud.333A Gaussian
blurring kernel is a generalization of the weighted sum in 3 from Algorithm 1.
The width and standard deviation of this kernel are treated as hyperparameters
that are determined via cross validation. We apply it only in this setting
because the convolution becomes only a 2D convolution, which is necessary in
the image space. We apply this convolution strategy in 1D space for each axis
of rotation. Consequently, the uncertainty-aware action decision is made over
the resulting uncertainty-aware heatmaps for both translation and rotation.

## References

  * Yang et al. [2023] S. Yang, O. Nachum, Y. Du, J. Wei, P. Abbeel, and D. Schuurmans, “Foundation models for decision making: Problems, methods, and opportunities,” _arXiv preprint arXiv:2303.04129_ , 2023.
  * Dasari et al. [2019] S. Dasari, F. Ebert, S. Tian, S. Nair, B. Bucher, K. Schmeckpeper, S. Singh, S. Levine, and C. Finn, “Robonet: Large-scale multi-robot learning,” _arXiv preprint arXiv:1910.11215_ , 2019.
  * Ebert et al. [2021] F. Ebert, Y. Yang, K. Schmeckpeper, B. Bucher, G. Georgakis, K. Daniilidis, C. Finn, and S. Levine, “Bridge data: Boosting generalization of robotic skills with cross-domain datasets,” _arXiv preprint arXiv:2109.13396_ , 2021.
  * Brohan et al. [2022] A. Brohan, N. Brown, J. Carbajal, Y. Chebotar, J. Dabis, C. Finn, K. Gopalakrishnan, K. Hausman, A. Herzog, J. Hsu _et al._ , “Rt-1: Robotics transformer for real-world control at scale,” _arXiv preprint arXiv:2212.06817_ , 2022.
  * Brohan et al. [2023] A. Brohan, N. Brown, J. Carbajal, Y. Chebotar, X. Chen, K. Choromanski, T. Ding, D. Driess, A. Dubey, C. Finn _et al._ , “Rt-2: Vision-language-action models transfer web knowledge to robotic control,” _arXiv preprint arXiv:2307.15818_ , 2023.
  * Shridhar et al. [2023] M. Shridhar, L. Manuelli, and D. Fox, “Perceiver-actor: A multi-task transformer for robotic manipulation,” in _Conference on Robot Learning_. PMLR, 2023, pp. 785–799.
  * Goyal et al. [2023] A. Goyal, J. Xu, Y. Guo, V. Blukis, Y.-W. Chao, and D. Fox, “Rvt: Robotic view transformer for 3d object manipulation,” _arXiv preprint arXiv:2306.14896_ , 2023.
  * Shridhar et al. [2022] M. Shridhar, L. Manuelli, and D. Fox, “Cliport: What and where pathways for robotic manipulation,” in _Conference on Robot Learning_. PMLR, 2022, pp. 894–906.
  * Ding et al. [2019] Y. Ding, C. Florensa, P. Abbeel, and M. Phielipp, “Goal-conditioned imitation learning,” _Advances in neural information processing systems_ , vol. 32, 2019.
  * Stepputtis et al. [2020] S. Stepputtis, J. Campbell, M. Phielipp, S. Lee, C. Baral, and H. Ben Amor, “Language-conditioned imitation learning for robot manipulation tasks,” _Advances in Neural Information Processing Systems_ , vol. 33, pp. 13 139–13 150, 2020.
  * Lynch and Sermanet [2020] C. Lynch and P. Sermanet, “Language conditioned imitation learning over unstructured data,” _arXiv preprint arXiv:2005.07648_ , 2020.
  * Reed et al. [2022] S. Reed, K. Zolna, E. Parisotto, S. G. Colmenarejo, A. Novikov, G. Barth-Maron, M. Gimenez, Y. Sulsky, J. Kay, J. T. Springenberg _et al._ , “A generalist agent,” _arXiv preprint arXiv:2205.06175_ , 2022.
  * Driess et al. [2023] D. Driess, F. Xia, M. S. Sajjadi, C. Lynch, A. Chowdhery, B. Ichter, A. Wahid, J. Tompson, Q. Vuong, T. Yu _et al._ , “Palm-e: An embodied multimodal language model,” _arXiv preprint arXiv:2303.03378_ , 2023.
  * Jiang et al. [2022] Y. Jiang, A. Gupta, Z. Zhang, G. Wang, Y. Dou, Y. Chen, L. Fei-Fei, A. Anandkumar, Y. Zhu, and L. Fan, “Vima: General robot manipulation with multimodal prompts,” _arXiv preprint arXiv:2210.03094_ , 2022.
  * Shah et al. [2023] D. Shah, A. Sridhar, N. Dashora, K. Stachowicz, K. Black, N. Hirose, and S. Levine, “Vint: A foundation model for visual navigation,” _arXiv preprint arXiv:2306.14846_ , 2023.
  * Lindemann et al. [2023] L. Lindemann, M. Cleaveland, G. Shim, and G. J. Pappas, “Safe planning in dynamic environments using conformal prediction,” _IEEE Robotics and Automation Letters_ , 2023.
  * Dixit et al. [2023] A. Dixit, L. Lindemann, S. X. Wei, M. Cleaveland, G. J. Pappas, and J. W. Burdick, “Adaptive conformal prediction for motion planning among dynamic agents,” in _Learning for Dynamics and Control Conference_. PMLR, 2023, pp. 300–314.
  * Lekeufack et al. [2023] J. Lekeufack, A. A. Angelopoulos, A. Bajcsy, M. I. Jordan, and J. Malik, “Conformal decision theory: Safe autonomous decisions from imperfect predictions,” _arXiv preprint arXiv:2310.05921_ , 2023.
  * Guo et al. [2017] C. Guo, G. Pleiss, Y. Sun, and K. Q. Weinberger, “On calibration of modern neural networks,” in _International conference on machine learning_. PMLR, 2017, pp. 1321–1330.
  * Shafer and Vovk [2008] G. Shafer and V. Vovk, “A tutorial on conformal prediction.” _Journal of Machine Learning Research_ , vol. 9, no. 3, 2008.
  * Ren et al. [2023] A. Z. Ren, A. Dixit, A. Bodrova, S. Singh, S. Tu, N. Brown, P. Xu, L. Takayama, F. Xia, J. Varley _et al._ , “Robots that ask for help: Uncertainty alignment for large language model planners,” _arXiv preprint arXiv:2307.01928_ , 2023.
  * Firoozi et al. [2023] R. Firoozi, J. Tucker, S. Tian, A. Majumdar, J. Sun, W. Liu, Y. Zhu, S. Song, A. Kapoor, K. Hausman _et al._ , “Foundation models in robotics: Applications, challenges, and the future,” _arXiv preprint arXiv:2312.07843_ , 2023.
  * James et al. [2020] S. James, Z. Ma, D. R. Arrojo, and A. J. Davison, “Rlbench: The robot learning benchmark & learning environment,” _IEEE Robotics and Automation Letters_ , vol. 5, no. 2, pp. 3019–3026, 2020.
  * Zeng et al. [2021] A. Zeng, P. Florence, J. Tompson, S. Welker, J. Chien, M. Attarian, T. Armstrong, I. Krasin, D. Duong, V. Sindhwani _et al._ , “Transporter networks: Rearranging the visual world for robotic manipulation,” in _Conference on Robot Learning_. PMLR, 2021, pp. 726–747.


11institutetext: Wellcome/EPSRC Centre for Interventional and Surgical
Sciences, University College London, Gower Street, London WC1E 6BT, UK
22institutetext: School of Biomedical Engineering & Imaging Sciences, King’s
College London, Strand, London WC2R 2LS, UK
22email<EMAIL_ADDRESS>
# An Unsupervised Approach to Ultrasound Elastography with End-to-end Strain
Regularisation

Rémi Delaunay 1122 0000-0002-0398-4995 Yipeng Hu 11 0000-0003-4902-0486 Tom
Vercauteren 1122 0000-0003-1794-0456

###### Abstract

Quasi-static ultrasound elastography (USE) is an imaging modality that
consists of determining a measure of deformation (i.e. strain) of soft tissue
in response to an applied mechanical force. The strain is generally determined
by estimating the displacement between successive ultrasound frames acquired
before and after applying manual compression. The computational efficiency and
accuracy of the displacement prediction, also known as time-delay estimation,
are key challenges for real-time USE applications. In this paper, we present a
novel deep-learning method for efficient time-delay estimation between
ultrasound radio-frequency (RF) data. The proposed method consists of a
convolutional neural network (CNN) that predicts a displacement field between
a pair of pre- and post-compression ultrasound RF frames. The network is
trained in an unsupervised way, by optimizing a similarity metric between the
reference and compressed image. We also introduce a new regularization term
that preserves displacement continuity by directly optimizing the strain
smoothness. We validated the performance of our method by using both
ultrasound simulation and in vivo data on healthy volunteers. We also compared
the performance of our method with a state-of-the-art method called OVERWIND
[17]. Average contrast-to-noise ratio (CNR) and signal-to-noise ratio (SNR) of
our method in 30 simulation and 3 in vivo image pairs are 7.70 and 6.95, 7 and
0.31, respectively. Our results suggest that our approach can effectively
predict accurate strain images. The unsupervised aspect of our approach
represents a great potential for the use of deep learning application for the
analysis of clinical ultrasound data.

###### Keywords:

Ultrasound elastography Time delay estimation convolutional neural network

##  1 Introduction

Ultrasound elastography (USE) is an imaging technique that enables the
characterization of tissue mechanical properties [21]. Since its introduction
in 1991 [19], strain imaging has shown usefulness in diagnostic applications
where pathological alterations induce modification of tissue stiffness, such
as lesion detection in liver disease [12] and tumour characterisation in the
thyroid [16], breast [7] and prostate cancer [3]. This work focuses on quasi-
static, free-hand strain elastography, where a time-varying axial compression
is applied with an ultrasound transducer to the targeted tissue [23]. The
tissue mechanical behavior is then determined by mapping the relative
deformation (i.e. strain) induced by manual compression (i.e. stress).
Various methods have been proposed over the years to measure the strain. The
main approach consists of determining the spatial displacement between a pair
of radio-frequency (RF) ultrasound image data, acquired before and after
applying an axial compression. The displacement estimation, also known as
time-delay estimation, has been historically performed by maximizing a
correlation function between local frame windows, either in the time or phase
domain [18, 20, 2, 22, 1]. More recently, different approaches have added a
regularization parameter to account for displacement discontinuity and improve
displacement estimates [15, 6]. Although these methods demonstrated the
ability to make an accurate displacement estimation, current techniques are
sensitive to noise and global decorrelation, i.e. the change of speckle
appearance due to out-of-plane motion. Furthermore, real-time imaging is an
important feature of elastography, and a trade-off is often made between the
precision of standard approaches and their computational cost.
The recent progress of learning-based methods in computer vision for optical
flow estimation have inspired new approaches for strain elastography [9].
Those methods demonstrated the ability of neural networks to exploit the
ultrasound high-frequency content and to robustly produce accurate
displacement estimates [14, 25]. Previous networks have been trained using
ultrasound simulation associated with ground truth displacement and strain
labels. Accurate ground-truth images can be difficult to obtain for quasi-
static elastography because the magnitude of applied stress is unknown.
However, real-world ultrasound data often exhibits complex speckle patterns
and echogenic features that can be quite challenging to replicate in
ultrasound simulation.
In this paper, we proposed an unsupervised method for time-delay estimation
that allows a neural network to be trained directly on clinical data and
predict tissue displacement. Unlike previous methods, our CNN training
procedure is performed without ground truth labels. Instead, the network
weights are optimized by minimizing a dissimilarity function between the pre-
compression image and warped compressed image. We also introduce a new
regularization term that preserves displacement continuity by directly
optimizing the smoothness of the strain prediction. We validated our method by
applying it on both real ultrasound data and simulations. The performance of
our method was evaluated by comparing our displacement field and strain
prediction with the ground truth labels and a gold standard USE technique
[17]. To the best of our knowledge, this is the first unsupervised deep-
learning method applied to strain imaging.

##  2 Methods

###  2.1 Problem Statement

Our method follows an approach similar to a standard image registration
framework, which aims to find a spatial transformation that maps a moving
image into the space of a reference image. In the case of a non-rigid image
registration solution, this transformation can be represented as a dense
displacement field (DDF). This transformation is generally optimized in an
iterative manner, by maximizing an objective function that measures the
similarity between the warped moving image and the reference image, and a
regularization parameter to ensure displacement continuity. From a learning-
based approach perspective, the image mapping is predicted by a neural network
instead of being directly optimized. An overview of the method is presented in
Fig. 1.

###  2.2 Displacement Estimation

Our network takes as input a pair of pre- and post-compression 2D RF frames,
here named $Pre$ and $Post$, and predicts a DDF $u$. The network parameters
are learned by minimizing a dissimilarity metric $L_{sim}$ between pairs of 2D
RF ultrasound data. Our training loss function also includes a regularization
term, $L_{reg}$, acting on the predicted displacement field $u$ and associated
with a weighting hyper-parameter $\alpha$, to ensure balance between the
likelihood and smoothness of the predicted transformation. The optimization
problem can be written as:

$\hat{\theta}=\arg\min_{\theta}\ [L_{sim}(Pre,Post;u)\ +\ \alpha\ .\
L_{reg}(u)]$ (1)

where $\theta$ represents the network parameters that are optimized through
stochastic gradient descent.

Figure 1: Overview of the method.

The dissimilarity metric corresponds to a negative local normalized cross-
correlation (NCC) which average the NCC score between sliding windows sampled
from the pre-compression image and the transformed post-compression image,
resampled with the predicted displacement field. The NCC between two local
image windows, $W_{pre}$ and $W_{post}$, with $i,j$ pixel components can be
written as:

$NCC=\frac{1}{N}\sum_{i,j}\frac{\left[W_{pre}(i,j)-\mu_{W_{pre}}\right]\times\left[W_{post}(i,j)-\mu_{W_{post}}\right]}{\sigma_{W_{pre}}\times\sigma_{W_{post}}}$
(2)

where $N$ is the number of pixels ($i,j$) and $\mu$ and $\sigma$ correspond to
the mean and standard deviation of the images, respectively.

The regularisation term consists of the L1-norm of the second spatial
derivatives of the predicted displacement field $u$. Given that the strain
modulus corresponds to the displacement gradient, minimizing its second
derivative allows to enforce the strain map smoothness. The regularisation
term can be written as:

$L_{reg}=\sum_{i,j}\ (|\ {\partial_{x}^{2}u_{i,j}}\ |+|\
{\partial_{y}^{2}u_{i,j}}\ |)$ (3)

where $\partial_{x}^{2}$ and $\partial_{y}^{2}$ are the second partial
derivatives in axial and lateral directions, respectively.

###  2.3 Implementation

####  2.3.1 Network

The architecture of our network was presented in a method for medical image
registration [8], and consists in an encoder-decoder CNN with skip
connections. The encoder part is composed of four down-sampling blocks, which
capture the hierarchical features necessary to establish correspondence
between the image pair. Each down-sampling block consists of two convolutional
layers with a residual network shortcut, batch normalization, and leaky
rectified linear unit (Leaky ReLU). Symmetrically, the decoder part is
composed of four up-sampling blocks that consists of an additive up-sampling
layer summed over a transpose convolutional layer. Finally, each up-sampling
block outputs a displacement field that is convoluted and resized to the input
size, then summed to output the predicted displacement field. The network was
implemented in TensorFlow using NiftyNet [4]. It was trained by using the Adam
optimiser, starting at a learning rate of 10-3, with a minibatch of 4. The
regularisation weight was set to $\alpha=2$ to ensure displacement continuity.

####  2.3.2 Strain Estimation

In USE, the strain estimates are obtained by computing the displacement field
gradient. However, direct differentiation of the displacement field is rarely
used because gradient operations generate a significant amount of noise in the
resulting strain map. We used the least-squares strain estimator (LSQSE) to
improve the elastogram signal-to-noise ratio [13]. The strain estimation was
also implemented in TensorFlow to facilitate efficient parallel computing. In
inference, the strain map prediction rate reach a total of 13 images per
second on a 12GB NVIDIA GTX-1080ti GPU.

##  3 Experiments

The performance of our method was evaluated on ultrasound simulations and in
vivo data. We compared our results with a state-of-the-art strain elastography
method called “tOtal Variation Regularization and WINDow-based time delay
estimation” (OVERWIND) [17]. The OVERWIND results were obtained by using the
publicly available MATLAB implementation, and default parameters were chosen
for the simulation experiment. The regularisation coefficients were manually
tuned for the in vivo data to globally maximize the NCC and Contrast-to-Noise
Ratio (CNR) scores of the three cases. The results on simulation were also
compared with strain estimates obtained by training our network with a
supervised loss function. The supervised loss function was used before for
time-delay estimation [25], and corresponded to the mean absolute difference
(MAE) between the network prediction and the ground truth labels. The quality
of the strain estimates were assessed with the CNR and Signal-to-Noise Ratio
(SNR), which are metrics commonly used in USE [24, 17].

####  3.0.1 Simulation Dataset

The Field-II software [11, 10] was used to generate the ultrasound images.
Each simulation consisted in a 3D rectangle of size 38x40 mm, containing a
cylindrical-shaped inclusion with a randomly assigned diameter (from 8 to 12
mm) and position. The speckle pattern typically observed in ultrasound imaging
was obtained by randomly assigning a total of 400,000 scatterers across each
digital phantom. The axial compression was assigned randomly to each phantom
and represented between 0.5% and 4% of the phantom total length. The Young’s
modulus of the inclusion was set to different values (i.e. 8, 15, 45 and 75
kPa) while the background was fixed to 25 kPa. Tissue displacements were
estimated by finite element method using the Partial Differential Equation
Toolbox from MATLAB, and were used to interpolate the scatterers position. A
total of 192 RF lines were simulated for each image, with probe central and
sampling frequency of 7 and 40 MHz, respectively. The background scatterers
were associated with a random intensity value from a Gaussian distribution, to
mimic homogeneous tissue. To increase the network’s robustness to noise and
image intensities, the inclusion scatterers intensities were either assigned
to 0 or similar to the background. Moreover, white Gaussian noise with random
signal power values, from 5 to 20 dBw, was added on each image to increase
robustness to noise. Finally, 160 ultrasound image pairs were simulated, where
100 were used for training and 30 for validation and testing, respectively.

####  3.0.2 In Vivo Dataset

We created our own in vivo dataset by collecting images of the arms in three
human volunteers. The RF data was acquired from a Cicada 128PX system equipped
with a 10 MHz linear probe from Cephasonics (Cephasonics Inc., USA). The
images were reconstructed using the delay-and-sum beamformer from SUPRA [5].
Experimental protocol consisted in acquiring sequences of images while slowly
applying an axial compression on the volunteer’s arm with the handheld
ultrasound probe. The in vivo dataset included 1300 image pairs for training
and 300 pairs for validation and testing. The three image pairs presented in
the results section were taken from the testing partition and exhibit one or
several blood vessels located in the arm.

###  3.1 Results

Figure 2 shows axial strain images estimated from three simulated image pairs
and obtained with the finite element method, OVERWIND and both supervised and
unsupervised models. The averaged CNR and SNR values of the entire testing
dataset (i.e., 30 image pairs) are displayed in table 1.

Figure 2: Comparison of three axial strain fields computed from ultrasound simulations by finite element method (i.e., ground truth), OVERWIND, the unsupervised and supervised network. OVERWIND’s regularisation parameters : $\alpha 1=\beta 1=20$ and $\alpha 2=\beta 2=8$. Table 1: Mean and standard deviation of SNR and CNR for the strain images of the simulation testing dataset obtained with finite element method, OVERWIND, the unsupervised and supervised network. |  | Ground truth
---
(mean std)
| Unsupervised
---
(mean std)
| Supervised
---
(mean std)
| OVERWIND
---
(mean std)
SNR | 7.51 (2.61) | 6.95 (2.54) | 7.79 (2.01) | 9.31 (3.51)
CNR | 9.15 (2.73) | 7.70 ( 3.8) | 4.22 (2.08) | 6.33 (3.6)

Figure 3 shows axial strains estimated from three different in vivo image
pairs. CNR and SNR values of each case can be found in table 2. The local NCC
values between the post-compression and resampled pre-compression images are
also displayed, to indicate the quality of the predicted displacement fields.
The mean of the three CNR values are 7 and 8.43 for our unsupervised model and
OVERWIND, respectively. The average SNR values are 0.31 and 0.7 for our method
and OVERWIND, respectively. Displacement field estimates from both the
simulation and in vivo experiments are available in the supplementary
material.

Figure 3: Comparison of axial strains estimated by our method and OVERWIND on three image pairs from the in vivo testing dataset . OVERWIND’s regularisation parameters : $\alpha 1=\beta 1=0.2$ and $\alpha 2=\beta 2=0.05$. Table 2: SNR and CNR of strain estimates from our unsupervised method and OVERWIND, and local NCC scores for the three image pair results from the in vivo dataset. | Unsupervised | OVERWIND
---|---|---
| Case 1 | Case 2 | Case 3 | Case 1 | Case 2 | Case 3
SNR | 0.14 | 0.07 | 0.1 | 0.12 | 0.35 | 0.23
CNR | 6.22 | 5.2 | 4.5 | 4.9 | 3.15 | 2.23
LNCC | 0.87 | 0.81 | 0.94 | 0.90 | 0.81 | 0.93

##  4 Discussion

In this work, we presented a new deep-learning approach for the estimation of
the displacement and strain maps between a pair of ultrasound RF data
undergoing an axial compression. We validated our method on both ultrasound
simulation and in vivo data . Our method is completely unsupervised and ground
truth images collected from finite element analysis were only used to assess
the performance of our method.

Our results on ultrasound simulation indicate that our method predicts strain
estimates with a significantly better CNR compared to the supervised network,
with 7.70 and 4.20 respectively. This suggests that our training loss
function, which includes a term that penalizes the strain smoothness, improves
the strain contrast. Our experiments showed comparable results to OVERWIND, a
state-of-the-art method which has, in terms of CNR, already outperformed a
previous classical approach [6] by 27.26%, 144.05%, and 49.90% on average in
simulation, phantom, and in-vivo data, respectively, as reported in [17]. In
addition, our method is fully automatic in inference while OVERWIND strain
estimation relies on the correct adjustment of its regularisation parameters.

Finally, the OVERWIND real-time performance had not been quantitatively
reported while our approach reached a strain prediction rate of about 13
frames per second on a 12GB NVIDIA GTX-1080ti GPU. The real-time inference of
our network and its ability to be trained without ground truth labels
represents a great potential for the use of learning-based methods in
ultrasound strain elastography.

####  4.0.1 Acknowledgement

This work was supported by the EPSRC [NS/A000049/1], [NS/A000050/1],
[EP/L016478/1] and the Wellcome Trust [203148/Z/16/Z], [203145/Z/16/Z]. Tom
Vercauteren is supported by a Medtronic/RAEng Research Chair [RCSRF1819/7/34].

## References

  * [1] Ara, S.R., Mohsin, F., Alam, F., Rupa, S.A., Awwal, R., Lee, S.Y., Hasan, M.K.: Phase-based direct average strain estimation for elastography. IEEE Transactions on Ultrasonics, Ferroelectrics, and Frequency Control 60(11), 2266–2283 (2013). https://doi.org/10.1109/TUFFC.2013.6644732
  * [2] Chen, H., Shi, H., Varghese, T.: Improvement of elastographic displacement estimation using a two-step cross-correlation method. Ultrasound in Medicine and Biology 33(1), 48–56 (2007). https://doi.org/10.1016/j.ultrasmedbio.2006.07.022
  * [3] Correas, J.M., Tissier, A.M., Khairoune, A., Khoury, G., Eiss, D., Hélénon, O.: Ultrasound elastography of the prostate: State of the art. Diagnostic and Interventional Imaging 94(5), 551–560 (2013). https://doi.org/10.1016/j.diii.2013.01.017
  * [4] Gibson, E., Li, W., Sudre, C., Fidon, L., Shakir, D.I., Wang, G., Eaton-Rosen, Z., Gray, R., Doel, T., Hu, Y., Whyntie, T., Nachev, P., Modat, M., Barratt, D.C., Ourselin, S., Cardoso, M.J., Vercauteren, T.: Niftynet: a deep-learning platform for medical imaging. Computer Methods and Programs in Biomedicine 158, 113–122 (2018). https://doi.org/10.1016/j.cmpb.2018.01.025
  * [5] Göbl, R., Navab, N., Hennersperger, C.: SUPRA: open-source software-defined ultrasound processing for real-time applications: A 2D and 3D pipeline from beamforming to B-mode. International Journal of Computer Assisted Radiology and Surgery 13(6), 759–767 (2018). https://doi.org/10.1007/s11548-018-1750-6
  * [6] Hashemi, H.S., Rivaz, H.: Global Time-Delay Estimation in Ultrasound Elastography. IEEE Transactions on Ultrasonics, Ferroelectrics, and Frequency Control 64(10), 1625–1636 (2017). https://doi.org/10.1109/TUFFC.2017.2717933
  * [7] Hiltawsky, K.M., Krüger, M., Starke, C., Heuser, L., Ermert, H., Jensen, A.: Freehand ultrasound elastography of breast lesions: Clinical results. Ultrasound in Medicine and Biology 27(11), 1461–1469 (2001). https://doi.org/10.1016/S0301-5629(01)00434-3
  * [8] Hu, Y., Modat, M., Gibson, E., Li, W., Ghavami, N., Bonmati, E., Wang, G., Bandula, S., Moore, C.M., Emberton, M., Ourselin, S., Noble, J.A., Barratt, D.C., Vercauteren, T.: Weakly-supervised convolutional neural networks for multimodal image registration. Medical Image Analysis 49, 1–13 (2018). https://doi.org/10.1016/j.media.2018.07.002
  * [9] Ilg, E., Mayer, N., Saikia, T., Keuper, M., Dosovitskiy, A., Brox, T.: Flownet 2.0: Evolution of optical flow estimation with deep networks. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 2462–2470 (2017). https://doi.org/10.1109/CVPR.2017.179
  * [10] Jensen, J.A.: FIELD: A program for simulating ultrasound systems. Medical and Biological Engineering and Computing 34, 351–352 (1996)
  * [11] Jensen, J.A., Svendsen, N.B.: Calculation of Pressure Fields from Arbitrarily Shaped, Apodized, and Excited Ultrasound Transducers. IEEE Transactions on Ultrasonics, Ferroelectrics, and Frequency Control 39(2), 262–267 (1992). https://doi.org/10.1109/58.139123
  * [12] Jeong, W.K., Lim, H.K., Lee, H.K., Jo, J.M., Kim, Y.: Principles and clinical application of ultrasound elastography for diffuse liver disease. Ultrasonography 33(3), 149–160 (2014). https://doi.org/10.14366/usg.14003
  * [13] Kallel, F., Ophir, J.: A least-squares strain estimator for elastography. Ultrasonic Imaging 19(3), 195–208 (1997). https://doi.org/10.1177/016173469701900303
  * [14] Kibria, M.G., Rivaz, H.: Gluenet: Ultrasound elastography using convolutional neural network. In: Stoyanov, D, Taylor, Z, Aylward, S, Tavares, JMRS, Xiao, Y, Simpson, A, et al. (eds.), MICCAI 2018, LNCS. vol. 11070, pp. 21–28. Springer International Publishing (2018). https://doi.org/10.1007/978-3-030-01045-4_3
  * [15] Kuzmin, A., Zakrzewski, A.M., Anthony, B.W., Lempitsky, V.: Multi-frame elastography using a handheld force-controlled ultrasound probe. IEEE Transactions on Ultrasonics, Ferroelectrics, and Frequency Control 62(8), 1486–1500 (2015). https://doi.org/10.1109/TUFFC.2015.007133
  * [16] Kwak, J.Y., Kim, E.K.: Ultrasound elastography for thyroid nodules: Recent advances. Ultrasonography 33(2), 75–82 (2014). https://doi.org/10.14366/usg.13025
  * [17] Mirzaei, M., Asif, A., Rivaz, H.: Combining Total Variation Regularization with Window-Based Time Delay Estimation in Ultrasound Elastography. IEEE Transactions on Medical Imaging 38(12), 2744–2754 (2019). https://doi.org/10.1109/TMI.2019.2913194
  * [18] Ophir, J., Cespedes, I., Garra, B., Ponnekanti, H., Huang, Y., Maklad, N.: Elastography: Ultrasonic imaging of tissue strain and elastic modulus in vivo. European Journal of Ultrasound 3(1), 49–70 (1996). https://doi.org/10.1016/0929-8266(95)00134-4
  * [19] Ophir, J., Céspedes, I., Ponnekanti, H., Yazdi, Y., Li, X.: Elastography: A quantitative method for imaging the elasticity of biological tissues. Ultrasonic Imaging 13(2), 111–134 (1991). https://doi.org/10.1177/016173469101300201
  * [20] Shi, H., Varghese, T.: Two-dimensional multi-level strain estimation for discontinuous tissue. Physics in Medicine and Biology 52(2), 389–401 (2007). https://doi.org/10.1088/0031-9155/52/2/006
  * [21] Sigrist, R.M., Liau, J., Kaffas, A.E., Chammas, M.C., Willmann, J.K.: Ultrasound elastography: Review of techniques and clinical applications. Theranostics 7(5), 1303–1329 (2017). https://doi.org/10.7150/thno.18650
  * [22] Varghese, T., Konofagou, E.E., Ophir, J., Alam, S.K., Bilgen, M.: Direct strain estimation in elastography using spectral cross-correlation. Ultrasound in Medicine and Biology 26, 1525–1537 (2000). https://doi.org/10.1016/S0301-5629(00)00316-1
  * [23] Varghese, T.: Quasi-Static Ultrasound Elastography. Ultrasound Clinics 4(3), 323–338 (2009). https://doi.org/10.1016/j.cult.2009.10.009
  * [24] Wang, J., Huang, Q., Zhang, X.: Ultrasound elastography based on the normalized cross-correlation and the pso algorithm. In: 2017 4th International Conference on Systems and Informatics (ICSAI). pp. 1131–1135. IEEE (2017). https://doi.org/10.1109/ICSAI.2017.8248455
  * [25] Wu, S., Gao, Z., Liu, Z., Luo, J., Zhang, H., Li, S.: Direct reconstruction of ultrasound elastography using an end-to-end deep neural network. In: Stoyanov, D, Taylor, Z, Aylward, S, Tavares, JMRS, Xiao, Y, Simpson, A, et al. (eds.), MICCAI 2018, LNCS. vol. 11070, pp. 374–382. Springer International Publishing (2018). https://doi.org/10.1007/978-3-030-00928-1_43


* (110) D. G. Madland, “Recent results in the development of a global medium-energy nucleon-nucleus optical-model potential,” in  _Proceedings of a Specialists’ Meeting on Preequilibrium Nuclear Reactions_, NEA/NEANDC(1988)245/U, pp. 103–110.