Learning Video-Conditioned Policy on Unlabelled Data with Joint Embedding Predictive Transformer
Authors: Hao Luo
Zongqing Lu
Venue: ICLR 2025 Poster (2025)
Keywords: Learning from Videos
Video-Conditioned Policy
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose the Joint Embedding Predictive Transformer (JEPT) to train a generalizable video-conditioned policy by leveraging both action-labeled and unlabeled expert videos.

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.

Published as a conference paper at ICLR 2025
LEARNING VIDEO-CONDITIONED POLICY ON UNLA- BELLED DATA WITH JOINT EMBEDDING PREDICTIVE TRANSFORMER
Hao Luo' and Zongqing Lu
1 School of Computer Science, Peking University BBeijing Academy of Artificial Intelligence
ABSTRACT
The rideo-conditioned policy takes prompt videos of the desired tasks as condition and regarded for prospective generalizability. Despite its promise, training a video-conditioned policy is non-trivial due the need for abundant demonstrations. In some tasks, the expert rollouts are merely available as videos, and costly and time-consuming efforts required annotate action labels. To address this, we explore training video-conditioned policy on a mixture demonstrations and unlabeled expert videos reduce reliance extensive manual annotation. We introduce the Joint Embedding Predictive Transformer (JEPT) to learn a video- conditioned policy through sequence modeling. JEPT 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 generalizable policy. JEPT outperforms baselines in the tasks without action-labeles data and unseen tasks. We also experimentally reveal the potential of JEPT simple visual priors injection approach to enhance the sideo-conditioned policy.
Action-Labeled Dataset
ApyAvy
(Ar. Ay.As)
(As.
obs. O
act.
Behavior Cloning
Action Predictor
grangtviee
demonstration
Unlabeled Dataset
AAx,Ay,4e)
prompt
prompt video
video
Jobs.
Joint Embedding
next abs.
Joint Embedding Predictor
Joint Embedding Encoder
JEPT
prompt
Action Predictor
(Ax, (dx,Ay,44 Ay,
Figure 1: Illustration of the BC and JEPT for visual-conditioned policy learning. Video-conditioned Behavior Cloning (BC) (right top) directly models the policy from the demonstrations. JEPT (bottom) decomposes imitation learning into visual transition prediction (solid arrows) and inverse dynamics learning (dashed arrows). BC is constrained to the action-laheled dataset, whereas JEPT enables visual transition prediction unlabeled data and leverages the mixture dataset.
"Correspondence Zongqing Lu <zongqing.loe lu www..@pk.uu..........mmm

---

Published as a conference paper at ICLR 2025
1 INTRODUCTION
Learning generalizable policies through multi-task imitation learning remains considerable chal- lenge. A generalizable policy should adapt seamlessly to novel tasks whose demonstrations are absent from the training dataset. In response to this, the video-conditioned policy (Jiang et al., 2023; Bahl et al., 2023; Shah et al., 2023) has garnered attention for its superior potential to generalize across tasks. A video-conditioned policy takes a prompt video depicting the task as condition and executes the desired task in the dynamics encountered. Utilizing a prompt video as a policy condition exhibits a better potential for generalization compared to other forms of task specifications, such as instructions or goals given in language (Brohan et al., 2023; Padalkar & et al., 2023; Jiang et al., 2019; Yenamandra et al., 2023) image (Bousmalis et al., 2024; Lee et al., 2020; Du et al., 2024; Nair et al., 2018) formats. The abundant information conveyed in videos offers sufficient policy guidance for novel tasks, thus circumventing the misleading caused by language ambiguities or the absence of process depiction in static images (Jain et al., 2024).
Although comprehensive information in prompt videos imbues video-conditioned policy with flex- ibility and generalization, learning such a policy is more than trivial. Imitating the tasks depicted in prompt videos requires models capable of both temporal reasoning to understand the task and fine- fine-grained control to replicate it. To this end, previous methods (Shah et al., 2023; Jain et al., 2024) typically conduct behavior cloning on large datasets of paired prompt videos and expert demonstra- tions. However, procuring expert demonstrations can be prohibitively costly and time-consuming. In the tasks where expert policies are hard to obtain, we can acquire the videos of the expert rollouts through human intervention, but these videos are devoid of action labels. Conducting behavior cloning in this scenario requires additional efforts to annotate the action labels. This raises the question: Can these expert videos' be effectively harnessed without additional annotation, thus allevi- ating the data collection burden? Accordingly, we explore the possibility of training a generalizable video-conditioned policy dataset containing both expert demonstrations and unlabeled expert videos paired with prompt videos.
Unlike methods that directly imitate expert demonstrations, behavior cloning is infeasible on the mixture dataset due to the absence of the action labels from the data samples with expert videos. Intuitively, we decompose the behavior cloning into two synergistic subtasks to fully exploit the mixture dataset of expert demonstrations and expert videos. Specifically, the process of video- conditioned imitation can be broken down into visual transition prediction and inverse dynamics learning. In the visual transition prediction, the model learns how the prompt videos should manifest in the dynamics of the tasks. In inverse dynamics learning, the model infers the actions required to realize the visual transition. By combining them, the model can learn to predict plausible future observations and then convert them into actions, enabling condiccond imitation.
This design potentially suits the setting of mixture dataset. On the one hand, visual transition prediction is more task-specific in visual imitation learning task and can be more sufficiently learned from both expert videos and demonstrations. On the other hand, although the inverse dynamics is merely embedded in the demonstrations, it remains universally applicable across tasks. Consequently, the expert videos could enhance the learning of visual transitions, while the inverse dynamics derived from the demonstrations could generalize to guide the execution of the tasks without action labels.
In light of these insights, we propose the Joint Embedding Predictive Transformer (JEPT) encap- sulate both visual transition and inverse dynamics for video-conditioned policy learning. Building on the framework of Decision Transformers (Chen et al., 2021; Lee et al., 2022; Furuta et al., 2022), we employ a video-conditioned Transformer-basee architecture to perform sequence modeling, serving as the policy. As shown in Figure 1, we modify the typical behavior cloning sequence into a two-step process. First, JEPT predicts the embeddings of the next observations, thus capturing visual transition. Then, conditioned on these predicted embeddings, JEPT predicts the corresponding actions, thereby learning inverse dynamics. Additionally, we integrate the Joint Embedding Predictive Architecture (JEPA) (LeCun, 2022) into the sideo-conditioned Transformer. Prior works (Assran et al., 2023; Bardes al., 2024) have demonstrated that JEPA is an effective visual representation learning ap- proach, compressing visual inputs into predictive features. By incorporating JEPA, we aim to learn
For simplicity and distinction, we use 'prompt videos' referring to videos depicting the desired task and used as the policy condition, while "expert videos' for the visual observation sequence of the expert rollouts.
2

---

Published as a conference paper at ICLR 2025
predictive visual representations of future observations, thereby improving the transfer of inverse dynamics between action-labeled and unlabeled data.
To evaluate the effectiveness of JEPT, we conduct experiments Meta- Mooa- (Yu et al., 2020a) and Robosuite (Zhu et al., 2020). In our Meta- World experiments. JEPT outperforms all the basclines regarding the average success rates of the tasks within the dataset and the unseen tasks. In the Robosuite experiments, we further validate the effectiveness of JEPT in handling tasks where there larger discrepancy between prompt videos and task dynamics. Additionally, we explore the potential of JEPT in injecting visual priors via adjusting the inputs of the joint embedding encoder. Our experiments on Meta-Wont Neta-World reveal that the optical flow priors can effectively improve the performance. These results demonstrate the effectiveness of JEPT in jointly leveraging the mixture dataset of expert demonstrations and expert videos to train a generalizable video-conditioned policy.
In summary, our contributions are as follows:
We propose a novel paradigm, JEPT, to learn video-conditioned policies using the dataset containing mixture of expert videos and demonstrations.
We design JEPT as a joint embedding predictive architecture to learn an abstract visual representation for learning a generalizable video-conditioned policy.
We explore the JEPT as an approach of visual priors injection and experimentally find that JEPT can effectively leverage visual prior knowledge and achieve better performance.
2 RELATED WORKS
Video Prompt Policy Learning. As a densely informative data form, video inherently contains abundant information for task completion when used as a task description. The detailed guidance of the task completion process in the video makes it possible to learn the policy generalizable to novel tasks beyond the training dataset. Some works (Finn et al., 2017; Duan et al., 2017: Yu et al., 2018) employ meta-learning methods to adapt the policy to novel tasks depicted by prompt video, which requires the similarity between the tasks. More recent works focus on learning policy via behavior cloning datasets of paired prompt videos and expert demonstrations. These works take the prompt videos as a condition input to the policy network and generalize to unseen tasks with corresponding prompt videos. Various auxiliary mechanisms have been explored achieve generalizable imitation, such as inverse dynamics prediction (Dasari & Gupta, 2021), cross-painiting (Chen et al., 2024), skill decomposition (Shin et al., 2024; 2023), hierarchical policy learning (Jain al., 2023), text-aligned representation (Jang et al., 2022), contrastive video encoder (Chane-Sanc et al., 2023) and observation- attentive representation (Jain et al., 2024). Some works (Sivakumar al., 2022) design object-centric decomposition for some specific manipulation tasks and achieve generalization within the same task category. The core of these methods remains direct behavior cloning, necessitating access to action-labeled demonstration data. Recent works (Jain et al., 2024; Shah et al., 2023; Jiang et al., 2023) construct large-scale datasets for video-conditioned policy learning, evaluating that large-scale and well well-aligned data effectively improve the rideo-conditioned policy learning in terms of generalization and success rates. However, these works require extensive demonstration annotation efforts, which is costly and time-consuming. Our work basically follows the video-conditioned policy learning setting. Still, we explore how the video-conditioned policy can learn effectively when the dataset is a mixture of expert videos and demonstrations paired with prompt videos.
Learning from Videos. Given the relative accessibility of video data compared to action-labeled demonstrations, some works explore leveraging unlabeled videos to aid in policy learning. Despite the absence of action labels, videos contain substantial decision- rich temporal information. A path for learning from videos involves self-supervised representation learning on vidcos, including masked autoencoder (Radosavovic et al., 2023; Xiao et al., 2022; Yang et al., 2024), temporal contrastive learning (Li et al., 2024; Nair et al., 2023), and video prediction (Seo et al., 2022; Luo et al., 2024). These methods aim to learn compact and informative representation from videos that aid in policy learning. Other approaches focus on learning reward functions (Escontrela et al., 2023; Yu et al., 2020b; Bruce et al., 2023) or value functions (Bobrin et al., 2024; Chang et al., 2022) to expedite online policy learning. Due to the absence of action labels, these methods require fine-tuning the downstream tasks for adaptation. Additionally, some research utilizes datasets combining videos and demonstrations, learning an inverse dynamics model (IDM) (Baker al., 2022; Schmeckpeper et al., 2021: Zheng et al., 2023; Zhang al., 2022; Kim
knowledge due to their
3

---

Published as a conference paper at ICLR 2025
Observation Tokens
her4
0gr
Rg
ac-1
Re+1
he
hyl+1
Rt
ht+3
(Optional) Action Tokens
Joint Embedding Tokens
Cideo-Conditioned Causal Predictor
Joint Embedding Encoder
E
prt
Prompt Video Encoder
0e'
Ot-1
O2
Ar++1
02'
O2
0c+1
p
Prompt Videos
Expert Videos Demonstrations
Figure 2: The architecture of the Joint Embedding Predictive Transformer (JEPT). The model consists of joint embedding encoder, prompt video encoder, and video-conditioned causal predictor. The demonstrations and the expert videos predicts the tokens causally. The visual embedding token of the next observation predicted to learn predictive joint embedding. We use t' == t k + 1 to denote the start of the context.
processed as sequence of tokens, and the causal predictor
al., 2023) latent inverse dynamics model (Schmidt & Jiang, 2024; Ye et al., 2022; Edwards et al., 2019) on the action-labeled data. The learned IDM can used to annotate the video data, facilitatiigg imitation learning derive policies. In these works, the decoupling of inverse dynamics learning and visual imitation may limit the generalizability of the IDM. In our work, however, the learning of visual transitions and inverse dynamics is conducted jointly, offering better generalization.
Joint Embedding Predictive Architecture. The Joint Embedding Predictive Architecture (JEPA) (LeCun, 2022) emerges as a promising representation learning framework. JEPA optimizes predictive loss on the encoded embeddings to learn an embedding space where the embeddings are predictable to each other. Unlike the generative architectures, JEPA optimizes the predictive loss on embeddings rather than raw data, potentially discarding extraneous information to learn more compact and meaningful representations. In order to train JEPA, it is a common practice to construct data pairs to predict and be predicted. JEPA is versatile framework. Various works employ different data pair construction methods enable JEPA to learn embedding spaces tailored to specific information, enhancing tasks such image classification (Assran et al., 2023), mask classification (Kim et al., 2024), video understanding (Bardes et al., 2024), and motion and content learning (Bardes al., 2023). Our work integrates JEPA with sequence modeling trajectories, providing a novel data pair construction method for JEPA. This extends JEPA to architectures akin to the Decision Transformers, aiming learn visual representations that facilitate inverse dynamics knowledge transfer.
3 METHODOLOGY
In this section, we describe how our Joint Embedding Predictive Transformer (JEPT) works on the mixture of expert demonstrations and expert videos paired with prompt videos.
3.1 PRELIMINARIES
We aim to learn video-conditioned policy across set of tasks. The environment is viewed as a collection of Partially Observable Markov Decision Process variants (POMDPs), with each task denote the state space, action space, transition function, observation space, and observation function, respectively, shared across all tasks. Specifically, V; comprises set of prompt videos depicting the task M,
represented by POMDP variant M, == (S, A. P, O, Ω, V.). Here, S,A.P, O,
Cideo-Conditioned Policy. During interaction with the environment, the agent receives a prompt video V e V. illustrating the desired task the start of each rollout. At each timestep the agent obtains a visual observation o, E O, derived from the current state 8, E S via the observation function Ω: 8 > O. Upon executing action a, A, the environment transitions to the subsequent
4

---

Published as a conference paper at ICLR 2025
state $1+1 according to the transition function P: S A S + [0. 1]. A iidoo-conditioned policy generates actions based on the visual observation context and the prompt video. Formally, a video-conditioned policy is denoted as π (0.10)=--+1a)) V). where 0 represents the visual observation context over the last l timesteps.
Mixture Dataset. We consider training the video-conditioned policy mixture dataset. Ideally, we collect three types of data for imitation learning: (1) the prompt video V == (v), U2...., UN) specifying the task with N frames, (2) the expert video O = (oo, o1...., OT), the visual observation sequence recorded alongside a T-ttep expert rollout of the task, and (3) the corresponding action sequence A = (ag, ai..... ar-1) executed during the same rollout. In the tasks where expert demonstrations are available, we assume access to the dataset Ddemo- and each data sample damoo E Demmo is in the form:
demo = ((v), 02. .2......N.. UN). (oo, 4o: 01,01,11111111 OT) m.).
Here, x² represents the expert policies, ~ denotes the rollout process, and indexes the dataset sample. Conversely, in the tasks where merely videos of expert rollouts are available, we denote this dataset as Dvia, and each data sample dvia E Dvia is in the form:
(2) Each data sample in Deemo consists of a prompt video V and a paired expert demonstration (O, A), while each data sample in Dvia consists of prompt video V and paired expert video O. In summary, our objective is to train the video-conditioned policy π (ar ((+1))) V) on a dataset mixed of these two kinds of data, denoted Demo U Dvid.
((1), 02, ",,...,, UN). (On, OT)...,o)~~T:;
3.2 MODEL ARCHITECTURE
Our JEPT employs Transformer-basee architecture to perform sequence modeling on the data, following the supervised learning paradigm of the Generalized Decision Transformers (Furuta et al., 2022). In this context, demonstrations are treated as sequences of observation and action tokens, where the sequence modeling serves as codoo-conditioned policy via predicting the corresponding action tokens. To perform the sequence modeling, JEPT comprises two kinds of modules: (1) visual encoders that process the high-dimensional visual input and (2) a causal predictor aggregating the visual representations to causally predict tokens in the sequence. The comprehensive architecture of JEPT is illustrated in Figure 2, and we elucidate each component in detail below.
3.2.1 VISUAL ENCODERS
Joint Embedding Encoder. A Joint Embedding Encoder is employed learn the visual representa- tion for each observation. In our design. the Joint Embedding Encoder captures information from both tinle-timestep and contextual levels. Structurally, the Joint Embedding Encoder comprises a spatial encoder for single-timestep encoding and a bi-directional temporal encoder for contextual encoding. Formally, the Joint Embedding Encoder Mots takes as input a visual observation context and outputs the visual representations:
h4) = Mobs (oi-R41, o1)
Each representation h2 consists of Nob. tokens. The visual embeddings of the next observation is predicted by the causal predictor and thus a joint embedding will be learned. Via learning joint embeddings predictive of the previous observation tokens, the encoder can compress the high- dimensional visual space into a low-dimensional space and discard extraneous information for effective and generalizable policy learning.
Prompt Video Encoder. Prompt Video Encoder encodes a prompt video into an embedding that serves reference for the policy. To specify the desired task. the embedding of the prompt video is used for an implicit understanding of task execution. Specifically, the Prompt Video Encoder comprises a per-frame encoder and Perceiver Resampler (Jaegle et al., 2022). The Perceiver Resampler aggregates frame embeddings from the per-frame encoder to produce an overall embedding Epm of the prompt video. Each representation Epn consists of Npm tokens. Epn abstracts task-relevant attributes from videos and serves as the policy condition. Formally, the Prompt Video Encoder outputs as:
Finm V pm (V1, U2,..., UN). UN)
5

