Few-Shot In-Context Imitation Learning via Implicit Graph Alignment
Authors: Vitalis Vosylius
Edward Johns
Venue: CoRL 2023 Poster (2023)
Keywords: Few-Shot Imitation Learning
Graph Neural Networks

Abstract: Consider the following problem: given a few demonstrations of a task across a few different objects, how can a robot learn to perform that same task on new, previously unseen objects? This is challenging because the large variety of objects within a class makes it difficult to infer the task-relevant relationship between the new objects and the objects in the demonstrations. We address this by formulating imitation learning as a conditional alignment problem between graph representations of objects. Consequently, we show that this conditioning allows for in-context learning, where a robot can perform a task on a set of new objects immediately after the demonstrations, without any prior knowledge about the object class or any further training. In our experiments, we explore and validate our design choices, and we show that our method is highly effective for few-shot learning of several real-world, everyday tasks, whilst outperforming baselines. Videos are available on our project webpage at https://www.robot-learning.uk/implicit-graph-alignment.

Few-Shot In-Context Imitation Learning via Implicit Graph Alignment
Vitalis Vosylius and Edward Johns
The Robot Learning Lab
Imperial College London United Kingdom
vitalis, . wow.yliiiit.iiimiiiaalll. ac.u
Consider the following problem: given a few demonstrations of a task across few different objects, how can robot learn to perform that same task on new, previously unseen objects? This challenging because the large variety of objects within class makes it difficult to infer the task-relevant lationship between the new objects and the objects in the demonstrations. We address this by formulating imitation learning lem between graph representations of objects. Consequently, we show that this conditioning allows for in-context learning, where robot can perform task on set of new objects immediately after the demonstrations, without any prior knowledge about the object class any further training. In our experi- ments, we explore and validate our design choices, and we show that our method is highly effective for few-shot learning of several real-world, everyday tasks, whilst outperforming baselines. Videos are available on our project webpage at attps://wwe robot-leerning uk ht//mpplicit. gragg-rpphae aeapphallgnnan alignment.
Abstract:
conditional alignment prob-
Keywords: Few-Shot Imitation Learning, Graph Neural Networks
A few dernonstrations of lask with different sets objects
Immediate lask completion by execution the found alignment trajectory
Implicit Graph Alignment
Novel objects
training dataset task-agnoste aligrments
Figure 1: Given few (e.g. three) demonstrations of new task across multiple object instances (left), robot can perform this same task with previously unseen objects (right), based on graph alignment energy model trained on a diverse dataset of task-agnostic object alignments (middle).
Introduction
For robot to perform manipulation task involving two objects, it must infer task-relevant parts of these objects and align them throughout the trajectory in a task-specific way. For example, the spout of teapot should be aligned with the opening of a mug when pouring tea. However, objects within a class typically have slightly different shapes, so inferring the parts which should be aligned for a given task is challenging due the diversity of object shapes. This becomes in an imitation learning setting, where this task-specific alignment should be inferred only from demonstrations, and ideally, without requiring any prior knowledge about the objects or their class.
challenging
We approach this challenge as a few-shot in-context imitation learning problem, and learn condi- tional distribution of alignments that can produce task-relevant alignments of novel objects from fuomjustt
7th Conference on Robot Learning (CoRI. 2023), Atlanta, USA.

---

