opentau.datasets.utils

Utility functions for dataset management, I/O, and validation.

This module provides a comprehensive set of utility functions for working with LeRobot datasets, including file I/O operations, metadata management, data validation, version compatibility checking, and HuggingFace Hub integration.

The module is organized into several functional areas:

  • Dictionary manipulation: Flattening/unflattening nested dictionaries

  • File I/O: JSON and JSONL reading/writing with automatic directory creation

  • Metadata management: Loading and saving dataset info, statistics, episodes,

    tasks, and advantages

  • Data validation: Frame and episode buffer validation with detailed error

    messages

Key Features:
  • Automatic serialization: Converts tensors and arrays to JSON-compatible

    formats.

  • Comprehensive validation: Validates frames and episodes.

  • Path management: Standard paths for dataset structure (meta/, data/).

Constants:

DEFAULT_CHUNK_SIZE: Maximum number of episodes per chunk (1000). ADVANTAGES_PATH, INFO_PATH, EPISODES_PATH, STATS_PATH: Standard paths.

Classes:
IterableNamespace: Namespace object supporting both dictionary iteration

and dot notation access.

Functions:
Dictionary manipulation:

flatten_dict: Flatten nested dictionaries with separator-based keys. unflatten_dict: Expand flattened keys into nested dictionaries. serialize_dict: Convert tensors/arrays to JSON-serializable format.

File I/O:

load_json, write_json: JSON file operations. load_jsonlines, write_jsonlines: JSONL operations.

Data validation:

validate_frame: Validate frame data against feature specifications. validate_episode_buffer: Validate episode buffer before adding.

(Note: Truncated for brevity, apply the same flat indentation to the rest)

Example

Load dataset metadata:

>>> info = load_info(Path("my_dataset"))
>>> stats = load_stats(Path("my_dataset"))
>>> episodes = load_episodes(Path("my_dataset"))

Validate a frame:

>>> features = {"state": {"dtype": "float32", "shape": (7,)}}
>>> frame = {"state": np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7])}
>>> validate_frame(frame, features)

Functions

append_jsonlines(data, fpath)

Append a single dictionary to a JSON Lines (JSONL) file.

backward_compatible_episodes_stats(stats, ...)

Create episode-level statistics from global statistics for backward compatibility.

cast_stats_to_numpy(stats)

Convert statistics dictionary values to numpy arrays.

check_timestamps_sync(timestamps, ...[, ...])

This check is to make sure that each timestamp is separated from the next by (1/fps) +/- tolerance to account for possible numerical error.

check_version_compatibility(repo_id, ...[, ...])

Check compatibility between a dataset version and the current codebase version.

create_branch(repo_id, *, branch[, repo_type])

Create a branch on a existing Hugging Face repo.

create_empty_dataset_info(codebase_version, ...)

Create an empty dataset info dictionary with default values.

create_lerobot_dataset_card([tags, dataset_info])

Keyword arguments will be used to replace values in src/opentau/datasets/card_template.md.

cycle(iterable)

The equivalent of itertools.cycle, but safe for Pytorch dataloaders.

dataset_to_policy_features(features)

Convert dataset features to policy feature format.

embed_images(dataset)

Embed image bytes into the dataset table before saving to parquet.

flatten_dict(d[, parent_key, sep])

Flatten a nested dictionary structure by collapsing nested keys into one key with a separator.

get_delta_indices_soft(...)

Returns soft indices (not necessarily integer) for delta timestamps based on the provided information.

get_episode_data_index(episode_dicts[, episodes])

Compute data indices for episodes in a flattened dataset.

get_hf_features_from_features(features)

Convert dataset features dictionary to HuggingFace Features object.

get_nested_item(obj, flattened_key[, sep])

Get a nested item from a dictionary-like object using a flattened key.

get_repo_branches(repo_id)

Return the branch names of a dataset repo on the Hub.

get_repo_versions(repo_id)

Returns available valid versions (branches and tags) on given repo.

get_safe_version(repo_id, version[, ...])

Resolve version to a concrete revision available on the repo.

hf_transform_to_torch(items_dict)

