One-Shot Imitation Learning with Invariance Matching for Robotic Manipulation

TL;DR

This work tackles one-shot generalization in robotic manipulation by proposing IMOP, a framework that transfers actions via matching invariant regions between a single demonstration and test scenes. IMOP replaces direct end-effector pose regression with a two-step process: predict invariant regions using a graph-attention network and solve a least-squares pose transfer to compute the test-state action pose, guided by correspondences. The approach yields state-of-the-art performance on 18 RLBench base tasks and strong generalization to 22 novel tasks without fine-tuning, including robustness to large shape variations and one-shot sim-to-real transfer. The results demonstrate that leveraging invariant geometry across scenes can substantially improve data efficiency and generalization in manipulation, with practical implications for robot autonomy in unstructured environments.

Abstract

Learning a single universal policy that can perform a diverse set of manipulation tasks is a promising new direction in robotics. However, existing techniques are limited to learning policies that can only perform tasks that are encountered during training, and require a large number of demonstrations to learn new tasks. Humans, on the other hand, often can learn a new task from a single unannotated demonstration. In this work, we propose the Invariance-Matching One-shot Policy Learning (IMOP) algorithm. In contrast to the standard practice of learning the end-effector's pose directly, IMOP first learns invariant regions of the state space for a given task, and then computes the end-effector's pose through matching the invariant regions between demonstrations and test scenes. Trained on the 18 RLBench tasks, IMOP achieves a success rate that outperforms the state-of-the-art consistently, by 4.5% on average over the 18 tasks. More importantly, IMOP can learn a novel task from a single unannotated demonstration, and without any fine-tuning, and achieves an average success rate improvement of $11.5\%$ over the state-of-the-art on 22 novel tasks selected across nine categories. IMOP can also generalize to new shapes and learn to manipulate objects that are different from those in the demonstration. Further, IMOP can perform one-shot sim-to-real transfer using a single real-robot demonstration.

Figures (8)

• Figure 1: Example of a pick-and-pour task executed by a real Kuka robot after observing a single demonstration and without any re-training. The correspondences between estimated invariant regions of the 3D point cloud in the demonstration and in the test scene are visualized in red lines. The invariant regions are predicted and matched by a neural network, trained offline.
• Figure 2: Examples of tasks, and object-level generalization in various trials. After training on base tasks, IMOP is evaluated on novel tasks that are substantially different from the base tasks. Every learned task is evaluated in multiple trials, each with different object layouts and orientations. For each novel task, only one recorded trajectory is given as a demonstration. We summarize the performance of IMOP on both base and novel tasks in Section \ref{['sec:exp-base']} and \ref{['sec:exp-novel']}. Furthermore, we study the ability of IMOP to generalize to new objects beyond those in the demonstrations in Section \ref{['sec:exp-beyond-demo']}.
• Figure 3: Overview of the proposed invariant region matching network and correspondence-based pose regression. Given the observed transition of the support state $s_i$, we build KNN graphs from the scene point cloud and apply graph self-attention (within $s_i$ and $s_i'$) and cross-attention (between $s_i$ and $s_i'$) layers. The invariant region $\mathcal{I}(s_i)$ is predicted as the set of activated points through a point-wise sigmoid over the support state $s_i$. Next, we apply graph cross-attention layers between the KNN graphs of $\mathcal{I}(s_i)$ and query state $s_j$. The correspondence matrix $C$ is predicted by matching the point-wise features $h_{\mathcal{I}(s_i)}$ and $h_{s_j}$. The action pose $T_j$ of query state $s_j$ is analytically solved from the action pose $T_i$ of support state $s_i$, correspondence matrix $C$, $P_{\mathcal{I}(s_i)}$ and $P_{s_j}$ (points in $\mathcal{I}(s_i)$ and $s_j$), as detailed in Section \ref{['sec:corr-pose-reg']}.
• Figure 4: Overview of the proposed state routing network. Given the query state $s_j$ and one-shot demonstration trajectory $\tau$, we extract the scene-level features for $s_j$ and each state in $\tau$. The scene features are concatenated with the corresponding low-dimensional robot states, and applied with cross-attention layers. The state with the strongest attention to the query state $s_j$ is selected as the support state $s_i$.
• Figure 5: Visualization of invariant regions $\mathcal{I}(s_i)$ as estimated by our invariant region matching network on selected base and novel tasks. The invariant regions on $s_i$ are highlighted with green masks and yellow bounding boxes. The end-effector position is marked with red dot. $\mathcal{I}(s_i)$ generally covers the target object of the action that is applied by the robot to transition from state $s_i$ to next state $s'_{i}$. For example, the invariant region of 'turn the lamp on' covers the 'lamp switch' area. The invariant region of 'pull the rope straight' covers one side of the rope and the rope target. It is worth noting that IMOP only takes point cloud as input and has no access to segmentation masks.

Theorems & Definitions (1)

• Definition 3.1