---

Published as a conference paper at ICLR 2025
3.2.2 CAUSAL PREDICTOR
Connecting the visual encoders, video-conditioned Causal Predictor integrates the visual repre- sentations of prompt videos and observations to conduct sequence modeling on the mixture dataset.
To effectively leverage the mixture dataset, the predictor is designed to conduct two subtasks of behavior cloning, visual transition prediction and inverse dynamics learning. The visual transition delineates how observations should evolve to accomplish tasks depicted by the prompt videos. For- mally, visual transition prediction approximates the distribution P' (ot+1 logz V). which is embedded (01+1/o<t: in both Daemo and Dvid. Through visual transition prediction, the model learns to align the visual observation sequence with the prompt videos. Inverse dynamics learning approximates the distri- bution P (aelor, 0t+1), reflecting the actions required to realize a given visual transition. Although the action labels are only available in Demmo P (a, Ot: 0t4-1) is shared across tasks due to the shared transition function P. Given the challenge of predicting in the raw visual observation space, we re- place o++1 with the joint embedding of the next observation hy+1 in seek of compressed embedding space conducive to knowledge transfer. Accordingly, the causal predictor is tasked with capturing P phe+1loge+ V) and (a. 0t, f4+11).
The input and output sequences of the causal predictor are designed to concurrently capture the visual transition (hh+1 V) and the inverse dynamics P (aylor, hy+1). As shown in Figure with the video representation Exut as the prompt, the causal predictor takes the trajectory sequence
(Ot-R+1, 2 ................................ Or, 14+1) as input to causally predict ( O h âz) respec- tively. Structurally, the predictor comprises a causal Transformer encoder // pred and two prediction heads, T dbs and act. Respectively, Tobs predicts the joint embedding of the next observation h++1 from the hidden states of observation tokens Og, and I aa predicts the action tokens a from the hidden states of h441- Formally, the predictor encodes the input sequence:
Z V pred (Eprss Oc- k+1, /1---42 Or, hx+1)
(5)
and predicts the tokens:
h1t1 = Tobs (z(0)..
â = Past (z(Ri)
(6)
Here, - denotes the hidden states of the corresponding tokens in Z, while he-1 and â represent the predicted tokens.
In this predictive form, the causal predictor alternatively models P (hh-1 Epn, higt, Oge) capture the visual transition and P (az] Ept: oct, 0ct, hę(+1) to capture the inverse dynamics. During inference, the predictor iteratively feeds the predicted joint embedding tokens back into the input sequence to predict tokens. By combining the two types of token prediction together, the causal predictor indeed works as a planning- based policy, which predicts the desired next observations first and then converts the plan into actions.
3.3 TRAINING PROCEDURE
With the designed sequence modeling form, JEPT can simultaneously capture visual transition and inverse dynamics. In this subsection, we describe how JEPT is optimized on the mixture dataset Daemo U Dvia. Specifically, two predictive losses, the visual transition loss and the inverse dynamics loss, are optimized to predict the joint embedding tokens and the action tokens, respectively.
Visual Transition Loss. For the joint embedding prediction, compute the average T.2 distance between the predicted joint embedding and the encoded joint embedding as the visual transition loss Lots to approximate the visual transition (he+1) hig, Oge). For both Demmo and Dvia. the joint embedding of the next observations is available. The visual transition loss Los. is defined as:
Lobs Devm
h
Inverse Dynamics Loss. For the action prediction, we compute the average Cross-Entropy between the predicted actions and the discrete ground-truth actions the inverse dynamics loss Lan to

---

Published as a conference paper at ICLR 2025
approximate the inverse dynamics P (ael Epnt Oct, het+1). The ground ground-truth action label is merely available in Deemo The inverse dynamics loss Lue is defined as:
Sat = E(V.O,A)
Σ
a, log d;
(8)
k
Via optimizing Lobs, JEPT learns to generate the visual transitions aligned with the prompt videos. By optimizing Last JEPT learns to convert the visual transitions into actions. Formally, all the components of JEPT are trained with the loss Liotal
Liotal Lobs + LdED
Lacts
(9)
where è is a hyperparameter balancing the two losses, and IdEDam. is an indicator function that equals when the data sample d is from Demmo and 0 otherwise. Considering that JEPT predicts the encoded joint embeddings instead of the raw visual observations, we adopt an alternative training procedure to separately optimize the Joint Embedding Encoder and other components in case of potential model collapse. Additionally, an exponential moving average (EMA) of the joint embedding used when the Joint Embedding Encoder is fixed target network stabilize the training, which is widely adapted in the previous JEPAs (Assran et al., 2023; Bardes et al., 2024). The overall process shown in Algorithm I in Appendix A.1.
EXPERIEMNTS
In this section, we evaluate the effectiveness of our proposed JEPT on the mixture dataset, where the model learns from action-labeled expert demonstrations and unlabeled expert videos. We evaluate JEPT and baselines on two simulated benchmarks, et al., 2020). Both benchmarks provide a variety of robotic manipulation tasks and are widely used for evaluating visual control tasks. Via the experiments, we aim to figure out: (I) whether JEPT effectively leverages the additional Dvid to improve the performance on unlabeled tasks, (2) whether the learned policy can generalize to unseen tasks, and (3) whether the joint embedding predictive mechanism in JEPT is essential for knowledge transfer.
World (Yu et al., 2020a) and RoboSuite (Zhu
4.1 EXPERIMENT SETUP
We first introduce the experimental setup, including the environments, dataset, metrics, and baselines used in our experiments. More details are available in Appendix B.
Environments. We replace the language task descriptions in the environments with prompt videos recorded in the environments as task specifications. In our Meta-World experiments, a Sawyer robot arm interacts with various objects. The prompt videos performing identical manipulations, which means there is no visual gap between the prompt videos and the expert videos. In our Robosuite experiments. Panda robot arm interacts with objects. The prompt videos are recorded by performing the same manipulation with various robot arms, including Panda, Sawyer, IIWA, and UR5e, introducing a visual gap between the prompt and expert videos.
recorded with the same robot arm
Dataset. We select 18 tasks for the Meta-World task set and 15 tasks for the Robosuite task set. In order to construct the mixture dataset Daemo Dvia and evaluate on the unseen tasks, we split each task set into three subsets: (1) Taemm: the tasks with the prompt videos and the paired expert demonstrations, (2) Tvia: the tasks with the prompt videos and the paired expert videos, and (3) Tunsen: the tasks with merely the prompt videos. For each task, construct expert policy in the vector state space and collect demonstrations or vidcos by running this expert policy in the rendered environments. Duemo is collected from the tasks in Taemos while D is collected from the tasks in Tvia- We also collect prompt videos for the tasks from Tunseen- which are not included in the mixture dataset.
Metrics. We evaluate the model performance terms of success rate, which is the percentage of successful episodes within 50 trials with different random seeds. The success rates of the tasks from Tasmo- Tvid and Tunsen reflect the model's 's ability of behavior cloning, learning from videos, and one-shot imitation learning, respectively. We also calculate the average success rates of seen and unseen tasks to evaluate the overall performance.
Baselines. The baselines chosen for comparison with JEPT in our setting are follows:
7

---

Published as a conference paper at ICLR 2025
Table 1: Success Rates (%) of JEPT and the basclines calculated from 50 trials for each task in Meta-Wort The average success rate of the tasks in Tuemo and Tvia are listed in the first two rows, while the individual success rate of the 4 tasks in T Tunseen are listed in the following rows. The average success rates of the seen and unseen tasks are listed in the last two row.
Task
Vid2Robet-D
Vid2Robot-M
BC+IDM
JEPT+MWM
DT
JEPT
T
51.8
46.3
49.5
44,5
61.3
51.3
T vil
4.7
28.7
8,3
21.3
13.0
31.7
Handle Press
22.0
8.0
18.0
6.0
28.0
Lever Pull
0.0
4.0
0.0
4.0
0.0
10.0
Plate Slide Back
4.0
8.0
0.0
2.0
0.0
14.0
Faucet Open
4.0
14.0
4.0
14.0
0.0
22.0
Seen Average
28.2
37.5
28.9
32,9
37.1
41,5
Unseen Average
4.5
12.0
3.0
9.5
1.5
18.5
Table 2: Success Rates (%) the ablations calculated from 50 trials for each task in Meta- World. We train JEPT on a dataset with various task numbers for Dvid collection. The average success rate of the tasks in Taemo is listed in the first row, while the individual success rate of the 4 tasks Tunscen are listed in the following rows. The average success rate of the unseen tasks calculated in the last row.
Task
None
2 Tasks
4 Tasks
Tasks (JEPT)
T Tamm
59.5
49.8
47.0
51.3
Handle Press
6.0
4.0
16.0
28.0
Lever Pull
0.0
0.0
4.0
10.0
Plate Slide Back
0.0
0.0
10.0
14.0
Faucet Open
20.0
22.0
Unseen Average
1.5
1.5
12.5
18.5
Vid2Robot (Jain et al., 2024); conducts behavior cloning and auxiliary representation learning on large dataset where all the data is labiln-labeled. We train two variants of Vid2Robot: Vid2Robot-D refers to Vid2Robot trained merely on DDemoo while Vid2Robot-M refers to Vid2Robot trained on Demmo U Diiia with behavior cloning loss masked from D
BC+IDM: Zheng et al. (2023); Kim et al. (2023) proposed a paradigm leveraging the mixture dataset for visual imitation learning. These methods separately learn inverse dynamics model for action annotation and conduct behavior cloning on the mixture dataset. We accommodate these methods video-conditioned imitation learning.
DT*: We remove the joint embedding predictive mechanism from JEPT, which indeed conducts behavior cloning with DT variant on DJemo-
JEPT+MWM (Sco al., 2023): We replace the joint embedding in JEPT with the representation learned with MWM as baseline. MWM provides a practical representation learning method for visual control tasks. Unlike the joint embedding JEPT, MWM learns generative Masked Auto-Encoder (MAE).
4.2 META-WORLD EXPERIMENTS
In the Meta-World experiments, we divide the 18 tasks into three subsets: Tacmo Trid and Tunseen respectively containing 8, 6 and 4 tasks.
Comparisons with Baselines. The task success rates of JEPT and the baselines are presented in Table 1, where JEPT surpasses all baselines in terms of average success rates. Regarding the success rates of tasks from Tacmo JEPT performs slightly worse than DT similar to the gap between Vid2Robot-N M and Vid2Robot-D. This may be side-effect of the more tasks integrated into the policy, leading to a decrease in the performance of individual tasks. Notably, JEPT exhibits a significant advantage in tasks from both Tria and 'Tuseen indicating that JEPT effectively leverages the additional Dvia to learn a more generalizable policy. Compared to Vid2Robot-M, the sequence modeling design in JEPT, extracting visual transition and inverse dynamics, might suit the mixture dataset more. Although BC+IDM also captures the inverse dynamics, JEPT integrates transition-awann representation learning with inverse dynamics learning the mixture dataset instead of IDM learned merely on D D which may capture more generalizable inverse dynamics knowledge and accounts for the better performance of JEPT. Additionally, the outperformance of JEPT over DT
8

---

Published as a conference paper at ICLR 2025
Embeddings from Vid2Robot-M
Embeddings from MWM
Embeddings from JEPT
Door Close
Drawer Close
Faucet Open
Figure 3: Visualization of the observation embeddings learned in JEPT (right), MWM (mid), and Vid2Robot-M (left). The observation embeddings are calculated in a rollout video from 'Door Close' in Taemo- 'Drawer Close' in Tvias and 'Facet Open' in Tunseen- We apply t-SNE on one whole embedding set and split the projected vectors according to the algorithms into three sub-figures.
and JEPT+MWM indicates that the predictive joint embedding in JEPT more suitable for learning on mixture datasets than non-specialized or alternative embedding mechanisms.
Ablation Study. We conduct an ablation study to investigate the impact of the size of Dvva. As is listed in Table 2, we train JEPT datasets where 0, 4 and tasks covered in Dvia- The success rates of tasks increases as the size of Tvid grows. The results indicates that the additional video data is benificial and JEPT can effectively leverage the data in Dvid to learn more generalizable video-conditioned policy. Additionally, when the D grows from 0 tasks to 6 tasks, the performance degradation on the seen tasks might result from higher integration of the policy, which stands with our observation in Table 1.
Visualizations. We also visualize the visual embeddings learned in JEPT, MWM, and Vid2Robot-M. Specifically, we select one task from each of Taeeoo Tyid, and Tunseen and calculate the observation embeddings of a rollout video from each task. We apply t-SNE to the whole set containing observation embeddings of the three videos from JEPT, MWM, and Vid2Robot to project them into a shared 2D space. As shown in Figure 3, the observation embeddings from Vid2Robot-M and MWM demonstrate severe out-of-distributiin phenomena among the action-labeled, unlabeled, and unseen tasks. The embeddings from JEPT exhibit greater similarity in the distribution, potentially contributing to its superior generalization.
4.3 ROBOSUITE EXPERIMENTS
We also conduct experiments the Robosuite, which offers variety of robotic arms for manip- ulation tasks. Thus we evaluate the performance of JEPT in the case where there a visual gap between prompt videos and expert videos. We split the 15 tasks into three subsets: Taemo Tvia. and Tunseens containing 6, 5, and 4 tasks, respectively. The results of these experiments are shown in Table 3. Although the success rates in these tasks are somewhat lower than those in the World experiments due to the increased difficulty. JEPT still outperforms the basclines across all tasks. Similarly to the world results, the success rate of JEPT in Taemo is lower than that of DT* possibly due to the integration of more tasks into single policy. However, JEPT still exceeds all the baselines in Tvis and 'Tussen: indicating that JEPT can effectively leverage D to learn more generalizable policies despite visual gaps.
4.4 VISUAL PRIOR INJECTION
Considering that visual priors learned from external datasets may enhance the generalization of the model due to their universality, we additionally explore the possibility of injecting such priors from the pre-trained models into JEPT. A directly injection approach is adopted in our practice. We simply replace the output of the single-timestep spatial encoder in the Joint Embedding Encoder with the output from the visual encoders of some pre-trained visual models. Specifically, we experiment on Meta- World with FlowFormer (Huang al., 2022). VideoMAE-v2 (Wang et al., 2023), Dino- v2 (Oquab et al., 2024), and SAM (Kirillov et al., 2023). representing visual priors related to optical flow, video reconstruction, depth estimation, and visual segmentation, respectively. As shown

---

Published as a conference paper at ICLR 2025
Table 3: Success Rates (%) of JEPT and the baselines calculated from 50 trials for each task in Robosuite. The average success rate of the tasks in Tdemo and Tvia are listed in the first two rows, while the individual success rate of the 4 tasks in Tunseen are listed the following rows. 'X T refers to the task where Panda robot arm performs the task T with the prompt video recorded with X robot arm. The average success rates the seen and unseen tasks are listed in the last two rows.
Task
Vid2Robot-D
Vid2Robot-M
BC+IDM
JEPT+MWM
DT
JEPT
T demo
26.8
20.7
26.7
24.3
36.3
27.2
T Tii
4.0
11.6
8.4
10.8
2.0
12.8
Panda Lift
4.0
16.0
6.0
22.0
8.0
38.0
Sawyer Lift
0.0
2.0
0.0
6.0
0.0
16.0
IIWA Lift
0.0
4.0
0.0
2.0
0.0
12.0
UR5e Lift
0.0
0.0
0.0
0.0
0.0
Seen Average
15.5
16.1
17.5
17.6
19.2
20.1
Unseen Average
1.0
5.5
1.5
7.5
2.0
18.5
Table 4: Success Rates (%) of JEPT injected with different visual priors calculated from 50 trials for each task in Meta- Meta-World tasks. The average success rate the tasks Taemo and Tvid are listed the first two rows, while the individual success rate of the 4 tasks in Tunseen listed in the following rows. The average success rate of the seen and unseen tasks are calculated in the last two row.
Task
JEPT-Flow
JEPT-ViddoMAE-22
JEPT+Dino-12
JEPT+SAM
Tamm
59.8
33.3
38.5
41.3
Tvid
35.3
11.7
9.3
19.0
Handle Press
36.0
8.0
8.0
16.0
Lever Pull
10.0
2.0
2.0
4.0
Plate Slide Back
6.0
0.0
0.0
2.0
Faucet Open
26.0
4.0
4.0
10.0
Seen Average
47.5
22.5
23.9
30.1
Unseen Average
19.5
3.5
3.5
8.0
in Table 4, JEPT+FlowFormer outperforms other visual priors and surpasses the original JEPT. This indicates that the optical flow priors are beneficial for visual control tasks and can generalize across tasks. Meanwhile, JEPT can effectively incorporate the optical flow priors to enhance model performance. 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.
CONCLUSION
In this paper, we propose the Joint Embedding Predictive Transformer (JEPT), a novel approach for video-conditioned policy learning. JEPT is designed to learn from mixture of expert demonstrations and expert vidcos paired with prompt videos, aiming to reduce the burden of action label annotation. To suit the mixture dataset, we decompose the indeo-conditioned policy learning into two subtasks: visual transition prediction and inverse dynamics learning. By jointly learning the two subtasks the sequence modeling, JEPT works a planning-base based policy. We implement JEPT the joint embedding predictive architecture to learn an abstract representation of visual observation, which aids in generalizing video-conditioned policy. Experimentally, we evaluate the effectiveness of JEPT on a series of visual control tasks. Additionally, we explore the JEPT as simple visual priors injection approach and find it valid in injecting optical flow knowledge.
extension of
Limitations. A series of experimental results of JEPT indicate that joint learning visual transition and inverse dynamics allow it to effectively leverage the mixture dataset and derive generalizable policy. Despite this, have also identified some limitations. Since our work focuses solving problems within datasets containing expert videos, our method does not include additional designs for the potential visual gap between the prompt videos and task dynamics. We have validated the effectiveness of JEPT in addressing a certain level of visual gap in Robosuite environments, but more design and validation are needed to apply JEPT to one-shot visual imitation with more discrepancy prompt videos such as human videos. In this regard, our attempt at injecting visual priors could be one beneficial approach.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by NSFC under Grant 62450001 and 62476008. The authors would like to thank the anonymous reviewers for their valuable comments and advice. The authors express special thanks to Xinrun Xu (v.xinrun gmail.com) for her assistance in the figures of this work.
REFERENCES
Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from images with a cännt-embddiig predictive architecture. In Proc. of Computer Vision and Pattern Recognition, 2023.
Shikhar Bahl, Russell Mendonea, Lili Chen, Unnat Jain, and Deepak Pathak. Affordances from human videos as versatile representation for robotics. In Proc. of Computer Vision and Pattern Recognition, 2023.
Bowen Baker, 1lge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. In Proc. of Neural Information Processing Systems, 2022.
Adrien Bardes, Jean Ponce, and Yann LeCun. Mc-jepa: A joint-embedding predictive architecture for self-supervised learning of motion and content features. arXiv preprint arXiv:2307. 2023.
Adrien Bardes, Quentin Garrido, Jean Ponce, Xinlei Chen, Michael Rabbat, Yann LeCun, Mido Assran, and Nicolas Ballas. V-JEPA: Latent video prediction for visual representation learning, 2024. URL https: http:::///ppeeeeeeeeee wttp:/iwwwweeefoouuneoeeea www.oorwwww.w..oooooooooo
Maksim Bobrin, Nazar Buzun, Dmitrii Krylov, and Dmitry V Dylov, Align your intents: Offline imitation learning via optimal transport. arXiv preprint 2xX.::2402..3337.. 2024.
Konstantinos Bousmalis, Giulia Vezzani, Dushyant Rao, Coline Manon Devin, Alex X. Lee, Maria Bauza Villalonga, Todor Davehev, Yuxiang Zhou, Agrim Gupta, Akhil Raju, Antoine Laurens, Claudio Fantacci, Valentin Dalibard, Martina Zambelli, Murilo Fernandes Martins, Rugile Pevceviciute, Michiel Blokzijl, Misha Denil, Nathan Batchelor, Thomas Lampe, Emilio Parisotto, Konrad Zolna, Scott Reed, Sergio Gómez Colmenarejo, Jonathan Scholz, Abbas Abdol- malcki, Oliver Groth, Jean- Baptitt Regli, Oleg Sushkov, Thomas Rothörl, Jose Enrique Chen, Yusuf Aytar, David Barker, Joy Ortiz, Martin Riedmiller, Jost Tobias Springenberg, Raia Hadsell, Francesco Nori, and Nicolas Heess. Robocat: A self-improving generalist agent for robotic manipulation. Transactions on Machine Learning Research, 2024.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choromanski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, et al. Rt-2: Vision-language-action models transfer web knowledge robotic control. arXiv preprint arXiv:2307.1588 2023.
Jake Bruce, Ankit Anand, Bogdan Mazoure, and Rob Fergus. Learning about progress from experts. In Proc. of International Conference on Learning Representations, 2023.
Elliot Chane-Sane Cordelia Schmid, and Ivan Laptev. Learning video-conditions policies for unseen manipulation tasks. In Proc. of International Conference on Roboties and Automation, 2023.
Matthew Chang, Arjun Gupta, and Saurabh Gupta. Learning value functions from undirected state-only experience. In Proc. of International Conference Learning Representations, 2022.
Lawrence Yunliang Chen, Kush Hari, Karthik Dharmarajan, Chenfeng Xu, Quan Vuong, and Ken Goldberg. Mirage: Cross-embodiment zero-shot policy transfer with cross- cross-nainting arXiv preprint 1924492402.19249, 2024.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Proc. of Neural Information Processing Systems, 2021.
11

---

Published as a conference paper at ICLR 2025
Sudeep Dasari and Abhinav Gupta. Transformers for one-shot visual imitation. In Proc. of Conference on Robot Learning, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In Proc. of International Conference on Learning Representations, 2021.
Yilun Du, Sherry Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Josh Tenenbaum, Dale Schuurmans, and Picter Abbeel. Learning universal policies via text-guided video generation. In Proc. of Neural Information Processing Systems, 2024.
Yan Duan, Marcin Andrychowicz. Bradly Stadie, OpenAl Jonathan Ho. Jonas Schneider, Ilya Sutskever, Picter Abbeel, and Wojeiech Zaremba. One-shot imitation learning. In Proc. of Neural Information Processing Systems, 2017.
Ashley Edwards, Himanshu Sahni, Yannick Schroecker, and Charles Isbell. Imitating latent policies from observation. In Proc. of International Conference on Machine Learning, 2019.
Alejandro Escontrela, Ademi Adeniji, Wilson Yan, Ajay Jain, Xue Bin Peng, Ken Goldberg, Young- woon Lee, Danijar Hafner, and Pieter Abbeel. Vidco prediction models as rewards for reinforcement learning. In Proc. of Neural Information Processing Systems, 2023.
Chelsea Finn, Tianhe Yu, Tianhao Zhang, Picter Abbeel, and Sergey Levine. One-shot visual imitation learning via meta-learning In Proc. of Conference Robot Learning, 2017.
Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. In Proc. of International Conference Learning Representations, 2022.
Zhaoyang Huang, Xiaoyu Shi, Chao Zhang, Qiang Wang, Ka Chun Cheung. Hongwei Qin, Jifeng Dai, and Hongsheng Li. Flowformer: transformer architecture for optical flow. Proc. of European Conference on Computer Vision, 2022.
Andrew Jaegle, Sebastian Borgeaud, Jean- Can-ttapistt Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al. Perceiver io: A general architecture for structured inputs & outputs. Proc. of International Conference on Learning Representations, 2022.
Vidhi Jain, Yixin Lin, Eric Undersander, Yonatan Bisk, and Akshara Rai. Transformers are adaptable task planners. In Proc. of Conference Robot Learning, 2023.
Vidhi Jain, Maria Attarian, Nikhil Joshi, Ayzaan Wahid, Danny Driess, Quan Vuong, Pannag R Sanketi, Pierre Sermanet, Stefan Welker, Christine Chan, et al. Vid2robot: End-to-end video- conditioned policy learning with cross-attention transformers. arXiv preprint 1294, 2024.
Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. Be-z: Zero-shot task generalization with robotic imitation learning. In Proc. of Conference on Robol Learning, 2022.
Yiding Jiang, Shixiang Shane Gu, Kevin P Murphy, and Chelsea Finn. Language as abstraction for hierarchical deep reinforcement learning. In Proc. of Neural Information Processing Systems, 2019.
Yunfan Jiang, Agrim Gupta, Zichen Zhang, Guanzhi Wang, Yongqiang Dou, Yanjun Chen, Li Fei-Fei, Anima Anandkumar, Yuke Zhu, and Linxi Fan. Vima: General robot manipulation with multimodal prompts. In Proc. of International Conference on Machine Learning, 2023.
Dong-Hon Hee Kim, Sungduk Cho, Hyeonwoo Cho, Chanmin Park, Jinyoung Kim, and Won Hwa Kim. Joint-embedding predictive architecture for self-supervised learning of mask classification architecture. arXiv preprint arXiv:2407. 10733, 2024.
12

---

Published as a conference paper at ICLR 2025
Moo Jin Kim, Jiajun Wu, and Chelsea Finn. Giving robots hand: Learning generalizable ma- nipulation with eye-in-hand human video demonstrations. arXiv preprint arXiv:2307.05959 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 Proc. of International Conference on Computer Vision, 2023.
Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. 2022.
Kuang-Huel Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadar- rama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. In Proc. of Neural Information Processing Systems, 2022.
Lisa Lee, Ben Eysenbach, Russ R Salakhutdinov, Shixiang Shane Gu, and Chelsea Finn. Weakly- supervised reinforcement learning for controllable behavior. In Proc. of Neural Information Processing Systems, 2020.
Siyuan Li, Shijie Han, Yingnan Zhao, By Liang, and Peng Liu, Auxiliary reward generation with transition distance representation learning. arXiv preprint rrXiv:2402.07412, 2024.
Hao Luo, Bohan Zhou, and Zongqing Lu. Pre-trained visual dynamics representations for efficient policy learning. In Proc. of European Conference on Computer Vision, 2024.
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. In Proc. of Conference on Robot Learning, 2023.
Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. In Proc. of Neural Information Processing Systems, 2018.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A universal visual representation for robot manipulation. In Proc. of Conference on Robot Learning. 2023.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy V. Vo. Mare Szaframiec, Vasil Khalidov, Pierre Fernandez, Daniel HAZIZA, Francisco Massa, Alaaeldin El-Nouby, Mido Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Po- Yao Huang, Shang- Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Herve Jegou, Julien Mairal, Patrick Labatut, Ar- mand Joulin, and Piotr Bojanowski. DINOv2: Learning robust visual features without supervision. Transactions on Machine Learning Research, 2024.
Abhishck Padalkar and et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint arXiv:2330.08800 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. In Proc. of Neural Information Processing Systems, 2019.
Ilija Radosavovic, Tete Xiao, Stephen James, Pieter Abbeel, Jitendra Malik, and Trevor Darrell. Real-world robot learning with masked visual pre-training. In Proc. of Conference Robot Learning, 2023.
Karl Schmeckpeper, Oleh Rybkin, Kostas Daniilidis, Sergey Levine, and Chelsea Finn. Reinforcement learning with videos: Combining offline observations with interaction. In Proc. of Conference on Robot Learning, 2021.
Dominik Schmidt and Minqi Jiang. Learning act without actions. In Proc. of International Conference on Learning Representations, 2024.
Younggyo Seo, Kimin Lee, Stephen L James, and Pieter Abbeel. Reinforcement learning with action-free pre-training from videos. In Proc. of International Conference on Machine Learning, 2022.
13

---

Published as a conference paper at ICLR 2025
Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control. In Proc. of Conference on Robot Learning. 2023.
Rutav Shah, Roberto Martín-Martín, and Yuke Zhu. MUTEX: Learning unified policies from multimodal task specifications, In Proc. of Conference on Robot Learning, 2023.
Sangwoo Shin, Daehee Lee, Minjong Yoo, Woo Kyung Kim, and Honguk Woo. One-shot imitation in non-stationary environment via multi-modal skill. In Proc. of International Conference on Machine Learning, 2023.
Sangwoo Shin, Minjong Yoo, Jcongwoo Lee, and Honguk Woo. Semtra: A semantic skill translator for eross domsin zero-shot policy adaptation. In Proc. of Association for the Advancement of Artificial Intelligence, 2024.
Aravind Sivakumar, Kenneth Shaw, and Deepak Pathak. Robotic telekinesis: Learning a robotic hand imitator by watching humans on youtube. In Proc. of Robotics: Science and Systems, 2022.
Limin Wang, Bingkun Huang, Zhiyu Zhao, Zhan Tong, Yinan He, Wang, Yali Wang, and Yu Qiao. Videomae v2: Scaling video masked autoencoders with dual masking. In Proc. of Computer Vision and Puttern Recognition, 2023.
Tete Xiao, Lija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv preprint arXfv:2033.06173 2022.
Jiange Yang, Bei Liu, Jianlong Fu, Bocheng Pan, Gangshan Wu, and Limin Wang. Spatiotemporal predictive pre-training for robotic motor control. arXiv preprint arXiv:2403. 053004 2024.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09643 2021.
Weirui Ye, Yunsheng Zhang, Picter Abbeel, and Yang Gao. Become a proficient player with limited data through watching pure videos. In Proc. of International Conference Learning Representations, 2022.
Sriram Yenamandra, Arun Ramachandran. Karmesh Yadav, Austin S. Wang, Mukul Khanna, Theophile Gervet, Tsung-Yen Yang. Vidhi Jain, Alexander Clegg. John M. Turner, Zsolt Kira, Manolis Savva, Angel X. Chang. Devendra Singh Chaplot, Dhruv Batra, Roozbeh Mottaghi, mobile manipulation. In Proc. of
Yonatan Bisk, and Chris Paxton. Homerobot: Conference on Robot Learning, 2023.
Tianhe Yu, Picter Abbeel, Sergey Levine, and Chelsea Finn. One-shot hierarchical imitation learning of compound visuomotor tasks. arXiv preprint 1rXiv:1810.1144. 2018
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Proc. of Conference on Robot Learning. 2020a.
Xingrui Yu, Yueming Lyu, and Ivor Tsang. Intrinsic reward driven imitation learning via generative model. In Proc. of International Conference on Machine Learning, 2020b.
Qihang Zhang, Zhenghao Peng, and Bolei Zhou. Learning to drive by watching youtube videos: Action-conditioned contrastive policy pretraining. In Proc. of European Conference on Computer Vision, 2022.
Qinqing Zheng, Mikael Henaff, Brandon Amos, and Aditya Grover. Semi-supervised offline rein- forcement learning with action-free trajectories. In Proc. of International Conference on Machine Learning, 2023.
Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martín-Martin, Abhishek Joshi, Soroush Nasiriany, and Yifeng Zhu. robosuite: A modular simulation framework and benchmark for robot learning. arXiv preprint arXiv:2009. 12293, 2020.
14

---

Published as a conference paper at ICLR 2025
A IMPLEMENTATIONS
We build our framework based on PyTorch (Paszke et al., 2019) and use the implementations of Transformer modules from the codebase x-transformers' We use the 3090 Nvidia GPU and i9- 12900K CPU for the JEPT training and testing. It takes around 10 hours to train JEPT on the Meta-World dataset and 16.5 hours train JEPT on the Robosuite dataset. The details of the implementations are as follows.
A.1 ALGORITHM
Algorithm Joint Embedding Predictive Transformer
I: Input: Mixture dataset Ddemo U 'ivid
2: Initialize: JEPT components ooss, pmn. 10 "pred, Tobs Faa with random weights
3: for 1,2....du
4:
for d E DDemo U Dvii do
Optimizing Joint Embedding Encoder
5: Encode V and O into Epn and (h+-k+1,.... he) as Equations and 4
6: Predict the joint embedding tokens and action tokens as Equations and
7: Compute Looual as Equation 9
8:
Update V with Liotal end for
9:
10:
for d E Ddemo U Dvid do
Optimizing other components
11:
Repeat steps in lines 5-7
Update 1/ prts V predt I Tobs I with Livaal
13: end for
end for
A.2 MODEL. DETAILS
Joint Embedding Encoder. To encode visual observations, we utilize pre-trained VIT (Dosovitskiy et al., 2021) variant, 'vit ba__base___pg p 224', the spatial visual encoder. To effectively compress the tokens of visual observations, a Perceiver-IO) (Jaegle et al., 2022) with a fixed number of learnable queries is employed to compress the visual tokens of single-frame observations. The VIT weights remain fixed during training, allowing the Perceiver-IO to function as an adaptor network. The compressed visual tokens are concatenated with representations of the self-state vector produced by an MLP to form the observation tokens. A bi-directional Transformer temporal encoder is used encode the observation tokens with context awareness. For visual prior injection approaches in Section 4.4, we simply replace the outputs of the fixed ViT with the embedding outputs of visual encoders pre-trained to capture visual priors. Notably, there are subtle distinctions input processing among these four visual prior injections. For FlowFormer, observation tokens at each temporal step are computed from adjacent frames. VideoMAE v2 replicates all video frames excluding the first to preserve the original temporal length. Dino v2 and SAM v2 both utilize single-frame images inputs. Moreover, these models diverge architectural design, pre-training datasets, and methodological approaches. Despite these variations, our visual prior experiments do not aim comparatively evaluate visual prior performance. Instead, our objective is to ascertain whether a joint embedding predictive of these representations with visual priors further compress visual information and enhance generalization.
basc_patchl6
To encode the prompt videos, utilize the same ViT variant the Joint Embedding Encoder for encoding the visual tokens. Two cascaded Perceiver-IOs are employed to aggregate these tokens. The first Perceiver compresses the tokens from single frame, while the second compresses the flattened tokens from the first layer to form the prompt video embeddings. During training, only the Perceiver-IOs are trained, the VIT weights remain fixed.
Prompt Video Encoder:
Video-Cone Cideo--oodiiionedd Causal Predictor. To process observations in expert videos and demonstrations, we employ the same strueture the Joint Embedding Encoder to encode single-timestep observations. Notably, the single-timestep encoder shares the weights in our implementation, and the EMAs are
15

