Spatiotemporal Predictive Pre-training for Robotic Motor Control

TL;DR

This work tackles the gap in robotic motor-control pre-training where static content priors predominate, proposing STP, a spatiotemporal predictive pre-training framework that uses asymmetric masking and decoupled dual decoders to learn both content and motion features from large-scale egocentric videos. By freezing the encoder for downstream behavior cloning and exploring post-pre-training and hybrid pre-training, STP demonstrates improved generalization to unseen, distractor-rich environments and better data efficiency across simulated and real-world tasks. Key findings include the superiority of STP over MAE and VC-1 baselines, the value of motion-focused pre-training, and nuanced insights into temporal-prediction conditioning and decoder architecture. Overall, STP offers a practical, scalable approach for robust, transferable vision-based robotic control.

Abstract

Robotic motor control necessitates the ability to predict the dynamics of environments and interaction objects. However, advanced self-supervised pre-trained visual representations in robotic motor control, leveraging large-scale egocentric videos, often focus solely on learning the static content features. This neglects the crucial temporal motion clues in human video, which implicitly contain key knowledge about interacting and manipulating with the environments and objects. In this paper, we present a simple yet effective robotic motor control visual pre-training framework that jointly performs spatiotemporal prediction with dual decoders, utilizing large-scale video data, termed as STP. STP adheres to two key designs in a multi-task learning manner. First, we perform spatial prediction on the masked current frame for learning content features. Second, we utilize the future frame with an extremely high masking ratio as a condition, based on the masked current frame, to conduct temporal prediction for capturing motion features. The asymmetric masking and decoupled dual decoders ensure that our image representation focusing on motion information while capturing spatial details. Extensive simulation and real-world experiments demonstrate the effectiveness and generalization abilities of STP, especially in generalizing to unseen environments with more distractors. Additionally, further post-pre-training and hybrid pre-training unleash its generality and data efficiency. Our code and weights will be released for further applications.

Figures (7)

• Figure 1: STP framework. Left: During pre-training, we sample the current frame and the future frame from the video clip, and carry out spatiotemporal predictive pre-training. Right: During extensive downstream motor control tasks evaluation, we freeze the pre-trained encoder to extract visual state representations and discard the decoders.
• Figure 2: Temporal decoder design. (a) Standard joint-self architecture. (b) The self-cross architecture.
• Figure 3: The real-world experiments setup and five evaluation demonstrations.
• Figure 4: Attention Visualization. We use the [CLS] token as query, average the attention of all heads at the last layer of the frozen ViT encoder, where the size of the attention value is directly proportional to the intensity of the yellow light. Top: MAE pre-training. Bottom: STP pre-training.
• Figure 5: Some examples of our STP prediction result on Ego4D videos. For each six tuple, we show the ground-truth (left), masked frames (middle), STP prediciton results (right), current frames (top), and future frames (bottom). We simply overlay the output with the visible patches to improve visual quality.