Instant Policy: In-Context Imitation Learning via Graph Diffusion
Authors: Vitalis Vosylius
Edward Johns
Venue: ICLR 2025 Oral (2025)
Keywords: In-context Imitation Learning
Robotic Manipulation
Graph Neural Networks
Diffusion Models
Primary area: applications to robotics, autonomy, planning
TL;DR: We formulate In-Context Imitation Learning as a diffusion-based graph generation problem and learn it using procedurally generated pseudo-demonstrations.

Abstract: Following the impressive capabilities of in-context learning with large transformers, In-Context Imitation Learning (ICIL) is a promising opportunity for robotics. We introduce Instant Policy, which learns new tasks instantly from just one or two demonstrations, achieving ICIL through two key components. First, we introduce inductive biases through a graph representation and model ICIL as a graph generation problem using a learned diffusion process, enabling structured reasoning over demonstrations, observations, and actions. Second, we show that such a model can be trained using pseudo-demonstrations – arbitrary trajectories generated in simulation – as a virtually infinite pool of training data. Our experiments, in both simulation and reality, show that Instant Policy enables rapid learning of various everyday robot tasks. We also show how it can serve as a foundation for cross-embodiment and zero-shot transfer to language-defined tasks.

Published as a conference paper at ICLR 2025
INSTANT POLICY: IN-CONTEXT IMITATION LEARNING VIA GRAPH DIFFUSION
Vitalis Vosylius and Edward Johns The Robot Learning Lab at Imperial College London vi aasls vosyl uus 90 mmperrrr r a 1l.1... ac. uk
ABSTRACT
Following the impressive capabilities of in-context learning with large transform- ers, In-Context Imitation Learning (ICIL) is a promising opportunity for robotics. We introduce Instant Policy, which learns new tasks instantly (without further training) from just one or two demonstrations, achieving ICIL through two key components. First, we introduce inductive biases through graph representation and model ICIL as a graph generation problem with a learned diffusion process, enabling structured reasoning over demonstrations, observations, and actions. Second, we show that such a model can be trained using pseudo-demoonstattions virtually infinite pool of
arbitrary trajectories generated in simulation training data. Simulated and real experiments show that Instant Policy enables rapid learning of various everyday robot tasks. We also show how it can serve foundation for cross- eross-embodiment and zero-shot transfer to language-defined tasks. atk/inntanttpolicyy.
Code and videos are available https:/wwwwrrboo-eeammii
Instant Policy
In-Context Graph Diffusion Model
Trained Pseudo Demonstrations
A few demonstrations of novel task defining the context.
Instantly execute closed-loop policy after the provided demonstrations.
Instant Policy acquires skills instantly after providing demos at test time. We model in- context imitation learning as graph-based diffusion process, trained using pseudoodemmonstrationss.
Figure 1:
INTRODUCTION
Robot policies acquired through Imitation Learning (IL) have recently shown impressive capabili- ties, but today's 's Behavioural Cloning (BC) methods still require hundreds or thousands of demon- strations per task (Zhao al.). Meanwhile, language and vision communities have shown that when large transformers are trained on sufficiently large and diverse datasets, we see the emergence of In- In-Context Learning (ICL) (Brown, 2020). Here, trained models can use test-time examples of novel task (the context), and instantly generalise to new instances of this task without updating the model's ' weights. This now offers a promising opportunity of In-Context Imitation Learning (ICIL) in robotics. To this end, we present Instant Policy, which enables new tasks to be learned instantly: after providing just one or two demonstrations, new configurations of that task can be performed immediately, without any further training. This is far time-efficient and convenient than BC methods, which require numerous demonstrations and hours of network training for each new task.
ICL in language and vision benefits from huge and readily available datasets, which do not exist for robotics. As such, we are faced with two primary challenges. 1) Given the limited available data, we need appropriate inductive biases in observation and action representations for efficient learning in

---

Published as a conference paper at ICLR 2025
3D space: 2) Given the inefficiency and cost of manually collecting roboties data, need means to easily collect training data in scalable way. In this work, we propose solutions to both these.
We address the first challenge by introducing a novel graph-based representation that integrates demonstrations, current point cloud observations, and the robot's actions within a unified graph space. We then cast ICIL as a diffusion-based graph generation problem, enabling demonstrations and observations to be interpreted effectively in order to predict the robot's actions. To address the second challenge, we observe that in traditional BC, the model's weights directly encode policies for a specific set of tasks, whereas in ICIL, the model's weights should encode a more general, task- agnostic ability to interpret and act upon the given context. Due to this, we found that we were able to train the model using pseudo demonstrations - sets of procedurally generated robot trajectories, but where each set of demonstrations for task semantically consistent. This approach allows us to generate virtually infinite training data by scaling up the simulated data.
Our experiments, with both simulated and real-worke tasks, show that Instant Policy can learn vari- ous everyday tasks, whilst achieving higher task success rates than state-of-the-art baselines trained on the same data. As an emergent ability, we also observed generalisation capabilities to object geometries unseen from the test-time demonstrations, Importantly, we found that performance im- proves as more data is generated and used for simultancous training, offering scalable opportunities with abundant simulated data. In our further experiments on downstream applications, Instant Pol- icy also achieves cross-embodiment transfer from human-hand demonstrations to rohot policies, and zero-shot transfer to language-de defined tasks without needing large language-annotra annotated datasets.
Our contributions are as follows: 1) We cast In-Context Imitation Learning as a diffusion-bart graph generation problem; 2) We show that this model can be trained using procedurally generated pseuudoddemonstrations::; 3) We evaluate in simulation and the real world across various everyday tasks, showing strong performance, encouraging scaling trends, and promising downstream uses.
2 RELATED WORK
In-Context Learning (ICL). ICL is emerging paradigm in machine learning which allows mod- els to adapt to new tasks using a small number of examples, without requiring explicit weight updates or retraining. Initially popularised in natural language processing with models like GPT-3 (Brown, 2020). ICL has been applied to enable robots to rapidly adapt to new tasks by using foundation models (Di Palo & Johns, 2024), finding consistent object alignments (Vosylius & Johns, 2023a), identifying invariant regions of the state space (Zhang & Boularias, 2024), and directly training policies aimed at task generalisation (Duan et al., 2017; Fu et al., 2024) or cross-embodiment trans- fer (Jang et al., 2022; Jain et al., 2024; Vogt et al., 2017). Despite these advancements, challenges remain in achieving generalisation to tasks unseen during training and novel object geometries. In- stant Policy addresses this by leveraging simulated pscudo-demonstratioos to generate abundant and diverse data, while its structured graph representation ensures that this data is utilised efficiently.
Diffusion Models. Diffusion models (Ho et al., 2020) have gamered significant attention across various domains, due to their ability to iteratively refine randomly sampled noise through a learned denoising process, ultimately generating high-quality samples from the underlying distribution. Ini- tially popularised for image generation (Ramesh et al., 2021), diffusion models have recently been applied to robotics. They have been utilised for creating image augmentations (Yu et al., 2023; Man- dlekar et al., 2023) to help robots adapt to diverse environments, generating 'imagined' goals (Kape- lyukh et al., 2023) or subgoals (Black et al., 2023) for guiding robotic policies, and learning precise control policies (Chi et al., 2023; Vosylius et al., 2024). In contrast, our work proposes novel use of diffusion models for graph generation, enabling structured learning of complex distributions.
Graph Neural Networks (GNNs). Graph Neural Networks (GNNs) allow learning on structured data using cessaag--asssing or attention-based strategies. These capabilities have been applied across a wide range of domains, including molecular chemistry Jha et al. (2022), social network analysis (Hu et al., 2021), and recommendation systems (Shi et al., 2018). In robotics, GNNs have been employed for obtaining reinforcement learning (RL) policies (Wang et al., 2018; Sferrazza et al., 2024), managing object rearrangement tasks (Kapelyukh & Johns, 2022), and learning affor- dance models for skill transfer Vosylius & Johns (2023b). In our work, we build on these foundations by studying structured graph representations for ICIL, enabling learning of the relationships between demonstrations, observations, and actions.
2

---