a few context examples, i.e. demonstrations. We propose to learn this distribution from point cloud observations using novel heterogeneous sraph-based energy model, whose graph structure with transformer-basee attention is capable of representing intricate relations between different parts of the objects, and whose implicit energy-based learning effective at capturing complex and multi- modal distributions. These design choices allow us to perform few-shot imitation learning on any arbitrary objects and over any arbitrary number of demonstrations, using just a single trained model.
But where can we obtain a large-scale dataset of task-relevant alignments between objects needed to train such a model? Our insight here is that simply by deforming pairs of arbitrary ShapeNet [1] objects in simulation, using an object augmentation function trained for a single object category of chairs [2], we can create sensible new instances of objects of an arbitrary category and align them consistently, providing us with the data required to train the aforementioned model. By training on objects from diverse categories and multiple alignments between them, we are able to predict task-relevant alignments for novel object pairs, all from the same trained model.
We evaluated our method both with simulation and real-world experiments and analysed its ability to align novel objects in a few-shot imitation learning setting, where each demonstration consists of one or more waypoints in a trajectory. Our real-world results show that for everyday tasks, such as pouring from a teapot into a mug, sweeping with a brush, and hanging cap onto peg, we achieve 80% task success rate under the following conditions: (1) only four demonstrations are provided, (2) the test objects are novel and unseen during the demonstrations, and (3) we assume no prior knowledge about the objects or class of objects. These results, when compared to the baselines, validate our method as an effective few-shot imitation in-context learning framework, capable of efficiently learning everyday robotic tasks from just few demonstrations.
2 Related Work
The field of imitation learning has seen various approaches that aim to solve robotic tasks based on single demonstration [3, 4, 5, 6]. However, these methods typically limited to scenarios where the objects used in the demonstrations and during deployment are either the same or sufficiently sim- ilar. This lack of flexibility prevents them from being applied to category-level manipulation tasks. Other methods tackle this limitation by leveraging category-level keypoints, which can be anno- tated [7, 8), predicted [9]. obtained through sulf--uuprrvsedd correspondence methods [10, 11, 12]. However, selecting such keypoints in a task-agnostic manner is challenging, thereby restricting the applicability of these approaches. Transporter Nets [13], utilises learned feature maps and template matching to identify pick-and-place poses, but is restricted to pick-and-place tasks and struggles to predict arbitrary 6DoF poses. We also train a function to predict the quality of an object alignment, but we do so in an unrestricted (3(3) space, enabling us to handle much more complex and arbitrary tasks defined by the demonstrations. Graph neural networks [14. 15. 16, 17] and implicit learning of functions [18, 19, 20, 21] have recently gained significant attention for robot manipulation, and we study a novel formulation of these ideas for few-shot imitation learning.
The most closely related work to ours comes from Relational Neural Descriptor Fields (R-NDF) [22, 23] and TAX-Pose [24], both of which can generalise to novel instances within the same object category. [22] achieves object alignment by matching descriptors extracted from a pre-trained occu- pancy network. However, it requires per-object category training and the annotation of task-relevant object parts. Furthermore, [22] averages descriptors from multiple demonstrations, which can pose challenges when dealing with significantly different objects or multimodal demonstrations. [24] em- ploys cross-object attention for estimating dense cross- cross-object correspondences. However, it requires separate network training for each new alignment between the objects, limiting its applicability. In contrast to previously described methods, our approach learns a general conditional distribution of alignments using an energy-based model, which can be applied to objects of unseen categories and arbitrary alignments without the need for additional training beyond the initial demonstrations.
3 Method
3.1
Problem Formulation & Overview
Problem Setting. We consider an imitation learning setting involving two objects, namely a grasped object (or a robot gripper) and target object, between which the interaction should occur. We do
2

---

not assume any prior knowledge about these objects, and denote their observations as OA, and OB. The alignment, or relative spatial configuration between these objects, is denoted as general representation A(OA, On), which should be independent of the global configurations of the objects, and focusses only on their relative alignment. As we will explain later, we use segmented point clouds as observations Oa and Og and devise heterogeneous graph to represent A(O OB).
Objective. Given an arbitrary task involving any two objects Oa and OB. our goal to infer tra- jectory of alignments between them that would complete the task. We denote this trajectory as a se- quence of alignments Atstt (Os, OB). where is the number of time steps in a trajectory. In general, to accomplish this, we would require knowledge of the task-conditions distribution p( p( 1::T task), 1
but defining this analytically is infeasible for all possible task and object combinations and from partial point cloud observations, Instead, we propose to approximate it using few samples from this distribution [Alemo (OA, O'))],", . by learning a conditional distribution po (o [A1:T demo): ),, parameterised by 0. Here, N denotes the number of demonstrations.
Deployment. When presented with novel objects Oa and Og, our objective is to perform task by sampling from the learned conditional distribution, eliminating the need for task-specific training or fine-tuning. We sample from po (AAesst transformations 71:2" which when applied to object A, yield alignments within the distribution. Assuming rigid grasps, we can directly apply this sequence of transformations to the robot's end- effector, causing it to move the grasped object along the task-relevant trajectory. In this work, we treat individual time steps in trajectory independently, hence, going forward, we will drop the superscript E, and refer to set of N demonstrations Ademo for conciseness.
OA, OB)l(AAtmaa: ).N) by finding a sequence of
Single demonstration, same objacts Unique Alggment
Singls demoratiaiion, different objects Many poss ble aligmmens
Multipie damorstrations, different objects Task Rei event aligmment
Demo
Test
Demo
Test
Demos
Test
Figure 2: Illustration motivating the need for multiple demonstrations across different sets of objects, in order to generalise to unseen objects when inferring rask-relevant alignments.
Multiple Demonstrations Requirement. When the objects during test time are the same as during the demonstrations, sampling from the learned distribution would yield alignments identical to the demonstrations (see Figure left). However, if they differ from those in the demonstrations, the task-relevant parts should be aligned consistently. Inferring relevant object parts for an arbitrary task is infeasible from single demonstration when objects used are different (see Figure mid- dle). Therefore, we use multiple demonstrations with different objects from the same categories and make an assumption that across the demonstrations, the relative alignment of the task-relevant parts remains consistent and can be inferred (see Figure 2 right). With increasing number of demon- strations and a broader range of object shapes, the conditional distribution should converge towards Dirac delta function 8(.). In cases where the demonstrations are multimodal or inconsistent, the distribution becomes multimodal, representing different possible ways of completing the task.
3.2 Generalisation Through Shape Augmentation
To learn a conditional distribution Pa(Ate (Adem)) which generalises to arbitrary alignments be- tween different instances of objects from the same category, we need a diverse dataset containing multiple objects with different geometries aligned in semantically consistent way. Obtaining such a dataset would require a huge effort in annotation or solving the underlying semantic object align- shape
ment problem in the first place. Instead, we propose to utilise augmentation and treat each deformation of the object as a new instance. The ability to gener- ate new instances of specific category of objects and having privileged information regarding the ground truth correspondences allows us to create an arbitrary number of alignments between the ob- jects with specific parts consistently aligned (see Figure 3). Although these deformations might not
3

---

represent realistic objects one might encounter in the real world. we hope that we can encapsulate the true underlying distribution of objects within the convex hull of the produced deformations.
Arbitrary objects Shapp Not
Instanoos of objects acquired through deformations
Relevant Parts, that maintain the same algnments.
Figure 3: Two arbitrary objects and their deformed variations aligned based on specific parts. Colour maps indicate dense correspondences, green and blue spheres highlight the parts used for alignment.
In practice, we create a diverse dataset of alignments, as follows. 1) Random Object Selection: We randomly select two objects from the ShapeNet dataset representing O^ and OB, and apply ran- dom (3) transformations to them, creating an arbitrary alignment between them. 2) Deformation Generation: To create variations of objects, we employ correspondence peeesoondnce-preseriing implicit gener- ative model called DIF-Net [2]. By deforming the objects, we obtain different instances of the same category, each with its own unique shape characteristics. 3) Alignment Based on Correspondences: For each original object, we randomly sample a point on its surface. Using the known correspon- dences between the original and deformed objects, we align the deformed objects based on the local neighbourhood of points around the sampled point. This alignment ensures that certain parts of the objects maintain the same relative pose consistently. 4) Rendering Point Clouds: Finally, we ren- der the point clouds of the aligned objects using simulated depth cameras, resulting in sample of consistently aligned object pairs. The resulting dataset D {(P), P/, \ LM contains samples, Jj
where each sample consists of N aligned object pairs. We utilíse these samples by using N alignments as context and the left-out alignment as the target alignment for those object.
3.3
Alignment Representation
Alignment Energy Model
Optimisation
Demo
Demo
Damn
Local Encoding
Groph Ceetton
Transfome
Teat
Aligrment
Twat Aligoment
Demo
Edeo toen:
CCnmucoaeee dee costooas arcsinng
Intultion: transformaion target objecte consistent with demonetrations by minimiiing energy.
Inputa
Local Geometry Embeddings
Graph Represeraton
Energy
Nignment
Figure 4: (Left) High-level structure of the model used to learn the alignment distribution. (Right) The learnt energy landscape 2D, and the graph alignment during different optimisation steps.
As hinted previously, we use point clouds PA and Pj as the initial representation of the objects. To learn the proposed distribution Po PAcsst Adtmm), we need to reason about the geometries of different parts of the objects and their relative alignments. Directly using dense point clouds would make this problem extremely computationally expensive. Therefore, A, ), that encodes different parts of the objects and the relative poses between them.
design a representation for
First, sample K points from each point cloud using Furthest Point Sampling Algorithm and encode local geometry around them using the Set Abstraction (SA) layer [25], to obtain a of fea- ture vectors J and their positions describing the whole object {Fi pp K. (see Figure 4). We utilise

---

Vector Neurons [26] which, together with the local nature of the SA layer, provide us with (3)- Equivariant embeddings F. To ensure that these embeddings describe the local geometry around the sampled point, we pre-train this network as an implicit occupancy network [27], reconstructing the encoded object. Note that each local embedding used reconstruct only a part of the object, reducing the complexity of the problem and allowing it to generalise more easily. We provide full description of the occupancy network used in Appendix B. To encode the relative alignment between the two objects based on the sets of feature and position pairs, we then construct a heterogeneous graph representation. This graph captures the relationships between different parts of the objects and fully deseribes their geometry. Each edge Eęj is assigned feature vector representing the relative position between the nodes. To model frequeryy changes in the relative position between the nodes, utilise positional encoding [28] to express edge features.
3.4 Learning the Conditional Alignment Distribution
We use our devised alignment representation A and learn the proposed conditional distribution Mp(AtesrlAdemm) by employing a novel heterogeneous graph-based transformer energy model de- noted as (Atestt Ademo) Intuitively, the energy-basee model should compare Atest with Ademo- and predict if alignments are consistent (low energy) not (high energy). Then, by minimising the predicted energy with respect Atest, we can find desired alignment of test objects.
Firstly, we combine the graphs describing Atest and Ademo by connecting nodes from demonstrated and test alignments with directional edges, equipped with learnable embeddings (see Figure 4). These connections effectively propagate information about the relative alignment of the objects dur- ing the demonstrations to the test observation. This enables the graph capture the alignment patterns observed in the context and makes it suitable for making predictions about whether the test observation is consistent with the demonstrations. To propagate the information through the graph in structured manner, we utilise graph transformer convolutions [29], which, for heterogenous graph, can be viewed as collection of cross-attention modules. We then learn (o Ademo) by employing an InfoNCE-style [30] loss function, and minimise the negative log- log-likeihhod of:
exp(-Egg Eg(Arest Ademo)) exp(- (op(-Eo(AAaa (Arsst Ademo)) + Eg Noeg exp( Eo (op(Aesst Ademo)
po (Atest]Ademar (Alest) N Nnex
Here Nneg is a number of counter-examples, used to approximate the, otherwise intractable, nor- malisation constant. We create these counter-examples by applying S3)(3 transformations to nodes in the graph describing OA, transforming both their position PA and feature vectors FA. Note that are able to do this due to the (E(33-equivarran property of our point cloud encodings. We obtain the transformations using a combination of random and Langevin Dynamics [31, 32] sampling. At inference, we sample from the learnt distribution by minimising the energy using gradient-based optimisation [33, 31] to find an SE(3) transformation T which, when applied to D. would result in an alignment between the objects which is consistent with the demonstrations (see Figure 4 right).
T = argmin Es (T x OA, TESE(3)
(2)
Equation 2 can be understood as follows: Adjust the pose of the object A iteratively until the en- ergy is minimised, indicating that the alignment between the objects A and B becomes consistent with the demonstrated alignments. In practice, we use separate networks for finding orientation and translation alignments, optimising them sequentially. In Appendix C, we provide full details regarding the training of the proposed energy model, and optimisation on the SE(3) manifold. Us- ing a graph neural network with transformer-based attention allows us to capture complex relational information between different parts of two objects, and dynamically handle an increasing number of demonstrations during inference. The energy-based approach allows us to learn complex and multi-modal distributions. It can capture the inherent variability and uncertainty in the alignment process, allowing for different ways of aligning objects. This flexibility crucial when multiple valid alignments exist, c.g. when dealing with symmetrical objects multi-moda modal demonstrations.
5

---

4 Experiments
To thoroughly assess the effectiveness of our method in learning a general alignment distribution and its practicality in real-world robotic tasks, we conducted experiments in two distinct settings: (1) a simulation setting, where we have access to ground truth information, and (2) a real-world setting, where we evaluate the performance of the complete robotic pipeline on everyday robotic tasks. A single model trained on 500K samples (5 consistent alignments each) is used for all the experiments, evaluating it by providing new demonstrations as context. Videos are available our project webpage at https: http::////wwrbbbto/eeeee/ee www.oooooontiiiieeoe uk/implicitt t- werplliiittggphh--ee al ignnent.
4.1 Exploring Generalisation Modes
Our first set of simulated experiments aims to answer the following question: Is our model capable of generalising and identifying lask-relevant alignments when presented with both novel objects from familiar categories and entirely novel object categories?
Experimental Procedure. We constructed five distinct evaluation datasets, each introducing an additional mode of generalisation requirement. The generation procedure follows the methodology outlined in Section 3.2, and only includes objects from non-symmetric categories in order to evaluate the orientation accuracy of found aliguments. The five evaluation datasets encompass the following samples: 1) Seen Alignments: This dataset comprises objects and their alignments that were seen during training. 2) Unseen Aliguments: In this dataset, objects that were seen during training are aligned in novel ways, introducing unfamiliar alignment configurations. 3) Unseen Object Instances: This dataset contains unseen instances of object categories that were encountered during training. 4) Unseen Object Categories: Here, objects are from categories that were unseen during training. 5) Multi-Modal Demonstrations: This dataset focuses on multimodal demonstrations of objects from unseen categories. This means that the demonstrations exhibit multiple modes, indicating different possible object alignments. Unseen, here refers to objects, alignments or categories, that have not been used during training. Every described dataset contains 1000 unique samples, each sample with 5 consistent object alignments, 4 of which are used as the context, and as ground truth label.
Evaluation Metrics. Having access to the ground truth alignment of the objects, report the translation and rotation errors of the predicted alignment, when starting point clouds are initialised in random (3) poses. In the case of multi-modal demonstrations, we report the error to the closest ground truth alignment. To put the later presented error metrics in perspective, object sizes in the previously described datasets range from 8 cm 35 cm.
Baselines. In our evaluation, we compare our method of Implicit Graph Alignment (IGA) against several baselines and its own variants to demonstrate its effectiveness. 1) ICP, where the target and grasped object point clouds are aligned to the closest match amongst the available demonstra- tions. It serves to validate the need to learn the alignment distribution. 2) TAX-Post Pose-DCC is a variant of TAX- TAX-Pose-GG [24]. where instcad of one-hot encoding, we use an average embedding of the demonstrations a conditional variable. For fair comparison, this is our attempt extend cross-object correspondence approach to be conditional on demonstrations (the original TAX-Pose- GC conditional a one-hot encoding of desired placement). 3) R-NDF is Relational Neural Descriptor Fields [22], which aims to validate our choice of using a heterogeneous graph represen- tation for the alignment of objects, instead of matching local descriptors of two objects as in R-NDF. 4) Ours DirReg), J is a variant of our method which, using our graph representation, directly regresses the transformation T instead of learning an energy function. 5) Ours(DFM is another variant of our method, which predicts (3) distance to the ground truth alignment and learns a distance field instead of an energy function. For fair comparison, our method and all baselines were trained on the same amount of data from diverse range of object categories.
(1) Seen Aligaments (deg)
(25 Unseen Aligaments Transicm)
<3s Unoeen Object Instunoos
(4) Llneeen Object Cauppries
(5) Multi-Maa
Demonstration
Trams (cm)
(dep)
Trans (cm)
Fransicm
Run (deg)
Thers (cm)
Roi (dog)
ICP
7.7
770
ski
69.4 250
16.
14.0
150
KNDDU
83112
123
811
218
92233
23.
98154
138=78
96.11 15.5
Pax-Pone-DCC DC
21849.
19
1964 10.4
11724
368
263498
467
213489
12634 18
239-94
125.1 W0.
Oues (Di)
+1224
27.1644
4.
275596
33233
31
6.7350
31.12 16.7
76.8 26.
Oues IDM
31212
423
17
1361
37213
43.4
4.420
11534014
10.
14911997
Ches ([CA)
21113
13
121
13
11.761
1.
147
30=13
Table 1: Translation and rotation errors (means and standard deviations) of alignments for different
modes of generalisation. Values of the best-performing method are presented in bold.

---

Results. Results shown in Table reveal that using non-learning based methods, such as ICP. and extending the cross-object correspondence approach be conditional, yield unsatisfactory perfor- mance. The sub-optimal performance of R-NDF can be attributed to the high diversity within object categories and instances, as well as the arbitrary alignments between them (as it primarily focuses on finding alignments between objects in close proximity). Other variants of our approach show promising results in terms of translation, but struggle to accurately determine the orientation. Our proposed method, using an implicit energy-based model, excels in both translation and orientation, and across different modes of generalisation.
4.2 Exploring Demonstration Diversity
Our second set of simulated experiments aims to answer the following question: Which is more important: the number of demonstrations, or the diversity of demonstrations?
Experimental Procedure. We explore this question by varying both the diversity and the number of demonstrations provided as context. We create 3 different evaluation datasets of demonstrations, each with increasing diversity of objects. We vary the diversity of the ob- jects by increasing the random scaling range, as well as the scale of the latent vector a, which controls the amount of deformations applied by [2]. For visualisation of the diversity of objects the created evaluation datasets see Appendix D.4.
LnW Bive diy Mesium Strerity figh Sivesity
12
10
Eroor cm
the DIF-Net model
Number Demonstrations
Figure 5: Translation error based the number of demonstrations for 3 different sets of diversities. a single model is used for all the evaluations. Analysis of Figure 5 reveals that the performance of our method does not improve when the number of demonstrations increases, unless the diversity among those demonstrations also increases. How- ever, an increase in diversity can substantially enhance the performance. Notably, the performance of our method remains relatively unchanged when provided with either 10 or 15 demonstrations, suggesting that the remaining errors are not attributable to insufficient information in the demon- strations. Instead, they likely stem from inherent errors of the learnt model or ambiguities present in the generated datasets.
Results. Translation errors are shown in Fig- ure (orientation errors follow a similar trend and can be seen in Appendix D.4). Note that
4.3 Real- World Experiments
Through our real-world experiments, we aim answer the following question: Can a robot learn real-world everyday tasks from just small number of human demonstrations?
Real-World Setup. Our experiments are performed using the Franka Emika robot, along with three RealSense D415 depth cameras. To obtain point clouds of the objects, the cameras first capture the point cloud of the target object with the robot not obstructing their view. Then, the robot moves its gripper to a position where the grasped object is visible to both cameras, after which the cameras capture point cloud of the grasped object.
Tasks. We evaluate our approach on six different tasks, which can be seen in Figure 6. 1) Grasping.
The goal is to grasp different pans by the handle. 2) Stacking. The goal is to stack two bowls. 3) Sweeping. The goal to sweep marbles into dustpan with brush. 4) Hanging. The goal is to hang cap onto deformable stand. 5) Inserting. The goal is to insert bottle into a shoe. 6) Pouring. The goal is to pour a marble into a mug. We provide success criteria in Appendix D.I.
Experimental Procedure. Trajectories for Grasping and Stacking tasks are defined by single waypoint at the grasping or placing locations and executed with a scripted controller, consistent with R-NDF [22] and TAX-Pose [24]. For these tasks, we provide 3 demonstrations. Trajectories for other tasks are defined by waypoints, which need to be reached sequentially. For these tasks, we provide 4 demonstrations each. Each demonstration is a unique combination of different grasped
7

