MaIL: Improving Imitation Learning with Selective State Space Models
Authors: Xiaogang Jia
Qian Wang
Atalay Donat
Bowen Xing
Ge Li
Hongyi Zhou
Onur Celik
Denis Blessing
Rudolf Lioutikov
Gerhard Neumann
Venue: CoRL 2024 (2024)
Keywords: Imitation Learning
Sequence Models
Denoising Diffusion Policies

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.

MalL: Improving Imitation Learning with Selective State Space Models
Xiaogang Jia Qian Wang"i Atalay Donat Bowen Xing Ge Lit
Hongyi Zhou' Onur Celik Denis Blessing' Rudolf Lioutikov Gerhard Neumann' I Autonomous Learning Robots, Karlsruhe Institute of Technology Intuitive Robots Lab, Karlsruhe Institute of Technology
Abstract:
This work presents Mamba Imitation Learning (MalL), a novel imitation learn- ing (IL) architecture that provides an alternative to otate-of-the-art (SoTA) Transformer-based policies. Mall. leverages Mamba, a state-space model de- signed to selectively focus on key features of the data. While Transformers are highly effective in data-rich environments due to their dense attention mecha- nisms, they struggle with smaller datasets, often leading overfitting or sub- optimal representation learning. In contrast, Mamba's architecture enhances rep- resentation learning efficiency by focusing on key features and reducing model complexity. This approach mitigates overfitting and enhances generalization, even when working with limited data. Extensive evaluations on the LIBERO benchmark demonstrate that MalL consistently outperforms Transformers on all LIBERO tasks with limited data and matches their performance when the full dataset is available. Additionally, MalL's effectiveness validated through its superior performance in three real robot experiments. Our code is available at https://github com ALRhub/MaIL.
Keywords: Imitation Learning, Sequence Models, Denoising Diffusion Policies
1 Introduction
Imitation learning (IL) [1] from human data has shown remarkable success in acquiring robot poli- cies 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 Trans- former 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 considered state-of-the-art for processing sequential data. Here, current approaches either use an decoder-only structure [5], or decoder-encoder architecture [6]. Which of these architectures excels often task dependent. The performance of transformers usually comes with large models that are difficult to train, especially domains where data is scarce. An alternative concept for handling a sequence of observations are state space models [12]. These models assume a linear rela- tionship between observations (embeddings) and are usually computationally more efficient. Recent approaches such Mamba [13], 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 MalL, a novel imitation learning policy architecture that uses Mamba as a backbone. MalL can be used standalone policy, or as part of more advanced processes such as
Equul contribution, correspondence to jia266163 gmail.com
8th Conference Robot Learning (CoRL. 2024), Munich, Germany.

---

