SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 4: Action Execution. Synthesized video plans and executed actions in the simulated environment. The two video plans roughly align with each other.
In actions to execute, we further compare with transactions to execute, similar to the goal-conditioned layer [25] (Image + TT). Finally, to highlight the immediate UniPi with learning a diffusion process that, as per future robot actions in the joint space (as opposed to [21, 27]). For both our method and each baseline, instructions using pretrained T5 embeddings. Note that learning baselines are not directly applicable as we do not consider environment.
Figure 5: Multitask Video Generation. Generated video plans on new test tasks in the multitask setting. UniPi is able to synthesize plan across a set of environments.
Frame Condition	Frame Consistency	Temporal Heirarchy	Place	Relation
No	No	No	13.2 ± 3.2	12.4 ± 2.4
Yes	No	No	52.4 ± 2.9	34.7 ± 2.6
Yes	Yes	No	53.2 ± 3.0	39.4 ± 2.8
Yes	Yes	Yes	59.1 ± 2.5	53.2 ± 2.0
Table 2: Task Completion Accuracy Ablations. Each component of UniPi improves its performance. Performance reported on the seen place and relation tasks.
language instruction specifies placing a block in relation to some other block as opposed to direct pick-and-place.
Combinatorial Generalization. In Table 1, we find that UniPi generalizes well to both seen and novel combinations of language prompts. We illustrate our action generation pipeline in Figure 4 and different generated video plans using our approach in Figure 3.
Ablations. In Table 2, we ablate UniPi on seen language instructions and in-relation-to tasks. Specifically, we study the effect of conditioning the video generative model on the first observation frame (frame condition), tiling the observed frame across timesteps (frame consistency) and super-resolving video generation across time (temporal hierarchy). In settings where frame consistency is not enforced, we provide a zeroed out image as context to the non-start frames in a video. We found that all components of UniPi
Figure 7: Adaptable Planning. By guiding test-time sampling towards an intermediate image, we can adapt our planning procedure to move a particular block.
are crucial for good performance. We found that frame conditioning and consistency enabled videos to be consistent with the observed image and temporal hierarchy enabled more detailed video plans.
Adaptability. We next assess the ability of UniPi to adapt at test time to new constraints. In Figure 7, we illustrate the ability to construct plans which color and move one particular block to a specified geometric relation.
4.2 Multi-Environment Transfer
We next evaluate the ability of UniPi to effectively learn across a set of different tasks and generalize, at test time, to a new set of unseen environments.
Setup. To measure multi-task learning and transfer, we use the suite of language guided manipulation tasks from [28]. We train our method using demonstrations across a set of 10 separate tasks from [28], and evaluate the ability of our approach to transfer to 3 different test tasks. Using a scripted oracle agent, we generate a set of 200k videos of language execution in the environment. We report the underlying accuracy in which each language instruction is completed. Details of this environment can be found in Appendix A.
Figure 6: High Fidelity Plan Generation. UniPi can generate high resolution video plans across different language prompts.
Model	Place Bowl	Pack Object	Pack Pair
State + Transformer BC	9.8 ± 2.6	21.7 ± 3.5	1.3 ± 0.9
Image + Transformer BC	5.3 ± 1.9	5.7 ± 2.1	7.8 ± 2.6
Image + TT	4.9 ± 2.1	19.8 ± 0.4	2.3 ± 1.6
Diffuser	14.8 ± 2.9	15.9 ± 2.7	10.5 ± 2.4
UniPi (Ours)	51.6 ± 3.6	75.5 ± 3.1	45.7 ± 3.7
Table 3: Task Completion Accuracy on Multitask Environment. UniPi generalizes well to new environments when trained on a set of different multi-task environments.
Baselines. We use the same baseline methods as in Section 4.1. While our environment setting is similar to that of [28], this method is not directly comparable to our approach, as CLIPort abstracts actions to the existing primitives of pick and place as opposed to using the joint space of a robot. CLIPort is also designed to solve the significantly simpler problem of inferring only the poses upon which to pick and place objects (with no easy manner to adapt to our setting).
Multitask Generalization. In Table 3 we present results of UniPi and baselines across new tasks. The UniPi approach is able to generalize and synthesize new videos and decisions of different language tasks, and can generate videos consisting of picking different kinds of objects and different colored objects. We further present video visualizations of our approach in Figure 5.
4.3 Real World Transfer
Finally we evaluate the extent to which UniPi can generalize to real world scenarios and construct complex behaviors by leveraging widely available videos on the internet.
Setup. Our training data consists of an internet-scale pretraining dataset and a smaller real-world robotic dataset. The pretraining dataset uses the same data as [19], which consists of 14 million video-text pairs, 60 million image-text pairs, and the publicly available LAION-400M image-text dataset. The robotic dataset is adopted from the Bridge dataset [29] with 7.2k video-text pairs, where we use the task IDs as texts. We partition the 7.2k video-text pairs into train (80%) and test (20%) splits. We pretrain UniPi on the pretraining dataset followed by finetuning on the train split of the Bridge data. Architectural details can be found in Appendix A.
Video Synthesis. We are particularly interested in the effect of pretraining on internet-scale video data that is not specific to robotics. We report the CLIP scores, FIDs, and FVDs (averaged across frames and computed on 32 samples) of UniPi trained on Bridge data, with and without pretraining. As shown in Table 4, UniPi with pretraining achieves significantly higher FID and FVD and a marginally better CLIP score than UniPi without pretraining, suggesting that pretraining on non-robot data helps with generating plans for robots. Interestingly, UniPi without pretraining often synthesizes plans that fail to complete the task (Figure 6), which is not well reflected in the CLIP
Figure 8: Pretraining Enables Combinatorial Generalization. Figure 9: Robustness to Background Change.  
tion. Internet pretraining enables UniPi to synthesize videos. UniPi learns to be robust to changes of underlying of tasks not seen in training. In contrast, a model trained background, such as black cropping or the addition from scratch incorrectly generates plans of different tasks. of photo-shopped objects.
Model (24x40)	CLIP Score ↑	FID ↓	FVD ↓	Success ↑
No Pretrain	24.43 ± 0.04	17.75 ± 0.56	288.02 ± 10.45	72.6%
Pretrain	24.54 ± 0.03	14.54 ± 0.57	264.66 ± 13.64	77.1%
Table 4: Video Generation Quality of UniPi on Real Environment. The use of existing data on the internet improves video plan predictions under all metrics considered.
score, suggesting the need for better generation metrics for control-specific tasks. To tackle the lack of such a metric, we develop a surrogate metric for evaluating task success from the generated videos. Specifically, we train a success classifier that takes in the last frame of a generated video and predicts whether the task is successful or not. We find that training a model from scratch achieves  $72.6\%$  success while finetuning from a pretrained model improves performance to  $77.1\%$ . In both settings, generated videos are able to successfully complete most tasks.
Generalization. We find that internet-scale pretraining enables UniPi to generalize to novel task commands and scenes in the test split not seen during training, whereas UniPi trained only on task-specific robot data fails to generalize. Specifically, Figure 8 shows the results of novel task commands that do not exist in the Bridge dataset. Additionally, UniPi is relatively robust to background changes such as black cropping or the addition of photo-chipped objects as shown in Figure 9.
Learning Generative Models of the World. Models trained to generate environment rewards and dynamics that can serve as "world models" for model-based reinforcement learning and planning have been recently scaled to large-scale architectures developed for vision and language [9, 25, 30, 31]. These works separate learning the world model from planning and policy learning, and arguably present a mismatch between the generative modeling objective of the world and learning optimal policies. Additionally, learning a world model requires the training data to be in a strict state-action-reward format, which is incompatible with the largely available datasets on the internet, such as YouTube videos. While methods such as VPT [32] can utilize internet-scale data through learning an inverse dynamics model to label unlabeled videos, an inverse dynamics model itself does not support model-based planning or reinforcement learning to further improve learned policies beyond imitation learning. Our text-conditioned video policies can be seen as jointly learning the world model and conducting hierarchical planning simultaneously, and is able to leverage widely available datasets that are not specifically designed for sequential decision making.
Diffusion Models for Decision Making. Diffusion models have recently been applied to different decision making problems [21, 27, 33, 34, 35, 36]. Most similar to this work, [21] trained an unconditional diffusion model to generate trajectories consisting of joint-based states and actions, and used a separately trained reward model to select generated plans. On the other hand, [27] trained a conditional diffusion model to guide behavior synthesis from desired rewards, constraints or agent skills. Unlike both works, which learn task-specific policies from scratch, our approach of text-condition video generation as a universal policy can leverage internet-scale knowledge to learn generalist agents that can be deployed to a variety of novel tasks and environments. Additionally, [37]
applied web-scale text-conditioned image diffusion to generate a goal image to condition a policy on, whereas our work uses video diffusion to learning universal policies directly.
Learning Generalist Agents. Inspired by the success of large-scale pretraining in vision and language domains, large-scale sequence and image models have recently been applied to learning generalist decision making agents [6, 26, 38]. However, these generalist agents can only operate under environments with the same state and action spaces (e.g., Atari games) [26, 38], or require studious tokenization [6] that might seem unnatural in scenarios where different environments have distinct state and actions spaces. Another downside of using customized tokens for control is the inability to directly utilize knowledge from pretrained vision and language models. Our approach, on the other hand, uses text and images as universal interfaces for policy learning so that the knowledge from pretrained vision and language models can be preserved. The choice of diffusion as opposed to autoregressive sequence modeling also enables long-term and hierarchical planning.
Learning Text-Conditioned Policies. There has been a growing amount of work using text commands as a way to learn multi-task and generalist control policies [39, 40, 24, 41, 42, 43]. Different from our framing of video-as-policies, existing work directly trains a language-conditioned control policy in the action space of some specific robot, leaving cross-morphology multi-environment learning of generalist agents as an unsolved problem. We believe this paper is the first to propose images as a universal state and action space to enable broad knowledge transfer across environments, tasks, and even between humans and robots.
We have demonstrated the utility of representing policies using text-conditioned video generation, showing that this enables effective combinatorial generalization, multi-task learning, and real world transfer. These positive results point to the broader direction of using generative models and the wealth of data on the internet as powerful tools to generate general-purpose decision making systems.
Limitations. Our current approach has several limitations. First, the underlying video diffusion process can be slow, it can take a minute to generate highly photorealistic videos. This slowness can be overcome by distilling the diffusion process into a faster sampling network [44], which in our initial experimentation resulted in a 16x speed-up. UniPi may further be sped up with by faster speed samplers in diffusion models. Second, the environments considered in this work are generally fully observed. In partially observable environments, video diffusion models might make hallucination of objects or movements that are unfaithful or not in the physical world. Integrating video models with semantic knowledge about the world may help resolve this issue, and the integration of UniPi with LLMs would be an interesting direction of future work.
[1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020. 1
[2] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022. 1
[3] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022. 1
[4] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pages 8748-8763. PMLR, 2021. 1
[39] Wenlong Huang, Fei Xia, Ted Xiao, Harris Chan, Jacky Liang, Pete Florence, Andy Zeng, Jonathan Thompson, Igor Mordatch, Yevgen Chebotar, et al. Inner monologue: Embodied reasoning through planning with language models. arXiv preprint arXiv:2207.05608, 2022. 10
[40] Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, et al. Do as i can, not as i say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01691, 2022. 10
[41] Corey Lynch and Pierre Sermanet. Language conditioned imitation learning over unstructured data. arXiv preprint arXiv:2005.07648, 2020. 10
[42] Parsa Mahmoudieh, Deepak Pathak, and Trevor Darrell. Zero-shot reward specification via grounded natural language. In ICLR 2022 Workshop on Generalizable Policy Learning in Physical World, 2022. 10
[43] Hao Liu, Lisa Lee, Kimin Lee, and Pieter Abbeel. Instruction-following agents with jointly pre-trained vision-language models. arXiv preprint arXiv:2210.13431, 2022. 10
[44] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512, 2022. 10
[45] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J. Fleet. Video diffusion models. arXiv preprint arXiv:2204.03458, 2022. 14
[46] Chitwan Sahara, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022. 14
A Architecture, Training, and Evaluation Details
A.1 Video Diffusion Training Details
We use the same base architecture and training setup as [45] which utilizes Video U-Net architecture with 3 residual blocks of 512 base channels and channel multiplier [1, 2, 4], attention resolutions [6, 12, 24], attention head dimension 64, and conditioning embedding dimension 1024. We use noise schedule log SNR with range [-20, 20]. We make modifications Video U-Net to support first-frame conditioning during training. Specifically, we replicate the first frame to be conditioned on at all future frame indices, and apply temporal super resolution model condition on the replicated first frame by concatenating the first frame channel-wise to the noisy data similar to [46]. We use temporal convolutions as opposed to temporal attention to mix frames across time, to maintain local temporal consistency across time, which has also been previously noted in [19]. We train each of our video diffusion models for 2M steps using batch size 2048 with learning rate 1e-4 and 10k linear warmup steps. We use 256 TPU-v4 chips for our first-frame conditioned generation model and temporal super resolution model.
We use T5-XXL [22] to process input prompts which consists of 4.6 billion parameters. For combinatorial and multi-task generalization experiments on simulated robotic manipulation, we train a first-frame conditioned video diffusion models on  $10 \times 48 \times 64$  videos (skipping every 8 frames) with 1.7B parameters and a temporal super resolution of  $20 \times 48 \times 64$  (skipping every 4 frames) with 1.7B parameters. The resolution of the videos are chosen so that the objects being manipulated (e.g., blocks being moved around) are clearly visible in the video. For the real world video results, we finetune the  $16 \times 40 \times 24$  (1.7B),  $32 \times 40 \times 24$  (1.7B),  $32 \times 80 \times 48$  (1.4B), and  $32 \times 320 \times 192$  (1.2B) temporal super resolution models pretrained on the data used by [19].
A.2 Inverse Dynamics Training Details
UniPi's inverse dynamics model is trained to directly predict the 7-dimensional controls of the simulated robot arm from an image observation mean squared error. The inverse dynamics model consists of a 3x3 convolutional layer, 3 layers of 3x3 convolutions with residual connection, a mean-pooling layer across all pixel locations, and an MLP layer of (128, 7) channels to predict the final controls. The inverse dynamics model is trained using the Adam optimizer with gradient norm clipped at 1 and learning rate 1e-4 for a total of 2M steps where linear warmup is applied to the first 10k steps.
A.3 Baselines Training Details
We describe the architecture details of various baselines below. The training details (e.g., learning rate, warm up, gradient clip) of each baseline follow those of the inverse dynamics model detailed above.
Transformer BC [6, 26]. We employ the same transformer architecture as the 10M model of [26] with 4 attention layers of 8 heads each and hidden size 512. We apply 4 layers of  $3 \times 3$  convolution with residual connection to extract image features, which, together with T5 text embeddings, are used as inputs to the transformer. We additionally experimented with vision transformer style linearization of the image patches similar to [26], but found the performance to be similar. We use a context length of 4 and skip every 4 frames similar to UniPi's inverse dynamics. We tried increasing the context length of the transformer to 8 but it did not help improve performance.
Transformer TT [25]. We use a similar transformer architecture as the Transformer BC baseline detailed above. Instead of predicting the immediate next control in the sequence as in Transformer BC, we predict the next 8 controls (skipping every 4 controls similar to other baselines) at the output layer. We have also tried autoregressively predicting the next 8 controls, but found the errors to accumulate quickly without additional discretization.
State-Based Diffusion [21]. For the state-based diffusion baseline, we use a similar architecture as UniPi's first-frame conditioned video diffusion, where instead of diffusing and generating future
image frames, we replicate future controls across different pixel locations and apply the same U-Net structure as UniPi to learn state-based diffusion models.
A.4 Details of the Combinatorial Planning Task
In the combinatorial planning tasks, we sample random 6 DOF poses for blocks, colored bowls, the final placement box. Blocks start off uncolored (white) and must be placed in a bowl to obtain a color. The robot then must manipulate and move the colored block to have the desired geometric relation in the placement box. The underlying action space of the agent corresponds to 6 joint values of robot plus a discrete contact action. When the contact action is active, the nearest block on the table is attached to the robot gripper (where for methods that predict continuous actions, we thresholded action prediction  $>0.5$  to correspond to contact). Given individual action predictions for different models, we simulate the next state of the environment by running the joint controller in Pybullet to try reach the predicted joint state (with a timeout of 2 seconds due to certain actions being physically infeasible). As only a subset of the video dataset contained action annotations, we trained the inverse-dynamics model on action annotations from 20k generated videos.
A.5 Details of the CLIPort Multi-Environment Task
In the CLIPort environment, we use the same action space as the combinatorial planning tasks and execute actions similarly using the built in joint controller in Pybullet. As our training data, we use a scripted agent on put-block-in-bowl-unseen-colors, packing-unseen-google-objects-seq, assembling-kits-seq-unseen-colors, stack-block-pyramid-seq-seen-colors, tower-of-hanoi-seq-seen-colors, assembling-kits-seq-seen-colors, tower-of-hanoi-seq-unseen-colors, stack-block-pyramid-seq-unseen-colors, packing-seen-google-objects-seq, packing-boxes-pairs-seen-colors, packing-seen-google-objects-group. As our test data, we used the environments put-block-in-bowl-seen-colors, packing-unseen-google-objects-group, packing-boxes-pairs-unseen-colors. We trained the inverse dynamics on action annotation across the 200k generated videos.
B.1 Additional Results on Combinatorial Generalization
Figure 10: Combinatorial Video Generation. Additional results on UniPi's generated videos for unseen language goals at test time.
B.2 Additional Results on Multi-Environment Transfer
Figure 11: Multitask Video Generation. Additional results on UniPi's generated video plans on different new tasks in the multitask setting.
B.3 Additional Results on Real-World Transfer
Figure 12: High Fidelity Plan Generation. Additional results on UniPi's high resolution video plans across different language prompts.