---

Grasping
Stacking
Sweeping
Hanging
Inserting
Pouring
Demo
Teat
Demo
Demo
Teot
Demo
Test
Demo
Demo
Test
Figure 6: Real-world tasks used to evaluate our method, and the corresponding objects used for providing demonstrations and testing.
and target objects. Notably, the Hanging task involves deformable stand and deformable cap, both of which are randomly deformed before each demonstration and test. We evaluate each task on different combinations / deformations of grasped and target test objects, repeating the evaluation times by starting the objects in different configurations, totalling 15 evaluations per task per method. No objects used to provide the demonstrations are used during the evaluation.
Results. As we can see from Table 2, ICP and R-NDF baselines can be suc- cessfully used to complete Grusping or Stacking tasks but struggle with more challenging ones, e.g. Hanging or Pouring. This be attributed to the relatively high task tolerances and limited variation in the object geome- tries between demonstrations and inference for these tasks. Our method, on the other hand, can consistently complete all the evaluated tasks from just few demonstrations, with the notable ception of a Sweeping task. We hypothesise that this due to major geometry differences between demonstration and test objects, that were not covered by the shape augmentation model.
Grosp
Stack
Sweep
Hang
Insert
Avg.
ICP
9715
14715
8715
0715
7/15
2715
44 31%
R-NDF
7715
14715
4715
3715
+ 29%
IGA
14/15
13715
8715
12715
11715
12715
80 15%
Table Success rates of our method and two best perform- ing baselines in a real-world setting.
5 Discussion
Limitations. We now share the main limitations of our method, in order of importance. 1) As with many similar works, we assume that segmented point clouds are available of sufficient observability to capture task-relevant object parts, and thus we require three depth cameras. However, our ex- periments showed that with real-world hardware and noisy observations, our method still highly effective for everyday tasks. 2) Our current implementation of gradient-based sampling is not yet suitable for reactive closed-loop control on today's computing hardware, as inference can be time- consuming (up to 15 seconds). 3) We model trajectories as sparse waypoints, which may not be suitable for tasks requiring more complex dynamics. However, in future work, our method could be extended to track much denser trajectory, including velocities. 4) We assume that objects remain rigidly grasped during a trajectory. Therefore, whilst we have shown that our method works with deformable objects (the Hanging task), we have not yet extended this framework for objects which may deform during a task. 5) Our current formulation models a task as the interaction of 2 objects, which might not always be the case. However, could extend our method to arbitrary number of objects, if Og is used to represent an observation of the entire environment. 6) Training of our model relies on data generated using heuristic method, which may introduce inaccuracies in the training samples and limit the model's 's precision.
Conclusions. We have proposed a method for few-shot imitation learning, which is able to perform real-world everyday tasks on previously unseen objects, from three or four demonstrations of that task on similar objects. We achieved this by designing novel graph alignment strategy based on an energy-based model and in-context learning, which was empirically shown to outperform alternative methods. Our experiments and videos highlight that this is also practical method which enables a robot to perform a task immediately after the demonstrations, without requiring any further data collection any prior knowledge about the object class, using just single trained network.
8

