Source-linked AI summary

Zero-Shot Robotic Manipulation with Pretrained Image-Editing Diffusion Models

Kevin Black, Mitsuhiko Nakamoto, Pranav Atreya, Homer Walke, Chelsea Finn, Aviral Kumar, Sergey Levine

arXiv:2310.10639v1cs.RO

TL;DR

Generalist robots need to manipulate novel objects and scenarios despite limited coverage in robot training data. SuSIE uses a finetuned image-editing diffusion model to propose visual subgoals for a low-level goal-conditioned policy, and reports state-of-the-art CALVIN results plus stronger real-world performance than several strong baselines. Its main limitation is that the separately trained subgoal model does not account for the low-level policy’s actual capabilities.

  • Problem

    Generalist robots must recognize and reason about novel objects and scenarios that may be absent from their training data.

  • Method

    SuSIE finetunes a pretrained image-editing model on video to generate language-conditioned future subgoal observations, then trains a robot-data goal-conditioned policy to reach them.

  • Results

    SuSIE achieves state-of-the-art CALVIN results and outperforms language-conditioned, privileged-goal, and RT-2-X baselines on real-world manipulation tasks.

  • Takeaways & Limitations

    Subgoal generation incorporates pretrained semantic and visual understanding into precise low-level robot control and supports generalization to novel real-world objects, distractors, and scenes.

  • Takeaways & Limitations

    The diffusion model and low-level policy are trained separately, so the subgoal model assumes data-reachable states are reachable by the policy; performance is often bottlenecked by the low-level policy.

Abstract

from arXiv · show

If generalist robots are to operate in truly unstructured environments, they need to be able to recognize and reason about novel objects and scenarios. Such objects and scenarios might not be present in the robot's own training data. We propose SuSIE, a method that leverages an image-editing diffusion model to act as a high-level planner by proposing intermediate subgoals that a low-level controller can accomplish. Specifically, we finetune InstructPix2Pix on video data, consisting of both human videos and robot rollouts, such that it outputs hypothetical future "subgoal" observations given the robot's current observation and a language command. We also use the robot data to train a low-level goal-conditioned policy to act as the aforementioned low-level controller. We find that the high-level subgoal predictions can utilize Internet-scale pretraining and visual understanding to guide the low-level goal-conditioned policy, achieving significantly better generalization and precision than conventional language-conditioned policies. We achieve state-of-the-art results on the CALVIN benchmark, and also demonstrate robust generalization on real-world manipulation tasks, beating strong baselines that have access to privileged information or that utilize orders of magnitude more compute and training data. The project website can be found at http://rail-berkeley.github.io/susie .

1 INTRODUCTION

SuSIE addresses the challenge of controlling robots around novel objects and scenarios by using pretrained image-editing models for semantic subgoal generation. A low-level goal-reaching policy then executes those subgoals, yielding strong benchmark and real-world performance.

  • Motivation: Generalist robots must recognize and reason about novel objects and scenarios beyond their training distribution.The motivation includes both semantic understanding and the physical ability to manipulate unfamiliar objects.
  • Motivation: Existing uses of pretrained vision-language models improve semantic information in robot learning, but their effect on low-level policy execution remains unclear.The paper distinguishes high-level visual generalization from precise policy execution.
  • Approach: SuSIE finetunes a pretrained image-editing model on video to generate hypothetical future frames from the current frame and a language task description.The generated frames serve as intermediate subgoals rather than direct actions.
  • Approach: A low-level goal-reaching policy executes actions toward each generated subgoal, allowing the high-level semantic model to transfer across data with different physical embodiments.The division separates semantic reasoning from precise robot actuation.
  • Results: SuSIE achieves state-of-the-art CALVIN results and outperforms prior methods, privileged-information oracles, and RT-2-X on real-world tasks involving novel objects, distractors, and scenes.RT-2-X uses a 55-billion-parameter model and an order of magnitude more robot data.

2 RELATED WORK

