robomimic.utils.dataset.SequenceDataset class extends PyTorch’s default
torch.utils.data.Dataset 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
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( hdf5_path=dataset_path, obs_keys=( # observations we want to appear in batches "robot0_eef_pos", "robot0_eef_quat", "image", "object", ), dataset_keys=( # can optionally specify more keys here if they should appear in batches "actions", "rewards", "dones", ), 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 hdf5_normalize_obs=False, filter_by_attribute=None, # can optionally provide a filter key here )
The absolute / relative path to the hdf5 file containing training demonstrations. See datasets page for the expected data structure.
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
Keys of non-observation data to read from a demonstration. Typically include
Length of the demonstration sub-sequence to fetch.
Whether to allow fetching subsequence that ends beyond the sequence. For example, given a demo of length 10 and
pad_seq_length=Trueallows the dataset object to access subsequence at
__get_item(index=5)__by repeating the last frame 5 times.
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
None: Always fetch data on-demand.
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.
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.