opentau.datasets.lerobot_dataset

LeRobot dataset implementation for robot learning data management.

This module provides the core dataset implementation for loading, creating, and managing robot learning datasets. It supports both loading existing datasets from the HuggingFace Hub or local disk, as well as creating new datasets for data recording.

The dataset structure consists of:

• Metadata: Info, statistics, tasks, and episode information stored as JSON

• Data files: Episode data stored as Parquet files organized by chunks

• Videos: Optional video files for camera observations stored as MP4 files

Key Features:

• Temporal alignment: Supports delta timestamps for temporal feature alignment, enabling sampling of 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 and standardization.

• Version compatibility: Automatic version checking and backward compatibility handling for datasets created with older format versions.

• Asynchronous image writing: Optional async image writer for high-frequency data recording without blocking the main process.

• Statistics management: Per-episode and aggregated statistics for data normalization, with automatic computation and aggregation.

• Video handling: Supports multiple video backends (torchcodec, pyav, video_reader) for efficient video encoding and decoding.

Classes:

DatasetMetadata

Base class for dataset metadata management.

LeRobotDatasetMetadata

Metadata manager for LeRobot datasets with Hub integration, version checking, and statistics loading.

VQADatasetMetadata

Metadata manager for vqa datasets.

BaseDataset

Base PyTorch Dataset class with common functionality.

LeRobotDataset

Main dataset class for robot learning data, supporting loading from Hub/local disk, temporal alignment, video/image handling, and data recording.

Functions:

retry_random_on_failure

Decorator to retry dataset item retrieval with random indices on failure.

Example

Load an existing dataset:

>>> dataset = LeRobotDataset(cfg, repo_id="my-robot-dataset")
>>> dataloader = DataLoader(dataset, batch_size=32)

Create a new dataset for recording:

>>> dataset = LeRobotDataset.create(
...     repo_id="my-new-dataset",
...     fps=30,
...     features={"state": {"shape": (7,), "dtype": "float32"}},
...     use_videos=True
... )

Functions

aggregate_selected_stats(meta[, episodes, ...])	Per-feature stats aggregated over the SELECTED episodes.
resolve_selected_episodes(all_episodes, ...)	Resolve the effective episode list for a dataset selection.
retry_random_on_failure(f)	Decorator to retry dataset item retrieval with random indices on failure.
suppress_control_mode_warning(repo_id)	Mark repo_id as already-warned so the missing-control_mode warning does not fire for it during subsequent LeRobotDataset.__init__ calls.

Classes

BaseDataset(cfg)	Base class for all robot learning datasets.
DatasetMetadata(*[, info, stats])	Base class for dataset metadata containing info and statistics.
LeRobotDataset(cfg, repo_id[, root, ...])	Main dataset class for loading and managing robot learning data.
LeRobotDatasetMetadata(repo_id[, root, ...])	Metadata manager for LeRobot datasets with Hub integration and version handling.
VQADatasetMetadata(*[, info, stats])	Metadata class for vqa datasets (vision-language datasets).

class opentau.datasets.lerobot_dataset.BaseDataset(cfg: TrainPipelineConfig)

Bases: Dataset

Base class for all robot learning datasets.

This abstract base class provides common functionality for both LeRobotDataset and VQADataset, including data format standardization, image processing, and vector padding. It ensures all datasets conform to a standard format regardless of their source or structure.

Key Features:

• Standard data format conversion: Maps dataset-specific feature names to standard names (camera0, camera1, state, actions, etc.)

• Image standardization: Resizes and pads images to target resolution while maintaining aspect ratio

• Vector padding: Pads state and action vectors to maximum dimensions

• Data type conversion: Converts floating-point tensors to bfloat16 for memory efficiency

• String normalization: Ensures prompts and responses have consistent newline formatting

Subclasses must implement:

• _get_feature_mapping_key(): Returns the key used for feature name mapping (e.g., “lerobot/aloha_mobile_cabinet”)

resolution

Target image resolution (height, width).

num_cams

Number of camera views in each sample.

max_state_dim

Maximum dimension for state vectors.

max_action_dim

Maximum dimension for action vectors.

action_chunk

Number of actions processed in a chunk.

Example

Create a custom dataset:

>>> class MyDataset(BaseDataset):
...     def _get_feature_mapping_key(self):
...         return "my-dataset"

__init__(cfg: TrainPipelineConfig)

static pad_vector(vector, new_dim)

Only the last dimension of the vector is padded to ‘new_dim’ with zeros.

resize_with_pad(img, width, height, pad_value=0) → Tensor

Resize an image to target dimensions with padding.

Maintains aspect ratio by resizing to fit within target dimensions, then pads on the left and top to reach exact target size.

Parameters:

• img – Input image tensor of shape (C, H, W) or (T, C, H, W).

• width – Target width.

• height – Target height.

• pad_value – Value to use for padding. Defaults to 0.

Returns:

Resized and padded image tensor of shape (C, height, width) or (T, C, height, width), matching the input rank.

Raises:

ValueError – If input is not 3D or 4D.

shallow_copy_with_dropout(*, enable_dropout: bool, enable_prompt_substitution: bool = False) → BaseDataset

Return a shallow copy that only diverges in the augmentation toggles.

Used by opentau.datasets.factory.make_dataset() to give the validation subset its own dataset instance so optional-key dropout and prompt substitution can be toggled independently of the training subset. The copy shares meta, hf_dataset, episode_data_index, cached segment tables, and every other instance attribute by reference — only enable_optional_key_dropout and enable_prompt_substitution diverge. If a future refactor adds an instance attribute that needs to diverge between train and val, it must be set here too.

class opentau.datasets.lerobot_dataset.DatasetMetadata(*, info: dict | None = None, stats: dict | None = None)

Bases: object

Base class for dataset metadata containing info and statistics.

info

Dictionary containing dataset information (features, fps, etc.).

stats

Dictionary containing dataset statistics for normalization.

repo_id

Repository ID of the dataset (set by subclasses).

__init__(*, info: dict | None = None, stats: dict | None = None)

property camera_keys: list[str]

Keys to access visual modalities (regardless of their storage method).

property features: dict[str, dict]

All features contained in the dataset.

property fps: int | None

Native frame rate of the dataset.

Returns None for datasets without a temporal axis (e.g. VQA image-text datasets). Subclasses with real frame-rate metadata (LeRobotDatasetMetadata) override this to return the int from info["fps"]. Downstream callers that emit fps as a sample key treat None as the pad signal so heterogeneous mixtures (VLA + VQA) stay batchable.

property image_keys: list[str]

Keys to access visual modalities stored as images.

property names: dict[str, list | dict]

Names of the various dimensions of vector modalities.

property shapes: dict

Shapes for the different features.

property video_keys: list[str]

Keys to access visual modalities stored as videos.

class opentau.datasets.lerobot_dataset.LeRobotDataset(cfg: TrainPipelineConfig, repo_id: str, root: str | Path | None = None, episodes: list[int] | None = None, excluded_episodes: list[int] | None = None, image_transforms: Callable | None = None, delta_timestamps: dict[str, ndarray | list[float]] | None = None, delta_timestamps_std: dict[str, ndarray | list[float]] | None = None, delta_timestamps_lower: dict[str, ndarray | list[float]] | None = None, delta_timestamps_upper: dict[str, ndarray | list[float]] | None = None, tolerance_s: float = 0.0001, revision: str | None = None, force_cache_sync: bool = False, download_videos: bool = True, video_backend: str | None = None, image_resample_strategy: str = 'nearest', vector_resample_strategy: str = 'nearest', standardize: bool = True, return_advantage_input: bool = False, skip_timestamp_check: bool = False, prompt_substitutions: dict[str, list[str]] | None = None, data_features_name_mapping: dict[str, str] | None = None)

Bases: BaseDataset

Main dataset class for loading and managing robot learning data.

This class provides a PyTorch Dataset interface for robot learning datasets stored in the LeRobot format. It supports loading from HuggingFace Hub or local disk, handles temporal alignment with delta timestamps, manages video and image data, and provides data recording capabilities.

The dataset structure consists of:

• Metadata: JSON files containing info, statistics, episodes, tasks

• Data files: Parquet files organized by chunks containing episode data

• Videos: Optional MP4 files for camera observations

Key Features:

• Hub integration: Automatic download from HuggingFace Hub with version compatibility checking

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