Prior work incorporates pretrained models into visual augmentation, language- or video-conditioned planning, and model-based control. SuSIE differs by generating only reachable visual waypoints and using an image-only low-level policy for zero-shot manipulation.

  • Pretrained semantic models: Prior approaches use pretrained models to augment scenes or initialize language- and image-based robot policies.These methods target semantic scene understanding or directly produce action sequences.
  • Pretrained semantic models: The paper reports that vision-language pretraining does not necessarily improve precise low-level execution, whereas SuSIE uses subgoals to steer motion.The comparison concerns object localization and gripper precision.
  • Planning and control: Text-conditioned low-level policies face a grounding problem that SuSIE avoids because its low-level controller receives images rather than language.Language is used by the high-level subgoal generator.
  • Planning and control: Unlike UniPi and HiP, SuSIE predicts the next waypoint instead of generating every frame in a complete rollout.The paper presents waypoint prediction as reducing the generative burden.
  • Model-based control: Classical model-based reinforcement-learning methods generally lack pretrained semantic knowledge and zero-shot generalization to new tasks.SuSIE also decouples its action-agnostic generative model from policy training.
  • Model-based control: Single-task transformer and diffusion planning methods require target-task training data, unlike SuSIE’s zero-shot planning approach.Their policies are trained using data from the task they must solve.

3 PRELIMINARIES AND PROBLEM STATEMENT

The problem setting is language-conditioned robotic control over new scenes and commands, with three data sources differing in language and action annotations. Evaluation measures task success from a new initial scene.

  • Problem setting: The task is to accomplish a manipulation behavior described by a novel language command.At test time, the command and scene may be new.
  • Data sources: Language-labeled video data contains task descriptions without robot actions, while language-labeled robot data contains both language labels and actions.These sources support different parts of the learning pipeline.
  • Data sources: Unlabeled robot data contains action trajectories without language annotations.It is exemplified by play data.
  • Trajectory representation: Each trajectory consists of image or state sequences and executed actions under standard Markov decision-process assumptions.Language labels describe the task completed by the trajectory.
  • Evaluation: At test time, methods receive a new scene and natural-language task description, and are evaluated by task success rate.The success rate measures completion starting from the new initial scene.

4 SUSIE: SUBGOAL SYNTHESIS VIA IMAGE EDITING

SuSIE separates language-guided subgoal synthesis from image-based goal reaching. A finetuned diffusion model proposes intermediate visual states, and a goal-conditioned policy repeatedly acts toward refreshed subgoals.

  • Overview: SuSIE decouples robot control into generating future subgoals and executing them with a low-level controller.This separation is intended to use semantic Internet knowledge without requiring the pretrained model to select low-level actions.
  • Phase I: Subgoal synthesis: The subgoal model is initialized from a text-guided image-editing model and finetuned on language-labeled video and robot trajectories.It models pθ(s_edited | s_orig, l), mapping an initial image and language label to a future subgoal image.
  • Phase I: Subgoal synthesis: The training distribution selects future subgoals between k_min and k_max steps ahead, balancing reachability with task progress.The distribution q(j | i) controls which future states are used as training subgoals.
  • Phase II: Goal reaching: The low-level controller is language-agnostic and is trained as a goal-conditioned behavioral-cloning policy to reach nearby future frames.It uses robot data and only needs to handle subgoals within the modeled temporal horizon, with k_delta providing overhead for imperfect predictions.
  • Test-time control: At test time, SuSIE samples a subgoal from the current image and language command, executes the policy for k_test steps, then refreshes the subgoal iteratively.This alternating procedure continues until the time limit.
  • Test-time control: Algorithm 1 specifies the alternating loop: sample a subgoal, predict and execute actions toward it, observe the robot state, and repeat.The procedure operates with a time limit and a subgoal-sampling interval.

5 EXPERIMENTAL EVALUATION

