# HI, FRB, what’s your z: The first FRB host galaxy redshift from radio
observations

M. Glowacki International Centre for Radio Astronomy Research (ICRAR), Curtin
University, Bentley, WA 6102, Australia  A. Bera  International Centre for
Radio Astronomy Research (ICRAR), Curtin University, Bentley, WA 6102,
Australia  K. Lee-Waddell International Centre for Radio Astronomy Research
(ICRAR), The University of Western Australia, 35 Stirling Hwy, Crawley, WA
6009, Australia  CSIRO Space and Astronomy, PO Box 1130, Bentley, WA 6102,
Australia  International Centre for Radio Astronomy Research (ICRAR), Curtin
University, Bentley, WA 6102, Australia  A. T. Deller  Centre for Astrophysics
and Supercomputing, Swinburne University of Technology, Hawthorn, VIC, 3122,
Australia  T. Dial  Centre for Astrophysics and Supercomputing, Swinburne
University of Technology, Hawthorn, VIC, 3122, Australia  K. Gourdji  Centre
for Astrophysics and Supercomputing, Swinburne University of Technology,
Hawthorn, VIC, 3122, Australia  M. Caleb  Sydney Institute for Astronomy
(SIfA), School of Physics, The University of Sydney, Camperdown NSW 2006,
Australia  L. Marnoch  School of Mathematical and Physical Sciences, Macquarie
University, NSW 2109, Australia  Australia Telescope National Facility, CSIRO
Space and Astronomy, P.O. Box 76, Epping, NSW 1710, Australia  Astrophysics
and Space Technologies Research Centre, Macquarie University, Sydney, NSW
2109, Australia  ARC Centre of Excellence for All-Sky Astrophysics in 3
Dimensions (ASTRO 3D), Australia  J. Xavier Prochaska University of California
- Santa Cruz 1156 High St. Santa Cruz, CA, USA 95064  Kavli IPMU (WPI), UTIAS,
The University of Tokyo, Kashiwa, Chiba 277-8583, Japan  Division of Science,
National Astronomical Observatory of Japan, 2-21-1 Osawa, Mitaka, Tokyo
181-8588, Japan  S. D. Ryder  School of Mathematical and Physical Sciences,
Macquarie University, NSW 2109, Australia  Astronomy, Astrophysics and
Astrophotonics Research Centre, Macquarie University, Sydney, NSW 2109,
Australia  R. M. Shannon  Centre for Astrophysics and Supercomputing,
Swinburne University of Technology, Hawthorn, VIC, 3122, Australia  N. Tejos
Instituto de Física, Pontificia Universidad Católica de Valparaíso, Casilla
4059, Valparaíso, Chile

###### Abstract

Identification and follow up observations of the host galaxies of fast radio
bursts (FRBs) not only help us understand the environments in which the FRB
progenitors reside, but also provide a unique way of probing the cosmological
parameters using the dispersion measures of FRBs and distances to their
origin. A fundamental requirement is an accurate distance measurement to the
FRB host galaxy, but for some sources viewed through the Galactic plane,
optical/NIR spectroscopic redshifts are extremely difficult to obtain due to
dust extinction. Here we report the first radio-based spectroscopic redshift
measurement for an FRB host galaxy, through detection of its neutral hydrogen
(Hi) 21-cm emission using MeerKAT observations. We obtain an Hi-based redshift
of $z=0.0357$ for the host galaxy of FRB 20230718A, an apparently non-
repeating FRB detected in the CRAFT survey and localized at a Galactic
latitude of –0.367∘. Our observations also reveal that the FRB host galaxy is
interacting with a nearby companion, which is evident from the detection of an
Hi bridge connecting the two galaxies. This result demonstrates the value of
Hi to obtain redshifts of FRBs at low Galactic latitudes and redshifts. Such
nearby FRBs whose dispersion measures are dominated by the Milky Way can be
used to characterise these components and thus better calibrate the remaining
cosmological contribution to dispersion for more distant FRBs that provide a
strong lever arm to examine the Macquart relation between cosmological DM and
redshift.

HI line emission — Fast radio bursts — Radio transient sources

##  1 Introduction

Fast radio bursts (FRBs), first identified by Lorimer et al. (2007), are
highly energetic radio pulses occurring on timescales of milliseconds. In
order to disentangle the various theories on FRB progenitors (see review by
Cordes & Chatterjee, 2019), FRBs need to be localised to their host galaxies,
and subsequently studied through follow-up observations. In addition, the
‘Macquart relation’ between the dispersion measure (DM) and redshift (Macquart
et al., 2020) has been shown to reveal the previously ‘missing baryons’ (hot
ionised gas in the intergalactic medium), can resolve the Universe’s large-
scale structure (Rafiei-Ravandi et al., 2021; Lee et al., 2022), and be used
for cosmological studies (James et al., 2022b; Baptista et al., 2023).
Spectroscopic redshift measurements of the FRB hosts are essential for such
studies.

However, spectroscopic redshift information of the FRB host galaxy candidates
is not always straightforward to obtain. Indeed, the first precisely localized
FRB (FRB 20121102A; Chatterjee et al., 2017) occurred in an intrinsically
faint galaxy along a modestly extincted sightline through the Galaxy. If
placed at $z>1$ or behind another magnitude of extinction, the identification
of the host would become prohibitively expensive for even the largest ground-
based telescopes. FRB hosts have been found at $z\sim 1$ (Ryder et al., 2023),
while Marnoch et al. (2023) presented follow-up of FRB 20210912A, which is
still without an identified host despite deep observations with the Very Large
Telescope (VLT) - even deeper optical time is required for these high-redshift
FRBs through limited telescope resources. Another scenario is an FRB localised
at low Galactic latitude, where dust extinction from the Galactic plane
hampers optical follow-up. This is an issue for FRB surveys such as the
Commensal Real-time ASKAP Fast Transients (CRAFT; Macquart et al., 2010;
Bannister et al., 2017) survey, which collects data simultaneously with other
scheduled observations, and so includes FRB searches during other major
science surveys pointed toward the Galactic plane, with the Australian Square
Kilometre Array Pathfinder telescope (ASKAP; Deboer et al., 2009; Hotan et
al., 2021). Other examples include the FAST telescope, where survey teams also
conduct commensal FRB searches, such as with the Galactic Plane Pulsar Survey
(GPPS; Han et al., 2021), and the Meer(more) TRAnsients and Pulsars (MeerTRAP;
Bezuidenhout et al., 2022). Through GPPS observations Zhou et al. (2023)
recently reported five FRBs all with Galactic latitude $b<4^{\circ}$.

Photometric redshifts are generally not accurate enough to be a reliable
substitute for optical spectroscopic measurements in the Macquart relation.
However, another alternative exists for FRBs at low Galactic latitudes:
spectral line transitions in wavelength ranges that are less affected by dust
extinction. For instance in the radio regime, the 21-cm (1420 MHz) line traces
neutral hydrogen (Hi), the star-forming fuel in galaxies. The 1665-1667 MHz
doublet for hydroxyl (OH) is typically associated with megamasers arising from
starburst activity within the host galaxy. Such transitions are not without
issues - for example, the strength of the 21-cm Hi transition decreases with
redshift and is hard to detect at $z>0.1$ due to both this effect and radio
frequency interference (RFI) from artificial satellites affecting the
$\sim$1150–1300 MHz regime ($0.09<z_{\rm HI}<0.23$), while the Hi radio beam
size is often much larger (tens of arcseconds) than the FRB localisation.
Nonetheless, it remains an avenue to obtain a spectroscopic redshift for FRB
host galaxies when optical follow-up is difficult. Hi has been detected and
studied in five FRB host galaxies to date, where four have shown signs of
strong asymmetry in the Hi global spectrum and/or disturbed Hi intensity maps
which were attributed to recent galaxy merger or interaction events
(Michałowski, 2021; Kaur et al., 2022; Lee-Waddell et al., 2023), and the
fifth with less signs of interaction (Glowacki et al., 2023). Thus far, no
non-detections of Hi have been reported in nearby FRB hosts, in line with the
finding by Gordon et al. (2023) that FRB host galaxies tend to be in star-
forming galaxies. Additionally, Hsu et al. (2023) presented asymmetric
profiles of molecular gas (CO) in the host galaxy of FRB 20180924B at
$z=0.3216$.

In this paper we present the first redshift for an FRB host galaxy solely
measured from Hi in emission. In Section 2 we describe the FRB localisation
and follow-up spectral-line observations in Hi, and present the Hi results in
Section 3.1. In Section 4 we briefly discuss the potential for future Hi
follow-up studies, and summarise our findings in Section 5.

##  2 Observations

###  2.1 FRB detection and localisation

During commensal observing with ASKAP, an FRB was successfully detected with a
signal-to-noise ratio S/N of 10.9 on 18 August 2023 UT 07:02:08. Through the
CRAFT Effortless Localisation and Enhanced Burst Inspection pipeline (CELEBI;
Scott et al., 2023), the FRB was localised to RA=08:32:38.804,
Dec=–40:27:06.33 (J2000), with a 1$\sigma$ uncertainty ellipse of $0\farcs 37$
in RA and $0\farcs 39$ in Dec, with a position angle of –2.3∘. The S/N of the
FRB in the post-processed image was 22.9. This FRB, henceforth FRB 20230718A,
was found to be coincident with the galaxy WISEA J083238.73–402705.3, with no
available redshift in the literature. The host galaxy is also seen in imaging
for the DECam Plane Survey (DECaPS2; Saydjari et al., 2023) as an apparent
faint red galaxy due to dust extinction (see Fig. 1). A preliminary analysis
of Probabilistic Association of Transients to their Hosts (PATH; Aggarwal et
al., 2021), without correcting for extinction nor removing any stars, lead to
a PATH posterior lower limit of 81%. Further analysis of the burst profile of
FRB 20230718 and the polarisation properties will be presented in Scott et
al., (in prep.).

The measured DM for FRB 20230718A from CELEBI is 476.6 pc cm-3, which
ordinarily would imply via the Macquart relation a host galaxy redshift
$z>0.3$. However, given the Galactic latitude of the host galaxy is –0.367∘,
the contribution of the Milky Way at this position is found to range from 393
or 421 pc cm-3 from the Baror/Prochaska111https://github.com/FRBs/ne2001 or
PyGEDM (Python Galactic electron density model; Price et al., 2021)
implementation for the NE2001 model (Cordes & Lazio, 2002, 2003), to as high
as 450 pc cm-3 from the YMW16 model (Yao et al., 2017). Hence the
extragalactic DM would be 26 $<$ DMEG $<$ 83 pc cm-3, consistent with a host
galaxy at $z<0.1$. We note that the cosmic DM DMcosmic = DMEG \- DMhost would
be smaller than the extragalactic DMEG when requiring the host galaxy DM $>0$
pc cm-3. High Galactic extinction ($>3$ mag in most bands; Schlafly &
Finkbeiner, 2011) is likely the reason why no optical spectroscopic redshift
was available for the galaxy WISEA J083238.73–402705.3, and motivated an
alternative approach for obtaining a spectroscopic redshift.

###  2.2 Searching for the host galaxy of FRB 20230718A in Hi-21 cm emission

We began by searching for archival Hi datasets. No detection was seen in the
Hi Parkes All Sky Survey (HIPASS; Barnes et al., 2001). We note that the
3$\sigma$ $M_{\rm HI}$ sensitivity in HIPASS is stated to be 106 $d_{\rm
Mpc}^{2}$ M⊙, which corresponded to $\sim$2.7$\times$1010 M⊙ at 163 Mpc (table
1 of Barnes et al. (2001)), higher than the Hi mass observed here (see Section
3.1). A search of the MeerKAT
archive222https://apps.sarao.ac.za/katpaws/archive-search hosted by the South
African Radio Astronomical Observatory (SARAO) indicated various pointings in
L-band which contained the FRB position by the ‘Legacy Survey of the Galactic
Plane’, project ID SSV-20180721-FC-01. Three observations with capture block
IDs of 1574801265, 1574911560, and 1574542866 were identified with the FRB
position within $0.5^{\circ}$ of the field centre of each pointing; two of the
observations had the same pointing for a combined 1 hr of on-source time. A
preliminary inspection of the data revealed weak Hi emission at the FRB
position. This motivated a deeper follow-up observation in L-band through
Director Discretionary Time (DDT). Observation 1697669539 was centred on the
FRB localisation, and carried out on 18–19 October 2023 with MeerKAT as part
of proposal ID DDT-20231017-MG-01 (see Table 1). The observation was taken in
L-band (856 MHz bandwidth centred at 1284 MHz) at 32K spectral resolution
(26.123 kHz wide channels, corresponding to 5.7 km s-1 channels at the
detected Hi emission frequency). The SARAO SDP continuum image quality report
gave a root mean square (RMS) noise of 11 $\mu$Jy. We simultaneously searched
for repeats bursts in realtime using the MeerTRAP single pulse search pipeline
as described in Caleb et al. (2022); Rajwade et al. (2022). No repeat bursts
were detected above a fluence threshold of 0.09 Jy ms for a 1 ms wide burst
around the DM of the FRB.

