opentau.datasets

Dataset management and processing utilities for robot learning and vision-language tasks.

This module provides a comprehensive toolkit for loading, creating, managing, and processing datasets for training vision-language-action (VLA) models. It supports both robot learning datasets (with actions and states) and vision-language vqa datasets (for multimodal understanding tasks).

The module is organized into several key components:

• Core Datasets: LeRobotDataset for robot learning data with support for temporal alignment, multi-modal data, and version compatibility.

• VQA Datasets: Vision-language datasets (CLEVR, COCO-QA, VSR) for training visual understanding without robot actions.

• Dataset Mixtures: WeightedDatasetMixture for combining multiple datasets with controlled sampling proportions.

• Data Processing: Utilities for statistics computation, image/video handling, transforms, and format standardization.

• Factory Functions: High-level functions for creating datasets and mixtures from configuration objects.

Key Features:

• HuggingFace Integration: Seamless loading from HuggingFace Hub with automatic version checking and backward compatibility.

• Temporal Alignment: Delta timestamps enable sampling features at different time offsets with optional Gaussian noise for data augmentation.

• Multi-modal Support: Handles images, videos, state vectors, actions, and text prompts with automatic format conversion.

• Weighted Sampling: Combine heterogeneous datasets with configurable sampling weights for balanced training.

• Standard Data Format: Unified data format across all datasets for consistent model input/output interfaces.

• Statistics Management: Automatic computation and aggregation of dataset statistics for normalization.

• Video Handling: Multiple video backends (torchcodec, pyav, video_reader) for efficient frame extraction and encoding.

• Asynchronous I/O: High-performance image writing for real-time data recording without blocking.

Main Modules:

• lerobot_dataset: Core dataset implementation for robot learning data.

• vqa: Vision-language vqa datasets (CLEVR, COCO-QA, VSR).

• dataset_mixture: Weighted combination of multiple datasets.

• factory: Factory functions for creating datasets from configurations.

• utils: Utility functions for I/O, metadata management, and validation.

• compute_stats: Statistics computation and aggregation utilities.

• transforms: Image transformation pipelines for data augmentation.

• video_utils: Video encoding, decoding, and metadata extraction.

• image_writer: Asynchronous image writing for high-frequency recording.

• sampler: Episode-aware sampling with boundary frame filtering.

• standard_data_format_mapping: Feature name and loss type mappings.

Example

Create a dataset mixture from configuration:

>>> from opentau.datasets.factory import make_dataset_mixture
>>> mixture = make_dataset_mixture(train_cfg)
>>> dataloader = mixture.get_dataloader()

Load a single dataset:

>>> from opentau.datasets.factory import make_dataset
>>> dataset = make_dataset(dataset_cfg, train_cfg)

Access vqa datasets:

>>> from opentau import available_vqa_datasets
>>> print(list(available_vqa_datasets.keys()))
['clevr', 'cocoqa', 'dummy', 'vsr']

Modules

backward_compatibility	Backward compatibility error handling for dataset format versions.
compute_stats	Statistics computation and aggregation for dataset features.
dataset_mixture	Weighted dataset mixture for combining multiple datasets with controlled sampling.
factory	Factory functions for creating datasets and dataset mixtures.
grounding	Utilities for encoding spatial outputs as PaliGemma-style location tokens.
image_writer	Asynchronous image writing utilities for high-frequency data recording.
lerobot_dataset	LeRobot dataset implementation for robot learning data management.
online_buffer	An online buffer for the online training loop in train.py
sampler	Episode-aware sampler for PyTorch DataLoader.
speed_percentiles	Per-task percentile bucketing for the speed optional key.
standard_data_format_mapping	Standard data format mapping for dataset feature names and loss types.
transforms	Image transformation utilities for data augmentation.
utils	Utility functions for dataset management, I/O, and validation.
video_utils	Video encoding, decoding, and information extraction utilities.
vqa	Vision-language vqa datasets for multimodal learning.