Project Page

Back to Parsimonious Latents:

Learning Task-Centric World Models from Visual Foundations

Minghao Fu   Fan Feng   Nicklas Hansen   Biwei Huang

University of California, San Diego

arXiv Paper Code Interactive Demo 16B Backbone

What latent space should a world model learn for planning and control beyond visual?

Abstract

World models enable agents to predict future dynamics conditioned on actions, making the choice of latent representation central to planning and control. Such representations are often either learned directly from pixels with limited semantic structure or inherited from frozen visual foundation models with excessive task-irrelevant detail, yielding state spaces that are poorly matched to downstream planning and control. This is especially challenging in reward-free offline settings, where the model must learn from fixed trajectories without reward supervision or online interaction.

To address this, we propose TC-WM, a framework for turning foundation-model embeddings into compact, task-sufficient world representations. The key design is to treat the pretrained embedding space as a semantic scaffold rather than as the final state space: TC-WM linearly projects high-dimensional visual embeddings into a compact latent as the dynamic space, aligns a subspace with the agent's physical state via contrastive learning, and reconstructs embeddings to preserve useful visual structure. This combines the generality of foundation features with the controllability of task-centric dynamics. Theoretically, we show that TC-WM suffices to identify the task-centric latent factors up to a simple transformation. Empirically, TC-WM enables test-time planning across diverse environments (e.g., Robomimic and D4RL), achieving better world-modeling quality and more precise control than state-of-the-art approaches.

Where the Task-Centric Latent Acts

TC-WM splits its compact latent into a task-centric part zs and a contextual part zc. The glow marks where a block drives the predicted change — pick a task and block to compare.

View
Task
Block
Threshold 0.12
Gamma 0.80
Gain 1.00
TC-WM · task-centricconcentrated
TC-WM · no alignmentdiffuse
DINO-WMdiffuse

Reading the glow. A saliency map: where advancing only that latent block one step changes the decode.

Why Visual Embeddings Aren't a Control State

Frozen foundation embeddings describe how a scene looks. Planning and control need a state: a compact handle on the factors an action can change. TC-WM recovers one with a single linear projection inside the embedding.

Comparison of world-model paradigms (a)-(c) and TC-WM (d)-(e)

Figure 1: Comparison of world-model paradigms. (a) Generative WM (MDN-RNN, IRIS) โ€” predict pixels directly. (b) Latent WM (TD-MPC, MuZero) โ€” predict in a learned latent. (c) Embedding WM (DINO-WM, V-JEPA) โ€” predict in a frozen foundation embedding. (d) TC-WM dynamics โ€” predict in a compact latent inside the embedding. (e) Task-centric structure โ€” align zs with proprioception, anchor zc via embedding reconstruction.

Method

TC-WM architecture: linear projection of joint visual and proprioception embeddings into a compact latent, with task-centric alignment, dynamics, and embedding reconstruction

Architecture. A frozen visual backbone is linearly projected into a compact latent; a designated subspace is aligned with proprioception via InfoNCE; a ViT predicts latent dynamics; a linear decoder reconstructs the embedding to prevent collapse. Under partial alignment, the task-centric block is identifiable up to an affine map.

Key Findings

01

Compact beats verbatim

A single linear projection of frozen DINOv2 embeddings gives the lowest latent-rollout error across nine tasks. Predicting in a parsimonious latent is more accurate than predicting in the raw embedding.

02

Task-centric planning wins

Aligning a latent subspace with proprioception is the only method that surpasses DINO-WM on every Robomimic task, giving more precise control.

03

The latent stays interpretable

Alignment concentrates the dynamic latent on the gripper and manipulated object. Without it, the influence spreads across the whole scene.

Robomimic highlight: success rate, latent-rollout MSE, linear probing on Lift and Can, and anti-collapse

Figure 2: Robomimic highlight. (1) Success rate. (2) Latent-rollout MSE. (3–4) Linear probes on Lift and Can. (5) TC-WM avoids the latent collapse seen when rolling out directly on foundation embeddings.

Results

World model prediction accuracy across nine environments

