SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yilun Du\*†, Mengjiao Yang \*‡§, Bo Dai ‡¶, Hanjun Dai ‡, Ofir Nachum ‡, Joshua B. Tenenbaum †, Dale Schuurmans ‡¶, Pieter Abbeel §  
MIT† Google DeepMind ‡ UC Berkeley § Georgia Tech ¶ University of Alberta  
https://universal-policy.github.io/
A goal of artificial intelligence is to construct an agent that can solve a wide variety of tasks. Recent progress in text-guided image synthesis has yielded models with an impressive ability to generate complex novel images, exhibiting combinatorial generalization across domains. Motivated by this success, we investigate whether such tools can be used to construct more general-purpose agents. Specifically, we cast the sequential decision making problem as a text-conditioned video generation problem, where, given a text-encoded specification of a desired goal, a planner synthesizes a set of future frames depicting its planned actions in the future, after which control actions are extracted from the generated video. By leveraging text as the underlying goal specification, we are able to naturally and combinatorially generalize to novel goals. The proposed policy-as-video formulation can further represent environments with different state and action spaces in a unified space of images, which, for example, enables learning and generalization across a variety of robot manipulation tasks. Finally, by leveraging pretrained language embeddings and widely available videos from the internet, the approach enables knowledge transfer through predicting highly realistic video plans for real robots<sup>2</sup>.
Building models that solve a diverse set of tasks has become a dominant paradigm in the domains of vision and language. In natural language processing, large pretrained models have demonstrated remarkable zero-shot learning of new language tasks [1, 2, 3]. Similarly, in computer vision, models such as those proposed in [4, 5] have shown remarkable zero-shot classification and object recognition capabilities. A natural next step is to use such tools to construct agents that can complete different decision making tasks across many environments.
However, training such agents faces the inherent challenge of environmental diversity, since different environments operate with distinct state action spaces (e.g., the joint space and continuous controls in MuJoCo are fundamentally different from the image space and discrete actions in Atari). Such diversity hampers knowledge sharing, learning, and generalization across tasks and environments. Although substantial effort has been devoted to encoding different environments with universal tokens in a sequence modeling framework [6], it is unclear whether such an approach can preserve the rich knowledge embedded in pretrained vision and language models and leverage this knowledge to transfer to downstream reinforcement learning (RL) tasks. Furthermore, it is difficult to construct reward functions that specify different tasks across environments.
In this work, we address the challenges in environment diversity and reward specification by leveraging video (i.e., image sequences) as a universal interface for conveying action and observation
* denotes equal contribution. Correspondence to yilundu@mit.edu and sherryy@berkeley.edu.
2See video visualizations at https://universal-policy.github.io.
Figure 1: Text-Conditional Video Generation as Universal Policies. Text-conditional video generations enables us to train general purpose policies on wide sources of data (simulated, real robots and YouTube) which may be applied to downstream multi-task settings requiring combinatorical language generalization, long-horizon planning, or internet-scale knowledge.
behavior in different environments, and text as a universal interface for expressing task descriptions. In particular, we design a video generator as a planner that sequentially conditions on a current image frame and a text passage describing a current goal (i.e., the next high-level step) to generate a trajectory in the form of an image sequence, after which an inverse dynamics model is used to extract the underlying actions from the generated video. Such an approach allows the universal nature of language and video to be leveraged in generalizing to novel goals and tasks across diverse environments. Specifically, we instantiate the text-conditioned video generation model using video diffusion. A set of underlying actions are then regressed from the synthesized frames and used to construct a policy to implement the planned trajectory. Since the language description of a task is often highly correlated with control actions, text-conditioned video generation naturally focuses generation on action relevant parts of the video. The proposed model, UniPi, is visualized in Figure 1.
We have found that formulating policy generation via text-conditioned video synthesis yields the following advantages:
Combinatorial Generalization. The rich combinatorial nature of language can be leveraged to synthesize novel combinatorial behaviors in the environment. This enables the proposed approach to rearrange objects to new unseen combinations of geometric relations, as shown in Section 4.1.
Multi-task Learning. Formulating action prediction as a video prediction problem readily enables learning across many different tasks. We illustrate in Section 4.2 how this enables learning across language-conditioned tasks and generalizing to new ones at test time without finetuning.
Action Planning. The video generation procedure corresponds to a planning procedure where a sequence of frames representing actions is generated to reach the target goal. Such a planning procedure is naturally hierarchical: a temporally sparse sequence of images toward a goal can first be generated, before being refined with a more specific plan. Moreover, the planning procedure is steerable, in the sense that the plan generation can be biased by new constraints introduced at test-time through test-time sampling. Finally, plans are produced in a video space that is naturally interpretable by humans, making action verification and plan diagnosis easy. We illustrate the efficacy of hierarchical sampling in Table 2 and steerability in Figure 7.
Internet-Scale Knowledge Transfer. By pretraining a video generation model on a large-scale text-video dataset recovered from the internet, one can recover a vast repository of "demonstrations" that aid the construction of a text-conditioned policy in novel environments. We illustrate how this enables the realistic synthesis of robot motion videos from language instructions in Section 4.3.
The main contribution of this work is to formulate text-conditioned video generation as a universal planning strategy from which diverse behaviors can be synthesized. While such an approach departs from typical policy generation in RL, where subsequent actions to execute are directly predicted from a current state, we illustrate that UniPi exhibits notable generalization advantages over traditional policy generation methods across a variety of domains.
We first motivate a new abstraction, the Unified Predictive Decision Process (UPDP), as an alternative to the Markov Decision Process (MDP) commonly used in RL, and then show an instantiation of a UPDP with diffusion models.
2.1 Markov Decision Process
The Markov Decision Process [7] is a broad abstraction used to formulate many sequential decision making problems. Many RL algorithms have been derived from MDPs with empirical success [8, 9, 10], but existing algorithms are typically unable to combinatorially generalize across different environments. Such difficulty can be traced back to certain aspects of the underlying MDP abstraction:
i) The lack of a universal state interface across different control environments. In fact, since different environments typically have separate underlying state spaces, one would need to construct a complex state representation to represent all environments, making learning difficult.
ii) The explicit requirement of a real-valued reward function in an MDP. The RL problem is usually defined as maximizing the accumulated reward in an MDP. However, in many practical applications, how to design and transfer rewards is unclear and different across environments.
iii) The dynamics model in an MDP is environment and agent dependent. Specifically, the dynamics model  $T(s'|s,a)$  that characterizes the transition between states  $(s,s')$  under action  $a$  is explicitly dependent to the environment and action space of the agent, which can be significantly different between different agents and tasks.
2.2 Unified Predictive Decision Process
These difficulties inspire us to construct an alternative abstraction for unified sequential decision making across diverse environments. Our abstraction, termed Unified Predictive Decision Process (UPDP), exploits images as a universal interface across environments, texts as task specifiers to avoid reward design, and a task-agnostic planning module separated from environment-dependent control to enable knowledge sharing and generalization.
Formally, we define a UPDP to be a tuple  $\mathcal{G} = \langle \mathcal{X},\mathcal{C},H,\rho \rangle$ , where  $\mathcal{X}$  denotes the observation space of images,  $\mathcal{C}$  denotes the space of textual task descriptions,  $H\in \mathcal{N}$  is a finite horizon length, and  $\rho (\cdot |x_0,c):\mathcal{X}\times \mathcal{C}\to \Delta (\mathcal{X}^H)$  is a conditional video generator. That is,  $\rho (\cdot |x_o,c)\in \Delta (\mathcal{X}^H)$  is a conditional distribution over  $H$ -step image sequences determined by the first frame  $x_0$  and the task description  $c$ . Intuitively,  $\rho$  synthesizes  $H$ -step image trajectories that illustrate possible paths for completing a target task  $c$ . For simplicity, we focus on finite horizon, episodic tasks.
Given a UPDP  $\mathcal{G}$ , we define a trajectory-task conditioned policy  $\pi(\cdot|\{x_h\}_{h=0}^H, c): \mathcal{X}^{H+1} \times \mathcal{C} \to \Delta(\mathcal{A}^H)$  to be a conditional distribution over  $H$ -step action sequences  $\mathcal{A}^H$ . Ideally,  $\pi(\cdot|\{x_h\}_{h=0}^H, c)$  specifies a conditional distribution of action sequences that achieves the given trajectory  $\{x_h\}_{h=0}^H$  in the UPDP  $\mathcal{G}$  for the given task  $c$ . To achieve such an alignment, we will consider an offline RL scenario where we have access to a dataset of existing experience  $\mathcal{D} = \{(x_i, a_i)_{i=0}^{H-1}, x_H, c\}_{j=1}^n$  from which both  $\rho(\cdot|x_0, c)$  and  $\pi(\cdot|\{x_h\}_{h=0}^H, c)$  can be estimated.
In contrast to an MDP, a UPDP directly models video-based trajectories and bypasses the need to specify a reward function beyond the textual task description. Since the space of video observations  $\mathcal{X}^H$  and task descriptions  $\mathcal{C}$  are both naturally shared across environments and easily interpretable by humans, any video-based planner  $\rho(\cdot|x_0,c)$  can be conveniently reused, transferred and debugged. Another benefit of a UPDP over an MDP is that UPDP isolates the video-based planning with  $\rho(\cdot|x_0,c)$  from the deferred action selection using  $\pi(\cdot|\{x_h\}_{h=0}^H,c)$ . This design choice isolates planning decisions from action-specific mechanisms, allowing the planner to be environment and agent agnostic.
UPDP can be understood as implicitly planning over an MDP and directly outputting an optimal trajectory based on the given instructions. Such a UPDP abstraction bypasses reward design, state extraction and explicit planning, and allows for non-Markovian modeling of an image-based state space. However, learning a planner in UPDP requires videos and task descriptions, whereas traditional MDPs do not require such data, so whether an MDP or UPDP is more suitable for a given task depends on what types of training data are available. Although the non-Markovian model and the requirement of video and text data induce additional difficulties in UPDP comparing to MDP, it is possible to leverage existing large text-video models that have been pretrained on massive, web-scale datasets to alleviate these complexities.
2.3 Diffusion Models for UPDP
Let  $\tau = [x_1, \ldots, x_H] \in \mathcal{X}^H$  denote a sequence of images. We leverage the significant recent advances in diffusion models for capturing the conditional distribution  $\rho(\tau | x_0, c)$ , which we will
leverage as a text and initial-frame conditioned video generator in a UPDP. We emphasize that the UPDP formulation is also compatible with other probabilistic models, such as a variational autoencoder [11], energy-based model [12, 13], or generative adversarial network [14]. For completeness we briefly cover the core formulation at a high-level, but defer details to background references [15].
We start with an unconditional model. A continuous-time diffusion model defines a forward process  $q_{k}(\tau_{k}|\tau) = \mathcal{N}(\cdot ;\alpha_{k}\tau ,\sigma_{k}^{2}I)$ , where  $k\in [0,1]$  and  $\alpha_{k},\sigma_{k}^{2}$  are scalars with predefined schedules. A generative process  $p(\tau)$  is also defined, which reverses the forward process by learning a denoising model  $s(\tau_k,k)$ . Correspondingly  $\tau$  can be generated by simulating this reverse process with an ancestral sampler [16] or numerical integration [17]. In our case, the unconditional model needs to be further adapted to condition on both the text instruction  $c$  and the initial image  $x_0$ . Denote the conditional denoiser as  $s(\tau_k,k|c,x_0)$ . We leverage classifier-free guidance [18] and use  $\hat{s} (\tau_k,k|c,x_0) = (1 + \omega)s(\tau_k,k|c,x_0) - \omega s(\tau_k,k)$  as the denoiser in the reverse process for sampling, where  $\omega$  controls the strength of the text and first-frame conditioning.
3 Decision Making with Videos
Next we describe the proposed approach UniPi in detail, which is a concrete instantiation of the diffusion UPDP. UniPi incorporates the two main components discussed in Section 2, as shown in Figure 2: (i) a diffusion model for the universal video-based planner  $\rho (\cdot |x_0,c)$ , which synthesizes videos conditioned on the first frame and task descriptions; and (ii) a task-specific action generator  $\pi (\cdot |\{x_h\}_{h = 0}^H,c)$ , which infers actions sequences from generated videos through inverse dynamics modeling.
3.1 Universal Video-Based Planner
Encouraged by the recent success of text-to-video models [19], we seek to construct a video diffusion module as the trajectory
planner, which can faithfully synthesize future image frames given an initial frame and textual task description. However, the desired planner departs from the typical setting in text-to-video models [20, 19] which normally generate unconstrained videos given a text description. Planning through video generation is more challenging as it requires models to both be able to generate constrained videos that start at a specified image, and then complete the target task. Moreover, to ensure valid action inference across synthesized frames in a video, the video prediction module needs to be able to track the underlying environment state across synthesized video frames.
Conditional Video Synthesis. To generate a valid and executable plan, a text-to-video model must synthesize a constrained video plan starting at an initial image that depicts the initial configuration of the agent and environment. One approach to solve this problem is to modify the underlying test-time sampling procedure of an unconditional model, by fixing the first frame of the generated video plan to always begin at the observed image, as done in [21]. However, we found that this performed poorly and led to subsequent frames in the video plan to deviate significantly from the original observed image. Instead, we found it more effective to explicitly train a constrained video synthesis model by providing the first frame of each video as explicit conditioning context during training.
Trajectory Consistency through Tiling. Existing text-to-video models typically generate videos where the underlying environment state changes significantly during the temporal duration [19]. To construct an accurate trajectory planner, it is important that the environment remain consistent across all time points. To enforce environment consistency in conditional video synthesis, we provide, as additional context, the observed image when denoising each frame in the synthesized video. In particular, we re-purpose a temporal super-resolution video diffusion architecture, and provide as context the conditioned visual observation tiled across time, as the opposed to a low temporal-resolution video for denoising at each timestep. In this model, we directly concatenate each
Figure 2: Given an input observation and text instruction, we plan a set of images representing agent behavior. Images are converted to actions using an inverse dynamics model.
intermediate noisy frame with the conditioned observed image across sampling steps, which serves as a strong signal to maintain the underlying environment state across time.
Hierarchical Planning. When constructing plans in high dimensional environments with long time horizons, directly generating a set of actions to reach a goal state quickly becomes intractable due to the exponential blow-up of the underlying search space. Planning methods often circumvent this issue by leveraging a natural hierarchy in planning. Specifically, planning methods first construct coarse plans operating on low dimensional states and actions, which may then be refined into plans in the underlying state and action spaces. Similar to planning, our conditional video generation procedure likewise exhibits a natural temporal hierarchy. We first generate videos at a coarse level by sparsely sampled videos ("abstractions") of the desired behavior along the time axis. Then we refine the videos to represent valid behavior in the environment by super-resolving videos across time. Meanwhile, coarse-to-fine super-resolution further improves consistency via interpolation between frames.
Flexible Behavioral Modulation. When planning a sequence of actions to a given sub-goal, one can readily incorporate external constraints to modulate the generated plan. Such test-time adaptability can be implemented by composing a prior  $h(\tau)$  during plan generation to specify desired constraints across the synthesized action trajectory [21], which is also compatible with UniPi. In particular, the prior  $h(\tau)$  can be specified using a learned classifier on images to optimize a particular task, or as a Dirac delta on a particular image to guide a plan towards a particular set of states. To train the text-conditioned video generation model, we utilize the video diffusion algorithm in [19], where pretrained language features from T5 [22] are encoded. Please see Appendix A for the underlying architecture and training details.
3.2 Task Specific Action Adaptation
Given a set of synthesized videos, we may train a small task-specific inverse-dynamics model to translate frames into a set of actions as described below.
Inverse Dynamics. We train a small model to estimate actions given input images. The training of the inverse dynamics is independent from the planner and can be done on a separate, smaller and potentially suboptimal dataset generated by a simulator.
Action Execution. Finally, we generate an action sequence given  $x_0$  and  $c$  by synthesizing  $H$  image frames and applying the learned inverse-dynamics model to predict the corresponding  $H$  actions. Inferred actions can then be executed via closed-loop control, where we generate  $H$  new actions after each step of action execution (i.e., model predictive control), or via open-loop control, where we sequentially execute each action from the initially inferred action sequence. For computational efficiency, we use an open-loop controller in all our experiments in this paper.
The focus of these experiments is to evaluate UniPi in terms of its ability to enable effective, generalizable decision making. In particular, we evaluate
(1) the ability to combinatorially generalize across different subgoals in Section 4.1,
(2) the ability to effectively learn and generalize across many tasks in Section 4.2,
(3) the ability to leverage existing videos on the internet to generalize to complex tasks in Section 4.3.
See experimental details in Appendix A. Additional results are given in Appendix B and videos in the supplement.
4.1 Combinatorial Policy Synthesis
First, we measure the ability of UniPi to combinatorially generalize to different language tasks.
Setup. To measure combinatorial generalization, we use the combinatorial robot planning tasks in [23]. In this task, a robot must manipulate blocks in an environment to satisfy language instructions, i.e., put a red block right of a cyan block. To accomplish this task, the robot must first pick up a white block, place it in the appropriate bowl to paint it a particular color, and then pick up and place the block in a plate so that it satisfies the specified relation. In contrast to [23] which uses
Model	Seen	Novel
Place	Relation	Place	Relation
State + Transformer BC [24]	19.4 ± 3.7	8.2 ± 2.0	11.9 ± 4.9	3.7 ± 2.1
Image + Transformer BC [24]	9.4 ± 2.2	11.9 ± 1.8	9.7 ± 4.5	7.3 ± 2.6
Image + TT [25]	17.4 ± 2.9	12.8 ± 1.8	13.2 ± 4.1	9.1 ± 2.5
Diffuser [21]	9.0 ± 1.2	11.2 ± 1.0	12.5 ± 2.4	9.6 ± 1.7
UniPi (Ours)	59.1 ± 2.5	53.2 ± 2.0	60.1 ± 3.9	46.1 ± 3.0
Table 1: Task Completion Accuracy in Combinatorial Environments. UniPi generalizes to seen and novel combinations of language prompts in Place (e.g., place X in Y) and Relation (e.g., place X to the left of Y) tasks.
Figure 3: Combinatorial Video Generation. Generated videos for unseen language goals at test time.
pre-programmed pick and place primitives for action prediction, we predict actions in the continuous robotic joint space for both the baselines and our approach.
We split the language instructions in this environment into two sets: one set of instructions (70%) that is seen during training, and another set (30%) that is only seen during testing. The precise locations of individual blocks, bowls, and plates in the environment are fully randomized in each environment iteration. We train the video model on 200k example videos of generated language instructions in the train set. Details of this environment can be found in Appendix A. We constructed demonstrations of videos in this task by using a scripted agent.
Baselines. We compare the proposed approach with three separate representative approaches. First, we compare to existing work that uses goal-conditioned transformers to learn across multiple environments, where goals can be specified as episode returns [26], expert demonstrations [6], or text and images [24]. To represent these baselines, we construct a transformer behavior cloning (BC) agent to predict the subsequent action to execute given the task description and either the visual observation (Image + Transformer BC) or the underlying robot joint state (State + Transformer BC). Second,
given that our approach regresses a sequence of actions to execute, we further compare with transformer models that regress a sequence of future actions to execute, similar to the goal-conditioned behavioral cloning of the Trajectory Transformer [25] (Image + TT). Finally, to highlight the importance of the video-as-policy approach, we compare UniPi with learning a diffusion process that, conditioned on an image observation, directly infers future robot actions in the joint space (as opposed to diffusing future image frames), corresponding to [21, 27]. For both our method and each baseline, we condition the policy on encoded language instructions using pretrained T5 embeddings. Note that in this setting, existing offline reinforcement learning baselines are not directly applicable as we do not have access to the reward functions in the environment.
Metrics. To compare UniPi with baselines, we measure final task completion accuracy across new instances of the environment and associated language prompts. We subdivide the evaluation along two axes: (1) whether the language instruction has been seen during training and (2) whether the