Mamba as a motion encoder for robotic imitation learning
Toshiaki Tsuji







Mamba as a motion encoder
 for robotic imitation learning

    Toshiaki Tsuji^1

^1Toshiaki Tsuji is with the 
Dept. Electrical Engineering, Electronics and Applied Physics, Saitama University, 255 Shimo-okubo, Saitama, 338-8570 
E-mail: tsuji@ees.saitama-u.ac.jp

    September 9, 2024
=============================================================================================================================================================================================================





§ ABSTRACT 
Recent advancements in imitation learning, particularly with the integration of LLM techniques, 
are set to significantly improve robots' dexterity and adaptability. 
In this study, we propose using Mamba, a state-of-the-art architecture with potential applications 
in LLMs, for robotic imitation learning, highlighting its ability to function as an encoder that 
effectively captures contextual information.
By reducing the dimensionality of the state space, Mamba operates similarly to an autoencoder.  
It effectively compresses the sequential information into state variables while preserving the essential 
temporal dynamics necessary for accurate motion prediction.
Experimental results in tasks such as Cup Placing and Case Loading demonstrate that 
despite exhibiting higher estimation errors, Mamba achieves superior success rates compared to 
Transformers in practical task execution. 
This performance is attributed to Mamba’s structure, which encompasses the state space model. 
Additionally, the study investigates Mamba’s capacity to serve as a real-time motion generator 
with a limited amount of training data.




§ INTRODUCTION

In the field of robotic manipulation, the application of imitation learning has been rapidly advancing in 
recent years <cit.>.
Conventional robot control methods relying on manual coding necessitate complex motion planning 
as well as accurate modeling of the environment and objects,
imitation learning offers the advantage of intuitively and efficiently teaching tasks to robots.
Furthermore, with the recent advancements in deep learning techniques, it is becoming possible to acquire 
generalized control policies adaptable to various environments in robotics.
So far, various architectures have been proposed, such as behavior cloning <cit.>,
inverse reinforcement learning <cit.>, and GAIL <cit.>.
There is an increase in multimodal imitation learning that extends modalities <cit.>
and models representing complex tasks using hierarchical architecture <cit.> has a potential to extend imitation learning to complicated tasks. 
Integration of force control and imitation learning is essential for enhancing performance of contact rich 
tasks <cit.>. 
Depending on the development of these imitation learning studies, 
there is potential for robots to attain human-level dexterity 
and adaptability even in tasks that have been considered difficult to automate.

This research aims to verify the significance of using deep state-space 
models (SSMs) <cit.>, particularly Mamba <cit.>,

as a model for robotic imitation learning.
The rapid development of LLMs has not only impacted language tasks but also influenced fields 
such as image and speech processing, drawing attention to their potential applications in robotics <cit.>.
In robotic imitation learning, LLMs are expected to contribute to mapping language instructions to physical 
actions <cit.>. 
Not only that, but they have also been applied to action generation and path planning due to their versatile knowledge 
representation and multimodal integration capabilities <cit.>.  
While many of these studies are based on Transformers with attention mechanisms,
Mamba is also a good candidate as a model that achieves equivalent long-term memory retention. 
As Mamba is based on a continuous SSM, it could exhibit better characteristics than Transformers
as a model for robotic imitation learning, which deals with sequences of continuous sensor-actuator responses. 
Addressing the challenges of limited training data and the need for quick responses, 
this research develops a model with simple configuration suitable for real-time motion generation in robots. 
We then examine how much capability this simplified model alone possesses as an IL motion generation model.
Previous studies have been designed on the premise of incorporating long sequences of states entirely into the model,
but to achieve both real-time responsiveness and adaptability to diverse environments, it is desirable for designers to be
able to handle and process past context and current state information separately.
Additionally, for data storage and real-time processing, it is preferable that each piece of information is appropriately compressed.
Therefore, focusing on Mamba's ability to be applied as an encoder that stores context as low-dimensional information,
we extend it into a model that assigns past contextual information to state variables.




§ RELATED STUDIES




  §.§.§ Imitation Learning

