PaperlayerView exact revision on arXiv

Source-linked AI summary

DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning

Gaoyue Zhou, Hengkai Pan, Yann LeCun, Lerrel Pinto

arXiv:2411.04983v2cs.ROcs.AI

TL;DR

DINO-WM addresses the challenge of generalizing offline world models beyond task-specific training and precomputed solutions. It predicts future DINOv2 patch embeddings and plans toward visual goals, achieving zero-shot solutions and higher success across varied environments, while requiring sufficient offline state-action coverage.

  • Problem

    Existing world models often use feed-forward, task-specific policies, while auxiliary information such as demonstrations, inverse models, or dense rewards can reduce offline generality.

  • Method

    DINO-WM predicts future DINOv2 pretrained spatial patch embeddings from offline trajectories and actions, enabling latent-space visual-goal planning without pixel reconstruction.

  • Results

    DINO-WM achieves zero-shot behavioral solutions and higher success across environment variations, improving arbitrary-goal success by 45% on average and LPIPS performance by 56% on the hardest tasks.

  • Takeaways & Limitations

    The results support task-agnostic visual reasoning and test-time planning across unseen configurations without expert demonstrations, reward modeling, or pre-learned inverse models.

  • Takeaways & Limitations

    DINO-WM assumes offline datasets with sufficient state-action coverage and ground-truth actions, and currently plans in action space rather than integrating hierarchical fine-grained control.

Abstract

from arXiv · show

The ability to predict future outcomes given control actions is fundamental for physical reasoning. However, such predictive models, often called world models, remains challenging to learn and are typically developed for task-specific solutions with online policy learning. To unlock world models' true potential, we argue that they should 1) be trainable on offline, pre-collected trajectories, 2) support test-time behavior optimization, and 3) facilitate task-agnostic reasoning. To this end, we present DINO World Model (DINO-WM), a new method to model visual dynamics without reconstructing the visual world. DINO-WM leverages spatial patch features pre-trained with DINOv2, enabling it to learn from offline behavioral trajectories by predicting future patch features. This allows DINO-WM to achieve observational goals through action sequence optimization, facilitating task-agnostic planning by treating goal features as prediction targets. We demonstrate that DINO-WM achieves zero-shot behavioral solutions at test time on six environments without expert demonstrations, reward modeling, or pre-learned inverse models, outperforming prior state-of-the-art work across diverse task families such as arbitrarily configured mazes, push manipulation with varied object shapes, and multi-particle scenarios.

1. Introduction

DINO-WM addresses the limits of task-specific and online world models by learning task-agnostic visual dynamics from offline trajectories and optimizing behavior at test time. It uses pretrained DINOv2 embeddings to support zero-shot goal reaching across varied environments.

  • Motivation: Existing feed-forward policies must encode solutions for possible tasks and scenarios before deployment, making broad generalization difficult.The supplied passage states that successful generalization requires exposure to similar scenarios during training, while learning every possible solution is infeasible or inefficient.
  • Motivation: Online world models require environment interaction and policy-specific data coverage, so new tasks can require retraining even within the same environment.The online setting continuously collects data, but the learned model is accurate only over the optimized policy’s coverage.
  • Motivation: Offline world models often depend on expert demonstrations, structured keypoints, pretrained inverse models, or dense rewards, reducing their generality.The paper frames the central question as whether alternative auxiliary information can preserve offline world-model generality.
  • Approach: DINO-WM predicts future DINOv2 patch embeddings from offline trajectories and uses model predictive control to optimize visual goal reaching at test time.Its latent-space formulation uses spatial and object-centric pretrained features rather than reconstructing raw observations during world-model training.
  • Results: 56% improvement over prior state-of-the-art work is reported on LPIPS metrics for the hardest tasks.The metric concerns visual reconstruction from trained decoders used to measure future world modeling quality.
  • Results: 45% average improvement over prior work is reported for reaching arbitrary goals on the hardest tasks.The result concerns behavior optimization using latent world models.
  • Results: DINO-WM generalizes across environment variations within task families and achieves higher success rates than prior work.Examples include different maze layouts and object shapes for navigation and manipulation.

2. Related Work

