Learning Video-Conditioned Policy on Unlabelled Data with Joint Embedding Predictive Transformer

We are thankful to the reviewer for their meticulous and in-depth review and for providing valuable insights. We are pleased to receive recognition for the effectiveness and the reasonability of our method. Next, we will respond to the weaknesses.

W1. Does BC+IDM train both the video encoder and the inverse dynamics model on demonstration videos? What if a fixed pretrained video encoder is used instead? This could ensure that the representations of demonstration and expert videos are in the same space, potentially aiding in the learning of generalizable actions.

In the BC+IDM method, there are two video encoders involved: one is used for processing trajectories, and the other is used to encode the prompt video. From the context, the reviewer is likely referring to the former. Indeed, when learning IDM, the video encoder that processes trajectories is trained on demonstration data. It is reasonable to replace it with a fixed pre-trained video encoder to reduce distribution shift and thus potentially help generalize actions. Most pre-trained video encoders use a kernel size of 2 for temporal compression, which results in the representation length being halved in the temporal dimension. Similar to our approach in visual prior injection with VideoMAE v2, we repeat all frames within the video except for the first frame to ensure consistent length. Using VideoMAE v2, experimental results are shown in Appendix C.2. After switching to a fixed video encoder, BC+IDM is denoted as BC+IDM$^\star$. There is a slight improvement in success rates on unlabeled tasks, but there is no significant gain on unseen tasks. In all, the improvement from using a fixed video encoder is limited.

W2. How does JEPA influence performance? Did the authors experiment with other pretrained models, including different architectures and pretrained data?

There might be a slight misunderstanding at this point. We do not use a pre-trained JEPA; instead, we extend JEPA to a Transformer architecture for decision-making tasks and train it on the data we collect.

W3. The comparison of different visual prior injection methods appears inconsistent. The four models rely on distinct visual signals as inputs, which has led the authors to conclude that optical flow is the most suitable prior. However, these models were trained on different datasets and under varying conditions (self-supervised or supervised). It is crucial to detail these differences, as they may influence the performance.

Thank the reviewer for pointing this out. In our visual prior experiments, the inputs for the four models are indeed inconsistent (single frame, two frames, multiple frames). Although all models use a visual Transformer architecture, there are differences in their structures (such as patch size, number of layers, and number of heads). This certainly warrants further explanation, and we have added this information in Appendix A.2.

Additionally, as we mentioned in the paper, "The poor performance of other visual priors may result from the inherent unsuitability of these priors for visual control tasks or the inadequacy of the injection approach, requiring further exploration in future work." We don't conclude that optical flow is the most suitable prior for video-conditioned policy learning. It just performs better in our naive injection design. These visual encoders have been trained on large datasets across various visual tasks and possess certain generalizable priors for handling visual inputs. Our visual prior experiments aim to determine whether a joint embedding that predicts these representations with visual priors can further compress visual information and enhance generalization. From the perspective of introducing more generalizable visual priors, using different priors to process visual inputs makes sense. However, successfully learning a predictive joint embedding of these priors and effectively improving performance may depend on selecting more refined injection methods and better training method designs. We only reveal the potential of JEPT in this aspect, and further detailed research on other visual prior injections is left for future work.

We appreciate the reviewer for the detailed and comprehensive reviews, and for offering insightful feedback. We are delighted to receive recognition for our method and experiments. Below, we will respond to the concerns raised by the reviewers.

W1. More qualititative samples might be also helpful to see the experiments in action, and to show how the rollout of actions and $h_t$ changes over time.

Thanks for pointing this out. We add additional visualizations of the actions and $h_t$ in the rollouts in Appendix C.3 (Fig 7 and 6).

W2. Lack of any real world examples is one weakness, it would be nice if the paper can show these nice properties of giving a simple video demo, and the robot can mimic it in real world setting, that would be great, and can also answer many questions regarding the usefulness of the work.

We sincerely thank the reviewer for highlighting this important point. We wholeheartedly agree that training on real-world datasets is vital and represents a significant challenge. Unfortunately, due to our limited experimental resources and the tight time constraints during the rebuttal period, we are unable to conduct this experiment at present. In our setting, the model is trained using aligned pairs of video prompts and expert video or demonstration data. For real-world environments, we do not yet have access to the necessary equipment to perform these experiments. Additionally, for simulated environments using real-world video prompts, we lack the aligned data needed between real-world video prompts and simulator environments. In fact, the experiments in Vid2Robot have demonstrated that simple designs can achieve visual imitation in the real world through training with a large amount of aligned data. We acknowledge the importance of real-world experiments and will address this in future research when we overcome these limitations. Thanks for the valuable feedback and understanding.

