Tsinghua University, Huawei and others proposed iVideoGPT: specializing in interactive world models

iVideoGPT meets the high interactivity needs of the world model.

Generative models have made significant progress in recent years, among which video generation is becoming a new frontier. An important application of these generative video models is to learn in an unsupervised manner on diverse Internet-scale data for building predictive world models. These world models are expected to accumulate common-sense knowledge about how the world works, allowing predictions of potential future outcomes based on the behavior of agents.

By leveraging these world models, reinforcement learning-based agents can imagine, reason, and plan within the world model, allowing them to perform tasks in the real world with just a few experiments. Acquire new skills more safely and efficiently.

Despite the fundamental connection between generative models and world models, there are still significant differences between the development of generative models for video generation and world models for agent learning. the gap. One of the main challenges is how to achieve the best balance between interactivity and scalability.

#In the field of model-based reinforcement learning, world models mainly use recurrent network architecture. This design facilitates interactive behavioral learning by allowing observations or latent states to be passed based on actions at each step. However, these models mostly focus on game or simulation environments, have simple data, and have limited ability to model large-scale and complex in-the-wild data.

In contrast, Internet-scale video generation models can synthesize realistic long videos that can be controlled with textual descriptions or future action sequences. While such models allow for high-level, long-term planning, their trajectory-level interactivity does not provide agents with sufficient granularity to effectively learn precise behaviors as a fundamental skill.

Researchers from Tsinghua University, Huawei Noah's Ark Laboratory, and Tianjin University proposed iVideoGPT (Interactive VideoGPT), which is an scalable autoregressive Transformer framework. It integrates multi-modal signals (visual observations, actions and rewards) into a series of tokens, enabling the agent to conduct interactive experiences by predicting the next token.

iVideoGPT uses novel compression tokenization technology to effectively discrete high-dimensional visual observations. Leveraging its scalable architecture, researchers were able to pre-train iVideoGPT on millions of human and robot operation trajectories, thereby establishing a versatile foundation that can be used as an interactive world model for a variety of downstream tasks. This research promotes the development of interactive universal world models.

• Paper address: https://arxiv.org/pdf/2405.15223
• Paper title: iVideoGPT: Interactive VideoGPTs are Scalable World Models

Method

In this part, the research team introduces a scalable world model architecture - iVideoGPT, which is extremely flexible and can integrate multi-modal information, including visual observations, actions, rewards and other potential inputs.

#The core of iVideoGPT includes a compression tokenizer for discretizing video frames, and an autoregressive transformer for predicting subsequent tokens. By pre-training on diverse video data, the model can acquire extensive world knowledge and then efficiently transfer to downstream tasks.

Architecture

Compression tokenization. Transformer performs particularly well at handling sequences of discrete tokens. VQGAN is a commonly used visual tokenizer that converts raw pixels into discrete tokens. The researchers proposed to use a new conditional VQGAN consisting of a dual encoder and decoder {(E_c, D_c), (E_p, D_p)} to tokenize the video.

As shown in Figure 3a, the initial context frame contains rich context information and is tokenized and reconstructed independently through N tokens:

In contrast, due to the temporal redundancy between context frames and future frames, only necessary change information, such as the position and pose of moving objects, needs to be encoded. The above process is achieved by using conditional encoders and decoders:

#The researchers implemented the conditional mechanism by using cross-attention between multi-scale feature maps. In general, the tokenizer is trained with the following goals:

The tokenization proposed in this study mainly has two benefits:

• First, it significantly reduces the sequence length of the tokenized video, which grows linearly with the number of frames, but the growth rate n is much smaller;
• Secondly, through conditional encoding, the transformer that predicts subsequent tokens can more easily maintain the temporal consistency of the context and focus on modeling the necessary dynamic information.

Interactive predictions for Transformer. After tokenization, the video is flattened into a series of tokens:

length is . Special slot tokens [S] are inserted to delineate frame boundaries and facilitate the fusion of additional low-dimensional modalities such as actions. As shown in Figure 3b, a GPT-like autoregressive transformer is used for interactive video prediction by generating next-tokens frame by frame. In this work, the team used the model size of GPT-2 but adapted the LLaMA architecture in order to take advantage of recent innovations in LLM architectures, such as rotational position embedding.

Pre-training

Large language models can be self-supervised through next-word prediction way to gain extensive knowledge from Internet texts. Similarly, the action-free video pre-training paradigm of the world model uses video prediction as the pre-training goal to provide Internet-scale supervision for the physical world knowledge that LLM lacks.

The researchers pre-trained iVideoGPT on this general goal, applying cross-entropy loss to predict subsequent video tokens:

Pre-training data. Although there are a large number of videos available on the Internet, due to computational limitations, the researchers pretrained iVideoGPT specifically for the field of robotic manipulation. They utilized a mix of 35 datasets from the Open X-Embodiment (OXE) dataset and the Something-Something v2 (SSv2) dataset, totaling 1.5 million trajectories.

Fine-tuning

action conditions and reward prediction. The team’s architecture is designed to flexibly integrate additional modalities to learn an interactive world model, as shown in Figure 3b. Actions are integrated via linear projection and added to slot token embeddings. For reward prediction, instead of learning a separate reward predictor, they added a linear head on the hidden state of the last token of each observation.

This multi-task learning method can enhance the model's attention to task-related information, thereby improving the prediction accuracy of control tasks. In addition to the cross-entropy loss of Equation (3), they also used the mean square error loss for reward prediction.

Tokenizer Adapt. The research team chose to update the complete model, including the tokenizer, to adapt to downstream tasks, and found this strategy to be more effective than parameter-efficient fine-tuning methods.

#There is little literature exploring the use of VQGAN tokenizer for domain-specific data. In this work, since tokenization decouples dynamic information from contextual conditions, it is assumed that although this model may encounter unseen objects in downstream tasks, such as different types of robots, the transformer learns from diverse scenarios Basic physics knowledge - such as movement and interaction - is shared.

This hypothesis is supported by experiments where they migrated iVideoGPT from mixed pre-trained data to the unseen BAIR dataset, where the pre-trained transformer can zero-sample Generalizing to predict natural motion requires only fine-tuning the tokenizer for unseen robot grippers (see Figure 7). This feature is particularly important for scaling GPT-like transformers to large sizes, enabling lightweight alignment across domains while keeping the transformer intact.

Experiment

As shown in Table 1, compared with the SOTA method, iVideoGPT Demonstrates competitive performance while enabling interactivity and scalability in its architecture. Although preliminary experiments were performed at the low resolution of 64×64, iVideoGPT can be easily extended to RoboNet’s 256×256.

See Figure 9 for qualitative results.

Figure 4 shows the success rate of iVideoGPT compared to the baseline model. iVideoGPT significantly outperforms all baselines on both RoboDesk tasks and achieves average performance comparable to the strongest model SVG'.

Figure 6 shows that the model-based algorithm not only improves the sample efficiency than the model-free algorithm, but also meets or exceeds the performance of DreamerV3.

The next study analyzes the zero-shot video prediction capabilities of large-scale pre-trained iVideoGPT on the unseen BAIR dataset. Interestingly, we observe in the second row of Figure 7 that iVideoGPT predicts the natural motion of a robot gripper without fine-tuning—albeit differently from the pre-training dataset. This shows that although the model has limited zero-shot generalization to completely unseen robots due to insufficient diversity in the pre-training data, it effectively separates scene context from motion dynamics. In contrast, using the adapted tokenizer, the non-fine-tuned Transformer successfully transfers the pre-trained knowledge and predicts the motion of the novel robot in the third row, providing similar perception to the fully fine-tuned Transformer in the fourth row. The quality and quantitative results are shown in Figure 8a.

For more results, please refer to the original paper.

The above is the detailed content of Tsinghua University, Huawei and others proposed iVideoGPT: specializing in interactive world models. For more information, please follow other related articles on the PHP Chinese website!