• Video/image handling: Supports both video files and individual images with automatic frame extraction and synchronization

• Episode filtering: Load specific episodes by index

• Data recording: Create new datasets and add episodes programmatically

• Statistics: Per-episode and aggregated statistics for normalization

Two Usage Modes:

1. Loading existing datasets: From local disk or HuggingFace Hub

2. Creating new datasets: Using the create() classmethod for data recording

cfg

Training pipeline configuration.

repo_id

Repository ID of the dataset.

root

Local root directory for the dataset.

meta

LeRobotDatasetMetadata instance containing all metadata.

hf_dataset

HuggingFace Dataset containing parquet data.

episodes

Dictionary mapping episode_index to episode info.

image_transforms

Optional image transforms to apply.

delta_timestamps_params

Processed delta timestamp parameters.

video_backend

Backend used for video decoding.

standardize

Whether to standardize data format.

Example

Load dataset from Hub:

>>> dataset = LeRobotDataset(cfg, repo_id="lerobot/aloha")
>>> dataloader = DataLoader(dataset, batch_size=32)

Load specific episodes:

>>> dataset = LeRobotDataset(
...     cfg,
...     repo_id="lerobot/aloha",
...     episodes=[0, 1, 2, 5, 10]
... )

Create new dataset for recording:

>>> dataset = LeRobotDataset.create(
...     cfg,
...     repo_id="my-new-dataset",
...     fps=30,
...     features={"state": {"dtype": "float32", "shape": (7,)}}
... )

__init__(cfg: TrainPipelineConfig, repo_id: str, root: str | Path | None = None, episodes: list[int] | None = None, excluded_episodes: list[int] | None = None, image_transforms: Callable | None = None, delta_timestamps: dict[str, ndarray | list[float]] | None = None, delta_timestamps_std: dict[str, ndarray | list[float]] | None = None, delta_timestamps_lower: dict[str, ndarray | list[float]] | None = None, delta_timestamps_upper: dict[str, ndarray | list[float]] | None = None, tolerance_s: float = 0.0001, revision: str | None = None, force_cache_sync: bool = False, download_videos: bool = True, video_backend: str | None = None, image_resample_strategy: str = 'nearest', vector_resample_strategy: str = 'nearest', standardize: bool = True, return_advantage_input: bool = False, skip_timestamp_check: bool = False, prompt_substitutions: dict[str, list[str]] | None = None, data_features_name_mapping: dict[str, str] | None = None)

Initialize LeRobotDataset.

2 modes are available for instantiating this class, depending on 2 different use cases:

1. Your dataset already exists:

   • On your local disk in the ‘root’ folder. This is typically the case when you recorded your dataset locally and you may or may not have pushed it to the hub yet. Instantiating this class with ‘root’ will load your dataset directly from disk. This can happen while you’re offline (no internet connection).

   • On the Hugging Face Hub at the address https://huggingface.co/datasets/{repo_id} and not on your local disk in the ‘root’ folder. Instantiating this class with this ‘repo_id’ will download the dataset from that address and load it, pending your dataset is compliant with codebase_version v2.0. If your dataset has been created before this new format, you will be prompted to convert it using our conversion script from v1.6 to v2.0, which you can find at lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py.

2. Your dataset doesn’t already exists (either on local disk or on the Hub): you can create an empty LeRobotDataset with the ‘create’ classmethod. This can be used for recording a dataset or port an existing dataset to the LeRobotDataset format.

In terms of files, LeRobotDataset encapsulates 3 main things:

• metadata:

  • info contains various information about the dataset like shapes, keys, fps etc.

  • stats stores the dataset statistics of the different modalities for normalization

  • tasks contains the prompts for each task of the dataset, which can be used for task-conditioned training.

• hf_dataset (from datasets.Dataset), which will read any values from parquet files.

• videos (optional) from which frames are loaded to be synchronous with data from parquet files.

A typical LeRobotDataset looks like this from its root path:

.
├── data
│   ├── chunk-000
│   │   ├── episode_000000.parquet
│   │   ├── episode_000001.parquet
│   │   ├── episode_000002.parquet
│   │   └── ...
│   ├── chunk-001
│   │   ├── episode_001000.parquet
│   │   ├── episode_001001.parquet
│   │   ├── episode_001002.parquet
│   │   └── ...
│   └── ...
├── meta
│   ├── episodes.jsonl
│   ├── info.json
│   ├── stats.json
│   └── tasks.jsonl
└── videos
    ├── chunk-000
    │   ├── observation.images.laptop
    │   │   ├── episode_000000.mp4
    │   │   ├── episode_000001.mp4
    │   │   ├── episode_000002.mp4
    │   │   └── ...
    │   ├── observation.images.phone
    │   │   ├── episode_000000.mp4
    │   │   ├── episode_000001.mp4
    │   │   ├── episode_000002.mp4
    │   │   └── ...
    ├── chunk-001
    └── ...

Note that this file-based structure is designed to be as versatile as possible. The files are split by episodes which allows a more granular control over which episodes one wants to use and download. The structure of the dataset is entirely described in the info.json file, which can be easily downloaded or viewed directly on the hub before downloading any actual data. The type of files used are very simple and do not need complex tools to be read, it only uses .parquet, .json and .mp4 files (and .md for the README).

Parameters:

• cfg (TrainPipelineConfig) – Training configuration object.

• repo_id (str) – This is the repo id that will be used to fetch the dataset. Locally, the dataset will be stored under root/repo_id.

• root (Path | None, optional) – Local directory to use for downloading/writing files. You can also set the HF_OPENTAU_HOME environment variable to point to a different location. Defaults to ‘~/.cache/huggingface/opentau’.

• episodes (list[int] | None, optional) – If specified, this will only load episodes specified by their episode_index in this list. Defaults to None.

• excluded_episodes (list[int] | None, optional) – Episode indices to drop. Takes precedence over episodes (an index present in both is excluded); when episodes is None the denylist is applied to the full episode list. Defaults to None.

• image_transforms (Callable | None, optional) – You can pass standard v2 image transforms from torchvision.transforms.v2 here which will be applied to visual modalities (whether they come from videos or images). Defaults to None.

• delta_timestamps (dict[list[float]] | None, optional) – Dictionary mapping feature names to lists of delta timestamps in seconds. For example, {'state': [0.0], 'action': [0, 0.5, 1.0]} means state is sampled at the current time and action is sampled at three offsets. A {feature}_is_pad boolean mask of the same length is added to the returned item. Defaults to None.

• delta_timestamps_std – (dict[list[float]] | None, optional): Per-feature standard deviation for the delta timestamps. Absent keys are treated as deterministic (std=0). Defaults to None.

• delta_timestamps_lower – (dict[list[float]] | None, optional): Per-feature lower bound for the delta timestamps. Defaults to None.

• delta_timestamps_upper – (dict[list[float]] | None, optional): Per-feature upper bound for the delta timestamps. Defaults to None.

• tolerance_s (float, optional) – Tolerance in seconds used to ensure data timestamps are actually in sync with the fps value. It is used at the init of the dataset to make sure that each timestamps is separated to the next by 1/fps +/- tolerance_s. This also applies to frames decoded from video files. Defaults to 1e-4.

• revision (str, optional) – An optional Git revision id which can be a branch name, a tag, or a commit hash. Defaults to current codebase version tag.

• download_videos (bool, optional) – Flag to download the videos. Note that when set to True but the video files are already present on local disk, they won’t be downloaded again. Defaults to True.

• video_backend (str | None, optional) – Video backend to use for decoding videos. Defaults to torchcodec when available int the platform; otherwise, defaults to ‘pyav’. You can also use the ‘pyav’ decoder used by Torchvision, which used to be the default option, or ‘video_reader’ which is another decoder of Torchvision.

• image_resample_strategy – str: Resampling strategy to use for image features. If ‘linear’, it will use linear interpolation between two immediate timestamps. If ‘nearest’, it will use nearest neighbor interpolation. Defaults to ‘nearest’.

• vector_resample_strategy – str: Resampling strategy to use for non-image features, such as action or state. If ‘linear’, it will use linear interpolation between two immediate timestamps. If ‘nearest’, it will use nearest neighbor interpolation. Defaults to ‘nearest’.

• standardize (bool, Optional) – Flag to enable standardization in __getitem__. Defaults to True.