---

Acknowledgments
The authors thank Norman Di Palo, Kamil Dreczkowski, Ivan Kapelyukh, Pietro Vitiello, Yifei Ren, Teyun Kwon, and Georgios Papagiannis for their valuable contributions through insightful discussions. Additionally, special thanks to Georgios Papagiannis for developing the robot controller used our real-world experiments.
References
[1] A. X. Chang, T. Funkhouser, L. Guibas, P. Hanrahan, Q. Huang, Z. Li, S. Savarese, M. Savva, S. Song, H. Su, al. Shapenet: An information-rich 3d model repository. arXiv preprini arXiv: 1512030012 2015.
[2] Y. Deng, J. Yang, and X. Tong. Deformed implicit field: Modeling 3d shapes with learned dense correspondence. In Proceedings of the IEEF/CVF Conference on Computer Vision and Puttern Recognition, pages 10286-10296, 2021.
[3] Z. Mandi, F. Liu, K. Lee, and P. Abbeel. Towards more generalizable one-shot visual imitation learning. In 2022 International Conference Robotics and Automation (ICRA), pages 2434- 2444. IEEF, 2022.
[4] E. Johns. Coarse-to-fine imitation learning: Robot manipulation from single demonstration. In 2021 IEEE international conference on robotics and automation (ICRA), pages 4613-4619. IEEE, 2021.
[5] E. Valassakis, G. Papagiannis, N. Di Palo, and E. Johns. Demonstrate once, imitate imme- diately (dome): Learning visual servoing for one-shot imitation learning. In 2022 IEEE/RS/ International Conference on Intelligent Robots and Systems (IROS), pages 8611-8621. IEEE, 2022.
[6] Vitiello, K. Dreczkowski, and E. Johns. One-shot imitation learning: A pose estimation perspective. In Conference Robot Learning, 2023.
[7] I.. Manuelli, W. Gao, P. Florence, and R. Tedrake. kpam: Keypoint affordances for category- level robotic manipulation. In Robotics Research: The 19th International Symposium ISRR, pages 132-157. Springer, 2022.
[8] W. Gao and R. Tedrake. kpam 2.0: Feedback control for category-leve robotic manipulation. IEEE Robotics and Automation Letters, 6(2):2962-296) 2021.
[9] Z. Qin, K. Fang. Y. Zhu, L. Fei-Fei, and S. Savarese. Keto: Learning keypoint representations for tool manipulation. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 7278-7285. IEEE, 2020.
[10] R. Florence, L. Manuelli, and R. Tedrake. Dense object nets: Learning dense visual object descriptors by and for robotic manipulation. arXiv preprint arXiv:1806.0887 08756, 2018.
[11] I.. Manuelli, Y. Li, P. Florence, and R. Tedrake. Keypoints into the future: Self-supervised cor- respondence in model-based reinforcement learning. arXiv preprint arXiv:2009.05085 2020.
[12] P. Florence, I.. Manuelli, and R. Tedrake. gegessupppriseeddccreepoddedee in visuomotor policy learning. IEEE Robotics and Automation Letters, 5(2):492-499, 2019.
[13] A. Zeng, P. Florence, J. Tompson, S. Welker, J. Chien, M. Attarian, T. Armstrong, I. Krasin, D. Duong, V. Sindhwani, et al. Transporter networks: Rearranging the visual world for robotic manipulation. In Conference on Robot Learning, pages 726-747. PMLR, 2021.
[14] V. Vosylius and E. Johns. Where to start? transferring simple skills to complex environments. In Conference on Robot Learning, pages 471-481 PMLR, 2023,
[15] M. Sieb, Z. Xian, Huang, O. Kroemer, and Fragkiadaki. Graph-structured visual imita- tion. In Conference on Robot Learning, pages 979-988 PMLR, 2020.

---