---

Published as a conference paper at ICLR 2025
calculated when producing the target joint embeddings. A causal Transformer is utilized, taking as input the sequence of all compressed tokens. A causal mask, where tokens attend only to preceding tokens and tokens from the same observations, ensures prediction causality. To predict the joint embedding tokens and action tokens, two prediction heads are applied to the causal Transformer's output. Specifically, an ML. predicts the joint embedding tokens, while cross-attention Transformer with learnable action queries for each action dimension predicts the action tokens.
A.3 BASELINE ADAPTATION
Vid2Robot We implement Vid2Robot following the original paper (Jain et al., 2024), codebase released. Similar to our structure, a spatial visual encoder followed by a Perceiver processes the spatiotemporal information of videos and observation sequences. We the same ViT variant as the Joint Embedding Encoder to encode the visual tokens of the prompt videos and observations, ensuring consistency in comparison. We align the Perceiver structure and the action prediction head in Vid2Robot with our framework. Since no text instruction is provided in our setting, we remove the auxiliary text-video contrastive loss in Vid2Robot.
MWM We adapt the released codebase of MWM (Seo et al., 2023) to our setting. The only difference the absence reward signal in our setting; thus, we remove the reward prediction loss.
DT We adapt our implementation of JEPT to perform sequence modeling in typical behavior cloning sequence ((ae)Epe 01, 01 .1.,)..,))))
A.4 HYPERPARAMETERE
In our implementation, the Perceiver-IO structure shared across the Joint Embedding Encoder, the Prompt Video Encoder, and the Indoo-Conditioned Causal Predictor. The hyperparameters for the Perceiver-IOs and other modules are detailed in Table 5. During training, Σpn is initially excluded in the causal predictor as a warm-up strategy. Furthermore, the hyperparameters for the training process are listed in Table 6.
B ENVIRONMENT SETUP DETAILS
B. ENVIRONMENT
Meta-World Environment. In the Meta- Wool-World environment, Sawyer robot arm is directed to execute variety of manipulation tasks. We employ 18 tasks originally from Meta- World to establish our video prompt visual control framework. As shown in Figure 4, for each task, video of the identical manipulation in the same environment, albeit with a different random seed, is provided as the task prompt, while the original task descriptions are omitted. We utilize the customized cameral perspective for visual observation and resize the image from the camera to 224 x 224 pixels. In addition to the visual observation, we assume the agent is aware of current self-state, which encompasses the joint angles, joint velocities, and end-effector position. vector representing these self-state details is supplied alongside the visual camera image to form the observation. For action prediction using Cross-Entropp the original continuous action space is discretized into bins for each action dimension. During the rollout of each task, the episode is truncated after 250 steps, with an action repeat 2.
Robosuite Environment. In the Robosuite environment, we select tasks involving a Panda robot arm performing various manipulation tasks. We employ two tasks originally from the Robosuite simulator and four tasks implemented by MimicGen (Mandlekar et al., 2023) within the Robosuite framework to create our video prompt visual task. As illustrated in Figure 4, a video depicting specific robot arm completing the same manipulation serves the task prompt. The robot in these prompt videos may be a Panda, Sawyer, IIWA, or UR5e, indicating potential visual gap between the prompt videos and the task dynamics. Two tasks are considered distinct if the robot arm in the prompt videos or the desired manipulation differs. Consequently, we construct 15 tasks from the 6 tasks of MimicGen. We utilize the 'agentview image' attribute of the state the visual image. Similar to the Meta- World environment, the visual image is resized to 224 x 224 pixels and accompanied by the self-state vector. The continuous action space is discretized into bins per dimension. During each
16

