A world model is a learned simulator. Give it the recent past — a few frames and the actions just taken — and it predicts what happens next. Once an agent can imagine the consequence of an action, it can plan: search over action sequences, keep the ones whose imagined futures look good, and act on them. This is the appeal of model-based control, and it is why world models have become a central topic in scaling up robot learning.
An ideal world model is not simply to make world models predict better, but to understand the latent space in which prediction should happen for planning and control. Most existing work improves world models through better tokenization, longer rollout horizons, or sharper visual prediction. In contrast, we ask what state representation a world model should operate on. A predictor is only useful for control if its latent state preserves the action-relevant factors needed for decision making. With an inappropriate state space, increasing model capacity does not necessarily improve planning; with a control-sufficient state space, even a compact dynamics model can support effective planning.
The dominant design choice today is to learn world-model dynamics either in the pixel space, as in video-generation-based simulators, or in the embedding space of a frozen visual foundation model, such as JEPA or DINO-WM. The latter has become an especially natural recipe: use a pretrained encoder to map each image into a dense field of transferable patch features, freeze it, and train the dynamics model directly in that feature space, and this approach generally can work well in practice. However, it makes the world model inherit the representational biases of the frozen encoder. These features may encode texture, lighting, background appearance, table patterns, wall color, and other visual factors that are useful for recognition or correspondence, but irrelevant to control and often invariant under the agent's actions.
A visual representation answers what the scene looks like. Control needs a latent that answers what matters for acting.
In this blog, we introduce TC-WM, a task-centric world model built on top of visual foundation representations. Rather than treating a pretrained embedding as the state space for dynamics learning, we use it as a semantic scaffold from which a more compact control-relevant state can be extracted. The visual encoder remains frozen, but prediction is not performed directly in its embedding space. Instead, TC-WM learns a single linear projection into a compact latent space, decomposes this latent into a control-centric component aligned with the agent's own physical state and a complementary context component, and learns dynamics in this projected space. The resulting representation removes much of the visual variation that is useful for recognition but unnecessary for control, while retaining the factors needed for planning.
01Why Visual Embeddings Are Not a State for Control
For a single frame, pre-trained visual embeddings (e.g., DINO) produce a dense grid of patch tokens, with each token carrying a high-dimensional visual descriptor. This representation is too broad for control: it preserves local appearance, object identity, texture, illumination, background structure, and other factors that are useful for recognition and correspondence. The object to be manipulated and the agent body are present in this tensor, but they are entangled together with many visual degrees of freedom that are only weakly related to the control problem.
For control, this creates a mismatch between the encoder's representation and the state required by the dynamics model. The variables that determine action effects are much lower-dimensional: the pose of the end-effector, the pose of the relevant object, their relative geometry, and contact or interaction state. These factors are spatially localized and causally affected by actions. In contrast, much of the frozen embedding describes factors that are either static or uncontrollable, such as background appearance, lighting, or texture. Learning dynamics over the full embedding therefore asks the world model to predict many dimensions that do not participate in the action-conditioned transition structure.
The mismatch also changes the geometry of the learned state space. In the frozen embedding space, two observations can be far apart because of visual changes that are irrelevant to the task, while two control-distinct configurations may be close if they share similar appearance. For planning, we want the opposite: distances in latent space should reflect differences in controllable task state. States should be close when the agent, object, and interaction configuration are equivalent for control, even if the rendered images differ in texture, lighting, or background. TC-WM addresses this by extracting such a state from the frozen representation, rather than retraining the visual encoder or predicting directly in its full feature space.
02One Simple Linear Projection, Not a New Encoder
A natural approach to this representation mismatch is to train or fine-tune a new encoder. TC-WM takes a different route for both practical and diagnostic reasons: we keep the visual encoder fixed and ask whether the control-relevant state can be recovered from its existing representation. The conjecture is that the frozen embedding already contains the variables needed for control, but does not organize them in the form required by a dynamics model. The issue is therefore not missing information, but entanglement: task-relevant factors are mixed with appearance, background, and other nuisance variation. TC-WM extracts these factors using the simplest transformation possible: a learned linear projection from the frozen embedding to a compact latent state.
Let \(x_t\) be the frozen visual embedding of the observation, concatenated with a trainable embedding of the agent's proprioceptive reading. We project it linearly to a compact latent,
and we learn the dynamics in \(z\) rather than in \(x\):
Two components make this design work. First, the compactness of \(z\) constrains what the dynamics model can carry forward. Instead of modeling every visual degree of freedom present in the frozen embedding, the model is forced to retain a smaller set of variables that are predictive under actions. Second, the projection is linear. This is not just an architectural simplification, but an explicit assumption about the pretrained representation: the visual encoder has already performed the hard nonlinear transformation from pixels to semantic features, so the remaining task is to identify the control-relevant subspace within those features.
We empirically ablated this design choice by comparing the learned linear projection with both a random projection and a nonlinear projection head. Replacing the learned linear map with a random projection hurts performance, indicating that the control-relevant subspace must be selected rather than arbitrarily compressed. Replacing it with a nonlinear head does not improve performance, suggesting that the useful operation is closer to subspace extraction than to learning another deep representation.
03Aligning Part of the Latent With the Physical States
A compact latent is not necessarily task-centric. The projection reduces dimensionality, but by itself it does not specify which directions should encode controllable state and which should encode contextual visual information. TC-WM introduces this structure using a signal that is available in embodied control without additional annotation: the agent's own proprioception.
We disentangle the latent into two parts, \(z = [z^s, z^c]\). The first, \(z^s\), is the control-centric slice. We ask it to agree with the agent's physical state \(s^p\): joint angles, end-effector pose, and the readout the robot reports about itself. We tie them with a contrastive objective that pulls a slice and its matching proprioceptive reading together and pushes mismatched pairs apart:
Hence, only \(z^s\) is constrained by proprioception, so this part of the latent is anchored to variables directly tied to the agent's body and action-conditioned motion. The remaining component, \(z^c\), is left unconstrained by this alignment and can encode task-relevant context not explained by proprioception, such as object configuration, scene structure, visual appearance, and other information needed for prediction.
This split introduces a second requirement: the latent must remain informative enough for dynamics learning. If the projection is optimized only toward compactness and proprioceptive alignment, it can discard visual information that is not immediately captured by \(z^s\) but is still needed to predict future observations and task outcomes. To prevent this, TC-WM adds a linear decoder that reconstructs the original embedding,
and penalizes the reconstruction error. The decoder is used only during training. Its role is to regularize the decomposition: information not captured by the proprioception-aligned slice should remain recoverable from the full latent, encouraging contextual factors to be represented in \(z^c\) rather than dropped.
The linear projection also makes the representation easier to analyze. Under the model assumptions, the projection identifies the task-centric subspace up to an affine transformation, rather than learning an unconstrained nonlinear code. This matters because the latent is not only optimized for downstream performance; it is intended to correspond to the causal factors relevant for control. In this sense, \(z^s\) is not merely a useful feature vector, but a structured handle on the controllable part of the state.
The ablations on Lift support the role of each component. Removing embedding reconstruction significantly degrades both planning performance and visual prediction, indicating that the compact latent has lost information needed by the dynamics model. Removing proprioceptive alignment mainly reduces task success while preserving visual quality, suggesting that the model can still predict plausible futures but no longer organizes the latent around control-relevant variables. Varying the dimensional split between \(z^s\) and \(z^c\) has a much smaller effect, indicating that the method is not highly sensitive to the exact allocation as long as the alignment and reconstruction constraints are present.
04Dynamics and Planning
With the latent defined, the rest is a world model. A Transformer predicts the next latent from a short history of latents and actions,
and a companion head predicts the next proprioceptive state, so the part of the latent tied to the body stays grounded as a rollout unfolds. Training minimizes a sum of four terms: latent dynamics, proprioceptive dynamics, the alignment loss, and the reconstruction regularizer:
The framework is trained offline, from fixed trajectories of images, actions, and proprioception, with no reward labels and no online interaction. At test time, the world model is frozen, and the controller is matched to the task. For goal-reaching (the navigation mazes and the planar push), we plan with the cross-entropy method (CEM), sampling action sequences. For locomotion, we train a policy with SAC against the learned dynamics. For high-dimensional manipulation, a latent diffusion planner (LDP) proposes trajectories that the world model then scores.
05Why We Should Choose TC-WM
We evaluate across nine offline visual-control environments in three regimes: navigation (Maze, Wall), manipulation (Lift, Can, Square, Push-T), and locomotion (Reacher, Cheetah, Hopper). The baselines include TD-MPC2, DreamerV3, MuZero, and DINO-WM.
The first thing a good state buys is predictability. Because the latent keeps what matters, and the decoder keeps it honest, multi-step rollouts stay accurate where the task lives, rather than drifting toward a blurry, plausible-looking mean.
Each result is an open-loop rollout: the original demonstration, TC-WM's embedding reconstruction, and its multi-step prediction rolled forward from the first frames.
TC-WM is competitive across goal-reaching and locomotion. The gap is widest exactly where the argument predicts: manipulation, where the controllable object occupies a small, moving part of a busy scene, and appearance-dominated latents waste the most capacity.
The most informative evidence comes from the perturbation analysis. We perturb the control-centric slice \(z^s\) and measure which image patches are most affected through the reconstruction or prediction pathway. The response is spatially concentrated on the gripper and the manipulated object, and it follows the object as it moves through the scene. This indicates that \(z^s\) primarily controls the parts of the representation associated with action-conditioned interaction. In contrast, perturbing \(z^c\) produces changes that are more distributed over the background and other contextual regions. These are factors that can affect visual prediction but are not directly tied to the agent's own motion. Thus, even though the only explicit supervision comes from proprioception, the learned split produces a meaningful separation: \(z^s\) captures the controllable interaction state, while \(z^c\) absorbs contextual visual variation.
06Scales to a Video World Model
We scale up the framework with Cosmos3-Nano base model, a 16B video foundation model, finetune it the task-centric way (we call the result TC-Cosmos3), and roll out real DROID robot episodes.
Every clip below is conditioned on the same ground-truth actions from the same start frame, so any divergence is the world model instead of the policy. Ground truth is the real episode; TC-Cosmos3 is Cosmos3-Nano finetuned with the TC-WM objectives. Each clip stitches three camera views: the wrist on top, two exterior views below.
In real-robot rollouts, TC-WM stays anchored to the gripper and the manipulated object, while allowing non-interacting parts of the scene to vary. This reflects the same disentanglement measured in simulation: control-relevant interaction factors remain tied to the agent and object, while background variation is absorbed as context. Compared with the original base model (Cosmos3-Nano), TC-WM is both more physically consistent and more task-centric.
07From Representation to Control
This work points toward a more physical and task-centric view of world modeling: instead of predicting all visual variation equally, the model should allocate capacity to the factors that govern interaction, control, and future action effects. In this sense, “back to parsimonious latents” should be taken literally. Although larger and more expressive representations often improve transfer in perception, acting can require the opposite: identifying the compact subspace that carries the task-relevant state and planning within it. The opening question — what latent space should a world model learn for planning and control? — therefore has a concrete answer: a latent space that is compact, physically grounded, and sufficient for the task.