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

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

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

Expands docs/source/api/configs.mdx from 5 documented classes to the
module's full public surface. Adds lerobot.configs to
check_docstrings.py's MODULES_TO_CHECK ratchet and removes the
module's ruff D-ignore. Docstrings and docstring-format changes only;
no behavioral changes.

.github/workflows/documentation-upload-pr.yml

@@ -33,7 +33,7 @@ jobs:
       github.event.workflow_run.event == 'pull_request' &&
       github.event.workflow_run.conclusion == 'success' &&
       github.repository == 'huggingface/lerobot'
-    uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@931031bf2b54aabb134ceb54980a6a2860a00f11  # main
+    uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@6108e850ae1cf2f71bb0815a600bcd50c39abfa7  # main
     with:
       package_name: lerobot
     secrets:

.github/workflows/documentation.yml

@@ -60,14 +60,10 @@ jobs:
       github.repository == 'huggingface/lerobot'
     permissions:
       contents: read
-    uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@931031bf2b54aabb134ceb54980a6a2860a00f11  # main
+    uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@6108e850ae1cf2f71bb0815a600bcd50c39abfa7  # main
     with:
       commit_sha: ${{ github.sha }}
       package: lerobot
-      # The shared workflow builds its venv with the runner's system Python, which is 3.10 on
-      # ubuntu-22.04. lerobot requires >=3.12, so without this the install fails during setup —
-      # before `pre_command` below ever runs. Added upstream in huggingface/doc-builder#808.
-      python_version: "3.12"
       # doc-builder ships a mock-deps registry entry for lerobot, so the reusable workflow takes its
       # "light install" path: `pip install ./lerobot --no-deps` plus a handful of real dependencies.
       # That is not enough to import lerobot — draccus runs `register_subclass` at import time and
@@ -96,12 +92,11 @@ jobs:
     permissions:
       contents: read
       pull-requests: write
-    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@931031bf2b54aabb134ceb54980a6a2860a00f11  # main
+    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@6108e850ae1cf2f71bb0815a600bcd50c39abfa7  # main
     with:
       commit_sha: ${{ github.event.pull_request.head.sha }}
       pr_number: ${{ github.event.number }}
       package: lerobot
       # See the comment on build_main_docs. The PR workflow passes its own `--version pr_<n>`, so no
       # additional_args are needed here.
-      python_version: "3.12"
       pre_command: uv pip install "./lerobot[dataset]"

docs/source/api/configs.mdx

@@ -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

docs/source/hardware_guide.mdx

