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: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-07 13:41:57

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.

Main paper: 
 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. 
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 inputindependent 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.Early imitation learning methods primarily focused on learning one-toone 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. 
preliminaries: 
 Here, we briefly review the basics of the Mamba model and explain the policy representations considered in this work. 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× 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 × in the Mamba block denotes matrix multiplication, and σ 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.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 ∈ R B×L×D and the output sequence y ∈ R B×L×D , where B, L, D refer to the batch size, sequence length and dimension respectively. A standard SSM with hidden state dimension N defines time-invariant parameters A ∈ R N×N , B ∈ R N×D , C ∈ R D×N and the time-step vector ∆ 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 C.1). 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,EQUATIONwhere B ∈ R B×L×D , C ∈ R B×L×D , and ∆ ∈ R B×L×D , Linear refers to linear projection layers and SoftPlus is a smooth approximation of ReLU. Then the output can be calculated viaEQUATIONwhere (A, B) are discretized [42] counterparts of (A, B) with time-steps ∆. 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. 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., π(a|s) = N (a|µ θ (s), σ 2 I). Maximizing the likelihood for the model parameters θ simplifies to a mean-square error (MSE) loss, that is,EQUATIONwhere the expectation over s, a is approximated using state-action pairs in the demonstration data.Denoising Diffusion Policies utilize a denoising function ϵ θ to sample from a Markov chain (a t ) T t=1EQUATIONstarting from a N ∼ N (0, I) to produce a noise-free action a 0 for a given observation s. The denoising function is trained to predict the source noise z 0 ∼ N (0, I) of a noisy action a t = √ ᾱt a 0 + z t by minimizing the lossEQUATIONwhere ᾱt = t j=1 α j . The expectation over t 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 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 µ θ of a conditional Gaussian distribution. Conversely, when using Denoising Diffusion Policies (DDPs), these architectures parameterize the denoising function ϵ θ . Given the straightforward nature of the former case, we focus on introducing these architectures in the context of DDPs.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 ϵ θ that takes a sequence of observations s k-K:k , noisy actions a t k:k+J and diffusion step t to generate a less noisy sequence of actions a t-1 k:k+J . The diffusion step is encoded using time-embedding TE. The observations are encoded using a ResNet-18, with shared weights across images from different time steps. An action encoder ENC a 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 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.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 encoderdecoder 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 Enc m is used to process the time embedding and state embedding, and the Mamba Decoder Dec m is used to process the noise embedding. Since the inputs to the E m and D m have different lengths of sequences, we propose to add learnable variables to complement each sequence. The method is outlined in Alg 2. 
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?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: Decoderonly 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.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 Table 1 : Performance on LIBERO benchmark with 20% data, 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.of robotic learning and manipulation capabilities. The task visualizations are in Figure 3 . More details can be found in Appendix B.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.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 . 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.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 7 8 . The corresponding results are presented in Table 2 -4.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. 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. 
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. 
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 benchmark and real robot experiments demonstrate that MaIL not only matches but also surpasses the performance of existing baselines with limited demonstrations, establishing it as a promising approach for IL tasks. 
 
 Appendix: A. In Table 1 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 experiments on these tasks using language embeddings, with access to 20% and 100% of the data. The results are presented in Tables 7 and 8 • 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.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.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.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.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) ∈ R D → y(l) ∈ R D , that maps x(l) into a hidden state h(l) ∈ R N , which then can be projected onto the output y(t), that is, EQUATIONwhere ∆ is the step size. The model then can be computed as either linear recurrence, that is,EQUATIONor through a global convolution * given by y = x * K, with K = (CB, CAB, . . . , CA k B).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.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 10 .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 9 . We list the training hyperparameters for the Transformer-based and Mamba-based policies in Table 10 . 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.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. 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.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.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.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.To better understand the advantages of MaIL over Transformers, we conducted a detailed analysis of the latent representations produced by both methods. Specifically, we compared the Encoder-Decoder architectures of MaIL and Transformers by visualizing their high-dimensional representations 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, MaIL's state representations are notably more structured across all trajectories compared to those generated by Transformers. More structured representations could potentially improve the generalization of MaIL, which explains the advantages of MaIL over Transformers. 