diffuser in the diffusion process. We implement MalL in two variants. In the decoder-only variant, Mall. 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 's formalism, an encoder-decoder variant is not straightforwardly implementable due to possibly varying input and output structures. We show how we can extend MalL 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-ari LIBERO [ [14] II. benchmarks with and without history input show that MalL achieves improved performance compared to similarly complex transformer-base IL. policies. MalL'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 re- quires 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 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 trans- formers. Mall. exploits the selective state space models present in Mamba and proposes 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 ar- chitectures, including restricted representation power with long-sequence modeling and slow train- ing times due 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 effi- ciency 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, 39J. Due to their strong generalization ability and rich representation power to capture multimodal action distribu- tions, they have become the SoTA in the imitation learning domain. Many of these models also utilize Transformers policies, leveraging their rich representation capabilities [6, 41, 39J. Instead of relying transformer backbone, MalL proposes novel architecture that is based on Mamba [13]. In the evaluations, we show the benefits of MalL, especially in scenarios with sequences of observation inputs over state-of-the-ar methods.
3 Preliminaries
Here, we briefly review the basics of the Mamba model and explain the policy representations sidered in this work.

---

Time Encoding
Linear Projection
Mamba Block
Actoos
M
Nx
SSM
Mamba
Conv
PE
Linear Projection
Linear Projection
TE
RewNett-8
Action Encoder
Noised Actions
Figure I: D-Ma: Mamba denoising architecture integrates ResNet-18 for state encoding and an ac- tion encoder for action encoding. The state sequence has a length of K, while the action sequence diffusion step has length of J. Before feeding the inputs into the Mamba module, positional encoding (PE) and time encoding (TE) enhance the imputs, where the 8ę and aş share the same po- sitional encoding. The mamba module has Nx mamba blocks, with a detailed structure [13] shown on the left. The outputs from the mamba module processed by a linear output layer, resulting the one-step denoising actions. The symbol x in the Mamba block denotes matrix multiplication, and the SiLU activation function.
Positionaa Encocing
Encoding
Denoked
Actions
M
Linear
N*
Aggregation
Mamba
Mamba
PE
PE
AV
AV
AV
AV
TV
SV
SV
SV
SV
PP
Learnable action
Leamable time
Learnable
Figure 2: ED-Ma: Differentt from the D-Ma model, ED-Ma contains Mamba encoder which is used to process the time embedding and state embedding, and 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 introduced to the decoder output for sequence alignment.
3.1 Mamba: Selective State-Space Models
Inspired by the Attention mechanism [15] in Transformers, Mamba [13] improves upon the Struc- tured 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 E BxXLx and the output sequence y E BBxLLD where B, L,, refer the batch size, sequence length and dimension respectively. standard SSM with hidden state dimension defines time-invariant parameters A RREX B R C E RDxH and the time-step vector A to map the inputs up x() to hidden state, which then can be projected to the output y(() (More details are in Appendix C.1). However, this property limits the SSM for contextual learning. Mamba implements the selection mechanism by making the SSM parameters function of the input, that is,
B Linear(x), C =Lineer((x)) A SoftPlus(Linear(z))
(1)
where E BBXXXD C E R2xLXD and A E BLLLD Linear refers to linear projection layers and SoftPlus is smooth approximation of ReLU. Then the output be calculated via
= SSM B, (x)).
where (A, B) are discretized [42] counterparts of (A, B) with time-steps A. As the time-varying model can only be calculated in recurrent way, Mamba further implements hardware- method to compute the selective SSM efficiently. Figure shows simplified Mamba block.

---

Algorithm D--M DDP Inference
Algorithm ED-Maa DDP Inference
Require: Observation 8g-Ke Kiks
Require: Observation Sk- Kik. Mamba Encoder Encm. Mamba
Mamba Decoder Decm
Decoder Decm: Position Embedding PE, Learnallc action variables agck+J, Icarnable state variables šд Sk-Kik, Icarnable
Action Encoder Enc..
Position Embedding PE
time variables {
1: a (/
1: a ~ NV I)
2: = T................. ......1 do
2: for T.... ......1 do
3: Er Linear(!)
(E(())
4:
E. PE RRsSNet ks))
4: E. PE( (ee(ReNet(Ss--x..)
5:
E. PE(Linear(a (a))))))
5: Ea PE (EE(Enaa(B) (nca(
6:
Em Dm((ca((e (e, , Es, ac:k+)))
Sequence alignment
6:
ak*++J
1-1
7:
E. Dmi(ca(f, By- Kik, Ea])
b. First decoder layer
Sinear(Decn)) Be, Ea))
agi+J Dm Ea) b. Remaining decoder layers
7: end for
9: end for
8: return acc++
10: return J
3.2 Policy Representations
In this work, policies (DDPs). For clarity, we focus the non-sequential case.
two policy representations: behavioral cloning (BC) and denoising diffusion
Behavioral Cloning assumes a parameterized conditional Gaussian distribution as policy repre- sentation, i.e., T (als)) = N (alluo(s), I)). Maximizing the likelihood for the model parameters simplifies to mean-square error (MSE) loss, that is,
LB (0) Ex [s))o(s)
where the expectation over s, is approximated using state-action pairs in the demonstration data. Denoising Diffusion Policies utilize denoising function Eg to sample from a Markov chain (a (a')?=
Otk ( 1 at v1 āi
co (taai,t,s)) t. + Varz', where z' ~N,,1.
(4)
starting from ay ~ (O, to produce a noise-free action a for a given observation s. The denoising function trained to predict the source noise z² N/ I) of a noisy action af vā,a" + zz by minimizing the loss
(D) Es,a,t [ale
aa, s)
(5)
where a, = =I1 Og.. The expectation over corresponds to a uniform sampling in {1....,T).
Mamba for Imitation Learning
In this section, we explain how we leverage Mamba for Imitation Learning (IL). Drawing inspiration Decoodrr (ED-Tr) Transformers, we propose from the successful Decoder-only (D-Tr) Mamba (D-Ma) and Encoder- Mamba (ED- two Mamba-based architectures: Decoder-only
Ma). These architectures serve the parameterization for the policy. Specifically, when employing Behavioral Cloning (BC), these architectures parameterize the mean μ4 of conditional Gaussian distribution. Conversely, when using Denoising Diffusion Policies (DDPs), these architectures pa- rameterize the denoising function éş. Given the straightforward nature of the former case, we focus on introducing these architectures in the context DDPs.
4.1 Decoder-Only Mamba
Similar to the Decoder-only transformer, use 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 denoising function ey that takes sequence of observations Sk-K:k- noisy

---

actions a. and diffusion step t to generate less noisy sequence of actions t-1 J. The diffusion step is encoded using time-embedding TE. The observations are encoded using ResNet-18, with shared weights across images from different time steps. An action encoder ENC, is used to tokenize the noisy action inputs. Additionally,, position embeddings 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 Dec,- 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 pro- cess complex input-output relationships, especially for cases where the input and output sequences differ in structure. However, Mamba does not provide such mechanism to support the encoder- decoder structure since the target and the source share the same sequence length. Here, propose 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 Encm is used to pro- cess the time embedding and state embedding, and the Mamba Decoder Decm is used to process the noise embedding. Since the inputs to the Em and Dm have different lengths of sequences, we propose add learnable variables 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 MalL in imitation learning. Our investigation focuses on the following key questions:
Q1) Can MalL achieve comparable superior performance to Transformers?
Q2) Can MalL utilize multi-modal inputs, such as language instructions?
Q3) How effectively do MalL handle sequential information in observations?
5.1 Baselines
In this work, we mainly aim to explore the potential of Mamba visual imitation learning com- pared to Transformer structures. Therefore, our experiments contain four architectures: Decoder- only Transformer (D-Tr), Encoder-Decoder Transformer (ED-Tr), Decoder-only Mamba (D-Ma), Sicoder-Decoder Mamba (ED-Ma). For fair comparison, we ResNet18 to encode visual in- puts for each method. For tasks that use language instructions, 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 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 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 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
5

---

