FIT: Future-Aware Imitation Transformer for Long-Horizon Robotic Manipulation

Abstract

Developing a generalist agent capable of performing multiple long-horizon tasks and solving novel ones remains an open challenge in robotic learning. Current imitation learning methods show promise in handling individual long-horizon tasks, but they struggle to scale to multitask scenarios and fail to execute new tasks. This work proposes a framework that learns a unified imitation policy capable of generalizing to diverse interaction-rich tasks and environmental variations. Our approach introduces a task decomposition strategy that leverages a large language model (LLM) and robot proprioception to infer agent–object interaction events from expert demonstrations. With these interaction events as input, we introduce the future-aware imitation transformer (FIT), a transformer-based policy that acquires interaction-aware skills. At every time step, FIT predicts the actions, the latent representations of the future states to anticipate the outcomes of actions, and a task embedding to maintain awareness of the long-horizon manipulation objective. We evaluated our method in simulated LIBERO, Franka kitchen, and real-world xARM6 environments. Experimental results demonstrate that our method outperforms all baselines, particularly on novel tasks, where existing methods fail to generalize.