Enhancing Visual Domain Robustness in Behaviour Cloning via Saliency-Guided Augmentation

AI summary

Overview

Research area: Robotics — vision-based behaviour cloning (BC), visual domain generalisation, and image data augmentation.

Technical level: Intermediate. Readers benefit from familiarity with imitation learning policies, image encoders, and data augmentation terminology, but the core idea (blend training images using a saliency map) is explained without requiring deep background.

Scope: The paper introduces RoboSaGA, a saliency-guided superimposition augmentation that blends in-domain robot demonstration images with out-of-domain images at per-pixel strength, and evaluates it against Random Crop, Colour Jitter, Guided Erase, Random Overlay, and SODA across simulated and real-world manipulation tasks and three policy architectures.

What This Paper Is About

Vision-based behaviour cloning trains robot policies from images of expert demonstrations, but those datasets are usually collected in limited visual conditions, so policies break when lighting, shadows, backgrounds, or distractor objects change. Standard image augmentations such as Random Crop and Colour Jitter are mild and do not cover these larger visual domain shifts, while superimposition augmentations borrowed from computer vision risk destroying task-critical content. The paper's goal is an augmentation that can blend in strong out-of-domain texture while protecting the parts of the image the policy actually depends on.

Key Contributions

1. RoboSaGA, a superimposition augmentation for vision-based BC that uses per-pixel saliency maps, back-propagated from the visual encoder's output features via FullGrad, as the blending matrix between the in-domain image and a randomly drawn out-of-domain image. Unlike KeepAugment, it uses adaptive per-pixel blending rather than fixed-size rectangular binary crops, and derives saliency from visual features rather than classification logits.

2. A saliency clipping mechanism (s^{v_i} := min(g^{v_i}, λ), with λ set to 0.8 across all experiments) that caps normalised saliency so task-critical regions still receive a subtle overlay, which the authors report is important for robustness to background changes.

3. A global saliency buffer that stores 8-bit single-channel saliency maps at reduced resolution and updates only the least recently updated β trajectories per batch, letting RoboSaGA avoid recomputing FullGrad for every augmented sample. The paper reports roughly 67 MB of graphics memory to store maps for 10,000 samples at 84×84, and a warm-up period (default 10 epochs) during which all-one saliency maps are returned.

4. A simulation and real-world evaluation protocol comparing RoboSaGA against Random Crop, Random Crop + Colour Jitter, SODA, Guided Erase, and Random Overlay across three policy families (BC-MLP, BC-RNN, and Diffusion Policy) and five tasks (Lift, Can, Square, Transport in simulation, Toy in the real world), using a performance-gap metric defined as the in-domain Random Crop baseline success rate minus the success rate under visual domain shift with the augmentation applied.

Main Findings

• Traditional augmentations fail on distractors and backgrounds. Averaged over all tested policies and tasks, Random Crop produced a mean gap of 0.60 for distractor and background variations, and Colour Jitter reduced this only from 0.60 to 0.55, an 8% relative improvement. Colour Jitter did help with plain lighting changes, reducing the average gap to 0.10, but its improvement was limited under shadows, where the gap was 0.33 versus less than 0.1 for overlay methods.

• Guided Erase and Random Overlay trade off in opposite directions. In the BC-MLP ablation over Lift, Can, and Square, Guided Erase was strongest against distractors (gap 0.06) but weak against background changes (0.55), while Random Overlay was more consistent across shifts but less effective against distractors (0.31). RoboSaGA reduced Random Overlay's distractor gap from 0.31 to 0.06 and Guided Erase's background gap from 0.55 to 0.13, with a mean gap of 0.10 versus 0.68 for Random Crop, 0.31 for Guided Erase, and 0.26 for Random Overlay.

• Random Overlay alone is a strong, low-cost baseline. Averaged over all policies, Random Overlay reduced the gap for distractor and background variations from Colour Jitter's 0.55 to 0.24 and for lighting and shadow variations from 0.21 to 0.06, and in real-world tests from 0.71 to 0.18.

• RoboSaGA improves on Random Overlay. It reduced the simulated mean gap to 0.14 and the real-world mean gap to 0.05, corresponding to relative improvements of 41.6% and 50.0% over Random Overlay on the two simulated variation groups and 72% in the real-world setting.

• Representation learning was not necessary here. Using SODA, a student-teacher architecture with a consistency loss, as a baseline, the authors found that direct training with augmented samples delivered comparable yet marginally improved performance, reducing the gap by an average of 0.04 under distractor and background variations and 0.02 under lighting and shadows.

• Gains transfer across policy types. Random Overlay, RoboSaGA, and SODA consistently improved over their Random Crop counterparts across BC-MLP, BC-RNN, and Diffusion Policy. The Diffusion Policy exhibited the least performance degradation.

• The saliency buffer trades almost nothing for speed. With the buffer, the average success rate over Lift, Can, and Square under visual domain shifts was 0.75 versus 0.74 without it, while per-batch augmentation processing time dropped from 0.70s to 0.22s.

• Saliency maps do not always match human intuition. During reaching, saliency highlighted the general region around the gripper and nut while the eye-in-hand view focused only on the handle; during lifting, saliency shifted away from the object; during insertion, the front camera overlooked the target peg and nut. The authors hypothesise this arises because vision-based BC integrates proprioception, observation history, and multiple views, and note that history-independent BC-MLP produced more human-interpretable saliency maps than history-dependent counterparts.

• Two reported weaknesses. BC-MLP and BC-RNN showed comparatively lower gains across all tested methods on the Transport task, which the authors attribute partly to its second-person view being dominated by two manipulators with little target information. Saliency computation with a ResNet18 encoder on 84×84 inputs also took about 1.5 times longer than training itself.