Get a transform function that convert items from Hugging Face dataset (pyarrow) to torch tensors.

is_valid_version(version)

Check if a version string is valid and can be parsed.

load_advantages(local_dir)

Load advantage values from the advantages.json file.

load_episodes(local_dir)

Load episodes from the episodes.jsonl file.

load_episodes_and_stats_v30(local_dir)

Load v3.0 episodes and per-episode stats, parsing the metadata shards once.

load_episodes_stats(local_dir)

Load episode statistics from the episodes_stats.jsonl file.

load_episodes_stats_v30(local_dir[, ...])

Reconstruct per-episode stats from the flattened stats/* columns of the v3.0 episodes parquet.

load_episodes_v30(local_dir[, episodes_table])

Load episodes from the v3.0 meta/episodes/**/*.parquet shards.

load_image_as_numpy(fpath[, dtype, ...])

Load an image file as a numpy array.

load_info(local_dir)

Load dataset info from the standard info.json file.

load_json(fpath)

Load JSON data from a file.

load_jsonlines(fpath)

Load JSON Lines (JSONL) data from a file.

load_stats(local_dir)

Load dataset statistics from the standard stats.json file.

load_tasks(local_dir)

Load tasks from the tasks.jsonl file.

load_tasks_v30(local_dir)

Load tasks from the v3.0 meta/tasks.parquet file.

serialize_dict(stats)

Serialize a dictionary containing tensors and arrays to JSON-serializable format.

unflatten_dict(d[, sep])

Unflatten a dictionary by expanding keys with separators into nested dictionaries.

validate_episode_buffer(episode_buffer, ...)

Validate that an episode buffer is properly formatted.

validate_feature_dtype_and_shape(name, ...)

Validate that a feature value matches its expected dtype and shape.

validate_feature_image_or_video(name, ...)

Validate that an image or video feature matches expected shape.

validate_feature_numpy_array(name, ...)

Validate that a numpy array feature matches expected dtype and shape.

validate_feature_string(name, value)

Validate that a feature value is a string.

validate_features_presence(actual_features, ...)

Validate that required features are present and no unexpected features exist.

validate_frame(frame, features[, ...])

Validate that a frame dictionary matches the expected features.

write_episode(episode, local_dir)

Write an episode entry to the episodes.jsonl file.

write_episode_stats(episode_index, ...)

Write episode statistics to the episodes_stats.jsonl file.

write_info(info, local_dir)

Write dataset info dictionary to the standard info.json file.

write_json(data, fpath)

Write data to a JSON file.

write_jsonlines(data, fpath)

Write data to a JSON Lines (JSONL) file.

write_stats(stats, local_dir)

Write dataset statistics to the standard stats.json file.

write_task(task_index, task, local_dir)

Write a task entry to the tasks.jsonl file.

Classes

IterableNamespace([dictionary])

A namespace object that supports both dictionary-like iteration and dot notation access.
class opentau.datasets.utils.IterableNamespace(dictionary: dict[str, Any] | None = None, **kwargs)[source]

Bases: SimpleNamespace

A namespace object that supports both dictionary-like iteration and dot notation access. Automatically converts nested dictionaries into IterableNamespaces.

This class extends SimpleNamespace to provide: - Dictionary-style iteration over keys - Access to items via both dot notation (obj.key) and brackets (obj[“key”]) - Dictionary-like methods: items(), keys(), values() - Recursive conversion of nested dictionaries

Parameters:

  • dictionary – Optional dictionary to initialize the namespace

  • **kwargs – Additional keyword arguments passed to SimpleNamespace

Examples

>>> data = {"name": "Alice", "details": {"age": 25}}
>>> ns = IterableNamespace(data)
>>> ns.name
'Alice'
>>> ns.details.age
25
>>> list(ns.keys())
['name', 'details']
>>> for key, value in ns.items():
...     print(f"{key}: {value}")
name: Alice
details: IterableNamespace(age=25)
__init__(dictionary: dict[str, Any] | None = None, **kwargs)[source]
items()[source]
keys()[source]
values()[source]
opentau.datasets.utils.append_jsonlines(data: dict, fpath: Path) → None[source]