Decision: 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.

Reviews: review: title: title: Replacing Transformers in IL with Mamba
paper_summary: paper_summary: **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.
main_review: main_review: ## Overview

The paper introduces the Mamba for Imitation Learning (MaIL), presenting a novel approach by incorporating the Mamba model with established IL techniques such as Behavior Cloning and Denoising Diffusion Policies. This integration is explored through two architectural variants: a decoder-only model and an encoder-decoder model. The study provides a comprehensive evaluation of these models through both simulation and real-world experiments. Results demonstrate that the Mamba architecture for Imitation Learning (MaIL) significantly outperforms traditional Transformer-based methods in IL, showcasing its effectiveness and adaptability in both simulated and practical settings.

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.


## Strengths

- The utilization of the Mamba model within the realm of imitation learning, while intuitive, proposes a compelling alternative to the prevailing Transformer-based methodologies.

- Extensive evaluations are conducted utilizing both simulated benchmarks (LIBERO) and real robotic tasks, effectively demonstrating the method's adaptability and efficiency across diverse scenarios.

- MaIL consistently surpasses Transformer-based approaches in various tasks, notably under conditions of limited data availability, a critical factor for practical robotic applications.

- Demonstrating proficiency in managing multi-modal inputs, such as images and language, MaIL proves to be highly effective for a wide range of complex and versatile robotic tasks.

## Weaknesses

- While the empirical evidence underscores MaIL's robust performance against Transformer-based architectures, the underlying reasons of such advantages remain ambiguous. A deeper exploration into why Mamba excels would enhance the understanding of its efficacy.

- A notable strength of Transformer architectures is their ability to manage long-horizon dependencies. The capability of MaIL to handle similar tasks remains uncertain. Testing MaIL on tasks that require extensive long-horizon reasoning could verify its potential in more complex scenarios.

- The Mamba architecture is noted for its speed advantage, particularly when dealing with very long sequences, as per the original Mamba paper. However, in typical imitation learning scenarios where sequence lengths are much shorter, the relevance of this speed benefit is questionable. Clarifying whether Mamba retains its efficiency advantage in standard imitation learning contexts would be valuable.

- MaIL's effectiveness seems to converge with that of Transformers when large datasets are employed, which may diminish its advantages in data-abundant situations. Given that Transformers excel at scaling with large volumes of data [1,2,3], further investigation into MaIL's performance with substantial datasets is warranted.

- The claim that "It is evident that Mamba-based models significantly outperform Transformers when data is scarce and perform comparably as the dataset size increases" is intriguing yet lacks detailed analysis or explanation. Providing a thorough examination or empirical backing for this observation would substantiate the claim and clarify the conditions under which MaIL is most effective.

[1] Brohan, Anthony, et al. "Rt-1: Robotics transformer for real-world control at scale." _arXiv preprint arXiv:2212.06817_ (2022).

[2] Brohan, Anthony, et al. "Rt-2: Vision-language-action models transfer web knowledge to robotic control." _arXiv preprint arXiv:2307.15818_ (2023).

[3] Kim, Moo Jin, et al. "OpenVLA: An Open-Source Vision-Language-Action Model." _arXiv preprint arXiv:2406.09246_ (2024).
questions: questions: - 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?
limitations: limitations: 3
review_summary: review_summary: 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.
title: title: Improving Imitation Learning via Mamba
paper_summary: paper_summary: 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.
main_review: main_review: **Strengths**

* The idea and motivation of the paper are solid. Mamba has recently received significant attention for matching Transformers' performance on some benchmarks with lower memory usage and better speed. It makes sense to study the performance of Mamba-based imitation policies in robotics.

* The authors have proposed both decoder-only and encoder-decoder architectures, similar to Transformers, and evaluated them in both behavior cloning and denoising diffusion policies. The authors propose a novel approach called Mamba Aggregation for the encoder-decoder variant.

* The proposed approach outperforms Transformer-based counterparts on 5 LIBERO suites, some including language modality, when having access to a small portion of datasets (20%).

* They have conducted real-world experiments on three different tasks, and in all of them, MaIL outperforms the baseline.


