PaperlayerView exact revision on arXiv

Source-linked AI summary

RoboEngine: Plug-and-Play Robot Data Augmentation with Semantic Robot Segmentation and Background Generation

Chengbo Yuan, Suraj Joshi, Shaoting Zhu, Hang Su, Hang Zhao, Yang Gao

arXiv:2503.18738v2cs.RO

TL;DR

Existing robot visual augmentation can require camera calibration or controlled setups, while collecting data across diverse scenes is costly. RoboEngine combines robot segmentation and fine-tuned background generation in a plug-and-play toolkit. With demonstrations from one scene, it generalizes manipulation to six new scenes and reports performance gains over two baselines.

  • Problem

    Robot visual augmentation methods can require prerequisites such as camera calibration or green-screen environments, while collecting data across diverse scenes is costly.

    Source: §I · p. 1

  • Method

    RoboEngine combines a robot-scene segmentation dataset and model with a fine-tuned, physics- and task-aware background-generation diffusion model in a plug-and-play toolkit.

    Source: §I · p. 2, C. Plug-and-Play RoboEngine Augmentation · p. 3

  • Results

    With demonstrations from one scene, RoboEngine generalizes manipulation to six entirely new scenes, improving performance by 210% over the no-augmentation baseline and 20% over the previous SoTA method.

    Source: §I · p. 2

  • Takeaways & Limitations

    The authors hope RoboEngine will help alleviate visual robustness issues in imitation learning and serve as a standard augmentation tool.

    Source: §V · p. 6

    RoboEngine does not handle temporal consistency between frames or consider multi-view or 3D augmentation.

    Source: §V · p. 6

Abstract

from arXiv · show

Visual augmentation has become a crucial technique for enhancing the visual robustness of imitation learning. However, existing methods are often limited by prerequisites such as camera calibration or the need for controlled environments (e.g., green screen setups). In this work, we introduce RoboEngine, the first plug-and-play visual robot data augmentation toolkit. For the first time, users can effortlessly generate physics- and task-aware robot scenes with just a few lines of code. To achieve this, we present a novel robot scene segmentation dataset, a generalizable high-quality robot segmentation model, and a fine-tuned background generation model, which together form the core components of the out-of-the-box toolkit. Using RoboEngine, we demonstrate the ability to generalize robot manipulation tasks across six entirely new scenes, based solely on demonstrations collected from a single scene, achieving a more than 200% performance improvement compared to the no-augmentation baseline. All datasets, model weights, and the toolkit are released https://roboengine.github.io/

I. INTRODUCTION

RoboEngine targets the prerequisites and limited realism and diversity of existing robot augmentation with a plug-and-play toolkit for physics- and task-aware scenes. In real-robot experiments, it generalizes manipulation from demonstrations in one scene to six new scenes, outperforming both no augmentation and a previous method.

I. INTRODUCTION

  • Collecting demonstrations across diverse scenes can mitigate visual disturbances in imitation learning, but the time and resource costs make that approach unscalable.Generative augmentation offers a more scalable way to produce diverse visual backgrounds for training.Source: §I · p. 1
  • Existing methods lack simultaneous convenience, realism, and diversity, with prerequisites such as green screens or camera calibration and limits on scene-wide changes.Random-image or texture changes may violate physical constraints and lead to degraded real-world performance due to distribution shifts.Source: §I · pp. 1–2
  • RoboEngine packages plug-and-play visual augmentation for physics- and task-aware robot scenes without prerequisite conditions.The toolkit is intended to let users generate backgrounds with just a few lines of code.Source: §I · p. 2
  • 210% and 20% performance improvements over the no-augmentation baseline and previous SoTA method, respectively, enabled generalization from one demonstration scene to six entirely new scenes.Performance continued to improve to some extent as the number of augmented demonstrations increased.Source: §I · p. 2
  • RoboSeg contributes 3,800 high-quality annotated robot-scene images and supports training a generalizable, high-quality robot segmentation model.The paper describes its model as the first generalizable “wire-level” robot segmentation model.Source: §I · p. 2
  • RoboEngine integrates the segmentation model with a fine-tuned, physics- and task-aware diffusion model for background generation.The toolkit is presented as a convenient augmentation interface similar to ColorJitter in computer vision.Source: §I · p. 2

