# Adaptive Policy Learning for Offline-to-Online Reinforcement Learning

Han Zheng1, Xufang Luo2, Pengfei Wei3,
Xuan Song4, Dongsheng Li2, Jing Jiang1 This work was done when Han Zheng was
interning with Microsoft Research Asia.Corresponding author. Email:
<EMAIL_ADDRESS>
###### Abstract

Conventional reinforcement learning (RL) needs an environment to collect fresh
data, which is impractical when online interactions are costly. Offline RL
provides an alternative solution by directly learning from the previously
collected dataset. However, it will yield unsatisfactory performance if the
quality of the offline datasets is poor. In this paper, we consider an
offline-to-online setting where the agent is first learned from the offline
dataset and then trained online, and propose a framework called Adaptive
Policy Learning for effectively taking advantage of offline and online data.
Specifically, we explicitly consider the difference between the online and
offline data and apply an adaptive update scheme accordingly, that is, a
pessimistic update strategy for the offline dataset and an optimistic/greedy
update scheme for the online dataset. Such a simple and effective method
provides a way to mix the offline and online RL and achieve the best of both
worlds. We further provide two detailed algorithms for implementing the
framework through embedding value or policy-based RL algorithms into it.
Finally, we conduct extensive experiments on popular continuous control tasks,
and results show that our algorithm can learn the expert policy with high
sample efficiency even when the quality of offline dataset is poor, e.g.,
random dataset.

## Introduction

Conventional online reinforcement learning (RL) methods (Haarnoja et al. 2018;
Fujimoto, van Hoof, and Meger 2018) usually learn from experiences generated
by online interactions with the environment. They are impractical in some
real-world applications, e.g., dialog (Jaques et al. 2019) and education
(Mandel et al. 2014), where interactions are costly. Furthermore, many offline
data have already been generated by one or more policies, but online RL agents
usually cannot utilize them directly (Fujimoto, Meger, and Precup 2019). In
contrast, supervised learning achieves remarkable successes across a range of
domains (Kotsiantis et al. 2007; Brown et al. 2020) by directly leveraging
existing large-scale datasets, such as ImageNet (Deng et al. 2009). To apply
such data-driven learning paradigm with RL objectives, many offline RL methods
have been proposed and aroused much attention recently (Levine et al. 2020).
They try to learn policies from static datasets which are pre-collected by
arbitrary policies.

Existing offline RL studies mainly focus on remedying the issue of
distribution mismatch or out-of-distribution (OOD) actions by employing a
pessimistic update scheme (Kumar et al. 2019, 2020) or in combination with
imitation learning (Fujimoto, Meger, and Precup 2019). However, when the
dataset is fixed and sub-optimal, it is nearly impossible for offline RL to
learn the optimal policy (Kidambi et al. 2020). Even worse, when a large
portion of actions or states is not covered within the training set
distribution, offline RL methods usually fail to learn good policies (Kumar et
al. 2020; Fu et al. 2020; Levine et al. 2020). Moreover, policy evaluation
under the offline RL setting is also challenging. Although some methods, such
as off-policy evaluation algorithms (Dann et al. 2014), are proposed for this
problem, they are still not ideal for the practical purpose.

Some recent works address the above issues by employing an offline-to-online
setting. Such methods (Lee et al. 2021; Nair et al. 2020) focus on pre-
training a policy using the offline dataset and fine-tuning the policy through
further online interactions. Even though their methods alleviate the above
problems, they have not considered the different advantages of offline and
online data and utilised both well. Specifically, offline data can prevent
agents from prematurely converging to sub-optimal policies thanks to the
potential data diversity, while online data can stabilize training and
accelerate convergence (Nair et al. 2020; Sutton and Barto 2011). They can
contribute to better policy learning in different ways. But prior methods only
stress one of them. For instance, Nair et al. (2020) employs a pessimistic
strategy to update the policy, which may be problematic for improving policy
performance during the online phase. Unlike Nair et al. (2020), Lee et al.
(2021) only selects near-on-policy data from the offline dataset during online
phase to alleviate the distribution shift in transition from offline to
online. Such a strategy ignores a large part of the offline dataset, whose
potential data diversity is important for learning better policies.

To tackle the above problems, in this paper, we propose that the advantage of
online and offline data should be emphasized in an adaptive way. First,
considering their different characteristics, separate updating strategies
should be employed for online and offline data, respectively. To do so, we
present a novel framework called Adaptive Policy Learning (APL) that takes
advantages of them effectively. The core idea is simple: when learning from
online data, the agent is updated in an optimistic way, and when learning from
the offline dataset, the agent is optimized by a pessimistic strategy. The
intuition behind this is that near-on-policy data can be collected via online
interactions, so an optimistic strategy is used here for better policy
improvement, and potentially useful offline datasets can be collected by
arbitrary policies, so a pessimistic strategy is used to exploit all these
data. In addition, to distinguish between offline and online data in a simple
way, we design a two-level replay buffer for the APL framework. We further
provide value-based and policy-based implementations for APL framework through
embedding state-of-the-art (SOTA) online and offline RL methods into the
framework. Experimental results demonstrate that our algorithm can learn
expert policies in most tasks with a high sample efficiency regardless of the
quality of offline datasets. More specifically, our algorithm performs better
by using only $20\%$ online interactions compared with the previous offline-
to-online method (Nair et al. 2020).

Our contributions can be summarized as below:

  * •

We propose a simple framework called Adaptive Policy Learning (APL) that
considers different advantages of offline and online data for policy learning,
and can effectively make use of them.

  * •

We further provide value-based and policy-based algorithms to implement the
APL framework, showing APL is compatible with various RL methods.

  * •

We test our algorithm on the popular continuous control tasks MuJoCo (Todorov,
Erez, and Tassa 2012) and compare it with some strong baselines. The results
clearly show that the APL framework is effective under the offline-to-online
setting.

## Related Work

Online RL In general, online RL algorithms can be divided into two categories,
i.e., on-policy and off-policy algorithms. On-policy methods (Schulman et al.
2015, 2017) update the policy using data collected by its current behavior
policy. As ignoring the logged data collected by its history behavior
policies, they usually have a lower sample efficiency than the off-policy
methods (Fujimoto, van Hoof, and Meger 2018; Chen et al. 2021a), which enable
the policy to learn from experience collected by history behavior policies.
However, they cannot learn from history trajectories collected by other
agents’ behavior policies (Fujimoto, Meger, and Precup 2019; Kumar et al.
2020). So, the need for huge online interaction makes online RL impractical
for some real-world applications, such as dialog agents (Jaques et al. 2019)
or the education system (Mandel et al. 2014).

Offline RL Offline RL algorithms assume the online environment is unavailable
and learn policies only from the pre-collected dataset. As the value
estimation error cannot be corrected using online interactions, these methods
utilize a pessimistic updating strategy to relieve the distribution mismatch
problem (Fujimoto, Meger, and Precup 2019; Kumar et al. 2019). Model-free
offline RL methods generally employ value or policy penalties to constrain the
updated policy close to the data collecting policy (Wu, Tucker, and Nachum
2019; Kumar et al. 2020; Fujimoto, Meger, and Precup 2019; He and Hou 2020;
Ghosh et al. 2022). Model-based methods use predictive models to estimate
uncertainties of states and then update the policy in a pessimistic way based
on them (Kidambi et al. 2020; Yu et al. 2020; Chen et al. 2021b). Those
offline RL methods cannot guarantee a good performance, especially when the
data quality is poor (Kumar et al. 2020). Besides, policy evaluation when the
online environment is unavailable is also challenging. Even though off-policy
evaluation (OPE) methods (Dann et al. 2014) present alternative solutions,
they are still far from perfect.

The above issues in online and offline RL motivate us to investigate the
offline-to-online setting.

Offline-to-online RL Some works focus on the mixed setting where the agent is
first learned from the offline dataset and then trained online. Nair et al.
(2020) propose an advantage-weighted actor-critic (AWAC) method that restricts
the policy to select actions close to those in the offline data by an implicit
constraint. Kostrikov, Nair, and Levine (2022) present IQL, implementing a
weighted behavioral cloning step for better online policy improvement, which
can also be used in offline-to-online setting. When online interactions are
available, such conservative designs may adversely affect the performance.
OFF2ON (Lee et al. 2021) employs a balanced replay scheme to address the
distribution shift issue. It uses offline data by only selecting near-on-
policy samples. Unlike these works, our method addresses all online and
offline data, and explicitly considers the difference between them by
adaptively applying optimistic and pessimistic updating schemes for online and
offline data, respectively. Besides, our framework is more flexible because it
can be easily applied to policy or value-based methods. Moreover, our
algorithm is much more robust to different hyper-parameters (see the
experiment section). Particularly, without careful tuning, our algorithm can
achieve better or comparable performance than prior methods, while OFF2ON uses
a specialized network architecture, and the hyper-parameters are fine-tuned.
Matsushima et al. (2021) focus on optimizing deployment efficiency, i.e., the
number of distinct data-collection policies used during learning, by employing
a behavior-regularized policy updating strategy. However, they ignore existing
offline dataset, and dose not focus on improving sample efficiency, while both
are addressed in our paper. Some works (Zhu et al. 2019; Vecerik et al. 2017;
Rajeswaran et al. 2018; Kim et al. 2013) can also learn from online
interactions and offline data. However, they need expert demonstrations
instead of any dataset, limiting their applicability.

## Preliminaries

