T2-z: Test-Time Optimization of Environment Latents for Robotic World Models

Chenyuan Zhou; Yining Hong; Huang Huang; Yunfan Jiang; Wenlong Huang; David Held; Vitor Campagnolo Guizilini; Yue Wang

cs.ROOct 15, 2026

T2-z: Test-Time Optimization of Environment Latents for Robotic World Models

Authors: Chenyuan Zhou, Yining Hong, Huang Huang, Yunfan Jiang, Wenlong Huang, David Held, Vitor Campagnolo Guizilini, Yue Wang,

Organizations: Stanford University · CMU · Toyota Research Institute · University of Southern California

Abstract

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) $\downarrow$	Physical Error (Final) $\downarrow$	LPIPS $\downarrow$	PSNR $\uparrow$	Action Success $\uparrow$
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

\times

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 $\bigcirc$ denoting training environments and $\triangle$ 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.

Task	Method	Physical Error (Mean) $\downarrow$	Physical Error (Final) $\downarrow$	LPIPS $\downarrow$	PSNR $\uparrow$	Action Success $\uparrow$
Door Close	Standard World Model	21.82	41.22	0.05700	29.96	22.22%
	LoRA TTA	17.35	32.14	0.04946	30.24	22.22%
	DINOv2 Context	21.12	38.90	0.05490	30.28	22.22%
	TTT-KQV	19.87	35.93	0.05164	30.95	22.22%
	T2-z (Ours)	7.95	13.47	0.03208	32.15	55.56%
Ball Friction	Standard World Model	28.36	64.42	0.03153	36.04	0.00%

Table 2: Real-world robotic task performance. T2-z outperforms baselines across real-world tasks.

Figure 5: Ablation studies. (a) Test-Time Optimization Learning-rate strategy. Constant Maximum/Minimum remove coarse-to-fine scheduling, while Ours

\times 10/\div 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) $\downarrow$	Physical Error (Final) $\downarrow$	LPIPS $\downarrow$	PSNR $\uparrow$	Action Success $\uparrow$
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\to 128\to 3072$ MLPs	Model blocks

Appendix

Table 4: Model components and optimization roles. Test-time adaptation updates only the environment latent.

Task	Physical envs. $\times$ episodes/env.	Frames	Stride	Batch	Budget
Friction	$35\times 10$	41	1	$8\times 4$	$2\times 24$
Gravity	$20\times 10$	61	1	$8\times 4$	$2\times 24$
Mass Collision	$20\times 9$	61	1	$8\times 8$	$2\times 24$
Light Switch	$12\times 10$	33	1	$8\times 8$	$2\times 24$
Mass Balance	$20\times 15$	41	2	$8\times 8$	$2\times 24$
Mass $\times$ Friction	$84\times 9$	61	1	$16\times 6$	$4\times 48$

Appendix

Table 5: Simulation training scale by physical environment. Diverse Backgrounds uses 30 scene backgrounds. Budget is H200 GPUs

\times

hours per method.

Task	Physical envs. / episodes	Frames	Stride	Native RGB	Input RGB	Batch	Budget
Ball Friction	8 / 240	61	1	$560\times 248$	$320\times 160$	$6\times 8$	$2\times 24$
Door Close	10 / 438	49	1	$640\times 480$	$256\times 192$	$6\times 8$	$2\times 24$
Stick Balance	9 / 421	33	3	$640\times 480$	$256\times 192$	$6\times 8$	$2\times 24$
Soft Pull	9 / 146	33	3	$512\times 256$	$320\times 160$	$10\times 6$	$2\times 24$

Appendix

Table 6: Real-world training scale by physical environment and native/model-input RGB resolution. Budget is H200 GPUs

\times

hours per method.

Task	$K$	Query videos
Friction	1	90
Gravity	1	90
Mass Collision	1	80
Light Switch	2	60
Mass Balance	2	130
Mass $\times$ Friction	1	80

Appendix

Table 7: Simulation evaluation counts: test-time interaction trajectories and query videos.