Table 1: Details of the MeerKAT observation for proposal ID DDT-20231017-MG-01. SBID | 1697669539
---|---
Phase centre | 08:32:38.804, –40:27:06.33
Bandpass calibrator | J1939-6342
Gain calibrator | J1744-5144
Channel width | 26.123 kHz (later rebinned by factor of 2)
On-source integration time | 4 hr
Robust weighting | 1
High spatial res cube |  UV range of 0 – $\sim$40k$\lambda$, radio beam of 19.3”$\times$15.2”, PA of 100∘, 2” pixels
|  2$\sigma$ $N_{\rm HI}$ sensitivity of 5.2$\times$1019 cm-2 in intensity map
Low spatial res cube |  UV range of 0 – $\sim$10k$\lambda$, radio beam of 38.6”$\times$24.9”, PA of 93∘, 5” pixels
|  2$\sigma$ $N_{\rm HI}$ sensitivity of 1.0$\times$1019 cm-2 in intensity map

Figure 1: Intensity, velocity, and dispersion (moments 0, 1, and 2) maps of
the MeerKAT Hi DDT observations, as measured by SoFiA. Top: intensity map
radio contours for the higher spatial resolution overlaid on DECam 3-colour
image. Contours are at column density multiples of 1, 2, 4, 8, … of
5.2$\times$1019 cm-2, where the lowest contour is at 2$\sigma$ significance.
The FRB localisation region is given by the cyan ellipse, enlarged by a factor
of 5 for visibility, and is coincident with a peak in Hi emission and a faint,
apparently red, optical galaxy (noting extinction is a factor). Bottom panels:
The higher spatial resolution spectral-line velocity and dispersion maps are
displayed.

Figure 2: Moment maps as in Fig. 1, for the lower resolution cube. Contours
are at column density multiples of 1, 2, 4, 8, … of 1.0$\times$1019 cm-2,
where the lowest contour is at 2$\sigma$ significance.

The raw data were transferred to the ilifu supercomputing cloud system and
reduced there. Bandpass, flux, and phase calibration, along with self-
calibrated continuum imaging, was performed using the processMeerKAT
pipeline333https://idia-pipelines.github.io/docs/processMeerKAT, which is
written in Python, uses a purpose-built CASA (McMullin et al., 2007)
Singularity container, and employs MPICASA (a parallelised form of CASA). Data
was at this stage rebinned by a factor of 2 (i.e. to 16K mode, 52.246 kHz-wide
channels). Model continuum visibility data were subtracted from the corrected
visibility data using the CASA task uvsub. A first-order polynomial fit to the
continuum was then calculated and subtracted using the CASA task uvcontsub for
all channels to remove residual continuum emission from the spectral line
data, with known Hi emission channels excluded from the continuum fit.
Finally, two spectral line cubes were created using tclean with robust = 1 and
cleaning to 3$\times$RMS (root mean square) levels, and variable uv-distances
for the uvrange parameter and pixel sizes with no uvtaper used (see Table 1).
The RMS per 52.246 kHz channel was fairly stable across the middle of the
band, with a per-channel noise of 0.18 and 0.16 mJy beam-1 at the Hi emission
for the FRB host galaxy for the low and high resolution cubes. All channels
were convolved to a common synthesized beam for each spectral-line cube, and
katbeam444https://github.com/ska-sa/katbeam used for primary beam correction.
The Source Finding Application 2 (SoFiA 2; Serra et al., 2015; Westmeier et
al., 2021) program was run on each cube.

##  3 Results

###  3.1 HI properties

Table 2: Measured properties of the host galaxy FRB 20230718A and its companion in order of Hi flux, Hi mass, Hi distance, Hi frequency $Freq_{\rm HI}$, Hi position, and radio continuum flux, as measured from the low spatial resolution cube (except for $Freq_{\rm HI}$ and Hi position, where measurements from the distinct SoFiA detections in the high resolution cube are used instead). Quantity | FRB host | Neighbour | Combined system
---|---|---|---
$S_{\rm HI}$ |  0.17 $\pm$ 0.02 Jy km s-1 |  0.45 $\pm$ 0.03 Jy km s-1 |  0.69 $\pm$ 0.03 Jy km s-1
$M_{\rm HI}$ |  9.7 $\pm\,$1.0$\times$108 M⊙ |  2.6 $\pm$ 0.2$\times$109 M⊙ |  4.0 $\pm$ 0.2$\times$109 M⊙
$D_{\rm HI}$ | 162.9 Mpc | 161.2 Mpc | 162.1 Mpc
$Freq_{\rm HI}$ | 1371.49 MHz | 1371.96 MHz | –
Hi Position | 08:32:38.642, –40:27:05.43 | 08:32:35.992, –40:26:04.56 | –
$S_{\rm cont}$ | 0.97 mJy | Host unclear | –

In Figures 1 and 2 we present Hi intensity maps, overlaid on a DECaPS2
3-colour image ($g,r,z$ bands), accompanied by the velocity and dispersion
maps for the high and low resolution Hi cubes respectively. Hi emission was
seen for both the FRB host galaxy with the peak overlapping with the FRB
localisation, and a neighbouring Hi-rich galaxy with no clear optical
counterpart (henceforth referred to as ‘neighbour’ or ‘neighbouring galaxy’).
Using the SoFiA measured Hi position for the neighbouring galaxy, the
projected separation distance between it and the FRB localisation is 55 kpc.
We also detect a Hi bridge connecting the two, at $>8\sigma$ sensitivity in
the low resolution cubes, or to a Hi column density of least 8$\times$1019
cm-2.

In Figure 3 we present Hi spectra for the putative FRB host galaxy, the
neighbouring galaxy, and the combined region including the Hi bridge, from the
low resolution cube. The presence of the Hi bridge, as well as a somewhat
lopsided Hi spectrum with more flux on the lower velocity side, shows that the
FRB host galaxy distribution is interacting with its neighbour. The
neighbouring galaxy appears to be regularly rotating with a lopsided double-
horned spectral-line profile. The dispersion map shows higher values towards
the south of the neighbouring galaxy rather than the centre. We do not detect
any other Hi sources along the line of sight to the FRB host galaxy, although
we do detect several other Hi galaxies in emission at a similar redshift
elsewhere in the cube (0.4–1.9 Mpc projected distances from the FRB host),
indicating this FRB host galaxy resides in a Hi-rich galaxy group.

Figure 3: Hi spectra for the FRB host galaxy (red), the neighbouring galaxy
(blue), and the total combined system including the Hi bridge (black). We note
the lack of any other spectroscopic redshift information to compare with here.

We summarise the following properties calculated from the Hi data and
ancillary data in Table 2, using the Hi redshift of $z~{}=~{}0.0357$
(corresponding to a distance of 162.8 Mpc assuming Planck cosmology and
$H_{\rm 0}$ = 67.7 km s-1 Mpc-1; Planck Collaboration et al., 2016). We note
that it is difficult to distinguish between the Hi associated with the FRB
host galaxy and the bridge connecting it to its neighbour. The companion is
more Hi massive than the FRB host by 2.7 times (or $\sim$65% of the combined
Hi mass).

An unresolved radio continuum source with a flux density of 0.97 mJy is
associated with the FRB and Hi emission at $\sim 5\sigma$ significance. The
WISE colours of $(W1-W2)=0.286\pm 0.075$ and $(W2-W3)=4.107\pm 0.092$ suggest
that the host of FRB 20230718A is unlikely to harbor a highly efficient
accreting AGN as it fits among the normal star-forming spiral region of the
Wide-field Infrared Survey Explorer (WISE) colour-colour diagram (see fig. 10
of Wright et al., 2010). The red $(W2-W3)$ colour is in part attributed to
Galactic extinction.

##  4 Discussion

###  4.1 HI environment

The FRB host galaxy is not isolated, given the clear detection of a Hi bridge
between it and the neighbour galaxy. Therefore, this would be the fifth FRB
host galaxy with a sign of galaxy interaction within the Hi observations to
date out of six published so far, akin to Michałowski (2021); Kaur et al.
(2022); Lee-Waddell et al. (2023). Owing to the higher dispersion seen in the
neighbour galaxy towards the south (FRB host galaxy) side rather than the
centre (Fig. 1, bottom-right panel), it is possible that the FRB host galaxy
is accumulating gas from the neighbour. This would align with the proposed
‘fast FRB channel model’ discussed in Michałowski (2021), where a galaxy
interaction event triggers star formation, leading to the birth of massive and
relatively short-lived stars that result in the creation of neutron stars or
magnetars, which are popular FRB progenitor candidates.

Additionally, at least seven other Hi emission detections within the same
MeerKAT observation have been made within a 5 MHz block containing the FRB
host galaxy Hi emission, identified from visual inspection and confirmed
through SoFiA. This FRB host hence resides in a Hi-rich galaxy group. In
cross-matching four optical galaxy group catalogues with the Arecibo Legacy
Fast ALFA Survey (ALFALFA; Haynes et al., 2018), Jones et al. (2020) found
that ALFALFA contained a far greater number of field galaxies rather than
group members, with no region of the environment parameter space explored with
group galaxies as the dominant population.

###  4.2 FRB host galaxy redshifts

This is the first demonstration of using a radio spectral transition to
measure the redshift of an FRB host galaxy. Currently, one limitation arises
from the diminishing intensity of the Hi 21-cm transition with redshift, where
only the most Hi massive galaxies would be visible with even a full night of
observation with MeerKAT above $z>0.23$. RFI limits Hi studies between
$0.09<z<0.23$. While OH megamaser emission seen through the 1665–1667 MHz
doublet spans a different redshift space not affected by RFI (e.g. RFI affects
the $0.27<z<0.44$ redshift space for this transition), such detections are
less plentiful compared to Hi as they are typically only associated with
starburst galaxies. Little more than 100 OH megamasers exist in the current
literature, although the advent of SKA pathfinders are poised to improve on
this space (Roberts et al., 2021).

Hi searches in absorption avoid the issue of a fall-off in intensity of the
transition with redshift seen for emission-line searches, but would require
either the FRB host to be a sufficiently bright radio galaxy to enable such a
search, or a search within the FRB signal itself. The latter is particularly
interesting as it would yield information on both the cold neutral hydrogen
and the hot ionised gas within the FRB sightline (see Fender & Oosterloo,
2015). No Hi absorption was seen from a preliminary analysis of the FRB
20230718A burst profile.

We briefly consider different applications from using the 21-cm Hi transition
for determining redshifts for FRB host galaxies. Potentially 10s to 100s of
FRBs a year will be found through the Galactic plane. The host galaxies of
these FRBs are both useful for specific scientific goals, and difficult to
obtain via optical spectroscopy. The first application of low Galactic
latitude FRBs, specifically those originating in nearby galaxies, is to assist
in characterising the Milky Way disk and halo DM models. Nearby FRBs have the
twin benefit that the hosts can most easily be identified in Hi emission,
while the uncertainty in the minimum non-Milky Way contribution can be very
small for low-DM bursts. Currently, halo models are uncertain - the DM
contributed by Milky Way halo DMhalo has been estimated to range 10–80 pc cm-3
(Prochaska & Zheng, 2019; Keating & Pen, 2020), and is often simply assumed to
be 50 pc cm-3 (e.g. James et al., 2022a). For FRB 20230718A, with DMEG ranging
between $\sim$26–83 pc cm-3, and assuming 10 pc cm-3 per 0.01 in redshift (as
done in James et al. (2022a) for the Macquart relation), DMcosmic would be
$\sim$35.7 pc cm-3, placing an upper limit of both DMhost and DMhalo combined
of $\sim$47.3 pc cm-3. Even assuming the lower end for DMhalo would leave a
fairly modest amount of DM contributed by the FRB host galaxy.