A. Visual Generalizable Imitation Learning

Imitation learning is a key approach to visuomotor robot control, but raw Diffusion Policy is sensitive to visual disturbances. The paper focuses on data augmentation as one strategy for improving robustness and generalization.

A. Visual Generalizable Imitation Learning

B. Robot Data Augmentation

Robot data collection is costly, while augmentation methods vary in how well they handle large visual changes and in their prerequisites. Generative methods can provide larger changes but still face practical and physical-reliability constraints.

B. Robot Data Augmentation

  • Collecting robot data is costly and labor-intensive, especially across diverse environments, motivating methods that synthesize diverse data at minimal cost.Source: B. Robot Data Augmentation · p. 2
  • Cropping and ColorJitter work for in-domain visual generalization but struggle with large visual changes.Source: B. Robot Data Augmentation · p. 2
  • Generative methods handle larger visual changes but are not plug-and-play when they require camera calibration, scene objects for inpainting, or scene-specific hyperparameter adjustment.The cited methods also offer limited visual variation that may fail under extreme conditions.Source: B. Robot Data Augmentation · p. 2
  • GreenAug can make more significant visual changes but requires a green screen, struggles in unstructured environments, and neglects physical reliability.The passage links neglect of physical reliability to out-of-distribution issues.Source: B. Robot Data Augmentation · p. 2

III. METHOD

RoboEngine generates robot masks without prerequisites and combines segmentation with fine-tuned diffusion-based background generation in a plug-and-play pipeline.

III. METHOD

  • RoboEngine's method introduces a robot-scene segmentation dataset and develops a generalizable, high-quality robot segmentation model to generate masks without prerequisites.Source: §III · p. 2
  • The method fine-tunes a diffusion model for background generation.Source: §III · p. 2
  • The segmentation and diffusion models are integrated into a seamless plug-and-play pipeline.Source: §III · p. 2

A. Task Definition

RoboEngine augments a robot demonstration dataset by separating robot, task-object, and background regions, replacing the background, and recombining the image while retaining its associated information.

A. Task Definition

  • RoboEngine augments each demonstration image by segmenting the robot, task-related objects, and background, then generating a new background and recombining the regions.The task-related object mask includes objects such as the mouse and pad in “Put Mouse to the Pad.”Source: A. Task Definition · p. 3
  • The augmented dataset replaces each image with its augmented version while retaining the associated information, including instructions, proprioception, and actions.The input dataset pairs images with this associated information.Source: A. Task Definition · pp. 2–3

B. RoboSeg Dataset and Robot Segmentation Model

RoboSeg supplies fine-grained robot-scene annotations for training Robo-SAM, addressing the difficulty of obtaining accurate robot masks with off-the-shelf models.

B. RoboSeg Dataset and Robot Segmentation Model

  • Off-the-shelf models such as SAM-V2 and LISA struggle to segment robots accurately, often necessitating green screens or camera calibration.Source: B. RoboSeg Dataset and Robot… · p. 3
  • RoboSeg contains 3,800 images sampled from over 35 robot datasets, with robot-main, robot-auxiliary, and task-object masks plus task instructions.The annotations include robot wires; each image also has 10 scene descriptions generated with GPT-4o.Source: B. RoboSeg Dataset and Robot… · p. 3
  • Fine-tuning language-conditioned EVF-SAM on RoboSeg produces Robo-SAM for high-quality open-world robot segmentation.The authors choose language conditioning because they find language provides additional visual cues that enhance the base model’s performance.Source: B. RoboSeg Dataset and Robot… · p. 3

C. Plug-and-Play RoboEngine Augmentation

RoboEngine combines robot and task-object segmentation with generated backgrounds in a toolkit that requires no prerequisite setup.

