Datasets

We hope our dataset is beneficial to your research in mobile manipulation! Below, we provide direct links to each of our dataset partitions and provide a step-by-step guide on how to reproduce our imitation learning policy results.

Dataset Download Links

Our dataset consists of four subsets:

• Expert: 110 successful demonstrations collected from an expert demonstrator (low horizon, low variance)

• Suboptimal: 110 successful demonstrations collected from a suboptimal demonstrator (longer horizon, higher variability)

• Generalization: 22 successful demonstrations collected from an expert demonstrator with key furniture object locations swapped (dishwasher and dresser swapped, sink and oven swapped)

• Sample: 5 demonstrations from the Expert subset, useful for prototyping and debugging locally without downloading the other huge subsets
  

All of our datasets are completely compatible with the robomimic framework and can be easily downloaded directly using the download_momart_datasets.py script.

Below are the links (with corresponding sizes) to download each subset for each task! Alternatively, we provide a helpful utility script for directly downloading the datasets (this assumes you have robomimic installed, see Getting Started section for more details):

• python -m robomimic.scripts.download_momart_datasets <ARGS> (Use --help flag to view download options)

Dataset Overview

Each of our dataset partitions consist of a single .hdf5 file containing demonstration data from a specific operator / environment configuration (expert / suboptimal / generalization) and task (setup table from dishwasher, unload dishwasher to dresser, ...). Note that our dataset structure is fully compatible with robomimic. We describe the structure of the .hdf5 file as follows ("Table Cleanup to Dishwasher" used as an example):

• table_cleanup_to_dishwasher.hdf5

  • mask

    • train: train data split indices

    • valid: validation data split indices

  • data

    • env_args (str): dictionary specifying environment configuration, in string form

    • total (int): Number of samples in this .hdf5 file

    • demo_0

      • num_samples (int): Number of samples in demo_0

      • actions (ndarray): (num_samples, 10) array corresponding to actions taken during demo_0 (10 = action space size)

      • dones (ndarray): (num_samples,) array corresponding to whether episode is completed (1 => completed, 0 => not completed)

      • rewards (ndarray): (num_samples,) rewards received during episode (0 everywhere)

      • states (ndarray): (num_samples, n_states) global state of the world at each timestep

      • task_successes (ndarray): (num_samples,) array corresponding to whether task is completed (1 => completed, 0 => not completed)

      • obs

        • rgb (ndarray): (num_samples, H, W, 3) RGB observations in range [0, 255] from head camera

        • rgb_wrist (ndarray): (num_samples, H, W, 3) RGB observations in range [0, 255] from wrist camera

        • depth (ndarray): (num_samples, H, W, 1) Normalized depth observations in range [0, 1] from head camera

        • depth_wrist (ndarray): (num_samples, H, W, 1) Normalized depth observations in range [0, 1] from wrist camera

        • proprio (ndarray): (num_samples, 10) Robot body observations corresponding to (head_joint_pos, grasped, eef_pos, eef_quat)

        • proprio_nav (ndarray): (num_samples, 2) Robot base observations corresponding to (linear vel magnitude, angular z vel)

        • scan (ndarray): (num_samples, 1, 100) Normalized LIDAR scan observations in range [0, 1]

        • gt_nav (ndarray): (num_samples, 4) Ground-truth navigation observations corresponding to (base xy pos, cos(z rot), sin(z rot))

        • object (ndarray): (num_samples, 3) Grouth-truth observations corresponding to (red bowl pos)

      • next_obs

        • <same as obs, but offset by 1 step>

    • demo_1
      ...

    • demo_2
      ...

...

Getting Started

Our datasets can be used immediately using the simulation platform iGibson and robot learning framework robomimic! Our hardware requirements are as follows:

• Linux machine

• Python 3.7

• NVIDIA GPU with >=11GB of VRAM
  

Proceed with the following steps:

1. Create conda environment

• conda create -y -n il_mm python=3.7

• conda activate il_mm

2. Install robomimic

• git clone https://github.com/ARISE-Initiative/robomimic.git

• cd robomimic

• pip install -e .

3. Install our branch of iGibson

• git clone https://github.com/StanfordVL/iGibson.git --single-branch --branch momart iGibson_momart --recursive

• cd iGibson_momart

• pip install -e .

• rm -r gibson2/data/* (We initially clear the modified assets, since we need to import the default ones in the following step)

• Download iGibson assets from HERE, extract, and move the assets folder into gibson2/data

• Download iGibson ig_dataset from HERE, extract, and move the ig_dataset folder into gibson2/data

• git reset --hard (make sure our modified variants of specific iGibson assets are being used)

4. Download datasets

• Download our mobile manipulation datasets(s) to the directory of your choice, assumed to be <DATASET_DIR>, using one of the following ways:

  • Direct dataset links listed above

  • python <ROBOMIMIC_DIR>/robomimic/scripts/download_momart_datasets.py <ARGS> (Use --help flag to view download options)

5. Visualize data

• python <ROBOMIMIC_DIR>/robomimic/scripts/playback_dataset.py --dataset <DATASET_DIR>/<NAME>.hdf5 --render <ARGS> (Use --help flag to view all options)

  • Plays back the dataset, optionally with actions, rendering, and /or saving to video

• python <ROBOMIMIC_DIR>/robomimic/scripts/get_dataset_info.py --dataset <DATASET_DIR>/<NAME>.hdf5 (Use --help flag to view all options)

  • Prints out metadata info from dataset

6. Training

• python <ROBOMIMIC_DIR>/robomimic/examples/train_bc_rnn.py --dataset_type momart --output <OUTPUT_FPATH> <ARGS> (Use --help flag to view all options)

  • Example training script using training configuration used in MOMART paper (useful for reproducing published results)

  • --debug flag for running a quick training session for sanity checking using a small sample dataset (~1.6GB)

  • --dataset <DATASET_DIR>/<NAME>.hdf5 for running actual training on the specified dataset

• python <ROBOMIMIC_DIR>/robomimic/scripts/train.py --config <CONFIG_FPATH> <ARGS> (Use --help flag to view all options)

  • General training script that uses robomimic configs

  • --debug flag for running a quick training session for sanity checking

  • --dataset <DATASET_DIR>/<NAME>.hdf5 for overriding the dataset specified in the config file