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, Li Fei-Fei, Jiajun Wu, Leonidas Guibas, Yejin Choi · N/A · 2026

Summary

Abstract Summary

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.

Key Points

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.

Method

We introduce T2-z ( fig. 2 & algorithm 1 ), a framework for adapting robotic world models to hidden environmental variation through direct optimization of an environment latent. Our key idea is to separate the general visual predictive capability of a world model from unobserved environment-specific information.

During training, we learn these two components through alternating optimization under a shared flow-matching objective.

Method — T2-z: Test-Time Optimization of Environment Latents for Robotic World Models — 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.

Experiments

We evaluate T2-z across 7 simulated settings and 4 real-world robotic tasks. In simulation, Friction requires predicting how far an object moves after a push under different surface friction; Gravity varies gravitational strength and requires predicting the object motion; Mass Collision varies colliding object mass and predicts the resulting collision motion.

Sources and Demonstrations

Figures are reproduced from Zhou et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.