[Paper Notes] T2-z: Test-Time Optimization of Environment Latents for Robotic World Models
Published:
This post supports English / 中文 switching via the site language toggle in the top navigation.
TL;DR
Two visually similar objects can move different distances under the same push. A video world model needs evidence about the environment’s dynamics to resolve that ambiguity. T2-z stores this evidence in a 32-dimensional environment latent. At deployment, it freezes the world model and optimizes the latent against one or two observed interaction trajectories, then reuses it to predict other actions. Training alternates between learning shared dynamics and fitting the environment latents, using the same flow-matching loss for both.
The mechanism is a form of learned system identification through prediction error. The most useful detail is how training prepares a small, optimizable conditioning space. Simply optimizing a full-width conditioning token performs poorly in the paper’s ablation.
Paper and sources
T2-z: Test-Time Optimization of Environment Latents for Robotic World Models, by Chenyuan Zhou, Yining Hong, Huang Huang, Yunfan Jiang, Wenlong Huang, David Held, Vitor Campagnolo Guizilini, Yue Wang, Li Fei-Fei, Jiajun Wu, Leonidas Guibas, and Yejin Choi. These notes follow the October 15, 2026 preprint, arXiv:2610.22364v1, including its PDF appendices. The authors also provide a project page, code repository, and dataset. All experimental results below are author-reported; the implementation implications are my reading of the method.
1. What an environment latent represents
The prediction model is
\[p_\theta(\mathbf y\mid\mathbf x,\mathbf a,\mathbf z_e),\]where $\mathbf x$ is the initial observation, $\mathbf a$ is an action sequence, and $\mathbf y$ is the future video. The shared parameters $\theta$ describe the predictive model; $\mathbf z_e\in\mathbb R^{32}$ carries information specific to environment $e$. Different episodes from the same environment share a latent while their observations and actions vary.
This last condition matters. Training does not need numerical labels for friction, mass, or spring stiffness, but it does maintain environment-indexed groups of interactions. The available grouping tells the learner which episodes should share an explanation. “No physical-parameter supervision” therefore still assumes that training data can be organized by environment (Appendix A).
The representation can also capture a discrete interaction rule. In Light Switch, the hidden variation determines which button interactions activate the lamp. The prediction objective can distinguish these rules without assigning each latent coordinate a predefined physical meaning.
2. Learn dynamics slowly, fit latents quickly
Both sets of variables use the flow-matching objective
\[\mathcal L_{\mathrm{FM}}(\theta,\mathbf z_e) =\mathbb E\left[ \left\|\mathbf v_\theta(\mathbf y_t,t\mid\mathbf x,\mathbf a,\mathbf z_e) -\mathbf u_t\right\|_2^2\right].\]Here $\mathbf y_t$ is the noisy/interpolated future representation and $\mathbf u_t$ is its target flow velocity. The initial conditioning frame is excluded from the prediction loss. Thus “prediction error” refers to the model’s training objective; it does not require a differentiable object tracker or a supplied friction estimate.
With latents fixed, model updates minimize this loss across environments. With the model fixed, each latent is updated using the interactions assigned to its environment. In compact gradient notation,
\[\begin{aligned} \theta &\leftarrow \theta-\eta_\theta\nabla_\theta \mathbb E_{e,\mathcal D_e}\mathcal L_{\mathrm{FM}}(\theta,\mathbf z_e), &&\text{latents fixed},\\ \mathbf z_e &\leftarrow\mathbf z_e-\eta_z\nabla_{\mathbf z_e} \mathbb E_{\mathcal D_e}\mathcal L_{\mathrm{FM}}(\theta,\mathbf z_e), &&\text{model fixed}. \end{aligned}\]The implementation begins with 300 model updates, then alternates 200 latent updates and 200 model updates. The learning rates are $10^{-5}$ for the shared model and $0.03$ for latents. Training latents start from $\mathcal U(0,1)$ and are clipped to $[0,1]$ after updates. Model blocks update the DiT, action encoder, and latent projectors; the video VAE stays frozen (Section 3 and Appendix A.2).
The distinction from joint optimization is measurable. On the Friction ablation, joint model–latent training gives mean physical error 7.03 and action success 68%; the alternating method gives 3.85 and 72%. This supports the training schedule in that setting, without establishing one universally optimal block length.
3. How 32 numbers condition a video model
The backbone is BWM built on Wan2.2-TI2V-5B, with DiT width 3072. It takes an initial RGB frame and aligned 14-dimensional robot actions. Text conditioning is disabled. Two separate MLPs project the environment latent from 32 → 128 → 3072, with LayerNorm before the first linear layer and GELU between the two linear layers.
One projector creates an environment token appended to the action conditioning sequence. Every video patch can attend to that token through cross-attention. The other creates a modulation vector added to the timestep and action embeddings before the existing time projection. That path changes adaptive-normalization shifts, scales, and residual gates, and also conditions the output head (Appendix A.1).
These learned projectors explain why the adaptation space deserves attention. At test time, gradients move 32 coordinates through a conditioning map learned across training environments. The Direct Token ablation instead optimizes a 3072-dimensional token and reports 94.50 mean physical error and 24% action success on Friction. The low-dimensional parameterization and its training procedure are substantive parts of the method.
4. Test-time interaction becomes reusable context
For a new environment, collect actual outcomes of randomly sampled actions:
\[\mathcal D_{\mathrm{test}}=\{(\mathbf x_i,\mathbf a_i,\mathbf y_i)\}_{i=1}^{N}.\]Freeze all shared weights, initialize a new latent, and optimize
\[\mathbf z_{\mathrm{test}}^{(s+1)}= \mathbf z_{\mathrm{test}}^{(s)}- \eta_z^{(s)}\nabla_{\mathbf z} \frac1N\sum_{i=1}^{N} \mathcal L_{\mathrm{FM}}(\theta^*,\mathbf z;\mathbf x_i,\mathbf a_i,\mathbf y_i) \bigg|_{\mathbf z=\mathbf z_{\mathrm{test}}^{(s)}}.\]The main evaluations use 40 updates, with learning rates 3.0, 1.5, 0.5, and 0.15, ten updates each. Query outcomes are excluded from adaptation. The resulting latent is reused for different candidate actions, whose predicted videos guide action selection (Algorithm 1 and Appendix A.2).
flowchart TD
A[New environment] --> B[Execute probe actions]
B --> C[Observed interaction trajectories]
C --> D[Flow-matching loss through frozen world model]
D --> E[Update only the 32-dimensional environment latent]
E --> D
E --> F[Reuse adapted latent]
G[Candidate action sequences] --> H[Predict videos with frozen world model]
F --> H
H --> I[Select action using predicted task outcome]
“One trajectory” means a sequence of interaction observations. Its information content depends on the action. Appendix E.1 makes this visible: increasing the count from one to two raises Light Switch success from 50.0% to 87.5% and Door Close success from 22.2% to 55.6%. Soft Pull stays at 81.8% for one, two, or four trajectories. More observations help when they resolve an ambiguity the first probe left open.
5. Results, with metric boundaries
Evaluation spans seven simulated settings and four real tasks. Simulation covers friction, gravity, collision mass, switch rules, mass balance, combined mass and friction, and friction across 30 visual backgrounds. Real tasks are Door Close, Ball Friction, Stick Balance, and Soft Pull. All adaptive baselines receive the same test-time interaction trajectories and use the same backbone (Tables 1–2).
| Evaluation | Standard world model | T2-z |
|---|---|---|
| Simulation: mean physical error, mixed task units | 12.8575 | 4.7971 |
| Simulation: mean action success | 37.9% | 74.2% |
| Real world: mean tracked-state error, pixels | 19.30 | 10.85 |
| Real world: mean action success | 26.47% | 70.80% |
The simulation average mixes object-position errors in pixels, lamp-score error, and tilt-angle error in degrees. It is a benchmark summary with heterogeneous units. Per-task comparisons are easier to interpret physically. Real-world tracking errors use native-image pixels, whose resolutions also vary by task.
Action success and visual accuracy can disagree. On Gravity, DINOv2 Context has lower mean physical error than T2-z (5.9300 vs 7.3408), while T2-z has higher action success (82.4% vs 76.5%). On Soft Pull, T2-z improves action success while its LPIPS is worse than the standard model. A reproduction should evaluate the action decisions directly, alongside video and state errors.
The data scale is also specific: Appendix B lists 35 × 10 episodes for simulated Friction and 84 × 9 for Mass × Friction. The four real tasks use 240, 438, 421, and 146 episodes for Ball Friction, Door Close, Stick Balance, and Soft Pull. Those adaptation results sit on top of a pretrained video backbone and task training; the test-time trajectory counts alone do not describe the total data requirement.
6. What I would carry into a new system
I would start with T2-z when a competent action-conditioned world model already exists and repeated episodes differ through hidden environmental properties. It gives a concrete way to fit a small context vector from outcomes, and to reuse that context across candidate actions. The Friction dimension sweep favors 32 dimensions over 8, 16, 64, and 128 under the tested protocol; that is a useful starting point, with retuning needed for a different task.
There are two practical boundaries. First, latent optimization still requires backpropagation through the frozen video model. Updating only 32 numbers does not make the computation proportional to 32; Appendix E.4 profiles 40 adaptation updates plus video generation on an H200. Second, latent-space PCA shows organization correlated with hidden factors, but each task has its own PCA basis. These visualizations support structured representations within the experiments; they do not establish calibrated recovery of physical parameters or a universal coordinate system.
My next experiment would compare random probes with actions selected to distinguish plausible environment explanations. The paper already shows that a second interaction can matter much more in one task than another. Choosing which interaction to collect could be as consequential as increasing the number of latent updates.