Problem
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, such as friction, mass, and human-defined rules. Most robotic world models assume full observability or rely on pre‑specified system parameters, which limits their ability to generalize to new environments or cope with unmodeled dynamics.
Method
We introduce T2‑Z, a test‑time optimization method that augments a robotic world model with a learnable environment latent. During training, T2‑Z alternates between two stages: (1) fixing the latent and learning general action‑conditioned dynamics across environments, and (2) freezing the world model and optimizing the latent to capture environmental variation purely from prediction errors. At test time, the world model remains frozen and a newly initialized environment latent is optimized from one or more interaction trajectories, adapting the model to previously unseen configurations.
Findings
Across seven simulated settings and four real‑world robotic tasks, T2‑Z reduces the object prediction error metric from 15.20 to 7.00 and raises action success from 33.73% to 72.96%. Ablations show that effective latents can be inferred from as little as a single test‑time interaction trajectory. The emergent latents capture structured representations of individual and compounded factors (e.g., friction and mass) across diverse visual scenes.
Limitations
The method is evaluated on a limited set of simulated and real‑world tasks; performance may vary with more complex dynamics or higher‑dimensional latent spaces. The approach assumes the world model architecture is sufficiently expressive to benefit from latent conditioning, and does not address potential optimization instability during test‑time latent search.
Practical Significance
T2‑Z provides a lightweight, online adaptation mechanism that enables robots to quickly adjust their internal dynamics model to new environments without requiring explicit system identification or retraining. This improves robustness and sample efficiency for model‑based reinforcement learning and control in uncertain, changing real‑world settings.