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

1 INTRODUCTION

• T2-Z alternates between learning shared action-conditioned dynamics with environment latents fixed and optimizing the latents through prediction error with the world model frozen.The compact latent has no supervision specifying which environmental factors to represent.Source: §1 · p. 2
• At test time, T2-Z freezes the world model and optimizes a newly initialized environment latent from one or more interactions in a previously unseen configuration.The optimized latent conditions predictions for comparing candidate action outcomes before execution.Source: §1 · p. 2
• Physical prediction error falls from 12.86 to 4.80 in simulation, while action success rises from 37.9% to 74.2% in simulation and 26.5% to 70.8% in the real world.T2-Z is evaluated across 7 simulated settings and 4 real-world robotic tasks.Source: §1 · p. 2
• Learned latents capture structured environmental factors across diverse visual scenes, and an ablation finds useful representations from a single test-time interaction trajectory.In multi-background setups, the representation captures environmental dynamics without becoming entangled with scene appearance.Source: §1 · pp. 2–3

2 RELATED WORKS

• Vanilla world models predict the same outcome across different environmental conditions, while T2-Z optimizes an environment latent from test-time interaction without predefined factors or a learned environment encoder.System-identification approaches such as RMA use predefined factor values and infer representations through a learned encoder.Source: Figure 2 · p. 3
• Robotic world models and their test-time adaptations predict futures from visual observations but do not account for environmental variation that is not directly observable.T2-Z separates that variation from the world model’s observable scene representation through environment-latent conditioning.Source: §2 · p. 3
• RMA-style methods prespecify environment representations using privileged parameters, whereas T2-Z learns representations implicitly through world-prediction error without ground-truth physical parameters or a learned history encoder.The related-work discussion says prespecified representations and history encoders limit flexibility under unknown, compounded, or out-of-distribution variation.Source: §2 · p. 3
• Test-time adaptation also includes policy, prompt, memory, retrieval, video-replanning, and predictive-model approaches.These approaches adapt different components rather than necessarily optimizing a dedicated environment representation.Source: §2 · p. 3

3.1 PROBLEM SETUP

• T2-Z augments an action-conditioned world model with an environment latent, and future-observation prediction is conditioned on the observation, action, and latent.The latent receives no supervision about which environmental factors it should represent.Source: §3.1 · p. 4

3.1.1 TRAINING

• In Stage I, the model parameters are optimized across trajectories from different environments and actions while latents remain fixed, encouraging shared predictive dynamics.Training alternates stages so improved latents condition later world-model optimization and the improving model supports later latent optimization.Source: §3.1.1 · p. 5
• In Stage II, the world model is frozen and only each environment’s latent is updated to match observed interaction outcomes, separating environment-specific variation from the fixed model.Systematically different outcomes drive latents toward different regions of latent space.Source: §3.1.1 · p. 5

3.1.2 TEST-TIME OPTIMIZATION

• At test time, T2-Z freezes the trained world model and optimizes a new latent against one or more observed interaction trajectories, then reuses it to predict unobserved actions.Direct optimization can incorporate varying numbers of test-time trajectories, and the cited ablation reports a single trajectory is sufficient in most settings.Source: §3.1.2 · pp. 5–6

4.1 EXPERIMENTAL SETUP

• The evaluation spans seven simulated settings and four real-world tasks, with physical prediction, video-generation quality, and downstream action success measured against adaptive and non-adaptive baselines.The baselines share the world-model backbone, and adaptive baselines receive the same test-time interactions as T2-Z.Source: §4.1 · pp. 6–7

4.2 SIMULATION EXPERIMENTS

• 36.9% to 76.6% in five-task average action success and 0.0469 to 0.0325 in average LPIPS; mean physical error also falls on Friction, Mass Collision, and Mass Balance.The reported physical-error changes are 11.06 to 3.85 pixels, 21.93 to 6.45 pixels, and 1.18° to 0.64°, respectively.Source: §4.2 · p. 7, Table 1 · p. 7
• For single-factor friction, learned latents form a one-dimensional structure corresponding closely to friction values, although ground-truth factor values are used only to annotate the visualization.PCA projects the latents into two dimensions for visualization; the factor values are not provided to the model.Source: §4.2 · p. 8, Figure 4 · p. 8
• 35.0% to 69.7% action success under jointly varying mass and friction, while the learned latent forms a two-dimensional structure corresponding to those factors.The model is not told the number, identities, or values of the underlying factors.Source: Table 1 · p. 7, §4.2 · p. 8, Figure 4 · p. 8
• Across 30 diverse visual backgrounds, latents organize consistently by the same underlying friction variation rather than separating by scene appearance.T2-Z outperforms the baselines across every reported metric in this setting.Source: §4.1 · p. 6, Table 1 · p. 7, §4.2 · p. 8

4.3 REAL-WORLD ROBOTIC EXPERIMENTS