Another application is using the distances derived from Hi measurements for
FRB hosts to help determine the distances of scattering screens to the FRB
source (Sammons et al., 2023), given that Galactic scintillation can be
assumed to be a dominating effect (e.g. Chawla et al., 2022). This in turn can
aid model comparison for FRB progenitors and their immediate environments,
alongside disentangling estimations of both DMhost and DMcosmic (Ocker et al.,
2022).

Nearby FRBs away from the Galactic plane could also be localised through the
Hi 21-cm transition. The number of FRBs localised will improve from an
expected increase in FRB detection rates, such as for CRAFT through its CRACO
(CRAFT COherent) upgrade to the real-time detection system. While most FRBs
detected by CHIME are not well constrained in their position to enable
redshift searches, new localisations for one-off and repeating bursts have
been reported (Bhardwaj et al., 2023; Ibik et al., 2023), while the CHIME
Outriggers project will detect and localize FRBs with 50 mas precision (Mena-
Parra et al., 2022). The 110-antenna Deep Synoptic Array (DSA-110) telescope
also recently shared localisations for 11 FRBs (Law et al., 2023). A lower
localisation precision is sufficient to obtain a host galaxy for nearby FRBs,
compared to higher redshift FRBs. Previously conducted large Hi surveys such
as HIPASS and ALFALFA will be searchable against any such FRB localisations,
as will the Widefield ASKAP L-band Legacy All-sky Blind surveY (WALLABY;
Koribalski et al., 2020) which is currently ongoing and resulted in the first
commensal detection of an FRB and the Hi in its host galaxy (Glowacki et al.,
2023). WALLABY offers an improved angular resolution and depth ($z<0.09$) to
ALFALFA and HIPASS, and so will enable association of nearby FRBs to hosts
without the need for dedicated optical follow-up. Probabilistic associations
of nearby FRBs through PATH could hence be applied to Hi data, rather than
optical datasets.

##  5 Conclusions

We have presented Hi follow-up observations from the MeerKAT radio telescope
of FRB 20230718A, which was localised to sub-arcsecond precision to a faint
dust extinct galaxy at low Galactic latitude. Hi in emission was detected for
the host galaxy, as well as a nearby companion galaxy with no clear optical
host in the available optical photometry. A Hi bridge is seen between the two
galaxies, indicating that the host of FRB 20230718A is interacting and
potentially undergoing enhanced star formation, which could have led to a
progenitor formed from a massive star, as seen in four other FRB host galaxies
detected in Hi.

From this observation we obtained a redshift of 0.0357, the first redshift
measurement for an FRB host galaxy not derived first from optical
observations. This demonstrates the ability to use the Hi 21-cm transition to
follow up FRB host galaxies for which it may be difficult to obtain optical
spectroscopic measurements, such as for other FRB hosts behind the Galactic
plane. Redshift information obtained through Hi provides an avenue to
including such localised FRBs in studies such as constraining the DM
contribution from the Milky Way and FRB host galaxy, FRB progenitor models,
and cosmological studies.

##  6 Acknowledgements

We thank Ben Stappers for assistance with the MeerTRAP single pulse search
pipeline. MG is supported by the Australian Government through the Australian
Research Council’s Discovery Projects funding scheme (DP210102103). RMS and
ATD acknowledge support through Australian Research Council Future Fellowship
FT190100155 and Discovery Project DP220102305. MC acknowledges support of an
ARC Discovery Early Career Research Award DE220100819 funded by the Australian
Government and the ARC Centre of Excellence for All Sky Astrophysics in 3
Dimensions (ASTRO 3D), through project number CE170100013. KG acknowledges
support through Australian Research Council Discovery Project DP200102243. LM
acknowledges the receipt of an MQ-RES scholarship from Macquarie University.
Authors JXP and NT, as members of the Fast and Fortunate for FRB Follow-up
team, acknowledge support from NSF grants AST-1911140, AST-1910471 and
AST-2206490.

This scientific work uses data obtained from Inyarrimanha Ilgari Bundara, the
CSIRO Murchison Radio-astronomy Observatory. We acknowledge the Wajarri Yamaji
People as the Traditional Owners and native title holders of the Observatory
site. CSIRO’s ASKAP radio telescope is part of the Australia Telescope
National Facility (https://ror.org/05qajvd42). Operation of ASKAP is funded by
the Australian Government with support from the National Collaborative
Research Infrastructure Strategy. ASKAP uses the resources of the Pawsey
Supercomputing Research Centre. Establishment of ASKAP, Inyarrimanha Ilgari
Bundara, the CSIRO Murchison Radio-astronomy Observatory and the Pawsey
Supercomputing Research Centre are initiatives of the Australian Government,
with support from the Government of Western Australia and the Science and
Industry Endowment Fund. We also thank the MRO site staff. The MeerKAT
telescope is operated by the South African Radio Astronomy Observatory, which
is a facility of the National Research Foundation, an agency of the Department
of Science and Innovation.

This work was performed on the OzSTAR national facility at Swinburne
University of Technology. The OzSTAR program receives funding in part from the
Astronomy National Collaborative Research Infrastructure Strategy (NCRIS)
allocation provided by the Australian Government, and from the Victorian
Higher Education State Investment Fund (VHESIF) provided by the Victorian
Government. We acknowledge the use of the ilifu cloud computing facility -
www.ilifu.ac.za, a partnership between the University of Cape Town, the
University of the Western Cape, Stellenbosch University, Sol Plaatje
University, the Cape Peninsula University of Technology and the South African
Radio Astronomy Observatory. The ilifu facility is supported by contributions
from the Inter-University Institute for Data Intensive Astronomy (IDIA - a
partnership between the University of Cape Town, the University of Pretoria
and the University of the Western Cape), the Computational Biology division at
UCT and the Data Intensive Research Initiative of South Africa (DIRISA). This
work was carried out using the data processing pipelines developed at the
Inter-University Institute for Data Intensive Astronomy (IDIA) and available
at https://idia-pipelines.github.io. IDIA is a partnership of the University
of Cape Town, the University of Pretoria and the University of the Western
Cape. This work made use of the CARTA (Cube Analysis and Rendering Tool for
Astronomy) software (DOI 10.5281/zenodo.3377984 – https://cartavis.github.io;
Comrie et al., 2021). This research has made use of the NASA/IPAC
Extragalactic Database (NED) which is operated by the Jet Propulsion
Laboratory, California Institute of Technology, under contract with the
National Aeronautics and Space Administration. This research made use of
hips2fits,555https://alasky.cds.unistra.fr/hips-image-services/hips2fits a
service provided by CDS.

## References

  * Aggarwal et al. (2021) Aggarwal, K., Budavári, T., Deller, A. T., et al. 2021, ApJ, 911, 95, doi: 10.3847/1538-4357/abe8d2
  * Bannister et al. (2017) Bannister, K. W., Shannon, R. M., Macquart, J. P., et al. 2017, ApJ, 841, L12, doi: 10.3847/2041-8213/aa71ff
  * Baptista et al. (2023) Baptista, J., Prochaska, J. X., Mannings, A. G., et al. 2023, arXiv e-prints, arXiv:2305.07022, doi: 10.48550/arXiv.2305.07022
  * Barnes et al. (2001) Barnes, D. G., Staveley-Smith, L., de Blok, W. J. G., et al. 2001, MNRAS, 322, 486, doi: 10.1046/j.1365-8711.2001.04102.x
  * Bezuidenhout et al. (2022) Bezuidenhout, M. C., Barr, E., Caleb, M., et al. 2022, MNRAS, 512, 1483, doi: 10.1093/mnras/stac579
  * Bhardwaj et al. (2023) Bhardwaj, M., Michilli, D., Kirichenko, A. Y., et al. 2023, arXiv e-prints, arXiv:2310.10018, doi: 10.48550/arXiv.2310.10018
  * Caleb et al. (2022) Caleb, M., Heywood, I., Rajwade, K., et al. 2022, Nature Astronomy, 6, 828, doi: 10.1038/s41550-022-01688-x
  * Chatterjee et al. (2017) Chatterjee, S., Law, C. J., Wharton, R. S., et al. 2017, Nature, 541, 58, doi: 10.1038/nature20797
  * Chawla et al. (2022) Chawla, P., Kaspi, V. M., Ransom, S. M., et al. 2022, ApJ, 927, 35, doi: 10.3847/1538-4357/ac49e1
  * Comrie et al. (2021) Comrie, A., Kuo-Song Wang, Shou-Chieh Hsu, et al. 2021, CARTA: The Cube Analysis and Rendering Tool for Astronomy, Zenodo, doi: 10.5281/ZENODO.3377984
  * Cordes & Chatterjee (2019) Cordes, J. M., & Chatterjee, S. 2019, ARA&A, 57, 417, doi: 10.1146/annurev-astro-091918-104501
  * Cordes & Lazio (2002) Cordes, J. M., & Lazio, T. J. W. 2002, arXiv e-prints, astro, doi: 10.48550/arXiv.astro-ph/0207156
  * Cordes & Lazio (2003) —. 2003, arXiv e-prints, astro, doi: 10.48550/arXiv.astro-ph/0301598
  * Deboer et al. (2009) Deboer, B. D. R., Gough, R. G., Bunton, J. D., et al. 2009, Proc. IEEE, 97, 1507, doi: 10.1109/JPROC.2009.2016516
  * Fender & Oosterloo (2015) Fender, R., & Oosterloo, T. 2015, MNRAS, 451, L75, doi: 10.1093/mnrasl/slv065
  * Glowacki et al. (2023) Glowacki, M., Lee-Waddell, K., Deller, A. T., et al. 2023, ApJ, 949, 25, doi: 10.3847/1538-4357/acc1e3
  * Gordon et al. (2023) Gordon, A. C., Fong, W.-f., Kilpatrick, C. D., et al. 2023, ApJ, 954, 80, doi: 10.3847/1538-4357/ace5aa
  * Han et al. (2021) Han, J. L., Wang, C., Wang, P. F., et al. 2021, Research in Astronomy and Astrophysics, 21, 107, doi: 10.1088/1674-4527/21/5/107
  * Haynes et al. (2018) Haynes, M. P., Giovanelli, R., Kent, B. R., et al. 2018, ApJ, 861, 49, doi: 10.3847/1538-4357/aac956
  * Hotan et al. (2021) Hotan, A. W., Bunton, J. D., Chippendale, A. P., et al. 2021, PASA, 38, e009, doi: 10.1017/pasa.2021.1
  * Hsu et al. (2023) Hsu, T.-Y., Hashimoto, T., Hatsukade, B., et al. 2023, MNRAS, 519, 2030, doi: 10.1093/mnras/stac3655
  * Ibik et al. (2023) Ibik, A. L., Drout, M. R., Gaensler, B. M., et al. 2023, arXiv e-prints, arXiv:2304.02638, doi: 10.48550/arXiv.2304.02638
  * James et al. (2022a) James, C. W., Prochaska, J. X., Macquart, J. P., et al. 2022a, MNRAS, 509, 4775, doi: 10.1093/mnras/stab3051
  * James et al. (2022b) James, C. W., Ghosh, E. M., Prochaska, J. X., et al. 2022b, MNRAS, 516, 4862, doi: 10.1093/mnras/stac2524
  * Jones et al. (2020) Jones, M. G., Hess, K. M., Adams, E. A. K., & Verdes-Montenegro, L. 2020, MNRAS, 494, 2090, doi: 10.1093/mnras/staa810
  * Kaur et al. (2022) Kaur, B., Kanekar, N., & Prochaska, J. X. 2022, ApJ, 925, L20, doi: 10.3847/2041-8213/ac4ca8
  * Keating & Pen (2020) Keating, L. C., & Pen, U.-L. 2020, MNRAS, 496, L106, doi: 10.1093/mnrasl/slaa095
  * Koribalski et al. (2020) Koribalski, B. S., Staveley-Smith, L., Westmeier, T., et al. 2020, Ap&SS, 365, 118, doi: 10.1007/s10509-020-03831-4
  * Law et al. (2023) Law, C. J., Sharma, K., Ravi, V., et al. 2023, arXiv e-prints, arXiv:2307.03344, doi: 10.48550/arXiv.2307.03344
  * Lee et al. (2022) Lee, K.-G., Ata, M., Khrykin, I. S., et al. 2022, ApJ, 928, 9, doi: 10.3847/1538-4357/ac4f62
  * Lee-Waddell et al. (2023) Lee-Waddell, K., James, C. W., Ryder, S. D., et al. 2023, PASA, 40, e029, doi: 10.1017/pasa.2023.27
  * Lorimer et al. (2007) Lorimer, D. R., Bailes, M., McLaughlin, M. A., Narkevic, D. J., & Crawford, F. 2007, Science, 318, 777, doi: 10.1126/science.1147532
  * Macquart et al. (2010) Macquart, J.-P., Bailes, M., Bhat, N. D. R., et al. 2010, PASA, 27, 272, doi: 10.1071/AS09082
  * Macquart et al. (2020) Macquart, J. P., Prochaska, J. X., McQuinn, M., et al. 2020, Nature, 581, 391, doi: 10.1038/s41586-020-2300-2
  * Marnoch et al. (2023) Marnoch, L., Ryder, S. D., James, C. W., et al. 2023, MNRAS, 525, 994, doi: 10.1093/mnras/stad2353
  * McMullin et al. (2007) McMullin, J. P., Waters, B., Schiebel, D., Young, W., & Golap, K. 2007, in Astronomical Society of the Pacific Conference Series, Vol. 376, Astronomical Data Analysis Software and Systems XVI, ed. R. A. Shaw, F. Hill, & D. J. Bell, 127
  * Mena-Parra et al. (2022) Mena-Parra, J., Leung, C., Cary, S., et al. 2022, AJ, 163, 48, doi: 10.3847/1538-3881/ac397a
  * Michałowski (2021) Michałowski, M. J. 2021, ApJ, 920, L21, doi: 10.3847/2041-8213/ac2b35
  * Ocker et al. (2022) Ocker, S. K., Cordes, J. M., Chatterjee, S., et al. 2022, ApJ, 931, 87, doi: 10.3847/1538-4357/ac6504
  * Planck Collaboration et al. (2016) Planck Collaboration, Ade, P. A. R., Aghanim, N., et al. 2016, A&A, 594, A13, doi: 10.1051/0004-6361/201525830
  * Price et al. (2021) Price, D. C., Flynn, C., & Deller, A. 2021, PASA, 38, e038, doi: 10.1017/pasa.2021.33
  * Prochaska & Zheng (2019) Prochaska, J. X., & Zheng, Y. 2019, MNRAS, 485, 648, doi: 10.1093/mnras/stz261
  * Rafiei-Ravandi et al. (2021) Rafiei-Ravandi, M., Smith, K. M., Li, D., et al. 2021, ApJ, 922, 42, doi: 10.3847/1538-4357/ac1dab
  * Rajwade et al. (2022) Rajwade, K. M., Bezuidenhout, M. C., Caleb, M., et al. 2022, MNRAS, 514, 1961, doi: 10.1093/mnras/stac1450
  * Roberts et al. (2021) Roberts, H., Darling, J., & Baker, A. J. 2021, ApJ, 911, 38, doi: 10.3847/1538-4357/abe944
  * Ryder et al. (2023) Ryder, S. D., Bannister, K. W., Bhandari, S., et al. 2023, Science, 382, 294, doi: 10.1126/science.adf2678
  * Sammons et al. (2023) Sammons, M. W., Deller, A. T., Glowacki, M., et al. 2023, MNRAS, 525, 5653, doi: 10.1093/mnras/stad2631
  * Saydjari et al. (2023) Saydjari, A. K., Schlafly, E. F., Lang, D., et al. 2023, ApJS, 264, 28, doi: 10.3847/1538-4365/aca594
  * Schlafly & Finkbeiner (2011) Schlafly, E. F., & Finkbeiner, D. P. 2011, ApJ, 737, 103, doi: 10.1088/0004-637X/737/2/103
  * Scott et al. (2023) Scott, D. R., Cho, H., Day, C. K., et al. 2023, arXiv e-prints, arXiv:2301.13484, doi: 10.48550/arXiv.2301.13484
  * Serra et al. (2015) Serra, P., Westmeier, T., Giese, N., et al. 2015, MNRAS, 448, 1922, doi: 10.1093/mnras/stv079
  * Westmeier et al. (2021) Westmeier, T., Kitaeff, S., Pallot, D., et al. 2021, MNRAS, 506, 3962, doi: 10.1093/mnras/stab1881
  * Wright et al. (2010) Wright, E. L., Eisenhardt, P. R. M., Mainzer, A. K., et al. 2010, AJ, 140, 1868, doi: 10.1088/0004-6256/140/6/1868
  * Yao et al. (2017) Yao, J. M., Manchester, R. N., & Wang, N. 2017, ApJ, 835, 29, doi: 10.3847/1538-4357/835/1/29
  * Zhou et al. (2023) Zhou, D. J., Han, J. L., Jing, W. C., et al. 2023, MNRAS, 526, 2657, doi: 10.1093/mnras/stad2769


# MaIL: Improving Imitation Learning with Mamba

Xiaogang Jia ${\color[rgb]{1,1,1}i}^{\dagger\ddagger}$ Qian Wang∗† Atalay
Donat† Bowen Xing‡ Ge Li†
Hongyi Zhou‡ Onur Celik† Denis Blessing† Rudolf Lioutikov‡ Gerhard Neumann†
† Autonomous Learning Robots, Karlsruhe Institute of Technology
‡ Intuitive Robots Lab, Karlsruhe Institute of Technology Equal contribution,
correspondence to<EMAIL_ADDRESS>
###### Abstract

This work introduces Mamba Imitation Learning (MaIL), a novel imitation
learning (IL) architecture that offers a computationally efficient alternative
to state-of-the-art (SoTA) Transformer policies. Transformer-based policies
have achieved remarkable results due to their ability in handling human-
recorded data with inherently non-Markovian behavior. However, their high
performance comes with the drawback of large models that complicate effective
training. While state space models (SSMs) have been known for their
efficiency, they were not able to match the performance of Transformers. Mamba
significantly improves the performance of SSMs and rivals against
Transformers, positioning it as an appealing alternative for IL policies. MaIL
leverages Mamba as a backbone and introduces a formalism that allows using
Mamba in the encoder-decoder structure. This formalism makes it a versatile
architecture that can be used as a standalone policy or as part of a more
advanced architecture, such as a diffuser in the diffusion process. Extensive
evaluations on the LIBERO IL benchmark and three real robot experiments show
that MaIL: i) outperforms Transformers in all LIBERO tasks, ii) achieves good
performance even with small datasets, iii) is able to effectively process
multi-modal sensory inputs, iv) is more robust to input noise compared to
Transformers.

