MaIL

Paper, Arxiv, Project Page, CoRL 2024

Xiaogang Jia*¹², Qian Wang*¹ Atalay Donat¹, Bowen Xing¹, Ge LI¹, Hongyi Zhou², Onur Celik¹, Denis Blessing¹, Rudolf Lioutikov², Gerhard Neumann¹

¹Autonomous Learning Robots, Karlsruhe Institute of Technology

²Intuitive Robots Lab, Karlsruhe Institute of Technology

* indicates equal contribution

This project encompasses the MaIL codebase, which includes the implementation of the Decoder-only Mamba and Encoder-Decoder Mamba for BC and DDPM models. Highlights of MaIL:

MaIL achieves better results compared to Transformer-based models on the LIBERO benchmark with 20% data.
MaIL can be used as a standalone policy or as a part of advanced methods like diffusion.
MaIL has a much more structured latent space compared to Transformer-based models.

Installation

To begin, clone this repository locally

git clone git@github.com:ALRhub/MaIL.git

Installing requirements

conda create -n mail python=3.8
conda activate mail

# adapt to your own cuda version if you need
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

pip install -r requirements.txt

Installing LIBERO Setup

git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git
cd LIBERO
pip install -e .

Installing Mamba

pip install mamba-ssm==1.2.0.post1

Usage

File System

MaIL
└── agents # model implementation
    └── models
    ...
└── configs # task configs and model hyper parameters
└── environments
        ...
    └── dataset # data saving folder and data process
        └── data
        ...
└── scripts # running scripts and hyper parameters
    └── aligning
    └── stacking
    ...
└── task_embeddings # language embeddings
└── simulation # task simulation
...

Reproduce the results

Train decoder-only mamba with BC on LIBERO-Spatial and LIBERO-Object tasks using 3 seeds

bash scripts/bc/libero_so_bc_mamba_dec.sh

Train encoder-decoder mamba with DDPM on LIBERO-Spatial and LIBERO-Object tasks using 3 seeds

bash scripts/3seed/libero_so_mamba_encdec.sh

Acknowledgements

The code of this repository relies on the following existing codebases:

[Mamba] https://github.com/state-spaces/mamba
[D3Il] https://github.com/ALRhub/d3il
[LIBERO] https://github.com/Lifelong-Robot-Learning/LIBERO

Citation

If you found the code usefull, please cite our work:

@inproceedings{
jia2024mail,
title={Ma{IL}: Improving Imitation Learning with Selective State Space Models},
author={Xiaogang Jia and Qian Wang and Atalay Donat and Bowen Xing and Ge Li and Hongyi Zhou and Onur Celik and Denis Blessing and Rudolf Lioutikov and Gerhard Neumann},
booktitle={8th Annual Conference on Robot Learning},
year={2024},
url={https://openreview.net/forum?id=IssXUYvVTg}
}