The robomimic.utils.dataset.SequenceDataset class extends PyTorch’s default to interface with our demonstration datasets. The class supports accessing demonstration sub-sequences (as opposed to individual states) by index and both on-demand fetching and in-memory caching. This page walks through the key concepts of the SequenceDataset interface. Please refer to the official PyTorch documentation and a short example on how to use the Dataset and DataLoader interfaces to build a training pipeline. Also, please see the released datasets section for guidance on downloading datasets that work out of the box with the repository.

Here is a sample dataset object:

dataset = SequenceDataset(
    obs_keys=(                      # observations we want to appear in batches
    dataset_keys=(                  # can optionally specify more keys here if they should appear in batches
    seq_length=10,                  # length-10 temporal sequences
    pad_seq_length=True,            # pad last obs per trajectory to ensure all sequences are sampled
    hdf5_cache_mode="all",          # cache dataset in memory to avoid repeated file i/o
    filter_by_attribute=None,       # can optionally provide a filter key here
  • hdf5_path

    • The absolute / relative path to the hdf5 file containing training demonstrations. See datasets page for the expected data structure.

  • obs_keys

    • A list of strings specifying which observation modalities to read from the dataset. This is typically read from the config file: our implementation pools observation keys from config.observation.modalities.obs.low_dim and config.observation.modalities.obs.rgb.

  • dataset_keys

    • Keys of non-observation data to read from a demonstration. Typically include actions, rewards, dones.

  • seq_length

    • Length of the demonstration sub-sequence to fetch.

  • pad_seq_length

    • Whether to allow fetching subsequence that ends beyond the sequence. For example, given a demo of length 10 and seq_length=10, setting pad_seq_length=True allows the dataset object to access subsequence at __get_item(index=5)__ by repeating the last frame 5 times.

  • hdf5_cache_mode

    • Optionally cache the dataset in memory for faster access. The dataset supports three caching modes: ["all", "low_dim", or None].

      • all: Load the entire dataset into the RAM. This mode minimizes data loading time but incurs the largest memory footprint. Recommended if the dataset is small or when working with low-dimensional observation data.

      • low_dim: Load only the low-dimensional observations into RAM. Always use this mode when possible as loading low-dim data incurs nontrivial overhead. Low-dim observations are specified at config.observation.modalities.obs.low_dim.

      • None: Always fetch data on-demand.

  • hdf5_normalize_obs

    • If True, normalize observations by computing the mean observation and std of each observation (in each dimension and modality), and normalizing unit mean and variance in each dimension.

  • filter_by_attribute

    • if provided, use the provided filter key to look up a subset of demonstrations to load. See the documentation on filter keys for more information.