opentau.datasets.factory

Factory functions for creating datasets and dataset mixtures.

This module provides factory functions to create individual datasets and weighted dataset mixtures from configuration objects. It handles the setup of delta timestamps, image transforms, and metadata configuration before instantiating datasets.

The factory supports two types of datasets:

  1. LeRobot datasets: Standard robot learning datasets loaded from HuggingFace repositories with configurable delta timestamps for temporal alignment.

  2. VQA datasets: Vision-language vqa datasets (CLEVR, COCO-QA, VSR, etc.) for multimodal learning tasks.

Key Features:

  • Delta timestamp resolution: Automatically configures temporal offsets for features.

  • Image transform support: Applies configurable image transformations during dataset creation.

  • Imagenet stats override: Optionally replaces dataset statistics with ImageNet normalization statistics for camera features.

  • VQA dataset registration: Supports extensible vqa dataset registration through side-effect imports.

Functions:
make_dataset: Creates a single dataset instance from a DatasetConfig,

handling delta timestamp setup, image transforms, and metadata configuration.

make_dataset_mixture: Creates a WeightedDatasetMixture from a

TrainPipelineConfig containing multiple dataset configurations.

resolve_delta_timestamps: Resolves delta timestamps configuration based

on TrainPipelineConfig settings, mapping features to temporal groups.

Constants:
IMAGENET_STATS: ImageNet normalization statistics (mean, std, min, max)

used for camera feature normalization when use_imagenet_stats is enabled.

Example

Create a single dataset:
>>> dataset = make_dataset(dataset_cfg, train_cfg, return_advantage_input=False)
Create a dataset mixture:
>>> mixture = make_dataset_mixture(train_cfg, return_advantage_input=False)
>>> dataloader = mixture.get_dataloader()

Functions

make_dataset(cfg, train_cfg[, ...])

Handles the logic of setting up delta timestamps and image transforms before creating a dataset.

make_dataset_mixture(cfg[, ...])

Creates a dataset mixture from the provided TrainPipelineConfig.

resolve_delta_timestamps(cfg, dataset_cfg, ...)

Resolves per-feature delta_timestamps based on TrainPipelineConfig.
opentau.datasets.factory.make_dataset(cfg: DatasetConfig, train_cfg: TrainPipelineConfig, return_advantage_input: bool = False) → BaseDataset | Tuple[BaseDataset, BaseDataset][source]

Handles the logic of setting up delta timestamps and image transforms before creating a dataset.

A train and validation dataset are returned if train_cfg.val_freq is greater than 0. The validation dataset is a subset of the train dataset, and is used for evaluation during training. The validation dataset is created by splitting the train dataset into train and validation sets based on the effective split ratio: the per-dataset cfg.val_split_ratio when set, otherwise the mixture-wide train_cfg.dataset_mixture.val_split_ratio (the per-dataset value None inherits the mixture default).

Parameters:

  • cfg (DatasetConfig) – A DatasetConfig used to create a LeRobotDataset.

  • train_cfg (TrainPipelineConfig) – A TrainPipelineConfig config which contains a DatasetConfig and a PreTrainedConfig.

  • return_advantage_input (bool) – Whether the created dataset includes advantage inputs including “success”, “episode_end_idx”, “current_idx”, “last_step”, “episode_index”, and “timestamp”. Defaults to False.

Raises:

  • ValueError – If exactly one of cfg.vqa and cfg.repo_id is not provided.

  • ValueError – If cfg.vqa is not a supported vqa dataset.

Returns:

A single dataset or a tuple of (train_dataset, val_dataset) if val_freq > 0.

Return type:

BaseDataset or Tuple[BaseDataset, BaseDataset]

opentau.datasets.factory.make_dataset_mixture(cfg: TrainPipelineConfig, return_advantage_input: bool = False) → WeightedDatasetMixture | Tuple[WeightedDatasetMixture, WeightedDatasetMixture][source]

Creates a dataset mixture from the provided TrainPipelineConfig.

Parameters:

  • cfg (TrainPipelineConfig) – The configuration containing the datasets to mix. If cfg.dataset_mixture.weights is None, each dataset is weighted by its length (cast to float).

  • return_advantage_input (bool) – Whether the datasets should return advantage inputs including “success”, “episode_end_idx”, “current_idx”, “last_step”, “episode_index”, and “timestamp”. Defaults to False.

Returns:

An instance of WeightedDatasetMixture containing the datasets, or a tuple of (train_mixture, val_mixture) if val_freq > 0.

Return type:

WeightedDatasetMixture or Tuple[WeightedDatasetMixture, WeightedDatasetMixture]

opentau.datasets.factory.resolve_delta_timestamps(cfg: TrainPipelineConfig, dataset_cfg: DatasetConfig, ds_meta: LeRobotDatasetMetadata) → tuple[dict[str, ndarray], dict[str, ndarray], dict[str, ndarray], dict[str, ndarray]][source]

Resolves per-feature delta_timestamps based on TrainPipelineConfig.

Parameters:
Returns:

A 4-tuple (mean, std, lower, upper) of dicts mapping feature names to lists of delta-timestamp values. Keys that appear only in mean will be filled with sensible defaults by LeRobotDataset.compute_delta_params.