M ASK V I T: M ASKED V ISUAL P RE -T RAINING FOR V IDEO P REDICTION
Anonymous authors
Paper under double-blind review
A BSTRACT
The ability to predict future visual observations conditioned on past observationsand motor commands can enable embodied agents to plan solutions to a variety oftasks in complex environments.
This work shows that we can create good videoprediction models by pre-training transformers via masked visual modeling.
Ourapproach, named MaskViT, is based on two simple design decisions.
First, formemory and training efficiency, we use two types of window attention: spatialand spatiotemporal.
Second, during training, we mask a variable percentage oftokens instead of a fixed mask ratio.
For inference, MaskViT generates all tokens via iterative refinement where we incrementally decrease the masking ratiofollowing a mask scheduling function.
On several datasets we demonstrate that
MaskViT outperforms prior works in video prediction, is parameter efficient, generates high-resolution videos ( 256 × 256 ) and can be easily adapted to performgoal-conditioned video prediction.
Further, we demonstrate the benefits of inference speedup (up to 512 × ) due to iterative decoding by using MaskViT forplanning on a real robot.
Our work suggests that we can endow embodied agentswith powerful predictive models by leveraging the general framework of maskedvisual modeling with minimal domain knowledge.
I NTRODUCTION
Evidence from neuroscience suggests that human cognitive and perceptual capabilities are supportedby a predictive mechanism to anticipate future events and sensory signals (Tanji & Evarts, 1976;
Wolpert et al., 1995).
Such a mental model of the world can be used to simulate, evaluate, and selectamong different possible actions.
This process is fast and accurate, even under the computationallimitations of biological brains (Wu et al., 2016).
Endowing robots with similar predictive capabilities would allow them to plan solutions to multiple tasks in complex and dynamic environments,e.g., via visual model-predictive control (Finn & Levine, 2017; Ebert et al., 2018).
Predicting visual observations for embodied agents is however challenging and computationallydemanding: the model needs to capture the complexity and inherent stochasticity of future eventswhile maintaining an inference speed that supports the robot’s actions.
Therefore, recent advances inautoregressive generative models, which leverage Transformers (Vaswani et al., 2017) for buildingneural architectures and learn good representations via self-supervised generative pretraining (Devlin et al., 2019), have not benefited video prediction or robotic applications.
We in particularidentify three technical challenges.
First, memory requirements for the full attention mechanismin
Transformers scale quadratically with the length of the input sequence, leading to prohibitivelylarge costs for videos.
Second, there is an inconsistency between the video prediction task and autoregressive masked visual pretraining – while the training process assumes partial knowledge ofthe ground truth future frames, at test time the model has to predict a complete sequence of futureframes from scratch , leading to poor video prediction quality (Yan et al., 2021; Feichtenhofer et al.,2022).
Third, the common autoregressive paradigm effective in other domains would be too slowfor robotic applications.
To address these challenges, we present Mask ed
Vi deo T ransformers (MaskViT): a simple, effectiveand scalable method for video prediction based on masked visual modeling.
Since using pixelsdirectly as frame tokens would require an inordinate amount of memory, we use a discrete variationalautoencoder (dVAE) (Van Den Oord & Vinyals, 2017; Esser et al., 2021) that compresses frames
VQGAN
Iteration Iteration
Iteration… …
t = T - 1 t = 1 t = t = T - 1 t = 1 t =t
=
T - 1 t = 1 t =(a) Training (b) Inference x L
Latent token
Masked latent token 3D window attention
Window T x (h’ // 4) x (w’ // 4) Window 1 x h’ x w’ (context) t
= T -1 t = 1 t =(context) t = T - 1 t = 1 t =(context) t
=
T - 1 t = 1 t =(context) (context) (context)
Bidirectional Window Transformer
A variable number of tokens in future frames are masked, and the network is trained to predict the masked tokens.
A block in MaskViT consists of two layers with window-restricted attention: spatial and spatiotemporal
Inference
Videos are generated via iterative refinement where we incrementally decrease the masking ratio following a mask scheduling function page into a smaller grid of visual tokens.
We opt for compression in the spatial (image) domain insteadof the spatiotemporal domain (videos), as preserving the correspondence between each original andtokenized video frame allows for flexible conditioning on any subset of frames – initial (past), final(goal), and possibly equally spaced intermediate frames.
However, despite operating on tokens,representing 16 frames at 256 tokens per frame still requires 4 , 096 tokens, incurring prohibitivememory requirements for full attention.
Hence, to further reduce memory, MaskViT is composed ofalternating transformer layers with non-overlapping window-restricted (Vaswani et al., 2017) spatialand spatiotemporal attention.
To reduce the inconsistency between the masked pretraining and the video prediction task and tospeed up inference, we take inspiration from non-autoregressive, iterative decoding methods in generative algorithms from other domains (Sohl-Dickstein et al., 2015;
Ho et al., 2020; Nichol & Dhariwal, 2021; Ghazvininejad et al., 2019; Chang et al., 2022).
We propose a novel iterative decodingscheme for videos based on a mask scheduling function that specifies, during inference, the numberof tokens to be decoded and kept at each iteration.
In contrast to autoregressive decoding, whichinvolves predicting tokens one by one, our iterative decoding scheme is faster as the number ofdecoding iterations is significantly less than the number of tokens.
A few initial tokens are predicted over multiple initial iterations, and then the majority of the remaining tokens can be predictedrapidly over the final few iterations.
This brings us closer to the ultimate video prediction task,where only the first frame is known and all tokens for other frames must be inferred.
To furtherclose the training-test gap, during training we mask a variable percentage of tokens, instead of usinga fixed masking ratio.
This simulates the different masking ratios MaskViT will encounter duringiterative decoding in the actual video prediction task.
Through experiments on several publicly available real-world video prediction datasets (Ebert et al.,2017; Geiger et al., 2013; Dasari et al., 2019), we demonstrate that MaskViT achieves competitive orstate-of-the-art results in a variety of metrics.
Moreover, MaskViT can predict considerably higherresolution videos ( 256 × 256 ) than previous methods.
We also show the flexibility of MaskViTby adapting it to predict goal-conditioned video frames.
In addition, thanks to iterative decoding,
MaskViT is up to 512 × faster than autoregressive methods, enabling its application for planningon a real robot (§ 4.5).
These results indicate that we can endow embodied agents with powerfulpredictive models by leveraging the advances in self-supervised learning in language and vision,without engineering domain-specific solutions.
R ELATED W ORK
Video prediction.
The video prediction task refers to the problem of generating videos conditionedon past frames (Ranzato et al., 2014; Lotter et al., 2016), possibly with an additional natural language description (Li et al., 2018;
Gupta et al., 2018; Pan et al., 2017; Wu et al., 2021b) and/ormotor commands (Finn et al., 2016; Villegas et al., 2019; Wu et al., 2021a; Babaeizadeh et al.,2021).
Multiple classes of generative models have been utilized to tackle this problem, such as Generative adversarial networks (GANs) (Clark et al., 2019; Tulyakov et al., 2018; Luc et al., 2020),
Variational Autoencoders (VAEs) (Villegas et al., 2019; Wu et al., 2021a; Babaeizadeh et al., 2021;2018; Denton & Fergus, 2018; Akan et al., 2021; 2022), invertible networks (Dorkenwald et al.,2021), autoregressive (Yan et al., 2021; Rakhimov et al., 2020; Nash et al., 2022) and diffusion (Hoet al., 2022; Voleti et al., 2022) models.
Our work focuses on predicting future frames conditionedon past frames or motor commands and belongs to the family of two-stage methods that first encodethe videos into a downsampled latent space and then use transformers to model an autoregressiveprior (Yan et al., 2021;
Rakhimov et al., 2020).
A common drawback of these methods is the largeinference time due to autoregressive generation.
MaskViT overcomes this issue by using an iterativedecoding scheme, which significantly reduces inference time.
Masked autoencoders.
Masked autoencoders are a type of denoising autoencoder (Vincent et al.,2008) that learn representations by (re)generating the original input from corrupted (i.e., masked)inputs.
Masked language modeling (MLM) was first proposed in BERT (Devlin et al., 2019) andhas revolutionized the field of natural language processing, especially when scaled to large datasetsand model sizes (Brown et al., 2020; Radford et al., 2019).
MLM with fixed and low masking ratiohas also been extended for multi-modal representation learning (Lu et al., 2019; Sun et al., 2019;
Zellers et al., 2021).
The success in NLP has also been replicated in vision by masking a fixedbut high ratio of patches of pixels (He et al., 2021;
Dosovitskiy et al., 2020) or masking tokensgenerated by a pretrained dVAE (Bao et al., 2022; Chen et al., 2020).
Recently, these works havealso been extended to video domains to learn good representations for action recognition (Tonget al., 2022; Feichtenhofer et al., 2022).
Unlike them, we apply masked visual modeling for videoprediction, and we use a variable masking ratio during training to reduce the difference betweenmasked pretraining and video prediction.
Another related line of work is leveraging good visualrepresentations learnt via self supervised learning methods (Laskin et al., 2020; Nair et al., 2022;
Parisi et al., 2022) including masked autoencoders (Xiao et al., 2022) for motor control.
Non-autoregressive decoding.
A key limitation of autoregressive decoding is that each token ispredicted sequentially.
Hence, the decoding time is proportional to the dimensionality of the data.
An appealing alternative is iterative decoding, where all tokens are predicted simultaneously andthen refined for a fixed number of steps independent of the data dimensionality (Gu et al., 2017;
Ghazvininejad et al., 2019; Saharia et al., 2020).
Ghazvininejad et al.
(2019) proposed a maskpredict decoding paradigm for machine translation by training a bidirectional decoder instead ofa causal decoder.
During inference, entire translated sequences are predicted in parallel withina constant number of steps.
Chen et al.
(2020) extend this framework to the image domain anduse mask scheduling functions during training and decoding.
We extend the framework of maskpredict (Ghazvininejad et al., 2019) and MaskGiT (Chen et al., 2020) for video prediction.
Wepre-train a bidirectional window transformer model via masked visual modeling where we maska variable percentage of tokens independent of the mask scheduling function.
Importantly, unlike
MaskGiT, instead of keeping the most confident tokens during each decoding step, we add temperature annealed Gumbel noise to the token confidence to produce diverse outputs.
M ASK V I T: M ASKED V IDEO T RANSFORMER
MaskViT is the result of a two-stage training procedure (Van Den Oord & Vinyals, 2017; Razaviet al., 2019):
First, we learn an encoding of the visual data that discretizes images into tokensbased on a discrete variational autoencoder (dVAE).
Next, we deviate from the common autoregressive training objective and pre-train a bidirectional transformer with window-restricted attention viamasked visual modeling (MVM).
In the following section, we describe our image tokenizer, bidirectional transformer, masked visual pre-training, and iterative decoding procedure.
Ground truth
Predicted
KITTI t =(context) t
=t =
t =t = t =t =
RoboNet
Ground truth
Predicted (context) t
=t =
t =t = t =t =
t =
3. L EARNING V ISUAL T OKENS
Videos contain too many pixels to be used directly as tokens in a transformer architecture.
Hence,to reduce dimensionality, we first train a VQ-VAE (Van Den Oord & Vinyals, 2017) for individualvideo frames so that we can represent videos as sequences of grids of discrete tokens.
VQ-VAE consists of an encoder E ( x ) that encodes an input image x ∈ R H
× W × 3 into a series of latent vectors.
The vectors are discretized through a nearest neighbour look up in a codebook of quantized embeddings, Z = { z k } Kk =1 ⊂ R n z .
A decoder D is trained to predict a reconstruction of the image, ˆ x , fromthe quantized encodings.
In our work, we leverage VQ-GAN (Esser et al., 2021), which improvesupon VQ-VAE by adding adversarial (Goodfellow et al., 2014) and perceptual losses (Johnson et al.,2016; Zhang et al., 2018).
Each video frame is individually tokenized into a 16 × 16 grid of tokens,regardless of their original resolution (Fig. 1, a, left).
Instead of using 3D extensions of VQ-VAEwhich perform spatiotemporal compression of videos (Yan et al., 2021), our per-frame compressionenables us to condition on arbitrary context frames: initial, final, and possibly intermediate ones.
[Figure]
Inspired by the success of masked language (Devlin et al., 2019) and image
(Bao et al., 2022; Heet al., 2021) modeling, and in the spirit of unifying methodologies across domains, we pre-train
MaskViT via MVM for video prediction.
Our pre-training task and masking strategy are straightforward: we keep the latent codes corresponding to context frames intact and mask a random numberof tokens corresponding to future frames.
The network is trained to predict masked latent codesconditioned on the unmasked latent codes.
Concretely, we assume access to input context frames for T c time steps, and our goal is to predict
T p frames during test time.
We first quantize the entire video sequence into latent codes Z ∈
R T × h × w .
Let Z p =
[ z i ]
Ni =1 denote the latent tokens corresponding to future video frames, where
N = T p × h × w .
Unlike prior work on MVM (Bao et al., 2022; He et al., 2021) that uses a fixedmasking ratio, we propose to use a variable masking ratio that reduces the gap between pre-trainingtask and inference leading to better evaluation results (see § 3.4).
Specifically, during training, foreach video in a batch, we first select a masking ratio r ∈ [0 . 5 , 1) and then randomly select andreplace ⌊ r · N ⌋ tokens in Z p with a [MASK] token.
The pre-training objective is to minimize thenegative log-likelihood of the visual tokens given the masked video as input:
[Equation]
where D is the training dataset, N M represents randomly masked positions, and Z Mp denotes theoutput of applying the mask to Z p , and Z c are latent tokens corresponding to context frames.
The
MVM training objective is different from the causal autoregressive training objective: L AR =− E x ∈D (cid:104) (cid:80) ∀ i ∈
Z p log p ( z i | z
j<i , Z c ) (cid:105) .
The key difference is that for MVM the conditional dependence is bidirectional : all masked tokens are predicted conditioned on all tokens.
3. B IDIRECTIONAL W INDOW T RANSFORMER
Transformer models composed entirely of global self-attention modules incur significant computeand memory costs, especially for video tasks.
To achieve more efficient modeling, we propose tocompute self-attention in windows, based on two types of non-overlapping configurations: 1) Spatial
Window (SW): attention is restricted to all the tokens within a subframe of size 1 × h × w (the firstdimension is time); 2) Spatiotemporal Window (STW): attention is restricted within a 3D windowof size T × h ′ × w ′ .
We sequentially stack the two types of window configurations to gain both localand global interactions in a single block (Fig. 1, a, center) that we repeat L times.
Surprisingly, wefind that a small window size of h ′ =
w ′ = 4 is sufficient to learn a good video prediction modelwhile significantly reducing memory requirements (Table 2b).
Note that our proposed block enjoysglobal interaction capabilities without requiring padding or cyclic-shifting like prior works (Liuet al., 2021a;b), nor developing custom CUDA kernels for sparse attention (Child et al., 2019) asboth window configurations can be instantiated via simple tensor reshaping.
3. I TERATIVE D ECODING
Decoding tokens autoregressively during inference is time-consuming, as the process scales linearlywith the number of tokens, and this can be prohibitively large (e.g., 4 , 096 for a video with 16frames and 256 tokens per frame).
Our video prediction training task allows us to predict futurevideo frames via a novel iterative non-autoregressive decoding scheme: inspired by the forwarddiffusion process in diffusion models (Ho et al., 2020; Nichol & Dhariwal, 2021) and the iterativedecoding in generative models e.g., mask-predict Ghazvininejad et al. (2019) and MaskGiT Changet al.
, we predict videos in T steps where T << N , the total number of tokens to predict.
Concretely, let γ ( t ) , where t ∈ { 0 T , 1 T , . .
, T − 1 T } , be a mask scheduling function (Fig. 3) thatcomputes the mask ratio for tokens as a function of the decoding steps.
We choose γ ( t ) such that itis monotonically decreasing with respect to t , and it holds that γ (0)
→ 1 and γ (1) → 0 to ensurethat our method converges.
At t = 0 , we start with Z =
[ Z c , Z p ] where all the tokens in Z pare [MASK] tokens.
At each decoding iteration, we predict all the tokens conditioned on all thepreviously predicted tokens.
For the next iteration, we mask out n = ⌈ γ ( tT ) N ⌉ tokens by keepingall the previously predicted tokens and the most confident token predictions in the current decodingstep.
We use the softmax probability as our confidence measure.
Empirically we found that selecting the most confident tokens while performing iterative decodinggenerates videos with little or no motion.
Intuitively, given only a few ( 1 − 2 ) context frames,selecting the most confident tokens results in videos where the context frames are copied across alltime steps.
Hence, we add temperature annealed Gumbel noise to the token confidence to encouragethe model to produce more diverse outputs.
Concretely, let C t ∈ R
N be the vector of softmaxprobabilities of the sampled tokens Z tp at decoding iteration t .
We select the most confident tokensto keep at iteration t from C tg where,
[Equation]
E XPERIMENTAL E VALUATION
In this section, we evaluate our method on three different datasets and compare its performancewith prior state-of-the-art methods, using four different metrics.
We also perform extensive ablationstudies of different design choices, and showcase that the speed improvements due to iterative decoding enable real-time planning for robotic manipulation tasks.
For qualitative results, see § B.2and videos on our project website.
[Table]
4. E XPERIMENTAL S ETUP
Implementation.
Our transformer model is a stack of L blocks, where each block consists of twotransformer layers with attention restricted to the window size of 1 × 16 × 16 (spatial window)and T × 4 × 4 (spatiotemporal window), unless otherwise specified.
We use learnable positionalembeddings, which are the sum of space and time positional embeddings.
See § A.1 for architecturedetails and hyperparameters.
Metrics.
We use four evaluation metrics to compare our method with prior work:
Fr ´ echet Video
Distance (FVD) (Unterthiner et al., 2018), Peak Signal-to-noise Ratio (PSNR), Structural Similarity Index Measure (SSIM) (Wang et al., 2004) and Learned Perceptual Image Patch Similarity(LPIPS)
(Zhang et al., 2018).
To account for the stochastic nature of video prediction, we followprior work (Babaeizadeh et al., 2021; Villegas et al., 2019) and report the best SSIM, PSNR, and
LPIPS scores over 100 trials for each video.
For FVD, we use all 100 with a batch size of 256.
Weonly conducted 1 trial per video for evaluating performance on the BAIR dataset (Ebert et al., 2017).
4. C OMPARISON WITH P RIOR W ORK
BAIR.
We first evaluate our model on the BAIR robot pushing dataset (Ebert et al., 2017), one ofthe most studied video modeling datasets (Fig. 7).
We follow the evaluation protocol of prior worksand predict 15 video frames given 1 context frame and no actions.
The lack of action conditioningmakes this task extremely challenging and tests the model’s ability to predict plausible future robottrajectories and object interactions.
MaskViT achieves similar performance to FitVid (Babaeizadehet al., 2021) while being more parameter efficient, and it outperforms all other prior works (Table 1).
Finally, to predict action-conditioned future frames, we linearly project the action vectors and addthem to Z .
As expected, action conditioning performs the best, with 25% improvement in FVD.
KITTI.
The KITTI dataset (Geiger et al., 2013) is a relatively small dataset of 57 training videos.
Wefollow the evaluation protocol of prior work (Villegas et al., 2019) and predict 25 video frames given5 context frames.
Compared to other datasets in our evaluation, KITTI is especially challenging,as it involves dynamic backgrounds, limited training data, and long-horizon predictions.
We usecolor jitter and random cropping data augmentation for training VQ-GAN and do not use any dataaugmentation for training the second stage.
Across all metrics, we find that MaskViT is significantlybetter than prior works while using fewer parameters (Table 1).
Training a transformer model withfull self-attention would require prohibitively large GPU memory due to the long prediction horizon( 30 × 16 × 16 = 7680 ).
However, MaskViT can attend to all tokens because its spatiotemporalwindows significantly reduce the size of the attention context.
We also report video predictionresults for the KITTI dataset at 256 × 256 resolution
Fig. 2, 8), a higher resolution that prior workwas not able to obtain.
RoboNet.
RoboNet (Dasari et al., 2019) is a large dataset of 15 million video frames of 7 differentrobotic arms interacting with objects and provides 5 dimensional robot action annotations.
Wefollow the evaluation protocol of prior work (Babaeizadeh et al., 2021) and predict 10 video framesgiven 2 context frames and future actions.
At 64 × 64 resolution, MaskViT is competitive but does
[Table]
efficient
[Table]
not outperform prior works (Table 1).
FVD of the VQ-GAN reconstructions is a lower bound for
MaskViT.
We found flicker artifacts in the VQ-GAN reconstructions, probably due to our use ofper-frame latents, resulting in a high FVD score of 121 for the VQ-GAN reconstructions.
MaskViTachieves FVD scores very close to this lower bound but performs worse than prior works due totemporally inconsistent VQ-GAN reconstructions.
Finally, we also report video prediction resultsfor the RoboNet dataset at 256 × 256 resolution
Fig. 2, 9).
4. F LEXIBLE C ONDITIONING
Planning for long horizon tasks via Visual MPC can quickly become computationally intractableas the search complexity grows exponentially with the number of planning steps.
Goal conditionedplanning can be especially helpful in reducing the search space by breaking down tasks into subgoals (Nasiriany et al., 2019;
Pertsch et al., 2020).
Methods like FitVid (Babaeizadeh et al.
, 2021)and prior two-stage methods (Yan et al., 2021), which rely on compression in both temporal and spatial domain, cannot predict goal conditioned video frames.
In contrast, we can easily adapt MaskViTto predict goal-conditioned video frames by including the last frame in Z c .
Goal conditioning significantly ( 18% FVD improvement) improves performance without requiring any architectural change(Table 1).
We believe our method could also be extended to perform video interpolation but leavethat investigation to future work.
4. A BLATION S TUDIES
We ablate MaskViT to understand the contribution of each design decision with the default settings:6 blocks, 768 embedding dimension, 1 × 16 × 16 spatial window, 16 × 4 × 4 spatiotemporal window, and variable masking ratio (Table 2).
See § B.1 and § B.2 for additional ablation studies anddiscussion.
Model hyperparameters.
We compare the effect of the number of blocks and the embeddingdimension in Table 2a.
We find that having a larger embedding dimension improves the performanceslightly, whereas increasing the number of blocks does not improve the performance.
Spatiotemporal window (STW).
An important design decision of MaskViT is the size of the STW(Table 2b).
We compare three different window configurations and MaskViT with full self-attentionin all layers.
Note that training a model with full self-attention requires using gradient checkpointing,which significantly increases the training time.
In fact, the STW size of 16 × 4 × 4 achieves betteraccuracy, while requiring 60% less memory and speeds up training time by 3 . 3 × .
Masking ratio during training.
We find that a fixed masking ratio results in poor video predictionperformance (Table 2c).
A large masking ratio until a maximum of 95% decreases FVD.
Furtherincrease in masking ratio significantly deteriorates the performance.
A variable masking ratio performs best, as it best approximates the different masking ratios encountered during inference.
Mask scheduling.
The choice of mask scheduling function and the number of decoding iterationsduring inference has a significant impact on image (Chang et al., 2022) and video generation quality(Fig. 3).
We compare three types of scheduling functions: concave (cosine, square, cubic, exponen-
[Table]
[Figure]
tial), linear, and convex (square root).
Concave functions performed significantly better than linearand convex functions.
Videos have a lot of redundant information due to temporal coherence, andconsequently with only 5% of unmasked tokens (Table 2c) the entire video can be correctly predicted.
The critical step is to predict these few tokens very accurately.
We hypothesize that concavefunctions perform better as they capture this intuition by slowly predicting the initial tokens overmultiple iterations and then rapidly predicting the majority of remaining tokens conditioned on the(more accurate) initial tokens in the final iterations.
Convex functions operate in an opposite mannerand thus perform significantly worse.
Across all functions, FVD improved with increased numbersof decoding steps until a certain point.
Further increasing the decoding steps did not improve FVD.
Additionally, we found that selecting the most confident tokens while performing iterative decoding led to video predictions with little or no motion.
Hence, we add temperature annealed Gumbelnoise to the token confidence to encourage the model to produce more diverse outputs (Fig. 3).
Empirically, we found that a temperature value of 4 . 5 works best across different datasets.
4. V ISUAL M ODEL P REDICTIVE C ONTROL WITH M ASK V I T ON A R EAL R OBOT
[Table]
In this section, we test how our generally designed framework performs whenapplied to real robotic tasks.
We evaluate the capability of our method onthe control of embodied agents throughexperimental evaluation on a Sawyerrobot arm
We first train our modelon the RoboNet dataset along with asmall collection of random interactiontrajectories in our setup.
We thenleverage MaskViT to perform visualmodel-predictive control, and evaluatethe robot’s performance on several manipulation tasks.
Setup and data collection.
We found that the VQGAN model trained only on RoboNet data didnot generalize to variations in lighting, and background, or to novel objects.
Therefore, we autonomously collect 120 K frames of additional data with a random policy to finetune VQGAN to oursetup by augmenting RoboNet.
We hypothesize that this dependency on domain-specific data couldbe substantially reduced with a large-scale pre-training of VQ-VAEs on internet-scale data as shownby prior work (Ramesh et al., 2021), but it is beyond the scope of our experiment.
At each timestep,the robot takes a 5 -dimensional action representing a desired change in end-effector pose:
[ x, y, z ]gripper position, θ yaw angle, and a binary gripper open/close command.
During data collection,the robot interacts with a diverse collection of random household objects (Fig. 4).
[Figure]
Model predictive control using MaskViT.
We evaluate the planning capabilities of MaskViT usingvisual foresight (Finn & Levine, 2017; Ebert et al., 2018) on a series of robotic pushing tasks.
Ourcontrol evaluation contains two task types: table setting , which involves pushing a bowlto a specified location, and sweeping , where objects are moved into an unseen dustpan.
Foreach task, the robot is given a 64 × 64 goal image and we perform planning based on MaskViTby optimizing a sequence of actions using the cross-entropy method (CEM)
(Boer et al., 2004),using the ℓ 2 pixel error between the last predicted image and the goal as the cost.
We comparewith two baseline methods: a FitVid model trained on combined RoboNet and the domain-specificdataset, and a random Gaussian policy.
We evaluate three variants of our model: one trained onthe combined RoboNet and domain-specific datasets ( all data ), one pretrained using RoboNetand then finetuned with domain-specific data ( finetuned ), and one trained only on RoboNet ( RNonly ).
See § A.2 for additional details and hyperparameters.
Results.
As shown in Table 4, our model improves performance compared to FitVid on real robotexperiments.
Our model achieves slightly better performance when trained jointly on all dataas opposed to pre-training on RoboNet and then finetuning on real robot data.
Qualitatively, wefound that domain-specific data significantly improves the scene fidelity and arm motion predictionaccuracy for our model, resulting in better planning performance.
We highlight two advantages of our method when applied to real robot tasks compared to priormethods.
The first advantage is the computation efficiency during inference time.
Prior two-stagemethods (Rakhimov et al., 2020; Yan et al., 2021) cannot be used for real-time planning due toextremely slow inference speed.
Furthermore, compared to autoregressive models, MaskViT isorders of magnitude more efficient as shown in Table 3.
Indeed, our method achieves ∼ 6 .
5 secondsper CEM iteration.
Second, compared to CNN-based architectures such as FitVid, which is the stateof-the-art method on this task, our method does not require any domain-specific inductive biases andallows for a more general inference procedure, without sacrificing any performance.
C ONCLUSION
In this work, we explore MaskViT, a simple but powerful and versatile method for video predictionthat leverages masked visual modeling as a pre-training task and transformers with window attention as a backbone for computation efficiency.
We showed that by masking a variable number oftokens during training, we can achieve competitive video prediction results.
Our iterative decodingscheme is significantly faster than autoregressive decoding and achieves state-of-the-art performancein planning for real robot manipulation tasks.
Limitations and future work.
While our results are encouraging, we found that using per framequantization can lead to flickering artifacts, especially in videos that have a static background likein RoboNet.
Although MaskViT is efficient in terms of memory and parameters, scaling up videoprediction, especially for scenarios that have significant camera motion (e.g., self-driving (Geigeret al., 2013) and egocentric videos (Grauman et al., 2021)) remains challenging.
Finally, an important future avenue of exploration is scaling up the complexity of robotic tasks (Srivastava et al.
,2022) integrating our video prediction method in more complex planning algorithms.
R EPRODUCIBILITY S TATEMENT
We use an open-source implementation of VQ-GAN ( https://github.com/CompVis/taming-transformers ) for all our experiments.
We use the default parameters specified inthe original implementation unless specified in Table A.1.
We use PyTorch (Paszke et al., 2019)1 .
7 library for implementing MaskViT. Architecture details and hyperparameters are described in§ 3 and Table A.1 respectively.
The inference parameters used for iterative decoding for differentdatasets are provided in Table A.1.
Finally, please refer to the detailed description of our real robotdata collection setup and experiments in § A.2.
A I MPLEMENTATION D ETAILSA. T RAINING M ASK V I T
VQ-GAN.
We train a VQ-GAN (Esser et al., 2021) model for each dataset which downsamples eachframe into 16 × 16 latent codes, i.e., by a factor of 4 for frames of size 64 × 64 frames and 16 forframes of size 256 × 256 .
Table 5 summarizes our settings for all three datasets.
Training VQ-GANwith discriminator loss can lead to instabilities.
Hence, as suggested by (Esser et al., 2021)
we start
GAN losses after the reconstruction loss has converged.
We also found that GAN losses were notalways helpful, especially at lower input resolutions for BAIR and RoboNet.
Transformer.
Our transformer model is a stack of L blocks, where each block consists of twotransformer layers with attention restricted to the window size of 1 × 16 × 16 (spatial window)and T × 4 × 4 (spatiotemporal window), unless otherwise specified.
We use learnable positionalembeddings, which are the sum of space and time positional embeddings.
Following (Liu et al.,2021a), we adopt relative position biases in our layers.
We use the Adam (Kingma & Ba, 2015)optimizer with linear warmup (Goyal et al., 2017) and a cosine decay learning rate schedule.
Table 5summarizes our settings for all three datasets.
Evaluation.
We find the optimal evaluation parameters by doing a grid search of the followingparameters: γ (cosine, square), temperature (3, 4.5) and decoding iterations depending on the prediction horizon length.
We use a top-p value of 0 .
95 for the BAIR dataset only.
Table 5 summarizesour evaluation settings.
[Table]
KITTI KITTI
[Table] [Table]
A. R EAL R OBOT E XPERIMENTS
Data collection.
Our robot setup consists of a Sawyer robot arm with a Logitech C922 PRO consumer webcam for recording video frames at 640 × 480 resolution.
All raw image observationsare center-cropped to 480 × 480 resolution before being resized to 64 × 64 for model training andcontrol in our experiments.
For the finetuning dataset, we autonomously collect 5000 trajectories of30 timesteps.
At each step the robot takes a 5-dimensional action representing a change in state ofthe end-effector: a delta translation in Cartesian space, [ x, y, z ] , for the gripper position in meters,change in θ yaw angle of the end-effector, and a binary gripper open/close command.
Followingthe action space used to collect the RoboNet dataset, the pitch and roll of the end-effector are keptfixed such that the gripper points with the fingers towards the table surface.
Each action withina trajectory is selected independently and each dimension of the action vector is independent ofthe others, sampled from a diagonal Gaussian distribution, except the gripper open/close commandthat closes automatically when the z -position of the end-effector reaches below a certain threshold to increase the rate of object interaction.
The random action distribution is parameterized by
N (0 , diag ([0 . 035 , 0 . 035 , 0 .
08 , π/ 18 , 2]) .
During data collection, we provide the robot with a diverse set of training objects to interact with.
During evaluation, we test on tasks which require therobot to manipulate bowls in one setting and to push training items into an unseen dustpan in another.
Visual-MPC.
Our control strategy is a visual MPC (Finn & Levine, 2017; Ebert et al., 2018)procedure.
Given a start and goal image I 0 , I g ∈ R 64 × 64 × 3 , the objective is to find an optimalsequence of actions to reach the goal observation from the start.
The planning objective can bewritten as: min a 1 ,a 2 ,...a H (cid:80)
Hi =1 c i ∥
ˆ f
( I 0 , a 1 , ...a H )
i − I g ∥ 22 ,
where ˆ f ( I 0 , a 1 , ...a H ) i represents thei th predicted frame by the learned video prediction model (MaskViT in our case), and c i are asequence of constant hyperparameters that determine the importance of the difference between thepredicted frame and the goal for each time step.
We use the cross-entropy method (CEM) (De Boer et al., 2005) to optimize a sequence of H = 10future actions for this objective.
In each planning iteration, we first sample M = 256 sequencesof random actions.
We then provide these sequences, together with two consecutive context frames(the previous and current step observations), and one context action (the action taken at the previousstep) to MaskViT. Action sequences are sampled according to a multivariate Gaussian distribution.
To bias action sampling towards smoother trajectories, the noise samples for actions in a givenrandom trajectory are correlated across time as in Nagabandi et al.
Specifically, given acorrelation coefficient hyperparameter β , we first compute u 1 i , u 2 i , ...
u Mi i.i.d ∼ N (0 , Σ i ) , where Σ i isthe variance of the action at timestep i in the current optimization iteration.
The noise at timestepi for the j th random trajectory, n ji , is then computed as a weighted combination of the new noisesample and the noise sample at the previous timestep, that is, n ji = (1 − β ) ∗ u ji + β ∗ n ji − 1 .
Afterall noise samples are computed, they are summed with means µ i for each timestep, which are alsoiteratively updated.
The final random trajectories are formed by rounding the elements in the lastaction dimension (gripper action) to − 1 or 1 , whichever is closer.
Next, we compare the predictions to the goal image by computing the ℓ 2 error and summing overtime as described by the objective above.
We weight the cost on the final timestep by 10 × , but stillinclude the costs on intermediate timesteps in the summation to encourage the robot to solve thetask quickly.
The best action sequences based on this score are used to refit the sampling distribution mean and variance for the next optimization iteration.
After K = 3 optimization iterations,we execute the best scoring action sequence on the robot for the first 3 steps before performingreplanning.
The robot uses a total of 15 steps to solve the task, including one initial action = [0 , 0 , − 0 .
1 , 0]which is executed at the beginning of every trajectory.
This ensures that at least two context imagesprovided for planning.
The planning hyperparameters are summarized in Table 6.
FitVid model training details.
For the FitVid comparison, we train FitVid on a combined datasetconsisting of RoboNet and our finetuning data, just as when training our model in the all datacase.
Like in the original paper, we train the model to predict 10 future frames given 2 contextframes, and use the Adam optimizer with a learning rate of 3 e − 4 .
To be consistent with the trainingsetup for our model, we do not use data augmentation.
We train for 245 K gradient steps using abatch size of 32 with 4 NVIDIA TITAN RTX GPUs.
[Table]
Evaluation.
We perform control evaluation on two categories of tasks: table setting andsweeping .
For each task, we test 5 different variations with 3 trials each.
A trial is consideredsuccessful if the center of the object of interest is within 8 cm of the goal position after the trajectoryis complete.
Model inference for real robot control is performed using 8 NVIDIA RTX 3090 GPUswith a batch size of 16 per GPU.
We use 5 decoding iterations, which yields a forward pass time ofapproximately 6 . 2 seconds for a batch of 256 samples.
B A DDITIONAL R ESULTSB. Q UANTITATIVE R ESULTS
Factorized attention.
An important decision decision of MaskViT is factorizing full-self attentioninto a block comprised of spatial and sptio-temporal window layers.
This design ensures that eachtoken can attend to all other tokens just as in a full-self-attention layer but without huge memoryrequirements.
In Table 8, we compare MaskViT with a Transformer model composed entirely of
STW layers.
For all window sizes, the FVD scores are worse than our default block design.
We notethat a higher spatial resolution helps, which showcases the importance of our block design.
VQ-GAN training dataset.
A key benefit of operating on individual frames is that we can potentially use VQ-VAEs trained on internet scale datasets (Ramesh et al., 2021).
As a step towards thatgoal, we train a VQ-GAN model jointly on all three datasets considered in this work, i.e. BAIR,
RoboNet and KITTI.
During training, we create a batch consisting of equal number of frames fromeach dataset.
We use the hyper-parameters for the RoboNet VQ-GAN model (Table 5).
As shownin Table 9, across the three data sets, the performance of MaskViT trained with a shared VQ-GANis better or similar to MaskViT trained with VQ-GAN specific to the dataset.
Real robot experiments.
Table 7 shows per-task success rates for our real-world robotic controlevaluation.
Each of our two task types (table setting, sweeping) has 5 variations, each of which involves different objects to push (unseen bowls for table setting, toys previously seen in the finetuningdata for sweeping) and different target locations.
B. Q UALITATIVE R ESULTS
Video prediction.
We present additional qualitative video prediction results for BAIR (Fig. 7),
KITTI (Fig.
8) and RoboNet (Fig. 9).
Real robot experiments.
Fig.
5 and Fig.
6 depict sample predictions for MaskViT ( all data )and MaskViT trained only on RoboNet ( RN only ) for two example control tasks.
We observethat with our model, the planner is able to find sequences of actions which bring the blue bowl orthe soft red hat close to the position specified in the goal image.
However, with a model which istrained only on the RoboNet dataset, planning fails.
We see qualitatively that the model trained in
[Table] [Table]
[Table]
Across all three datasets the performance of
MaskViT trained with a shared VQ-GAN is better or similar to MaskViT trained with dataset specific VQ-GAN.
RoboNet-only, even when solely performing reconstruction of the first two context images usingthe VQ-GAN component, is unable to reconstruct the background and robot arm with high fidelity.
Despite the diversity of the RoboNet dataset, finetuning on domain-specific data is still required toproduce reasonable predictions in our setting.
Discussion.
We performed extensive experiments comparing different design choices: model size,window design, masking ratio, and mask scheduling functions.
Qualitatively, we found that allmodel sizes and different window designs work equally well, as also indicated by similar FVDscores.
However, when our backbone consisted only of spatio-temporal windows, we found theperformance to be much worse (Table 8).
A key failure mode was lack of global consistency inthe predicted video frames, e.g. the robot arm would generally disappear after the first few frames.
We believe that this failure could be attributed to the lack of global interactions i.e. a token couldno longer attend to all other tokens.
Models trained with a low masking ratio or without addingannealed noise had a high FVD score because they showed very little or no motion of the robotarm.
Finally, in comparison with other works, we found that the biggest qualitative difference wasin the KITTI dataset.
KITTI presents a unique challenge, as we have to predict 25 video frames.
Asshown in Fig.
2 and Fig.
8, MaskViT can predict future video frames without perceptually visibledegradation, whereas prior work predict frames that degrade rapidly.
B. L IMITATIONS AND F UTURE W ORK
Although MaskViT addresses important limitations of current work in video prediction, we believethe field has not yet witnessed progress similar to text-to-image generation (Ramesh et al., 2021).
Key areas of improvement are short-term temporal consistency, long-term temporal consistency,increasing the prediction horizon, and moving beyond narrow domain-specific datasets to “ in-thewild ” settings like everyday egocentric videos (Grauman et al., 2021).
One possible way to improve short-term consistency is incorporating temporal consistency losses for
VQ-GAN training.
Despite using a bidirectional window attention mechanism to reduce memoryrequirements, we are still limited to predicting 25 frames in the future.
Possible ways to predict longhorizon videos include training multiple networks at different temporal resolutions and incorporatingadvances in increasing the context length of Transformers.
Finally, scaling both data and modelparameters might be necessary to move to more realistic and complicated everyday scenes.
t = 2 t =
S t a r t
i ng i m age
V i s
ua li
z a t i on o f be s t p r ed i c t ed p l an a t t = 1 f ound a ft e r
C E M
i t e r a t i on s
G oa l i m age
M ask
V i T ( a ll d a t a ) r e c on s t r u c t ed c on t e x t t =t =
t =
V i s
ua li
z a t i on o f be s t p r ed i c t ed p l an a t t = 4 f ound a ft e r
C E M
i t e r a t i on s r e c on s t r u c t ed c on t e x t t =t =
V i s
ua li
z a t i on o f be s t p r ed i c t ed p l an a t t = 1 f ound a ft e r
C E M
i t e r a t i on s
M ask
V i T (
R obo
N e t on l y )
r e c on s t r u c t ed c on t e x t t =t =
F u ll t a sk e x e c u t i on
F i gu r e li t = 2 t =
S t a r t
i ng i m age
V i s
ua li
z a t i on o f be s t p r ed i c t ed p l an a t t = 1 f ound a ft e r
C E M
i t e r a t i on s
G oa l i m age
M ask
V i T ( a ll d a t a ) r e c on s t r u c t ed c on t e x t t =t =
t =
V i s
ua li
z a t i on o f be s t p r ed i c t ed p l an a t t = 4 f ound a ft e r
C E M
i t e r a t i on s r e c on s t r u c t ed c on t e x t t =t =
V i s
ua li
z a t i on o f be s t p r ed i c t ed p l an a t t = 1 f ound a ft e r
C E M
i t e r a t i on s
M ask
V i T (
R obo
N e t on l y )
r e c on s t r u c t ed c on t e x t t =t =
F u ll t a sk e x e c u t i on
F i gu r e 6 :
V i s u a li za ti on s f o r t a s k sweep : push hat to bottomrig ht corner .
G r ound t r u t h t =
P r ed i c t ed t =t =t =t =t = t = t =
G r ound t r u t h t =
P r ed i c t ed t =t =t =t =t = t = t =
G r ound t r u t h t =
P r ed i c t ed t =t =t =t =t = t = t =
G r ound t r u t h t =
P r ed i c t ed t =t =t =t =t = t = t =
( c on t e x t ) t =( c on t e x t ) t =( c on t e x t ) t =( c on t e x t ) t =
F i gu r e 7 :
Q u a li t a t i v e re s u l t s :
B A I
R ( ac ti on fr ee , × 64 )
G r ound t r u t h t =
P r ed i c t ed t =t =t = t = t = t = t =
G r ound t r u t h
P r ed i c t ed
G r ound t r u t h
P r ed i c t ed
G r ound t r u t h
P r ed i c t ed ( c o n t e x t ) t
=t =
t = t = t =t =
t = t = t = t =
t = 2t =t =
t =
( c on t e x t ) t =t =
t =t =
t = t = t = t =
( c on t e x t ) t =t =
t = t =
t =t = t = t =
t = t =
t = 2t =t =
t =
( c o n t e x t )
F i gu r e 8 :
Q u a li t a t i v e re s u l t s :
K I
TT I (× 256 )
G r ound t r u t h t =
P r ed i c t ed t =t =t =t =t =t =t =
G r ound t r u t h
P r ed i c t ed
G r ound t r u t h
P r ed i c t ed
G r ound t r u t h
P r ed i c t ed ( c o n t e x t ) t
=t =
t = t = t =t =
t =t =
t =t = t =t =
t =
t =
( c on t e x t ) t =t =
t =t = t =t =
t =
t =( c on t e x t ) t =t =
t = t =
t =t = t =t =
t =t = t =t =
t =
t =
( c o n t e x t )
F i gu r e 9 :
Q u a li t a t i v e re s u l t s :
R o b o N e t (× 256 )