In RL, interactions between the agent and environment are usually modelled
using Markov decision process (MDP)
$\left(\mathcal{S},\mathcal{A},p_{\text{M}},r,\gamma\right)$, with state space
$\mathcal{S}$ (state $\mathbf{s}\in\mathcal{S}$), action space $\mathcal{A}$
(action $\mathbf{a}\in\mathcal{A}$). At each discrete time step, the agent
takes an action $\mathbf{a}$ based on the current state $\mathbf{s}$, and the
state changes into $\mathbf{s}^{\prime}$ according to the transition dynamics
$p_{\text{M}}\left(\mathbf{a}^{\prime}\mid\mathbf{s},\mathbf{a}\right)$, and
the agent receives a reward $r\in\mathbb{R}$. The agent’s objective is to
maximize the return, which is defined as
$R_{t}=\sum_{n=0}^{\infty}\gamma^{n}r_{t+n}$, where $t$ is the time step, and
$\gamma\in[0,1)$ is the discounted factor. The mapping from $\mathbf{s}$ to
$\mathbf{a}$ is denoted by the stochastic policy
$\pi:\mathbf{a}\sim\pi(\cdot|\mathbf{s})$. Policy can be stochastic or
deterministic, and we use the stochastic from in this paper for generality.
Each policy $\pi$ have a corresponding action-value function
$Q^{\pi}(\mathbf{s},\mathbf{a})=\mathbb{E}_{\pi}\left[R_{t}\mid\mathbf{s},\mathbf{a}\right]$,
which is the expected return following the policy after taking action
$\mathbf{a}$ in state $\mathbf{s}$. The action-value function of policy $\pi$
can be updated by the Bellman operator $\mathcal{T}^{\pi}$:

$\mathcal{T}^{\pi}Q(\mathbf{s},\mathbf{a})=\mathbb{E}_{\mathbf{s}^{\prime}}\left[r+\gamma
Q\left(\mathbf{s}^{\prime},\pi\left(\mathbf{s}^{\prime}\right)\right)\right]$
(1)

Q-learning (Sutton and Barto 2011) directly learns the optimal action-value
function
$Q^{*}(\mathbf{s},\mathbf{a})=\max_{\pi}Q^{\pi}(\mathbf{s},\mathbf{a})$, and
such Q-function can be modelled using neural networks (Mnih et al. 2015).

In principle, off-policy methods, such as Q-learning, can utilize experiences
collected by any policy, and thus they usually maintain a replay buffer
$\mathcal{B}$ to store and repeatedly learn from experiences collected by
behavior policies (Agarwal, Schuurmans, and Norouzi 2020). Such capability
also enables off-policy methods in the offline setting by storing offline data
into the buffer $\mathcal{B}$ and not updating the buffer during learning
since no further interactions are available here (Levine et al. 2020).
However, this simple adjusting generally results in poor performance due to
bootstrapping errors from OOD actions, especially when the dataset is not
diverse (Kumar et al. 2020; Agarwal, Schuurmans, and Norouzi 2020; Fujimoto,
Meger, and Precup 2019), and this is also the problem tackled in most offline
RL works.

This paper focuses on the offline-to-online setting, where the agent learns
from both the offline and online datasets. We use off-policy methods in the
online part because they can utilize data more effectively and generally have
higher sample efficiency than on-policy ones. Thus, if no additional marks,
online RL methods refer to off-policy algorithms in the rest of this paper.

Figure 1: Learning curves for the online agent initialized with the D4RL (Fu
et al. 2020) dataset hopper-medium-replay-v0. Scores are averaged over five
random seeds and the shaded areas represent the standard deviation. The
normalized score of 100 is the average returns of a domain-specific expert
while normalized score of 0 corresponds to the average returns of an agent
taking uniformly random actions across the action space.

## Methodology

In this section, we first give an illustrative example to explain our
motivation. Then, we introduce the proposed APL framework, trying to couple
online and offline RL in an adaptive way. Finally, we present detailed
algorithms for implementing APL framework.

An illustrative example We test the SOTA offline and online RL methods under
the offline-to-online setting, i.e., Conservative Q-learning (CQL) (Kumar et
al. 2020) and Randomized Ensembled Double Q-learning (REDQ) (Chen et al.
2021a), respectively. Specifically, we first pre-train the agent with the
offline dataset for $100K$ steps. Then, the agent is fine-tuned online by
alternately conducting the interaction and updating process. The agent
interacts with the environment for $1K$ steps, and is updated for $10K$ steps.
The total online interaction steps are around $100K$.

The offline-to-online setting described above should be a favourable one for
policy learning because both diverse offline data and online interactions are
available. Therefore, we expect the following results: the initial offline
training with the fixed dataset will provide relatively good but not expert
performance for the agent, and then the agent can be further improved by the
online process since online interaction data can be obtained. At last, we may
acquire a well-performed policy, with a normalized score close to or better
than $100$ (i.e., the performance of an expert).

Nonetheless, the experiment results in Figure 1 do not show what we expected.
Specifically, starting points of the two curves show that the offline
algorithm CQL and online algorithm REDQ can obtain normalized scores greater
than $0$ (i.e., the performance of a random policy). This indicates that by
learning from the offline dataset, the agent can obtain a relatively good
starting point, as expected. And these starting scores (below $30$) are also
far from the expert score, which leaves a large room for further promotion in
the online phase. However, both algorithms have trouble in the online process.
REDQ suffers from severe instability, and in the end, the policy almost
degenerates into a random one. Although the CQL agent can keep stable during
learning, its improvement is very limited, which indicates it is ineffective
for leveraging online interaction.

These results indicate that a pure online RL algorithm may be problematic for
handling offline and online data in a single training process. Furthermore,
the pure offline RL algorithm cannot effectively use valuable online
interaction data due to its conservative updating strategy.

### Adaptive Policy Learning Framework

In this subsection, we present the Adaptive Policy Learning (APL) framework,
trying to tackle the above problem by adaptively leveraging online and offline
data. The underlying idea is simple and can be described as follows. When
online interactions are available, we try to obtain near-on-policy data from
them and choose an optimistic updating strategy since these data reflect the
true situation of the current policy. By contrast, we use a more pessimistic
updating strategy when data are sampled from the offline dataset. In this way,
we can make full use of both online and offline data and explicitly consider
their differences by separately applying suitable updating schemes. It is
worth stressing that being optimistic here means we rely on an online RL
update strategy which is optimistic compared with offline RL methods, while
being pessimistic means an offline RL update strategy.

Then, we give a general formalization for the above idea. Since this idea can
be applied to both value-based and policy-based methods, we use $C$ to
represent a policy or state-action value function, and define a unified
updating rule as follows:

$\displaystyle
C^{k+1}\leftarrow\mathcal{F}\left(\mathbb{A}(C^{k})+\mathcal{W}(\mathbf{s},\mathbf{a})\mathbb{B}(C^{k})\right),$
(2)

where $k$ is the times of iterative updates. For the right-hand side, the
first term $\mathbb{A}(C)$ stands for the optimistic updating strategy. It is
the common learning target used in online RL, such as the Bellman error for
value-based methods. The second term $\mathbb{B}(C)$ stands for the
pessimistic updating strategy. It is a penalty term to make the learned agent
take actions close to the dataset or take conservative actions. Besides, to
express the idea of adaptive learning, we apply a weight function
$\mathcal{W}(\mathbf{s},\mathbf{a})$ to the penalty term $\mathbb{B}(C)$.
Specifically, when we use near-on-policy data,
$\mathcal{W}(\mathbf{s},\mathbf{a})$ will be a small value, and the updating
relies more on $\mathbb{A}(C)$, leading to a relatively optimistic updating
strategy. On the contrary, when we use the offline data,
$\mathcal{W}(\mathbf{s},\mathbf{a})$ will be a large value, and the updating
strategy is relatively pessimistic. In addition, we use $\mathcal{F}$ to
denote the general updating operator, such as $\operatorname{argmin}$ or
$\operatorname{argmax}$. Using this general operator, the updating rule for
both value and policy-based methods can be unified in Eq. 2. Note that we only
introduce the general framework here, and detailed algorithms are presented in
the next sections.

Online-Offline Replay Buffer We now introduce a simple but effective online-
offline replay buffer (OORB) to distinguish between near-on-policy and offline
data. OORB consists of two buffers. One is the online buffer that collects the
near-on-policy data generated via online interactions. To ensure the data in
the online buffer is near-on-policy, we set this buffer to be small, and newly
collected online data are stored in it by following the first-in-first-out
rule. The other is the offline buffer containing previously collected offline
datasets which can come from arbitrary policies, and all data generated via
interactions in the online phase.

Data are sampled from OORB following a Bernoulli distribution. With a
probability $p$, data are sampled from the online buffer, and with probability
$1-p$, they are sampled from the offline buffer. The effect of $p$ on the
final performance is tested via ablation studies in the experiment section.

Algorithm 1 Greedy Conservative Q-ensemble Learning

// Online interaction steps in each iteration $T_{\text{on}}$, updating steps
in each iteration $T_{\text{off}}$, initial offline training steps
$T_{\text{initial}}$

// OORB threshold $p$, OORB starting sampling size $T_{\text{s}}$

Input: Total online interaction steps $S_{T}$

Initialize online buffer $B_{\text{on}}$ to empty, offline buffer
$B_{\text{off}}\leftarrow\text{offline dataset}$

Initialize policy $\pi$, Q-functions $Q_{i},i\in N$

Set the updating step counter $t$

Set the total online interaction step counter $S_{\text{on}}\leftarrow 0$

Train the agent for $T_{\text{initial}}$ steps using the offline dataset

repeat

$t\leftarrow 0$

Interact with the environment for $T_{\text{on}}$ steps

Store collected experiences to both online buffer $B_{\text{on}}$ and offline
buffer $B_{\text{off}}$

$S_{\text{on}}\leftarrow S_{\text{on}}+T_{\text{on}}$

for $t<T_{\text{off}}$ do

Sample a random value $p_{s}\sim\mathbb{U}(0,1)$

