HGDpolicy: Hyperbolic Geometry-Guided Dual-Space Diffusion Policy for Few-shot Contextual Imitation Learning in Robot Manipulation

Robotic imitation learning aims to rapidly adapt to novel manipulation tasks from only a few human demonstrations. Existing contextual imitation learning methods typically organize demonstrations, observations, and actions as graphs and employ diffusion models for action generation. However, sparse demonstrations require the policy to infer latent subtask relations and transition patterns beyond local state-action correspondences, while Euclidean graph representations are limited in capturing such hierarchical structures. Hyperbolic geometry offers a natural alternative, yet its integration with graph diffusion remains challenging: hierarchical task information must be reliably extracted from few demonstrations and effectively transferred from non-Euclidean representations to anisotropic action generation in Euclidean space. To address these challenges, we propose HGDPolicy, a hyperbolic geometry-guided dual-space diffusion policy for contextual robotic imitation learning. HGDPolicy first introduces transition-aware hyperbolic graph modeling to capture task hierarchies and state-transition patterns from sparse demonstrations. It further develops feasible-action geometric perception, which encodes the learned hyperbolic task structure as an action diffusion prior and dynamically adjusts the center and uncertainty of feasible action distributions according to the current task context. In this way, hierarchical task structure directly modulates the action diffusion process, rather than merely serving as an auxiliary representation. Experiments on nine RLBench generalization tasks show that HGDPolicy achieves a higher average task success rate than the compared baselines, reaching 59.4%, compared with 51.9% for the best-performing baseline. Real-world robotic experiments further demonstrate rapid adaptation to novel tasks with only a few demonstrations and generalization to unseen objects without additional data collection.