Published as a conference paper at ICLR 2025
3 INSTANT POLICY
3.1 OVERVIEW & PROBLEM FORMULATION
Overview, We address the problem of In- In-Context Imitation Learning, where the goal is for the robot complete novel task immediately after the provided demonstrations. At test time, one or a few demos of a novel task are provided to define the context, which our trained Instant Policy network interprets together with the current point cloud observation, and infers robot actions suitable for closed-loop reactive control (see Figure 1). This enables instantaneous policy acquisition, without extensive real-world data collection or training. We achieve this through a structured graph repre- sentation (Section 3.2), learned diffusion process (Section 3.3), and an abundant source of diverse simulated meeudo-dmmonstrationn (Section 3.4).
Problem Formulation. We express robot actions a as end-effector displacements TEA E SE(3) (which, when time-scaled, correspond to velocities), along with binary open-close commands for the gripper, dg E {0, 1}. Such actions move the robot's gripper from frame E to new frame A and change its binary state accordingly. Our observations, or at timestep t, consist of seg- mented point clouds pf, the current end-effector pose in the world frame W, TivE SE((3), and a binary gripper state sy E {0, 13. Formally, our goal is to find a probability distribution, p(a ot, {(oj a:y) 1))1). from which robot actions can be sampled and executed. Here, denotes the action prediction horizon, while and N represent the demonstration length and the number of demonstrations, respectively. For conciseness, from now onwards we refer to the demon- strations, which define the task at test time and are not used during training, as context C, and the action predictions
Analytically defining such distribution is infeasible, therefore aim to learn it from simulated decudd-demonsttattoon using a novel graph-based diffusion process.
3.2 GRAPH REPRESENTATION
To learn the described conditional probability of actions, we first need to choose a suitable repre- sentation that would capture the key elements of the problem and introduce appropriate inductive biases. We propose a heterogeneous graph that jointly expresses context, current observation, and future actions, capturing complex relationships between the robot and the environment and ensuring that the relevant information is aggregated and propagated in a meaningful manner. This graph is constructed using segmented point cloud observations, as shown Figure 2.
State Representation
Joint Demo, Current State and Future Actions Representation
Context
Current State
Action
Action T
Sequence of States
Dome
Demo N
Figure 2: (Left) A local graph, Gr, representing the robot's state (blue nodes) and local geometries of the objects (green nodes). (Right) A graph representing 2 demos (3 waypoints each), the current state, and 2 future actions. Edge colours represent different edge types in a heterogeneous graph.
Local Representation. The core building block of our representation is the observation at time step t, which we express as a local graph g (P', Ti sg) (Figure 2, left). First, we sample M points from dense point cloud p' using the Furthest Point Sampling Algorithm and encode lo- cal geometry around them with Set Abstraction (SA) layers (Qi et al., 2017), obtaining feature vectors F and positions p (Fi, ) = o( (P). The separately pre-trained Φ, implicit oc-
3

---

Published as a conference paper at ICLR 2025
eupancy network (Mescheder al., 2019), ensures these features describe local geometries (details in Appendix A). We then represent the gripper's s state in the same format, {Fr -. Pg]i=1
16 by rigidly transforming key points Pkp on the end-effector, Ps Twæ Pkp and assigning them embeddings that encode node distinction and gripper state information F g = [fg. Øy (sy)]] Finally,
link scene and gripper nodes with directional edges and assign edge attributes e representing relative positions in Cartesian space. To increase the precision and capture high-frequency changes in the positions of the described nodes, we represent edges as Eij = (sin(2' (pj p;)); cos cos(22((() (i)).-...nnn222 (P) pi)), cos(22-1mm)))))m (p; - pt))), similar to Zhou et al. (2023).
Context Representation. While gj captures the environment state, sequence of such graphs, interleaved with actions, defines a trajectory within context C' (Figure 2, middle). We perform this interleaving by linking gripper nodes across time to represent their relative movement (red edges) and connecting all demonstration gripper nodes to the current ones to propagate relevant information (grey edges). This enables the graph efficiently handle any number of demos, regardless of length, whilst ensuring that the number of edges grows linearly. The result, Ge ( enables a structured flow of information between the context and the current observation.
Action Representation. To express future actions = (TEA: ag) within the graph repre- sentation, we construct local graphs as if the actions were executed and the gripper moved: (Pt, x ag). This allows 'imagining' spatial implications of actions. Thus, the ac- tions are fully described within the positions and the features of the nodes these local graphs. To represent actions as relative movements from the current gripper pose, we then add edges between current and future gripper nodes with position- position-based embeddings that represent relative movement between subsequent timesteps. These edges propagate the information from the cur- rent observation (and indirectly the context) to the nodes representing the actions. The final graph, (GP (g) (M), aggregates relevant information from the context and the current ob- servation and propagates it to nodes representing the actions, enabling effective reasoning about the robot actions by ensuring the observations and actions are expressed the same graph space.
3.3 LEARNING ROBOT ACTION VIA GRAPH DIFFUSION
To utilise our graph representation effectively, we frame ICIL as graph generation problem and learn a distribution over previously described graphs (g) using a diffusion model, depicted in Figure 3. This approach involves forward and backward Markov-chain processes, where the graph altered and reconstructed in each phase. At test time, the model iteratively updates only the parts of the graph representing robot actions, implicitly modelling the desired conditional action probability.
Training
Deployment
Noisy Actions
Intuition: Nodes representing the gripper iltrrlivelyy they represent desired actione.
Input: Joint graph mepresentation constructed using noisy actions.
Graph Transformer
(Gripper nodes:
of
Output Denoising direction for nodas represenning the actions.
diffusion process
During the dirusion process
end the dirlusion process
Figure 3: (Left) High-level structure of the network used to train graph-based diffusion model. (Right) Position of gripper action nodes during the denoising process for one of the predicted actions.
Training. Training diffusion model includes, firstly, the forward process, where noise is itera- tively added to the samples extracted from the underlying data distribution g(). In this phase, we construct noise-an natterd graph by adding noise to the robot actions according to Ho et al. (2020):
a(gh | (N'(" v -Baaa (,.)). k= 1,...., K
4

---

Published as a conference paper at ICLR 2025
Here, N represents the normal distribution, 3x the variance schedule, and K the total number of diffusion steps. This process gradually transitions the ground truth graph representation into graph constructed using actions sampled from Gaussian distribution.
Inversely, in the reverse diffusion process, the aim is to reconstruct the original data sample, in case the graph, from its noise-altered state, utilising a parameterised model po g*). Intuitively, such a model needs to learn how the gripper nodes representing the robot actions should be adjusted, such that the whole graph moves closer to the true data distribution g()). Formally, the parameterised model learns a denoising process of actions using our graph representation ((a) as:
- = G 7E0
(ro((""))) .)) + N
(2)
Here, Ea(.) can be interpreted as effectively predicting the gradient field, based on which a single noisy gradient descent step taken (Chi et al., 2023). As we represent actions as collections of nodes with their associated positions p and features, that depend on the binary gripper actions dg. such a gradient field has two components Eo pVp Vagl". As we will discuss later, Vag can be used directly in the diffusion process, while a set Vp predictions is an over-parameteriiatioon of a gradient direction the SE(3) manifold, and additional steps need to be used to compute precise denoising update. However, this can result in a large translation dominating small rotation, and vice versa, preventing precisely learning both components well. To address this, represent the denoising directions as a combination of centre- of- mass movement and rotation around it, ef- fectively decoupling the translation and rotation predictions while remaining in Cartesian space as (V, V = [ A ReA X Pkp R X pkpl" with Vğ = Vit + Vý, representing flow (red arrows in Figure 3. left). Here, tEA E R and REA E SO(3) define the SE(() transformation representing actions TEA- Thus learn Ey by making per-node predictions ex IR and optimis- ing the variational lower bound of the data likelihood which has been shown (Ho et al., 2020) to be equivalent minimising (SE(s* ((0(()) (g As our parameterised model, we use heterogeneous graph transformer, which updates features of each node in the graph, Fi, as (Shi et al., 2020):
F.= = WIFF Σ att, (W2F, + Wsess): att,, = softmar (WaF)) (WaF, + Wsesj) jEN(()
Vd
Here, W, represent learnable weights. Equipping our model with such a structured attention mech- anism allows for selective and informative information aggregation which is propagated through the graph in a meaningful way. while ensuring that memory and computational complexity scales linearly with increasing context length (both N and L). More details can be found in Appendix C.
Deployment. During test time, we create the graph representation using actions sampled from the normal distribution, together with the current observation and the demonstrations the context. We then make predictions describing how gripper nodes should be adjusted, and update the positions of these nodes by taking denoising step according the DDIM (Song et al., 2020):
1 ak-1 an (p' vakpg)
p9 - /ak-IPs +
(4)
Here, P9 = + Ap, Apr. This leaves us with two sets of points y k-1 and p that implicitly represent gripper poses at denoising time steps & - and k. As we know the ground truth corre- spondences between them, we can extract an (I) transformation that would align them using a Singular Value Decomposition (SVD) (Arun et al., 1987) as:
T arg
||p
TE- x 1112
Finally, the R-1 is calculated by applying calculated transformation Tk-1,k to a². Note that for gripper opening and closing actions utilising Equation 4 directly is sufficient. This process is re- peated times until the graph that is in distribution is generated and, as a byproduct, final a actions are extracted, allowing us sample from the initially described distribution p(a. or, C).
5

---

Published as a conference paper at ICLR 2025
3.4 AN INFINITE POOL OF DATA
Now that we can learn the conditional distribution of actions, we need to answer the question of where a sufficiently large and diverse dataset will come from, ensure that the learned model can be used for a wide range of real-world tasks. With In-Context Learning, the model does not need to encode task-specific policies into its weights. Thus it is possible simulate 'arbitrary but consistent" trajectories as training data. Here, consistent means that while the trajectories differ, they 'perform' the same type of pseudo-task at semantic level. We call such trajectories pseudo-demonstin ggeudo-demonstrationss
Data Generation. Firstly, to ensure generalisation across object geometries, we populate simu- lated environment using diverse range of objects from the ShapeNet dataset (Chang et al., 2015). We then create pseudo-tasks by randomly sampling object-centric waypoints near on the objects, that the robot needs to reach in sequence. Finally, by virtually moving the robot gripper between them and occasionally mimicking rigid grasps by attaching objects to the gripper. create pscudo- demonstrations trajectories that resemble various manipulation tasks. Furthermore, randomising the poses of the objects and the gripper, allows us to create many -esnttddmoosttrttionn performing the same pseudo-task resulting in the data that we use to train our In-Context model.
Pseudo-Task
Psoudo-Task
Figure 4: Examples of the simulated trajectories dreuddoddmoonttratioons for pseudo-tasks.
In practice, facilitate more efficient learning of common skills, we bias sampling towards way- points resembling tasks like grasping or pick-and-place Note that as the environment dynamics and task specifications, such as feasible grasps, are defined as context at inference, we do not need to ensure that these trajectories are dynamically even kinematically feasible. In theory, with enough randomisation, the convex hull the generated trajectories would encapsulate all the possible test- time tasks. More information about the data generation process can be found in Appendix D.
Data Usage. During training, we sample N decudo-demonstrationn for a given pseudo-task, using N - 1 define the context while the model learns to predict actions for the Nth. Although pseudo- demonstrations are the primary training data, our approach can integrate additional data sources in the same format, allowing the model adapt specific settings handle noisier observations.
EXPERIMENTS
We conducted experiments in two distinct settings: 1) simulation with RLBench (James et al., 2020) and ground truth segmentations, and 2) real-world everyday tasks. Our experiments study perfor- mance relative baselines, understand the effect of different design choices, to reveal the scaling trends, and to showcase applicability in cross- cooos-embodiment and modality transfer. Videos are avail- able on our anonymous webpage at https://www.robbl www.oooo--oooaooinngggggo wwwamiing.uiinstanttt... uk/instant-policyy
Experimental Setup. In all our experiments, unless explicitly stated otherwise, we use a single model perform various manipulation tasks by providing N=2 demos, expressed L=10 waypoints as context and predict T=8 future actions. We train this model for 2.5M optimisation steps using pseudo-dent demoosttttos that are continuously generated in parallel, which roughly equivalent to using 700K unique trajectories. During training, we randomise the number of demos in context be- tween and When we discuss integrating additional training data beyond pseudoodemonnstration refer to models fine-tuned for additional 100K optimisation steps using 50/50 mix of pscudo- demonstrations and new data. For more information, please refer to Appendix E.
Baselines. We compare Instant Policy to 3 baselines which also enable In-Context Imitation Learn- ing, namely: BC-Z (Jang et al., 2022), Vid2Robot (Jain et al., 2024). and a GPT2-style model (Rad- ford et al., 2019). BC-Z combines latent embedding of the demonstrations with the current observa- tion and uses an MLP-based model to predict robot actions, Vid2Robot utilises a Perceiver Resam- pler (Jaegle al., 2021) and cross-attention integrate information from the context and currett

---

Published as a conference paper at ICLR 2025
observation, and GPT2 uses causal self-attention predict the next tokens in the sequence, which in our case are robot actions. For a fair comparison, we implemented all baselines by adapting them to work with point cloud observation using the same pre- rre-traind encoder, and all have roughly the same number trainable parameters. Additionally, all components that rely on language-annote such auxiliary losses, removed because our generated pseudo-demonstration do not have such information. To highlight this, we add an asterisk to these methods when discussing results.
data,
4.1 SIMULATED EXPERIMENTS
The aim of our first set of experiments is two-fold: 1) to evaluate the effectiveness of Instant Pol- icy in performing various manipulation tasks by comparing state-of-the-art baselines, and 2) to investigate the role the training data plays in generalising to unseen scenarios. We use a standard RLBench setup using the Franka Emika Panda and test Instant Policy (IP) and the baselines on 24 tasks, 100 rollouts each, randomising the poses of the objects in the environment each time. Addi- tionally, we test models trained using only pseudo-demonstration
(PD only) and a combination of pseudo-demonstrations and 20 demonstrations for each of 12 out of the 24 RLBench tasks (PD++),
Results Discussion. The first notable observation from the results, presented in Table 1, is that all methods achieve non-zero success rates when only pseudo-demoonstration
used and can perform well on at least the simpler tasks. This indicates that these aecudo-demonstrationn are a powerful source of limitless data for In-Context Imitation Learning. Our second observation is that incorpo- rating additional demonstrations from the same domain can greatly boost the performance, helping with generalisation to unseen tasks and novel object poses. Our third observation is that Instant Policy achieves significantly higher success rates than the baselines, showing the importance of our graph representation and its ability to interpret the context. We further demonstrate this by visual- ising attention weights (Figure 5), which reveal the model's 's ability to understand the task's 's current stage and identify the relevant information in the context. We discuss fallure cases in Appendix G.
Policy
Policy
Closs jur
out
Tuilet down
0.33
Close laptop
mileswaae
0.2/0.87
0.76/11 1.30
Close
0100
0.9970.5
Phone
1.00
0.5170.50
0.8/00.51
10.55
Open jar
0.12/0338
01170.30
0.22/0.51
Lan
1,007 1.00
0.82/033
0.90/0:01
10.94
Toilet woal up
0.341.000
063/063
Take umbrella out
0.88 /091
042/0064
0.90/0.00
0.75/0.89
Meat off grill
0.77/0.9
0.75/064
0.7670.35
0.0/0.30
biock
0.75 1,00
0.10/.114
0.12/0.16
10.1/0166
Open microwave
0.23/0.56
0.00/0.13
0.80/0.02
0.000.00
Push button
00.00/ 1.00
0.75 27.51
0.5 0,53
0.80/00.1
Paper roll
0.70/0..5
0.32 11.53
0.29/0.55
0.26/0.4 0.48
Basketball hoop
0.56/0.9
0.02 0.0270099
0000 0.00
0.05 10.07
Put zubish bin
0.97/0..9
0.11 1.11
0.12/0.14
0.18/0.17
Meat on grill
0,78/ 1.00
0.64 70.89
0.51 0..77
0.55/0.81
Pur umbeelia
0.31 /0.37
0.35 0.33
0.41 70.29
0.28/0.39
Flip swich
0.40
0. 0.15/0633
0..16
0..1/
Lamp
0.42/0.41
0,00//0000
0.00
0.00/ 6.00
Averae (PD++) Seen
8.07
0.59
0.60
0.65
Averae (PD++) Unseen
0.82
0.43
0.37
0.31
Averge (PD only) All
0.71
0.36
0:38
0.34
Table 1: Success rates for Instant Policy and baselines on 24 tasks. 100 rollouts for each (trained
using only pseudo-demonstrations / with additional demos from the 12 RLBench tasks on the left).
Task Snapahots
1.0
Attention Weight
Demonstration
Start
Grasping
Placing
Moowing graspig poir
0.0
Moving
Figure S: Attention weights visualised on sub- sub-graph edges two different timesteps the phone- on-base task, showing the model's 's ability to track task progress and aggregate relevant information.
4.2 INSTANT POLICY DESIGN CHOICES & SCALING TRENDS
Our next set of experiments investigates the impact of various hyperparameters on the performance of method, focusing design choices requiring model re-training, inference parameters that

---

Published as a conference paper at ICLR 2025
alter model behaviour at test time, and scaling trends model capacity and training time increase. For the design choices and inference parameters, we calculate the average change in success rate on 24 unseen RLBench tasks, with respect to the base model used in the previous set of experiments, while for the scaling trends, we report validation loss on a hold-out set of pseudo-demonssrations to how well it can capture the underlying data distribution.
Design Choices
Inference Parameters
Action Mode
A%
Diffusion
4%
Prediction Horizon (T)
A%
# Diffusion Steps (K)
A%
Demo
A%
# Demos
A%
Mode
Length (L)
(N)
Ap
-15
Flow
0
-52
-16
1
-71
-12
(Ap,, Ap.)
Sample
-13
5
-26
(Ar, Ag)
-37
Noise
-7
8
0
10
0
(Ar, A6)
-21
No Diff
-29
16
-4
8
0
15
Table 2: Performance change of ablation variants when compared to the base model.
Design Choices. We now examine the following: action mode, diffusion mode, and the prediction horizon. For action modes, we compare our proposed parameterisation, which decouples translation and rotation, against an approach without such decoupling, and more conventional approaches like combining translation with quaternion or angle-axis representations. For diffusion mode, we evalu- ate predicting flow versus added noise, direct sample, and omitting diffusion, regressing the actions directly. Lastly, we assess the impact of predicting different numbers of actions. The results, shown in Table (left), show that these choices greatly influence performance. Decoupling translation and rotation in Cartesian space allows for precise low-level action learning. The diffusion process is vital for capturing complex action distributions, with predicting flow showing the best results. Fi- nally, predicting multiple actions is helpful, but this also increases computational complexity. For a detailed discussion of other design choices, including unsuccessful ones, please refer to Appendix H.
Inference Parameters. Using a diffusion with a flcxible representation that handles arbitrary con- text lengths allows us to adjust model performance inference. We investigate the impact of the number diffusion steps, the demonstration length, and the number of demonstrations in the con- text, as shown in Table (right). Results show that with just two denoising steps, good perfor- mance can be achieved. Demonstration length is critical; it must be dense enough to convey how the task should be solved, this information is not encoded in the model weights. This evident when only the final goal is provided (demonstration length = 1), leading to poor performance. However, extending beyond certain point shows minimal improvement, as the RLBench tasks can often be described by just few waypoints, For more complex tasks, dense demonstrations would be crucial. Finally, performance improves with multiple demonstrations, though using more than two seems to be unnecessary. This is because two demonstrations are sufficient to disambiguate the task when generalising only over object poses. However, as we will show in other experiments, this does not hold when the test objects differ in geometry from those in the demonstrations.
Scaling Trends. The ability to continuously generate training data in simulation allows our model's performance to be limited only by available compute (training time) and model ca- pacity (number of trainable parameters). To as- sess how these factors influence approach, we trained three model variants with different numbers of parameters and evaluated them af-
0.04
69M
117M
178M
Validation Loss
0.01
ter varying numbers of optimisation steps (Fig- ure 6). The results show that the model's ability to capture the data distribution (as re- flected by decreasing validation loss) scales well with both training time and model com- plexity. This offers some promise that scaling compute alone could enable the development of high-performing models for robot manipu- lation. Qualitatively, we observed similar performance trend unseen RLBench tasks. With increased training, we see an increase in performance, However, it plateaus eventually. Similarly, by increasing the model's capacity from 69M to 1 117M, the success reached before plateauing
0.005
Took
200K
4D0K
BOOK
LBM
2.5M
Optimisation
Steps
Figure 6: Validation loss curves for three different model sizes.
8

---

Published as a conference paper at ICLR 2025
increases significantly. However, further increasing the number of trainable parameters 178M results in only minor, insignificant improvements in performance. This suggests the need for more diverse and representative data. Such data could come from available robotics datasets the gener- ation of pseudo-demoonstraioos that more closely resemble real tasks.
4.3 REAL-WORLD EXPERIMENTS
In real-world experiments, we evaluate our method's ability to learn everyday tasks and generalise to novel objects, unseen in both the training data and the context. We use a Sawyer robot with Robotiq 2F-85 gripper and two external RealSense D415 depth cameras. We obtain segmentation by seeding the XMem++ (Bekuzarov et al., 2023) object tracker with initial results from SAM (Kirillov et al., 2023), and we provide demonstrations using kinesthetic teaching. To help the model handle imperfect segmentations and noisy point clouds more effectively, we further co-fine-tunce the model used our previous experiments using demos from 5 tasks not included in the evaluation.
Insert Paper Roll
Open Airryer
Flip Bottle
Stack Bowls
Kinckk over Creeper
Kettle on Stove
Close Coffee Maclina
Hang Cable
Open Box
Turn Tap Right
Turn Tap Laft
Take Rose Cut
Push Cars Togeher
Pick up Kettle
Cleaa Rox
Open Cash Register
Figure 7: The 16 tasks used in our real-world evaluation.
model's ability to tackle various tasks in the real world, tested it and the baselines on 16 everyday tasks (Figure 7). We evaluated all methods using 10 rollouts, randomising the poses of the objects in the environment each time. From the results (Ta- ble 3), we can see that Instant Policy is able to complete various everyday tasks from just a couple of demonstrations with high success rate, outperforming the baselines by a large margins.
Tasks. To evaluate
Open
Flip Bottle
Kettle
Hang Cahle
Paper Roll
Airfryer
Bowls
Creeper
Stove
Coffee Machine
Instant Policy
9710
107
8710
BC.Z*
// 10
5/10
0/ 10
2/
5/ 10
1/ 10
1/ 10
0/10
Vid2Rohot*
3/10
7/
4/10
GPT2*
1/
6/10
1/
4/
5/ 10
5/ 1/10
5/ 10
1/10
Open
Box
Right
Turn Tap
Left
Tap
Rose Out
Push Cans
Together
Kettle
up
Close
Box
Open
Register
Average, %
Instant Policy
8/10
10/10
10710
9710
5710
10/
107
107
88.75
BC-Z*
8/10
2/10
3/10
07 10
2/
10/
7/
8/10
34.38
Vid2Reber*
10
4/10
3/10
10
1/
10
7/
40.63
GPT2*
10
5/10
5/10
0/10
0/10
10/
5/10
7/10
36.88
Table 3: Real-world success rates for Instant Policy and the baselines, with 10 rollouts each.
Generalisation to Novel Objects. While all of our previous experiments focused on evaluat- ing method's performance on the same ob- jects used in the demonstrations, here we aim its ability to generalise novel object geometries at test time. We do so by provid- ing demonstrations (i.e., defining the context) with different sets of objects from the same se- mantic category, and testing on a different ob- ject from that same category. For the evalua- tion, we use four different tasks (Figure 8), each with six sets of objects (four for the demon- strations/context and two for evaluation). We evaluate method with different number
Open Box
Fold Half
Demo
Test
Place a mug on plate Dema
Unplug Charger
Teut
Demo
Text
Figure 8: Objects used in the generalisation ex- periment (numbers indicate their usage stage).

---

completely Published as a conference paper at ICLR 2025
of demonstrations in the context, randomising the poses of the test objects during each roll- out (5 rollouts for each unseen object set). The results, presented in Table 4, show that, with an increasing number of demonstrations across different objects, the performance novel object geometries increases. This indicates that Instant Policy is capable of selectively ag- gregating and interpolating the information present in the context to disambiguate the task and the parts of the objects that are relevant to it. It is important to note that this is an emer- gent behaviour, as we never trained our model with objects from different geometries across the context, and is enabled by the graph representation and structured cross-attention mechanism.
4.3.1 DOWNSTREAM APPLICATIONS
Cross-embodiment transfer. Since our model uses segmented point clouds and defines the robot state by the end-effector pose and grip- per state, different embodiments can be used to define the context and roll out the policy, pro- vided the mapping between them is known. We demonstrate this by using human-hand demon- strations with a handcrafted mapping the gripper state, allowing us to transfer the policy directly to the robot. qualitative experiments, we achieve similar success rates on simple tasks, like pick-and and-elace, compared to kinesthetic teach- ing. However, for more complex tasks, this approach limited by the handcrafted mapping. See our webpage for video examples and refer to Appendix I for more information about this mapping.
Open
Fold
Mug Plate
Unplug Charger
Average,
N
Box
in Half
1
2/ 10
7/10 10
7/10
0/10
40
2
5/ 10
8/ 10
8/10
0/10
52.5
3
10/ 10
10/10
9/10
5/10
85
4.
10/10
10/10
9/10
7/10
90
Table 4: Success rates of Instant Policy with a different number of demonstrations (N), enabling generalisation to novel object geometries.
Modality change. While obtaining policy immediately after demonstrations is powerful and efficient tool, it still requires human effort in providing those demonstrations. We can circumvent this by exploiting the bottleneck our trained model, which holds the information about the context and the current observation needed to predict actions. This information is aggregated in the gripper nodes of the current observations. If we approximate this bottleneck representation using different modalities, such as language, be achieved with a smaller, language-annotater dataset and a contrastive objective. Using language- annotated trajectories from RLBenchh and rollout data from previous experiments, qualitatively demonstrate zero-shot task completion based solely on language commands. For more details, see Appendix J, and for videos, visit our webpage https:/wwww.robo--eemmnn hus.iiintiiiiittttooioiio
bypass using demonstrations context altogether. This can
5 DISCUSSION
Limitations. While Instant Policy demonstrates strong performance, it has several limitations. First, like many similar approaches, we assume the availability of segmented point clouds for sufficient observability. Second, point cloud observations lack colour or other semantically rich information. Third, our method focuses relatively short-horizon tasks where the Markovian assumption holds. Fourth, Instant Policy sensitive to the quality and downsampling of demonstrations at inference. Fifth, it does not address collision avoidance or provide end-to-end control of the full configuration space of the robot arm. Finally, it lacks the precision needed for tasks with extremely low tolerances or rich contact dynamics. However, we believe many of these limitations can be addressed primarily through improvements in generation of the pseudo-demonstratiios such as accounting for colli- sions, incorporating long-horizon tasks, and by improving the graph representation with additional features from vision models, force information, or past observations.
Conclusions. In this work, we introduced Instant Policy, a novel framework for In In-Context Imitation Learning that enables immediate robotic skill acquisition following one or two test-time demonstra- tions. This is compelling alternative paradigm to today's widespread behavioural cloning methods, which require hundreds or thousands of demonstrations. We showed that our novel graph structure enables data from demonstrations, current observations, and future actions, to be propagated ef- fectively via novel graph diffusion process. Importantly, Instant Policy can be trained with only demoosstttossstti generated in simulation, providing a virtually unlimited data source that is constrained only by available computational resources. Experiments showed strong performance relative to baselines, the ability to learn everyday real-world manipulation tasks, generalisation to novel object geometries, and encouraging potential for further downstream applications.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
K Somani Arun, Thomas S Huang, and Steven D Blostein. Least-squares fitting of two 3-d point sets. IEEE Transactions on pattern analysis and machine intelligence, (5):698-700, 1987.
Jimmy Lei Ba. Layer normalization. arXiv preprint urXiv:160700645 2016.
Maksym Bekuzarov, Ariana Bermudez, Joon- Joon-Young Lee, and Hao Li. Xmem++: Production-level video segmentation from few annotated frames. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 635-644, 2023.
Kevin Black, Mitsuhiko Nakamoto, Pranav Atreya, Homer Walke, Chelsea Finn, Aviral Kumar, and Sergey Levine. Zero-shot robotic manipulation with pretrained image-editing diffusion models. arXiv preprint arXiv:2310. 10639, 2023.
Tom B Brown. Language models few-shot learners. arXiv preprint arXiv:2005.14165. 2020.
Angel X Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, et al. Shapenet: An information-rich 3d model repository. arXiv preprint 1512.0:512.03000 2015.
Cheng Chi, Siyuan Feng. Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shu- Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137 2023.
Norman Di Palo and Edward Johns. Keypoint action tokens enable in-context imitation learning in robotics. arXiv preprint arXiv:2403.19557 195788 2024.
Yan Duan, Marcin Andrychowicz, Bradly Stadie, OpenAl Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojeiech Zaremba. One-shot imitation learning. Advances in neural information processing systems, 30, 2017.
Letian Fu, Huang Huang, Gaurav Datta, Lawrence Yunliang Chen, William Chung-Ho Panitch, Fangchen Liu, Hui Li, and Ken Goldberg. In-context imitation learning via next-token prediction. arXiv preprint arXiv:2408. 12940,...88, 15980, 2024.
Dan Hendrycks and Kevin Gimpel.
Gaussian error linear units (gelus).
arXiv preprint
S55i::6606.04111 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogb-Isc: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2103.094300 2021.
Andrew Jaegle, Sebastian Borgcaud, Jean- Baptiste Alayrac, Carl Doersch, Catalin lonescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al. Perceiver io: A general architecture for structured inputs & outputs. arXiv preprint 107.v:210..14799 2021.
Vidhi Jain, Maria Attarian, Nikhil Joshi. Ayzaan Wahid, Danny Driess. Quan Vuong, Pannag Sanketi, Pierre Sermanet, Stefan Welker, Christine Chan, et al. Vid2robot: End-to-end video- conditioned policy learning with cross-attention transformers. arXiv preprint 2rXiv:2403.12943 2024.
Stephen James, Zicong Ma, David Rovick Arrojo, and Andrew Davison. Rlbench: The robot learning benchmark & learning environment, IEEE Robotics and Automation Letters, 5(2):3019 3026, 2020.
Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. Bc-z: Zero-shot task generalization with robotic imitation learning. In Confer- ence on Robot Learning, pp. 991-1002. PMLR, 2022.
Kanchan Jha, Sriparna Saha, and Hiteshi Singh. Prediction of protein interaction using graph neural networks. Scientific Reports, 12(1):8366) 2022.

---

Published as a conference paper at ICLR 2025
Ivan Kapelyukh and Edward Johns. My house, my rules: Learning tidying preferences with graph neural networks. In Conference on robot learning. pp. 740-749. PMLR, 2022.
Ivan Kapelyukh, Vitalis Vosylius, and Edward Johns. Dall-e-bot: Introducing web-scale diffusion models to roboties. IEEE Robotics and Automation Letters, 8(7):3956 3963, 2023.
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. In Proceed- ings of the IEEE/CVF International Conference on Computer Vision, pp. 4015 2023.
James J Kuffner and Steven M La Valle. Rit-connect: An efficient approach single-query path planning. In Proceedings 2000 ICRA. Millennium Conference. IEEE International Conference on Robotics and Automation. Symposia Proceedings (Cat. No. 00CH37065), volume 2, pp. 995- 1001. IEEE, 2000.
I Loshchilov. Decoupled weight decay regularization. arXiv preprint arXiv: 17X1v:1711..100111 2017.
Camillo Lugaresi, Jiugiang Tang, Hadon Nash, Chris McClanahan, Esha Uboweja, Michael Hays, Fan Zhang, Chuo-Ling Chang, Ming Guang Yong. Juhyun Lee, et al. Mediapipe: A framework for building perception pipelines. arXiv preprini arXiv:1906.088 72, 2019.
Ajay Mandlekar, Soroush Nasiriany, Bowen Wen, Iretiayo Akinola, Yashraj Narang, Linxi Fan, Yuke Zhu, and Dieter Fox. Mimicgen: A data generation system for scalable robot learning using human demonstrations. arXiv preprint arXiv:2310.175596 2023.
Matthew Matl. Pyrender. https: / ttt.:/ggthhh...co.......
hptpl//yyyedde//ee 2019.
Lars Mescheder, Michael Oechsle, Michael Niemeyer, Sebastian Nowozin, and Andreas Geiger. Oc- cupancy networks: Learning 3d reconstruction in function space. In Proceedings the IEEF/CVF conference on computer vision and pattern recognition, pp. 4460-4470, 2019.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1):99-106 2021.
Georgios Papagiannis, Norman Di Palo, Pietro Vitiello, and Edward Johns. R+ x: Retrieval and execution from everyday human videos. arXiv preprint arXiv:2407. 12907,, 2024.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019. URL http: http://paperrss neurips cc/ppprr 9015 1155---ytrch-eee ipprrtt vv- styl 1e hee-hiigh wew.hhh-errformneeeeeeeeeeeeeeee leannng pdf.
Charles Ruizhongtai Qi, Li Yi, Hao Su, and Leonidas J Guibas. Pointnet++:: Deep hierarchical fea- ture learning on point sets in a metric space. Advances in neural information processing systems, 30, 2017.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International conference on machine learning, pp. 8821-8831. Pmlr, 2021.
N Reimers. Sentence-ber bert: Sentence embeddings using siamese bert-networks. arXiv preprint arXiv: 10084, 2019,
Careeoo Sferrazza, Dun-Ming Huang, Fangchen Liu, Jongmin Lee, and Pieter Abbeel. Body trans- former: Leveraging robot embodiment for policy learning. arXiv preprint arXiv:2408.06316 2024.
12

---

Published as a conference paper at ICLR 2025
Chuan Shi, Binbin Hu, Wayne Xin Zhao, and S Yu Philip. Heterogeneous information network embedding for recommendation. IEEE transactions on knowledge and data engineering, 31(2): 357-370, 2018.
Yunsheng Shi, Zhengjie Huang, Shikun Feng, Hui Zhong, Wenjin Wang, and Yu Sun. Masked label prediction: Unified message passing model for semi-supervised classification. arXiv preprint arXiv:2009.03509 2020.
Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Perceiver-actor, A multi-task transformer for robotic manipulation. In Conference on Robot Learning, pp. 785-799. PMLR, 2023.
Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502. 2020.
Julen Urain, Niklas Funk, Jan Peters, and Georgia Chalvatzaki. Se (i)-diffusionfield::: Learning smooth cost functions for joint grasp and motion optimization through diffusion. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 5923-5930. IEEE, 2023.
David Vogt, Simon Stepputtis, Steve Grehl, Bernhard Jung, and Heni Ben Amor. A system for learning continuous human-robot interactions from human-human demonstrations. In 2017 IEEE International Conference Robotics and Automation (ICRA), pp. 2882-2889) IEEE, 2017.
Vitalis Vosylius and Edward Johns. Few-shot in-context imitation learning via implicit graph align- ment. arXiv preprint arXiv:23100 12238, 2023a.
Vitalis Vosylius and Edward Johns. Where to start? transferring simple skills to complex environ- ments. In Conference on Robot Learning, pp. 471-481. PMLR, 2023b.
Vitalis Vosylius, Younggyo Seo, Jafar Uruç, and Stephen James. Render and diffuse: Aligning image and action spaces for diffusion-based behaviour cloning, arXiv preprint arXiv:2405.18190 2024.
Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In International conference on learning representations, 2018.
Tianhe Yu. Ted Xiao, Austin Stone, Jonathan Tompson, Anthony Brohan, Su Wang. Jaspiar Singh, Clayton Tan, Jodilyn Peralta, Brian Ichter, al. Scaling robot learning with semantically imag- ined experience. arXiv preprint arXiv:2302.115533 115500022.155,,, 2023.
Xinyu Zhang and Abdeslam Boularias. One-shot imitation learning with invariance matching for robotic manipulation. arXiv preprint arXiv:2405. 131778 2024.
Tony Z Zhao, Jonathan Tompson, Danny Driess, Pete Florence, Seyed Kamyar Seyed Ghasemipour, Chelsea Finn, and Ayzaan Wahid. Aloha unleashed: A simple recipe for robot dexterity. In 8th Annual Conference on Robot Learning.
Allan Zhou, Moo Jin Kim, Lirui Wang, Pete Florence, and Chelsea Finn. Nerf in the palm of your hand: Corrective augmentation for robotics via novel-view synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 17907-17917, 2023.
13

---

Published as a conference paper at ICLR 2025
APPENDIX
A GEOMETRY ENCODER
Here, we describe the local geometry encoder used to represent observations of the environment set of nodes. Formally, the local encoder encodes the dense point cloud into set of feature vectors together with their associated positions as: {Fi, "'''" = & P). Here, each feature J describes the local geometry around the point p'. We ensure this by pre-training an occupancy network (Mescheder et al., 2019), that consists of an encoder de. which embeds local point clouds, and a decoder de which given this embedding and query point is tasked determine whether the query lays on the surface of the object: (e(de(P), - [0, 1]. The high-level structure of our occupancy network can be seen in Figure 9. Note that each local embedding is used to reconstruct only part of the object, reducing the complexity of the problem and allowing it to generalise more easily.
Occupancy Net Decoder
PointNet Encoder
Figure 9: High-level structure of the occupancy network.
We parameterise de network composed of 2 Set Abstraction layers (Qi et al., 2017) enhanced with Nerf-like sine/cosine embeddings (Mildenhall et al., 2021). It samples M centroids from the dense point cloud and embeds the local geometries around them into feature vectors of size 512. Instead of expressing positions of points relative to the sampled centroids p¹ Pj p¹ E IR3 we express them as (sin(2 (p; pt)), cos(2 (2) _ sin(2' (p; p:)), (29 p;))), enabling the model capture high-frequency changes the position of the dense points and capture the local geometry more precisely. We parametrise Vie as an eight-layer MLP with residual connections, that uses the same Nerf-like embeddings to represent the position of the query point. We use objects from a diverse ShapeNet (Chang et al., 2015) dataset to generate the training data needed to train the occupancy network. For training Instant Policy, we do not use the decoder and keep the encoder frozen.
B TRAINING
Training our diffusion model involves a forward and backward Markov chain diffusion process, which is outlined in Equations 1 and 2. Intuitively, we add noise to the ground truth robot actions and learn how to remove this noise in the graph space (see Figure 10).
In practice, training includes 4 main steps: 1) noise is added to the ground truth actions, 2) noisy ac- tions are used to construct our graph representation, 3) the network predicts how nodes representing robot actions need to be adjusted to effectively remove the added noise, and 4) the prediction and ground truth labels are used to calculate the loss function, and weights of the network are updated accordingly.
To add noise to the action expressed as (TEA E SE(3), ag E R, we first project TEA to se(3) using Logmap, normalise the resulting vectors, add the noise described by Ho et al. (2020), unnormalise the result and extract the noisy end endeeffector transformation T EA using Expmap. Such process can be understood adding noise SE(3) transformation in its tangential space. We
14

---

Published as a conference paper at ICLR 2025
Sample data
Alter with noise
Learn to remove graph space
Noisy Actions
Add noise sampled from norman distrbutiin actions
Demonstrations
Graph Transformer
Current State
Future Actione
Noisy Actions
Figure 10: High-level overview of the training process. (Left) A data point is sampled from the dataset. (Middle) Noise is added the ground truth actions. (Right) Using demonstrations, current observation and noisy actions, graph representation is constructed, which is used to predict, how to remove the added noise in the graph space.
can do this because around actions (end-effector displacements) are sufficiently small. For bigger displacements, unnormalised noise should be projected onto the SIE(3 manifold directly, done by Vosylius & Johns (2023a) and Urain et al. (2023). Adding noise to real-valued gripper actions can be done directly using the process described by Ho al. (2020).
C NETWORK ARCHITECTURE
Here describe the neural network used learn the denoising process on graphs, enabling us to generate graphs G and implicitly model the conditional action probability. Our parametrised neural network takes the constructed graph representation as input and predicts the gradient field for each gripper node representing the actions: (g)).. These predictions are then used in the diffusion process allowing to iteratively update the graph and ultimately extract desired low-level robot actions. In practice, for computational efficiency and more controlled information propagation, we are using three separate networks o, é and v. updating relevant parts of the graph in sequence as:
-CG (G.(O (G
6
Here, a operates on local subgraphs Gi and propagates initial information about the point cloud ob- servation the gripper nodes, jo additionally propagates information through the demonstrated tra- jectories and allows all the relevant information from the context to be gathered at the gripper nodes of the current subgraph. Finally, propagates information to nodes the graph representing the tions. Using such a structured and controlled propagation of information through the graph, together with the learnable attention mechanism described in Equation 3, allows the model continuously aggregate relevant information from the context and make accurate predictions about the actions. Additionally, it also results in a clear and meaningful bottleneck in the network with all the relevant information from the context aggregated in specific set of nodes (G(G.((r )) ())) This bottleneck representation could be used for retrieval or as shown in our experiments, to switch modalities, for example to language, via smaller annotated dataset and contrastive objective.
Each of the three separate networks is heterogeneous graph transformer (Equation 3) with lay- ers and hidden dimension of size 1024 (16 heads, each with 64 dimensions). As we using heterogeneous graphs, each node and edge type are processed with separate learnable weights and aggregated via summation produce all node-wise embeddings. This can be understood as set of cross-attention mechanisms, each responsible for processing different parts of the graph representa- tion. We layer normalisation layers (Ba, 2016) between every attention layer and add additional
15

---

Published as a conference paper at ICLR 2025
residual connections to ensure good propagation of gradients throughout the network. Finally, fea- tures of the nodes representing robot actions are processed with a 2-layer MLP equipped with GeLU activations (Hendrycks & Gimpel, 2016) to produce the per-node denoising directions.
D DATA GENERATION
Our data generation process, firstly, includes populating a scene with objects with which the robot will interact. We do so by sampling two objects from the ShapeNet dataset and placing them ran- domly on plane. Next, we define pseudo-task by sampling a sequence of waypoints on or near those objects. The number of these waypoints is also randomly selected to be between 2 and 6, inherently modelling various manipulation tasks. We assign gripper state, mimicking the rigid robotic grasp and release. We then sample a starting pose for the gripper, where initialise mesh of Robotiq 2F-85 gripper. By moving the gripper between the aforementioned waypoints and attaching or detaching the closest object to when the gripper state changes, we create a pseudo-demonstration. To further increase the diversity in pseudo-tasks, we use different interpolation strategies between the waypoints (e.g. linear, cubic or interpolating while staying on spherical manifold). We record gripper poses and segmented point cloud obser- vations using PyRender (Matl, 2019) and three simulated depth cameras. We ensure that the spacing between the subsequent spaces is constant and uniform (Iem and 3 degrees, same as used for the normalisation of actions). Moving objects to different poses, choosing a different starting gripper pose and repeating the process results in several pseudo-demonstrationn for the same pseudo-tar -takk, which we use to train our In In-Context model. As mentioned in Section 3.4, we do not need to ensure that these generated trajectories are dynamically or even kinematically feasible, as the environment dynamics and task specifications, such as feasible grasp, defined as context at inference.
more waypoints to change the
Bias Sampling. To facilitate more efficient learning of common skills, we bias the sampling to favour waypoints resembling common tasks such grasping or pick-and-place. This does not require creating dynamically feasible trajectories but rather involves designing sampling strategies for waypoints that loosely approximate these tasks. For instance, by selecting specific part of an object, moving the simulated gripper to that location, and closing the gripper, we can simulate a grasping task, even if the grasp itself is not physically feasible. We design such sampling strategies for common tasks, such as grasping, pick-and-place, opening or closing. Pseudo-demoonstration are generated using these strategies for half of the samples, while the rest use completely random waypoints.
Data Augmentation. To facilitate the emergence of recovery behaviour of the learnt policies, we further augment the generated trajectories. Firstly, for 30% of the trajectories, we add local distur- bances associated with actions that would bring the robot back the reference trajectory, similarly to how it is done by Zhou et al. (2023). Secondly, for 10% of the data points, we purposely change the gripper's 's open-close state. This, we found to be crucial, as, without it, the policy would never try to re-grasy grasp an object after initially closing the gripper.
E IMPLEMENTATION DETAILS
Here we discuss the implementation details of the Instant Policy, which we found to be important in making the method perform well.
Demo Processing. Although our network can handle an arbitrary number of demonstrations of any length, we downsample the demo trajectories fixed length (L 10 in our experiments). First, we record demonstrations at rate of 25Hz and 10Hz in simulation and the real world, respectively. The lower rate in the real world caused by the simultaneous segmentation of objects of interest by Xmem++ (Bekuzarov et al., 2023). We then include the start and end of the trajectories and all the waypoints where the open-close state of the gripper changed. We then include the waypoints in the trajectory, where the gripper sufficiently slowed down, indicating important trajectory stages (similar to Shridhar et al. (2023)). Finally, if the current number of the trajectory waypoint less than L, add intermediate waypoints between the already extracted ones.
Data Augmentation. To achieve robust policies, we found that it is crucial to randomise current observations and subsequent actions during training. The network can easily overfit to binary gripper
16

---

Published as a conference paper at ICLR 2025
states (if it is open, just keep it open, and if it is closed, just keep it closed). To tackle this, we, with a 10% probability, flip the current gripper state used as an input to the model. This greatly increased generation
the robustness of the resulting policies. Additionally, during the pseudo-demonstrattee process, we added local perturbations to the pose of the gripper (adjusting point cloud observations and actions accordingly), further increasing robustness and enabling recovery behaviour.
Normalisation. We normalise all the outputs of our mode to be [-1,1], a step that we found to be crucial. To this end, we manually define the maximum end-effector displacement between subsequent action predictions to be no more than lem in translation and 3 deg in rotation and clamp the noisy actions to be within this range. Thus the flow prediction is capped be most twice the size of this range. Knowing this, we normalise Vp, and Vp, to be between -1 and 1 independently, enabling efficient network training. For the gripper opening-closing actions, this can done easily as they are expressed as binary states {0, }.. We do not normalise the position of the point cloud observations but rather rely on the sine/cosine embeddings, a strategy that we found to be sufficient.
Point Cloud Representation. We use segmented point cloud observations of objects of interest in the environment as our observations. These segmented point clouds do not include the points on the robot or other static objects such as the table or distractors, In practice, We downsample the point clouds to contain 2048 points and express them in the end-effector frame as TEW X P to achieve stronger spatial generalisation capabilities. These point clouds are then processed with geometry encoder, described in Section A, producing M = 16 nodes used to construct our devised graph representation.
Action Denoising. When updating TEA during our denoising process, we use calculated Tx,x-1 (as described in Section 3.3) and calculate the transformation representing end-effector actions during the denoising process T TrAA-1 x T These actions then used to construct graph representation that is used in the next denoising step. In practice, because we express point cloud observations in the end-effector frame, we apply the inverse of these actions to the M points repre- senting the scene and construct local graphs of actions as G (TTA FA x Pt, Tw p. ag). As there are no absolute positions in the graph, this equivalent to applying the actions the gripper pose Tú Tive. but it allows us to recompute the geometry embeddings of the point clouds at their new pose, better matching the ones from the demonstrations.
Training, We trained our model using AdamW (Loshchilov. 2017) optimiser with 1e learn- ing rate for 2.5M optimisation steps (approx. 5 days on a single NVIDIA GeForce RTX 3080-ti) followed by a 50K steps learning rate cool-down period. For efficient training, we used float16 precision and compiled our models using torch compile capabilities (Paszke et al., 2019). Training data in the form of pseudo-demonsttattinn were continuously generated during training, replacing the older sample to ensure that the model did not see the same data point several times, preventing overfitting.
F SIMULATION EXPERIMENTAL SETUP
Here, we describe the changes we made to standard RLBench setup (James et al., 2020) when conducting our experiments. 1) We generated all the demonstrations (for the context and for those used during training as described in Section 4) using only Cartesian Space planning - we disregarded all demonstrations that were generated using RRT-based motion planner (Kuffner & Valle, 2000). We did so to ensure that the demonstrations did not have arbitrary motions that would not be captured by our observations of segmented point clouds and end-effeetor poses. 2) We restricted the orientations of the objects in the environment to be within /33/3, π /33. We did so to match the distribution of object poses to the one present in generated pseudo-demone decudo-demonstrationn,, It also ensured that most tasks could be solved without complex motions requiring motion planners.
G FAILURE CASES
Here discuss the observed failure modes of Instant Policy during experiments. Given dif- ferent setups and assumptions, we do so for each of our experiments independently. However, the discussed failure modes are shared across the experiments.
17

---

Published as a conference paper at ICLR 2025
Simulated Experiments. During simulated experiments using RLBench (James et al., 2020), we observed several common failure modes of Instant Policy. First of all, tasks such as Open Mi- crowave Put Umbrella into Rack require extremely high precision in action predictions, oth- erwise, the inaccurate dynamics of the simulator will prevent the task from being completed. As such, sometimes the handle of the microwave would slip from the gripper, the umbrella would fly off when in contact with the robot. Second, tasks such as Flipping Switch or Pushing a Button terminate immediately after the task condition is met. As we predict actions of not doing anything at the end of the trajectory, this resulted in the policy stopping before the task is fully completed state virtually the same as the desired one. Moreover, our generated nneudo-demonstration do not include any collision avoidance, which has proven to be problem for tasks such as Turning the Lamp On, where the robot occasionally collides with the lamp by moving in a straight line to- wards the button. Finally, other failure modes usually included policy stalling at a certain point or oscillating between two configurations. We hypothesise that such behaviour is caused by conflicting information in the provided demonstrations and violating the Markovian assumption. In the future, this could be addressed by incorporating past observations into the graph representation.
Real-World Tasks. By far, the most common failure mode in real-world experiments was the segmentation failure caused by several occlusions. Additionally, imperfect segmentation sometimes included parts of the robot the table, causing the policy perform irrelevant actions. This also sometimes degraded the quality of the demonstrations by including irrelevant points (and thus nodes in the constructed graph). Moreover, we observed that the overall quality of the demonstrations, in terms of smoothness and clearly directed motions, had major impact on the performance of Instant Policy. If recorded demonstrations included inconsistent and arbitrary motions, information in the context was conflicting, resulting in the policy stalling or oscillating. Finally, other observed failure cases mainly involved policy not completing the task due to the lack of precision.
Generalisation to Novel Geometries. When evaluating Instant Policy using objects unseen neither during training nor demonstration at inference, policy sometimes just mimicked the motion observed during the demonstrations without achieving the desired outcome. With an increasing number of diversity demos in the context, such behaviour was minimised. However, some tasks (c.g. placing mug plate) were completed mainly due to the high tolerance of the task rather than true generalisation capabilities.
eoooddmmbodiment transfer exper- iments were caused by incorrect mapping of hand poses to end-effector poses and an insufficient field of view in observations. This caused the robot to occasionally miss the precise grasp- ing locations, closing the gripper at stages where it was not intended, and, in general, resulted in demonstrations of poorer quality.
Transfer. The main failure cases during
Modality Transfer. Replacing demonstrations with language descriptions of the task yielded promising results in our qualitative experiments. However, the observed behaviour was sometimes mismatched with the object geometries in the environment. For instance, the policy would execute appropriate motions (e.g.. pushing or closing) but incorrect locations relative to the objects. This issue likely stems from object features containing only geometric information without any seman- tic context. Incorporating additional features from vision foundation models into the point cloud observations and expanding the language-annote annotated dataset could help address this limitation.
H THINGS THAT DID NOT WORK
Here we discuss various design choices we considered before settling for the approach, described in Section 3.
Fully-Con Connnected Graph. Initially, we experimented with fully connected graph, which effectively acts as a transformer with dense self-attention. While the attention mechanism should, in theory, leam the structure relevant to the task, this approach failed to produce good results, even for simple tasks.
One Big Network. Instead of using three separate networks in sequence (as described in Section C), we experimented with a single larger network, which led to a significant drop in performance. We hypothesize that this is because, early on, the nodes lack sufficient information to reason about ac-
18

---

Published as a conference paper at ICLR 2025
tions, causing much of the computation to be wasted and potentially resulting in conflicting learning signals.
More Gripper Nodes. We express the robot state as set of six nodes in the graph. In theory, we can use an arbitrary number (>> 3) of such nodes, allowing more flexible aggregation of relevant information. We experimented with different numbers of such nodes and observed minimal changed in performance, while the computational requirements increased significantly.
No Pre-trained Geometry Encoder. During the training of Instant Policy, keep the geometry encoder frozen. We experimented with training this model from scratch end-to-end, as well as fine- tuning it. Training from scratch did not work at all, while fine-tuning resulted in significantly worse performance. We also experimented with larger sizes of the encoder and saw no improvement, indicating that the geometry information was already well represented.
Homogeneous Graph. Instead of using a heterogenous graph transformer, which processes differ- ent types of nodes and edges using separate sets of learnable weights, we tried using a homogeneous variant with distinct embeddings added to the nodes and edges. This approach resulted in signifi- cantly worse performance, given the same number of trainable parameters. This indicates that by not sharing the same weights, different parts of the network can better focus on aggregating and interpreting relevant information, resulting more efficient learning.
Predicting Waypoints. Initially, we tried predicting spare waypoints instead of low-level actions, e.g. velocities, that progress the execution of a task. We found, that because of these waypoints represent larger end-effector displacements, predicting them with high precision was challenging. Intuitively, this is the result of the increased action space that, when normalised, needs to be repre- sented in interval [-1.
Larger Learning Rates. For our experiments, we used a relatively small learning rate of le To speed up the training process, tried increasing it. However, with increased learning rate found the training process to be unstable, resulting in large gradients and increasing training loss. We also tried using several different optimisers, using AdamW (Loshchilov, 2017) resulting in the best performance.
I CROSS CROSS-EMB -MMOEIEEBBDD TRANSFER
As described in Section 1.3.1, our approach allows us to provide demonstrations using one embod- iment (e.g. using human hands) and instantly deploy a policy on robot, given a mapping between different embodiments is known. This is because our observations are composed of segmented point clouds that do not include points on the robot and its end-effector pose. Thus, by mapping the pose of a human hand to the robot's end-effector pose, we can effectively obtain the same observations. In our experiments, we achieve this mapping using hand keypoint detector from Mediapipe (Lugaresi et al., 2019) and manually designing a mapping between these key points and the corresponding robot's end-effector pose. We model the position of the end-effector be represented by the mid- way position between the index finger and the thumb and estimate the orientation using an additional point on the palm of the hand. this way, we effectively overparametrise the SE(3) pose of the hand, modelled as parallel gripper, using set of positions. This allows us to complete simple tasks, such as grasping or pick-and-place. However, for more precise tasks, such crude mapping can be insufficient. It could be addressed by using more elaborate mappings between human hands and robot grippers, for example, as done by Papagiannis et al. (2024).
J MODALITY TRANSFER
Using our graph representation together with network architecture, discussed in Section C, results in a clear information bottleneck with all the relevant information from the context aggregated in a specific set of nodes (o (G( ))). Information present in the nodes of the graph representing the current information holds all the necessary information to compute precise robot action appropriate for the current situation, and a trained vi()) has the capacity to do it. We exploit this bottleneck and learn to approximate it using the current observation and a language description of a task and utilise a frozen v.) to compute the desired robot actions in the same way as done when the context includes demonstrations. We learn this approximation using the local graph rep-
19

---

Published as a conference paper at ICLR 2025
resentation of the current observation 9j and a language embedding of the task description fiangs produced by Sentence- BER (Reimers, 2019). We use graph transformer architecture, similar to the one used to learn o, and incorporate fiang train this network using a language-annotated dataset comprising demonstrations from RLBench and rollouts from experiments, along with contrastive objective. At inference, language description of a task and, based on the current observation, compute the embeddings of the aforementioned bottleneck. We then use compute robot actions that are executed closed- looped, allowing for zero-shot generalisation to tasks described by language. Although showing dataset, further improvements could
additional type of node in the graph. We
provide
promising performance using only small language-annotatee be achieved by incorporating semantic information into the observation, using a variational learning framework and expanding the dataset size. We leave these investigations for future work.
20


Reviews:
[{"review_id": "iPb73FJhC2", "reviewer": "Reviewer_LPk7", "summary": "This paper introduces the Instant Policy method for in-context imitation learning. Instant Policy represents the state as a graph computed from a point cloud and object segmentation. Edges in this graph connect gripper nodes from the current observation to gripper nodes in the in-context demonstrations. Instant Policy generates actions through a graph diffusion process that produces a gripper state in the graph. The model is trained with \"pseudo-task\" data of a gripper interacting with simulated objects in arbitrary ways. Experiments show that Instant Policy can perform tabletop manipulation tasks from two demonstrations better than ICIL baselines. Instant Policy can also work with cross-embodiment demonstrations and language instructions.", "questions": "1. Perhaps I missed it, but what does the \"*\" in Table 1 next to the baseline numbers indicate?\n\n1. Why are orientations of objects in the environment fixed to be within some range (L872)? Why not increase the range of rotations in the pseudo-demonstrations or have the policy generalize to these new rotations?\n\n1. Do the objects in RLBench tasks overlap with the objects used in the pseudo-demonstrations? This is important for assessing the generalization abilities of Instant Policy trained only on the pseudo-demonstrations.\n\n1. What exactly does \"plateaus at varying levels\" on L430 mean? Does this mean the performance does not exceed the numbers in Table 1, even for larger policy sizes?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. Instant Policy empirically outperforms ICIL baselines in 12 RLBench tasks by a large margin. This indicates that the graph representation and graph diffusion process from Instant Policy provides meaningful improvements over prior approaches. \n\n1. The pseudo-demonstrations are an effective way to scale data collection for robot manipulation tasks in ICIL. This data source is scalable because it simulates interactions with arbitrary objects without the requirement that they are dynamically or kinematically feasible. The value of this data is confirmed by all methods achieving some non-zero ICIL performance when trained with this pseudo-demonstration source.\n\n1. The architecture of Instant Policy is ablated in Table 2, confirming that the action and diffusion mode are important for performance. \n\n1. The paper shows scaling results for Instant Policy in Figure 6. \n\n1. Instant Policy outperforms baselines on real-world robot experiments. These results show that the Instant Policy assumptions of point cloud and object segmentation hold for real-world scenarios where these sensors will be noisy. \n\n1. The paper shows several interesting extensions of Instant Policy. Table 4 shows Instant Policy generalizes to new objects in the real-world. The graph state representation of Instant Policy allows it to work with human demonstrations for the context as well. Finally, language instructions can also replace the provided demonstration context to work with a language-specified tasks.\n\n1. Supplementary G provides useful transparency about architecture and training decisions that did not work, which will be useful for researchers in the field.\n\n1. Qualitative results are provided in the supplementary section's analysis of the failure modes, the supplementary website, and in Figure 5. \n\n1. The paper and code provide sufficient details for reproducing the results.", "weaknesses": "1. It is unclear how to extend the Instant Policy to partially observable settings. Instant Policy relies on all the objects in the scene being represented via nodes in the local graph. Even with addressing partial observability by appending previous local graphs, Instant Policy still would struggle with object clutter or occluded objects in containers. The graph state representation that enables Instant Policy to perform well in ICIL in simple tabletop tasks could be a hindrance when scaling to more complex task settings.\n\n1. The pseudo-demonstration generation lacks collision avoidance behaviors. As the trajectories are not necessarily feasible and are just with a single object, they demonstrate a simple interaction motion. Such data does not show how to avoid other obstacles, such as reaching around another object or into a tight space. This limitation is also confirmed in the RLBench failure cases in L891, which mention that the robot gripper typically moves in a straight line toward objects. Based on this and the previously mentioned weakness, the Instant Policy method seems limited to tabletop manipulation where the scene is fully observable and the gripper only needs to move in a straight line to interact with objects. \n\n1. Instant Policy has two major components: the graph structure and the graph diffusion process. The baselines do not use this same diffusion process or graph structure. The paper does not investigate which of these components of graph diffusion or local graph structure is more important. A simple next-action prediction architecture operating over the graph structure input, like the GPT2 baseline, may be sufficient for good performance like Instant Policy. \n\n1. The method assumes object segmentation to form the local graph. However, this is not a major weakness since the paper demonstrates this works in the real world with noisy observations, so this assumption is likely possible in many realistic settings.\n\n1. The paper uses a fixed number of demonstration steps for all demonstrations being trained on (L819). While Table 2 shows the method is robust to the demonstration length, to scale to more diverse demonstration scenarios, the method must also handle demonstrations of varying lengths where demonstrations with more complex behavior require more steps. \n\nMinor:\n1. L88 mentions that Instant Policy achieves \"true generalization to unseen tasks\". However, the paper never defines what is meant by true generalization and why prior approaches lack this capability.\n\n1. The definition of the term \"pseudo-demonstrations\" was delayed until late in the paper in Section 3.4. I suggest adding some high-level explanation of what's meant by \"pseudo-demonstrations\" earlier in the paper since this term is often referenced earlier in the paper.", "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "p0x0e9nfIf", "reviewer": "Reviewer_WsTz", "summary": "The paper presents Instant Policy, a method aimed at improving In-Context Imitation Learning (ICIL) by enabling task learning from just one or two demonstrations without additional training. It introduces a graph-based representation to effectively combine demonstrations and observations, framing ICIL as a diffusion-based graph generation problem. The authors also explore training with procedurally generated pseudo-demonstrations, which allows for scalable data generation. The results suggest that this approach can perform well across various tasks in both simulation and real-world contexts, showing potential for scalability and generalization.", "questions": "1. About baselines: Do the baseline models also learn from pseudo-demonstrations during the training phase? In the inference phase, is the accessed context modality video, images, or a data structure similar to INSTANT POLICY (e.g., point clouds, points)?\n2. About training dataset: What types of data were used to train the models tested in simulation versus those tested in the real world? Did both use pseudo-demonstrations generated from ShapeNet?\n3. About related works: What is the main difference between your work and One-Shot / Few-Shot Imitation Learning approaches, such as [1]?\n4. About inference: I noticed that the model predicts T=8 future actions, and I'm curious whether the model completes the task based on the provided demonstration by making just one prediction.\n\n[1] Zhang, Xinyu, and Abdeslam Boularias. \"One-Shot Imitation Learning with Invariance Matching for Robotic Manipulation.\" arXiv preprint arXiv:2405.13178 (2024).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The main contribution of this work is the Instant Policy framework for In-Context Imitation Learning, which enables immediate robotic skill acquisition from just one or two demonstrations using a novel graph diffusion process. \n2. The motivation is well-founded. Large Language Models (LLMs) demonstrate impressive capabilities in in-context learning. The paper analyzes the underlying reasons behind the emergence of these capabilities. It also addresses the challenges of acquiring such abilities in the field of robotics and proposes solutions for each challenge identified.\n3. The experimental results and demonstrations indicate that the model possesses certain in-context learning capabilities. It also showcases interesting abilities in downstream applications, including the capacity to learn from demonstrations even after cross-embodiment transfer and modality change during the reasoning phase.", "weaknesses": "1. The writing in the paper could be improved, particularly in the model training section, which is somewhat unclear. Figure 3's left image does not effectively illustrate this process.\n2. The collection and use of pseudo-demonstrations are not well explained, especially regarding how the model utilizes them and the differences in their usage compared to the RLBench dataset or real-world datasets. For details, see the question section.", "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OOKBpeL64x", "reviewer": "Reviewer_ibtP", "summary": "In the presented work, the authors introduce Instant Policy, a novel approach for in-context imitation learning (ICIL). The key contribution of this work is that demonstrations are represented as a graph, including the observations and actions trained from pseudo-demonstrations generated procedurally in simulation. Furthermore, given a single demonstration of the target task, the proposed method is able to generate a suitable control policy by re-framing the ICIL problem as a diffusion-based graph generation problem. The presented work is interesting for the perspective of not requiring data for the target manipulation tasks but is instead able to use an approach similar to learning from play.", "questions": "- How are sub-goals formulated in the diffusion process? Is it just some horizon the diffusion model is able to generate \"in one go\"?\n- How many pseudo-demonstrations are used to train the initial model across how many objects?\n- Transferring the policy from simulation to the real robot seemed to have required only 25 examples (5 rollouts across 5 tasks). How was performance assessed, and would more samples result in significantly better results?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The graph-based diffusion process is interesting, particularly how the robot's motion is generated through the iterative diffusion process.\n- The graph diffusion process in section 3.3 is well described, and generally, the paper is relatively easy to read.\n- The results of the presented method are convincing, consistently outperforming the chosen baselines.\n- The modality change in section 4.3.1 is very interesting and should be expanded on (the cross-embodiment transfer could be placed in the appendix instead)", "weaknesses": "- In section 3.2, it is unclear how actions are encoded. The graph seems to be encoding relative spatial relationships between key-points on task-relevant objects and the gripper, however, the \"Action Representation\" section is intuitively unclear. How are the position-based embeddings learned with respect to the actions? What exactly are the actions?\n- Given that demonstrations are encoded as a sequence of relative spatial relationships, why are action nodes needed? Wouldn't it be possible to just \"condition\" the motion on the current observation with respect to the demonstrated graphs and optimize some form of minimal deformation of the graph? (similar to Vogt, David, et al. \"A system for learning continuous human-robot interactions from human-human demonstrations.\" _2017 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2017.) (this might be a candidate for some discussion in the related work)\n- The authors claim that, unlike BC, their method does not require much data; however, generating pseudo-demonstrations in simulation is also an expensive process. How does generating pseudo-demonstrations compare to alternative approaches, e.g., this one (Lynch, Corey, et al. \"Learning latent plans from play.\" _Conference on robot learning_. PMLR, 2020.)? (this might be a candidate for some discussion in the related work)\n- The generalization to novel objects within the same semantic group seems to stem from the model's ability to generalize due to the selection of only a few key points in the point cloud, making the \"novelty\" of the semantically similar objects questionable. How does the model handle novel objects from semantic groups that were never part of the training data, which would be a common occurrence in household tasks?", "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZnygU6ntqJ", "reviewer": "Reviewer_Wb4x", "summary": "The paper presents a novel method for rapid robot policy learning. The proposed \"Instant Policy\" framework allows robots to instantly learn tasks from just one or two demonstrations by modeling in-context imitation learning (ICIL) as a graph-based diffusion process. The approach uses graph representations for demonstrations, observations, and actions, combined with a learned diffusion process, enabling efficient and scalable task learning. It can be trained using pseudo-demonstrations, offering a virtually infinite source of training data, and achieves high success rates on various tasks in both simulated and real-world environments.", "questions": "See weakness.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents a substantial amount of experiments in both simulator and real-world, showcasing the effectiveness of the proposed method.\n2. The proposed method demonstrates good performance in adapting to unseen tasks.\n3. The idea of using pseudo-demonstrations to ensure that the learned network can be used for a wide range of real-world tasks is interesting.", "weaknesses": "1. I think there is a notation mistake in Line 163 $\\{\\mathcal{F}_g^i,p_g^i\\}^6_{i=1}$.\n2. Since the method consists of three kinds of graphs, the ablation study is missing.\n3. The method section is hard to follow, it will be better with a clearer illustration figure.", "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": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)