opentau.datasets.dataset_mixture

Weighted dataset mixture for combining multiple datasets with controlled sampling.

This module provides functionality to combine multiple PyTorch datasets into a single weighted mixture, enabling training on heterogeneous datasets with controlled sampling proportions. It supports hierarchical sampling strategies that efficiently handle large-scale dataset combinations while maintaining memory efficiency.

The module implements a two-level sampling approach:

1. Dataset-level sampling: Selects which dataset to sample from based on specified weights.

2. Sample-level sampling: Uniformly samples within the selected dataset.

This hierarchical approach avoids expensive multinomial sampling over millions of individual samples by operating at the dataset level, making it scalable for large dataset mixtures.

Key Features:

• Weighted sampling: Control relative sampling frequency of different datasets through configurable weights.

• Memory-efficient sampling: Hierarchical sampler processes samples in chunks to minimize memory overhead.

• Metadata aggregation: Automatically aggregates and standardizes metadata from multiple datasets, including statistics normalization and feature name mapping.

• Format standardization: Converts dataset-specific feature formats to a common standard format, handling vector padding and missing cameras.

Classes:

WeightedDatasetMixture: Main class for combining multiple datasets with

weighted sampling. Creates concatenated datasets and provides DataLoader with hierarchical sampling.

HierarchicalSampler: Custom PyTorch sampler that implements two-level

weighted sampling (dataset selection, then uniform sample selection).

DatasetMixtureMetadata: Aggregates metadata from multiple datasets,

standardizes feature names, pads vectors, and combines statistics.

Functions:

pad_vector: Pads the last dimension of a vector to a target size with zeros.

Example

Create a dataset mixture with two datasets resampled to a shared 30 Hz:

>>> datasets = [dataset1, dataset2]
>>> weights = [0.7, 0.3]  # 70% from dataset1, 30% from dataset2
>>> mixture = WeightedDatasetMixture(cfg, datasets, weights, action_freq=30.0)
>>> dataloader = mixture.get_dataloader()

Mixed-frequency mixture (no resampling) — each dataset is sampled at its own native fps, so a single batch can contain samples drawn at different rates:

>>> mixture = WeightedDatasetMixture(cfg, datasets, weights, action_freq=None)

Functions

compute_norm_key(robot_type, control_mode, ...)	Compute the normalization-head key for a dataset.
pad_vector(vector, new_dim)	Pad the last dimension of a vector to a target size with zeros.

Classes

DatasetMixtureMetadata(cfg, metadatas, ...)	Per-(robot_type, control_mode) normalization metadata for a mixture.
HierarchicalSampler(dataset_lengths, ...[, ...])	With-replacement sampler for a ConcatDataset that first samples a dataset according to dataset_probs, and then samples uniformly within that dataset.
WeightedDatasetMixture(cfg, datasets, ...)	A class to combine multiple PyTorch Datasets and create a DataLoader that samples from them according to specified weightings.

class opentau.datasets.dataset_mixture.DatasetMixtureMetadata(cfg: TrainPipelineConfig, metadatas: List[DatasetMetadata], dataset_weights: List[float], dataset_names: List[str] | None = None, name_maps: List[dict[str, str] | None] | None = None)

Bases: object

Per-(robot_type, control_mode) normalization metadata for a mixture.

Each underlying dataset’s stats are normalised into the standard data format (feature renaming, state/action padding to cfg.max_state_dim / cfg.max_action_dim, missing-camera zero placeholders). Datasets sharing the same (robot_type, control_mode) are then grouped into a single norm head whose stats are sample-count-pooled via aggregate_stats(). The policy’s Normalize / Unnormalize layers stack one row per norm head and use a per-sample index (chosen via dataset_to_norm_index) to select the right row.

Datasets whose (robot_type, control_mode) pair is missing — empty, None, whitespace, or the "unknown" sentinel — fall back to keying by the dataset’s own deduplicated mixture name, giving them a private head. See compute_norm_key().

per_dataset_stats

One entry per underlying dataset, parallel to dataset_names. Used by aggregated_action_stats() and kept for diagnostic / back-compat consumers.

dataset_names

Ordered deduplicated mixture-level names (matches WeightedDatasetMixture._make_dataset_names output).

dataset_name_to_index

{name: i} reverse lookup over dataset_names (per-dataset axis, NOT the norm-head axis).

per_norm_key_stats

One entry per unique norm head, parallel to norm_keys. This is what the policy’s stacked Normalize / Unnormalize buffers consume.

norm_keys

Ordered deduplicated norm-head identifiers ("<robot_type>::<control_mode>" or fallback dataset name).

norm_key_to_index

{norm_key: row} reverse lookup over norm_keys — the norm-head axis on the policy.

dataset_to_norm_index

{dataset_name: norm_head_row} — the per-sample mapping consumed by _TaggedDataset at training time and by the policy at inference.

norm_key_to_dataset_names

{norm_key: [dataset_name, ...]} — operator diagnostic showing which datasets share each head.

__init__(cfg: TrainPipelineConfig, metadatas: List[DatasetMetadata], dataset_weights: List[float], dataset_names: List[str] | None = None, name_maps: List[dict[str, str] | None] | None = None)

aggregated_action_stats() → dict[str, ndarray]

