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:
- Create conda environment
conda create -y -n il_mm python=3.7
conda activate il_mm
- Install robomimic
cd robomimic
pip install -e .
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)
- 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)
- 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
- 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