# Denoising Heat-inspired Diffusion with Insulators for Collision Free Motion
Planning

Junwoo Chang∗1, Hyunwoo Ryu∗1, Jiwoo Kim1, Soochul Yoo1
Joohwan Seo2, Nikhil Prakash2, Jongeun Choi†1,2, Roberto Horowitz2
1Yonsei University, 2University of California, Berkeley
<EMAIL_ADDRESS>
{joohwan_seo, nikhilps<EMAIL_ADDRESS>
∗Contributed Equally. $\quad{}^{\dagger}$Corresponding author

###### Abstract

Diffusion models have risen as a powerful tool in robotics due to their
flexibility and multi-modality. While some of these methods effectively
address complex problems, they often depend heavily on inference-time obstacle
detection and require additional equipment. Addressing these challenges, we
present a method that, during inference time, simultaneously generates only
reachable goals and plans motions that avoid obstacles, all from a single
visual input. Central to our approach is the novel use of a _collision-
avoiding diffusion kernel_ for training. Through evaluations against behavior-
cloning and classical diffusion models, our framework has proven its
robustness. It is particularly effective in multi-modal environments,
navigating toward goals and avoiding unreachable ones blocked by obstacles,
while ensuring collision avoidance.

##  1 Introduction

Diffusion probabilistic models have been recognized for multi-modality in
various fields [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]. While motion planning [12,
13, 14, 15] and manipulation [16, 17, 18, 19] often present complex
challenges, these diffusion models provide effective solutions [20, 21, 22,
23, 24, 25, 23, 26, 27], highlighting their growing significance in robotics.
In particular, Urain et al. [27] introduced a novel method that utilizes
guided diffusion to generate both goals and motions. Compared to the process
where goals are generated and motions are planned separately, this integrated
approach is more efficient and has the strength of generating only reachable
goals. However, this approach necessitates comprehensive knowledge of
obstacles, potentially requiring equipment such as LIDAR during the inference
time, which can be costly and time-consuming.

To overcome this issue, we introduce an alternative approach that utilizes a
diffusion model with only a single visual input, visualized in Fig.2. Rather
than relying on diffusion guidance to direct collision avoidance, our method
incorporates this behavior into the motion planning process. The core of our
approach is leveraging a diffusion model with a collision-avoiding diffusion
kernel. Drawing inspiration from the behavior in which heat avoids insulators,
similar to collision avoidance motions, we have adapted the modified heat
equation to build our diffusion kernel that mimic behaviors of heat dispersion
with insulating obstacles. As a consequence, our approach operates efficiently
on just a top-down view image during the inference, making it well-suited for
real-world scenarios.

We summarize our contributions as: (1) We introduce an approach leveraging
collision-avoiding diffusion kernel within the diffusion model, presenting a
clear distinction from Gaussian distributions, as seen in Fig.1. (2) Our
method offers an end-to-end motion planning technique that, using only a top-
down view image as input during inference, ensures collision avoidance and
generates only reachable goals. This approach eliminates the need for detailed
object state information and potential auxiliary equipment, making the process
of inferring obstacle avoidance more efficient. Our experiments highlight the
model’s capability to generate reachable goals while avoiding collisions and
exhibiting multi-modality against baseline methods.

Figure 1: Collision avoiding diffusion kernel. The comparison clearly shows
the differences between two diffusion kernels in an environment with
obstacles. (a) The collision-avoiding diffusion kernel moves without invading
any obstacles. (b) In contrast, the Gaussian diffusion kernel often runs into
obstacles, indicating a higher risk of collisions.

##  2 Background

### Problem statement.

Let’s define the state space in a $2D$ environment,
$\mathcal{X}\subset\mathbb{R}^{2}$. The subset of $\mathcal{X}$ where
obstacles exist is $\mathcal{X}_{obs}\subset\mathcal{X}$. The feasible space
is $\mathcal{X}_{free}=\mathcal{X}\backslash\mathcal{X}_{obs}$. The primary
objective is to generate the goal as well as the state sequence
$\tau=\\{x_{1},...,x_{T}\\}$ where each $x_{t}\in\mathcal{X}_{free}$. Given
only a top-down view map image, $y$, our method aims to generate the sequence
$\tau$ that reaches the goal distribution, $p_{g}(x|y)$. We solve the problem
by leveraging the concepts of the diffusion model.

### Score-based model learning and sampling.

Diffusion models work by repeatedly denoising a noisy distribution [28, 29].
Such models are able to enable the generation of sequence $\tau$ using a
visual input. We can express $p_{t}(x|y)$ as the distribution at a diffusion
step, conditioned on the given image, and $p_{0}(x|y)$ as the distribution
centered around the goal state with minimal variance, equivalent to
$p_{g}(x|y)$. Using anealed Langevin dynamics, we can generate the sequence
$\tau$, provided $p_{t}(x|y)$ excluding any $x\in\mathcal{X}_{obs}$. Training
involves adding noise to the goal distribution using a diffusion kernel,
denoted as $p_{0t}(x_{t}|x_{0},y)$. This process gradually transforms the
original distribution into a noised one given by $p_{t}(x_{t}|y)=\int
p_{0t}(x_{t}|x_{0},y)p_{0}(x_{0}|y)dx_{0}$. The model’s goal is to accurately
estimate the score $s_{\theta^{*}}(x_{t})=\nabla_{x_{t}}\log p_{t}(x_{t}|y)$,
similarly to Welling and Teh [30]. After training, the model can generate
states from the perturbed distribution using annealed Langevin dynamics
sampling [29]. With a sufficiently small stepsize and repeated samplings, it
is proven that the ultimately generated $\tilde{x}$ belongs to the goal
distribution $p_{0}(x|y)$ [29].

##  3 Goal generation with collision-avoiding diffusion kernel

### Collision avoiding diffusion kernel.

To utilize the method in Section 2, we need a diffusion kernel that guarantees
collision avoidance. While both intuition and empirical evidence suggest that
a Gaussian diffusion kernel falls short, the collision-avoiding diffusion
kernel offers a viable solution, as demonstrated in Fig.1. Our approach,
inspired by heat dispersion behavior, adopts the distribution by solving the
heat equation, which can be simply denoted by:

$\frac{\partial u}{\partial t}=\nabla\cdot\left(K(x)\nabla u\right)$ (1)

Figure 2: Architecture and overview of our method. The model processes the
visual input $y$ and time $t$ to produce an output vector field, which is then
bilinearly interpolated with the input state $x_{t}$ to obtain the score
function at $x_{t}$. To determine the next state leading to the goal, we
employ annealed Langevin dynamics sampling.

The presented equation is a partial differential equation with position-
dependent heat conductivity, $K(x)$. For our purposes, the edges of the map
and obstacles act as perfect insulators, making $K(x)$ equal to zero.
Elsewhere, a uniform heat conductivity $K(x)=\alpha$ prevails. We empirically
observed that the traditional equation was not compatible with our method and
was also time-consuming. To address these challenges, we utilized a modified
equation and solved it using the explicit method, detailed in Appendix C. To
avoid potential confusion, we represent the dispersion time step $t$ as $k$,
which is proportionate to the diffusion step, $t:k\propto t$. The resulting
distribution is considered as a collision-avoiding diffusion kernel, defined
as $p_{0t}:=u^{k}$ and $p_{0}:=u^{h}$, where $h$ is a small value of $k$.

### Training with the collision avoiding diffusion kernel.

Our model is then trained accordingly as:

$L(\theta)=\mathbb{E}_{t}\left(\lambda(t)\mathbb{E}_{x_{0}}\mathbb{E}_{x_{t}|x_{0},y_{obs}}[\|s_{\theta}(x_{t},t,y)-\nabla_{x_{t}}\log
p_{0t}(x_{t}|x_{0},y_{obs})\|^{2}_{2}]\right)$ (2)

where $y_{obs}$ is the corresponding obstacle image mask. We chose
$\lambda(t)=f(\sigma_{t})^{2}$, with $\sigma_{t}=\sqrt{2k}$ denoting the
standard deviation (std) of the Gaussian heat kernel. The function $f$
stabilizes $\sigma_{t}$ once it equals half the image size, approximating the
std of the distribution confined by the map boundaries. Leveraging the U-net
structure from Ho et al. [28], our network incorporates an added bilinear
interpolation step, enabling score estimation via score vector field from $y$
and interpolation at $x_{t}$.

### Generating states for motion planning.

Given a randomly generated map image and an initial state $x_{T}$, we
iteratively produce a sequence of states $x_{t}\in\mathcal{X}_{free}$. This
generation is achieved by addressing the reverse SDE with our trained score
model, leveraging the annealed Langevin dynamics sampling (details in Appendix
B):

$\tilde{x}_{t-1}=\tilde{x}_{t}+\frac{\alpha_{t}}{2}s_{\theta}(x_{t},t,y)+\sqrt{\alpha_{t}}z_{t},~{}~{}~{}~{}z_{t}\sim\mathcal{N}(0,I)$
(3)

where $\alpha_{t}=\epsilon\cdot\frac{\lambda(t)}{\lambda(T)}$. In practice, we
set parameters as $\epsilon=0.0008,~{}T=10$, and performed 100 annealing
iterations. Every generated state belongs to $p_{t}(x|y)$ [29], progressively
approaching the goal distribution, $p_{0}(x|y)$, with collision avoidance.

Figure 3: Experiment results of our method. In the $64\times 64$ input image,
black areas indicate obstacles, red dots illustrate states originating from
the initial distribution, and green apples mark the goals. (a) The first row
demonstrates an experiment with two multi-modal goals generated. It proves the
multi-modality of our model by moving the states toward each goal. (b) The
second row presents a similar setup, but with one goal being unreachable. It
only generates the reachable goal without any collisions with the obstacle.

##  4 Experiments

We assessed our method’s performance using a benchmark experiment on a noisy
map from Pearce et al. [22], setting it against three established baseline
models for comparison. The first model is expert-data behavior cloned (BC).
The second and third, named Gaussian diffusion and Gaussian diffusion + RRT*,
utilize a Vanilla diffusion model for goal generation, with the latter also
incorporating RRT* for planning.

Table 1: Experiment results. Success rates and KL divergence for $10$K samples, aggregated from $100$ samples across $100$ episodes. Whenever samples encountered obstacles or if there was a failure in trajectory generation, they remained stationary at their current location. | Uni-modality | Multi-modality | With unreachable goal
---|---|---|---
Models |  Success rate |  KL diverg. |  Success rate |  KL diverg. |  Success rate |  KL diverg.
BC |  100. |  6.44 |  0. |  13.32 |  0. |  14.83
Gaussian diffusion |  34.23 |  10.37 |  57.52 |  11.91 |  30.23 |  12.32
Gaussian diff + RRT* |  98.81 |  6.36 |  97.31 |  6.34 |  50.31 |  23.01
Ours |  98.5 |  6.05 |  99.98 |  6.45 |  100. |  6.46

All models under comparison employ the same U-net architecture and are
described in Appendix D. Our evaluation criteria encompass Uni-modality,
Multi-modality, and With unreachable goal generation. Uni-modality measures
performance in reaching a single goal, Multi-modality assesses success in
multi-goal scenarios, and With unreachable goal generation evaluates the
model’s ability to avoid generating unreachable goals. During the initial
experiment, a single goal exists, and two goals in the next experiment. In the
final experiment, an obstacle is added near one of the goals, making it
unreachable.

Figure 4: Graphical analysis of experiment results. (a): success rate of
reaching goal distributions. (b) KL divergence between the goal distributions
and the empirical distribution. Both indicates that our method shows robust
and good performance in all scenarios, compared to baselines.

### Analysis.

In Table 1, we detail the success rate, the count of states that successfully
reach any goal distributions, and the KL divergence, a metric measuring the
divergence between the empirical distribution of generated samples and ground-
truth goal distribution. The BC model performed well in a single-goal
environment. However, it lacked multi-modality, as a consequence of its
limitations from training with MSE loss, as noted in Pearce et al. [22]. The
Gaussian diffusion lagged in both goal attainment and divergence, as it
doesn’t account for collision avoidance. The Gaussian diffusion + RRT*
performed well in both the Uni-modality and Multi-modality tests. Yet, it
mistakenly generated the unreachable goal in the last experiment, performing
about $50\%$ success rate. In contrast, our method consistently showcased
strong performance across all experiments. As illustrated in Fig.4, our method
consistently maintains high success rate and low KL divergence, indicating a
reliable performance in achieving the desired goal distributions. The results
highlight our method’s potential performance in real-world tasks. The visual
results are displayed in Fig.3 and Appendix A.

##  5 Conclusion

We have introduced an end-to-end method with a conditioned diffusion model to
generate reachable goals without collisions with just a visual input. The core
of our method is the utilization of a heat-inspired collision-avoiding
diffusion kernel within a diffusion model to ensure collision avoidance.
Furthermore, it can effectively avoid generating unreachable goals. Thus, we
presented a framework that is user-friendly in the sense that it does not
require inference time collision checking or any auxiliary equipment. Future
work involves extending this approach to accommodate high-dimensional sensory
inputs in the field of robotics.

## Acknowledgments and Disclosure of Funding

This work was supported by the National Research Foundation of Korea (NRF)
grants funded by the Korea government (MSIT) (No.RS-2023-00221762 and No.
2021R1A2B5B01002620). This work was also supported by the Korea Institute of
Science and Technology (KIST) intramural grants (2E31570), and a Berkeley
Fellowship.

## References

  * Song et al. [2023] Jiaming Song, Qinsheng Zhang, Hongxu Yin, Morteza Mardani, Ming-Yu Liu, Jan Kautz, Yongxin Chen, and Arash Vahdat.  Loss-guided diffusion models for plug-and-play controllable generation.  2023\.
  * Meng et al. [2023] Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans.  On distillation of guided diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 14297–14306, 2023.
  * Kawar et al. [2023] Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang, Tali Dekel, Inbar Mosseri, and Michal Irani.  Imagic: Text-based real image editing with diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 6007–6017, 2023.
  * De Bortoli [2022] Valentin De Bortoli.  Convergence of denoising diffusion models under the manifold hypothesis.  _arXiv preprint arXiv:2208.05314_ , 2022.
  * Rombach et al. [2022] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 10684–10695, 2022.
  * Gao et al. [2020] Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma.  Learning energy-based models by diffusion recovery likelihood.  _arXiv preprint arXiv:2012.08125_ , 2020.
  * Somepalli et al. [2023] Gowthami Somepalli, Vasu Singla, Micah Goldblum, Jonas Geiping, and Tom Goldstein.  Understanding and mitigating copying in diffusion models.  _arXiv preprint arXiv:2305.20086_ , 2023.
  * Liu et al. [2023] Nan Liu, Yilun Du, Shuang Li, Joshua B Tenenbaum, and Antonio Torralba.  Unsupervised compositional concepts discovery with text-to-image generative models.  _arXiv preprint arXiv:2306.05357_ , 2023.
  * Chung et al. [2023] Hyungjin Chung, Jeongsol Kim, Sehui Kim, and Jong Chul Ye.  Parallel diffusion models of operator and image for blind inverse problems.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 6059–6069, 2023.
  * Trippe et al. [2022] Brian L Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi Jaakkola.  Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem.  _arXiv preprint arXiv:2206.04119_ , 2022.
  * Yim et al. [2023] Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola.  Se (3) diffusion model with application to protein backbone generation.  _arXiv preprint arXiv:2302.02277_ , 2023.
  * Kavraki et al. [1996] Lydia E Kavraki, Petr Svestka, J-C Latombe, and Mark H Overmars.  Probabilistic roadmaps for path planning in high-dimensional configuration spaces.  _IEEE transactions on Robotics and Automation_ , 12(4):566–580, 1996.
  * LaValle [1998] Steven LaValle.  Rapidly-exploring random trees: A new tool for path planning.  _Research Report 9811_ , 1998.
  * Ni and Qureshi [2022] Ruiqi Ni and Ahmed H Qureshi.  Ntfields: Neural time fields for physics-informed robot motion planning.  _arXiv preprint arXiv:2210.00120_ , 2022.
  * Ni and Qureshi [2023] Ruiqi Ni and Ahmed H Qureshi.  Progressive learning for physics-informed neural motion planning.  _arXiv preprint arXiv:2306.00616_ , 2023.
  * Seo et al. [2023] Joohwan Seo, Nikhil Potu Surya Prakash, Xiang Zhang, Changhao Wang, Jongeun Choi, Masayoshi Tomizuka, and Roberto Horowitz.  Robot manipulation task learning by leveraging se (3) group invariance and equivariance.  _arXiv preprint arXiv:2308.14984_ , 2023.
  * Kim et al. [2023] Jiwoo Kim, Hyunwoo Ryu, Jongeun Choi, Joohwan Seo, Nikhil Potu Surya Prakash, Ruolin Li, and Roberto Horowitz.  Robotic manipulation learning with equivariant descriptor fields: Generative modeling, bi-equivariance, steerability, and locality.  In _RSS 2023 Workshop on Symmetries in Robot Learning_ , 2023.
  * Ryu et al. [2022] Hyunwoo Ryu, Hong-in Lee, Jeong-Hoon Lee, and Jongeun Choi.  Equivariant descriptor fields: Se (3)-equivariant energy-based models for end-to-end visual robotic manipulation learning.  _arXiv preprint arXiv:2206.08321_ , 2022.
  * Simeonov et al. [2022] Anthony Simeonov, Yilun Du, Andrea Tagliasacchi, Joshua B Tenenbaum, Alberto Rodriguez, Pulkit Agrawal, and Vincent Sitzmann.  Neural descriptor fields: Se (3)-equivariant object representations for manipulation.  In _2022 International Conference on Robotics and Automation (ICRA)_ , pages 6394–6400. IEEE, 2022.
  * Janner et al. [2022] Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine.  Planning with diffusion for flexible behavior synthesis.  _arXiv preprint arXiv:2205.09991_ , 2022.
  * Xiao et al. [2023] Wei Xiao, Tsun-Hsuan Wang, Chuang Gan, and Daniela Rus.  Safediffuser: Safe planning with diffusion probabilistic models.  _arXiv preprint arXiv:2306.00148_ , 2023.
  * Pearce et al. [2023] Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, et al.  Imitating human behaviour with diffusion models.  _arXiv preprint arXiv:2301.10677_ , 2023.
  * Pearce and Zhu [2022] Tim Pearce and Jun Zhu.  Counter-strike deathmatch with large-scale behavioural cloning.  In _2022 IEEE Conference on Games (CoG)_ , pages 104–111. IEEE, 2022.
  * Carvalho et al. [2022] Joao Carvalho, Mark Baierl, Julen Urain, and Jan Peters.  Conditioned score-based models for learning collision-free trajectory generation.  In _NeurIPS 2022 Workshop on Score-Based Methods_ , 2022.
  * Carvalho et al. [2023] Joao Carvalho, An T Le, Mark Baierl, Dorothea Koert, and Jan Peters.  Motion planning diffusion: Learning and planning of robot motions with diffusion models.  _arXiv preprint arXiv:2308.01557_ , 2023.
  * Ryu et al. [2023] Hyunwoo Ryu, Jiwoo Kim, Junwoo Chang, Hyun Seok Ahn, Joohwan Seo, Taehan Kim, Jongeun Choi, and Roberto Horowitz.  Diffusion-edfs: Bi-equivariant denoising generative modeling on se (3) for visual robotic manipulation.  _arXiv preprint arXiv:2309.02685_ , 2023.
  * Urain et al. [2023] Julen Urain, Niklas Funk, Jan Peters, and Georgia Chalvatzaki.  Se (3)-diffusionfields: Learning smooth cost functions for joint grasp and motion optimization through diffusion.  In _2023 IEEE International Conference on Robotics and Automation (ICRA)_ , pages 5923–5930. IEEE, 2023.
  * Ho et al. [2020] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in neural information processing systems_ , 33:6840–6851, 2020.
  * Song and Ermon [2019] Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _Advances in neural information processing systems_ , 32, 2019.
  * Welling and Teh [2011] Max Welling and Yee W Teh.  Bayesian learning via stochastic gradient langevin dynamics.  In _Proceedings of the 28th international conference on machine learning (ICML-11)_ , pages 681–688, 2011.
  * Gerald [2004] Curtis F Gerald.  _Applied numerical analysis_.  Pearson Education India, 2004.

## Appendix

##  Appendix A Benchmark experiment results

Figure 5: Performance evaluation of the BC model in terms of multi-modality.
Black areas denote obstacles, red dots represent states initiated from the
initial distribution, and green apples mark the goals. (a) The first row
presents a scenario with one reachable goal generated. (b) The second row
displays a test similar to the first but with multi-modal goals generated. It
indicates that the behavior-cloning method performs well in uni-modal
scenarios, but lacks the multi-modality. Figure 6: Performance evaluation of
the Gaussian model in sampling reachable goals. In the visual representation,
black areas denote obstacles, blue dots illustrate states generated by the
Gaussian model from the initial distribution, and green apples mark the goals.
(a) The first row demonstrates an experiment with two reachable multi-modal
goals generated. (b) The second row presents a scenario similar to the first,
but with one unreachable goal. We use this method directly for planning or add
another motion planner to move the states toward the generated goals. The
result proves that the Gaussian diffusion model performs the multi-modality,
but cannot assure reachable goal generation or collision avoidance.

##  Appendix B Detailed implementation

To address convergence speed issues, Ryu et al. [26] modifies the original
Langevin dynamics sampling equation as follows:

$\tilde{x}_{t-1}=\tilde{x}_{t}+\frac{\epsilon}{2}s_{\theta}(x_{t},t)t^{k_{1}}+t^{\frac{k_{1}+k_{2}}{2}}z_{t},~{}~{}~{}~{}z_{t}\sim\mathcal{N}(0,I)$
(4)

This approach exhibits impressive performance in control tasks, particularly
in state sampling without collisions, mirroring our scenario. While they
modify $\alpha_{t}$ to $t^{k_{1}}$ for their purposes, we opt for
$\alpha_{t}^{k_{1}}$ for clarity, making the temperature term adapt as
$\alpha_{t}^{k_{2}}$. Given that in our configuration
$\alpha_{t}=\epsilon\cdot\frac{\lambda(t)}{\lambda(T)}$ and
$\alpha_{t}\rightarrow 0$ as $t\rightarrow 0$, our adapted Langevin dynamics
sampling equation can be expressed as:

$\tilde{x}_{t-1}=\tilde{x}_{t}+\frac{\alpha_{t}^{k_{1}}}{2}s_{\theta}(x_{t},t,y)+\alpha_{t}^{\frac{k_{1}+k_{2}}{2}}z_{t},~{}~{}~{}~{}z_{t}\sim\mathcal{N}(0,I)$
(5)

In practice, we choose $k_{1}=0.6$ and $k_{2}=0.4$ for faster convergence.

##  Appendix C Explicit method and revised heat equation

It is able to solve the heat equation by utilizing the standard explicit
method [31]. Let’s discretize $x,y,t$ as: $x_{i}=i\Delta x$, $y_{j}=j\Delta
y$, and $k=n\Delta k$. For numerical stability, it is essential to satisfy
$\Delta k\leq\frac{\Delta x^{2}}{4\alpha}$, under the assumption that $\Delta
x=\Delta y$. By choosing $\Delta x=1$ to apply in $2D$ pixel space and $\Delta
k=\frac{1}{4\alpha}$, the numerical approach can be denoted as:

$u^{n+1}_{i,j}=u^{n}_{i,j}+K_{i,j}\left(u^{n}_{i+1,j}+u^{n}_{i-1,j}+u^{n}_{i,j+1}+u^{n}_{i,j-1}-4u^{n}_{i,j}\right)$
(6)

The given equation has limitations: it is time consuming since the heat value
increases at a slow rate, leading to slow dispersion. Furthermore, empirical
analysis suggests it is not appropriate for our method. Therefore, we have
made modifications to the equation:

$u^{n+1}_{i,j}=u^{n}_{i,j}+K(i,j)\left(u^{n}_{i+1,j}+u^{n}_{i-1,j}+u^{n}_{i,j+1}+u^{n}_{i,j-1}-Vu^{n}_{i,j}\right)$
(7)

where $V$ denotes the number of valid neighboring states, specifically
excluding insulators such as map edges and obstacles. After computing, the
result is smoothed using a Gaussian filter to mitigate numerical inaccuracies.
Subsequently, it is normalized by dividing by its sum, transforming it into a
collision-avoiding diffusion kernel. We set $k$ in a range from 12.5 to
3612.5, with an exponential increase aligning with the progression of $t$.
Here, we can interpret $p_{0t}$ as the heat distribution at time step $k$.

##  Appendix D Baseline pseudo codes

Algorithm 1 Behavior Cloning model sampling

1:input image $y$, stationary number $t$, sampling timesteps $N$

2:Initial states $x_{0}$

3:for $i=0$ to $N$ do

4: $dx\leftarrow BC(y,t,x_{i})$

5: $x_{i+1}\leftarrow x_{i}+dx$ $\triangleright$ Motion Planning return
$x_{N}$

Algorithm 2 Gaussian Diffusion sampling

1:input image $y$, step size $\epsilon$

2:Initial states $x_{T}$

3:for $t=T$ to $1$ do $\triangleright$ Motion Planning with Each State
Generated

4: $s_{\theta}\leftarrow Gaussian(y,t,x_{t})$

5: $x_{t-1}\leftarrow x_{t}+\frac{\epsilon}{2}s_{\theta}+\sqrt{\epsilon}z_{t}$
$\triangleright$ $z_{t}\sim N(0,I)$ return $x_{0}$

Algorithm 3 Gaussian Diffusion + RRT* sampling

1:input image $y$, step size $\epsilon$

2:Initial states $x_{T}$

3:for $t=T$ to $1$ do $\triangleright$ Goal Sampling

4: $s_{\theta}\leftarrow Gaussian(y,t,x_{t})$

5: $x_{t-1}\leftarrow x_{t}+\frac{\epsilon}{2}s_{\theta}+\sqrt{\epsilon}z_{t}$
$\triangleright$ $z_{t}\sim N(0,I)$

6:$x_{0}\leftarrow RRT^{*}(x_{0},x_{T})$ $\triangleright$ Motion Planning
toward the Generated Goal

7:return $x_{0}$


11institutetext: CuraCloud corporation, U.S.A22institutetext: Department of
Computer Science, Chengdu University of Information Technology, P.R. China,
610225

# Robust Multimodal Image Registration using Deep Recurrent Reinforcement
Learning ††thanks: Supported in part by the National Natural Science
Foundation of China under Grant 61602065, Sichuan province Key Technology
Research and Development project under Grant 2017RZ0013, Scientific Research
Foundation of the Education Department of Sichuan Province under Grant
No.17ZA0062; J201608 supported by Chengdu University of Information and
Technology (CUIT) Foundation for Leaders of Disciplines in Science, project
KYTZ201610 supported by the Scientific Research Foundation of CUIT.

Shanhui Sun 11 0000-0001-9841-8592 Jing Hu 22 0000-0003-0921-0592 Mingqing Yao
22 0000-0002-5966-5440 Jinrong Hu 22 0000-0001-7732-8141 Xiaodong Yang 22
0000-0002-0973-2537 Qi Song 11 Xi Wu The first two authors contributed equally
to this paper22 0000-0002-7659-1631

###### Abstract

The crucial components of a conventional image registration method are the
choice of the right feature representations and similarity measures. These two
components, although elaborately designed, are somewhat handcrafted using
human knowledge. To this end, these two components are tackled in an end-to-
end manner via reinforcement learning in this work. Specifically, an
artificial agent, which is composed of a combined policy and value network, is
trained to adjust the moving image toward the right direction. We train this
network using an asynchronous reinforcement learning algorithm, where a
customized reward function is also leveraged to encourage robust image
registration. This trained network is further incorporated with a lookahead
inference to improve the registration capability. The advantage of this
algorithm is fully demonstrated by our superior performance on clinical MR and
CT image pairs to other state-of-the-art medical image registration methods.

###### Keywords:

Multimodal image registration Reinforcement learning reward function lookahead
inference.

##  1 Introduction

Image registration is a basic yet important pre-process in many applications
such as remote sensing, computer-assisted surgery and medical image analysis
and processing. In the context of brain registration, for instance, accurate
alignment of the brain boundary and corresponding structures inside the brain
such as hippocampus are crucial for monitoring brain cancer development [1].
Although extensive efforts have been made over three decades, image
registration remains an open problem given the complexity clinical situations
faced by the algorithms.

The core of image registration is to seek a spatial transformation that
establishes pixel/voxel correspondence between a pair of fixed and moving
images under rotation, scaling and translation transformation condition.
Conventionally, the mapping between two images is obtained by minimizing an
objective function with regard to some similarity criterion [1]. Therefore,
two factors are fundamental: image feature representations and similarity
measure [2]. Common image features include gradient, edge, geometric shape and
contour, image skeleton, landmark, response of Gabor filter, or the intensity
histogram. Recently, local invariant features have been widely applied to
image registration [3]. As for similarity measure, sum-of-squared-differences,
correlation coefficient, correlation ratio and mutual information are commonly
used [4]. Three-dimensional extensions of these similarity measures have also
been proposed to facilitate medical image registration.

Regarding the complexity of visual appearance of an image in presence of
noise, outliers, bias field distortion, and spatially varying intensity
distortion, defining the above-mentioned two factors appropriately is a
challenging task. Although extensive studies have been carried out on the
design of these two factors, they are somewhat designed manually and cannot be
quite adaptable in the wide range of image modalities [1, 4, 5]. To this end,
inspired by the success of neural architectures and deep learning for their
strong semantic understanding, several works have generated image feature or
similarity metric from scratch by using a deep convolutional neural network,
without attempting to use any hand-crafted formations [6, 7, 8, 9, 10].

Figure 1: Varying human anatomy in 2D multimodal image registration. Top: MR
images (fixed image). Bottom: CT images (moving image). The goal of image
registration is to align the moving image to the fixed image.

Recently, another type of approaches has also been proposed with focus on
predicting registration parameters directly [5]. High parameter dimensionality
and the non-linearity between image appearance and registration parameters
render predictions non-trivial. More specifically, high dimensionality of data
and parameter space (e.g. nine degree-of-freedom) challenge 3D medical image
registration; while the challenges of 2D medical image registration are huge
variability in appearances and shapes as indicated in Fig. 1.

Within the notion of the third type, we cast image registration as a
sequential decision-making framework, in which both feature representations
and similarity metric are implicitly learned from deep neural networks. An
overview of our framework is shown in Fig. 2. An artificial agent explores the
space of spatial transformation and determines the next best transformation
parameter at each time step. Unlike other similar image registration methods
that either use pre-trained neural network or learn the greedy exploration
process in a supervised manner [11, 12], our method explores the searching
space freely by using a multithread actor-critic scheme (A3C). Our major
contributions are:

  * •

We present a novel reinforcement learning (RL) framework for image
registration utilizing a combined policy and value network. This actor-critic
network can explore transformation parameter spaces freely, thereby avoiding
the local minima when a pose is far from initial displacement.

  * •

To the best of our knowledge, this is the first RL-based image registration
method that handle similarity transformation. To cope with the transformation
parameter unit discrepancy, we introduce a new reward function driven by
landmark error. Experiments also suggest that using such reward function helps
convergence.

  * •

A Monte Carlo rollout strategy is performed as a look-ahead guidance in the
test phase, whose “terminating” prediction is not trivial due to unknown
terminal states.

Figure 2: Overview of our proposed framework. The upper part illustrates a
sequential decision making driven by the agent who interacts with environment.
The environment produces next observation image. The left bottom part depicts
registration agent using CNN-LSTM network. The right bottom shows action
types.

##  2 Related Work

###  2.1 Reinforcement learning

In a classical reinforcement learning (RL) framework, an agent interacts with
an environment in a sequence of observations, rewards and actions. At each
time-step $t$, the agent selects an action $a_{t}$ from a set of pre-defined
actions, so as to maximize the future rewards $R_{t}$. The action is then
passed to the environment and modifies its internal state and gives out
immediate reward $r_{t}$ of this input action. This mechanism makes RL
suitable for decision-making problems. The most impressive work is from
DeepMind [13], who has been able to train an agent to play the game of Go and
achieve superhuman performance. Other research fields in computer vision also
witness the success of RL [14, 15, 16, 17, 18]. For instance, an active object
tracking system was proposed recently based on an actor-critic reinforcement
learning model [16], and this kind of actor-critic model was also validated to
be effective in an image captioning scenario [17].

In general, there are two types of RL methods: on-policy and off-policy. The
former includes Sarsa, n-step method and actor-critic method and the latter
includes Q-learning. In this paper, the framework of actor-critic is adopted.
More specifically, the policy function is called an actor, which takes actions
based on the current policy
$\pi\left({{a_{t}}\left|{{s_{t}};\theta}\right.}\right)$; the value function
is called a critic, which serves as a baseline to evaluate the quality of the
action by returning the state value $V\left({{s_{t}};{\theta_{v}}}\right)$ for
the current state under policy $\pi$. To reduce the variance in policy
gradient, an advantage function
$A\left({{a_{t}};{s_{t}}}\right)={R_{t}}-V\left({{s_{t}};{\theta_{v}}}\right)$
is used for action $a_{t}$ at state $s_{t}$, where the expected future return
$R_{t}$ is calculated as a discounted sum of future rewards up to $T$ time
steps with discount factor $\gamma\in\left({0,1}\right]$:
${R_{t}}=\sum\nolimits_{t^{\prime}=t}^{t+T-1}{{\gamma^{t^{\prime}-t}}{r_{t^{\prime}}}}$.
Parameters $\theta$ of policy network and parameters ${\theta_{v}}$ of value
network are updated as follows:

$\theta\leftarrow\theta+\tau{\nabla_{\theta}}\log\pi\left({{a_{t}}\left|{{s_{t}}}\right.;\theta}\right)\left({{R_{t}}-V\left({{s_{t}};{\theta_{v}}}\right)}\right)+\beta{\nabla_{\theta}}H\left({\pi\left({{s_{t}};\theta}\right)}\right)$
(1)
${\theta_{v}}\leftarrow{\theta_{v}}-\frac{1}{2}\tau{\nabla_{{\theta_{v}}}}{\left({{R_{t}}-V\left({{s_{t}};{\theta_{v}}}\right)}\right)^{2}}$
(2)

where $\tau$ is the learning rate, $H\left(\cdot\right)$ is the entropy and
$\beta$ is a regularization factor.

Among the actor-critic frameworks, A3C [19] is the one that employs
asynchronous parallel threads. Multiple threads run at the same time with
unrelated copies of the environment, each generating its own sequences of
training samples and maintaining an independent environment-agent interaction.
The network parameters are shared across the threads and updated every $T$
time steps asynchronously using Eqs. (1) and (2) in each thread. A3C is
reported to be fast yet stable [16].

###  2.2 Image registration

Early image registration processes can be roughly classified into two types:
intensity-based methods and feature-based methods. The former, focusing on the
image’s gray spaces, maximizes the similarity between pixel intensities to
determine the alignment between two images. Cross correlation and mutual
information are considered to be gold standard similarity measures for this
kind of methods. The latter extracts salient invariant image features and uses
the correlation between those features to determine the optimal alignment. The
main difficulty of early registration methods comes from the great variability
of image when captured by different physical principles and environments,
which translates in the lack of a general rule for images to be represented
and compared. Motivated by the great advancement in computer vision triggered
by deep learning technologies like convolutional neural networks (CNNs),
several researchers have proposed to apply such techniques to the field of
image registration. For example, Wu et al. [20] combined CNN with independent
subspace analysis, the learned image features were then used to replace the
handcrafted features in HAMMER registration model. Simonovsky et al. [7]
employed CNNs to estimate a similarity cost between two patches from differing
modalities.

Recently, another way to formulate image registration problem is to directly
predict transformation parameters [5]. Based on large deformation
diffeomorphic metric mapping (LDDMM) registration framework, Yang et al. [21]
designed a deep encoder-decoder network to initialize the momentum value for
each pixel and then evolved over time to obtain the final transformation. Miao
et al. [22] proposed a CNN-based regression approach or 2D/3D images, in which
the CNNs are trained with artificial examples generated by manually adapting
the transformation parameters for the input training data. However, this
method requires a good initialization in proximity to the true poses. Liao et
al. [11] formulated image registration as a sequential action learning problem
and leveraged supervised learning to greedily choose alignment action at each
exploration step for the sake of huge parameter searching space and the
comparatively limited training data. However, the greedy searching may not be
global optimal. Ma et al. [23], the most relevant work to ours, extended work
[11] via Q-learning framework trained using reinforcement learning. Although
this work manages to search the transformation parameter space freely, a huge
amount of state-action histories have to be saved during training. This
becomes challenging when extended to 3D medical image registration.

##  3 Method

###  3.1 Problem formulation

Let ${{\bf{I}}_{f}}$ be a fixed image and ${{\mathbf{I}}_{m}}$ be a moving
image. The task of image registration is to estimate the best spatial
similarity transformation ${T_{t}}$ from ${{\mathbf{I}}_{m}}$ to
${{\bf{I}}_{f}}$. ${T_{t}}$ is parameterized by 4 parameters with 2
translations [${t_{x}}$, ${t_{y}}$], one rotation $\alpha$ and one scaling
$s$:

${T_{t}}\left({{t_{x}},{t_{y}},s,\alpha}\right)=\left[{\begin{array}[]{*{20}{c}}{s\cos\alpha}&{-s\sin\alpha}&{{t_{x}}}\\\
{s\sin\alpha}&{s\cos\alpha}&{{t_{y}}}\end{array}}\right]$ (3)

As illustrated in Fig. 2, we formulate this image registration problem as a
sequential decision making process that at each time step, the agent decides
which variable in ${T_{t}}$ should be altered so that the moving image can be
sequentially aligned to the fixed image. This process is modeled as a Markov
Decision Process (MDP) with $\left({S,A,{r_{t}},\gamma}\right)$, where $S$ is
a set of states and $A$ is a set of actions, $r_{t}$ is the reward function
the agent receives when taking a specific action at a specific state, $\gamma$
is the discount factor that controls the importance of future rewards. In
contrast to previous work [11] and [23], we solved MDP in the framework of
actor-critic using a new state approximation network, a novel reward function
and a novel inference procedure. We focus on learning the policy function
$\pi$ via deep reinforcement learning. More specifically, a new deep neural
network is devised as a non-linear function approximator for $\pi$, where
action $a_{t}$ at time $t$ can be drawn by $a_{t}\sim\pi(s_{t};\theta)$, where
$s_{t}$ is composed of ${\bf{I}}_{m}$ and ${\bf{I}}_{f}$ at time $t$.

Table 1: Details of the combined policy and value network structure. C16-8-S4 represents 16 feature maps of convolution $8\times 8$ kernel with stride 4. FC256 indicates fully connected layer with 256 output dimensions. LSTM256 indicates that the dimension of all the elements (hidden state and input node) in the LSTM unit is 256. Layer | 1 | 2 | 3 | 4 | 5 | 6
---|---|---|---|---|---|---
Parameter | C16-8-S4 | C16-4-S2 | C32-4-MS-2 | FC256 | LSTM256 | Policy (8)/Value (1)

###  3.2 State Space Approximation and Action Definition

In work [11] and [23], state space is approximated using CNN and thereby
neighbouring frames are not considered in the process of decision making. In
our work, we utilize CNN to encode image states and a long short-term memory
(LSTM) recurrent neural network to encode hidden states between neighbouring
frames. Fig. 3 illustrates overview of the proposed network architecture. The
input of the network is a concatenation of ${{\mathbf{I}}_{m}}$ and
${{\bf{I}}_{f}}$. The LSTM layer has two output (FC) layers: policy function,
$\pi(s_{t};\theta)$ and state value function, $v(s_{t};\theta_{v})$ . Table 1
illustrates the details of our neural network. Each convolutional layer as
well as the fully connected layer are followed by a eLU activation function.
The action space consists of 8 candidate transformations that lead to the
change of $\pm 1$ pixel in translations, $\pm{1^{\circ}}$ for rotation, and
$\pm 0.05$ for scaling. At time $t$, after the selection of action $a_{t}$ by
the agent, the transformation matrix $T_{t+1}$ becomes ${a_{t}}\circ{T_{t}}$.

Figure 3: The architecture of the proposed neural network. The network weights
are shared for each LSTM unroll step and LSTM unroll $T$ times in our work.

###  3.3 Reward Definition

Given a certain state, a reward is generated to reflect the value of current
selected action. The definition of reward is somehow a tricky concept, as it
mimics human learning. A wrong definition of reward will lead, with a high
probability, to unsatisfactory learning results [24]. It is a natural
intuition for most deep learning-based image registration methods that the
reward function should encourage the transformation matrix, generated from the
sequential actions made by the agent, to be closed to the ground truth.
Therefore, the reward function in work [11] and [23] is inversely proportional
to Euclidean distance between these two matrices.

However, considering scaling effect causes parameter unit discrepancy in our
framework, since the scaling is in a comparatively small unit, and its
parameter change is only $\pm 0.05$ whereas the parameter change for
translation and rotation is $\pm 1$. This implies that the Euclidean distance
for scaling in two transformation matrices can be small, despite the rotation
and translation being large. Therefore, a large reward is still obtained even
scaling is not well registered. Although using a scalar to weight the
contribution between scaling and other two parameters can mitigate this
effect, the value of scalar should be carefully tuned. Instead of fine-tuning
a specific weight function, we alternatively using landmark error to define
the reward function. In our method, the key-points selected by scale invariant
feature transform (SIFT) features [25] are viewed as landmarks. These landmark
reference set ${{\mathbf{p}}_{G}}$ are computed from the ground truth of
moving image. Afterwards, they are warped (${{\mathbf{\tilde{p}}}_{G}}$) using
the perturbation transformation matrix. Then, for each action $a_{t}$, the
warped landmarks are transformed back using $T_{t+1}$. The Euclidean distance
$D$ between the transformed landmarks and the corresponding ground truth is
used to define the reward for action $a_{t}$:

${r_{t}}=-D=-\frac{1}{{{}^{\\#}\left\\{{{{\mathbf{p}}_{G}}}\right\\}}}\sum\limits_{i}{{{\left\|{{p_{i}}-{{\tilde{p}}_{i}}\circ{T_{t+1}}}\right\|}_{2}}},{p_{i}}\in{{\mathbf{p}}_{G}},{\tilde{p}_{i}}\in{{\mathbf{\tilde{p}}}_{G}}$
(4)

where ${p_{i}}$ and ${\tilde{p}_{i}}$ are the landmark points, $\circ$ denotes
the align operator, ${}^{\\#}\left\\{{}\right\\}$ calculates the number of
points. In addition, if $D$ is smaller than a threshold, we assume that the
terminal is triggered, and a terminal reward is set in this situation.

###  3.4 Training using Reinforcement Learning

We train the agent in an end-to-end manner using reinforcement learning
algorithm similar to A3C proposed in [19]. However, rather than running
different copies of a single game in [19], each thread runs with a different
image pair in our work. The episode starts from a pair of images and ends with
a terminal state (image pair aligned) or reaches a maximum episode length
($M_{e}$) in our work. Since
$\pi\left({\cdot\left|{{s_{t}};\theta}\right.}\right)$ and
$V\left({{s_{t}};{\theta_{v}}}\right)$ are combined in our method, parameters
of the agent are represented by
$\Theta=\left\\{{\theta,{\theta_{v}}}\right\\}$. Gradients of $\Theta$ are
backpropagated from the actor-critic outputs to the lower-level layers. Note
that different training episodes with similar initial poses would cause over-
fitting, since small change of actions make two successive states strongly
correlated. To handle this problem, we randomly perturb initial moving image
at every training episode.

###  3.5 Training using Supervised Learning

For better comparison, we also train the neural network (to more specific, the
policy network) via supervised learning. Under this strategy, the agent at
each time step selects the “optimal” action ${a_{t}}$ that has the minimal
distance for the new transformation ${a_{t}}\circ{T_{t}}$ to the ground truth
transformation ${T_{g}}$ among all the legal actions. Recall that in
reinforcement learning, the action $a_{t}$ is selected according to the output
probability distribution over all legal actions. Therefore, in supervised
learning strategy, only the shortest path for action space is explored.

To train the neural network, the policy network
$\pi\left({\cdot\left|{{s_{t}};\theta}\right.}\right)$ is trained by
minimizing cross entropy loss:
$-\sum\nolimits_{t=1}^{T}{\log\pi\left({{a_{t}}\left|{{s_{t}}}\right.;\theta}\right)}$
and the value network $V\left({{s_{t}};{\theta_{v}}}\right)$ is trained by
minimizing the mean squared loss
${\left({{R_{t}}-V\left({{s_{t}};{\theta_{v}}}\right)}\right)^{2}}$.

###  3.6 Inference with Monte Carlo Rollout

For training, the agent learns a registration policy that maps the current
state to the optimal action. For inference in test, the agent applies the
learned policy to approach the correct alignment. The stopping criteria in
test phase is based on value network prediction $v_{t}\geq{trs}$, where
$v_{t}$ is a predict value at time $t$ and $trs$ is a value threshold close to
the terminal reward. However, it is not trivial to find a good $trs$ from the
critic network as a termination since terminal state is unknown in the test
phase, so this value can only be determined empirically on training data with
new random perturbations. But, we often observe that either transformation
parameters jitter around specific values or the terminal is not reachable in
the test phase. To handle this challenge, a Monte Carlo (MC) method is
proposed in the test phase to simulate multiple searching paths, so as to
predict a better action. Note that we only perform MC rollout when $v_{t}$
reaches $trs$; otherwise the action is the one that with a highest policy
probability.

To perform MC rollout, a number $N_{mc}$ of trajectories are simulated from
the state $s_{t}$, whose $v_{t}\geq{trs}$, with a fixed searching depth
$D_{mc}$. For each trajectory $i$, the agent at first randomly selects an
action, then subsequent actions $a_{t^{\prime}}$ are chosen following
$a_{t^{\prime}}\sim\pi(s_{t^{\prime}};\theta)$ and at the same time obtained
the corresponding values by value network. Suppose that the trajectory $i$
finally gets a transformation matrix with parameters
$\left[{{t_{xi}},{t_{yi}},{s_{i}},{\alpha_{i}}}\right]$. Moreover,
accumulating the values along trajectory $i$ obtains the total path value
$V_{i}=\sum_{ri=t}^{t+D_{mc}}v_{ri}$, where $ri$ is an explored node at
trajectory $i$. We computed the final parameters each based on weighted
average:

$\left[{{t_{x}},{t_{y}},s,\cos{\alpha},\sin{\alpha}}\right]=\frac{1}{{\sum\nolimits_{i=1}^{N_{mc}}{{V_{i}}}}}\sum\nolimits_{i=1}^{N_{mc}}{{V_{i}}\times\left[{{t_{xi}},{t_{yi}},{s_{i}},\cos{\alpha_{i}},\sin{\alpha_{i}}}\right]},$
(5)

where $\alpha$ can be computed from $\cos{\alpha}$ and $\sin{\alpha}$. The
idea behind this formula is that it forms a Monte Carlo importance sampling
along different trajectories, and the empirical mean results an approximation
to the expected value.

##  4 Experiment and Results

####  4.0.1 Dataset

In this paper, we apply our proposed framework to multi-modality image
registration problem. The dataset used in our experiment contains 100 paired
axis view MR images and CT images each from 99 patients diagnosed as
nasopharyngeal carcinoma. The original CT and MR images have different
resolutions, i.e. CT has the resolution of $0.84\times 0.84\times 3$ mm and MR
has the resolution of $1\times 1\times 1$ mm. In term of image preprocessing,
these CT images were resampled to an isotropic resolution of 1mm. In addition,
since the ground truth alignment of these two data modalities is unfortunately
not easily obtainable, the standard of alignment is estimated using an off-
the-shelf toolbox Elastix [28] for the sake of efficiency. More specifically,
the MR image is treated as a fixed image while the CT image is treated as a
moving image. Pre-registration for all MR and CT pairs for the same patient
was carried out in 3D by Elastix with standard parameters. Taking into
consideration that adjacent slices are strongly correlated, only six slices
out of 100 are selected for each modality to exhibit considerable variation
between images, which ends up with 594 MR-CT pairs for our experiments.

We randomly selected 79 patients as the training data resulting 474 image
pairs, and 20 patients as testing data resulting 120 image pairs. In training
stage, at each episode, the moving images are randomly perturbed with
translation $[-25,25]$ pixels with a step $1$ pixel for $x$ and $y$ axes
respectively, a rotation $[-{30^{\circ}},{30^{\circ}}]$ with a step
$1^{\circ}$, and a scaling factor $[0.75,1.25]$ with a step $0.05$. The
training procedure thus sees $160000$ initial image pairs ($20000$ training
episodes $\times{8}$ threads). For testing, two data sets were generated using
different moving image perturbation distribution. Each moving image has 64
perturbations. As a result, we have 7680 image pairs in each data set.
“$E_{1}$” is a data set consisting of moving images that are perturbed using
the same perturbation priors as those used for generating training data.
“$E_{2}$” is a data set consisting of moving images that are perturbed using a
larger range: translation $[-30,30]$ pixels for $x$ and $y$ axes respectively,
a rotation $[-{45^{\circ}},{45^{\circ}}]$, and a scaling factor $[0.75,1.25]$.
All the training images and testing images are resized to $168\times 168$ in
our experiments.

####  4.0.2 Methods used for comparison

To evaluate the effectiveness of our method, we compare it with a SIFT-based
image registration method [25], a pure SL method (named as ”pure SL”) which
contains a feedforward network with CT and MRI as inputs and transformation
matrix as output (used a same CNN network architecture as in our RL model), as
well as several variations of the proposed framework so as to demonstrate our
contributions in agent training, reward function and lookahead inference.
Recall that in our method the policy and value networks are trained by RL
(named as “RL”), a landmark error based method is proposed for reward function
(named as “LME”), and Monte Carlo rollout strategy is used for prediction in
the last step (named as “MC”). Three corresponding substitutes are also used
in later experiments, including agent trained with SL (named as “SL”), reward
function calculated by the transformation matrix distance (named as “matrix”),
and no lookahead inference. In general, the methods used in our experiments
are: [SIFT, pure SL, RL-matrix, SL-matrix, RL-matrix-MC, SL-matrix-MC, RL-LME,
SL-LME, RL-LME-MC (the proposed method), SL-LME-MC]. Note that in all
experiments, target registration error (TRE) is used as a quantitative
measure, which is the distance after registration between corresponding points
not used in transformation matrix calculation [26].

####  4.0.3 Training and testing details

We trained the agent with 8 asynchronous threads, and with the Adam optimizer
at an initial learning rate of 0.0001, $\gamma=0.99$ and $\beta=0.1$. Training
episodes had the maximum length of 20000 cycles, and the maximum length
($M_{e}$) of each episode is 500 steps. Each episode is terminated whenever
the agent reaches a distance terminal (in terms of landmark error or
transformation matrix error) or the maximum episode length. The distance
threshold for terminal is 1 for both kinds of errors, and the terminal reward
in RL training is 10. The boostrapping length (the network weights updating
frequency) in training procedure is 30. Each thread randomly selects a new
training image pair at every two episodes. Stopping thresholds $trs$ are RL-
LME: 10, RL-LME-MC: 9, SL-LME: -0.05 and SL-LME-MC: -0.1. We also applied
these thresholds to their “matrix” counterparts. For MC rollout methods,
simulation number $N_{mc}=20$ and searching depth $D_{mc}=10$.

####  4.0.4 Evaluation and Results

We evaluated the models and inference variants (w/o MC) on datasets $E_{1}$
and $E_{2}$, respectively. Table 2 summarizes the quantitative results. It is
clear that all the deep-learning based image registration methods
significantly outperform the method using SIFT. It is important to note that
our method is far better than pure SL, which makes sense because the distance
between the moving image and fixed image is so large that a direct regression
learned by the latter is not feasible. The improved performance by SL-matrix,
SL-matrix-MC and SL-LME suggests that using discrete action step by step is
beneficial for the ”bad” registration scenario. Table 2 also proves that these
three components, including model “RL”, LME-based reward function and
predication strategy, all contribute to the good performance of the proposed
method. Furthermore, the proposed method is also compared with Elastix, which
is used as a gold standard to pre-align the original MR-CT pairs for our
experiments. Although we have found that this software can obtain a very
accurate alignment for each image pair, the deformation between original MR-CT
is in fact not severe. Is Elastix still a gold standard for our method when it
applies to a more challenging data, i.e. the testing dataset used in our
experiment, is worth looking into. Therefore, we performed a 2D version
Elastix registration on $E_{1}$ and $E_{2}$. Note that 2D Elastix may yield
registration results closer to ground truth, since the same cost function is
used as the one used for ground truth generation. Such a comparison is made
and summarized in Table 2. A visual comparison is also presented in Fig. 4. We
can see that Elastix fails to register images with large deformation whereas
our method is more robust to the challenging cases.

Table 2: Evaluation of methods on different datasets ($E_{1}$ and $E_{2}$) in terms of TRE. The statistics mean ($\mu$), standard deviation ($std$), median ($50^{th}$) and $90\%$ percentile ($90^{th}$) are performed. The bold numbers indicate the best performer in that column. Note that 2D Elastix results may exist bias due to it was performed using the same cost function as ground truth generation algorithm. | $E_{1}$ | $E_{2}$
---|---|---
$\mu\pm{std}$ | $50^{th}$ | $90^{th}$ | $\mu\pm{std}$ | $50^{th}$ | $90^{th}$
SIFT | 11.79$\pm$9.81 | 9.04 | 19.99 | 12.15$\pm$9.89 | 9.13 | 20.06
Elastix | 1.31$\pm$1.06 | 0.91 | 2.46 | 2.33$\pm$1.86 | 1.07 | 5.33
Pure SL | 4.59$\pm$3.00 | 3.72 | 8.28 | 6.62$\pm$5.11 | 5.11 | 13.12
SL-matrix | 2.29$\pm$2.92 | 1.81 | 3.77 | 2.99$\pm$4.72 | 2.04 | 4.83
SL-matrix-MC | 2.88$\pm$4.52 | 2.12 | 4.39 | 3.58$\pm$13.20 | 2.20 | 5.17
SL-LME | 2.60$\pm$6.62 | 1.34 | 3.15 | 3.06$\pm$7.12 | 1.41 | 4.35
SL-LME-MC | 3.10$\pm$41.67 | 1.09 | 3.03 | 3.16$\pm$20.69 | 1.19 | 4.20
RL-matrix | 1.59$\pm$1.61 | 1.32 | 2.96 | 1.95$\pm$3.55 | 1.40 | 3.41
RL-matrix-MC | 1.31$\pm$1.64 | 1.00 | 2.38 | 1.81$\pm$8.00 | 1.07 | 3.21
RL-LME | 1.66$\pm$1.56 | 1.30 | 3.12 | 1.91$\pm$2.39 | 1.41 | 3.61
RL-LME-MC (Ours) | 1.39$\pm$1.45 | 0.98 | 2.59 | 1.66$\pm$2.30 | 1.08 | 3.17
Figure 4: Examples of registration results. (a) Fixed image, (b) initial
moving image, (c) ground truth poses, results of (d) Elastix (e) SL-matrix,
(f) RL-LME-MC (Ours), (g) Pure SL.

##  5 Discussion

From the above experiments, it is clear that the proposed RL-LME-MC image
registration framework is more powerful and robust than the other methods,
which validates the effectiveness of the advocated three components. But,
besides these three factors, the network architecture (CNN+LSTM) is also worth
looking into. In this section, the importance of LSTM is firstly discussed.
After that, the above-mentioned three components of our proposed are
respectively analyzed. Last, we also talk about the difference between our
method and Elastix.

###  5.1 Network architecture analysis

Although there are several image registration methods that have been proposed
so far, to the best of our knowledge, only one work that handles registration
problem using RL solely. A Dueling Network framework was used in their work,
focusing on learning the action value for the agent. Our method, however,
alternatively uses A3C to learn both action probability and state value for
the agent. In this way, there is no need for our method to use a replay-memory
that hinders the future application of RL to 3D image registration in case of
high dimensionality. But, lacking replay-memory requires our neural network
should be elegant enough to obtain a good feature representation that
facilitates neural network training.

A CNN-LSTM network is proposed in this paper, where the LSTM is leveraged to
learn the information (hidden state) among consecutive frames that are
correlated in fact. A natural question is “how important is LSTM for our
method”. To answer it, we trained a new neural network that substitutes LSTM
with FC128 (named as “CNN-LME”). Fig. 5(a) and (b) compares this new network
with our proposed one in terms of agent’s speed reaching terminals and
cumulative reward per episode. Clearly, incorporating LSTM into CNN
facilitates a more efficient and stable convergence. We believe such
efficiency attributes to the capability of LSTM unit that takes into account
historical states when producing output. In addition, CNN-LME-MC results
$2.61\pm{3.19}$ pixels on $E_{1}$ and $3.52\pm{5.70}$ pixels on $E_{2}$ in
terms of TRE (mean$\pm$standard deviation). This result is worsen than our
approach.

###  5.2 Impacts of three components in our method

####  5.2.1 Model RL

To assess the merits of RL, it is compared with a SL-based network using the
same architecture. Although there are several reasonable doubts about the
assert that RL surpasses SL, we have found in our study that using RL achieved
higher registration accuracy than all the cases using SL, as reflected in
Table 2. It should be noted that the advantage for RL is more obvious when
handling $E_{2}$, the perturbation of which is larger than the training data.
This implies that by exploring more path enables RL a very strong learning
ability and generalization ability.

####  5.2.2 Reward function

Reward function is a fundamental factor in RL. Although in Table 2, no
significant improvement of quantitative score is observed by using LME over
matrix. It is even worse, in certain cases, that using LME obtained lower
scores (see the third and fourth rows in Table 2 for details). However, we
have found in our pilot study that for the model CNN-LME, if its reward
function is changed into matrix, this new model (CNN-matrix) fails to converge
and in fact very unstable within the same number of training steps, as shown
in Fig. 5(c) and (d). Although so far the mechanism behind this phenomenon is
uncertain, it, on the hand, proves that defining a proper reward function can
somewhat compensate for a weak network architecture.

####  5.2.3 Lookahead

According to Table 2, it is slightly better for a RL neural network than its
SL counterpart to be embedded with a prediction step. This is possibly due to
the reason that the action selected at each time step in RL is towards the
goal of maximizing the future return. Incorporating a prediction step to RL
would reinforce such intendancy. On the contrary, SL follows a greedy strategy
that only considers the local optimal action. Hence, using a lookahead
strategy implies averaging actions from future time steps, which would be
highly likely to worsen the “optimal” situation compared with using the action
at one time step. In addition, for “MC”, we made an early stopping (smaller
$ths$) for compensating Monte Carlo computation expense. However, if the agent
is stopped ”too” early, using MC rollout may not reach an optimal solution.
Due to this, an increased standard deviation can be observed at method RL-
matrix-MC on $E_{2}$ data set. But, relaxing the threshold will address this
problem.

Figure 5: Learning curves for different models. Steps per episode when agents
reach terminals for (a) RL-LME vs CNN-LME and (c) CNN-matrix vs CNN-LME.
Cumulative rewards per episode for (b) RL-LME vs CNN-LME and (d) CNN-matrix vs
CNN-LME.

###  5.3 Comparison with Elastix

Elastix [28] is a state-of-the-art medical image registration software that
has been widely used in several research projects and it is still under active
development. Both 2D and 3D image registration frameworks are provided, only
dissimilar at image processing part. Based on the observation that Elastix
performs well when the pose of moving image is in proximity to the fixed
image’s, its 3D version is used to generate the ground truth for our
experiment. In comparison, the 2D Elastix could provide us a good reference to
understand our method, although its results would be biased since it has a
similar framework to the 3D algorithm. On data set $E_{1}$, 2D Elastix
outweighs others. However, on data set $E_{2}$, where we observed the pose of
moving image is far from fixed image’s, i.e. a large part of image missed, or
in other words, a large portion is out of field of view, Elastix algorithm
fails to fulfill the alignment task, as reflected in experiment results (Fig.
4 and Table 2). Quantitatively, Elastix performs much worse than our approach
in terms of $90\%$ percentile on this large range testing data set. This is
because Elastix is basically a least square optimization, and thereby is not
robust to outliers and lack of high-level semantics understanding (e.g. parts,
shapes and temporal context) which is crucial for robust alignment. In
contrast, our proposed method performed similar on $E_{1}$ in terms of mean,
standard deviation, median and $90\%$ percentile statistics of TRE, but much
better on $E_{2}$. This implies that our method is able to perform semantic
understanding as well as extrapolation although perturbation range is not
included in the training prior.

##  6 Conclusion

In this paper, we present a new learning paradigm in the context of image
registration based on reinforcement learning. Different from previous work
that also build upon reinforcement learning, our method devised a combined
policy network and value network to respectively generate the action
probability and state value, without the need of extra storage on exploration
history. To learn this network, we use a state-of-the-art reinforcement
learning approach A3C with a novel landmark error based reward function. A
Monte Carlo rollout prediction step is also proposed to further improvement.
We evaluate our method with image slices from MR and CT scans, and find that
our method achieves state-of-the-art performance on standard benchmark.
Detailed analyses on our method are also conducted to understand its merits
and properties. Our future work includes the sub-pixel image registration,
extension to 3D registration, as well as the generalization to other computer
vision problems.

## References

  * [1] de Vos, B.D., Berendsen, F.F., Viergever, M.A., Staring, M., Išgum, I.:  End-to-end unsupervised deformable image registration with a convolutional neural network (2017)
  * [2] CY, Z., HC, Z., JF, L., SJ, S., B, L.:  Multimodal image matching based on multimodality robust line segment descriptor.  Neurocomputing 177 (2016) 290–303
  * [3] Liao, S., Chung, A.C.S.:  Feature based nonrigid brain mr image registration with symmetric alpha stable filters.  IEEE Transactions on Medical Imaging 29 (2010) 106–119
  * [4] Razlighi, Q.R., Kehtarnavaz, N., Yousefi, S.:  Evaluating similarity measures for brain image registration.  Journal of visual communication and image representation 27 (2013) 977–987
  * [5] Litjens, G., Kooi, T., Bejnordi, B.E., Arindra, A., Setio, A., Ciompi, F., Ghafoorian, M., Laak, J.A.W.M.V.D., Ginneken, B.V., Clara, I.S.:  A survey on deep learning in medical image analysis.  Medical Image Analysis 42 (2017) 60–88
  * [6] Uzunova, H., Wilms, M., Handels, H., Ehrhardt, J.:  Training cnns for image registration from few samples with model-based data augmentation.  In: International Conference on Medical Image Computing and Computer-Assisted Intervention, Springer (2017) 223–231
  * [7] Simonovsky, M., Gutiérrez-Becker, B., Mateus, D., Navab, N., Komodakis, N.:  A deep metric for multimodal registration.  In: International Conference on Medical Image Computing and Computer-Assisted Intervention, Springer (2016) 10–18
  * [8] Wu, G., Kim, M., Wang, Q., Gao, Y., Liao, S., Shen, D.:  Unsupervised deep feature learning for deformable registration of mr brain images.  In: International Conference on Medical Image Computing and Computer-Assisted Intervention, Springer (2013) 649–656
  * [9] Sokooti, H., de Vos, B., Berendsen, F., Lelieveldt, B.P., Išgum, I., Staring, M.:  Nonrigid image registration using multi-scale 3d convolutional neural networks.  In: International Conference on Medical Image Computing and Computer-Assisted Intervention, Springer (2017) 232–239
  * [10] Cao, X., Yang, J., Zhang, J., Nie, D., Kim, M., Wang, Q., Shen, D.:  Deformable image registration based on similarity-steered cnn regression.  In: International Conference on Medical Image Computing and Computer-Assisted Intervention, Springer (2017) 300–308
  * [11] Liao, R., Miao, S., de Tournemire, P., Grbic, S., Kamen, A., Mansi, T., Comaniciu, D.:  An artificial agent for robust image registration.  In: AAAI. (2017) 4168–4175
  * [12] Krebs, J., Mansi, T., Delingette, H., Zhang, L., Ghesu, F.C., Miao, S., Maier, A.K., Ayache, N., Liao, R., Kamen, A.:  Robust non-rigid registration through agent-based action learning.  In: International Conference on Medical Image Computing and Computer-Assisted Intervention, Springer (2017) 344–352
  * [13] Silver, D., Huang, A., Maddison, C.J., Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., et al.:  Mastering the game of go with deep neural networks and tree search.  nature 529 (2016) 484–489
  * [14] Bellver, M., Giró-i Nieto, X., Marqués, F., Torres, J.:  Hierarchical object detection with deep reinforcement learning.  arXiv preprint arXiv:1611.03718 (2016)
  * [15] Caicedo, J.C., Lazebnik, S.:  Active object localization with deep reinforcement learning.  In: Computer Vision (ICCV), 2015 IEEE International Conference on, IEEE (2015) 2488–2496
  * [16] Luo, W., Sun, P., Mu, Y., Liu, W.:  End-to-end active object tracking via reinforcement learning.  arXiv preprint arXiv:1705.10561 (2017)
  * [17] Ren, Z., Wang, X., Zhang, N., Lv, X., Li, L.J.:  Deep reinforcement learning-based image captioning with embedding reward.  arXiv preprint arXiv:1704.03899 (2017)
  * [18] Tan, B., Xu, N., Kong, B.:  Autonomous driving in reality with reinforcement learning and image translation.  arXiv preprint arXiv:1801.05299 (2018)
  * [19] Mnih, V., Badia, A.P., Mirza, M., Graves, A., Lillicrap, T., Harley, T., Silver, D., Kavukcuoglu, K.:  Asynchronous methods for deep reinforcement learning.  In: International Conference on Machine Learning. (2016) 1928–1937
  * [20] Wu, G., Kim, M., Wang, Q., Gao, Y., Liao, S., Shen, D.:  Unsupervised deep feature learning for deformable registration of mr brain images.  In: International Conference on Medical Image Computing and Computer-Assisted Intervention, Springer (2013) 649–656
  * [21] Yang, X., Kwitt, R., Niethammer, M.:  Fast predictive image registration.  In: Deep Learning and Data Labeling for Medical Applications.  Springer (2016) 48–57
  * [22] Miao, S., Wang, Z.J., Liao, R.:  A cnn regression approach for real-time 2d/3d registration.  IEEE transactions on medical imaging 35 (2016) 1352–1363
  * [23] Ma, K., Wang, J., Singh, V., Tamersoy, B., Chang, Y.J., Wimmer, A., Chen, T.:  Multimodal image registration with deep context reinforcement learning.  In: International Conference on Medical Image Computing and Computer-Assisted Intervention, Springer (2017) 240–248
  * [24] Zhao, D., Hu, Z., Xia, Z., Alippi, C., Zhu, Y., Wang, D.:  Full-range adaptive cruise control based on supervised adaptive dynamic programming.  Neurocomputing 125 (2014) 57–67
  * [25] Lowe, D.G.:  Distinctive image features from scale-invariant keypoints.  International journal of computer vision 60 (2004) 91–110
  * [26] Fitzpatrick, J.M., West, J.B.:  The distribution of target registration error in rigid-body point-based registration.  IEEE transactions on medical imaging 20 (2001) 917–927
  * [27] Kocsis, L., Szepesvári, C.:  Bandit based monte-carlo planning.  In: European conference on machine learning, Springer (2006) 282–293
  * [28] Stefan Klein and Marius Staring and Keelin Murphy and Max A. Viergever and Josien P.W. Pluim:  elastix: a toolbox for intensity-based medical image registration.  IEEE Transactions on Medical Imaging 29 (2010) 196 – 205
  * [29] Hava Lester and Simon R. Arridgem:  A survey of hierarchical non-linear medical image registration.  Pattern Recognition 32 (1999) 129 – 149