Single mixture-wide action stats (mean/std/min/max/count).

Backwards-compat helper for the rare consumers that genuinely need a single set of action stats across the whole mixture — currently only fit_fast_tokenizer.py, which fits one BPE codec over a global action range. Most callers should consume per_dataset_stats / dataset_names directly.

property features: dict[str, dict]

Return standard data format

class opentau.datasets.dataset_mixture.HierarchicalSampler(dataset_lengths: List[int], dataset_probs: List[float], num_samples: int, *, generator: Generator | None = None, seed: int | None = None, chunk_size: int = 262144)

Bases: Sampler[int]

With-replacement sampler for a ConcatDataset that first samples a dataset according to dataset_probs, and then samples uniformly within that dataset. This avoids multinomial over a huge number of categories (over 2^24) by operating at the dataset level.

__init__(dataset_lengths: List[int], dataset_probs: List[float], num_samples: int, *, generator: Generator | None = None, seed: int | None = None, chunk_size: int = 262144)

class opentau.datasets.dataset_mixture.WeightedDatasetMixture(cfg: TrainPipelineConfig, datasets: List[BaseDataset], dataset_weights: List[float], action_freq: float | None)

Bases: object

A class to combine multiple PyTorch Datasets and create a DataLoader that samples from them according to specified weightings.

__init__(cfg: TrainPipelineConfig, datasets: List[BaseDataset], dataset_weights: List[float], action_freq: float | None)

Initializes the WeightedDatasetMixture.

Parameters:

• cfg (TrainPipelineConfig) – Configuration for the training pipeline.

• datasets (List[Dataset]) – A list of PyTorch Dataset objects.

• dataset_weights (List[float]) – A list of weights corresponding to each dataset. These determine the relative sampling frequency.

• action_freq (Optional[float]) – Common action frequency (Hz) the mixture’s datasets are resampled to. None means no resampling — each dataset is sampled at its native fps, so a single batch may mix samples from sources running at different rates (mixed-frequency training). Stored as informational state and forwarded to downstream consumers (e.g. BaseDataset._action_freq); not used arithmetically here.

get_combined_val_dataloader() → DataLoader | None

Create one deterministic sequential DataLoader over the whole mixture.

Unlike get_per_dataset_dataloaders() (one loader per dataset), this returns a single loader over the concatenated mixture so that, under accelerator.prepare, every rank’s shard is full even when individual validation subsets have fewer frames than world_size — the per-dataset loaders leave most ranks idle on tiny subsets and stack that idle time across datasets. Each sample still carries its dataset_index / dataset_repo_id provenance (injected by _TaggedDataset), so the validation loop can disaggregate metrics per (dataset, control_mode) from the batch rather than relying on homogeneous per-dataset loaders. shuffle=False + drop_last=False make the pass order-deterministic (seed-independent) and score every sample exactly once.

Returns:

A single DataLoader over the mixture, or None when the mixture is empty (mirrors the empty-dataset skip in get_per_dataset_dataloaders()).

get_dataloader() → DataLoader

Create and return a PyTorch DataLoader with weighted sampling.

Uses HierarchicalSampler to first sample a dataset according to weights, then uniformly sample within that dataset.

Returns:

DataLoader configured for weighted hierarchical sampling.

Raises:

ValueError – If no non-empty dataset has a positive sampling weight.

get_per_dataset_dataloaders() → dict[str, DataLoader]

Create one sequential DataLoader per underlying dataset.

Intended for per-dataset evaluation (e.g. per-dataset validation loss), where each dataset should be iterated exactly once rather than mixed via weighted hierarchical sampling. Empty datasets are skipped.

Returns:

Mapping from dataset_name to its DataLoader.

opentau.datasets.dataset_mixture.compute_norm_key(robot_type: str | None, control_mode: str | None, fallback_name: str) → tuple[str, bool]

Compute the normalization-head key for a dataset.

Datasets that share the same (robot_type, control_mode) are expected to share normalization stats because they share the units, axis count, and physical ranges of the proprio / action vectors. When either tag is missing, fall back to keying by the dataset’s own name so the dataset still gets a head — it just won’t share one with anything else.

A value is treated as missing when it is None, empty after strip(), or matches “unknown” case-insensitively (the sentinel that LeRobotDatasetMetadata.control_mode returns when info.json[“control_mode”] is absent — and the typical stand-in for missing robot_type).

Parameters:

• robot_type – From meta.info[“robot_type”] (after overrides).

• control_mode – From meta.info[“control_mode”] (after overrides).

• fallback_name – The dataset’s deduplicated mixture-level name, used as the head key when fallback fires.

Returns:

A (norm_key, fallback_fired) pair. norm_key is “<robot_type>::<control_mode>” (preserving the original casing of each tag) on the happy path and fallback_name otherwise; fallback_fired is True iff the fallback path was taken.

opentau.datasets.dataset_mixture.pad_vector(vector: ndarray, new_dim: int) → ndarray

Pad the last dimension of a vector to a target size with zeros.

Parameters:

• vector – Input numpy array to pad.

• new_dim – Target size for the last dimension.

Returns:

Padded array with the last dimension expanded to new_dim. If the vector already has the target dimension, returns it unchanged.