Instant Policy: In-Context Imitation Learning via Graph Diffusion

Regarding the description of robot state representation and its notation.

Thank you for thoroughly checking our notations and raising your concerns. However, we believe that the notation is consistent and correct. We represent the gripper’s state as a collection of the positions of key points on the end-effector with their associated feature vectors. A set of positions is an over-parametrization of an SE(3) pose, while distinct feature vectors allow us to distinguish them while also holding information about the binary open-close gripper state. Thus a set of these points and feature vectors fully describe the state of the end-effector.

Regarding the graph representation and its ablations.

Although our graph representation consists of three separate parts describing different components of In-Context Imitation Learning (demonstrations, current observation and future actions), together they form a single graph that is used by our model. In Appendix G (Things that Did not Work), we describe other graph structures that failed to produce good results. Namely, a fully connected graph and a homogeneous graph. We did not include these results in the ablations section as the results were not very informative and left it as a qualitative discussion in the appendix.

Regarding improvements in describing our method.

Thank you for suggesting that the paper could be improved and made easier to understand. We have updated parts of the paper that could cause confusion (namely, the action representation, training procedure, use of the generated pseudo-demonstrations, and scaling trends) and added an additional figure to better convey the training procedure to the appendix. We also added new videos and visualisation on our anonymous webpage (https://sites.google.com/view/instant-policy/), that can help to understand our work better. If there are any particular parts of the paper that would benefit from further improvements, please let us know.

Regarding how actions are represented.

Thank you for pointing out how our paper could be improved and made easier to understand. In general, we represent a robot action as the relative transformation of the end-effector that would move the gripper from its current pose to a new desired pose and a binary open-close gripper action. These transformations are restricted to, at most, moving the gripper 1cm and rotating it by 3 degrees between subsequent timesteps, and, when time-scaled, can be interpreted as end-effector velocities.

To incorporate such actions into our graph representation, we transform the nodes representing the gripper, moving them to the pose they would reach if the actions were executed, and set their features based on binary gripper actions. Edges between the nodes representing the current observation and those representing the actions are equipped with positional embeddings that encode the gripper's movement between timesteps. These embeddings are not learned but are expressed as a combination of sine and cosine functions, similar to the approach used in the original NeRF paper. This integration allows low-level robot actions to be represented in the graph, with the end-effector's motion captured through the gripper nodes' positions and binary gripper actions encoded in their features. We have updated our paper to convey this more clearly.

At inference, our goal is to move the nodes representing robot actions to their desired position and adjust their features such that they correspond to correct binary gripper actions. This is done using a diffusion process, starting from randomly initialised actions, as described by Equations 4 and 5. Please see the newly added visualisations of the diffusion process on our anonymised webpage at https://sites.google.com/view/instant-policy/.

Regarding the action nodes and alternative approaches.

That is an excellent point and a good suggestion for the discussion in the related work. For completeness, we have included the work suggested by the reviewer in the related work section.

Our approach can be viewed as optimising the deformation of the graph, but we deform parts of the graph representing actions and do so using a learnt diffusion process, rather than analytically. There are several reasons for this. First of all, we represent robot actions as displacements from the current configuration (i.e. velocities) rather than the desired robot configurations, allowing us to directly execute the predicted actions which results in reactive and robust manipulation policies. Therefore, to represent this relative movement, we include both, the current observation and the actions in the graph representation. Second, our context is composed of robot trajectories rather than a single desired pose. As the reviewer suggests, this problem could be tackled by reaching relative spatial relationships in sequence. However, it would require a module for understanding what the next goal state in a trajectory should be or rely on open-loop execution. Instead, we allow the model to learn at what stage of the task it currently is and what points in the trajectory are relevant (please see the newly added visualisation in our anonymised webpage). Furthermore, our approach allows us to easily integrate multiple demonstrations into the context, enabling generalisation to novel object geometries.

Thus to answer the original question of why action nodes are needed: they express low-level robot actions instead of specific desired key poses, they allow us to simultaneously reason about and predict multiple actions, and circumvent the need for manually identifying the stage of the task (as it is done by nodes representing the current observation).

Regarding the effort needed to generate pseudo-demonstrations.

Generating pseudo-demonstrations is an extremely cheap and fast process when compared to collecting the demonstrations in the real world or handcrafting tasks in simulation. Play data that the reviewer refers to is easier to obtain than purposely collecting demonstrations. However, it still requires significant human time and effort, even though it might not be as mundane as doing the same task over and over again to collect a dataset of demonstrations. Additionally, utilising play data to train an In-Cotext model is not straightforward, as we need multiple trajectories performing semantically the same task. Pseudo-demonstrations, on the other hand, can be generated procedurally without much human effort and they do not need to be dynamically feasible as during deployment proper demonstrations will be used to define real tasks. Thus, even computational requirements are substantially smaller due to not needing to run physics simulations.

Thank you for your in-depth response to my initial review. I have taken a look at it as well as the reviews and responses from the other reviewers.

Please incorporate some intuition regarding the action nodes into the final paper. I am convinced by your argument about pseudo-demonstration and novel objects, so I have increased my score to accept.

Regarding the partially observable settings.

We agree that sometimes occluded objects make the problem more challenging, and we acknowledged this in the limitations section of the original paper. In our real-world experiments, we observed strong performance even with partial point clouds, which are sometimes severely occluded by the robot's gripper. We hypothesise that this is mainly due to the use of local geometry features, which allow the model to identify specific parts of the objects and act accordingly, even if the whole object is not observed.

As the reviewer suggests, incorporating past observations could help further mitigate the impact of partial observability, which would be a simple extension of our graph representation. Additionally, while we currently rely on segmented point clouds, our framework could readily be extended to use point clouds of the entire scene. Randomising distractor objects in the scene during pseudo-demonstration generation would add robustness against occlusions. Finally, integrating additional semantic features into the observations would further increase the model's ability to reason about partial observations.

Thus, although partial observability and clutter can cause issues for the current implementation of Instant Policy, the overall framework has the potential to be extended to mitigate these issues.

Regarding collision avoidance.

We agree with the reviewer that collision avoidance is not addressed in the current implementation of Instant Policy, and we point it out in the limitations section. First, we would like to note that the movement in a straight line that the reviewer mentions is observed only when the robot is approaching the object, and non-linear motions are performed while completing various tasks, such as opening a box.

The limitation of not accounting for collision avoidance itself, however, is not a limitation of the proposed framework in general, but rather of the specific implementation in our experiments. Our formulation of the problem and the graph representation allows us to use dense point clouds of the whole scene, which would provide sufficient information to reason about collision avoidance. Similarly, pseudo-demonstration generation could then be extended to use motion planning to avoid randomised obstacles when approaching the objects of interest. Finally, although in this work, we express the robot state as the pose of the end-effector, the graph representation can be extended to incorporate different parts of the robot, allowing control of the full configuration of the robot through known robot kinematics.

Regarding a baseline using the next-action prediction architecture operating over the graph structure input.

In our ablation experiments, we do study the influence that the diffusion process has. One of the ablation variants of Instant Policy, described in Section 4 ("No Diff" in Table 2), does not use diffusion but rather predicts the next action directly, as the reviewer suggests. This comparison allows us to estimate the added value of using diffusion on graphs rather than relying solely on the graph representation. As shown in Table 2, we see a 29% drop in the success rate, indicating that the diffusion process is indeed a crucial part of the method. We attribute this boost in performance when diffusion is used to two factors. First, the diffusion objective is overall more expressive, enabling capturing complex multi-modal distributions. Second, iteratively adjusting the nodes in the graph representing the actions ensures that observations and actions are expressed in the same graph space, and the model does not need to learn complex and highly non-linear mapping between these spaces. Intuitively, such an approach allows the model to 'imagine' spatial implications of considered actions, as nodes are updated as if the actions were executed.

Regarding the assumption of using object segmentation.

While the current implementation of Instant Policy assumes access to segmentations, the overall framework could be easily extended to utilise dense point clouds of the entire scene. And as the reviewer points out, the proposed method already demonstrated good performance in real-world settings using noisy observations.

Thank you for the response. The response sufficiently addresses my listed weaknesses, so I raise my score to accept.

Regarding the description of the training procedure.

Thank you for pointing out how our paper could be improved and made easier to understand. The training process can be mainly understood by looking at Equations 1 and 2. In short, training includes 4 main steps: 1) noise is added to the ground truth actions, 2) noisy actions 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. We have added a section in the appendix, describing the training procedure and included an additional figure that more clearly illustrates the training process.