if $p_{s}<p$ and $S_{\text{on}}>T_{\text{s}}$ then

Sample a batch of $(\mathbf{s},\mathbf{a})$ from online buffer $B_{\text{on}}$

Set the $\mathcal{W}(\mathbf{s},\mathbf{a})$ to 0

else

Sample a batch of $(\mathbf{s},\mathbf{a})$ from offline buffer
$B_{\text{off}}$

Set the $\mathcal{W}(\mathbf{s},\mathbf{a})$ to 1

end if

Update the Q-functions $Q_{i},i\in N$ by Eq. 6

Update the policy $\pi$ by Eq. 7

$t\leftarrow t+1$

end for

until $S_{\text{on}}>S_{T}$

### Value-Based Implementation

We first present an action-value based implementation for APL framework by
incorporating CQL (Kumar et al. 2020) as the value-based offline RL method and
an ensemble online RL algorithm REDQ (Chen et al. 2021a). We name the
implementation Greedy-Conservative Q-ensemble Learning (GCQL).

Specifically, $C$ in Eq. 2 is the state-action value $Q$ here, and we use the
updating function in REDQ as the optimistic strategy. Since REDQ is an
ensemble method which uses a set of Q-functions, we use $i$ as the index of
the Q-functions, and the size of the set is $N$. Hence,

$\displaystyle\mathbb{A}(Q^{k}_{i})=\mathbb{E}_{\mathbf{s},\mathbf{a},\mathbf{s^{\prime}}\sim{\text{OORB}},\mathbf{a^{\prime}}\sim\pi^{k}(\cdot\mid
s^{\prime})}\left[\left(Q^{k}_{i}(\mathbf{s},\mathbf{a})-\mathcal{B}^{\pi}\hat{Q}^{k}(\mathbf{s^{\prime}},\mathbf{a^{\prime}})\right)^{2}\right],$
(3)

where $k$ is the times of iterative updates same as Eq. 2, and
“$\sim{\text{OORB}}$” represents sampling data from OORB. Operator
$\mathcal{B}^{\pi}\hat{Q}^{k}(\mathbf{s^{\prime}},\mathbf{a^{\prime}})$ is

$r+\gamma\min_{i\in\mathcal{M}}\hat{Q}^{k}_{i}\left(\mathbf{s^{\prime}},\mathbf{a^{\prime}}\right),\quad\mathbf{a^{\prime}}\sim\pi^{k}\left(\cdot\mid\mathbf{s^{\prime}}\right).$
(4)

Here, $\pi^{k}$ is the current policy, and $\hat{Q}$ denotes a target $Q$
function for stabilizing the learning process (Mnih et al. 2015). Following
REDQ’s design, we randomly select two Q-functions from a set of them for
ensemble, with $\mathcal{M}$ representing the set of selected indexes.

Then, we adopt the conservative regularizer in CQL as the second term in Eq.
2. For every Q-function in the set:

$\displaystyle\mathbb{B}(Q^{k}_{i})=\alpha\mathbb{E}_{\mathbf{s}\sim{\text{OORB}}}\left[\log\sum_{\mathbf{a^{\prime}}}\exp(Q^{k}_{i}(\mathbf{s},\mathbf{a^{\prime}}))-\mathbb{E}_{\mathbf{a}\sim{\text{OORB}}}[Q^{k}_{i}(\mathbf{s},\mathbf{a})]\right],$
(5)

where action $\mathbf{a^{\prime}}$ is sampled from the current policy, i.e.,
$\mathbf{a^{\prime}}\sim\pi^{k}(\cdot|s)$. Hence, the overall updating
function of $Q$ is:

$Q^{k+1}_{i}=\arg\min_{Q^{k}_{i}}\left\\{\mathbb{A}(Q^{k}_{i})+\mathcal{W}(\mathbf{s},\mathbf{a})\mathbb{B}(Q^{k}_{i})\right\\}$
(6)

Finally, the update function of policy is shown as follows:

$\displaystyle\pi^{k+1}=\arg\max_{\pi^{k}}\mathbb{E}_{\mathbf{a}\sim\pi^{k}\left(\cdot\mid\mathbf{s}\right)}\left[\mathbb{E}_{i\in
N}\left[Q_{i}^{k}\left(\mathbf{s},\mathbf{a}\right)\right]-\alpha\log\pi^{k}\left(\mathbf{a}\mid\mathbf{s}\right)\right].$
(7)

### Policy-Based Implementation

Besides value-based methods, APL framework can also be easily implemented with
policy-based methods. Here we take TD3BC (Fujimoto and Gu 2021) as an example,
and name our algorithm Greedy-Conservative TD3BC (GCTD3BC). Specifically,
TD3BC uses the policy update step in TD3, and adds a behavior cloning (BC)
term to regularize the policy. Therefore, GCTD3BC takes the policy update
function in TD3 as the optimistic term, and takes the BC regularizer as the
pessimistic term. So, the policy update rule in GCTD3BC is

$\displaystyle\pi^{k+1}=\underset{\pi^{k}}{\operatorname{argmax}}\mathbb{E}_{(\mathbf{s},\mathbf{a})\sim\text{OORB}}\left[\mathbb{A}(\pi^{k})+\mathcal{W}(\mathbf{s},\mathbf{a})\mathbb{B}(\pi^{k})\right],$
(8)

where $\mathbb{A}(\pi^{k})=\lambda Q^{k}(\mathbf{s},\pi(\mathbf{s}))$ is the
value function in TD3, and
$\mathbb{B}(\pi^{k})=-(\pi^{k}(\mathbf{s})-\mathbf{a})^{2}$ is the BC
regularizer. Hyper-parameter $\lambda$ is used to adjust two terms.

We now introduce how $\mathcal{W}(\mathbf{s},\mathbf{a})$ is set in both
implementations. We use a very simple method: when data is sampled from the
online buffer, $\mathcal{W}(\mathbf{s},\mathbf{a})$ is set to 0, otherwise 1.
This is why we use the term Greedy-Conservative to describe our algorithms.
Specifically, we greedily exploit the near-on-policy data with the online RL
scheme without any conservative regularizer. And we conservatively exploit the
offline data by employing the offline RL methods. Formally, this strategy can
be explained as below:

$\mathcal{W}(\mathbf{s},\mathbf{a})\leftarrow\begin{cases}0&~{}\text{if}~{}(\mathbf{s},\mathbf{a})\sim\text{online
buffer}\\\ 1&~{}\text{otherwise}\end{cases}$ (9)

To be more clear, we summarize GCQL in Algorithm 1, and explain the main steps
as follows. Firstly, we learn from the existing offline data for
$T_{\text{initial}}$ steps to leverage them. To make good use of the offline
data, we usually set $T_{\text{initial}}$ to a large value, e.g., $100K$.
Secondly, we begin the following interleaving learning process. We conduct the
online interaction for $T_{\text{on}}$ steps and store newly collected
experiences to OORB, and then we update the agent for $T_{\text{off}}$ steps.
For higher sample efficiency, we set the number of online interaction steps
$T_{\text{on}}$ to a small value, e.g., $1K$, and the number of updating steps
$T_{\text{off}}$ to a value larger than $T_{\text{on}}$, e.g., $10K$. If the
batch of data for updating the policy and Q-functions comes from the online
buffer, $\mathcal{W}(\mathbf{s},\mathbf{a})$ is set to 0, otherwise 1.

Algorithm GCTD3BC is similar to Algorithm 1 except that GCTD3BC uses the
action-value objective in TD3 (Fujimoto, van Hoof, and Meger 2018) and the
policy updating rule is defined in Eq. 8.

Environment | GCQL | GCTD3BC | CQL | REDQ_ON | REDQ | TD3_ON | TD3BC | AWAC | OFF2ON | IQL
---|---|---|---|---|---|---|---|---|---|---
walker2d-r | 31$\pm$27 | 5$\pm$3 | 7$\pm$9 | 71$\pm$11 | 5$\pm$3 | 7$\pm$2 | 6$\pm$3 | 12 | 20$\pm$13 | 7$\pm$3
hopper-r | 58$\pm$30 | 35$\pm$22 | 10$\pm$1 | 78$\pm$37 | 2$\pm$1 | 10$\pm$2 | 11$\pm$0 | 63 | 81$\pm$21 | 10$\pm$2
halfcheetah-r | 101$\pm$2 | 69$\pm$8 | 46$\pm$4 | 59$\pm$2 | 32$\pm$1 | 39$\pm$1 | 35$\pm$3 | 53 | 85 $\pm$3 | 28$\pm$7
walker2d-m | 94$\pm$6 | 90$\pm$7 | 83$\pm$1 | 71$\pm$11 | 2$\pm$3 | 7$\pm$2 | 79$\pm$2 | 80 | 89$\pm$2 | 51$\pm$13
hopper-m | 83$\pm$11 | 99$\pm$3 | 70$\pm$23 | 78$\pm$37 | 3$\pm$1 | 10$\pm$2 | 80$\pm$13 | 91 | 59 $\pm$9 | 42$\pm$9
halfcheetah-m | 66$\pm$3 | 62$\pm$2 | 25$\pm$8 | 59$\pm$2 | 46$\pm$1 | 39$\pm$1 | 43$\pm$1 | 41 | 58$\pm$2 | 40$\pm$0
walker2d-me | 93$\pm$12 | 102$\pm$2 | 105$\pm$1 | 71$\pm$11 | 12$\pm$3 | 7$\pm$2 | 110$\pm$3 | 78 | 101$\pm$24 | 58$\pm$22
hopper-me | 110$\pm$1 | 110$\pm$2 | 109$\pm$5 | 78$\pm$37 | 40$\pm$15 | 10$\pm$2 | 110$\pm$0 | 112 | 82 $\pm$21 | 72$\pm$16
halfcheetah-me | 102$\pm$1 | 103$\pm$2 | 92$\pm$2 | 59$\pm$2 | 9$\pm$3 | 39$\pm$1 | 98$\pm$2 | 41 | 100$\pm$1 | 38$\pm$17
walker2d-mr | 97$\pm$16 | 90$\pm$9 | 57$\pm$5 | 71$\pm$11 | 13$\pm$2 | 7$\pm$2 | 60$\pm$5 | - | 71$\pm$32 | 30$\pm$13
hopper-mr | 72$\pm$20 | 87$\pm$11 | 37$\pm$4 | 78$\pm$37 | 0$\pm$1 | 10$\pm$2 | 38$\pm$1 | - | 60$\pm$23 | 31$\pm$10
halfcheetah-mr | 62$\pm$2 | 53$\pm$1 | 50$\pm$0 | 59$\pm$2 | 28$\pm$25 | 39$\pm$1 | 47$\pm$0 | - | 57$\pm$1 | 42$\pm$2
Total | 969$\pm$131 | 905$\pm$71 | 691$\pm$63 | 624$\pm$200 | 192$\pm$59 | 224$\pm$20 | 717$\pm$33 | - | 863$\pm$152 | 449$\pm$114