> Keywords: Imitation Learning, Sequence Models, Denoising Diffusion Policies

##  1 Introduction

Imitation learning (IL) [1] from human data has shown remarkable success in
acquiring robot policies that can solve complex tasks [2, 3, 4, 5, 6]. As
human behavior is inherently multi-modal and non-Markovian, prior research has
demonstrated great benefits in using historical information [7], predicting a
sequence of actions instead of a single action [3, 5] and using methods that
can model multi-modal distributions [5, 6, 8]. Recent works [5, 6] therefore
base their policies on Transformer models to effectively handle sequences of
observations. The Transformer’s self-attention and cross-attention mechanisms
have led to remarkable results in various domains [9, 10, 11] and are
considered state-of-the-art for processing sequential data. Here, current
approaches either use an decoder-only structure [5], or a decoder-encoder
architecture [6]. Which of these architectures excels is often task dependent.
The performance of transformers usually comes with large models that are
difficult to train, especially in domains where data is scarce. An alternative
concept for handling a sequence of observations are state space models [12].
These models assume a linear relationship between observations (embeddings)
and are usually computationally more efficient. Recent approaches such as
Mamba [13], a selective state space model, rigorously improve the performance
of state space models and rival against transformers in many tasks. Due to its
properties in inference speed, memory usage, and efficiency, Mamba is an
appealing model for IL policies.

This work proposes MaIL, a novel imitation learning policy architecture that
uses Mamba as a backbone. MaIL can be used as a standalone policy, or as part
of more advanced processes such as a diffuser in the diffusion process. We
implement MaIL in two variants. In the decoder-only variant, MaIL processes
the noised actions and the observation features [5] together with the time
embedding of the diffusion process and outputs the denoised actions. However,
due to Mamba’s formalism, an encoder-decoder variant is not straightforwardly
implementable due to possibly varying input and output structures. We show how
we can extend MaIL to an encoder-decoder variant by extending the inputs with
learnable action, state, and time embedding variables such that learning can
still be done efficiently. We show that this model works more efficiently for
multi-modal inputs such as image and language input. Our extensive evaluations
on the state-of-the-art LIBERO [14] IL benchmarks with and without history
input show that MaIL achieves improved performance compared to similarly
complex transformer-based IL policies. MaIL’s high performance is further
confirmed on three challenging real robot experiments.

##  2 Related Works

Sequence Models. In recent years, Transformers [15, 10, 16] have become the
leading approach in handling long-term dependencies in sequential data. The
self-attention mechanism in Transformers allows processing sequences in
parallel, effectively addressing the limitations of RNN in sequential data
processing [17, 18, 19, 20]. However, structured state space models [12, 21,
22, 13] provide an appealing alternative to Transformers. While transformers
scale quadratically in the sequence length, structured state space models
scale linearly [13]. Earlier approaches [12, 23, 24] rely on a convolutional
formulation of the learnable matrices that allows training in parallel and
only requires sequential processing during inference. However, this
formulation requires time and input-independent state space matrices harming
the performance. In contrast, recent works [13] rely on associative scans that
also allow for parallel computations but additionally allow for input-
dependent learnable matrices [13]. For a more detailed review, we refer the
reader to [25]. This work leverages the state-of-the-art method Mamba [13]
that showed competing performance compared to transformers. MaIL exploits the
selective state space models present in Mamba and proposes a novel
architecture for improved performance in imitation learning, especially in
sequential inputs.

Imitation Learning (IL). Early imitation learning methods primarily focused on
learning one-to-one mappings between state-action pairs. Despite demonstrating
promising results in tasks like autonomous driving [26] and robot control [27,
8, 28], these methods overlooked the rich temporal information contained in
history. Subsequent approaches incorporated RNNs to encode observation
sequences [29, 30, 31, 32], demonstrating that the utilization of historical
observations can enhance model performance. However, these methods suffer from
the inherent limitations of RNN-based architectures, including restricted
representation power with long-sequence modeling and slow training times due
to their unsuitability for large-scale parallelization. With the raise of
Transformers in NLP and CV domains, modern imitation learning methods have
adopted Transformers as their backbone [33, 34, 35]. Transformers can model
long sequences while maintaining training efficiency through parallelized
sequence processing. This trend extends to IL with multi-modal sensory inputs
[36, 37, 38, 39], where Transformers encode both image and language sequences.
Recently, Diffusion Models has demonstrated superiority in imitation learning
[5, 40, 6, 41, 39]. Due to their strong generalization ability and rich
representation power to capture multimodal action distributions, they have
become the SoTA in the imitation learning domain. Many of these models also
utilize Transformers as policies, leveraging their rich representation
capabilities [6, 41, 39]. Instead of relying on a transformer backbone, MaIL
proposes a novel architecture that is based on Mamba [13]. In the evaluations,
we show the benefits of MaIL, especially in scenarios with sequences of
observation inputs over state-of-the-art methods.

##  3 Preliminaries

Figure 1:  D-Ma: Mamba denoising architecture integrates ResNet-18 for state
encoding and an action encoder for action encoding. The state sequence has a
length of $K$, while the action sequence at diffusion step $t$ has a length of
$J$. Before feeding the inputs into the Mamba module, positional encoding (PE)
and time encoding (TE) enhance the inputs, where the $s_{k}$ and $a_{k}$ share
the same positional encoding. The mamba module has N$\times$ mamba blocks,
with a detailed structure [13] shown on the left. The outputs from the mamba
module are processed by a linear output layer, resulting the one-step
denoising actions. The symbol $\bm{\times}$ in the Mamba block denotes matrix
multiplication, and $\bm{\sigma}$ the SiLU activation function. Figure 2: ED-
Ma: Different from the D-Ma model, ED-Ma contains a Mamba encoder which is
used to process the time embedding and state embedding, and a Mamba decoder
which is used to process the noisy actions. In order to aggregate the
information from encoder and decoder, learnable action variables are
introduced to the encoder input and learnable time variables and state
variables are introduced to the decoder output for sequence alignment.

Here, we briefly review the basics of the Mamba model and explain the policy
representations considered in this work.

###  3.1 Mamba: Selective State-Space Models