Task	Tracked quantity	Object-centric error
Friction	Pushed-block centroid	ADE/FDE
Gravity	Blue projectile centroid	ADE/FDE
Mass Collision	Red struck-object centroid	ADE/FDE
Light Switch	Yellow-pixel fraction in lamp region	Mean/final absolute error (pp)
Mass Balance	Balance-bar tilt angle	Mean/final angle error ( ∘ )
Mass $\times$ Friction	Red struck-object centroid	ADE/FDE

Appendix

Table 8: Simulation object-centric quantities extracted from the predicted and ground-truth RGB videos. ADE/FDE are in image pixels.

Panel of Figure 7 — Figure 7: Representative ground-truth frames from the simulated tasks. Diverse Backgrounds uses the Friction push task with varied scene backgrounds.

Task	$K$	Query videos
Ball Friction	1	80
Door Close	2	100
Stick Balance	2	45
Soft Pull	1	24

Appendix

Table 9: Real-world evaluation counts: test-time interaction trajectories and query videos.

Task	Tracked quantity	Object-centric error
Ball Friction	Visible ball center	ADE/FDE
Door Close	Outer door-edge position	ADE/FDE
Stick Balance	Two end markers, averaged	ADE/FDE
Soft Pull	Dragged object’s visible center	ADE/FDE

Appendix

Table 10: Real-world object-centric measurements from the main camera. ADE/FDE are reported in native-image pixels.

Panel of Figure 8 — 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) $\downarrow$			Action success (%) $\uparrow$
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 $\downarrow$	FDE $\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	Success $\uparrow$
$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 $\times$ 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

Panel of Figure 12 — Figure 12: Friction: selected test actions.

Panel of Figure 13 — Figure 13: Gravity: selected test actions.

Panel of Figure 14 — Figure 14: Mass Collision: selected test actions.

Panel of Figure 15 — Figure 15: Light Switch: selected test actions.

Panel of Figure 16 — Figure 16: Mass Balance: selected test actions.

Panel of Figure 17 — Figure 17: Mass $\times$ Friction: selected test actions.

Panel of Figure 18 — Figure 18: Diverse Backgrounds: selected test actions.

Panel of Figure 19 — Figure 19: Ball Friction: selected test actions.

Panel of Figure 20 — Figure 20: Door Close: selected test actions.

Panel of Figure 21 — Figure 21: Stick Balance: selected test actions.

Figure 22: Soft Pull: selected test actions.

World Models for Robotics

World models for robotics predict how scenes and robot states change under actions to support planning, control, and learning.

See all 380 papers

RoboJEPA: Scaling Robotic Latent World Models
Oct 7, 2026Artem Zholus, Nicolas Beltran-Velez, Jianhao Yuan +9World Models Latent World Models
FAIR at Meta · Chandar Research Lab · Mila - Quebec AI Institute +2
LaLa-World: Scaling Latent World Models with Natural Language Actions
Oct 20, 2026Junha Song, Geonmo Gu, Jeongeun Park +5World Model Learning Latent World Models
NAVER AI Lab · Georgia Institute of Technology · Brown University +1
Cross-Embodiment Robot Foundation World Models with Latent Actions
Oct 7, 2026Huang Huang, Sriram Yenamandra, Arjun Majumdar +9World Models for Robotics Action-Conditioned World Models
Department of Computer Science, Stanford University, US · Work partially done while at Meta FAIR Robotics.

Oct 7, 2026J/K move · Enter open · S save

RoboJEPA: Scaling Robotic Latent World Models

Artem Zholus, Nicolas Beltran-Velez, Jianhao Yuan, Sarath Chandar+8

FAIR at Meta · Chandar Research Lab · Mila - Quebec AI Institute+2

World Models Latent World Models World Model-Based Planning World Model Evaluation Joint-Embedding Predictive Architecture Cross-Embodiment Robot Learning

Open paper AskPDF arXiv

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.

Abstract

World Models for Robotics

Explore similar work

RoboJEPA: Scaling Robotic Latent World Models

LaLa-World: Scaling Latent World Models with Natural Language Actions

Cross-Embodiment Robot Foundation World Models with Latent Actions