Learning Universal Policies via Text-Guided Video Generation

In plain words

Most robot controllers pick the next motor command directly from what the robot sees. This paper does something different: it asks a video generator to imagine the next few seconds of the task, as a short clip that starts from the current camera image and follows a text instruction such as "put a red block right of a cyan block". A second, small model then looks at the imagined frames and works out which motor commands would produce them.

The appeal is that text and video are shared across very different robots and simulators, so one planner can be reused, and the imagined video can be inspected by a person before anything moves. The authors report that this helps the robot handle instructions it was not trained on, and that a video model pretrained on general internet clips can produce plausible plans for a real robot arm.

TL;DR

UniPi plans with a text-conditioned video diffusion model. It is trained to continue an observed first frame into a video of the task, and an inverse-dynamics model turns the generated frames into joint-space actions. The paper reports higher task completion than state-based diffusion and transformer behavior cloning on simulated manipulation, and, for a real robot, a benefit from pretraining on internet video. All numbers are author-reported.

Problem

Standard reinforcement learning formulates a task as a Markov decision process with a reward function and a dynamics model that depend on each environment's state and action space. The authors argue this makes knowledge hard to share across environments, because state spaces differ, rewards are hard to specify and transfer, and the dynamics are tied to one agent §2.1. Their alternative, the Unified Predictive Decision Process (UPDP), treats images as the shared observation space and text as the task specifier, and separates planning (a video generator) from action selection (an inverse-dynamics-style policy) §2.2. The method is motivated by the need to combinatorially generalize to new language goals and to train across many environments §1.

The paper assumes an offline setting: a dataset of observation-action sequences with a text description per episode, from which both the planner and the policy are estimated §2.2. Without video-and-text data, the approach does not apply.

Method

The method has two trained parts. The planner is a diffusion model that generates an image sequence conditioned on the first observed frame and the task text §3.1. The policy is a task-specific inverse-dynamics model that maps generated frames to actions §3.2.

Planner: conditional video synthesis. The planner is a text-to-video diffusion model adapted to start from a given image. Instead of fixing the first frame at test time, as Diffuser-style sampling does, the authors train the model with the first frame as explicit conditioning §3.1. Two other choices are part of the design:

1. Tiling the observed frame across time. The observed image is repeated at every frame index and concatenated with each noisy frame at every denoising step, which keeps the environment consistent across the video §3.1 app. A.1.
2. Coarse-to-fine temporal hierarchy. A coarse video is generated at sparse time steps, then a temporal super-resolution model fills in the frames §3.1 app. A.1.

The text is encoded with pretrained T5 features §3.1 app. A.1. Sampling uses classifier-free guidance, with the denoiser formula written inline in §2.3 (not numbered in the source).

# Guided denoiser for one reverse-diffusion step  [§2.3]
# w sets the strength of text and first-frame conditioning
def guided_denoise(noisy_video, noise_level, text, first_frame):
    # Denoiser with text and first frame as conditioning  [§2.3]
    conditional = denoiser(noisy_video, noise_level, text, first_frame)
    # Denoiser with the conditioning dropped  [§2.3]
    unconditional = denoiser(noisy_video, noise_level)
    # Extrapolate away from the unconditional prediction  [§2.3]
    return (1 + w) * conditional - w * unconditional

Denoiser input with tiling. How the first frame enters the video U-Net, as described in the appendix:

# Build the denoiser input from noisy frames and the observed frame  [app. A.1]
def denoiser_input(noisy_frames, first_frame):
    # Repeat the first observed frame at every future frame index  [app. A.1]
    tiled = repeat_over_time(first_frame, num_frames)
    # Concatenate channel-wise with each noisy frame  [app. A.1] [§3.1]
    return concat_channels(noisy_frames, tiled)

Policy: inverse dynamics and open-loop execution. An inverse-dynamics model is trained separately, on a smaller and possibly suboptimal dataset from a simulator, to regress the actions behind an observed image §3.2 app. A.2. At run time the planner produces H frames, the inverse-dynamics model predicts H actions, and the actions are executed in order without re-planning. The authors use this open-loop mode in all experiments for speed, though a closed-loop (model predictive control) mode is described §3.2.

# Plan a video, infer actions, execute open-loop  [§3.2] [§3]
# Input: observed image and task text; Output: none (acts in the env)
def plan_and_act(observed_image, task_text):
    # Encode the text with pretrained T5 features  [§3.1] [app. A.1]
    text_features = t5_encode(task_text)
    # Generate H frames conditioned on the observed image  [§3.1]
    video_plan = sample_video(observed_image, text_features)
    # Infer one action per planned frame  [§3.2] [app. A.2]
    actions = [inverse_dynamics(frame) for frame in video_plan]
    # Execute the whole sequence open-loop, without re-planning  [§3.2]
    for action in actions:
        env_step(action)