The experiments test SuSIE’s zero-shot generalization and precision across simulated and real-world manipulation settings, including novel environments, objects, and language commands. SuSIE outperforms prior methods in simulation and achieves the strongest real-world performance, with results attributing gains to subgoal guidance, pretrained video data, and improved low-level precision.

  • Comparisons: The experiments compare SuSIE with language-conditioned policies and other pretrained vision-language approaches to assess generalization, precision, and the value of Internet and non-robot video data.The real-world comparison includes RT-2-X, a 55-billion-parameter model trained on more than 1.1 million robot trajectories.
  • Experimental scenarios: The evaluation tests zero-shot manipulation in CALVIN and across three real-world scenes with progressively harder semantic generalization.The real-world scenes vary in visual departure from training data and mixtures of seen and unseen objects; CALVIN trains on A, B, and C and tests on D.
  • Simulation results: SuSIE achieves superior zero-shot CALVIN performance when training on environments A, B, and C and testing on environment D, significantly outperforming the UniPi reimplementation.Performance is measured by success rates for completing language-instruction chains.
  • Real-world results: SuSIE achieves the best real-world performance across scenes, including uniquely consistent bell-pepper grasping in Scene B.It outperforms RT-2-X despite RT-2-X using substantially more robot and Internet data.
  • Real-world results: In Scene C, SuSIE remains best while RT-2-X is close behind because the objects are easier to grasp and low-level precision matters less.Observed failures for both methods were generally imprecise grasps or early dropping rather than semantic misunderstanding; toothpaste was the hardest object.
  • Precision and ablations: SuSIE outperforms the privileged Oracle GCBC baseline on average, indicating that generated subgoals improve both semantic generalization and low-level precision.Oracle GCBC still cannot reliably grasp the bell pepper in Scene B because of insufficient low-level precision.
  • Pretraining and video data: Internet-scale pretrained weights and Something-Something cotraining improve subgoal quality and out-of-distribution generalization, especially for tasks involving the wooden bowl.BridgeData-only models hallucinate background items or misunderstand out-of-distribution tasks more often, while full-model subgoals yield a small average execution improvement.

6 DISCUSSION AND FUTURE WORK

SuSIE uses diffusion-generated image subgoals to guide low-level goal-conditioned control, improving semantic generalization and precision across simulation and real-world tasks. Its remaining bottleneck is the separately trained low-level policy, which may not match the diffusion model’s assumed reachability.

  • Discussion: SuSIE generates image subgoals to guide a low-level goal-conditioned policy, improving zero-shot generalization and overall policy precision.The subgoal model imports semantic benefits from pretraining while providing fine-grained guidance to low-level control.
  • Discussion: Figure 4 compares Bridge + Smthn-Smthn, Bridge Only, and From Scratch models on in-distribution and unseen out-of-distribution scenes.The comparison emphasizes the roles of Internet-scale pretraining and video cotraining in subgoal quality.
  • Discussion: SuSIE improves over prior techniques on CALVIN and generalization and low-level control on real-world manipulation tasks.In real-world experiments, it outperforms language-conditioned behavioral cloning, a privileged goal-conditioned policy, and RT-2-X.
  • Future work: The diffusion model and low-level policy are trained separately, so the diffusion model assumes data-reachable subgoals are reachable by the policy.The authors hypothesize that modeling the low-level policy’s capabilities could improve performance.

A IMPLEMENTATION DETAILS

The implementation appendix documents goal-sampling hyperparameters in Table 5.

  • Implementation details: Table 5 documents the goal-sampling hyperparameters used in the experiments.The passage provides no further hyperparameter values.

A.1 SUSIE IMPLEMENTATION DETAILS

