Source-linked AI summary
Scaling Robot Learning with Semantically Imagined Experience
Tianhe Yu, Ted Xiao, Austin Stone, Jonathan Tompson, Anthony Brohan, Su Wang, Jaspiar Singh, Clayton Tan, Dee M, Jodilyn Peralta, Brian Ichter, Karol Hausman, Fei Xia
TL;DR
Robot learning needs diverse data, but collecting large real-world datasets requires substantial human effort or engineering. ROSIE uses text-guided diffusion and inpainting to augment existing robotic data with novel objects, tasks, backgrounds, and distractors. Real-world experiments show improved unseen-task performance, robustness to novel scenes and distractors, and success-detection generalization.
Problem
Robot policies generalize less than large-scale vision-language models, while collecting diverse real-world robot data is difficult to scale.
Method
ROSIE parses novel instructions, selects image regions, and uses text-guided diffusion inpainting to augment existing robotic episodes without additional robot data collection.
Results
ROSIE-trained policies solve unseen tasks, become more robust to distractors and backgrounds, and improve out-of-distribution success detection.
Takeaways & Limitations
Semantically imagined image experience provides a way to expand robotics datasets and improve learned-model robustness without extra real-world data collection.
Takeaways & Limitations
ROSIE augments object and scene appearance but does not generate new motions, and per-frame augmentation can reduce temporal consistency.
Abstract
from arXiv · showhide
Recent advances in robot learning have shown promise in enabling robots to perform a variety of manipulation tasks and generalize to novel scenarios. One of the key contributing factors to this progress is the scale of robot data used to train the models. To obtain large-scale datasets, prior approaches have relied on either demonstrations requiring high human involvement or engineering-heavy autonomous data collection schemes, both of which are challenging to scale. To mitigate this issue, we propose an alternative route and leverage text-to-image foundation models widely used in computer vision and natural language processing to obtain meaningful data for robot learning without requiring additional robot data. We term our method Robot Learning with Semantically Imagened Experience (ROSIE). Specifically, we make use of the state of the art text-to-image diffusion models and perform aggressive data augmentation on top of our existing robotic manipulation datasets via inpainting various unseen objects for manipulation, backgrounds, and distractors with text guidance. Through extensive real-world experiments, we show that manipulation policies trained on data augmented this way are able to solve completely unseen tasks with new objects and can behave more robustly w.r.t. novel distractors. In addition, we find that we can improve the robustness and generalization of high-level robot learning tasks such as success detection through training with the diffusion-based data augmentation. The project's website and videos can be found at diffusion-rosie.github.io
1 Introduction
Robot policies remain less general than large-scale vision-language models because diverse robot data spanning motor skills, objects, and visual domains is difficult to collect. ROSIE addresses this bottleneck by using text-guided diffusion and inpainting to create semantically novel robotic experiences from existing data.
- Diverse robot data is limited across motor skills, objects, and visual domains, constraining policy generalization.
- Collecting large real-world datasets requires either engineering-heavy autonomous schemes or laborious human teleoperation.One cited effort used 17 months and 13 robots for 130k demonstrations; another used 7 robots and 16 months for 800k autonomous episodes.
- Traditional image augmentations do not provide novel robotic experiences for mastering new skills or generalizing to semantically new environments.
- Text-guided diffusion models augment task semantics through natural-language interfaces and generate photorealistic objects zero-shot from internet-scale data.
- Replacing objects in existing demonstrations can create semantically labeled tasks while preserving task-relevant scene content, such as a grasp.For example, a green chip bag near an orange can become a yellow chip bag near a peach.
- ROSIE parses novel instructions, identifies regions to alter, and uses inpainting while leaving the rest of each image untouched.The resulting augmentations include novel tasks, distractors, and semantically meaningful backgrounds.
2 Related Work
Prior robot-learning scaling approaches rely on simulation, teleoperation, autonomous collection, or scripting. ROSIE instead directly applies text-conditioned diffusion to real-world data, generating novel tasks and distractors without depth or manually specified masks.
- Robot-learning data has been scaled through simulation, teleoperation, autonomous collection, and scripted behaviors.
- ROSIE directly augments real-world data with diffusion-based photorealistic image manipulation rather than changing simulated parameters.
- Unlike related methods, ROSIE generates novel distractors and tasks without depth, manually provided masks, or object meshes.
- ROSIE automatically selects inpainting regions with text guidance and uses text-guided diffusion to produce realistic augmentations.
3 Preliminaries
The paper uses diffusion models for image inpainting and learns vision-language-conditioned robot policies through imitation learning. Policies map image observations and language instructions to action distributions using behavioral cloning.
- Diffusion models iteratively denoise Gaussian noise into distribution samples and can fill masked image regions through inpainting.
- The robot dataset contains episodes of observations, actions, and language instructions describing the demonstrated task.
- Behavioral cloning minimizes the negative log-likelihood of demonstrated actions to learn a policy conditioned on image observations and language.
- The RT-1 architecture uses FiLM-conditioned EfficientNet, TokenLearner, and a Transformer to output actions.
4 Robot Learning with Semantically Imagened Experience (ROSIE)
ROSIE scales robot data generation through semantic image augmentation: it localizes regions, edits them with text-guided inpainting, and trains policies on the resulting episodes. The pipeline supports targeted object, distractor, and background changes while preserving trajectory consistency and enables manipulation in unseen settings.
- Pipeline: ROSIE augments existing robotic episodes by adding semantically different circumstances and applying the changes consistently across every trajectory frame.The pipeline first modifies the instruction, identifies a relevant region, and inpaints it throughout the episode.
- Augmentation Examples: ROSIE can replace in-hand objects, backgrounds, and distractors with novel content while retaining the manipulation trajectory and actions.Examples include microfiber cloths replacing a chip bag and unseen kitchen surfaces replacing the regular tabletop or drawer.
- Augmentation Region Localization: The pipeline localizes target or distractor regions with open-vocabulary segmentation before applying image edits.OWL-ViT detects instructed objects, tables, and passthrough objects so sampled masks avoid existing objects.
- Augmentation Text Proposal: ROSIE uses hand-engineered or LLM-proposed prompts to specify novel objects and semantically meaningful distractors.Hand-engineered prompts target out-of-training-distribution objects, while GPT-3 proposes objects with detailed visual descriptions.
- Diffusion Model for Text-Guided Inpainting: Imagen Editor performs masked, text-guided inpainting that generates photorealistic content while preserving unmasked image regions.The method is agnostic to the inpainting model, but Imagen Editor is used for high-resolution realistic augmentations and object-oriented masks.
5 Experiments
The experiments evaluate whether ROSIE can learn unseen manipulation skills, improve robustness to distractors and backgrounds, and support high-level success detection. Across these settings, ROSIE-generated augmentations improve generalization and robustness over the specified RT-1 baselines.
- Experimental setup: The evaluation uses ∼130k robot demonstrations covering 744 language instructions and compares original training with text and image augmentation variants.Experiments address new skills, visual robustness, and success detection.
- RQ1: Learning new skills: ROSIE fine-tuned RT-1 outperforms pre-trained RT-1 and instruction-augmented RT-1 on tasks involving unseen containers and generated manipulation objects.The experiments include moving objects near containers, placing objects into containers, grasping unknown deformable objects, and placing objects into an unseen sink.
- RQ1: Learning new skills: 75%: ROSIE outperforms pre-trained RT-1 and instruction-augmented RT-1 by at least 75% when placing objects into novel containers.The comparison is reported for the novel-container placement experiment.
- RQ1: Learning new skills: ROSIE learns to place objects into a real kitchen sink despite no real-world training data for that task.The sink scenario is generated by transforming drawer-placement episodes.
- RQ2: Robustifying manipulation policies: 115% improvement: ROSIE improves robustness across varying table textures and backgrounds, outperforming RT-1 NoAug in 7 of 8 settings.ROSIE performs similarly to NoAug in the remaining scenario.
- RQ2: Robustifying manipulation policies: ROSIE-generated distractors improve chip-bag placement into drawers when unseen objects are already present.The authors interpret this as mitigating premature termination caused by the model mistaking the first frame for a successful task state.
- RQ3: Success detection: Increasing ROSIE augmentation makes success detectors increasingly robust on difficult cluttered out-of-distribution drawer scenarios in F1 score.The prior detector without augmentation struggles in this setting, whereas ROSIE obtains reasonable performance.
6 Societal Impact
The paper highlights risks from using text-guided image generation for data augmentation and describes controls intended to reduce them.
- Risk management: ROSIE modifies local image patches with narrowly scoped semantic labels and recommends inspecting and filtering generated data before training.These practices are presented as responsible-AI measures for reducing risks from text-guided image generation.
7 Discussion, Future Work, and Conclusion
The paper concludes that ROSIE expands robotics datasets without real-world data collection and improves robustness across language-conditioned manipulation tasks. It also identifies limits involving motion generation, temporal consistency, and computational cost.
- Conclusion: ROSIE uses off-the-shelf text-guided image generation to expand robotics datasets without additional real-world data collection.The system generates new instructions and corresponding image-editing prompts.
- Conclusion: The method produces semantically meaningful image augmentations that support more robust learned models in out-of-distribution scenes.The paper validates ROSIE on varied language-conditioned manipulation tasks.
- Limitations and future work: ROSIE augments object and scene appearance but does not generate new motions, and per-frame augmentation can reduce temporal consistency.The authors note no observed performance drop for their Robotics Transformer architecture, while identifying temporal consistency as future work.
- Limitations and future work: Diffusion-based image augmentation is computationally heavy, limiting the ability to perform augmentation on the fly.The paper identifies computational efficiency as a future direction.
A.1 Implementation Details and Hyperparameters
Implementation trains RT-1 with a fixed mixture of original and ROSIE-generated episodes, using OWL-ViT for mask selection, LLM prompting for semantic edits, and Imagen Editor for inpainting.
- Policy training: RT-1 is fine-tuned for 85k steps at 1×10^-6 using a 1:1 mixture of 130k original and ROSIE-generated episodes.The pretrained RT-1 has 35M parameters and was trained for 315k steps at 1×10^-4 before fine-tuning.
- Mask generation: OWL-ViT filters low-confidence region and passthrough masks and selects the highest-scoring mask when multiple detections exist.The masks identify augmentation regions and objects that should remain untouched.
- Prompt generation: A 1-shot LLM prompt generates semantically meaningful region, passthrough-object, and inpainting prompts for the augmentation pipeline.The recipe is applied across the experimental tasks.
- Inpainting: Imagen Editor uses a 64×64 base checkpoint and a 256×256 super-resolution model to produce image augmentations.The models are run directly for inference during inpainting.
- Evaluation: Evaluation uses 10 rollouts for novel-container movement and unseen cloth grasping, 8 for novel-container placement, and 5 per can for sink placement.These rollout counts are specified per method.
- Compute: Generating 1k episodes requires 1 TPU-hour for OWL-ViT segmentation and 2 hours on four TPUs for Imagen Editor augmentation.Policy training uses 16 TPUs for 1 day.
B Examples of Augmentations
ROSIE generates semantically realistic objects and distractors for robot-learning data, including novel containers and clutter, while learned attention can focus on augmented objects.
- Novel containers: ROSIE generates novel containers, such as those used for moving objects near or placing objects into unseen containers.Figure 10 visualizes inpainted novel containers used in experiments described in Section 5.1.
- Distractors: ROSIE augments scenes with distractors including coke cans on tables and chip bags inside open drawers.These augmentations correspond to experiments in Section 5.2.
- Attention: Attention visualizations show some RT-1 heads focusing on augmented objects, suggesting the augmentations affect learned visual attention.Bright regions in the overlay indicate higher spatial attention.
- Visual realism: ROSIE-generated objects are described as semantically realistic and typically include realistic shading on tables or inside drawers.The authors identify this realism as beneficial for training manipulation policies on augmented data.
C Failure Cases of Generated Prompts and Images
ROSIE has failure cases in both prompt generation and image augmentation: zero-shot prompting can produce unusable prompts, while irregular masks can yield incomplete or incorrect objects.
- Prompt failures: Few-shot prompting is required because zero-shot LLM prompts can hallucinate unusable augmentation instructions.The reported zero-shot response incorrectly describes picking up the can near the sink rather than replacing the scene content.
- Generated-object failures: ROSIE sometimes generates the wrong container, producing bowl-shaped objects instead of woven baskets or glass mason jars.These failures are shown in Figure 12.
- Mask-related failures: Irregular masks around in-hand objects can prevent ROSIE from generating requested replacements in full shape.Examples include incomplete blue microfiber cloths or yellow rubber ducks with part of the original green chip bag remaining.
- Potential improvement: The authors suggest fine-tuning Imagen Editor on robotic manipulation data could substantially improve generation results.They also state that imperfect generations may mainly create instruction-image misalignment without greatly harming policy results.