---

Published as a conference paper at ICLR 2025
Table 5: The hyperparameters of the JEPT Modules.
Single-Frand Frame Perceiver-10
num latents
64
num queries
embedding channels
512
attention heads
encoder cross-attettion layers
encoder self-attention layers
decoder cross- aatsentton layers
attention feedforward dim
1024
attention dropout
0.1
Temproal Perceiver-IO
num latents
128
num queries
32
embedding channels
512
attention heads
encoder cross-attention layers encoder self-attentio layers
decoder cross-attention layers
attention feedforward dim
892
attention dropout
0.08
Bi-directional Tempral Encoder
joint embedding dim
256
attention layers
attention heads
attention dim
512
attention feedforward dim attention dropout
1024
0.1
attention activation
GeLU
Causal Transformer
Max prompt length prompt tokens num Npm
250
32
obs tokens num Nobs
context length k embedding channels attention heads attention layers
584
attention feedforward dim
1168
attention dropout
0.05
Attention Activation
GeLU
Action Prediction Head I lact
Cross-attention layers
Cross-attention heads
Cross-attention dim
256
Feedforward dim
512
activation
GeLU
Joint Embedding Prediction Head l'obs
layers
hidden dim activation
256
LeakyRelu
17

---

Published as a conference paper at ICLR 2025
Table 6: The hyperparameters of the training process.
Training Hyperparameters
Value
Minibatch size
90
Optimizer
AdamW
Joint Embedding Encoder Learning Rate
3e-5
Prompt Video Encoder Learning Rate
7e-5
Jideo-Conditioned Causal Predictor Learning Rate
7e-5
Weight Decay
le-4
Max Gradient Clip
1.0
Warm Up Steps
200
Loss Weight
8
prompt videos
expert vidoos/dttmmmoosstattionn
Meta-World
robosuite
Figure 4: The video prompt task setting of the WorldWorld and Robosuite. In Meta-World, the task specification is video of the same manipulation in the same environment with different random seed. In Robosuite, the task specification is video of certain robot arm executing the same manipulation, and the visual gap between the prompt videos and the task dynamics likely.
task rollout, the episode is truncated after 200 steps, although there is no restriction on episode length within the Robosuite environment.
B.2 DATASET
Task Subset. In both Meta-Worke and Robosuite, tasks are manually categorized into three subsets: those with expert demonstrations and prompt videos, Taomo: those with expert videos and prompt videos, Twai and those with only prompt videos, Tumsean The specific task divisions are detailed in Table 7 and 8. For Pera-World experiments, 500 expert videos/demonssraationss
collected
by executing expert policies on tasks from Taemo U T For Robosuite experiments, 1000 expert videos/demonstrations are gathered for each task in the dataset. For the tasks in Tinsenn 10 prompt videos are sampled for trials with different random seeds on each to assess the one-shot imitation capability of the learned policy. In the ablation study of task numbers covered in Trid. we use the 'Button Press Topdown Wall' and 'Door Open' tasks for the 2-task variant. Additionally, we include the 'Reach Wall' and 'Drawer Close" tasks for the 4-task variant.
Expert Policy. To construct the data used for training and evaluation, we develop expert policies for individual tasks within the vector state space, collecting both prompt videos and corresponding expert demonstrations. Initially, expert policies are formulated in the continuous action space. In Meta- World, policies are trained using PPO for each task, utilizing the vector state and reward function provided by the environment, We integrate rule-based policies from the GitHub repository with PPO policies to form expert policies, executing rule-based policies with a probability of 0.85 and PPO policies with a probability of 05 ensure diversity. For tasks in Robosuite, expert policies are learne through behavior cloning on the expert demonstrations from MimicGen. Additionally, expert policies for 'Stack" and 'Lift' tasks trained using DrQ-v2 (Yarats et al., 2021) with no demonstrations available in Mimicgen. To collect expert rollouts in the discrete action space, continuous expert policies are used to generate actions, which are then replaced with the nearest discrete actions in the bins for interaction with the environment. A performance drop is noted in the discrete action space compared to the continuous action space. Specifically, the average success rate for expert policies in
18

