Predicting how the world responds to an action requires more than what can be seen: robots must account for environmental factors that are neither directly observable nor explicitly specified, from physical properties such as friction and mass to human-defined environmental rules. We introduce T2-Z for test-time optimization of environment latents, which augments a robotic world model with a learnable environment latent to condition world predictions. This compact latent captures underlying environmental variation without supervision of what environmental factors exist or what values they take. During training, T2-Z alternates between two stages: a world-model optimization stage that fixes the latents and learns general action-conditioned dynamics across environments; and an environment-latent optimization stage that freezes the world model and optimizes the conditioning latents, allowing latent structure to emerge purely through world prediction. At test time, the world model remains frozen and a newly initialized environment latent is optimized from one or more test-time interactions, adapting the model to previously unseen environmental configurations. Across 7 simulated settings and 4 real-world robotic tasks, T2-Z reduces object prediction error metric from 15.20 to 7.00 and raises action success from 33.73% to 72.96%. Its emergent latents capture structured representations of individual and compounded factors across diverse visual scenes. Ablations further show that T2-Z can efficiently infer environment latents from as little as a single test-time interaction trajectory.
Figures & tables
Figure 1: T2-z for Test-Time Environment Latent Optimization. (a) Given the same observation and action, T2-z directly optimizes environment latents to capture different environmental conditions ( e.g., ball mass), leading the world model to predict different outcomes. (b) Without knowing how many factors exist or what they represent, PCA of the learned environment latents reveals a two-dimensional structure that closely reflects the underlying factor distribution. (c) We evaluate T2-z on real-world robotic tasks spanning diverse unobserved environmental factors. (d) Across diverse visual scenes, the learned latents reveal a unified structure that reflects the same underlying environmental factor without being entangled with visual appearance.
Figure 2: Comparison with existing approaches. (a) Vanilla world models do not explicitly represent unobserved environmental variation, producing the same prediction under different environmental conditions. (b) System-identification approaches such as RMA learn environment representations from predefined factors and their ground-truth values, then infer them from interaction history through a learned encoder. (c) T2-z requires neither predefined environmental factors nor a learned environment encoder, instead directly optimizing an emergent environment latent from test-time interaction through prediction error.
Figure 3: T2-z framework. During training, T2-z alternates between two stages: (1) world-model optimization, which learns shared action-conditioned dynamics with environment latents fixed, and (2) environment-latent optimization, which optimizes the latents through prediction error with the world model frozen. (3) At test time, the world model remains frozen while a new environment latent is directly optimized from interaction and then reused to predict the outcomes of different actions.
Task
Method
Physical Error (Mean) ↓
Physical Error (Final) ↓
LPIPS ↓
PSNR ↑
Action Success ↑
Mass Balance
Standard World Model
1.1844
4.5571
0.0244
32.969
40.0%
LoRA TTA
1.7975
7.3105
0.0631
24.895
50.0%
DINOv2 Context
1.2636
4.9480
0.0239
33.263
30.0%
TTT-KQV
2.5075
10.1759
0.0312
32.114
30.0%
T2-z (Ours)
0.6403
2.0842
0.0240
32.814
80.0%
Friction
Standard World Model
11.0626
24.4116
0.0667
31.113
32.0%
Table 1: Simulation performance across environmental variations. For Friction, Gravity, Mass Collision, Mass × Friction, and Diverse Scenes, physical error is measured in pixels; for Light Switch, it is lamp-score error; and for Mass Balance, it is bar-tilt error in degrees. The average is computed over all seven tasks.
Figure 4: PCA of learned environment latents. We visualize learned environment latents using PCA onto x-y axes, with ◯ denoting training environments and △ denoting test environments. Ground-truth factors (denoted by color or size) are shown only for visualization and are never provided to the model. (a–b) show that structured latent representations emerge in both simulated and real-world single-factor settings. (c) shows that a unified latent structure emerges across diverse visual backgrounds, indicating that the representation does not become entangled with scene appearance. (d) shows that a two-dimensional latent structure emerges when mass and friction vary jointly, with x-axis primarily corresponding to friction and y-axis to mass, despite the model never knowing how many underlying factors are present or what they represent.
Figure 5: Ablation studies. (a) Test-Time Optimization Learning-rate strategy. Constant Maximum/Minimum remove coarse-to-fine scheduling, while Ours ×10/÷10 scale our schedule uniformly. (b) Ablating the test-time optimization budget shows that performance peaks at 40 steps. (c,d) Effect of the number of test-time interaction trajectories. Bold markers denote our chosen trajectory count, after which performance largely saturates.
Method
Physical Error (Mean) ↓
Physical Error (Final) ↓
LPIPS ↓
PSNR ↑
Action Success ↑
Joint Model-Latent Training
7.03
15.66
0.0657
31.203
68%
Direct Token (3072-D)
94.50
80.49
0.7626
15.218
24%
32-D Random [−0.05,0.05]
4.81
9.73
0.0495
32.172
60%
Alternating Latent/Model (100/100)
5.41
11.50
0.0484
32.208
56%
Alternating Latent/Model (400/400)
4.31
8.80
0.0507
32.430
68%
T2-z
3.85
7.59
0.0495
31.976
72%
Table 3: Additional Ablation study. Joint Model–Latent Training optimizes the world model and latent together without separate stages; Direct Token directly optimizes the full-dimensional conditioning token; 32-D Random [−0.05,0.05] changes the latent initialization range; and Alternating Latent/Model change the number of updates before switching between the two optimization stages, compared with our default 200/200.
Figure 6: Qualitative results in simulation and the real world. Within each task, we apply the same actions across different environments, which produce different ground-truth outcomes due to underlying environmental variation. T2-z closely matches these environment-dependent outcomes.
Appendix figures & tables25 assets
Supplementary material from the paper’s appendix.
Appendix
Component
Initialization or structure
Training update
Video DiT
BWM-initialized Wan2.2 backbone
Model blocks
Action encoder
BWM-initialized 14-D action pathway
Model blocks
Video VAE
Wan2.2 video encoder/decoder
Frozen
Environment table
Learnable 32-D latents indexed by environment
Latent blocks
Latent projectors
Two 32→128→3072 MLPs
Model blocks
Appendix
Table 4: Model components and optimization roles. Test-time adaptation updates only the environment latent.
Task
Physical envs. × episodes/env.
Frames
Stride
Batch
Budget
Friction
35×10
41
1
8×4
2×24
Gravity
20×10
61
1
8×4
2×24
Mass Collision
20×9
61
1
8×8
2×24
Light Switch
12×10
33
1
8×8
2×24
Mass Balance
20×15
41
2
8×8
2×24
Mass × Friction
84×9
61
1
16×6
4×48
Appendix
Table 5: Simulation training scale by physical environment. Diverse Backgrounds uses 30 scene backgrounds. Budget is H200 GPUs × hours per method.
Task
Physical envs. / episodes
Frames
Stride
Native RGB
Input RGB
Batch
Budget
Ball Friction
8 / 240
61
1
560×248
320×160
6×8
2×24
Door Close
10 / 438
49
1
640×480
256×192
6×8
2×24
Stick Balance
9 / 421
33
3
640×480
256×192
6×8
2×24
Soft Pull
9 / 146
33
3
512×256
320×160
10×6
2×24
Appendix
Table 6: Real-world training scale by physical environment and native/model-input RGB resolution. Budget is H200 GPUs × hours per method.
Table 10: Real-world object-centric measurements from the main camera. ADE/FDE are reported in native-image pixels.
Figure 8: Physical designs of the four real-world tasks. The Ball task changes the number of steel beads inside the ball (not visible in the video frame), and the Door task uses a spring-loaded latch.
Physical error (mean) ↓
Action success (%) ↑
Task
K=1
K=2
K=4
K=1
K=2
K=4
Mass Collision
6.452
6.534
5.946
61.1
66.7
72.2
Light Switch
8.26
3.43
3.41
50.0
87.5
87.5
Door Close
11.337
7.954
6.889
22.2
55.6
55.6
Soft Pull
18.710
18.363
18.060
81.8
81.8
81.8
Appendix
Table 11: Numerical results for Figure 5(c,d) at K=1,2,4 test-time interaction trajectories. Bold marks the selected K . Physical error is in pixels except for Light Switch (lamp-score error in percentage points).
Latent dim.
ADE ↓
FDE ↓
PSNR ↑
SSIM ↑
LPIPS ↓
Success ↑
8
8.92
19.26
31.293
0.9474
0.0646
44%
16
7.63
15.90
29.557
0.9470
0.0614
64%
32 (ours)
3.85
7.59
31.976
0.9513
0.0495
72%
64
6.86
13.98
31.692
0.9504
0.0551
64%
128
7.08
14.51
31.372
0.9483
0.0581
64%
Appendix
Table 12: Environment-latent dimensionality on Friction under the same K=1 evaluation protocol (ten held-out physical environments; 90 query videos). Object ADE/FDE are main-view centroid errors in pixels; video metrics use the concatenated main and wrist views. Action success is evaluated on 25 reachable target decisions. Best values are bold.
Figure 9: Latent-space PCA across simulation (a–d) and real-world tasks (e–h). Circles denote training-time latents and triangles denote test-time-adapted latents. Color indicates the underlying physical setting; point size additionally indicates friction in Mass × Friction. These factor annotations are used only for visualization and are not provided to the model. The Soft Pull test-time latents shown here use four test-time interaction trajectories.
Figure 10: Training resource profile. Time is the median compute-only optimizer duration over ten measured updates after two warmup updates, excluding dataset reads; memory is peak allocated GPU memory. Each update uses six loss-bearing video chunks.
Figure 11: Compute-only inference profile on one H200 GPU, excluding model and sample loading and MP4 writing. Top: processing of one test-time interaction trajectory plus generation of the same five query videos. Bottom: interaction processing time for K=1,2,4 with the update counts held fixed within each method. The bottom panel does not compare query-generation time at K>1 .
Figure 12: Friction: selected test actions.
Figure 13: Gravity: selected test actions.
Figure 14: Mass Collision: selected test actions.
Figure 15: Light Switch: selected test actions.
Figure 16: Mass Balance: selected test actions.
Figure 17: Mass × Friction: selected test actions.
Figure 18: Diverse Backgrounds: selected test actions.
Latent world models have shown a remarkable ability to predict future states and to plan in the real world. In practice, however, we lack a principled way to estimate how their capabilities scale with model size, data, and compute, an open problem that slows progress in the field. In this work we present RoboJEPA, a world model based on the Joint Embedding Predictive Architecture (JEPA) and trained on a large-scale dataset spanning 12 robotic embodiments. We show that RoboJEPA's imagination error, the error of its latent rollouts, follows a second-order power law in compute, allowing us to predict model quality well beyond the scale at which the law is fit. We further show that downstream robotic planning performance improves predictably with compute, and that imagination error is strongly correlated with it, making it a reliable proxy for real-robot evaluation. Finally, we demonstrate that latent world models can be deployed zero-shot as robotic agents, planning toward a single goal image to solve tasks requiring long-horizon planning on real hardware. We release all model checkpoints together with our training and robot deployment code. To our knowledge, this is the first work to establish scaling laws for multi-embodiment robotic world models trained on real robot data, and RoboJEPA, at 8B parameters, is the largest JEPA predictor model trained to date.
Artem Zholus, Nicolas Beltran-Velez, Jianhao Yuan +9
FAIR at Meta · Chandar Research Lab · Mila - Quebec AI Institute +2
This paper shows that JEPA-style latent world models (LWMs) can learn across diverse domains and data sources by adopting the language-conditioned interface of video generation models. We establish this approach as an effective and efficient basis for planning with imagination and downstream policy fine-tuning. To this end, we introduce LaLa-World, a Language-conditioned Latent World model that learns from simulation and robot trajectories together with broader image and video data. Two designs make this work. First, we describe both numerical actions and visual changes in language, enabling joint training of a single future predictor across heterogeneous datasets. Second, we build the predictor on a language model, allowing it to interpret language conditions while predicting future latent states. The resulting predictor achieves higher average MPC planning performance than existing LWMs trained separately for each domain or task. Adding image and video data further improves average planning performance and robustness to action phrasing. In a matched comparison with an RAEv2-based video generation model baseline, LaLa-World reaches comparable planning performance with 6× less training cost and 13.9× faster planning. Beyond planning, the pretrained LaLa-World serves as an initialization for LaLa-Policy, which jointly predicts actions and future visual representations. LaLa-Policy performs competitively with Cosmos Policy on LIBERO and RoboCasa, with up to 5.5× faster inference.
Junha Song, Geonmo Gu, Jeongeun Park +5
NAVER AI Lab · Georgia Institute of Technology · Brown University +1
The diversity of robot embodiments and action spaces makes it challenging to build robot world models that generalize across different embodiments. We introduce the Latent Action-Conditioned Robot World Model (LAC-WM), which operates within a learned unified latent action space shared across diverse embodiments. This unified action space improves the world model's performance when adapted to previously unseen robot embodiments. We compare LAC-WM with an Explicit Action-Conditioned World Model (EAC-WM), which conditions on explicit motion labels. Our results show that explicit action conditioning leads to disjoint action representations across embodiments, limiting downstream performance when adapting to new robots. We evaluate both models on dexterous manipulation tasks and a modified LIBERO benchmark. LAC-WM improves downstream performance over EAC-WM by up to 46.7% on dexterous manipulation and 11.7% on LIBERO. Crucially, the unified latent action space allows LAC-WM's downstream performance to scale positively with the number of embodiments used during pretraining. In contrast, the disjoint action space in EAC-WM leads to decreased performance as the number of pretraining embodiments increases. These results highlight the importance of a unified action space for efficient cross-embodiment learning, addressing a key challenge in robotics. Project website: https://lacwm.github.io/
Huang Huang, Sriram Yenamandra, Arjun Majumdar +9
Department of Computer Science, Stanford University, US · Work partially done while at Meta FAIR Robotics.