Spatiotemporal Predictive Pre-training for Robotic Motor Control

C1: The high costs associated with evaluating real-world tasks using different random seeds make it challenging to report variances. However, assessing the impact of multiple random seeds in simulated tasks could provide more reliable statistical insights. As shown in Table 1, STP's performance improvement over baselines is marginal (STP 63.7 vs. VC-1 61.8, and STP-L/16(Post PT) 78.4 vs. MAE-L/16(Post PT) 76.7). Given the inherent stochastic nature of imitation learning and reinforcement learning, evaluations across multiple episodes and various random seeds are crucial to validate the proposed methods effectively.

R1: Thanks for your suggestions. To address your concern, we first report results with multiple random seeds, then clarify some confusion points in Tab 1, and finally report extra experiment results with more evaluation.

(1) We rerun the experiments and report the mean and standard deviation, demonstrating the stability of our experiments.

For the same BC seed, our policy training is fully reproducible, with the only uncertainty being the slight difference that still exist in MuJoCo rendering even with the same policy and evaluation seed. Therefore, we further rerun our paper's STP-B and MAE-B baseline twice during limited time, obtaining three results and their mean and variance, as shown below. These demonstrate the stability of our experiments. In addition, the rendering in Trifinger is not subject to randomness, hence the results are fully reproducible.

(2) On confusion of Tab 1.: we use the same pre-training data to conduct a fair comparison between MAE and STP.

Our STP (EgoClip) utilize less pre-training data than the publicly VC-1 (full Ego4D+MNI). Under the same pre-training data and number of training epochs, our STP-B outperforms MAE-B baseline (63.7 vs. 59.6), which shows a significant improvement of 4.1. Note that VC-1 shares the same technique with MAE only with different pre-training data.

(3) About the performance improvement of ViT-Large (PT).

We analyze that under the post-pre-training and ViT-Large setting, the representation might be overfitted to the target domain, and the performance improvement tends to saturate. In the future, we may resort to larger pre-training data.

(4) Our initial evaluation scale was enormous, hence the performance improvement is solid and stable.

The Number of evaluation episodes = Number of tasks × Number of BC seeds × Number of evaluation seeds × Number of camera views. Therefore, the total number of our evaluation episodes is 5×3×25×1 + 5×3×50×2 + 5×3×25×1 + 2×3×25×1 + 2×1×25×1 = 2450.

(5) We add the evaluation results from RLbench, demonstrating the generality of our improvements.

In addition, we have added a performance comparison between STP and VC-1 on 20 randomly selected tasks in RLBench. We use the advanced RVT [6] as the policy module, training a multi-task policy with only 20 demonstrations per task. Each task is evaluated over 25 episodes, and the results are in the PDF file; our STP (52.0) outperforms VC-1 (43.2) by 8.8, further proving the effectiveness of our STP.

[6] Rvt: Robotic view transformer for 3d object manipulation.

C2: The authors could strengthen their paper by highlighting these related works, demonstrating awareness of existing methods, situating their contributions and highlighting the differences between STP and these baselines.

R2: Thank you for your suggestion. We will further strengthen the discussion of these related works and highlight our STP's differences in the revised paper.

Our STP has distinct differences from these works that consider the temporal movements in terms of objectives and techniques. R3M, LIV, and DecisionNCE capture the alignment between language instructions and task progression through contrastive learning for representation learning or reward learning. VIP focuses on self-supervised contrastive learning to pretrain a reward function. Voltron performs language-guided reconstruction for multimodal video representation.

In contrast, our STP is a purely self-supervised method through mask-based generative modeling for spatiotemporal prediciton, using asymmetric masking and decoder architecture to jointly capture content and motion features. Our method does not require language descriptions. Our generative pre-training paradigm differentiates STP to these contrastive methods from the technique perspective.

In summary, our STP has the following advantages: It is more general (learning standard image representation instead of focusing on a specific application), scalable (self-supervised learning and plain ViT), efficient (high masking ration), and effective (asymmetric masking and decoder architecture design ensure joint learning of content and motion features).

I agree with the authors that this paper conducts more comprehensive ablations and evaluations to support the effectiveness of each design choices. In this sense, I am very happy to increase my score. Meanwhile, considering the similarity to Voltron (the authors provide more detailed techinique insights about the differences to voltron, but the high-level differences are like what I pointed out). I decide to increase to a 5 (boardline accept).