Table 1: Averaged normalized score over last three evaluation iterations and 5
random seeds. $\pm$ captures the standard deviation over seeds. The highest
performing scores are in bold. The results of AWAC are taken from their paper
(Nair et al. 2020). REDQ_ON and TD3_ON are pure online methods without any
offline pre-training. We rerun them using implementations from authors to
ensure identical evaluation process. Suffix “-r” = random-v0 dataset, “-m” =
medium-v0 dataset, “-me” = medium-expert-v0, and “-mr” = medium-replay-v0. All
learning curves are showed in Appendix.

## Experiments

In this section, we design experiments to verify the effectiveness of our
framework from three perspectives: (1) the performance compared with
competitive baselines; (2) ablation studies to test the effect of key
components used in our methods; (3) the influence of different hyper-
parameters.

Tasks All experiments were done on the continuous control task set MuJoCo
(Todorov, Erez, and Tassa 2012), and the offline dataset comes from the
popular offline RL benchmark D4RL (Fu et al. 2020). Here, we test three tasks,
i.e., walker2d, hopper and halfcheetah, and each task takes four different
kinds of offline dataset, which are random-v0, medium-v0, medium-replay-v0 and
medium-expert-v0.

Settings We set $T_{\text{on}}$ in Algorithm 1 to $1K$. To better exploit the
offline dataset, we set $T_{\text{initial}}$ and $T_{\text{off}}$ to $100K$
and $10K$, respectively. For OORB, we set $p=0.5$ for GCQL and $p=0.1$ for
GCTD3BC, and $T_{\text{s}}$ to $10K$ for both of them. The size of online and
offline buffer is set to $20K$ and $3M$, respectively. We input $S_{T}$ as
$100K$. The above configurations keep the same across all tasks, datasets and
methods.

Figure 2: Ablation studies on Halfcheetah. GCQL_WG: GCQL without the greedy
updating strategy. GCQL_WGO: GCQL without the greedy updating strategy and
online buffer. Results on four tasks are averaged over three random seeds.

Compared Methods  We compare our algorithms, i.e., GCQL and GCTD3BC, with 8
competitive baselines, i.e., CQL, REDQ_ON, REDQ, TD3_ON, TD3BC, AWAC (Nair et
al. 2020), OFF2ON (Lee et al. 2021) and IQL (Kostrikov, Nair, and Levine
2022). REDQ_ON denotes that we rerun the officially released codes by the
authors (Chen et al. 2021a) without changing hyper-parameters, and do not use
any offline data to pre-train for this online method. We use the “_ON” to
indicates the agent is trained purely online without offline pre-training, and
methods without “_ON” are in the offline-to-online setting. Similarly, TD3_ON
is a pure online agent trained with TD3. For a fair comparison, we also
include a re-implemented version in GCQL, i.e. REDQ, as a baseline, where the
number of Q-functions is set to 5 for computational efficiency (as in GCQL)
and pre-training with offline data is leveraged. The results of AWAC are
directly taken from their paper. For OFF2ON, since the authors did not release
the code for pre-training, we contact the authors and use pre-trained agents
provided by them, and then use their codes for the rest of training. For other
methods, results are obtained by rerunning their codes under the offline-to-
online setting. We use default hyper-parameters for all methods without
further tuning.

### Overall Performance

We list scores in Table 1, and include all learning curves in Appendix. First,
it is clear that our methods GCQL and GCTD3BC perform much better than
baselines, well demonstrating the effectiveness of them. Second, we notice
that our methods are more robust comparing with baselines. Specifically,
offline RL methods (i.e., CQL and TD3BC) can obtain higher scores than pure
online methods (i.e., REDQ_ON and TD3_ON) when the dataset is not random, but
they have poor performance when the dataset is random. In contrast, offline-
to-online methods (e.g., GCQL and OFF2ON) can perform well regardless of
random dataset or not. Note that OFF2ON is fine-tuned. For instance, the
critic’s network architecture is different from that in the original CQL
paper. Our methods use default hyper-parameters in the original paper without
fine-tuning, and can perform better than those fine-tuned methods. Third, it
is surprising that single-agent method GCTD3BC have higher scores and faster
learning speed than ensemble method GCQL on some tasks, such as hopper-
medium-v0. All results shows APL can benefit more from the online and offline
data and gain high sample efficiency. Also, the general APL framework can be
successfully applied to the various RL methods.

Figure 3:  GCQL-i2e5 and GCQL-i5e4 mean that $T_{\text{initial}}$ is $2e5$ and
$5e4$, respectively. And GCQL-o2e4 and GCQL-o5e3 mean $T_{\text{off}}$ is set
to corresponding values. Results are averaged over three random seeds. Figure
4: Learning curves for agents with different OORB threshold $p$. GCQL-X means
that $p$ is set to X. Results are averaged over three random seeds.

### Ablation Study on Key Components in APL

To investigate the effect of key components in our method, we conduct the
following ablation studies on GCQL. The main differences between APL framework
and prior offline RL methods is the adaptive update schema and online buffer
in OORB. To this end, we design two variants of GCQL to investigate such two
components’ effect. GCQL_WG: We remove the adaptive update schema by fixing
$\mathcal{W}$ to 1, resulting in the variant of GCQL without greedy updating
strategy. GCQL_WGO: We further remove the online buffer, resulting in the
variant of GCQL without greedy updating strategy and online buffer.

It is easy to deduce that when the dataset includes the expert data, all
variants perform well. The underlying reason is obvious as offline RL methods
can perform quite well when data quality is good. However, for other datasets
with worse quality (e.g., medium-replay, medium and random), online buffer and
greedy update schema play an important role in boosting the performance and
stabilizing the learning process. Results show that for these worse datasets,
both GCQL_WG and GCQL_WGO have a clear performance drop compared with GCQL.
Specifically, it is observed that, for medium datasets, both CQL and GCQL_WGO
suffer a serious stability issue, which verifies that the online buffer indeed
stabilizes the learning process. Besides, GCQL_WG performs better than
GCQL_WGO, which indicates that the online buffer can not only stabilize the
learning process but also improve the performance. Moreover, GCQL still
outperforms GCQL_WG by a large margin on the sub-optimal dataset, which
verifies the efficiency of our greedy update scheme. In summary, greedy update
schema and online buffer are both crucial for improving the sample efficiency
in GCQL.

### Analysis on Hyper-parameters

As we do not use careful fine-tuning techniques (e.g., grid search) for hyper-
parameters in our methods, one may wonder how the hyper-parameters affect the
performance. To this end, we conduct experiments to investigate their
influence on GCQL. Here, we use the most complicated one among three tasks,
i.e., walker2d, in this experiment, and we analyze initial offline training
steps $T_{\text{initial}}$, updating steps in each iteration $T_{\text{off}}$
and OORB threshold $p$ in the section.

Learning curves for analyzing $T_{\text{initial}}$ and $T_{\text{off}}$ are
shown in Figure 3. We try a larger and smaller value than the default one to
test its impacts. Specifically, $T_{\text{initial}}$ is tested with $2e5$ and
$5e4$, and $T_{\text{off}}$ is tested with $2e4$ and $5e3$. We can see that
the performance of GCQL is insensitive to the $T_{\text{initial}}$ and
$T_{\text{off}}$, especially for datasets with not poor quality, e.g.,
datasets except for the random one. Besides, Figure 4 shows curves for
analyzing $p$, which have bigger impact on the performance than
$T_{\text{initial}}$ and $T_{\text{off}}$. $p$ is tested with $0.3$, $0.4$,
$0.5$, $0.6$ and $0.7$, and results show that methods with higher $p$ perform
better on the medium and medium-replay datasets, while those with lower $p$
perform better on the other datasets. Particularly, only $p=0.5$ can achieve a
clear performance improvement in both the medium and random datasets. Overall,
the default $p=0.5$ is the most appropriate setting, suggesting that taking
the online and offline data equally important may be the best option for GCQL
in most cases.

## Conclusion and Future Work

In this paper, we propose an Adaptive Policy Learning (APL) framework for
offline-to-online reinforcement learning. In APL, the advantages of online and
offline data are considered in an adaptive way, so that they are well-utilized
for policy learning. Furthermore, a value-based and a policy-based algorithm
are provided for implementing APL framework. Finally, we conduct comprehensive
experiments and results demonstrate that our methods can obtain best sample
efficiency in the offline-to-online setting comparing with several competitive
baselines. In the future, we will continue to further improve the robustness
of APL. For example, we will try to minimize the impact of offline dataset’s
quality on the performance. At last, we hope this work could bridge the gap
between offline and online RL.