---

Published as a conference paper at ICLR 2025
Table 7: The division of tasks in our experiments.
Task Subset
Tasks
Button Press, Button Press Topdown, Handle Pull, Reach,
Taemo
Door Close, Window Open, Plate Slide, Faucet Close
Button Press Topdown Wall, Door Open, Reach Wall, Drawer Close, Plate Slide Back Side, Window Close
Tvil
Tunseen
Handle Press, Lever Pull, Plate Slide Back, Faucet Open
Table 8: The division tasks in our Robosuite experiments. 'X refers to the task where the Panda robot arm performs the task T with the prompt video recorded with X robot arm.
Task Subset
Tasks
Panda Stack, Sawyer Stack, IIWA Stack, UR5e Stack,
Taemo
Three Pieces Assembly, Coffee
Panda Square, Sawyer Square, IIWA Square, UR5e Square, Stack Three
Tvid
T Tunseen
Panda Lift, Sawyer Lift, IIWA Lift, UR5e Lift
Meta-World World decreases from 91.5% to 85.3% disteriicttiination,, and in Robosuite, it declines from 78.3% to 7%. For prompt video collection, record visual observation sequences of the expert rollouts, while for expert videos or demonstrations, both visual observation sequences and self-state vector sequences are recorded.
C ADDITIONAL EXPERIMENTS
C.I VISUAL PRIOR INJECTION ON ROBOSUITE
In order to supplement the effects of various visual priors on JEPT in different enviromments, conduct experiments on visual prior injection within the RoboSuite environment. The implementation of visual prior injection is consistent with the experiments condueted in Meta-World. The results, as shown in Table 9, indicate that the performance of JEPT can be slightly enhanced by utilizing priors based on optical flow process the visual inputs. This suggests that optical flow might be a suitable prior for enhancing JEPT's performance across different environments. However, all of these visual encoders have been trained on large datasets across numerous visual tasks and possess certain generalizable priors for handling visual inputs. Our experiments with visual priors aim to ascertain whether joint embedding predictive of these representations, when combined with visual priors, can further compress visual information and improve generalization. From the perspective of introducing more generalizable visual priors, these different priors for processing visual inputs should be logical. However, the success of learning predictive joint embeddings of these priors and effectively enhancing performance may depend on selecting more refined injection methods and better training designs. We merely highlight JEPT's potential in this domain, and further detailed research on other visual prior injections is left for future exploration.
C.2 ADDITIONAL ABLATIONS
We conduct additional ablation studies to explore the influence of the hyperparameter e in JEPT. The value of is varied within the set (0.3, 1, 5, 8, 12), and the performance of JEPT is assessed in Meta-World. The results are presented in Table 10. When e remains within a reasonable range, JEPT exhibits similar performance. However, when is excessively small large, JEPT's performance declines significantly. This underscores the critical role the hyperparameter in JEPT's s efficacy. In other experiments, set 8, as it yields the optimal performance in Meta- World.
To thoroughly investigate the impact of the ratio between D Ddemo and D on model performance, we conduct additional ablation experiments. Specifically, beyond the original 14-task training dataset,
19

