Pretrained Features as Teachers, Not States: Compact Latent World Models for Visual Planning
Organizations: Delft University of Technology (TU Delft)
Abstract
A key design choice in visual world models is what information the latent state is encouraged to preserve. Models that learn a latent space with only a latent-prediction objective adapt to environment dynamics, but the predictive objective alone admits collapsed solutions and provides no explicit constraint on the visual information retained by the latent. Models built directly on pretrained visual features inherit a strong visual prior, but tie prediction and planning to a frozen state space. We propose Anchor-WM, which instead uses pretrained visual features as a teacher for learning a compact latent state rather than as the world-model state itself. Anchor-WM learns this compact latent with action-conditioned dynamics, while distilling the dense DINO representations through a training-only distillation head. This encourages the compact latent to retain visual information captured by the pretrained representation while discouraging representational collapse. At inference, the DINO encoder and distillation head are discarded, allowing prediction and planning to operate entirely in the learned latent space. Anchor-WM achieves the best average success rate across the four evaluated environments while enabling substantially faster planning than world models operating directly on DINO patch features. These results show that pretrained visual features can effectively guide world-model learning without defining the state space used for prediction and planning.
Figures & tables
| Method | Latent Size | TwoRoom | Reacher | PushT | Cube | Average |
|---|---|---|---|---|---|---|
| Predictive | 1 192 | 35.2 2.2 | 10.6 5.3 | 3.2 1.1 | 46.5 2.9 | 23.9 |
| PLDM | 1 192 | 93.0 2.0 | 69.5 5.8 | 71.8 6.1 | 51.2 1.6 | 71.4 |
| LeWM | 1 192 | 86.1 1.8 | 78.4 5.7 | 90.4 4.3 | 64.8 2.1 | 79.9 |
| DINO-WM | 196 384 | 99.6 0.9 | 67.2 1.8 | 65.2 6.6 | 74.4 2.6 | 76.6 |
| DINO-WM (CLS) | 1 384 | 59.3 1.7 | 45.2 3.4 | 37.8 5.6 | 50.2 3.8 | 48.1 |
| Anchor-WM (ours) | 1 192 | 99.6 0.9 | 82.4 4.3 | 84.8 5.4 | 72.0 4.2 | 84.7 |
| Method | TwoRoom | Reacher | PushT | Cube | Average |
|---|---|---|---|---|---|
| LeWM | 79.2 3.6 | 95.2 1.8 | 41.6 7.1 | 44.8 8.3 | 65.2 |
| DINO-WM | 96.0 2.5 | 90.0 5.7 | 32.0 2.8 | 53.2 5.0 | 67.8 |
| Anchor-WM (ours) | 96.4 1.7 | 96.4 3.3 | 46.4 3.6 | 45.2 2.7 | 71.1 |
| Method | TwoRoom | Reacher | PushT | Cube | Average |
|---|---|---|---|---|---|
| LeWM | 62.8 8.2 | 56.4 8.4 | 10.8 3.0 | 64.0 4.0 | 48.5 |
| DINO-WM | 98.0 2.4 | 60.8 7.0 | 62.8 4.8 | 76.4 3.6 | 74.5 |
| Anchor-WM (ours) | 99.2 1.1 | 47.2 10.8 | 33.2 4.8 | 68.0 6.5 | 61.9 |
| Property | Model | Linear MSE | Linear | MLP MSE | MLP |
|---|---|---|---|---|---|
| Joint Position | LeWM | ||||
| DINO-WM | |||||
| Anchor-WM (ours) | |||||
| Joint Velocity | LeWM | ||||
| DINO-WM | |||||
| Anchor-WM (ours) |
| Ablation | Setting | TwoRoom | Reacher | PushT | Cube |
|---|---|---|---|---|---|
| Distillation Objective | MSE | 99.6 0.9 | 82.4 4.3 | ||
| Cosine | 89.2 1.8 | 78.6 3.1 | |||
| Distillation Target | Patch | 99.6 0.9 | 82.4 4.3 | ||
| Global | 83.8 2.4 | 54.7 3.7 | |||
| DINO Guidance | Distillation | 99.6 0.9 | 82.4 4.3 | ||
| Alignment | 50.0 5.7 | 48.1 2.3 |
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
| Property | Model | Linear MSE | Linear | MLP MSE | MLP |
|---|---|---|---|---|---|
| Agent Location | LeWM | ||||
| DINO-WM | |||||
| Anchor-WM (ours) | |||||
| Block Location | LeWM | ||||
| DINO-WM | |||||
| Anchor-WM (ours) |
| Hyperparameter | Value |
|---|---|
| Encoder | ViT-Tiny (patch size 14), from scratch |
| Input resolution | |
| Latent dimension | 192 (CLS token) |
| Predictor | Causal Transformer, 6 layers |
| Attention heads / head dim | 16 / 64 |
| Predictor MLP dimension | 2048 |
World Model-Based Planning
World model-based planning uses models of how environments change to compare imagined actions and choose goal-directed behavior.
See all 245 papers