Append a single dictionary to a JSON Lines (JSONL) file.

Creates parent directories if they don’t exist. Appends the data as a new line to the existing file.

Parameters:

  • data – Dictionary to append as a new line.

  • fpath – Path to the JSONL file (will be created if it doesn’t exist).

opentau.datasets.utils.backward_compatible_episodes_stats(stats: dict[str, dict[str, ndarray]], episodes: Iterable[int]) → dict[int, dict[str, dict[str, ndarray]]][source]

Create episode-level statistics from global statistics for backward compatibility.

In older dataset versions, statistics were stored globally rather than per-episode. This function creates per-episode statistics by assigning the same global stats to each episode.

Parameters:

  • stats – Global statistics dictionary.

  • episodes – List of episode indices.

Returns:

Dictionary mapping episode_index to the same statistics dictionary.

opentau.datasets.utils.cast_stats_to_numpy(stats) → dict[str, dict[str, ndarray]][source]

Convert statistics dictionary values to numpy arrays.

Flattens the dictionary, converts all values to numpy arrays, then unflattens to restore the original structure.

Parameters:

stats – Dictionary with statistics (values may be lists or other types).

Returns:

Dictionary with the same structure but all values as numpy arrays.

opentau.datasets.utils.check_timestamps_sync(timestamps: ndarray, episode_indices: ndarray, episode_data_index: dict[str, ndarray], fps: int, tolerance_s: float, raise_value_error: bool = True) → bool[source]

This check is to make sure that each timestamp is separated from the next by (1/fps) +/- tolerance to account for possible numerical error.

Parameters:

  • timestamps (np.ndarray) – Array of timestamps in seconds.

  • episode_indices (np.ndarray) – Array indicating the episode index for each timestamp.

  • episode_data_index (dict[str, np.ndarray]) – A dictionary that includes ‘to’, which identifies indices for the end of each episode.

  • fps (int) – Frames per second. Used to check the expected difference between consecutive timestamps.

  • tolerance_s (float) – Allowed deviation from the expected (1/fps) difference.

  • raise_value_error (bool) – Whether to raise a ValueError if the check fails.

Returns:

True if all checked timestamp differences lie within tolerance, False otherwise.

Return type:

bool

Raises:

ValueError – If the check fails and raise_value_error is True.

opentau.datasets.utils.check_version_compatibility(repo_id: str, version_to_check: str | Version, current_version: str | Version, enforce_breaking_major: bool = True) → None[source]

Check compatibility between a dataset version and the current codebase version.

Parameters:

  • repo_id – Repository ID of the dataset.

  • version_to_check – Version of the dataset to check.

  • current_version – Current codebase version.

  • enforce_breaking_major – If True, raise error for major version mismatches. Defaults to True.

Raises:

BackwardCompatibilityError – If the dataset version is too old (major version mismatch).

opentau.datasets.utils.create_branch(repo_id, *, branch: str, repo_type: str | None = None) → None[source]

Create a branch on a existing Hugging Face repo. Delete the branch if it already exists before creating it.

opentau.datasets.utils.create_empty_dataset_info(codebase_version: str, fps: int, robot_type: str | None, features: dict, use_videos: bool) → dict[source]

Create an empty dataset info dictionary with default values.

Parameters:

  • codebase_version – Version of the codebase used to create the dataset.

  • fps – Frames per second used during data collection.

  • robot_type – Type of robot used (can be None).

  • features – Dictionary of feature specifications.

  • use_videos – Whether videos are used for visual modalities.

Returns:

Dictionary containing dataset metadata with initialized counters and paths.

opentau.datasets.utils.create_lerobot_dataset_card(tags: list | None = None, dataset_info: dict | None = None, **kwargs) → DatasetCard[source]

Keyword arguments will be used to replace values in src/opentau/datasets/card_template.md. Note: If specified, license must be one of https://huggingface.co/docs/hub/repositories-licenses.

opentau.datasets.utils.cycle(iterable)[source]