• return_advantage_input (bool, Optional) – Flag to return advantage inputs (“success”, “episode_end_idx”, “current_idx”, “last_step”, “episode_index”, “timestamp”, ). Defaults to False. Ignored if standardize is False.

• skip_timestamp_check (bool, Optional) – If True, bypass the load-time check_timestamps_sync call (a warning is logged). Frame-to-frame spacing is NOT verified against 1/fps; downstream delta_timestamps lookups may sample unintended frames. Defaults to False. Does not affect the record-time check inside add_episode.

• prompt_substitutions (dict[str, list[str]] | None, optional) – Mapping from an on-disk task string (exact match against meta.tasks) to a non-empty list of non-empty substitute prompts. During __getitem__ a matching sample’s task is ALWAYS replaced by a uniform random draw from its list (gated by self.enable_prompt_substitution); unmapped tasks pass through unchanged. Keys matching no on-disk task string raise a ValueError at init. Defaults to None.

• data_features_name_mapping (dict[str, str] | None, optional) – This instance’s standard-role -> dataset-column mapping. When set it wins over the process-global DATA_FEATURES_NAME_MAPPING registry (see BaseDataset._get_name_map()), so two mixture entries sharing a repo_id and control_mode can declare different mappings (e.g. two camera views of the same repo) without clobbering each other. None (default) falls back to the registry.

add_frame(frame: dict) → None

This function only adds the frame to the episode_buffer. Apart from images — which are written in a temporary directory — nothing is written to disk. To save those frames, the ‘save_episode()’ method then needs to be called.

Video features listed in deferred_video_keys (set via create()) may be omitted from the frame; their observations will be attached later via attach_video().

attach_video(episode_index: int, video_key: str, input_video_path: str | Path, overwrite: bool = False, vcodec: str = 'libsvtav1', pix_fmt: str = 'yuv420p', g: int | None = 2, crf: int | None = 30, start_time: float | None = None) → Path

Attach a pre-recorded MP4 video to an episode with deferred video observations.

The source video is resampled to the dataset’s FPS and trimmed so it contains exactly as many frames as the episode. The resulting file is placed in the standard video path for the dataset.

This method is meant to be called after save_episode() for episodes whose video observations were deferred (see the deferred_video_keys parameter of create()).

After attaching videos for all deferred keys and episodes, you should call update_video_info() to update the dataset metadata with the actual video properties (resolution, codec, etc.).

Parameters:

• episode_index – Index of the episode to attach the video to.

• video_key – The video feature key (e.g. "observation.images.top").

• input_video_path – Path to the source MP4 file.

• overwrite – Whether to overwrite an existing video at the target path.

• vcodec – Video codec for re-encoding. Defaults to “libsvtav1”.

• pix_fmt – Pixel format. Defaults to “yuv420p”.

• g – GOP size. Defaults to 2.

• crf – Constant Rate Factor. Defaults to 30.

• start_time – Optional start offset in seconds into the source video. Useful when only a portion of the recording overlaps with the episode data. Defaults to None (start from the beginning).

Returns:

Path to the written video file inside the dataset.

Raises:

• ValueError – If video_key is not a declared video feature, or the episode index does not exist.

• FileNotFoundError – If input_video_path does not exist.

clear_episode_buffer() → None

static compute_delta_params(mean: dict[str, ndarray | list[float]] | None, std: dict[str, ndarray | list[float]] | None, lower: dict[str, ndarray | list[float]] | None, upper: dict[str, ndarray | list[float]] | None) → tuple[dict[str, ndarray], dict[str, ndarray], dict[str, ndarray], dict[str, ndarray]]

Process the parameters mean, std, lower and upper for delta timestamps.

Delta timestamps are computed dynamically in __getitem__ with clip(dT, lower, upper) where dT ~ N(mean, std^2). Each parameter is a dictionary mapping feature names to sequences of floats.

For example, mean = {"state": [0.0], "action": [0.0, 0.5, 1.0]} means the state feature will be sampled at t + 0.0 and the action feature will be sampled at three offsets t, t+0.5 and t+1.0.

Keys present in std / lower / upper but absent from mean are ignored. Keys absent from std / lower / upper but present in mean receive sensible defaults (0 for std, -inf / +inf for bounds).