SuSIE fine-tunes InstructPix2Pix on video and robot data, then uses diffusion-based goal-reaching control with specified sampling and guidance settings.

  • SuSIE implementation: InstructPix2Pix is fine-tuned for 40k steps with AdamW, learning rate 1e-4, batch size 1024, and EMA decay 0.999.Training uses 800 warmup steps, weight decay 0.01, and a single v4-64 TPU pod for 17 hours.
  • SuSIE implementation: Real-world training mixes 60% Something-Something video data with 40% BridgeData V2 robot data, while simulation uses CALVIN exclusively.The data mixture differs between simulation and real-world experiments.
  • SuSIE implementation: Inference uses image guidance weight 2.5, text guidance weight 7.5, and 50 DDIM sampling steps.These settings govern image-editing diffusion at test time.
  • SuSIE implementation: The goal-reaching policy stacks observation and goal images, encodes them with ResNet-50, and conditions a diffusion process modeling the action distribution.Diffusion-based policies are used to capture multimodality in robot data.
  • SuSIE implementation: Temporal ensembling is used instead of receding-horizon resampling because switching between sampled action modes produced jerky behavior.The passage contrasts the two action-sequence execution strategies.

A.2.1 MOO

This appendix describes MOO’s object-mask pipeline and the UniPi comparison, including implementation choices and observed limitations.

  • A.2.1 MOO: MOO conditions a language policy on a four-channel image containing an object mask produced from OWL-ViT bounding boxes.The target object and destination are extracted from language before detection.
  • A.2.1 MOO: BridgeData V2’s diverse annotations require GPT-3.5-Turbo-Instruct to extract target entities before OWL-ViT mask generation.Ground-truth masks are manually supplied at test time through clicks on the initial camera image.
  • A.2.1 MOO: Training-time MOO bounding boxes are sometimes incorrect because BridgeData V2 is unstructured, potentially confusing the policy.Figure 5 shows examples of the noisy extracted boxes.
  • A.2.1 MOO: UniPi generates full frame sequences with video diffusion and extracts actions using an inverse dynamics model applied to adjacent frames.The predicted actions are executed before the process repeats from a later generated frame.
  • A.2.1 MOO: The real-world UniPi video model does not reliably produce high-quality generations, although its simulation quality suffices for CALVIN tasks.The real-world generation examples are deferred to Appendix C.

B EXPERIMENTAL SETUP

The experiments evaluate SuSIE on CALVIN and real-world manipulation tasks using repeated trials, language-instruction chains, and varied object placements. Additional CALVIN comparisons use manually reset oracle goal states and different trajectory counts.

  • CALVIN: In CALVIN, each trial gives the policy 360 timesteps to complete a simulator-selected chain of 5 language instructions.The next instruction is selected based on environment affordances after the previous instruction is completed.
  • CALVIN: CALVIN results are averaged over 100 trajectories, while UniPi video-model results use 25 trajectories because video generation is computationally costly.The short 10-timestep generation horizon also requires frequent regeneration.
  • Oracle comparison: Oracle GCBC comparisons use simulator ground-truth goal images obtained by manually resetting internal states, and report results over 200 evaluation trajectories.This procedure is feasible only for selected CALVIN tasks listed in Table 3.
  • Real world: Real-world evaluation covers 9 tasks with 10 trials per task, varying object positions and orientations while keeping evaluation configurations consistent across methods.In Scene A, the plate, eggplant, and carrot positions and orientations change every trial, while the pot remains in the drying rack.

C QUALTITATIVE EXAMPLES OF VIDEO GENERATION

Qualitative examples show that competing video-generation models often produce blurry, hallucinated, or physically inconsistent outputs, whereas SuSIE examples visualize subgoals and successful CALVIN rollouts. These generation failures can confuse or degrade low-level control.

  • UniPi (HiP): UniPi (HiP) fails to generate a clear robot gripper and produces consistently blurry, visually poor videos.Because of this failure, the model was not evaluated on the real robot.
  • UniPi (Ours): UniPi (Ours) succeeds in a few trials, but most generated videos contain hallucination and physical inconsistency that confuse the low-level policy.These generation problems led to poor performance in the real-world evaluation.
  • SuSIE: SuSIE qualitative examples visualize generated subgoals alongside trajectory rollouts for a successful 5-instruction CALVIN language chain.The figure presents both the intermediate subgoals and the resulting robot behavior.
Loading 2310.10639v1…