Compare commits

..

1 Commits

Author SHA1 Message Date
CarolinePascal 8353f84f52 docs(configs): bring src/lerobot/configs/ to 100% docstring coverage
Documents every remaining public class/function across configs/: the
types.py feature/normalization enums and PolicyFeature, the
PreTrainedConfig and RewardModelConfig base classes (full Args: blocks,
abstract property/method contracts, from_pretrained overrides), the
TrainPipelineConfig and EvalPipelineConfig top-level pipeline configs,
default.py's DatasetConfig/WandBConfig/EvalConfig/PeftConfig/JobConfig,
dataset.py's DatasetRecordConfig, video.py's VideoEncoderConfig, the
accelerator.py ActivationCheckpointingMode enum, and parser.py's CLI
argument-parsing helpers.

Converts existing inline `#` field comments on config dataclasses to
machine-checked Args: blocks (matching each field against its actual
constructor signature), and corrects a couple of stale/misattributed
comments found along the way (EvalPipelineConfig's misplaced field
comment, PreTrainedConfig's phantom normalization_mapping field).

Expands docs/source/api/configs.mdx from 5 documented classes to the
module's full public surface. Adds lerobot.configs to
check_docstrings.py's MODULES_TO_CHECK ratchet and removes the
module's ruff D-ignore. Docstrings and docstring-format changes only;
no behavioral changes.
2026-08-07 11:57:10 +02:00
33 changed files with 997 additions and 883 deletions
+62
View File
@@ -10,14 +10,26 @@ itself with `@register_subclass("name")` and is then selectable by that name on
[[autodoc]] lerobot.configs.train.TrainPipelineConfig
## EvalPipelineConfig
[[autodoc]] lerobot.configs.eval.EvalPipelineConfig
## PreTrainedConfig
[[autodoc]] lerobot.configs.PreTrainedConfig
## RewardModelConfig
[[autodoc]] lerobot.configs.rewards.RewardModelConfig
## DatasetConfig
[[autodoc]] lerobot.configs.DatasetConfig
## DatasetRecordConfig
[[autodoc]] lerobot.configs.DatasetRecordConfig
## EvalConfig
[[autodoc]] lerobot.configs.EvalConfig
@@ -25,3 +37,53 @@ itself with `@register_subclass("name")` and is then selectable by that name on
## WandBConfig
[[autodoc]] lerobot.configs.WandBConfig
## PeftConfig
[[autodoc]] lerobot.configs.PeftConfig
## JobConfig
[[autodoc]] lerobot.configs.JobConfig
## Feature types
[[autodoc]] lerobot.configs.FeatureType
[[autodoc]] lerobot.configs.PipelineFeatureType
[[autodoc]] lerobot.configs.NormalizationMode
[[autodoc]] lerobot.configs.PolicyFeature
[[autodoc]] lerobot.configs.RTCAttentionSchedule
## Video encoding
[[autodoc]] lerobot.configs.VideoEncoderConfig
[[autodoc]] lerobot.configs.RGBEncoderConfig
[[autodoc]] lerobot.configs.DepthEncoderConfig
[[autodoc]] lerobot.configs.encoder_config_from_video_info
## Distributed training
[[autodoc]] lerobot.configs.parallelism.ParallelismConfig
[[autodoc]] lerobot.configs.parallelism.ContextParallelConfig
[[autodoc]] lerobot.configs.accelerator.AcceleratorConfig
[[autodoc]] lerobot.configs.accelerator.FSDPConfig
[[autodoc]] lerobot.configs.accelerator.DDPConfig
[[autodoc]] lerobot.configs.accelerator.GradientAccumulationConfig
[[autodoc]] lerobot.configs.accelerator.CompileConfig
[[autodoc]] lerobot.configs.accelerator.ActivationCheckpointingConfig
[[autodoc]] lerobot.configs.accelerator.ActivationCheckpointingMode
-34
View File
@@ -21,37 +21,3 @@ See [Using LeRobotDataset](../lerobot-dataset-v3) for the format and the common
## StreamingLeRobotDataset
[[autodoc]] lerobot.datasets.StreamingLeRobotDataset
## EpisodeAwareSampler
[[autodoc]] lerobot.datasets.sampler.EpisodeAwareSampler
## Editing a dataset
Functions in `lerobot.datasets.dataset_tools` for editing an existing `LeRobotDataset` on disk: adding,
removing, or modifying features; splitting, merging, or deleting episodes; re-encoding video; and
recomputing statistics. Each returns a new dataset rather than mutating the source in place.
[[autodoc]] lerobot.datasets.dataset_tools.add_features
[[autodoc]] lerobot.datasets.dataset_tools.remove_feature
[[autodoc]] lerobot.datasets.dataset_tools.modify_features
[[autodoc]] lerobot.datasets.dataset_tools.modify_tasks
[[autodoc]] lerobot.datasets.dataset_tools.delete_episodes
[[autodoc]] lerobot.datasets.dataset_tools.split_dataset
[[autodoc]] lerobot.datasets.dataset_tools.merge_datasets
[[autodoc]] lerobot.datasets.dataset_tools.recompute_stats
[[autodoc]] lerobot.datasets.dataset_tools.reencode_dataset
[[autodoc]] lerobot.datasets.dataset_tools.convert_image_to_video_dataset
## Aggregating datasets
[[autodoc]] lerobot.datasets.aggregate.aggregate_datasets
+1 -1
View File
@@ -439,8 +439,8 @@ ignore = [
"src/lerobot/async_inference/**" = ["D"]
"src/lerobot/cameras/**" = ["D"]
"src/lerobot/common/**" = ["D"]
"src/lerobot/configs/**" = ["D"]
"src/lerobot/data_processing/**" = ["D"]
"src/lerobot/datasets/**" = ["D"]
"src/lerobot/distributed/**" = ["D"]
"src/lerobot/envs/**" = ["D"]
"src/lerobot/jobs/**" = ["D"]
+1 -2
View File
@@ -12,8 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Public API for lerobot configuration types and base config classes.
"""Public API for lerobot configuration types and base config classes.
NOTE: TrainPipelineConfig, EvalPipelineConfig, and TrainRLServerPipelineConfig
are intentionally NOT re-exported here to avoid circular dependencies
+7
View File
@@ -171,6 +171,13 @@ class CompileConfig:
class ActivationCheckpointingMode(str, Enum):
"""The activation-checkpointing strategy applied to FSDP wrap units.
**Attributes**:
- **NONE** -- No activation checkpointing.
- **FULL** -- Checkpoint every wrap unit.
"""
NONE = "none"
FULL = "full"
+51 -31
View File
@@ -23,56 +23,76 @@ from .video import DepthEncoderConfig, RGBEncoderConfig, depth_encoder_defaults,
@dataclass
class DatasetRecordConfig:
# Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).
"""Shared dataset recording configuration used by both `lerobot-record` and `lerobot-rollout`.
Args:
repo_id (`str`, *optional*, defaults to `""`): Dataset identifier. By convention it should match
`'{hf_username}/{dataset_name}'` (e.g. `lerobot/test`).
single_task (`str`, *optional*, defaults to `""`): A short but accurate description of the task performed during the
recording (e.g. `"Pick the Lego block and drop it in the box on the right."`).
root (`str | Path | None`, *optional*): Root directory where the dataset will be stored (e.g.
`'dataset/path'`). If `None`, defaults to `$HF_LEROBOT_HOME/repo_id`.
fps (`int`, *optional*, defaults to 30): Limit the frames per second.
episode_time_s (`int | float`, *optional*, defaults to 60): Number of seconds for data recording
for each episode.
reset_time_s (`int | float`, *optional*, defaults to 60): Number of seconds for resetting the
environment after each episode.
num_episodes (`int`, *optional*, defaults to 50): Number of episodes to record.
video (`bool`, *optional*, defaults to `True`): Encode frames in the dataset into video.
push_to_hub (`bool`, *optional*, defaults to `True`): Upload dataset to the Hugging Face Hub.
private (`bool | None`, *optional*): If `True`, upload as private; if `None`, defer to the org
default on the Hub (only affects orgs).
tags (`list[str] | None`, *optional*): Add tags to your dataset on the Hub.
num_image_writer_processes (`int`, *optional*, defaults to 0): Number of subprocesses handling the
saving of frames as PNG. Set to 0 to use threads only; set to >=1 to use subprocesses, each
using threads to write images. The best number of processes and threads depends on your
system. We recommend 4 threads per camera with 0 processes. If fps is unstable, adjust the
thread count. If still unstable, try using 1 or more subprocesses.
num_image_writer_threads_per_camera (`int`, *optional*, defaults to 4): Number of threads writing
the frames as png images on disk, per camera. Too many threads might cause unstable
teleoperation fps due to the main thread being blocked. Not enough threads might cause low
camera fps.
video_encoding_batch_size (`int`, *optional*, defaults to 1): Number of episodes to record before
batch encoding videos. Set to 1 for immediate encoding (default behavior), or higher for
batched encoding.
rgb_encoder (`RGBEncoderConfig`, *optional*): Video encoder settings for camera MP4s (codec,
quality, GOP, etc.). Tuned via CLI nested keys, e.g. `--dataset.rgb_encoder.vcodec=h264`.
depth_encoder (`DepthEncoderConfig`, *optional*): Video encoder settings for depth-map MP4s (codec,
quality, GOP, etc.). Tuned via CLI nested keys.
streaming_encoding (`bool`, *optional*, defaults to `False`): Enable streaming video encoding:
encode frames in real-time during capture instead of writing PNG images first. Makes
`save_episode()` near-instant. More info in the documentation:
https://huggingface.co/docs/lerobot/streaming_video_encoding
encoder_queue_maxsize (`int`, *optional*, defaults to 30): Maximum number of frames to buffer per
camera when using streaming encoding. ~1s buffer at 30fps. Provides backpressure if the encoder
can't keep up.
encoder_threads (`int | None`, *optional*): Number of threads per encoder instance. `None` means
auto (codec default). Lower values reduce CPU usage; maps to `'lp'` (via `svtav1-params`) for
libsvtav1 and `'threads'` for h264/hevc.
no_stamp (`bool`, *optional*, defaults to `False`): Skip appending the date-time tag to `repo_id`,
keeping the user-provided name as-is (e.g. self-managed versioned names intended for a later
`lerobot-edit-dataset merge`).
"""
repo_id: str = ""
# A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.")
single_task: str = ""
# Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
root: str | Path | None = None
# Limit the frames per second.
fps: int = 30
# Number of seconds for data recording for each episode.
episode_time_s: int | float = 60
# Number of seconds for resetting the environment after each episode.
reset_time_s: int | float = 60
# Number of episodes to record.
num_episodes: int = 50
# Encode frames in the dataset into video
video: bool = True
# Upload dataset to Hugging Face hub.
push_to_hub: bool = True
# If True, upload as private; if None, defer to the org default on the Hub (only affects orgs).
private: bool | None = None
# Add tags to your dataset on the hub.
tags: list[str] | None = None
# Number of subprocesses handling the saving of frames as PNG. Set to 0 to use threads only;
# set to ≥1 to use subprocesses, each using threads to write images. The best number of processes
# and threads depends on your system. We recommend 4 threads per camera with 0 processes.
# If fps is unstable, adjust the thread count. If still unstable, try using 1 or more subprocesses.
num_image_writer_processes: int = 0
# Number of threads writing the frames as png images on disk, per camera.
# Too many threads might cause unstable teleoperation fps due to main thread being blocked.
# Not enough threads might cause low camera fps.
num_image_writer_threads_per_camera: int = 4
# Number of episodes to record before batch encoding videos
# Set to 1 for immediate encoding (default behavior), or higher for batched encoding
video_encoding_batch_size: int = 1
# Video encoder settings for camera MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys,
# e.g. ``--dataset.rgb_encoder.vcodec=h264`` (see ``RGBEncoderConfig``).
rgb_encoder: RGBEncoderConfig = field(default_factory=rgb_encoder_defaults)
# Video encoder settings for depth-map MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys.
depth_encoder: DepthEncoderConfig = field(default_factory=depth_encoder_defaults)
# Enable streaming video encoding: encode frames in real-time during capture instead
# of writing PNG images first. Makes save_episode() near-instant. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding
streaming_encoding: bool = False
# Maximum number of frames to buffer per camera when using streaming encoding.
# ~1s buffer at 30fps. Provides backpressure if the encoder can't keep up.
encoder_queue_maxsize: int = 30
# Number of threads per encoder instance. None = auto (codec default).
# Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
encoder_threads: int | None = None
# Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
# (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
no_stamp: bool = False
def stamp_repo_id(self) -> None:
+128 -58
View File
@@ -27,35 +27,65 @@ logger = logging.getLogger(__name__)
@dataclass
class DatasetConfig:
# You may provide a list of datasets here. `train.py` creates them all and concatenates them. Note: only data
# keys common between the datasets are kept. Each dataset gets and additional transform that inserts the
# "dataset_index" into the returned item. The index mapping is made according to the order in which the
# datasets are provided.
"""A dataset to train on. `TrainPipelineConfig.dataset` may be a list of these, concatenated together.
Only data keys common between multiple datasets are kept. Each dataset gets an additional transform
that inserts the `"dataset_index"` into the returned item, with the index mapping made according to
the order in which the datasets are provided.
Args:
repo_id (`str`): The Hub repo ID (or local dataset name, if `root` is set) to load.
repo_type (`str`, *optional*, defaults to `"dataset"`): Hub repository type: `"dataset"` (the
default) or `"bucket"` for an HF Storage Bucket streamed over `hf://buckets/`. Buckets are
streaming-only, so `"bucket"` requires `streaming=True`.
root (`str | None`, *optional*): Root directory for a concrete local dataset tree (e.g.
`'dataset/path'`). If `None`, local datasets are looked up under `$HF_LEROBOT_HOME/repo_id` and
Hub downloads use a revision-safe cache under `$HF_LEROBOT_HOME/hub`.
episodes (`list[int] | None`, *optional*): Episode indices to include. If `None`, all episodes are
used.
exclude_episodes (`list[int] | None`, *optional*): Episode indices to drop (e.g. corrupt or
heterogeneous ones). Applied on top of `episodes`.
image_transforms (`ImageTransformsConfig`, *optional*): Image augmentation settings applied at load
time.
revision (`str | None`, *optional*): Hub revision (commit hash, branch, or tag) to load.
use_imagenet_stats (`bool`, *optional*, defaults to `True`): Whether to use ImageNet normalization
statistics for visual features instead of the dataset's own.
video_backend (`str`, *optional*): The video decoding backend to use.
return_uint8 (`bool`, *optional*, defaults to `False`): When `True`, RGB video frames are returned
as `uint8` tensors (0-255) instead of `float32` (0.0-1.0). This reduces memory and speeds up
DataLoader IPC. The training pipeline handles the conversion.
depth_output_unit (`str`, *optional*, defaults to `"mm"`): Physical unit depth maps are dequantized
to at load time: `"mm"` (millimeters) or `"m"` (metres). Has no effect on datasets without depth
cameras.
streaming (`bool`, *optional*, defaults to `False`): Stream the dataset instead of downloading it
locally.
eval_split (`float`, *optional*, defaults to 0.0): Fraction of episodes held out per task for
offline evaluation (0.0 = disabled).
"""
repo_id: str
# Hub repository type: "dataset" (default) or "bucket" for an HF Storage Bucket streamed over
# hf://buckets/. Buckets are streaming-only, so "bucket" requires streaming=true.
repo_type: str = "dataset"
# Root directory for a concrete local dataset tree (e.g. 'dataset/path'). If None, local datasets are
# looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
root: str | None = None
episodes: list[int] | None = None
# Episode indices to drop (e.g. corrupt or heterogeneous ones). Applied on top of `episodes`.
exclude_episodes: list[int] | None = None
image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
revision: str | None = None
use_imagenet_stats: bool = True
video_backend: str = field(default_factory=get_safe_default_video_backend)
# When True, RGB video frames are returned as uint8 tensors (0-255) instead of float32 (0.0-1.0).
# This reduces memory and speeds up DataLoader IPC. The training pipeline handles the conversion.
return_uint8: bool = False
# Physical unit depth maps are dequantized to at load time: "mm" (millimeters) or "m" (metres).
# Has no effect on datasets without depth cameras.
depth_output_unit: str = DEFAULT_DEPTH_UNIT
streaming: bool = False
# Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
eval_split: float = 0.0
def __post_init__(self) -> None:
"""Validate `repo_type`/`streaming`/`depth_output_unit`/`eval_split`/`episodes`/`exclude_episodes`.
Raises:
ValueError: If `repo_type` isn't `"dataset"` or `"bucket"`; if `repo_type="bucket"` is combined
with `streaming=False` or a nonzero `eval_split`; if `depth_output_unit` isn't a recognized
unit; if `eval_split` is outside `[0.0, 1.0)`; or if `episodes` contains negative or
duplicate indices.
"""
if self.repo_type not in ("dataset", "bucket"):
raise ValueError(f"repo_type must be 'dataset' or 'bucket', got {self.repo_type!r}")
if self.repo_type == "bucket" and not self.streaming:
@@ -92,35 +122,63 @@ class DatasetConfig:
@dataclass
class WandBConfig:
"""Weights & Biases logging settings for `lerobot-train`.
Args:
enable (`bool`, *optional*, defaults to `False`): Whether to log this run to Weights & Biases.
disable_artifact (`bool`, *optional*, defaults to `False`): Set to `True` to disable saving an
artifact despite `save_checkpoint=True`.
project (`str`, *optional*, defaults to `"lerobot"`): The WandB project to log to.
entity (`str | None`, *optional*): The WandB entity (team or username) to log under.
notes (`str | None`, *optional*): Notes attached to the WandB run.
run_id (`str | None`, *optional*): An existing WandB run id to resume logging into.
mode (`str | None`, *optional*): WandB mode: `"online"`, `"offline"`, or `"disabled"`. Defaults to
`"online"`.
add_tags (`bool`, *optional*, defaults to `True`): If `True`, save the training configuration as
tags on the WandB run.
"""
enable: bool = False
# Set to true to disable saving an artifact despite training.save_checkpoint=True
disable_artifact: bool = False
project: str = "lerobot"
entity: str | None = None
notes: str | None = None
run_id: str | None = None
mode: str | None = None # Allowed values: 'online', 'offline' 'disabled'. Defaults to 'online'
add_tags: bool = True # If True, save configuration as tags in the WandB run.
mode: str | None = None
add_tags: bool = True
@dataclass
class EvalConfig:
"""Settings for the periodic in-training simulation-environment evaluation.
Args:
n_episodes (`int`, *optional*, defaults to 50): Number of episodes to run per evaluation.
batch_size (`int`, *optional*, defaults to 0): The number of environments to use in a
`gym.vector.VectorEnv`. `0` auto-tunes based on available CPU cores and `n_episodes`.
use_async_envs (`bool`, *optional*, defaults to `True`): Whether to use asynchronous environments
(multiprocessing). Automatically downgraded to a `SyncVectorEnv` when `batch_size` is 1.
recording (`bool`, *optional*, defaults to `False`): Whether to record eval rollouts as a LeRobot
dataset on disk.
recording_repo_id (`str | None`, *optional*): If set, push recorded eval datasets to the Hub under
this repo id (one repo per task, suffixed by task and env index). Requires `recording=True`.
recording_private (`bool`, *optional*, defaults to `False`): Whether the pushed recording
repositories should be private.
"""
n_episodes: int = 50
# `batch_size` specifies the number of environments to use in a gym.vector.VectorEnv.
# Set to 0 for auto-tuning based on available CPU cores and n_episodes.
batch_size: int = 0
# `use_async_envs` specifies whether to use asynchronous environments (multiprocessing).
# Defaults to True; automatically downgraded to SyncVectorEnv when batch_size=1.
use_async_envs: bool = True
# Whether to record eval rollouts as a LeRobot dataset on disk.
recording: bool = False
# If set, push recorded eval datasets to the Hub under this repo id (one repo per task,
# suffixed by task and env index). Requires recording=true.
recording_repo_id: str | None = None
# Whether the pushed recording repositories should be private.
recording_private: bool = False
def __post_init__(self) -> None:
"""Validate `recording_repo_id`/`recording`, and resolve/cap `batch_size`.
Raises:
ValueError: If `recording_repo_id` is set without `recording=True`.
"""
if self.recording_repo_id is not None and not self.recording:
raise ValueError("eval.recording_repo_id requires eval.recording=true.")
if self.batch_size == 0:
@@ -141,54 +199,66 @@ class EvalConfig:
@dataclass
class PeftConfig:
# PEFT offers many fine-tuning methods, layer adapters being the most common and currently also the most
# effective methods so we'll focus on those in this high-level config interface.
"""PEFT (parameter-efficient fine-tuning) settings, e.g. LoRA adapters.
PEFT offers many fine-tuning methods, layer adapters being the most common and currently also the
most effective methods so we'll focus on those in this high-level config interface.
Args:
target_modules (`list[str] | str | None`, *optional*): Either a string (module name suffix or
`'all-linear'`), a list of module name suffixes, or a regular expression describing module
names to target with the configured PEFT method. Some policies have a default value for this
so that you don't *have* to choose which layers to adapt, but it might still be worthwhile
depending on your case.
full_training_modules (`list[str] | None`, *optional*): Names/suffixes of modules to fully
fine-tune and store alongside adapter weights. Useful for layers that are not part of a
pre-trained model (e.g., action state projections). Depending on the policy this defaults to
layers that are newly created in pre-trained policies. If you're fine-tuning an already trained
policy you might want to set this to `[]`. Corresponds to PEFT's `modules_to_save`.
method_type (`str`, *optional*, defaults to `"LORA"`): The PEFT (adapter) method to apply to the
policy. Needs to be a valid PEFT type.
init_type (`str | None`, *optional*): Adapter initialization method. Look at the specific PEFT
adapter documentation for defaults.
r (`int`, *optional*, defaults to 16): We expect that all PEFT adapters are in some way doing
rank-decomposition, therefore this parameter specifies the rank used for the adapter. In
general a higher rank means more trainable parameters and closer to full fine-tuning.
lora_alpha (`int | None`, *optional*): Alpha parameter for LoRA scaling (`scaling = lora_alpha /
r`). In general, a higher alpha means stronger adaptation signal. If `None`, the PEFT library
defaults to `alpha=8`, which may dampen high-rank adapters. Common values are `r` (`alpha ==
rank`) or `2*r`.
"""
# Either a string (module name suffix or 'all-linear'), a list of module name suffixes or a regular expression
# describing module names to target with the configured PEFT method. Some policies have a default value for this
# so that you don't *have* to choose which layers to adapt but it might still be worthwhile depending on your case.
target_modules: list[str] | str | None = None
# Names/suffixes of modules to fully fine-tune and store alongside adapter weights. Useful for layers that are
# not part of a pre-trained model (e.g., action state projections). Depending on the policy this defaults to layers
# that are newly created in pre-trained policies. If you're fine-tuning an already trained policy you might want
# to set this to `[]`. Corresponds to PEFT's `modules_to_save`.
full_training_modules: list[str] | None = None
# The PEFT (adapter) method to apply to the policy. Needs to be a valid PEFT type.
method_type: str = "LORA"
# Adapter initialization method. Look at the specific PEFT adapter documentation for defaults.
init_type: str | None = None
# We expect that all PEFT adapters are in some way doing rank-decomposition therefore this parameter specifies
# the rank used for the adapter. In general a higher rank means more trainable parameters and closer to full
# fine-tuning.
r: int = 16
# Alpha parameter for LoRA scaling (scaling = lora_alpha / r).
# In general, a higher alpha means stronger adaptation signal.
# If None, the PEFT library defaults to alpha=8, which may dampen high-rank adapters.
# Common values are r (alpha == rank) or 2*r.
lora_alpha: int | None = None
@dataclass
class JobConfig:
# Where training runs. None (omitted) or "local" runs on this machine.
# Any other value is an HF Jobs flavor and submits the run to HF Jobs.
# List available flavors + pricing with `hf jobs hardware` command.
"""Where and how a training run executes: locally, or dispatched to an HF Jobs flavor.
Args:
target (`str | None`, *optional*): Where training runs. `None` (omitted) or `"local"` runs on this
machine. Any other value is an HF Jobs flavor and submits the run to HF Jobs. List available
flavors and pricing with the `hf jobs hardware` command.
image (`str`, *optional*, defaults to `"huggingface/lerobot-gpu:latest"`): Runtime image for the
remote job (ignored for local runs).
timeout (`str | None`, *optional*, defaults to `"2d"`): Max wall-clock for the remote job as an HF
Jobs duration string (e.g. `"2h"`). HF Jobs itself defaults to `"2d"`; we pass an explicit,
generous cap instead. Set a smaller value to fail fast, or a larger one for long runs.
detach (`bool`, *optional*, defaults to `False`): Submit and exit instead of streaming the job logs
in the foreground.
tags (`list[str]`, *optional*): Extra tags attached to the HF job and to any dataset this run
pushes to the Hub. A `"lerobot"` tag is always added; e.g. `--job.tags '["lelab"]'` adds more.
"""
target: str | None = None
# Runtime image for the remote job (ignored for local runs).
image: str = "huggingface/lerobot-gpu:latest"
# Max wall-clock for the remote job as an HF Jobs duration string (e.g. "2h").
# Defaults to "2d": We pass an explicit, generous cap instead. Set a smaller
# value to fail fast, or a larger one for long runs.
timeout: str | None = "2d"
# Submit and exit instead of streaming the job logs in the foreground.
detach: bool = False
# Extra tags attached to the HF job and to any dataset this run pushes to the
# Hub. A "lerobot" tag is always added; e.g. --job.tags '["lelab"]' adds more.
tags: list[str] = field(default_factory=list)
# Two entry points to the same predicate: the staticmethod tests a raw target string
+21 -6
View File
@@ -28,21 +28,36 @@ logger = getLogger(__name__)
@dataclass
class EvalPipelineConfig:
# Either the repo ID of a model hosted on the Hub or a path to a directory containing weights
# saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch
# (useful for debugging). This argument is mutually exclusive with `--config`.
"""The top-level configuration for `lerobot-eval`, parsed by draccus from CLI flags and/or a YAML file.
Args:
env (`envs.EnvConfig`): The simulation environment to evaluate the policy in.
eval (`EvalConfig`, *optional*): Number of episodes, batching, and recording settings for the
evaluation run.
policy (`PreTrainedConfig | None`, *optional*): Loaded from `--policy.path`, either the repo ID of
a model hosted on the Hub or a path to a directory containing weights saved using
`PreTrainedPolicy.save_pretrained`. If not provided, the policy is initialized from scratch
(useful for debugging).
output_dir (`Path | None`, *optional*): Where to save evaluation outputs.
job_name (`str | None`, *optional*): A name for the run.
seed (`int | None`, *optional*, defaults to 1000): Seed used for the evaluation environments.
rename_map (`dict[str, str]`, *optional*): Rename map for the observation, to override the image
and state keys.
trust_remote_code (`bool`, *optional*, defaults to `False`): Explicit consent to execute remote
code from the Hub (required for Hub environments).
"""
env: envs.EnvConfig
eval: EvalConfig = field(default_factory=EvalConfig)
policy: PreTrainedConfig | None = None
output_dir: Path | None = None
job_name: str | None = None
seed: int | None = 1000
# Rename map for the observation to override the image and state keys
rename_map: dict[str, str] = field(default_factory=dict)
# Explicit consent to execute remote code from the Hub (required for hub environments).
trust_remote_code: bool = False
def __post_init__(self) -> None:
"""Resolve `--policy.path` into a loaded config, and derive `job_name`/`output_dir` when unset."""
# HACK: We parse again the cli args here to get the pretrained path if there was one.
policy_path = parser.get_path_arg("policy")
if policy_path:
@@ -75,5 +90,5 @@ class EvalPipelineConfig:
@classmethod
def __get_path_fields__(cls) -> list[str]:
"""This enables the parser to load config from the policy using `--policy.path=local/dir`"""
"""This enables the parser to load config from the policy using `--policy.path=local/dir`."""
return ["policy"]
+30 -19
View File
@@ -96,6 +96,7 @@ def get_cli_overrides(field_name: str, args: Sequence[str] | None = None) -> lis
def parse_arg(arg_name: str, args: Sequence[str] | None = None) -> str | None:
"""Return the value of `--{arg_name}=value` or `--{arg_name} value` in `args` (`sys.argv[1:]` if `None`)."""
if args is None:
args = sys.argv[1:]
option = f"--{arg_name}"
@@ -115,7 +116,7 @@ def parse_plugin_args(plugin_arg_suffix: str, args: Sequence[str]) -> dict[str,
Args:
plugin_arg_suffix (str): The suffix to identify plugin-related arguments.
cli_args (Sequence[str]): A sequence of command-line arguments to parse.
args (`Sequence[str]`): A sequence of command-line arguments to parse.
Returns:
dict: A dictionary containing the parsed plugin arguments where:
@@ -156,7 +157,7 @@ def load_plugin(plugin_path: str) -> None:
registered with their parents using the `register_subclass` decorator.
Args:
plugin_path (str): The Python package path to the plugin (e.g. "mypackage.plugins.myplugin")
plugin_path (str): The Python package path to the plugin, e.g. "mypackage.plugins.myplugin".
Raises:
PluginLoadError: If the plugin cannot be loaded due to import errors or if the package path is invalid.
@@ -180,6 +181,7 @@ def load_plugin(plugin_path: str) -> None:
) from e
def iter_namespace(ns_pkg: ModuleType) -> Iterable[ModuleInfo]:
"""Iterate the direct submodules of `ns_pkg`, yielding their fully-qualified names."""
return pkgutil.iter_modules(ns_pkg.__path__, ns_pkg.__name__ + ".")
try:
@@ -192,6 +194,7 @@ def load_plugin(plugin_path: str) -> None:
def get_path_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
"""Return `--{field_name}.path`'s value from CLI `args`, or from a YAML/JSON config if not on the CLI."""
result = parse_arg(f"{field_name}.{PATH_KEY}", args)
if result is None:
result = _config_path_args.get(field_name)
@@ -199,21 +202,24 @@ def get_path_arg(field_name: str, args: Sequence[str] | None = None) -> str | No
def get_yaml_overrides(field_name: str) -> list[str]:
"""Return the CLI-style overrides extracted from `field_name`'s YAML/JSON config path block, if any."""
return _config_yaml_overrides.get(field_name, [])
def get_type_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
"""Return `--{field_name}.type`'s value from CLI `args` (`sys.argv[1:]` if `None`)."""
return parse_arg(f"{field_name}.{draccus.CHOICE_TYPE_KEY}", args)
def _register_scoped_actions(
wrapper: Wrapper, parser: SuppressingArgumentParser, cli_args: Sequence[str]
) -> None:
"""Like draccus's own Wrapper.register_actions, but for a ChoiceType field only recurses into
the already-selected subclass (per CLI `.type` args), instead of every registered choice.
"""Like draccus's own Wrapper.register_actions, but only recurses a ChoiceType field's subclasses.
This mirrors draccus 0.11.x's internal wrapper traversal because its public parser eagerly registers
every choice before parsing the command line. Keep this in sync when updating draccus.
Only the already-selected subclass (per CLI `.type` args) is recursed into, instead of every
registered choice. This mirrors draccus 0.11.x's internal wrapper traversal because its public
parser eagerly registers every choice before parsing the command line. Keep this in sync when
updating draccus.
"""
if isinstance(wrapper, ChoiceWrapper):
group = parser.add_argument_group(title=wrapper.title, description=wrapper.description)
@@ -253,8 +259,10 @@ def _register_scoped_actions(
def print_scoped_help(config_class: type, cli_args: Sequence[str]) -> None:
"""Prints --help output scoped to the choices already resolved on the CLI (e.g. --env.type=pusht),
instead of draccus's default of expanding every registered subclass of every ChoiceType field."""
"""Prints --help output scoped to the choices already resolved on the CLI (e.g. --env.type=pusht).
Instead of draccus's default of expanding every registered subclass of every ChoiceType field.
"""
parser = SuppressingArgumentParser(formatter_class=SimpleHelpFormatter)
parser.add_argument(
f"--{draccus.utils.CONFIG_ARG}", type=str, help="Path for a config file to parse with draccus"
@@ -264,6 +272,7 @@ def print_scoped_help(config_class: type, cli_args: Sequence[str]) -> None:
def filter_arg(field_to_filter: str, args: Sequence[str] | None = None) -> list[str]:
"""Return `args` with `--{field_to_filter}` (and its value, if separate) removed."""
if args is None:
return []
option = f"--{field_to_filter}"
@@ -285,12 +294,11 @@ def filter_arg(field_to_filter: str, args: Sequence[str] | None = None) -> list[
def filter_path_args(fields_to_filter: str | list[str], args: Sequence[str] | None = None) -> list[str]:
"""
Filters command-line arguments related to fields with specific path arguments.
"""Filters command-line arguments related to fields with specific path arguments.
Args:
fields_to_filter (str | list[str]): A single str or a list of str whose arguments need to be filtered.
args (Sequence[str] | None): The sequence of command-line arguments to be filtered.
args (Sequence[str] | None, *optional*): The sequence of command-line arguments to be filtered.
Defaults to None.
Returns:
@@ -380,19 +388,22 @@ def extract_path_fields_from_config(config_path: str, path_fields: list[str]) ->
def wrap(config_path: Path | None = None) -> Callable[[F], F]:
"""
HACK: Similar to draccus.wrap but does three additional things:
- Will remove '.path' arguments from CLI in order to process them later on.
- If a 'config_path' is passed and the main config class has a 'from_pretrained' method, will
initialize it from there to allow to fetch configs from the hub directly
- Will load plugins specified in the CLI arguments. These plugins will typically register
their own subclasses of config classes, so that draccus can find the right class to instantiate
from the CLI '.type' arguments
"""HACK: Similar to draccus.wrap but does three additional things.
- Will remove '.path' arguments from CLI in order to process them later on.
- If a 'config_path' is passed and the main config class has a 'from_pretrained' method, will
initialize it from there to allow to fetch configs from the hub directly
- Will load plugins specified in the CLI arguments. These plugins will typically register
their own subclasses of config classes, so that draccus can find the right class to instantiate
from the CLI '.type' arguments
"""
def wrapper_outer(fn: F) -> F:
"""Wrap `fn` so its first argument is resolved from the CLI/config instead of passed directly."""
@wraps(fn)
def wrapper_inner(*args: Any, **kwargs: Any) -> Any:
"""Build `fn`'s config argument from the CLI/config file (unless already given), then call `fn`."""
argspec = inspect.getfullargspec(fn)
argtype = argspec.annotations[argspec.args[0]]
if len(args) > 0 and type(args[0]) is argtype:
+88 -21
View File
@@ -39,50 +39,64 @@ logger = getLogger(__name__)
@dataclass
class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: ignore[misc,name-defined] #TODO: draccus issue
"""
Base configuration class for policy models.
"""Base configuration class for policy models.
Every concrete policy config also declares a `normalization_mapping: dict[str, NormalizationMode]`
field (mapping a `FeatureType` name, e.g. `"STATE"`/`"VISUAL"`, to the `NormalizationMode` to apply),
with a policy-specific default — not declared here since it has no sensible shared default.
Args:
n_obs_steps: Number of environment steps worth of observations to pass to the policy (takes the
current step and additional steps going back).
input_features: A dictionary defining the PolicyFeature of the input data for the policy. The key represents
the input data name, and the value is PolicyFeature, which consists of FeatureType and shape attributes.
output_features: A dictionary defining the PolicyFeature of the output data for the policy. The key represents
the output data name, and the value is PolicyFeature, which consists of FeatureType and shape attributes.
normalization_mapping: A dictionary that maps from a str value of FeatureType (e.g., "STATE", "VISUAL") to
a corresponding NormalizationMode (e.g., NormalizationMode.MIN_MAX)
n_obs_steps (`int`, *optional*, defaults to 1): Number of environment steps worth of observations
to pass to the policy (takes the current step and additional steps going back).
input_features (`dict[str, PolicyFeature] | None`, *optional*): A dictionary defining the
`PolicyFeature` of the input data for the policy. The key represents the input data name, and
the value is a `PolicyFeature`, which consists of `type` and `shape` attributes. Can be set to
`None`/`null` in order to infer those values from the dataset.
output_features (`dict[str, PolicyFeature] | None`, *optional*): A dictionary defining the
`PolicyFeature` of the output data for the policy, with the same key/value semantics as
`input_features`.
device (`str | None`, *optional*): The torch device, e.g. `"cuda"`, `"cuda:0"`, `"cpu"`, or `"mps"`.
If unset or unavailable, `__post_init__` auto-selects one.
use_amp (`bool`, *optional*, defaults to `False`): Whether to use Automatic Mixed Precision for
training and evaluation, with automatic gradient scaling. Auto-disabled by `__post_init__`
when AMP isn't available on `device`.
use_peft (`bool`, *optional*, defaults to `False`): Whether the policy employed PEFT for training.
push_to_hub (`bool`, *optional*, defaults to `True`): Whether to push the policy to the Hugging Face
Hub after training.
repo_id (`str | None`, *optional*): The Hub repo ID to push to. Required when `push_to_hub` is
`True`.
private (`bool | None`, *optional*): Whether to upload to a private repository on the Hugging Face
Hub.
tags (`list[str] | None`, *optional*): Tags to add to the policy on the Hub.
license (`str | None`, *optional*): The license to add to the policy on the Hub.
pretrained_path (`Path | None`, *optional*): Either the repo ID of a model hosted on the Hub or a
path to a directory containing weights saved using `PreTrainedPolicy.save_pretrained`. If not
provided, the policy is initialized from scratch.
pretrained_revision (`str | None`, *optional*): Hub revision (commit hash, branch, or tag) to pin
the pretrained model version.
"""
n_obs_steps: int = 1
# `input_features` can be set to None/null in order to infer those values from the dataset.
input_features: dict[str, PolicyFeature] | None = field(default_factory=dict)
output_features: dict[str, PolicyFeature] | None = field(default_factory=dict)
device: str | None = None # e.g. "cuda", "cuda:0", "cpu", or "mps"
# `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
# automatic gradient scaling is used.
device: str | None = None
use_amp: bool = False
# Whether the policy employed PEFT for training.
use_peft: bool = False
push_to_hub: bool = True # type: ignore[assignment] # TODO: use a different name to avoid override
repo_id: str | None = None
# Upload on private repository on the Hugging Face hub.
private: bool | None = None
# Add tags to your policy on the hub.
tags: list[str] | None = None
# Add tags to your policy on the hub.
license: str | None = None
# Either the repo ID of a model hosted on the Hub or a path to a directory containing weights
# saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch.
pretrained_path: Path | None = None
# Optional Hub revision (commit hash, branch, or tag) to pin the pretrained model version.
pretrained_revision: str | None = None
def __post_init__(self) -> None:
"""Auto-select `device` when unset/unavailable, and disable `use_amp` when AMP isn't available on it."""
if not self.device or not is_torch_device_available(self.device):
auto_device = auto_select_torch_device()
logger.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
@@ -97,6 +111,7 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
@property
def type(self) -> str:
"""The policy's registered `draccus.ChoiceRegistry` name (e.g. `"act"`, `"diffusion"`)."""
choice_name = self.get_choice_name(self.__class__)
if not isinstance(choice_name, str):
raise TypeError(f"Expected string from get_choice_name, got {type(choice_name)}")
@@ -105,32 +120,52 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
@property
@abc.abstractmethod
def observation_delta_indices(self) -> list | None: # type: ignore[type-arg] #TODO: No implementation
"""Offsets, relative to the current step, of the observation timesteps the policy consumes.
`None` means only the current step is used.
"""
raise NotImplementedError
@property
@abc.abstractmethod
def action_delta_indices(self) -> list | None: # type: ignore[type-arg] #TODO: No implementation
"""Offsets, relative to the current step, of the action timesteps the policy predicts/consumes.
`None` means only the current step is used.
"""
raise NotImplementedError
@property
@abc.abstractmethod
def reward_delta_indices(self) -> list | None: # type: ignore[type-arg] #TODO: No implementation
"""Offsets, relative to the current step, of the reward timesteps the policy consumes.
`None` means only the current step is used.
"""
raise NotImplementedError
@abc.abstractmethod
def get_optimizer_preset(self) -> OptimizerConfig:
"""Return this policy's default `OptimizerConfig`, used when `use_policy_training_preset` is set."""
raise NotImplementedError
@abc.abstractmethod
def get_scheduler_preset(self) -> LRSchedulerConfig | None:
"""Return this policy's default `LRSchedulerConfig`, or `None` if it uses no scheduler."""
raise NotImplementedError
@abc.abstractmethod
def validate_features(self) -> None:
"""Check that `input_features`/`output_features` contain what this policy requires.
Raises:
ValueError: If a required feature is missing or has an unexpected shape/type.
"""
raise NotImplementedError
@property
def robot_state_feature(self) -> PolicyFeature | None:
"""The input `PolicyFeature` for the robot's proprioceptive state (`observation.state`), if any."""
if not self.input_features:
return None
for ft_name, ft in self.input_features.items():
@@ -140,6 +175,7 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
@property
def env_state_feature(self) -> PolicyFeature | None:
"""The input `PolicyFeature` of type `FeatureType.ENV` (environment state), if any."""
if not self.input_features:
return None
for _, ft in self.input_features.items():
@@ -149,12 +185,14 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
@property
def image_features(self) -> dict[str, PolicyFeature]:
"""All input features of type `FeatureType.VISUAL`, keyed by feature name."""
if not self.input_features:
return {}
return {key: ft for key, ft in self.input_features.items() if ft.type is FeatureType.VISUAL}
@property
def action_feature(self) -> PolicyFeature | None:
"""The output `PolicyFeature` for the action (`action`), if any."""
if not self.output_features:
return None
for ft_name, ft in self.output_features.items():
@@ -182,6 +220,35 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
revision: str | None = None,
**policy_kwargs: Any,
) -> T:
"""Download a policy's `config.json` from the Hub (or read it locally) and parse it.
The concrete policy config subclass is resolved from the serialized `"type"` tag (e.g. `"act"`,
`"diffusion"`) rather than being fixed by `cls`, so calling this on the `PreTrainedConfig` base
class works for any registered policy type.
Args:
pretrained_name_or_path (`str | Path`): Either the `repo_id` of the config hosted on the Hub,
or a path to a directory containing a `config.json` saved via `.save_pretrained`.
force_download (`bool`, *optional*, defaults to `False`): Whether to force (re-)downloading
the files from the Hub, overriding the existing cache.
resume_download (`bool | None`, *optional*): Deprecated; ignored by the underlying Hub client.
proxies (`dict[Any, Any] | None`, *optional*): A dictionary of proxy servers to use by protocol
or endpoint.
token (`str | bool | None`, *optional*): The token to use as HTTP bearer authorization for
remote files. By default, uses the token cached by `huggingface-cli login`.
cache_dir (`str | Path | None`, *optional*): Path to the folder where cached files are stored.
local_files_only (`bool`, *optional*, defaults to `False`): If `True`, avoid downloading the
file and return the path to the local cached file if it exists.
revision (`str | None`, *optional*): Revision on the Hub: a branch name, git tag, or commit id.
Defaults to the latest commit on `main`.
policy_kwargs: Forwarded as CLI-style overrides via `policy_kwargs["cli_overrides"]`
(a list of `--key=value` strings applied on top of the loaded config); any other keys are
ignored.
Raises:
FileNotFoundError: If `config.json` isn't found locally or on the Hub.
ValueError: If `config.json` has no `"type"` field, or its value isn't a registered policy type.
"""
model_id = str(pretrained_name_or_path)
config_file: str | None = None
if Path(model_id).is_dir():
+57 -7
View File
@@ -43,31 +43,45 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
"""Base configuration for reward models.
Args:
input_features: A dictionary defining the PolicyFeature of the input data for the reward. The key represents
the input data name, and the value is PolicyFeature, which consists of FeatureType and shape attributes.
output_features: A dictionary defining the PolicyFeature of the output data for the reward. The key represents
the output data name, and the value is PolicyFeature, which consists of FeatureType and shape attributes.
input_features (`dict[str, PolicyFeature]`, *optional*): A dictionary defining the `PolicyFeature`
of the input data for the reward. The key represents the input data name, and the value is a
`PolicyFeature`, which consists of `type` and `shape` attributes.
output_features (`dict[str, PolicyFeature]`, *optional*): A dictionary defining the `PolicyFeature`
of the output data for the reward, with the same key/value semantics as `input_features`.
device (`str | None`, *optional*): The torch device, e.g. `"cuda"`, `"cuda:0"`, `"cpu"`, or `"mps"`.
If unset or unavailable, `__post_init__` auto-selects one.
pretrained_path (`str | None`, *optional*): Either the repo ID of a model hosted on the Hub or a
path to a directory containing weights saved using `.save_pretrained`. If not provided, the
reward model is initialized from scratch.
pretrained_revision (`str | None`, *optional*): Optional Hub revision, e.g. a commit hash, branch,
or tag, to pin the pretrained reward model version.
push_to_hub (`bool`, *optional*, defaults to `False`): Whether to push the reward model to the
Hugging Face Hub after training.
repo_id (`str | None`, *optional*): The Hub repo ID to push to. Required when `push_to_hub` is
`True`.
license (`str | None`, *optional*): The license to add to the reward model on the Hub.
tags (`list[str] | None`, *optional*): Tags to add to the reward model on the Hub.
private (`bool | None`, *optional*): Whether to upload to a private repository on the Hugging Face
Hub.
"""
# Reuses PolicyFeature
input_features: dict[str, PolicyFeature] = field(default_factory=dict)
output_features: dict[str, PolicyFeature] = field(default_factory=dict)
device: str | None = None
pretrained_path: str | None = None
# Optional Hub revision (commit hash, branch, or tag) to pin the pretrained reward model version.
pretrained_revision: str | None = None
push_to_hub: bool = False
repo_id: str | None = None
# Hub metadata
license: str | None = None
tags: list[str] | None = None
private: bool | None = None
def __post_init__(self) -> None:
"""Auto-select `device` when unset or unavailable."""
if not self.device or not is_torch_device_available(self.device):
auto_device = auto_select_torch_device()
logger.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
@@ -75,6 +89,7 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
@property
def type(self) -> str:
"""The reward model's registered `draccus.ChoiceRegistry` name."""
choice_name = self.get_choice_name(self.__class__)
if not isinstance(choice_name, str):
raise TypeError(f"Expected string from get_choice_name, got {type(choice_name)}")
@@ -82,14 +97,17 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
@property
def observation_delta_indices(self) -> list | None: # type: ignore[type-arg]
"""`None`: reward models consume only the current observation timestep."""
return None
@property
def action_delta_indices(self) -> list | None: # type: ignore[type-arg]
"""`None`: reward models consume only the current action timestep."""
return None
@property
def reward_delta_indices(self) -> list | None: # type: ignore[type-arg]
"""`None`: reward models consume only the current reward timestep."""
return None
def get_optimizer_preset(self) -> OptimizerConfig | None:
@@ -97,9 +115,14 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
return None
def get_scheduler_preset(self) -> LRSchedulerConfig | None:
"""Default LR scheduler for this reward model. `None` here; overridden by subclasses that need one."""
return None
def validate_features(self) -> None:
"""Check that `input_features`/`output_features` contain what this reward model requires.
No-op here; overridden by subclasses that have required features.
"""
pass
def _save_pretrained(self, save_directory: Path) -> None:
@@ -122,6 +145,33 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
revision: str | None = None,
**reward_kwargs: Any,
) -> T:
"""Download a reward model's `config.json` from the Hub (or read it locally) and parse it.
The concrete reward-model config subclass is resolved from the serialized `"type"` tag, so
calling this on the `RewardModelConfig` base class works for any registered reward-model type.
Args:
pretrained_name_or_path (`str | Path`): Either the `repo_id` of the config hosted on the Hub,
or a path to a directory containing a `config.json` saved via `.save_pretrained`.
force_download (`bool`, *optional*, defaults to `False`): Whether to force (re-)downloading
the files from the Hub, overriding the existing cache.
resume_download (`bool | None`, *optional*): Deprecated; ignored by the underlying Hub client.
proxies (`dict[Any, Any] | None`, *optional*): A dictionary of proxy servers to use by protocol
or endpoint.
token (`str | bool | None`, *optional*): The token to use as HTTP bearer authorization for
remote files. By default, uses the token cached by `huggingface-cli login`.
cache_dir (`str | Path | None`, *optional*): Path to the folder where cached files are stored.
local_files_only (`bool`, *optional*, defaults to `False`): If `True`, avoid downloading the
file and return the path to the local cached file if it exists.
revision (`str | None`, *optional*): Revision on the Hub: a branch name, git tag, or commit id.
Defaults to the latest commit on `main`.
reward_kwargs: Forwarded as CLI-style overrides via `reward_kwargs["cli_overrides"]`
(a list of `--key=value` strings applied on top of the loaded config); any other keys are
ignored.
Raises:
FileNotFoundError: If `config.json` isn't found locally or on the Hub.
"""
model_id = str(pretrained_name_or_path)
config_file: str | None = None
if Path(model_id).is_dir():
+121 -32
View File
@@ -108,75 +108,119 @@ def _migrate_legacy_rabc_fields(config: dict[str, Any]) -> dict[str, Any] | None
@dataclass
class TrainPipelineConfig(HubMixin):
"""The top-level configuration for `lerobot-train`, parsed by draccus from CLI flags and/or a YAML file.
Args:
dataset (`DatasetConfig`): The dataset(s) to train on.
env (`envs.EnvConfig | None`, *optional*): The simulation environment to periodically evaluate the
policy in (see `env_eval_freq`). Required when `env_eval_freq > 0`.
policy (`PreTrainedConfig | None`, *optional*): The policy to train. Mutually exclusive with
`reward_model`.
reward_model (`RewardModelConfig | None`, *optional*): The reward model to train instead of a
policy. Mutually exclusive with `policy`.
output_dir (`Path | None`, *optional*): Where to save all of the run outputs. If you run another
training session with the same value its contents will be overwritten unless `resume` is set.
job_name (`str | None`, *optional*): A name for the run.
resume (`bool`, *optional*, defaults to `False`): Resume a previous run. Pass `--config_path`
pointing at either a local checkpoint's `train_config.json` or a Hub repo id holding
`checkpoints/<step>/` subtrees (the latest checkpoint is downloaded and resumed from). When
resuming, the default behavior is to use the configuration from the checkpoint, regardless of
what's provided with the training command at the time of resumption (CLI `--*` flags still
override).
seed (`int | None`, *optional*, defaults to 1000): Seed used for training (e.g. model
initialization, dataset shuffling) and for the evaluation environments.
cudnn_deterministic (`bool`, *optional*, defaults to `False`): Use deterministic cuDNN algorithms
for reproducibility. Disables `cudnn.benchmark` and may reduce training speed by ~10-20 percent.
num_workers (`int`, *optional*, defaults to 4): Number of workers for the dataloader.
batch_size (`int`, *optional*, defaults to 8): The training batch size.
prefetch_factor (`int`, *optional*, defaults to 4): Number of batches loaded in advance by each
dataloader worker.
persistent_workers (`bool`, *optional*, defaults to `True`): Keep dataloader worker processes alive
between epochs.
dataloader_multiprocessing_context (`str | None`, *optional*, defaults to `"spawn"`): DataLoader
worker start method. `"spawn"` is safer than `"fork"` with non-fork-safe libs (PyAV /
torchcodec / ffmpeg), but adds some worker-startup time per run since workers re-import modules
instead of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
when appropriate, or set it to `None` to use Python's platform default.
steps (`int`, *optional*, defaults to 100000): Total number of training steps.
env_eval_freq (`int`, *optional*, defaults to 20000): Run the policy in the simulation environment
every N steps to measure reward/success (0 = disabled).
log_freq (`int`, *optional*, defaults to 200): Logging frequency, in steps.
eval_steps (`int`, *optional*, defaults to 0): Compute eval loss on held-out episodes every N steps
(0 = disabled). Requires `eval_split > 0`.
max_eval_samples (`int`, *optional*, defaults to 0): Cap on total eval samples, split uniformly
across tasks (0 = use all held-out data).
tolerance_s (`float`, *optional*, defaults to 0.0001): Maximum timestamp difference tolerated when
loading dataset frames, in seconds.
save_checkpoint (`bool`, *optional*, defaults to `True`): Whether to save checkpoints during
training.
save_freq (`int`, *optional*, defaults to 20000): Save a checkpoint every `save_freq` training
iterations and after the last training step. A non-positive value disables periodic saving,
keeping only the final checkpoint.
checkpoint_format (`CheckpointFormat`, *optional*, defaults to `CheckpointFormat.SAFETENSORS`):
Model-artifact format inside checkpoints; non-default values require a sharded run.
use_policy_training_preset (`bool`, *optional*, defaults to `True`): Use the policy's own
optimizer/scheduler presets when `optimizer`/`scheduler` aren't explicitly set.
optimizer (`OptimizerConfig | None`, *optional*): The optimizer to use. Falls back to the policy's
preset when `use_policy_training_preset` is `True`.
scheduler (`LRSchedulerConfig | None`, *optional*): The learning-rate scheduler to use. Falls back
to the policy's preset when `use_policy_training_preset` is `True`.
parallelism (`ParallelismConfig`, *optional*): Process topology: `dp_replicate` / `dp_shard` for HSDP
and context-parallel degree placeholders.
accelerator (`AcceleratorConfig`, *optional*): Execution runtime handed to the Accelerator: mixed
precision, gradient accumulation, FSDP/DDP tuning knobs, compile & activation-checkpointing
placeholders.
eval (`EvalConfig`, *optional*): Settings for the periodic simulation-environment evaluation.
wandb (`WandBConfig`, *optional*): Weights & Biases logging settings.
peft (`PeftConfig | None`, *optional*): PEFT (e.g. LoRA) settings, when fine-tuning with adapters
instead of full-parameter training.
job (`JobConfig`, *optional*): Where to run training: locally (default), or an HF Jobs flavor.
save_checkpoint_to_hub (`bool`, *optional*, defaults to `False`): Push each saved checkpoint to the
Hub (`policy.repo_id`) as it is written, not just the final model (useful to monitor progress
mid-run). The final model is pushed regardless. Works the same locally and remotely.
sample_weighting (`SampleWeightingConfig | None`, *optional*): Sample weighting configuration (e.g.
for RA-BC training).
rename_map (`dict[str, str]`, *optional*): Rename map for the observation, to override the image
and state keys.
"""
dataset: DatasetConfig
env: envs.EnvConfig | None = None
policy: PreTrainedConfig | None = None
reward_model: RewardModelConfig | None = None
# Set `dir` to where you would like to save all of the run outputs. If you run another training session
# with the same value for `dir` its contents will be overwritten unless you set `resume` to true.
output_dir: Path | None = None
job_name: str | None = None
# Set `resume` to true to resume a previous run. Pass `--config_path` pointing at either a local
# checkpoint's train_config.json or a Hub repo id holding `checkpoints/<step>/` subtrees (the
# latest checkpoint is downloaded and resumed from). Note that when resuming, the default behavior
# is to use the configuration from the checkpoint, regardless of what's provided with the training
# command at the time of resumption (CLI `--*` flags still override).
resume: bool = False
# `seed` is used for training (eg: model initialization, dataset shuffling)
# AND for the evaluation environments.
seed: int | None = 1000
# Set to True to use deterministic cuDNN algorithms for reproducibility.
# This disables cudnn.benchmark and may reduce training speed by ~10-20 percent.
cudnn_deterministic: bool = False
# Number of workers for the dataloader.
num_workers: int = 4
batch_size: int = 8
prefetch_factor: int = 4
persistent_workers: bool = True
# DataLoader worker start method. "spawn" is safer than "fork" with
# non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
# worker-startup time per run since workers re-import modules instead
# of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
# when appropriate, or set it to `null` to use Python's platform default.
dataloader_multiprocessing_context: str | None = "spawn"
steps: int = 100_000
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
env_eval_freq: int = 20_000
log_freq: int = 200
# Compute eval loss on held-out episodes every N steps (0 = disabled). Requires eval_split > 0.
eval_steps: int = 0
# Cap on total eval samples, split uniformly across tasks (0 = use all held-out data).
max_eval_samples: int = 0
tolerance_s: float = 1e-4
save_checkpoint: bool = True
# Checkpoint is saved every `save_freq` training iterations and after the last training step.
# A non-positive value disables periodic saving, keeping only the final checkpoint.
save_freq: int = 20_000
# Model-artifact format inside checkpoints; non-default values require a sharded run.
checkpoint_format: CheckpointFormat = CheckpointFormat.SAFETENSORS
use_policy_training_preset: bool = True
optimizer: OptimizerConfig | None = None
scheduler: LRSchedulerConfig | None = None
# Process topology: dp_replicate / dp_shard (HSDP) and context-parallel degree placeholders.
parallelism: ParallelismConfig = field(default_factory=ParallelismConfig)
# Execution runtime handed to the Accelerator: mixed precision, gradient accumulation,
# FSDP/DDP tuning knobs, compile & activation-checkpointing placeholders.
accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig)
eval: EvalConfig = field(default_factory=EvalConfig)
wandb: WandBConfig = field(default_factory=WandBConfig)
peft: PeftConfig | None = None
# Where to run training (local default, or an HF Jobs flavor). See JobConfig.
job: JobConfig = field(default_factory=JobConfig)
# Push each saved checkpoint to the Hub (policy.repo_id) as it is written, not
# just the final model (useful to monitor progress mid-run). Optional; the
# final model is pushed regardless. Works the same locally and remotely.
save_checkpoint_to_hub: bool = False
# Sample weighting configuration (e.g., for RA-BC training)
sample_weighting: SampleWeightingConfig | None = None
# Rename map for the observation to override the image and state keys
rename_map: dict[str, str] = field(default_factory=dict)
checkpoint_path: Path | None = field(init=False, default=None)
@@ -262,6 +306,21 @@ class TrainPipelineConfig(HubMixin):
self.reward_model.pretrained_path = str(policy_dir)
def validate(self) -> None:
"""Resolve pretrained sources and cross-field defaults, and fail fast on invalid combinations.
Called by draccus after parsing. Resolves `--policy.path`/`--reward_model.path`/`resume` into a
loaded config, derives `job_name` and `output_dir` when unset, and applies the policy's
optimizer/scheduler presets when `use_policy_training_preset` is `True`.
Raises:
ValueError: On an unsupported `dataloader_multiprocessing_context`, neither `policy` nor
`reward_model` configured, a `rename_map` without a pretrained checkpoint, an unsplit
dataset with `eval_steps > 0`, a missing `repo_id` when pushing to the Hub, or
`save_checkpoint_to_hub` without `policy.repo_id` — or (see `_validate_distributed`) an
unsupported distributed-training combination.
FileExistsError: If `output_dir` already exists and `resume` is `False`.
NotImplementedError: If `dataset.repo_id` is a list (multi-dataset training).
"""
available_contexts = multiprocessing.get_all_start_methods()
if (
self.dataloader_multiprocessing_context is not None
@@ -389,6 +448,7 @@ class TrainPipelineConfig(HubMixin):
return ["policy", "reward_model"]
def to_dict(self) -> dict[str, Any]:
"""Encode the config to a plain, JSON-serializable dictionary (via `draccus.encode`)."""
return draccus.encode(self) # type: ignore[no-any-return] # because of the third-party library draccus uses Any as the return type
def _save_pretrained(self, save_directory: Path) -> None:
@@ -409,6 +469,35 @@ class TrainPipelineConfig(HubMixin):
revision: str | None = None,
**kwargs: Any,
) -> "TrainPipelineConfig":
"""Download a run's `train_config.json` from the Hub (or read it locally) and parse it.
Falls back to the latest checkpoint's config when the repo has no root `train_config.json` (a repo
of periodic checkpoints from an interrupted run), so a resume can start straight from
`--config_path=<repo>`. Legacy RA-BC fields in a JSON config are migrated to the current
`sample_weighting` schema.
Args:
pretrained_name_or_path (`str | Path`): Either the `repo_id` of the run hosted on the Hub, or
a path to a directory containing a `train_config.json` saved via `.save_pretrained`.
force_download (`bool`, *optional*, defaults to `False`): Whether to force (re-)downloading
the files from the Hub, overriding the existing cache.
resume_download (`bool | None`, *optional*): Deprecated; ignored by the underlying Hub client.
proxies (`dict[Any, Any] | None`, *optional*): A dictionary of proxy servers to use by protocol
or endpoint.
token (`str | bool | None`, *optional*): The token to use as HTTP bearer authorization for
remote files. By default, uses the token cached by `huggingface-cli login`.
cache_dir (`str | Path | None`, *optional*): Path to the folder where cached files are stored.
local_files_only (`bool`, *optional*, defaults to `False`): If `True`, avoid downloading the
file and return the path to the local cached file if it exists.
revision (`str | None`, *optional*): Revision on the Hub: a branch name, git tag, or commit id.
Defaults to the latest commit on `main`.
kwargs: Forwarded as CLI-style overrides via `kwargs["cli_args"]` (a list of `--key=value`
strings applied on top of the loaded config); any other keys are ignored.
Raises:
FileNotFoundError: If `train_config.json` isn't found locally, on the Hub, or on any checkpoint
within the Hub repo.
"""
model_id = str(pretrained_name_or_path)
config_file: str | None = None
if Path(model_id).is_dir():
+48
View File
@@ -18,6 +18,18 @@ from enum import Enum
class FeatureType(str, Enum):
"""The category of data a `PolicyFeature` represents.
**Attributes**:
- **STATE** -- A robot/environment proprioceptive state vector.
- **VISUAL** -- An image or video feature.
- **ENV** -- Environment-provided state, distinct from robot proprioception (e.g. simulation
environment state).
- **ACTION** -- An action vector.
- **REWARD** -- A scalar reward.
- **LANGUAGE** -- A natural-language feature (e.g. task instruction tokens).
"""
STATE = "STATE"
VISUAL = "VISUAL"
ENV = "ENV"
@@ -27,11 +39,28 @@ class FeatureType(str, Enum):
class PipelineFeatureType(str, Enum):
"""Which side of a processor pipeline a feature belongs to.
**Attributes**:
- **ACTION** -- The feature is part of the action space.
- **OBSERVATION** -- The feature is part of the observation space.
"""
ACTION = "ACTION"
OBSERVATION = "OBSERVATION"
class NormalizationMode(str, Enum):
"""The normalization strategy applied to a feature by a `NormalizerProcessorStep`.
**Attributes**:
- **MIN_MAX** -- Scale to `[-1, 1]` using the feature's min/max statistics.
- **MEAN_STD** -- Center and scale to unit variance using the feature's mean/std statistics.
- **IDENTITY** -- Leave the feature unchanged.
- **QUANTILES** -- Scale to `[-1, 1]` using the feature's 1st/99th percentile statistics.
- **QUANTILE10** -- Scale to `[-1, 1]` using the feature's 10th/90th percentile statistics.
"""
MIN_MAX = "MIN_MAX"
MEAN_STD = "MEAN_STD"
IDENTITY = "IDENTITY"
@@ -41,11 +70,30 @@ class NormalizationMode(str, Enum):
@dataclass
class PolicyFeature:
"""Describes one entry of a policy's input/output feature space.
Args:
type (`FeatureType`): The category of the feature.
shape (`tuple[int, ...]`): The feature's shape, excluding the batch dimension.
"""
type: FeatureType
shape: tuple[int, ...]
class RTCAttentionSchedule(str, Enum):
"""The prefix-attention weighting schedule used by the Real-Time Chunking (RTC) policy.
Controls how much weight is given to the previous action chunk's prediction versus the new one,
over the overlap region between consecutive chunks.
**Attributes**:
- **ZEROS** -- No prefix attention: weight is 1.0 before `start`, then 0.0.
- **ONES** -- Full prefix attention: weight is 1.0 up to `end`, then 0.0.
- **LINEAR** -- Linearly ramps the weight down from 1.0 to 0.0 between `start` and `end`.
- **EXP** -- Like `LINEAR`, but with an exponential (rather than linear) decay curve.
"""
ZEROS = "ZEROS"
ONES = "ONES"
LINEAR = "LINEAR"
+34 -17
View File
@@ -84,18 +84,33 @@ DEPTH_ENCODER_INFO_FIELD_NAMES: frozenset[str] = frozenset({"depth_min", "depth_
@dataclass
class VideoEncoderConfig:
"""Video encoder configuration."""
"""Video encoder configuration.
vcodec: str = "libsvtav1" # Video codec name. "auto" picks a hardware codec if available, else libsvtav1.
pix_fmt: str = "yuv420p" # Pixel format (e.g. yuv420p).
g: int | None = 2 # GOP size (keyframe interval).
crf: int | float | None = 30 # Quality level. Lower means better quality and larger files.
preset: int | str | None = None # Speed/quality preset. Accepted values are codec-specific.
fast_decode: int = 0 # Fast-decode tuning. Accepted values are codec-specific, 0 disables it.
Args:
vcodec (`str`, *optional*, defaults to `"libsvtav1"`): Video codec name. `"auto"` picks a hardware
codec if available, else `libsvtav1`.
pix_fmt (`str`, *optional*, defaults to `"yuv420p"`): Pixel format (e.g. `yuv420p`).
g (`int | None`, *optional*, defaults to 2): GOP size (keyframe interval).
crf (`int | float | None`, *optional*, defaults to 30): Quality level. Lower means better quality
and larger files.
preset (`int | str | None`, *optional*): Speed/quality preset. Accepted values are codec-specific.
fast_decode (`int`, *optional*, defaults to 0): Fast-decode tuning. Accepted values are
codec-specific; 0 disables it.
video_backend (`str`, *optional*, defaults to `"pyav"`): Encoding backend. Only `"pyav"` is
currently supported.
extra_options (`dict[str, Any]`, *optional*): Extra codec options merged last, e.g. `{"tune":
"film"}`.
"""
vcodec: str = "libsvtav1"
pix_fmt: str = "yuv420p"
g: int | None = 2
crf: int | float | None = 30
preset: int | str | None = None
fast_decode: int = 0
# TODO(CarolinePascal): add torchcodec support + find a way to unify the
# two backends (encoding and decoding).
video_backend: str = "pyav" # Encoding backend. Only "pyav" is currently supported.
# Extra codec options merged last, e.g. {"tune": "film"}.
video_backend: str = "pyav"
extra_options: dict[str, Any] = field(default_factory=dict)
# Source-data channel count this encoder is expected to handle. ``None``
@@ -104,6 +119,7 @@ class VideoEncoderConfig:
_DEFAULT_CHANNELS: ClassVar[int | None] = None
def __post_init__(self) -> None:
"""Resolve `vcodec` (e.g. `"auto"`), apply the libsvtav1 default preset, and validate the config."""
self.resolve_vcodec()
# Empty-constructor ergonomics: ``VideoEncoderConfig()`` must "just work".
if self.preset is None and self.vcodec == "libsvtav1":
@@ -112,9 +128,7 @@ class VideoEncoderConfig:
@classmethod
def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]:
"""Parse the ``video.*`` keys of a feature ``info`` block into
constructor kwargs.
"""
"""Parse the ``video.*`` keys of a feature ``info`` block into constructor kwargs."""
video_info = video_info or {}
kwargs: dict[str, Any] = {}
@@ -147,6 +161,7 @@ class VideoEncoderConfig:
Args:
encoders: List of encoder names to detect. If a string, it is converted to a list.
Returns:
List of available encoder names. If the video backend is not "pyav", returns an empty list.
"""
@@ -211,6 +226,7 @@ class VideoEncoderConfig:
opts: dict[str, Any] = {}
def set_if(key: str, value: Any) -> None:
"""Set `opts[key]` to `value` (stringified if `as_strings`), unless `value` is `None`."""
if value is not None:
opts[key] = value if not as_strings else str(value)
@@ -302,9 +318,10 @@ class DepthEncoderConfig(VideoEncoderConfig):
@classmethod
def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]:
"""Layer the depth-specific tuning (``depth_min`` / ``depth_max`` /
``shift`` / ``use_log``) on top of the base parser. Missing keys
fall back to the class defaults.
"""Layer the depth-specific tuning on top of the base parser.
Adds ``depth_min`` / ``depth_max`` / ``shift`` / ``use_log``. Missing keys fall back to the
class defaults.
"""
kwargs = super()._kwargs_from_video_info(video_info)
video_info = video_info or {}
@@ -328,8 +345,8 @@ def encoder_config_from_video_info(video_info: dict | None) -> VideoEncoderConfi
otherwise.
Args:
video_info: A feature's ``info`` dict as persisted in ``info.json``,
or ``None`` (treated as an empty dict).
video_info (`dict | None`): A feature's ``info`` dict as persisted in ``info.json``, or ``None``
(treated as an empty dict).
Returns:
A :class:`DepthEncoderConfig` for depth features, otherwise a
+53 -89
View File
@@ -58,38 +58,12 @@ type ChunkFile = tuple[int, int]
class IndexState(TypedDict):
"""The current write cursor for a non-video (parquet) output stream during aggregation.
**Attributes**:
- **chunk** (`int`) -- The chunk index currently being written to.
- **file** (`int`) -- The file index, within `chunk`, currently being written to.
- **src_to_dst** (`dict[ChunkFile, ChunkFile]`, *optional*) -- Maps each source dataset's
`(chunk, file)` to the destination `(chunk, file)` its rows were merged into.
"""
chunk: int
file: int
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
class VideoIndex(TypedDict):
"""The current write cursor for a video output stream during aggregation.
**Attributes**:
- **chunk** (`int`) -- The chunk index currently being written to.
- **file** (`int`) -- The file index, within `chunk`, currently being written to.
- **latest_duration** (`float`) -- The duration, in seconds, appended to the current destination
file so far.
- **episode_duration** (`float`) -- The duration, in seconds, of the episode currently being
concatenated.
- **src_to_offset** (`dict[ChunkFile, float]`, *optional*) -- Maps each source `(chunk, file)` to
the time offset, in seconds, at which it was appended into its destination file.
- **src_to_dst** (`dict[ChunkFile, ChunkFile]`, *optional*) -- Maps each source `(chunk, file)` to
the destination `(chunk, file)` its video was concatenated into.
- **dst_file_durations** (`dict[ChunkFile, float]`, *optional*) -- The final total duration, in
seconds, of each completed destination file.
"""
chunk: int
file: int
latest_duration: float
@@ -106,7 +80,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
"""Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.
Args:
all_metadata (`list`): List of `LeRobotDatasetMetadata` objects to merge.
all_metadata: List of LeRobotDatasetMetadata objects to merge.
Returns:
dict: A dictionary of merged video feature info.
@@ -152,7 +126,7 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[i
Video encoder info is not considered for validation but is merged during aggregation in ``merge_video_feature_info_for_aggregate``.
Args:
all_metadata (`list`): List of `LeRobotDatasetMetadata` objects to validate.
all_metadata: List of LeRobotDatasetMetadata objects to validate.
Returns:
tuple: A tuple containing (fps, robot_type, features) from the first metadata.
@@ -161,6 +135,7 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[i
ValueError: If any metadata has different fps, robot_type, or features
than the first metadata in the list.
"""
fps = all_metadata[0].fps
robot_type = all_metadata[0].robot_type
features = all_metadata[0].features
@@ -189,13 +164,14 @@ def update_data_df(
previously aggregated data in the destination dataset.
Args:
df (`DataFrame`): DataFrame containing the data to be updated.
src_meta (`LeRobotDatasetMetadata`): Source dataset metadata.
dst_meta (`LeRobotDatasetMetadata`): Destination dataset metadata.
df: DataFrame containing the data to be updated.
src_meta: Source dataset metadata.
dst_meta: Destination dataset metadata.
Returns:
pd.DataFrame: Updated DataFrame with adjusted indices.
"""
df["episode_index"] = df["episode_index"] + dst_meta.info.total_episodes
df["index"] = df["index"] + dst_meta.info.total_frames
@@ -221,15 +197,16 @@ def update_meta_data(
to correctly map source file indices to their destination locations.
Args:
df (`DataFrame`): DataFrame containing the metadata to be updated.
dst_meta (`LeRobotDatasetMetadata`): Destination dataset metadata.
meta_idx (`IndexState`): Dictionary containing current metadata chunk and file indices.
data_idx (`IndexState`): Dictionary containing current data chunk and file indices.
videos_idx (`VideoIndexState`): Dictionary containing current video indices and timestamps.
df: DataFrame containing the metadata to be updated.
dst_meta: Destination dataset metadata.
meta_idx: Dictionary containing current metadata chunk and file indices.
data_idx: Dictionary containing current data chunk and file indices.
videos_idx: Dictionary containing current video indices and timestamps.
Returns:
pd.DataFrame: Updated DataFrame with adjusted indices and timestamps.
"""
df["meta/episodes/chunk_index"] = df["meta/episodes/chunk_index"] + meta_idx["chunk"]
df["meta/episodes/file_index"] = df["meta/episodes/file_index"] + meta_idx["file"]
@@ -390,21 +367,15 @@ def aggregate_datasets(
4. Finalizing the aggregated dataset with proper statistics
Args:
repo_ids (`list`): List of repository IDs for the datasets to aggregate.
aggr_repo_id (`str`): Repository ID for the aggregated output dataset.
roots (`list[pathlib.Path] | None`, *optional*): List of root paths for the source
datasets.
aggr_root (`pathlib.Path | None`, *optional*): Root path for the aggregated dataset.
data_files_size_in_mb (`int | None`, *optional*): Maximum size for data files in MB. Falls
back to `DEFAULT_DATA_FILE_SIZE_IN_MB` when not set.
video_files_size_in_mb (`int | None`, *optional*): Maximum size for video files in MB. Falls
back to `DEFAULT_VIDEO_FILE_SIZE_IN_MB` when not set.
chunk_size (`int | None`, *optional*): Maximum number of files per chunk. Falls back to
`DEFAULT_CHUNK_SIZE` when not set.
concatenate_videos (`bool`, *optional*, defaults to `True`): When `False`, keep one mp4 per
source file instead of packing into shards.
concatenate_data (`bool`, *optional*, defaults to `True`): When `False`, keep one parquet per
source file instead of packing into shards.
repo_ids: List of repository IDs for the datasets to aggregate.
aggr_repo_id: Repository ID for the aggregated output dataset.
roots: Optional list of root paths for the source datasets.
aggr_root: Optional root path for the aggregated dataset.
data_files_size_in_mb: Maximum size for data files in MB (defaults to DEFAULT_DATA_FILE_SIZE_IN_MB)
video_files_size_in_mb: Maximum size for video files in MB (defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB)
chunk_size: Maximum number of files per chunk (defaults to DEFAULT_CHUNK_SIZE)
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
concatenate_data: When False, keep one parquet per source file instead of packing into shards.
"""
logger.info("Start aggregate_datasets")
@@ -487,14 +458,12 @@ def aggregate_videos(
Creates new video files when size limits are exceeded.
Args:
src_meta (`LeRobotDatasetMetadata`): Source dataset metadata.
dst_meta (`LeRobotDatasetMetadata`): Destination dataset metadata.
videos_idx (`VideoIndexState`): Dictionary tracking video chunk and file indices.
video_files_size_in_mb (`float`): Maximum size for video files in MB.
chunk_size (`int`): Maximum number of files per chunk.
concatenate_videos (`bool`, *optional*, defaults to `True`): When `False`, keep one mp4 per
source file instead of packing into shards.
src_meta: Source dataset metadata.
dst_meta: Destination dataset metadata.
videos_idx: Dictionary tracking video chunk and file indices.
video_files_size_in_mb: Maximum size for video files in MB (defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB)
chunk_size: Maximum number of files per chunk (defaults to DEFAULT_CHUNK_SIZE)
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
Returns:
dict: Updated videos_idx with current chunk and file indices.
"""
@@ -612,13 +581,12 @@ def aggregate_data(
have multiple data files (e.g., from a previous merge operation).
Args:
src_meta (`LeRobotDatasetMetadata`): Source dataset metadata.
dst_meta (`LeRobotDatasetMetadata`): Destination dataset metadata.
data_idx (`IndexState`): Dictionary tracking data chunk and file indices.
data_files_size_in_mb (`float`): Maximum size for data files in MB.
chunk_size (`int`): Maximum number of files per chunk.
concatenate_data (`bool`, *optional*, defaults to `True`): When `False`, keep one parquet per
source file instead of packing into shards.
src_meta: Source dataset metadata.
dst_meta: Destination dataset metadata.
data_idx: Dictionary tracking data chunk and file indices.
data_files_size_in_mb: Maximum size for data files in MB.
chunk_size: Maximum number of files per chunk.
concatenate_data: When False, keep one parquet per source file instead of packing into shards.
Returns:
dict: Updated data_idx with current chunk and file indices.
@@ -692,11 +660,11 @@ def aggregate_metadata(
and writes them to the destination with proper file rotation.
Args:
src_meta (`LeRobotDatasetMetadata`): Source dataset metadata.
dst_meta (`LeRobotDatasetMetadata`): Destination dataset metadata.
meta_idx (`IndexState`): Dictionary tracking metadata chunk and file indices.
data_idx (`IndexState`): Dictionary tracking data chunk and file indices.
videos_idx (`VideoIndexState`): Dictionary tracking video indices and timestamps.
src_meta: Source dataset metadata.
dst_meta: Destination dataset metadata.
meta_idx: Dictionary tracking metadata chunk and file indices.
data_idx: Dictionary tracking data chunk and file indices.
videos_idx: Dictionary tracking video indices and timestamps.
Returns:
dict: Updated meta_idx with current chunk and file indices.
@@ -759,22 +727,18 @@ def append_or_create_parquet_file(
from becoming too large. Handles both regular parquet files and those containing images.
Args:
df (`DataFrame`): DataFrame to write to the parquet file.
src_path (`Path`): Path to the source file, used for size estimation.
idx (`IndexState`): Dictionary containing current `chunk` and `file` indices.
max_mb (`float`): Maximum allowed file size in MB before rotation.
chunk_size (`int`): Maximum number of files per chunk before incrementing the chunk index.
default_path (`str`): Format string for generating file paths.
contains_images (`bool`, *optional*, defaults to `False`): Whether the data contains images
requiring special handling.
aggr_root (`pathlib.Path | None`, *optional*): Root path for the aggregated dataset.
hf_features (`datasets.features.features.Features | None`, *optional*): HuggingFace Features
schema used for proper image typing.
concatenate (`bool`, *optional*, defaults to `True`): When `False`, always rotate to a new
file instead of appending to the current one.
one_row_group_per_episode (`bool`, *optional*, defaults to `False`): Whether to emit one
parquet row group per episode. Set to `True` for data parquet files; left `False` for the
episodes-metadata parquet, which already has one row per episode.
df: DataFrame to write to the parquet file.
src_path: Path to the source file (used for size estimation).
idx: Dictionary containing current 'chunk' and 'file' indices.
max_mb: Maximum allowed file size in MB before rotation.
chunk_size: Maximum number of files per chunk before incrementing chunk index.
default_path: Format string for generating file paths.
contains_images: Whether the data contains images requiring special handling.
aggr_root: Root path for the aggregated dataset.
hf_features: Optional HuggingFace Features schema for proper image typing.
concatenate: When False, always rotate to a new file instead of appending to the current one.
one_row_group_per_episode: True for DATA parquet (emit one row group per episode); False for
the episodes-metadata parquet (already one row per episode).
Returns:
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
@@ -838,8 +802,8 @@ def finalize_aggregation(
aggregated statistics from all source datasets.
Args:
aggr_meta (`LeRobotDatasetMetadata`): Aggregated dataset metadata.
all_metadata (`list`): List of all source dataset metadata objects.
aggr_meta: Aggregated dataset metadata.
all_metadata: List of all source dataset metadata objects.
"""
logger.info("write tasks")
write_tasks(aggr_meta.tasks, aggr_meta.root)
+28 -62
View File
@@ -28,22 +28,16 @@ DEFAULT_QUANTILES = [0.01, 0.10, 0.50, 0.90, 0.99]
class RunningQuantileStats:
"""Maintains running statistics for batches of vectors.
"""
Maintains running statistics for batches of vectors, including mean,
standard deviation, min, max, and approximate quantiles.
Includes mean, standard deviation, min, max, and approximate quantiles. Statistics are computed per
feature dimension and updated incrementally
Statistics are computed per feature dimension and updated incrementally
as new batches are observed. Quantiles are estimated using histograms,
which adapt dynamically if the observed data range expands.
"""
def __init__(self, quantile_list: list[float] | None = None, num_quantile_bins: int = 5000):
"""Initialize empty running statistics.
Args:
quantile_list: Quantiles to track (e.g. `0.01` for the 1st percentile). Defaults to
`DEFAULT_QUANTILES` (1st, 10th, 50th, 90th, 99th percentiles).
num_quantile_bins: Number of histogram bins used to estimate quantiles.
"""
self._count = 0
self._mean = None
self._mean_of_squares = None
@@ -210,9 +204,8 @@ def estimate_num_samples(
dataset_len: int, min_num_samples: int = 100, max_num_samples: int = 10_000, power: float = 0.75
) -> int:
"""Heuristic to estimate the number of samples based on dataset size.
The power controls the sample growth relative to dataset size. Lower the power for less number of
samples.
The power controls the sample growth relative to dataset size.
Lower the power for less number of samples.
For default arguments, we have:
- from 1 to ~500, num_samples=100
@@ -228,24 +221,11 @@ def estimate_num_samples(
def sample_indices(data_len: int) -> list[int]:
"""Return evenly-spaced indices into a sequence of length `data_len`, sized by `estimate_num_samples`."""
num_samples = estimate_num_samples(data_len)
return np.round(np.linspace(0, data_len - 1, num_samples)).astype(int).tolist()
def auto_downsample_height_width(img: np.ndarray, target_size: int = 150, max_size_threshold: int = 300):
"""Downsample a `(C, H, W)` image by integer striding if either dimension exceeds `max_size_threshold`.
Args:
img (`np.ndarray`): Input image in `(C, H, W)` layout to potentially downsample.
target_size (`int`, *optional*, defaults to 150): Approximate size, in pixels, that the
largest side should be reduced to.
max_size_threshold (`int`, *optional*, defaults to 300): Size, in pixels, above which the
largest side of `img` triggers downsampling.
Returns:
`img` unchanged, or strided down so its largest side is roughly `target_size`.
"""
_, height, width = img.shape
if max(width, height) < max_size_threshold:
@@ -257,16 +237,6 @@ def auto_downsample_height_width(img: np.ndarray, target_size: int = 150, max_si
def sample_images(image_paths: list[str]) -> np.ndarray:
"""Load and downsample a sampled subset of `image_paths` into a single `uint8` array.
Args:
image_paths (`list[str]`): Paths of all images for the episode/feature, from which a
subset is sampled (see `sample_indices`).
Returns:
A `(N, C, H, W)` `uint8` array of the sampled, downsampled images (see
`auto_downsample_height_width`), where `N` is chosen by `sample_indices`.
"""
sampled_indices = sample_indices(len(image_paths))
images = None
@@ -439,7 +409,6 @@ def _compute_basic_stats(
Args:
array: Reshaped array ready for statistics computation
sample_count: Number of samples represented in the data
quantile_list: Quantiles to fill with the mean value. Defaults to `DEFAULT_QUANTILES`.
Returns:
Dictionary with basic statistics and quantiles set to mean values
@@ -478,14 +447,13 @@ def get_feature_stats(
- Global: axis=None computes statistics over entire array
Args:
array (`np.ndarray`): Input data array with a shape appropriate for the specified `axis`.
axis (`int | tuple[int, ...] | None`): Axis or axes along which to compute statistics:
`(0, 2, 3)` for image data (batch, channels, height, width), `0` or `(0,)` for
vector/tabular data (samples, features), `(1,)` to compute across features, or
`None` for global statistics over the entire array.
keepdims (`bool`): If `True`, reduced axes are kept as dimensions of size 1.
quantile_list (`list[float] | None`, *optional*): Quantiles to compute (e.g. `0.01` for
the 1st percentile). Defaults to `DEFAULT_QUANTILES` when not provided.
array: Input data array with shape appropriate for the specified axis
axis: Axis or axes along which to compute statistics
- (0, 2, 3): For image data (batch, channels, height, width)
- 0 or (0,): For vector/tabular data (samples, features)
- (1,): For computing across features
- None: For global statistics over entire array
keepdims: If True, reduced axes are kept as dimensions with size 1
Returns:
Dictionary containing:
@@ -528,13 +496,10 @@ def compute_episode_stats(
- Strings: Skipped (no statistics computed)
Args:
episode_data (`dict[str, list[str] | np.ndarray]`): Mapping from feature name to its data
for the episode: a list of file paths for `image`/`video` features, or a numpy array
for numerical features.
features (`dict`): Dataset feature metadata, keyed by feature name, describing each
feature's `dtype` and shape.
quantile_list (`list[float] | None`, *optional*): Quantiles to compute (e.g. `0.01` for
the 1st percentile). Defaults to `DEFAULT_QUANTILES` when not provided.
episode_data: Dictionary mapping feature names to data
- For images/videos: list of file paths
- For numerical data: numpy arrays
features: Dictionary describing each feature's dtype and shape
Returns:
Dictionary mapping feature names to their statistics dictionaries.
@@ -665,6 +630,7 @@ def aggregate_stats(stats_list: list[dict[str, dict]]) -> dict[str, dict[str, np
- new_mean = (mean of all data, weighted by counts)
- new_std = (std of all data)
"""
_assert_type_and_shape(stats_list)
data_keys = {key for stats in stats_list for key in stats}
@@ -724,16 +690,16 @@ def compute_relative_action_stats(
statistics suitable for normalization.
Args:
hf_dataset (`datasets.Dataset`): The underlying HuggingFace dataset, must expose
`"action"`, `"observation.state"`, and `"episode_index"` columns.
features (`dict`): Dataset feature metadata; must contain `"action"` with a `"shape"`
entry and optionally `"names"`.
chunk_size (`int`): Number of consecutive frames per action chunk.
exclude_joints (`list[str] | None`, *optional*): Joint names whose dimensions should
remain absolute instead of being converted to relative actions.
num_workers (`int`, *optional*, defaults to 0): Number of parallel threads used for
computation. Values `<= 1` run single-threaded; NumPy releases the GIL so threads
give real parallelism here.
hf_dataset: The underlying HuggingFace dataset with "action",
"observation.state", and "episode_index" columns.
features: Dataset feature metadata (must contain "action" with "shape"
and optionally "names").
chunk_size: Number of consecutive frames per action chunk.
exclude_joints: Joint names whose dimensions should remain absolute
(not converted to relative actions).
num_workers: Number of parallel threads for computation. Values ≤1
mean single-threaded. Numpy releases the GIL so threads give
real parallelism here.
Returns:
Statistics dict with keys "mean", "std", "min", "max", "q01", …, "q99".
+3 -4
View File
@@ -529,9 +529,9 @@ class LeRobotDatasetMetadata:
return self.info.video_files_size_in_mb
def get_task_index(self, 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.
"""
Given a task in natural language, returns its task_index if the task already exists in the dataset,
otherwise return None.
"""
if task in self.tasks.index:
return int(self.tasks.loc[task].task_index)
@@ -774,7 +774,6 @@ class LeRobotDatasetMetadata:
}
def __repr__(self):
"""A short summary: repo ID, total episode/frame counts, and feature keys."""
feature_keys = list(self.features)
return (
f"{self.__class__.__name__}({{\n"
+1 -3
View File
@@ -293,9 +293,7 @@ class DatasetReader:
return result
def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict[str, torch.Tensor]:
"""Decode the requested per-camera frame timestamps from `ep_idx`'s videos.
Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
"""Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
in the main process (e.g. by using a second Dataloader with num_workers=0). It will result in a
Segmentation Fault.
"""
+77 -86
View File
@@ -118,11 +118,10 @@ def delete_episodes(
consistent with its own metadata.
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset.
episode_indices (`list`): List of episode indices to delete.
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the edited dataset
will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/repo_id`.
repo_id (`str | None`, *optional*): Identifier for the edited dataset.
dataset: The source LeRobotDataset.
episode_indices: List of episode indices to delete.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
"""
if not episode_indices:
raise ValueError("No episodes to delete")
@@ -186,11 +185,10 @@ def split_dataset(
output split stays consistent with its own metadata.
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset to split.
splits (`dict`): Either a dict mapping split names to episode indices, or a dict mapping
split names to fractions (must sum to <= 1.0).
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the split
datasets will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/repo_id`.
dataset: The source LeRobotDataset to split.
splits: Either a dict mapping split names to episode indices, or a dict mapping
split names to fractions (must sum to <= 1.0).
output_dir: Root directory where the split datasets will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id.
Examples:
Split by specific episodes
@@ -282,14 +280,11 @@ def merge_datasets(
This is a wrapper around the aggregate_datasets functionality with a cleaner API.
Args:
datasets (`list`): List of LeRobotDatasets to merge.
output_repo_id (`str`): Identifier for the merged dataset.
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the merged dataset
will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/output_repo_id`.
concatenate_videos (`bool`, *optional*, defaults to `True`): When `False`, keep one mp4 per
source file instead of packing them into shards.
concatenate_data (`bool`, *optional*, defaults to `True`): When `False`, keep one parquet file
per source file instead of packing them into shards.
datasets: List of LeRobotDatasets to merge.
output_repo_id: Merged dataset identifier.
output_dir: Root directory where the merged dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/output_repo_id.
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
concatenate_data: When False, keep one parquet per source file instead of packing into shards.
"""
if not datasets:
raise ValueError("No datasets to merge")
@@ -332,14 +327,11 @@ def modify_features(
regardless of how many features are being added or removed.
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset.
add_features (`dict[str, tuple[numpy.ndarray | torch.Tensor | collections.abc.Callable, dict]] | None`, *optional*):
Dict mapping feature names to `(feature_values, feature_info)` tuples.
remove_features (`str | list[str] | None`, *optional*): Feature name(s) to remove. Can be a
single string or a list.
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the edited dataset
will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/repo_id`.
repo_id (`str | None`, *optional*): Identifier for the edited dataset.
dataset: The source LeRobotDataset.
add_features: Optional dict mapping feature names to (feature_values, feature_info) tuples.
remove_features: Optional feature name(s) to remove. Can be a single string or list.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
Returns:
New dataset with features modified.
@@ -438,11 +430,10 @@ def add_features(
copies the dataset once regardless of how many features are being added.
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset.
features (`dict`): Dictionary mapping feature names to `(feature_values, feature_info)` tuples.
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the edited dataset
will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/repo_id`.
repo_id (`str | None`, *optional*): Identifier for the edited dataset.
dataset: The source LeRobotDataset.
features: Dictionary mapping feature names to (feature_values, feature_info) tuples.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
Returns:
New dataset with all features added.
@@ -476,12 +467,10 @@ def remove_feature(
"""Remove features from a LeRobotDataset.
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset.
feature_names (`str | list[str]`): Name(s) of features to remove. Can be a single string or
a list.
output_dir (`str | pathlib.Path | None`, *optional*): Root directory where the edited dataset
will be stored. If not specified, defaults to `$HF_LEROBOT_HOME/repo_id`.
repo_id (`str | None`, *optional*): Identifier for the edited dataset.
dataset: The source LeRobotDataset.
feature_names: Name(s) of features to remove. Can be a single string or list.
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
Returns:
New dataset with features removed.
@@ -960,7 +949,7 @@ def _copy_and_reindex_episodes_metadata(
def _write_parquet(df: pd.DataFrame, path: Path, meta: LeRobotDatasetMetadata) -> None:
"""Write DataFrame to parquet.
"""Write DataFrame to parquet
This ensures images are properly embedded and the file can be loaded correctly by HF datasets.
"""
@@ -1468,14 +1457,13 @@ def modify_tasks(
- meta/info.json (total_tasks)
Args:
dataset (`LeRobotDataset`): The source LeRobotDataset to modify.
new_task (`str | None`, *optional*): Default task applied to any episode not covered by
`episode_tasks` or a matching `task_replacements` entry.
episode_tasks (`dict[int, str] | None`, *optional*): Dict mapping episode indices to task
strings. Takes precedence over both `task_replacements` and `new_task`.
task_replacements (`dict[str, str] | None`, *optional*): Dict mapping existing task strings to
new ones. Applied to episodes whose current task matches a key. Every key must be an
existing task.
dataset: The source LeRobotDataset to modify.
new_task: Default task applied to any episode not covered by `episode_tasks` or a
matching `task_replacements` entry.
episode_tasks: Optional dict mapping episode indices to task strings. Takes precedence
over both `task_replacements` and `new_task`.
task_replacements: Optional dict mapping existing task strings to new ones. Applied to
episodes whose current task matches a key. Every key must be an existing task.
At least one of `new_task`, `episode_tasks`, or `task_replacements` must be provided.
@@ -1606,19 +1594,19 @@ def recompute_stats(
"""Recompute stats.json from scratch by iterating all episodes.
Args:
dataset (`LeRobotDataset`): The LeRobotDataset to recompute stats for.
skip_image_video (`bool`, *optional*, defaults to `True`): If `True`, only recompute stats for
numeric features (action, state, etc.) and keep existing image/video stats unchanged.
relative_action (`bool`, *optional*, defaults to `False`): If `True`, compute action stats in
relative space by iterating all valid action chunks and subtracting the current state.
This matches the normalization distribution the model sees during training with
`use_relative_actions=True`.
relative_exclude_joints (`list[str] | None`, *optional*): Joint names to exclude from relative
conversion when `relative_action=True`. These dims keep absolute stats.
chunk_size (`int`, *optional*, defaults to 50): Action chunk size used for relative stats
computation. Should match `policy.chunk_size`. Only used when `relative_action=True`.
num_workers (`int`, *optional*, defaults to 0): Number of parallel threads for relative action
stats computation. Values <=1 mean single-threaded. Only used when `relative_action=True`.
dataset: The LeRobotDataset to recompute stats for.
skip_image_video: If True (default), only recompute stats for numeric features
(action, state, etc.) and keep existing image/video stats unchanged.
relative_action: If True, compute action stats in relative space by
iterating all valid action chunks and subtracting the current state.
This matches the normalization distribution the model sees during
training with ``use_relative_actions=True``.
relative_exclude_joints: Joint names to exclude from relative conversion when
relative_action=True. These dims keep absolute stats.
chunk_size: Action chunk size used for relative stats computation. Should match
``policy.chunk_size``. Only used when ``relative_action=True``.
num_workers: Number of parallel threads for relative action stats computation.
Values 1 mean single-threaded. Only used when ``relative_action=True``.
Returns:
The same dataset with updated stats.
@@ -1721,22 +1709,24 @@ def convert_image_to_video_dataset(
LeRobot dataset structure with videos stored in chunked MP4 files.
Args:
dataset (`LeRobotDataset`): The source LeRobot dataset with images.
output_dir (`pathlib.Path | None`, *optional*): Root directory where the converted dataset will
be stored. When `None`, defaults to `$HF_LEROBOT_HOME/repo_id`.
repo_id (`str | None`, *optional*): Identifier for the converted dataset.
rgb_encoder (`lerobot.configs.video.RGBEncoderConfig | None`, *optional*): Video encoder settings
applied to RGB cameras. When `None`, `rgb_encoder_defaults` is used.
depth_encoder (`lerobot.configs.video.DepthEncoderConfig | None`, *optional*): Video encoder
settings applied to depth-map cameras, including the quantization parameters persisted to
the dataset metadata. When `None`, `depth_encoder_defaults` is used.
episode_indices (`list[int] | None`, *optional*): Episode indices to convert. When `None`, all
episodes are converted.
num_workers (`int`, *optional*, defaults to 4): Number of threads for parallel processing.
max_episodes_per_batch (`int | None`, *optional*): Maximum episodes per video batch, to bound
memory use. `None` means no limit.
max_frames_per_batch (`int | None`, *optional*): Maximum frames per video batch, to bound memory
use. `None` means no limit.
dataset: The source LeRobot dataset with images.
output_dir: Root directory where the converted dataset will be stored. When
``None``, defaults to ``$HF_LEROBOT_HOME/repo_id``. Equivalent to
``new_root`` in ``EditDatasetConfig``.
repo_id: Converted dataset identifier. Equivalent to ``new_repo_id`` in
``EditDatasetConfig``.
rgb_encoder: Video encoder settings applied to RGB cameras. When ``None``,
:func:`~lerobot.configs.video.rgb_encoder_defaults` is used.
depth_encoder: Video encoder settings applied to depth-map cameras, including
the quantization parameters persisted to the dataset metadata. When
``None``, :func:`~lerobot.configs.video.depth_encoder_defaults` is used.
episode_indices: Episode indices to convert. When ``None``, all episodes are
converted.
num_workers: Number of threads for parallel processing.
max_episodes_per_batch: Maximum episodes per video batch, to bound memory use.
``None`` means no limit.
max_frames_per_batch: Maximum frames per video batch, to bound memory use.
``None`` means no limit.
Returns:
A new :class:`LeRobotDataset` with images encoded as videos.
@@ -1976,17 +1966,18 @@ def reencode_dataset(
Videos are re-encoded in-place and the video information in ``info.json`` is refreshed.
Args:
dataset (`LeRobotDataset`): An existing :class:`LeRobotDataset` whose videos will be re-encoded.
rgb_encoder (`lerobot.configs.video.RGBEncoderConfig | None`, *optional*): Target encoder
configuration applied to every RGB video file. If `None`, re-encoding is skipped for RGB
videos.
depth_encoder (`lerobot.configs.video.DepthEncoderConfig | None`, *optional*): Target encoder
configuration applied to every depth video file. If `None`, re-encoding is skipped for depth
videos. Quantization parameters will not override the ones in the current dataset.
encoder_threads (`int | None`, *optional*): Per-encoder thread count forwarded to
`reencode_video`. `None` lets the codec decide.
num_workers (`int | None`, *optional*): Number of parallel processes. `None` or `0` means
sequential (no multiprocessing); `1+` spawns a `ProcessPoolExecutor`.
dataset: An existing :class:`LeRobotDataset` whose videos will be
re-encoded.
rgb_encoder: Target encoder configuration applied to every RGB video
file. If ``None``, re-encoding is skipped for RGB videos.
depth_encoder: Target encoder configuration applied to every depth video
file. If ``None``, re-encoding is skipped for depth videos.
Quantization parameters will not override the ones in the current dataset.
encoder_threads: Per-encoder thread count forwarded to
:func:`reencode_video`. ``None`` lets the codec decide.
num_workers: Number of parallel processes. ``None`` or ``0`` means
sequential (no multiprocessing); ``1+`` spawns a
:class:`~concurrent.futures.ProcessPoolExecutor`.
Returns:
The same :class:`LeRobotDataset` instance with its metadata updated
+2 -1
View File
@@ -200,7 +200,8 @@ class DatasetWriter:
self.image_writer.save_image(image=image, fpath=fpath, compress_level=compress_level)
def add_frame(self, frame: dict) -> None:
"""Add a single frame to the current episode buffer.
"""
Add a single frame to the current episode buffer.
Apart from images written to a temporary directory, nothing is written to disk
until ``save_episode()`` is called.
+19 -41
View File
@@ -13,7 +13,9 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Depth encoding/decoding helpers for :class:`DepthEncoderConfig`."""
"""
Depth encoding/decoding helpers for :class:`DepthEncoderConfig`.
"""
import math
from typing import Literal
@@ -90,24 +92,13 @@ def quantize_depth(
``depth_min``, ``depth_max``, and ``shift`` are always in **metres**.
Args:
depth (`numpy.ndarray[tuple[int, ...], numpy.dtype[numpy.uint16]] | numpy.ndarray[tuple[int, ...], numpy.dtype[numpy.float32]] | torch.Tensor`): Depth
map to quantize. A `torch.Tensor` is moved to CPU before conversion.
depth_min (`float`, *optional*, defaults to 0.01): Depth, in metres, mapped to quantum
`0`.
depth_max (`float`, *optional*, defaults to 10.0): Depth, in metres, mapped to quantum
`DEPTH_QMAX`.
shift (`float`, *optional*, defaults to 3.5): Depth shift, in metres, used in log mode.
Must satisfy `depth_min + shift > 0`.
use_log (`bool`, *optional*, defaults to `True`): If `True`, quantize in log space, which
allocates more quanta to near-range depth.
pix_fmt (`str`, *optional*, defaults to `"gray12le"`): Pixel format used to build the
`av.VideoFrame` when `video_backend="pyav"`.
video_backend (`str | None`, *optional*, defaults to `"pyav"`): Video backend used for
encoding. When `"pyav"`, returns an `av.VideoFrame`; otherwise returns the raw
`uint16` array.
input_unit (`Literal`, *optional*, defaults to `"auto"`): Input unit policy: `"auto"`
infers the unit from `depth`'s dtype, while `"mm"` or `"m"` force millimetres or
metres respectively.
depth: Depth map; ``torch.Tensor`` is moved to CPU for conversion.
depth_min: Depth (metres) at quantum ``0``.
depth_max: Depth (metres) at quantum :data:`DEPTH_QMAX`.
shift: Depth shift (metres); used in log mode. Must satisfy ``depth_min + shift > 0``.
use_log: If ``True`` (default), quantize in log space.
video_backend: Video backend to use for encoding. Defaults to "pyav".
input_unit: Input unit policy (``"auto"``, ``"mm"``, ``"m"``).
Returns:
``numpy.ndarray``, ``dtype=uint16``, same shape as ``depth``, values in
@@ -183,28 +174,15 @@ def dequantize_depth(
Output layout is determined by ``output_channel_last``.
Args:
quantized (`numpy.ndarray[tuple[int, ...], numpy.dtype[numpy.uint16]] | av.video.frame.VideoFrame | torch.Tensor`): 12-bit
codes in `[0, DEPTH_QMAX]`, as a numpy array, `av.VideoFrame`, or `torch.Tensor`
(any integer or float dtype).
depth_min (`float`, *optional*, defaults to 0.01): Depth, in metres, mapped to quantum
`0`. Must match the value passed to `quantize_depth`.
depth_max (`float`, *optional*, defaults to 10.0): Depth, in metres, mapped to quantum
`DEPTH_QMAX`. Must match the value passed to `quantize_depth`.
shift (`float`, *optional*, defaults to 3.5): Depth shift, in metres, used in log mode.
Must match the value passed to `quantize_depth`.
use_log (`bool`, *optional*, defaults to `True`): If `True`, invert the log-space mapping
used by `quantize_depth`. Must match the encoding call.
pix_fmt (`str`, *optional*, defaults to `"gray12le"`): Pixel format used to extract the
plane data when `quantized` is an `av.VideoFrame`.
output_unit (`Literal`, *optional*, defaults to `"mm"`): `"mm"` returns `uint16`
millimetres, clipped to `[0, 65535]`, when returning a numpy array, or `float32`
millimetres when `output_tensor=True`. `"m"` returns `float32` metres in
`[depth_min, depth_max]`.
output_tensor (`bool`, *optional*, defaults to `True`): If `True`, return a
`torch.Tensor` instead of a numpy array.
output_channel_last (`bool`, *optional*, defaults to `False`): If `True`, add the
restored singleton channel dimension as the last axis instead of the third-to-last
axis.
quantized: 12-bit codes in ``[0, DEPTH_QMAX]``. ``np.ndarray``,
``av.VideoFrame``, or ``torch.Tensor`` (any integer or float dtype).
depth_min, depth_max, shift, use_log: Same as :func:`quantize_depth` (metres).
pix_fmt: Pixel format used to extract the plane from an ``av.VideoFrame``.
output_unit: ``"mm"`` returns ``uint16`` millimetres (rint, clip
``[0, 65535]``) when returning a numpy array, or ``float32`` mm when
``output_tensor=True``. ``"m"`` returns ``float32`` metres in
``[depth_min, depth_max]``.
output_tensor: If True, return a ``torch.Tensor`` instead of a numpy array.
Returns:
Depth map in the requested unit and dtype.
+8 -21
View File
@@ -101,10 +101,10 @@ def create_empty_dataset_info(
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, *optional*): The type of robot used, if any.
chunks_size (int | None, *optional*): Max files per chunk directory. Defaults to ``DEFAULT_CHUNK_SIZE``.
data_files_size_in_mb (int | None, *optional*): Max parquet file size in MB. Defaults to ``DEFAULT_DATA_FILE_SIZE_IN_MB``.
video_files_size_in_mb (int | None, *optional*): Max video file size in MB. Defaults to ``DEFAULT_VIDEO_FILE_SIZE_IN_MB``.
robot_type (str | None): The type of robot used, if any.
chunks_size (int | None): Max files per chunk directory. Defaults to ``DEFAULT_CHUNK_SIZE``.
data_files_size_in_mb (int | None): Max parquet file size in MB. Defaults to ``DEFAULT_DATA_FILE_SIZE_IN_MB``.
video_files_size_in_mb (int | None): Max video file size in MB. Defaults to ``DEFAULT_VIDEO_FILE_SIZE_IN_MB``.
Returns:
DatasetInfo: A typed dataset information object with initial metadata.
@@ -170,7 +170,7 @@ def check_delta_timestamps(
deltas in seconds.
fps (int): The frames per second of the dataset.
tolerance_s (float): The allowed tolerance in seconds.
raise_value_error (bool, *optional*, defaults to `True`): If True, raises an error on failure.
raise_value_error (bool): If True, raises an error on failure.
Returns:
bool: True if all deltas are valid, False otherwise.
@@ -219,18 +219,6 @@ def get_delta_indices(delta_timestamps: dict[str, list[float]], fps: int) -> dic
def validate_frame(frame: dict, features: dict) -> None:
"""Check that `frame` has a `"task"` key and matches `features` (minus auto-populated defaults).
Args:
frame (`dict`): The frame to validate, mapping feature names to their values, as passed by the
caller to `add_frame`.
features (`dict`): The dataset's feature specification, mapping feature names to their dtype and
shape metadata.
Raises:
ValueError: If `frame` is missing `"task"`, or has missing/extra features, or a feature's dtype
or shape doesn't match its definition in `features`.
"""
# DEFAULT_FEATURES (timestamp, frame_index, episode_index, index, task_index) are
# auto-populated by the recording pipeline (add_frame / save_episode) and must not
# be supplied by the caller. Excluding them here means any frame dict that contains
@@ -287,7 +275,7 @@ def validate_feature_dtype_and_shape(
Args:
name (str): The name of the feature.
feature (dict): The feature specification from the LeRobot features dictionary.
value (`numpy.ndarray | PIL.Image.Image | str`): The value of the feature to validate.
value: The value of the feature to validate.
Returns:
str: An error message if validation fails, otherwise an empty string.
@@ -349,7 +337,7 @@ def validate_feature_image_or_video(
Args:
name (str): The name of the feature.
expected_shape (list[str]): The expected shape, e.g. (C, H, W) or (H, W, C).
value (`numpy.ndarray | PIL.Image.Image`): The image or video frame data to validate.
value: The image data to validate.
Returns:
str: An error message if validation fails, otherwise an empty string.
@@ -395,8 +383,7 @@ def validate_feature_language(name: str, value) -> str:
Args:
name (str): The name of the feature.
value (`Any`): The value supplied for the language feature. Only checked for being `None`; any
other value is dropped with a warning.
value: The value to validate.
Returns:
str: Always an empty string — language values are non-fatal.
+6 -22
View File
@@ -27,10 +27,7 @@ logger = logging.getLogger(__name__)
def safe_stop_image_writer(func):
"""Decorator: on an exception from `func`, stop the `dataset` kwarg's image writer before re-raising."""
def wrapper(*args, **kwargs):
"""Call `func`; on any exception, stop `kwargs["dataset"].writer.image_writer` before re-raising."""
try:
return func(*args, **kwargs)
except BaseException:
@@ -129,7 +126,8 @@ def save_kwargs_for_path(fpath: Path, compress_level: int) -> dict:
def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1):
"""Saves a NumPy array or PIL Image to a file.
"""
Saves a NumPy array or PIL Image to a file.
This function handles both NumPy arrays and PIL Image objects, converting
the former to a PIL Image before saving. It includes error handling for
@@ -140,7 +138,7 @@ def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level
Args:
image (np.ndarray | PIL.Image.Image): The image data to save.
fpath (Path): The destination file path for the image.
compress_level (int, optional, *optional*, defaults to 1): The compression level for the saved
compress_level (int, optional): The compression level for the saved
image, as used by PIL.Image.save(). Defaults to 1.
Refer to: https://github.com/huggingface/lerobot/pull/2135
for more details on the default value rationale.
@@ -165,7 +163,6 @@ def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path, compress_level
def worker_thread_loop(queue: queue.Queue):
"""Pop `(image_array, fpath, compress_level)` items from `queue` and write each until a `None` sentinel."""
while True:
item = queue.get()
if item is None:
@@ -177,7 +174,6 @@ def worker_thread_loop(queue: queue.Queue):
def worker_process(queue: queue.Queue, num_threads: int):
"""Run `num_threads` `worker_thread_loop` threads against `queue` and block until they all exit."""
threads = []
for _ in range(num_threads):
t = threading.Thread(target=worker_thread_loop, args=(queue,))
@@ -189,9 +185,9 @@ def worker_process(queue: queue.Queue, num_threads: int):
class AsyncImageWriter:
"""This class abstracts away the initialisation of processes or/and threads.
It saves images on disk asynchronously, which is critical to control a robot and record data
"""
This class abstract away the initialisation of processes or/and threads to
save images on disk asynchronously, which is critical to control a robot and record data
at a high frame rate.
When `num_processes=0`, it creates a threads pool of size `num_threads`.
@@ -204,15 +200,6 @@ class AsyncImageWriter:
"""
def __init__(self, num_processes: int = 0, num_threads: int = 1):
"""Start the thread or process pool.
Args:
num_processes: Number of worker subprocesses. `0` uses threads only (in this process).
num_threads: Number of writer threads per process (or in this process, if `num_processes=0`).
Raises:
ValueError: If both `num_threads` and `num_processes` are non-positive.
"""
self.num_processes = num_processes
self.num_threads = num_threads
self.queue = None
@@ -243,18 +230,15 @@ class AsyncImageWriter:
def save_image(
self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
):
"""Enqueue `image` to be written to `fpath` asynchronously; returns immediately."""
if isinstance(image, torch.Tensor):
# Convert tensor to numpy array to minimize main process time
image = image.cpu().numpy()
self.queue.put((image, fpath, compress_level))
def wait_until_done(self):
"""Block until every enqueued image has been written to disk."""
self.queue.join()
def stop(self):
"""Signal all worker threads/processes to exit and wait for them to join. No-op if already stopped."""
if self._stopped:
return
+10 -21
View File
@@ -46,7 +46,6 @@ from .utils import (
def get_parquet_file_size_in_mb(parquet_path: str | Path) -> float:
"""Return the uncompressed size, in megabytes, of a parquet file's column data (from its metadata)."""
metadata = pq.read_metadata(parquet_path)
total_uncompressed_size = 0
for row_group in range(metadata.num_row_groups):
@@ -58,24 +57,20 @@ def get_parquet_file_size_in_mb(parquet_path: str | Path) -> float:
def get_hf_dataset_size_in_mb(hf_ds: Dataset) -> int:
"""Return the in-memory (Arrow buffer) size of a Hugging Face `Dataset`, in megabytes."""
return hf_ds.data.nbytes // (1024**2)
def load_nested_dataset(
pq_dir: Path, features: datasets.Features | None = None, episodes: list[int] | None = None
) -> Dataset:
"""Find parquet files in provided directory {pq_dir}/chunk-xxx/file-xxx.parquet.
Convert parquet files to pyarrow memory mapped in a cache folder for efficient RAM usage, then
concatenate all pyarrow references to return HF Dataset format.
"""Find parquet files in provided directory {pq_dir}/chunk-xxx/file-xxx.parquet
Convert parquet files to pyarrow memory mapped in a cache folder for efficient RAM usage
Concatenate all pyarrow references to return HF Dataset format
Args:
pq_dir (`Path`): Directory containing parquet files.
features (`datasets.features.features.Features | None`, *optional*): Features schema used to ensure
consistent loading of complex types like images.
episodes (`list[int] | None`, *optional*): List of episode indices to filter. Uses PyArrow
predicate pushdown for efficiency.
pq_dir: Directory containing parquet files
features: Optional features schema to ensure consistent loading of complex types like images
episodes: Optional list of episode indices to filter. Uses PyArrow predicate pushdown for efficiency.
"""
paths = sorted(pq_dir.glob("*/*.parquet"))
if len(paths) == 0:
@@ -88,7 +83,6 @@ def load_nested_dataset(
def get_parquet_num_frames(parquet_path: str | Path) -> int:
"""Return the number of rows in a parquet file, read from its metadata (no data is loaded)."""
metadata = pq.read_metadata(parquet_path)
return metadata.num_rows
@@ -124,7 +118,6 @@ def embed_images(dataset: datasets.Dataset) -> datasets.Dataset:
def write_info(info: DatasetInfo, local_dir: Path) -> None:
"""Write dataset info metadata to its standard file path (the inverse of `load_info`)."""
write_json(info.to_dict(), local_dir / INFO_PATH)
@@ -183,14 +176,12 @@ def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]] | None:
def write_tasks(tasks: pandas.DataFrame, local_dir: Path) -> None:
"""Write the task-prompt table to its standard parquet file path (the inverse of `load_tasks`)."""
path = local_dir / DEFAULT_TASKS_PATH
path.parent.mkdir(parents=True, exist_ok=True)
tasks.to_parquet(path)
def load_tasks(local_dir: Path) -> pandas.DataFrame:
"""Load the task-prompt table from its standard file path, indexed by task string."""
tasks = pd.read_parquet(local_dir / DEFAULT_TASKS_PATH)
tasks.index.name = "task"
return tasks
@@ -198,13 +189,12 @@ def load_tasks(local_dir: Path) -> pandas.DataFrame:
def write_episodes(episodes: Dataset, local_dir: Path) -> None:
"""Write episode metadata to a parquet file in the LeRobot v3.0 format.
This function writes episode-level metadata to a single parquet file.
Used primarily during dataset conversion (v2.1 → v3.0) and in test fixtures.
Args:
episodes (`Dataset`): Hugging Face `Dataset` containing the episode metadata.
local_dir (`Path`): Root directory where the dataset is stored.
episodes: HuggingFace Dataset containing episode metadata
local_dir: Root directory where the dataset will be stored
"""
episode_size_mb = get_hf_dataset_size_in_mb(episodes)
if episode_size_mb > DEFAULT_DATA_FILE_SIZE_IN_MB:
@@ -220,7 +210,6 @@ def write_episodes(episodes: Dataset, local_dir: Path) -> None:
def load_episodes(local_dir: Path) -> datasets.Dataset:
"""Load episode metadata, excluding per-episode `stats/*` columns (for faster access to the rest)."""
episodes = load_nested_dataset(local_dir / EPISODES_DIR)
# Select episode features/columns containing references to episode data and videos
# (e.g. tasks, dataset_from_index, dataset_to_index, data/chunk_index, data/file_index, etc.)
@@ -236,9 +225,9 @@ def load_image_as_numpy(
Args:
fpath (str | Path): Path to the image file.
dtype (np.dtype, *optional*, defaults to `float32`): The desired data type of the output array. If floating,
dtype (np.dtype): The desired data type of the output array. If floating,
pixels are scaled to [0, 1]. Only used for RGB images.
channel_first (bool, *optional*, defaults to `True`): If True, converts the image to (C, H, W) format.
channel_first (bool): If True, converts the image to (C, H, W) format.
Otherwise, it remains in (H, W, C) format.
Returns:
+4 -23
View File
@@ -44,14 +44,6 @@ logger = logging.getLogger(__name__)
class LeRobotDataset(torch.utils.data.Dataset):
"""A PyTorch `Dataset` over episodic robot data: per-frame state/action tensors, optional videos.
Backed by parquet files (`data/`) for tabular observation/action/reward data, optional video files
(`videos/`) for image observations, and a `meta/` directory holding `info.json` (shapes, keys, fps),
`stats.json` (normalization statistics), and per-episode metadata. See `__init__`'s docstring for the
on-disk layout, and `create()` for building a new (empty) dataset from scratch.
"""
def __init__(
self,
repo_id: str,
@@ -76,7 +68,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
*,
token: str | bool | None = None,
):
"""2 modes are available for instantiating this class, depending on 2 different use cases.
"""
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
@@ -175,9 +168,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
conversion. This works for both image-backed and video-backed observations and can later be
updated with `set_image_transforms()` or cleared with `clear_image_transforms()`.
Defaults to None.
delta_timestamps (dict[list[float]] | None, optional): Per-feature timestamp offsets (in
seconds, relative to a frame's own timestamp) of additional frames to return alongside it.
Defaults to None.
delta_timestamps (dict[list[float]] | None, optional): _description_. 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
@@ -194,11 +185,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
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.
return_uint8 (bool, optional): For RGB videos, whether to return raw uint8 frames instead of
the default float32 frames normalized to [0, 1]. Defaults to False.
depth_output_unit (str, optional): Physical unit depth maps are dequantized to at load time:
"mm" (millimeters) or "m" (metres). Has no effect on datasets without depth cameras.
Defaults to "mm".
batch_encoding_size (int, optional): Number of episodes to accumulate before batch encoding videos.
Set to 1 for immediate encoding (default), or higher for batched encoding. Defaults to 1.
rgb_encoder (RGBEncoderConfig | None, optional): Video encoder settings for cameras
@@ -409,7 +395,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
@property
def hf_dataset(self) -> datasets.Dataset:
"""The underlying Hugging Face Dataset object."""
"""The underlying Hugging Face Dataset object"""
self.reader = self._ensure_reader()
if self.reader.hf_dataset is None:
self.reader.load_and_activate()
@@ -554,7 +540,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
return self.hf_dataset[idx]
def __repr__(self):
"""A short summary: repo ID, selected episode/sample counts, and feature keys."""
feature_keys = list(self.features)
return (
f"{self.__class__.__name__}({{\n"
@@ -753,10 +738,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
during capture instead of writing images first.
encoder_queue_maxsize: Max buffered frames per camera when using
streaming encoding.
video_files_size_in_mb: Max video file size in MB. Defaults to
``DEFAULT_VIDEO_FILE_SIZE_IN_MB``.
data_files_size_in_mb: Max parquet file size in MB. Defaults to
``DEFAULT_DATA_FILE_SIZE_IN_MB``.
Returns:
A new :class:`LeRobotDataset` in write mode.
+3 -28
View File
@@ -51,20 +51,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
*,
token: str | bool | None = None,
):
"""Construct a `LeRobotDataset` for each `repo_id` and concatenate them.
Args:
repo_ids: The Hub repo IDs (or local dataset names, if `root` is set) to load.
root: Root directory containing the underlying datasets. Defaults to `$HF_LEROBOT_HOME`.
episodes: Optional mapping from `repo_id` to the episode indices to load from it.
image_transforms: Transform applied to visual observations in each underlying dataset.
delta_timestamps: Passed through to each underlying `LeRobotDataset`.
tolerances_s: Optional mapping from `repo_id` to its timestamp tolerance, in seconds. Defaults
to `1e-4` for every dataset.
download_videos: Whether to download video files for each underlying dataset.
video_backend: The video decoding backend to use.
token: Hugging Face Hub authentication token.
"""
super().__init__()
self.repo_ids = repo_ids
self.root = Path(root) if root else HF_LEROBOT_HOME
@@ -154,7 +140,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
@property
def features(self) -> datasets.Features:
"""The union of all underlying datasets' features (minus `disabled_features`)."""
features = {}
for dataset in self._datasets:
features.update(
@@ -201,26 +186,17 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
@property
def tolerance_s(self) -> float:
"""Tolerance in seconds used to discard loaded frames when their timestamps aren't close enough.
Only used when `delta_timestamps` is provided or when loading video frames from mp4 files.
"""Tolerance in seconds used to discard loaded frames when their timestamps
are not close enough from the requested frames. It is only used when `delta_timestamps`
is provided or when loading video frames from mp4 files.
"""
# 1e-4 to account for possible numerical error
return 1 / self.fps - 1e-4
def __len__(self):
"""The total number of frames across all underlying datasets."""
return self.num_frames
def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
"""Return the frame at `idx`, resolved to the underlying dataset it falls in.
Adds a `"dataset_index"` key identifying which underlying dataset the frame came from, and drops
any `disabled_features` keys.
Raises:
IndexError: If `idx` is out of bounds.
"""
if idx >= len(self):
raise IndexError(f"Index {idx} out of bounds.")
# Determine which dataset to get an item from based on the index.
@@ -243,7 +219,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
return item
def __repr__(self):
"""A summary: repo IDs, sample/episode counts, media type, fps, camera keys, and transforms."""
return (
f"{self.__class__.__name__}(\n"
f" Repository IDs: '{self.repo_ids}',\n"
+12 -17
View File
@@ -26,13 +26,12 @@ from lerobot.utils.feature_utils import hw_to_dataset_features
def create_initial_features(
action: RobotAction | None = None, observation: RobotObservation | None = None
) -> dict[PipelineFeatureType, dict[str, Any]]:
"""Creates the initial features dict for the dataset from action and observation specs.
"""
Creates the initial features dict for the dataset from action and observation specs.
Args:
action (`dict[str, typing.Any] | None`, *optional*): A dictionary of action feature names to their
types/shapes.
observation (`dict[str, typing.Any] | None`, *optional*): A dictionary of observation feature names
to their types/shapes.
action: A dictionary of action feature names to their types/shapes.
observation: A dictionary of observation feature names to their types/shapes.
Returns:
The initial features dictionary structured by PipelineFeatureType.
@@ -47,14 +46,12 @@ def create_initial_features(
# Helper to filter state/action keys based on compiled regex patterns.
def should_keep(key: str, patterns: tuple[re.Pattern] | None) -> bool:
"""Return `True` if `patterns` is `None` or any pattern in it matches `key`."""
if patterns is None:
return True
return any(pat.search(key) for pat in patterns)
def strip_prefix(key: str, prefixes_to_strip: tuple[str]) -> str:
"""Remove the first prefix in `prefixes_to_strip` that `key` starts with, if any."""
for prefix in prefixes_to_strip:
if key.startswith(prefix):
return key[len(prefix) :]
@@ -76,22 +73,20 @@ def aggregate_pipeline_dataset_features(
exclude_images: bool = False,
patterns: Sequence[str] | None = None,
) -> dict[str, dict]:
"""Aggregates and filters pipeline features to create a dataset-ready features dictionary.
"""
Aggregates and filters pipeline features to create a dataset-ready features dictionary.
This function transforms initial features using the pipeline, categorizes them as action or observations
(image or state), filters them based on `exclude_images` and `patterns`, and finally
formats them for use with a Hugging Face LeRobot Dataset.
Args:
pipeline (`DataProcessorPipeline`): The processor pipeline to apply to `initial_features`.
initial_features (`dict`): A dictionary of raw feature specs for actions and observations, keyed by
`PipelineFeatureType`.
use_videos (`bool`, *optional*, defaults to `True`): Controls the storage dtype for image features.
If `True`, images are stored as `"video"`; if `False`, they are stored as `"image"`.
exclude_images (`bool`, *optional*, defaults to `False`): If `True`, image features are dropped
entirely from the output.
patterns (`collections.abc.Sequence[str] | None`, *optional*): A sequence of regex patterns used to
filter action and state features.
pipeline: The DataProcessorPipeline to apply.
initial_features: A dictionary of raw feature specs for actions and observations.
use_videos: Controls the storage dtype for image features. If True, images are stored as "video"; if False, they are stored as "image".
exclude_images: If True, image features are dropped entirely from the output.
patterns: A sequence of regex patterns to filter action and state features.
Image features are not affected by this filter.
Returns:
A dictionary of features formatted for a Hugging Face LeRobot Dataset.
+4 -5
View File
@@ -41,11 +41,10 @@ def write_u16_plane(plane: av.video.plane.VideoPlane, src: np.ndarray, fill_valu
leave the padding untouched.
Args:
plane (`VideoPlane`): Destination 16-bit plane to copy into.
src (`ndarray`): Source image, shape `(height, width)`, dtype `uint16`.
fill_value (`int | None`, *optional*): If given, every pixel of the plane
(including the row padding) is set to this value first, so the padding
holds clean data instead of garbage.
plane: Destination 16-bit plane.
src: Source image, shape ``(height, width)``, dtype ``uint16``.
fill_value: If given, every pixel (padding included) is set to this first, so the
padding holds clean data instead of garbage.
"""
height, width = src.shape
stride_u16 = plane.line_size // np.dtype(np.uint16).itemsize
+1 -17
View File
@@ -55,8 +55,7 @@ class EpisodeAwareSampler:
seed: int = 0,
absolute_to_relative_idx: dict[int, int] | None = None,
):
"""Build the sampler from per-episode `[from, to)` frame-index boundaries.
"""
Args:
dataset_from_indices: Start index of each episode in the dataset.
dataset_to_indices: End index of each episode in the dataset.
@@ -65,13 +64,6 @@ class EpisodeAwareSampler:
drop_n_last_frames: Frames to drop from the end of each episode.
shuffle: Whether to shuffle the indices.
seed: Seed the permutation is derived from (together with the epoch).
absolute_to_relative_idx: Optional mapping from absolute dataset frame index to the relative
index actually yielded (e.g. when the sampler is used over a filtered subset).
Raises:
ValueError: If `drop_n_first_frames`/`drop_n_last_frames` is negative, if
`dataset_from_indices`/`dataset_to_indices` have different lengths, or if no episode has
any frames remaining after dropping.
"""
if drop_n_first_frames < 0:
raise ValueError(f"drop_n_first_frames must be >= 0, got {drop_n_first_frames}")
@@ -124,15 +116,12 @@ class EpisodeAwareSampler:
return [self._frame_index(k) for k in range(self._num_frames)]
def set_epoch(self, epoch: int) -> None:
"""Set the epoch the next `__iter__` call will use, without consuming an auto-advance."""
self._epoch = epoch
def state_dict(self) -> dict:
"""Return `{"epoch": ..., "start_index": ...}`, enough to resume mid-epoch sample-exactly."""
return {"epoch": self._epoch, "start_index": self._start_index}
def load_state_dict(self, state: dict) -> None:
"""Restore the epoch and within-epoch offset from a `state_dict()`-produced dict."""
self._epoch = state["epoch"]
self._start_index = state["start_index"]
@@ -151,10 +140,6 @@ class EpisodeAwareSampler:
return absolute_idx
def __iter__(self) -> Iterator[int]:
"""Yield frame indices for the current epoch (from `set_epoch`/`load_state_dict`), then advance it.
Shuffled if `self.shuffle`, using a permutation seeded from `(seed, epoch)`.
"""
# Advance epoch state eagerly, not on first consumption of the generator.
epoch, start = self._epoch, self._start_index
self._epoch += 1
@@ -171,7 +156,6 @@ class EpisodeAwareSampler:
yield self._frame_index(k)
def __len__(self) -> int:
"""The total number of frames across the sampled episodes (full length, even mid-resume)."""
return self._num_frames
+33 -49
View File
@@ -44,13 +44,17 @@ from .video_utils import (
class LookBackError(Exception):
"""Exception raised when trying to look back in the history of a Backtrackable object."""
"""
Exception raised when trying to look back in the history of a Backtrackable object.
"""
pass
class LookAheadError(Exception):
"""Exception raised when trying to look ahead in the future of a Backtrackable object."""
"""
Exception raised when trying to look ahead in the future of a Backtrackable object.
"""
pass
@@ -60,10 +64,11 @@ class _ShardExhaustedError(Exception):
class Backtrackable[T]:
"""Wrap any iterator/iterable so you can step back up to `history` items and look ahead.
"""
Wrap any iterator/iterable so you can step back up to `history` items
and look ahead up to `lookahead` items.
Looking ahead is bounded by `lookahead` items. This is useful for streaming datasets where you need
to access previous and future items
This is useful for streaming datasets where you need to access previous and future items
but can't load the entire dataset into memory.
Example:
@@ -93,16 +98,6 @@ class Backtrackable[T]:
__slots__ = ("_source", "_back_buf", "_ahead_buf", "_cursor", "_history", "_lookahead")
def __init__(self, iterable: Iterable[T], *, history: int = 1, lookahead: int = 0):
"""Wrap `iterable`, buffering up to `history` past items and `lookahead` future items.
Args:
iterable: The iterable to wrap.
history: How many past items `prev()`/`peek_back()` can reach. Must be `>= 1`.
lookahead: How many future items `peek_ahead()` can reach. Must be `> 0`.
Raises:
ValueError: If `history < 1` or `lookahead <= 0`.
"""
if history < 1:
raise ValueError("history must be >= 1")
if lookahead <= 0:
@@ -116,11 +111,9 @@ class Backtrackable[T]:
self._lookahead = lookahead
def __iter__(self) -> "Backtrackable[T]":
"""Return `self`; `Backtrackable` is its own iterator."""
return self
def __next__(self) -> T:
"""Return the next item, consuming from the back buffer first if `prev()` stepped back."""
# If we've stepped back, consume from back buffer first
if self._cursor < 0: # -1 means "last item", etc.
self._cursor += 1
@@ -135,9 +128,9 @@ class Backtrackable[T]:
return item
def prev(self) -> T:
"""Step one item back in history and return it.
Raises `LookBackError` if already at the oldest buffered item.
"""
Step one item back in history and return it.
Raises IndexError if already at the oldest buffered item.
"""
if len(self._back_buf) + self._cursor <= 1:
raise LookBackError("At start of history")
@@ -146,15 +139,17 @@ class Backtrackable[T]:
return self._back_buf[self._cursor]
def peek_back(self, n: int = 1) -> T:
"""Look `n` items back (n=1 == previous item) without moving the cursor."""
"""
Look `n` items back (n=1 == previous item) without moving the cursor.
"""
if n < 0 or n + 1 > len(self._back_buf) + self._cursor:
raise LookBackError("peek_back distance out of range")
return self._back_buf[self._cursor - (n + 1)]
def peek_ahead(self, n: int = 1) -> T:
"""Look `n` items ahead (n=1 == next item) without moving the cursor.
"""
Look `n` items ahead (n=1 == next item) without moving the cursor.
Fills the ahead buffer if necessary.
"""
if n < 1:
@@ -174,9 +169,9 @@ class Backtrackable[T]:
return self._ahead_buf[n - 1]
def history(self) -> list[T]:
"""Return a copy of the buffered history (most recent last).
The list length is at most the `history` argument passed at construction.
"""
Return a copy of the buffered history (most recent last).
The list length `history` argument passed at construction.
"""
if self._cursor == 0:
return list(self._back_buf)
@@ -185,12 +180,14 @@ class Backtrackable[T]:
return list(self._back_buf)[: self._cursor or None]
def can_peek_back(self, steps: int = 1) -> bool:
"""Check if we can go back `steps` items without raising a `LookBackError`."""
"""
Check if we can go back `steps` items without raising an IndexError.
"""
return steps < len(self._back_buf) + self._cursor
def can_peek_ahead(self, steps: int = 1) -> bool:
"""Check if we can peek ahead `steps` items.
"""
Check if we can peek ahead `steps` items.
This may involve trying to fill the ahead buffer.
"""
if self._lookahead > 0 and steps > self._lookahead:
@@ -278,8 +275,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
episodes (list[int] | None, optional): If specified, this will only load episodes specified by
their episode_index in this list.
image_transforms (Callable | None, optional): Transform to apply to image data.
delta_timestamps (dict[list[float]] | None, optional): Per-feature timestamp offsets (in
seconds, relative to a frame's own timestamp) of additional frames to return alongside it.
tolerance_s (float, optional): Tolerance in seconds for timestamp matching.
revision (str, optional): Git revision id (branch name, tag, or commit hash).
force_cache_sync (bool, optional): Flag to sync and refresh local files first.
@@ -289,8 +284,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
seed (int, optional): Reproducibility random seed.
rng (np.random.Generator | None, optional): Random number generator.
shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
return_uint8 (bool, optional): For RGB videos, whether to return raw uint8 frames instead of
the default float32 frames normalized to [0, 1].
depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm").
Defaults to "mm".
repo_type: "dataset" (default) or "bucket" to stream from an HF Storage Bucket
@@ -390,17 +383,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
@property
def num_frames(self):
"""The total number of frames in the dataset."""
return self.meta.total_frames
@property
def num_episodes(self):
"""The total number of episodes in the dataset."""
return self.meta.total_episodes
@property
def fps(self):
"""The dataset's recording frame rate."""
return self.meta.fps
@property
@@ -425,11 +415,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
# could be used with a ThreadPoolExecutor to run `make_frame` (especially video decoding)
# in parallel, feeding a queue from which this iterator will yield processed items.
def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
"""Yield frames via reservoir-buffered random sampling across shards, streaming indefinitely.
Samples a random shard, then a random frame from a fixed-size buffer refilled from that shard, so
no full shuffle or shard is ever fully materialized in memory.
"""
if self.video_decoder_cache is None:
self.video_decoder_cache = VideoDecoderCache()
@@ -507,7 +492,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
return dict.fromkeys(self.meta.video_keys, [start_ts])
def _make_padding_camera_frame(self, camera_key: str):
"""Variable-shape padding frame for the given camera key, shaped (H, W, C)."""
"""Variable-shape padding frame for given camera keys, given in (H, W, C)"""
return torch.zeros(self.meta.info.features[camera_key]["shape"]).permute(-1, 0, 1)
def _get_video_frame_padding_mask(
@@ -538,7 +523,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
return padding_mask
def make_frame(self, dataset_iterator: Backtrackable) -> Generator:
"""Makes a frame starting from a dataset iterator."""
"""Makes a frame starting from a dataset iterator"""
try:
item = next(dataset_iterator)
except StopIteration as e:
@@ -631,13 +616,12 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
return query_timestamps
def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict:
"""Decode the requested per-camera frame timestamps from `ep_idx`'s videos.
Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
"""Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
in the main process (e.g. by using a second Dataloader with num_workers=0). It will result in a
Segmentation Fault. This probably happens because a memory reference to the video loader is created in
the main process and a subprocess fails to access it.
"""
item = {}
for video_key, query_ts in query_timestamps.items():
root = self.meta.url_root if self.streaming and not self.streaming_from_local else self.root
@@ -680,9 +664,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
"""Get frames with delta offsets using the backtrackable iterator.
Args:
dataset_iterator (Backtrackable): The backtrackable iterator to peek/step through for delta
frames.
current_item (dict): Current item from the iterator.
ep_idx (int): Episode index.
Returns:
tuple: (query_result, padding) - frames at delta offsets and padding info.
@@ -787,7 +770,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
return query_result, padding
def _validate_delta_timestamp_keys(self, delta_timestamps: dict[list[float]]) -> None:
"""Validate that all keys in delta_timestamps correspond to actual features in the dataset.
"""
Validate that all keys in delta_timestamps correspond to actual features in the dataset.
Raises:
ValueError: If any delta timestamp key doesn't correspond to a dataset feature.
+13 -41
View File
@@ -63,21 +63,11 @@ hub_api.create_tag("{repo_id}", tag="_version_", repo_type="dataset")
"""
class CompatibilityError(Exception):
"""Base class for errors raised when a dataset's `codebase_version` doesn't match this install."""
...
class CompatibilityError(Exception): ...
class BackwardCompatibilityError(CompatibilityError):
"""Raised when a dataset was saved with an older, unsupported `codebase_version`."""
def __init__(self, repo_id: str, version: packaging.version.Version):
"""Build the error message pointing the user at the v2.1-to-v3.0 conversion script.
Raises:
NotImplementedError: If `version` isn't the one supported legacy version (2.1).
"""
if version.major == 2 and version.minor == 1:
message = V30_MESSAGE.format(repo_id=repo_id, version=version)
else:
@@ -88,10 +78,7 @@ class BackwardCompatibilityError(CompatibilityError):
class ForwardCompatibilityError(CompatibilityError):
"""Raised when a dataset was saved with a newer `codebase_version` than this install supports."""
def __init__(self, repo_id: str, version: packaging.version.Version):
"""Build the error message pointing the user at upgrading their `lerobot` install."""
message = FUTURE_MESSAGE.format(repo_id=repo_id, version=version)
super().__init__(message)
@@ -202,12 +189,6 @@ class DatasetInfo:
tools: list[dict] | None = None
def __post_init__(self) -> None:
"""Coerce feature shapes from list to tuple, and validate `fps`/`chunks_size`/file-size fields.
Raises:
ValueError: If `fps`, `chunks_size`, `data_files_size_in_mb`, or `video_files_size_in_mb` isn't
positive.
"""
# Coerce feature shapes from list to tuple — JSON deserialisation
# returns lists, but the rest of the codebase expects tuples.
for ft in self.features.values():
@@ -258,11 +239,6 @@ class DatasetInfo:
# Once all callers have been migrated to attribute access, remove these.
# ---------------------------------------------------------------------------
def __getitem__(self, key: str):
"""Deprecated dict-style read; use attribute access instead.
Raises:
KeyError: If `key` isn't a field on this class.
"""
import warnings
warnings.warn(
@@ -277,7 +253,6 @@ class DatasetInfo:
raise KeyError(key) from err
def __setitem__(self, key: str, value) -> None:
"""Deprecated dict-style write; use attribute assignment instead."""
import warnings
warnings.warn(
@@ -315,7 +290,6 @@ def has_legacy_hub_download_metadata(root: Path) -> bool:
def update_chunk_file_indices(chunk_idx: int, file_idx: int, chunks_size: int) -> tuple[int, int]:
"""Advance to the next `(chunk_idx, file_idx)`, rolling over to a new chunk once `chunks_size` is hit."""
if file_idx == chunks_size - 1:
file_idx = 0
chunk_idx += 1
@@ -381,7 +355,7 @@ def check_version_compatibility(
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, *optional*, defaults to `True`): If True, raise an error on major version mismatch.
enforce_breaking_major (bool): If True, raise an error on major version mismatch.
Raises:
BackwardCompatibilityError: If the dataset version is from a newer, incompatible
@@ -408,9 +382,9 @@ def get_repo_versions(repo_id: str, *, token: str | bool | None = None) -> list[
Args:
repo_id (str): The repository ID on the Hugging Face Hub.
token (`str | bool | None`, *optional*): Authentication token used for Hub requests. Pass a string
token, `True` to require the locally stored token, `False` to disable authentication, or `None`
to use the Hugging Face Hub default.
token: Authentication token used for Hub requests. Pass a string token,
``True`` to require the locally stored token, ``False`` to disable
authentication, or ``None`` to use the Hugging Face Hub default.
Returns:
list[packaging.version.Version]: A list of valid versions found.
@@ -440,7 +414,7 @@ def get_safe_version(
Args:
repo_id (str): The repository ID on the Hugging Face Hub.
version (str | packaging.version.Version): The target version.
token (`str | bool | None`, *optional*): Authentication token forwarded to the Hub version lookup.
token: Authentication token forwarded to the Hub version lookup.
Returns:
str: The safe version string (e.g., "v1.2.3") to use as a revision.
@@ -487,7 +461,7 @@ def create_branch(repo_id: str, *, branch: str, repo_type: str | None = None) ->
Args:
repo_id (str): The ID of the repository.
branch (str): The name of the branch to create.
repo_type (str | None, *optional*): The type of the repository (e.g., "dataset").
repo_type (str | None): The type of the repository (e.g., "dataset").
"""
api = HfApi()
@@ -512,12 +486,10 @@ def create_lerobot_dataset_card(
https://huggingface.co/docs/hub/repositories-licenses.
Args:
tags (list | None, *optional*): A list of tags to add to the dataset card.
dataset_info (DatasetInfo | None, *optional*): The dataset's info object, which will
tags (list | None): A list of tags to add to the dataset card.
dataset_info (DatasetInfo | None): The dataset's info object, which will
be displayed on the card.
kwargs (`Any`, *optional*): Values used to replace placeholders in the card template, e.g. `license`, which
must be a valid license identifier from
https://huggingface.co/docs/hub/repositories-licenses.
**kwargs: Additional keyword arguments to populate the card template.
Returns:
DatasetCard: The generated dataset card object.
@@ -552,12 +524,10 @@ def create_lerobot_dataset_card(
def is_float_in_list(target, float_list, threshold=1e-6):
"""Return `True` if `float_list` contains a value within `threshold` of `target`."""
return any(abs(target - x) <= threshold for x in float_list)
def find_float_index(target, float_list, threshold=1e-6):
"""Return the index of the first value in `float_list` within `threshold` of `target`, or -1."""
for i, x in enumerate(float_list):
if abs(target - x) <= threshold:
return i
@@ -565,7 +535,9 @@ def find_float_index(target, float_list, threshold=1e-6):
def safe_shard(dataset: datasets.IterableDataset, index: int, num_shards: int) -> datasets.Dataset:
"""Safe shards the dataset."""
"""
Safe shards the dataset.
"""
shard_idx = min(dataset.num_shards, index + 1) - 1
return dataset.shard(num_shards, index=shard_idx)
+70 -124
View File
@@ -61,18 +61,19 @@ def decode_video_frames(
return_uint8: bool = False,
is_depth: bool = False,
) -> torch.Tensor:
"""Decodes video frames using the specified backend.
"""
Decodes video frames using the specified backend.
Args:
video_path (Path): Path to the video file.
timestamps (list[float]): List of timestamps to extract frames.
tolerance_s (float): Allowed deviation in seconds for frame retrieval.
backend (str, optional, *optional*): Backend to use for decoding. Defaults to "torchcodec" when available
backend (str, optional): Backend to use for decoding. Defaults to "torchcodec" when available
in the platform; otherwise, defaults to "pyav". The legacy value "video_reader" is
accepted for one release as an alias for "pyav" and will be removed in a future version.
return_uint8 (bool, *optional*, defaults to `False`): For RGB videos, if True return raw uint8 frames without float32 normalization.
return_uint8 (bool): For RGB videos, if True return raw uint8 frames without float32 normalization.
This reduces memory for DataLoader IPC; normalization can be done on GPU afterward.
is_depth (bool, *optional*, defaults to `False`): Set to True if the video is a depth map (1 channel, uint12).
is_depth (bool): Set to True if the video is a depth map (1 channel, uint12).
Returns:
torch.Tensor: Decoded frames (RGB: float32 in [0,1] by default, or uint8 if return_uint8=True, Depth: uint12).
@@ -123,16 +124,14 @@ def decode_video_frames_pyav(
video can be adjusted at encoding time to trade off decoding speed against file size.
Args:
video_path (`pathlib.Path | str`): Path to the video file.
timestamps (`list`): List of timestamps, in seconds, to extract frames for.
tolerance_s (`float`): Allowed deviation in seconds between a queried timestamp
and the closest decoded frame.
log_loaded_timestamps (`bool`, *optional*, defaults to `False`): Whether to log
every decoded frame's timestamp at INFO level.
return_uint8 (`bool`, *optional*, defaults to `False`): For RGB videos, whether to
return raw uint8 frames instead of the default float32 frames normalized to [0, 1].
is_depth (`bool`, *optional*, defaults to `False`): Whether the video is a depth map
(1 channel, uint12).
video_path: Path to the video file.
timestamps: List of timestamps (in seconds) to extract frames for.
tolerance_s: Allowed deviation in seconds between a queried timestamp and the closest
decoded frame.
log_loaded_timestamps: When True, log every decoded frame's timestamp at INFO level.
return_uint8: For RGB videos, if True return raw uint8 frames (C, H, W).
Otherwise, return float32 in [0, 1] range.
is_depth: Set to True if the video is a depth map (1 channel, uint12).
Returns:
torch.Tensor of shape (len(timestamps), C, H, W).
@@ -266,31 +265,15 @@ class VideoDecoderCache:
ever opened until the process exits).
Args:
max_size (`int | None | object`, *optional*, defaults to `<unset>`): Maximum
number of decoders to retain. `None` disables eviction and restores legacy unbounded
behaviour. The sentinel default defers to the value of `LEROBOT_VIDEO_DECODER_CACHE_SIZE`
if set, otherwise `DEFAULT_DECODER_CACHE_SIZE`.
max_size: Maximum number of decoders to retain. ``None`` disables
eviction and restores legacy unbounded behaviour. Defaults to the
value of ``LEROBOT_VIDEO_DECODER_CACHE_SIZE`` if set, otherwise
:data:`DEFAULT_DECODER_CACHE_SIZE`.
"""
class _UnsetSentinel:
"""Singleton marker distinguishing "not passed" from an explicit `None` `max_size`.
Has a fixed `__repr__` (unlike a bare `object()`) so it renders identically across
processes, which keeps the class docstring's `defaults to` clause stable.
"""
def __repr__(self) -> str:
"""Return `"<unset>"`, a stable placeholder for docstrings/logging."""
return "<unset>"
_SENTINEL: ClassVar[object] = _UnsetSentinel()
_SENTINEL: ClassVar[object] = object()
def __init__(self, max_size: int | None | object = _SENTINEL):
"""Create the cache. See the class docstring for `max_size`.
Raises:
ValueError: If `max_size` is neither `None` nor a positive integer.
"""
if max_size is VideoDecoderCache._SENTINEL:
max_size = _default_max_cache_size()
if max_size is not None and max_size <= 0:
@@ -300,7 +283,6 @@ class VideoDecoderCache:
self._lock = Lock()
def __contains__(self, video_path: object) -> bool:
"""Return `True` if `video_path` (as `str`) has a cached decoder."""
with self._lock:
return str(video_path) in self._cache
@@ -356,7 +338,7 @@ class VideoDecoderCache:
class FrameTimestampError(ValueError):
"""Helper error to indicate the retrieved timestamps exceed the queried ones."""
"""Helper error to indicate the retrieved timestamps exceed the queried ones"""
pass
@@ -375,16 +357,11 @@ def decode_video_frames_torchcodec(
"""Loads frames associated with the requested timestamps of a video using torchcodec.
Args:
video_path (`pathlib.Path | str`): Path to the video file.
timestamps (`list`): List of timestamps, in seconds, to extract frames for.
tolerance_s (`float`): Allowed deviation in seconds between a queried timestamp
and the closest decoded frame.
log_loaded_timestamps (`bool`, *optional*, defaults to `False`): Whether to log
every decoded frame's timestamp at INFO level.
decoder_cache (`lerobot.datasets.video_utils.VideoDecoderCache | None`, *optional*): Decoder
cache to fetch the `VideoDecoder` from. Uses the module-level default cache if `None`.
return_uint8 (`bool`, *optional*, defaults to `False`): For RGB videos, whether to
return raw uint8 frames instead of the default float32 frames normalized to [0, 1].
video_path: Path to the video file.
timestamps: List of timestamps to extract frames.
tolerance_s: Allowed deviation in seconds for frame retrieval.
log_loaded_timestamps: Whether to log loaded timestamps.
decoder_cache: Optional decoder cache instance. Uses default if None.
Note: Setting device="cuda" outside the main process, e.g. in data loader workers, will lead to CUDA initialization errors.
@@ -474,19 +451,19 @@ def encode_video_frames(
RGB frames are encoded directly.
Args:
imgs_dir (`pathlib.Path | str`): Directory containing the frames to encode, named
`frame-000000` onwards (`.png` for RGB, `.tiff` for depth).
video_path (`pathlib.Path | str`): Output path for the encoded `.mp4` file.
fps (`int`): Frame rate of the output video.
video_encoder (`lerobot.configs.video.VideoEncoderConfig | None`, *optional*): Encoder settings
(codec, pixel format, quality, ...). When `None`, `rgb_encoder_defaults` is used. Pass a
`DepthEncoderConfig` to encode depth frames.
encoder_threads (`int | None`, *optional*): Per-encoder thread count forwarded to the codec.
`None` lets the codec decide.
log_level (`int | None`, *optional*, defaults to 24): libav log level to set while encoding,
or `None` to leave the current logging configuration unchanged.
overwrite (`bool`, *optional*, defaults to `False`): When `False` and `video_path` already
exists, skip encoding and log a warning. When `True`, re-encode and replace the existing file.
imgs_dir: Directory containing the frames to encode, named ``frame-000000``
onwards (``.png`` for RGB, ``.tiff`` for depth).
video_path: Output path for the encoded ``.mp4`` file.
fps: Frame rate of the output video.
video_encoder: Encoder settings (codec, pixel format, quality, ...). When
``None``, :func:`rgb_encoder_defaults` is used. Pass a
:class:`~lerobot.configs.video.DepthEncoderConfig` to encode depth frames.
encoder_threads: Per-encoder thread count forwarded to the codec. ``None``
lets the codec decide.
log_level: libav log level to set while encoding, or ``None`` to leave the
current logging configuration unchanged.
overwrite: When ``False`` and ``video_path`` already exists, skip encoding and
log a warning. When ``True``, re-encode and replace the existing file.
"""
if video_encoder is None:
video_encoder = rgb_encoder_defaults()
@@ -575,21 +552,16 @@ def reencode_video(
"""Re-encode a video file, optionally trimming it to ``[start_time_s, end_time_s)``.
Args:
input_video_path (`pathlib.Path | str`): Existing video file to read.
output_video_path (`pathlib.Path | str`): Path for the re-encoded file.
video_encoder (`lerobot.configs.video.VideoEncoderConfig | None`, *optional*): Encoder
configuration. Defaults to `rgb_encoder_defaults`.
encoder_threads (`int | None`, *optional*): Optional thread count forwarded to
`VideoEncoderConfig.get_codec_options`.
log_level (`int | None`, *optional*, defaults to 24): libav log level while encoding,
or `None` to leave logging unchanged.
overwrite (`bool`, *optional*, defaults to `False`): When `False` and `output_video_path`
already exists, skip and log a warning.
start_time_s (`float | None`, *optional*): When set, trim the output to start at this
timestamp, in seconds.
end_time_s (`float | None`, *optional*): When set, trim the output to end at this
timestamp, in seconds, exclusive.
input_video_path: Existing video file to read.
output_video_path: Path for the re-encoded file.
video_encoder: Encoder configuration. Defaults to :func:`rgb_encoder_defaults`.
encoder_threads: Optional thread count forwarded to :meth:`VideoEncoderConfig.get_codec_options`.
log_level: libav log level while encoding, or ``None`` to leave logging unchanged. Defaults to WARNING.
overwrite: When ``False`` and ``output_video_path`` already exists, skip and log a warning.
start_time_s: When set, trim the output to start at this timestamp (seconds).
end_time_s: When set, trim the output to end at this timestamp (seconds, exclusive).
"""
video_encoder = video_encoder or rgb_encoder_defaults()
if (start_time_s is not None and start_time_s < 0) or (end_time_s is not None and end_time_s < 0):
@@ -679,26 +651,25 @@ def concatenate_video_files(
overwrite: bool = True,
compatibility_check: bool = False,
):
"""Concatenate multiple video files into a single video file using pyav.
"""
Concatenate multiple video files into a single video file using pyav.
This function takes a list of video input file paths and concatenates them into a single
output video file. It uses ffmpeg's concat demuxer with stream copy mode for fast
concatenation without re-encoding.
Args:
input_video_paths (`list`): Ordered list of input video file paths to concatenate.
output_video_path (`Path`): Path to the output video file.
overwrite (`bool`, *optional*, defaults to `True`): Whether to overwrite the output
video file if it already exists.
compatibility_check (`bool`, *optional*, defaults to `False`): Whether to check that
the input videos share the same height, width, fps, codec, and pixel format
before concatenating.
input_video_paths: Ordered list of input video file paths to concatenate.
output_video_path: Path to the output video file.
overwrite: Whether to overwrite the output video file if it already exists. Default is True.
compatibility_check: Whether to check if the input videos are compatible. Default is False.
Note:
- Creates a temporary directory for intermediate files that is cleaned up after use.
- Uses ffmpeg's concat demuxer which requires all input videos to have the same
codec, resolution, and frame rate for proper concatenation.
"""
output_video_path = Path(output_video_path)
if output_video_path.exists() and not overwrite:
@@ -796,17 +767,6 @@ class _CameraEncoderThread(threading.Thread):
stop_event: threading.Event,
encoder_threads: int | None = None,
):
"""Set up the thread; frames are only consumed once `start()` is called.
Args:
video_path: Output MP4 path.
fps: Output frame rate.
video_encoder: Codec/quality settings; `DepthEncoderConfig` selects depth-map encoding.
frame_queue: Queue this thread reads `(frame, ...)` items from.
result_queue: Queue the final stats are pushed to once encoding finishes.
stop_event: Set by the caller to signal this thread to stop early.
encoder_threads: Number of threads passed to the codec, if it supports one.
"""
super().__init__(daemon=True)
self.video_path = video_path
self.fps = fps
@@ -818,10 +778,6 @@ class _CameraEncoderThread(threading.Thread):
self.encoder_threads = encoder_threads
def run(self) -> None:
"""Encode frames from `frame_queue` to `video_path` until a stop sentinel or `stop_event`.
Pushes the accumulated `RunningQuantileStats` to `result_queue` once encoding finishes.
"""
from .compute_stats import RunningQuantileStats, auto_downsample_height_width
container = None
@@ -942,8 +898,7 @@ class StreamingVideoEncoder:
queue_maxsize: int = 30,
encoder_threads: int | None = None,
):
"""Create the manager; per-camera encoder threads are started lazily on first frame.
"""
Args:
fps: Frames per second for the output videos.
rgb_encoder: Video encoder settings applied to all RGB cameras.
@@ -1150,9 +1105,11 @@ class StreamingVideoEncoder:
@dataclass
class VideoFrame:
# TODO(rcadene, lhoestq): move to Hugging Face `datasets` repo
"""Provides a type for a dataset containing video frames.
"""
Provides a type for a dataset containing video frames.
Example:
```python
data_dict = [{"image": {"path": "videos/episode_0.mp4", "timestamp": 0.3}}]
features = {"image": VideoFrame()}
@@ -1164,7 +1121,6 @@ class VideoFrame:
_type: str = field(default="VideoFrame", init=False, repr=False)
def __call__(self):
"""Return the pyarrow struct type backing this feature, as required by `datasets.Features`."""
return self.pa_type
@@ -1179,11 +1135,6 @@ with warnings.catch_warnings():
def get_audio_info(video_path: Path | str) -> dict:
"""Read audio-stream metadata (channels, codec, bit rate, sample rate, etc.) from a video file.
Returns:
A dict of `"audio.*"` keys, or `{"has_audio": False}` if `video_path` has no audio stream.
"""
# Set logging level
logging.getLogger("libav").setLevel(av.logging.WARNING)
@@ -1222,13 +1173,13 @@ def get_video_info(
"""Build the ``video.*`` / ``audio.*`` info dict persisted in ``info.json``.
Args:
video_path (`pathlib.Path | str`): Path to the encoded video file to probe.
video_encoder (`lerobot.configs.video.VideoEncoderConfig | None`, *optional*): If provided,
record the exact encoder settings used to encode this video. Stream-derived values take
precedence encoder fields are only written for keys not already populated from the
video file itself. When a `DepthEncoderConfig` is passed, the depth quantization
parameters (`depth_min` / `depth_max` / `shift` / `use_log`) are recorded so frames can
be dequantized on read.
video_path: Path to the encoded video file to probe.
video_encoder: If provided, record the exact encoder settings used to encode this
video. Stream-derived values take precedence encoder fields are only written for keys
not already populated from the video file itself. When a
:class:`~lerobot.configs.video.DepthEncoderConfig` is passed, the depth
quantization parameters (``depth_min`` / ``depth_max`` / ``shift`` /
``use_log``) are recorded so frames can be dequantized on read.
Returns:
The ``video.*`` / ``audio.*`` info dict, including ``is_depth_map`` which is
@@ -1276,10 +1227,11 @@ def get_video_info(
def get_video_duration_in_s(video_path: Path | str) -> float:
"""Get the duration of a video file in seconds using PyAV.
"""
Get the duration of a video file in seconds using PyAV.
Args:
video_path (`pathlib.Path | str`): Path to the video file.
video_path: Path to the video file.
Returns:
Duration of the video in seconds.
@@ -1297,7 +1249,8 @@ def get_video_duration_in_s(video_path: Path | str) -> float:
class VideoEncodingManager:
"""Context manager that ensures proper video encoding and data cleanup even if exceptions occur.
"""
Context manager that ensures proper video encoding and data cleanup even if exceptions occur.
This manager handles:
- Batch encoding for any remaining episodes when recording interrupted
@@ -1305,23 +1258,16 @@ class VideoEncodingManager:
- Removing empty image directories
Args:
dataset (`LeRobotDataset`): The LeRobotDataset instance.
dataset: The LeRobotDataset instance
"""
def __init__(self, dataset):
"""Store the `LeRobotDataset` this manager will finalize/clean up on exit."""
self.dataset = dataset
def __enter__(self):
"""Return `self`; no setup is needed on entry."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Finalize the dataset, cancelling pending videos and cleaning up interrupted-episode files.
Runs unconditionally (even if `exc_type` is set), so partial/interrupted recordings still leave
a consistent dataset on disk.
"""
writer = self.dataset.writer
if writer is not None:
if exc_type is not None and writer._streaming_encoder is not None:
+1 -1
View File
@@ -60,7 +60,7 @@ PATH_TO_LEROBOT = PATH_TO_REPO / "src" / "lerobot"
# Modules whose public objects are checked. Add a module here once its docstrings follow the standard.
MODULES_TO_CHECK = [
"lerobot.robots",
"lerobot.datasets",
"lerobot.configs",
]
# Objects that do not yet follow the standard, so the check can be green from day one. Removing an entry