The related work situates DINO-WM among model-based learning, generative world models, and pretrained visual representations. It emphasizes latent-space modeling for more precise planning and control while avoiding the computational burden of pixel-space video generation.

  • Model-based Learning: Model-based learning studies dynamics models and future-state prediction across control, planning, robotics, and vision-based embodied-agent learning.The passage identifies a broad literature spanning classical control and recent vision-based applications.
  • Generative Models as World Models: Generative world models predict videos conditioned on text or high-level action sequences for self-driving, control, and general-purpose video generation.The supplied passage describes these as initiatives built around large-scale foundation models.
  • DINO-WM: DINO-WM uses pretrained DINOv2 embeddings of image frames, then plans through the learned model with model predictive control toward a target observation.The figure caption presents this pipeline as training from pretrained embeddings, test-time behavior optimization, and improved performance over prior world models.
  • Generative Models as World Models: Pixel-space diffusion video generation is computationally expensive for test-time optimization, motivating latent-space world models for planning and control.The supplied passage specifically connects this computational cost to restricting methods such as model predictive control.
  • Pretrained Visual Representations: Pretrained visual representations provide compact features that capture spatial and semantic information for downstream tasks.The related-work passage lists pretrained image and video representation families including ImageNet, I-JEPA, DINO, V-JEPA, R3M, and MVP.

3. DINO World Models

DINO-WM learns task-agnostic visual dynamics from offline trajectories in a latent space built from frozen DINOv2 patch features. At test time, it optimizes action sequences toward RGB-image goals using model predictive control and cross-entropy search.

  • Overview: DINO-WM learns task-agnostic world models from precollected offline trajectories, using latent states rather than raw observations.The model is designed to support visual reasoning and control at test time without task-specific online retraining.
  • Decoder for interpretability: The optional decoder reconstructs images for visualization and interpretability, while remaining independent of transition-model training and planning.Avoiding pixel reconstruction during planning reduces computational costs.
  • Observation model: The frozen DINOv2 observation model converts each image into spatial patch embeddings z_t ∈ R^N×E, providing task- and environment-independent visual features.The representation is intended to capture spatial information useful for navigation and manipulation.
  • Transition model: The transition model predicts the next latent state from H previous latent states and actions, using causal frame-level attention for autoregressive prediction.Patch vectors within one observation are treated collectively to model global structure and temporal dependencies.
  • Transition model: Training uses teacher forcing and latent consistency losses over trajectory segments, without reconstructing pixel images.The action encoder maps actions to higher-dimensional representations before conditioning predictions.
  • Visual planning: At test time, DINO-WM uses MPC with CEM to optimize action sequences by minimizing latent-space MSE between predicted and goal states.Although gradient descent is possible because the model is differentiable, CEM empirically outperforms it in the experiments.

4. Experiments

The experiments evaluate DINO-WM and offline world-model baselines across six control suites, testing zero-shot planning, encoder choices, generalization, and visual prediction quality. DINO-WM performs especially strongly on manipulation and unseen-configuration settings while producing accurate long-horizon predictions.

  • Experimental setup: Experiments span six suites covering maze navigation, tabletop pushing, robotic-arm control, and deformable-object manipulation.Tasks use RGB observations and randomly sampled goal observations from arbitrary initial states.
  • Baselines: DINO-WM is compared with IRIS, DreamerV3, TD-MPC2, and an action-conditioned AVDC variant using offline datasets and model-predictive control.The comparisons omit reward and task information for the offline world-model baselines.
  • Zero-shot planning: 45% average improvement in success on the hardest tasks demonstrates DINO-WM’s performance for reaching arbitrary goals through latent-world-model planning.The evaluation measures success on Maze, Reach, PushT, and Wall, and Chamfer Distance on Rope and Granular.
  • Zero-shot planning: DINO-WM matches state-of-the-art world models on simpler Wall and PointMaze environments and significantly outperforms prior work on manipulation tasks requiring contact and object-dynamics inference.The authors attribute TD-MPC2’s poor performance to learning latent representations without reward signals.
  • Pre-trained representations: Patch-based visual representations outperform single-vector encoders as environment complexity increases, where precise control and spatial understanding are required.Single global vectors from R3M, ResNet, and DINO CLS lose spatial details needed for manipulation.
  • Novel configurations: DINO-WM generalizes to unseen layouts and object configurations, achieving the lowest Chamfer Distance in GranularRandom despite out-of-distribution particle counts.PushObj remains challenging because training included only four object shapes.
  • Visual prediction: Action-conditioned AVDC produces visually realistic but physically implausible futures and diverges from ground truth during long-horizon prediction.These failures make it insufficient for accurate task planning despite realistic image generation.
  • Visual prediction: 56% improvement on LPIPS for the hardest tasks indicates stronger predicted visual futures than prior state-of-the-art models.DINO-WM’s open-loop rollouts are visually indistinguishable from ground truth in the reported Push-T and Granular visualizations.

5. Conclusion