Figure 3: LIBERO benchmark suites with total 130 tasks in five different scenes.
w/o language
language
Policy
Backbone
LIBERO-Object
LIBERO-Spatill
LIBERO-Long
LIBERO-Goal
LIBERO-90
BC-HI
D-Tr
0.358
0.257
0.222 00.77
0.362
0.311
ED-Tr
0.480=0.027
0.333
0.265+0.013
0.325=#0.000
0.268+0.050
D-Ma
0.55022:062
0.410 o58
0.287:20.0m 037
0.48220.00 000
0.559_0.021
ED-Ma
O. 0.618=0.000 058
0.302 og9
0.275sa0.22
0.473=00.011
0.514+0.008
BC-H5
D-Tr
0.286:00n440
0.243:0.032
0.315==0.006
0,28440.012
ED-Tr
0.417:00.30
0,287=0.0pm n90
0.248:
0.338
0.300
D-Ma
0.765220.0 115
0.372 024
0.252
0.43:00029
0.529
BD-Ma
0.743:20:020
029
0.2622:0
0.317=20106
0.599 0308
DDP-HI
D-Tr
0.667:20:058
0.473 g:10
0.330220.02m 0.20
0.300
0.181 20
ED-Tr
0. 0.680.00.27
0.420=
0.373
0.478::0.083
0.309
D-Ma
0.728:00082
0 0.492:0..om 20.0000000
0,3433=0.030
0.510=0.043
1.4
ED-Ma
0.71320:00
0.483
0.395 oss
0.565
0.6033
DDP-HS
D-Tr
(.. 0.76220.070
0.43522..0 0:1
0.353:20.03
0.315220.58
0.161 021
RD-Tr
0.81 220006
0.415200000 097
0.348:20.019
0.51820.014
0.437 046
D-Ma
0.815 004
0. 0.538 a .44
0.387:10.40
0.388
0.352
RD-Ma
0.778220:00
0.513
0.417:20.025
0.563:20.070
Table 1: Performance on LIBERO benchmark with 20% data, where "w/o language" indicates that do not use language instructions and "w/ language" means language tokens generated from a pre-trained CLIP model, HI and H5 refer to using current state and 5 steps historical states respectively.
of robotic learning and manipulation capabilities. The task visualizations are in Figure 3. More details can found in Appendix B.
Evaluation Protocol We compared each method across five LIBERO task suites separately. In- stead of using full demonstrations, 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 parameear for both Transformer and Mamba, making sure they have 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.
Main Results. We report the main results in Table 1. Our Mamba-based architectures, D-Ma and ED-Ma, significantly outperform Transformer-basee methods across all LIBERO task suites based on the BC policy. Specifically, Mamba-based models achieve nearly 30% improvement in success rate in LIBERO-Object and LIBERO-90. When using the DDP policy, our models consis- tently surpass the Transformer baselines, with performance improvements exceeding 5% in most tasks. These results confirm Q1, demonstrating that MalL achieves superior performance com- pared to Transformers. To address Q2, we compare MalL with Transformers LIBERO- Goal and LIBERO-90, using additional language embeddings as inputs. We observe significant improvements with Mamba-based methods in these tasks, indicating that MalL effectively leverages multi-modal inputs. Given that most recent visual imitation learning works use historical observations as inputs, we evaluated the methods with and 5 historical observations. We found that historical informa- tion does not always enhance performance. Only in LIBERO-Ohject do the H5 models outperform the H1 models, while in other tasks, H5 models achieve similar worse results. Mamba-based H5 models again perform consistently better than Transformer based models, which indicates that MalL is able to capture sequential observation features effectively, answering Q3.

---

0.77
0.8
0.5
0..
0.7
0.4
0.6
0.3
11.2
0.3
0.3
occllsion sale
1055
2005
00.5
100%
1050
2055
1005
(a) DDP-H5: object
(h) DDP-H5: spatial
(c) BC-HI: 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.
Ablation on Observation Occlusions In order further understand the sequential learning ability of Transformer and Mamba, randomly mask out areas of images and test the model's perfor- mance 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 MalL performs well with only 20% of the demonstrations, 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. More experimental results can be found in Appendix Table 6. It is evident that Mamba-based models significantly outperform Transformers when data is scarce and perform comparably as the dataset size increases.
5.3 Real Robot Evaluation
We designed three challenging tasks based on the 7DoF Franka Panda Robot, uti- lizing visual inputs for the model. Two cameras, positioned at different angles in front the robot, provide the visual data. One image is cropped and resized (128, 256, 3). while the other is resized to (256, 256,3), 3). The entire setup is vi- sualized in Figure 5. These images are stacked at each timestep to form the ob- servations. 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 presented in Table 2-4.
Right side
Camera
Figure Real robot setup. We use visual inputs from two cameras as observations for policies.
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 demonstrations for each individual tasks. During 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 one ma- nipulation 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 blue block into 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
7

---

task demands precision control which makes it challenging, we fix the LEGO box position and only randomly initialize the state the block.
CupStacking. The robot required 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 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 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.
Two stage
Banona
Carret
Stage
Stage22
ED Ma
0.55
0.70
0.70
0.55
Figure 6: PickPlacing (front camera view)
Table 2: Pick-Place Success rates
Stage-1
Stage-2
ED-T:
0.45
0.33
ED-Ma
Figure 7: Inserting (right camera view)
Table Inserting success rates
DDP-HI
Stack-1
Stack-2
ED-T:
0.4
ED-Ma
0.8
0.55
Figure CupStacking (front camera view)
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 fair comparison, we used the same initial states for both Transformer and Mamba evaluations. From the results, Mamba- Mamba-based methods achieve comparable results with Transformer models.
Limitations
While MalL demonstrates excellent performance with smaller dataset sizes, its advantage becomes less pronounced as dataset scales up. When trained on larger datasets, Mall. achieves results comparable to, but not surpassing, those of Transformer models. Additionally, Mamba is designed 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 MalL,, 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, MalL offers a competitive alternative to the traditionally large and complex Transformer-based policies. The in- troduction of Mamba in an encoder-decoder structure enhances its versatility, making it suitable for standalone
well integration into advanced architectures like diffusion processes. Exten- sive evaluations the LIBERO benchmark and real robot experiments demonstrate that MalL not only matches but also surpasses the performance of existing baselines with limited demonstrations, establishing it as promising approach for IL tasks.