[16] I. Kapelyukh and E. Johns. My house, my rules: Learning tidying preferences with graph neural networks. In Conference on Robot Learning, pages 740-749. PMLR, 2022.
[17] C. Wang, Y. Zhang, X. Zhang, Z. Wu, X. Zhu, S. Jin, T. Tang, and M. Tomizuka. Offline- online learning of deformation model for cable manipulation with graph neural networks. IEEE Robotics and Automation Letters, 7(2):5544-5555) 2022.
[18] D. Driess, J.-S. Ha, M. Toussaint, and R. Tedrake. Learning models as functionals of signed- distance fields for manipulation planning. In Conference Robot Learning, pages 245-255. PMIR, 2022.
[19] J.-S. Ha, D. Driess, and M. Toussaint. Learning neural implicit functions as object representa- tions for robotic manipulation. arXiv preprint arX:v:2112204812,, 2021.
[20] J. Urain, N. Funk, G. Chalvatzaki, and Peters. Se (3)-diffusionnfells:: Learning cost functions for joint grasp and motion optimization through diffusion. arXiv preprint arXiv:2209.03855, 2022.
[21] I. Huang, Narang, R. Bajesy, F. Ramos, T. Hermans, and D. Fox. Defgraspnets: Grasp planning on 3d fields with graph neural nets. arXiv preprint arXiv:2303. 2023.
[22] A. Simconov, Y. Du, Y-C. Lin, A. R. Garcia, L. P. Kaclbling, T. Lozano-Pérre and P. Agrawal. Se (3)-cquivriiant relational rearrangement with neural descriptor fields. In Conference Robot Learning, pages 835-846. PMLR, 2023.
[23] A. Simeonov, Y. Du, A. Tagliasacchi, J. B. Tenenbaum, A. Rodriguez, P. Agrawal, and V. Sitz- mann. Neural descriptor fields: Se (3)-equivariant object representations for manipulation. In 2022 International Conference Robotics and Automation (ICRA), pages 2022.
IEEE,
[24] C. Pan, B. Okorn, H. Zhang, B. Eisner, and D. Held. Tax-pose: Task-specific pross-pose esti- mation for robot manipulation. In Conference on Robot Learning, pages 1783-1792. PMLR, 2023.
[25] C. R. Qi, L. Yi, H. Su, and L. J. Guibas. Pointnet++: Deep hierarchical feature learning on point sets in metric space. Advances neural information processing systems, 30, 2017.
[26] C. Deng, O. Litany, Y. Duan, A. Poulenard, A. Tagliasacchi, and L. J. Guibas. Vector neu- rons: A general framework for (3)-equivariant networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 12200-12209, 2021.
[27] L. Mescheder, M. Oechsle, M. Niemeyer, S. Nowozin, and A. Geiger. Occupancy networks: Leaming 3d reconstruction in function space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4460 4470, 2019.
[28] B. Mildenhall, P. P. Srinivasan, M. Tancik, J. T. Barron, R. Ramamoorthi, and R. Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1):99-106, 2021.
[29] Y. Shi, Z. Huang, S. Feng, H. Zhong, W. Wang, and Sun. Masked label prediction: Unified message passing model for semi- semi-supervised classification. arXiv preprint arXiv:2009.03509 2020.
[30] A. v. d. Oord, L.,, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807.03748, 2018
[31] Y. Du and I. Mordatch. Implicit generation and modeling with energy based models. Advances Neural Information Processing Systems, 32, 2019.
[32] P. Florence, C. Lynch, A. Zeng, O. A. Ramirez, A. Wahid, L. Downs, A. Wong, J. Lee, I. Mor- datch, and Tompson. Implicit behavioral cloning. In Conference on Robot Learning, pages 158-168, PMLR, 2022.
10

---

[33] M. Welling and Y. W. Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681- 688, 2011.
[34] C. R. Qi H. Su, K. Mo, and L. J. Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 652-660, 2017.
[35] D. Hendrycks and K. Gimpel. Gaussian enror linear units (gelus). arXiv preprint 1rXiv:1606.08415, 2016.
[36] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. 17Xiv:1711.05101, 2017.
arXiv preprint

---

Appendix
A Framework Overview
To learn proposed conditional distribution of alignments energy- mergy-basedd model, utilise various techniques to create the diverse dataset needed, pre-train an implicit occupancy network to capture local geometries of objects, and create heterogeneous graph representation of object alignments. We discuss these techniques and their implementation details in the following sections, while the overview can be in Figure 7.
We utilise ShapeNet two-fold: learn SE(3)-Equivariant embeddings representing the local geometry of objects, and 2) creale large consistent alignments between two objects.
Veeseer Neunt
Ossupancy Decoder
Encoder
Point cloud of an object separated into different local parts.
SE(J)-Equivarian local genmetry ambeddinns
Combined reconstruction of the original abject
Augnment
DIF-Nol
ShapeNet
Random reletive alignment between two objects
Deformed objects that retain specific relative alignment
Diverse Detaset of Consistent Aligmmett Between Objects
b
We then utilise the created dataset and trained local encoder to leam conditional distribution of alignments as implicit energy function whose minima represent the desired alignment between objects.
Demo
caea"
Demo
Demo N
Transforint
Naum PantNet
Creatoon
Alirmmnt
Test Algemmnt
Heterageneous Graph
Inputs
Local Geometry Embeddings
Representation
Energy
Figure An overview the techniques used in the design of our proposed framework.
B Local Encoder Network
The representation of object alignment A(O Og) as a heterogeneous graph is a crucial step in effectively capturing the relationship between the two objects. To achieve this, begin by encoding the segmented point clouds of the objects as sets of feature and position pairs. The underlying assumption is that each feature vector can effectively represent the local geometry of specific part of the object. By treating these feature vectors nodes in a graph and connecting them with edges that represent their relative positions, we create a graph representation that enables the network to focus on the specific parts of the objects. By adopting this graph- based approach, we are able to shift the network's attention towards local information and individual parts the objects, rather than relying solely on the global geometry of the entire objects. This localised representation facilitates more precise and targeted reasoning about the alignment between the objects, leading to improved performance in capturing the complex relationships and relative positions between the object parts.
12

---

