A Universal World Model Learned from Large Scale and Diverse Videos

Hanchen Cui ¹ and Yang Gao ^1,2

¹ Shanghai Qi Zhi Institute ² Tsinghua University

Paper Code Poster BibTeX

>200 h In-the-wild video data pretraining Large-scale, diverse videos

≈1.7 min Robotic data for adaptation 2,000 frame pairs at 20 Hz

0 labels Action-free pretraining Learn latent actions directly from videos

Learn from diverse videos. Adapt to new environments.

Pretrain a world model on action-free, in-the-wild videos,
then specialize it with a small amount of in-domain data.

In-the-wild video data trains a universal world model. In-domain robot data adapts it to a specific world model, connected to an agent. — Large-scale video pretraining provides a universal starting point. In-domain fine-tuning specializes the dynamics, and an action adapter connects the learned action space to environment controls.

Learn without action labels

Extract compact latent actions directly from pairs of neighboring video frames.

Capture diverse dynamics

Pretrain on everyday videos from Something-Something-v2 and Ego4D.

Adapt with less data

Improve predictions in a new domain with just 2,000 paired frames.

Abstract

World models play a crucial role in model-based reinforcement learning (RL) by providing predictive representations of an agent in an environment and enabling the agent to reason about the future and make more informed decisions. However, there are still two main problems limiting the applications of world models. First, current methods typically train world models using only massive domain-specific data, making it challenging to generalize to unseen scenarios or adapt to changes in the environment. Second, it is difficult to define actions when world models are trained using in-the-wild videos.

In this work, we tackle these two problems by learning a general-purpose world model from a diverse and large-scale real-world video dataset with extracted latent actions. Specifically, our approach leverages a pretrained vision encoder to project images of two adjacent frames into states, extracts latent actions into a low-dimensional space based on vector quantization, and learns a dynamics function using these latent actions.

Results show that the proposed generic world model can successfully extract latent actions from arbitrary neighboring frames when testing on in-the-wild video datasets. Furthermore, fine-tuning on only a small amount of in-domain data can significantly improve the accuracy of the generic world model when adapting to unseen environments.

Universal World Model framework

Learn compact latent actions from videos, then align them with environment-specific actions.

Top: adjacent video frames are encoded and passed through an action extractor, discrete codebook, and dynamics model. Bottom: an action adapter maps environment actions into the same latent action space. — **Video pretraining → action-conditioned adaptation.** A frozen visual encoder provides state representations. The action extractor and vector-quantized codebook capture transitions; the dynamics model predicts future states. An action adapter subsequently maps environment actions to the learned latent action space.

Implementation details

Visual representation

Encode the observations

A frozen MAE encoder and decoder provide visual representations and image reconstruction for 224 × 224 observations.

Latent action learning

Compress what changes

A 4-layer ViT action extractor and a VQ-VAE with 1,024 codes of 64 dimensions create a compact information bottleneck. A 2-layer ViT learns dynamics.

Action adaptation

Connect to real controls

An action adapter takes the current state and ground-truth environment action and learns to match the extracted latent action.

Video frames are sampled at 0.2-second intervals. During pretraining, the visual encoder and decoder remain frozen while the action extractor, codebook, and dynamics network are optimized.

Dynamics in the wild

Future-frame prediction on Something-Something-v2 and Ego4D,
conditioned on latent actions extracted from neighboring frames.

Five Something-Something-v2 examples and five Ego4D examples, each showing two observations, their reconstructions, and the predicted next frame. — **Learning transitions from observation.** The world model predicts a future frame from the current state and the extracted latent action. Reconstruction and prediction blur comes from the pretrained MAE decoder. On smaller screens, swipe horizontally or select Enlarge to inspect the examples.

Explore individual examples

Current observed frame, SSv2 example 1 — O _t Current observation

Reconstruction of current frame, SSv2 example 1 — R _t Reconstruction

Next observed frame, SSv2 example 1 — O _t+1 Next observation

EXAMPLE

LPIPS to next frame 0.193 vs. 0.280 to current ↓ LOWER IS BETTER

Static examples from the paper, not a live model. Reconstruction blur comes from the pretrained MAE decoder. LPIPS uses equally blurred observations.

Perceptual prediction results

Across all ten paper examples, predictions are closer to the next observation
than to the current observation. Lower LPIPS is better.

LPIPS between the prediction and each observation. Lower is better. All ten examples are closer to the next observation.
Dataset	Compared with	Example 1	Example 2	Example 3	Example 4	Example 5
SSv2	Current O _t	0.280	0.437	0.350	0.323	0.416
SSv2	Next O _t+1	0.193	0.240	0.224	0.201	0.153
Ego4D	Current O _t	0.470	0.378	0.540	0.427	0.397
Ego4D	Next O _t+1	0.370	0.247	0.385	0.232	0.284

Source: Table 1, supplied paper. Scores describe these examples rather than dataset-wide averages.

Adapt with 2k frame pairs

Small-data fine-tuning in a new environment.

Four robot examples comparing observed frames and predictions. The fine-tuned model produces clearer predictions than the unadapted pretrained model. — **A little domain data goes a long way.** Fine-tuning on 2,000 paired frames improves visual quality, with results qualitatively comparable to training on in-domain data from scratch.

W _pre Pretrained W _tfs From scratch W _ft Fine-tuned

Connect to environment actions

Aligning latent actions with robot controls.

Three Robosuite robot examples with enlarged gripper regions, comparing observations and predictions conditioned on latent actions or ground-truth actions. — **From latent actions to real controls.** An action adapter makes predictions conditional on environment-defined actions. Enlarged gripper regions reveal fine motion details.

R _LA Latent actions R _GT Ground-truth actions

Future work. The paper evaluates visual prediction and perceptual similarity. Integrating the model into a model-based reinforcement learning algorithm is left to future work.

BibTeX

If you find this research useful, please cite the paper.

Download paper

@inproceedings{cui2023universal,
  title={A universal world model learned from large scale and diverse videos},
  author={Cui, Hanchen and Gao, Yang},
  booktitle={NeurIPS 2023 Foundation Models for Decision Making Workshop},
  year={2023}
}