C1: It is unclear where the diverse image data for STP trained with Ego+I in Table 1 is obtained from.

R1: Sorry for the confusion. "STP trained with Ego+I" means that we perform a hybrid pre-training using EgoClip and ImageNet data. Specifically, we first initialize ViT with the ImageNet-MAE weight. During the pre-training process, for the image data from ImageNet, we conduct MAE pre-training; for the video data from EgoClip, we conduct STP pre-training. This resultes in a 0.5 performance improvement, indicating that our STP can also benefit from more diverse image data.

C2: The real-world experiments seem limited with only two real-world tasks where the MAE also performs reasonably well.

R2: Our STP has achieved significant advantages in the pouring task (45% -> 65%). It can more accurately align with the moving bowl and the pot. In addition, although MAE and STP have a same success rate in picking tasks, STP tends to execute grasping in a better position. Some works such as [2] have demonstrated that real-world environments and simulation settings yield similar conclusions, hence we have not carried out more and costly real-world evaluations. We will release our STP weight in the future, hoping that it could contribute to the community and be applied to more real-world environments and tasks.

[2] What do we learn from a large-scale study of pre-trained visual representations in sim and real environments?

C3: The authors must include comparisons with prior works using MAE for spatiotemporal learning [1].

R3: Thanks for your suggestion. Our STP differs from MAE-ST as STP executes asymmetric masking and decoder architecture design for decoupled spatial and temporal prediction on the 2D image model. On the contrary, MAE-ST jointly performs spatiotemporal reconstruction to pre-train the 3D video model, processing the temporal dimension and spatial dimension symmetrically.

We pre-trained a 4-frame MAE-ST based on the EgoClip dataset, and the results are shown below. We believe that the poorer performance of MAE-ST is due to the gap in temporal interaction between upstream and diverse downstream environments, which lead to a significant risk of cumulative error in the paradigm of imitation learning.

C1: This paper should highlight its differences with R3M, VIP, and V-PTR.

R1: Thank you for your suggestion. We will highlight these differences in the revised paper. Our STP has distinct differences from these works in objectives and techniques.

VIP and V-PTR respectively pre-train the value function through contrastive learning and TD learning, focusing on visual reward functions and value function in RL. R3M pre-trains image representation using time-contrastive learning and video-language alignment.

In contrast, our STP performs self-supervised, masking-based generative modeling, which is more efficient (high masking ratio), and scalable (self-supervised learning and simple backbone of ViT allows for large-scale pre-training). Our generative pre-training paradigm differentiates us to these contrastive methods. The superior performance of STP demonstrates the advantage of generative pre-training against them.

C2: The performance gap is not significant. The slight performance difference may be due to hyperparameter selection and randomness, as the paper did not provide error bars over multiple seeds.

R2: First, we clarify some confusions on the result in Tab 1.

(1) In a fair comparison, our improvement is significant.

Our STP (EgoClip) utilize less pre-training data than the VC-1 (full Ego4D+MNI). Under the same pre-training data and the number of training epochs, our STP-B outperforms MAE-B baseline (63.7 vs. 59.6), which shows a significant improvement of 4.1. Note that VC-1 shares the same technique with MAE only with different pre-training data.

(2) Our initial evaluation scale was enormous, hence the performance improvement is stable.

The Number of evaluation episodes = Number of tasks × Number of BC seeds × Number of evaluation seeds × Number of camera views. The total number of our evaluation episodes is 5×3×25×1 + 5×3×50×2 + 5×3×25×1 + 2×3×25×1 + 2×1×25×1 = 2450.

Second, we rerun the experiments with multiple seeds and report the mean and standard deviation.

During pre-training and BC, we did not deliberately select hyperparameters. The pre-training hyperparameters of STP follow MAE, and all representations adhere to the same setting during BC. For the same BC seed, our policy training is fully reproducible, with the only uncertainty being the slight difference that still exist in MuJoCo rendering even with the same policy and evaluation seed. Therefore, we further rerun our paper's STP-B and MAE-B baseline twice during rebuttal, obtaining three results and their mean and variance, as shown below. These demonstrate the stability of our STP. Additionally, the rendering in Trifinger is not subject to randomness, hence the results are fully reproducible.