To construct the local features, we first utilise the Furthest Point Sampling (FPS) algorithm to sample K points on the surface of the point cloud (8 our implementation), These sampled points serve as the centre positions p; for the subsequent calculation of local embeddings. Next, we group all the points in the original point cloud according to their closest centroid and re-centre them around their respective centroids. This grouping process results in K different point clouds, each representing a distinct part the object. To encode these local point clouds, we employ a shared PointNet model. This model takes each local point cloud input and generates feature vector F; that describes the local neighbourhood around each centroid. Our PointNet model consists of an eight- layer MLP (Multi-Layer Perceptron) with skip connections, serving the backbone for our local encoder. To introduce -qquivariie to the features, we incorporate Vector Neurons [26] into network. This approach, described in the Vector Neurons paper, helps ensure that the features maintain equivariance with respect to rotations in three-dimensional space. Overall, the local encoder comprises approximately 1.7 million trainable parameters, allowing it capture and encode the relevant local information from the point clouds.
Initial Point Cloud
Set of Local Embeddings
Reconstructed Point Cloud
a)
Entodier
Decadder
b)
Excaddr
Decoooer
Figure 8: Examples of of our trained auto-encoder when reconstructing a pan (a), and a bowl (b). Blue point clouds represent initial point cloud observations, yellow points represent sampled cen- troids, and red and green points represent network prediction made for that point, occupied (red) and not occupied (green).
To enforce, that the local embeddings indeed encode the local geometry, we pre-train them as an implicit occupancy network [27]. where decoder is given query point and local feature embed- ding and is asked to determine whether a query point lies on the surface of the encoded part of the object (F(F', .q) + [o, 1). Decoder is implemented as PointNet Model [34] with GeLU activation functions [35] (without Vector Neurons).
We utilise positional encoding [28] and express edge features (sin(2" (sin) (oss22TTT)...... sin(2"-1 cos(22-11pp)) where Py is the position of the query point, and is the number of frequencies used. In our experiments, we set L 10.
q
To train the occupancy network as an auto-encoder (as presented in Figure 8), a synthetic dataset is generated, consisting of point clouds for randomly sampled ShapeNet objects and corresponding labelled query points obtained using the PySDF library. This dataset comprises a total of 100,000 samples. During the training process, two NVIDIA RTX 2080ti GPUs were utilised for compu- tational acceleration. The training duration spanned a period of approximately days. We used AdamW [36] optimiser and scheduler our learning rate using the Cosine Annealing scheduler.
13

---

C Energy Based Model
To learn our proposed alignment distribution pe (A A Aemmo) we employ an energy- based ap- proach and model the distributions as:
Eg(A Ademo)
po (Atestl Ademo)
2 ((Aeest,
Here Z is a normalising constant. In practice, approximate this, otherwise intractable constant using counter-examples and minimise the negative log-likelihood of Equation
C.1 Architecture
We using heterogeneous graphs constructed using features described in Section B to represent the alignment between two objects G((F)) PAR. (F represented as relative positions between nodes expressed using positional encoding as:
PB)) Edges in the graph in the alignment
Cij = (sin(2"T(p; - pi)), cos(co - -p))ppp)----. sin((2-1 - pp)),, cos(22-) -p))) (4)
In our base model, we use L = 6. Nodes the demonstration and test alignment graphs are connected with direction edges equipped with learnable embeddings, effectively propagating in- formation about the demonstration alignments to the test alignment graph. Note, that we are using heterogeneous graphs, meaning different edges (connecting nodes from the same object, target and grasped objects, and connecting demonstration and test graphs) have different types and will be processed with separate learnable parameters. Finally, to make predictions based on the connected graphs, we add an additional type of node to the graph, which aggregates the information from the test alignment graph. This Node can be seen a Class token, and each of the considered alignments in batch (number of counter-exampee + 1) is connected to separate Classs node.
Having the designed graph structure, as a collection of cross-attention modules. These modules facilitate message passing hetween nodes in the graph, taking into account the specific types of nodes and edges in heterogeneous graph representation. For a specific type of nodes and edges in the graph, the message passing and attention mechanism can be expressed as:
graph transformer convolutions, which can be viewed
(WaF,) + Woess)
J W Σ
an (W2F, + Weeg); aij softmax JEN(()
Vd
(5)
Embedding from the Energy node (or Class token) is then processed with small MLP to produce the predicted energy.
Our base model is comprised of graph transformer convolutions with 4 multi- head attention heads, each with dimension of 64. Final MLP is composed of 2 layers (with dimensions 256) and GeLU activation functions [35]. The full model contains around 5.7M trainable parameters.
C.2 Training
To train our proposed energy model, we first need to create alignments of test objects used as counter-examples ( We do so by creating copies of Gtest ((F) P' K (F and applying (3) to the nodes in the graph describing the grasped object. Note that demonstra- tion alignment graphs not need be copied, as they are connected the test alignment graph with directional edges, propagating information one way.
14

---

To actually transform the nodes in the graph corresponding to the grasped object, both, the position and the feature vectors need be transformed. Given transformation Traises and it's corresponding rotation matrix Rooise we update the graph nodes as:
AA. 11" = nooise x PA. 112 FA Rnoise X FA
66)
Note that this is possible because of the use of SE(()-equivariant htt-eruuiiarraaeeebeddiinig described in Section B. During training, we create 256 different Âtest alignments per batch to approximate 2, each created using a unique noose- All the counter-examples alternative graph alignments are created directly on a GPU, facilitating an efficient training phase.
To calculate set of transformations Toise use combination of Langevin Dynamics (described in Section C.4) and uniform sampling at different scales. We start the training with uniform sam- pling in ranges of [-0.8, 0.8) metres for translation and [-π, .| radians for rotation. After N number of optimisation steps (10K in our implementation), we incorporate Langevin Dynamics sampling which we perform every 5 optimisation steps. During this phase. we also reduce the uniform sam- pling range to [-0.1,0.] 0.1] metres for translation and [- r/44 radians for rotation. Although creating negative samples using only Langevin Dynamics is sufficient, in practice, we found that our described sampling strategy leads to faster convergence and more stable training for this specific application of energy-basee models.
All models were trained on a single NVIDIA GeForce 3080Ti GPU for approximately I day.
During the training of the proposed energy model, several important tricks were employed to ensure stability, efficiency, and smoothness of the energy landscapes for effective gradient- based optimi- sation. These tricks contribute to the overall training process and facilitate the convergence of the model. The following tricks were identified as particularly significant:
L, Regularisation: To prevent the logits from diverging towards positive or negative in- finity, a small 7.2 regularisation term is added to the loss function. This regularisation term helps to control the magnitude of the logits and maintain stability during training.
Spectral Normalisation: Spectral normalisation is applied to all layers of the network. normalises the weights of the network using the spectral norm (the largest singular value) of each weight matrix. In our case, energy landscapes that were learnt without using spectral norms were unusable for gradient-based optimisation.
L2 Gradient Penalties: Gradient penalties are applied to the feature vectors of edges connecting grasped and target objects. This technique imposes an L2 regularisation on the gradients, penalising large changes in the input space. By doing so, the energy landscape becomes smoother and more amenable to gradient-based optimisation.
Pre-training on a Subset of the Data: When dealing with a large and diverse dataset, is beneficial to initialise the network by pre-training it on smaller subset of the training data. This pre-training process allows the gradients to flow in regions of the loss-function landscape that would otherwise be relatively flat. As result, the network can start from better initialisation point, accelerating the training process. In the specific case mentioned, pre-training approximately 1,000 samples saved approximately 70% of the total training time.
Dividing search space between multiple models: We train two models that are exactly the same, but one is trained without rotations, while the other one is trained with rotations only. It reduces the search space each model needs to generalise over and prevents conflicting gradients during inference.
Mix negative sampling strategy: It allows the model to initially learn coarse energy landscape (using random sampling), which later refined using small perturbations and Langevin sampling.
Regular checkpoint evaluation: It allows to detect training instabilities and select stable checkpoint with a smooth energy landscape suitable for gradient-based optimisation.
Although training energy models using a contrastive learning approach can be challenging and un- stable, the discussed training techniques, together with the use of a joint graph representation, that
15

---

mitigates the need for the network to lear highly non-linear mapping between observation and action spaces. lead stable and efficient training. Figure 9 shows typical training logs of our model including training loss and maximum and minimum predicted energies throughout training.
energy max expects,
energy mm
40
-20
15
-30
10
304
2006
506
1094
506
103c
Figure 9: Training logs of energy model (translation). Please note that the sudden jump in the training loss is due to the discussed change in the negative sampling strategy.
C.3 Scaling the Number Demonstrations
As discussed previously, nodes in the graph that belong to the same demonstration are densely connected using different types of edges. Nodes from demonstrated and test alignments are con- nected with directional edges (there are edge connections between demonstrated alignments themselves). Because of this, we can connect demonstrated alignments to any number of paral- lel test samples and make energy predictions for all of them in single forward pass. Therefore, the number of nodes in the graph grows linearly with the number of demonstrations, but the number of edges also grows linearly. More precisely, the number of edges the graph grows as N M x where N is the number of demonstrations, M is the number of parallel samples, and K is the num- ber of nodes per demonstration. Figure 10 shows graphs depicting how the number of nodes, the number of edges, memory usage and time required for single forward pass depends the number of demonstrations used.
2500
2 Fa rallll Sample Paraile! Sample
Number of Edges bx1000)
800
2 Parallel Sample Sample Sample
Sample
Number of Noues
2000
200
500
Number Demonstrations
Demonstrations
100
2 Parallel Sample
Sample
500
2 Perallel Sample
Sample
Forward Pass Time (ms)
Sample
Memery Usage (MB)
Sample
450
400
350
300
Number Demonstrations
Number Demonstrations
Figure 10: Graphs showing how various metrics of heterogeneous graph representation scale
with the number of demonstrations used.
16

---

C.4 Inference Optimisation
Assuming learnt, previously described, energy-based model, our goal at inference is use it to sample from the conditional distribution pa (AAtestl Ademo). We can not directly sample alignments of objects Atst. but we can compute SE(3) transformation T, that when applied to the grasped object would result in an alignment between the objects that are within the distribution po(-).
To solve Equation 2, utilise iterative gradient-hased approach (Langevin Dynamics sampling). Each iteration step k in the optimisation process updates the nodes of the graph alignment represen- tation corresponding to the grasped object as:
X
P k+1 117
The update x Tnoise(e")) IA,, Fx+1 update x Rnoiso x F
Here, ex ~ N E Re and Thoise is calculated using exponential mapping to project it to
(3))) as (t) Expmap(e). In practice, calculate T Tupdate SEE(3) (and Rupdate
(3))), ), we first transform the appropriate nodes in the graph using an identity transformation T₃ E
(3) and calculate its gradient using pack-propagation as V,, Ea x 0% On), Adema) E
R. Finally, Tupdate is calculated by taking an exponential mapping of VT, E.(.) as T Expmap(V E-))).
update
17

---

D Experiments
D.1 Task Definitions
We evaluate our approach on six different tasks: 1) Grasping. The goal is to grasp different pans by the handle, where success means the pan is lifted by the gripper. 2) Stacking. The goal is stack two bowls, where success means one bowl remains inside the other bowl. 3) Sweeping. The goal is to sweep marbles into a dustpan with a brush, where success means that out of the 3 marbles end up in the dustpan. 4) Hanging. The goal is to hang cap onto a deformable stand, where success means the cap rests on the stand. 5) Inserting. The goal is to insert a bottle into a shoe, where success means the bottle stays upright in the shoe. 6) Pouring. The goal is to pour marble into mug, where success means the marble ends up in the mug.
D.2 Implementation Details for Real-Worke Experiments
Obtaining Point Clouds: to obtain point cloud observation for our real-world experiments we are using 3 calibrated RealSense D415 depth cameras. To calibrate the extrinsics, we are using a standard charuco board and calibration implementation from the open-cv library. Typical calibration errors that we observe are in the range of 2 - 5 mm. To combine observations from different cameras, we first projected separate point cloud observations to common frame of reference (frame of the robot base) using calibrated extrinsics. We then concatenated the point clouds and used voxel downsampling (voxel size of 5 mm) remove the redundant information. We first capture the point cloud of the workspace using the three cameras (with the gripper empty) to obtain the observation of the target object. We then place the object in the robot's gripper. The robot then moves the end- effector to the pre-defined position in which the grasped object is seen from 2 external cameras, and the observation of the grasped object is captured. If the task does not involve the grasped object (e.g. grasping). we are using the point cloud the gripper obtained in simulation. Typical point clouds that we obtained using such procedure can be seen in Figure 11.
Gresping
Stacking
Sweeping
Henging
Inserting
Pouring
Figure 11:
Recording Demonstrations: We represent the demonstration trajectories series of waypoints of alignments between two objects. To obtain these waypoints of alignments of point clouds using kinesthetic teaching, we assume rigid grasps and utilise the forward kinematics of the robot. When the grasped object is moved to the desired waypoint, we record the pose of the end-effector and apply transformation to the point cloud of the grasped object, such that it is transformed to the current pose. This mitigates the need to re-capture point cloud observations at each waypoint. Finally, all the point clouds recorded transformed in the robot's end-effector frame. In this way, the predicted transformation can be directly applied to the robot's end-effector, reaching the predicted waypoint.
Executing the Predicted Trajectory: In this work, we model the trajectories as series of sparse waypoints that need to reached in sequence. Therefore, we use a standard position controller that interpolates the path between the waypoints and reaches them with constant velocity.
18