Regarding how the model utilises pseudo-demonstrations and how they differ from other sources of data.

Pseudo-demonstrations serve as a primary source of training data for our model. We generate sets of demonstrations performing the same type of pseudo-task (e.g. 3 pseudo-demonstrations for each pseudo-task). Then, to train our model, we use a subset of these demonstrations (e.g. 2 out of 3 demonstrations) to define the context and predict actions from the pseudo-demonstration that was left out when defining the context. Pseudo-demonstrations are never used during inference to define the context.

RLBench and real-world data, on the other hand, can be used either as an additional source of training data or to define the context at inference. For training, it follows the same format as the pseudo-demonstrations. A few demonstrations are collected performing the same task, and some are used to define the context while inferring actions from the other. When used during inference, they are solely used to define the task and the context. We have updated the explanation regarding the use of pseudo-demonstrations and how they differ from other sources of data in the updated paper.

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)?

We adapted all the baselines to operate on point cloud observations using the same pre-trained encoder used for Instant Policy, to have the same number of trainable parameters as our model, and trained them on the same data (removing parts of the baselines that depend on language annotated data, as pseudo-demonstrations do not have such information). We did this to ensure a fair comparison between the methods, avoiding using different modalities and data sources. We indicate this by adding an asterisk to the baselines and describing it on line 325 in the original paper and line 328 in the updated one.

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?

Thank you for pointing out a part of the paper that was not initially clear. All models were trained primarily using pseudo-demonstrations. In simulation, we tested two versions of the models: one trained using only pseudo-demonstrations and another trained with a combination of pseudo-demonstrations and a small amount of data from RLBench tasks. In the real-world setting, the model trained on both pseudo-demonstrations and RLBench data was further fine-tuned by adding a small amount of real-world data (five demonstrations for five tasks, none of which were used during evaluation).

In summary, pseudo-demonstrations allow us to generate arbitrary amounts of training data, while incorporating small amounts of in-domain data (e.g., from RLBench or the real world) improves model performance by adjusting to, for example, imperfect segmentations and noisy observations.

What is the main difference between your work and One-Shot / Few-Shot Imitation Learning approaches, such as [1]?

While works, such as the one reviewer points out, rely on some sort of formulation of one or few-shot imitation learning as a pose estimation problem, our work directly reasons about the demonstration trajectories and predicts low-level robot actions, such as end-effector velocities. In this way, we allow the model to reason about what stage of a task it is at and what the appropriate actions should be, resulting in reactive and robust manipulation policies (please see newly added videos on our anonymised webpage at https://sites.google.com/view/instant-policy/). Additionally, our method allows us to easily integrate additional demonstration trajectories, enabling generalisation to novel object geometries. For completeness, we added the work the reviewer suggested to the related work section.