## References

  * Agarwal, Schuurmans, and Norouzi (2020) Agarwal, R.; Schuurmans, D.; and Norouzi, M. 2020.  An optimistic perspective on offline reinforcement learning.  In _ICML_.
  * Brown et al. (2020) Brown, T. B.; Mann, B.; Ryder, N.; Subbiah, M.; Kaplan, J.; Dhariwal, P.; Neelakantan, A.; Shyam, P.; Sastry, G.; Askell, A.; et al. 2020.  Language models are few-shot learners.  _arXiv preprint arXiv:2005.14165_.
  * Chen et al. (2021a) Chen, X.; Wang, C.; Zhou, Z.; and Ross, K. W. 2021a.  Randomized Ensembled Double Q-Learning: Learning Fast Without a Model.  _ICLR_ , abs/2101.05982.
  * Chen et al. (2021b) Chen, X.-H.; Yu, Y.; Li, Q.; Luo, F.-M.; Qin, Z.; Shang, W.; and Ye, J. 2021b.  Offline Model-based Adaptable Policy Learning.  _Advances in Neural Information Processing Systems_ , 34: 8432–8443.
  * Dann et al. (2014) Dann, C.; Neumann, G.; Peters, J.; et al. 2014.  Policy evaluation with temporal differences: A survey and comparison.  _Journal of Machine Learning Research_.
  * Deng et al. (2009) Deng, J.; Dong, W.; Socher, R.; Li, L.-J.; Li, K.; and Fei-Fei, L. 2009.  ImageNet: A large-scale hierarchical image database.  In _2009 IEEE Conference on Computer Vision and Pattern Recognition_ , 248–255.
  * Fu et al. (2020) Fu, J.; Kumar, A.; Nachum, O.; Tucker, G.; and Levine, S. 2020.  D4rl: Datasets for deep data-driven reinforcement learning.  _arXiv preprint arXiv:2004.07219_.
  * Fujimoto and Gu (2021) Fujimoto, S.; and Gu, S. S. 2021.  A Minimalist Approach to Offline Reinforcement Learning.  arXiv:2106.06860.
  * Fujimoto, Meger, and Precup (2019) Fujimoto, S.; Meger, D.; and Precup, D. 2019.  Off-Policy Deep Reinforcement Learning without Exploration.  In _ICML_.
  * Fujimoto, van Hoof, and Meger (2018) Fujimoto, S.; van Hoof, H.; and Meger, D. 2018.  Addressing Function Approximation Error in Actor-Critic Methods.  In _International Conference on Machine Learning_.
  * Ghosh et al. (2022) Ghosh, D.; Ajay, A.; Agrawal, P.; and Levine, S. 2022.  Offline rl policies should be trained to be adaptive.  In _International Conference on Machine Learning_ , 7513–7530. PMLR.
  * Haarnoja et al. (2018) Haarnoja, T.; Zhou, A.; Abbeel, P.; and Levine, S. 2018.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  _arXiv preprint arXiv:1801.01290_.
  * He and Hou (2020) He, Q.; and Hou, X. 2020.  POPO: Pessimistic Offline Policy Optimization.  _ArXiv_ , abs/2012.13682.
  * Jaques et al. (2019) Jaques, N.; Ghandeharioun, A.; Shen, J. H.; Ferguson, C.; Lapedriza, A.; Jones, N.; Gu, S.; and Picard, R. 2019.  Way off-policy batch deep reinforcement learning of implicit human preferences in dialog.  _arXiv preprint arXiv:1907.00456_.
  * Kidambi et al. (2020) Kidambi, R.; Rajeswaran, A.; Netrapalli, P.; and Joachims, T. 2020.  MOReL : Model-Based Offline Reinforcement Learning.  _NeurIPS_ , abs/2005.05951.
  * Kim et al. (2013) Kim, B.; massoud Farahmand, A.; Pineau, J.; and Precup, D. 2013.  Learning from Limited Demonstrations.  In _NeurIPS_.
  * Kostrikov, Nair, and Levine (2022) Kostrikov, I.; Nair, A.; and Levine, S. 2022.  Offline reinforcement learning with implicit q-learning.  _International Conference on Learning Representations_.
  * Kotsiantis et al. (2007) Kotsiantis, S. B.; Zaharakis, I.; Pintelas, P.; et al. 2007.  Supervised machine learning: A review of classification techniques.  _Emerging artificial intelligence applications in computer engineering_ , 160(1): 3–24.
  * Kumar et al. (2019) Kumar, A.; Fu, J.; Soh, M.; Tucker, G.; and Levine, S. 2019.  Stabilizing off-policy q-learning via bootstrapping error reduction.  In _Advances in Neural Information Processing Systems_ , 11761–11771.
  * Kumar et al. (2020) Kumar, A.; Zhou, A.; Tucker, G.; and Levine, S. 2020.  Conservative Q-Learning for Offline Reinforcement Learning.  _NeurIPS_.
  * Lee et al. (2021) Lee, S.; Seo, Y.; Lee, K.; Abbeel, P.; and Shin, J. 2021.  Offline-to-Online Reinforcement Learning via Balanced Replay and Pessimistic Q-Ensemble.  _arXiv preprint arXiv:2107.00591_.
  * Levine et al. (2020) Levine, S.; Kumar, A.; Tucker, G.; and Fu, J. 2020.  Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems.  _ArXiv_ , abs/2005.01643.
  * Mandel et al. (2014) Mandel, T.; Liu, Y.-E.; Levine, S.; Brunskill, E.; and Popovic, Z. 2014.  Offline policy evaluation across representations with applications to educational games.  In _AAMAS_ , volume 1077.
  * Matsushima et al. (2021) Matsushima, T.; Furuta, H.; Matsuo, Y.; Nachum, O.; and Gu, S. 2021.  Deployment-efficient reinforcement learning via model-based offline optimization.  _ICLR_.
  * Mnih et al. (2015) Mnih, V.; Kavukcuoglu, K.; Silver, D.; Rusu, A. A.; Veness, J.; Bellemare, M. G.; Graves, A.; Riedmiller, M.; Fidjeland, A. K.; Ostrovski, G.; et al. 2015\.  Human-level control through deep reinforcement learning.  _nature_ , 518(7540): 529–533.
  * Nair et al. (2020) Nair, A.; Dalal, M.; Gupta, A.; and Levine, S. 2020.  Accelerating Online Reinforcement Learning with Offline Datasets.  _CoRR_ , abs/2006.09359.
  * Rajeswaran et al. (2018) Rajeswaran, A.; Kumar, V.; Gupta, A.; Vezzani, G.; Schulman, J.; Todorov, E.; and Levine, S. 2018.  Learning complex dexterous manipulation with deep reinforcement learning and demonstrations.  _Robotics: Science and Systems_.
  * Schulman et al. (2015) Schulman, J.; Levine, S.; Abbeel, P.; Jordan, M.; and Moritz, P. 2015.  Trust region policy optimization.  In _International conference on machine learning_ , 1889–1897. PMLR.
  * Schulman et al. (2017) Schulman, J.; Wolski, F.; Dhariwal, P.; Radford, A.; and Klimov, O. 2017.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_.
  * Sutton and Barto (2011) Sutton, R. S.; and Barto, A. G. 2011.  Reinforcement learning: An introduction.
  * Todorov, Erez, and Tassa (2012) Todorov, E.; Erez, T.; and Tassa, Y. 2012.  Mujoco: A physics engine for model-based control.  In _2012 IEEE/RSJ International Conference on Intelligent Robots and Systems_ , 5026–5033. IEEE.
  * Vecerik et al. (2017) Vecerik, M.; Hester, T.; Scholz, J.; Wang, F.; Pietquin, O.; Piot, B.; Heess, N.; Rothörl, T.; Lampe, T.; and Riedmiller, M. 2017.  Leveraging demonstrations for deep reinforcement learning on robotics problems with sparse rewards.  _arXiv preprint arXiv:1707.08817_.
  * Wu, Tucker, and Nachum (2019) Wu, Y.; Tucker, G.; and Nachum, O. 2019.  Behavior Regularized Offline Reinforcement Learning.  _ArXiv_ , abs/1911.11361.
  * Yu et al. (2020) Yu, T.; Thomas, G.; Yu, L.; Ermon, S.; Zou, J.; Levine, S.; Finn, C.; and Ma, T. 2020.  MOPO: Model-based Offline Policy Optimization.  _NeurIPS_ , abs/2005.13239.
  * Zhu et al. (2019) Zhu, H.; Gupta, A.; Rajeswaran, A.; Levine, S.; and Kumar, V. 2019.  Dexterous manipulation with deep reinforcement learning: Efficient, general, and low-cost.  In _2019 International Conference on Robotics and Automation (ICRA)_ , 3651–3657. IEEE.


11institutetext: University of Oxford, UK
11email<EMAIL_ADDRESS>
22institutetext: Mianyang Visual Object Detection and Recognition Engineering
Center, China

# VMambaMorph: a Multi-Modality Deformable Image Registration Framework based
on Visual State Space Model with Cross-Scan Module

Ziyang Wang 11 Jian-Qing Zheng 11 Chao Ma 22 Tao Guo

###### Abstract

