ydy0615/lerobot-clone
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-06-03 12:21:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore(docs): update doctrines pipeline files (#1872)

Browse Source 
* docs(processor): update docstrings batch_processor

* docs(processor): update docstrings device_processor

* docs(processor): update docstrings tokenizer_processor

* update docstrings processor_act

* update docstrings for pipeline_features

* update docstrings for utils

* update docstring for processor_diffusion

* update docstrings factory

* add docstrings to pi0 processor

* add docstring to pi0fast processor

* add docstring classifier processor

* add docstring to sac processor

* add docstring smolvla processor

* add docstring to tdmpc processor

* add docstring to vqbet processor

* add docstrings to converters

* add docstrings for delta_action_processor

* add docstring to gym action processor

* update hil processor

* add docstring to joint obs processor

* add docstring to migrate_normalize_processor

* update docstrings normalize processor

* update docstring normalize processor

* update docstrings observation processor

* update docstrings rename_processor

* add docstrings robot_kinematic_processor

* cleanup rl comments

* add docstring to train.py

* add docstring to teleoperate.py

* add docstrings to phone_processor.py

* add docstrings to teleop_phone.py

* add docstrings to control_utils.py

* add docstrings to visualization_utils.py

---------

Co-authored-by: Pepijn <pepijn@huggingface.co>
This commit is contained in:
Steven Palma
2025-09-08 18:44:15 +02:00
committed by GitHub
parent d32006440c
commit af9ddcf9a2
33 changed files with 2325 additions and 519 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
src/lerobot/datasets/pipeline_features.py
View File

@@ -27,14 +27,30 @@ def aggregate_pipeline_dataset_features(
use_videos: bool = True,
patterns: Sequence[str] | None = None,
) -> dict[str, dict]:
"""
Aggregates the pipeline's features and returns a features dict ready for the dataset,
filtered to only those keys matching any of the given patterns (for action/state only).
"""Aggregates and filters dataset features based on a data processing pipeline.
- `initial_features`: raw camera specs, e.g. {"front": (h,w,c), ...}
- `use_videos`: whether to treat image features as video streams
- `patterns`: regexes to filter action & state features; images are included
whenever use_videos=True, regardless of patterns.
This function determines the final structure of dataset features after applying a series
of processing steps defined in a pipeline. It starts with an initial set of hardware
features (e.g., camera image shapes), transforms them using the pipeline, and then
filters the results.
Image features are controlled by the `use_videos` flag, while action and state features
can be selectively included by matching their keys against the provided regex `patterns`.
The final output is formatted to be compatible with Hugging Face Datasets feature dictionaries.
Args:
pipeline (DataProcessorPipeline): The data processing pipeline that defines all
feature transformations.
initial_features (dict[str, Any]): A dictionary of initial hardware features, where
keys are feature names and values are their shapes or types (e.g., camera resolutions).
use_videos (bool): If `True`, includes image/video features in the output. Defaults to `True`.
patterns (Sequence[str] | None): An optional sequence of regular expression patterns.
Only action and state keys that match at least one pattern will be included. If `None`,
all action and state keys are kept. Defaults to `None`.
Returns:
dict[str, dict]: A dictionary representing the final dataset features, structured for
use with `datasets.Features`.
"""
import re

584
src/lerobot/datasets/utils.py
View File

@@ -75,13 +75,20 @@ DEFAULT_FEATURES = {
def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict:
"""Flatten a nested dictionary structure by collapsing nested keys into one key with a separator.
"""Flatten a nested dictionary by joining keys 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}
Example:
>>> dct = {"a": {"b": 1, "c": {"d": 2}}, "e": 3}
>>> print(flatten_dict(dct))
{'a/b': 1, 'a/c/d': 2, 'e': 3}
Args:
d (dict): The dictionary to flatten.
parent_key (str): The base key to prepend to the keys in this level.
sep (str): The separator to use between keys.
Returns:
dict: A flattened dictionary.
"""
items = []
for k, v in d.items():
@@ -94,6 +101,20 @@ def flatten_dict(d: dict, parent_key: str = "", sep: str = "/") -> dict:
def unflatten_dict(d: dict, sep: str = "/") -> dict:
"""Unflatten a dictionary with delimited keys into a nested dictionary.
Example:
>>> flat_dct = {"a/b": 1, "a/c/d": 2, "e": 3}
>>> print(unflatten_dict(flat_dct))
{'a': {'b': 1, 'c': {'d': 2}}, 'e': 3}
Args:
d (dict): A dictionary with flattened keys.
sep (str): The separator used in the keys.
Returns:
dict: A nested dictionary.
"""
outdict = {}
for key, value in d.items():
parts = key.split(sep)
@@ -107,6 +128,16 @@ def unflatten_dict(d: dict, sep: str = "/") -> dict:
def get_nested_item(obj: DictLike, flattened_key: str, sep: str = "/") -> Any:
"""Access an item in a nested dictionary using a flattened key.
Args:
obj (DictLike): The nested dictionary-like object.
flattened_key (str): A key with parts separated by `sep`.
sep (str): The separator used in the flattened key.
Returns:
Any: The value from the nested dictionary.
"""
split_keys = flattened_key.split(sep)
getter = obj[split_keys[0]]
if len(split_keys) == 1:
@@ -119,6 +150,19 @@ def get_nested_item(obj: DictLike, flattened_key: str, sep: str = "/") -> Any:
def serialize_dict(stats: dict[str, torch.Tensor | np.ndarray | dict]) -> dict:
"""Serialize a dictionary containing tensors or numpy arrays to be JSON-compatible.
Converts torch.Tensor, np.ndarray, and np.generic types to lists or native Python types.
Args:
stats (dict): A dictionary that may contain non-serializable numeric types.
Returns:
dict: A dictionary with all values converted to JSON-serializable types.
Raises:
NotImplementedError: If a value has an unsupported type.
"""
serialized_dict = {}
for key, value in flatten_dict(stats).items():
if isinstance(value, (torch.Tensor, np.ndarray)):
@@ -133,6 +177,17 @@ def serialize_dict(stats: dict[str, torch.Tensor | np.ndarray | dict]) -> dict:
def embed_images(dataset: datasets.Dataset) -> datasets.Dataset:
"""Embed image bytes into the dataset table before saving to Parquet.
This function prepares a Hugging Face dataset for serialization by converting
image objects into an embedded format that can be stored in Arrow/Parquet.
Args:
dataset (datasets.Dataset): The input dataset, possibly containing image features.
Returns:
datasets.Dataset: The dataset with images embedded in the table storage.
"""
# Embed image bytes into the table before saving to parquet
format = dataset.format
dataset = dataset.with_format("arrow")
@@ -142,38 +197,94 @@ def embed_images(dataset: datasets.Dataset) -> datasets.Dataset:
def load_json(fpath: Path) -> Any:
"""Load data from a JSON file.
Args:
fpath (Path): Path to the JSON file.
Returns:
Any: The data loaded from the JSON file.
"""
with open(fpath) as f:
return json.load(f)
def write_json(data: dict, fpath: Path) -> None:
"""Write data to a JSON file.
Creates parent directories if they don't exist.
Args:
data (dict): The dictionary to write.
fpath (Path): The path to the output JSON file.
"""
fpath.parent.mkdir(exist_ok=True, parents=True)
with open(fpath, "w") as f:
json.dump(data, f, indent=4, ensure_ascii=False)
def load_jsonlines(fpath: Path) -> list[Any]:
"""Load data from a JSON Lines file.
Args:
fpath (Path): Path to the JSON Lines file.
Returns:
list[Any]: A list of objects loaded from the file.
"""
with jsonlines.open(fpath, "r") as reader:
return list(reader)
def write_jsonlines(data: dict, fpath: Path) -> None:
"""Write a list of dictionaries to a JSON Lines file.
Creates parent directories if they don't exist.
Args:
data (dict): The list of dictionaries to write.
fpath (Path): The path to the output JSON Lines file.
"""
fpath.parent.mkdir(exist_ok=True, parents=True)
with jsonlines.open(fpath, "w") as writer:
writer.write_all(data)
def append_jsonlines(data: dict, fpath: Path) -> None:
"""Append a dictionary to a JSON Lines file.
Creates parent directories if they don't exist.
Args:
data (dict): The dictionary to append.
fpath (Path): The path to the JSON Lines file.
"""
fpath.parent.mkdir(exist_ok=True, parents=True)
with jsonlines.open(fpath, "a") as writer:
writer.write(data)
def write_info(info: dict, local_dir: Path):
"""Write dataset info metadata to its standard file path.
Args:
info (dict): The dataset information dictionary.
local_dir (Path): The root directory of the dataset.
"""
write_json(info, local_dir / INFO_PATH)
def load_info(local_dir: Path) -> dict:
"""Load dataset info metadata from its standard file path.
Also converts shape lists to tuples for consistency.
Args:
local_dir (Path): The root directory of the dataset.
Returns:
dict: The dataset information dictionary.
"""
info = load_json(local_dir / INFO_PATH)
for ft in info["features"].values():
ft["shape"] = tuple(ft["shape"])
@@ -181,16 +292,40 @@ def load_info(local_dir: Path) -> dict:
def write_stats(stats: dict, local_dir: Path):
"""Serialize and write dataset statistics to their standard file path.
Args:
stats (dict): The statistics dictionary (can contain tensors/numpy arrays).
local_dir (Path): The root directory of the dataset.
"""
serialized_stats = serialize_dict(stats)
write_json(serialized_stats, local_dir / STATS_PATH)
def cast_stats_to_numpy(stats) -> dict[str, dict[str, np.ndarray]]:
"""Recursively cast numerical values in a stats dictionary to numpy arrays.
Args:
stats (dict): The statistics dictionary.
Returns:
dict: The statistics dictionary with values cast to numpy arrays.
"""
stats = {key: np.array(value) for key, value in flatten_dict(stats).items()}
return unflatten_dict(stats)
def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]]:
"""Load dataset statistics and cast numerical values to numpy arrays.
Returns None if the stats file doesn't exist.
Args:
local_dir (Path): The root directory of the dataset.
Returns:
A dictionary of statistics or None if the file is not found.
"""
if not (local_dir / STATS_PATH).exists():
return None
stats = load_json(local_dir / STATS_PATH)
@@ -198,6 +333,13 @@ def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]]:
def write_task(task_index: int, task: dict, local_dir: Path):
"""Write a single task to the tasks metadata file.
Args:
task_index (int): The index of the task.
task (dict): The task description dictionary.
local_dir (Path): The root directory of the dataset.
"""
task_dict = {
"task_index": task_index,
"task": task,
@@ -206,6 +348,16 @@ def write_task(task_index: int, task: dict, local_dir: Path):
def load_tasks(local_dir: Path) -> tuple[dict, dict]:
"""Load tasks from the tasks metadata file.
Args:
local_dir (Path): The root directory of the dataset.
Returns:
A tuple containing:
- A dictionary mapping task index to task description.
- A dictionary mapping task description to task index.
"""
tasks = load_jsonlines(local_dir / TASKS_PATH)
tasks = {item["task_index"]: item["task"] for item in sorted(tasks, key=lambda x: x["task_index"])}
task_to_task_index = {task: task_index for task_index, task in tasks.items()}
@@ -213,15 +365,36 @@ def load_tasks(local_dir: Path) -> tuple[dict, dict]:
def write_episode(episode: dict, local_dir: Path):
"""Write a single episode's metadata to the episodes metadata file.
Args:
episode (dict): The episode metadata dictionary.
local_dir (Path): The root directory of the dataset.
"""
append_jsonlines(episode, local_dir / EPISODES_PATH)
def load_episodes(local_dir: Path) -> dict:
"""Load episode metadata from the episodes metadata file.
Args:
local_dir (Path): The root directory of the dataset.
Returns:
dict: A dictionary mapping episode index to episode metadata.
"""
episodes = load_jsonlines(local_dir / EPISODES_PATH)
return {item["episode_index"]: item for item in sorted(episodes, key=lambda x: x["episode_index"])}
def write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path):
"""Write statistics for a single episode to the episode stats file.
Args:
episode_index (int): The index of the episode.
episode_stats (dict): The statistics for the episode.
local_dir (Path): The root directory of the dataset.
"""
# We wrap episode_stats in a dictionary since `episode_stats["episode_index"]`
# is a dictionary of stats and not an integer.
episode_stats = {"episode_index": episode_index, "stats": serialize_dict(episode_stats)}
@@ -229,6 +402,14 @@ def write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path
def load_episodes_stats(local_dir: Path) -> dict:
"""Load per-episode statistics from the episode stats file.
Args:
local_dir (Path): The root directory of the dataset.
Returns:
dict: A dictionary mapping episode index to its statistics dictionary.
"""
episodes_stats = load_jsonlines(local_dir / EPISODES_STATS_PATH)
return {
item["episode_index"]: cast_stats_to_numpy(item["stats"])
@@ -239,12 +420,35 @@ def load_episodes_stats(local_dir: Path) -> dict:
def backward_compatible_episodes_stats(
stats: dict[str, dict[str, np.ndarray]], episodes: list[int]
) -> dict[str, dict[str, np.ndarray]]:
"""Create a per-episode stats dictionary from a global stats dictionary.
This is used for backward compatibility with older datasets that only had global stats.
Args:
stats (dict): The global dataset statistics.
episodes (list[int]): A list of episode indices.
Returns:
dict: A dictionary mapping each episode index to the global stats.
"""
return dict.fromkeys(episodes, stats)
def load_image_as_numpy(
fpath: str | Path, dtype: np.dtype = np.float32, channel_first: bool = True
) -> np.ndarray:
"""Load an image from a file into a numpy array.
Args:
fpath (str | Path): Path to the image file.
dtype (np.dtype): The desired data type of the output array. If floating,
pixels are scaled to [0, 1].
channel_first (bool): If True, converts the image to (C, H, W) format.
Otherwise, it remains in (H, W, C) format.
Returns:
np.ndarray: The image as a numpy array.
"""
img = PILImage.open(fpath).convert("RGB")
img_array = np.array(img, dtype=dtype)
if channel_first: # (H, W, C) -> (C, H, W)
@@ -255,10 +459,19 @@ def load_image_as_numpy(
def hf_transform_to_torch(items_dict: dict[torch.Tensor | None]):
"""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].
"""Convert a batch from a Hugging Face dataset to torch tensors.
This transform function converts items from Hugging Face dataset format (pyarrow)
to torch tensors. Importantly, images are converted from PIL objects (H, W, C, uint8)
to a torch image representation (C, H, W, float32) in the range [0, 1]. Other
types are converted to torch.tensor.
Args:
items_dict (dict): A dictionary representing a batch of data from a
Hugging Face dataset.
Returns:
dict: The batch with items converted to torch tensors.
"""
for key in items_dict:
first_item = items_dict[key][0]
@@ -273,6 +486,14 @@ def hf_transform_to_torch(items_dict: dict[torch.Tensor | None]):
def is_valid_version(version: str) -> bool:
"""Check if a string is a valid PEP 440 version.
Args:
version (str): The version string to check.
Returns:
bool: True if the version string is valid, False otherwise.
"""
try:
packaging.version.parse(version)
return True
@@ -286,6 +507,18 @@ def check_version_compatibility(
current_version: str | packaging.version.Version,
enforce_breaking_major: bool = True,
) -> None:
"""Check for version compatibility between a dataset and the current codebase.
Args:
repo_id (str): The repository ID for logging purposes.
version_to_check (str | packaging.version.Version): The version of the dataset.
current_version (str | packaging.version.Version): The current version of the codebase.
enforce_breaking_major (bool): If True, raise an error on major version mismatch.
Raises:
BackwardCompatibilityError: If the dataset version is from a newer, incompatible
major version of the codebase.
"""
v_check = (
packaging.version.parse(version_to_check)
if not isinstance(version_to_check, packaging.version.Version)
@@ -303,7 +536,14 @@ def check_version_compatibility(
def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
"""Returns available valid versions (branches and tags) on given repo."""
"""Return available valid versions (branches and tags) on a given Hub repo.
Args:
repo_id (str): The repository ID on the Hugging Face Hub.
Returns:
list[packaging.version.Version]: A list of valid versions found.
"""
api = HfApi()
repo_refs = api.list_repo_refs(repo_id, repo_type="dataset")
repo_refs = [b.name for b in repo_refs.branches + repo_refs.tags]
@@ -316,9 +556,22 @@ def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> str:
"""
Returns the version if available on repo or the latest compatible one.
Otherwise, will throw a `CompatibilityError`.
"""Return the specified version if available on repo, or the latest compatible one.
If the exact version is not found, it looks for the latest version with the
same major version number that is less than or equal to the target minor version.
Args:
repo_id (str): The repository ID on the Hugging Face Hub.
version (str | packaging.version.Version): The target version.
Returns:
str: The safe version string (e.g., "v1.2.3") to use as a revision.
Raises:
RevisionNotFoundError: If the repo has no version tags.
BackwardCompatibilityError: If only older major versions are available.
ForwardCompatibilityError: If only newer major versions are available.
"""
target_version = (
packaging.version.parse(version) if not isinstance(version, packaging.version.Version) else version
@@ -360,6 +613,17 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) ->
def get_hf_features_from_features(features: dict) -> datasets.Features:
"""Convert a LeRobot features dictionary to a `datasets.Features` object.
Args:
features (dict): A LeRobot-style feature dictionary.
Returns:
datasets.Features: The corresponding Hugging Face `datasets.Features` object.
Raises:
ValueError: If a feature has an unsupported shape.
"""
hf_features = {}
for key, ft in features.items():
if ft["dtype"] == "video":
@@ -387,6 +651,14 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
def _validate_feature_names(features: dict[str, dict]) -> None:
"""Validate that feature names do not contain invalid characters.
Args:
features (dict): The LeRobot features dictionary.
Raises:
ValueError: If any feature name contains '/'.
"""
invalid_features = {name: ft for name, ft in features.items() if "/" in name}
if invalid_features:
raise ValueError(f"Feature names should not contain '/'. Found '/' in '{invalid_features}'.")
@@ -395,6 +667,22 @@ def _validate_feature_names(features: dict[str, dict]) -> None:
def hw_to_dataset_features(
hw_features: dict[str, type | tuple], prefix: str, use_video: bool = True
) -> dict[str, dict]:
"""Convert hardware-specific features to a LeRobot dataset feature dictionary.
This function takes a dictionary describing hardware outputs (like joint states
or camera image shapes) and formats it into the standard LeRobot feature
specification.
Args:
hw_features (dict): Dictionary mapping feature names to their type (float for
joints) or shape (tuple for images).
prefix (str): The prefix to add to the feature keys (e.g., "observation"
or "action").
use_video (bool): If True, image features are marked as "video", otherwise "image".
Returns:
dict: A LeRobot features dictionary.
"""
features = {}
joint_fts = {key: ftype for key, ftype in hw_features.items() if ftype is float}
cam_fts = {key: shape for key, shape in hw_features.items() if isinstance(shape, tuple)}
@@ -427,6 +715,20 @@ def hw_to_dataset_features(
def build_dataset_frame(
ds_features: dict[str, dict], values: dict[str, Any], prefix: str
) -> dict[str, np.ndarray]:
"""Construct a single data frame from raw values based on dataset features.
A "frame" is a dictionary containing all the data for a single timestep,
formatted as numpy arrays according to the feature specification.
Args:
ds_features (dict): The LeRobot dataset features dictionary.
values (dict): A dictionary of raw values from the hardware/environment.
prefix (str): The prefix to filter features by (e.g., "observation"
or "action").
Returns:
dict: A dictionary representing a single frame of data.
"""
frame = {}
for key, ft in ds_features.items():
if key in DEFAULT_FEATURES or not key.startswith(prefix):
@@ -440,6 +742,21 @@ def build_dataset_frame(
def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFeature]:
"""Convert dataset features to policy features.
This function transforms the dataset's feature specification into a format
that a policy can use, classifying features by type (e.g., visual, state,
action) and ensuring correct shapes (e.g., channel-first for images).
Args:
features (dict): The LeRobot dataset features dictionary.
Returns:
dict: A dictionary mapping feature keys to `PolicyFeature` objects.
Raises:
ValueError: If an image feature does not have a 3D shape.
"""
# TODO(aliberts): Implement "type" in dataset features and simplify this
policy_features = {}
for key, ft in features.items():
@@ -471,11 +788,19 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
def combine_feature_dicts(*dicts: dict) -> dict:
"""
Merge LeRobot grouped feature dicts.
"""Merge LeRobot grouped feature dicts.
- For 1D numeric specs (dtype not image/video/string) with "names": we merge the names and recompute the shape.
- For others (observation.images.*), last one wins (if they are identical).
- For others (e.g. `observation.images.*`), the last one wins (if they are identical).
Args:
*dicts: A variable number of LeRobot feature dictionaries to merge.
Returns:
dict: A single merged feature dictionary.
Raises:
ValueError: If there's a dtype mismatch for a feature being merged.
"""
out: dict = {}
for d in dicts:
@@ -521,6 +846,18 @@ def create_empty_dataset_info(
use_videos: bool,
robot_type: str | None = None,
) -> dict:
"""Create a template dictionary for a new dataset's `info.json`.
Args:
codebase_version (str): The version of the LeRobot codebase.
fps (int): The frames per second of the data.
features (dict): The LeRobot features dictionary for the dataset.
use_videos (bool): Whether the dataset will store videos.
robot_type (str | None): The type of robot used, if any.
Returns:
dict: A dictionary with the initial dataset metadata.
"""
return {
"codebase_version": codebase_version,
"robot_type": robot_type,
@@ -541,6 +878,18 @@ def create_empty_dataset_info(
def get_episode_data_index(
episode_dicts: dict[dict], episodes: list[int] | None = None
) -> dict[str, torch.Tensor]:
"""Calculate the start and end indices for each episode in a flattened dataset.
Args:
episode_dicts (dict): A dictionary mapping episode index to episode metadata,
which must contain a "length" key.
episodes (list[int] | None): An optional list of episode indices to consider.
If None, all episodes are used.
Returns:
dict: A dictionary with "from" and "to" keys, containing torch tensors
with the start and end indices for each episode.
"""
episode_lengths = {ep_idx: ep_dict["length"] for ep_idx, ep_dict in episode_dicts.items()}
if episodes is not None:
episode_lengths = {ep_idx: episode_lengths[ep_idx] for ep_idx in episodes}
@@ -560,16 +909,19 @@ def check_timestamps_sync(
tolerance_s: float,
raise_value_error: bool = True,
) -> bool:
"""
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 if timestamps are separated by (1/fps) +/- tolerance.
This check ensures that consecutive timestamps within an episode are spaced
correctly, accounting for possible numerical errors. It ignores the boundaries
between episodes.
Args:
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',
episode_data_index (dict): 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.
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.
@@ -577,7 +929,8 @@ def check_timestamps_sync(
bool: True if all checked timestamp differences lie within tolerance, False otherwise.
Raises:
ValueError: If the check fails and `raise_value_error` is True.
ValueError: If `timestamps` and `episode_indices` shapes do not match, or if
the check fails and `raise_value_error` is True.
"""
if timestamps.shape != episode_indices.shape:
raise ValueError(
@@ -628,9 +981,23 @@ def check_timestamps_sync(
def check_delta_timestamps(
delta_timestamps: dict[str, list[float]], fps: int, tolerance_s: float, raise_value_error: bool = True
) -> bool:
"""This will check if all the values in delta_timestamps are multiples of 1/fps +/- tolerance.
This is to ensure that these delta_timestamps added to any timestamp from a dataset will themselves be
actual timestamps from the dataset.
"""Check if delta timestamps are multiples of 1/fps +/- tolerance.
This ensures that adding these delta timestamps to any existing timestamp in
the dataset will result in a value that aligns with the dataset's frame rate.
Args:
delta_timestamps (dict): A dictionary where values are lists of time
deltas in seconds.
fps (int): The frames per second of the dataset.
tolerance_s (float): The allowed tolerance in seconds.
raise_value_error (bool): If True, raises an error on failure.
Returns:
bool: True if all deltas are valid, False otherwise.
Raises:
ValueError: If any delta is outside the tolerance and `raise_value_error` is True.
"""
outside_tolerance = {}
for key, delta_ts in delta_timestamps.items():
@@ -656,6 +1023,15 @@ def check_delta_timestamps(
def get_delta_indices(delta_timestamps: dict[str, list[float]], fps: int) -> dict[str, list[int]]:
"""Convert delta timestamps in seconds to delta indices in frames.
Args:
delta_timestamps (dict): A dictionary of time deltas in seconds.
fps (int): The frames per second of the dataset.
Returns:
dict: A dictionary of frame delta indices.
"""
delta_indices = {}
for key, delta_ts in delta_timestamps.items():
delta_indices[key] = [round(d * fps) for d in delta_ts]
@@ -664,9 +1040,17 @@ def get_delta_indices(delta_timestamps: dict[str, list[float]], fps: int) -> dic
def cycle(iterable):
"""The equivalent of itertools.cycle, but safe for Pytorch dataloaders.
"""Create a dataloader-safe cyclical iterator.
See https://github.com/pytorch/pytorch/issues/23900 for information on why itertools.cycle is not safe.
This is an equivalent of `itertools.cycle` but is safe for use with
PyTorch DataLoaders with multiple workers.
See https://github.com/pytorch/pytorch/issues/23900 for details.
Args:
iterable: The iterable to cycle over.
Yields:
Items from the iterable, restarting from the beginning when exhausted.
"""
iterator = iter(iterable)
while True:
@@ -677,8 +1061,14 @@ def cycle(iterable):
def create_branch(repo_id, *, branch: str, repo_type: str | None = None) -> None:
"""Create a branch on a existing Hugging Face repo. Delete the branch if it already
exists before creating it.
"""Create a branch on an existing Hugging Face repo.
Deletes the branch if it already exists before creating it.
Args:
repo_id (str): The ID of the repository.
branch (str): The name of the branch to create.
repo_type (str | None): The type of the repository (e.g., "dataset").
"""
api = HfApi()
@@ -696,9 +1086,20 @@ def create_lerobot_dataset_card(
dataset_info: dict | None = None,
**kwargs,
) -> DatasetCard:
"""
Keyword arguments will be used to replace values in src/lerobot/datasets/card_template.md.
Note: If specified, license must be one of https://huggingface.co/docs/hub/repositories-licenses.
"""Create a `DatasetCard` for a LeRobot dataset.
Keyword arguments are used to replace values in the card template.
Note: If specified, `license` must be a valid license identifier from
https://huggingface.co/docs/hub/repositories-licenses.
Args:
tags (list | None): A list of tags to add to the dataset card.
dataset_info (dict | None): The dataset's info dictionary, which will
be displayed on the card.
**kwargs: Additional keyword arguments to populate the card template.
Returns:
DatasetCard: The generated dataset card object.
"""
card_tags = ["LeRobot"]
@@ -730,19 +1131,16 @@ def create_lerobot_dataset_card(
class IterableNamespace(SimpleNamespace):
"""
A namespace object that supports both dictionary-like iteration and dot notation access.
Automatically converts nested dictionaries into IterableNamespaces.
"""A namespace object that supports both dictionary-like iteration and dot notation.
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
This class extends `SimpleNamespace` to provide dictionary-style iteration,
access to items via brackets (`obj["key"]`), and dictionary-like methods
(`items()`, `keys()`, `values()`). Nested dictionaries are recursively
converted to `IterableNamespace` objects.
Args:
dictionary: Optional dictionary to initialize the namespace
**kwargs: Additional keyword arguments passed to SimpleNamespace
dictionary (dict, optional): A dictionary to initialize the namespace with.
**kwargs: Additional keyword arguments to initialize the namespace.
Examples:
>>> data = {"name": "Alice", "details": {"age": 25}}
@@ -756,10 +1154,16 @@ class IterableNamespace(SimpleNamespace):
>>> for key, value in ns.items():
... print(f"{key}: {value}")
name: Alice
details: IterableNamespace(age=25)
details: <__main__.IterableNamespace object at ...>
"""
def __init__(self, dictionary: dict[str, Any] = None, **kwargs):
"""Initialize the IterableNamespace.
Args:
dictionary (dict, optional): Dictionary to populate the namespace.
**kwargs: Keyword arguments to populate the namespace.
"""
super().__init__(**kwargs)
if dictionary is not None:
for key, value in dictionary.items():
@@ -769,22 +1173,46 @@ class IterableNamespace(SimpleNamespace):
setattr(self, key, value)
def __iter__(self) -> Iterator[str]:
"""Return an iterator over the keys of the namespace."""
return iter(vars(self))
def __getitem__(self, key: str) -> Any:
"""Allow bracket-style access to attributes.
Args:
key (str): The name of the attribute.
Returns:
Any: The value of the attribute.
"""
return vars(self)[key]
def items(self):
"""Return a view of the namespace's (key, value) pairs."""
return vars(self).items()
def values(self):
"""Return a view of the namespace's values."""
return vars(self).values()
def keys(self):
"""Return a view of the namespace's keys."""
return vars(self).keys()
def validate_frame(frame: dict, features: dict):
"""Validate a single data frame against the dataset's feature specification.
Checks for missing/extra features, and validates the dtype and shape of each
provided feature.
Args:
frame (dict): The data frame to validate.
features (dict): The LeRobot features dictionary for the dataset.
Raises:
ValueError: If the frame does not match the feature specification.
"""
expected_features = set(features) - set(DEFAULT_FEATURES)
actual_features = set(frame)
@@ -799,6 +1227,15 @@ def validate_frame(frame: dict, features: dict):
def validate_features_presence(actual_features: set[str], expected_features: set[str]):
"""Check for missing or extra features in a frame.
Args:
actual_features (set[str]): The set of feature names present in the frame.
expected_features (set[str]): The set of feature names expected in the frame.
Returns:
str: An error message string if there's a mismatch, otherwise an empty string.
"""
error_message = ""
missing_features = expected_features - actual_features
extra_features = actual_features - expected_features
@@ -814,6 +1251,19 @@ def validate_features_presence(actual_features: set[str], expected_features: set
def validate_feature_dtype_and_shape(name: str, feature: dict, value: np.ndarray | PILImage.Image | str):
"""Validate the dtype and shape of a single feature's value.
Args:
name (str): The name of the feature.
feature (dict): The feature specification from the LeRobot features dictionary.
value: The value of the feature to validate.
Returns:
str: An error message if validation fails, otherwise an empty string.
Raises:
NotImplementedError: If the feature dtype is not supported for validation.
"""
expected_dtype = feature["dtype"]
expected_shape = feature["shape"]
if is_valid_numpy_dtype_string(expected_dtype):
@@ -829,6 +1279,17 @@ def validate_feature_dtype_and_shape(name: str, feature: dict, value: np.ndarray
def validate_feature_numpy_array(
name: str, expected_dtype: str, expected_shape: list[int], value: np.ndarray
):
"""Validate a feature that is expected to be a numpy array.
Args:
name (str): The name of the feature.
expected_dtype (str): The expected numpy dtype as a string.
expected_shape (list[int]): The expected shape.
value (np.ndarray): The numpy array to validate.
Returns:
str: An error message if validation fails, otherwise an empty string.
"""
error_message = ""
if isinstance(value, np.ndarray):
actual_dtype = value.dtype
@@ -846,6 +1307,18 @@ def validate_feature_numpy_array(
def validate_feature_image_or_video(name: str, expected_shape: list[str], value: np.ndarray | PILImage.Image):
"""Validate a feature that is expected to be an image or video frame.
Accepts `np.ndarray` (channel-first or channel-last) or `PIL.Image.Image`.
Args:
name (str): The name of the feature.
expected_shape (list[str]): The expected shape (C, H, W).
value: The image data to validate.
Returns:
str: An error message if validation fails, otherwise an empty string.
"""
# Note: The check of pixels range ([0,1] for float and [0,255] for uint8) is done by the image writer threads.
error_message = ""
if isinstance(value, np.ndarray):
@@ -862,12 +1335,35 @@ def validate_feature_image_or_video(name: str, expected_shape: list[str], value:
def validate_feature_string(name: str, value: str):
"""Validate a feature that is expected to be a string.
Args:
name (str): The name of the feature.
value (str): The value to validate.
Returns:
str: An error message if validation fails, otherwise an empty string.
"""
if not isinstance(value, str):
return f"The feature '{name}' is expected to be of type 'str', but type '{type(value)}' provided instead.\n"
return ""
def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features: dict):
"""Validate the episode buffer before it's written to disk.
Ensures the buffer has the required keys, contains at least one frame, and
has features consistent with the dataset's specification.
Args:
episode_buffer (dict): The buffer containing data for a single episode.
total_episodes (int): The current total number of episodes in the dataset.
features (dict): The LeRobot features dictionary for the dataset.
Raises:
ValueError: If the buffer is invalid.
NotImplementedError: If the episode index is manually set and doesn't match.
"""
if "size" not in episode_buffer:
raise ValueError("size key not found in episode_buffer")