The equivalent of itertools.cycle, but safe for Pytorch dataloaders.

See https://github.com/pytorch/pytorch/issues/23900 for information on why itertools.cycle is not safe.

opentau.datasets.utils.dataset_to_policy_features(features: dict[str, dict]) → dict[str, PolicyFeature][source]

Convert dataset features to policy feature format.

Maps dataset features to policy feature types (VISUAL, ENV, STATE, ACTION) based on feature names and data types.

Parameters:

features – Dictionary mapping feature names to feature specifications.

Returns:

Dictionary mapping feature names to PolicyFeature objects.

Raises:

ValueError – If a visual feature doesn’t have 3 dimensions.

opentau.datasets.utils.embed_images(dataset: Dataset) → Dataset[source]

Embed image bytes into the dataset table before saving to parquet.

Converts the dataset to arrow format, embeds image storage, and restores the original format.

Parameters:

dataset – HuggingFace dataset containing images.

Returns:

Dataset with embedded image bytes, ready for parquet serialization.

opentau.datasets.utils.flatten_dict(d: dict, parent_key: str = '', sep: str = '/') → dict[source]

Flatten a nested dictionary structure by collapsing nested keys into one key with a separator.

For example:

>>> dct = {"a": {"b": 1, "c": {"d": 2}}, "e": 3}
>>> print(flatten_dict(dct))
{"a/b": 1, "a/c/d": 2, "e": 3}
opentau.datasets.utils.get_delta_indices_soft(delta_timestamps_info: tuple[dict[str, ndarray], dict[str, ndarray], dict[str, ndarray], dict[str, ndarray]], fps: int) → dict[str, ndarray][source]

Returns soft indices (not necessarily integer) for delta timestamps based on the provided information. Soft indices are computed by sampling from a normal distribution defined by the mean and standard deviation and clipping the values to the specified lower and upper bounds. Note: Soft indices can be converted to integer indices by either rounding or interpolation.

opentau.datasets.utils.get_episode_data_index(episode_dicts: dict[int, dict], episodes: list[int] | None = None) → tuple[dict[str, Tensor], dict[int, int]][source]

Compute data indices for episodes in a flattened dataset.

Calculates start and end indices for each episode in a concatenated dataset, and creates a mapping from episode index to position in the episodes list.

Parameters:

  • episode_dicts – Dictionary mapping episode_index to episode info dicts containing ‘length’ keys.

  • episodes – Optional list of episode indices to include. If None, uses all episodes from episode_dicts.

Returns:

  • episode_data_index: Dictionary with ‘from’ and ‘to’ tensors indicating start and end indices for each episode.

  • ep2idx: Dictionary mapping episode_index to position in the episodes list.

Return type:

Tuple of (episode_data_index, ep2idx)

opentau.datasets.utils.get_hf_features_from_features(features: dict) → Features[source]

Convert dataset features dictionary to HuggingFace Features object.

Maps feature types and shapes to appropriate HuggingFace feature types (Image, Value, Sequence, Array2D, Array3D, Array4D, Array5D).

Parameters:

features – Dictionary mapping feature names to feature specifications with ‘dtype’ and ‘shape’ keys.

Returns:

HuggingFace Features object compatible with the dataset library.

Raises:

ValueError – If a feature shape is not supported (more than 5 dimensions).

opentau.datasets.utils.get_nested_item(obj: DictLike, flattened_key: str, sep: str = '/') → Any[source]

Get a nested item from a dictionary-like object using a flattened key.

Parameters:

  • obj – Dictionary-like object to access.

  • flattened_key – Flattened key path (e.g., “a/b/c”).

  • sep – Separator used in the flattened key. Defaults to “/”.

Returns:

The value at the nested path specified by the flattened key.

Example

>>> dct = {"a": {"b": {"c": 42}}}
>>> get_nested_item(dct, "a/b/c")
42
opentau.datasets.utils.get_repo_branches(repo_id: str) → list[str][source]

Return the branch names of a dataset repo on the Hub.

Parameters:

repo_id – Repository ID of the dataset on the Hugging Face Hub.

Returns:

The names of every branch ref on the repo (e.g. ["main"]); tags are not included.

opentau.datasets.utils.get_repo_versions(repo_id: str) → list[Version][source]

Returns available valid versions (branches and tags) on given repo.

opentau.datasets.utils.get_safe_version(repo_id: str, version: str | Version, read_ceiling: str | Version = 'v3.0', allow_branch_fallback: bool = False) → str[source]

Resolve version to a concrete revision available on the repo.

Returns the requested version if the repo carries it, otherwise the latest compatible one. read_ceiling is the newest format the loader can read (default READ_CODEBASE_VERSION): when the requested version is unavailable but the repo only carries a newer-yet-still-readable format (e.g. a v3.0-only dataset while the default target is v2.1), the newest such version up to the ceiling is selected instead of raising ForwardCompatibilityError.

When the repo carries no codebase-version tag at all and allow_branch_fallback is set, the repo’s default branch (main, then master) is returned instead of raising RevisionNotFoundError — this lets untagged Hub datasets load when the caller did not pin a revision.

Parameters:

  • repo_id – Repository ID of the dataset on the Hugging Face Hub.

  • version – Desired codebase version (e.g. "v2.1").

  • read_ceiling – Newest format the loader can read. Defaults to READ_CODEBASE_VERSION.

  • allow_branch_fallback – When True and the repo carries no version tags, fall back to the main/master branch instead of raising. Set by callers only when no explicit revision was requested.

Returns:

either f"v{version}" for a resolved version, or a branch name ("main"/"master") on fallback.

Return type:

A revision usable as a Hub revision

Raises:

  • RevisionNotFoundError – The repo has no version tags and either allow_branch_fallback is False or it has no main/master branch.

  • BackwardCompatibilityError – Only older, incompatible major versions exist.

  • ForwardCompatibilityError – Only newer versions beyond the ceiling exist.

opentau.datasets.utils.hf_transform_to_torch(items_dict: dict[str, list])[source]

Get a transform function that convert items from Hugging Face dataset (pyarrow) to torch tensors. Importantly, images are converted from PIL, which corresponds to a channel last representation (h w c) of uint8 type, to a torch image representation with channel first (c h w) of float32 type in range [0,1].

Columns whose values cannot be converted to tensors (e.g. dict/struct annotation columns such as language_persistent) are not model inputs and are dropped from the returned dict rather than crashing the dataloader.

opentau.datasets.utils.is_valid_version(version: str) → bool[source]

Check if a version string is valid and can be parsed.

Parameters:

version – Version string to validate.

Returns:

True if the version string is valid, False otherwise.

opentau.datasets.utils.load_advantages(local_dir: Path) → dict | None[source]

Load advantage values from the advantages.json file.

Advantages are keyed by (episode_index, timestamp) tuples in the JSON file as comma-separated strings, which are converted to tuple keys.

Parameters:

local_dir – Root directory of the dataset containing meta/advantages.json.

Returns:

Dictionary mapping (episode_index, timestamp) tuples to advantage values, or None if the file doesn’t exist.

opentau.datasets.utils.load_episodes(local_dir: Path) → dict[int, dict][source]

Load episodes from the episodes.jsonl file.

Parameters:

local_dir – Root directory of the dataset containing meta/episodes.jsonl.

Returns:

Dictionary mapping episode_index to episode information dictionary.

opentau.datasets.utils.load_episodes_and_stats_v30(local_dir: Path) → tuple[dict[int, dict], dict[int, dict[str, dict[str, ndarray]]]][source]

Load v3.0 episodes and per-episode stats, parsing the metadata shards once.

The episode-metadata shards are read a single time (one Arrow table) and split into the load_episodes_v30() and load_episodes_stats_v30() results. The result is cached per local_dir (keyed by a content signature), so a mixture that lists the same repo under several configs parses its metadata once instead of once per config.

Parameters:

local_dir – Dataset root directory.

opentau.datasets.utils.load_episodes_stats(local_dir: Path) → dict[int, dict[str, dict[str, ndarray]]][source]

Load episode statistics from the episodes_stats.jsonl file.

