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

Authors: Hao Luo, Zongqing Lu
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Learning from Videos, Video-Conditioned Policy

Abstract:
The video-conditioned policy takes prompt videos of the desired tasks as a condition and is regarded for its prospective generalizability. Despite its promise, training a video-conditioned policy is non-trivial due to the need for abundant demonstrations. In some tasks, the expert rollouts are merely available as videos, and costly and time-consuming efforts are required to annotate action labels. To address this, we explore training video-conditioned policy on a mixture of demonstrations and unlabeled expert videos to reduce reliance on extensive manual annotation. We introduce the Joint Embedding Predictive Transformer (JEPT) to learn a video-conditioned policy through sequence modeling. JEPT is designed to jointly learn visual transition prediction and inverse dynamics. The visual transition is captured from both demonstrations and expert videos, on the basis of which the inverse dynamics learned from demonstrations is generalizable to the tasks without action labels. Experiments on a series of simulated visual control tasks evaluate that JEPT can effectively leverage the mixture dataset to learn a generalizable policy. JEPT outperforms baselines in the tasks without action-labeled data and unseen tasks. We also experimentally reveal the potential of JEPT as a simple visual priors injection approach to enhance the video-conditioned policy.

Decision: accept (poster)
Reviewer scores: 8, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730099782517

thread_type:
official_review

<thread_object>
creator: reviewer_zv2n

<content>
summary:
This paper introduces the Joint Embedding Predictive Transformer (JEPT), a novel approach for video-conditioned policy learning that leverages both expert demonstrations and expert videos paired with prompt videos to reduce the need for action label annotation. To effectively learn from the mixture datasets, JEPT decomposes the learning task into two subtasks: visual transition prediction and inverse dynamics learning. By jointly training these subtasks, JEPT enhances the generalization of the video-conditioned policy. Experimental results demonstrate that JEPT outperforms baseline methods.

strengths:
1. The proposed JEPT framework for learning from mixed datasets is clear and reasonable.
2. Experiments on the Meta-World and Robosuite benchmarks demonstrate the effectiveness of JEPT.
3. Results indicate that visual priors can be integrated into JEPT to improve generalization.

weaknesses:
1. 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.
2. How does JEPA influence performance? Did the authors experiment with other pretrained models, including different architectures and pretrained data?
3. 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.

questions:
1. 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?
2. How does JEPA influence performance? Did the authors experiment with other pretrained models, including different architectures and pretrained data?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730596187597

thread_type:
official_review

<thread_object>
creator: reviewer_grhe

<content>
summary:
the paper propose Joint Embedding Predictive Transformer (JEPT), where the model can be conditioned on some video based demonstations, and the causal transformer also learns to predict the next action and observation (in pixels) space. The prompt video encoder, can be used to give demonstactions just from a video, this is great because giving an intertion is hard with language, and even babies do this from visual demonstations. The JEPA style video model, encodes causally the state and actions, and the joint encoder maps actions to states and there are jointly optimized during training. The paper evaulates these models on various benchmarks, and shows good performance.

strengths:
The main stength of this paper is to able to prompt the robot from simple visual demonstrations. This has very nice benifits from not having any dependcy from langauges and from development physcology it make sense, as babies learn/ and repeat actions by watching others.  

The paper has through experiments on meta world and robosuit, and in most of the cases the performance is good compared to previous works. 

Ablations also clearly shows the benifits of joint embedding model.

weaknesses:
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. 

Lack of any real world examples is one weakness, it would be nice if the paper can show these nice propeties 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 usefullness of the work. 

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

questions:
please look at my strengths and weakness sections, and if you can adress the weakness section, i am happy to change my ratings.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730958312738

thread_type:
official_review

<thread_object>
creator: reviewer_yxfu

<content>
summary:
A training strategy for video-conditioned policy generation is proposed in this manuscript. Considering the challenge of annotating action labels, the author propose to split the behavior cloning task into two sub tasks to reduce the dependencies on annoatated data. The visual transition prediction task predicts the observation embedding of next time step in the feature space, and the inverse dynamics leanring task estimates the current action by taking the consecutive two observations of current time step. The training of the visual transition prediction needs only the prompt videos and expert videos without action annotation. The training of inverse dynamics learning requires the prompt video, the expert video and the corresponding annotations. This training strategy is quite similar to the popular strategy pre-training (visual transition prediction) + supervised learning (inverse dynamics learning) . The proposed method is evaluated on two benchmark datasets, the Meta-World task and the Robosuite task, and the experimental results show the effectiveness of the proposed mtehod.

