Key Idea:  We formulate In-Context Imitation Learning as a graph generation problem and use procedurally generated pseudo-demonstrations as a main source of training data, achieving manipulation skill acquisition instantaneously after the provided 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.

Here, you can see a person providing a single demonstration, after which Instant Policy is deployed, completing various everyday manipulation tasks. Instant Policy is also robust to introduced distractor objects and perturbations - the person moving the objects in the environment during the policy execution.

Graph Representation

To achieve efficient and generalisable In-Context Imitation Learning, we first need to choose a suitable representation 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 in the figure below.

In-Context Imitation Learning as a Graph Generation Problem

An Infinite Pool of Data

Unlike traditional Behavioural Cloning, which encodes tasks directly into the weights of a neural network, In-Context Imitation Learning defines tasks through the provided context, allowing the model to interpret and act based on a few examples. Leveraging this paradigm shift, we introduce pseudo-demonstrations—arbitrary yet semantically consistent, procedurally generated robot trajectories—as a virtually limitless source of training data. Here, consistent means that while the generated trajectories differ, they `perform' the same type of pseudo-task at a semantic level. Some of these generated pseudo-demonstrations can be used as context while inferring actions associated with the left-out one, making them suitable for In-Context Learning. This approach allows abundant data generation with minimal human effort, creating new opportunities for efficient robot learning. Examples of the pseudo-demonstrations for different pseudo-tasks can be seen in the visualisation below.

Robust and Generalisable Policies

Deploying Instant Policy on a real robot using just one or a few demonstrations results in robust and generalisable closed-loop manipulation policies! Below, you can see Instant Policy being deployed with external perturbations introduced by either moving the objects in the environments or the robot itself.

Understanding the Context

One might ask where this robustness comes from. To investigate this, we visualise the attention weights on the edges of the graph that are responsible for interpreting and aggregating relevant information from the context for placing a phone on a base task in simulation (see video below). We can see that Instant Policy is able to understand at what stage of a task it is at and focus on relevant parts of the context -- initially focusing on the grasping point in the demonstration, and once the phone is grasped, shifting the attention towards the placing point.

Emergent Generalisation Properties

The ability to add an arbitrary number of demonstrations to the context together with a structured graph representation allows Instant Policy to selectively aggregate and interpolate the information present in the context to disambiguate the task and the parts of the objects that are relevant to it. This results in an emergent generalisation capability to generalise to different object geometries from the ones used during the demonstrations. From the video below, we can see that Instant Policy initially fails to open a box with a different geometry at deployment. However, when more demonstrations with different boxes are provided, it is capable of disambiguating the task and generalising to different boxes at test time.

The generalisation to novel object geometries is also present in different types of tasks, such as folding as seen in the videos below.

Cross-Embodiment Transfer

Since our model uses segmented point clouds and defines the robot state by the end-effector pose and gripper state, different embodiments can be used to define the context and roll out the policy, provided the mapping between them is known. We define such a mapping between human hands and robot gripper using Mediapipe and transfer the policy directly to the robot while the demonstrations (i.e. the context) are provided using human hands.

Modality Change to Language

We can further use our trained model as a foundation for completing various manipulation tasks defined by language zero-shot, completely circumventing the need for demonstrations at test time. We achieve this using a smaller language-annotated dataset and approximating the information bottleneck of our model, which holds the information about the context and the current observation needed to predict robot actions. Using this information approximation and deploying Instant Policy allows us to define the context using language and instantaneously deploy the resulting policy, as shown in the video below.