[논문 리뷰] MaIL: Improving Imitation Learning with Mamba

Xiaogang Jia, Qian Wang외 8명|arXiv (Cornell University)|2024. 06. 12.

Tactile and Sensory InteractionsNeuroscience인용 수 3

한 줄 요약

MaIL는 기존 트랜스포머 기반 백본을 선택적 상태공간모델인 Mamba로 대체함으로써 계산 비용을 감소시키면서도 뛰어난 성능을 달성하는 새로운 이민러닝 아키텍처를 제안한다. Mamba를 인코더-디코더 구조에 적응시켜 LIBERO 벤치마크와 실제 로봇 작업에서 트랜스포머 기반 정책보다 뛰어난 성능을 내며, 특히 소규모 데이터셋과 노이즈가 있는 입력에서 뛰어난 성능을 보인다.

ABSTRACT

This work presents Mamba Imitation Learning (MaIL), a novel imitation learning (IL) architecture that provides an alternative to state-of-the-art (SoTA) Transformer-based policies. MaIL leverages Mamba, a state-space model designed to selectively focus on key features of the data. While Transformers are highly effective in data-rich environments due to their dense attention mechanisms, they can struggle with smaller datasets, often leading to overfitting or suboptimal representation learning. In contrast, Mamba's architecture enhances representation 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 MaIL consistently outperforms Transformers on all LIBERO tasks with limited data and matches their performance when the full dataset is available. Additionally, MaIL's effectiveness is validated through its superior performance in three real robot experiments. Our code is available at https://github.com/ALRhub/MaIL.

연구 동기 및 목표

• 이민러닝에서 대규모 트랜스포머 기반 정책의 높은 계산 비용과 학습 난이도 문제를 해결한다.
• 선택적 상태공간모델인 Mamba의 효율성과 장기 시퀀스 모델링 능력을 활용하여 학습 및 추론의 효율성을 향상시킨다.
• 다중모달 입력을 위한 인코더-디코더 아키텍처로 Mamba를 적응시켜 이민러닝에서 효과적인 시퀀스 모델링을 가능하게 한다.
• 특히 데이터가 부족하고 노이즈가 많은 조건에서 Mamba 기반 정책이 트랜스포머를 초월할 수 있음을 입증한다.
• 실제 벤치마크 환경과 실제 로봇 배포를 통해 제안된 아키텍처의 실용성을 검증한다.

제안 방법

• 이민러닝의 백본으로 Mamba를 채택하여 트랜스포머 기반 정책을 대체함으로써 계산 효율성을 향상시킨다.
• 행동, 상태, 시간 단계에 대한 학습 가능한 임bedding을 도입하여 Mamba를 인코더-디코더 구조로 확장할 수 있는 새로운 수식 체계를 설계한다.
• Mamba 기반 정책을 디노이징 헤드로 삼아 디퓨전 과정에 통합함으로써 노이즈가 섞인 관측값으로부터 정제된 행동을 예측할 수 있도록 한다.
• 관측 특징, 시간 임베딩, 행동 시퀀스를 인과적으로 처리하는 시퀀스 모델링 목적함수를 사용해 모델을 훈련한다.
• 이미지와 언어 임베딩 등의 다중모달 입력을 Mamba 처리 전에 입력 시퀀스에 융합함으로써 처리 능력을 향상시킨다.
• 독립형 정책 및 디퓨전 기반 설정 양쪽 모두에서 모델을 적용함으로써 배포 및 통합의 유연성을 확보한다.

Figure 1: D-Ma: Mamba denoising architecture integrates ResNet-18 for state encoding and an action encoder for action encoding. The state sequence has a length of $K$ , while the action sequence at diffusion step $t$ has a length of $J$ . Before feeding the inputs into the Mamba module, positional e

실험 결과

연구 질문

• RQ1 Mamba 기반 모델은 추론 효율성을 유지하거나 향상시키면서도 이민러닝에서 트랜스포머 기반 정책보다 뛰어난 성능을 낼 수 있는가?
• RQ2 제안된 Mamba의 인코더-디코더 적응이 시각 및 언어와 같은 다중모달 관측 입력에 대해 얼마나 잘 일반화되는가?
• RQ3 데이터가 부족하고 노이즈가 많은 조건에서 Mamba 기반 정책은 트랜스포머보다 뛰어난 성능을 유지하는가?
• RQ4 Mamba 기반 정책은 디노이징 디퓨전 정책 같은 디퓨전 기반 정책 학습 프레임워크에 효과적으로 통합될 수 있는가?
• RQ5 다양한 데이터셋 크기에서 학습 및 추론 비용을 고려할 때 Mamba와 트랜스포머 간의 성능 격차는 어떠한가?

주요 결과

• MaIL은 LIBERO 벤치마크의 모든 15개 작업에서 트랜스포머 기반 정책를 초월하여, 픽플레이스, 이중단계 픽플레이스, 삽입, 컵스택킹 작업에서 높은 성공률을 기록했다.
• 픽플레이스 작업에서 MaIL은 바나나 배치에 대해 0.55의 성공률을 기록했으며, 트랜스포머 기준 0.45를 초월했고, 당근 배치에선 0.70의 성공률을 기록해 트랜스포머의 0.25를 상회했다.
• 이중단계 픽플레이스 작업에서 MaIL은 0.45의 성공률을 기록했고, 트랜스포머의 0.20을 상회하여 복잡한 순차적 조작 작업에서 뛰어난 성능을 보였다.
• 삽입 작업에서 MaIL은 1단계에 0.55의 성공률, 2단계에 0.20의 성공률을 기록했으며, 트랜스포머의 0.40과 0.30을 각각 초월했다.
• 컵스택킹 작업에서 MaIL은 1단계에 0.80의 성공률, 2단계에 0.55의 성공률을 기록해 트랜스포머의 0.60과 0.40을 각각 상회했다.
• MaIL은 입력 노이즈에 대해 강인하며, 예상치 못한 물체 방향성에 대해서도 일반화 능력을 보였으며, 특히 트랜스포머가 성능을 떨어뜨리는 저데이터 환경에서 뛰어난 성능을 발휘했다.

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 a