OSVI-WM: One-Shot Visual Imitation for Unseen Tasks using World-Model-Guided Trajectory Generation
Authors: Raktim Gautam Goswami
Prashanth Krishnamurthy
Yann LeCun
Farshad Khorrami
Venue: NeurIPS 2025 poster ()
Keywords: One-Shot Visual Imitation Learning
World Model
Robot Learning
Primary area: reinforcement_learning

Abstract: Visual imitation learning enables robotic agents to acquire skills by observing expert demonstration videos. In the one-shot setting, the agent generates a policy after observing a single expert demonstration without additional fine-tuning. Existing approaches typically train and evaluate on the same set of tasks, varying only object configurations, and struggle to generalize to unseen tasks with different semantic or structural requirements. While some recent methods attempt to address this, they exhibit low success rates on hard test tasks that, despite being visually similar to some training tasks, differ in context and require distinct responses. Additionally, most existing methods lack an explicit model of environment dynamics, limiting their ability to reason about future states. To address these limitations, we propose a novel framework for one-shot visual imitation learning via world-model-guided trajectory generation. Given an expert demonstration video and the agent’s initial observation, our method leverages a learned world model to predict a sequence of latent states and actions. This latent trajectory is then decoded into physical waypoints that guide the agent’s execution. Our method is evaluated on two simulated benchmarks and three real-world robotic platforms, where it consistently outperforms prior approaches, with over 30% improvement in some cases.