Early imitation learning methods primarily focused on learning one-to-one mappings between state-action pairs.
Despite showing promising results in tasks such as autonomous driving <cit.> and robot control <cit.>,
these methods overlooked the rich temporal information contained in histories.
Therefore, numerous studies have successfully incorporated action sequences using RNNs <cit.>.
LSTMs <cit.> have been widely applied in many tasks where long-term dependencies need to be 
considered.
Seq2Seq models adopting LSTM as both encoder and decoder allow for optimization of generated trajectories in latent space,
enabling dimensionality reduction of optimization problems, as their latent variables compress sequence information <cit.>.
However, limitations in memory retention have been pointed out as potentially problematic when applied to very long sequences <cit.>.
Transformers are widely applied in recent imitation learning due to their ability to model longer sequences
than LSTM while maintaining training efficiency through parallelization of sequence processing.
Having developed in language and image processing, multimodal imitation learning using Transformers to encode
both image and language sequences is being researched <cit.> and there are also 
studies in robotic applications, too <cit.>.
As mentioned above, the Transformer is a powerful tool for robot imitation learning due to its high memory retention performance. 
We propose a model that has the same level of memory retention as the Transformer and is suitable for improving robot task performance.





  §.§.§ SSM

SSMs are also promising candidates to solve the issue in robotic imitation learning.
SSMs have long been applied to time series information prediction <cit.>.
Recently, there have been many attempts to combine SSMs with deep learning.
Watter et al. proposed an SSM combining locally linear state transition models
with nonlinear observation models using VAE <cit.>.
Gu et al. proposed HIPPO, a model applying SSMs to discrete neural networks <cit.>,
and extended it to S4 <cit.> and Mamba <cit.>.
These can achieve efficiency and performance surpassing Transformers in processing long-range dependencies.
It has long been known that CNNs are effective in extracting features necessary for motion generation from environmental information,
and as Mamba is reported to possess these characteristics, it is promising as an IL policy for robots,
with examples of application to robot imitation learning already emerging <cit.>.
Jia et al. demonstrated a method to extend Mamba to an encoder-decoder variant by augmenting input with learnable action,
state, and time embedding variables, enabling efficient learning and applying it to diffusion processes <cit.>.
Previous researches applying Mamba to robots have all been designed assuming complex contexts with large amounts of data,
resulting in large network structures with stacked layers of residual mechanisms.
This paper focuses on the fact that the context of robot manipulation tasks can be often represented by smaller number of words
compared to general sentences, and reduces the dimensionality of the state space. Considering learning with a small amount of
training data, we reduce the number of Mamba layers. We construct a model with such a low-dimensional simplified configuration
and examine how much capability it possesses as an IL motion generation model on its own.
Additionally, the dimensionality reduction of the state space means that the amount of contextual information can be reduced.
Utilizing this feature, we improve the model to generate actions using only the compressed contextual information and recent
state sequences.



  §.§.§ Autoencoder

The model adopted in this research can be interpreted as utilizing Mamba as an autoencoder, given its low-dimensional SSM configuration. 
This approach bears relevance to previous works employing autoencoders in robotics, such as <cit.>. 
Some of these studies also have architectures that provide inputs to latent variables <cit.>. 
The most significant difference of our study to these is the CNN and the selection mechanism. 
These features are expected to enable selective extraction of crucial information while mitigating performance degradation in long sequences.

The model adopted in this research can be interpreted as using Mamba as an autoencoder, as it sets the dimensions of Mamba's
SSM low. Therefore, it has relevance to previous works using autoencoders in robotics 
such as .  
Some of them also have an architecture to give inputs to latent variables
In our study, the most significant difference lies in having CNN and a selection mechanism.
These characteristics are expected to enable selective extraction of only important information and suppress performance
degradation in long sequences.



§ PROPOSED METHOD



 §.§ Architecture of State Space Model

The machine learning model implemented in this research is shown in Fig.<ref>(a). 
Here, x_t and y_t are input and output vectors at time t, respectively, and T is length of the sequence. 
Mamba traditionally inputs sequence information vectorized by a tokenizer into multiple overlapping Mamba layers, and then converts it back to words for output. In this research, the tokenizer is replaced with a linear layer because continuous 
data is given for input.
Additionally, the input dimension to the Mamba block is reduced by this linear layer. The computation results of the low-dimensional Mamba block are then increased to the same dimension as the input through a linear layer for output.
To predict future trajectories, the output time is advanced compared to the input. Also, to avoid performance degradation due to padding in the Convolution layer, the input layer is made longer by the kernel length. While not strictly identical, the input and output are nearly equivalent to an autoencoder structure.
The configuration of the Mamba block is shown in Fig. <ref>. 

