This repo demonstrates how to convert an existing DROID dataset into RLDS format for X-embodiment experiment integration.
First create a conda environment using the provided environment.yml file (use environment_ubuntu.yml
or environment_macos.yml
depending on the operating system you're using):
conda env create -f environment_ubuntu.yml
Then activate the environment using:
conda activate rlds_env
If you want to manually create an environment, the key packages to install are tensorflow
,
tensorflow_datasets
, matplotlib
, plotly
, cv2
and wandb
.
During teleoperation, DROID stores camera data in ZED's .svo
data format. For conversion to RLDS we first need
to convert these into .mp4
files. The easiest is to do this directly on the computer that is running the DROID
data collection program and controlling the robot, since it already has the necessary ZED python API
installed.
On the robot machine, run the svo_to_mp4.py
script. Afterward, you should see .mp4
video files in your data directory -- check that the videos look as expected.
Each video shows both stereo camera views, horizontally concatenated. Next, copy the full data folder to the machine
on which you want to run RLDS conversion below. If you're running RLDS conversion directly on the robot machine, you
do not need to copy the data.
You can modify the provided example to convert your own data. Follow the steps below:
1.Modify Dataset Splits: The function _split_paths()
determines the splits of the generated dataset (e.g. training, validation etc.).
If your dataset defines a train vs validation split, please provide the corresponding filenames, e.g.
by pointing the crawler
to the corresponding folders (i.e. the data folder you copied from the robot machine above).
If your dataset does not define splits, put all file names into the train
split.
2.Modify Dataset Conversion Code: We currently assume a fixed dummy language instruction for the full dataset. Please modify to use one or multiple appropriate instructions.
That's it! You're all set to run dataset conversion. Run the following:
pip3 install -e .
cd droid
tfds build --overwrite
The command line output should finish with a summary of the generated dataset (including size and number of samples).
Please verify that this output looks as expected and that you can find the generated tfrecord
files in ~/tensorflow_datasets/droid
.
Note 1: This will overwrite any existing dataset, so make sure to copy previously converted data to a new place or
add a --data_dir=<your_new_path>
to the build
command to write the output to a different folder.
Note 2: Conversion uses multi-threading. You can adjust the number of workers and number of episodes held in memory in parallel
via the N_WORKERS
and MAX_PATHS_IN_MEMORY
values in droid.py
.
If you would like to submit your data to the Open X-Embodiment dataset to help the community, you can add a citation and
license as described below. Please submit your data via the enrollment form and send a short email indicating your submission
to open-x-embodiment@googlegroups.com
.
3.Provide Dataset Description: Next, add a bibtex citation for your dataset in CITATIONS.bib
and add a short description
of your dataset in README.md
inside the dataset folder. You can also provide a link to the dataset website and please add a
few example trajectory images from the dataset for visualization.
4.Add Appropriate License: Please add an appropriate license to the repository.
To verify that the data is converted correctly, please run the data visualization script from the base directory:
python3 visualize_dataset.py droid
This will display a few random episodes from the dataset with language commands and visualize action and state histograms per dimension.
Note, if you are running on a headless server you can modify WANDB_ENTITY
at the top of visualize_dataset.py
and
add your own WandB entity -- then the script will log all visualizations to WandB.