@@ -62,10 +62,7 @@ Reference data points on a 4×H100 80 GB cluster (`accelerate launch --num_proce
 | `smolvla`   | 27m 49s    |      0.312 |           0.011 | ~80%     | `--policy.path=lerobot/smolvla_base`, `freeze_vision_encoder=false`, `train_expert_only=false`                                 |
 | `pi05`      | 3h 41m     |      2.548 |           0.014 | ~95%     | `--policy.pretrained_path=lerobot/pi05_base`, `gradient_checkpointing=true`, `dtype=bfloat16`, vision encoder + expert trained |
 
-Training logs separate the full iteration into `dataloading_s` (`next(dl_iter)`), `preprocessing_s`
-(image conversion and the policy pipeline), and `update_s` (the optimizer update). `step_s` covers all
-three and drives `samples_per_s`. The benchmark above predates this split, so its `dataloading_s` includes
-preprocessing.
+The `dataloading_s` vs. `update_s` ratio is the diagnostic that matters: when `dataloading_s` approaches `update_s`, more GPUs stop helping — your dataloader is the bottleneck and you should look at `--num_workers`, image resolution, and disk speed before adding compute.
 
 ### Schedule and checkpoints
 

docs/source/molmoact2.mdx

@@ -242,17 +242,6 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
   --repo-id=your_dataset
 ```
 
-Recording, resuming, and merging aggregate quantiles from per-episode summaries, so `meta/stats.json` ends up holding a conservative envelope (`min` for `q <= 50`, `max` for `q > 50`) rather than whole-dataset quantiles. To estimate the latter, scan every episode with a running histogram:
-
-```bash
-python src/lerobot/scripts/augment_dataset_quantile_stats.py \
-  --repo-id=your_dataset \
-  --overwrite \
-  --skip-images
-```
-
-`--skip-images` keeps the existing image statistics and avoids video decoding when only `STATE`/`ACTION` need recomputing, and `--root` reads a local dataset instead of the Hub. These values are histogram estimates, subject to discretization and rebinning error, so they can differ from the conservative ones — which changes MolmoAct2's normalized targets and therefore its loss scale. Statistics already saved inside an existing checkpoint are not affected.
-
 Alternatively, train MolmoAct2 with mean/std normalization:
 
 ```bash

docs/source/pi05.mdx

@@ -127,17 +127,6 @@ lerobot-edit-dataset \
 
 Or keep the dataset as-is and pass `--policy.normalization_mapping='{"ACTION": "MEAN_STD", "STATE": "MEAN_STD", "VISUAL": "IDENTITY"}'`.
 
-Recording, resuming, and merging aggregate quantiles from per-episode summaries, so `meta/stats.json` ends up holding a conservative envelope (`min` for `q <= 50`, `max` for `q > 50`) rather than whole-dataset quantiles. To estimate the latter, scan every episode with a running histogram:
-
-```bash
-python src/lerobot/scripts/augment_dataset_quantile_stats.py \
-    --repo-id=your_dataset \
-    --overwrite \
-    --skip-images
-```
-
-`--skip-images` keeps the existing image statistics and avoids video decoding when only `STATE`/`ACTION` need recomputing, and `--root` reads a local dataset instead of the Hub. These values are histogram estimates, subject to discretization and rebinning error, so they can differ from the conservative ones — which changes π₀.₅'s normalized targets and therefore its loss scale. Statistics already saved inside an existing checkpoint are not affected.
-
 ### Training Command Example
 
 The same finetune with the VLM frozen: less memory, at some cost in success rate. Swap `--dataset.repo_id` for your own dataset.

docs/source/policy_diffusion_README.md

@@ -2,25 +2,6 @@
 
 https://diffusion-policy.cs.columbia.edu
 
-## Training
-
-The reference implementation maintains an exponential moving average (EMA) of the policy weights during training and evaluates the EMA weights. To reproduce this behavior, enable the trainer's EMA shadow:
-
-```bash
-lerobot-train \
-  --policy.type=diffusion \
-  --ema.enable=true \
-  ...
-```
-
-Checkpoints then contain a directly loadable copy of the EMA weights next to the live ones, e.g. for evaluation:
-
-```bash
-lerobot-eval --policy.path=outputs/train/.../checkpoints/last/pretrained_model_ema ...
-```
-
-The EMA decay schedule (`--ema.inv_gamma`, `--ema.power`, ...) defaults to the reference implementation's values. For a constant decay instead of the warmup schedule (e.g. to match openpi's pi0/pi05 training), set `--ema.decay=0.99`.
-
 ## Citation
 
 ```bibtex

docs/source/policy_pi05_README.md

@@ -59,22 +59,6 @@ When `use_relative_actions=true`, the training script automatically:
 
 ---
 
-## EMA of the policy weights
-
-OpenPI maintains an exponential moving average of the weights during training (`ema_decay=0.99` by default) and keeps the EMA copy for inference. To reproduce this with the LeRobot trainer, enable the EMA shadow with a constant decay:
-
-```bash
-python -m lerobot.scripts.lerobot_train \
-  --policy.type=pi05 \
-  --dataset.repo_id=your_org/your_dataset \
-  --ema.enable=true \
-  --ema.decay=0.99
-```
-
-Checkpoints then contain a directly loadable copy of the EMA weights in `pretrained_model_ema/` next to the live ones. Note that the shadow is a full extra copy of the parameters on the GPU. Like OpenPI (which disables EMA in its LoRA configs), EMA is not supported together with PEFT adapters.
-
----
-
 ## Citation
 
 If you use this work, please cite both **OpenPI** and the π₀.₅ paper:

docs/source/writing_docstrings.mdx

@@ -26,7 +26,7 @@ def send_action(self, action: RobotAction, rate_hz: float = 30.0) -> RobotAction
         action (`dict[str, float]`):
             Target values keyed by motor name, e.g. `{"shoulder_pan.pos": 0.0}`. Keys must match the
             robot's action features.
-        rate_hz (`float`, *optional*, defaults to 30.0):
+        rate_hz (`float`, *optional*, defaults to `30.0`):
             Control loop frequency.
 
     Returns:
@@ -60,7 +60,7 @@ description, then the sections.
 ### The `Args:` line is machine-parsed
 
 ```
-name (`type`, *optional*, defaults to X):
+name (`type`, *optional*, defaults to `X`):
     Description, indented on its own line.
 ```
 
@@ -81,10 +81,6 @@ fails. Omit the clause entirely for required parameters:
 Types go in backticks. Use `*optional*` with no `defaults to` when the default is `None` or is otherwise not
 worth restating.
 
-The default value itself follows one rule, and the checker rewrites to match it: **numbers are bare,
-everything else is backticked** — `defaults to 30`, `defaults to 1e-05`, but `` defaults to `True` ``,
-`` defaults to `"socketcan"` ``. Booleans count as "everything else", not as numbers.
-
 ### `Returns:` is type-first
 
 One indented line, type first, then a colon, then the description:
@@ -183,33 +179,7 @@ Add files containing runnable examples to `utils/documentation_tests.txt`.
 
 Put examples on the three to five genuine entry points of a module. Examples on trivial accessors are noise.
 
-## Four patterns you will hit constantly
-
-### Constructor parameters go on the class
-
-**doc-builder renders a class from its class docstring and never reads `__init__.__doc__`.** An `Args:`
-block written on `__init__` is dropped from the page entirely — the parameter still appears in the rendered
-signature, but with no description beside it.
-
-Document constructor parameters in an `Args:` block on the **class** docstring:
-
-```python
-class SOFollower(Robot):
-    """A single SO-family follower arm.
-
-    Args:
-        config (`SOFollowerRobotConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
-    """
-
-    def __init__(self, config: SOFollowerRobotConfig):
-        super().__init__(config)
-```
-
-`__init__` then needs no docstring at all — `D107` is disabled repo-wide for exactly this reason. The
-payoff is not only that the parameters render: an `Args:` block on the class is checked against
-`inspect.signature(cls)` by `make check-docstrings`, so it cannot silently drift from the constructor. The
-same block on `__init__` is checked by nothing.
+## Three patterns you will hit constantly
 
 ### Config dataclasses
 

pyproject.toml

@@ -413,11 +413,6 @@ ignore = [
     # rendered page. Coverage of the things that do get rendered is enforced by interrogate instead.
     "D100",
     "D104",
-    # D107: `__init__` docstrings. doc-builder renders a class from its *class* docstring and never reads
-    # `__init__.__doc__`, so anything documented there is dropped from the page. Constructor parameters
-    # belong in an `Args:` block on the class, where they render and where `make check-docstrings`
-    # validates them against the signature.
-    "D107",
 ]
 
 [tool.ruff.lint.per-file-ignores]
@@ -444,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

@@ -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
@@ -22,7 +21,7 @@ Import them directly: ``from lerobot.configs.train import TrainPipelineConfig``
 """
 
 from .dataset import DatasetRecordConfig
-from .default import DatasetConfig, EMAConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
+from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
 from .recipe import MessageTurn, TrainingRecipe, load_recipe
 from .types import (
@@ -57,7 +56,6 @@ __all__ = [
     # Config classes
     "DatasetRecordConfig",
     "DatasetConfig",
-    "EMAConfig",
     "EvalConfig",
     "JobConfig",
     "MessageTurn",

src/lerobot/configs/accelerator.py

@@ -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

@@ -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

@@ -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:
@@ -140,108 +198,67 @@ class EvalConfig:
 
 
 @dataclass
-class EMAConfig:
-    """Exponential moving average (EMA) of the policy weights.
+class PeftConfig:
+    """PEFT (parameter-efficient fine-tuning) settings, e.g. LoRA adapters.
 
-    Standard practice for diffusion-style policies (Chi et al. 2023, "Diffusion Policy", section V.D):
-    the reference implementation enables it in every config and evaluates the EMA weights. Off by
-    default here because it keeps a second full copy of the parameters in memory.
+    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.
 
-    The decay follows the warmup schedule from diffusers' `EMAModel`:
-    `decay_t = 1 - (1 + t / inv_gamma) ** -power`, clamped to `[min_decay, max_decay]`.
-    The defaults mirror the reference implementation. Alternatively, set `decay` for a constant
-    decay at every step, as used by openpi for pi0/pi05 (`ema_decay=0.99`).
+    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`.
     """
 
-    enable: bool = False
-    # Constant decay coefficient (openpi-style, e.g. 0.99 for pi0/pi05). When set, the warmup
-    # schedule below is bypassed and the shadow uses this decay at every step.
-    decay: float | None = None
-    # Number of optimizer steps during which the shadow stays a hard copy of the live weights.
-    update_after_step: int = 0
-    # Warmup schedule parameters (see class docstring).
-    inv_gamma: float = 1.0
-    power: float = 0.75
-    min_decay: float = 0.0
-    max_decay: float = 0.9999
-    # Evaluate the EMA weights (instead of the live ones) during periodic env eval.
-    # Offline eval-loss (--eval_steps) always uses the live weights: it runs on every rank
-    # while the EMA shadow only lives on the main process.
-    use_for_eval: bool = True
-
-    def __post_init__(self) -> None:
-        if not (0.0 <= self.min_decay <= self.max_decay <= 1.0):
-            raise ValueError(
-                "Expected 0 <= ema.min_decay <= ema.max_decay <= 1, got "
-                f"min_decay={self.min_decay} and max_decay={self.max_decay}."
-            )
-        if self.inv_gamma <= 0:
-            raise ValueError(f"ema.inv_gamma must be positive, got {self.inv_gamma}.")
-        if self.power <= 0:
-            raise ValueError(f"ema.power must be positive, got {self.power}.")
-        if self.update_after_step < 0:
-            raise ValueError(f"ema.update_after_step must be >= 0, got {self.update_after_step}.")
-        if self.decay is not None:
-            if not 0.0 <= self.decay <= 1.0:
-                raise ValueError(f"ema.decay must be in [0, 1], got {self.decay}.")
-            # Keep the literals in sync with the field defaults above.
-            if self.min_decay != 0.0 or self.max_decay != 0.9999:
-                raise ValueError(
-                    "ema.decay (constant decay) and ema.min_decay/ema.max_decay (schedule clamp) are "
-                    "mutually exclusive: set one or the other."
-                )
-
-
-@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.
-
-    # 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

@@ -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

@@ -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,8 +388,8 @@ 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:
+    """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
@@ -391,8 +399,11 @@ def wrap(config_path: Path | None = None) -> Callable[[F], F]:
     """
 
     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

@@ -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

@@ -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

@@ -35,7 +35,7 @@ from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint
 from lerobot.utils.sample_weighting import SampleWeightingConfig
 
 from . import parser
-from .default import DatasetConfig, EMAConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
+from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
 from .rewards import RewardModelConfig
 
@@ -108,77 +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)
-    # Maintain an EMA shadow of the policy weights during training (see EMAConfig).
-    ema: EMAConfig = field(default_factory=EMAConfig)
     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)
 
@@ -264,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
@@ -391,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:
@@ -411,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

@@ -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

@@ -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

src/lerobot/datasets/compute_stats.py

@@ -613,15 +613,8 @@ def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, d
         for q_key in quantile_keys:
             if all(q_key in s for s in stats_ft_list):
                 quantile_values = np.stack([s[q_key] for s in stats_ft_list])
-                # Exact global quantiles cannot be recovered from quantile summaries.
-                # Keep a conservative envelope of the available estimates: min
-                # for lower quantiles and max for upper quantiles. The resulting
-                # values are bounds across the inputs, not global quantile estimates.
-                q_percent = int(q_key[1:])
-                if q_percent <= 50:
-                    aggregated[q_key] = np.min(quantile_values, axis=0)
-                else:
-                    aggregated[q_key] = np.max(quantile_values, axis=0)
+                weighted_quantiles = quantile_values * counts
+                aggregated[q_key] = weighted_quantiles.sum(axis=0) / total_count
 
     return aggregated
 

src/lerobot/distributed/factory.py

@@ -33,7 +33,11 @@ if TYPE_CHECKING:
     from lerobot.policies.pretrained import PreTrainedPolicy
 
 # Env vars through which `accelerate launch --config_file` (or a stray shell) would configure
-# accelerate behind the config system's back, making train_config.json lie about what ran.
+# accelerate behind the config system's back. Plugin `__post_init__`s read these silently as
+# field fallbacks (ACCELERATE_DYNAMO_* enables torch.compile through the default
+# TorchDynamoPlugin; ACCELERATE_GRADIENT_ACCUMULATION_STEPS overrides the explicitly passed
+# value inside Accelerator.__init__), which would make train_config.json lie about what ran.
+_ACCELERATE_ENV_PREFIXES = ("FSDP_", "PARALLELISM_CONFIG_", "ACCELERATE_DYNAMO_")
 _ACCELERATE_ENV_VARS = (
     "ACCELERATE_USE_FSDP",
     "ACCELERATE_USE_PARALLELISM_CONFIG",
@@ -55,7 +59,11 @@ def guard_against_env_interference() -> None:
     """
     if os.environ.get(_ENV_OVERRIDE):
         return
-    offending = sorted(name for name in _ACCELERATE_ENV_VARS if name in os.environ)
+    offending = sorted(
+        name
+        for name in os.environ
+        if name in _ACCELERATE_ENV_VARS or name.startswith(_ACCELERATE_ENV_PREFIXES)
+    )
     if offending:
         raise RuntimeError(
             f"Accelerate-configuring environment variables are set: {', '.join(offending)}. "

src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py

@@ -29,17 +29,18 @@ logger = logging.getLogger(__name__)
 
 
 class BiOpenArmFollower(BimanualMixin, Robot):
-    """A bimanual pair of OpenArm follower arms driven as one robot.
-
-    Args:
-        config (`BiOpenArmFollowerConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
-    """
+    """A bimanual pair of OpenArm follower arms driven as one robot."""
 
     config_class = BiOpenArmFollowerConfig
     name = "bi_openarm_follower"
 
     def __init__(self, config: BiOpenArmFollowerConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`BiOpenArmFollowerConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         super().__init__(config)
         self.config = config
 

src/lerobot/robots/bi_rebot_b601_follower/bi_rebot_b601_follower.py

@@ -33,16 +33,18 @@ class BiRebotB601Follower(BimanualMixin, Robot):
 
     Composes two single-arm :class:`RebotB601Follower` instances. Observation and
     action keys of each arm are namespaced with a ``left_`` / ``right_`` prefix.
-
-    Args:
-        config (`BiRebotB601FollowerConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
     """
 
     config_class = BiRebotB601FollowerConfig
     name = "bi_rebot_b601_follower"
 
     def __init__(self, config: BiRebotB601FollowerConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`BiRebotB601FollowerConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         super().__init__(config)
         self.config = config
 

src/lerobot/robots/bi_so_follower/bi_so_follower.py

@@ -29,17 +29,18 @@ logger = logging.getLogger(__name__)
 
 
 class BiSOFollower(BimanualMixin, Robot):
-    """A bimanual pair of [SO follower arms](https://github.com/TheRobotStudio/SO-ARM100) by TheRobotStudio.
-
-    Args:
-        config (`BiSOFollowerConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
-    """
+    """A bimanual pair of [SO follower arms](https://github.com/TheRobotStudio/SO-ARM100) by TheRobotStudio."""
 
     config_class = BiSOFollowerConfig
     name = "bi_so_follower"
 
     def __init__(self, config: BiSOFollowerConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`BiSOFollowerConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         super().__init__(config)
         self.config = config
 

src/lerobot/robots/earthrover_mini_plus/robot_earthrover_mini_plus.py

@@ -81,11 +81,6 @@ class EarthRoverMiniPlus(Robot):
     - Linear and angular velocity control
     - Battery and orientation telemetry
 
-    Args:
-        config (`EarthRoverMiniPlusConfig`):
-            The robot's configuration. Its `sdk_url` points at the Frodobots SDK server; there is no
-            serial port, since control goes over HTTP.
-
     **Attributes**:
         - **config** -- Robot configuration
         - **sdk_base_url** -- URL of the Frodobots SDK server (default: http://localhost:8000)
@@ -95,6 +90,11 @@ class EarthRoverMiniPlus(Robot):
     name = "earthrover_mini_plus"
 
     def __init__(self, config: EarthRoverMiniPlusConfig):
+        """Initialize EarthRover Mini Plus robot.
+
+        Args:
+            config: Robot configuration including SDK URL
+        """
         super().__init__(config)
         self.config = config
         self.sdk_base_url = "http://localhost:8000"

src/lerobot/robots/hope_jr/hope_jr_arm.py

@@ -39,16 +39,18 @@ class HopeJrArm(Robot):
 
     The arm and the hand are separate robots; pair this with [`~robots.hope_jr.HopeJrHand`] for a full
     limb. See [`~robots.Robot`] for the contract every method here implements.
-
-    Args:
-        config (`HopeJrArmConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
     """
 
     config_class = HopeJrArmConfig
     name = "hope_jr_arm"
 
     def __init__(self, config: HopeJrArmConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`HopeJrArmConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         super().__init__(config)
         self.config = config
         self.bus = FeetechMotorsBus(

src/lerobot/robots/hope_jr/hope_jr_hand.py

@@ -63,16 +63,18 @@ class HopeJrHand(Robot):
 
     Each hand is its own robot, so a two-handed setup uses two of these with different `side` values. See
     [`~robots.Robot`] for the contract every method here implements.
-
-    Args:
-        config (`HopeJrHandConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
     """
 
     config_class = HopeJrHandConfig
     name = "hope_jr_hand"
 
     def __init__(self, config: HopeJrHandConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`HopeJrHandConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         super().__init__(config)
         self.config = config
         self.bus = FeetechMotorsBus(

src/lerobot/robots/koch_follower/koch_follower.py

@@ -41,16 +41,18 @@ class KochFollower(Robot):
       wrist-to-elbow expansion, developed by Alexander Koch from
       [Tau Robotics](https://tau-robotics.com).
     - [Koch v1.1](https://github.com/jess-moss/koch-v1-1), developed by Jess Moss.
-
-    Args:
-        config (`KochFollowerConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
     """
 
     config_class = KochFollowerConfig
     name = "koch_follower"
 
     def __init__(self, config: KochFollowerConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`KochFollowerConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         super().__init__(config)
         self.config = config
         norm_mode_body = MotorNormMode.DEGREES if config.use_degrees else MotorNormMode.RANGE_M100_100

src/lerobot/robots/lekiwi/lekiwi.py

@@ -46,16 +46,18 @@ class LeKiwi(Robot):
     commands for the wheels.
 
     To drive one of these from another machine, use [`~robots.lekiwi.LeKiwiClient`].
-
-    Args:
-        config (`LeKiwiConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
     """
 
     config_class = LeKiwiConfig
     name = "lekiwi"
 
     def __init__(self, config: LeKiwiConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`LeKiwiConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         super().__init__(config)
         self.config = config
         norm_mode_body = MotorNormMode.DEGREES if config.use_degrees else MotorNormMode.RANGE_M100_100

src/lerobot/robots/lekiwi/lekiwi_client.py

@@ -36,16 +36,18 @@ class LeKiwiClient(Robot):
     Presents the same [`~robots.Robot`] interface as [`~robots.lekiwi.LeKiwi`], but every observation and
     action crosses a ZMQ connection to the host process running on the robot. Calibration stays on the
     robot, so this class does not perform it.
-
-    Args:
-        config (`LeKiwiClientConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
     """
 
     config_class = LeKiwiClientConfig
     name = "lekiwi_client"
 
     def __init__(self, config: LeKiwiClientConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`LeKiwiClientConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         import zmq
 
         self._zmq = zmq

src/lerobot/robots/lekiwi/lekiwi_host.py

@@ -40,13 +40,15 @@ class LeKiwiHost:
 
     Runs on the robot's own computer, receiving actions on one socket and publishing observations on
     another.
+    """
+
+    def __init__(self, config: LeKiwiHostConfig):
+        """Bind the command and observation sockets.
 
         Args:
             config (`LeKiwiHostConfig`):
                 Ports, loop frequency and watchdog settings for the host.
         """
-
-    def __init__(self, config: LeKiwiHostConfig):
         self.zmq_context = zmq.Context()
         self.zmq_cmd_socket = self.zmq_context.socket(zmq.PULL)
         self.zmq_cmd_socket.setsockopt(zmq.CONFLATE, 1)

src/lerobot/robots/omx_follower/omx_follower.py

@@ -39,16 +39,18 @@ class OmxFollower(Robot):
     """The [OpenMANIPULATOR-X](https://github.com/ROBOTIS-GIT/open_manipulator) follower arm.
 
     Developed by Woojin Wie and Junha Cha at [ROBOTIS](https://ai.robotis.com/).
-
-    Args:
-        config (`OmxFollowerConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
     """
 
     config_class = OmxFollowerConfig
     name = "omx_follower"
 
     def __init__(self, config: OmxFollowerConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`OmxFollowerConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         super().__init__(config)
         self.config = config
         norm_mode_body = MotorNormMode.DEGREES if config.use_degrees else MotorNormMode.RANGE_M100_100

src/lerobot/robots/openarm_follower/openarm_follower.py

@@ -41,16 +41,18 @@ class OpenArmFollower(Robot):
 
     Uses Damiao motors in MIT control mode. See [`~robots.Robot`] for the contract every method here
     implements.
-
-    Args:
-        config (`OpenArmFollowerConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
     """
 
     config_class = OpenArmFollowerConfig
     name = "openarm_follower"
 
     def __init__(self, config: OpenArmFollowerConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`OpenArmFollowerConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         super().__init__(config)
         self.config = config
 

src/lerobot/robots/reachy2/robot_reachy2.py

@@ -73,17 +73,18 @@ REACHY2_VEL = {
 
 
 class Reachy2Robot(Robot):
-    """[Reachy 2](https://www.pollen-robotics.com/reachy/), the humanoid by Pollen Robotics.
-
-    Args:
-        config (`Reachy2RobotConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
-    """
+    """[Reachy 2](https://www.pollen-robotics.com/reachy/), the humanoid by Pollen Robotics."""
 
     config_class = Reachy2RobotConfig
     name = "reachy2"
 
     def __init__(self, config: Reachy2RobotConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`Reachy2RobotConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         require_package("reachy2_sdk", extra="reachy2")
         super().__init__(config)
 

src/lerobot/robots/rebot_b601_follower/rebot_b601_follower.py

@@ -60,16 +60,18 @@ class RebotB601Follower(Robot):
 
     Motor communication is handled by the ``motorbridge`` package over a CAN bus,
     reached either through a Damiao serial bridge or a SocketCAN adapter.
-
-    Args:
-        config (`RebotB601FollowerRobotConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
     """
 
     config_class = RebotB601FollowerRobotConfig
     name = "rebot_b601_follower"
 
     def __init__(self, config: RebotB601FollowerRobotConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`RebotB601FollowerRobotConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         require_package("motorbridge", extra="rebot")
         super().__init__(config)
         self.config = config

src/lerobot/robots/robot.py

@@ -41,11 +41,6 @@ class Robot(abc.ABC):
     ...     robot.send_action(action)
     ```
 
-    Args:
-        config (`RobotConfig`):
-            The robot's configuration. Its `id` and `calibration_dir` decide where calibration is
-            read from and written to.
-
     **Attributes**:
         - **config_class** (`type[RobotConfig]`) -- The expected configuration class for this robot.
         - **name** (`str`) -- The unique robot name used to identify this robot type.
@@ -56,6 +51,13 @@ class Robot(abc.ABC):
     name: str
 
     def __init__(self, config: RobotConfig):
+        """Set up identity and calibration paths, loading an existing calibration file if there is one.
+
+        Args:
+            config (`RobotConfig`):
+                The robot's configuration. Its `id` and `calibration_dir` decide where calibration is
+                read from and written to.
+        """
         self.robot_type = self.name
         self.id = config.id
         self.calibration_dir = (

src/lerobot/robots/so_follower/so_follower.py

@@ -44,10 +44,6 @@ class SOFollower(Robot):
     Actions and observations are keyed `"<motor>.pos"`; cameras named in the config appear in observations
     under their own keys. See [`~robots.Robot`] for the contract every method here implements.
 
-    Args:
-        config (`SOFollowerRobotConfig`):
-            The robot's configuration. Its `port` and `cameras` determine what is connected.
-
     Example:
         ```python
         >>> from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig
@@ -62,6 +58,12 @@ class SOFollower(Robot):
     name = "so_follower"
 
     def __init__(self, config: SOFollowerRobotConfig):
+        """Build the robot from its configuration.
+
+        Args:
+            config (`SOFollowerRobotConfig`):
+                The robot's configuration. Its `port` and `cameras` determine what is connected.
+        """
         super().__init__(config)
         self.config = config
         # choose normalization mode depending on config if available

src/lerobot/robots/unitree_g1/g1_kinematics.py

@@ -24,16 +24,17 @@ logger = logging.getLogger(__name__)
 
 
 class WeightedMovingFilter:
-    """A fixed-length weighted moving average over recent samples, used to smooth IK solutions.
+    """A fixed-length weighted moving average over recent samples, used to smooth IK solutions."""
+
+    def __init__(self, weights, data_size=14):
+        """Set up the filter.
 
         Args:
-        weights (`Sequence[float]`):
+            weights:
                 Per-sample weights, newest first. Their length sets the window size.
             data_size (`int`, *optional*, defaults to 14):
                 Number of values in each sample.
         """
-
-    def __init__(self, weights, data_size=14):
         self._window_size = len(weights)
         self._weights = np.array(weights)
         self._data_size = data_size
@@ -75,14 +76,15 @@ class WeightedMovingFilter:
 
 
 class G1_29_ArmIK:  # noqa: N801
-    """Inverse kinematics for the G1's two arms, solved together as one optimisation problem.
+    """Inverse kinematics for the G1's two arms, solved together as one optimisation problem."""
+
+    def __init__(self, unit_test=False):
+        """Build the arm model and the IK solver.
 
         Args:
             unit_test (`bool`, *optional*, defaults to `False`):
                 Whether to run in test mode, which visualises the solution instead of driving a robot.
         """
-
-    def __init__(self, unit_test=False):
         import casadi
         import pinocchio as pin
         from huggingface_hub import snapshot_download

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -136,6 +136,13 @@ class UnitreeG1(Robot):
     `is_simulation=True` to drive a MuJoCo model instead of the physical robot.
 
     See [`~robots.Robot`] for the contract every method here implements.
+    """
+
+    config_class = UnitreeG1Config
+    name = "unitree_g1"
+
+    def __init__(self, config: UnitreeG1Config):
+        """Build the robot and, if one is configured, its locomotion controller.
 
         Args:
             config (`UnitreeG1Config`):
@@ -145,11 +152,6 @@ class UnitreeG1(Robot):
         Raises:
             ImportError: If the `unitree_g1` extra is not installed.
         """
-
-    config_class = UnitreeG1Config
-    name = "unitree_g1"
-
-    def __init__(self, config: UnitreeG1Config):
         require_package("unitree-sdk2py", extra="unitree_g1", import_name="unitree_sdk2py")
         super().__init__(config)
 

src/lerobot/robots/unitree_g1/unitree_sdk2_socket.py

@@ -44,28 +44,30 @@ class LowStateMsg:
     """
 
     class MotorState:
-        """Motor state data for a single joint.
+        """Motor state data for a single joint."""
+
+        def __init__(self, data: dict[str, Any]) -> None:
+            """Build one motor's state from a deserialized JSON frame.
 
             Args:
                 data (`dict[str, Any]`):
                     The motor's entry from the robot's state message.
             """
-
-        def __init__(self, data: dict[str, Any]) -> None:
             self.q: float = data.get("q", 0.0)
             self.dq: float = data.get("dq", 0.0)
             self.tau_est: float = data.get("tau_est", 0.0)
             self.temperature: float = data.get("temperature", 0.0)
 
     class IMUState:
-        """IMU sensor data.
+        """IMU sensor data."""
+
+        def __init__(self, data: dict[str, Any]) -> None:
+            """Build the IMU state from a deserialized JSON frame.
 
             Args:
                 data (`dict[str, Any]`):
                     The IMU's entry from the robot's state message.
             """
-
-        def __init__(self, data: dict[str, Any]) -> None:
             self.quaternion: list[float] = data.get("quaternion", [1.0, 0.0, 0.0, 0.0])
             self.gyroscope: list[float] = data.get("gyroscope", [0.0, 0.0, 0.0])
             self.accelerometer: list[float] = data.get("accelerometer", [0.0, 0.0, 0.0])
@@ -145,7 +147,10 @@ def ChannelFactoryInitialize(domain_id: int = 0, config: Any = None) -> None:  #
 
 
 class ChannelPublisher:
-    """ZMQ-based publisher that sends commands to the robot server.
+    """ZMQ-based publisher that sends commands to the robot server."""
+
+    def __init__(self, topic: str, msg_type: type) -> None:
+        """Bind the publisher to a topic.
 
         Args:
             topic (`str`):
@@ -153,8 +158,6 @@ class ChannelPublisher:
             msg_type (`type`):
                 The message class this topic carries.
         """
-
-    def __init__(self, topic: str, msg_type: type) -> None:
         self.topic = topic
         self.msg_type = msg_type
 
@@ -172,7 +175,10 @@ class ChannelPublisher:
 
 
 class ChannelSubscriber:
-    """ZMQ-based subscriber that receives state from the robot server.
+    """ZMQ-based subscriber that receives state from the robot server."""
+
+    def __init__(self, topic: str, msg_type: type) -> None:
+        """Bind the subscriber to a topic.
 
         Args:
             topic (`str`):
@@ -180,8 +186,6 @@ class ChannelSubscriber:
             msg_type (`type`):
                 The message class this topic carries.
         """
-
-    def __init__(self, topic: str, msg_type: type) -> None:
         self.topic = topic
         self.msg_type = msg_type
 

src/lerobot/scripts/augment_dataset_quantile_stats.py

@@ -25,11 +25,6 @@ quantile statistics (q01, q10, q50, q90, q99) in their metadata. This script:
 3. If missing, computes quantile statistics for all features
 4. Updates the dataset metadata with the new quantile statistics
 
-Statistics are accumulated into a single running histogram per feature across
-all episodes rather than aggregating per-episode quantile summaries. The
-resulting quantiles are histogram approximations, subject to discretization and
-range-rebinning error; image/video frames are sampled by default.
-
 Usage:
 
 ```bash
@@ -39,7 +34,9 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
 """
 
 import argparse
+import concurrent.futures
 import logging
+import os
 from pathlib import Path
 
 import numpy as np
@@ -52,10 +49,11 @@ from lerobot.datasets import (
     CODEBASE_VERSION,
     DEFAULT_QUANTILES,
     LeRobotDataset,
+    aggregate_stats,
     get_feature_stats,
     write_stats,
 )
-from lerobot.datasets.compute_stats import RunningQuantileStats, sample_indices
+from lerobot.datasets.compute_stats import sample_indices
 from lerobot.utils.utils import init_logging
 
 
@@ -81,25 +79,20 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
     return False
 
 
-def collect_episode_arrays(
-    dataset: LeRobotDataset,
-    episode_idx: int,
-    use_sampling: bool = True,
-    skip_images: bool = False,
-) -> dict[str, tuple[np.ndarray, int]]:
-    """Collect one episode's frames per feature, flattened to (num_samples, dim).
+def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
+    """Process a single episode and return its statistics.
 
     Args:
         dataset: The LeRobot dataset
-        episode_idx: Index of the episode to read
-        use_sampling: If True, sub-sample image/video frames to bound memory.
-            If False, use every frame (higher memory).
-        skip_images: If True, skip image/video features entirely.
+        episode_idx: Index of the episode to process
+        use_sampling: If True, sub-sample image/video frames per episode to bound
+            memory. If False, use every frame (exact, higher memory).
 
     Returns:
-        Mapping of feature name to that episode's values and the number of frames
-        they came from (which differs from the row count for image features).
+        Dictionary containing episode statistics
     """
+    logging.info(f"Computing stats for episode {episode_idx}")
+
     start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
     end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
 
@@ -109,9 +102,7 @@ def collect_episode_arrays(
     # numeric columns are cheap, so read them in full (exact).
     image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")]
     numeric_keys = [
-        k
-        for k in dataset.features
-        if dataset.features[k]["dtype"] not in ("image", "video", "string", "language")
+        k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
     ]
 
     collected_data: dict[str, list] = {}
@@ -123,7 +114,7 @@ def collect_episode_arrays(
             collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]]
 
     # Image/video features: decode only a sampled subset of frames.
-    if image_keys and not skip_images:
+    if image_keys:
         sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len))
         for offset in sampled_offsets:
             item = dataset[start_idx + offset]
@@ -131,82 +122,87 @@ def collect_episode_arrays(
                 if key in item:
                     collected_data.setdefault(key, []).append(item[key])
 
-    episode_arrays: dict[str, tuple[np.ndarray, int]] = {}
+    ep_stats = {}
     for key, data_list in collected_data.items():
+        if dataset.features[key]["dtype"] == "string":
+            continue
+
         data = torch.stack(data_list).cpu().numpy()
         if dataset.features[key]["dtype"] in ["image", "video"]:
             if data.dtype == np.uint8:
                 data = data.astype(np.float32) / 255.0
-            # (N, C, H, W) -> (N * H * W, C) so quantiles are computed per channel.
-            channels = data.shape[1]
-            values = data.transpose(0, 2, 3, 1).reshape(-1, channels)
+
+            axes_to_reduce = (0, 2, 3)
+            keepdims = True
         else:
-            values = data.reshape(-1, data.shape[-1]) if data.ndim > 1 else data.reshape(-1, 1)
-        episode_arrays[key] = (values, len(data_list))
+            axes_to_reduce = 0
+            keepdims = data.ndim == 1
 
-    return episode_arrays
+        ep_stats[key] = get_feature_stats(
+            data, axis=axes_to_reduce, keepdims=keepdims, quantile_list=DEFAULT_QUANTILES
+        )
+
+        if dataset.features[key]["dtype"] in ["image", "video"]:
+            ep_stats[key] = {
+                k: v if k == "count" else np.squeeze(v, axis=0) for k, v in ep_stats[key].items()
+            }
+
+    return ep_stats
 
 
-def compute_quantile_stats_for_dataset(
-    dataset: LeRobotDataset,
-    use_sampling: bool = True,
-    skip_images: bool = False,
-) -> dict[str, dict]:
-    """Compute whole-dataset statistics with one running histogram per feature.
+def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
+    """Compute quantile statistics for all episodes in the dataset.
 
     Args:
         dataset: The LeRobot dataset to compute statistics for
         use_sampling: If True, sub-sample image/video frames per episode to bound
-            memory. If False, use every frame (higher memory).
-        skip_images: If True, skip image/video features and leave their stats untouched.
+            memory. If False, use every frame (exact, higher memory).
 
     Returns:
-        Dictionary containing statistics with histogram-based global quantile estimates
+        Dictionary containing aggregated statistics with quantiles
 
     Note:
-        Episodes are accumulated sequentially because the running accumulators are
-        shared across all of them.
+        Video decoding operations are not thread-safe, so we process episodes sequentially
+        when video keys are present. For datasets without videos, we use parallel processing
+        with ThreadPoolExecutor for better performance.
     """
     logging.info(f"Computing quantile statistics for dataset with {dataset.num_episodes} episodes")
 
-    running_stats: dict[str, RunningQuantileStats] = {}
-    frame_counts: dict[str, int] = {}
-    row_counts: dict[str, int] = {}
-    # Kept only while a feature has a single row, so it can still be finalized.
-    single_row_arrays: dict[str, np.ndarray] = {}
+    episode_stats_list = []
+    has_videos = len(dataset.meta.video_keys) > 0
 
+    if has_videos:
+        logging.info("Dataset contains video keys - using sequential processing for thread safety")
         for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
-        episode_arrays = collect_episode_arrays(
-            dataset, episode_idx, use_sampling=use_sampling, skip_images=skip_images
-        )
-        for key, (array, num_frames) in episode_arrays.items():
-            running_stats.setdefault(key, RunningQuantileStats()).update(array)
-            frame_counts[key] = frame_counts.get(key, 0) + num_frames
-            row_counts[key] = row_counts.get(key, 0) + len(array)
-            if row_counts[key] < 2:
-                single_row_arrays[key] = array
+            ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
+            episode_stats_list.append(ep_stats)
     else:
-                single_row_arrays.pop(key, None)
+        logging.info("Dataset has no video keys - using parallel processing for better performance")
+        max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))
 
-    if not running_stats:
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+            future_to_episode = {
+                executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
+                for episode_idx in range(dataset.num_episodes)
+            }
+
+            episode_results = {}
+            with tqdm(total=dataset.num_episodes, desc="Processing episodes") as pbar:
+                for future in concurrent.futures.as_completed(future_to_episode):
+                    episode_idx = future_to_episode[future]
+                    ep_stats = future.result()
+                    episode_results[episode_idx] = ep_stats
+                    pbar.update(1)
+
+        for episode_idx in range(dataset.num_episodes):
+            if episode_idx in episode_results:
+                episode_stats_list.append(episode_results[episode_idx])
+
+    if not episode_stats_list:
         raise ValueError("No episode data found for computing statistics")
 
-    aggregated_stats: dict[str, dict] = {}
-    for key, accumulator in running_stats.items():
-        if row_counts[key] < 2:
-            # Histograms need at least two samples; mirror get_feature_stats' basic-stats path.
-            stats = get_feature_stats(single_row_arrays[key], axis=0, keepdims=False)
-        else:
-            stats = accumulator.get_statistics()
-        if dataset.features[key]["dtype"] in ["image", "video"]:
-            # Image stats are stored as (C, 1, 1) to broadcast over height and width.
-            stats = {k: v if k == "count" else v[:, np.newaxis, np.newaxis] for k, v in stats.items()}
-        # `get_feature_stats` counts frames, not the per-channel rows the accumulator sees.
-        stats["count"] = np.array([frame_counts[key]])
-        aggregated_stats[key] = stats
-
-    logging.info(f"Computed global histogram statistics for {len(aggregated_stats)} features")
-    return aggregated_stats
+    logging.info(f"Aggregating statistics from {len(episode_stats_list)} episodes")
+    return aggregate_stats(episode_stats_list)
 
 
 def augment_dataset_with_quantile_stats(
@@ -214,7 +210,6 @@ def augment_dataset_with_quantile_stats(
     root: str | Path | None = None,
     overwrite: bool = False,
     use_sampling: bool = True,
-    skip_images: bool = False,
 ) -> None:
     """Augment a dataset with quantile statistics if they are missing.
 
@@ -223,8 +218,7 @@ def augment_dataset_with_quantile_stats(
         root: Local root directory for the dataset
         overwrite: Overwrite existing quantile statistics if they already exist
         use_sampling: If True, sub-sample image/video frames per episode to bound
-            memory. If False, use every frame (higher memory).
-        skip_images: If True, skip image/video features and keep their existing stats
+            memory. If False, use every frame (exact, higher memory).
     """
     logging.info(f"Loading dataset: {repo_id}")
     dataset = LeRobotDataset(
@@ -238,13 +232,7 @@ def augment_dataset_with_quantile_stats(
 
     logging.info("Dataset does not contain quantile statistics. Computing them now...")
 
-    new_stats = compute_quantile_stats_for_dataset(
-        dataset, use_sampling=use_sampling, skip_images=skip_images
-    )
-
-    if skip_images and dataset.meta.stats:
-        for key, feature_stats in dataset.meta.stats.items():
-            new_stats.setdefault(key, feature_stats)
+    new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
 
     logging.info("Updating dataset metadata with new quantile statistics")
     dataset.meta.stats = new_stats
@@ -288,15 +276,10 @@ def main():
         "--no-sampling",
         action="store_true",
         help=(
-            "Compute stats over every frame (higher memory). By default, "
+            "Compute stats over every frame (exact, higher memory). By default, "
             "image/video frames are sub-sampled per episode to bound memory."
         ),
     )
-    parser.add_argument(
-        "--skip-images",
-        action="store_true",
-        help="Skip image/video features and preserve their existing stats",
-    )
 
     args = parser.parse_args()
     root = Path(args.root) if args.root else None
@@ -308,7 +291,6 @@ def main():
         root=root,
         overwrite=args.overwrite,
         use_sampling=not args.no_sampling,
-        skip_images=args.skip_images,
     )
 
 

src/lerobot/scripts/lerobot_train.py

@@ -76,7 +76,7 @@ from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_proces
 from lerobot.policies.factory import ProcessorConfigKwargs
 from lerobot.rewards import make_reward_pre_post_processors
 from lerobot.utils.collate import lerobot_collate_fn
-from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR
+from lerobot.utils.constants import TRAINING_STATE_DIR
 from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 from lerobot.utils.random_utils import set_seed
@@ -95,20 +95,6 @@ else:
 
 from .lerobot_eval import eval_policy_all
 
-EMA_STATE_FILENAME = "ema_state.pt"
-
-
-@contextmanager
-def _ema_weights(ema: Any, policy: PreTrainedPolicy) -> Iterator[None]:
-    """Temporarily swap the EMA shadow weights into `policy`, restoring the live ones on exit."""
-    params = list(policy.parameters())
-    ema.store(params)
-    ema.copy_to(params)
-    try:
-        yield
-    finally:
-        ema.restore(params)
-
 
 @contextmanager
 def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
@@ -606,65 +592,6 @@ def train(cfg: TrainPipelineConfig):
     dl_iter = cycle(dataloader)
     policy.train()
 
-    # EMA shadow of the policy weights (Chi et al. 2023, Diffusion Policy, section V.D). The shadow
-    # lives on the main process only, which is safe under DDP where every rank holds identical
-    # weights after each gradient sync. diffusers is imported lazily so the base training path does
-    # not depend on it.
-    ema = None
-    if cfg.ema.enable:
-        if parallel_dims.is_sharded:
-            raise NotImplementedError(
-                "--ema.enable=true is not supported with sharded training (FSDP2/HSDP/CP): the "
-                "parameters are sharded across ranks. Use a replicated (DDP) or single-GPU run."
-            )
-        if cfg.peft is not None:
-            raise NotImplementedError("--ema.enable=true is not supported together with PEFT adapters.")
-        require_package("diffusers", extra="diffusion")
-        if is_main_process():
-            from diffusers.training_utils import EMAModel  # noqa: PLC0415
-
-            # A constant --ema.decay is expressed through the schedule clamp: with
-            # min_decay == max_decay, the warmup curve is pinned to that value at every step.
-            min_decay = cfg.ema.min_decay if cfg.ema.decay is None else cfg.ema.decay
-            max_decay = cfg.ema.max_decay if cfg.ema.decay is None else cfg.ema.decay
-            ema = EMAModel(
-                accelerator.unwrap_model(policy).parameters(),
-                decay=max_decay,
-                min_decay=min_decay,
-                update_after_step=cfg.ema.update_after_step,
-                use_ema_warmup=True,
-                inv_gamma=cfg.ema.inv_gamma,
-                power=cfg.ema.power,
-            )
-            ema.to(device)
-            if cfg.ema.decay is not None:
-                logging.info(
-                    "EMA enabled: decay=%g (constant), update_after_step=%d, use_for_eval=%s",
-                    cfg.ema.decay,
-                    cfg.ema.update_after_step,
-                    cfg.ema.use_for_eval,
-                )
-            else:
-                logging.info(
-                    "EMA enabled: max_decay=%g, inv_gamma=%g, power=%g, update_after_step=%d, use_for_eval=%s",
-                    cfg.ema.max_decay,
-                    cfg.ema.inv_gamma,
-                    cfg.ema.power,
-                    cfg.ema.update_after_step,
-                    cfg.ema.use_for_eval,
-                )
-            if cfg.checkpoint_path is not None:
-                ema_path = cfg.checkpoint_path / TRAINING_STATE_DIR / EMA_STATE_FILENAME
-                if ema_path.exists():
-                    ema.load_state_dict(torch.load(ema_path, map_location=device, weights_only=True))
-                    logging.info("Resumed EMA shadow from %s", ema_path)
-                else:
-                    logging.warning(
-                        "Resuming with --ema.enable=true but %s is missing; "
-                        "restarting the shadow from the current weights.",
-                        ema_path,
-                    )
-
     train_metrics = {
         # Per-rank loss reflects only one shard of the global batch; mean recovers the loss the
         # data-parallel group is actually optimizing. grad_norm and lr are already identical on
@@ -675,10 +602,9 @@ def train(cfg: TrainPipelineConfig):
         "lr": AverageMeter("lr", ":0.1e"),
         # Report the slowest rank for bottleneck-style timings so multi-GPU runs surface the
         # true straggler instead of rank 0's view.
-        "dataloading_s": AverageMeter("data_s", ":.3f", reduction="max"),
-        "preprocessing_s": AverageMeter("prep_s", ":.3f", reduction="max"),
         "update_s": AverageMeter("updt_s", ":.3f", reduction="max"),
-        "step_s": AverageMeter("step_s", ":.3f", reduction="max"),
+        "dataloading_s": AverageMeter("data_s", ":.3f", reduction="max"),
+        # Derived from the post-reduce max step time; set once per log window on the main rank.
         "samples_per_s": AverageMeter("smp/s", ":.0f"),
     }
     if torch.cuda.is_available():
@@ -708,15 +634,13 @@ def train(cfg: TrainPipelineConfig):
         )
 
     for _ in range(step, cfg.steps):
-        step_start = time.perf_counter()
+        start_time = time.perf_counter()
         batch = next(dl_iter)
-        preprocessing_start = time.perf_counter()
-        train_tracker.dataloading_s = preprocessing_start - step_start
         for cam_key in dataset.meta.camera_keys:
             if cam_key in batch and batch[cam_key].dtype == torch.uint8:
                 batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0
         batch = preprocessor(batch)
-        train_tracker.preprocessing_s = time.perf_counter() - preprocessing_start
+        train_tracker.dataloading_s = time.perf_counter() - start_time
 
         train_tracker, _ = update_policy(
             train_tracker,
@@ -728,13 +652,6 @@ def train(cfg: TrainPipelineConfig):
             lr_scheduler=lr_scheduler,
             sample_weighter=sample_weighter,
         )
-        train_tracker.step_s = time.perf_counter() - step_start
-
-        # Pull one optimizer step of the live weights into the EMA shadow (main process only).
-        # The shadow tracks optimizer updates, not micro-batches: gate on the sync step under
-        # gradient accumulation.
-        if ema is not None and accelerator.sync_gradients:
-            ema.step(accelerator.unwrap_model(policy).parameters())
 
         # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
         # increment `step` here.
@@ -751,8 +668,11 @@ def train(cfg: TrainPipelineConfig):
             # Collective reduce must run on every rank, before the main-process gate below.
             train_tracker.reduce_across_ranks()
             if is_main_process():
-                if train_tracker.step_s.avg > 0:
-                    train_tracker.samples_per_s = samples_per_step / train_tracker.step_s.avg
+                # Cluster-wide throughput, derived from the already-reduced (max) step time so it
+                # reflects the slowest rank — which is what actually gates the next iteration.
+                step_time = train_tracker.update_s.avg + train_tracker.dataloading_s.avg
+                if step_time > 0:
+                    train_tracker.samples_per_s = samples_per_step / step_time
                 logging.info(train_tracker)
                 if wandb_logger:
                     # Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the
@@ -763,9 +683,6 @@ def train(cfg: TrainPipelineConfig):
                     if sample_weighter is not None:
                         weighter_stats = sample_weighter.get_stats()
                         wandb_log_dict.update({f"sample_weighting/{k}": v for k, v in weighter_stats.items()})
-                    if ema is not None and ema.cur_decay_value is not None:
-                        wandb_log_dict["ema/decay"] = ema.cur_decay_value
-                        wandb_log_dict["ema/step"] = ema.optimization_step
                     wandb_logger.log_dict(wandb_log_dict, step)
             train_tracker.reset_averages()
 
@@ -810,17 +727,6 @@ def train(cfg: TrainPipelineConfig):
                 accelerator=accelerator,
             )
             if is_main_process():
-                if ema is not None:
-                    # Save the shadow for exact resume, plus a directly loadable copy of the EMA
-                    # weights (lerobot-eval --policy.path=<checkpoint>/pretrained_model_ema).
-                    torch.save(ema.state_dict(), checkpoint_dir / TRAINING_STATE_DIR / EMA_STATE_FILENAME)
-                    unwrapped_policy = accelerator.unwrap_model(policy)
-                    ema_dir = checkpoint_dir / f"{PRETRAINED_MODEL_DIR}_ema"
-                    with _ema_weights(ema, unwrapped_policy):
-                        unwrapped_policy.save_pretrained(ema_dir)
-                        cfg.save_pretrained(ema_dir)
-                        preprocessor.save_pretrained(ema_dir)
-                        postprocessor.save_pretrained(ema_dir)
                 update_last_checkpoint(checkpoint_dir)
                 if cfg.save_checkpoint_to_hub:
                     push_checkpoint_to_hub(
@@ -836,18 +742,10 @@ def train(cfg: TrainPipelineConfig):
             if is_main_process():
                 step_id = get_step_identifier(step, cfg.steps)
                 logging.info(f"Eval policy at step {step}")
-                eval_policy_model = accelerator.unwrap_model(policy)
-                # Evaluate the EMA weights when enabled: the swap happens only on the main
-                # process (the other ranks wait at the barrier below) and is exactly undone
-                # afterwards, so the live weights stay in sync across ranks.
-                use_ema_for_eval = ema is not None and cfg.ema.use_for_eval
-                if use_ema_for_eval:
-                    logging.info("Evaluating the EMA weights")
-                weights_cm = _ema_weights(ema, eval_policy_model) if use_ema_for_eval else nullcontext()
-                with weights_cm, _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
+                with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
                     eval_info = eval_policy_all(
                         envs=eval_env,  # dict[suite][task_id] -> vec_env
-                        policy=eval_policy_model,
+                        policy=accelerator.unwrap_model(policy),
                         env_preprocessor=env_preprocessor,
                         env_postprocessor=env_postprocessor,
                         preprocessor=preprocessor,
@@ -906,25 +804,6 @@ def train(cfg: TrainPipelineConfig):
             peft_model=unwrapped if peft_model is not None else None,
         )
 
-        # The push above ships the live weights; when EMA is on, the weights that were
-        # evaluated are the shadow, so push those too under a sibling `<repo_id>-ema` repo.
-        # The shadow lives on the main process only, so this is rank-0-only by construction.
-        # Non-fatal: the live model is already up if this fails.
-        if ema is not None:
-            ema_repo_id = f"{active_cfg.repo_id}-ema"
-            orig_repo_id = unwrapped.config.repo_id
-            try:
-                unwrapped.config.repo_id = ema_repo_id
-                with _ema_weights(ema, unwrapped):
-                    unwrapped.push_model_to_hub(cfg, dataset_meta=dataset.meta)
-                preprocessor.push_to_hub(ema_repo_id)
-                postprocessor.push_to_hub(ema_repo_id)
-                logging.info("Pushed EMA weights to %s", ema_repo_id)
-            except Exception as exc:  # noqa: BLE001
-                logging.warning("Failed to push EMA weights to %s: %s", ema_repo_id, exc)
-            finally:
-                unwrapped.config.repo_id = orig_repo_id
-
     # Properly clean up the distributed process group
     accelerator.wait_for_everyone()
     accelerator.end_training()

tests/datasets/test_augment_quantile_stats.py

@@ -12,11 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from types import SimpleNamespace
-
 import numpy as np
 import pytest
-import torch
 
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 
@@ -27,9 +24,7 @@ from lerobot.scripts.augment_dataset_quantile_stats import (
 
 
 def _numeric_keys(dataset):
-    return [
-        k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string", "language")
-    ]
+    return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
 
 
 def _image_keys(dataset):
@@ -107,112 +102,3 @@ def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory)
     )
     stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
     assert has_quantile_stats(stats)
-
-
-class FakeHFDataset:
-    """Minimal stand-in exposing the column slicing used by the augment script."""
-
-    def __init__(self, columns: dict[str, list]):
-        self._columns = columns
-
-    def select_columns(self, keys):
-        return FakeHFDataset({key: self._columns[key] for key in keys})
-
-    def __getitem__(self, index):
-        return {key: values[index] for key, values in self._columns.items()}
-
-
-def test_compute_quantile_stats_skips_language_features():
-    class FakeDataset:
-        num_episodes = 1
-        features = {
-            "action": {"dtype": "float32"},
-            "observation.language": {"dtype": "language"},
-        }
-        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
-        hf_dataset = FakeHFDataset(
-            {
-                "action": [[0.0], [1.0]],
-                "observation.language": [
-                    [{"role": "user", "content": "pick"}],
-                    [{"role": "assistant", "content": "done"}],
-                ],
-            }
-        )
-
-    stats = compute_quantile_stats_for_dataset(FakeDataset())
-
-    assert set(stats) == {"action"}
-
-
-def test_compute_quantile_stats_skip_images_avoids_decoding():
-    class FakeDataset:
-        num_episodes = 1
-        features = {
-            "action": {"dtype": "float32"},
-            "observation.images.cam": {"dtype": "video"},
-        }
-        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
-        hf_dataset = FakeHFDataset({"action": [[0.0], [1.0]]})
-
-        def __getitem__(self, index):
-            raise AssertionError(f"video frame {index} was decoded despite skip_images=True")
-
-    stats = compute_quantile_stats_for_dataset(FakeDataset(), skip_images=True)
-
-    assert set(stats) == {"action"}
-
-
-def test_compute_quantile_stats_handles_single_frame():
-    class FakeDataset:
-        num_episodes = 1
-        features = {"action": {"dtype": "float32"}}
-        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 1}])
-        hf_dataset = FakeHFDataset({"action": [[5.0, 7.0]]})
-
-    stats = compute_quantile_stats_for_dataset(FakeDataset())
-
-    np.testing.assert_array_equal(stats["action"]["count"], np.array([1]))
-    for key in ("min", "max", "mean", "q01", "q10", "q50", "q90", "q99"):
-        np.testing.assert_allclose(stats["action"][key], np.array([5.0, 7.0]))
-
-
-def test_compute_quantile_stats_image_count_uses_frames():
-    frames = [torch.zeros(3, 2, 2), torch.ones(3, 2, 2)]
-
-    class FakeDataset:
-        num_episodes = 1
-        features = {"observation.images.cam": {"dtype": "video"}}
-        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
-        hf_dataset = FakeHFDataset({})
-
-        def __getitem__(self, index):
-            return {"observation.images.cam": frames[index]}
-
-    stats = compute_quantile_stats_for_dataset(FakeDataset(), use_sampling=False)
-    image_stats = stats["observation.images.cam"]
-
-    np.testing.assert_array_equal(image_stats["count"], np.array([2]))
-    assert image_stats["mean"].shape == (3, 1, 1)
-    np.testing.assert_allclose(image_stats["mean"], np.full((3, 1, 1), 0.5))
-
-
-def test_compute_quantile_stats_accumulates_across_episodes():
-    values = [[float(value)] for value in range(100)] + [[float(value)] for value in range(1000, 1010)]
-
-    class FakeDataset:
-        num_episodes = 2
-        features = {"action": {"dtype": "float32"}}
-        meta = SimpleNamespace(
-            episodes=[
-                {"dataset_from_index": 0, "dataset_to_index": 100},
-                {"dataset_from_index": 100, "dataset_to_index": 110},
-            ]
-        )
-        hf_dataset = FakeHFDataset({"action": values})
-
-    stats = compute_quantile_stats_for_dataset(FakeDataset())
-
-    np.testing.assert_array_equal(stats["action"]["count"], np.array([110]))
-    expected_q90 = np.percentile(np.asarray(values), 90, axis=0)
-    np.testing.assert_allclose(stats["action"]["q90"], expected_q90, atol=0.1)

tests/datasets/test_compute_stats.py

@@ -688,7 +688,7 @@ def test_compute_episode_stats_string_features_skipped():
 
 
 def test_aggregate_feature_stats_with_quantiles():
-    """Test aggregating feature stats that include quantiles uses conservative bounds."""
+    """Test aggregating feature stats that include quantiles."""
     stats_ft_list = [
         {
             "min": np.array([1.0]),
@@ -697,9 +697,6 @@ def test_aggregate_feature_stats_with_quantiles():
             "std": np.array([2.0]),
             "count": np.array([100]),
             "q01": np.array([1.5]),
-            "q10": np.array([2.0]),
-            "q50": np.array([5.0]),
-            "q90": np.array([9.0]),
             "q99": np.array([9.5]),
         },
         {
@@ -709,21 +706,22 @@ def test_aggregate_feature_stats_with_quantiles():
             "std": np.array([2.5]),
             "count": np.array([150]),
             "q01": np.array([2.5]),
-            "q10": np.array([3.0]),
-            "q50": np.array([6.0]),
-            "q90": np.array([11.0]),
             "q99": np.array([11.5]),
         },
     ]
 
     result = aggregate_feature_stats(stats_ft_list)
 
-    # Lower quantiles use min; upper quantiles use max, regardless of counts.
-    np.testing.assert_allclose(result["q01"], np.array([1.5]), atol=1e-6)
-    np.testing.assert_allclose(result["q10"], np.array([2.0]), atol=1e-6)
-    np.testing.assert_allclose(result["q50"], np.array([5.0]), atol=1e-6)
-    np.testing.assert_allclose(result["q90"], np.array([11.0]), atol=1e-6)
-    np.testing.assert_allclose(result["q99"], np.array([11.5]), atol=1e-6)
+    # Should preserve quantiles
+    assert "q01" in result
+    assert "q99" in result
+
+    # Verify quantile aggregation (weighted average)
+    expected_q01 = (1.5 * 100 + 2.5 * 150) / 250  # ≈ 2.1
+    expected_q99 = (9.5 * 100 + 11.5 * 150) / 250  # ≈ 10.7
+
+    np.testing.assert_allclose(result["q01"], np.array([expected_q01]), atol=1e-6)
+    np.testing.assert_allclose(result["q99"], np.array([expected_q99]), atol=1e-6)
 
 
 def test_aggregate_stats_mixed_quantiles():
@@ -880,60 +878,3 @@ def test_fixed_quantiles_always_computed():
         for q_key in expected_quantiles:
             assert q_key in episode_stats[key]
             assert episode_stats[key][q_key].shape == (features[key]["shape"][0],)
-
-
-def test_aggregate_stats_incremental_resume():
-    """Verify conservative bounds remain associative across incremental additions."""
-    # Start with episode 1 stats (narrow distribution)
-    ep1_stats = {
-        "action": {
-            "min": np.array([-10.0, -5.0]),
-            "max": np.array([10.0, 5.0]),
-            "mean": np.array([0.0, 0.0]),
-            "std": np.array([3.0, 1.5]),
-            "count": np.array([500]),
-            "q01": np.array([-9.0, -4.5]),
-            "q99": np.array([9.0, 4.5]),
-        },
-    }
-
-    # Episode 2: wider distribution on dim 0
-    ep2_stats = {
-        "action": {
-            "min": np.array([-30.0, -5.0]),
-            "max": np.array([40.0, 6.0]),
-            "mean": np.array([5.0, 0.5]),
-            "std": np.array([15.0, 2.0]),
-            "count": np.array([100]),
-            "q01": np.array([-25.0, -4.0]),
-            "q99": np.array([35.0, 5.5]),
-        },
-    }
-
-    # First aggregation: ep1 + ep2 (simulates save_episode for ep2)
-    cumulative = aggregate_stats([ep1_stats, ep2_stats])
-
-    # q01 should take min (conservative lower bound)
-    np.testing.assert_allclose(cumulative["action"]["q01"], np.array([-25.0, -4.5]))
-    # q99 should take max (conservative upper bound)
-    np.testing.assert_allclose(cumulative["action"]["q99"], np.array([35.0, 5.5]))
-
-    # Episode 3: even wider on dim 1
-    ep3_stats = {
-        "action": {
-            "min": np.array([-8.0, -20.0]),
-            "max": np.array([8.0, 25.0]),
-            "mean": np.array([0.0, 3.0]),
-            "std": np.array([2.0, 8.0]),
-            "count": np.array([50]),
-            "q01": np.array([-7.0, -18.0]),
-            "q99": np.array([7.0, 22.0]),
-        },
-    }
-
-    # Second aggregation: cumulative + ep3 (simulates save_episode for ep3)
-    cumulative2 = aggregate_stats([cumulative, ep3_stats])
-
-    # Bounds should widen monotonically
-    np.testing.assert_allclose(cumulative2["action"]["q01"], np.array([-25.0, -18.0]))
-    np.testing.assert_allclose(cumulative2["action"]["q99"], np.array([35.0, 22.0]))

tests/distributed/test_parallel_dims_and_factory.py

@@ -79,10 +79,13 @@ class TestParallelDims:
 
 
 class TestEnvGuard:
-    # Silent config overrides inside accelerate itself — the guard must catch them.
+    # ACCELERATE_DYNAMO_*/ACCELERATE_GRADIENT_ACCUMULATION_STEPS are silent config overrides
+    # inside accelerate itself — the guard must catch them too.
     _POISON = (
         "ACCELERATE_USE_FSDP",
-        "ACCELERATE_USE_PARALLELISM_CONFIG",
+        "FSDP_VERSION",
+        "PARALLELISM_CONFIG_DP_SHARD_SIZE",
+        "ACCELERATE_DYNAMO_BACKEND",
         "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
     )
 
@@ -99,7 +102,7 @@ class TestEnvGuard:
             guard_against_env_interference()
 
     def test_override_acknowledges(self, monkeypatch):
-        monkeypatch.setenv("ACCELERATE_USE_FSDP", "true")
+        monkeypatch.setenv("FSDP_VERSION", "2")
         monkeypatch.setenv(_ENV_OVERRIDE, "1")
         guard_against_env_interference()
 

tests/policies/vla_jepa/test_action_head.py

@@ -16,6 +16,7 @@ from conftest import (
     make_config,
     set_seed_all,
 )  # noqa: E402
+
 from lerobot.policies.vla_jepa.action_head import (  # noqa: E402
     VLAJEPAActionHead,
 )

tests/policies/vla_jepa/test_configuration.py

@@ -3,8 +3,8 @@
 from __future__ import annotations
 
 import pytest
-
 from conftest import ACTION_DIM, ACTION_HORIZON, IMAGE_SIZE, NUM_VIDEO_FRAMES, STATE_DIM, make_config
+
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE

tests/policies/vla_jepa/test_vla_jepa.py

@@ -32,6 +32,7 @@ from conftest import (  # noqa: E402
     make_train_batch,
     set_seed_all,
 )
+
 from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig  # noqa: E402
 from lerobot.policies.vla_jepa.modeling_vla_jepa import VLAJEPAPolicy  # noqa: E402
 from lerobot.utils.constants import ACTION  # noqa: E402

tests/training/test_ema.py

@@ -1,314 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# 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.
-
-"""Tests for the opt-in EMA shadow maintained by the training pipeline (--ema.enable=true)."""
-
-import draccus
-import numpy as np
-import pytest
-import torch
-
-pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
-
-from lerobot.configs.default import EMAConfig
-from lerobot.configs.train import TrainPipelineConfig
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR
-
-DUMMY_REPO_ID = "dummy/repo"
-DUMMY_STATE_DIM = 6
-DUMMY_ACTION_DIM = 6
-IMAGE_SIZE = 32
-N_EPISODES = 2
-EPISODE_LENGTH = 12
-
-
-def test_ema_config_defaults_match_the_reference():
-    cfg = EMAConfig()
-    assert not cfg.enable
-    assert cfg.inv_gamma == 1.0
-    assert cfg.power == 0.75
-    assert cfg.update_after_step == 0
-
-
-@pytest.mark.parametrize(
-    "kwargs",
-    [
-        {"min_decay": 0.5, "max_decay": 0.1},
-        {"max_decay": 1.5},
-        {"min_decay": -0.1},
-        {"inv_gamma": 0.0},
-        {"power": -1.0},
-        {"update_after_step": -1},
-        {"decay": 1.5},
-        {"decay": -0.1},
-        {"decay": 0.99, "min_decay": 0.5},
-        {"decay": 0.99, "max_decay": 0.9},
-    ],
-)
-def test_ema_config_rejects_invalid_values(kwargs):
-    with pytest.raises(ValueError):
-        EMAConfig(**kwargs)
-
-
-def test_ema_config_cli_parsing():
-    cfg = draccus.parse(
-        TrainPipelineConfig,
-        None,
-        args=[
-            f"--dataset.repo_id={DUMMY_REPO_ID}",
-            "--ema.enable=true",
-            "--ema.power=0.8",
-            "--ema.update_after_step=10",
-        ],
-    )
-    assert cfg.ema.enable
-    assert cfg.ema.power == 0.8
-    assert cfg.ema.update_after_step == 10
-
-
-def test_ema_config_cli_parsing_constant_decay():
-    cfg = draccus.parse(
-        TrainPipelineConfig,
-        None,
-        args=[
-            f"--dataset.repo_id={DUMMY_REPO_ID}",
-            "--ema.enable=true",
-            "--ema.decay=0.99",
-        ],
-    )
-    assert cfg.ema.enable
-    assert cfg.ema.decay == 0.99
-
-
-def test_ema_constant_decay_pins_the_schedule():
-    """min_decay == max_decay clamps the warmup curve to a constant (how --ema.decay is implemented)."""
-    pytest.importorskip("diffusers")
-    from diffusers.training_utils import EMAModel
-
-    model = torch.nn.Linear(4, 4)
-    ema = EMAModel(
-        model.parameters(), decay=0.99, min_decay=0.99, use_ema_warmup=True, inv_gamma=1.0, power=0.75
-    )
-    # The first update is a hard copy (decay 0); every one after uses the constant decay.
-    for step in range(1, 6):
-        ema.step(model.parameters())
-        if step > 1:
-            assert ema.cur_decay_value == 0.99
-
-
-def test_ema_weights_context_swaps_and_restores():
-    pytest.importorskip("diffusers")
-    from diffusers.training_utils import EMAModel
-
-    from lerobot.scripts.lerobot_train import _ema_weights
-
-    torch.manual_seed(0)
-    model = torch.nn.Linear(4, 4)
-    ema = EMAModel(model.parameters(), decay=0.9999, use_ema_warmup=True, inv_gamma=1.0, power=0.75)
-
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-    for _ in range(3):
-        model(torch.randn(2, 4)).sum().backward()
-        optimizer.step()
-        optimizer.zero_grad()
-        ema.step(model.parameters())
-
-    live = [p.detach().clone() for p in model.parameters()]
-    with _ema_weights(ema, model):
-        swapped = [p.detach().clone() for p in model.parameters()]
-    restored = list(model.parameters())
-
-    assert any(not torch.equal(a, b) for a, b in zip(live, swapped, strict=True))
-    assert all(torch.equal(a, b.detach()) for a, b in zip(live, restored, strict=True))
-
-
-def make_dummy_dataset(tmp_path):
-    features = {
-        "action": {"dtype": "float32", "shape": (DUMMY_ACTION_DIM,), "names": None},
-        "observation.state": {"dtype": "float32", "shape": (DUMMY_STATE_DIM,), "names": None},
-        "observation.images.top": {
-            "dtype": "image",
-            "shape": (IMAGE_SIZE, IMAGE_SIZE, 3),
-            "names": ["height", "width", "channel"],
-        },
-    }
-    root = tmp_path / "_dataset"
-    dataset = LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=30, features=features, root=root)
-    rng = np.random.default_rng(0)
-    for ep_idx in range(N_EPISODES):
-        for _ in range(EPISODE_LENGTH):
-            dataset.add_frame(
-                {
-                    "action": rng.standard_normal(DUMMY_ACTION_DIM).astype(np.float32),
-                    "observation.state": rng.standard_normal(DUMMY_STATE_DIM).astype(np.float32),
-                    "observation.images.top": rng.integers(
-                        0, 255, size=(IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8
-                    ),
-                    "task": f"task_{ep_idx}",
-                }
-            )
-        dataset.save_episode()
-    dataset.finalize()
-    return root
-
-
-def make_train_config(root, output_dir, steps, ema_enable, ema_decay=None):
-    from lerobot.configs.default import DatasetConfig
-    from lerobot.policies.factory import make_policy_config
-
-    policy_config = make_policy_config(
-        "diffusion",
-        device="cpu",
-        push_to_hub=False,
-        n_obs_steps=2,
-        horizon=8,
-        n_action_steps=4,
-        drop_n_last_frames=0,
-        down_dims=(32, 64),
-        diffusion_step_embed_dim=32,
-        spatial_softmax_num_keypoints=8,
-        num_inference_steps=2,
-        pretrained_backbone_weights=None,
-        use_group_norm=True,
-    )
-    cfg = TrainPipelineConfig(
-        dataset=DatasetConfig(repo_id=DUMMY_REPO_ID, root=str(root)),
-        policy=policy_config,
-        output_dir=output_dir,
-        steps=steps,
-        batch_size=2,
-        num_workers=0,
-        seed=42,
-        log_freq=0,
-        env_eval_freq=0,
-        save_freq=2,
-        ema=EMAConfig(enable=ema_enable, decay=ema_decay),
-    )
-    cfg.optimizer = policy_config.get_optimizer_preset()
-    cfg.scheduler = policy_config.get_scheduler_preset()
-    # The config is built in-process, so skip the CLI-oriented validation.
-    cfg.validate = lambda: None
-    return cfg
-
-
-def load_safetensors(path):
-    from safetensors.torch import load_file
-
-    return load_file(path)
-
-
-def test_train_diffusion_with_ema_checkpoint_and_resume(tmp_path):
-    pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
-    pytest.importorskip("diffusers", reason="diffusers is required (install lerobot[diffusion])")
-    from lerobot.scripts.lerobot_train import EMA_STATE_FILENAME, train
-
-    root = make_dummy_dataset(tmp_path)
-    output_dir = tmp_path / "_output"
-
-    cfg = make_train_config(root, output_dir, steps=4, ema_enable=True)
-    train(cfg)
-
-    checkpoint_dir = output_dir / "checkpoints" / "000004"
-    ema_state_path = checkpoint_dir / TRAINING_STATE_DIR / EMA_STATE_FILENAME
-    ema_model_dir = checkpoint_dir / f"{PRETRAINED_MODEL_DIR}_ema"
-
-    # The shadow state is saved for resume and tracks every optimizer step.
-    assert ema_state_path.exists()
-    ema_state = torch.load(ema_state_path, weights_only=True)
-    assert ema_state["optimization_step"] == 4
-
-    # A directly loadable EMA model is saved next to the live one, with different weights.
-    live_weights = load_safetensors(checkpoint_dir / PRETRAINED_MODEL_DIR / "model.safetensors")
-    ema_weights = load_safetensors(ema_model_dir / "model.safetensors")
-    assert set(live_weights) == set(ema_weights)
-    assert any(not torch.equal(live_weights[k], ema_weights[k]) for k in live_weights)
-
-    from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
-
-    policy = DiffusionPolicy.from_pretrained(str(ema_model_dir))
-    assert isinstance(policy, DiffusionPolicy)
-
-    # Resuming picks the shadow up where it left off instead of restarting it.
-    resume_cfg = make_train_config(root, output_dir, steps=6, ema_enable=True)
-    resume_cfg.resume = True
-    resume_cfg.checkpoint_path = checkpoint_dir
-    train(resume_cfg)
-
-    resumed_state = torch.load(
-        output_dir / "checkpoints" / "000006" / TRAINING_STATE_DIR / EMA_STATE_FILENAME,
-        weights_only=True,
-    )
-    assert resumed_state["optimization_step"] == 6
-
-
-def test_train_with_constant_ema_decay(tmp_path):
-    pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
-    pytest.importorskip("diffusers", reason="diffusers is required (install lerobot[diffusion])")
-    from lerobot.scripts.lerobot_train import EMA_STATE_FILENAME, train
-
-    root = make_dummy_dataset(tmp_path)
-    output_dir = tmp_path / "_output"
-
-    cfg = make_train_config(root, output_dir, steps=2, ema_enable=True, ema_decay=0.99)
-    train(cfg)
-
-    ema_state = torch.load(
-        output_dir / "checkpoints" / "000002" / TRAINING_STATE_DIR / EMA_STATE_FILENAME,
-        weights_only=True,
-    )
-    # The constant decay is implemented by pinning the schedule clamp to that value.
-    assert ema_state["decay"] == 0.99
-    assert ema_state["min_decay"] == 0.99
-    assert ema_state["optimization_step"] == 2
-
-
-def test_train_with_ema_and_gradient_accumulation(tmp_path):
-    """The shadow tracks optimizer steps, not micro-batches, under gradient accumulation."""
-    pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
-    pytest.importorskip("diffusers", reason="diffusers is required (install lerobot[diffusion])")
-    from lerobot.scripts.lerobot_train import EMA_STATE_FILENAME, train
-
-    root = make_dummy_dataset(tmp_path)
-    output_dir = tmp_path / "_output"
-
-    cfg = make_train_config(root, output_dir, steps=4, ema_enable=True)
-    cfg.accelerator.gradient_accumulation.steps = 2
-    train(cfg)
-
-    ema_state = torch.load(
-        output_dir / "checkpoints" / "000004" / TRAINING_STATE_DIR / EMA_STATE_FILENAME,
-        weights_only=True,
-    )
-    # 4 micro-batches / 2 accumulation steps = 2 optimizer updates.
-    assert ema_state["optimization_step"] == 2
-
-
-def test_train_without_ema_writes_no_ema_files(tmp_path):
-    pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
-    pytest.importorskip("diffusers", reason="diffusers is required (install lerobot[diffusion])")
-    from lerobot.scripts.lerobot_train import EMA_STATE_FILENAME, train
-
-    root = make_dummy_dataset(tmp_path)
-    output_dir = tmp_path / "_output"
-
-    cfg = make_train_config(root, output_dir, steps=2, ema_enable=False)
-    train(cfg)
-
-    checkpoint_dir = output_dir / "checkpoints" / "000002"
-    assert (checkpoint_dir / PRETRAINED_MODEL_DIR / "model.safetensors").exists()
-    assert not (checkpoint_dir / TRAINING_STATE_DIR / EMA_STATE_FILENAME).exists()
-    assert not (checkpoint_dir / f"{PRETRAINED_MODEL_DIR}_ema").exists()

tests/training/test_multi_gpu.py

@@ -21,10 +21,8 @@ This module tests multi-GPU training functionality with accelerate.
 These tests are designed to run on machines with 2+ GPUs and are executed
 in the nightly CI workflow.
 
-The tests launch `lerobot-train` through `accelerate launch` in a subprocess to properly test the
-distributed training environment. Accelerate is used as a plain launcher only: the topology comes
-from `--parallelism.*` flags, never from an accelerate YAML config (see
-`lerobot.distributed.factory.guard_against_env_interference`).
+The tests automatically generate accelerate configs and launch training
+with subprocess to properly test the distributed training environment.
 """
 
 import os
@@ -60,25 +58,73 @@ def download_dataset(repo_id, episodes):
     print(f"Dataset {repo_id} downloaded successfully")
 
 
-def run_accelerate_training(config_args, num_processes=4):
+def _write_multi_gpu_config(f, num_processes):
+    f.write("compute_environment: LOCAL_MACHINE\n")
+    f.write("distributed_type: MULTI_GPU\n")
+    f.write("mixed_precision: 'no'\n")
+    f.write(f"num_processes: {num_processes}\n")
+    f.write("use_cpu: false\n")
+    f.write("gpu_ids: all\n")
+    f.write("downcast_bf16: 'no'\n")
+    f.write("machine_rank: 0\n")
+    f.write("main_training_function: main\n")
+    f.write("num_machines: 1\n")
+    f.write("rdzv_backend: static\n")
+    f.write("same_network: true\n")
+
+
+def _write_fsdp_config(f, num_processes):
+    # FSDP1 with FULL_SHARD (ZeRO-3-equivalent) and FULL_STATE_DICT, matching
+    # docs/source/multi_gpu_training.mdx. ACT's repeated transformer blocks are the wrap units;
+    # fsdp_use_orig_params is required because LeRobot builds the optimizer before prepare().
+    f.write("compute_environment: LOCAL_MACHINE\n")
+    f.write("distributed_type: FSDP\n")
+    f.write("mixed_precision: 'no'\n")
+    f.write(f"num_processes: {num_processes}\n")
+    f.write("use_cpu: false\n")
+    f.write("gpu_ids: all\n")
+    f.write("machine_rank: 0\n")
+    f.write("main_training_function: main\n")
+    f.write("num_machines: 1\n")
+    f.write("rdzv_backend: static\n")
+    f.write("same_network: true\n")
+    f.write("fsdp_config:\n")
+    f.write("  fsdp_version: 1\n")
+    f.write("  fsdp_sharding_strategy: FULL_SHARD\n")
+    f.write("  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP\n")
+    f.write("  fsdp_transformer_layer_cls_to_wrap: ACTEncoderLayer,ACTDecoderLayer\n")
+    f.write("  fsdp_use_orig_params: true\n")
+    f.write("  fsdp_state_dict_type: FULL_STATE_DICT\n")
+
+
+def run_accelerate_training(config_args, num_processes=4, temp_dir=None, distributed_type="MULTI_GPU"):
     """
     Helper function to run training with accelerate launch.
 
-    `accelerate launch` is used as a plain launcher (no `--config_file`): it only sets the
-    rendezvous env vars, and the layout — DDP by default, FSDP with `--parallelism.dp_shard` —
-    comes from `config_args`.
-
     Args:
         config_args: List of config arguments to pass to lerobot_train.py
         num_processes: Number of processes (GPUs) to use
+        temp_dir: Temporary directory for outputs
+        distributed_type: "MULTI_GPU" (DDP) or "FSDP" — selects the generated accelerate config.
 
     Returns:
         subprocess.CompletedProcess result
     """
+
+    config_path = Path(temp_dir) / "accelerate_config.yaml"
+
+    # Write YAML config
+    with open(config_path, "w") as f:
+        if distributed_type == "FSDP":
+            _write_fsdp_config(f, num_processes)
+        else:
+            _write_multi_gpu_config(f, num_processes)
+
     cmd = [
         "accelerate",
         "launch",
-        f"--num_processes={num_processes}",
+        "--config_file",
+        str(config_path),
         "-m",
         "lerobot.scripts.lerobot_train",
     ] + config_args
@@ -127,7 +173,7 @@ class TestMultiGPUTraining:
                 "--num_workers=0",
             ]
 
-            result = run_accelerate_training(config_args, num_processes=4)
+            result = run_accelerate_training(config_args, num_processes=4, temp_dir=temp_dir)
 
             # Check that training completed successfully
             assert result.returncode == 0, (
@@ -170,7 +216,7 @@ class TestMultiGPUTraining:
                 "--num_workers=0",
             ]
 
-            result = run_accelerate_training(config_args, num_processes=2)
+            result = run_accelerate_training(config_args, num_processes=2, temp_dir=temp_dir)
 
             assert result.returncode == 0, (
                 f"Training failed:\nSTDOUT:\n{result.stdout}\n\nSTDERR:\n{result.stderr}"
@@ -200,12 +246,11 @@ class TestMultiGPUTraining:
 
     def test_fsdp_optimizer_save_and_resume(self):
         """
-        Test that FSDP saves the sharded optimizer state and can resume from it.
+        Test that FSDP saves the (gathered) optimizer state and can resume from it.
 
-        Trains a few steps under FSDP2 (`--parallelism.dp_shard=2`), verifies the DCP optimizer
-        shards are written next to the rest of the training state, then resumes from the
-        checkpoint for more steps and checks it completes without shape/key errors in the
-        resharding optimizer load path.
+        Trains a few steps under FSDP, verifies the gathered optimizer state is written next to the
+        rest of the training state, then resumes from the checkpoint for more steps and checks it
+        completes without shape/key errors in the FSDP optimizer load path.
         """
         # Pre-download dataset to avoid race conditions
         download_dataset("lerobot/pusht", episodes=[0])
@@ -220,7 +265,6 @@ class TestMultiGPUTraining:
                 "--policy.device=cuda",
                 "--policy.push_to_hub=false",
                 f"--output_dir={output_dir}",
-                "--parallelism.dp_shard=2",
                 "--batch_size=4",
                 "--steps=10",
                 "--env_eval_freq=-1",
@@ -230,33 +274,34 @@ class TestMultiGPUTraining:
                 "--num_workers=0",
             ]
 
-            result = run_accelerate_training(config_args, num_processes=2)
+            result = run_accelerate_training(
+                config_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP"
+            )
             assert result.returncode == 0, (
                 f"FSDP training failed:\nSTDOUT:\n{result.stdout}\n\nSTDERR:\n{result.stderr}"
             )
 
-            # Under sharding the optimizer state is written as DCP shards (proves the save
-            # collective ran); the model artifact stays a gathered model.safetensors at the
-            # default --checkpoint_format=safetensors.
-            checkpoint_dir = output_dir / "checkpoints" / "last"
-            training_state_dir = checkpoint_dir / "training_state"
-            optimizer_shards = training_state_dir / "optimizer_0"
-            assert optimizer_shards.is_dir(), f"FSDP optimizer shards not saved in {training_state_dir}"
-            assert any(optimizer_shards.iterdir()), f"FSDP optimizer shard dir is empty: {optimizer_shards}"
-            assert (checkpoint_dir / "pretrained_model" / "model.safetensors").exists(), (
-                f"Gathered model weights not saved in {checkpoint_dir}"
+            # The gathered optimizer state must be written under FSDP (proves the save collective ran),
+            # in the same safetensors format as single-GPU training.
+            training_state_dir = output_dir / "checkpoints" / "last" / "training_state"
+            optimizer_state = training_state_dir / "optimizer_state.safetensors"
+            optimizer_param_groups = training_state_dir / "optimizer_param_groups.json"
+            assert optimizer_state.exists(), f"FSDP optimizer state not saved in {training_state_dir}"
+            assert optimizer_param_groups.exists(), (
+                f"FSDP optimizer param groups not saved in {training_state_dir}"
             )
 
-            # Resume from the checkpoint for more steps. A successful run proves the DCP optimizer
-            # load accepts the saved state and reshards it without shape/key errors. The topology
-            # is restored from train_config.json, so --parallelism.* is not repeated here.
-            resume_config = checkpoint_dir / "pretrained_model" / "train_config.json"
+            # Resume from the checkpoint for more steps. A successful run proves load_fsdp_optimizer
+            # accepts the saved state and reshards it without shape/key errors.
+            resume_config = output_dir / "checkpoints" / "last" / "pretrained_model" / "train_config.json"
             resume_args = [
                 f"--config_path={resume_config}",
                 "--resume=true",
                 "--steps=20",
             ]
-            resume_result = run_accelerate_training(resume_args, num_processes=2)
+            resume_result = run_accelerate_training(
+                resume_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP"
+            )
             assert resume_result.returncode == 0, (
                 f"FSDP resume failed:\nSTDOUT:\n{resume_result.stdout}\n\nSTDERR:\n{resume_result.stderr}"
             )

tests/utils/test_logging_utils.py

@@ -197,7 +197,7 @@ def test_metrics_tracker_reduce_across_ranks_invokes_all_reduce(monkeypatch):
     tracker.reduce_across_ranks()
 
     assert captured["op"] == logging_utils.dist.ReduceOp.MAX
-    assert torch.allclose(captured["values"], torch.tensor([0.4], device=captured["values"].device))
+    assert torch.allclose(captured["values"], torch.tensor([0.4]))
     assert tracker.update_s.avg == pytest.approx(0.9)
     # Metrics without a reduction stay untouched.
     assert tracker.loss.avg == 1.0

utils/check_docstrings.py

@@ -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