MTIL: Encoding Full History with Mamba for Temporal Imitation Learning

TL;DR

Imitation learning for long-horizon tasks is hampered by the Markov assumption and the quadratic cost of attention-based architectures when processing full observation histories. MTIL leverages the Mamba-2 State Space Model to compress the entire observation sequence into a recurrent state $h_t$ and conditions the policy on $(o_t,h_t)$ to predict a chunk of actions $\hat{a}_{t:t+K-1}$. The authors provide a theoretical framing of MTIL as learning an implicit dynamical system, a history-driven world model, and validate it across ACT, LIBERO, Robomimic, and real-world dual-arm tasks, where MTIL significantly outperforms SOTA baselines. The results demonstrate that full temporal context is essential for disambiguating perceptual ambiguities and that the proposed architecture yields computationally feasible long-horizon imitation on high-dimensional observations.

Abstract

Standard imitation learning (IL) methods have achieved considerable success in robotics, yet often rely on the Markov assumption, which falters in long-horizon tasks where history is crucial for resolving perceptual ambiguity. This limitation stems not only from a conceptual gap but also from a fundamental computational barrier: prevailing architectures like Transformers are often constrained by quadratic complexity, rendering the processing of long, high-dimensional observation sequences infeasible. To overcome this dual challenge, we introduce Mamba Temporal Imitation Learning (MTIL). Our approach represents a new paradigm for robotic learning, which we frame as a practical synthesis of World Model and Dynamical System concepts. By leveraging the linear-time recurrent dynamics of State Space Models (SSMs), MTIL learns an implicit, action-oriented world model that efficiently encodes the entire trajectory history into a compressed, evolving state. This allows the policy to be conditioned on a comprehensive temporal context, transcending the confines of Markovian approaches. Through extensive experiments on simulated benchmarks (ACT, Robomimic, LIBERO) and on challenging real-world tasks, MTIL demonstrates superior performance against SOTA methods like ACT and Diffusion Policy, particularly in resolving long-term temporal ambiguities. Our findings not only affirm the necessity of full temporal context but also validate MTIL as a powerful and a computationally feasible approach for learning long-horizon, non-Markovian behaviors from high-dimensional observations.

Figures (7)

• Figure 1: Overview of the Mamba Temporal Imitation Learning (MTIL) architecture. Multi-modal inputs (images via DINOv2, state) are fused and processed by sequential Mamba-2 blocks, updating the recurrent state $h_t$ which encodes history. At each step $t$ across the entire trajectory, MTIL predicts an action chunk $\hat{a}_{t:t+K-1}$ (current plus $K-1$ future steps). This is supervised via L2 loss against ground truth actions $a_{t:t+K-1}$ from the demonstration (using last action for padding when near trajectory end). The historical context embedded in $h_t$ enables temporally coherent, long-horizon action generation.
• Figure 3: Dual UR3 experimental setup with four cameras (Top: Kinect; Side: D435i; Wrists: D405) and custom grippers.
• Figure 4: Sequential Insertion Success Rates (%), averaged over 50 roll-outs.
• Figure 5: Coordinated Pouring Success Rates (%), averaged over 50 roll-outs.
• Figure : (a) Learning Curves