DreamX: Video Diffusion Model is a Universal Latent Action Learner

DreamX: Video Diffusion Model as Universal Latent Action Learner§

DreamX introduces a method to learn a universal action representation from video by repurposing a pretrained video diffusion model. The key insight is that video diffusion models implicitly capture spatiotemporal dynamics that are consistent across different executions of the same action but invariant to embodiment and appearance changes. By extracting compact latent actions from these models, DreamX achieves state-of-the-art performance on latent action benchmarks and enables zero-shot action transfer across embodiments.

Core Methodology§

The approach leverages a pretrained video diffusion model (e.g., a latent video diffusion model) to encode video frames into a compact latent space. Specifically, each latent video frame is compressed by a factor of 300× (from 1,200 tokens to 4 tokens) using the model’s encoder. This yields a sequence of latent codes that represent the essential dynamics of the video.

Given a video $x_{1:T}$, the encoder $\mathcal{E}$ produces a latent sequence $z_{1:T} = \mathcal{E}(x_{1:T})$, where each $z_t \in \mathbb{R}^{4}$. The diffusion model’s generative prior ensures that these latents capture action-relevant information while discarding visual nuisances (e.g., background, object appearance). The latent actions are then defined as the differences between consecutive latents: $a_t = z_{t+1} - z_t$. These latent actions serve as a shared action interface across embodiments.

Training and Implementation§

DreamX is trained in two stages:

1. Video Diffusion Pretraining: A video diffusion model is pretrained on a large-scale video dataset (e.g., something like WebVid or Kinetics) to learn spatiotemporal features.
2. Latent Action Extraction: The encoder of the pretrained model is frozen, and latent actions are extracted from videos. Optionally, a lightweight action decoder can be trained to map latent actions back to video frames for reconstruction, ensuring the latents retain sufficient information.

Importantly, the method does not require action labels; it is self-supervised. The latent space is then used for downstream tasks: zero-shot action transfer (e.g., applying a latent action from a human video to a robot) and policy learning (e.g., learning a policy in the compact latent action space, reducing computational cost compared to pixel-level diffusion policies).

Experimental Results§

On LARY-Bench, DreamX achieves state-of-the-art performance among latent action models. It supports any-to-any video transfer and zero-shot action transfer across diverse embodiments (human and robot) and environments. Furthermore, policy learning in the latent action space reduces computational cost while maintaining effectiveness.

Code Illustration§

Below is a simplified PyTorch snippet illustrating the latent action extraction and policy learning pipeline:

import torch
import torch.nn as nn

class LatentActionExtractor(nn.Module):
    def __init__(self, video_diffusion_model):
        super().__init__()
        self.encoder = video_diffusion_model.encoder  # pretrained, frozen
        self.compress = nn.Linear(1200, 4)  # compress 1200 tokens to 4 tokens per frame

    def forward(self, video):
        # video: (B, T, C, H, W)
        B, T, C, H, W = video.shape
        # encode each frame independently (or use temporal encoder)
        frames = video.view(B*T, C, H, W)
        latents = self.encoder(frames)  # (B*T, 1200)
        latents = self.compress(latents)  # (B*T, 4)
        latents = latents.view(B, T, 4)
        # compute latent actions as differences
        actions = latents[:, 1:] - latents[:, :-1]  # (B, T-1, 4)
        return actions

# Policy learning in latent action space
class LatentPolicy(nn.Module):
    def __init__(self, obs_dim, action_dim=4):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_dim, 256),
            nn.ReLU(),
            nn.Linear(256, action_dim)
        )
    def forward(self, obs):
        return self.net(obs)

In practice, the encoder is a 3D convolutional or transformer-based video encoder from a diffusion model, and the compression may be implemented via pooling or a learned projection. The policy is trained with behavior cloning or reinforcement learning in the latent action space.

Conclusion§

DreamX demonstrates that video diffusion models can serve as powerful latent action learners, providing a universal action representation that generalizes across embodiments and enables efficient policy learning. This work bridges generative video modeling and embodied AI, offering a scalable path to action understanding from unlabeled video.