• 19.30/38.24 to 10.85/16.66 pixels in average mean/final physical error, 0.05194 to 0.04341 LPIPS, and 26.47% to 70.80% action success on four real-world tasks.On Stick Balance, success rises from 20.00% to 83.33%.Source: §4.3 · p. 8, Table 2 · p. 9
• The real-world experiments report larger improvements over the Standard World Model than simulation, amid unknown environmental factors and additional uncontrolled interaction variation.Source: §4.3 · pp. 8–9

4.4 ABLATIVE & QUALITATIVE STUDIES

• Ablations favor the coarse-to-fine learning-rate schedule; performance peaks at 40 optimization steps and largely saturates beyond the selected interaction-trajectory count.T2-Z uses one trajectory on seven tasks and two on four tasks; qualitative comparisons show it matches environment-dependent outcomes where baselines often produce incorrect configurations.Source: §4.4 · pp. 9–10, Figure 5 · p. 9, Figure 6 · p. 10

5 CONCLUSION

• Across 7 simulation and 4 real-world tasks, T2-Z outperforms baselines in video prediction and action selection.The passage does not provide numerical comparisons.Source: §5 · p. 10
• T2-Z alternates world-model optimization for action-conditioned dynamics with environment-latent optimization for hidden factors.The two optimization stages address dynamics and environmental variation, respectively.Source: §5 · p. 10
• For a new environment, T2-Z freezes the world model and optimizes only the environment latent.Source: §5 · p. 10

• Each episode supplies an initial RGB frame and aligned 14-D actions; the model predicts subsequent video frames, with the latent augmenting both action-conditioning paths.The initial frame enters the video-latent stream, while text conditioning is disabled.Source: §A.1 · p. 16
• Episodes in the same environment share a freely optimized 32-D latent, without an observation-to-latent encoder or supervision from physical parameters.Two separate projectors map the latent to model-width vectors.Source: §A.1 · p. 16
• Each projector uses two linear layers with an internal 128-D width and produces a D-dimensional output.The token and modulation projectors do not share weights and have no biases.Source: §A.1 · p. 16
• The environment latent conditions both cross-attention and per-patch modulation while preserving BWM’s original action pathways.The environment token joins action tokens in cross-attention; its modulation vector is combined with action and timestep vectors.Source: §A.1 · p. 16

A.2 TRAINING AND TEST-TIME ADAPTATION

• Training alternates 200-update latent blocks, which freeze the shared model, and 200-update model blocks, which freeze the latent table.An initial 300-update model block precedes this cycle.Source: §A.2 · p. 16, Table 4 · p. 17
• The inherited flow-matching video-prediction loss excludes the initial conditioning frame.The model block updates the DiT, action encoder, and two projectors; the latent block updates the environment table.Source: §A.2 · p. 16, Table 4 · p. 17
• For a held-out environment, test-time adaptation optimizes a 32-D latent from interaction trajectories while keeping all shared weights frozen.The main evaluations use 40 latent updates with a coarse-to-fine learning-rate schedule; the resulting latent is reused across query actions.Source: §A.2 · p. 17

B.1 TRAINING DATA AND RESOURCE BUDGETS

• Tables 5 and 6 report physical-environment and recorded-episode counts, while Table 5 also reports GPU-hour budgets and Diverse Backgrounds uses 30 scene backgrounds.Training batches are expressed as environments multiplied by action chunks per environment per optimizer update.Source: §B.1 · p. 17, Table 5 · p. 17, Table 6 · p. 17

B.2 SIMULATOR AND REAL-WORLD SETUP

• Simulation uses a LIBERO-based setup with a Franka Panda arm; real-world tasks use a Franka arm and fixed Intel RealSense D435i camera.Both settings use 14-channel actions; simulation videos are 224 × 224 at 20 fps, while real-world inputs are letterboxed to task-specific sizes.Source: §B.2 · p. 18

C.1 SHARED PROTOCOL AND METRICS

• All methods receive the same query inputs and candidate actions, and adapting methods receive the same test-time trajectories; no method sees query outcomes before prediction.Action-selection success is determined from the executed action’s actual outcome, while video quality includes PSNR, SSIM, and LPIPS.Source: §C.1 · p. 18
• Simulation evaluation tables count test-time interaction trajectories and query videos; K denotes trajectories per test environment.A query video is one action-conditioned video chunk evaluated at test time.Source: Table 7 · p. 18, §C.1 · p. 18

C.2 SIMULATED TASKS

• Simulated tasks vary friction, gravity, collision mass, lamp-control rules, balance mass ratios, and combined mass–friction conditions, with task-specific state tracking and success criteria.Friction and Diverse Backgrounds also test whether the representation follows dynamics across visual changes.Source: §C.2 · pp. 18–20
• Diverse Backgrounds varies scene appearance while retaining the Friction push task, testing whether the representation follows dynamics rather than background.Success is based on reaching the requested displacement region.Source: Figure 7 · p. 19, §C.2 · p. 20