Inspired by the Attention mechanism [15] in Transformers, Mamba [13] improves
upon the Structured State-Space Sequence Models (SSMs) [12] by using a
selective scan operator, to propagate or forget information over time,
allowing it to filter out relevant features. Specifically, we denote the input
sequence $x\in\mathbb{R^{\mathtt{B}\times\mathtt{L}\times\mathtt{D}}}$ and the
output sequence $y\in\mathbb{R^{\mathtt{B}\times\mathtt{L}\times\mathtt{D}}}$,
where $\mathtt{B}$, $\mathtt{L}$, $\mathtt{D}$ refer to the batch size,
sequence length and dimension respectively. A standard SSM with hidden state
dimension $\mathtt{N}$ defines time-invariant parameters
$\mathbf{A}\in\mathbb{R}^{\mathtt{N}\times\mathtt{N}}$,
$\mathbf{B}\in\mathbb{R}^{\mathtt{N}\times\mathtt{D}}$,
$\mathbf{C}\in\mathbb{R}^{\mathtt{D}\times\mathtt{N}}$ and the time-step
vector $\mathbf{\Delta}$ to map the inputs up to $x(l)$ to a hidden state,
which then can be projected to the output $y(l)$ (More details are in Appendix
A). However, this property limits the SSM for contextual learning. Mamba
implements the selection mechanism by making the SSM parameters a function of
the input, that is,

$\displaystyle\quad\mathbf{B}=\text{Linear}(x),\ \mathbf{C}=\text{Linear}(x),\
\mathbf{\Delta}=\text{SoftPlus}(\text{Linear}(x)),$ (1)

where $\mathbf{B}\in\mathbb{R}^{\mathtt{B}\times\mathtt{L}\times\mathtt{D}}$,
$\mathbf{C}\in\mathbb{R}^{\mathtt{B}\times\mathtt{L}\times\mathtt{D}}$, and
$\mathbf{\Delta}\in\mathbb{R}^{\mathtt{B}\times\mathtt{L}\times\mathtt{D}}$,
Linear refers to linear projection layers and SoftPlus is a smooth
approximation of ReLU. Then the output can be calculated via

$\displaystyle\quad
y=\text{SSM}(\mathbf{\overline{A},\mathbf{\overline{B}},C})(x),$ (2)

where $(\mathbf{\overline{A}},\mathbf{\overline{B}})$ are discretized [42]
counterparts of $(\mathbf{{A}},\mathbf{{B}})$ with time-steps
$\mathbf{\Delta}$. As the time-varying model can only be calculated in a
recurrent way, Mamba further implements a hardware-aware method to compute the
selective SSM efficiently. Figure 1 shows a simplified Mamba block.

###  3.2 Policy Representations

In this work, we use two policy representations: behavioral cloning (BC) and
denoising diffusion policies (DDPs). For clarity, we focus on the non-
sequential case.

Behavioral Cloning assumes a parameterized conditional Gaussian distribution
as policy representation, i.e.,
$\pi(\mathbf{a}|\mathbf{s})=\mathcal{N}(\mathbf{a}|\mu_{\theta}(\mathbf{s}),\sigma^{2}\textbf{I})$.
Maximizing the likelihood for the model parameters $\theta$ simplifies to a
mean-square error (MSE) loss, that is,

$\mathcal{L}_{\text{BC}}(\theta)=\mathbb{E}_{\mathbf{s},\mathbf{a}}\left[\|\mu_{\theta}(\mathbf{s})-\mathbf{a}\|^{2}_{2}\right],$
(3)

where the expectation over $\mathbf{s},\mathbf{a}$ is approximated using
state-action pairs in the demonstration data.

Denoising Diffusion Policies utilize a denoising function
$\bm{\epsilon}_{\theta}$ to sample from a Markov chain
$(\mathbf{a}^{t})_{t=1}^{T}$

$\mathbf{a}^{t-1}=\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{a}^{t}-\frac{1-\alpha_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\bm{\epsilon}_{\theta}(\mathbf{a}^{t},t,\mathbf{s})\right)+\sqrt{\alpha_{t}}\mathbf{z}^{t},\quad\text{where}\quad\mathbf{z}^{t}\sim\mathcal{N}(\mathbf{0},\textbf{I}),$
(4)

starting from $\mathbf{a}_{N}\sim\mathcal{N}(\mathbf{0},\textbf{I})$ to
produce a noise-free action $\mathbf{a}^{0}$ for a given observation
$\mathbf{s}$. The denoising function is trained to predict the source noise
$\mathbf{z}^{0}\sim\mathcal{N}(\mathbf{0},\textbf{I})$ of a noisy action
$\mathbf{a}^{t}=\sqrt{\bar{\alpha}_{t}}\mathbf{a}^{0}+\mathbf{z}^{t}$ by
minimizing the loss

$\mathcal{L}_{\text{DPM}}(\theta)=\mathbb{E}_{\mathbf{s},\mathbf{a},t}\left[\|\bm{\epsilon}_{\theta}(\mathbf{a}^{t},t,\mathbf{s})-\mathbf{z}^{0}\|^{2}_{2}\right],$
(5)

where $\bar{\alpha}_{t}=\prod_{j=1}^{t}\alpha_{j}$. The expectation over $t$
corresponds to a uniform sampling in $\\{1,\dots,T\\}$.

##  4 Mamba for Imitation Learning

In this section, we explain how we leverage Mamba for Imitation Learning (IL).
Drawing inspiration from the successful Decoder-only (D-Tr) and Encoder-
Decoder (ED-Tr) Transformers, we propose two Mamba-based architectures:
Decoder-only Mamba (D-Ma) and Encoder-Decoder Mamba (ED-Ma). These
architectures serve as the parameterization for the policy. Specifically, when
employing Behavioral Cloning (BC), these architectures parameterize the mean
$\mu_{\theta}$ of a conditional Gaussian distribution. Conversely, when using
Denoising Diffusion Policies (DDPs), these architectures parameterize the
denoising function $\bm{\epsilon}_{\theta}$. Given the straightforward nature
of the former case, we focus on introducing these architectures in the context
of DDPs.

Algorithm 1 D-Ma DDP Inference

1:Observation $\mathbf{s}_{k-K:k}$,

2: Mamba Decoder $\mathbf{D_{m}}$,

3: Action Encoder $\mathbf{E_{\mathbf{a}}}$,

4: Position Embedding $\mathbf{PE}$

5:$\mathbf{a}^{T}_{k:k+J}\sim\mathcal{N}(\mathbf{0},\textbf{I})$

6:for $t=T,\dotsc,1$ do

7: $E_{t}=\mathbf{TE}(t)$

8: $E_{\mathbf{s}}=\mathbf{PE}(\mathbf{ResNet}(\mathbf{s}_{k-K:k}))$

9:
$E_{\mathbf{a}}=\mathbf{PE}(\mathbf{E_{\mathbf{a}}}(\mathbf{a}^{t}_{k:k+J}))$

10: $\mathbf{a}^{t-1}_{k:k+J}=$

11: $\mathbf{Linear}(\mathbf{D_{m}}(E_{t},E_{\mathbf{s}},E_{\mathbf{a}}))$

12:end for

13:return $\mathbf{a}_{k:k+J}$

Algorithm 2 ED-Ma DDP Inference

1:Observation $\mathbf{s}_{k-K:k}$, Mamba Encoder $\mathbf{E_{m}}$, Mamba
Decoder $\mathbf{D_{m}}$, Position Embedding $\mathbf{PE}$, Learnable action
variables $\hat{\mathbf{a}}_{k:k+J}$, learnable state variables
$\hat{\mathbf{s}}_{k-K:k}$, learnable time variables $\hat{t}$

2:$\mathbf{a}^{T}_{k:k+J}\sim\mathcal{N}(\mathbf{0},\textbf{I})$

3:for $t=T,\dotsc,1$ do

4: $E_{t}=\mathbf{TE}(t)$

5: $E_{s}=\mathbf{PE}(\mathbf{ResNet}(\mathbf{s}_{k-K:k}))$

6: $E_{a}=\mathbf{PE}(\mathbf{Linear}(\mathbf{a}^{t}_{k:k+J}))$

7: $E_{s}=\mathbf{E_{m}}(\textrm{cat}(E_{t},E_{s},\hat{\mathbf{a}}_{k:k+J}))$
$\triangleright$ Sequence alignment

8:
$E_{a}=\mathbf{D_{m}}(\textrm{cat}(\hat{t},\hat{\mathbf{s}}_{k-K:k},E_{a}))$
$\triangleright$ First decoder layer

9: $\mathbf{a}^{t-1}_{k:k+J}=\mathbf{D_{m}}(E_{s}+E_{a})$ $\triangleright$
Remaining decoder layers

10:end for

11:return $\mathbf{a}_{k:k+J}$

###  4.1 Decoder-Only Mamba

Similar to the Decoder-only transformer, we use a Mamba block to process the
inputs. An overview of the Decoder-only Mamba architecture is shown in Figure
1. The Decoder-only Mamba for DDPs is designed to learn a denoising function
$\bm{\epsilon}_{\theta}$ that takes a sequence of observations
$\mathbf{s}_{k-K:k}$, noisy actions $\mathbf{a}^{t}_{k:k+J}$ and diffusion
step $t$ to generate a less noisy sequence of actions
$\mathbf{a}^{t-1}_{k:k+J}$. The diffusion step is encoded using time-embedding
$\mathbf{TE}$. The observations are encoded using a $\mathbf{ResNet}$-18, with
shared weights across images from different time steps. An action encoder
$\mathbf{ENC}_{\mathbf{a}}$ is used to tokenize the noisy action inputs.
Additionally, position embeddings $\mathbf{PE}$ are applied to both the
observations and actions. The time embedding, state embedding, and action
embedding will then be inputted into the Mamba Decoder $\mathbf{Dec}_{m}$. The
Mamba Decoder is implemented by stacking multiple mamba blocks with residual
connections and Layer Normalization. The full inference routine is shown in
Algorithm 1.

###  4.2 Encoder-Decoder Mamba

Compared to the Decoder-only Transformer containing only self-attention
mechanisms, the Encoder-Decoder Transformer with cross-attention is a more
flexible and effective design to process complex input-output relationships,
especially for cases where the input and output sequences differ in structure.
However, Mamba does not provide such a mechanism to support the encoder-
decoder structure since the target and the source share the same sequence
length. Here, we propose a novel approach called Mamba Aggregation which is
used to design the encoder-decoder version of Mamba. The visualization can be
found in Figure 2. The Mamba Encoder $\mathbf{Enc_{m}}$ is used to process the
time embedding and state embedding, and the Mamba Decoder $\mathbf{Dec_{m}}$
is used to process the noise embedding. Since the inputs to the
$\mathbf{E_{m}}$ and $\mathbf{D_{m}}$ have different lengths of sequences, we
propose to add learnable variables to complement each sequence. The method is
outlined in Alg 2.

##  5 Experiments

We conducted extensive experiments on both simulation benchmarks and real
robot setups to verify the effectiveness of MaIL in imitation learning. Our
investigation focuses on the following key questions:

Q1) Can MaIL achieve comparable or superior performance to Transformers?
Q2) Can MaIL utilize multi-modal inputs, such as language instructions?
Q3) How effectively do MaIL handle sequential information in observations?

###  5.1 Baselines

In this work, we mainly aim to explore the potential of Mamba in visual
imitation learning compared to Transformer structures. Therefore, our
experiments contain four architectures: Decoder-only Transformer (D-Tr),
Encoder-Decoder Transformer (ED-Tr), Decoder-only Mamba (D-Ma), Encoder-
Decoder Mamba (ED-Ma). For a fair comparison, we use ResNet18 to encode visual
inputs for each method. For tasks that use language instructions, we use the
pre-trained CLIP model [43] to get the corresponding language embedding which
is used in training and inference for all methods. Based on the above setting,
we implement the following imitation learning policies:

Behavior Cloning (BC) We implement a vanilla behavior cloning policy trained
with MSE loss with both Transformer and Mamba structures.

Denoising Diffusion Policies (DDP) Based on the same structures in BC, we
further implement a diffusion policy using a discrete denoising process [44].
We use 16 diffusion time steps for training and sampling for each
architecture.

###  5.2 Simulation Evaluation

LIBERO [14]: The evaluation is conducted using the LIBERO benchmark, which
encompasses five distinct task suites: LIBERO-Spatial, LIBERO-Object, LIBERO-
Goal, LIBERO-Long, and LIBERO-90. Each task suite comprises 10 tasks along
with 50 human demonstrations except for LIBERO-90 which contains 90 tasks with
50 demonstrations. Each task suite is designed to test different aspects of
robotic learning and manipulation capabilities. The task visualizations are in
Figure 3. More details can be found in Appendix C.