Methodology in Plain English

The policy network already contains a visual encoder that turns camera images into features. RoboSaGA asks that encoder which pixels it cares about, by running a FullGrad saliency computation backward from the feature output to the input image. Those saliency values are normalised, capped at 0.8, and used as a per-pixel mixing weight: where saliency is high, the original robot image is kept; where it is low, a random out-of-domain image shows through. This makes augmentation aggressive in task-trivial areas such as tabletop and background, and gentle in task-critical areas.

To avoid paying the FullGrad cost on every sample, saliency maps are cached in a buffer. Each batch updates the buffer for only 10% of the trajectories, storing maps as 8-bit single-channel images at reduced resolution, while 50% of trajectories in the batch are augmented using cached maps. The first 10 epochs use all-one saliency maps so the encoder can develop task-specific knowledge before strong augmentation begins. In-domain and out-of-domain images are drawn from the same training batch pipeline; the out-of-domain pool comprises 5,000 MSCOCO images plus 1,000 synthetic images with plain, gradient, grid, chess, and Perlin patterns, all subject to random rotation and brightness adjustment.

Experiments use four RoboMimic simulation tasks (Lift, Can, Square, Transport) with proficient human demonstrations and one real-world Toy pick-and-place task. Random Crop is the default baseline in every experiment, using a crop size of 90% of the input resized back to original resolution; Colour Jitter baselines are obtained by fine-tuning Random Crop baselines for 50 additional epochs with PyTorch's ColorJitter at 0.2 for brightness, contrast, saturation, and hue. In simulation, results are averaged over the top three checkpoints from 600 training epochs saved every 50 epochs with 50 rollouts per checkpoint (Lift trains for 200 epochs); in the real world, the final checkpoint after 500 epochs is used with 20 trials per data point.

Why This Matters

Research impact. The paper shows that a cheap, task-prior-free augmentation signature — overlay of random images — is surprisingly strong in behaviour cloning, and that a light amount of task-aware shaping on top of it buys further robustness. It also provides a direct comparison against representation-learning-based augmentation (SODA) and finds action-labelled BC does not obviously need it, which is a useful counterpoint to practice in image-based reinforcement learning. The saliency-visualisation analysis raises a question about whether policy-driven saliency is a reliable proxy for task relevance in multi-sensory policies.

Real-world applications:

• Warehouse or factory pick-and-place cells where lighting, shadows from passing objects, and changing tabletop surfaces shift between data collection and deployment.
• Household or service robots operating on cluttered surfaces, where distractor objects and background textures vary unpredictably across homes.
• Agricultural or outdoor manipulation, where natural lighting and background variation are large and out-of-domain imagery is easy to source but task-relevant imagery is expensive.
• Lab automation and other settings where demonstrations are scarce and expensive, making augmentation preferable to collecting new data under every visual condition.

Industry relevance. The method is designed as a drop-in change to existing training pipelines: no structural changes to the policy, no additional learning objectives, and no extra learnable modules. The out-of-domain images come from an existing dataset plus synthetic patterns, and the saliency buffer keeps the overhead bounded, which matters for teams training multi-camera diffusion policies on limited hardware. The real-world demonstration uses a 3 degrees-of-freedom pick-and-place task with BC-MLP, BC-RNN, and Diffusion Policy, indicating the approach is not tied to one architecture.

Future Directions

• Reduce saliency computation cost. The authors suggest approximating saliency at mid-feature layers, exploring alternative saliency extractors, or using knowledge distillation, where a smaller encoder is trained without augmentation and saliency is computed in later epochs to build an offline buffer.
• Improve performance on long-horizon, multi-subtask tasks. BC-MLP and BC-RNN showed comparatively lower gains on Transport across all tested methods; understanding why policies differ in distilling task-relevant information from images with high task-related pixel density is left open.
• Broaden real-world task coverage. The real-world evaluation is a concentrated 3-degrees-of-freedom pick-and-place task; the authors state that future work could extend the findings to a broader range of tasks.
• Reconcile saliency with human notions of task relevance. The mismatch between saliency maps and human expectations raises the question of which views, history lengths, and sensory inputs drive saliency, and whether the process could be steered toward more interpretable task focus. The paper notes that history-independent BC-MLP produces more human-interpretable maps and points to examples in its appendix.

Target Audience

Robotics and imitation-learning researchers working on visual generalisation for manipulation policies; practitioners training behaviour cloning policies who need robustness to lighting, shadow, background, and distractor variation without collecting new demonstrations; and computer vision researchers interested in how saliency-guided superimposition augmentations behave when task-critical semantics, multi-view inputs, and proprioception must be preserved rather than classification logits.

Authors’ abstract

In vision-based behavior cloning (BC), conventional image augmentations such as Random Crop and Color Jitter often fall short under substantial visual domain shifts, including changes in shadows, distractors, and backgrounds. Superimposition-based augmentations, which blend in-domain and out-of-domain images, have shown promise for improving generalization in computer vision, but their suitability for BC remains uncertain because task-critical semantics, spatiotemporal relationships, and agent-target interactions must be preserved. To address this, we introduce RoboSaGA, a Saliency-Guided Augmentation method within the superimposition family tailored for vision-based BC. RoboSaGA dynamically adjusts augmentation intensity at the pixel level using policy-driven saliency, enabling aggressive augmentation in task-irrelevant regions while preserving task-critical information. It integrates seamlessly into existing architectures without requiring structural modifications or additional learning objectives. Experiments in both simulated and real-world settings show that RoboSaGA preserves in-domain performance while substantially improving robustness to visual domain shifts, including distractor and background changes, as well as lighting and shadow variations. Code is available at https://github.com/Zheyu-Zhuang/RoboSaGA.