strengths:
1. The proposed method aims to introduce prompt video as guidance for task imitation learning, which is interesting.
2. The involvement of visual transition prediction for observation embedding prediction is good. It can not only help reduce the cost of action annotation, but also make the task easier to learn. The prediction of observation embedding in the next time step in feature space is easier than the direct embedding extraction from the raw observation of the next time step. However, the premise the embedding space is well learned.
3. The proposed method is extensively evaluated and achieves good performance.

weaknesses:
1. 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 apprecaited;
* 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 leanring, which can help utilize the mixture dataset, and reduce the dependency on action annoataion. 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 limitted;

2. 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.

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

questions:
1. How to handle the heterogeneity issue? Is the action space designed for only one robot configuration?
2. 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 the ratio influence the final performance? An involvement of the ablation study on the ratio would be appreciated.
3. Is the hyperparameter $c$ in Eq.9 also influenced by the ratio of $D_{demo}$ and $D_{vid}$?
4. 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}$.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731046069189

thread_type:
official_review

<thread_object>
creator: reviewer_4tqo

<content>
summary:
Video-conditioned policy learning often requires extensive action-labeled demonstrations, which are costly and time-consuming to acquire.  To address this problem, a method named Joint Embedding Predictive Transformer (JEPT) is introduced in this paper. JEPT jointly
learns visual transition prediction, which predicts the next visual state in a sequence based on the current observation and prompt video with unlabelled data, and inverse dynamics, which learns to infer actions that cause the transition between two states given the labeled data. The model is evaluated on two benchmarks, Meta-World and Robosuite, and is compared with recent state-of-the-art methods. Experimental results show that JEPT achieves superior performances.

strengths:
- The paper is novel and addresses real-world problems. More specifically, decomposing the behavior cloning into visual transition prediction and inverse dynamics learning is novel, prior approaches mostly rely on labeled data or direct behavior cloning, while JEPT makes use of unlabelled data and largely improved generalization ability and performance.
- The experiments are solid and thorough. JEPT is compared with recent state-of-the-art methods. The robustness analysis and visual prior injection studies are interesting and inspiring.

weaknesses:
- 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.
- 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.

questions:
Please refer to the Weaknesses section.
</content>
</thread_object>

</thread>


Metareview:
metareview:
(a) Scientific Claims and Findings

The paper introduces the Joint Embedding Predictive Transformer (JEPT), an approach for video-conditioned policy learning that reduces the need for action-labeled demonstrations. JEPT decomposes the learning task into visual transition prediction and inverse dynamics learning, leveraging both labeled and unlabeled data to enhance generalization. The model is evaluated on Meta-World and Robosuite benchmarks, showing better performance compared to state-of-the-art methods. Reviewers highlight the method's ability to utilize unlabelled data effectively and its potential to improve generalization in video-conditioned policy learning.

(b) Strengths

Reviewer 4tQo acknowledges the novelty of decomposing behavior cloning into two sub-tasks, which improves generalization and performance. yXfU finds the introduction of prompt video guidance and visual transition prediction interesting and effective in reducing annotation costs. grHE appreciates the ability to prompt robots from visual demonstrations, aligning with developmental psychology insights. Zv2N commends the clear and reasonable framework of JEPT and its demonstrated effectiveness on benchmarks.

(c) Weaknesses

The reviewers identify several weaknesses. 4tQo notes the model's sensitivity to visual priors and the lack of real-world dataset validation. yXfU suggests further polishing for clarity and questions the handling of heterogeneity among robots. grHE points out the absence of real-world examples and the need for more qualitative samples. Zv2N raises concerns about the training of video encoders and the consistency of visual prior injection methods, suggesting that differences in training conditions may affect performance.

(d) Decision Reasons

The AC agrees with the reviewer's recommendation to accept the paper. The final decision to accept the paper is based on its novel approach to video-conditioned policy learning and the strong empirical results demonstrated on benchmark datasets. The method's ability to leverage unlabelled data and improve generalization is a contribution, as highlighted by reviewers 4tQo and yXfU. While there are concerns about the model's sensitivity to visual priors and the lack of real-world validation, the overall strengths in innovation, experimental validation, and potential impact on the field outweigh these weaknesses. The paper's contributions to reducing annotation costs and enhancing policy learning make it a valuable addition to the conference.