Figure 3: LIBERO benchmark suites with total 130 tasks in five different
scenes.

Evaluation Protocol We compared each method across five LIBERO task suites
separately. Instead of using full demonstrations, we utilized only 20% of the
demonstrations for each sub-task, amounting to 100 trajectories per task
suite, except for LIBERO-90, which contains 900 trajectories. We tune the
hyper-parameters for both Transformer and Mamba, making sure they have a
similar amount of parameters. All models were trained for 50 epochs, and we
used the last checkpoint for evaluation. Following the official LIBERO
benchmark settings, we performed 20 rollouts for each sub-task, totaling 200
evaluations per task suite, except for LIBERO-90, which includes 1800
evaluations. We report the average success rate for each task suite over 3
seeds.

|  | w/o language | w/ language
---|---|---|---
Policy | Backbone | LIBERO-Object | LIBERO-Spatial | LIBERO-Long | LIBERO-Goal | LIBERO-90
BC-H1 | D-Tr | $0.358\scriptstyle{\pm 0.086}$ | $0.257\scriptstyle{\pm 0.030}$ | $0.222\scriptstyle{\pm 0.067}$ | $0.362\scriptstyle{\pm 0.034}$ | $0.311\scriptstyle{\pm 0.055}$
| ED-Tr | $0.480\scriptstyle{\pm 0.027}$ | $0.333\scriptstyle{\pm 0.029}$ | $0.265\scriptstyle{\pm 0.015}$ | $0.325\scriptstyle{\pm 0.020}$ | $0.268\scriptstyle{\pm 0.050}$
| D-Ma | $0.550\scriptstyle{\pm 0.062}$ | $\mathbf{0.410\scriptstyle{\pm 0.058}}$ | $\mathbf{0.287\scriptstyle{\pm 0.037}}$ | $\mathbf{0.482\scriptstyle{\pm 0.066}}$ | $\mathbf{0.559\scriptstyle{\pm 0.021}}$
| ED-Ma | $\mathbf{0.618\scriptstyle{\pm 0.058}}$ | $0.352\scriptstyle{\pm 0.069}$ | $0.275\scriptstyle{\pm 0.023}$ | $0.473\scriptstyle{\pm 0.031}$ | $0.514\scriptstyle{\pm 0.008}$
BC-H5 | D-Tr | $0.489\scriptstyle{\pm 0.170}$ | $0.286\scriptstyle{\pm 0.040}$ | $0.243\scriptstyle{\pm 0.052}$ | $0.315\scriptstyle{\pm 0.036}$ | $0.284\scriptstyle{\pm 0.012}$
| ED-Tr | $0.417\scriptstyle{\pm 0.300}$ | $0.287\scriptstyle{\pm 0.095}$ | $0.248\scriptstyle{\pm 0.021}$ | $0.338\scriptstyle{\pm 0.024}$ | $0.300\scriptstyle{\pm 0.012}$
| D-Ma | $\mathbf{0.765\scriptstyle{\pm 0.115}}$ | $0.372\scriptstyle{\pm 0.024}$ | $0.252\scriptstyle{\pm 0.072}$ | $\mathbf{0.43\scriptstyle{\pm 0.029}}$ | $0.529\scriptstyle{\pm 0.094}$
| ED-Ma | $0.743\scriptstyle{\pm 0.020}$ | $\mathbf{0.428\scriptstyle{\pm 0.029}}$ | $\mathbf{0.262\scriptstyle{\pm 0.047}}$ | $0.317\scriptstyle{\pm 0.105}$ | $\mathbf{0.599\scriptstyle{\pm 0.019}}$
DDP-H1 | D-Tr | $0.667\scriptstyle{\pm 0.058}$ | $0.473\scriptstyle{\pm 0.040}$ | $0.330\scriptstyle{\pm 0.029}$ | $0.300\scriptstyle{\pm 0.350}$ | $0.181\scriptstyle{\pm 0.024}$
| ED-Tr | $0.680\scriptstyle{\pm 0.027}$ | $0.420\scriptstyle{\pm 0.022}$ | $0.373\scriptstyle{\pm 0.042}$ | $0.478\scriptstyle{\pm 0.033}$ | $0.399\scriptstyle{\pm 0.030}$
| D-Ma | $\mathbf{0.728\scriptstyle{\pm 0.032}}$ | $\mathbf{0.492\scriptstyle{\pm 0.060}}$ | $0.343\scriptstyle{\pm 0.059}$ | $0.510\scriptstyle{\pm 0.043}$ | $0.405\scriptstyle{\pm 0.139}$
| ED-Ma | $0.713\scriptstyle{\pm 0.008}$ | $0.483\scriptstyle{\pm 0.031}$ | $\mathbf{0.395\scriptstyle{\pm 0.035}}$ | $\mathbf{0.565\scriptstyle{\pm 0.023}}$ | $\mathbf{0.603\scriptstyle{\pm 0.017}}$
DDP-H5 | D-Tr | $0.762\scriptstyle{\pm 0.070}$ | $0.435\scriptstyle{\pm 0.043}$ | $0.353\scriptstyle{\pm 0.033}$ | $0.315\scriptstyle{\pm 0.058}$ | $0.161\scriptstyle{\pm 0.024}$
| ED-Tr | $0.812\scriptstyle{\pm 0.065}$ | $0.415\scriptstyle{\pm 0.097}$ | $0.348\scriptstyle{\pm 0.019}$ | $0.518\scriptstyle{\pm 0.014}$ | $0.437\scriptstyle{\pm 0.046}$
| D-Ma | $\mathbf{0.815\scriptstyle{\pm 0.004}}$ | $\mathbf{0.538\scriptstyle{\pm 0.044}}$ | $0.387\scriptstyle{\pm 0.040}$ | $0.388\scriptstyle{\pm 0.027}$ | $0.352\scriptstyle{\pm 0.096}$
| ED-Ma | $0.778\scriptstyle{\pm 0.005}$ | $0.513\scriptstyle{\pm 0.024}$ | $\mathbf{0.417\scriptstyle{\pm 0.023}}$ | $\mathbf{0.563\scriptstyle{\pm 0.070}}$ | $\mathbf{0.455\scriptstyle{\pm 0.081}}$

Table 1: Performance on LIBERO benchmark, where ”w/o language” indicates that
we do not use language instructions and ”w/ language” means we use language
tokens generated from a pre-trained CLIP model, H1 and H5 refer to using
current state and 5 steps historical states respectively.

Main Results. We report the main results in Table 1. Our Mamba-based
architectures, D-Ma and ED-Ma, significantly outperform Transformer-based
methods across all LIBERO task suites based on the BC policy. Specifically,
Mamba-based models achieve nearly a 30% improvement in success rate in LIBERO-
Object and LIBERO-90. When using the DDP policy, our models consistently
surpass the Transformer baselines, with performance improvements exceeding 5%
in most tasks. These results confirm Q1, demonstrating that MaIL achieves
superior performance compared to Transformers. To address Q2, we compare MaIL
with Transformers on LIBERO-Goal and LIBERO-90, using additional language
embeddings as inputs. We observe significant improvements with Mamba-based
methods in these tasks, indicating that MaIL effectively leverages multi-modal
inputs. Given that most recent visual imitation learning works use historical
observations as inputs, we evaluated the methods with 1 and 5 historical
observations. We found that historical information does not always enhance
performance. Only in LIBERO-Object do the H5 models outperform the H1 models,
while in other tasks, H5 models achieve similar or worse results. Mamba-based
H5 models again perform consistently better than Transformer-based models,
which indicates that MaIL is able to capture sequential observation features
effectively, answering Q3.

Ablation on Observation Occlusions In order to further understand the
sequential learning ability of Transformer and Mamba, we randomly mask out
areas of images and test the model’s performance drop. The results are
reported in Figure 4. While for zero occlusions, the Transformer architectures
can perform en par with Mamba, adding occlusions degenerates the performance
of transformers more quickly, indicating that Mamba can better extract the
important information out of the history sequence.

Ablation on Dataset Size Given that MaIL performs well with only 20% of the
demonstrations, we are interested in evaluating its scalability with
increasing dataset size. We compared Mamba-based models with Transformer
models using the BC policy on the LIBERO-Spatial task. The results are
presented in Figure 4. It is evident that Mamba-based models significantly
outperform Transformers when data is scarce and perform comparably as the
dataset size increases.

(a) DDP-H5: object

(b) DDP-H5: spatial

(c) BC-H1: spatial

(d) BC-H5: spatial

Figure 4: Ablation study for observation occlusions and dataset size. (a) and
(b) study the influence of observation occlusions on Libero-Object and LIBERO-
Spatial respectively. (c) and (d) show the influence of the amount of data on
LIBERO-spatial.

###  5.3 Real Robot Evaluation

Figure 5: Real robot setup. We use visual inputs from two cameras as
observations for policies.

We designed three challenging tasks based on the 7DoF Franka Panda Robot,
utilizing visual inputs for the model. Two cameras, positioned at different
angles in front of the robot, provide the visual data. One image is cropped
and resized to $(128,256,3)$, while the other is resized to $(256,256,3)$. The
entire setup is visualized in Figure 5. These images are stacked at each
timestep to form the observations. We excluded the robot states from the
inputs, as previous studies have reported that including them can lead to poor
performance [7]. The action space is 8-dimensional, encompassing the joint
positions and the gripper state. The tasks settings detailed below and shown
in Figure 6-8. The corresponding results are presented in Table 2-4.

Pick-Place We first evaluate our policies using a pick place task. This task
contains two individual sub-tasks: place the banana on the plate or place the
carrot in the bowl. We collect 30 demonstrations for each individual tasks.
During the evaluation, the objects are initialized with similar positions but
different orientations.

Two-stage Pick-Place. As an enhanced version, this task concatenates the two
stages as one manipulation sequence: place both the banana and carrot
sequentially to their target areas. Due to the longer motion horizon and more
stochasticity, the task complexity significantly increases. The dataset
contains 100 demonstrations with objects randomly initialized.

Inserting. The robot is required to insert a blue block into a LEGO box. We
split this task into two stages: 1) Pick up the block and move it to the box
2) Insert the block inside the box. Since this task demands precision control
which makes it challenging, we fix the LEGO box position and only randomly
initialize the state of the block.

CupStacking. The robot is required to stack three cups of different sizes.
This task contains two stages: 1) Pick up the green cup and place it in the
blue cup and 2) Pick up the yellow cup and place it in the green cup. Three
cups are aligned almost at the same line at the beginning, then we randomly
add variances to each of them. During the evaluation, we use the same
initialized positions to evaluate the models and we record the success rate at
each stage.

Figure 6: PickPlacing (front camera view)

DDP-H1 |  sPick Places | Two-stage Pick-Place
---|---|---
|  sBananas |  s Carrots | Stage-1 | Stage-2
ED-Tr | $0.45$ | $0.25$ | $0.35$ | $\mathbf{0.2}$
ED-Ma ss | $\mathbf{0.55}$ | $\mathbf{0.70}$ | $\mathbf{0.45}$ | $0.1$

Table 2: Pick-Place Success rates

Figure 7: Inserting (right camera view)

DDP-H1 |  ssssssss Stage-1 ssssssss |  sssss Stage-2 ssss
---|---|---
ED-Tr | $0.4$ | $\mathbf{0.3}$
ED-Ma | $\mathbf{0.55}$ | $0.2$

Table 3: Inserting success rates

Figure 8: CupStacking (front camera view)

DDP-H1 |  ssssssss Stack-1 ssssssss |  sssssStack-2ssss
---|---|---
ED-Tr | $0.6$ | $0.4$
ED-Ma ss | $\mathbf{0.8}$ | $\mathbf{0.55}$

Table 4: CupStacking success rates

We compared the ED-Tr with ED-Ma using DDP-H1 model. We trained each method
for 100 epochs (convergent) and evaluated the model using the final
checkpoint. For each task, we performed 20 rollouts with different initial
states for the objects. To ensure a fair comparison, we used the same initial
states for both Transformer and Mamba evaluations. From the results, Mamba-
based methods achieve comparable results with Transformer models.

##  6 Limitations