---

Published as a conference paper at ICLR 2025
incorporate 6 additional tasks from Meta-World to collect training data (Handle Pull Side, Door Lock, Drawer Open, Plate Slide Side, Dial Tum, Handle Press Side). We maintain a constant total amount of training data while varying the ratio of demo to Tvia- For various ratios, experiment with different hyperparameters e, and the results are presented in Table 11. Given the variability of TTeeo and Tvia. report the average success rate Tunseen under different settings. The results indicate that within a certain range, Dvia effectively compensates for the lack of data with action labels. However, when demo is insufficient, it may adequately learn, making IDM difficult to generalize to unseen tasks. Moreover, the results suggest that when Dvid is reduced, larger is often required achieve better performance.
Considering that employing fixed video encoder mitigates the distribution shift in video repre- sentation between action-labeled and unlabeled data, thereby enhancing the generalization of the inverse dynamics model in BC+IDM, we conducted experiments on Meta-Worke using pre-trained video encoder, denoted as BC+IDM'. Most pre-trained video encoders utilize a kernel size of for temporal compression, halving the representation length 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 maintain consistent length and accurately decode the actions. Employing VideoMAE v2, the experimental results are presented Table 12. There is a slight improvement in success rates unlabeled tasks, but no gain on unseen tasks. Overall, the improvement from using fixed video encoder is limited.
C.3 ADDITIONAL VISUALIZATION
We conduct a series of visualizations within the Meta-World environment to reveal the distribution of some variants within the rollouts. Specifically, process 150 videos for each of the 18 different tasks using a prompt video encoder to encode them into Fxrt. We flatten these embeddings and apply t-SNE to reduce the dimensionality to two dimensions. As is shown in Figure 5, there clear distinctions between different tasks, suggesting that JEPT effectively captures and represents the unique features of each task. Additionally, observe 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.
Furthermore, also apply t-SNE visualize the actions and joint embeddings of the 150 trajectories for each task. Each point in the visualization is color-coded with varying transparency to denote its temporal position within the trajectory, with points closer the end being less transparent. The visualizations are shown in Figures 7 and 6. From this analysis, we find that the overall distributions of actions and joint embeddings across different tasks are similar, which partially explains the generalization capability of JEPT. Additionally, there are specific concentrated regions the distribution for each task, which likely arise from the unique nature each task.
Table Success Rates (%) of JEPT injected with different visual priors calculated from 50 trials for each task in RoboSuite tasks. The average success rate of the tasks in Taemo and Tvid are listed in the first two rows, while the individual success rate of the 4 tasks Tunscen are listed in the following rows. The average success rate of the seen and unseen tasks are calculated in the last
Task
JEPT+VideoMAE-w JEPT+FlowFormee
JEPT+Dino-v2
JEPT+SAM
JEPT
T Temo Tsid
27.7
20.7
24.3
26.7
27.3
13.2
11.6
10.8
8.4
12.8
Panda Lift
34.0
14.0
Sawyer Lift
18.0
0.0
ITWA Lift
8.0
8.0
0.0
0.0
12.0
UR5e Lift
14.0
0.0
0.0
0.0
8.0
Seen Average
20.4
16.1
17.6
17.5
20.1
Unseen Average
18.5
5.5
5.0
3.5
18.5
20

