Enhancing Visual Domain Robustness in Behaviour Cloning via Saliency-Guided Augmentation, CoRL 2024 [Paper]

This is the official implementation of RoboSaGA, a data augmentation technique designed to improve generalization against changes in lighting, shadows, the presence of distractors, and variations in table textures and backgrounds in the context of behavior cloning.

📈 The following updates have been made to RoboSaGA and the baselines compared to what was described in the paper:

• Average Pooling: Replaces Spatial Softmax for better saliency. Spatial Softmax introduces a diamond-shaped gradient prior, as its gradient computation depends on the constant multipliers introduced by coordinate grids.

✅ Test Training and Evaluation scripts for BC-MLP, and BC-RNN

🔳 Test training and evaluation scripts for Diffusion Policy.

To install Mujoco200, run the following commands:

chmod +x install_mujoco200.sh
./install_mujoco200.sh

Create the Conda environment using the provided environment.yaml file:

conda env create -f environment.yaml

Activate

conda activate robosaga

Datasets are stored within the data folder.

1. Unzip the backgrounds.zip file, which contains all out-of-domain images for data augmentation.

   mkdir data
   unzip backgrounds.zip -d data/

2. Download the robomimic datasets from this link.
3. Unzip the robomimic dataset to the corresponding path. For example, for the transport task, the proficient human (ph) dataset should be stored under ./data/robomimic/transport/ph.

The default config files for RoboSaGA are included in the robomimic/robomimic/robosaga_configs directory with the format <policy_type>_<aug_method>.json.

Training appears to be faster for BC-MLP and BC-RNN with the original robomimic environment compared to the counterpart integrated in Diffusion Policy (although Diffusion Policy is also capable of training BC-MLP and BC-RNN).

Navigate to the robomimic directory if you are not already there:

cd robomimic

Train with the corresponding config file. For example, to train the square task with BC-RNN using RoboSaGA, run:

python robomimic/scripts/train.py --config robomimic/robosaga_configs/transport/bc_rnn_saga.json

To be added

RoboSaGA supports evaluation under the following visual domain shifts:

• lighting
• lighting_and_shadow
• distractors
• backgrounds
• Combinations of the above

1. Navigate to the robomimic directory if you are not already there:

   cd robomimic

2. Run the evaluation script with your chosen settings. For example, to evaluate with distractors and background shifts, run:

   python robomimic/scripts/eval_visual_domain_shifts.py -e <PATH_TO_THE_EXPERIMENT> --vds distractors backgrounds --video

   Replace <PATH_TO_THE_EXPERIMENT> with the path to your trained model or experiment directory.

   The --vds flag specifies the visual domain shifts to evaluate, and the --video flag generates evaluation videos for visualization.

To be added

This project would not have been possible without the following amazing works and their code (listed alphabetically):

• Diffusion Policy [paper] [code]
• Fullgrad [paper] [code]
• Robomimic [paper] [code]

@inproceedings{zhuang2024enhancing,
  title = {Enhancing Visual Domain Robustness in Behaviour Cloning via Saliency-Guided Augmentation},
  author = {Zheyu Zhuang and Ruiyu Wang and Nils Ingelhag and Ville Kyrki and Danica Kragic},
  booktitle = {8th Annual Conference on Robot Learning},
  year = {2024},
  url = {https://openreview.net/forum?id=CskuWHDBAr}
}