---

D.3 Discussion on Failure Cases
The failure cases our method were mainly observed when the objects used have significantly different geometry from the ones to provide demonstrations. This is evident from the performance of our method on the Sweeping Task, where the brushes used to provide demonstrations are signif- icantly different from the ones used during testing (see Figure 12 top). The shape augmentation model we used to create the dataset of consistent alignments can not create such differences in augmented objects, making them significantly out of distribution for our trained model. Figure 12 bottom shows types of deformations that could be created from the demonstration objects. Note that such failure cases could be mitigated by using more powerful shape augmentation models coming from the fastly moving 3D vision community.
Objects used to provide demonstratons
Objects used during inference
Deformations
Deformations
Figure 12: Top- point cloud observations of brushes used for the Sweeping Task, Bottom types of shape deformations produced by DIF-Net [2] from the objects used for demonstrations.
D.4 Exploring Demonstration Diversity
To explore the importance of diversity in the demonstrations, created 3 different evaluation datasets, each with increasing diversity of objects. We vary the diversity of the objects by increasing the random scaling range, as well the scale of the latent vector a, which controls the amount of deformations applied by the DIF-Net model [2]. Examples of point cloud samples from the created datasets can be seen in Figure 13.
19

---

Context
Test
Low Diversity
Medium Diversity
High Diversity
Figure 13: Random samples from the 3 different datasets used for the data diversity experiment.
Green point cloud represent object A, while red point cloud represent object B.
Low Diversity Medium Diversity High Diversity
30
20
Emrr,
15
5
10 Number of Demonstrations
Figure 14: Rotational error based on the number of demonstrations for 3 different sets of diversities.
20