---

Published as a conference paper at ICLR 2025
Table 10: Success Rates (%) of JEPT and the ablations on calculated from 50 trials for each task in Meta- World. The average success rates of the tasks in Taemo and Tvia are listed in the first two rows, while the individual success rates of the tasks in Tunseen are listed the following rows. The average success rates of the and unseen tasks are listed in the last two row.
Task
0.3
1
5
8(EEPT)
10
12
T demo
38.5
33.0
47.0
51.3
52.5
54.5
T vid
19.0
17.7
24.3
31.7
27.7
12.3
Handle Press
4.0
12.0
24.0
28.0
14.0
6.0
Lever Pull
0.0
0.0
8.0
10.0
12.0
0.0
Plate Slide Back
2.0
8.0
6.0
14.0
16.0
2.0
Faucet Open
0.0
2.0
16.0
22.0
20.0
8.0
Seen Average
28.8
25.3
35.7
41.5
40.1
33.4
Unseen Average
1.5
5.5
13.5
18.5
15.5
4.0
Table 11: Average Success Rates (%) of JEPT variants the 4 tasks in Tunseen in Meta-World experiments. The columns represent the variants with different ratios of Taemo and T vids while the represent the variants with different values of hyperparameter
T demo
T vid
(1)
(19)
(3) + (17)
(5)
(15)
(8) + (12)
(10) (10)
(12)
(8)
0.0
0.0
8.0
5.0
3.0
9.0
5
0.0
1.0
5.5
7.5
10.0
8.5
8
1.5
2.0
3.5
8.5
18.0
18.5
10
0,0
0.0
6.0
10.0
16.0
10.0
C 12
0.5
2.5
8.0
19.5
14.5
9.5
16
0.0
5.0
12.5
12.0
10.0
10.5
C 20
1.0
0.0
10.0
11.0
8.0
4.0
Table 12: Success Rates (%) of JEPT and two baselines calculated from 50 trials for each task in Keta-World. The average success rate of the tasks in Taemo and Twa are listed in the first rows, while the individual success rate of the 4 tasks in Tunseen are listed in the following rows. The average success rates of the seen and unseen tasks are listed in the last two row.
Task
BC+IDM
BC+IDM
JEPT
T demo
49.5
47.5
51.3
T vid
8.3
14.7
31.7
Handle Press
8.0
7.0
28.0
Lever Pull
0.0
0.0
10.0
Plate Slide Back
0.0
0.0
14.0
Faucet Open
4.0
3.0
22.0
Seen Average
28.9
31.1
41.5
Unseen Average
3.0
2.5
18.5
21

