Dataset Card for CAST (Counterfactual Augmentation with Synthetic Trajectories)

For more information, see:

• Our website
• Our paper
• Our code

This repo includes three files:

• cast_counterfactual_dataset.tar.gz : language-annotated counterfactual trajectories
• cast_filtered_dataset.tar.gz : language-annotated trajectories from original GNM dataset
• atomic_datasets.tar.gz : grouped and annotated atomic skill trajecories ('turn left', 'turn right', 'go forward', etc.)

Uses

This dataset is intended to be used to train steerable VLAs. Code for training with this dataset is available here(cast-vla).

Getting started

To get started, download the datasets or the dataset you'd like to use

hf download catglossop/CAST-dataset --repo-type=dataset

or use download button next to the file name.

Then, the files can be unzipped

tar -xvf <target tar.gz file>

For use with our training code and training on TPUs, the data should then be uploaded to a GCS bucket

gsutil cp -r <your path> gs://<your bucket>

The data is stored in the RLDS format as tfrecords. The data format stored using this tfds features dict format.

tfds.features.FeaturesDict({
                'steps': tfds.features.Dataset({
                    'observation': tfds.features.FeaturesDict({
                        'image': tfds.features.Image(
                            shape=(128, 128, 3),
                            dtype=np.uint8,
                            encoding_format='png',
                            doc='Main camera RGB observation.',
                        ),
                        'state': tfds.features.Tensor(
                            shape=(3,),
                            dtype=np.float64,
                            doc='Robot state, consists of [2x position, 1x yaw]',
                        ),
                        'position': tfds.features.Tensor(
                            shape=(2,),
                            dtype=np.float64,
                            doc='Robot position',
                        ),
                        'yaw': tfds.features.Tensor(
                            shape=(1,),
                            dtype=np.float64,
                            doc='Robot yaw',
                        ),
                        'yaw_rotmat': tfds.features.Tensor(
                            shape=(3, 3),
                            dtype=np.float64,
                            doc='Robot yaw rotation matrix',
                        ),
                    }),
                    'action': tfds.features.Tensor(
                        shape=(2,),
                        dtype=np.float64,
                        doc='Robot action, consists of 2x position'
                    ),
                     'action_angle': tfds.features.Tensor(
                        shape=(3,),
                        dtype=np.float64,
                        doc='Robot action, consists of 2x position, 1x yaw',
                    ),
                    'discount': tfds.features.Scalar(
                        dtype=np.float64,
                        doc='Discount if provided, default to 1.'
                    ),
                    'reward': tfds.features.Scalar(
                        dtype=np.float64,
                        doc='Reward if provided, 1 on final step for demos.'
                    ),
                    'is_first': tfds.features.Scalar(
                        dtype=np.bool_,
                        doc='True on first step of the episode.'
                    ),
                    'is_last': tfds.features.Scalar(
                        dtype=np.bool_,
                        doc='True on last step of the episode.'
                    ),
                    'is_terminal': tfds.features.Scalar(
                        dtype=np.bool_,
                        doc='True on last step of the episode if it is a terminal step, True for demos.'
                    ),
                   'language_instruction': tfds.features.Tensor(
                        shape=(10,),
                        dtype=tf.string,
                        doc='Language Instruction.'
                    ),
                }),
                'episode_metadata': tfds.features.FeaturesDict({
                    'file_path': tfds.features.Text(
                        doc='Path to the original data file.'
                    ),
                    'episode_id': tfds.features.Scalar(
                        dtype=tf.int32,
                        doc='Episode ID.'
                    ),
                    'normalization_fctor': tfds.features.Scalar(
                        dtype=tf.float32,
                        doc='Normalization factor.'
                    ),
                }),
            }))

Some specific things to note:

• The language instruction field includes up to ten alternative instructions for a trajectory. Empty strings are used as padding (see our code)).
• As we have counterfactual data, there are actions without corresponding image observations. Placeholder white images are used in place making it easy to filter them out later (see our code)

Converting the CAST dataset to your own version

To start converting the CAST dataset to your own dataset version, it is easiest to load the dataset and go through each example, reformatting or saving data as needed. Using the script provided in this repo, fill in your own data conversion code for each sample and run the script as follows

python convert_cast_dataset.py --dataset_name <dataset_name> --data_dir <local or GCS path to data> --output_dir <desired output dir> <--use_gcs if using GCS>

Dataset Creation

For information of the autolabeling, please read our paper and see our code.

Source Data

This dataset is build off the the GNM data mixture:

• SACSoN
• Cory Hall
• Go Stanford
• SCAND
• RECON
• Seattle
• Tartan Drive