While MaIL demonstrates excellent performance with smaller dataset sizes, its
advantage becomes less pronounced as the dataset scales up. When trained on
larger datasets, MaIL achieves results comparable to, but not surpassing,
those of Transformer models. Additionally, Mamba is designed to be fast and
efficient for large-scale sequences. However, in the context of imitation
learning policies where sequences are relatively short, the inference time for
Transformers is similar to that of Mamba. This reduces the performance
efficiency advantage of Mamba in these scenarios.

##  7 Conclusion

In conclusion, this work presents MaIL, a novel imitation learning (IL) policy
architecture that bridges the gap between efficiency and performance in
handling sequences of observations. By leveraging the strengths of state space
models and rigorously improving upon them, MaIL offers a competitive
alternative to the traditionally large and complex Transformer-based policies.
The introduction of Mamba in an encoder-decoder structure enhances its
versatility, making it suitable for standalone use as well as integration into
advanced architectures like diffusion processes. Extensive evaluations on the
LIBERO IL benchmark and real robot experiments demonstrate that MaIL not only
matches but also surpasses the performance of existing baselines, establishing
it as a promising approach for IL tasks.

## References

  * Osa et al. [2018] T. Osa, J. Pajarinen, G. Neumann, J. A. Bagnell, P. Abbeel, J. Peters, et al.  An algorithmic perspective on imitation learning.  _Foundations and Trends® in Robotics_ , 7(1-2):1–179, 2018.
  * Jia et al. [2024] X. Jia, D. Blessing, X. Jiang, M. Reuss, A. Donat, R. Lioutikov, and G. Neumann.  Towards diverse behaviors: A benchmark for imitation learning with human demonstrations.  _arXiv preprint arXiv:2402.14606_ , 2024.
  * Zhao et al. [2023] T. Z. Zhao, V. Kumar, S. Levine, and C. Finn.  Learning fine-grained bimanual manipulation with low-cost hardware.  _arXiv preprint arXiv:2304.13705_ , 2023.
  * Florence et al. [2022] P. Florence, C. Lynch, A. Zeng, O. A. Ramirez, A. Wahid, L. Downs, A. Wong, J. Lee, I. Mordatch, and J. Tompson.  Implicit behavioral cloning.  In _Conference on Robot Learning_ , pages 158–168. PMLR, 2022.
  * Chi et al. [2023] C. Chi, S. Feng, Y. Du, Z. Xu, E. Cousineau, B. Burchfiel, and S. Song.  Diffusion policy: Visuomotor policy learning via action diffusion.  In _Proceedings of Robotics: Science and Systems (RSS)_ , 2023.
  * Reuss et al. [2023] M. Reuss, M. Li, X. Jia, and R. Lioutikov.  Goal conditioned imitation learning using score-based diffusion policies.  In _Robotics: Science and Systems_ , 2023.
  * Mandlekar et al. [2021] A. Mandlekar, D. Xu, J. Wong, S. Nasiriany, C. Wang, R. Kulkarni, L. Fei-Fei, S. Savarese, Y. Zhu, and R. Martín-Martín.  What matters in learning from offline human demonstrations for robot manipulation.  _arXiv preprint arXiv:2108.03298_ , 2021.
  * Blessing et al. [2024] D. Blessing, O. Celik, X. Jia, M. Reuss, M. Li, R. Lioutikov, and G. Neumann.  Information maximizing curriculum: A curriculum-based approach for learning versatile skills.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Dosovitskiy et al. [2021] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=YicbFdNTTy.
  * Devlin et al. [2019] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.
  * Wen et al. [2023] Q. Wen, T. Zhou, C. Zhang, W. Chen, Z. Ma, J. Yan, and L. Sun.  Transformers in time series: A survey, 2023.
  * Gu et al. [2021] A. Gu, K. Goel, and C. Re.  Efficiently modeling long sequences with structured state spaces.  In _International Conference on Learning Representations_ , 2021.
  * Gu and Dao [2023] A. Gu and T. Dao.  Mamba: Linear-time sequence modeling with selective state spaces.  _arXiv preprint arXiv:2312.00752_ , 2023.
  * Liu et al. [2024] B. Liu, Y. Zhu, C. Gao, Y. Feng, Q. Liu, Y. Zhu, and P. Stone.  Libero: Benchmarking knowledge transfer for lifelong robot learning.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Vaswani et al. [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Achiam et al. [2023] J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_ , 2023.
  * Rumelhart et al. [1986] D. E. Rumelhart, G. E. Hinton, and R. J. Williams.  Learning internal representations by error propagation, parallel distributed processing, explorations in the microstructure of cognition, ed. de rumelhart and j. mcclelland. vol. 1. 1986.  _Biometrika_ , 71:599–607, 1986.
  * Hochreiter and Schmidhuber [1997] S. Hochreiter and J. Schmidhuber.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * Cho et al. [2014] K. Cho, B. Van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio.  Learning phrase representations using rnn encoder-decoder for statistical machine translation.  _arXiv preprint arXiv:1406.1078_ , 2014.
  * Chung et al. [2014] J. Chung, C. Gulcehre, K. Cho, and Y. Bengio.  Empirical evaluation of gated recurrent neural networks on sequence modeling.  _arXiv preprint arXiv:1412.3555_ , 2014.
  * Zhu et al. [2024] L. Zhu, B. Liao, Q. Zhang, X. Wang, W. Liu, and X. Wang.  Vision mamba: Efficient visual representation learning with bidirectional state space model.  _arXiv preprint arXiv:2401.09417_ , 2024.
  * Smith et al. [2023] J. T. Smith, A. Warrington, and S. Linderman.  Simplified state space layers for sequence modeling.  In _The Eleventh International Conference on Learning Representations_ , 2023.
  * Gupta et al. [2022] A. Gupta, A. Gu, and J. Berant.  Diagonal state spaces are as effective as structured state spaces.  _Advances in Neural Information Processing Systems_ , 35:22982–22994, 2022.
  * Hasani et al. [2022] R. Hasani, M. Lechner, T.-H. Wang, M. Chahine, A. Amini, and D. Rus.  Liquid structural state-space models.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Wang et al. [2024] X. Wang, S. Wang, Y. Ding, Y. Li, W. Wu, Y. Rong, W. Kong, J. Huang, S. Li, H. Yang, et al.  State space model for new-generation network alternative to transformers: A survey.  _arXiv preprint arXiv:2404.09516_ , 2024.
  * Pomerleau [1988] D. A. Pomerleau.  Alvinn: An autonomous land vehicle in a neural network.  _Advances in neural information processing systems_ , 1, 1988.
  * Schaal [1996] S. Schaal.  Learning from demonstration.  _Advances in neural information processing systems_ , 9, 1996.
  * Taranovic et al. [2022] A. Taranovic, A. G. Kupcsik, N. Freymuth, and G. Neumann.  Adversarial imitation learning with preferences.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Sun et al. [2017] W. Sun, A. Venkatraman, G. J. Gordon, B. Boots, and J. A. Bagnell.  Deeply aggrevated: Differentiable imitation learning for sequential prediction.  In _International conference on machine learning_ , pages 3309–3318. PMLR, 2017.
  * Wang et al. [2019] B. Wang, E. Adeli, H.-k. Chiu, D.-A. Huang, and J. C. Niebles.  Imitation learning for human pose prediction.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 7124–7133, 2019.
  * Du et al. [2021] C. Du, Z. Wang, A. A. Malcolm, and C. L. Ho.  Imitation learning for autonomous driving based on convolutional and recurrent neural networks.  In _2021 International Conference on High Performance Big Data and Intelligent Systems (HPBD &IS)_, pages 256–260. IEEE, 2021.
  * Lynch et al. [2020] C. Lynch, M. Khansari, T. Xiao, V. Kumar, J. Tompson, S. Levine, and P. Sermanet.  Learning latent plans from play.  In _Conference on robot learning_ , pages 1113–1132. PMLR, 2020.
  * Shafiullah et al. [2022] N. M. Shafiullah, Z. Cui, A. A. Altanzaya, and L. Pinto.  Behavior transformers: Cloning $k$ modes with one stone.  _Advances in neural information processing systems_ , 35:22955–22968, 2022.
  * Cui et al. [2022] Z. J. Cui, Y. Wang, N. M. M. Shafiullah, and L. Pinto.  From play to policy: Conditional behavior generation from uncurated robot data.  _arXiv preprint arXiv:2210.10047_ , 2022.
  * Lee et al. [2024] S. Lee, Y. Wang, H. Etukuru, H. J. Kim, N. M. M. Shafiullah, and L. Pinto.  Behavior generation with latent actions.  _arXiv preprint arXiv:2403.03181_ , 2024.
  * Shridhar et al. [2023] M. Shridhar, L. Manuelli, and D. Fox.  Perceiver-actor: A multi-task transformer for robotic manipulation.  In _Conference on Robot Learning_ , pages 785–799. PMLR, 2023.
  * Wu et al. [2023] H. Wu, Y. Jing, C. Cheang, G. Chen, J. Xu, X. Li, M. Liu, H. Li, and T. Kong.  Unleashing large-scale video generative pre-training for visual robot manipulation.  _arXiv preprint arXiv:2312.13139_ , 2023.
  * Shah et al. [2023] R. Shah, R. Martín-Martín, and Y. Zhu.  Mutex: Learning unified policies from multimodal task specifications.  _arXiv preprint arXiv:2309.14320_ , 2023.
  * Reuss et al. [2024] M. Reuss, Ö. E. Yağmurlu, F. Wenzel, and R. Lioutikov.  Multimodal diffusion transformer: Learning versatile behavior from multimodal goals.  In _First Workshop on Vision-Language Models for Navigation and Manipulation at ICRA 2024_ , 2024.  URL https://openreview.net/forum?id=Pt6fLfXMRW.
  * Scheikl et al. [2024] P. M. Scheikl, N. Schreiber, C. Haas, N. Freymuth, G. Neumann, R. Lioutikov, and F. Mathis-Ullrich.  Movement primitive diffusion: Learning gentle robotic manipulation of deformable objects.  _IEEE Robotics and Automation Letters_ , 2024.
  * Sridhar et al. [2023] A. Sridhar, D. Shah, C. Glossop, and S. Levine.  Nomad: Goal masked diffusion policies for navigation and exploration.  _arXiv preprint arXiv:2310.07896_ , 2023.
  * Zhang et al. [2023] M. Zhang, K. K. Saab, M. Poli, T. Dao, K. Goel, and C. Ré.  Effectively modeling time series with simple discrete state spaces, 2023.
  * Radford et al. [2021] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al.  Learning transferable visual models from natural language supervision.  In _International conference on machine learning_ , pages 8748–8763. PMLR, 2021.
  * Ho et al. [2020] J. Ho, A. Jain, and P. Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Lin et al. [2021] Y. Lin, A. S. Wang, G. Sutanto, A. Rai, and F. Meier.  Polymetis.  https://facebookresearch.github.io/fairo/polymetis/, 2021.

##  Appendix A State Space Models

Inspired by the Attention mechanism [15] in Transformers, Mamba [13] improves
upon the Structured State-Space Sequence (S4) Model by using a selective scan
operator, to propagate or forget information over context, allowing it to
filter out relevant features. Specifically, the S4 model is inspired by a
continuous system, $x(l)\in\mathbb{R^{\mathtt{D}}}\mapsto
y(l)\in\mathbb{R^{\mathtt{D}}}$, that maps $x(l)$ into a hidden state
$h(l)\in\mathbb{R}^{\mathtt{N}}$, which then can be projected onto the output
$y(t)$, that is,

$y(l)=\mathbf{C}h(l)\quad\text{with}\quad
h^{\prime}(l)=\mathbf{A}h(l)+\mathbf{B}x(l),$ (6)

where $\mathbf{A}\in\mathbb{R}^{\mathtt{N}\times\mathtt{N}}$ are the evolution
parameters, $\mathbf{B}\in\mathbb{R}^{\mathtt{N}\times\mathtt{D}}$ and
$\mathbf{C}\in\mathbb{R}^{\mathtt{D}\times\mathtt{N}}$ are the projection
parameters. S4 transforms the parameters $(\mathbf{A},\mathbf{B})$ to discrete
parameters $(\mathbf{\overline{A}},\mathbf{\overline{B}})$ via

$\displaystyle\mathbf{\overline{A}}$
$\displaystyle=\exp{(\mathbf{\Delta}\mathbf{A})}\quad\text{and}\quad\mathbf{\overline{B}}$
$\displaystyle=(\mathbf{\Delta}\mathbf{A})^{-1}(\exp{(\mathbf{\Delta}\mathbf{A})}-\mathbf{I})\cdot\mathbf{\Delta}\mathbf{B},$
(7)

where $\mathbf{\Delta}$ is the step size. The model then can be computed as
either linear recurrence, that is,

$\displaystyle\quad y_{t}$
$\displaystyle=\mathbf{C}h_{t},\quad\text{with}\quad h_{t}$
$\displaystyle=\mathbf{\overline{A}}h_{t-1}+\mathbf{\overline{B}}x_{t},$ (8)

or through a global convolution $*$ given by

$\displaystyle\mathbf{y}$
$\displaystyle=\mathbf{x}*\mathbf{\overline{K}},\quad\text{with}\quad\mathbf{\overline{K}}$
$\displaystyle=(\mathbf{C}\mathbf{\overline{B}},\mathbf{C}\mathbf{\overline{A}}\mathbf{\overline{B}},\dots,\mathbf{C}\mathbf{\overline{A}}^{k}\mathbf{\overline{B}}).$
(9)

##  Appendix B Transformer Architecture

Here we depict two transformer-based architectures in diffusion policy: a
decoder-only model (Figure 9) and an encoder-decoder model (Figure 10). Both
architectures leverage the strengths of transformer models to effectively
handle sequential data and capture long-range dependencies.

Figure 9: Decoder-only learning block. This architecture integrates ResNet-18
for state encoding and an action encoder for actions with horizon $J$, with
both components feeding into a self-attention mechanism. Positional encoding
(PE) and time encoding (TE) enhance the inputs. Ultimately, the output of
self-attention is fed to a linear output layer to predict future actions.
Figure 10: Encoder-decoder learning block.This figure illustrates the
architecture of an encoder-decoder transformer block designed for policy
learning. In the encoder, states are encoded using ResNet-18, enhanced by time
encoding (TE) and positional encoding (PE), and processed through self-
attention. The decoder then utilizes self-attention on the encoded actions and
employs cross-attention to integrate the encoded states from the encoder.
Ultimately, the output of cross-attention is fed to a linear output layer to
predict future actions.

##  Appendix C Additional Task Details

###  C.1 LIBERO benchmark

  * •

LIBERO-Spatial contains various objects and different spatial relationships
among those objects within a consistent layout. This suite challenges the
robot’s spatial understanding and its ability to navigate and manipulate
objects based on their spatial configurations.

  * •

LIBERO-Object suite has tasks involving different objects on the same layout,
requiring the robot to pick-place a unique object at a time. It emphasizes the
policy’s ability to accurately recognize various types of objects.

  * •

LIBERO-Goal suite maintains the same objects and fixed spatial relationships
but varies the task goals. This suite challenges the robot’s ability to
understand and achieve different objectives despite the consistent arrangement
of objects.

  * •

LIBERO-Long suite focuses on long-horizon tasks, evaluating the robot’s
ability to perform and manage extended sequences of actions over a prolonged
period.

  * •

LIBERO-90 contains 90 short-horizon tasks with significant diversity. This
suite includes a wide variety of object categories, layouts, and goals,
providing a comprehensive evaluation of the robot’s adaptability and
versatility in handling diverse scenarios.

###  C.2 Real Robot Tasks

####  C.2.1 CupStacking

In the CupStacking task, there are three cups with differing sizes which the
robot needs to stack inside one another. The large cup is on the left, the
small cup is on the right and the middle-sized cup is in the middle. In the
training data, only one of the two modes is covered: the middle-sized cup is
put into the large cup first, and the small cup is stacked afterwards. The
horizontal positions of the cups are fixed, but their vertical positions can
vary.

####  C.2.2 Inserting

In the Inserting task, there is a LEGO box with a hole in the middle and
another block which can fit in the hole of the LEGO box. The goal in this task
is to grab the block and insert it in the hole of the LEGO box. The position
and orientation of the LEGO box are fixed, but the position and orientation of
the block can vary inside a specific area on the table.

####  C.2.3 PickPlacing

In the PickPlacing task, there are two different bowls, a banana and a carrot.
The goal is to pick the banana and the carrot and place them inside the bowls.
The banana is always placed in the left bowl, and the carrot is placed in the
right bowl. The positions and orientations of the bowls are fixed, but the
positions and orientations of the banana and the carrot can vary inside a
specific area on the table.

##  Appendix D Model Details

###  D.1 Parameter Comparison

For all Mamba-based policies, we fix the expansion factor at 2. The parameters
exhibit slight variance when using the same number of layers due to
differences in the choice of hidden state values and convolution widths.
Therefore, at the same layer depth, we use the average parameter values to
account for these variations. We list the comparison of parameters for all
models in Table 5.

We also evaluate the inference time on a local PC equipped with an RTX 2060
GPU, using a batch size of 32 to ensure all models are assessed under the same
conditions in Table 5.

Table 5: Parameters comparison

Method | Layers | Params | Inference time | Inference time 5history
---|---|---|---|---
BC-Dec | 6 | 23.6M | 0.03043 | 0.05906
| 8 | 24.0M | 0.03305 | 0.06306
BC-EncDec | 4 + 4 | 24.2M | 0.03230 | 0.06368
| 4 + 6 | 24.7M | 0.03505 | 0.06563
BC-Mamba | 10 | 23.6M | 0.02603 | 0.05340
| 12 | 23.8M | 0.02797 | 0.05477
| 16 | 24.3M | 0.02855 | 0.05614
BC-MambaDec | 6 + 6 | 23.8M | 0.02752 | 0.05465
| 6 + 8 | 24.0M | 0.02833 | 0.05733
| 6 + 10 | 24.2M | 0.02974 | 0.05862
| 8 + 10 | 24.4M | 0.03074 | 0.06007
| 10 + 10 | 24.6M | 0.03190 | 0.06112
DDPM-Dec | 6 | 23.7M | 0.10300 | 0.14074
| 8 | 24.1M | 0.12882 | 0.16428
DDPM-EncDec | 4 + 4 | 24.3M | 0.14443 | 0.17995
| 4 + 6 | 24.8M | 0.17244 | 0.21022
DDPM-Mamba | 10 | 23.7M | 0.11185 | 0.14384
| 12 | 23.9M | 0.12616 | 0.16145
| 16 | 24.2M | 0.15699 | 0.18236
DDPM-MambaCross | 6 + 6 | 23.8M | 0.12653 | 0.15984
| 6 + 8 | 24.0M | 0.14025 | 0.17350
| 6 + 10 | 24.3M | 0.15977 | 0.19398
| 8 + 10 | 24.5M | 0.16883 | 0.20553
| 10 + 10 | 24.7M | 0.18332 | 0.22571

###  D.2 Training Details

We list the training hyperparameters for the Transformer-based and Mamba-based
policies in Table 6. To ensure a fair comparison, we tune the hyperparameters
of both policies at the same level.

The policies are trained using the human expert demonstrations provided by
LIBERO, where we only take 10 demonstrations for each task in main experiment.

All models are trained on a cluster equipped with 4 A100 GPUs, with a batch
size of 256, over 50 epochs using 3 different seeds. Finally, we calculate the
average success rate across these 3 seeds.

Table 6: Hyperparameter in Simulated Experiments

Methods / Parameters | libero_object | libero_spatial | libero_long | libero_goal | libero_90
---|---|---|---|---|---
history = 5 |  |  |  |  |
DDPM-Mamba |  |  |  |  |
Number of Mamba Layer | 10 | 16 | 16 | 12 | 16
Hidden State Dimension | 16 | 8 | 16 | 8 | 16
Convolution Width | 4 | 4 | 4 | 2 | 4
Expand Factor | 2 | 2 | 2 | 2 | 2
DDPM-MambaCross |  |  |  |  |
Number of Encoder Layer | 8 | 10 | 8 | 10 | 8
Number of Decoder Layer | 10 | 10 | 10 | 10 | 10
Hidden State of Encoder | 8 | 8 | 8 | 8 | 8
Hidden State of Decoder | 8 | 16 | 8 | 16 | 8
Encoder Convolution Width | 4 | 4 | 2 | 2 | 2
Decoder Convolution Width | 2 | 2 | 2 | 2 | 2
Expand Factor | 2 | 2 | 2 | 2 | 2
DDPM-Decoder |  |  |  |  |
Number of Attention Block | 8 | 8 | 8 | 8 | 8
Number of Attention Head | 4 | 4 | 4 | 4 | 4
DDPM-EncDec |  |  |  |  |
Number of Encoder Block | 4 | 4 | 4 | 4 | 4
Number of Decoder Block | 6 | 6 | 6 | 6 | 6
Number of Encoder Head | 4 | 4 | 4 | 4 | 4
Number of Decoder Head | 4 | 4 | 4 | 4 | 4
history = 1 |  |  |  |  |
DDPM-Mamba |  |  |  |  |
Number of Mamba Layer | 10 | 16 | 16 | 12 | 16
Hidden State Dimension | 16 | 8 | 16 | 8 | 16
Convolution Width | 4 | 4 | 4 | 2 | 4
Expand Factor | 2 | 2 | 2 | 2 | 2
DDPM-MambaCross |  |  |  |  |
Number of Encoder Layer | 10 | 10 | 10 | 10 | 8
Number of Decoder Layer | 8 | 10 | 10 | 8 | 10
Hidden State of Encoder | 16 | 16 | 8 | 16 | 8
Hidden State of Decoder | 8 | 8 | 16 | 16 | 8
Encoder Convolution Width | 2 | 4 | 2 | 2 | 2
Decoder Convolution Width | 4 | 2 | 4 | 4 | 2
Expand Factor | 2 | 2 | 2 | 2 | 2
DDPM-Decoder |  |  |  |  |
Number of Attention Block | 8 | 8 | 8 | 8 | 8
Number of Attention Head | 4 | 4 | 4 | 4 | 4
DDPM-EncDec |  |  |  |  |
Number of Encoder Block | 4 | 4 | 4 | 4 | 4
Number of Decoder Block | 6 | 6 | 6 | 6 | 6
Number of Encoder Head | 4 | 4 | 4 | 4 | 4
Number of Decoder Head | 4 | 4 | 4 | 4 | 4

###  D.3 Ablation on Observation Occlusions

For the image-based model, recent observations are not significantly different
from the current observation. To evaluate the influence of historical data and
compel the model to focus more on past observations, we employ occlusion
techniques on the observations. This involves randomly masking portions of the
images like the Figure 11. The occlusion rate determines the extent to which
the images are masked.

Figure 11:  Visualization of Data Augmentation

##  Appendix E Real Robot Details

###  E.1 Setup

The current setup consists of two robots and two computers. One of the
computers, the robot PC, runs a real-time OS to control the robots reliably
using Polymetis (E.2), which is a real-time Pytorch controller manager for
robots. The robot PC uses a PD controller with the desired joint position as a
configurable parameter. The robot PC hosts a server from which the desired
joint position and the gripper of the robot can be configured. The other PC,
the workstation, runs the models and sends the desired joint position and
gripper state to the robot PC at each time step. The workstation also has
access to two cameras, which are used to capture images as input for the
models.

###  E.2 Polymetis

Polymetis [45] is a software framework designed to improve the control and
interaction of robotic systems. It connects high-level decision-making with
low-level hardware control, allowing robots to perform complex tasks like
assembly and inspection accurately. Key features include real-time motion
planning, adaptive control, and sensory feedback. Polymetis supports various
robotic platforms and sensors, making it useful for both research and
industry. Its simplicity, detailed documentation, and strong API make
Polymetis a valuable tool for enhancing robotic capabilities and intelligence.

###  E.3 Data Collection

Teleoperation is used to collect data for all the real robot tasks, where the
leader robot is controlled by a human and the follower robot follows the
leader robot as shown in Figure 12. The objects are put in front of the
follower robot and the cameras do not see the leader robot or the human. The
current joint state of the leader robot is sent to the follower robot as the
desired joint state. The state of the gripper is considered to be binary,
either closed or open. A threshold is set for the gripper of the leader robot;
if the current width is below the threshold, the gripper of the follower robot
closes, otherwise it opens.

Figure 12: Data collection scene for real robot

###  E.4 Evaluation

For the evaluation, using the output of the model sometimes activate the
security mechanism of the robot because it violates a certain constraint. To
solve this issue, a trajectory is generated between the current joint position
and the output of the model. The points of this trajectory are then given to
the robot at each time step instead of the raw output of the model. The length
of this trajectory varies depending on how far away the output of the model is
from the current robot state.