Parameters:

local_dir – Root directory of the dataset containing meta/episodes_stats.jsonl.

Returns:

Dictionary mapping episode_index to statistics dictionary (with numpy arrays).

opentau.datasets.utils.load_episodes_stats_v30(local_dir: Path, episodes_table: Table | None = None) → dict[int, dict[str, dict[str, ndarray]]][source]

Reconstruct per-episode stats from the flattened stats/* columns of the v3.0 episodes parquet. Returns the same shape as load_episodes_stats() so aggregate_stats accepts v3.0 stats identically: count is (1,) int and image/video stats are (C, 1, 1).

Each stats/* column is materialized to numpy column-wise – flatten the Arrow list column once and reshape to (num_episodes, -1) – instead of a per-cell Python coercion, which is the dominant cost when loading large v3.0 episode tables. A ragged column (non-uniform per-episode length) falls back to the per-cell path.

Parameters:

  • local_dir – Dataset root directory.

  • episodes_table – Optional pre-read Arrow table so the shards are parsed once across episodes + stats (see load_episodes_and_stats_v30()).

opentau.datasets.utils.load_episodes_v30(local_dir: Path, episodes_table: Table | None = None) → dict[int, dict][source]

Load episodes from the v3.0 meta/episodes/**/*.parquet shards.