It features a residual network with skip connections that perform identity mapping 
(see the bottom arrow in Fig.<ref>), similar to the model in <cit.>, 
but modified for application in robot trajectory generation. 
Previous studies applying Mamba to robots <cit.> have made Mamba blocks 
handle features enhanced by encoders or signal processing processes on input and output data. 
In contrast, this research adopts a simpler model configuration consisting solely of Mamba blocks.

The input consists of 16-axis time series data: 
8-axis angles and 8-axis torques corresponding to joint motors of the robot. 


To simplify the model, the 3 degrees of freedom of posture are assumed to be fixed in this study. 
The internal structure of the residual network consists of convolutional layers, SSM, and FNN layers. 
The state variable h_t and output y_t of SSM at time t are derived using the 
following equations.

    {[ h_t = A̅h_t-1 + B̅x_t;            y_t = Ch_t ].





    A̅   =   exp (ΔA), 
    B̅   =    (Δ A)^-1(exp (Δ A) - I) ·Δ B.

In these equations, A, B, and C represent the state transition matrix, input matrix, 
and output matrix, respectively, with the superscript bars indicating the parameters of the discretized equations. 
x is the input to SSM, and the subscript t denotes the time step. 
The parameters B, C, and Δ are variable and adjusted by an FNN through a Selection Mechanism. 







 §.§ Design of the proposed method

The configuration of SSM is based on <cit.>, while the dimensionality of the SSM is reduced 
based on the assumption that the variation in robot movements is limited compared to the highly diverse patterns in natural language. 
In our study, the SSM dimensionality was set to 4, achieving characteristics similar to an autoencoder.
In addition, this study shows that the performance of motion generation is sufficient even in an extremely simplified configuration. 
The number of input and output linear layers n_o, n_i and the number of Mamba blocks n  are 1, 3 and 1, respectively. 

While Mamba implements functionality akin to the Attention mechanism 
by incorporating a gated MLP (gMLP) within its block,  
this study excludes gMLP due to its tendency to cause overfitting and unstable learning. 
Most of the continuous-time theories in control and signal processing are based on the assumption that all past responses influence 
the present, so this modification is not unreasonable for imitation learning in robotics. 






The residual network outputs the estimated values of 16-axis position and force data N samples ahead. 
In this study, the angles and torques of each joint are divided into 8 axes each. 
During training, the time series data of position and force are inputted for each sample, 
and the estimated position and force data are compared with the post-sample training data to obtain the error. 
The loss is calculated from this error, and the model is trained through backpropagation.

Mamba has convolutional layers close to the input, and SSM recursively uses the low-dimensional information of 
past histories as input for the next step, thus exhibiting features of both CNN and RNN. 
Also, previous studies have provided the entire input sequence at once, while the model in this paper, 
as shown in Fig. <ref>(b), is structured to provide the input sequence only for a certain period. 
Here, w is the length of the short period input. 
State variables generated in past steps are input to reflect the context before the input period in motion generation. 
Since state variables are low-dimensional compared to input, the amount of information to be preserved is reduced.

It is often a problem in RNNs that their performance degrades in tasks that require consideration of long-term dependencies. 
This challenge can be overcome by setting the matrix A appropriately. 
As shown in Eq. (<ref>), matrix A is the main factor to determine the time constant of memory retention. 
When the matrix A is learned as a neural network parameter, the time constant of memory retention 
is learned to minimize the error function. 
If designers want to assign the memory time constant arbitrarily, matrix A can be given as a fixed value. 
Although there is a trade-off between the length of memory retention and 
the frequency bands of extractable features, 
it is possible to balance both by comprehensively setting high and low time constants 
since the eigenvalues of matrix A can be set multiple times.



§ EVALUATION




 §.§ Robot Configuration

Training data was collected through bilateral control of an 8-degree-of-freedom robot, as shown in Fig. <ref>.  
The performance of imitation learning was evaluated using the follower robot. 
The robot employed in this study was the CRANE-X7 by RT Corporation, which features a 7-axis manipulator 
with a 1-axis gripper at the end. A single control PC was used to calculate inputs for both robots, performing 
bilateral control for each axis in joint space 
The control PC communicated with a machine learning PC equipped with a GPU via UDP, with operating cycles 
of 100 ms and 2 ms, respectively. 
The kernel size of the convolutional layers was set to 4, the loss function used was MSE, and the learning rate was set to 0.001. To prevent overfitting, training was halted when the loss function appeared to start increasing. 





 §.§ Evaluation Method

This paper presents offline and online evaluations of motion imitation learning, respectively.
The offline test aimed to quantitatively evaluate the model's trajectory prediction performance considering long-term dependencies. 
Motions were generated using a model trained with training data obtained from the experimental system, and errors in the predicted 
output were evaluated when test data from the experimental system was given as input.
We evaluated the model on three tasks of increasing complexity: simple up-and-down movements, cup placement after grasping, and case loading.
The up-and-down motion was first repeated for 10 seconds.  
Next, the task was repeated twice and then stopped.
In the former trial, the next trajectory can be predicted based only on the short-term memory, while in the latter trial, long-term memory 
is required to predict the stopping after repeated movements.
In addition, we tested cup placing and case loading as tasks that simulate actual work.
In the former task, the robot grasped an empty 7-cm-diameter, 9-cm-high cup placed in a fixed position and evaluated whether 
the cup was placed in a specified 12-cm-square area without tipping it over.
The latter evaluated whether the subject could grasp a paper box 21 cm long, 3 cm wide, and 8 cm high, and place it in an upright position in a 22-cm-width case.
Since the clearances between the case and the box is 1 cm, the box often contacts the case during training, requiring the instructor to adjust the position accordingly.
Therefore, the task has training data with a higher variability of motion.
The robot in this study was not equipped with a vision system and reproduced a situation in which the position of the object was almost 
identified by placing the object in position with an error of 5 mm or less.
Hyperparameters of LSTM and Transformer were determined with reference to the studies of 
 <cit.>, respectively, which have the closest control system configuration.
For all the tasks, the number of trials for training and test data were 18 and 6, respectively.


 §.§ Offline Evaluation

The errors of each model are compared in the Table <ref>.
Here, in order to have the same evaluation as that used for the error function, the error is based on the integration of information in different units, angle (rad) and torque (Nm).
The Transformer model with full period input consistently demonstrated superior performance across all tasks, 
whereas the Transformer with short period input showed performance degradation due to the unavailability of past history information.
This is supported by the fact that the degradation of the short period input transformer is relatively low for repetitive up-and-down movements, which do not require historical information.
Although the RMSE of Mamba was not as small as that of the Transformer with full period input, it was confirmed that a similar performance can be obtained with only short period input and contextual state variables.



Next, we examined the differences when the diagonal elements of matrix  A were learned by a neural network versus when they were fixed.
Fig. <ref> shows the time responses of the state variable h_t for Mamba with w=20 under different A matrices. 
To confirm whether feature values can be obtained from state variables as work progresses, arrows were added to indicate timings of work progression such as posture alignment, approaching the cup, grasping, and releasing.
In all examples, the diagonal elements were given values multiplied by 0.4 from the maximum value to ensure uniform variation between axes of A.
For the autoencoder to perform sufficiently well, it is desirable for the SSM, which corresponds to latent variables, to be dispersed throughout the space. Also, to extract task features with high performance, it is desirable for state variables to fluctuate in conjunction with characteristic input signals.
When the elements of A were small, it was confirmed that the SSM response values varied greatly over time. However, there was significant bias in the spatial distribution, and it was confirmed that the slope of the state variables was almost uniform, with changes dependent on input being small compared to the overall response values.
When the elements of A were large, the SSM values did not change significantly depending on time. In the results where the maximum values of A elements were 0.5 and 2.5,  
the state variables were widely distributed in space and changed as the work progressed.
In other words, to extract low-dimensional features from input information into the state space, the elements of matrix  
A need to be given appropriately to achieve suitable time constants.
Fig. <ref> compares the RMSE when the diagonal elements of matrix  A are given as fixed values.
When matrix  A is too large or too small, the prediction accuracy of robot motion deteriorates, but when the elements of 
 A are set within an appropriate range, high prediction accuracy is obtained.
The red dotted line represents the average prediction accuracy of 0.451±0.216 when the diagonal elements of matrix 
A were learned as parameters of the neural network. When appropriate fixed values were given, the prediction accuracy 
was better compared to when adjusted by machine learning.
This result suggests the advantage of giving matrix A as fixed values rather than learning it with a neural network, 
especially when the amount of training data is limited, because it results in a model with lower degrees of freedom.








 §.§ On-line evaluation

Table 2 shows the success rates when each task was performed 20 times on the actual machine.

To mitigate the substantial oscillations observed in the Transformer's output
a first-order LPF (low-pass filter) with a time constant of 0.3 seconds 
was applied to the output to suppress vibrations, and the filtered values were given to the robot as command values to improve 
the task success rate. As a result of this adjustment, the Transformer with full period input showed the smallest output RMSE and demonstrated 
high success rates in actual machine tasks. However, a certain percentage of failures occurred in both Cup Placing and Case Loading tasks. 
The main factors for failure were dropping of grasped objects due to vibration and unexpected contact with the environment 
that was not present during training. These issues were more pronounced in the Transformer with short period input, resulting 
in further decreased success rates. 

On the other hand, when performing tasks using Mamba, there was almost no vibration, resulting in a 100% success rate for 
Cup Placing, which has a large tolerance for position deviations. In Case Loading, there were two instances where improper contact 
with the case prevented accurate placement in the desired location, but the overall success rate was higher than that of the Transformer. 
This is thought to be due to the minimal vibration in Mamba's output, despite its RMSE being larger than the Transformer's.
Transformers are known to generally have weak inductive bias, which often brings positive effects. 
However, in robot motion generation, this characteristic does not guarantee output continuity and tends to cause vibrations. 
The reason why Mamba does not generate vibrations is that its SSM has a structure that integrates state variables, 
incorporating the regularity of continuous output changes into the model. 
This functions as an inductive bias, enabling efficient learning from data. 
In robot trajectory generation, continuity based on dynamic constraints is crucial, so the inductive bias provided by SSMs 
like Mamba can have a positive effect on robot imitation learning.

Furthermore, to evaluate real-time motion generation capabilities, Mamba's computation time was reduced to 25ms for experiment. 
As a result of the sampling time being reduced to one-fourth of the original, motions were generated four times faster. 
Although this acceleration increased vibrations due to dynamics inconsistencies and raised the failure rate, 
the overall decrease in success rate was minimal, confirming Mamba's real-time adaptability and stability.








§ LIMITATION

The model proposed in this research generates motion based on state variables and 
short-term sequences rather than using all sequences as input. 
While this approach effectively reduces the network structure and the scale of data 
that needs to be retained, it may impact the Selection Mechanism. 
When all sequences are inputted, the Selection Mechanism captures long-term dependencies 
across all inputs and filters out unnecessary information. 
In the proposed model, information is filtered based on the values of state variables given 
as past context and the input from a limited period. 
The experimental results in this paper suggest that this does not hinder task performance; 
however, the impact may vary depending on the task being applied, necessitating further investigation.

This study presents a modified version of Mamba, excluding gMLP, which has demonstrated high performance in experiment. 
However, it remains unclear to what extent the exclusion 
of gMLP offers advantages. The tasks addressed in this research are relatively simple and can be realized 
within a context similar to System 1 of dual-process theory, which requires a less complicated reasoning. 
For tasks requiring more advanced reasoning, akin to System 2, it is likely that gMLP could be necessary. 
Therefore, it is essential to investigate the relationship between the amount of training data, the anticipated 
modalities, the tasks to be imitated, and the appropriate model design, for developing design principles. 



§ CONCLUSION

This study demonstrates that Mamba offers significant 
advantages for robotic imitation learning. 

By reducing the dimensionality of the state space, Mamba functions effectively as an autoencoder, 
compressing past contextual information into low-dimensional latent variables. 
This characteristic enables efficient learning and robust performance in motion generation.

The experiments revealed that, despite a higher RMSE compared to Transformers, 
Mamba consistently achieved higher success rates in robotic tasks, owing to its ability to suppress 
vibrations and maintain continuity in output trajectories. This highlights the importance of inductive 
bias in robotic systems, where the integration of state variables within a low-dimensional SSM 
framework contributes to more reliable and stable task execution.

Furthermore, Mamba's ability to generate real-time motions with reduced computation time underscores 
its potential for deployment in dynamic environments where rapid and adaptive responses are critical. 
These findings suggest that SSMs like Mamba, with their built-in structural advantages, 
are well-suited for advancing the capabilities of robot imitation learning, providing a strong foundation for future developments in this field.





 
IEEEtran 