Image registration, a critical process in medical imaging, involves aligning
different sets of medical imaging data into a single unified coordinate
system. Deep learning networks, such as the Convolutional Neural Network
(CNN)-based VoxelMorph, Vision Transformer (ViT)-based TransMorph, and State
Space Model (SSM)-based MambaMorph, have demonstrated effective performance in
this domain. The recent Visual State Space Model (VMamba), which incorporates
a cross-scan module with SSM, has exhibited promising improvements in modeling
global-range dependencies with efficient computational cost in computer vision
tasks. This paper hereby introduces an exploration of VMamba with image
registration, named VMambaMorph. This novel hybrid VMamba-CNN network is
designed specifically for 3D image registration. Utilizing a U-shaped network
architecture, VMambaMorph computes the deformation field based on target and
source volumes. The VMamba-based block with 2D cross-scan module is redesigned
for 3D volumetric feature processing. To overcome the complex motion and
structure on multi-modality images, we further propose a fine-tune recursive
registration framework. We validate VMambaMorph using a public benchmark brain
MR-CT registration dataset, comparing its performance against current state-
of-the-art methods. The results indicate that VMambaMorph achieves competitive
registration quality. The code for VMambaMorph with all baseline methods is
available on GitHub.

https://github.com/ziyangwang007/VMambaMorph

###### Keywords:

Image Registration Visual Mamba (VMamba) Space State Model

##  1 Introduction

Predicting for non-rigid anatomical shifts, deformable image registration
stands as a pivotal component of medical image analysis [22]. This vital
process ensures the spatial alignment of images prior to their examination.
Registration is instrumental in medical imaging, enabling the juxtaposition of
images obtained over time (for longitudinal studies), from varied scanner
technologies (multi-modal registration), or among different individuals for
broader research studies (inter-subject registration).

The essence of image registration lies in deducing the spatial transformation,
${\phi}:\mathbb{R}^{n}\to\mathbb{R}^{n}$, which is articulated through
parameters that define either a linear spatial transformation (like rigid or
affine transformations) or a continuum of movements (or displacements)
indicative of deformation, represented by
$\phi[\textbf{{x}}]\in\mathbb{R}^{d}$ at a specific coordinate
$\textbf{{x}}\in\mathbb{Z}^{d}$ within a target image $\textbf{{I}}^{\rm
tgt}\in\mathbb{R}^{{n}}$, originating from a source image $\textbf{{I}}^{\rm
src}\in\mathbb{R}^{n}$. Here, $n$ denotes the dimensions of a 3D image as
$n={H\times W\times T}$, with $d,T,H,W$ corresponding to the dimensionality,
thickness, height, and width of the image, respectively. Traditionally
approached as an optimization challenge, image registration endeavors to
minimize a dissimilarity metric $\mathcal{D}$ alongside a regularization term
$\mathcal{R}$:

$\hat{\phi}=\underset{\phi}{\mathrm{argmin}}{\big{(}\mathcal{L}_{\rm
dissim}(\phi(\textbf{{I}}^{\rm src}),~{}\textbf{{I}}^{\rm
tgt})+\lambda\mathcal{L}_{\rm smooth}(\phi)\big{)}}$ (1)

In this formulation, $\phi$ signifies the hypothesized spatial transformation,
while $\lambda$ is the regularization term’s weight.

Figure 1: The Recursive Registration Framework of VMambaMorph.

While traditional registration techniques are capable of calculating
diffeomorphic displacement fields with considerable accuracy, their
substantial computational requirements and extended processing durations
hinder their practicality for real-time applications. Over the last decade,
registration methods powered by deep learning [1] have risen as efficient
alternatives. These techniques deliver rapid registration capabilities and, in
some instances, achieve accuracy comparable to that of conventional methods:

${\phi}=\mathcal{R}(\textbf{{I}}^{\rm src},\textbf{{I}}^{\rm
tgt};{\textbf{{w}}})\ $ (2)

Here, $\mathcal{R}$ represents the deep learning model. This model is refined
through a training process aimed at minimizing the loss function, exemplified
by Equation (1).

In a typical learning-based framework, the source and target images/volumes
are concatenated and fed into a single-branch network volume, which is then
processed through a UNet-like architecture to compute the deformation field.

Voxelmorph [1] is an early deep learning-based deformable image registration,
using classic convolutional network structure, UNet [19]. However, the
convolutional network structure suffers from a limited receptive field, thus
causing a limited capture range of deformation between two images.

The attention mechanism, as proposed by Vaswani et al. [24], is designed to
capture long-range dependencies in contextual sequential features. It has
since been adapted for computer vision in the form of the Vision Transformer
[6], achieving a global receptive field capable of addressing the issue of
limited capture ranges. Consequently, it has been employed in multiple studies
for both rigid image registration [34, 35] and deformable image registration
[31, 21, 37]. However, the attention mechanism is also known to incur high
computational complexity. Several previous efforts [15, 36] have investigated
progressively coarse-to-fine image registration strategies to reduce
computational demands while retaining the benefits of extensive capture range
for deformation. Nevertheless, these methods also present challenges in
preserving discontinuity in the predicted deformation [36].

In this landscape, the emerging Mamba framework, named SSM, which is designed
for efficient feature processing in sequential features, offers a compelling
alternative [7, 8, 9]. The initial iterations of Mamba have been succeeded by
Visual Mamba [16], which incorporates a selective scan module within the Mamba
architecture, enhancing its performance in computer vision tasks. This
advancement has led to a slew of Mamba-based networks being proposed for a
variety of applications [18, 30, 29, 25, 38, 26]. MambaMorph [11] is one such
innovation, utilizing a hybrid Mamba-CNN network as the registration module to
compute deformation fields. Despite these developments, there remains an
untapped potential in the latest Mamba techniques, especially considering the
differences between sequential and visual features.

In this paper, we explore the application of Visual Mamba [16] to deformable
image registration, which has been further developed to efficiently extract
visual features from images and volumes. The contributions of this research
are outlined as follows:

  1. 1.

To the best of our knowledge, this represents the first investigation into the
use of the Visual State Space Model (VMamba) [16] for medical image
registration tasks.

  2. 2.

Inspired by the recent success of VMamba [16], the 2D image-based Visual State
Space (VSS) block is redesigned for 3D volumetric feature processing.

  3. 3.

VSS blocks are integrated with a conventional CNN-based U-shaped network to
serve as the registration module.

  4. 4.

To overcome the challenge of complex motion and structure in multi-modality
image registration, a recursive registration framework is adopted. In
addition, a fine-grained feature extractor is utilized before registration.

  5. 5.

The proposed registration method, VMambaMorph, is validated on a public
benchmark MR-CT dataset, achieving performance that surpasses existing state-
of-the-art networks for deformable image registration.

Figure 2: The Architecture of VMambaMorph, and the Details of Visual State
Space Block.

##  2 Method

The proposed VMamaMorph-based recursive registration framework is sketched in
Figure 1 and the network structure of VMambaMorph is shown in 2, consisting of
a two-branch feature extractor module, registration module, and STN. A pair of
target volume and source volume are respectively denoted as $\textbf{{I}}^{\rm
tgt}$ and $\textbf{{I}}^{\rm src}$, and their corresponding segmentation is
denoted as $\textbf{{S}}^{\rm tgt}$ and $\textbf{{S}}^{\rm src}$. The
registration module is $\mathcal{R}(\cdot,\cdot)$ and the feature extractor is
as $\mathcal{F}(\cdot)$. The deformation field is denoted as $\phi$. The
losses include $\mathcal{L}_{\rm dice}(\cdot,\cdot)$ and $\mathcal{L}_{\rm
smooth}(\cdot)$ following [1, 11].

###  2.1 Visual State Space Model

The CNN-based networks have been established as robust mechanisms for local
feature extraction [17, 19, 1, 32], and ViT-based networks are strength in
global-range dependency modeling [20, 3, 27, 5, 31]. Although hybrids of CNN
and ViT have been explored [4, 5, 28], they still incur significant
computational costs due to ViT’s quadratic complexity. In this landscape, the
emerging Mamba framework, named SSM, offers a compelling alternative. The
initial iterations of Mamba have been succeeded by Visual Mamba [16], which
incorporates a selective scan module within the Mamba architecture, enhancing
its performance in computer vision tasks. The conventional SSM is developed as
a linear time-invariant system function to map $x(t)\in\mathbb{R}\mapsto
y(t)\in\mathbb{R}$ through a hidden state $h(t)\in\mathbb{R}^{N}$, given
$A\in\mathbb{C}^{N\times N}$ as the evolution parameter,
$B,C\in\mathbb{C}^{N}$ as the projection parameters for a state size $N$, and
skip connection $D\in\mathbb{C}^{1}$. The process is illustrated in Equation 3
known as linear ordinary differential equations.

$h^{\prime}(t)=Ah(t)+Bx(t),y(t)=Ch(t)+Dx(t)$ (3)

Considering the input $x_{k}\in\mathbb{R}^{L\times D}$ a sampled vector within
the signal flow of length $L$, the discrete version of this linear model can
be transformed by zeroth-order hold rule given a timescale parameter
$\Delta\in\mathbb{R}^{D}$ [10, 12].

$h_{t}=\overline{A}h_{k-2}+\overline{B}x_{k}$ (4)

$y_{t}={C}h_{k}+\overline{D}x_{k}$ (5)

$\overline{A}=e^{\Delta A}$ (6)

$\overline{B}=(e^{\Delta A}-I)A^{-1}B$ (7)

$\overline{C}=C$ (8)

where $B,C\in\mathbb{R}^{D\times N}$. The approximation of $\overline{B}$
refined using first-order Taylor series $\overline{B}=\left(e^{\Delta
A}-I\right)A^{-1}B\approx\left(\Delta A\right)\left(\Delta A\right)^{-1}\Delta
B=\Delta B$. The Visual Mamba further introduce Cross-Scan Module then
integrate convolutional operations into the VSS block, which is detailed in
[8, 16]. We have redesigned the original VSS network blocks for the encoder of
registration module for 3D image processing considering the volumetric medical
image registration.

###  2.2 Recursive Registration Framework

