Technical Summary: Back to Parsimonious Latents: Learning Task-Centric World Models from Visual Foundations

Problem Statement

World models enable agents to predict future dynamics conditioned on actions, a capability central to planning and control. However, existing approaches face a fundamental trade-off in latent representation:

1. Pixel-based models (e.g., Dreamer, MuZero) learn dynamics directly from pixels or compact latents shaped by reward signals. While effective, they often incur high computational costs or require online interaction/reward supervision to shape the latent space.
2. Foundation embedding models (e.g., DINO-WM, V-JEPA) leverage frozen pre-trained visual encoders to predict directly in a high-dimensional semantic space. While these offer rich semantic priors and zero-shot generalization, the embeddings contain excessive task-irrelevant details (textures, lighting, background) that do not affect control.

In reward-free offline settings, where agents must learn from fixed trajectories without reward supervision, these issues are exacerbated. Using raw foundation embeddings as the state space forces the model to waste capacity on irrelevant variations, leading to poor modeling efficiency, reduced control interpretability, and potential representation collapse, particularly in high-dimensional action spaces (e.g., 7-DoF robotic arms). The core challenge is to extract a compact, task-sufficient latent representation from these rich but noisy foundation embeddings without online interaction.

Methodology: TC-WM

The authors propose TC-WM (Task-Centric World Model), a framework that treats pretrained foundation embeddings not as the final state space, but as a "semantic scaffold" from which task-relevant dynamics are distilled.

Core Architecture

TC-WM operates on a pipeline of visual embeddings ($x_t$) and proprioceptive signals ($s^p_t$):

1. Joint Embedding: A frozen visual encoder (e.g., DINOv2) maps images to visual features, while a trainable linear layer encodes proprioception. These are concatenated to form a joint embedding $x_t$.
2. Latent Encoder (Projection): A linear encoder projects the high-dimensional joint embedding into a compact latent state $z_t$. Crucially, $z_t$ is factorized into two subspaces:
   • $z^s_t$ (State-relevant): Aligned with the physical state.
   • $z^c_t$ (Complementary): Retains residual semantic structure (appearance, background) to satisfy reconstruction constraints.
3. Latent Dynamics: A Transformer-based predictor ($f_\theta$) models temporal transitions in the compact latent space ($\hat{z}_{t+1} = f_\theta(z_{t-H:t}, a_{t-H:t})$).
4. Task-Centric Alignment: The model employs a contrastive objective (InfoNCE) to align the $z^s_t$ subspace with the proprioceptive state $s^p_t$. This actively forces the encoder to separate task-relevant physical factors from nuisance variations.
5. Embedding Decoder: A linear decoder reconstructs the original visual embedding from $z_t$. This acts as a regularizer, ensuring the compact latent remains anchored to the foundation manifold and preventing degenerate solutions.

Training Objective

The total loss combines four components:

• Latent Dynamics Loss ($L^z_{dyn}$): Minimizes prediction error in the latent space.
• TC Dynamics Loss ($L^s_{dyn}$): Predicts future proprioception directly from history to enforce task-centric dynamics.
• Alignment Loss ($L_{align}$): Contrastive loss aligning $z^s_t$ with $s^p_t$.
• Reconstruction Loss ($L_{rec}$): Reconstructs the visual embedding from $z_t$ to preserve visual structure.

Test-Time Planning

The learned task-centric latent space supports diverse planning strategies:

• Low-dimensional control (Navigation): Model Predictive Control (MPC) using the Cross-Entropy Method (CEM) to optimize action sequences in latent space.
• High-dimensional manipulation: Latent Diffusion Planning (LDP) combined with an Inverse Dynamics Model (IDM) to generate executable actions from latent trajectories.

Key Contributions

1. Task-Centric Decomposition: The paper introduces a method to decompose high-dimensional foundation embeddings into a compact, task-sufficient latent space by explicitly separating physical state dynamics from semantic nuisance variables.
2. Active Identifiability: Unlike prior works that passively assume mechanism diversity for identifiability, TC-WM uses contrastive alignment with proprioception to actively induce distinct conditionals for different physical states, theoretically guaranteeing the recovery of task-centric factors up to an affine transformation.
3. Reward-Free Offline Learning: The framework achieves high-quality world modeling and planning in offline settings without reward signals, relying solely on visual observations and proprioceptive side-information.
4. Theoretical Guarantee: The authors provide a theorem proving that under mild assumptions (contextual dynamics, latent-observed variability, and sufficient alignment), the learned compact latent is identifiable and recovers the true physical state up to a reparameterization.

Experimental Results

The authors evaluate TC-WM across 9 offline benchmarks (Maze, Wall, Push-T, Lift, Can, Square, Reacher, Cheetah, Hopper) against baselines including TD-MPC2, DreamerV3, MuZero, and DINO-WM.

• World Model Quality: TC-WM achieves the lowest latent rollout prediction error (MSE) across nearly all tasks, particularly in manipulation environments where accurate long-horizon dynamics are critical. While image-space reconstruction is competitive, the primary gain is in the stability and accuracy of the latent dynamics.
• Planning Performance:
  • In navigation (low-dimensional actions), TC-WM matches the performance of strong baselines like TD-MPC2.
  • In manipulation (high-dimensional, contact-rich actions), TC-WM significantly outperforms all baselines. For example, on the Robomimic "Lift" task, TC-WM achieves a 54% success rate compared to 29% for DINO-WM and 31% for TD-MPC2.
• Representation Analysis:
  • Linear Probing: Linear regressors trained on TC-WM's latent space recover proprioceptive and object states with higher $R^2$ scores than baselines, confirming the latent is physically grounded.
  • Anti-Collapse: TC-WM demonstrates better utilization of the latent space (higher effective rank) compared to DINO-WM, which suffers from representation collapse when using raw embeddings directly.
  • Robustness: TC-WM maintains performance under visual perturbations (Gaussian noise, color jitter) better than DINO-WM, with smaller drops in SSIM.

Significance and Claims

The paper claims that foundation embeddings are best utilized as a semantic scaffold rather than a direct state representation. By projecting these embeddings into a task-centric subspace aligned with physical signals, agents can achieve:

• Parsimonious Latents: A reduction in dimensionality that filters out irrelevant visual details while preserving the dynamics necessary for control.
• Improved Control: Enhanced planning performance in complex, high-dimensional environments where standard foundation models fail due to noise and irrelevant variations.
• Theoretical Soundness: A formal justification for why linear projections and contrastive alignment are sufficient to recover task-relevant factors from frozen visual encoders.

The authors position TC-WM as a bridge between the generalization power of foundation models and the controllability required for physical agents, specifically addressing the gap in reward-free offline learning scenarios.