---

Acknowledgments
This work was supported by funding from the pilot program Core Informatics of the Helmholtz As- sociation (HGF). NS and GN were supported by the Carl Zeiss Foundation under the project JuBot (Jung Bleiben mit Robotern). Xiaogang Jia acknowledges the support from the China Scholar- ship Council (CSC). The authors acknowledge support by the state of Baden- Württemberg through bwHPC, well as the HoreKa supercomputer funded by the Ministry of Science, Research and the Arts Baaden-Württebeerr and by the German Federal Ministry of Education and Research.
References
[1] Osa, J. Pajarinen, Neumann, A. Bagnell, P. Abbeel, J. Peters, al. algorithmic perspective imitation learning. Foundations and Trends@ in Robotics, 7(1-2):1-179, 2018.
[2] 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.
[3] T. Z. Zhao, V. Kumar, S. Levine, and C. Finn. Learning fine tne-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:2304. 1 13705, 2023.
[4] Florence, C. Lynch, A. Zeng, O. A. Ramirez, A. Wahid, L. Downs, A. Wong, J. Lee, I. Mor- datch, and Tompson. Implicit behavioral cloning. In Conference on Robot Learning, pages 158-168. PMLR, 2022.
[5] C. Chi, S. Feng, Du. Z. Xu, E. Cousineau, Burchfiel, and S. Song, Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
[6] M. Reuss, M. Li, X. Jia, and R. Lioutikov. Goal conditioned imitation learning using score- based diffusion policies. In Robotics: Science and Systems, 2023.
[7] A. Mandlekar, D. Xu, J. Wong, S. Nasiriany, C. Wang, R. Kulkarni, L. Fei-Fei, S. Savarese, Y. Zhu, and R. Martin-Martin. What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXiv:2108. 032998,,, 2021.
[8] 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.
[9] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. De- hghani, 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 Con- ference on Learning Representations, 2021. URL VicbFaNTTy.
met/forum?id=
[10] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.
[11] Q. Wen, T. Zhou, C. Zhang, W. Chen, Z. Ma, Yan, and L. Sun. Transformers time series: survey, 2023.
[12] A. Gu, K. Goel, and C. Re. Efficiently modeling long sequences with structured state spaces. In International Conference Learning Representations, 2021.
[13] A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint arXiv:2312.00752. 2023.

---

[14] B. Liu, Y. Zhu, C. Gao. Y. Feng, Q. Liu, Y. Zhu, and P. Stone. Libero: Benchmarking knowl- edge transfer for lifelong rohot learning. Advances in Neural Information Processing Systems, 36, 2024.
[15] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. Gomez, E. Kaiser, and I. Polo- sukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
[16] Achiam, Adler, S. Agarwal, L. Ahmad, Akkaya, F. L. Aleman, D. Almeida, Altenschmidt, S. Altman, S. Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
[17] 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 meclelland. vol. 1. 1986. Biometrika, 71:599-607, 1986.
[18] S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural computation, 98): 1780, 1997.
[19] K. Cho, B. Van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Ben- gio. Learning phrase representations using rnn encoder-decoder for statistical machine trans- lation. arXiv preprint 140iv:1406.1 2014.
[20] J. Chung, C. Gulcehre, K. Cho, and Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv: 1412.3555, 2014.
[21] L. Zhu, 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.
[22] T. Smith, A. Warrington, and Linderman, Simplified state space layers for sequence modeling. In The Eleventh International Conference Learning Representations, 2023.
[23] 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.
[24] R. Hasani, M. Lechner, T.-H. Wang, M. Chahine, A. Amini, and D. Rus. Liquid structural state-space models. In The Eleventh International Conference Learning Representations, 2022.
[25] 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 1rXiv:2404.09516, 2024.
[26] D. A. Pomerleau. Alvinn: An autonomous land vehicle a neural network. Advances in neural information processing systems, 1. 1988.
[27] S. Schaal. Learning from demonstration. Advances in neural information processing systems,
9, 1996.
[28] A. Taranovic, Kupesik, N. Freymuth, and G. Neumann. Adversarial imitation learning with preferences. In The Eleventh International Conference on Learning Representattonss 2022.
[29] W. Sun, A. Venkatraman, G. J. Gordon, B. Boots, and J. Bagnell. Deeply aggrevated: Dif- ferentiable imitation learning for sequential prediction. In International conference machine learning, pages 3309-3318 PMLR, 2017.
10

---