C.3 REAL-WORLD TESTS

• Real-world Ball Friction changes bead load, Door Close changes latch spring strength, and both define success by a task-specific target outcome.The ball must reach a marked target; the door must close and latch.Source: §C.3 · p. 20, Figure 8 · p. 21
• Stick Balance varies end weights and tests whether the stick maintains balance; trials with no predicted balancing action count as failures.The model selects among predicted balancing actions, and one is chosen uniformly at random for execution.Source: §C.3 · pp. 20–21
• Soft Pull varies tube stiffness and curvature, testing whether predicted rotation thresholds separate no-slide from slide outcomes in two real-robot trials.The tests rotate to the predicted threshold minus and plus 6°; both outcomes must match their target conditions for success.Source: §C.3 · p. 21

D BASELINE IMPLEMENTATION DETAILS

• All four baselines use task-level training data and per-method GPU budgets, with shared query inputs, candidate actions, success criteria, and test-time trajectories for adaptive methods.DINOv2 Context and TTT-KQV differ in how sampled trajectories are used during training.Source: D BASELINE IMPLEMENTATION DETAILS · p. 21

D.1 STANDARD WORLD MODEL

• The Standard World Model fine-tunes the BWM-initialized DiT and action encoder on pooled task trajectories without an environment latent or test-time adaptation.It uses AdamW with a 10−5 model learning rate, 0.01 weight decay, 100 warmup updates, and gradient clipping at 0.5.Source: §D.1 · p. 21

D.2 LORA TTA

• LoRA TTA adds rank-8 adapters to the DiT attention projections, freezes the backbone, and optimizes reset adapters on each environment’s interaction trajectories.Adapters remain fixed across that environment’s query actions; the Friction configuration uses 50 updates at learning rate 10−4.Source: §D.2 · pp. 21–22

D.3 DINOV2 CONTEXT

• DINOv2 Context jointly trains a trajectory-context head with the BWM components, then infers test-time context in a forward pass without gradient updates.A frozen DINOv2 encoder supplies visual features, and the context head is trained using query-video flow-matching loss.Source: §D.3 · p. 22

D.4 TTT-KQV

• TTT-KQV inserts gated MLP modules into the BWM and learns fast weights from chunked interaction streams using per-token binding losses.At inference, it replays observed interactions at each denoising timestep; each query gets an independent update–read branch, with no query ground truth in the update.Source: §D.4 · p. 22

E.1 TEST-TIME INTERACTION-COUNT ABLATION

• Table 11 reports Figure 5(c,d) results for K = 1, 2, and 4 test-time interaction trajectories, with physical error units specified for Light Switch.The caption identifies the selected K but supplies no result values.Source: Table 11 · p. 23

E.2 LATENT DIMENSIONALITY

• Table 12 evaluates environment-latent dimensionality on Friction using K = 1, ten held-out environments, 90 query videos, and 25 action decisions.The caption identifies object-error, video-metric, and action-success measures but gives no values.Source: Table 12 · p. 23

E.3 LATENT-SPACE PCA ACROSS TASKS

• The latent-space PCA is fit separately for each task using centered training latents; test-time latents are projected into the same task-specific basis.Physical-setting labels are used only for visualization, not given to the model.Source: §E.3 · p. 23, Figure 9 · p. 24

E.4 COMPUTATIONAL EFFICIENCY

• Training profiles report median compute-only optimizer time per update and peak GPU memory; inference timing separates interaction processing from query generation.The training profile uses six loss-bearing chunks per update, while the inference profile uses 25 denoising steps and five query chunks.Source: §E.4 · p. 23, Figure 10 · p. 25, Figure 11 · p. 26
• At K = 1, T2-Z takes 18.09 s for latent adaptation and five query generations, versus 32.73 s for LoRA TTA and 37.57 s for cached TTT-KQV.These are compute-only inference times under the five-query comparison protocol.Source: §E.4 · p. 25, Figure 11 · p. 26

F ADDITIONAL QUALITATIVE RESULTS

• Figures 12–22 compare predictions for selected held-out actions, using the same initial observation and candidate action across methods.Each task shows two test environments and four actions per environment, except Soft Pull, which shows two actions per environment.Source: F ADDITIONAL QUALITATIVE RESULTS · p. 26
• The simulated examples include Friction, Gravity, Mass Collision, Light Switch, Mass Balance, Mass × Friction, and Diverse Backgrounds.The figure captions identify these as selected test actions.Source: Figure 12 · p. 27, Figure 13 · p. 28, Figure 14 · p. 29
• The remaining examples cover Ball Friction, Door Close, Stick Balance, and Soft Pull.The figure captions identify these as selected test actions.Source: Figure 19 · p. 34, Figure 20 · p. 35, Figure 21 · p. 36