#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""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()
"""
import copy
import logging
from typing import List, Optional, Tuple, Union
import numpy as np
import torch
# NOTE: Don't delete; imported for side effects.
import opentau.datasets.vqa.clevr # noqa: F401
import opentau.datasets.vqa.cocoqa # noqa: F401
import opentau.datasets.vqa.dummy # noqa: F401
import opentau.datasets.vqa.vsr # noqa: F401
from opentau import available_vqa_datasets
from opentau.configs.default import DatasetConfig
from opentau.configs.train import TrainPipelineConfig
from opentau.datasets.dataset_mixture import WeightedDatasetMixture
from opentau.datasets.lerobot_dataset import (
BaseDataset,
LeRobotDataset,
LeRobotDatasetMetadata,
suppress_control_mode_warning,
)
from opentau.datasets.standard_data_format_mapping import DATA_FEATURES_NAME_MAPPING, feature_mapping_key
from opentau.datasets.transforms import ImageTransforms
from opentau.datasets.utils import DeltaTimestampInfo
IMAGENET_STATS = {
"min": [[[0.0]], [[0.0]], [[0.0]]], # (c,1,1)
"max": [[[1.0]], [[1.0]], [[1.0]]], # (c,1,1)
"mean": [[[0.485]], [[0.456]], [[0.406]]], # (c,1,1)
"std": [[[0.229]], [[0.224]], [[0.225]]], # (c,1,1)
}
def _apply_metadata_overrides(dataset: BaseDataset, dataset_cfg: DatasetConfig) -> None:
"""Apply ``robot_type`` / ``control_mode`` overrides from a DatasetConfig.
The overrides are written through to ``dataset.meta.info`` so downstream
consumers (``meta.control_mode``, ``_emit_optional_keys``) observe the
overridden value. ``None`` means "do not override"; any string value
(including ``""``) is applied.
"""
if dataset_cfg.robot_type is not None:
dataset.meta.info["robot_type"] = dataset_cfg.robot_type
if dataset_cfg.control_mode is not None:
dataset.meta.info["control_mode"] = dataset_cfg.control_mode
# `LeRobotDataset.__init__` caches `self.control_mode = self.meta.control_mode`
# before this override fires, so refresh the attribute when present.
if hasattr(dataset, "control_mode"):
dataset.control_mode = dataset_cfg.control_mode
[docs]
def resolve_delta_timestamps(
cfg: TrainPipelineConfig, dataset_cfg: DatasetConfig, ds_meta: LeRobotDatasetMetadata
) -> DeltaTimestampInfo:
"""Resolves per-feature delta_timestamps based on TrainPipelineConfig.
Args:
cfg (TrainPipelineConfig): The TrainPipelineConfig to read delta_indices from.
dataset_cfg (DatasetConfig): The dataset configuration.
ds_meta (LeRobotDatasetMetadata): The dataset from which features and fps are used to build
delta_timestamps against.
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``.
"""
delta_timestamps: dict[str, list[float]] = {}
action_freq = cfg.dataset_mixture.action_freq
# Mixed-frequency training: `action_freq=None` opts out of resampling.
# Substituting `ds_meta.fps` makes every delta-timestamp land exactly on
# this dataset's native frame boundaries, so nearest-neighbor sampling is
# a no-op and consecutive frames are returned unchanged.
if action_freq is None:
action_freq = ds_meta.fps
if dataset_cfg.repo_id is None:
raise ValueError("dataset_cfg.repo_id must not be None when resolving delta timestamps.")
if dataset_cfg.data_features_name_mapping is not None:
# This entry's own mapping — the registry may hold another entry's
# mapping when two entries share a repo_id and control_mode (see
# BaseDataset._get_name_map for the fetch-time counterpart).
name_map = dataset_cfg.data_features_name_mapping
else:
# Runs before `_apply_metadata_overrides`, so prefer the config's control_mode
# override, falling back to the on-disk value, to resolve dual-split columns.
control_mode = (
dataset_cfg.control_mode if dataset_cfg.control_mode is not None else ds_meta.control_mode
)
mkey = feature_mapping_key(dataset_cfg.repo_id, control_mode)
name_map = DATA_FEATURES_NAME_MAPPING[
mkey if mkey in DATA_FEATURES_NAME_MAPPING else dataset_cfg.repo_id
]
reverse_name_map = {v: k for k, v in name_map.items()}
for key in ds_meta.features:
if key not in reverse_name_map:
continue # only process camera, state, and action features
standard_key = reverse_name_map[key]
if (
standard_key == "actions"
and cfg.policy is not None
and cfg.policy.action_delta_indices is not None
):
delta_timestamps[key] = [i / action_freq for i in cfg.policy.action_delta_indices]
elif "camera" in standard_key or standard_key == "state":
n_obs = cfg.dataset_mixture.n_obs_history
if n_obs is not None:
interval = getattr(cfg.policy, "history_interval", 1)
delta_timestamps[key] = [-(n_obs - 1 - i) * interval / action_freq for i in range(n_obs)]
else:
delta_timestamps[key] = [0.0]
dt_mean = {k: np.array(v) for k, v in delta_timestamps.items()}
return dt_mean, {}, {}, {}
def _subset_meta(meta: LeRobotDatasetMetadata) -> LeRobotDatasetMetadata:
"""Per-subset metadata copy that shares read-only per-episode structures.
Returns a shallow copy of ``meta`` that shares ``episodes`` /
``episodes_stats`` (and every other attribute) by reference, deep-copying
only the small aggregated ``stats`` dict so the train/val halves can hold
independent normalization stats without cross-contamination.
``random_split`` divides a dataset by frame index, so the per-episode
``episodes`` / ``episodes_stats`` are identical across the two subsets and
are read-only on the training path (only the dataset-*creation*
``save_episode`` / stats-conversion paths write them). A blanket
``deepcopy(meta)`` cloned them per subset: for a large per-episode-stats
dataset listed under many mixture configs that is tens of GB of needless
copies per rank, and — worse — copy-on-write surface that blows up host RAM
once forked across dataloader workers (every worker touches the clones).
Sharing them by reference removes both costs. Mirrors the share-by-reference
contract of :meth:`BaseDataset.shallow_copy_with_dropout`.
Args:
meta: The source dataset metadata to copy for a train/val subset.
"""
subset_meta = copy.copy(meta)
subset_meta.stats = copy.deepcopy(meta.stats)
return subset_meta
[docs]
def make_dataset(
cfg: DatasetConfig,
train_cfg: TrainPipelineConfig,
return_advantage_input: bool = False,
) -> Union[BaseDataset, Tuple[BaseDataset, BaseDataset]]:
"""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).
Args:
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:
BaseDataset or Tuple[BaseDataset, BaseDataset]: A single dataset or a tuple of (train_dataset, val_dataset) if val_freq > 0.
"""
image_transforms = ImageTransforms(cfg.image_transforms) if cfg.image_transforms.enable else None
if isinstance(cfg.vqa, str) + isinstance(cfg.repo_id, str) != 1:
raise ValueError("Exactly one of `cfg.vqa` and `cfg.repo_id` should be provided.")
if isinstance(cfg.vqa, str):
ds_cls = available_vqa_datasets.get(cfg.vqa)
if ds_cls is None:
raise ValueError(
f"Unknown vqa dataset '{cfg.vqa}'. Supported datasets are: {available_vqa_datasets.keys()}"
)
# TODO support dataset-specific arg / kwargs
dataset = ds_cls(train_cfg)
elif isinstance(cfg.repo_id, str):
ds_meta = LeRobotDatasetMetadata(cfg.repo_id, root=cfg.root, revision=cfg.revision)
dt_mean, dt_std, dt_lower, dt_upper = resolve_delta_timestamps(train_cfg, cfg, ds_meta)
# Suppress the "missing control_mode" warning when the user is
# providing an explicit override — they already know it's missing.
# Ordering invariant: this MUST run before `LeRobotDataset(...)` below;
# once `__init__` emits the warning the suppression is a no-op.
if cfg.control_mode is not None:
suppress_control_mode_warning(cfg.repo_id)
# Per-dataset values win over the mixture-wide default; `None` means
# "inherit". See `DatasetConfig` / `DatasetMixtureConfig` docstrings.
effective_tolerance = (
cfg.tolerance_s if cfg.tolerance_s is not None else train_cfg.dataset_mixture.tolerance_s
)
effective_skip = (
cfg.skip_timestamp_check
if cfg.skip_timestamp_check is not None
else train_cfg.dataset_mixture.skip_timestamp_check
)
dataset = LeRobotDataset(
train_cfg,
cfg.repo_id,
root=cfg.root,
episodes=cfg.episodes,
excluded_episodes=cfg.excluded_episodes,
delta_timestamps=dt_mean,
delta_timestamps_std=dt_std,
delta_timestamps_lower=dt_lower,
delta_timestamps_upper=dt_upper,
tolerance_s=effective_tolerance,
image_transforms=image_transforms,
revision=cfg.revision,
video_backend=cfg.video_backend,
image_resample_strategy=train_cfg.dataset_mixture.image_resample_strategy,
vector_resample_strategy=train_cfg.dataset_mixture.vector_resample_strategy,
return_advantage_input=return_advantage_input,
skip_timestamp_check=effective_skip,
prompt_substitutions=cfg.prompt_substitutions,
data_features_name_mapping=cfg.data_features_name_mapping,
)
else:
raise ValueError("Exactly one of `cfg.vqa` and `cfg.repo_id` should be provided.")
_apply_metadata_overrides(dataset, cfg)
# TODO vqa datasets implement stats in original feature names, but camera_keys are standardized names
if (
not isinstance(cfg.vqa, str)
and isinstance(cfg.repo_id, str)
and "dummy" not in cfg.repo_id
and cfg.use_imagenet_stats
):
if dataset.meta.stats is None:
dataset.meta.stats = {}
for key in dataset.meta.camera_keys:
for stats_type, stats in IMAGENET_STATS.items():
if key not in dataset.meta.stats:
dataset.meta.stats[key] = {}
dataset.meta.stats[key][stats_type] = np.array(stats, dtype=np.float32)
if train_cfg.val_freq > 0:
# Per-dataset value wins over the mixture-wide default; `None` means
# "inherit". Mirrors the `tolerance_s` / `skip_timestamp_check`
# resolution above. See `DatasetConfig` / `DatasetMixtureConfig` docs.
effective_val_split = (
cfg.val_split_ratio
if cfg.val_split_ratio is not None
else train_cfg.dataset_mixture.val_split_ratio
)
val_size = int(len(dataset) * effective_val_split)
train_size = len(dataset) - val_size
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
# Share the large, read-only per-episode metadata (`episodes`,
# `episodes_stats`) across the two halves; deep-copy only the small
# aggregated `stats`. A blanket `deepcopy(meta)` here is tens of GB of
# needless per-rank copies for a large per-episode-stats dataset and
# becomes copy-on-write surface that OOMs the host across dataloader
# workers. See `_subset_meta`.
train_dataset.meta = _subset_meta(dataset.meta) # type: ignore[assignment]
val_dataset.meta = _subset_meta(dataset.meta) # type: ignore[assignment]
# Subset wraps the same underlying dataset by reference, so the
# training and validation halves would share every instance attribute
# — including the optional-key dropout and prompt-substitution flags.
# Give the val subset its own shallow copy whose only divergent
# attributes are those toggles. See
# ``BaseDataset.shallow_copy_with_dropout`` for the contract on what
# stays shared.
val_dataset.dataset = dataset.shallow_copy_with_dropout( # type: ignore[attr-defined]
enable_dropout=train_cfg.dataset_mixture.val_enable_optional_key_dropout,
enable_prompt_substitution=train_cfg.dataset_mixture.val_enable_prompt_substitution,
)
return train_dataset, val_dataset # type: ignore[return-value]
return dataset
logger = logging.getLogger(__name__)
def _resolve_weights(
configured_weights: Optional[List[float]], datasets: list, label: str = "datasets"
) -> List[float]:
"""Return explicit weights or infer them from dataset lengths.
Args:
configured_weights: User-provided weights, or None to infer.
datasets: The list of datasets whose lengths are used when
``configured_weights`` is None.
label: Human-readable label used in the log message
(e.g. "train" or "val").
Returns:
A list of float weights, one per dataset.
"""
if configured_weights is not None:
return configured_weights
weights = [float(len(ds)) for ds in datasets]
logger.info("No explicit weights provided; inferring %s weights from dataset lengths: %s", label, weights)
return weights
def _validate_metadata_requirements(cfg: TrainPipelineConfig, datasets: list, label: str) -> None:
"""Raise if the mixture requires non-empty robot_type / control_mode and
any dataset (after overrides) still has an empty value.
"""
require_robot = cfg.dataset_mixture.require_non_empty_robot_type
require_control = cfg.dataset_mixture.require_non_empty_control_mode
if not (require_robot or require_control):
return
dataset_cfgs = cfg.dataset_mixture.datasets
# Invariant from `make_dataset_mixture`: each dataset_cfg appends exactly
# one entry to `datasets` (and at most one to `val_datasets`). Assert
# rather than silently skipping so a future refactor that breaks the
# invariant doesn't quietly bypass the require_non_empty_* checks.
assert len(dataset_cfgs) == len(datasets), (
f"dataset_cfgs ({len(dataset_cfgs)}) and {label} datasets ({len(datasets)}) "
"must be 1:1; cannot validate metadata requirements."
)
bad: list[str] = []
for dc, ds in zip(dataset_cfgs, datasets, strict=True):
info = ds.meta.info
identifier = dc.repo_id or dc.vqa or type(ds).__name__
if require_robot and not (info.get("robot_type") or ""):
bad.append(f"{identifier}: robot_type is empty")
if require_control and not (info.get("control_mode") or ""):
bad.append(f"{identifier}: control_mode is empty")
if bad:
raise ValueError(
"DatasetMixtureConfig requires non-empty metadata fields, but the "
f"following {label} datasets are missing values after overrides:\n - "
+ "\n - ".join(bad)
+ "\nSet `DatasetConfig.robot_type` / `DatasetConfig.control_mode` "
"on the offending dataset(s) to provide an override."
)
[docs]
def make_dataset_mixture(
cfg: TrainPipelineConfig, return_advantage_input: bool = False
) -> Union[WeightedDatasetMixture, Tuple[WeightedDatasetMixture, WeightedDatasetMixture]]:
"""Creates a dataset mixture from the provided TrainPipelineConfig.
Args:
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:
WeightedDatasetMixture or Tuple[WeightedDatasetMixture, WeightedDatasetMixture]: An instance of WeightedDatasetMixture containing the datasets, or a tuple of (train_mixture, val_mixture) if val_freq > 0.
"""
datasets = []
val_datasets = []
for dataset_cfg in cfg.dataset_mixture.datasets:
res = make_dataset(dataset_cfg, cfg, return_advantage_input=return_advantage_input)
if isinstance(res, tuple):
datasets.append(res[0])
val_datasets.append(res[1])
else:
datasets.append(res)
_validate_metadata_requirements(cfg, datasets, label="train")
if val_datasets:
_validate_metadata_requirements(cfg, val_datasets, label="val")
train_weights = _resolve_weights(cfg.dataset_mixture.weights, datasets, label="train")
train_mixture = WeightedDatasetMixture(cfg, datasets, train_weights, cfg.dataset_mixture.action_freq)
if val_datasets:
val_weights = _resolve_weights(cfg.dataset_mixture.weights, val_datasets, label="val")
val_mixture = WeightedDatasetMixture(cfg, val_datasets, val_weights, cfg.dataset_mixture.action_freq)
return train_mixture, val_mixture
return train_mixture