FoAM: Foresight-Augmented Multi-Task Imitation Policy for Robotic Manipulation

Abstract: Multi-task imitation learning (MTIL) has shown significant potential in robotic manipulation by enabling agents to perform various tasks using a single policy. This simplifies the policy deployment and enhances the agent's adaptability across different scenarios. However, key challenges remain, such as maintaining action reliability (e.g., avoiding abnormal action sequences that deviate from nominal task trajectories) and generalizing to unseen tasks with a few expert demonstrations. To address these challenges, we introduce the Foresight-Augmented Manipulation Policy (FoAM), a novel MTIL policy that pioneers the use of multi-modal goal condition as input and introduces a foresight augmentation in addition to the general action reconstruction. FoAM enables the agent to reason about the visual consequences (states) of its actions and learn more expressive embedding that captures nuanced task variations. Extensive experiments on over 100 tasks in simulation and real-world settings demonstrate that FoAM significantly enhances MTIL policy performance, outperforming state-of-the-art baselines by up to 41% in success rate. Meanwhile, we released our simulation suites, including a total of 10 scenarios and over 80 challenging tasks designed for manipulation policy training and evaluation. See the project homepage projFoAM.github.io for project details.