SlowGuess/NeurIPS_2023_OCR huggingface dataset:
4.3 Generalization of Video Prediction Rewards
Prior works in the space of large, pretrained generative models have shown a powerful ability to generalize beyond the original training data, demonstrated by their ability to produce novel generations (e.g., unseen text captions for image or video generation [33, 45, 40]) as well as learn more generalizable models from limited finetuning data [8, 46]. Both capabilities provide promising directions for extending large video generation models to VIPER, where we can leverage text-video models to capture priors specified by text commands, or finetune on a small set of demonstrations to learn a task-specific prior that better generalizes to novel states.
Figure 7: Sampled video predictions for in distribution reference videos (Train) and an OOD arm/task combination (OOD). The video model displays cross-embodiment generalization to arm/task combination not observed in the training data. Video model generalization can enable specifying new tasks where no reference data is available.
In this section, we seek to understand how this generalization can be used to learn more general reward functions. We train a model on two datasets of different robot arms, and evaluate the cross-embodiment generalization capabilities of the model. Specifically, we gather demonstrations for 23 tasks on the Rethink Robotics Sawyer Arm, and demonstrations for 30 tasks on the Franka Panda robotic arm, where only 20 tasks are overlapping between arms. We then train a task-conditioned autoregressive video
Figure 8: (Left) Training curve for RL agent trained with VIPER on OOD task. (Right) Task-conditional likelihood for reference and random trajectory for an OOD task.
model on these demonstration videos and evaluate the video model by querying unseen arm/task combinations, where a single initial frame is used for open loop predictions.
Sample video model rollouts for in distribution training tasks and an OOD arm/task combination are shown in Figure 7. Even though the video model was not directly trained on demonstrations of the Franka Panda arm to solve the saucepan task in RLBench, it is able to generate reasonable trajectories for the arm and task combination. Figure 8 further validates this observation by assigning higher likelihood to expert videos (Ref) compared to random robot trajectories (Rand). We observe that these generalization capabilities also extend to downstream RL, where we use our trained video model with VIPER to learn a policy for the Franka Robot arm to solve an OOD task without requiring demos for that specific task and arm combination. Figure 21 further extends this analysis to include more in distribution and out of distribution tasks. These results demonstrate a promising direction for future work in applying VIPER to larger scale video models that will be able to better generalize and learn desired behaviors only through a few demonstrations.
4.4 Impact of Data Quality and Quantity
Learning from sub-optimal data is an important feature of learned reward models, as large-scale video data may often contain suboptimal solutions to a given task. As such, we evaluate VIPER's ability to learn rewards from sub-optimal data.
We train video models with suboptimal (good) data, which has  $50 - 75\%$  returns of the expert data. In Figure 9, VIPER learns the suboptimal behaviors provided in the suboptimal video data. This suboptimal VIPER can be still useful if combined with a sparse task reward, comparable to an agent learned from dense task rewards.
Additionally, we evaluate how VIPER performs under different video dataset sizes. As shown in Figure 10, VIPER can learn a meaningful reward function only with one expert trajectory, although adding more videos quickly improves the performance of VIPER, with diminishing returns after 50 trajectories.
Figure 9: Atari performance with VIPER models trained on suboptimal data.
Figure 10: DMC performance with VIPER trained on different dataset sizes, where N is the number of expert trajectories. *Original dataset size.
In this section, we study the contributions of various design decisions when using video prediction models as rewards: (1) how to weight the exploration objective, (2) which video model to use, and (3) what context length to use. Ablation studies are performed across DMC tasks.
Figure 11: Effect of exploration reward term, video model choice, and context length on downstream RL performance. An equally weighted exploration reward term and longer video model context leads to improved performance. MaskGIT substantially underperforms VideoGPT as a choice of video model. The BYOL model performs moderately due to the deterministic architecture not properly handling multi-modality.
Exploration objective As discussed in subsection 3.2, an exploration objective may help the RL agent learn a distribution of behaviors that closely matches the distribution of the original data and generalizes, while preventing locally optimal solutions. In Figure 11, we ablate the  $\beta$  parameter introduced in Equation 6 using a VideoGPT-like model as the VIPER backbone.  $\beta = 0$  corresponds to no exploration objective, whereas  $\beta = 1$  signifies equal weighting between  $r^{\mathrm{VIPER}}$  and  $r^{\mathrm{exp}}$ . We ablate the exploration objective using the Plan2Explore [35] reward, which provides the agent with a reward proportional to the disagreement between an ensemble of one-step dynamics models. Using no exploration objective causes the policy's behavior to collapse, while increasing the weight of the exploration objective leads to improved performance.
Video model Although our experiments focus on using an autoregressive VideoGPT-like model to compute likelihoods, VIPER generally allows for any video model that supports computing conditional likelihoods or implicit densities.
Figure 11 shows additional ablations replacing our VideoGPT model with a similarly sized MaskGIT [6] model, where frames are modeled with MaskGIT over space, and autoregressive over time. MaskGIT performs substantially worse than the VideoGPT model, which is possibly due to noisier likelihoods from parallel likelihood computation. In addition, while VideoGPT only requires 1 forward pass to compute likelihood, MaskGIT requires as many as 8 forward passes on the sequence of tokens for a frame, resulting in an approximately  $8 \times$  slowdown in reward computation, and  $2.5 \times$  in overall RL training.
Finally, we evaluate the performance of a video model that computes implicit densities using a negative distance metric between online predictions and target encodings with a recurrent Bootstrap Your Own Latent (BYOL) architecture [12, 11]. We refer the reader to Appendix A.1 for more details about the implementation. While BYOL outperforms MaskGIT in Figure 11, its deterministic recurrent architecture is unable to predict accurate embeddings more than a few steps into the future. This limits the ability of the learned reward function to capture temporal dynamics. We observe that agents trained with BYOL often end up in a stationary pose, which achieves high reward from the BYOL model.
Context length The context length  $k$  of the video model is an important choice in determining how much history to incorporate into the conditional likelihood. At the limit where  $k = 0$ , the reward is the unconditional likelihood over each frame. Such a choice would lead to stationary behavior on many tasks. Figure 11 shows that increasing the context length can help improve performance when leveraging a VideoGPT-based model for downstream RL tasks, subject to diminishing returns. We hypothesize that a longer context length may help for long-horizon tasks.
This paper presents Video Prediction Rewards (VIPER), a general algorithm that enables agents to learn behaviors from videos of experts. To achieve this, we leverage a simple pre-trained autoregressive video model to provide rewards for a reinforcement learning agent. These rewards are parameterized as the conditional probabilities of a particular frame given a context past of frames. In addition, we include an entropy maximization objective to ensure that the agent learns diverse behaviors which match the video model's trajectory distribution. VIPER succeeds across 3 benchmarks and 28 tasks, including the DeepMind Control Suite, Atari, and the Reinforcement Learning Benchmark. We find that simple data augmentation techniques can dramatically improve the effectiveness of VIPER. We also show that VIPER generalizes to out-of-distribution tasks for which no demonstrations were provided. Moreover, VIPER can learn reward functions for a wide variety of tasks using a single video model, outperforming AMP [30], which suffers from mode collapse as the diversity of the expert videos increases.
Limitations of our work include that VIPER is trained using in-domain data of expert agents, which is not a readily-available data source in the real world. Additionally, video prediction rewards trained on stochastic data may lead the agent to prefer states that minimize the video model's uncertainty, which may lead to sub-optimal behaviors. This challenge may be present for cases where either the environment is stochastic or the demonstrator provides noisy demonstrations. Moreover, selecting the right trade-off between VQCode size and context length has a significant impact on the quality of the learned rewards, with small VQCodes failing to model important components of the environment (e.g., the ball in Atari Breakout) and short context lengths leading to myopic behaviors which results in poor performance. Exploring alternative video model architectures for VIPER would likely be a fruitful research direction.
To improve the generalization capabilities of VIPER, larger pre-trained video models are necessary. Future work will explore how fine-tuning or human preferences can be used to improve video prediction rewards. Using text-to-video models remain another interesting direction for future work in extracting text-conditioned task-specific priors from a pretrained video model. In addition, extending this line of work to leverage Video Diffusion Models [18] as video prediction rewards may lead to interesting outcomes.
This work was supported in part by an NSF Fellowship, NSF NRI #2024675, ONR MURI N00014-22-1-2773, Komatsu, and the Vanier Canada Graduate Scholarship. We also thank Google TPU Research Cloud and Cirrascale (https://cirrascale.com) for providing compute resources.
[1] Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference on Machine Learning, 2004.
[2] Bowen Baker, Ilge Akkaya, Peter Zhokhov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. In Advances in Neural Information Processing Systems, 2022.
[3] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
[4] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540, 2016.
[5] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations, 2019.
[6] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11315-11325, 2022.
[7] Soumith Chintala, Emily Denton, Martin Arjovsky, and Michael Mathieu. How to train a gan? tips and tricks to make gans work. https://github.com/soumith/ganhacks, 2020.
[8] Yilun Dai, Mengjiao Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Josh Tenenbaum, Dale Schuurmans, and Pieter Abbeel. Learning universal policies via text-guided video generation. arXiv preprint arXiv:2302.00111, 2023.
[9] Andreas Doerr, Christian Daniel, Martin Schiegg, Nguyen-Tuong Duy, Stefan Schaal, Marc Toussaint, and Trimpe Sebastian. Probabilistic recurrent state-space models. In International conference on machine learning, pages 1280-1289. PMLR, 2018.
[10] Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12873-12883, 2021.
[11] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. In Advances in neural information processing systems, volume 33, pages 21271-21284, 2020.
[12] Zhaohan Guo, Shantanu Thakoor, Miruna Píslar, Bernardo Avila Pires, Florent Altché, Corentin Tallec, Alaa Saade, Daniele Calandriello, Jean-Bastien Grill, Yunhao Tang, et al. Byol-explore: Exploration by bootstrapped prediction. In Advances in neural information processing systems, volume 35, pages 31855-31870, 2022.
[13] Agrim Gupta, Stephen Tian, Yunzhi Zhang, Jiajun Wu, Roberto Martin-Martin, and Li Fei-Fei. Maskvit: Masked visual pre-training for video prediction. arXiv preprint arXiv:2206.11894, 2022.
[14] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pages 2555-2565. PMLR, 2019.
[50] Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarej, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang. Task-relevant adversarial imitation learning. In Conference on Robot Learning, 2020.
A.1 Bootstrap Your Own Latent Details
Model Architecture. The BYOL model used in this work derives from the recurrent BYOL-Explore architecture proposed in [12], with a few modifications. Namely Guo et al. propose BYOL-Explore, a multi-step predictive latent world model. The BYOL-Explore architecture trains an online network using targets generated by an exponential moving average target network on future observations. Observations  $o_t$  are first encoded into an observation representation  $f_{\theta}(o_t)$ . Online predictions are computed by then encoding the observation representation using a closed-loop RNN  $h_{\theta}^{c}(f_{\theta}(o_t))$ . Guo et al. also pass actions into the closed loop RNN, but we wish to learn an action-free reward mechanism that can be trained solely from videos. At each timestep, the carry  $b_t$  from the closed loop RNN is then fed into an open-loop RNN  $h_{\theta}^{o}$  which predicts open-loop representations  $K$  steps into the future:  $(b_{t,k} \in \mathcal{R}^{\mathcal{M}})_{k=1}^{K-1}$ , where  $b_{t,k} = h_{\theta}^{o}(b_{t,k-1})$ . The original BYOL-Explore architecture feeds actions  $a_{t+k+1}$  as inputs to the open loop cell  $h_{\theta}^{o}$  as well, but we opt to omit these inputs for the same reason outlined above.
Given the deterministic RNN architecture used to predict online targets, omitting the action conditioning may lead to poor performance in multi-modal or stochastic environments. A more principled approach would utilize a probabilistic recurrent state space model [9, 14] to account for the multimodality of futue states. We leave this approach to future work.
Finally, a predictor head  $g_{\theta}(b_{t,k})$  outputs online predictions. We refer the reader to [12] for a figure of the outlined architecture.
The target network is an observation encoder  $f_{\phi}$  whose parameters are an exponential moving average of  $f_{\theta}$ . The loss function for the online network is then defined as:
$$
\mathcal {L} _ {\text {B Y O L - E x p l o r e}} (\theta , t, k) = \left\| \frac {g _ {\theta} (b _ {t , k})}{\| g _ {\theta} (b _ {t , k}) \| _ {2}} - \operatorname {s g} \left(\frac {f _ {\phi} (o _ {t + k})}{\| f _ {\phi} (o _ {t + k}) \| _ {2}}\right) \right\| _ {2} ^ {2},
$$
$$
\mathcal {L} _ {\mathrm {B Y O L - E x p l o r e}} (\theta) = \frac {1}{B (T - 1)} \sum_ {t = 0} ^ {T - 2} \frac {1}{K (t)} \sum_ {k = 1} ^ {K (t)} \mathcal {L} _ {\mathrm {B Y O L - E x p l o r e}} (\theta , t, k),
$$
where  $K(t) = \min(K, T - 1 - t)$  is the valid open-loop horizon for a trajectory of length  $T$  and  $\mathbf{sg}$  is the stop-gradient operator.
Computing Rewards. The uncertainty associated with the transition  $(o_{t},a_{t},o_{t + 1})$  is the sum of the corresponding prediction losses:
$$
\ell_ {t} = \sum_ {p + q = t + 1} \mathcal {L} _ {\mathrm {B Y O L - E x p l o r e}} (\theta , j, p, q),
$$
where  $0 \leq p \leq T - 2$ ,  $1 \leq q \leq K$  and  $0 \leq t \leq T - 2$ . This accumulates all the losses corresponding to the latent dynamics model uncertainties relative to the observation  $o_{t + 1}$ . For our purposes, we calculate rewards as  $r_t = -\ell_t$ , which is equivalent to negating the reward used in the original BYOL-Explore formulation, and can be interpreted as negating the exploration objective to ensure that the agent stays close to the trajectory distribution found in the original dataset.
A.2 Video Model Hyperparameters
Table 1: Hyperparameters and training details for all VQ-GAN models
DMC	Atari	RLBench (single + multi task)
Input size	64 × 64 × 3	64 × 64 × 3	64 × 64 × 3
Latent size	8 × 8	16 × 16	16 × 16
β (commitment loss coefficient)	0.25	0.25	0.25
Batch size	128	128	128
Learning rate	10-4	10-4	10-4
Learning rate schedule	constant	constant	constant
Training steps	200k	200k	200k
Base channels	128	128	128
Ch mult	[1, 1, 2, 2]	[1, 2, 2]	[1, 2, 2]
Num res blocks	1	1	2
Codebook size	1024	1024	1024
Codebook dimension	64	64	64
Perceptual loss weight	0.1	0.1	0.1
Disc base features	32	32	32
Disc gradient penalty weight	108	108	108
Disc max hidden feature size	512	512	512
Disc mbstd group size	4	4	4
Disc mbstd num features	1	1	1
Disc loss weight	0.1	0.1	0.1
Table 2: Hyperparameters and training details for all VideoGPT and MaskGit models
DMC	Atari	RLBench (single task)	RLBench (multi-task)
Sequence length	16	16	4	4
Frame skip	1	1	4	4
Latent size	8 × 8	16 × 16	16 × 16	16 × 16
Number of classes	16	8	n/a	34
Batch size	32	32	32	32
Learning rate	10-4	10-4	10-4	10-4
Learning rate schedule	constant	constant	constant	constant
Training steps	500k	500k	500k	500k
Adam (β1,β2)	(0.9, 0.999)	(0.9, 0.999)	(0.9, 0.999)	(0.9, 0.999)
Hidden dim	256	512	256	512
Feedforward dim	1024	2048	1024	2048
Number of heads	8	8	8	8
Number of layers	8	8	8	8
Dropout	0	0	0	0
A.3 Reinforcement Learning Hyperparameters
Table 3: Hyperparameters and training details for DreamerV3
DMC	Atari	RLBench (single task)	RLBench (multi-task)
General
Replay Capacity (FIFO)	10^6	10^6	10^6	10^6
Start learning (prefill)	0	0	0	0
Batch Size	16	16	16	16
Batch length	64	64	64	64
MLP Size	2 × 512	4 × 1024	4 × 1024	4 × 1024
Activation	LayerNorm + swish
World Model
RSSM Size	512	4096	4096	4096
Number of Latents	32	32	32	32
Classes per Latent	32	32	32	32
KL Balancing	0.667	0.667	0.667	0.667
Actor Critic
Imagination Horizon	15	15	15	15
Discount	0.995	0.995	0.995	0.995
Return Lambda	0.95	0.95	0.95	0.95
Target Update Interval	50	50	50	50
All Optimizers
Gradient Clipping	100	100	100	100
Learning Rate	10^-4	10^-4	10^-4	10^-4
Adam epsilon	10^-6	10^-6	10^-6	10^-6
Plan2Explore
Ensemble size	10	10	10	10
Table 4: Hyperparameters and training details for DrQ
DMC
Actor Critic
MLP Size	256 × 256
CNN Features	32 × 64 × 128 × 256
CNN Filters	3 × 3 × 3 × 3
CNN Strides	2 × 2 × 2 × 2
Learning Rate	3e-4
Discount	0.99
Number of Critics	2
τ	0.005
Init Temperature	0.1
Augmentations	Random crop
Random Network Distillation
RND CNN Features	32 × 64 × 64
RND MLP Size	512 × 512
RND learning rate	3e-4
RND Exploration Weight	1
A.4 Environment Training Curves
Figure 12: Results across 6 Reinforcement Learning Benchmark tasks, with mean and standard deviation computed across 3 seeds.
Figure 13: Results across 7 Atari tasks. Scores computed using Human-Normalized Mean across 3 seeds.
Figure 14: Results across 15 DeepMind Control Suite tasks, with mean and standard deviation computed across 3 seeds. AMP runs were stopped early due to poor performance.
A.5 Training Environments
Figure 15: A single autoregressive video model is trained on 30 tasks for the Franka Panda and 23 tasks for the ReThink Robotics Sawyer, using  $16 \times 16$  VQCodes and a context length of 4, with frame skip 4.
Figure 16: A single task-conditioned autoregressive video model is trained on 15 DeepMind Control Tasks, using  $8 \times 8$  VQCodes and a context length of 16.
Figure 17: A single autoregressive video model is trained on 7 Atari tasks, using  $16 \times 16$  VQCodes and a context length of 16.
A.6 Visualizing Video Prediction Uncertainty
Video model uncertainty can be visualized by upsampling the conditional likelihoods over VQCodes. Since VQCodes tend to model local features, this provides a useful tool for analyzing which regions of the image the video model is uncertain about. We visualize video prediction uncertainty for three environments below:
(a) Reference Trajectory
(b) Random Trajectory
Figure 18: Uncertainty visualized for a reference trajectory (top) and a random trajectory (bottom). Brighter values correspond to higher likelihoods. For the random trajectory, the video model assigns lower likelihoods to regions of the image containing the ball. This is especially true for random trajectories, where the video model, trained solely on expert trajectories, cannot accurately predict what happens when the agent plays sub-optimally.
Figure 19: Uncertainty visualized for a random trajectory from the dmc cartpole balance task. Brighter values correspond to higher likelihoods.
(b) Random Trajectory
Figure 20: Uncertainty visualized for a reference trajectory (top) and a random trajectory (bottom). Brighter values correspond to higher likelihoods. Notice the high uncertainty over the position of objects on the table for the first frame. This corresponds to the unconditional likelihood (with no context). Since the position of the saucepan is randomized for every episode, the video model assigns some uncertainty to this initial configuration. For the random trajectory, there is high uncertainty assigned to the position of the arm since the video model has not seen trajectories that display poor behavior.
A.7 VIPER Out-of-distribution Analysis
Figure 21: Predicted  $r^{\mathrm{VIPER}}$  on various RLBench tasks. We evaluate 10 trajectories for each task. *Training data for the VIPER model.