[30] B. Wang, E. Adeli, H.-k. Chiu, D.-A. Huang, and J. C. Niebles. Imitation learning for hu- man pose prediction. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 7124-7133, 2019.
[31] 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.
[32] C. Lynch, M. Khansari, T. Xiao, V. Kumar, J. Tompson, S. Levine, and Sermanet. Learning latent plans from play. In Conference on robot learning, pages 1113-1132. PMLR, 2020.
[33] N. M. Shafiullah, Cui, A. A. Altanzaya, and L. Pinto. Behavior transformers: Cloning modes with one stone. Advances in neural information processing systems, 35:22955-22968, 2022.
[34] 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.
[35] S. Lee, Y. Wang, H. Etukuru, H. Kim, N. M.M. Shafiullah, and L. Pinto. Behavior generation with latent actions. arXiv preprint arXiv:2403.03188, 2024.
[36] 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.
[37] H. Wu, Y. Jing, C. Cheang, G. Chen, Xu, X. Li, M. Liu, H. Li, and T. Kong. Unleash- ing large-scale video generative pre-training for visual robot manipulation. arXiv preprint arXiv:2312.13139, 2023.
[38] R. Shah, R. Martin-Martin, and Y. Zhu. Mutex: Learning unified policies from multimodal task specifications. arXiv preprint arXiv:2309. 2023.
[39] M. Reuss, Ö. E. Yagmurlu, F. Wenzel, and R. Lioutikov. Multimodal diffusion transformer: Learning versatile behavior from multimodal goals. In First Workshop on Vision-Langure Models for Navigation and Manipulation at ICRA 2024, 2024. URL htt.s: httptttss://pppeeeeiiiei net/forrumid=PPtELL..
[40] 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.
[41] A. Sridhar, D. Shah, C. Glossop, and S. Levine. Nomad: Goal masked diffusion policies for navigation and exploration. arXiv preprint arXiv:2310.07896 2023.
[42] M. Zhang, Saab, Poli, Dao, K. Goel, and Ré. Effectively modeling time series with simple discrete state spaces, 2023.
[43] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervi- sion. In International conference on machine learning, pages 8748-8763, PMLR, 2021.
[44] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-688 2020.
[45] Y. Lin, A. S. Wang, G. Sutanto, A. Rai, and F. Meier. Polymetis. htt../////r/oooomeeitieee 2021.
https: //
facebookresearchh
[46] L. Van der Maaten and G. Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.

---

A More Experimental Results
A.1 Comparison on Different Checkpoints
Evaluating multiple checkpoints can provide important insights into our model's generalization ca- pabilities. Here, we conducted additional experiments assessing the model after 40, 50, and 60 epochs. The results presented in Table suggest that the model neither overfits nor underfits the data at these intervals.
w/o language
w/ language
DDP-H5
Backbone
LIBERO-Objee
LIBERO-S LIBERO-Spatial
LIBERO-Lomd
LIBERO-God
Epoch 40
ED-Tr
0.770
0.415
0.345a0.044
0.502
ED-Ma
0.802:400300
0.478
0.355:+..22
0.615:0.028
Epoch 50
ED-Tr
0. 0..58
0.412++0.80
0.333++0.023
0.497
ED-Ma
0. 051
0.500=0.0st
0.385=+0.025
0.580:20.041
Epoch
ED-Tr
0.767:10.045
0.448
0.3.18+0..00
0.505:20.036
ED-Ma
0.805#0..00
0.497
0.390=0 028
0.610=0.035
Max
ED-Tr
0.770=
0.448+0.000
0.348
0.505
ED-Ma
0.807
0.500=0.08
0.390
0.611:00008
Table Comparison on different checkpoints based on DDP-H5 methods.
A.2 Comparison on Full Demonstrations
To investigate how Mall and Transformers scale with dataset size, we evaluated both models on the LIBERO benchmark with access to the full dataset. Due to time constraints, we focused solely on the Encoder-Decodee architecture for both MalL and Transformers. We conducted experiments across different checkpoints, evaluating each model at three checkpoints and reporting their best performance. The results are presented Table 6:
w/o language
w/ language
Policy
Backbone
LIBERO-Object
LIBERO Spatial
LIBERO-Long
LIBERO-Goal
BC-HI
ED-Tr
0.847#0.00
0.658=0.051
0.658=0.03
0.668
ED-Ma
0.862=0.088
0..93
0.64220.011
00.00
BC-H5
ED- Tr
0.910=0.018
0.733:0..04
0.632
0. 0.772
ED-Ma
0.907410.019
0.697=0.062
0.687=0.038
0.833=00.015
DDP-HI
ED-Tr
0.918=0.0m 028
0.707=2.035
0.683:20.015
0.788:20.054
ED-Ma
0.913
0.725
0.700
0.818=0.011
DDP-H5
ED-Tr
0.952=0.0m 003
0.745=2.018
0.731+0.081
0.873:10.00
ED-Ma
0.932+4.029
0.758
0.772=0.004
0.893
Table 6: Performance on LIBERO benchmark with access to the full dataset.
A.3 Experiments with Language Instructions
In Table of our manuscript, we did not utilize language instructions for the LIBERO-Object, LIBERO-Spatial, and LIBERO-Long tasks. To gain further insights, we conducted additional periments these tasks using language embeddings, with access to 20% and 100% of the data. The results presented in Tables 7 and 8:
12

---