---

Published as a conference paper at ICLR 2025
Window Open
Button Press Topdown
Reach
Handle Pull
Door Close
Plate Slide Back
Button Press
Handle Press
Lever Pull
Window Close
Faucet Open
Plate Slide
Dour Open
Reach
Drawer Close
Faucet Close
Plate Slide Back
Button Press Topdown
Figure 5: Visualization of the prompt video embeddings learned in JEPT. The prompt video embed- dings are calculated from 150 videos each task in Meta- World. We apply t-SNE on the whole embedding set.
22

---

Published as a conference paper at ICLR 2025
Window Open
Handle Pull
Button Press
Window Close
Door Open
Faucet Close
Button Press Topdown
Door Close
Handle Press
Faucet Open
Reach
Plate Slide Back Side
Reach Wall
Plate Slide Back
Lever Pull
Plate Slide
Drawer Close
Button Press Topdown Wall
Figure 6: Visualization of the joint embeddings learned in JEPT. The joint embeddings are calculated from 150 trajectories of each task in Meta-World, We apply t-SNE on one whole embedding set and split the projected vectors according to the tasks into 18 sub-figures. The transparency of the points reflects the temporal positions of the corresponding embeddings, with points closer to the end being less transparent.
23

---

Published as a conference paper at ICLR 2025
Window Open
Handle Pull
Button Press
Window Close
Door Open
Fancet Close
Button Press Topdown
Door Close
Handle Press
Faucet Open
Reach
Plate Slide Back Side
Reach Wall
Plate Slide Back
Lever Pull
Plate Slide
Drawer Close
Button Press Topdown Wall
Figure 7: Visualization of the actions learned in JEPT. The actions collected from 150 trajectories of each task in Meta-World We apply t-SNE on one whole embedding set and split the projected vectors according the tasks into 18 sub-figures. The transparency of the points reflects the temporal positions of the corresponding embeddings, with points closer to the end being less transparent.
24


Reviews:
[{"review_id": "FGMsUkYYVY", "reviewer": "Reviewer_4tQo", "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\nlearns 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.", "questions": "Please refer to the Weaknesses section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n- 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.\n- 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.", "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": "u8mcxsCQAg", "reviewer": "Reviewer_yXfU", "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.", "questions": "1. How to handle the heterogeneity issue? Is the action space designed for only one robot configuration?\n2. 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.\n3. Is the hyperparameter $c$ in Eq.9 also influenced by the ratio of $D_{demo}$ and $D_{vid}$?\n4. 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}$.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The proposed method aims to introduce prompt video as guidance for task imitation learning, which is interesting.\n2. 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.\n3. 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\n\n* 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;\n* 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;\n\n2. 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.\n\n3. 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.", "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": "XbiFeIQQG9", "reviewer": "Reviewer_grHE", "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.", "questions": "please look at my strengths and weakness sections, and if you can adress the weakness section, i am happy to change my ratings.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.  \n\nThe paper has through experiments on meta world and robosuit, and in most of the cases the performance is good compared to previous works. \n\nAblations 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. \n\nLack 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. \n\nIt 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.", "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": "H4cFE36T2S", "reviewer": "Reviewer_Zv2N", "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.", "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?\n2. How does JEPA influence performance? Did the authors experiment with other pretrained models, including different architectures and pretrained data?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed JEPT framework for learning from mixed datasets is clear and reasonable.\n2. Experiments on the Meta-World and Robosuite benchmarks demonstrate the effectiveness of JEPT.\n3. 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.\n2. How does JEPA influence performance? Did the authors experiment with other pretrained models, including different architectures and pretrained data?\n3. 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.", "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 (Poster)