W3. It would be also nice, if the paper can address how the patch tokens are encoded, and how they are handled during the autoregressive stage of the policy.

Thank the reviewer for pointing this out. In the main text, we use the term 'spatial encoder' as a shorthand to refer to this component. We have provided an explanation in the appendix. Specifically, we use a Vision Transformer (ViT) to encode the visual input and flatten all the patch tokens from the ViT to feed into the Perceiver IO, which compresses the number of tokens using learnable queries. The structure of the Perceiver IO is the same as in the original paper, with some hyperparameters modified. During the autoregressive phase, the compressed visual tokens are used as observation tokens for the current timestep and are input into the Video-Conditioned Causal Predictor. In the causal predictor, the causal mask is applied only temporally, allowing patch tokens within the same timestep to attend to each other. We have highlighted these parts in green in the Appendix.

We express our gratitude to the reviewer for the meticulous and thorough review, and for providing valuable insights. We are so pleased to receive recognition for the method and the experiments. Below, we will respond to the weaknesses and issues raised by the reviewers:

W1: The paper is not easy to follow, especially for the researcher not in the robotics field. A further polishment of the manuscript would be great

• The authors are suggested to clarify each term such as prompt video (which is used as a whole global representation in the method), expert video. A simple explanation would be appreciated;
• Based on my understanding, the key design of the proposed method is the split of behavior cloning into two sub-tasks the visual transition prediction and the inverse dynamics learning, which can help utilize the mixture dataset, and reduce the dependency on action annotation. However, the authors did not discuss much on why this design works and why it is suitable for the situation when the annotation of actions is limited;

Thanks for pointing this out. To more clearly present our setting, we have replaced the first figure with a clearer one and made some adjustments in the text. The core design of the algorithm works for two reasons: on one hand, visual transition prediction leverages the visual transition information in action-labeled and unlabeled expert videos; on the other hand, the generalization of inverse dynamics knowledge allows for inferring reasonable actions on unlabeled tasks. For a visual imitation learning task, visual transition prediction is the more task-specific part, requiring more data to learn, whereas inverse dynamics is more task-unified and can generalize even with limited action annotation. We have also made slight adjustments to emphasize this.

W2: From the view of pretraining, a well pre-trained network needs tremendous unlabelled data to fit a good network. For video/image understanding, the requiremenet of large-scale data is easy to fullfill. However for the robotics task, the large-scale data is relatively diffcult to acquire.

W3: Considering the diversity of robot configuration, the action state or the description of action in each robot is different. This paper does not mention how to handle the heterogeneity issue among robots, which limits the scenario span of the proposed method.

Q1. How to handle the heterogeneity issue? Is the action space designed for only one robot configuration?

These three pieces of feedback seem to be related. From a data perspective, large-scale data for decision-making problems is emerging, such as Open-X and Vid2Robot. Our work, in fact, reduces the burden of constructing these large-scale datasets by using unlabeled data to expand the dataset for video-conditioned policy. When training on large-scale data, the heterogeneity issue absolutely needs to be taken into consideration. We appreciate the reviewer for pointing this out, which is very enlightening. In our setting, only one action space is included. For the mixed data from multiple action spaces, the impact on visual transition prediction is smaller because it focuses more on aligning with the prompt video at the visual level. The inverse dynamics part, which directly involves multiple action spaces, can be handled more crudely by merging them into a large action space and continuing to use a shared network for learning, or more finely by equipping each action space with a decoder head to achieve mapping to different spaces. Of course, this is only a possible extension of our framework for data with multiple action spaces, and the specific effectiveness needs to be experimentally verified in future work. Due to our limited experimental resources and the time constraints during the rebuttal period, we are currently unable to conduct experiments in this aspect. We do thank the reviewer for the very enlightening suggestion. Our current work focuses on incorporating unlabeled data into video-conditioned policy, and we will further consider expanding our training dataset from this perspective.

Q2. What is the ratio of $D_{demo}$ and $D_{vid}$ in the mixture dataset? From the view of model pretraining, does the $D_{vid}$ require much more samples than $D_{demo}$? How does the ratio influence the final performance? An involvement of the ablation study on the ratio would be appreciated.

Q3. Is the hyperparameter $c$ in Eq.9 also influenced by the ratio of $D_{demo}$ and $D_{vid}$?