w/ language
Policy
Backhone
LIBERO-Ohject
LIBERO-Spatial
LIBERO-Long
BC-HI
ED-Tr
0.596+0.123
0.474
0.316=0.035
ED-Ma
0.0892
0..08
0.408
BC-H5
ED-Tr
0.714440.036
0.426=10.03
0.278
ED-Ma
0.005
0.395
0.303
DDP-HI
ED-Tr
0.7611:0.00
0.001
0.396:2.031
ED-Ma
0.781=0.045
0.576=0.022
0 0.499=0.008
DDP-H5
ED-Tr
0.864=0.025
0.618:10.053
0.431
ED-Ma
0.854+0.000
0.609
0.460=0.035
Table 7: Performance on LIBERO benchmark with 20% data.
w/ language
Policy
Backhone
LIBERO-Ohject
LIBERO-Spatial
LIBERO-Loong
BC-HI
ED-Tr
0.803+0.013
0.718+0.020
0.724:0.072
ED-Ma
00.59 0.776
0.770=0.054
BC-H5
ED-Tr
0.916+0.027 0.800+0.015
0.764=0.039
ED-Ma
0.929=0.017
0.838
0.758
DDP-HI
ED-Tr
0.916=0.02z
0.813
0.763:10:037
ED-Ma
0.001
0.743:10.043
0.786
DDP-H5
ED-Tr
0.955=0.0m
0.836+0.045
0.826:10.035
ED-Ma
0.943+0.050
0.883:+0.00
0 0.830=0.00m
Table 8: Performance on LIBERO benchmark with full data.
B
Additional Task Details
B.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 's spatial understanding and its ability to navigate and manipulate objects based on their spatial configurations.
LIBERO-Objec Object suite has tasks involving different objects the same layout, requiring the robot to pick-place unique object at 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 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-99 contains 90 short-horizon tasks with significant diversity. This suite includes a wide variety of object categories, layouts, and goals, providing comprehensive evaluation of the robot's adaptability and versatility in handling diverse scenarios.
B.2 Real Robot Tasks
B.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 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 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 vary.
13

---