The paper concludes that DINO-WM models visual dynamics in latent space without pixel reconstruction and supports task-agnostic test-time planning. It also identifies data coverage, action availability, and low-level control as important boundaries for future work.

  • Contribution: DINO-WM models visual dynamics in latent space without requiring pixel-space reconstruction.The method uses pretrained visual representations and supports planning toward observational goals.
  • Conclusion: DINO-WM captures environmental dynamics, generalizes to unseen configurations, and generates zero-shot downstream solutions through test-time planning.The conclusion frames this as progress toward task-agnostic world modeling, reasoning, and control.
  • Limitations and future work: DINO-WM assumes offline datasets with sufficient state-action coverage, which is difficult to obtain in highly complex environments.The authors suggest combining the method with exploration and model updates as new experiences become available.
  • Limitations and future work: The method relies on ground-truth agent actions and currently plans in action space rather than integrating hierarchical high-level planning with low-level control.The paper proposes hierarchical extensions for more fine-grained control tasks.

Impact Statement

The paper aims to facilitate task-agnostic world models and their applications.

  • Impact: The paper targets learning and applying task-agnostic world models.The authors do not identify a specific societal consequence requiring emphasis.

A.1. Environments and Dataset Generation

The paper evaluates six environments spanning navigation, rigid and deformable manipulation, and robotic control, using offline random-action trajectories for training and varied goals or configurations for testing.

  • Navigation: Maze training uses 2000 fully random trajectories, while Wall uses 1920 trajectories for fixed layouts and 10240 for varied configurations.
  • Robotic control: Reacher training uses 3000 trajectories of 100 steps, with testing requiring the entire arm to match arbitrary target poses.
  • Rigid manipulation: Push-T requires both the pusher and T-block to reach target locations within 25 steps, testing precise contact-rich dynamics.
  • Deformable manipulation: Rope manipulation trains on 1000 random 20-step trajectories and tests varied goal configurations from diverse initial states.
  • Particle manipulation: Granular manipulation uses about one hundred particles, with 1000 training trajectories from one initial configuration and testing across varied shapes and particle arrangements.

A.2. Environment families for testing generalization

Generalization is tested by changing layouts, object shapes, and material configurations while retaining task families represented during training.

  • Spatial layouts: WallRandom tests navigation with non-overlapping wall and door positions not seen during training.
  • Object shapes: PushObj trains on four block shapes and evaluates on two unseen Tetris-like or plus-shaped blocks.
  • Material configurations: GranularRandom varies particle count and requires gathering all particles into a square at a randomly sampled location.
  • Visualization: Visualizations of the generalization environments are provided in Figure 5.

A.3. Pretraining features

The feature comparisons distinguish pretrained visual representations by whether they preserve global or spatial information for world-model inputs.

  • Compared representations: R3M uses a ResNet-18 pretrained on real-world human manipulation videos.
  • Compared representations: ImageNet uses a ResNet-18 pretrained on the ImageNet-1K dataset.
  • DINOv2 features: DINOv2 provides Patch and CLS embeddings, with CLS described as a one-dimensional vector encoding global image information.
  • Spatial information: Patch-based representations retain spatial information that global feature vectors from R3M, ResNet, and DINO CLS lose for manipulation tasks.

A.4.1. SCALING LAWS OF DINO-WM

The appendix describes DINO-WM’s scaling behavior, training ablations, latent-space planning procedures, and comparisons with action-conditioned generative models and alternative world models.

  • Scaling behavior: As PushT training data increases from 200 to 18500 trajectories, prediction quality and planned behavior improve significantly.The authors attribute this trend to capturing more diverse dynamics and environmental nuances.
  • Attention-mask ablation: The causal-attention-mask ablation compares models across varying history lengths h on PushT.
  • Reconstruction-loss ablation: Training the predictor with decoder-propagated reconstruction loss performs reasonably on PushT but falls slightly short of independently trained prediction.The comparison supports separating feature learning from reconstruction objectives.
  • Latent-space planning: Planning encodes current and goal images, predicts latent trajectories under candidate actions, and minimizes final-state MSE to the goal latent.The cost is C = ∥ẑ_T − z_g∥2, with latent transitions generated recursively by the predictor.
  • CEM and MPC: CEM samples action sequences, selects the lowest-cost candidates, updates their distribution, and executes the first k actions before replanning.
  • Planning variants: Gradient descent directly optimizes actions through backpropagation, while MPC uses receding-horizon CEM and empirically outperforms gradient descent.
  • Generative-model comparison: Action-conditioned AVDC diverges from ground-truth observations during long-term open-loop prediction, limiting its accuracy for task planning.
  • Evaluation and implementation: DINO-WM’s visual prediction quality is evaluated with SSIM and LPIPS, while Table 10 reports inference, rollout, and CEM planning times.
Loading 2411.04983v2…