Emits the same {episode_index: {...}} shape as load_episodes(), so get_episode_data_index, the per-episode caches, and the file-path accessors are unchanged. Each value additionally carries the v3.0 file-mapping columns. The flattened stats/* columns are dropped here and loaded separately by load_episodes_stats_v30().

Columns are read column-wise (Arrow to_pylist per column, then zipped into per-episode rows) instead of DataFrame.to_dict(orient="records"), which is materially faster on large episode tables.

Parameters:

  • local_dir – Dataset root directory.

  • episodes_table – Optional pre-read Arrow table so the shards are parsed once across episodes + stats (see load_episodes_and_stats_v30()).

opentau.datasets.utils.load_image_as_numpy(fpath: str | ~pathlib.Path, dtype: ~numpy.dtype[~typing.Any] | None | type[~typing.Any] | ~numpy._typing._dtype_like._SupportsDType[~numpy.dtype[~typing.Any]] | str | tuple[~typing.Any, int] | tuple[~typing.Any, ~typing.SupportsIndex | ~collections.abc.Sequence[~typing.SupportsIndex]] | list[~typing.Any] | ~numpy._typing._dtype_like._DTypeDict | tuple[~typing.Any, ~typing.Any] = <class 'numpy.float32'>, channel_first: bool = True) → ndarray[source]

Load an image file as a numpy array.

Parameters:

  • fpath – Path to the image file.

  • dtype – Data type for the array. Defaults to np.float32.

  • channel_first – If True, return array in (C, H, W) format; otherwise (H, W, C). Defaults to True.

Returns:

Image as numpy array. If dtype is floating point, values are normalized to [0, 1]. Otherwise, values are in [0, 255].

opentau.datasets.utils.load_info(local_dir: Path) → dict[source]

Load dataset info from the standard info.json file.

Converts feature shapes from lists to tuples for consistency.

Parameters:

local_dir – Root directory of the dataset containing meta/info.json.

Returns:

Dataset info dictionary with feature shapes as tuples.

Raises:

ValueError – If meta/info.json is missing, empty, or invalid.

opentau.datasets.utils.load_json(fpath: Path) → Any[source]

Load JSON data from a file.

Parameters:

fpath – Path to the JSON file.

Returns:

Parsed JSON data (dict, list, or primitive type).

Raises:

  • ValueError – If the file is empty or contains only whitespace.

  • json.JSONDecodeError – If the file contains invalid JSON.

opentau.datasets.utils.load_jsonlines(fpath: Path) → list[Any][source]

Load JSON Lines (JSONL) data from a file.

Parameters:

fpath – Path to the JSONL file.

Returns:

List of dictionaries, one per line in the file.

opentau.datasets.utils.load_stats(local_dir: Path) → dict[str, dict[str, ndarray]] | None[source]

Load dataset statistics from the standard stats.json file.

Parameters:

local_dir – Root directory of the dataset containing meta/stats.json.

Returns:

Dictionary with statistics as numpy arrays, or None if the file doesn’t exist.

opentau.datasets.utils.load_tasks(local_dir: Path) → tuple[dict, dict][source]

Load tasks from the tasks.jsonl file.

Parameters:

local_dir – Root directory of the dataset containing meta/tasks.jsonl.

Returns:

  • tasks_dict: Dictionary mapping task_index to task description.

  • task_to_index_dict: Dictionary mapping task description to task_index.

Return type:

Tuple of (tasks_dict, task_to_index_dict)

opentau.datasets.utils.load_tasks_v30(local_dir: Path) → tuple[dict, dict][source]

Load tasks from the v3.0 meta/tasks.parquet file.

v3.0 replaces tasks.jsonl with a parquet indexed by the task string (named task) holding an integer task_index column. Returns the same (tasks, task_to_task_index) contract as load_tasks().

opentau.datasets.utils.serialize_dict(stats: dict[str, Tensor | ndarray | dict]) → dict[source]

Serialize a dictionary containing tensors and arrays to JSON-serializable format.

Converts torch.Tensor and np.ndarray to lists, and np.generic to Python scalars. The dictionary structure is preserved through flattening and unflattening.

Parameters:

stats – Dictionary containing statistics with tensor/array values.

Returns:

Dictionary with serialized (list/scalar) values in the same structure.

Raises:

NotImplementedError – If a value type is not supported for serialization.

opentau.datasets.utils.unflatten_dict(d: dict, sep: str = '/') → dict[source]

Unflatten a dictionary by expanding keys with separators into nested dictionaries.

Parameters:

  • d – Dictionary with flattened keys (e.g., {“a/b”: 1, “a/c/d”: 2}).

  • sep – Separator used to split keys. Defaults to “/”.

Returns:

{“b”: 1, “c”: {“d”: 2}}}).

Return type:

Nested dictionary structure (e.g., {“a”

Example

>>> dct = {"a/b": 1, "a/c/d": 2, "e": 3}
>>> print(unflatten_dict(dct))
{"a": {"b": 1, "c": {"d": 2}}, "e": 3}
opentau.datasets.utils.validate_episode_buffer(episode_buffer: dict, total_episodes: int, features: dict, deferred_features: set[str] | None = None) → None[source]

Validate that an episode buffer is properly formatted.

Checks that required keys exist, episode_index matches total_episodes, buffer is not empty, and all features are present.

Parameters:

  • episode_buffer – Dictionary containing episode data to validate.

  • total_episodes – Total number of episodes already in the dataset.

  • features – Dictionary of expected feature specifications.

  • deferred_features – Optional set of feature names whose data will be provided later and may be absent from the episode buffer.

Raises:

  • ValueError – If the buffer is missing required keys, is empty, or has mismatched features.

  • NotImplementedError – If episode_index doesn’t match total_episodes.

opentau.datasets.utils.validate_feature_dtype_and_shape(name: str, feature: dict, value: ndarray | Image | str) → str[source]

Validate that a feature value matches its expected dtype and shape.

Routes to appropriate validation function based on feature type.

Parameters:

  • name – Name of the feature being validated.

  • feature – Feature specification dictionary with ‘dtype’ and ‘shape’ keys.

  • value – Actual value to validate.

Returns:

Error message string (empty if validation passes).

Raises:

NotImplementedError – If the feature dtype is not supported.

opentau.datasets.utils.validate_feature_image_or_video(name: str, expected_shape: list[str], value: ndarray | Image) → str[source]

Validate that an image or video feature matches expected shape.

Supports both channel-first (C, H, W) and channel-last (H, W, C) formats.

Parameters:

  • name – Name of the feature being validated.

  • expected_shape – Expected shape as [C, H, W].

  • value – Actual image/video value (PIL Image or numpy array).

Returns:

Error message string (empty if validation passes).

Note

Pixel value range validation ([0,1] for float, [0,255] for uint8) is performed by the image writer threads, not here.

opentau.datasets.utils.validate_feature_numpy_array(name: str, expected_dtype: str, expected_shape: list[int], value: ndarray) → str[source]

Validate that a numpy array feature matches expected dtype and shape.

Parameters:

  • name – Name of the feature being validated.

  • expected_dtype – Expected numpy dtype as a string.

  • expected_shape – Expected shape as a list of integers.

  • value – Actual numpy array to validate.

Returns:

Error message string (empty if validation passes).

opentau.datasets.utils.validate_feature_string(name: str, value: str) → str[source]

Validate that a feature value is a string.

Parameters:

  • name – Name of the feature being validated.

  • value – Actual value to validate.

Returns:

Error message string (empty if validation passes).

opentau.datasets.utils.validate_features_presence(actual_features: set[str], expected_features: set[str], optional_features: set[str]) → str[source]

Validate that required features are present and no unexpected features exist.

Parameters:

  • actual_features – Set of feature names actually present.

  • expected_features – Set of feature names that must be present.

  • optional_features – Set of feature names that may be present but aren’t required.

Returns:

Error message string (empty if validation passes).

opentau.datasets.utils.validate_frame(frame: dict, features: dict, deferred_features: set[str] | None = None) → None[source]

Validate that a frame dictionary matches the expected features.

Checks that all required features are present, no unexpected features exist, and that feature types and shapes match the specification.

Parameters:

  • frame – Dictionary containing frame data to validate.

  • features – Dictionary of expected feature specifications.

  • deferred_features – Optional set of feature names whose data will be provided later (e.g. video observations attached after episode recording). These features are treated as optional during validation.

Raises:

ValueError – If the frame doesn’t match the feature specifications.

opentau.datasets.utils.write_episode(episode: dict, local_dir: Path) → None[source]

Write an episode entry to the episodes.jsonl file.

Parameters:

  • episode – Episode dictionary containing episode_index, tasks, length, etc.

  • local_dir – Root directory of the dataset where meta/episodes.jsonl will be written.

opentau.datasets.utils.write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path) → None[source]

Write episode statistics to the episodes_stats.jsonl file.

Serializes tensors and arrays in the stats before writing.

Parameters:

  • episode_index – Index of the episode.

  • episode_stats – Dictionary containing statistics for the episode (may contain tensors/arrays).

  • local_dir – Root directory of the dataset where meta/episodes_stats.jsonl will be written.

opentau.datasets.utils.write_info(info: dict, local_dir: Path) → None[source]

Write dataset info dictionary to the standard info.json file.

Parameters:

  • info – Dataset info dictionary to write.

  • local_dir – Root directory of the dataset where meta/info.json will be written.

opentau.datasets.utils.write_json(data: dict, fpath: Path) → None[source]

Write data to a JSON file.

Creates parent directories if they don’t exist. Uses 4-space indentation and allows non-ASCII characters.

Parameters:

  • data – Dictionary or other JSON-serializable data to write.

  • fpath – Path where the JSON file will be written.

opentau.datasets.utils.write_jsonlines(data: dict, fpath: Path) → None[source]

Write data to a JSON Lines (JSONL) file.

Creates parent directories if they don’t exist. Writes each item in the data iterable as a separate line.

Parameters:

  • data – Iterable of dictionaries to write (one per line).

  • fpath – Path where the JSONL file will be written.

opentau.datasets.utils.write_stats(stats: dict, local_dir: Path) → None[source]

Write dataset statistics to the standard stats.json file.

Serializes tensors and arrays to JSON-compatible format before writing.

Parameters:

  • stats – Dictionary containing dataset statistics (may contain tensors/arrays).

  • local_dir – Root directory of the dataset where meta/stats.json will be written.

opentau.datasets.utils.write_task(task_index: int, task: dict, local_dir: Path) → None[source]

Write a task entry to the tasks.jsonl file.

Parameters:

  • task_index – Integer index of the task.

  • task – Task description dictionary.

  • local_dir – Root directory of the dataset where meta/tasks.jsonl will be written.