Training of the video model (as stated). Two diffusion models are trained per setting: a first-frame-conditioned video model and a temporal super-resolution model. The architecture follows the video diffusion work of Ho et al. (2022) with temporal convolutions rather than temporal attention app. A.1.

Hyperparameters

Values are as stated in the paper or its appendix. Where the paper is silent, the table says so.

Results

Author-reported. All results are task completion accuracy in the authors' simulated or real environments, with the ± values as printed in the source tables (the paper does not define them).

Combinatorial generalization (simulated robot, 70/30 split of instructions; Place = "put X in Y", Relation = "put X left of Y") tab. 1 §4.1:

Ablation on seen Place and Relation tasks tab. 2 §4.1:

Multi-environment transfer (10 training tasks, 3 unseen test tasks, CLIPort-style suite) tab. 3 §4.2:

Real-world video generation (Bridge data, 24x40 model, 32 samples for FID/FVD, success from a separately trained success classifier on the last frame) tab. 4 §4.3:

The paper reports generalization to novel task commands that are absent from the Bridge data and some robustness to background edits, shown only as figures §4.3.

Paper vs code

No official code.

Searched, in order: the paper's linked project page (universal-policy.github.io, which returned no content on two fetches); GitHub repository search for unipi, universal policies video and universal-policy; a search for unipi restricted to the first author's account and to google-deepmind; and the first page of the first author's public repositories (truncated, no UniPi-like name). No repository belonging to the authors was found.

Pitfalls

• First-frame conditioning is trained, not applied at test time only. The paper says test-time clamping of the first frame performed poorly and trains the model with the first frame as context instead §3.1. Re-implementations that only clamp the first frame deviate from the reported method.
• Frame-count notation is undefined. "10x48x64" and "20x48x64" with "frame skip 8" and "frame skip 4" are given in app. A.1, but the paper does not state the axis order or the exact subsampling rule. Read them as (frames, height, width) only after confirming with the authors.
• Inverse-dynamics input is ambiguous. app. A.2 describes the model as predicting controls "from an image observation", while §3.2 applies it to a generated video of H frames, and the baseline text uses "context length 4 and skip every 4 frames similar to the inverse dynamics". Whether the model takes a single frame or a short frame stack is not stated explicitly.
• Unconditional branch of guidance. The guidance formula in §2.3 writes the unconditional denoiser as s(τ_k, k), dropping the text. Whether the first frame is also dropped in that branch is not stated.
• Guidance and sampling hyperparameters are missing. The guidance weight ω, the sampler, step count, noise schedule and conditioning dropout rate are not stated in the paper or appendix. Any reimplementation has to choose them.
• Open-loop execution. All experiments run open-loop §3.2, so errors in the generated plan are not corrected during execution. The paper does not report closed-loop results.
• Reported differences from prior baselines. The authors predict continuous joint actions for their baselines, unlike the pick-and-place primitives of the PDSketch environment §4.1. Numbers are not directly comparable with prior work on that benchmark.
• Text and figures disagree on the real-world metrics. §4.3 says pretraining gives "significantly higher FID and FVD", but tab. 4 marks FID and FVD as lower-is-better and the pretrained row has the lower values. Read the table.
• ± values are undefined. The tables print ± for every result, but the paper does not say whether these are standard deviations over seeds, episodes or instructions, nor how many trials were run.
• Success classifier for real-world results. The 72.6% and 77.1% success rates depend on a classifier the authors trained on the last generated frame §4.3. The classifier's own accuracy is not reported.
• Compute. Training used 256 TPU-v4 chips for the first-frame and temporal super-resolution models app. A.1. The paper says the diffusion process "can take a minute to generate highly photorealistic videos" and that distilling it into a faster sampler gave a 16x speed-up in initial experiments, which is not part of the reported method §5.
• Partially observable settings are not covered. The authors note that video models may hallucinate objects or motion in partially observable environments §5.

Lineage

• Builds on: video diffusion models of Ho et al. (2022), including the video U-Net and temporal super-resolution setup app. A.1; classifier-free guidance of Ho and Salimans (2022) §2.3; the T5 text encoder §3.1 app. A.1.
• Compared against: Diffuser, a state-based diffusion planner of Janner et al. (2022) §4.1; transformer behavior cloning of Lee et al. (2022) and Brohan et al. (2022) app. A.3; Trajectory Transformer of Janner et al. (2021) app. A.3; CLIPort §4.2; PDSketch of Mao et al. (2022) §4.1.
• Superseded by: not known from the content of this wiki.