OSVI-WM: One-Shot Visual Imitation for Unseen Tasks using World-Model-Guided Trajectory Generation
Raktim Gautam Goswami" Prashanth Krishnamurthy' Yann LeCun Farshad Khorrami
New York University Tandon School of Engineering New York University Courant Institute of Mathematical Sciences "Meta-FAIR
Abstract
Visual imitation learning enables robotic agents to acquire skills by observing expert demonstration videos. the one-shot setting, the agent generates a policy after observing single expert demonstration without additional fine-tuning. Existing approaches typically train and evaluate on the same set of tasks, varying only object configurations, and struggle to generalize to unseen tasks with different semantic structural requirements. While some recent methods attempt to address this, they exhibit low success rates hard test tasks that, despite being visually similar some training tasks, differ in context and require distinct responses. Additionally, most existing methods lack explicit model of environment dynamics, limiting their ability to reason about future states. To address these limitations, we propose novel framework for one-shot visual imitation learning via world-model-guide trajectory generation. Given an expert demonstration video and the agent's 'sent's initial observation, our method leverages a learned world model to predict sequence of latent states and actions. This latent trajectory is then decoded into physical waypoints that guide the agent's execution. Our method evaluated on two simulated benchmarks and three real-world robotic platforms, where it consistently outperforms prior approaches, with over 30% improvement some cases. The code is available at httpss//gghuu..mm//kiigggssssssu
Introduction
Imitation the sincerest of flattery.
Charles Caleb Coltor
Intelligent beings like humans capable of learning wide range of skills by observing and imitating others [50]. Even toddlers, with basic understanding of the world's dynamics, solve complex tasks by watching expert demonstrations. Inspired by this natural learning process, principles have been
applicd in robotics 13. 26 35, 622 where expert demonstrations, typically collected through human teleoperation, used teach agents to perform tasks through imitation. Imitation learning finds uses across diverse domains, including medical robotics, collaborative robotics, and industrial automation.
In successful imitation, the agent discovers the skill demonstrated by the expert, adapts it to its own embodiment, and executes it within its environment (58). the one-shot visual imitation (OSVI) setting, the agent must derive policy from single expert demonstration video, without any additional training [12]. Notably, the expert and agent may differ in embodiment [[O]; in fact, both prior work 43 37 and our approach have utilized human demonstration videos as input. Existing methods, however, often rely on the strong assumption that the training and testing tasks are nearly identical, typically differing only in object locations. As a result, generalization
Lacon: Or Many Things in Few Words, 1820.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Expert Demonstration Video
foresees future observations in the latent space using learned world model.
The expert shows 'what' the task and 'how' to do it.
Agent's Observation
Predicted Latent Trajectory
World-Model Model-Guided Trajectory Generation
Latent Space
Low-Level Control
Waypoint Controller
Waypoints
Decoder
Figure OSVI-WM infers the task from the expert demonstration and, along with the agent's observation "foresees" future latent states using a world-model-guidee trajectory generation module. The predicted trajectory is decoded into physical waypoints for control.
to semantically or structurally different tasks remains limited. MOSAIC [34 and AWDA 5 were among the first to investigate generalization to unseen tasks. In particular, AWDA (5) demonstrated that the performance of standard methods like [T, 60) degrades severely when evaluated on tasks that differ even marginally from those seen during training. Although AWDA outperforms earlier methods [10 60), it exhibits low success rates on test tasks that are visually similar but contextually different, such opening closing a sliding window. Moreover, incorporates an explicit model of the environment's dynamics, limiting their ability to reason about long-term consequences plan into the future.
of these approaches
Mathematically modeling the environment using raw camera images challenging due to the presence of diverse objects and their complex interactions. To address this, neural networks have been used to learn world models that predict how the environment will evolve response to the agent's 's actions. These models support applications in reinforcement learning 488 21
28 39
23 56 and behavior
cloning g 37). Inspired by such advances, we propose OSVI WM (Fig D. a framework for one-shot visual imitation for unseen tasks using world-modee model-guided trajectory generation. Given expert demonstration video and the agent's 's initial observation, OSVI-WX WW encodes each image into a shared latent space representing the environment state for the respective image. These embeddings recursively processed through learnable action and world models to predict a trajectory of future latent states, allowing the agent plan into the future. Predicting actions in the latent space also helps mitigate the multi-modality of real-world action distributions. The predicted latent trajectory is then decoded into physical waypoints
5 using pooling and Multi-Layer Perceptron (MLP) layers.
The predicted waypoints guide the agent's execution and provide supervision for end-to-end training. To enhance the world model's ability to capture environment dynamics, auxiliary loss is applied on the predicted latent trajectory during training. Further, OSVI- WM supports re- re-planning: if the agent doesn't t accurately complete the task by following the predicted waypoints, it takes and plans a fresh set of waypoints based on its current state. Notably, unlike methods that rely on
observation
large-scale pretraining extensive demonstration data 37 3 55). OSVI-WM is trained directly
within the task domain. In summary, our main contributions are:
An efficient end-to-end imitation learning architecture trained solely on in-domain data, without requiring large-scale pretraining.
A novel world-model--e guided trajectory generation module tailored for OSVI unseen tasks.
Robustness enhancement at test time by using a waypoint controller with re-planning.
Extensive experiments in both simulated and real-world settings, demonstrating that OSVI-WM outperforms existing methods on unseen tasks.

---

2 Related Works
One-Shot Imitation Learning: Imitation learning, or learning from demonstrations, is broadly categorized into inverse reinforcement learning I1 40] and behavior cloning [46
31 7 12, 42).
While the field has been extensively studied for years, a comprehensive review is beyond the scope Our method falls under behavior cloning, specifically in the one-shot visual imitation (OSVI) setting. In standard behavior cloning, expert demonstrations and corresponding expert actions are available to train the agent. In contrast, OSVI [12 provides only a video of the expert performing the task, without access to the underlying actions. One of the earliest works on OS [12] introduced the problem setting and proposed a soft-attention-based learning framework for block stacking. Since then, a number of have improved performance in this setting. While ego-centric setups (e.g.. DINOBot (TI have been explored, we focus on methods using external camera views. One such work, T-OSVI [O), uses transformers 49 model long-range dependencies in demonstrations, and like many other works 22 51 60). it supports expert-agent embodiment mismatch.
of this paper; we, therefore, refer readers to existing surveys [3] 35 62]
methods [1] 10 63 27
While these approaches have shown promising results, they typically evaluated on the same tasks that were used for training with minor variations, such as changes in object positions or quantities. MOSAIC [34] proposed a more challenging setup with completely unseen test tasks, building on which, AWDA 5 demonstrated that existing methods like T-OSVI [T0] and DAML (60] perform poorly in such settings. AWDA improved generalization through attributed waypoints, demonstration augmentation, and image mixup, but achieved low success rates on hard tasks requiring different responses despite visual similarity. While IMOP [63] used point clouds to aid generalization, we focus solely on monocular RGB images, which are more accessible and require no specialized hardware. Notably, none of these approaches models environment dynamics, limiting their ability to reason about future states. OSVI-WM addresses this gap through a world-modellguidee trajectory generation module that enables future-aware planning and improved generalization unseen tasks.
World Model in Robotics:
Intelligent beings like animals are believed to possess internal models of the world that support control and planning for task execution [52
Inspired by this, In reinforcement learning (RL), several works 17 28 29 18 39 56 6 44 [4] have focused on learning world models from agent rollouts. These models are then used as environments for RL policy training, often also predicting rewards based on the agent's actions. Most prior works train world models directly in the pixel space, i.e., conditioned on the current image and action, the model predicts the next-frame image, often also integrated with diffusion-based frameworks. These world models, however, requires modeling low-level pixel details that are often irrelevant for downstream tasks (2). To address this, recent methods [24] 23 36 30 learn world models in latent embedding space, allowing for more efficient and abstract representations of the environment
53
16 38
researchers have increasingly explored the use of world models in robotics 45
Beyond RL, world models have been explored for broader robotic learning tasks. For instance, SWIM 37 pre-trains a world model from human videos and fine-tunes it on robot data in an unsupervised manner. GR1 [54] and (59 propose pre-training strategies to improve downstream robotic motor control. Most of these approaches, however, rely on large-scale internet data or massive pretraining datasets, leading to high computational demands. DynaMo
(8] in contrast, proposed in-domain latent dynamics pretraining for the encoder, which was used to train downstream policies. Similarly, DINO- WWO [64] used the pre-trained DINOv2 44 model (with frozen weights) to build in-context world models and applied model predictive control based optimization for task execution. Inspired by these advances, our proposed framework, OSVI-WM, learns world model entirely from in-domain data. It predicts latent trajectories based on an expert demonstration and the agent's 's current observation. Unlike prior work, OSVI-WM, including its world model, is trained end-to-end, without any pretraining, simplifying the pipeline and reducing computational complexity.
3 Our Method
We formulate the problem following prior work in OSVI [O0 34 5]. Let T= (Ti, .Tk....TT))) denote a set of tasks, partitioned into disjoint training (Trrasn) and test (Tinst) sets. In contrast to [10 and consistent with [ our setup uses different training and test tasks, rather than minor variations of the same task. Each training task consists of multiple sequences of expert and agent trajectories. The expert trajectory comprises demonstration vidco frames {E....., EN ). while
Problem Formulation:
3

---

OSVI-WM: One-Shot Visual Imitation for Unseen Tasks
using World-Model- Trajectory Generation
Expert Demonstratio
Soft Spatial Pool
Fy
Temporal Attn Pool
Sinusoidal P.E.
World-Model-Guidee
Trajectory Generation
Waypoint Predictor
ew
ResNet Encoder
mTT... (1)
Ti+1 Wir (ir.0...... Ci (1,
WF Wa
W0
Implemented using
Predicted Waypoints
71
Causal Transformers
Agent Trajectory
Soft Spatial Pool
Loss on the predicted latent trajectory
Soft DTW Loss
Stop Grad
L1 loss
Ground Truth Waypoints
Ra
TM
This block implemented only during training.
Figure 2: OSVI-WM: The expert demonstration frames Ei...., EN and the agent's current observa- tion Ri are encoded into a latent space using a ResNet encoder. A world-model-guides trajectory generation module predicts future latent states 72, ...... : T'M which are decoded into physical waypoints for control. During training, supervision is applied on the predicted waypoints and latent states.
the agent trajectories are represented by observation-state pairs (R. where R = (R....., R&d denotes the observation frames and S denotes the corresponding agent states. Within each task, all expert and agent trajectories correspond to variations (e.g., different object configurations) of the same high-level task. The model, trained Terain. evaluated on Trest, which contains only expert demonstration videos (no agent rollouts). The performance is measured by recording the success rate of the agent in completing these test tasks. The key challenges in this setup are: (a) test tasks Teeat differ from the training tasks Tirain: (b) the expert and the agent may have different embodiments; (c) there is direct alignment between expert and agent states or actions; and (d) avoid reliance on costly large-scale data, the method should be trained in-domain.
Given expert demonstration video frames E Ey and the agent's initial observation Ri, OSVI-WM (Fig. 2 encodes them using shared ResNet-18 (25) encoder to obtain latent states for the expert and agent. These are passed the world-model-guiee
Method Overview:
trajectory generation module (Sec. 3.T to predict a trajectory of future latent states, which is pooled, first spatially, then temporally, into a compact representation. Using this as input, an MLP-based waypoint predictor generates physical waypoints (Sec. B.2). Loss functions are applied to the predicted waypoints and the latent states predicted by the world model (Sec. 3.3)
3.1 World-Model- Trajectory Generation
The expert demonstration and the agent's initial observa- tion frames are each processed through ResNet-18 25 encoder, producing feature maps of shape (F, W) be- fore the encoder's s global pooling layer, where = 512
World W
Ti+1
Agent States { 7
Model
Concat along
Concat along time dimension
feature dimension
and (H, W) depend the input resolution. These fea- tures are flattened to produce latent states for the expert frames (ei...., eN) and the agent's observation (r1), each frame with shape (F x H x W'. A latent state is com- pact, abstract representation of the environment that en- codes key features such as the configuration of objects and dynamic elements. A latent action similarly repre- sents an abstract transformation that drives state transi- tions. Operating in this latent space further helps mitigate the multi-modality of real-world action distributions. The sooidd-odel-guidee trajectory generation module consists an action model и and world model W. At timestep predicts latent action using the agent's 's current and past latent states along with the
{eNs ...... e1 Expert Demonstration
Gi+1
Action Model
π
Figure 3: WM-Guided Trajectory Gen- eration: Starting with the agent's ini- tial observation and the expert demon- stration, future states are recursively pre- dicted using action and world models.

---

expert demonstration's 's latent states. This action is used by W to predict the agent's 's next latent state:
a, (rr............................. T1,
T++1= W(r,, a............................a.. 01); i= (l..... M 1}. (1)
This multi-step recursive process (Fig. 3) generates all future latent states of the agent (r2,... ras), forming the generated trajectory. The world model (WM) essentially learns the environment's dynamics to predict the next latent state based on the current latent state and latent action.
Within the action model at each step i, the expert states (ex,..., ei and agent states {ri,.... } are concatenated along the temporal dimension, resulting in shape (N + F x H x W). This is processed through causal transformer blocks, and the output corresponding to the final temporal index is used as the latent action az. For the world model W. the latent actions {a...... a1} are concatenated with the corresponding agent states (Tay.... T1 along the feature dimension obtain shape (i, 2F x H x W). This sequence is processed through a causal transformer, and the output at the last time step is projected to shape (F x H W) produce the next state r++1-
3.2 Predicting Waypoints
The predicted trajectory of future states (72; . TTM) is reshaped back to shape (F, H, W) at each timestep to retain spatial structure. A spatial pooling module (TO) then applies 2D softmax across each feature map and computes an expected 2D coordinate over a
Spatial and Temporal Pooling:
[-1,1] grid for each feature dimension. These coordinates are concatenated, producing a spatially pooled output of shape 2F for each of the M timesteps (F2,... FMM) Sinusoidal positional embeddings are then added along the temporal dimension, and attention pooling module with learned query summarizes the sequence into a single -F-dimensional vector.
The pooled vector is passed through two-layer MLP to predict 5 attributed waypoints, each of dimension 4. Following [ each waypoint encodes the end-effector's 3D position in the camera's 's coordinate frame and binary gripper state (open/closed). These positions are transformed into the robot's coordinate frame using a known camera-to-robot transformation. waypoint controller, using inverse kinematics of the robot, computes low-level control commands for execution. To ensure accurate grasping, especially in the final centimeters before contact, we follow [5] and use end-effector-moumii depth camera correct residual pose errors.
Attributed Waypoint Prediction:
To improve robustness during execution, OSVI-WM supports waypoint re- pl-aannnnng which activates if the agent fails to complete the task using the predicted waypoints. In simulation, success or failure is automatically determined from environment configurations. Although we do not use re-planning on real-world benchmarks in this work, success (or completion) signals could be obtained via learned reward/value models [24 or iision-languagg-basee completion detec- tors 332 (77 In OSVI-WM, once all waypoints are executed, the agent captures new observation, which, combined with the original expert demonstration, generates updated waypoints. Notably OSVI-WM outperforms existing baselines even without re-planning (Fig.
Waypoint Re-Planning:
3.3 OSVI-WM Training
Loss Functions: We jointly optimize the world model and the waypoint predictions. To supervise the world model, ground-trute -ruth agent observations Re.... Rys from the training set are encoded into latent states using our encoder, followed by the spatial pooling layer to produce fz..... fM. An Ln loss is applied between these and the predicted latent states after spatial pooling (F....., FFM): Cum M EMe I|A, f.ll. To prevent model collapse and ensure stable training, we apply stop-gradient to ....... AM, preventing backpropagation through the ground-truth latent states. For waypoint supervision, we interpolate the ground-truth agent trajectory into a denser sequence and use the differentiable soft dynamic time- warping loss [9 (Cadtw). Although only 5 waypoints are used at inference, we follow (5) and predict multiple sets (1 to 5 waypoints) during training, applying separate losses for each set to improve supervision across different trajectory granularities. The overall training objective is C = Ludtw + ( where a(7) is a scheduling factor that balances the two loss terms based on the training iteration For the Meta-World [6]] benchmark experiments in Sec. 4 a(T) is initialized to I and decays exponentially to 0.05 by the end of training. For Human-Franka-PPe (T) 10; for all other datasets, keep ax (T( 1.
Training Settings: Each expert demonstration is downsampled to N 10 frames, and each agent trajectory to M = 6 frames. Following prior work [10, 5), we pair every expert demonstration
5

---

Button Press
Pick Place Wall
Window Open
Door Unlock
(a) Mcta-World Test Tasks
(b) Pick-and-Plact
Figure Simulation Environments: Test tasks are different from the ones used for training. Additionally, Pick-and-PP Place uses different embodiments for expert and agent.
Agent
Expert
Agent
Target
Target
Expert
Agent
Target
Target
Turget
Target
(a) Two-Franka-PP
(b) Human-Franka-PP
(c) Hurman Hurnan-Franka-P Franka-Pusi
Figure 5: Real-World Environments: (a) Pick-and-place setup with expert (gray) and agent (white) Franka arms mounted at different locations. (b) Similar setup with human expert and Franka agent. (c) Box push setup with human expert and Franka agent.
with every agent trajectory within the same task, resulting in combinatorial dataset expansion. We train for 10 epochs in simulation and 100 in real-world settings using Adam (33) optimizer with one-cycle learning rate scheduler [47 (start: 0.0002, end: 2 x 10 10-"), and batch size of 128. We apply asymmetric demonstration image mixup [5] regularization to help the model focus on task-relevant cues rather than video semantics. Although training
use early stopping for Meta- World to avoid overfitting and report results using the overall best rest-performing checkpoint. For simulation runs, used RTX A4000 GPU with 128 GB RAM and Intel i9 CPU, and for real-world, we used an RTX 2080 Ti, 64 GB RAM, and Intel i7 CPU.
for 10 epochs,
4 Experiments
Our experiments are designed to address the following key questions: (a) How well does OSVI-WM generalize novel, unseen tasks? (b) Can OSVI-WM be used for real-world robotic tasks? (c) What the impact of the world-model-guided trajectory generation module on overall performance? (d)
Which components of OSVI-WM are most critical its success?
4.1 Eyaluation Settings
Environments: We evaluate OSVI-WM on simulation (Fig. and real-world benchmarks (Fig.
(a) Meta-World 61 simulation benchmark of 50 manipulation tasks, split into 46 for training and 4 for testing following [ 'est tasks are different from train tasks and are divided into 'easy' tasks (Button-Press- V2, Tick-Place-Wall-V22 which differ from training tasks due to the presence/absence of distractor, and 'hard' tasks Window-Oper Handow-Open-V22 Door- Unlock-Y V2, which, despite being visually similar to Window-Close-V V2 and Door-Lock- V2 from the training set, respectively, differ in context and require distinct responses.
(b) Pick-and-Place [10] simulated benchmark involving 4 objects and targets, yielding 16 tasks. Following [5] we use 14 tasks for training and 2 for testing. This benchmark features different embodiments: the expert
Sawyer arm, and the agent uses Franka arm.
(c) Two-Franka-PP (Fig. [l): A real- world setup with two Franka arms mounted at different tabletop locations, with the gray arm the expert and the white arm the agent. The tasks are similar to Pick-and-Plas use 3 tasks for training and for testing. The difference in mounting location between the expert and agent introduces additional imitation complexity.
above, with two objects and two targets (4 tasks total). We
(d) Mannan-Fraana-PPe PP (Fig 5p): Similar to Two-Franka-Fin -nnv, but the expert is human arm and the agent Franka robot. This greater embodiment mismatch increases difficulty. As before, we use tasks for training and I for testing.

---

Table Success rates (in %) comparison on the Meta-World [6] and Pick-and- Place [0] simulation benchmarks. Best results are highlighted. We also report if a method uses additional training data.
Method
Add. Data
Pick Place
Meta- World
Easy
Hard
All
60 DAML
No
4
8
T-OSVI 10
No
10
50
7
28.5
AWDA
S
No
100
74
11
42.5
AWDA
S
Yes
100
73
30
51.5
OSVI-WM (Ours)
No
100
96
71.5
83.8
Table 2:
Real-World experiments:
Success rates and execution breakdowns (in %6 are reported.
T-OSVI* denotes T-OSVI [00] aided with end-effector depth sensing for improved grasping.
Two-Franka-PP
Hluman-Franka-1
Human-Franka- Push
Method
Breakdown
Breakdown
Breakdown
Overall Success
Reach
Grasp Obj.
Overall
Reach
Grasp
Overall
Reach
Obj.
Success
Obj.
Obj.
Success
Obj.
T-OSVI 10 T-OSVI*
52
0
0
0
0
16
96
10
4
52
36
0
0
0
AWDA
5
88
88
88
92
92
92
76
100
OSVI-WM (Ours)
96
96
96
92
92
92
100
100
Human-Frain Franaa--Faana-Push (Fig. This features two colored boxes on tabletop. The task
involves selecting a box and pushing it either forward or backward (four tasks total). We train on three and test on the held-out one. The expert is human and agent a Franka arm.
In the pick-and-place tasks above, sequences from test object-targe pairs are excluded from training, While this may appear to be an extension of training tasks, most prior methods require training on all pairs to succeed. Generalizing to an unseen pair requires the model to truly follow the expert, rather than rely solely on scene semantics. Additional environmental details are the appendix.
Baselines: As noted in Sec. and Sec. 2 most one-shot visual imitation methods are evaluated on tasks that are minor variations of training tasks. AWDA (5) is the only method that uses the same sensors and environment settings as ours and evaluates on unseen test tasks, making it our primary baseline. We also compare against two widely used methods: DAML 660 and T-OSVI [O.. In real-world experiments, T-OSVI VI fails to achieve successful grasps. To mitigate this, we augment it with the same depth- based gripper correction used in OSVI-W WM and AWDA, where an end-effector- mounted depth camera corrects residual pose errors. We refer to this variant as T-OSVI* and report its results Table 2 For OSVI-WM. re-planning used only on the MetaWorld 'hard' tasks, as we observe strong performance even without it on all other benchmarks.
4.2 Results
How well does OSVI-WM generalize novel, unseen tasks? We compare OSVI-WM against baseline methods in Ta the Pick-and-Place [10] and Meta- World [6]] benchmarks. OSVI-WM evaluated over 100 rollouts per task with varied object configurations, and success rates are reported. Baseline results are reported from 5 OSVI-WM outperforms prior methods, achieving success rates up to 30% higher than AWDA even when AWDA uses additional training data. OSVI-WM's ability to predict future latent states enables effective planning, resulting in strong performance even on the World 'hard' tasks. Failure cases in this setting often occurred in edge cases where objects, such windows or doors, were positioned far from the agent, making manipulation challenging due to depth ambiguity in external RGB observations. Nonetheless, OSVI-WN achieves 71.5% success rate on these tasks, outperforming all baselines.
Can OSVI-WM be used for real-world robotic tasks? We evaluate OSVI-WM and baselines on the real-world setups, each tested over 25 rollouts using the same 25 object config ioniigurations across methods for fair comparison (Table These setups introduce added complexity: in Two-Franka-PPe -P, the expert and agent robots are mounted at different locations, while in HaimannFranna-PPP and Human-Frame Push, the expert a human. Despite these challenges, OSVI-WM achieves consistently high success rates, showing strong real-world applicability. The table also breaks performance into sub-tasks like
7

---

( Expert Demonstration
(b) Agent Imitation Using OSVI-WM
Figure 6: Real- World Qualitative Example: The human expert demonstrates the test task of picking and placing a can in the green target, which the agent successfully imitates using OSVI-WM.
Success Rate (%)
100
Success Rate (%)
100
75
75
50
25
25
Button
Pick-Place Window
Door
Average
Window
Door Unlock
Average
Press
Wall
Open
Unlock
Open
Model
Loss
AWDD
No
With
Figure 7: Ablation study on the world model (WM) and WM loss. Using both achieves the best performance.
Figure 8: Ablation Study on OSVI-
WM's Re-Planning (RP).
reaching and grasping. In Two-Franka-PP, T-OSVI [10 reaches the object in 52% of trials but fails to grasp it. T-OSVI* (i.e., T-OSVI with grasp correction) improves grasping 36% but completes the full task only once (4%). AWDA succeeds in 88% of the trials, failing due to incorrect object selection. OSVI-WM performs best, succeeding 96% of the trials. -aannan---aaa-PPPP challenging due to the embodiment gap. In this setting, T-OSVI and T-OSVI* fail to reach the correct object, while AWDA and OSVI-WM both succeed 92 % the trials. Human Franan--aanaaaa- -ush, all methods consistently reach the correct object. However, due to the low margin for error in this setting, AWDA and T-OSVI often collided with the object during fine manipulation, leading to reduced success rates of 76% and 16,, respectively. OSVI-WM executed these movements more accurately, achieving 100% success. Since there no grasping involved this task, we do not run T-OSVI*.
more
We present a qualitative example of OSVI-WM in the Human- Franka-PP real-world setup. In Fig. 6h. a human expert demonstrates the task of placing a can on a green target. Using only this demonstration, the agent (white Franka arm) successfully completes the task (Fig. [6p) despite changes object locations and embodiment differences. This task was excluded from training, illustrating OSVI- WM''s ability to generalize to unseen tasks by following the expert's 's intent rather than relying solely on visual semanties. This example showcases OSVI-WM effectiveness under all key conditions outlined in our problem setup: unseen tasks, embodiment mismatch, unaligned trajectories, and in-domain training. More visualizations on different tasks are in the appendix.
What is the impact of the world-model-guidee trajectory generation module? We assess the impact of this module through two ablations on the World dataset. The first removes the world model loss Lumi the second removes the world-model-guidee trajectory generation module entirely, replacing it with single action model that predicts an action vector from the expert demonstration and the agent's initial observation. This vector is passed through the pooling and waypoint prediction modules. As shown in Fig. removing the world model yields the lowest performance, while adding it without Zum provides limited gains. The loss Lur is applied to future latent states predicted by the world model and independent of ground-truth waypoint supervision. acts as an auxiliary dynamics-consistency regularizer and encourages robust, task-agnostic representations, thus improving generalization to unseen tasks. Hence, combining the world model with fuom achieves the best performance, underscoring the importance of both the module and its training signal.

---

(a) Example Ground Truth Trajectory
2
(b) OSVI WM Generated Latent Trajectory Decoded to Images
Figure 9: (a) An example ground truth rollout of the Button-Press- V2 task from Meta- World 66]. (b) Latent trajectory generated using OSVI-WM's world model decoded into images for visualization.
In Fig we visualize the predicted trajectory generated by OSVI- WM's world model on the Button- Press-V2 V2 task [61]. To convert the latent states into images, we train a Transpose Convolution-banl image predictor. Images from the task are encoded using OSVI WM'ss trained encoder and spatial pooling layer, and the image predictor is trained to predict back the pixels. During evaluation, we generate the future latent states (f'...., Fd)) using the expert demonstration and initial agent frame, and decode them into images (Fig(9p) using the image predictor. A ground truth rollout shown for comparison in Fig. Dh. Two key observations emerge: (a) the predicted trajectory captures meaningful, taskccompgtiing states, effectively summarizing the task; (b) the latent states are informative enough to reconstruct interpretable images. As expected, later states appear more pixelated due to accumulated errors in recursive latent state predictions. Despite this, the decoded trajectory clearly reflects the full task execution and produces accurate downstream performance.
Which components of OSVI-WM are most critical to its success? As shown above, the world-model-guidee trajectory generation module, along with supervision from WM loss fuom- is critical to OSVI-WM's success. Re-planning further boosts performance, particu- larly on the challenging Meta-Wosd World 'hard' tasks. Fig. g isolates the impact of re-planning, showing approximate 30-point gain. For comparison,
also report the average success rate of AWDA [5] (with additional training data). Notably, even without re-planning, OSVI- WM outperforms AWDA. Table presents further ablation studies evaluating key architectural choices. Removing stop-gree graddgaddent in the WM loss significantly degrades performance, dropping Meta-World accuracy by 7% and Pick-and-Place by 80%, clearly indicating model collapse. Applying spatial pooling before the world-model-guidd trajectory generation block introduces an in- formation bottleneck, resulting in a 85% drop in Pick-and-Place and an 4% drop Meta-World Switching from multi-step recursive to single- step trajectory prediction has no ef- fect on Pick-and-Place but lowers Meta-World performance by over 15% due to Meta- World''''''s greater complexity. Together, these results show that OSVI-WM's effectiveness stems from key components: (a) the world-model-guided tra- jectory generation with WM loss, (b) re-planning, (c) stop-gradient in WM supervision, (d) spatial pooling after the world model, and (e) multi-step trajectory prediction. Regularization strategies like asymmetric demonstration image mixup [5] further improve generalization.
Table Ablation results (avg. success rates) on the Pick- Place [TO] and Meta-World [6]]. evaluating the impact of differ- ent architectural choices by individually disabling them.
Stop Grad in WM Loss after WM Prediction Pick-Place Meta-Wort
Spatial Pool Multi-Step Benchmarks
20
76.5
15
79.5
X
100
67.3
100
83.8
Division of train and test tasks: The choice of in the Meta- Meta-World simulation environment. During training, the model learns two key components: inferring tasks from the expert and acquiring the skills needed to complete these tasks. Hence, careful consideration required when choosing training tasks to ensure meaningful generalization to unseen benchmark. tasks. To further investigate this, we conducted two ablation studies on the Press-Topdoww- Topdown-Waa In the first, we excluded all button-related tasks (Butoon-Pess: Topsoo,, Button Press- Button- Button-Presss- Wall) from training and evaluated on Button-Press. The model still achieved a 97%
task division affects performance, especially

---

success rate, nearly identical to when the button-related tasks were included. This is because button pressing is relatively kinematically less complicated and the required skills were already acquired from other training tasks. In the second ablation, we removed Window-Closs from the training set and evaluated on Window-Open. Here, performance decreased from 95% to 80%..hy The sliding window tasks here involve more complex physical dynamics. So the absence of related training data led to noticeable drop in generalization. If the training set included a very large variety of tasks with diverse motions and skills, the model would potentially perform well on any task, even if no similar tasks were present in the training set. However, this approach would require large out-of-contex( datasets, significant training resources, and extended training times. In OSVI-WM, follow the train-test split proposed in AWDA [5] for consistency with prior work.
Conclusion
We introduced OSVI-WM, a novel framework for one-shot visual imitation on unseen tasks. Given single expert demonstration and the agent's initial observation, OSVI-WM encodes them into latent states and uses world-model-guider approach to predict a trajectory of future latent states. This trajectory is decoded into physical waypoints that control the robot. Extensive experiments in both simulation and real-world settings demonstrate OSVI-WM's strong performance, and ablation studies confirm the importance of key design choices in the framework. By generalizing effectively from single demonstration to unseen tasks, OSVI-WM takes a step toward enabling autonomous robotic systems capable of operating in complex, real-world domains such as industrial automation and medical assistance, without constant human supervision.
Limitations: While OSVI-WM demonstrates promising results, there are limitations that open it only predicts 3D positions and a gripper state, without modeling the end-effector's orientation. This may limit its applicability to tasks requiring rotational actions, such as screwing or side-grasping. The current formulation focuses on single-task execution; extending the framework to sequential or multi-task settings is an interesting direction for future work. The division of train and test task can potentially influence the performance of OSVI-WM as detailed in Section 4 Further, although the test tasks are unseen during training, they belong to the same family of manipulators as the training data. While these limitations do not hinder current performance, they highlight promising directions for enhancing the framework in more complex task settings.
avenues for future work. Like prior work [] 34),
Acknowledgements
This paper supported in part by the Army Research Office under grant number W911NF21-1-0155: by the New York University Abu Dhabi (NYUAD) Center for Artificial Intelligence and Robotics, funded by Tamkeen under the NYUAD Research Institute Award CG010, and by NSF under grant number 2208189.
References
[1] Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the International Conference on Machine Learning, page 1, 2004.
[2] Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from images with a joint- embedding predictive architecture. In Proceedings of the Conference Computer Vision and Pattern Recognition, pages 15619-15629, 2023.
[3] Homanga Bharadhwaj. Abhinav Gupta, Shubham Tulsiani, and Vikash Kumar. Zero-shot robot manipulation from passive human videos. arXiv preprint arXiv:2302.02011, 2023.
[4] Vishal Bhutani, Anima Majumder, Madhu Vankadari, Samrat Dutta, Aaditya Asati, and Swagat Kumar. Attentive one-shot meta-imitation learning from visual demonstration. Proceedings of the International Conference Robotics and Automation, pages 8584-8590. IEEE, 2022.
10

---

[5] Matthew Chang and Saurabh Gupta. One-shot visual imitation via attributed waypoints and demonstration augmentation. In Proceedings of the International Conference on Robotics and Automation, pages 5055-5062. IEFE, 2023.
[6] Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202..00. 09480, 2022.
[7] Zichen Jeff Cui, Yibin Wang, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. From play policy: Conditional behavior generation from uncurated robot data. In Proceedings of the International Conference on Learning Representations, 2023.
[8] Zichen Jeff Cui, Hengkai Pan, Aadhithya Iyer, Siddhant Haldar, and Lerrel Pinto. Dynamo: In- domain dynamics pretraining for visuo-motor control. A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang, editors. Proceedings of the Advances in Neural Information Processing Systems, volume 37, pages 33933 -3936.3 Curran Associates, Inc., 2024.
[9] Marco Cuturi and Mathieu Blondel. Soft-dtw: differentiable loss function for time-series. In Proceedings of the International Conference on Machine Learning, pages 894-903. PMLR, 2017.
[10] Sudeep Dasari and Abhinav Gupta. Transformers for one-shot visual imitation. In Jens Kober, Fabio Ramos, and Claire Tomlin, editors, Proceedings of the Conference on Robot Learning, volume 155 of Proceedings of Machine Learning Reseurch, pages 2071-2084. PMLR, 16-18 Nov 2021.
[11] Norman Di Palo and Edward Johns. Dinobot: Robot manipulation via retrieval and alignment with vision foundation models. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pages 2798-2805. IEEE, 2024.
[12] Yan Duan, Marein Andrychowicz. Bradly Stadie, OpenAl Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojeiech Zaremba. One-shot imitation learning. In I. Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Proceedings of the Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017
[13] Chrisantus Eze and Christopher Crick. Learning by watching: review of video-based learning approaches for robot manipulation. arXiv preprint arXiv:2402.07127, 2024.
[14] Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. In Proceed- ings of the International Conference on Robotics and Automation, pages 7786-2793, IEEE, 2017.
[15] Chelsea Finn, Tianhe Yu, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot visual imitation learning via meta-learning In Proceedings of the Conference on robot learning, pages 357-368. PMLR, 2017.
[16] J Randall Flanagan and Alan M Wing. The role of internal models in motion planning and control: evidence from grip force adjustments during movements of hand-held loads. Journal of Neuroscience, 17(4)1119--152, 1997.
[17] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Ccsa-Bianchh, and R. Garnett, editors, Proceedings of the Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
[18) Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Proceedings of the International Conference on Machine Learning, pages 2555-2565. PMLR, 2019.
[19] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In Proceedings of the International Conference on Learning Representations, 2020.

---

[20] Danijar Hafner. Timothy P Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In Proceedings of the International Conference on Learning Representations, 2021.
[21] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301 2.K1::2300.0104,, 2023.
[22] Siddhant Haldar, Zhuoran Peng, and Lerrel Pinto. Baku: An efficient transformer for multi-task policy learning. In Proceedings of the Annual Conference on Neural Information Processing Systems, 2024.
[23] Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpe2: Scalable, robust world models for continuous control. In Proceedings of the International Conference on Learning Representations, 2024.
[24] Nicklas A Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predictive control. In Proceedings of the International Conference on Machine Learning, pages 8387-8406. PMLR, 2022.
[25] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, In Proceedings of the Conference on Computer Vision and Pattern Recognition. pages 770-778, 2016.
[26] Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning:
A survey of learning methods. ACM Comput. Surv., 50(2), April 2017. ISSN 0360-0300.
[27] Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. Bc-z: Zero-shot task generalization with robotic imitation learning. In Proceedings of the Conference on Robot Learning, pages 991-1002. PMLR, 2022.
[28] Eukasz. Kaiser, Mohammad Babaeizadeh, Piotr Mitos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model based reinforcement learning for atari. In Proceedings of the International Conference on Learning Representations, 2020.
[29] Kushal Kedia, Prithwish Dan, Angela Chao, Maximus Adrian Pace, and Sanjiban Choudhury. One-shot imitation under mismatched execution. arXiv preprint arXiv:2409.06615, 2024.
[30] Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. Advances in Neural Information Processing Systems, 33:741-752, 2020.
[31] Seungjae Lee, Yibin Wang. Haritheja Etukuru, H Jin Kim, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. Behavior generation with latent actions. In Proceedings of the International Conference on Machine Learning, pages
PMLR, 2024.
[32] Haichao Liu, Sikai Guo, Pengfei Mai, Jiahang Cao, Haoang Li, and Jun Ma. RobodexvIm: Vi- sual language model-enabled task planning and motion control for dexterous robot manipulation. arXiv preprint arXiv:2503.0/6111 2025.
[33] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05100 2017.
[34] Zhao Mandi, Fangchen Liu, Kimin Lee, and Pieter Abbeel. Towards more generalizable one- shot visual imitation learning. In Proceedings of the International Conference on Robotics and Automation, pages 2434-2444. IEEE, 2022.
[35] Robert McCarthy, Daniel CH Tan, Dominik Schmidt, Fernando Acero, Nathan Herr, Yilun Du, Thomas G Thuruthel, and Zhibin Li. Towards generalist robot learning from internet video: A survey, arXiv preprint arXiv:24044.966 2024.
[36] Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak. Dis- covering and achieving goals via world models. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Proceedings of the Advances in Neural Information Processing Systems, volume 34, pages 24379-24391. Curran Associates, Inc., 2021.
12

---

[37] Russell Mendonca, Shikhar Bahl, and Deepak Pathak. Structured World Models from Human Videos. In Proceedings of Robotics: Science and Systems, Daegu, Republic of Korea, July 2023. doi: (0.15607/RSSS.223..XXX.
[38] R Chris Miall and Daniel M Wolpert. Forward models for physiological motor control. Neural networks, 98): 126-1129, 1996.
[39] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficien world
models. In Proceedings of the International Conference on Learning Representationss, 2023.
[40] Andrew Ng, Stuart Russell, et al. Algorithms for inverse reinforcement learning. In Proceed- ings of the International Conference Machine Learning, volume 1, page 2, 2000.
[41] Maxime Oquab, Timothée Dareet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. Transactions on Machine Learning Research Journal, pages 1-31, 2024.
[42] Deepak Pathak, Parsa Mahmoudich, Guanghao Luo, Pulkit Agrawal, Dian Chen, Yide Shentu, Evan Shelhamer, Jitendra Malik, Alexei A Efros, and Trevor Darrell. Zero-shot visual imitation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition workshops, pages 2050-2053, 2018.
[43] Ilija Radosavovic, Tete Xiao, Stephen James, Picter Abbeel, Jitendra Malik, and Trevor Darrell. Real-world robot learning with masked visual pre-training. In Karen Liu, Dana Kulic, and Jeff Ichnowski, editors, Proceedings of the Conference on Robot Learning, volume 205 of Proceedings of Muchine Learning Reseurch, pages 416-426. PMLR, 14-18 Dec 2023.
[44] Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. In Proceedings of the International Conference on Learning Representations, 2023.
[45] Ryo Sakagami, Florian S Lay, Andreas Dömel, Martin J Schuster, Alin Albu-Schäffer, and Freek Stulp. Robotic world models- conielssscccneepuliaaatiooo review, and engineering best practices. Frontiers in Robotics and AI, 10: 10:1253049, 2023.
[46] Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya (Arty) Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning k modes with one stone. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Proceedings of the Advances in Neural Information Processing Systems, volume 35, pages 22955-22968. Curran Associates, Inc., 2022.
[47] Leslie N Smith and Nicholay Topin. Super-convergence: Very fast training of neural networks using large learning rates. In Artificial intelligence and machine learning for multi-domain operations applications, volume 1 11006, pages 369-386. SPIE, 2019.
[48]
Richard Sutton and Andrew Barto. Reinforcement Learning: An Introduction. MIT Press, 2 edition, 2018.
[49] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[50] Stefan Vogt and Roland Thomaschke. From visuo-motor interactions to imitation learning: behavioural and brain imaging studies. Journal of Sports Sciences, 25(5):497-517, 2007.
[51] Chen Wang, Linxi Fan, Jiankai Sun, Ruohan Zhang, Li Fei-Fei, Danfei Xu, Yuke Zhu, and Anima Anandkumar. Mimicplay: Long-horizon imitation learning by watching human play. In Jie Tan, Mare Toussaint, and Kourosh Darvish, editors, Proceedings of the Conference Robot Learning, volume 229 of Proceedings of Machine Learning Research, pages 201-221. PMLR, 06-09 Nov 2023.
[52] Daniel M Wolpert, Zoubin Ghahramani, and Michael Jordan, An internal model for sensori- motor integration. Science, 26905232) 1995.
13

---

[53] Daniel M Wolpert, R Chris Miall, and Mitsuo Kawato. Internal models in the cerebellum. Trends in cognitive sciences, (0))338-3477 1998.
[54] Hongtao Wu, Ya Jing, Chilam Cheang, Guangzeng Chen, Jiafeng Xu, Xinghang Li, Minghuan Liu, Hang Li, and Tao Kong. Unleashing large-scale video generative pre-training for visual robot manipulation. In Proceedings of the International Conference on Learning Representa- tions, 2024.
[55] Jialong Wu, Haoyu Ma, Chaoyi Deng, and Mingsheng Long. Pre-training contextualized world models with in-the-wild videos for reinforcement learning. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and Levine, editors. Proceedings of the Advances in Neural Information Processing Systems, volume 36, pages 39719-39743. Curran Associates, Inc., 2023.
[56] Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. Day- dreamer: World models for physical robot learning. In Proceedings of the Conference on Robot Learning, pages
PMLR, 2023.
[57] Chuyan Xiong, Chengyu Shen, Xiaoqi Li, Kaichen Zhou, Jiaming Liu, Ruiping Wang, and Hao Dong. Autonomous interactive correction milm for robust robotic manipulation. In 8th Annual Conference Robot Learning, 2024.
[58] Mengda Xu, Zhenjia Xu, Cheng Chi, Manuela Veloso, and Shuran Song. Xskill: Cross embodiment skill discovery. In Proceedings of the Conference on Robot Learning, pages 3536-3555 PMLR, 2023.
[59] Jiange Yang, Bei Liu, Jianlong Fu, Bocheng Pan, Gangshan Wu, and Limin Wang. Spa- tiotemporal predictive pre-training for robotic motor control. arXiv preprint 2xXiv:2403.05344 2024.
[60] Tianhe Yu, Chelsea Finn, Annie Xie, Sudeep Dasari, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot imitation from observing humans via domain addaii-adaptive meta-learning. arXiv preprint arXiv: 1802.01557 2018
[61] 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 Proceedings of the Conference on Robot Learning, pages 1094-1 1000. PMLR, 2020.
[62] Maryam Zare, Parham M Kebria, Abbas Khosravi, and Saeid Nahavandi. A survey of imitation learning: Algorithms, recent developments, and challenges. IEEE Transactions Cybernetics, 2024.
[63] Xinyu Zhang and Abdeslam Boularias. One-shot imitation learning with invariance matching for robotic manipulation. arXiv preprint arXiv:2405. 13178, 2024.
[64] Gaoyue Zhou, Hengkai Pan, Yann LeCun, and Lerrel Pinto. Dino-wm: World models on pre-trained visual features enable zero-shot planning. arXiv preprint arXiv:2411.04400 2024.
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly state the paper's core contribution: framework for one-shot visual imitation for unseen tasks via world- cele-mooell-aaae based trajectory synthesis. This has been described in detail in Sections [ [ and 3 and supported with
experiments on benchmark environments demonstrating generalization to unseen tasks (Section 4) aligning well with the stated contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The paper discusses the limitations in the second paragraph of Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
15

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the suppiemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The paper provides sufficient detail Sections 3 and 4 to reproduce the main experiments, with additional implementation specifics in the appendix. Code will be released upon acceptance, further supporting reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
16

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
The paper uses publicly available simulation benchmarks. The code and real-world experimental data will be made publicly available through the Github repository linked in the abstract.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to rum to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips //ipsscc/ppuiii htce/ublic/aaee htt::/uiiss///eee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
The paper provides comprehensive training and evaluation details, including data splits, hyperparameters, and optimization settings in Sections 3.3 and with additional specifics in the appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We have provided extensive experiments showing the benefits of the proposed approach. Since we have empirically found the model's output waypoints to be highly repeatable for the same environment configuration, error bars are not explicitly provided. In Section E we report mean success rates over multiple rollouts with different environ- ment configurations and initialization showing the performance improvements over prior basclines.
17

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the eπror bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (c.g. negative error rates).
show in tables
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
The computer resources used are provided in Section 3.3 of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httppppe//eurrpp...cpppp///////p
Answer: [Yes]
Justification:
The paper follows the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
18

---

The paper discusses, in the introduction and conclusion, how the proposed methodology for one-shot imitation learning can positively impact society by enabling autonomous robotic systems in areas such as medical assistance and industrial automation. Since the paper addresses foundational research, there are no known negative societal impacts.
Justification:
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All assets used in the paper are properly cited, the license and terms of use explicitly mentioned and properly respected.
Guidelines:
19

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooel as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
20

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does involve LLMs as any important, original, or non standardd components.
Please refer to our LLM policy for what should or should not be described.
21


Reviews:
[{"review_id": "kjBpKwJT09", "reviewer": "Reviewer_9img", "summary": "This paper tackles poor generalization in one-shot visual imitation learning by proposing a framework guided by a learned world model. Instead of direct policy learning, the method plans a trajectory of states and actions in the world model's latent space, which is then decoded into physical waypoints for execution. This approach is designed to better handle unseen tasks with different underlying contexts. Evaluated across simulated and three real-world platforms, the method reportedly demonstrates consistent and significant improvements over prior art, with over a 30% performance gain in some scenarios.", "questions": "- In the simulated experiments, it seems that how to divide the train and test tasks will dramatically affect the experimental results. For instance, the Button-Press task shares very similar action behaviors with the Button-Press-Topdown task. So what are the considerations behind this split? How are the similar tasks in the trainset different from the testing tasks?\n- What are the specific settings of Human-Franka tasks? What is the camera setup? Are the objects fully randomized on the tabletop?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n- Technically solid method pipeline. One-shot transfer to unseen tasks is a challenging topic.\n- Comprehensive experiments conducted on 2 sim and 3 real-world benchmarks.\n- Conducting Human-to-Robot skill transfer, where the embodiment gap is quite large.\n\n## Weaknesses\n- Lack of comparison against one-shot imitation learning methods through alignment and open-loop replay, e.g., DINOBot.\n- I am not familiar with world-model-based methods for manipulation. While Human-to-Robot skill transfer is impressive to me, the studied manipulation tasks are generally simple, i.e., top-down grasping with low precision requirements. Given that BC-based or VLA-based policies have enabled very dexterous manipulation behaviors, I believe a discussion is necessary to illustrate the high-level pros and cons of this world-model-based approach when compared with more prevalent and straightforward BC-based methods.\n- Please refer to the questions for more minor weaknesses.\n\nDINOBot: Robot Manipulation via Retrieval and Alignment with Vision Foundation Models. ICRA'24.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed review and for the constructive comments. We further appreciate your recognition of the strengths of the paper. In the following, we address each of the concerns (weaknesses) and questions.\n\n**Concern 1: Lack of comparison against one-shot imitation learning methods through alignment and open-loop replay, e.g., DINOBot.**\n\n**Response**: DINOBot (and similar pixel alignment and open-loop replay methods) perform one-shot imitation by using a gripper-mounted camera to visually compare the current scene to demonstration frames. It retrieves the closest visual match and estimates pixel-level alignment to grip the object and execute a trajectory. While its use of foundation models for visual retrieval followed by pixel-level alignment is unique and effective without further training, its problem setup differs from ours. DINOBot relies solely on an ego-centric gripper camera, whereas OSVI-WM uses external third-person cameras. Consequently, DINOBot has shortcomings against external camera based methods as discussed below. Moreover, adapting DINOBot’s retrieval and pixel alignment framework for external cameras is not trivial. Therefore, we restricted our comparisons to one-shot visual imitation methods that use external third-person camera views. The main shortcomings we observed when attempting to use DINOBot are detailed below.\n1. To conduct retrieval and pixel alignment, DINOBot requires the target object to be visible to the gripper camera at all times. In our experiments, the target is often initially outside the gripper’s field of view, making such assumptions invalid.\n2. In multi-object settings like our pick-and-place tasks, which involve four randomly arranged objects and four targets, both the current observation and demonstration images contain multiple objects with different configurations. As a result, there is no single, consistent visual transformation that aligns the current scene’s pixels to that of the demonstration. DINOBot’s reliance on clear visual correspondence for a single object breaks down in such cluttered and compositional environments.\n\nWe will include this discussion on DINOBot and similar works in the updated manuscript to clarify the differences.\n\n**Concern 2: I am not familiar with world-model-based methods for manipulation. While Human-to-Robot skill transfer is impressive to me, the studied manipulation tasks are generally simple, i.e., top-down grasping with low precision requirements. Given that BC-based or VLA-based policies have enabled very dexterous manipulation behaviors, I believe a discussion is necessary to illustrate the high-level pros and cons of this world-model-based approach when compared with more prevalent and straightforward BC-based methods.**\n\n**Response**: Thank you for your suggestion. We will include a discussion on the pros and cons of our world-model-based approach compared to straightforward BC-based or VLA-based approaches in the updated manuscript. Generally, both world-model-based and BC-based (or VLA-based) approaches ultimately aim to execute tasks by predicting actions. However, the key distinction is that world-model-based approaches also learn the world dynamics independently of the task. This discourages the model from overly fixating on action prediction and encourages learning more generic features through the additional training signal provided by the world model loss. Most existing methods, whether world-model-based or BC-based, typically focus on training and testing within the same set of tasks, as detailed in our manuscript. Many of these methods have demonstrated good results in complex dexterous manipulation tasks, as you mentioned. However, when faced with out-of-distribution tasks, like those in our examples, existing BC-based methods such as [5, 10, 57] often show low success rates, even on top-down grasping tasks. Additionally, our experience indicates that even recent VLA-based robot foundation models like Pi0 [A] and Groot [B] require in-context retraining for new tasks for accurate manipulation, despite being pre-trained on large datasets.\n\n**Question 1-A: In the simulated experiments, it seems that how to divide the train and test tasks will dramatically affect the experimental results. For instance, the Button-Press task shares very similar action behaviors with the Button-Press-Topdown task.**\n\n**Response**: How one divides the train-test task will indeed impact results, particularly in the Meta-World simulation environment. In OSVI-WM, to be consistent with prior work, we adopt the train-test split proposed in AWDA [5] for Meta-World. During training, the model learns two key components: inferring tasks from the expert and acquiring the skills needed to complete these tasks. Hence, careful consideration is required when choosing training tasks to ensure meaningful generalization to unseen tasks. To explore this further, we conducted two new ablation studies:\n* Ablation 1: Button Press Tasks Removed from Training. We removed all training tasks related to button pressing (Button-Press-Topdown, Button-Press-Topdown-Wall, Button-Press-Wall) and evaluated the model on the Button-Press task. The model still achieved a high success rate of 97%, nearly identical to when the button-related tasks were included. This is because button pressing is relatively kinematically less complicated and the required skills were already acquired from other training tasks.\n* Ablation 2: Window Tasks Removed from Training. We removed the Window-Close task from training and evaluated on Window-Open. In this case, performance dropped from 95% to 80%. The sliding window tasks here involve more complex physical dynamics. So the absence of related training data led to a noticeable drop in generalization. \n\nIf the training set included a very large variety of tasks with diverse motions and skills, the model would potentially perform well on any task, even if no similar tasks were present in the training set. However, this approach would require large out-of-context datasets, significant training resources, and extended training times. \n\n**Question 1-B: So what are the considerations behind this split?**\n\n**Response**: Based on Ablation 1 and Ablation 2, we could, in principle, construct numerous train-test splits and get different results. However, to maintain consistency with prior work, we adopt the split proposed in AWDA [5] for Meta-World. \n\n**Question 1-C: How are the similar tasks in the trainset different from the testing tasks?**\n\n**Response**: As mentioned in Section 4.1 of the manuscript and detailed further in Section C of the appendix, the similar tasks in the training set differ from those in the testing set in Meta-World either through the introduction or omission of distractors such as walls (for the ‘easy’ tasks), or by requiring a different set of motions (for the ‘hard’ tasks), which can only be executed correctly by following the expert demonstration.\n\n**Question 2: What are the specific settings of Human-Franka tasks? What is the camera setup? Are the objects fully randomized on the tabletop?**\n\n**Response**: In the Human-Franka tasks (Human-Franka-PP and Human-Franka-Push), the Franka agent is positioned on a tabletop, with a camera fixed near the table's edge at a height that provides a clear view of the table and the objects, as depicted in Figs. 5b and 5c. We utilize an Intel RealSense D455 camera with known intrinsic and extrinsic (camera-to-robot transformation) parameters, capturing images at 640x480 resolution. Additionally, an Intel RealSense D435 is mounted on the robot's end-effector to assist with grasping, as detailed in Section 3.2 of the manuscript. Object locations are fully randomized during both training and testing. For testing, to ensure fairness across all methods, we use the same 25 random object configurations across all methods.\n\n**References**\n\n[A] Black, Kevin, Noah Brown, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai et al. \"$\\pi_0 $: A Vision-Language-Action Flow Model for General Robot Control.\" arXiv preprint arXiv:2410.24164 (2024).\n\n[B] Bjorck, Johan, Fernando Castañeda, Nikita Cherniadev, Xingye Da, Runyu Ding, Linxi Fan, Yu Fang et al. \"Gr00t n1: An open foundation model for generalist humanoid robots.\" arXiv preprint arXiv:2503.14734 (2025)."}, {"review_id": "WZIvf5OcPw", "reviewer": "Reviewer_DoHi", "summary": "This paper introduces OSVI-WM, a novel framework for one-shot visual imitation learning (OSVI) designed to generalize to unseen tasks. The core problem it addresses is that existing OSVI methods often fail when test tasks are semantically or structurally different from training tasks. Also, the agent and the expert have different embodiments and there is no direct alignment between their states or actions. The proposed solution is a model that, given a single expert demonstration video and the robot's initial observation, uses a learned action model (to infer the intention of the task shown by the expert and output a latent action) and a learned world model to plan a sequence of future latent agent states corresponding to the task. The sequence is then decoded into a series of physical waypoints for the robot's controller to imitate and execute. The entire framework, including the world model, is trained end-to-end on in-domain data without requiring massive external video pretraining. The authors demonstrate through experiments on simulated and real-world robotic platforms that OSVI-WM significantly outperforms prior methods, particularly on challenging, unseen tasks.", "questions": "1. The action model is an inverse dynamic model (IDM) conditioned on given tasks, and the world model is a corresponding forward dynamics model (FDM), in the context of latent action learning [1-3]. Could you elaborate on the benefit of this architecture, versus a single, larger model that directly predicts the next latent states $r_{i+1}$ from the history and the expert demonstration? Other design choices also include, directly learning a seq2seq-style connector, to translate the expert latent state sequence into corresponding agent space [4, 5]. Does the factorization of action model and dynamics model encourage learning a better representation or extracting the task semantics than other approaches? \n2. The world model here is conditioned on the expert videos (tasks), so it is more like a goal-conditioned future state predictor rather than pure dynamics model. While effective, I think the framing could be more precise.\n3. The process of transforming the generated waypoints into low-level robot actions requires further clarification (Attributed Waypoint Prediction in Section 3.2). For example, how does the waypoint controller computes low-level control commands for execution using inverse kinematics of the robot?\n4. The current method assumes the availability of aligned image-state pairs in the agent sequence. Could the authors discuss how the proposed approach could be adapted to a setting where only image-action trajectories are provided, where proprioceptive states of the agent are replaced by robot actions?\n\n[1] Ye W, Zhang Y, Abbeel P, et al. Become a proficient player with limited data through watching pure videos[C]//The Eleventh International Conference on Learning Representations. 2022.\\\n[2] Schmidt D, Jiang M. Learning to act without actions[J]. arXiv preprint arXiv:2312.10812, 2023.\\\n[3] Bruce J, Dennis M D, Edwards A, et al. Genie: Generative interactive environments[C]//Forty-first International Conference on Machine Learning. 2024.\\\n[4] Mazzaglia P, Verbelen T, Dhoedt B, et al. GenRL: Multimodal-foundation world models for generalization in embodied agents[J]. Advances in neural information processing systems, 2024, 37: 27529-27555.\\\n[5] Wang Y, Yu R, Wan S, et al. FOUNDER: Grounding Foundation Models in World Models for Open-Ended Embodied Decision Making[C]//Forty-second International Conference on Machine Learning. 2025.", "limitations": "The authors should be more precise about the boundaries of their method's generalization capabilities. Acknowledging that the \"unseen\" tasks are still \"in-domain\" (i.e., from the same family of manipulations as the training data) would provide a more nuanced and accurate picture of the current work's scope.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper tackles a crucial and difficult problem in robotics: one-shot generalization to new tasks from a single video. Moving beyond simple variations in object positions to tasks that are semantically distinct is a significant step towards more capable and flexible robots.\n2. The technical quality of the paper is high. The architecture is well-designed and the experimental evaluation is exceptionally thorough, including challenging human-to-robot imitation with embodiment mismatch. The comprehensive ablation studies are also convincing.\n3. The core idea lies in learning a meta-action-model to infer the semantics of the given task and output high-level latent actions, and integrating a world model to generate future latent trajectory plans given the actions and the expert video, which is then decoded it into low-level waypoints. This is a well-motivated approach.\n\nWeaknesses:\n1. The re-planning mechanism boosts performance on hard tasks, but it need manual assessment in real-world autonomy, requiring a human in the loop. An automated failure detection mechanism is needed for re-planning in real-world application.\n2. The paper's main claim is generalization to unseen tasks, but the \"unseen\" test tasks in are semantic or combinatorial variations of tasks seen during training. This is a valid and challenging setup, but it leaves open the question of how the model would perform on tasks that are truly out-of-domain (e.g., training on pick-place tasks and testing on a wiping or drawer-opening task). The generalization claim, while supported by the experiments, may be limited to new tasks within the same family of manipulations.", "was_revised": true, "final_justification": "I would like to thank the authors for their detailed response. I maintain my positive score.", "rebuttal": "We thank the reviewer for the detailed review, acknowledging the strengths of the paper, and for the constructive comments. We address each of the concerns (weaknesses) and questions in the following.\n\n**Concern 1: The re-planning … application.**\n\n**Response**: OSVI-WM uses re-planning only for Meta-World 'hard' tasks (Section 4.1). Since these are simulation tasks, success is automatically verified by the environment configuration. Although we do not apply re-planning in real-world experiments, enabling autonomous re-planning would require automated task success detection. Some prior works have explored vision-language models [A,B] or reward/value networks [22] for this purpose. We will include a discussion on this in the updated manuscript. Notably, even without re-planning, OSVI-WM surpasses the performance of existing methods (Fig. 8).\n\n**Concern 2: The paper's … manipulations.**\n\n**Response**: During training, the model essentially learns two components: to infer tasks from the expert and to acquire the skills necessary to complete these tasks. In completely out-of-domain tasks, the model might still grasp the task context from the expert but may not learn the required skills if the training and testing tasks involve completely different skill sets. For instance, pick-and-place tasks involve skills like reaching, grasping, moving, and placing objects, whereas wiping requires maintaining contact with a surface while moving the end-effector. This skill disparity can lead to failures in testing if the model is trained only on pick-and-place tasks. To further support this, we conducted two additional ablations:\n* Ablation 1: We removed all training tasks related to button pressing (Button-Press-Topdown, Button-Press-Topdown-Wall, Button-Press-Wall) and evaluated the model on the Button-Press task. The model still achieved a high success rate of 97%, nearly identical to when the button-related tasks were included. This is because button pressing is relatively kinematically less complicated and the required skills were already acquired from other training tasks.\n* Ablation 2: We removed the Window-Close task from training and evaluated on Window-Open. In this case, performance dropped from 95% to 80%. The sliding window tasks here involve more complex physical dynamics. So the absence of related training data led to a noticeable drop in generalization.\n\nHowever, if the training set included a very large variety of tasks with diverse motions and skills, the model would potentially perform well on any task, even if no similar tasks were present in the training set. Achieving this would require large datasets, high training resources, and training times. We will add a discussion on this in the updated manuscript to clarify how the generalization scales to different unseen tasks. \n\n**Question 1-A: The action model is an inverse dynamic model (IDM) conditioned on given tasks, and the world model is a corresponding forward dynamics model (FDM), in the context of latent action learning [1-3].**\n\n**Response**: The world model functions like a forward dynamics model, while the action model resembles an inverse dynamics model, with the distinction that OSVI-WM’s action model is conditioned on expert videos rather than the agent’s future states ($a_i = \\pi(r_i, \\dots, r_1, e_N, \\dots, e_1)$). \n\n**Question 1-B: Could you elaborate … approaches?**\n\n**Response**: Directly predicting the next latent state using only ground truth actions can cause the model to overfit to those actions, resulting in the model only learning features necessary for such action prediction. This dependency often results in poorer generalization to unseen test tasks (Fig. 7). In contrast, incorporating a world model loss provides auxiliary signals that help the model learn environment dynamics independently of any tasks. While a seq2seq-style connector directly learns the ground truth states, the factorization between the action and dynamics (world) model offers key advantages. The dynamics model depends only on the latent actions and current agent states, not directly on the goals, thus enhancing the model's interpretability and modularity and encouraging better representation learning. Additionally, if the goal representation changes (e.g., to text), only the actions model needs retraining, leaving the world model intact.\n\n**Question 2: The world model … precise**\n\n**Response**: OSVI-WM’s world model is conditioned solely on the latent actions and the agent’s states, rather than directly on the expert videos ($r_{i+1} = \\mathcal{W}(r_i, a_i, \\dots, r_1, a_1)$). The action model, however, is conditioned on the expert videos ($a_i = \\pi(r_i, \\dots, r_1, e_N, \\dots, e_1)$). By making this separation, we get key advantages as detailed in the response to your Question 1-B. We will make this difference clear in the updated manuscript.\n\n**Question 3: The process of … robot?**\n\n**Response**: The generated waypoints represent the end-effector’s 3D positions in the camera’s coordinate frame and binary gripper states (Section 3.2). These positions are transformed into the robot frame using a known camera-to-robot transformation. We denote these 3D positions in the robot’s frame as $p_k = (x_k, y_k, z_k)$, where $k = \\{1, 2, \\ldots, q\\}$. The desired rotation $\\mathrm{rot}_k$ of the end-effector (which remains constant in our case) is concatenated to $p_k$ to obtain 6D desired pose $H_k = (x_k, y_k, z_k, \\mathrm{rot}_k)$. We use this desired pose with the controller as detailed below:\n* Simulation Experiments: Obtaining the joint angles $\\theta_k$ involves solving the inverse kinematics problem , $\\mathrm{F.K.}(\\theta_k) = H_k$. Given that the robots used in our experiments have more than 6 degrees of freedom, the inverse kinematics does not have a unique solution. Therefore, we utilize the following optimization problem:\n$$\\min_{\\theta_k} \\left\\| H_k - \\mathrm{F.K.}(\\theta_k) \\right\\|$$\nwhere $\\theta_k$ is the robot's joint angles and $\\mathrm{F.K.}$ is the forward kinematics. To avoid jerky motion from large pose jumps between waypoints, we interpolate sub-waypoints between the current and target waypoint poses. The controller tracks these sub-waypoints, while the gripper follows the predicted gripper state.\n\n* Real-World Experiments: We employ a differential inverse kinematics approach, optimizing for joint velocity by solving: \n$$\\min_{\\dot{\\theta}_k} \\left\\| J \\dot{\\theta}_k - K_p e \\right\\|^2 + \\lambda  \\left\\|(I - J^\\dagger J)(\\dot{\\theta}_k - \\dot{\\theta}_b) \\right\\|^2$$\n$J$: Jacobian of the end-effector; $\\dot{\\theta}_k$: joint velocity; $K_p$: proportionality constant; $e$: pose error between the current and target poses. The second term in the equation above ensures that the joint velocities remain close to a nominal joint velocity $\\dot{\\theta}_b$ (please refer to [C] for more details). The resulting velocities are passed to the Franka robot, running an internal proprietary controller at approximately 4 kHz. As in simulation, sub-waypoints ensure smooth motion, and the gripper executes the predicted open/close action.\n\n\n\n**Question 4: The current … actions?**\n\n**Response**: In most robotics datasets [58,10,D,E,F,G,H] and real data collected in the lab, access to the robot’s proprioceptive states, typically joint angles and angular velocities, is standard. These states can be converted into end-effector poses and velocities via the robot’s forward kinematics. In OSVI-WM, we utilize these end-effector poses to generate ground-truth waypoints, which serve as target actions for training the model as detailed in the manuscript. However, in the uncommon scenario where proprioceptive states are unavailable, robotic datasets may represent actions as the difference between successive joint angles (or end-effector poses). We can still reconstruct the full joint trajectory by cumulatively summing these deltas starting from a known initial joint configuration (which is feasible with current-day robot hardware, which provides reasonable sensor resolutions; methods to reconstruct estimates of the robot pose from external cameras are also available). Using forward kinematics, we can then convert the reconstructed joint angles at each time step into end-effector poses, which are subsequently used to generate ground-truth waypoints for training OSVI-WM. This reconstruction-based approach allows us to maintain the overall framework and learning objectives of OSVI-WM even in the absence of direct proprioceptive state data.\n\n**Limitations: The authors … scope.**\n\n**Response**: Thank you for this feedback. We will include a discussion of this in the limitation section, along the lines of the response to your concern 2.\n\n**References**\n\n[A] Liu, Haichao, et al. \"Robodexvlm: Visual language model-enabled task planning and motion control for dexterous robot manipulation.\" arXiv preprint arXiv:2503.01616 (2025).\n\n[B] Xiong, Chuyan, et al. \"Autonomous interactive correction MLLM for robust robotic manipulation.\" In 8th Annual Conference on Robot Learning. 2024.\n\n[C] Lynch, Kevin M., and Frank C. Park. Modern robotics. Cambridge University Press, 2017.\n\n[D] Nasiriany, Soroush, et al. \"Robocasa: Large-scale simulation of everyday tasks for generalist robots.\" arXiv preprint arXiv:2406.02523 (2024).\n\n[E] Abhishek Gupta, et al.. “Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning”. arXiv preprint arXiv:1910.11956, 2019. 2, 5, 17\n\n[F] Vikash Kumar. “Manipulators and Manipulation in high dimensional spaces”. PhD thesis,\nUniversity of Washington, Seattle, 2016.\n\n[G] Kumar, Vikash, et al. \"Robohive: A unified framework for robot learning.\" Advances in Neural Information Processing Systems 36 (2023): 44323-44340.\n\n[H] Khazatsky, Alexander, et al. \"Droid: A large-scale in-the-wild robot manipulation dataset.\" arXiv preprint arXiv:2403.12945 (2024)."}, {"review_id": "4XKHjUS94g", "reviewer": "Reviewer_5LPE", "summary": "This paper presents OSVI-WM, a novel one-shot visual imitation learning framework based on world-model-guided trajectory generation, designed to address the limited generalization capability of existing methods on unseen tasks. OSVI-WM integrates expert demonstration videos with the agent's initial observation to predict sequences of latent states and actions through a learned world model, which are subsequently decoded into physical waypoints for task execution. Extensive experiments conducted across both simulated and real-world robotic platforms demonstrate the effectiveness of the proposed method, with results showing that OSVI-WM significantly outperforms existing approaches, particularly on \"hard\" tasks that exhibit semantic or structural differences. Furthermore, the framework incorporates waypoint re-planning to enhance execution robustness. \nThe key contributions of this work include: (1) an end-to-end architecture requiring no large-scale pretraining; (2) a dedicated world-model-guided trajectory generation module for one-shot visual imitation; (3) improved test-time robustness through re-planning; and (4) comprehensive experimental validation across diverse environments.", "questions": "1.Generalization to Unseen Positions and Objects: The paper states that OSVI-WM can generalize to new tasks with different object positions and across different embodiments, but it does not provide a detailed explanation of how the model achieves this capability. For example, how does the model ensure the accuracy of action prediction when encountering object placements or new object categories that were not seen during training? Additionally, compared to baseline methods, does the robustness of OSVI-WM to changes in object positions stem from the dynamic modeling of the world model? Further clarification is needed.\n2.Visualization and Analysis of Latent States and Actions: The paper demonstrates the visualization of latent states through image reconstruction (Figure 9), but lacks an in-depth analysis of the visualization of latent actions and their relationship with states. For example, do latent actions correspond to the movements of the robotic arm and changes in the interacted objects? How do changes in latent states drive action generation? It is suggested to supplement the visualization and analyze the correlation between the two to enhance interpretability.\n3.The Significance and Efficiency of Replanning: Replanning is described as “generating new waypoints when the task is not completed,” which seems equivalent to repeated execution until success. If so, this strategy may reduce efficiency in practical applications due to multiple attempts. The triggering conditions for replanning (such as timeouts or deviation thresholds) and its comparison with the success rate of single planning need to be explained. For example, Figure 8 shows that replanning improves performance, but it does not quantify its time cost. If replanning is merely a repeated call to the same module, then its significance is limited.\n4.Redundancy in Latent States: Although latent states are designed to be “compact” representations, the features extracted by the Res-Net encoder may contain a large amount of background information unrelated to the task (such as texture or lighting). How can the world model focus on task-relevant features? Moreover, if there is redundancy in the latent space, will it reduce the efficiency of trajectory generation?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths：\n1.Novel Integration of World Models: OSVI-WM is the first to introduce a world model into one-shot visual imitation learning, establishing an explicit environment dynamics model. By predicting latent states and generating trajectories, it significantly enhances the model’s ability to reason about future states.\n2.Superior Generalization Performance: Experimental results demonstrate that OSVI-WM outperforms existing methods on unseen tasks in both simulation and real-world robotic platforms, achieving over 30% improvement in success rates on challenging tasks with semantic or structural variations.\n3.Strong Cross-Embodiment Adaptability: The method performs robustly across different embodiments (e.g., human-to-robot imitation) without requiring additional training data or large-scale pretraining, lowering the barrier for real-world deployment.\n4.Enhanced Robustness via Replanning: The waypoint re-planning mechanism dynamically adjusts the execution strategy upon failure, further improving task completion rates.\n\n\nWeaknesses：\n1.Lack of End-Effector Orientation Control: The framework's limitation to predicting only 3D positions and gripper states restricts its applicability in tasks that require precise end-effector orientation control. Many real-world manipulation tasks, such as screwing, drilling, or side-grasping, necessitate accurate orientation control.\n2.Limited Handling of Dynamic Environments: Although the world model aims to capture environment dynamics, its ability to handle highly dynamic or uncertain environments remains unexplored. In scenarios with multiple moving objects, sudden changes in the environment, or significant variations in object properties, the model's predictions may become less reliable,.\n3.Redundancy in Latent States: Although latent states are designed to be “compact” representations, the features extracted by the Res-Net encoder may contain a large amount of redundant information unrelated to the task (such as texture or lighting). This redundant information can interfere with the model's learning and prediction of key task features, reducing the model's efficiency and accuracy. Therefore, further research is needed on how to remove this redundant information from the latent states to improve the model's performance and generalization ability.\n4.Limited Theoretical Analysis: The paper lacks a rigorous theoretical analysis of the framework's properties.", "was_revised": true, "final_justification": "I would like to thank the authors for their response. After consideration of the clarifications provided and the additional analyses presented in the response, I maintain my original assessment and scoring of this manuscript.", "rebuttal": "We thank the reviewer for the valuable comments. Some of the weaknesses listed, such as Limited Handling of Dynamic Environments, Redundancy in Latent States, and Limited Theoretical Analysis, are indeed active research areas within world models. OSVI-WM beats all previous results on one-shot visual imitation. Additionally, the modularity of our framework will facilitate integration of future advancements in the above three research areas. Below, we address each of the concerns (weaknesses) and questions in detail.\n\n**Concern 1: Lack of End-Effector Orientation Control.**\n\n**Response**: As noted in the limitations paragraph in Section 5 of the manuscript, OSVI-WM currently predicts only 3D positions and gripper states, aligning with existing works like [10,5,32]. While our current tasks do not demand precise orientation control, the framework is designed to be flexible. It will be extended in future work to include orientation control by incorporating end-effector orientation data into the waypoints during training. This adaptability ensures that OSVI-WM can be tailored to more complex tasks like screwing, drilling, or side-grasping.\n\n**Concern 2: Limited Handling of Dynamic Environments**\n\n**Response**: The trained world model is generally good at handling object movements directly influenced by the agent's actions, such as a button pressed by a robot (Fig. 9). Even in cases where external agents move objects important to the task, the world model can usually dynamically adapt to such changes appropriately. For other dynamic elements that do not directly affect the agent, like moving people or other robots, probabilistic modeling or optimization over environment-dependent latent variables can be employed. However, there is no widely accepted method for modeling such dynamic movements, and it is an active area of research in world models. Existing works in robotic control [32, 5, 10, 57, 61, 8] typically focus on static environments. While our current approach aligns with this trend, future work could explore integrating more sophisticated models to better handle dynamic and uncertain environments.\n\n**Concern 3: Redundancy in Latent States**\n\n**Response**: As the reviewer correctly notes, latent states are intended to provide compact representations of the environment that focus on the most relevant features. To reduce the influence of irrelevant factors such as texture or lighting, we train the world model directly in the latent embedding space rather than in the image pixel space. Specifically, the loss function (Section 3.3) is applied on the latent state embeddings instead of the decoded image pixels, where cues like texture and lighting tend to dominate. This approach is consistent with recent advances in self-supervised learning [2] and robotics [61, 8], where training in latent space has been shown to better capture task-relevant features while avoiding overemphasis on pixel-level patterns like brightness or texture.\n\nWhile quantifying redundancy in latent representations remains an active area of research, we argue that for the purposes of the current work, some degree of redundancy in a larger embedding space is not detrimental to performance. This is because the latent embeddings are subsequently passed through pooling and MLP-based waypoint predictor layers, where the MLP learns to attend to task-relevant features. The primary requirement for the latent representation is to effectively disentangle image features. As long as the latent space captures the underlying structure of the scene, the downstream MLP can leverage this disentanglement to generate accurate waypoint predictions. Most learning-based methods, including competing approaches like [5, 32, 10, 57, 60], utilize encoders such as ResNet to extract features. Although redundancy is possible in all of them, it is not widely recognized as a major factor in reducing performance.\n\n**Concern 4: Limited Theoretical Analysis**\n\n**Response**: We acknowledge the importance of theoretical analysis in understanding the properties of any framework. However, providing a rigorous theoretical analysis for large learning-based methods, such as OSVI-WM, is inherently challenging due to their complex and data-driven nature. Consequently, similar to other works in this area, we have focused on empirical validation to demonstrate the effectiveness and robustness of the various design choices in our framework (Section 4.2).\n\n**Question 1: Generalization to Unseen Positions and Objects**\n\n**Response**: As outlined in the ablation studies in Section 4.2 of the manuscript, OSVI-WM’s generalization capabilities arise from several key components: (a) world-model-guided trajectory generation with WM loss, (b) re-planning, (c) stop-gradient in WM supervision, (d) spatial pooling after the world model, and (e) multi-step trajectory prediction. Many one-shot visual imitation methods struggle with generalization because the model plans actions based on the objects in the image while underutilizing the expert demonstration. To discourage this, we employ the world model to predict the next latent states alongside the predicted waypoints. The WM loss applied to these predicted latent states is independent of the waypoints, providing a regularizing effect. Additionally, the use of asymmetric demonstration mixup enhances generalization. We will incorporate these clarifications into the updated manuscript for further clarification.\n\n**Question 2: Visualization and Analysis of Latent States and Actions**\n\n**Response**: The latent actions in our model do indeed correspond to the movements of the robotic arm. We qualitatively analyze this relationship through the following experiments:\n* Effect of Latent Actions on Latent States: We observe that when a zero vector is used as the latent action, the predicted latent state remains unchanged, indicating no movement. In contrast, when a random latent action vector is applied, the visualized latent state exhibits small, unstructured movements of the robotic arm. This qualitatively confirms that latent actions have a direct and controllable influence on latent states.\n* Effect of Latent States on Action Generation: To assess how changes in latent states drive the generation of latent actions, we perform an intervention during multi-step prediction in Eq. 1. Specifically, we manually reset a future predicted latent state $r_j$​ (where $1<j<M−1$) to be equal to the initial state $r_1$​, and feed this into the action encoder. We find that the latent actions adapt accordingly, attempting to recover the trajectory and guide the latent state back toward the originally planned sequence. This demonstrates that the action generation mechanism is sensitive to changes in the latent state and aims to maintain trajectory consistency.\n\nWhile we are unable to include these new visualizations in the rebuttal due to submission guidelines, we will incorporate them into the final version of the paper to enhance interpretability and provide clearer insights into the interaction between latent states and actions.\n\n**Question 3: The Significance and Efficiency of Replanning**\n\n**Response**: When task execution fails (detected by the simulation environment configuration), OSVI-WM triggers re-planning by taking a new observation and generating a fresh sequence of waypoints based on the robot's current state. The replanning mechanism in OSVI-WM is, therefore, not equivalent to simple repeated execution. Instead of open-loop retries, replanning is a closed-loop process that takes a new observation of the current robot state and generates fresh waypoints based on that feedback. This closed-loop replanning does incur additional time cost. For the results shown in Figure 8, the time comparison is as follows:\n* Without re-planning: Each episode runs for a maximum of 500 steps (~40 seconds on our system with one RTX A4000 GPU, 128 GB RAM, and an Intel i9 CPU).\n* With re-planning: New waypoints are generated every 500 steps, for up to 2000 steps (~160 seconds) total.\n\nDespite this overhead, replanning substantially improves success rates. Notably, in our paper, replanning is only applied for the ‘hard’ Meta-World tasks. Even without replanning, OSVI-WM achieves a success rate of ~42% on these tasks (71.5% after re-planning), which is already higher than all existing baselines (Fig. 8). Additionally, it is important to highlight that many baseline methods, such as [10], are also inherently closed-loop but still perform worse than OSVI-WM, further emphasizing the strength of our approach beyond just repeated planning. We will include a discussion on the time cost of replanning in the revised manuscript.\n\n**Question 4: Redundancy in Latent States**\n\n**Response**: As discussed in our response to your Concern 3, we explicitly address the issue of unrelated background information, such as texture or lighting, by training the world model in the latent embedding space instead of the image pixel space. This design choice ensures that the model focuses on higher-level, task-relevant representations, rather than low-level visual features. Recent works in self-supervised learning and robotics [2, 8, 61] have similarly adopted this approach to improve generalization and reduce sensitivity to irrelevant visual cues.\n\nSimilar to the response to your Concern 3, we claim that, while some redundancy in latent features may exist, the primary requirement for the latent representation is to effectively disentangle image features. As long as the latent space captures the underlying structure of the scene, the downstream MLP can leverage this disentanglement to generate accurate waypoint predictions. Most learning-based methods, including competing approaches like [5, 32, 10, 57, 60], utilize encoders such as ResNets to extract features. Although redundancy is possible in all of them, it is not widely recognized as a major factor in reducing performance."}, {"review_id": "A7HRiT1pfa", "reviewer": "Reviewer_9Vzo", "summary": "This paper proposes a one-shot visual imitation learning method for unseen robotic manipulation tasks. By taking a single expert demonstration as context information, this model generates a sequence of waypoints to complete novel tasks. To enhance the framework's ability for modeling dynamic environments, this paper introduces a world-model-guided trajectory generation proxy loss. The effectiveness of this method is validated through extensive experiments in both simulation and the real world.", "questions": "1. Why does the world model loss bring benefits to the one-shot imitation learning on novel scenarios? For example, the Human-Franka-PP task only contains 3 training manipulation tasks; the world model would overfit to these tasks, so why would it have better generalization ability for novel tasks?\n\n2. For dynamic world modeling, various methods exist to reconstruct subsequent frames. Could you provide a comparison with other techniques, like VQ-VAEs and diffusion models, to further assess how well the representation captures dynamic scenes?\n\n3. How do different types of similarity—such as semantic similarity and visual similarity—between the training dataset and validation tasks affect learning outcomes? Furthermore, the impact of various related tasks within the training dataset on overall performance remains underexplored.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\n1. This paper provides sufficient real-world experiments to prove the effectiveness.\n2. This paper is well-organized, and the proposed method is clearly explained.\n3. It is impressive to imitate human videos for a novel manipulation task.\n\n## Weaknesses\n\n1. Although the experimental results indicate that incorporating the world model loss is beneficial, it remains unclear why this leads to improvements on unseen tasks. Training a model on a small dataset is prone to overfitting, which would be expected to hinder its generalization capabilities on novel tasks. More detailed qualitative and quantitative analyses of the world model loss should be provided.\n\n2. The similarity between training scenarios and novel validation scenarios would significantly influence the performance of one-shot imitation learning systems. However, this lacks a comprehensive analysis of this relationship.", "was_revised": true, "final_justification": "I think this paper could be accepted due to the impressive one-shot learning performance.", "rebuttal": "We thank the reviewer for the valuable comments. We are happy that you found our human-franka experiments impressive. We address each of the concerns (weaknesses) and questions in the following.\n\n**Concern 1: Although the experimental results indicate that incorporating the world model loss is beneficial, it remains unclear why this leads to improvements on unseen tasks. Training a model on a small dataset is prone to overfitting, which would be expected to hinder its generalization capabilities on novel tasks. More detailed qualitative and quantitative analyses of the world model loss should be provided.**\n\n**Response**: Training a model on a small dataset can indeed lead to overfitting if not properly regularized. Most existing one-shot visual imitation methods are trained end-to-end to predict actions, often relying heavily on image semantics while underutilizing the expert demonstrations. These methods typically emphasize features crucial for action prediction alone. While regularization techniques like asymmetric demonstration image mixup, as used in AWDA [5], have shown improvements in generalization by disrupting this correlation, we enhance this further by incorporating the world model loss. This loss $\\mathcal{L}_{wm}$ on the future latent states predicted by the world model is independent of the ground truth actions (actions being waypoints in our case). This additional supervision helps the model focus on learning robust representations of environment dynamics, which improves generalization to unseen tasks.\n\n**Concern 2: The similarity between training scenarios and novel validation scenarios would significantly influence the performance of one-shot imitation learning systems. However, this lacks a comprehensive analysis of this relationship.**\n\n**Response**: Thank you for the suggestion. We agree that the similarity between training scenarios and novel validation scenarios can impact the performance of one-shot imitation learning systems. During training, the model learns two key components: inferring tasks from the expert and acquiring the skills needed to complete these tasks. Generally, the higher the similarity between the skills required in the training tasks and those in the test tasks, the better the generalization. To explore this further, we conducted two new ablation studies:\n* Ablation 1: Button Press Tasks Removed from Training. We removed all training tasks related to button pressing (Button-Press-Topdown, Button-Press-Topdown-Wall, Button-Press-Wall) and evaluated the model on the Button-Press task. The model still achieved a high success rate of 97%, nearly identical to when the button-related tasks were included. This is because button pressing is relatively kinematically less complicated and the required skills were already acquired from other training tasks.\n* Ablation 2: Window Tasks Removed from Training. We removed the Window-Close task from training and evaluated on Window-Open. In this case, performance dropped from 95% to 80%. The sliding window tasks here involve more complex physical dynamics. So the absence of related training data led to a noticeable drop in generalization.\n\nWe will include these results and a corresponding discussion in the updated manuscript to better clarify this relationship.\n\n**Question 1: Why does the world model loss bring benefits to the one-shot imitation learning on novel scenarios? For example, the Human-Franka-PP task only contains 3 training manipulation tasks; the world model would overfit to these tasks, so why would it have better generalization ability for novel tasks?**\n\n**Response**: As detailed in the response to your Concern 1, the world model loss provides an additional regularizing training signal, encouraging the model to focus on general features necessary for predicting the next states of the environment, alongside predicting the waypoints. This impact is more evident on larger datasets like Meta-World (as shown in Fig. 7 of the manuscript) compared to smaller datasets like Human-Franka-PP. Increasing the amount of data would certainly help mitigate overfitting. Nonetheless, OSVI-WM achieves high success rates on the novel tasks within Human-Franka-PP, Human-Franka-Push, and Two-Franka-PP (Table 2 of the manuscript), outperforming competing methods.\n\n**Question 2: For dynamic world modeling, various methods exist to reconstruct subsequent frames. Could you provide a comparison with other techniques, like VQ-VAEs and diffusion models, to further assess how well the representation captures dynamic scenes?**\n\n**Response**: VQ-VAE and diffusion models have proven effective for modeling dynamic scenes and have been successfully extended to world models (e.g., [A, B]). Pixel-space world models such as those based on VQ-VAE and diffusion techniques are often computationally intensive to train and deploy. Given the scope and constraints of this rebuttal, we provide a qualitative comparison with these methods based on insights from existing literature. We compare OSVI-WM’s world model against VQ-VAE and diffusion-based world models along several key dimensions:\n1. World model representation: As discussed in Section 2 of the manuscript, world models can represent the environment’s state either in pixel space or in a latent embedding space. VQ-VAE and diffusion-based world models usually operate in the pixel space, where future frames are generated or reconstructed directly as images. In contrast, OSVI-WM represents the world in a latent embedding space, optimized specifically for downstream prediction tasks such as waypoint planning. Several recent works ([2, 23, 22, 29, 61]) have shown that latent-space representations are more effective for learning meaningful abstractions, as they avoid modeling low-level pixel details that are often irrelevant for downstream tasks.\n2. Dynamic world: World models are generally designed to capture how environments evolve in response to agent actions. However, they differ in how they handle other dynamic elements (e.g., moving obstacles, background agents). Diffusion models can implicitly model such non-agent-dependent dynamics due to their probabilistic nature, though this comes with significant computational cost. OSVI-WM, in its current form, does not explicitly model dynamics unrelated to the agent. However, it can be extended to do so by introducing an additional environment-dependent learnable latent token, enabling the representation of external scene changes.\n3. Computational efficiency: OSVI-WM is more efficient, as it bypasses costly pixel-level reconstruction. VQ-VAE and especially diffusion models are, typically, more resource-intensive due to decoder use and iterative generation steps.\n4. Planning Efficiency: OSVI-WM’s latent embeddings are directly usable for downstream tasks like planning and control. In fact, some recent works (e.g., [61,C,D]) have shown that learned latent states can be effectively combined with model-based controllers without any further training. In contrast, VQ-VAE and diffusion outputs often require an additional encoder to derive planning-compatible representations, adding overhead.\n5. Interpretibility and visualization: VQ-VAE and diffusion models offer more intuitive visualization since their outputs are in pixel space. OSVI-WM requires a separate decoder to visualize latent states, making interpretation less direct.\n6. State-Action space: VQ-VAE uses a discrete latent space, while OSVI-WM and diffusion models operate in continuous spaces making them more suitable for continuous control.\n\nWe will include a brief discussion of the above in the revised manuscript.\n\n**Question 3: How do different types of similarity—such as semantic similarity and visual similarity—between the training dataset and validation tasks affect learning outcomes? Furthermore, the impact of various related tasks within the training dataset on overall performance remains underexplored.**\n\n**Response**: We conducted new ablation studies (as detailed in the response to your Concern 2) showing that skill similarity, rather than purely visual or semantic similarity, has a stronger impact on generalization. When core skills needed for the test task are present in the training set, performance remains high even if the tasks differ visually or semantically. We'll include these insights and results in the updated manuscript.\n\n**References**\n\n[A] Robine, Jan, Tobias Uelwer, and Stefan Harmeling. \"Smaller world models for reinforcement learning.\" Neural Processing Letters 55, no. 8 (2023): 11397-11427.\n\n[B] Ding, Zihan, Amy Zhang, Yuandong Tian, and Qinqing Zheng. \"Diffusion world model: Future modeling beyond step-by-step rollout for offline reinforcement learning.\" arXiv preprint arXiv:2402.03570 (2024).\n\n[C] Bar, Amir, Gaoyue Zhou, Danny Tran, Trevor Darrell, and Yann LeCun. \"Navigation world models.\" In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 15791-15801. 2025.\n\n[D] Bai, Yutong, Danny Tran, Amir Bar, Yann LeCun, Trevor Darrell, and Jitendra Malik. \"Whole-Body Conditioned Egocentric Video Prediction.\" arXiv preprint arXiv:2506.21552 (2025)."}]

Author rebuttal:

Meta-review: This submission presents OSVI-WM, a one-shot imitation framework for robot manipulation tasks. The framework uses a world model (WM) to guide the generation of a trajectory (a sequence of latent states and actions) based on a (human) video demonstration. The trajectory is then decoded into a sequence of waypoints that the robot tracks during execution. The approach claims to exhibit higher generalization from train to test task distributions.

All reviewers were initially weakly on the accept side and had questions about the details of the method, the experiment setup, train and test task distributions, and the mechanism by which generalization arises in the proposed method. In response, the authors have allayed the reviewers' concerns by providing additional explanations and, importantly, doing additional valuable ablation studies. The reviewers didn't have any major specific questions after the rebuttal but also didn't increase their "weak accept" scores.

The metareviewer thinks that from the standpoint of solving robot manipulation the proposed approach falls somewhat short: it essentially does motion planning but doesn't exhibit advanced object manipulation, currently lacking precise gripper control. However, this weakness is offset by this submission's well-executed and well-analyzed use of WMs to enable generalization. While the approach does need additional work to be competitive in manipulating objects, the robot manipulation community will benefit from the insights provided by OSVI-WM regarding WM use in this domain.
Decision: Accept (poster)