This is an excellent aspect of the ablation study, and we thank the reviewer for pointing it out. In the original ablation, we adjusted the number of tasks in $D_{vid}$​ to change its size, which alters the amount of data used for training to some extent. To address this, we conduct supplementary experiments, adding data from 6 more tasks in Metaworld, adjusting the ratio of $T_{vid}$​ and $T_{demo}$​ across a total of 20 training tasks, and testing different values of hyperparameter $c$ in Eq.9 under different ratios. We report the average success rate on unseen tasks, and the results are updated in Appendix C.2.

When the number of  $T_{demo}$​ is within a certain range (5-12), the ratio has little impact on success rate. When the number of  $T_{demo}$​ is too small, it is difficult to learn inverse knowledge sufficiently from the extremely limited action annotations, thus affecting the success rate. Additionally, when the proportion of $D_{vid}$​ increases, a larger $c$ may be needed to achieve good performance.

Q4. Considering $E_{prt}$ is a global representation of the prompt video, it would be interesting to see how $E_{prt}$ of videos from different action categories distribute in the feature space. Ideally, there should be high differentiation among the $E_{prt}$ from different action categories such that the model will not be misled by the guidance from $E_{prt}$.

We encoded the video prompts from 18 Metaworld tasks using the video encoder learned by JEPT, flattened these $E_{prt}$​, and visualized them using t-SNE. The figure is put in Appendix C.3 (Figure 5). Indeed, clear differences among different tasks can be observed, indicating that the model can learn distinctions between different tasks to avoid misleading. Additionally, we observe a degree of similarity between related tasks, such as ‘Faucet Close’ and ‘Faucet Open’, as well as ‘Window Open’, ‘Window Close’, and ‘Drawer Close’. This similarity reflects the inherent relationships between tasks, which may facilitate knowledge transfer during the learning process.

We are grateful to the reviewer for the meticulous and thorough review, and for providing insightful feedback. We are pleased to receive recognition for the comments such as 'novel', 'address real-world problems', and 'solid and thorough experiments'. Below, we will respond to the weaknesses raised by the reviewers:

W1: The model’s performance appears sensitive to the choice of visual priors. It is unclear how to select desired priors across different environments, or when facing a new task how to choose the prior.

To better illustrate the impact of visual prior injection across different environments, we conduct additional experiments on robosuite, replicating the visual prior injection method used in the meta-world experiment. The results, updated in Appendix C.1 and listed below, suggest that using visual encoders trained for optical flow estimate tasks still yields better results. This somewhat indicates that using optical flow can be a choice for different environments in our injection implementation. As we have mentioned in the paper, 'The poor performance of other visual priors may result from the inherent unsuitability of these priors for visual control tasks or the inadequacy of the injection approach, requiring further exploration in future work.' These visual encoders have been trained on large amounts of data across many visual tasks, and possess certain generalizable priors in handling visual inputs. Our visual prior experiments aim to determine whether a joint embedding predictive of these representations with visual priors can further compress visual information and enhance generalization. From the perspective of introducing more generalizable visual priors, these different priors for processing visual inputs ought to make sense. However, learning predictive joint embedding of these priors and effectively improving performance may depend on the choice of more refined injection methods and better design of training methods. We only reveal the potential of JEPT in this aspect, and further detailed research on other visual prior injections is left for future work.

W2: The paper primarily validates JEPT on simulated datasets. The paper can be improved by adding real-world datasets for more complex environments, strengthening the claims of generalizability and practical applicability.

We sincerely thank the reviewer for highlighting this important point. We fully acknowledge the critical role that training on real-world datasets plays in advancing our field. Unfortunately, due to our current limitations in experimental resources and the time constraints of the rebuttal period, we are unable to conduct such experiments at this time. Our current model is trained using aligned pairs of video prompts and expert video or demonstration data. For real-world environments, we lack the necessary equipment to run experiments. Similarly, for simulated environments using real-world video prompts, we lack the aligned data necessary to bridge real-world video prompts with simulator environments. Our research focuses on introducing a training method that integrates demonstrations with unlabeled expert videos. We recognize that utilizing extensive real-world video data could significantly enhance our work, and we are eager to explore this avenue in future studies.

Dear Reviewer 4tQo,

We sincerely appreciate your invaluable feedback, which plays a crucial role in enhancing the quality of our work.

As the rebuttal deadline approaches, we want to check whether our response addresses your concerns. If there are any additional questions or issues you would like us to address, please let us know. We would be grateful if you could provide stronger comments to support our work after reviewing our response.

Thank you once again for your support and guidance.

Best regards,

Authors of Submission6280