**Weaknesses**

* Although the paper is well-written and easy to follow, it can be misleading at times. In the abstract, the caption of Table 1, and the conclusion, the authors claim MaIL outperforms Transformers in all LIBERO tasks without clarifying that the evaluations are conducted with all models having access to only 20% of the dataset. 

* The ablation on dataset size is not sufficiently informative. The authors compare their approach with the baseline using the full dataset on only two tasks from the five task suites and only with the BC variants. A comprehensive comparison table, similar to Table 1, using the full dataset across all tasks, would have been more insightful.

* Although one of the main contributions of the paper is the Mamba Aggregation approach in the Encoder-Decoder architecture, the authors do not elaborate much on their inspirations and design choices for this.
questions: questions: * How do Transformer-based baselines compare with your approach when having access to the full LIBERO dataset?
* 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?
limitations: limitations: 3
review_summary: review_summary: 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.
title: title: Novel architecture but results not convincing
paper_summary: paper_summary: 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.
main_review: main_review: Strengths:
The paper is well-written with the contributions explained clearly throughout the paper. The usage of this architecture is novel and can potentially challenge the state-of-the-art Transformer backbones for visuomotor imitation learning.

Weaknesses:
1. I am a bit worried about the quantitative results in this paper. Success rates shown in Table 1 on the LIBERO multi-task benchmark are mostly and roughly in the range 0.3 to 0.5, which I am not convinced is a good success rate on this benchmark, and seemingly at par or even worse than that in the LIBERO’s paper itself. MaIL and the considered transformer architecture’s success rates are especially low on the LIBERO-Long suite of tasks, which does not seem to convey that MaIL seems to solve the problem with processing long sequences of observations efficiently. 
2. The authors mentioned in Section 5.2 that all models were trained for 50 epochs - can the authors clarify what they considered as an epoch, i.e. did it comprise of all possible sequences that can be chunked out of the dataset of demos and if there was any overlap between obs-action sequences the model was trained on. I think it would also be best for reproducibility if the authors mentioned the number of gradient steps the models were trained on.
3. Given that the authors present results having evaluated just one checkpoint, I would like to understand if the model fit/overfit/underfit the data at that checkpoint - something that I believe could be overlooked if multiple checkpoints were evaluated.
4. Why was BC-RNN not considered as a baseline?
questions: questions: Please refer to the weaknesses mentioned above.
limitations: limitations: 3
review_summary: review_summary: 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.
title: title: Good paper
paper_summary: paper_summary: 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.
main_review: main_review: My opinion is positive. The paper explores the utilization of Mamba, a novel state space model which is an alternative to Transformers. This work proposes using Mamba in imitation learning with selective state space models (MaIL). They propose two possible architectures of MaIL: an encoder-only Mamba denoising architecture (D-Ma) and a Mamba encoder-decoder architecture (ED-Ma).

In addition:
- The manuscript is well-organized and written, and the goals and contributions of the authors are clearly stated.
- The research questions they want to answer with their experiments are clearly stated and discussed.
- The figures and algorithms used in the paper complement the text and help to understand the authors' explanations.
- The appendix section includes valuable information regarding the transformer architecture, ablation studies, model, and experimental details.
- The authors clearly explain the limitations of their method when the dataset scales up and the advantage compared to Transformers is less noticeable.

However, I recommend that the authors expand the discussion of the experimental results, especially the cases in which D-Ma is outperformed by ED-Ma in LIBERO-Long and the possible reasons for this. 
Also, the authors did not sufficiently evaluate Mamba's ability to handle long-term contextual information. They only tested H5 conditions.
Considering that Mamba is an SSM, the obtained state space representations are better analyzed qualitatively.
questions: questions: - 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.
- Why is the performance of ED-Ma lower than ED-Tr in stage 2 of the Two-stage Pick-Place and Inserting tasks?
- 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?
- 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%.
- In Figure 4 (b), the label is DDP-H5. Is this correct? Should it not be DDP-H5: object?
- Could you have additional discussion over the application of the proposed method to BC with longer context?
- Could you analyze the representations of the formed state space qualitatively?
limitations: limitations: 3
review_summary: review_summary: 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.