Finally, we add extra evaluation results from RLbench, demonstrating the generality of our improvements.

In addition, we added a performance comparison between STP-B and VC-1 base on 20 randomly selected tasks in RLBench. We use the advanced RVT [5] as the policy, training a multi-task policy with only 20 demonstrations per task. Each task is evaluated over 25 episodes, and the results are in the PDF file; our STP (52.0) outperforms VC-1 (43.2) by 8.8, further proving the effectiveness of our STP.

C3: VIP and V-PTR should be included as baselines.

R3: Thanks for your suggestions. Since V-PTR has not released weight, we report the results for VIP and LIV [4] as baseline comparison. It is worth noting that both VIP and LIV utilize ResNet50 as their backbone, making a direct comparison with ViT-B is unfair. Additionally, the difference in feature dimensions between ResNet50 and ViT-B (2048 vs. 768) results in a discrepancy in the number of trainable MLPs policy parameters. Therefore, to enable a a more fair comparison, we construct the policy with an equivalent number of parameters for STP-B, which we refer to as STP†, and the results demonstrate the superior performance of our STP.

C4: What is the evaluation protocol for downstream tasks? Does all evaluation use an expert dataset and perform imitation learning to learn a policy?

R4: We follow a single-task setup due to without task condiiton. Yes, for each task, we learn a single policy based on the representation. In extra RLbench experiments, we use the multi-task setup.

C5: How did you collect the dataset for real-world experiments?

R5: As shown in A.5, following the approach in [6], we collect robot data using a VR tele-operation setup.

[4] LIV: Language-Image Representations and Rewards for Robotic Control. [5] Rvt: Robotic view transformer for 3d object manipulation [6] Openvr: Teleoperation for manipulation.

C1: The major concern is the novelty of the previous methods, considering several related papers that leverage human data and visuals pretraining to downstream tasks have been proposed [1-3].

R1: Thanks for your comments. Our STP exhibits some essential differences or advantages with these works as follows.

As for [2] and [3], our STP has different motivation and techniques. Both [2] and [3] perform a history-aware policy pre-training through language-driven video prediction, where [2] uses VQ-VAE and video diffusion techniques, while [3] employs auto-regressive GPT-style techniques based on frozen ViT representation. In contrast, our STP performs image representation pre-training through joint spatiotemporal prediction in a self-supervised manner, using asymmetric masking and decoder architecture design. Representation pre-training is orthogonal to these two methods.

As for [1], it pre-trains visual representations by predicting the transition frame and detecting the interaction object. It requires language and bounding box annotations, using only 93K video clips for pre-training. Instead, our STP is a purely self-supervised method without language or object annotations. Our STP is a much simpler design without the multi-frame causality modeling, multimodal token aggregator, and multiheaded attention pooling from [1], and only uses only ViT backbone. The following fair comparison (both use ViT-B) indicate that our STP achieves stronger performance, thanks to the scalability (masking and self-supervised learning enable lager scale data) of our proxy task, the asymmetric masking strategy, and specific decoder architecture design for spatiotemporal prediction.

In summary, our STP makes orthogonal contributions to [2] and [3]. Compared to [1], it has the following advantages: It is more general (learning standard ViT representation), scalable (self-supervised learning), efficient (high masking ration), and effective (asymmetric masking and decoder architecture design for spatiotemporal prediction). Finally, thank you for your reminder, and we will strengthen these distinctions in the revised paper.

C2: The experiment only contains imitation learning experiments in downstream tasks, while the reinforcement learning framework with sub-optimal data is not considered.

R2: Thanks for your suggestion. In rebuttal, we select the Panda-Door and Panda-TwoArmPegInHole tasks from the Robosuite [4] simulation environment for reinforcement learning evaluation. We employ DrQ-v2 as our RL algorithm and compare the results of VC-1 (ViT-B) and our STP (ViT-B) as frozen visual representations. Due to the large fluctuations in success rate, we report the maximum reward value under 200,000 steps, and the results preliminarily verify the effectiveness of our STP within the RL framework.

[4] A modular simulation framework and benchmark for robot learning.