Search papers, labs, and topics across Lattice.
This paper introduces Task-Centric World Models (TC-WM), a method to learn compact, task-relevant latent spaces for planning and control from pretrained visual foundation models in offline settings. TC-WM projects high-dimensional visual embeddings into a compact latent space, aligns a subspace with the agent's physical state via contrastive learning, and reconstructs embeddings to preserve useful visual structure. Experiments on Robomimic and D4RL demonstrate that TC-WM achieves better world-modeling quality and more precise control than existing approaches by combining the generality of foundation models with task-centric dynamics.
Foundation model embeddings contain so much irrelevant detail that projecting them into a task-centric latent space actually *improves* world model quality and control performance.
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 underlying 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.
