StructPolicy: Robust Imitation Learning Policy Guided by Structure Map

Zehao Du, Jiude Wei, Cewu Lu, Jianhua Sun

Keywords: Imitation Learning, Robot Manipulation

Abstract:

Imitation Learning (IL) offers an effective approach for robot manipulation by learning a mapping function from visual inputs to actions.
However, this paradigm suffers from domain discrepancies and is highly sensitive to distribution shifts, as the mapping function inherent in IL tends to overfit task-irrelevant visual noise, thereby compromising robustness. To address this, we propose StructPolicy, a lightweight module that guides the robot to acquire structural knowledge by constructing a stable and task-relevant Structure Map.
The structure map is composed of task-relevant object structures and reveals the topological cues essential for accomplishing manipulation tasks.
By filtering out visually distracting noise and retaining only the structural attributes, it effectively guides the robot's policy learning toward robust manipulation.
We begin by introducing a general object structure representation based on atomic geometric primitives, enabling flexible composition and scalability to arbitrary objects.
Building on this, we design StructGen, a module that automatically constructs structure maps from visual observations.
Finally, we design StructTransformer, which employs hierarchical attention over structure vertices to extract map features from the structure map for action prediction.
We extensively evaluate StructPolicy across 3 IL models, 3 different simulators, and 50 diverse manipulation tasks.
In addition, through ablation studies under various visual changes, including visual noise, camera viewpoint shifts, light intensity, and background color variations, we demonstrate that StructPolicy consistently improves robustness against distribution shifts.
Results demonstrate consistent and significant performance improvements across all tasks, validating the effectiveness and robustness of StructPolicy. Our code will be released publicly to facilitate reproducibility and further research.

Primary Area: applications to robotics, autonomy, planning

Submission Number: 733