C. Plug-and-Play RoboEngine Augmentation

  • Robo-SAM segments the robot and EVF-SAM segments task-related objects conditioned on their names in the task instruction.The generation step uses the resulting masks to create a physics- and task-aware background.Source: C. Plug-and-Play RoboEngine Augmentation · p. 3
  • The RoboSeg dataset provides high-quality, fine-grained semantic segmentation annotations across diverse robots and environments.Source: Figure 2 · p. 3
  • The background-generation model randomly selects scene descriptions from a pool collected from RoboSeg.Source: C. Plug-and-Play RoboEngine Augmentation · p. 3
  • RoboEngine integrates its segmentation and background-generation models in a toolkit that needs no prerequisite setup and supports augmentation with a few lines of code.Source: C. Plug-and-Play RoboEngine Augmentation · p. 3

A. Segmentation Results

Robo-SAM outperforms the compared segmentation baselines on both test sets, while RoboEngine is the only augmentation method reported to combine visual diversity with physics feasibility. Real-robot evaluation trains from one scene and tests across six new scenes.

A. Segmentation Results

  • More than 0.12 GIoU separates Robo-SAM from every baseline on both the Test and Zero-shot sets.The table compares segmentation performance using GIoU; the accompanying figure reports that only Robo-SAM yields usable masks for downstream augmentation.Source: Table I · p. 4, A. Segmentation Results · p. 4, Figure 3 · p. 4

B. Augmentation Baseline and Results

  • RoboEngine is the only compared augmentation method reported to provide both high visual diversity and physics feasibility.The compared methods include generated backgrounds, inpainting, ImageNet backgrounds, random textures, and no augmentation.Source: B. Augmentation Baseline and Results · p. 4

C. Real Robot Setting

  • Demonstrations collected in one scene are tested across four new scenes for Fold Towel and two for Put Mouse, with similar table heights to focus on visual generalization.The real-robot comparison reports behavior score and success rate across six new scenes.Source: C. Real Robot Setting · p. 4, Table II · p. 5

D. Policy Evaluation Results

RoboEngine outperforms the best baselines in policy evaluation, while scaling tests find gains from more augmented demonstrations that eventually slow. The conclusion also identifies temporal consistency and multi-view or 3D augmentation as limitations.

D. Policy Evaluation Results

  • 20% over the best baselines and 210% over the no-augmentation baseline are the reported performance improvements for RoboEngine.The paper reports average scores and success rates across tasks in Table II, with per-scene results in Appendix V-C.Source: D. Policy Evaluation Results · p. 5
  • Inpainting performance depends on task-irrelevant object masks and hyperparameters, often produces small visual changes, and has the highest time cost.The authors do not recommend this method for practical use.Source: D. Policy Evaluation Results · p. 5
  • Background, ImageNet, and Texture all improve substantially over No aug; Texture performs slightly better, while Texture and ImageNet are faster than RoboEngine.The paper recommends RoboEngine when compute is available and Texture or ImageNet when resources are limited.Source: D. Policy Evaluation Results · p. 5

D. Details of Scaling Trend Experiment

E. Performance Scaling Trend

  • Mixing original and augmented data improves performance over 1×, but not over 2× augmentation alone; more augmented demonstrations help with diminishing gains toward a bottleneck.The experiment evaluates Fold Towel (Finish) across three novel scenes; Figure 5 reports un-normalized average behavior score.Source: E. Performance Scaling Trend · p. 6, Figure 5 · p. 6

V. CONCLUSION

  • RoboEngine combines Robo-SAM, a fine-tuned physics- and task-aware background-generation model, and an out-of-the-box toolkit; it does not handle temporal consistency or multi-view and 3D augmentation.The conclusion mentions video diffusion and depth estimation followed by re-rendering as possible ways to address those limitations.Source: §V · p. 6

A. Details of Model Training

  • Robo-SAM is fine-tuned for 30 epochs on 3,562 training and 92 validation images, while BackGround-Diffusion is fine-tuned for 100 epochs using all RoboSeg images.The reported learning rates are 1e-5 for Robo-SAM and 5e-3 for BackGround-Diffusion; both use batch size 32.Source: A. Details of Model Training · p. 6

B. Evaluation Criteria; C. Raw Evaluation Results

Loading 2503.18738v2…