classmethod create(repo_id: str, fps: int, root: str | Path | None = None, robot_type: str | None = None, features: dict | None = None, use_videos: bool = True, tolerance_s: float = 0.0001, image_writer_processes: int = 0, image_writer_threads: int = 0, video_backend: str | None = None, image_resample_strategy: str = 'nearest', vector_resample_strategy: str = 'nearest', standardize: bool = True, skip_video_stats: bool = False, deferred_video_keys: set[str] | None = None) → LeRobotDataset

Create a LeRobot Dataset from scratch in order to record data.

Parameters:

deferred_video_keys – Optional set of video feature keys whose image observations will be provided later via attach_video() instead of being passed to add_frame(). When set, these keys are omitted from frame validation, episode video encoding, and post-save video existence checks. After all episodes are recorded, call attach_video() for each episode to supply an MP4 that will be resampled to the dataset FPS and trimmed to the episode length.

create_episode_buffer(episode_index: int | None = None) → dict

create_hf_dataset() → Dataset

Create an empty HuggingFace dataset with the correct features.

Returns:

Empty dataset with features matching the dataset specification.

download_episodes(download_videos: bool = True) → None

Downloads the dataset from the given ‘repo_id’ at the provided version. If ‘episodes’ is given, this will only download those episodes (selected by their episode_index). If ‘episodes’ is None, the whole dataset will be downloaded. Already-present files in ‘local_dir’ are not re-downloaded.

download_files(files: list[str]) → None

Fetch the files from files that are missing on disk, one hf_hub_download each.

snapshot_download(allow_patterns=<thousands of explicit paths>) spends many minutes GIL-held and I/O-idle inside filter_repo_objects, whose fnmatch loop is O(repo_files x patterns) — long enough to trip the NCCL watchdog while sibling ranks wait at the _sync broadcast. hf_hub_download targets each file by exact path, skipping the filter entirely; a thread pool overlaps the network round-trips (the GIL is released during I/O).

encode_episode_videos(episode_index: int, skip_keys: set[str] | None = None) → dict

Use ffmpeg to convert frames stored as png into mp4 videos. Note: encode_video_frames is a blocking call. Making it asynchronous shouldn’t speedup encoding, since video encoding with ffmpeg is already using multithreading.

Parameters:

• episode_index – Index of the episode to encode.

• skip_keys – Optional set of video keys to skip (e.g. deferred video keys).

encode_videos() → None

Use ffmpeg to convert frames stored as png into mp4 videos. Note: encode_video_frames is a blocking call. Making it asynchronous shouldn’t speedup encoding, since video encoding with ffmpeg is already using multithreading.

property features: dict[str, dict]

property fps: int

Frames per second used during data collection.

get_episodes_file_paths() → list[str]

Get file paths for all selected episodes.

Returns paths for both parquet data files and video files (if applicable) for all episodes in the dataset.

Returns:

List of file paths for episode data and videos.

property hf_features: Features

Features of the hf_dataset.

load_hf_dataset() → Dataset

hf_dataset contains all the observations, states, actions, rewards, etc.