World-model prediction. Lowest latent-prediction error on nearly all tasks; competitive image reconstruction. Lower is better.

Planning performance under CEM and LDP

Planning. CEM on Maze / Wall / Push-T / Cheetah / Hopper; LDP on Lift / Can / Square. Only method that surpasses DINO-WM on every LDP task.

Clean observationclean input
TC-WM rollout from cleanTC-WM pred (clean)
Perturbed observationnoisy input
TC-WM rollout from perturbedTC-WM pred (noisy)

Lift  ·  Gaussian noise + color jitter

Unseen perturbations. Under additive Gaussian noise and per-channel color jitter never seen during training, TC-WM's open-loop rollout preserves cube position and manipulator silhouette at near-clean fidelity.

Open-Loop Rollout

Lift originaloriginal
Lift reconstructionreconstruction
Lift TC-WM rolloutrollout

Lift

Can originaloriginal
Can reconstructionreconstruction
Can TC-WM rolloutrollout

Can

Square originaloriginal
Square reconstructionreconstruction
Square TC-WM rolloutrollout

Square

Wall originaloriginal
Wall reconstructionreconstruction
Wall TC-WM rolloutrollout

Wall

Maze originaloriginal
Maze reconstructionreconstruction
Maze TC-WM rolloutrollout

Maze

Cheetah originaloriginal
Cheetah reconstructionreconstruction
Cheetah TC-WM rolloutrollout

Cheetah

Analysis

Architecture and projection-head ablations on Robomimic

Architecture / projection ablations on Robomimic. Visual encoder / latent-source choice (a) and projection-head ablation (b), each showing success rate on the left and mirrored SSIM on the right. RP denotes randomly projected DINOv2 embeddings. Linear projection is the best; DINOv2 and DINOv3 dominate as foundation encoders.

Loss component ablation on Lift

Loss components on Lift. Removing embedding reconstruction collapses both planning and visual fidelity; removing proprioceptive supervision lowers success rate while preserving SSIM; the latent split dimension matters little.

Environments

Nine offline visual-control tasks: navigation, locomotion, and manipulation. We use images, actions, and proprioception during training, without reward.

Maze trajectory Maze
Wall trajectory Wall
Push-T trajectory Push-T
Robomimic Lift trajectory Lift
Robomimic Can trajectory Can
Robomimic Square trajectory Square
DMC Reacher trajectory Reacher
DMC Cheetah trajectory Cheetah
DMC Hopper trajectory Hopper

Scaling to a Video World Model

The method is not specific to DINOv2 or to simulation. We finetune Cosmos3-Nano, a 16B video foundation model, with the task-centric objective — denoted TC-Cosmos3 — and evaluate it on real DROID robot episodes. Each clip is conditioned on identical ground-truth actions from a common start frame, so any divergence reflects the world model rather than the policy: Ground truth · TC-Cosmos3 · Cosmos3-Nano.

Put the blocks into the bowl, then wipe the table

Ground truthTC-Cosmos3Cosmos3-Nano

Remove the cables from the plastic bag and put them on the table

Ground truthTC-Cosmos3Cosmos3-Nano

Plug the charger head into the adapter

Ground truthTC-Cosmos3Cosmos3-Nano

Put the black bottle and the sponge on the top shelf

Ground truthTC-Cosmos3Cosmos3-Nano

Remove the screwdriver from the rack and put it in the cup

Ground truthTC-Cosmos3Cosmos3-Nano

Real-robot video rollouts. Each panel stitches three camera views (wrist on top, two exterior below). TC-Cosmos3 stays anchored to the gripper and the manipulated object while letting non-interacting background vary — the same disentanglement measured in simulation, now carried by a 16B video model on real hardware.

BibTeX

@article{fu2026tcwm,
  title  = {Back to Parsimonious Latents: Learning Task-Centric World Models from Visual Foundations},
  author = {Fu, Minghao and Feng, Fan and Hansen, Nicklas and Huang, Biwei},
  journal= {arXiv preprint arXiv:2605.25620},
  year   = {2026}
}