Reviews:
[{"review_id": "xSZuj9Rdxr", "reviewer": "Reviewer_r1Xw", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This paper looks at the problem of generalizing robotics tasks to new objects in an imitation learning setting. The problem is formulated as inferring a trajectory of alignments between a grasped object and a target object. The proposed approach involves learning a task-conditioned alignment distribution, which is modeled as a graph-based transformer energy model. The dataset for training is created using ShapeNet.\nExperimental evaluation looks at generalization with novel objects and novel object categories, the effect of diversity in demonstration, and real-world experiments.", "summary_of_recommendation": "The paper introduces a novel method for generalizing to new objects. The presented idea has been effectively validated using several experiments, under certain assumptions (as listed in the Limitations section). Overall, this is a good proof-of-concept, and the research direction can be further explored by the community.", "technical_quality": "Very Good:  Convincing.", "clarity_of_presentation": "Very Good:  Well organized and clear.", "potential_impact": "", "robotics_focus": "Sufficient demonstration on hardware", "quality_of_limitations": "Limitations are addressed clearly", "questions_for_rebuttal": "1. Is the same conditional alignment distribution used for both simulated and real-world experiments? If not, how is the training data for real-world tasks generated?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Strengths:\n1. The main idea of the paper, i.e. modeling the relative alignment of objects conditioned on the task using a graph-based transformer energy model, is interesting. While certain assumptions are made, as detailed in the Limitations section, the work serves as a proof-of-concept for the approach, which is worth exploring further.\n2. The experiments are thorough, and look at various research questions. Specifically, the authors experiment with 5 scenarios for generalization: seen alignments, unseen alignments, unseen object instances, unseen object categories, multimodal demos. The results show that the proposed approach is effective across all scenarios. The second set of experiments shows that diverse demos are required to improve generalization. Finally, real-world experiments on 6 tasks shows that the proposed approach is effective.\n3. The paper is well-written and is easy to follow.\n\nWeaknesses:\n1. The title of the paper could be a little misleading, given that the term \"in-context learning\" is used primarily for LLMs these days.", "extra_scores": {"impact": "2: Unsure, difficult to assess."}, "extra_text": {}}, {"review_id": "lxcPUU5_ph", "reviewer": "Reviewer_R4fh", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This paper proposes a new method for few-shot imitation learning named implicit graph alignment, which generates an alignment for two test-time objects conditioned on a few demonstrations. The alignment, which is a SE3 transformation, is generated by optimizing a learned energy function, whose input is a heterogenous graph built between the test-time and demonstration objects. The energy function is trained in simulation using shape augmentation via a InfoNCE loss such that test-time alignments that matches the demonstration alignments have low energy. The method is evaluated in simulation and on 6 real-world tasks and achieves lower alignment error compared to previous methods, and shows good generalization towards novel object instances/categoriess and spatial configurations. ", "summary_of_recommendation": "See as above. I am open to increase the score if my concerns are addressed during the rebuttal period. ", "technical_quality": "Very Good:  Convincing.", "clarity_of_presentation": "Very Good:  Well organized and clear.", "potential_impact": "", "robotics_focus": "Sufficient demonstration on hardware", "quality_of_limitations": "Limitations are addressed clearly", "questions_for_rebuttal": "I have 2 major questions:\n\n- It is well known that training with Info-NCE loss to optimize the energy function can be instable. Did the authors encounter similar issues with this, and if so, how is this solved?\n- How can the method scale up with the number of demonstrations? It seems that the number of graph nodes would grow linearly with the number of demonstrations, and the number of edges would grow quadratically (the paper did not mention how the edges are constructed for the graph, so a clarification would be needed). This seems to make the method impractical for a large number of demonstrations.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Strengths:\n- The idea and algorithm of learning an implicit alignment by optimizing an energy function based on a graph built between the test-time objects and demonstration objects is quite novel.\n- The experiments evaluations are well done in both simulation and in the real-world, and the method is shown to perform better than prior methods, with better generalization abilities.\n- The paper is clearly written and well-motivated.\n\nWeakness:\n- See Questions for rebuttal. ", "extra_scores": {"impact": "4: The work contains interesting new ideas that will potentially have major impact in robotics or machine learning."}, "extra_text": {}}, {"review_id": "hz_wWCIlauq", "reviewer": "Reviewer_4Z1U", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "The work presents a new method for few-shot imitation learning of object manipulation tasks via aligning object-level graph representations. Specifically, the system focuses on one-robot arm manipulation of one object in relation to another object using <15 demonstrations and 3 depth cameras for inputs. The main novelty of this work, compared to prior art in R-NDF and TAX-Pose, is using an energy-based model / graph which can better generalize to unseen object categories and alignments compared to the former occupancy networks. Generalization is further achieved using correspondence-preserving shape augmentation (procedurally). Experiments in simulation show that the model generalizes better across seen/unseen alignments and objects. Experiments in real show show outperformance against baselines.\n", "summary_of_recommendation": "Overall, the work is convincing in demonstrating few-shot imitation learning for a broad set of objects and tasks which is well supported by experiments, ablations, and the method. Given the visuals and videos, it is clear that the results demonstrated have promise on real hardware and are relevant for addressing potentially important robotics problems (grapsing, pouring, stacking, sweeping), especially given the formulation designed around invariance to object domain and specific detail. \n", "technical_quality": "Excellent: Very convincing and thorough.", "clarity_of_presentation": "Very Good:  Well organized and clear.", "potential_impact": "", "robotics_focus": "Sufficient demonstration on hardware", "quality_of_limitations": "Limitations are addressed clearly", "questions_for_rebuttal": "- Were experiments run with multiple random seeds for training?\n- Could confidence intervals for the real-world experiments also be provided, to show statistical significance?\n", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Strengths:\n- The work demonstrates few shot imitation learning of object pair alignment which they show generalizes to novel, unseen objects and clear evaluations describing performance gains upon prior work\n- The explanation for the task and training / deployment time pipelines are well written\n- The method is clear in how it approaches novel object / pose generalization, with the augmentation based pipeline for synthesizing training data and energy model which does not make assumptions about specific object categories / types\n- The experiment procedures in sim/real are described well and are convincing in showing generalization capabilities (especially beyond prior works), the few-shot / sample requirements and tradeoffs, and real world performance\n\nWeaknesses:\n- Though the work claims that it can be used to solve several / generic robotic tasks, all tasks are relatively similar in being a correct alignment problem between an object grasping in a robotic gripper and a static object. Specifically, some of the tasks feel artificial for this setting, for example in the sweeping task, it appears that the marbles are directly placed within the path between the dustpan and brush, rather than using perception / reasoning to clean up all marbles.\n- Figure 4 and Appendix Figure 1 are appreciated, though a more detailed figure presenting the global view of the proposed framework / system would make the method more clear, given the broad set of works leveraged / described in Sections 3.2 to 3.4, e.g. DIF-Net, Further Point Sampling, Set Abstraction, Vector Neurons, the occupancy network, positional encodings, and the transformer energy model.\n\nMisc:\n- On line 87, missing word: “defining this analytically is infeasible [for] all possible task and …”\n", "extra_scores": {"impact": "4: The work contains interesting new ideas that will potentially have major impact in robotics or machine learning."}, "extra_text": {}}, {"review_id": "12PLnsXGVLj", "reviewer": "Reviewer_UYjn", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This paper introduces a solution for imitation learning that takes a minimal number of demonstrations. The idea is learn graph representations to enable the model to conduct in-context learning when executing tasks with a new set of objects, thereby eliminating the need for additional training.\n\nThe method takes point cloud inputs and employs ShapeNet as its training data source. An object augmentation function is utilized for data augmentation. To construct the dataset, two random objects are chosen and deformed using a generative model, while their relative poses remain unchanged. The graph representation is used to get an energy score between multiple objects pairs. Then a series of transformations was learned to minimize this score. \n\nBoth simulated and real-world experiments were conducted. Tasks in real-world experiments included grasping, stacking, sweeping, hanging, pouring, and inserting.", "summary_of_recommendation": "The proposed method appears to have a strong potential in generalizing to novel objects and scenes, a capability that is crucial in many real-world applications. While the methodology could benefit from additional clarity in certain aspects such as noise tolerance levels and the script controller implementation for real-world experiments, it already provides a solid foundation. However, it should be noted that my assessment is made without a deep expertise in this direction. \n\nI've read the rebuttal and my concerns have been addressed. The additional result on failure cases looks interesting and convincing.", "technical_quality": "Good: Generally convincing, but some additional experiments or analyses needed.", "clarity_of_presentation": "Fair: Major revision needed to improve organization or clarity.", "potential_impact": "", "robotics_focus": "Sufficient demonstration on hardware", "quality_of_limitations": "Limitations are addressed clearly", "questions_for_rebuttal": "- In the limitations section, the authors claim that the model can effectively handle noisy observations. However, the extent of this noise tolerance remains unclear. The level of noise the method can accommodate needs further explanation.\n- The implementation of the script controller for real-world experiments is not clearly detailed.\n- How are point clouds obtained from the three cameras in these real-world scenarios?\n- It would be helpful if the authors could provide visualizations highlighting the types of unseen objects the model fails to handle effectively. A comparison detailing how these unseen objects differ from the ones that the model has seen could provide valuable insights.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper introduces a unique approach aimed at achieving demonstrated alignments between two objects. The problem is interesting and the proposed method is novel. However, it lacks a comprehensive description of its implementation, with several key details missing (see questions below)", "extra_scores": {"impact": "4: The work contains interesting new ideas that will potentially have major impact in robotics or machine learning."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduced an algorithm that can make imitation learning policies generalize to new objects at test time. The key idea is to learn an implicit alignment between the test-time objects and the demonstration objects by optimizing an implicit energy function. The experiments demonstrated the effective few-shot learning capabilities of this method on real-world manipulation tasks. This paper received positive ratings from all four reviewers. The reviewers appreciated the novelty of the proposed approach and the real-world experimental validations. Some clarification questions have been raised and properly addressed by the rebuttal.

However, one lingering issue is the choice of real-world tasks. As 4Z1U correctly pointed out, "all tasks are relatively similar," and "some of the tasks feel artificial." It remains to be seen whether this approach could scale up to handle the increased complexity of real-world, everyday tasks. Nonetheless, considering the contributions made in this paper and the unanimous support from the reviewers, the AC believes that this paper has passed the bar of acceptance at CoRL and recommends accepting it for this conference.
Decision: Accept (Poster)