Addressing the intricate challenge presented by complex motion and structural
variations in multi-modality image registration, we adopt a recursive
registration framework structure aimed at achieving coarse-to-fine
registration, as depicted in Fig. 1. This initiative is inspired by [33, 37].
In our proposed method, the residual transformation $\varphi_{k}$,
representing the alignment discrepancy between the target image
$\textbf{{I}}^{\rm tgt}$ and the successively warped source feature map
$\phi_{k-1}(\textbf{{I}}^{\rm src})$ from the previous $k-1$ registration
level, is computed by the operator $\mathcal{R}$ and progressively refined
through compositional accumulation:

$\left\\{\begin{array}[]{cc}\phi_{k}=\phi_{k-1}\circ\varphi_{k}\\\
\varphi_{k}=\mathcal{R}(\phi_{k-1}(\textbf{{I}}^{\rm src}),\textbf{{I}}^{\rm
tgt})\end{array}\right.,~{}~{}0<k\leq K$ (9)

where $\circ$ signifies the composition of two deformation fields, with
$\phi_{0}$ being initialized without deformation field. The proposed
VMambaMorph network, as illustrated in Fig. 2, integrates a weight-sharing,
enhanced fine-grained feature extractor and a U-shaped VMamba-CNN hybrid
registration module to accurately estimate the deformation field $\psi_{k}$,
and then to refine the final prediction $\phi_{k}$. Following the optimal
setting validated by [36], the recursive number is set to ${K}_{\rm t}=2$ in
the training process, and varying recursive number ${K}_{\rm i}>1$ is compared
in the inference stage to optimize the registration performance.

###  2.3 Enhanced Fine-grained Feature Extractor

The existing methods, such as [1], generally combine the two volumes into a
unified entry to estimate the voxel-wise spatial correspondences. While this
approach may be adequate for registering volumes of the same modality, which
naturally exhibit similar visual characteristics, it often proves inadequate
for multi-modality registrations characterized by significant differences in
appearance. In such cases, it becomes crucial to extract matching features
from corresponding anatomical regions across the divergent volumes.

Following MambaMorph [11], a refined yet uncomplicated UNet architecture that
features a single down-sampling step for efficient feature extraction. This
modification better preserves the local characteristics of images. This, in
turn, enhances the detail captured for the task at hand. VMambaMorph leverages
a simultaneous learning approach to perform both feature extraction and
registration, employing weight-sharing among feature extractors within our
proposed model.

###  2.4 Training Setting

The training aim of image registration is to minimize the Loss $\mathcal{L}$
based on two labels $\textbf{{S}}^{\rm tgt}$ and $\textbf{{S}}^{\rm src}$,
which is formulated in Equation 10.

$\min\limits_{\textbf{{w}}}\mathcal{L}_{\rm dice}(\phi(\textbf{{S}}^{\rm
tgt}),\textbf{{S}}^{\rm src})+\lambda_{\rm s}\mathcal{L}_{smooth}(\phi)$ (10)

where the dense deformation vector field $\phi$ is calculated by:

$\phi=\mathcal{R}(\textbf{{F}}^{\rm tgt},\textbf{{F}}^{\rm src};\textbf{{w}})$
(11)

with the feature maps are calculated by:

$\textbf{{F}}^{\\{\rm tgt,src\\}}=\mathcal{F}(\textbf{{I}}^{\\{\rm
tgt,src\\}};\textbf{{w}})$ (12)

and $\phi(\cdot)$ denotes the warping operation and is implemented by STN
[14]. $\mathcal{L}_{dice}(\cdot,\cdot)$ is the supervised loss based on dice
similarity coefficient, and $\mathcal{L}_{smooth}(\cdot)$ denotes the same
smooth loss as VoxelMorph [1, 11].

##  3 Experiments

Implementation Details: Our code has been developed under Ubuntu 20.04 in
Python 3.10.13 using Pytorch 1.13.1 and CUDA 12.1 using a single Nvidia
GeForce RTX 3090 GPU, and Intel(R) Intel Core i9-10900K. The runtimes averaged
around 15 hours. VMambaMorph and all baseline methods are trained for 300
epoches, the batch size is set to 1, the optimizer is Adam, and the learning
rate is initially set to 0.001. The setting is also applied to all other
baseline methods directly without any further modification.

Dataset: A public benchmark brain MR-CT registration dataset named SR-Reg is
utilized in this paper [11]. The SR-Reg dataset consists of 180 well-aligned,
skull-stripped and intensity-rectified MR-CT pairs, with corresponding
segmentation labels, and it is mainly based on the SynthRAD 2023 dataset [23,
2, 13]. The original size of each volume is with $192\times 208\times 176$
voxels with a resolution of $1\times 1\times 1mm^{3}$. More details can be
found at [11]. Considering the input size requirement of VMamba, and
computational cost on a single GPU, we have resized the volume to $128\times
128\times 128$, and randomly select 150 cases for training, 10 cases for
validation, and the rest 20 cases for testing.

Table 1: Quantitative comparison of VMambaMorph methods against other baseline methods on the SR-Reg testing set. Model | FeatEx | Dice | HD95 | P|J∅|≤0 | Mem | Par | Time
---|---|---|---|---|---|---|---
|  | (%) | vox | (%) | (Gb) | (Mb) | (sec)
Initial | - | 62.42$\pm$3.29 | 3.73$\pm$0.41 | - | - | - | -
VoxelMorph [1] | ✗ | 71.27$\pm$1.95 | 2.11$\pm$0.25 | 0.13$\pm$0.02 | 1.22 | 0.09 | 0.01
VoxelMorph [1] | ✓ | 76.98$\pm$1.27 | 1.69$\pm$0.14 | 0.08$\pm$0.02 | 1.98 | 0.16 | 0.04
TransMorph [5] | ✗ | 77.87$\pm$1.99 | 1.71$\pm$0.19 | 0.66$\pm$0.05 | 2.81 | 14.29 | 0.06
TransMorph [5] | ✓ | 82.31$\pm$1.73 | 1.39$\pm$0.16 | 0.58$\pm$0.04 | 3.27 | 14.57 | 0.08
MambaMorph [11] | ✗ | 79.17$\pm$1.78 | 1.60$\pm$0.18 | 0.39$\pm$0.02 | 2.47 | 7.31 | 0.06
MambaMorph [11] | ✓ | 81.92$\pm$2.10 | 1.45$\pm$0.24 | 0.34$\pm$0.00 | 3.24 | 7.59 | 0.09
VMambaMorph | ✓ | 82.94$\pm$2.01 | 1.35$\pm$0.18 | 1.04$\pm$0.05 | 3.25 | 9.64 | 0.10

Metrics: We use the the percentage of mean dice-coefficients (Dice), 95%
Hausdorff distance (HD95) to measure the registration accuracy, percentage of
non-positive Jacobian determinant (P|J∅|≤0) to evaluate the diffeomorphic
property of deformation field. The inference time (Time) with the unit of
second (s), GPU memory cost (Mem) with the unit of Gigabyte (Gb) and the
amount of parameters of network with the unit Mb is used to evaluate the
proposed VMamabaMorph with other baseline methods.

Comparison with State-of-the-Art: To ensure a fair and thorough evaluation of
VMambaMorph, we have selected a range of baseline methods for comparison: the
purely CNN-based VoxelMorph [1], the Transformer-CNN hybrid TransMorph [5],
and the recent Mamba-CNN hybrid MambaMorph [11].

Table 2: Ablation study of VMambaMorph methods for feature extractor (FeatEx) and recursive training ($K_{\rm t}$) or inference ($K_{\rm i}$). Model | $K_{\rm t}$ | $K_{\rm i}$ | FeatEx | Dice | HD95 | P|J∅|≤0 | Mem | Par | Time
---|---|---|---|---|---|---|---|---|---
|  |  |  | (%) | vox | (%) | (Gb) | (Mb) | (sec)
Initial | - | - | - | 62.42$\pm$3.29 | 3.73$\pm$0.41 | - | - | - | -
VMambaMorph | 1 | 1 | ✗ | 78.46$\pm$2.13 | 1.69$\pm$0.22 | 0.33$\pm$0.01 | 2.94 | 9.35 | 0.07
VMambaMorph | 1 | 1 | ✓ | 82.49$\pm$1.99 | 1.38$\pm$0.18 | 0.35$\pm$0.01 | 3.25 | 9.64 | 0.10
VMambaMorph | 1 | 2 | ✓ | 82.84$\pm$2.05 | 1.44$\pm$0.20 | 1.05$\pm$0.01 | 3.93 | 9.64 | 0.20
VMambaMorph | 1 | 3 | ✓ | 82.80$\pm$2.06 | 1.47$\pm$0.18 | 1.71$\pm$0.07 | 3.93 | 9.64 | 0.29
VMambaMorph | 1 | 4 | ✓ | 82.62$\pm$2.11 | 1.50$\pm$0.18 | 2.29$\pm$0.09 | 3.93 | 9.64 | 0.36
VMambaMorph | 2 | 1 | ✓ | 82.47$\pm$2.01 | 1.38$\pm$0.18 | 0.35$\pm$0.01 | 3.93 | 9.64 | 0.10
VMambaMorph | 2 | 2 | ✓ | 82.94$\pm$2.01 | 1.35$\pm$0.18 | 1.04$\pm$0.05 | 3.93 | 9.64 | 0.19
VMambaMorph | 2 | 3 | ✓ | 82.85$\pm$2.07 | 1.47$\pm$0.18 | 1.69$\pm$0.06 | 3.93 | 9.64 | 0.28
VMambaMorph | 2 | 4 | ✓ | 82.65$\pm$2.16 | 1.51$\pm$0.17 | 2.26$\pm$0.08 | 3.93 | 9.64 | 0.40
Figure 3: The Training History of VoxelMorph, TransMorph, MambaMorph, and
VMambaMorph with and without the Feature Extractors.

Training History: Figure 3 depicts the training progression for all image
registration methods, with the X-axis representing the training epoch and the
Y-axis denoting the loss observed on the training and validation sets. It is
important to highlight that the incorporation of the fine-grained feature
extractor substantially reduces and stabilizes the training loss. Notably,
VMambaMorph demonstrates the most consistent performance, maintaining the
lowest level of loss on the validation set throughout the training process.
This evidence underscores the efficacy of VMambaMorph in learning from the
data.

Results: The quantitative results from a direct comparison of VMambaMorph with
other baseline methods across all evaluation metrics are reported in Table 1.
Our proposed method and the top-performing results are highlighted in bold. It
is noteworthy that VMambaMorph achieves 16%, 6%, and 1% higher Dice scores
compared with VoxelMorph, TransMorph, and MambaMorph, respectively.

Ablation Study: To validate the effectiveness of the fine-grained feature
extractor, we conducted an ablation study. The performance impact of the
feature extractor (denoted as FeatEx) is evaluated in conjunction with all
methods in Table 2. The deployment of FeatEx is indicated with a ✓, while its
absence is marked with an ✗. All methods demonstrate improvement when the
feature extractor is utilized. To validate the proposed recursive registration
framework, the bottom of Table 2 further explores the affect of varying
recursive numbers in training and inference.

##  4 Conclusion

In this paper, we introduce VMambaMorph, a multi-modality deformable image
registration framework based on the Visual State Space Model. The architecture
of VMambaMorph is developed with a hybrid registration module with 3D VMamba-
based network blocks and 3D CNN-based network blocks. Following the recent
advances in feature learning for multi-modality registration, VMambaMorph
incorporates the recursive registration framework with a pair of fine-grained
feature extractor to achieve high-level spatial features more effectively. Our
validation on a public MR-CT registration dataset indicates that VMambaMorph
surpasses recent methodologies such as MambaMorph, as well as established
techniques like VoxelMorph and TransMorph. This paves the way for future
applications of VMamba in broader computer vision tasks.

## References

  * [1] Balakrishnan, G., Zhao, A., Sabuncu, M.R., Guttag, J., Dalca, A.V.: Voxelmorph: a learning framework for deformable medical image registration. IEEE transactions on medical imaging 38(8), 1788–1800 (2019)
  * [2] Billot, B., Greve, D.N., Puonti, O., Thielscher, A., Van Leemput, K., Fischl, B., Dalca, A.V., Iglesias, J.E., et al.: Synthseg: Segmentation of brain mri scans of any contrast and resolution without retraining. Medical image analysis 86, 102789 (2023)
  * [3] Cao, H., Wang, Y., Chen, J., Jiang, D., Zhang, X., Tian, Q., Wang, M.: Swin-unet: Unet-like pure transformer for medical image segmentation. In: European conference on computer vision. pp. 205–218. Springer (2022)
  * [4] Chen, J., etc: Transunet: Transformers make strong encoders for medical image segmentation. arXiv preprint arXiv:2102.04306 (2021)
  * [5] Chen, J., Frey, E.C., He, Y., Segars, W.P., Li, Y., Du, Y.: Transmorph: Transformer for unsupervised medical image registration. Medical image analysis 82, 102615 (2022)
  * [6] Dosovitskiy, A., etc: An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 (2020)
  * [7] Gu, A.: Modeling Sequences with Structured State Spaces. Stanford University (2023)
  * [8] Gu, A., Dao, T.: Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 (2023)
  * [9] Gu, A., Goel, K., Ré, C.: Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396 (2021)
  * [10] Gu, A., Johnson, I., Goel, K., Saab, K., Dao, T., Rudra, A., Ré, C.: Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems 34, 572–585 (2021)
  * [11] Guo, T., Wang, Y., Meng, C.: Mambamorph: a mamba-based backbone with contrastive feature learning for deformable mr-ct registration. arXiv preprint arXiv:2401.13934 (2024)
  * [12] Gupta, A., Gu, A., Berant, J.: Diagonal state spaces are as effective as structured state spaces. Advances in Neural Information Processing Systems 35, 22982–22994 (2022)
  * [13] Hoopes, A., Mora, J.S., Dalca, A.V., Fischl, B., Hoffmann, M.: Synthstrip: skull-stripping for any brain image. NeuroImage 260, 119474 (2022)
  * [14] Jaderberg, M., Simonyan, K., Zisserman, A., et al.: Spatial transformer networks. Advances in neural information processing systems 28 (2015)
  * [15] Kang, M., Hu, X., Huang, W., Scott, M.R., Reyes, M.: Dual-stream pyramid registration network. Medical Image Analysis 78, 102379 (2022)
  * [16] Liu, Y., Tian, Y., Zhao, Y., Yu, H., Xie, L., Wang, Y., Ye, Q., Liu, Y.: Vmamba: Visual state space model. arXiv preprint arXiv:2401.10166 (2024)
  * [17] Long, J., Shelhamer, E., Darrell, T.: Fully convolutional networks for semantic segmentation. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 3431–3440 (2015)
  * [18] Ma, J., Li, F., Wang, B.: U-mamba: Enhancing long-range dependency for biomedical image segmentation. arXiv preprint arXiv:2401.04722 (2024)
  * [19] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: International Conference on Medical image computing and computer-assisted intervention. pp. 234–241. Springer (2015)
  * [20] Shi, J., He, Y., Kong, Y., Coatrieux, J.L., Shu, H., Yang, G., Li, S.: Xmorpher: Full transformer for deformable medical image registration via cross attention. In: International Conference on Medical Image Computing and Computer-Assisted Intervention. pp. 217–226. Springer (2022)
  * [21] Song, X., Guo, H., Xu, X., Chao, H., Xu, S., Turkbey, B., Wood, B.J., Wang, G., Yan, P.: Cross-Modal Attention for MRI and Ultrasound Volume Registration. In: International Conference on Medical Image Computing and Computer-Assisted Intervention. pp. 66–75. Springer (2021)
  * [22] Sotiras, A., Davatzikos, C., Paragios, N.: Deformable medical image registration: A survey. IEEE Transactions on Medical Imaging 32(7), 1153–1190 (2013)
  * [23] Thummerer, A., van der Bijl, E., Galapon Jr, A., Verhoeff, J.J., Langendijk, J.A., Both, S., van den Berg, C.N.A., Maspero, M.: Synthrad2023 grand challenge dataset: Generating synthetic ct for radiotherapy. Medical physics 50(7), 4664–4674 (2023)
  * [24] Vaswani, A., etc: Attention is all you need. In: Advances in neural information processing systems (2017)
  * [25] Wang, Z., Ma, C.: Semi-mamba-unet: Pixel-level contrastive cross-supervised visual mamba-based unet for semi-supervised medical image segmentation. arXiv preprint arXiv:2402.07245 (2024)
  * [26] Wang, Z., Ma, C.: Weak-mamba-unet: Visual mamba makes cnn and vit work better for scribble-based medical image segmentation. arXiv preprint arXiv:2402.10887 (2024)
  * [27] Wang, Z., Ma, C.: Dual-contrastive dual-consistency dual-transformer: A semi-supervised approach to medical image segmentation. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 870–879 (2023)
  * [28] Wang, Z., Yang, C.: Mixsegnet: Fusing multiple mixed-supervisory signals with multiple views of networks for mixed-supervised medical image segmentation. Engineering Applications of Artificial Intelligence 133, 108059 (2024)
  * [29] Wang, Z., et al.: Mamba-unet: Unet-like pure visual mamba for medical image segmentation. arXiv preprint arXiv:2402.05079 (2024)
  * [30] Xing, Z., Ye, T., Yang, Y., Liu, G., Zhu, L.: Segmamba: Long-range sequential modeling mamba for 3d medical image segmentation. arXiv preprint arXiv:2401.13560 (2024)
  * [31] Zhang, Y., Pei, Y., Zha, H.: Learning dual transformer network for diffeomorphic registration. In: International Conference on Medical Image Computing and Computer-Assisted Intervention. pp. 129–138. Springer (2021)
  * [32] Zhang, Z., Li, S., Wang, Z., Lu, Y.: A novel and efficient tumor detection framework for pancreatic cancer via ct images. In: 2020 42nd Annual International Conference of the IEEE Engineering in Medicine & Biology Society (EMBC). pp. 1160–1164. IEEE (2020)
  * [33] Zhao, S., Dong, Y., Chang, E.I., Xu, Y., et al.: Recursive cascaded networks for unsupervised medical image registration. In: Proceedings of the IEEE/CVF international conference on computer vision. pp. 10600–10610 (2019)
  * [34] Zheng, J.Q., Lim, N.H., Papież, B.W.: D-net: Siamese based network for arbitrarily oriented volume alignment. In: International Workshop on Shape in Medical Imaging. pp. 73–84. Springer (2020)
  * [35] Zheng, J.Q., Lim, N.H., Papież, B.W.: Accurate volume alignment of arbitrarily oriented tibiae based on a mutual attention network for osteoarthritis analysis. Computerized Medical Imaging and Graphics 106, 102204 (2023)
  * [36] Zheng, J.Q., Wang, Z., Huang, B., Lim, N.H., Papież, B.W.: Residual aligner-based network (ran): Motion-separable structure for coarse-to-fine discontinuous deformable registration. Medical Image Analysis 91, 103038 (2024)
  * [37] Zheng, J.Q., Wang, Z., Huang, B., Vincent, T., Lim, N.H., Papież, B.W.: Recursive deformable image registration network with mutual attention. In: Annual Conference on Medical Image Understanding and Analysis. pp. 75–86. Springer (2022)
  * [38] Zhu, L., Liao, B., Zhang, Q., Wang, X., Liu, W., Wang, X.: Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint arXiv:2401.09417 (2024)