admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-08-08 17:39:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

docs(configs): bring src/lerobot/configs/ to 100% docstring coverage

Browse Source 
Documents every remaining public class/function across configs/: the
types.py feature/normalization enums and PolicyFeature, the
PreTrainedConfig and RewardModelConfig base classes (full Args: blocks,
abstract property/method contracts, from_pretrained overrides), the
TrainPipelineConfig and EvalPipelineConfig top-level pipeline configs,
default.py's DatasetConfig/WandBConfig/EvalConfig/PeftConfig/JobConfig,
dataset.py's DatasetRecordConfig, video.py's VideoEncoderConfig, the
accelerator.py ActivationCheckpointingMode enum, and parser.py's CLI
argument-parsing helpers.

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

Expands docs/source/api/configs.mdx from 5 documented classes to the
module's full public surface. Adds lerobot.configs to
check_docstrings.py's MODULES_TO_CHECK ratchet and removes the
module's ruff D-ignore. Docstrings and docstring-format changes only;
no behavioral changes.
This commit is contained in:
CarolinePascal
2026-08-07 11:57:10 +02:00
parent 741005d719
commit 8353f84f52
14 changed files with 649 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/api/configs.mdx
+62
View File
@@ -10,14 +10,26 @@ itself with `@register_subclass("name")` and is then selectable by that name on
[[autodoc]] lerobot.configs.train.TrainPipelineConfig
## EvalPipelineConfig
[[autodoc]] lerobot.configs.eval.EvalPipelineConfig
## PreTrainedConfig
[[autodoc]] lerobot.configs.PreTrainedConfig
## RewardModelConfig
[[autodoc]] lerobot.configs.rewards.RewardModelConfig
## DatasetConfig
[[autodoc]] lerobot.configs.DatasetConfig
## DatasetRecordConfig
[[autodoc]] lerobot.configs.DatasetRecordConfig
## EvalConfig
[[autodoc]] lerobot.configs.EvalConfig
@@ -25,3 +37,53 @@ itself with `@register_subclass("name")` and is then selectable by that name on
## WandBConfig
[[autodoc]] lerobot.configs.WandBConfig
## PeftConfig
[[autodoc]] lerobot.configs.PeftConfig
## JobConfig
[[autodoc]] lerobot.configs.JobConfig
## Feature types
[[autodoc]] lerobot.configs.FeatureType
[[autodoc]] lerobot.configs.PipelineFeatureType
[[autodoc]] lerobot.configs.NormalizationMode
[[autodoc]] lerobot.configs.PolicyFeature
[[autodoc]] lerobot.configs.RTCAttentionSchedule
## Video encoding
[[autodoc]] lerobot.configs.VideoEncoderConfig
[[autodoc]] lerobot.configs.RGBEncoderConfig
[[autodoc]] lerobot.configs.DepthEncoderConfig
[[autodoc]] lerobot.configs.encoder_config_from_video_info
## Distributed training
[[autodoc]] lerobot.configs.parallelism.ParallelismConfig
[[autodoc]] lerobot.configs.parallelism.ContextParallelConfig
[[autodoc]] lerobot.configs.accelerator.AcceleratorConfig
[[autodoc]] lerobot.configs.accelerator.FSDPConfig
[[autodoc]] lerobot.configs.accelerator.DDPConfig
[[autodoc]] lerobot.configs.accelerator.GradientAccumulationConfig
[[autodoc]] lerobot.configs.accelerator.CompileConfig
[[autodoc]] lerobot.configs.accelerator.ActivationCheckpointingConfig
[[autodoc]] lerobot.configs.accelerator.ActivationCheckpointingMode
pyproject.toml
-1
View File
@@ -439,7 +439,6 @@ 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/configs/__init__.py
+1 -2
View File
@@ -12,8 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Public API for lerobot configuration types and base config classes.
"""Public API for lerobot configuration types and base config classes.
NOTE: TrainPipelineConfig, EvalPipelineConfig, and TrainRLServerPipelineConfig
are intentionally NOT re-exported here to avoid circular dependencies
src/lerobot/configs/accelerator.py
+7
View File
@@ -171,6 +171,13 @@ class CompileConfig:
class ActivationCheckpointingMode(str, Enum):
"""The activation-checkpointing strategy applied to FSDP wrap units.
**Attributes**:
- **NONE** -- No activation checkpointing.
- **FULL** -- Checkpoint every wrap unit.
"""
NONE = "none"
FULL = "full"
src/lerobot/configs/dataset.py
+51 -31
View File
@@ -23,56 +23,76 @@ from .video import DepthEncoderConfig, RGBEncoderConfig, depth_encoder_defaults,
@dataclass
class DatasetRecordConfig:
# Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).
"""Shared dataset recording configuration used by both `lerobot-record` and `lerobot-rollout`.
Args:
repo_id (`str`, *optional*, defaults to `""`): Dataset identifier. By convention it should match
`'{hf_username}/{dataset_name}'` (e.g. `lerobot/test`).
single_task (`str`, *optional*, defaults to `""`): A short but accurate description of the task performed during the
recording (e.g. `"Pick the Lego block and drop it in the box on the right."`).
root (`str | Path | None`, *optional*): Root directory where the dataset will be stored (e.g.
`'dataset/path'`). If `None`, defaults to `$HF_LEROBOT_HOME/repo_id`.
fps (`int`, *optional*, defaults to 30): Limit the frames per second.
episode_time_s (`int | float`, *optional*, defaults to 60): Number of seconds for data recording
for each episode.
reset_time_s (`int | float`, *optional*, defaults to 60): Number of seconds for resetting the
environment after each episode.
num_episodes (`int`, *optional*, defaults to 50): Number of episodes to record.
video (`bool`, *optional*, defaults to `True`): Encode frames in the dataset into video.
push_to_hub (`bool`, *optional*, defaults to `True`): Upload dataset to the Hugging Face Hub.
private (`bool | None`, *optional*): If `True`, upload as private; if `None`, defer to the org
default on the Hub (only affects orgs).
tags (`list[str] | None`, *optional*): Add tags to your dataset on the Hub.
num_image_writer_processes (`int`, *optional*, defaults to 0): Number of subprocesses handling the
saving of frames as PNG. Set to 0 to use threads only; set to >=1 to use subprocesses, each
using threads to write images. The best number of processes and threads depends on your
system. We recommend 4 threads per camera with 0 processes. If fps is unstable, adjust the
thread count. If still unstable, try using 1 or more subprocesses.
num_image_writer_threads_per_camera (`int`, *optional*, defaults to 4): Number of threads writing
the frames as png images on disk, per camera. Too many threads might cause unstable
teleoperation fps due to the main thread being blocked. Not enough threads might cause low
camera fps.
video_encoding_batch_size (`int`, *optional*, defaults to 1): Number of episodes to record before
batch encoding videos. Set to 1 for immediate encoding (default behavior), or higher for
batched encoding.
rgb_encoder (`RGBEncoderConfig`, *optional*): Video encoder settings for camera MP4s (codec,
quality, GOP, etc.). Tuned via CLI nested keys, e.g. `--dataset.rgb_encoder.vcodec=h264`.
depth_encoder (`DepthEncoderConfig`, *optional*): Video encoder settings for depth-map MP4s (codec,
quality, GOP, etc.). Tuned via CLI nested keys.
streaming_encoding (`bool`, *optional*, defaults to `False`): Enable streaming video encoding:
encode frames in real-time during capture instead of writing PNG images first. Makes
`save_episode()` near-instant. More info in the documentation:
https://huggingface.co/docs/lerobot/streaming_video_encoding
encoder_queue_maxsize (`int`, *optional*, defaults to 30): Maximum number of frames to buffer per
camera when using streaming encoding. ~1s buffer at 30fps. Provides backpressure if the encoder
can't keep up.
encoder_threads (`int | None`, *optional*): Number of threads per encoder instance. `None` means
auto (codec default). Lower values reduce CPU usage; maps to `'lp'` (via `svtav1-params`) for
libsvtav1 and `'threads'` for h264/hevc.
no_stamp (`bool`, *optional*, defaults to `False`): Skip appending the date-time tag to `repo_id`,
keeping the user-provided name as-is (e.g. self-managed versioned names intended for a later
`lerobot-edit-dataset merge`).
"""
repo_id: str = ""
# A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.")
single_task: str = ""
# Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
root: str | Path | None = None
# Limit the frames per second.
fps: int = 30
# Number of seconds for data recording for each episode.
episode_time_s: int | float = 60
# Number of seconds for resetting the environment after each episode.
reset_time_s: int | float = 60
# Number of episodes to record.
num_episodes: int = 50
# Encode frames in the dataset into video
video: bool = True
# Upload dataset to Hugging Face hub.
push_to_hub: bool = True
# If True, upload as private; if None, defer to the org default on the Hub (only affects orgs).
private: bool | None = None
# Add tags to your dataset on the hub.
tags: list[str] | None = None
# Number of subprocesses handling the saving of frames as PNG. Set to 0 to use threads only;
# set to ≥1 to use subprocesses, each using threads to write images. The best number of processes
# and threads depends on your system. We recommend 4 threads per camera with 0 processes.
# If fps is unstable, adjust the thread count. If still unstable, try using 1 or more subprocesses.
num_image_writer_processes: int = 0
# Number of threads writing the frames as png images on disk, per camera.
# Too many threads might cause unstable teleoperation fps due to main thread being blocked.
# Not enough threads might cause low camera fps.
num_image_writer_threads_per_camera: int = 4
# Number of episodes to record before batch encoding videos
# Set to 1 for immediate encoding (default behavior), or higher for batched encoding
video_encoding_batch_size: int = 1
# Video encoder settings for camera MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys,
# e.g. ``--dataset.rgb_encoder.vcodec=h264`` (see ``RGBEncoderConfig``).
rgb_encoder: RGBEncoderConfig = field(default_factory=rgb_encoder_defaults)
# Video encoder settings for depth-map MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys.
depth_encoder: DepthEncoderConfig = field(default_factory=depth_encoder_defaults)
# Enable streaming video encoding: encode frames in real-time during capture instead
# of writing PNG images first. Makes save_episode() near-instant. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding
streaming_encoding: bool = False
# Maximum number of frames to buffer per camera when using streaming encoding.
# ~1s buffer at 30fps. Provides backpressure if the encoder can't keep up.
encoder_queue_maxsize: int = 30
# Number of threads per encoder instance. None = auto (codec default).
# Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
encoder_threads: int | None = None
# Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
# (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
no_stamp: bool = False
def stamp_repo_id(self) -> None:
src/lerobot/configs/default.py
+128 -58
View File
@@ -27,35 +27,65 @@ logger = logging.getLogger(__name__)
@dataclass
class DatasetConfig:
# You may provide a list of datasets here. `train.py` creates them all and concatenates them. Note: only data
# keys common between the datasets are kept. Each dataset gets and additional transform that inserts the
# "dataset_index" into the returned item. The index mapping is made according to the order in which the
# datasets are provided.
"""A dataset to train on. `TrainPipelineConfig.dataset` may be a list of these, concatenated together.
Only data keys common between multiple datasets are kept. Each dataset gets an additional transform
that inserts the `"dataset_index"` into the returned item, with the index mapping made according to
the order in which the datasets are provided.
Args:
repo_id (`str`): The Hub repo ID (or local dataset name, if `root` is set) to load.
repo_type (`str`, *optional*, defaults to `"dataset"`): Hub repository type: `"dataset"` (the
default) or `"bucket"` for an HF Storage Bucket streamed over `hf://buckets/`. Buckets are
streaming-only, so `"bucket"` requires `streaming=True`.
root (`str | None`, *optional*): Root directory for a concrete local dataset tree (e.g.
`'dataset/path'`). If `None`, local datasets are looked up under `$HF_LEROBOT_HOME/repo_id` and
Hub downloads use a revision-safe cache under `$HF_LEROBOT_HOME/hub`.
episodes (`list[int] | None`, *optional*): Episode indices to include. If `None`, all episodes are
used.
exclude_episodes (`list[int] | None`, *optional*): Episode indices to drop (e.g. corrupt or
heterogeneous ones). Applied on top of `episodes`.
image_transforms (`ImageTransformsConfig`, *optional*): Image augmentation settings applied at load
time.
revision (`str | None`, *optional*): Hub revision (commit hash, branch, or tag) to load.
use_imagenet_stats (`bool`, *optional*, defaults to `True`): Whether to use ImageNet normalization
statistics for visual features instead of the dataset's own.
video_backend (`str`, *optional*): The video decoding backend to use.
return_uint8 (`bool`, *optional*, defaults to `False`): When `True`, RGB video frames are returned
as `uint8` tensors (0-255) instead of `float32` (0.0-1.0). This reduces memory and speeds up
DataLoader IPC. The training pipeline handles the conversion.
depth_output_unit (`str`, *optional*, defaults to `"mm"`): Physical unit depth maps are dequantized
to at load time: `"mm"` (millimeters) or `"m"` (metres). Has no effect on datasets without depth
cameras.
streaming (`bool`, *optional*, defaults to `False`): Stream the dataset instead of downloading it
locally.
eval_split (`float`, *optional*, defaults to 0.0): Fraction of episodes held out per task for
offline evaluation (0.0 = disabled).
"""
repo_id: str
# Hub repository type: "dataset" (default) or "bucket" for an HF Storage Bucket streamed over
# hf://buckets/. Buckets are streaming-only, so "bucket" requires streaming=true.
repo_type: str = "dataset"
# Root directory for a concrete local dataset tree (e.g. 'dataset/path'). If None, local datasets are
# looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
root: str | None = None
episodes: list[int] | None = None
# Episode indices to drop (e.g. corrupt or heterogeneous ones). Applied on top of `episodes`.
exclude_episodes: list[int] | None = None
image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
revision: str | None = None
use_imagenet_stats: bool = True
video_backend: str = field(default_factory=get_safe_default_video_backend)
# When True, RGB video frames are returned as uint8 tensors (0-255) instead of float32 (0.0-1.0).
# This reduces memory and speeds up DataLoader IPC. The training pipeline handles the conversion.
return_uint8: bool = False
# Physical unit depth maps are dequantized to at load time: "mm" (millimeters) or "m" (metres).
# Has no effect on datasets without depth cameras.
depth_output_unit: str = DEFAULT_DEPTH_UNIT
streaming: bool = False
# Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
eval_split: float = 0.0
def __post_init__(self) -> None:
"""Validate `repo_type`/`streaming`/`depth_output_unit`/`eval_split`/`episodes`/`exclude_episodes`.
Raises:
ValueError: If `repo_type` isn't `"dataset"` or `"bucket"`; if `repo_type="bucket"` is combined
with `streaming=False` or a nonzero `eval_split`; if `depth_output_unit` isn't a recognized
unit; if `eval_split` is outside `[0.0, 1.0)`; or if `episodes` contains negative or
duplicate indices.
"""
if self.repo_type not in ("dataset", "bucket"):
raise ValueError(f"repo_type must be 'dataset' or 'bucket', got {self.repo_type!r}")
if self.repo_type == "bucket" and not self.streaming:
@@ -92,35 +122,63 @@ class DatasetConfig:
@dataclass
class WandBConfig:
"""Weights & Biases logging settings for `lerobot-train`.
Args:
enable (`bool`, *optional*, defaults to `False`): Whether to log this run to Weights & Biases.
disable_artifact (`bool`, *optional*, defaults to `False`): Set to `True` to disable saving an
artifact despite `save_checkpoint=True`.
project (`str`, *optional*, defaults to `"lerobot"`): The WandB project to log to.
entity (`str | None`, *optional*): The WandB entity (team or username) to log under.
notes (`str | None`, *optional*): Notes attached to the WandB run.
run_id (`str | None`, *optional*): An existing WandB run id to resume logging into.
mode (`str | None`, *optional*): WandB mode: `"online"`, `"offline"`, or `"disabled"`. Defaults to
`"online"`.
add_tags (`bool`, *optional*, defaults to `True`): If `True`, save the training configuration as
tags on the WandB run.
"""
enable: bool = False
# Set to true to disable saving an artifact despite training.save_checkpoint=True
disable_artifact: bool = False
project: str = "lerobot"
entity: str | None = None
notes: str | None = None
run_id: str | None = None
mode: str | None = None # Allowed values: 'online', 'offline' 'disabled'. Defaults to 'online'
add_tags: bool = True # If True, save configuration as tags in the WandB run.
mode: str | None = None
add_tags: bool = True
@dataclass
class EvalConfig:
"""Settings for the periodic in-training simulation-environment evaluation.
Args:
n_episodes (`int`, *optional*, defaults to 50): Number of episodes to run per evaluation.
batch_size (`int`, *optional*, defaults to 0): The number of environments to use in a
`gym.vector.VectorEnv`. `0` auto-tunes based on available CPU cores and `n_episodes`.
use_async_envs (`bool`, *optional*, defaults to `True`): Whether to use asynchronous environments
(multiprocessing). Automatically downgraded to a `SyncVectorEnv` when `batch_size` is 1.
recording (`bool`, *optional*, defaults to `False`): Whether to record eval rollouts as a LeRobot
dataset on disk.
recording_repo_id (`str | None`, *optional*): If set, push recorded eval datasets to the Hub under
this repo id (one repo per task, suffixed by task and env index). Requires `recording=True`.
recording_private (`bool`, *optional*, defaults to `False`): Whether the pushed recording
repositories should be private.
"""
n_episodes: int = 50
# `batch_size` specifies the number of environments to use in a gym.vector.VectorEnv.
# Set to 0 for auto-tuning based on available CPU cores and n_episodes.
batch_size: int = 0
# `use_async_envs` specifies whether to use asynchronous environments (multiprocessing).
# Defaults to True; automatically downgraded to SyncVectorEnv when batch_size=1.
use_async_envs: bool = True
# Whether to record eval rollouts as a LeRobot dataset on disk.
recording: bool = False
# If set, push recorded eval datasets to the Hub under this repo id (one repo per task,
# suffixed by task and env index). Requires recording=true.
recording_repo_id: str | None = None
# Whether the pushed recording repositories should be private.
recording_private: bool = False
def __post_init__(self) -> None:
"""Validate `recording_repo_id`/`recording`, and resolve/cap `batch_size`.
Raises:
ValueError: If `recording_repo_id` is set without `recording=True`.
"""
if self.recording_repo_id is not None and not self.recording:
raise ValueError("eval.recording_repo_id requires eval.recording=true.")
if self.batch_size == 0:
@@ -141,54 +199,66 @@ class EvalConfig:
@dataclass
class PeftConfig:
# PEFT offers many fine-tuning methods, layer adapters being the most common and currently also the most
# effective methods so we'll focus on those in this high-level config interface.
"""PEFT (parameter-efficient fine-tuning) settings, e.g. LoRA adapters.
PEFT offers many fine-tuning methods, layer adapters being the most common and currently also the
most effective methods so we'll focus on those in this high-level config interface.
Args:
target_modules (`list[str] | str | None`, *optional*): Either a string (module name suffix or
`'all-linear'`), a list of module name suffixes, or a regular expression describing module
names to target with the configured PEFT method. Some policies have a default value for this
so that you don't *have* to choose which layers to adapt, but it might still be worthwhile
depending on your case.
full_training_modules (`list[str] | None`, *optional*): Names/suffixes of modules to fully
fine-tune and store alongside adapter weights. Useful for layers that are not part of a
pre-trained model (e.g., action state projections). Depending on the policy this defaults to
layers that are newly created in pre-trained policies. If you're fine-tuning an already trained
policy you might want to set this to `[]`. Corresponds to PEFT's `modules_to_save`.
method_type (`str`, *optional*, defaults to `"LORA"`): The PEFT (adapter) method to apply to the
policy. Needs to be a valid PEFT type.
init_type (`str | None`, *optional*): Adapter initialization method. Look at the specific PEFT
adapter documentation for defaults.
r (`int`, *optional*, defaults to 16): We expect that all PEFT adapters are in some way doing
rank-decomposition, therefore this parameter specifies the rank used for the adapter. In
general a higher rank means more trainable parameters and closer to full fine-tuning.
lora_alpha (`int | None`, *optional*): Alpha parameter for LoRA scaling (`scaling = lora_alpha /
r`). In general, a higher alpha means stronger adaptation signal. If `None`, the PEFT library
defaults to `alpha=8`, which may dampen high-rank adapters. Common values are `r` (`alpha ==
rank`) or `2*r`.
"""
# Either a string (module name suffix or 'all-linear'), a list of module name suffixes or a regular expression
# describing module names to target with the configured PEFT method. Some policies have a default value for this
# so that you don't *have* to choose which layers to adapt but it might still be worthwhile depending on your case.
target_modules: list[str] | str | None = None
# Names/suffixes of modules to fully fine-tune and store alongside adapter weights. Useful for layers that are
# not part of a pre-trained model (e.g., action state projections). Depending on the policy this defaults to layers
# that are newly created in pre-trained policies. If you're fine-tuning an already trained policy you might want
# to set this to `[]`. Corresponds to PEFT's `modules_to_save`.
full_training_modules: list[str] | None = None
# The PEFT (adapter) method to apply to the policy. Needs to be a valid PEFT type.
method_type: str = "LORA"
# Adapter initialization method. Look at the specific PEFT adapter documentation for defaults.
init_type: str | None = None
# We expect that all PEFT adapters are in some way doing rank-decomposition therefore this parameter specifies
# the rank used for the adapter. In general a higher rank means more trainable parameters and closer to full
# fine-tuning.
r: int = 16
# Alpha parameter for LoRA scaling (scaling = lora_alpha / r).
# In general, a higher alpha means stronger adaptation signal.
# If None, the PEFT library defaults to alpha=8, which may dampen high-rank adapters.
# Common values are r (alpha == rank) or 2*r.
lora_alpha: int | None = None
@dataclass
class JobConfig:
# Where training runs. None (omitted) or "local" runs on this machine.
# Any other value is an HF Jobs flavor and submits the run to HF Jobs.
# List available flavors + pricing with `hf jobs hardware` command.
"""Where and how a training run executes: locally, or dispatched to an HF Jobs flavor.
Args:
target (`str | None`, *optional*): Where training runs. `None` (omitted) or `"local"` runs on this
machine. Any other value is an HF Jobs flavor and submits the run to HF Jobs. List available
flavors and pricing with the `hf jobs hardware` command.
image (`str`, *optional*, defaults to `"huggingface/lerobot-gpu:latest"`): Runtime image for the
remote job (ignored for local runs).
timeout (`str | None`, *optional*, defaults to `"2d"`): Max wall-clock for the remote job as an HF
Jobs duration string (e.g. `"2h"`). HF Jobs itself defaults to `"2d"`; we pass an explicit,
generous cap instead. Set a smaller value to fail fast, or a larger one for long runs.
detach (`bool`, *optional*, defaults to `False`): Submit and exit instead of streaming the job logs
in the foreground.
tags (`list[str]`, *optional*): Extra tags attached to the HF job and to any dataset this run
pushes to the Hub. A `"lerobot"` tag is always added; e.g. `--job.tags '["lelab"]'` adds more.
"""
target: str | None = None
# Runtime image for the remote job (ignored for local runs).
image: str = "huggingface/lerobot-gpu:latest"
# Max wall-clock for the remote job as an HF Jobs duration string (e.g. "2h").
# Defaults to "2d": We pass an explicit, generous cap instead. Set a smaller
# value to fail fast, or a larger one for long runs.
timeout: str | None = "2d"
# Submit and exit instead of streaming the job logs in the foreground.
detach: bool = False
# Extra tags attached to the HF job and to any dataset this run pushes to the
# Hub. A "lerobot" tag is always added; e.g. --job.tags '["lelab"]' adds more.
tags: list[str] = field(default_factory=list)
# Two entry points to the same predicate: the staticmethod tests a raw target string
src/lerobot/configs/eval.py
+21 -6
View File
@@ -28,21 +28,36 @@ logger = getLogger(__name__)
@dataclass
class EvalPipelineConfig:
# Either the repo ID of a model hosted on the Hub or a path to a directory containing weights
# saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch
# (useful for debugging). This argument is mutually exclusive with `--config`.
"""The top-level configuration for `lerobot-eval`, parsed by draccus from CLI flags and/or a YAML file.
Args:
env (`envs.EnvConfig`): The simulation environment to evaluate the policy in.
eval (`EvalConfig`, *optional*): Number of episodes, batching, and recording settings for the
evaluation run.
policy (`PreTrainedConfig | None`, *optional*): Loaded from `--policy.path`, either the repo ID of
a model hosted on the Hub or a path to a directory containing weights saved using
`PreTrainedPolicy.save_pretrained`. If not provided, the policy is initialized from scratch
(useful for debugging).
output_dir (`Path | None`, *optional*): Where to save evaluation outputs.
job_name (`str | None`, *optional*): A name for the run.
seed (`int | None`, *optional*, defaults to 1000): Seed used for the evaluation environments.
rename_map (`dict[str, str]`, *optional*): Rename map for the observation, to override the image
and state keys.
trust_remote_code (`bool`, *optional*, defaults to `False`): Explicit consent to execute remote
code from the Hub (required for Hub environments).
"""
env: envs.EnvConfig
eval: EvalConfig = field(default_factory=EvalConfig)
policy: PreTrainedConfig | None = None
output_dir: Path | None = None
job_name: str | None = None
seed: int | None = 1000
# Rename map for the observation to override the image and state keys
rename_map: dict[str, str] = field(default_factory=dict)
# Explicit consent to execute remote code from the Hub (required for hub environments).
trust_remote_code: bool = False
def __post_init__(self) -> None:
"""Resolve `--policy.path` into a loaded config, and derive `job_name`/`output_dir` when unset."""
# HACK: We parse again the cli args here to get the pretrained path if there was one.
policy_path = parser.get_path_arg("policy")
if policy_path:
@@ -75,5 +90,5 @@ class EvalPipelineConfig:
@classmethod
def __get_path_fields__(cls) -> list[str]:
"""This enables the parser to load config from the policy using `--policy.path=local/dir`"""
"""This enables the parser to load config from the policy using `--policy.path=local/dir`."""
return ["policy"]
src/lerobot/configs/parser.py
+30 -19
View File
@@ -96,6 +96,7 @@ def get_cli_overrides(field_name: str, args: Sequence[str] | None = None) -> lis
def parse_arg(arg_name: str, args: Sequence[str] | None = None) -> str | None:
"""Return the value of `--{arg_name}=value` or `--{arg_name} value` in `args` (`sys.argv[1:]` if `None`)."""
if args is None:
args = sys.argv[1:]
option = f"--{arg_name}"
@@ -115,7 +116,7 @@ def parse_plugin_args(plugin_arg_suffix: str, args: Sequence[str]) -> dict[str,
Args:
plugin_arg_suffix (str): The suffix to identify plugin-related arguments.
cli_args (Sequence[str]): A sequence of command-line arguments to parse.
args (`Sequence[str]`): A sequence of command-line arguments to parse.
Returns:
dict: A dictionary containing the parsed plugin arguments where:
@@ -156,7 +157,7 @@ def load_plugin(plugin_path: str) -> None:
registered with their parents using the `register_subclass` decorator.
Args:
plugin_path (str): The Python package path to the plugin (e.g. "mypackage.plugins.myplugin")
plugin_path (str): The Python package path to the plugin, e.g. "mypackage.plugins.myplugin".
Raises:
PluginLoadError: If the plugin cannot be loaded due to import errors or if the package path is invalid.
@@ -180,6 +181,7 @@ def load_plugin(plugin_path: str) -> None:
) from e
def iter_namespace(ns_pkg: ModuleType) -> Iterable[ModuleInfo]:
"""Iterate the direct submodules of `ns_pkg`, yielding their fully-qualified names."""
return pkgutil.iter_modules(ns_pkg.__path__, ns_pkg.__name__ + ".")
try:
@@ -192,6 +194,7 @@ def load_plugin(plugin_path: str) -> None:
def get_path_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
"""Return `--{field_name}.path`'s value from CLI `args`, or from a YAML/JSON config if not on the CLI."""
result = parse_arg(f"{field_name}.{PATH_KEY}", args)
if result is None:
result = _config_path_args.get(field_name)
@@ -199,21 +202,24 @@ def get_path_arg(field_name: str, args: Sequence[str] | None = None) -> str | No
def get_yaml_overrides(field_name: str) -> list[str]:
"""Return the CLI-style overrides extracted from `field_name`'s YAML/JSON config path block, if any."""
return _config_yaml_overrides.get(field_name, [])
def get_type_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
"""Return `--{field_name}.type`'s value from CLI `args` (`sys.argv[1:]` if `None`)."""
return parse_arg(f"{field_name}.{draccus.CHOICE_TYPE_KEY}", args)
def _register_scoped_actions(
wrapper: Wrapper, parser: SuppressingArgumentParser, cli_args: Sequence[str]
) -> None:
"""Like draccus's own Wrapper.register_actions, but for a ChoiceType field only recurses into
the already-selected subclass (per CLI `.type` args), instead of every registered choice.
"""Like draccus's own Wrapper.register_actions, but only recurses a ChoiceType field's subclasses.
This mirrors draccus 0.11.x's internal wrapper traversal because its public parser eagerly registers
every choice before parsing the command line. Keep this in sync when updating draccus.
Only the already-selected subclass (per CLI `.type` args) is recursed into, instead of every
registered choice. This mirrors draccus 0.11.x's internal wrapper traversal because its public
parser eagerly registers every choice before parsing the command line. Keep this in sync when
updating draccus.
"""
if isinstance(wrapper, ChoiceWrapper):
group = parser.add_argument_group(title=wrapper.title, description=wrapper.description)
@@ -253,8 +259,10 @@ def _register_scoped_actions(
def print_scoped_help(config_class: type, cli_args: Sequence[str]) -> None:
"""Prints --help output scoped to the choices already resolved on the CLI (e.g. --env.type=pusht),
instead of draccus's default of expanding every registered subclass of every ChoiceType field."""
"""Prints --help output scoped to the choices already resolved on the CLI (e.g. --env.type=pusht).
Instead of draccus's default of expanding every registered subclass of every ChoiceType field.
"""
parser = SuppressingArgumentParser(formatter_class=SimpleHelpFormatter)
parser.add_argument(
f"--{draccus.utils.CONFIG_ARG}", type=str, help="Path for a config file to parse with draccus"
@@ -264,6 +272,7 @@ def print_scoped_help(config_class: type, cli_args: Sequence[str]) -> None:
def filter_arg(field_to_filter: str, args: Sequence[str] | None = None) -> list[str]:
"""Return `args` with `--{field_to_filter}` (and its value, if separate) removed."""
if args is None:
return []
option = f"--{field_to_filter}"
@@ -285,12 +294,11 @@ def filter_arg(field_to_filter: str, args: Sequence[str] | None = None) -> list[
def filter_path_args(fields_to_filter: str | list[str], args: Sequence[str] | None = None) -> list[str]:
"""
Filters command-line arguments related to fields with specific path arguments.
"""Filters command-line arguments related to fields with specific path arguments.
Args:
fields_to_filter (str | list[str]): A single str or a list of str whose arguments need to be filtered.
args (Sequence[str] | None): The sequence of command-line arguments to be filtered.
args (Sequence[str] | None, *optional*): The sequence of command-line arguments to be filtered.
Defaults to None.
Returns:
@@ -380,19 +388,22 @@ def extract_path_fields_from_config(config_path: str, path_fields: list[str]) ->
def wrap(config_path: Path | None = None) -> Callable[[F], F]:
"""
HACK: Similar to draccus.wrap but does three additional things:
- Will remove '.path' arguments from CLI in order to process them later on.
- If a 'config_path' is passed and the main config class has a 'from_pretrained' method, will
initialize it from there to allow to fetch configs from the hub directly
- Will load plugins specified in the CLI arguments. These plugins will typically register
their own subclasses of config classes, so that draccus can find the right class to instantiate
from the CLI '.type' arguments
"""HACK: Similar to draccus.wrap but does three additional things.
- Will remove '.path' arguments from CLI in order to process them later on.
- If a 'config_path' is passed and the main config class has a 'from_pretrained' method, will
initialize it from there to allow to fetch configs from the hub directly
- Will load plugins specified in the CLI arguments. These plugins will typically register
their own subclasses of config classes, so that draccus can find the right class to instantiate
from the CLI '.type' arguments
"""
def wrapper_outer(fn: F) -> F:
"""Wrap `fn` so its first argument is resolved from the CLI/config instead of passed directly."""
@wraps(fn)
def wrapper_inner(*args: Any, **kwargs: Any) -> Any:
"""Build `fn`'s config argument from the CLI/config file (unless already given), then call `fn`."""
argspec = inspect.getfullargspec(fn)
argtype = argspec.annotations[argspec.args[0]]
if len(args) > 0 and type(args[0]) is argtype:
src/lerobot/configs/policies.py
+88 -21
View File
@@ -39,50 +39,64 @@ logger = getLogger(__name__)
@dataclass
class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: ignore[misc,name-defined] #TODO: draccus issue
"""
Base configuration class for policy models.
"""Base configuration class for policy models.
Every concrete policy config also declares a `normalization_mapping: dict[str, NormalizationMode]`
field (mapping a `FeatureType` name, e.g. `"STATE"`/`"VISUAL"`, to the `NormalizationMode` to apply),
with a policy-specific default — not declared here since it has no sensible shared default.
Args:
n_obs_steps: Number of environment steps worth of observations to pass to the policy (takes the
current step and additional steps going back).
input_features: A dictionary defining the PolicyFeature of the input data for the policy. The key represents
the input data name, and the value is PolicyFeature, which consists of FeatureType and shape attributes.
output_features: A dictionary defining the PolicyFeature of the output data for the policy. The key represents
the output data name, and the value is PolicyFeature, which consists of FeatureType and shape attributes.
normalization_mapping: A dictionary that maps from a str value of FeatureType (e.g., "STATE", "VISUAL") to
a corresponding NormalizationMode (e.g., NormalizationMode.MIN_MAX)
n_obs_steps (`int`, *optional*, defaults to 1): Number of environment steps worth of observations
to pass to the policy (takes the current step and additional steps going back).
input_features (`dict[str, PolicyFeature] | None`, *optional*): A dictionary defining the
`PolicyFeature` of the input data for the policy. The key represents the input data name, and
the value is a `PolicyFeature`, which consists of `type` and `shape` attributes. Can be set to
`None`/`null` in order to infer those values from the dataset.
output_features (`dict[str, PolicyFeature] | None`, *optional*): A dictionary defining the
`PolicyFeature` of the output data for the policy, with the same key/value semantics as
`input_features`.
device (`str | None`, *optional*): The torch device, e.g. `"cuda"`, `"cuda:0"`, `"cpu"`, or `"mps"`.
If unset or unavailable, `__post_init__` auto-selects one.
use_amp (`bool`, *optional*, defaults to `False`): Whether to use Automatic Mixed Precision for
training and evaluation, with automatic gradient scaling. Auto-disabled by `__post_init__`
when AMP isn't available on `device`.
use_peft (`bool`, *optional*, defaults to `False`): Whether the policy employed PEFT for training.
push_to_hub (`bool`, *optional*, defaults to `True`): Whether to push the policy to the Hugging Face
Hub after training.
repo_id (`str | None`, *optional*): The Hub repo ID to push to. Required when `push_to_hub` is
`True`.
private (`bool | None`, *optional*): Whether to upload to a private repository on the Hugging Face
Hub.
tags (`list[str] | None`, *optional*): Tags to add to the policy on the Hub.
license (`str | None`, *optional*): The license to add to the policy on the Hub.
pretrained_path (`Path | None`, *optional*): Either the repo ID of a model hosted on the Hub or a
path to a directory containing weights saved using `PreTrainedPolicy.save_pretrained`. If not
provided, the policy is initialized from scratch.
pretrained_revision (`str | None`, *optional*): Hub revision (commit hash, branch, or tag) to pin
the pretrained model version.
"""
n_obs_steps: int = 1
# `input_features` can be set to None/null in order to infer those values from the dataset.
input_features: dict[str, PolicyFeature] | None = field(default_factory=dict)
output_features: dict[str, PolicyFeature] | None = field(default_factory=dict)
device: str | None = None # e.g. "cuda", "cuda:0", "cpu", or "mps"
# `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
# automatic gradient scaling is used.
device: str | None = None
use_amp: bool = False
# Whether the policy employed PEFT for training.
use_peft: bool = False
push_to_hub: bool = True # type: ignore[assignment] # TODO: use a different name to avoid override
repo_id: str | None = None
# Upload on private repository on the Hugging Face hub.
private: bool | None = None
# Add tags to your policy on the hub.
tags: list[str] | None = None
# Add tags to your policy on the hub.
license: str | None = None
# Either the repo ID of a model hosted on the Hub or a path to a directory containing weights
# saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch.
pretrained_path: Path | None = None
# Optional Hub revision (commit hash, branch, or tag) to pin the pretrained model version.
pretrained_revision: str | None = None
def __post_init__(self) -> None:
"""Auto-select `device` when unset/unavailable, and disable `use_amp` when AMP isn't available on it."""
if not self.device or not is_torch_device_available(self.device):
auto_device = auto_select_torch_device()
logger.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
@@ -97,6 +111,7 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
@property
def type(self) -> str:
"""The policy's registered `draccus.ChoiceRegistry` name (e.g. `"act"`, `"diffusion"`)."""
choice_name = self.get_choice_name(self.__class__)
if not isinstance(choice_name, str):
raise TypeError(f"Expected string from get_choice_name, got {type(choice_name)}")
@@ -105,32 +120,52 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
@property
@abc.abstractmethod
def observation_delta_indices(self) -> list | None: # type: ignore[type-arg] #TODO: No implementation
"""Offsets, relative to the current step, of the observation timesteps the policy consumes.
`None` means only the current step is used.
"""
raise NotImplementedError
@property
@abc.abstractmethod
def action_delta_indices(self) -> list | None: # type: ignore[type-arg] #TODO: No implementation
"""Offsets, relative to the current step, of the action timesteps the policy predicts/consumes.
`None` means only the current step is used.
"""
raise NotImplementedError
@property
@abc.abstractmethod
def reward_delta_indices(self) -> list | None: # type: ignore[type-arg] #TODO: No implementation
"""Offsets, relative to the current step, of the reward timesteps the policy consumes.
`None` means only the current step is used.
"""
raise NotImplementedError
@abc.abstractmethod
def get_optimizer_preset(self) -> OptimizerConfig:
"""Return this policy's default `OptimizerConfig`, used when `use_policy_training_preset` is set."""
raise NotImplementedError
@abc.abstractmethod
def get_scheduler_preset(self) -> LRSchedulerConfig | None:
"""Return this policy's default `LRSchedulerConfig`, or `None` if it uses no scheduler."""
raise NotImplementedError
@abc.abstractmethod
def validate_features(self) -> None:
"""Check that `input_features`/`output_features` contain what this policy requires.
Raises:
ValueError: If a required feature is missing or has an unexpected shape/type.
"""
raise NotImplementedError
@property
def robot_state_feature(self) -> PolicyFeature | None:
"""The input `PolicyFeature` for the robot's proprioceptive state (`observation.state`), if any."""
if not self.input_features:
return None
for ft_name, ft in self.input_features.items():
@@ -140,6 +175,7 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
@property
def env_state_feature(self) -> PolicyFeature | None:
"""The input `PolicyFeature` of type `FeatureType.ENV` (environment state), if any."""
if not self.input_features:
return None
for _, ft in self.input_features.items():
@@ -149,12 +185,14 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
@property
def image_features(self) -> dict[str, PolicyFeature]:
"""All input features of type `FeatureType.VISUAL`, keyed by feature name."""
if not self.input_features:
return {}
return {key: ft for key, ft in self.input_features.items() if ft.type is FeatureType.VISUAL}
@property
def action_feature(self) -> PolicyFeature | None:
"""The output `PolicyFeature` for the action (`action`), if any."""
if not self.output_features:
return None
for ft_name, ft in self.output_features.items():
@@ -182,6 +220,35 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
revision: str | None = None,
**policy_kwargs: Any,
) -> T:
"""Download a policy's `config.json` from the Hub (or read it locally) and parse it.
The concrete policy config subclass is resolved from the serialized `"type"` tag (e.g. `"act"`,
`"diffusion"`) rather than being fixed by `cls`, so calling this on the `PreTrainedConfig` base
class works for any registered policy type.
Args:
pretrained_name_or_path (`str | Path`): Either the `repo_id` of the config hosted on the Hub,
or a path to a directory containing a `config.json` saved via `.save_pretrained`.
force_download (`bool`, *optional*, defaults to `False`): Whether to force (re-)downloading
the files from the Hub, overriding the existing cache.
resume_download (`bool | None`, *optional*): Deprecated; ignored by the underlying Hub client.
proxies (`dict[Any, Any] | None`, *optional*): A dictionary of proxy servers to use by protocol
or endpoint.
token (`str | bool | None`, *optional*): The token to use as HTTP bearer authorization for
remote files. By default, uses the token cached by `huggingface-cli login`.
cache_dir (`str | Path | None`, *optional*): Path to the folder where cached files are stored.
local_files_only (`bool`, *optional*, defaults to `False`): If `True`, avoid downloading the
file and return the path to the local cached file if it exists.
revision (`str | None`, *optional*): Revision on the Hub: a branch name, git tag, or commit id.
Defaults to the latest commit on `main`.
policy_kwargs: Forwarded as CLI-style overrides via `policy_kwargs["cli_overrides"]`
(a list of `--key=value` strings applied on top of the loaded config); any other keys are
ignored.
Raises:
FileNotFoundError: If `config.json` isn't found locally or on the Hub.
ValueError: If `config.json` has no `"type"` field, or its value isn't a registered policy type.
"""
model_id = str(pretrained_name_or_path)
config_file: str | None = None
if Path(model_id).is_dir():
src/lerobot/configs/rewards.py
+57 -7
View File
@@ -43,31 +43,45 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
"""Base configuration for reward models.
Args:
input_features: A dictionary defining the PolicyFeature of the input data for the reward. The key represents
the input data name, and the value is PolicyFeature, which consists of FeatureType and shape attributes.
output_features: A dictionary defining the PolicyFeature of the output data for the reward. The key represents
the output data name, and the value is PolicyFeature, which consists of FeatureType and shape attributes.
input_features (`dict[str, PolicyFeature]`, *optional*): A dictionary defining the `PolicyFeature`
of the input data for the reward. The key represents the input data name, and the value is a
`PolicyFeature`, which consists of `type` and `shape` attributes.
output_features (`dict[str, PolicyFeature]`, *optional*): A dictionary defining the `PolicyFeature`
of the output data for the reward, with the same key/value semantics as `input_features`.
device (`str | None`, *optional*): The torch device, e.g. `"cuda"`, `"cuda:0"`, `"cpu"`, or `"mps"`.
If unset or unavailable, `__post_init__` auto-selects one.
pretrained_path (`str | None`, *optional*): Either the repo ID of a model hosted on the Hub or a
path to a directory containing weights saved using `.save_pretrained`. If not provided, the
reward model is initialized from scratch.
pretrained_revision (`str | None`, *optional*): Optional Hub revision, e.g. a commit hash, branch,
or tag, to pin the pretrained reward model version.
push_to_hub (`bool`, *optional*, defaults to `False`): Whether to push the reward model to the
Hugging Face Hub after training.
repo_id (`str | None`, *optional*): The Hub repo ID to push to. Required when `push_to_hub` is
`True`.
license (`str | None`, *optional*): The license to add to the reward model on the Hub.
tags (`list[str] | None`, *optional*): Tags to add to the reward model on the Hub.
private (`bool | None`, *optional*): Whether to upload to a private repository on the Hugging Face
Hub.
"""
# Reuses PolicyFeature
input_features: dict[str, PolicyFeature] = field(default_factory=dict)
output_features: dict[str, PolicyFeature] = field(default_factory=dict)
device: str | None = None
pretrained_path: str | None = None
# Optional Hub revision (commit hash, branch, or tag) to pin the pretrained reward model version.
pretrained_revision: str | None = None
push_to_hub: bool = False
repo_id: str | None = None
# Hub metadata
license: str | None = None
tags: list[str] | None = None
private: bool | None = None
def __post_init__(self) -> None:
"""Auto-select `device` when unset or unavailable."""
if not self.device or not is_torch_device_available(self.device):
auto_device = auto_select_torch_device()
logger.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
@@ -75,6 +89,7 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
@property
def type(self) -> str:
"""The reward model's registered `draccus.ChoiceRegistry` name."""
choice_name = self.get_choice_name(self.__class__)
if not isinstance(choice_name, str):
raise TypeError(f"Expected string from get_choice_name, got {type(choice_name)}")
@@ -82,14 +97,17 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
@property
def observation_delta_indices(self) -> list | None: # type: ignore[type-arg]
"""`None`: reward models consume only the current observation timestep."""
return None
@property
def action_delta_indices(self) -> list | None: # type: ignore[type-arg]
"""`None`: reward models consume only the current action timestep."""
return None
@property
def reward_delta_indices(self) -> list | None: # type: ignore[type-arg]
"""`None`: reward models consume only the current reward timestep."""
return None
def get_optimizer_preset(self) -> OptimizerConfig | None:
@@ -97,9 +115,14 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
return None
def get_scheduler_preset(self) -> LRSchedulerConfig | None:
"""Default LR scheduler for this reward model. `None` here; overridden by subclasses that need one."""
return None
def validate_features(self) -> None:
"""Check that `input_features`/`output_features` contain what this reward model requires.
No-op here; overridden by subclasses that have required features.
"""
pass
def _save_pretrained(self, save_directory: Path) -> None:
@@ -122,6 +145,33 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
revision: str | None = None,
**reward_kwargs: Any,
) -> T:
"""Download a reward model's `config.json` from the Hub (or read it locally) and parse it.
The concrete reward-model config subclass is resolved from the serialized `"type"` tag, so
calling this on the `RewardModelConfig` base class works for any registered reward-model type.
Args:
pretrained_name_or_path (`str | Path`): Either the `repo_id` of the config hosted on the Hub,
or a path to a directory containing a `config.json` saved via `.save_pretrained`.
force_download (`bool`, *optional*, defaults to `False`): Whether to force (re-)downloading
the files from the Hub, overriding the existing cache.
resume_download (`bool | None`, *optional*): Deprecated; ignored by the underlying Hub client.
proxies (`dict[Any, Any] | None`, *optional*): A dictionary of proxy servers to use by protocol
or endpoint.
token (`str | bool | None`, *optional*): The token to use as HTTP bearer authorization for
remote files. By default, uses the token cached by `huggingface-cli login`.
cache_dir (`str | Path | None`, *optional*): Path to the folder where cached files are stored.
local_files_only (`bool`, *optional*, defaults to `False`): If `True`, avoid downloading the
file and return the path to the local cached file if it exists.
revision (`str | None`, *optional*): Revision on the Hub: a branch name, git tag, or commit id.
Defaults to the latest commit on `main`.
reward_kwargs: Forwarded as CLI-style overrides via `reward_kwargs["cli_overrides"]`
(a list of `--key=value` strings applied on top of the loaded config); any other keys are
ignored.
Raises:
FileNotFoundError: If `config.json` isn't found locally or on the Hub.
"""
model_id = str(pretrained_name_or_path)
config_file: str | None = None
if Path(model_id).is_dir():
src/lerobot/configs/train.py
+121 -32
View File
@@ -108,75 +108,119 @@ def _migrate_legacy_rabc_fields(config: dict[str, Any]) -> dict[str, Any] | None
@dataclass
class TrainPipelineConfig(HubMixin):
"""The top-level configuration for `lerobot-train`, parsed by draccus from CLI flags and/or a YAML file.
Args:
dataset (`DatasetConfig`): The dataset(s) to train on.
env (`envs.EnvConfig | None`, *optional*): The simulation environment to periodically evaluate the
policy in (see `env_eval_freq`). Required when `env_eval_freq > 0`.
policy (`PreTrainedConfig | None`, *optional*): The policy to train. Mutually exclusive with
`reward_model`.
reward_model (`RewardModelConfig | None`, *optional*): The reward model to train instead of a
policy. Mutually exclusive with `policy`.
output_dir (`Path | None`, *optional*): Where to save all of the run outputs. If you run another
training session with the same value its contents will be overwritten unless `resume` is set.
job_name (`str | None`, *optional*): A name for the run.
resume (`bool`, *optional*, defaults to `False`): Resume a previous run. Pass `--config_path`
pointing at either a local checkpoint's `train_config.json` or a Hub repo id holding
`checkpoints/<step>/` subtrees (the latest checkpoint is downloaded and resumed from). When
resuming, the default behavior is to use the configuration from the checkpoint, regardless of
what's provided with the training command at the time of resumption (CLI `--*` flags still
override).
seed (`int | None`, *optional*, defaults to 1000): Seed used for training (e.g. model
initialization, dataset shuffling) and for the evaluation environments.
cudnn_deterministic (`bool`, *optional*, defaults to `False`): Use deterministic cuDNN algorithms
for reproducibility. Disables `cudnn.benchmark` and may reduce training speed by ~10-20 percent.
num_workers (`int`, *optional*, defaults to 4): Number of workers for the dataloader.
batch_size (`int`, *optional*, defaults to 8): The training batch size.
prefetch_factor (`int`, *optional*, defaults to 4): Number of batches loaded in advance by each
dataloader worker.
persistent_workers (`bool`, *optional*, defaults to `True`): Keep dataloader worker processes alive
between epochs.
dataloader_multiprocessing_context (`str | None`, *optional*, defaults to `"spawn"`): DataLoader
worker start method. `"spawn"` is safer than `"fork"` with non-fork-safe libs (PyAV /
torchcodec / ffmpeg), but adds some worker-startup time per run since workers re-import modules
instead of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
when appropriate, or set it to `None` to use Python's platform default.
steps (`int`, *optional*, defaults to 100000): Total number of training steps.
env_eval_freq (`int`, *optional*, defaults to 20000): Run the policy in the simulation environment
every N steps to measure reward/success (0 = disabled).
log_freq (`int`, *optional*, defaults to 200): Logging frequency, in steps.
eval_steps (`int`, *optional*, defaults to 0): Compute eval loss on held-out episodes every N steps
(0 = disabled). Requires `eval_split > 0`.
max_eval_samples (`int`, *optional*, defaults to 0): Cap on total eval samples, split uniformly
across tasks (0 = use all held-out data).
tolerance_s (`float`, *optional*, defaults to 0.0001): Maximum timestamp difference tolerated when
loading dataset frames, in seconds.
save_checkpoint (`bool`, *optional*, defaults to `True`): Whether to save checkpoints during
training.
save_freq (`int`, *optional*, defaults to 20000): Save a checkpoint every `save_freq` training
iterations and after the last training step. A non-positive value disables periodic saving,
keeping only the final checkpoint.
checkpoint_format (`CheckpointFormat`, *optional*, defaults to `CheckpointFormat.SAFETENSORS`):
Model-artifact format inside checkpoints; non-default values require a sharded run.
use_policy_training_preset (`bool`, *optional*, defaults to `True`): Use the policy's own
optimizer/scheduler presets when `optimizer`/`scheduler` aren't explicitly set.
optimizer (`OptimizerConfig | None`, *optional*): The optimizer to use. Falls back to the policy's
preset when `use_policy_training_preset` is `True`.
scheduler (`LRSchedulerConfig | None`, *optional*): The learning-rate scheduler to use. Falls back
to the policy's preset when `use_policy_training_preset` is `True`.
parallelism (`ParallelismConfig`, *optional*): Process topology: `dp_replicate` / `dp_shard` for HSDP
and context-parallel degree placeholders.
accelerator (`AcceleratorConfig`, *optional*): Execution runtime handed to the Accelerator: mixed
precision, gradient accumulation, FSDP/DDP tuning knobs, compile & activation-checkpointing
placeholders.
eval (`EvalConfig`, *optional*): Settings for the periodic simulation-environment evaluation.
wandb (`WandBConfig`, *optional*): Weights & Biases logging settings.
peft (`PeftConfig | None`, *optional*): PEFT (e.g. LoRA) settings, when fine-tuning with adapters
instead of full-parameter training.
job (`JobConfig`, *optional*): Where to run training: locally (default), or an HF Jobs flavor.
save_checkpoint_to_hub (`bool`, *optional*, defaults to `False`): Push each saved checkpoint to the
Hub (`policy.repo_id`) as it is written, not just the final model (useful to monitor progress
mid-run). The final model is pushed regardless. Works the same locally and remotely.
sample_weighting (`SampleWeightingConfig | None`, *optional*): Sample weighting configuration (e.g.
for RA-BC training).
rename_map (`dict[str, str]`, *optional*): Rename map for the observation, to override the image
and state keys.
"""
dataset: DatasetConfig
env: envs.EnvConfig | None = None
policy: PreTrainedConfig | None = None
reward_model: RewardModelConfig | None = None
# Set `dir` to where you would like to save all of the run outputs. If you run another training session
# with the same value for `dir` its contents will be overwritten unless you set `resume` to true.
output_dir: Path | None = None
job_name: str | None = None
# Set `resume` to true to resume a previous run. Pass `--config_path` pointing at either a local
# checkpoint's train_config.json or a Hub repo id holding `checkpoints/<step>/` subtrees (the
# latest checkpoint is downloaded and resumed from). Note that when resuming, the default behavior
# is to use the configuration from the checkpoint, regardless of what's provided with the training
# command at the time of resumption (CLI `--*` flags still override).
resume: bool = False
# `seed` is used for training (eg: model initialization, dataset shuffling)
# AND for the evaluation environments.
seed: int | None = 1000
# Set to True to use deterministic cuDNN algorithms for reproducibility.
# This disables cudnn.benchmark and may reduce training speed by ~10-20 percent.
cudnn_deterministic: bool = False
# Number of workers for the dataloader.
num_workers: int = 4
batch_size: int = 8
prefetch_factor: int = 4
persistent_workers: bool = True
# DataLoader worker start method. "spawn" is safer than "fork" with
# non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
# worker-startup time per run since workers re-import modules instead
# of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
# when appropriate, or set it to `null` to use Python's platform default.
dataloader_multiprocessing_context: str | None = "spawn"
steps: int = 100_000
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
env_eval_freq: int = 20_000
log_freq: int = 200
# Compute eval loss on held-out episodes every N steps (0 = disabled). Requires eval_split > 0.
eval_steps: int = 0
# Cap on total eval samples, split uniformly across tasks (0 = use all held-out data).
max_eval_samples: int = 0
tolerance_s: float = 1e-4
save_checkpoint: bool = True
# Checkpoint is saved every `save_freq` training iterations and after the last training step.
# A non-positive value disables periodic saving, keeping only the final checkpoint.
save_freq: int = 20_000
# Model-artifact format inside checkpoints; non-default values require a sharded run.
checkpoint_format: CheckpointFormat = CheckpointFormat.SAFETENSORS
use_policy_training_preset: bool = True
optimizer: OptimizerConfig | None = None
scheduler: LRSchedulerConfig | None = None
# Process topology: dp_replicate / dp_shard (HSDP) and context-parallel degree placeholders.
parallelism: ParallelismConfig = field(default_factory=ParallelismConfig)
# Execution runtime handed to the Accelerator: mixed precision, gradient accumulation,
# FSDP/DDP tuning knobs, compile & activation-checkpointing placeholders.
accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig)
eval: EvalConfig = field(default_factory=EvalConfig)
wandb: WandBConfig = field(default_factory=WandBConfig)
peft: PeftConfig | None = None
# Where to run training (local default, or an HF Jobs flavor). See JobConfig.
job: JobConfig = field(default_factory=JobConfig)
# Push each saved checkpoint to the Hub (policy.repo_id) as it is written, not
# just the final model (useful to monitor progress mid-run). Optional; the
# final model is pushed regardless. Works the same locally and remotely.
save_checkpoint_to_hub: bool = False
# Sample weighting configuration (e.g., for RA-BC training)
sample_weighting: SampleWeightingConfig | None = None
# Rename map for the observation to override the image and state keys
rename_map: dict[str, str] = field(default_factory=dict)
checkpoint_path: Path | None = field(init=False, default=None)
@@ -262,6 +306,21 @@ class TrainPipelineConfig(HubMixin):
self.reward_model.pretrained_path = str(policy_dir)
def validate(self) -> None:
"""Resolve pretrained sources and cross-field defaults, and fail fast on invalid combinations.
Called by draccus after parsing. Resolves `--policy.path`/`--reward_model.path`/`resume` into a
loaded config, derives `job_name` and `output_dir` when unset, and applies the policy's
optimizer/scheduler presets when `use_policy_training_preset` is `True`.
Raises:
ValueError: On an unsupported `dataloader_multiprocessing_context`, neither `policy` nor
`reward_model` configured, a `rename_map` without a pretrained checkpoint, an unsplit
dataset with `eval_steps > 0`, a missing `repo_id` when pushing to the Hub, or
`save_checkpoint_to_hub` without `policy.repo_id` — or (see `_validate_distributed`) an
unsupported distributed-training combination.
FileExistsError: If `output_dir` already exists and `resume` is `False`.
NotImplementedError: If `dataset.repo_id` is a list (multi-dataset training).
"""
available_contexts = multiprocessing.get_all_start_methods()
if (
self.dataloader_multiprocessing_context is not None
@@ -389,6 +448,7 @@ class TrainPipelineConfig(HubMixin):
return ["policy", "reward_model"]
def to_dict(self) -> dict[str, Any]:
"""Encode the config to a plain, JSON-serializable dictionary (via `draccus.encode`)."""
return draccus.encode(self) # type: ignore[no-any-return] # because of the third-party library draccus uses Any as the return type
def _save_pretrained(self, save_directory: Path) -> None:
@@ -409,6 +469,35 @@ class TrainPipelineConfig(HubMixin):
revision: str | None = None,
**kwargs: Any,
) -> "TrainPipelineConfig":
"""Download a run's `train_config.json` from the Hub (or read it locally) and parse it.
Falls back to the latest checkpoint's config when the repo has no root `train_config.json` (a repo
of periodic checkpoints from an interrupted run), so a resume can start straight from
`--config_path=<repo>`. Legacy RA-BC fields in a JSON config are migrated to the current
`sample_weighting` schema.
Args:
pretrained_name_or_path (`str | Path`): Either the `repo_id` of the run hosted on the Hub, or
a path to a directory containing a `train_config.json` saved via `.save_pretrained`.
force_download (`bool`, *optional*, defaults to `False`): Whether to force (re-)downloading
the files from the Hub, overriding the existing cache.
resume_download (`bool | None`, *optional*): Deprecated; ignored by the underlying Hub client.
proxies (`dict[Any, Any] | None`, *optional*): A dictionary of proxy servers to use by protocol
or endpoint.
token (`str | bool | None`, *optional*): The token to use as HTTP bearer authorization for
remote files. By default, uses the token cached by `huggingface-cli login`.
cache_dir (`str | Path | None`, *optional*): Path to the folder where cached files are stored.
local_files_only (`bool`, *optional*, defaults to `False`): If `True`, avoid downloading the
file and return the path to the local cached file if it exists.
revision (`str | None`, *optional*): Revision on the Hub: a branch name, git tag, or commit id.
Defaults to the latest commit on `main`.
kwargs: Forwarded as CLI-style overrides via `kwargs["cli_args"]` (a list of `--key=value`
strings applied on top of the loaded config); any other keys are ignored.
Raises:
FileNotFoundError: If `train_config.json` isn't found locally, on the Hub, or on any checkpoint
within the Hub repo.
"""
model_id = str(pretrained_name_or_path)
config_file: str | None = None
if Path(model_id).is_dir():
src/lerobot/configs/types.py
+48
View File
@@ -18,6 +18,18 @@ from enum import Enum
class FeatureType(str, Enum):
"""The category of data a `PolicyFeature` represents.
**Attributes**:
- **STATE** -- A robot/environment proprioceptive state vector.
- **VISUAL** -- An image or video feature.
- **ENV** -- Environment-provided state, distinct from robot proprioception (e.g. simulation
environment state).
- **ACTION** -- An action vector.
- **REWARD** -- A scalar reward.
- **LANGUAGE** -- A natural-language feature (e.g. task instruction tokens).
"""
STATE = "STATE"
VISUAL = "VISUAL"
ENV = "ENV"
@@ -27,11 +39,28 @@ class FeatureType(str, Enum):
class PipelineFeatureType(str, Enum):
"""Which side of a processor pipeline a feature belongs to.
**Attributes**:
- **ACTION** -- The feature is part of the action space.
- **OBSERVATION** -- The feature is part of the observation space.
"""
ACTION = "ACTION"
OBSERVATION = "OBSERVATION"
class NormalizationMode(str, Enum):
"""The normalization strategy applied to a feature by a `NormalizerProcessorStep`.
**Attributes**:
- **MIN_MAX** -- Scale to `[-1, 1]` using the feature's min/max statistics.
- **MEAN_STD** -- Center and scale to unit variance using the feature's mean/std statistics.
- **IDENTITY** -- Leave the feature unchanged.
- **QUANTILES** -- Scale to `[-1, 1]` using the feature's 1st/99th percentile statistics.
- **QUANTILE10** -- Scale to `[-1, 1]` using the feature's 10th/90th percentile statistics.
"""
MIN_MAX = "MIN_MAX"
MEAN_STD = "MEAN_STD"
IDENTITY = "IDENTITY"
@@ -41,11 +70,30 @@ class NormalizationMode(str, Enum):
@dataclass
class PolicyFeature:
"""Describes one entry of a policy's input/output feature space.
Args:
type (`FeatureType`): The category of the feature.
shape (`tuple[int, ...]`): The feature's shape, excluding the batch dimension.
"""
type: FeatureType
shape: tuple[int, ...]
class RTCAttentionSchedule(str, Enum):
"""The prefix-attention weighting schedule used by the Real-Time Chunking (RTC) policy.
Controls how much weight is given to the previous action chunk's prediction versus the new one,
over the overlap region between consecutive chunks.
**Attributes**:
- **ZEROS** -- No prefix attention: weight is 1.0 before `start`, then 0.0.
- **ONES** -- Full prefix attention: weight is 1.0 up to `end`, then 0.0.
- **LINEAR** -- Linearly ramps the weight down from 1.0 to 0.0 between `start` and `end`.
- **EXP** -- Like `LINEAR`, but with an exponential (rather than linear) decay curve.
"""
ZEROS = "ZEROS"
ONES = "ONES"
LINEAR = "LINEAR"
src/lerobot/configs/video.py
+34 -17
View File
@@ -84,18 +84,33 @@ DEPTH_ENCODER_INFO_FIELD_NAMES: frozenset[str] = frozenset({"depth_min", "depth_
@dataclass
class VideoEncoderConfig:
"""Video encoder configuration."""
"""Video encoder configuration.
vcodec: str = "libsvtav1" # Video codec name. "auto" picks a hardware codec if available, else libsvtav1.
pix_fmt: str = "yuv420p" # Pixel format (e.g. yuv420p).
g: int | None = 2 # GOP size (keyframe interval).
crf: int | float | None = 30 # Quality level. Lower means better quality and larger files.
preset: int | str | None = None # Speed/quality preset. Accepted values are codec-specific.
fast_decode: int = 0 # Fast-decode tuning. Accepted values are codec-specific, 0 disables it.
Args:
vcodec (`str`, *optional*, defaults to `"libsvtav1"`): Video codec name. `"auto"` picks a hardware
codec if available, else `libsvtav1`.
pix_fmt (`str`, *optional*, defaults to `"yuv420p"`): Pixel format (e.g. `yuv420p`).
g (`int | None`, *optional*, defaults to 2): GOP size (keyframe interval).
crf (`int | float | None`, *optional*, defaults to 30): Quality level. Lower means better quality
and larger files.
preset (`int | str | None`, *optional*): Speed/quality preset. Accepted values are codec-specific.
fast_decode (`int`, *optional*, defaults to 0): Fast-decode tuning. Accepted values are
codec-specific; 0 disables it.
video_backend (`str`, *optional*, defaults to `"pyav"`): Encoding backend. Only `"pyav"` is
currently supported.
extra_options (`dict[str, Any]`, *optional*): Extra codec options merged last, e.g. `{"tune":
"film"}`.
"""
vcodec: str = "libsvtav1"
pix_fmt: str = "yuv420p"
g: int | None = 2
crf: int | float | None = 30
preset: int | str | None = None
fast_decode: int = 0
# TODO(CarolinePascal): add torchcodec support + find a way to unify the
# two backends (encoding and decoding).
video_backend: str = "pyav" # Encoding backend. Only "pyav" is currently supported.
# Extra codec options merged last, e.g. {"tune": "film"}.
video_backend: str = "pyav"
extra_options: dict[str, Any] = field(default_factory=dict)
# Source-data channel count this encoder is expected to handle. ``None``
@@ -104,6 +119,7 @@ class VideoEncoderConfig:
_DEFAULT_CHANNELS: ClassVar[int | None] = None
def __post_init__(self) -> None:
"""Resolve `vcodec` (e.g. `"auto"`), apply the libsvtav1 default preset, and validate the config."""
self.resolve_vcodec()
# Empty-constructor ergonomics: ``VideoEncoderConfig()`` must "just work".
if self.preset is None and self.vcodec == "libsvtav1":
@@ -112,9 +128,7 @@ class VideoEncoderConfig:
@classmethod
def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]:
"""Parse the ``video.*`` keys of a feature ``info`` block into
constructor kwargs.
"""
"""Parse the ``video.*`` keys of a feature ``info`` block into constructor kwargs."""
video_info = video_info or {}
kwargs: dict[str, Any] = {}
@@ -147,6 +161,7 @@ class VideoEncoderConfig:
Args:
encoders: List of encoder names to detect. If a string, it is converted to a list.
Returns:
List of available encoder names. If the video backend is not "pyav", returns an empty list.
"""
@@ -211,6 +226,7 @@ class VideoEncoderConfig:
opts: dict[str, Any] = {}
def set_if(key: str, value: Any) -> None:
"""Set `opts[key]` to `value` (stringified if `as_strings`), unless `value` is `None`."""
if value is not None:
opts[key] = value if not as_strings else str(value)
@@ -302,9 +318,10 @@ class DepthEncoderConfig(VideoEncoderConfig):
@classmethod
def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]:
"""Layer the depth-specific tuning (``depth_min`` / ``depth_max`` /
``shift`` / ``use_log``) on top of the base parser. Missing keys
fall back to the class defaults.
"""Layer the depth-specific tuning on top of the base parser.
Adds ``depth_min`` / ``depth_max`` / ``shift`` / ``use_log``. Missing keys fall back to the
class defaults.
"""
kwargs = super()._kwargs_from_video_info(video_info)
video_info = video_info or {}
@@ -328,8 +345,8 @@ def encoder_config_from_video_info(video_info: dict | None) -> VideoEncoderConfi
otherwise.
Args:
video_info: A feature's ``info`` dict as persisted in ``info.json``,
or ``None`` (treated as an empty dict).
video_info (`dict | None`): A feature's ``info`` dict as persisted in ``info.json``, or ``None``
(treated as an empty dict).
Returns:
A :class:`DepthEncoderConfig` for depth features, otherwise a
utils/check_docstrings.py
+1
View File
@@ -60,6 +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.configs",
]
# Objects that do not yet follow the standard, so the check can be green from day one. Removing an entry