mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 8353f84f52 |
@@ -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
|
||||
|
||||
@@ -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
@@ -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"]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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
@@ -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():
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,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".
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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.
|
||||
"""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user