Loads the per-episode parquet files via load_dataset(“parquet”, …), which builds a memory-mapped Arrow cache under $HF_HOME/datasets/. The cache costs disk (~1-5x the parquet, compression-dependent — see #277) and nothing prunes it, so a multi-hundred-GB mixture needs that much extra disk provisioned. In exchange the loaded dataset is genuinely memory-mapped (resident pages are file-backed and reclaimable), so RAM stays bounded by the OS page cache rather than the dataset size. That is essential here: 8 ranks each load the full mixture, and the multi-hundred-GB video repos would otherwise OOM the node. A hand-rolled pa_ds.to_table() + Dataset(table) (or streaming to a self-written Arrow IPC file + Dataset.from_file) was tried and both materialised into anonymous RAM instead of mmapping — load_dataset routes through HF’s ParquetDatasetBuilder, whose Arrow cache Dataset.from_file does memory-map correctly.

Schema is inferred from the parquet files themselves; it is intentionally not validated against meta/info.json. So a parquet/info.json mismatch now loads silently — the parquet’s own schema wins and info.json is no longer authoritative. A mismatch between the parquet files of a single dataset still fails, but as a load_dataset concatenation error rather than the old explicit feature-cast error.

Files are passed in sorted-episode order so the hf_dataset row layout stays aligned with episode_data_index[“from”/”to”] — see the sorted(episodes) note in __init__. self.episodes is always a concrete list here: __init__ backfills it with the full episode list before this method is ever called.

property num_episodes: int

Number of episodes selected.

property num_frames: int

Number of frames in selected episodes.

pull_from_repo(allow_patterns: list[str] | str | None = None, ignore_patterns: list[str] | str | None = None) → None

push_to_hub(branch: str | None = None, tags: list | None = None, license: str | None = 'apache-2.0', tag_version: bool = True, push_videos: bool = True, private: bool = False, allow_patterns: list[str] | str | None = None, upload_large_folder: bool = False, **card_kwargs) → None

save_episode(episode_data: dict | None = None) → None

This will save to disk the current episode in self.episode_buffer.

Parameters:

episode_data (dict | None, optional) – Dict containing the episode data to save. If None, this will save the current episode in self.episode_buffer, which is filled with ‘add_frame’. Defaults to None.

Note

When deferred_video_keys are configured, the corresponding video features are excluded from validation, encoding, and existence checks. Use attach_video() after saving episodes to supply the video files.

start_image_writer(num_processes: int = 0, num_threads: int = 4) → None

stop_image_writer() → None

Whenever wrapping this dataset inside a parallelized DataLoader, this needs to be called first to remove the image_writer in order for the LeRobotDataset object to be pickleable and parallelized.

update_video_info() → None

Update video metadata from the first episode’s video files.

Call this after attaching all deferred videos to populate the info field of each video feature with actual video properties (resolution, codec, etc.) and persist the updated metadata to disk.

class opentau.datasets.lerobot_dataset.LeRobotDatasetMetadata(repo_id: str, root: str | Path | None = None, revision: str | None = None, force_cache_sync: bool = False)

Bases: DatasetMetadata

Metadata manager for LeRobot datasets with Hub integration and version handling.

This class manages all metadata for LeRobot datasets, including dataset info, statistics, episodes, tasks, and advantages. It handles loading from local disk or HuggingFace Hub, version compatibility checking, and provides utilities for accessing dataset files and information.

The class automatically handles:

• Loading metadata from local disk or downloading from HuggingFace Hub

• Version compatibility checking and automatic version resolution

• Backward compatibility with older dataset formats (v2.0 vs v2.1)

• Episode and task management

• Statistics aggregation (per-episode and global)

repo_id

Repository ID of the dataset on HuggingFace Hub.

root

Local root directory where the dataset is stored.

revision

Git revision (branch/tag/commit) of the dataset.

info

Dictionary containing dataset information (features, fps, paths, etc.).

stats

Aggregated statistics dictionary (mean, std, min, max, count).

episodes_stats

Per-episode statistics dictionary.

episodes

Dictionary mapping episode_index to episode information.

tasks

Dictionary mapping task_index to task descriptions.

task_to_task_index

Reverse mapping from task description to task_index.

advantages

Dictionary mapping (episode_index, timestamp) to advantage values.

Example

Load metadata from Hub:

>>> meta = LeRobotDatasetMetadata("lerobot/aloha_mobile_cabinet")
>>> print(f"Total episodes: {meta.total_episodes}")

Create new dataset metadata:

>>> meta = LeRobotDatasetMetadata.create(
...     repo_id="my-dataset",
...     fps=30,
...     features={"state": {"dtype": "float32", "shape": (7,)}}
... )

__init__(repo_id: str, root: str | Path | None = None, revision: str | None = None, force_cache_sync: bool = False)

add_task(task: str)

Given a task in natural language, add it to the dictionary of tasks.

property chunks_size: int

Max number of episodes per chunk.

property control_mode: str

Control-mode label from info.json. Defaults to ‘unknown’ when absent.

classmethod create(repo_id: str, fps: int, root: str | Path | None = None, robot_type: str | None = None, features: dict | None = None, use_videos: bool = True) → LeRobotDatasetMetadata

Creates metadata for a LeRobotDataset.

property data_path: str

Formattable string for the parquet files.

property fps: int

Frames per second used during data collection.

get_data_file_path(ep_index: int) → Path

Get the file path for a specific episode’s parquet data file.

Parameters:

ep_index – Episode index.

Returns:

Path to the parquet file for the episode.

get_episode_chunk(ep_index: int) → int

Get the chunk index for a given episode index.

Episodes are grouped into chunks for efficient storage.

Parameters:

ep_index – Episode index.

Returns:

Chunk index containing this episode.

get_task_index(task: str) → int | None

Given a task in natural language, returns its task_index if the task already exists in the dataset, otherwise return None.

get_video_file_path(ep_index: int, vid_key: str) → Path

Get the file path for a specific episode’s video file.

Parameters:

• ep_index – Episode index.

• vid_key – Video key/name (e.g., “camera0”).

Returns:

Path to the video file for the episode.

load_metadata() → None

Load dataset metadata from disk.

Loads info, tasks, episodes, statistics, and advantages from the dataset root directory. Handles version compatibility checks.

pull_from_repo(allow_patterns: list[str] | str | None = None, ignore_patterns: list[str] | str | None = None) → None

property robot_type: str | None

Robot type used in recording this dataset.

save_episode(episode_index: int, episode_length: int, episode_tasks: list[str], episode_stats: dict[str, dict]) → None

property total_chunks: int

Total number of chunks (groups of episodes).

property total_episodes: int

Total number of episodes available.

property total_frames: int

Total number of frames saved in this dataset.

property total_tasks: int

Total number of different tasks performed in this dataset.

update_video_info(skip_keys: set[str] | None = None) → None

Update video metadata from the first episode’s video files.

Warning: this function writes info from first episode videos, implicitly assuming that all videos have been encoded the same way. Also, this means it assumes the first episode exists.

Parameters:

skip_keys – Optional set of video keys to skip (e.g. deferred video keys whose files don’t exist yet).

property video_path: str | None

Formattable string for the video files.

class opentau.datasets.lerobot_dataset.VQADatasetMetadata(*, info: dict | None = None, stats: dict | None = None)

Bases: DatasetMetadata

Metadata class for vqa datasets (vision-language datasets).

Inherits fps -> None from DatasetMetadata since VQA samples have no temporal axis. BaseDataset._emit_optional_keys() sees this and emits fps_is_pad=True for VQA samples in a mixture, keeping every batch row schema-aligned with the LeRobot samples.

opentau.datasets.lerobot_dataset.aggregate_selected_stats(meta: LeRobotDatasetMetadata, episodes: list[int] | None = None, excluded_episodes: list[int] | None = None) → dict[str, dict[str, ndarray]]

Per-feature stats aggregated over the SELECTED episodes.

Falls back to meta.stats (the full-dataset aggregate) when no selection is active, or when the dataset predates per-episode stats (v2.0, which has nothing to recompute from). Mirrors the propagation done in LeRobotDataset.__init__ so off-line consumers (the FAST/BPE tokenizer fit, the norm-distribution diagnostic) normalize over the same episodes the policy trains on.

opentau.datasets.lerobot_dataset.resolve_selected_episodes(all_episodes: Iterable[int], episodes: list[int] | None, excluded_episodes: list[int] | None, repo_id: str = '') → list[int] | None

Resolve the effective episode list for a dataset selection.

The selection is (episodes if not None else all_episodes) minus excluded_episodes. excluded_episodes takes precedence: an index present in both is dropped. Returns a sorted list, or None when no selection is active (episodes is None and excluded_episodes is empty), meaning “use every episode”.

Raises:

ValueError – if the denylist removes every selected episode.

opentau.datasets.lerobot_dataset.retry_random_on_failure(f)

Decorator to retry dataset item retrieval with random indices on failure.

When a dataset item fails to load, this decorator will retry with random indices up to _total_rand_attempts times before raising an error.

Parameters:

f – The __getitem__ method to wrap.

Returns:

Wrapped function that retries on failure.

opentau.datasets.lerobot_dataset.suppress_control_mode_warning(repo_id: str) → None

Mark repo_id as already-warned so the missing-control_mode warning does not fire for it during subsequent LeRobotDataset.__init__ calls.

Intended for callers that supply an explicit control_mode override and therefore already know the on-disk metadata is missing the field. Must be invoked before the LeRobotDataset is constructed for repo_id; once __init__ has emitted the warning, further calls are a no-op.