B.2.2 Inserting
In the Inserting task, there is LEGO box with hole in the middle and another block which can fit in the hole of the LEGO box. The goal in this task grab the block and insert it in the hole the LEGO box. The position and orientation of the LEGO box are fixed, but the position and orientation of the block can vary inside specific area the table.
B.2.3 PickPlacing
In the PickPlacing task, there are two different bowls, a banana and carrot. The goal is to pick the banana and the carrot and place them inside the bowls. The banana is always placed 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.
Architectures
C.1
State Space Models
Inspired the Attention mechanism [15] in Transformers, Mamba [13] improves upon the Struc- tured State-Space Sequence (S4) Model by using selective operator, to propagate forget information over context, allowing it to filter out relevant features. Specifically, the S4 model is inspired by a continuous system, x() E R² + v(1) R,, that maps x(l) into hidden state h(l) E IR", which then can be projected onto the output yt), that is,
yl) Ch(l) with h'(1) Ah(1) Bx((),
where A E RRx are the evolution parameters, E R and C E DR are the projection parameters. S4 transforms the parameters (A, B) to discrete parameters (A, B) via
A exp AA) and B
(AA) I) AB,
where is the step size. The model then can be computed as either linear recurrence, that is,
V Ch with h4 Ah + Bz,,
(8)
or through global convolution given by
y == K, with = (CB, CAB, CA"B). B).
(9)
C.2 Transformers
Here we depict two transformer-based architectures in diffusion policy: a decoder-only model (Fig- ure 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.
D
Model Details
D.1 Parameter Comparison
For all Mamba-batt based policies, we fix the expansion factor at 2. The parameters exhibit slight vari- ance when using the same number of layers due to differences in the choice 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 models in Table 10.
We also evaluate the inference time local PC equipped with RTX 2060 GPU, using batch size of 32 to ensure all models are assessed under the same conditions in Table 9,
14

---

PE Positional Encoding
k
TE Time Encoder
Linear Output Layer
Nx
Self-Attention
TE
ResNet-18
Action Encoder
BK-K
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 linear output layer to predict future actions.
Encoding
TC
Time Encoder
Linear Output Layer
Cross-Attention
Nx
Nx
Self Self-Attention
Self-Attention
PE
PE
PE
PE
PE
PE
PE
PE
PE
TE
ResNet-18
Action Encoder
K
...
S2-2
8
a
Figure 10: Encoder-decoder learning bhock,This figure illustrates the architecture of 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-altention is fed to linear output layer to predict future actions.
D.2 Training Details
We list the training hyperparameters for the Transformer-basse and Mamba-based policies in Table 10. To ensure 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 only take 10 demonstrations for each task in main experiment.
All models are trained on cluster equipped with A100 GPUs, with a batch size of 256, over 50 epochs using 3 different seeds. Finally, caleulate the average success across these 3 seeds.
D.3 Ablation on Observation Occlusions
For the amage-based model, recent observations are not significantly different from the current ob- servation. To evaluate the influence of historical data and compel the model focus more on past observations, employ occlusion techniques on the observations. This involves randomly mask- ing portions of the images like the Figure 11. The occlusion rate determines the extent to which the images masked.
15

---

Method
Layers
Params
Inference time (HI)
Inference time (H5)
BC
D-Tr
23.6M
0,03043
0.05906
8
24.0M
0,03305
0.06306
ED-Tr
4
24.2M
0,03230
0.06368
4+6
24.7M
0.03505
0.06563
D-Ma
10
23.6M
0.02603
0.05340
12
23.8M
0.02797
0.05477
16
24.3M
0.02855
0.05614
ED-Ma
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.061 12
DDP
D-Tr
6
23.7M
0.10300
0.14074
8
24.1M
0.12882
0.16428
ED-Tr
4+4
24.3M
0.14443
0.17995
4+6
24.8M
0.17244
0.21022
D-Ma
10
23.7M
0.11185
0.14384
12
23.9M
0.12616
0.16145
16
24.2M
0.15699
0.18236
ED-Ma
6+6
23.8M
0.12653
0.15984
6+8
24.0M
0.14025
0.17350
10
24.3M
0.15977
0.19398
8 10
24.5M
0.16883
0.20553
10 10
24.7M
0.18332
0.22571
Table 9: Inference Time Comparison, H1 and H5 refer to using current state and 5 steps historical states respectively.
Figure 11: Visualization of Observation Occlusions
16

---

Number of Encoder Layer Methods Parameters
libero_ohject
lihero_spatial
lihero_long
lihero_goal
lihero_90
history
D-Ma
Number of Mamba Layer
10
16
16
12
16
Hidden State Dimension
16
8
16
Convolution Width
4
Expand Factor
2
2
2
ED-Ma
10
8
10
Number of Decoder Layer
10
Hidden State of Encoder
8
8
8
Hidden State of Decoder
16
8
16
8
Encoder Convolution Width
2
Decoder Convolution Width
2
2
2
Expand Factor
2
2
2
2
D-Tr
Number of Attention Block
8
8
Number of Attention Head 4
4
4
4
ED>-Tr
Number of Encoder Block
4
4
4
Number of Decoder Block
6
6
Number of Encoder Head
4
4
4
4
Number of Decoder Head
4
4
4
history
D-Ma
Number of Mamba Layer
10
16
16
12
16
Hidden State Dimension
16
8
16
Convolution Width
Expand Factor
2
2
2
ED-Ma
Number of Encoder Layer
10
10
10
Number of Decoder Layer
10
8
Hidden State of Encoder
16
16
8
16
8
Hidden State of Decoder
16
16
8
Encoder Convolution Width
2
Decoder Convolution Width
4
4
4
2
Expand Factor
2
2
2
2
D-Tr
Number of Attention Block
8
8
Number of Attention Head 4
4
4
4
ED-Tr
Number of Encoder Block
4
4
4
Number of Decoder Block
6
6
6
Number of Encoder Head
4
4
4
4
Number of Decoder Head
4
4
4
Table 10: Hyperparameter in Simulated Experiments
17

---

leader
follower robot
Figure 12: Data collection scene for real robot
E
Real Robot Details
E.1
Setup
The current setup consists of two robots and two computers. One of the computers, the robot PC, rums real-time OS to control the robots reliably using Polymetis (E.2), which is 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 server from which the desired joint position and the gripper of the robot be configured. The other PC, the workstation, the desired joint position and gripper state the robot PC each time step. The workstation also has access to two cameras, which are used to capture images as input for the models.
the models and sends
E.2 Polymetis
Polymetis [45] is software framework designed 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 mo- tion planning, adaptive control, and sensory feedback. Polymetis supports various robotic platforms and sensors, making 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 the desired joint state. The state of the gripper is considered to be binary, either closed open. A threshold is 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.
18

---

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 the robot each time step instead of the output of the model. The length of this trajectory varies depending how far away the output of the model is from the current robot state.
F Analysis on the State Space Representations
To better understand the advantages of MalL over Transformers, conducted detailed analysis of the latent representations produced by both methods. Specifically, we compared the Encoder- Decoder architectures of MalL and Transformers by visualizing their nigh-dimensiona represen- tations using t-SNE [46]. We employed BC-based models trained on the LIBERO-Object dataset using full demonstrations and rolled out the models across entire trajectories. The visualizations represent the latent spaces before the final action prediction layer (linear). Figure 13 displays the t-SNE results for four different trajectory representations.
As shown in the figure, MalL's state representations are notably more structured across all trajecto- ries compared to those generated by Transformers. More structured representations could potentially improve the generalization of MalL, which explains the advantages of MalL over Transformers.
19

---

MaiL
Transformers Representations
Comparent
componars
(a) Trajectory
MaIL Representations
Transtormers Repr Representations
5PE compontree
Frogeo
1."
1-5MC Comparen:
componertl
(b) Trajectory
MaiL Representations
Transformers Representations
Truediy Trggesso
< cor componee2
1CNC comparent
ENNC componatt
(c) Trajectory 3
Mail Representations
Transformers Representations
emppnort
regeoaiion
componert1
(d) Trajectory
Figure 13: Visualization of latent representations different trajectories.
20


Reviews:
[{"review_id": "ZBt85NHOGW", "reviewer": "Reviewer_k2cG", "summary": "", "questions": "", "limitations": "", "rating": 2, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 2, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "**Main Idea:** Employ the Mamba architecture for Imitation Learning (IL).  **Main Contributions:** - Integrates the Mamba model with Behavior Cloning and Denoising Diffusion Policies. - Furthermore, the study introduces two architectural variants: decoder-only and encoder-decoder. - Both simulation and real-world experiments validate the stronger performance of MaIL over Transformer-based approaches in IL.", "summary_of_recommendation": "Overall, this paper presents a timely \"A+B\" style contribution but lacks a clear explanation of the underlying mechanisms that drive its success. A more detailed exploration into why MaIL works would greatly enhance the depth and impact of the findings.", "technical_quality": 2, "clarity_of_presentation": 3, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "- Why do the authors emphasize MaIL's effective handling of multi-modal sensory input? I would assume this capability also extends to Transformer models, correct?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yR4Voa6nEP", "reviewer": "Reviewer_u9Yp", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This paper presents Mamba Imitation Learning (MaIL) as an efficient alternative to Transformer-based policies, leveraging the Mamba state space model.  MaIL leverages Mamba's advantages and matches Transformer performance with better speed, memory usage, and efficiency.  The authors offer two variants: - Decoder-only: Denoises actions using noised inputs and time embeddings. - Encoder-decoder: Handles multi-modal inputs with learnable embeddings.  MaIL outperforms Transformers on LIBERO benchmarks with 20% of the dataset, processes multi-modal inputs, and is robust to noise.  The authors demonstrate the effectiveness of MaIL in three real robot tasks.", "summary_of_recommendation": "Given that this paper explores the benefits of the recent Mamba architecture in the context of imitation learning, both in simulation and real-world setups, I find it to be a valuable contribution to the field. Therefore, I recommend accepting this paper.", "technical_quality": 3, "clarity_of_presentation": 3, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "* How do Transformer-based baselines compare with your approach when having access to the full LIBERO dataset?\n* Could you please elaborate more on the Aggregation approach? How are the state, action, and time variables learned? What are the advantages and disadvantages of this approach?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zz4yo69VVv", "reviewer": "Reviewer_mQZR", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 4, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "The paper presents a new architecture for handling sequences of observations in imitation learning, leveraging Mamba as the backbone. In their experiments, the authors try to assess this architecture’s performance for visuomotor imitation learning compared to a Transformer backbone, and the ability to handle multi-modal instructions such as language instructions. The authors present results on the LIBERO multi-task benchmark in simulation as well as on a real robot.", "summary_of_recommendation": "I believe the paper takes a stride into an important direction of discovering a better backbone for visuomotor imitation learning. However, I recommend a weak reject for this paper as I am not convinced that the quantitative results in simulation suggest that this architecture is sufficiently superior than the Transformer backbone. As I mentioned in my review, I believe the evaluation protocol is suboptimal and takeaways might change had evaluation been done differently. Additionally, there are also some missing baselines.", "technical_quality": 4, "clarity_of_presentation": 4, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "Please refer to the weaknesses mentioned above.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CKdUaWAzoT", "reviewer": "Reviewer_Vk9K", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This paper introduces MaIL (Mamba Imitation Learning), a novel approach that utilizes Mamba, a computationally efficient alternative to transformer-based models, for imitation learning. MaIL is evaluated against Behavior Cloning and Denoising Diffusion Policies, demonstrating superior performance. Extensive experiments in both simulation and real-world environments confirm that MaIL not only matches but often surpasses existing baseline methods.", "summary_of_recommendation": "My opinion is positive, and I would accept this manuscript for publication as it explores an exciting and novel approach of using Mamba in Imitation Learning. However, I suggest the authors to improve the discussion of the results of their experiments.", "technical_quality": 3, "clarity_of_presentation": 4, "potential_impact": 3, "robotics_focus": 3, "quality_of_limitations": 3, "questions_for_rebuttal": "- How are the performances of D-Ma and ED-Ma affected by the historical states? In the experiments, ED-Ma outperforms the others in BC-H1; however, when historical data is added in BC-H5, D-Ma outperforms ED-Ma.\n- Why is the performance of ED-Ma lower than ED-Tr in stage 2 of the Two-stage Pick-Place and Inserting tasks?\n- In the ablation studies of BC-H1 and BC-H5, at 60% and 100%, the success rates of ED-Ma and D-Ma were very similar to those of D-Tr and ED-Tr. What could be the reason for this?\n- Why, in the ablation studies BC-H5: spatial at 100%, is the success rate of D-Ma lower than the other methods? Its success rate at 100% seems to be the same as at 60%.\n- In Figure 4 (b), the label is DDP-H5. Is this correct? Should it not be DDP-H5: object?\n- Could you have additional discussion over the application of the proposed method to BC with longer context?\n- Could you analyze the representations of the formed state space qualitatively?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Strengths:
- Insight that object level tokenization needs less data is useful
- Aligns text level representations like Llama with object level representations
- Long tail scenario is a data hungry regime, achieves few shot generalization

Weaknesses:
- Performance highly depends on pretrained tokenizer, and the latter needs lots of data. Inference latency might be slower given larger models?
- Trajectory end to end performance not better than PARA-driver. This is a major weakness and is not well explained. 
- No experiments on hardware


Reviewer k2cG: Lack of explanation for MaIL's superior performance, its ability to handle long-horizon dependencies, and the relevance of Mamba's speed advantage in imitation learning. Also, the reviewer sought clarification on the emphasis on multi-modal input handling and the claim about MaIL's convergence with Transformers on large datasets.The authors provided motivations for MaIL's advantages, conducted experiments demonstrating its ability to handle long-horizon dependencies, and offered inference speed comparisons. They clarified their claims and conducted additional experiments with the full LIBERO dataset. They also justified their emphasis on MaIL's multi-modal handling capabilities.

Reviewer u9Yp: The paper could be misleading by claiming MaIL's outperformance without specifying the low-data setting. The ablation on dataset size was insufficient, and the Mamba Aggregation approach needed more elaboration. The authors acknowledged the potential misinterpretation and planned to correct it. They conducted additional experiments with the full dataset and provided a detailed explanation of the Mamba Aggregation approach.

Reviewer mQZR: The reviewer had concerns about the low success rates on the LIBERO benchmark, the definition of an epoch, whether the model fit the data well, and the absence of BC-RNN as a baseline. The authors attributed the low success rates to using a fraction of the dataset and provided results with the full dataset. They clarified the definition of an epoch, provided the number of gradient steps, conducted experiments with multiple checkpoints, and explained the exclusion of BC-RNN while providing results with an LSTM-based model.

Reviewer Vk9K: The reviewer suggested expanding the discussion of results, evaluating Mamba's ability to handle long-term context, and qualitatively analyzing state space representations. They also sought clarification on specific results and figure labels. The authors provided further discussion, conducted experiments with extended history lengths, and committed to visualizing state space representations. They also clarified figure labels and addressed specific questions about the results.
Decision: Accept