docs(optim): write the API reference docstrings

Starts Wave 2. Takes src/lerobot/optim/ to 100% public docstring coverage. Fixes dataclass Args: field
order to match the real generated __init__ signature (base-class fields keep their position even when
redeclared by a subclass). Adds docs/source/api/optim.mdx, which didn't exist before — needs a
_toctree.yml entry from whoever owns that file, see PR description.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

docs/source/api/configs.mdx

@@ -10,26 +10,14 @@ 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
@@ -37,53 +25,3 @@ 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/api/optim.mdx

@@ -0,0 +1,93 @@
+# Optimization
+
+`OptimizerConfig` and `LRSchedulerConfig` are the base configuration classes for the optimizers and learning
+rate schedulers used during training. `TrainPipelineConfig` composes one of each; see
+[`~optim.factory.make_optimizer_and_scheduler`] for how they are built from a policy's parameters.
+
+## make_optimizer_and_scheduler
+
+[[autodoc]] lerobot.optim.factory.make_optimizer_and_scheduler
+
+## OptimizerConfig
+
+[[autodoc]] lerobot.optim.OptimizerConfig
+    - type
+    - builds_multiple_optimizers
+    - default_choice_name
+    - build
+
+## AdamConfig
+
+[[autodoc]] lerobot.optim.AdamConfig
+    - build
+
+## AdamWConfig
+
+[[autodoc]] lerobot.optim.AdamWConfig
+    - build
+
+## SGDConfig
+
+[[autodoc]] lerobot.optim.SGDConfig
+    - build
+
+## MultiAdamConfig
+
+Builds a dictionary of Adam optimizers, one per parameter group — used when a policy needs separate
+optimizers for different components (e.g. actor/critic/temperature in SAC).
+
+[[autodoc]] lerobot.optim.MultiAdamConfig
+    - builds_multiple_optimizers
+    - build
+
+## XVLAAdamWConfig
+
+[[autodoc]] lerobot.optim.XVLAAdamWConfig
+    - build
+
+## save_optimizer_state
+
+[[autodoc]] lerobot.optim.save_optimizer_state
+
+## load_optimizer_state
+
+[[autodoc]] lerobot.optim.load_optimizer_state
+
+## LRSchedulerConfig
+
+[[autodoc]] lerobot.optim.LRSchedulerConfig
+    - type
+    - build
+
+## DiffuserSchedulerConfig
+
+[[autodoc]] lerobot.optim.DiffuserSchedulerConfig
+    - build
+
+## VQBeTSchedulerConfig
+
+[[autodoc]] lerobot.optim.VQBeTSchedulerConfig
+    - build
+
+## ConstantWithWarmupSchedulerConfig
+
+[[autodoc]] lerobot.optim.schedulers.ConstantWithWarmupSchedulerConfig
+    - build
+
+## CosineAnnealingWithWarmupSchedulerConfig
+
+[[autodoc]] lerobot.optim.schedulers.CosineAnnealingWithWarmupSchedulerConfig
+    - build
+
+## CosineDecayWithWarmupSchedulerConfig
+
+[[autodoc]] lerobot.optim.CosineDecayWithWarmupSchedulerConfig
+    - build
+
+## save_scheduler_state
+
+[[autodoc]] lerobot.optim.save_scheduler_state
+
+## load_scheduler_state
+
+[[autodoc]] lerobot.optim.load_scheduler_state

pyproject.toml

@@ -439,6 +439,7 @@ 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"]
@@ -446,7 +447,6 @@ ignore = [
 "src/lerobot/jobs/**" = ["D"]
 "src/lerobot/model/**" = ["D"]
 "src/lerobot/motors/**" = ["D"]
-"src/lerobot/optim/**" = ["D"]
 "src/lerobot/policies/**" = ["D"]
 "src/lerobot/processor/**" = ["D"]
 "src/lerobot/rewards/**" = ["D"]

src/lerobot/configs/__init__.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Public API for lerobot configuration types and base config classes.
+"""
+Public API for lerobot configuration types and base config classes.
 
 NOTE: TrainPipelineConfig, EvalPipelineConfig, and TrainRLServerPipelineConfig
 are intentionally NOT re-exported here to avoid circular dependencies

src/lerobot/configs/accelerator.py

@@ -171,13 +171,6 @@ 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,76 +23,56 @@ from .video import DepthEncoderConfig, RGBEncoderConfig, depth_encoder_defaults,
 
 @dataclass
 class DatasetRecordConfig:
-    """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`).
-    """
-
+    # Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).
     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,65 +27,35 @@ logger = logging.getLogger(__name__)
 
 @dataclass
 class DatasetConfig:
-    """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).
-    """
-
+    # 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.
     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:
@@ -122,63 +92,35 @@ 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
-    add_tags: bool = True
+    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.
 
 
 @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:
@@ -199,66 +141,54 @@ class EvalConfig:
 
 @dataclass
 class PeftConfig:
-    """PEFT (parameter-efficient fine-tuning) settings, e.g. LoRA adapters.
-
-    PEFT offers many fine-tuning methods, layer adapters being the most common and currently also the
-    most effective methods so we'll focus on those in this high-level config interface.
-
-    Args:
-        target_modules (`list[str] | str | None`, *optional*): Either a string (module name suffix or
-            `'all-linear'`), a list of module name suffixes, or a regular expression describing module
-            names to target with the configured PEFT method. Some policies have a default value for this
-            so that you don't *have* to choose which layers to adapt, but it might still be worthwhile
-            depending on your case.
-        full_training_modules (`list[str] | None`, *optional*): Names/suffixes of modules to fully
-            fine-tune and store alongside adapter weights. Useful for layers that are not part of a
-            pre-trained model (e.g., action state projections). Depending on the policy this defaults to
-            layers that are newly created in pre-trained policies. If you're fine-tuning an already trained
-            policy you might want to set this to `[]`. Corresponds to PEFT's `modules_to_save`.
-        method_type (`str`, *optional*, defaults to `"LORA"`): The PEFT (adapter) method to apply to the
-            policy. Needs to be a valid PEFT type.
-        init_type (`str | None`, *optional*): Adapter initialization method. Look at the specific PEFT
-            adapter documentation for defaults.
-        r (`int`, *optional*, defaults to 16): We expect that all PEFT adapters are in some way doing
-            rank-decomposition, therefore this parameter specifies the rank used for the adapter. In
-            general a higher rank means more trainable parameters and closer to full fine-tuning.
-        lora_alpha (`int | None`, *optional*): Alpha parameter for LoRA scaling (`scaling = lora_alpha /
-            r`). In general, a higher alpha means stronger adaptation signal. If `None`, the PEFT library
-            defaults to `alpha=8`, which may dampen high-rank adapters. Common values are `r` (`alpha ==
-            rank`) or `2*r`.
-    """
+    # 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 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.
-    """
-
+    # 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.
     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,36 +28,21 @@ logger = getLogger(__name__)
 
 @dataclass
 class EvalPipelineConfig:
-    """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).
-    """
-
+    # 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`.
     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:
@@ -90,5 +75,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,7 +96,6 @@ 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}"
@@ -116,7 +115,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.
-        args (`Sequence[str]`): A sequence of command-line arguments to parse.
+        cli_args (Sequence[str]): A sequence of command-line arguments to parse.
 
     Returns:
         dict: A dictionary containing the parsed plugin arguments where:
@@ -157,7 +156,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.
@@ -181,7 +180,6 @@ 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:
@@ -194,7 +192,6 @@ 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)
@@ -202,24 +199,21 @@ 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 only recurses a ChoiceType field's subclasses.
+    """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.
 
-    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.
+    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)
@@ -259,10 +253,8 @@ 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"
@@ -272,7 +264,6 @@ 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}"
@@ -294,11 +285,12 @@ 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, *optional*): The sequence of command-line arguments to be filtered.
+        args (Sequence[str] | None): The sequence of command-line arguments to be filtered.
             Defaults to None.
 
     Returns:
@@ -388,22 +380,19 @@ def extract_path_fields_from_config(config_path: str, path_fields: list[str]) ->
 
 
 def wrap(config_path: Path | None = None) -> Callable[[F], F]:
-    """HACK: Similar to draccus.wrap but does three additional things.
-
-    - Will remove '.path' arguments from CLI in order to process them later on.
-    - If a 'config_path' is passed and the main config class has a 'from_pretrained' method, will
-      initialize it from there to allow to fetch configs from the hub directly
-    - Will load plugins specified in the CLI arguments. These plugins will typically register
-        their own subclasses of config classes, so that draccus can find the right class to instantiate
-        from the CLI '.type' arguments
+    """
+    HACK: Similar to draccus.wrap but does three additional things:
+        - Will remove '.path' arguments from CLI in order to process them later on.
+        - If a 'config_path' is passed and the main config class has a 'from_pretrained' method, will
+          initialize it from there to allow to fetch configs from the hub directly
+        - Will load plugins specified in the CLI arguments. These plugins will typically register
+            their own subclasses of config classes, so that draccus can find the right class to instantiate
+            from the CLI '.type' arguments
     """
 
     def wrapper_outer(fn: F) -> F:
-        """Wrap `fn` so its first argument is resolved from the CLI/config instead of passed directly."""
-
         @wraps(fn)
         def wrapper_inner(*args: Any, **kwargs: Any) -> Any:
-            """Build `fn`'s config argument from the CLI/config file (unless already given), then call `fn`."""
             argspec = inspect.getfullargspec(fn)
             argtype = argspec.annotations[argspec.args[0]]
             if len(args) > 0 and type(args[0]) is argtype:

src/lerobot/configs/policies.py

@@ -39,64 +39,50 @@ 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.
-
-    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.
+    """
+    Base configuration class for policy models.
 
     Args:
-        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: 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 = 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
+    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.
     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}'.")
@@ -111,7 +97,6 @@ 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)}")
@@ -120,52 +105,32 @@ 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():
@@ -175,7 +140,6 @@ 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():
@@ -185,14 +149,12 @@ 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():
@@ -220,35 +182,6 @@ 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,45 +43,31 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     """Base configuration for reward models.
 
     Args:
-        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.
+    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.
     """
 
+    # 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}'.")
@@ -89,7 +75,6 @@ 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)}")
@@ -97,17 +82,14 @@ 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:
@@ -115,14 +97,9 @@ 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:
@@ -145,33 +122,6 @@ 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

@@ -108,119 +108,75 @@ def _migrate_legacy_rabc_fields(config: dict[str, Any]) -> dict[str, Any] | None
 
 @dataclass
 class TrainPipelineConfig(HubMixin):
-    """The top-level configuration for `lerobot-train`, parsed by draccus from CLI flags and/or a YAML file.
-
-    Args:
-        dataset (`DatasetConfig`): The dataset(s) to train on.
-        env (`envs.EnvConfig | None`, *optional*): The simulation environment to periodically evaluate the
-            policy in (see `env_eval_freq`). Required when `env_eval_freq > 0`.
-        policy (`PreTrainedConfig | None`, *optional*): The policy to train. Mutually exclusive with
-            `reward_model`.
-        reward_model (`RewardModelConfig | None`, *optional*): The reward model to train instead of a
-            policy. Mutually exclusive with `policy`.
-        output_dir (`Path | None`, *optional*): Where to save all of the run outputs. If you run another
-            training session with the same value its contents will be overwritten unless `resume` is set.
-        job_name (`str | None`, *optional*): A name for the run.
-        resume (`bool`, *optional*, defaults to `False`): Resume a previous run. Pass `--config_path`
-            pointing at either a local checkpoint's `train_config.json` or a Hub repo id holding
-            `checkpoints/<step>/` subtrees (the latest checkpoint is downloaded and resumed from). When
-            resuming, the default behavior is to use the configuration from the checkpoint, regardless of
-            what's provided with the training command at the time of resumption (CLI `--*` flags still
-            override).
-        seed (`int | None`, *optional*, defaults to 1000): Seed used for training (e.g. model
-            initialization, dataset shuffling) and for the evaluation environments.
-        cudnn_deterministic (`bool`, *optional*, defaults to `False`): Use deterministic cuDNN algorithms
-            for reproducibility. Disables `cudnn.benchmark` and may reduce training speed by ~10-20 percent.
-        num_workers (`int`, *optional*, defaults to 4): Number of workers for the dataloader.
-        batch_size (`int`, *optional*, defaults to 8): The training batch size.
-        prefetch_factor (`int`, *optional*, defaults to 4): Number of batches loaded in advance by each
-            dataloader worker.
-        persistent_workers (`bool`, *optional*, defaults to `True`): Keep dataloader worker processes alive
-            between epochs.
-        dataloader_multiprocessing_context (`str | None`, *optional*, defaults to `"spawn"`): DataLoader
-            worker start method. `"spawn"` is safer than `"fork"` with non-fork-safe libs (PyAV /
-            torchcodec / ffmpeg), but adds some worker-startup time per run since workers re-import modules
-            instead of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
-            when appropriate, or set it to `None` to use Python's platform default.
-        steps (`int`, *optional*, defaults to 100000): Total number of training steps.
-        env_eval_freq (`int`, *optional*, defaults to 20000): Run the policy in the simulation environment
-            every N steps to measure reward/success (0 = disabled).
-        log_freq (`int`, *optional*, defaults to 200): Logging frequency, in steps.
-        eval_steps (`int`, *optional*, defaults to 0): Compute eval loss on held-out episodes every N steps
-            (0 = disabled). Requires `eval_split > 0`.
-        max_eval_samples (`int`, *optional*, defaults to 0): Cap on total eval samples, split uniformly
-            across tasks (0 = use all held-out data).
-        tolerance_s (`float`, *optional*, defaults to 0.0001): Maximum timestamp difference tolerated when
-            loading dataset frames, in seconds.
-        save_checkpoint (`bool`, *optional*, defaults to `True`): Whether to save checkpoints during
-            training.
-        save_freq (`int`, *optional*, defaults to 20000): Save a checkpoint every `save_freq` training
-            iterations and after the last training step. A non-positive value disables periodic saving,
-            keeping only the final checkpoint.
-        checkpoint_format (`CheckpointFormat`, *optional*, defaults to `CheckpointFormat.SAFETENSORS`):
-            Model-artifact format inside checkpoints; non-default values require a sharded run.
-        use_policy_training_preset (`bool`, *optional*, defaults to `True`): Use the policy's own
-            optimizer/scheduler presets when `optimizer`/`scheduler` aren't explicitly set.
-        optimizer (`OptimizerConfig | None`, *optional*): The optimizer to use. Falls back to the policy's
-            preset when `use_policy_training_preset` is `True`.
-        scheduler (`LRSchedulerConfig | None`, *optional*): The learning-rate scheduler to use. Falls back
-            to the policy's preset when `use_policy_training_preset` is `True`.
-        parallelism (`ParallelismConfig`, *optional*): Process topology: `dp_replicate` / `dp_shard` for HSDP
-            and context-parallel degree placeholders.
-        accelerator (`AcceleratorConfig`, *optional*): Execution runtime handed to the Accelerator: mixed
-            precision, gradient accumulation, FSDP/DDP tuning knobs, compile & activation-checkpointing
-            placeholders.
-        eval (`EvalConfig`, *optional*): Settings for the periodic simulation-environment evaluation.
-        wandb (`WandBConfig`, *optional*): Weights & Biases logging settings.
-        peft (`PeftConfig | None`, *optional*): PEFT (e.g. LoRA) settings, when fine-tuning with adapters
-            instead of full-parameter training.
-        job (`JobConfig`, *optional*): Where to run training: locally (default), or an HF Jobs flavor.
-        save_checkpoint_to_hub (`bool`, *optional*, defaults to `False`): Push each saved checkpoint to the
-            Hub (`policy.repo_id`) as it is written, not just the final model (useful to monitor progress
-            mid-run). The final model is pushed regardless. Works the same locally and remotely.
-        sample_weighting (`SampleWeightingConfig | None`, *optional*): Sample weighting configuration (e.g.
-            for RA-BC training).
-        rename_map (`dict[str, str]`, *optional*): Rename map for the observation, to override the image
-            and state keys.
-    """
-
     dataset: DatasetConfig
     env: envs.EnvConfig | None = None
     policy: PreTrainedConfig | None = None
     reward_model: RewardModelConfig | None = None
+    # Set `dir` to where you would like to save all of the run outputs. If you run another training session
+    # with the same value for `dir` its contents will be overwritten unless you set `resume` to true.
     output_dir: Path | None = None
     job_name: str | None = None
+    # Set `resume` to true to resume a previous run. Pass `--config_path` pointing at either a local
+    # checkpoint's train_config.json or a Hub repo id holding `checkpoints/<step>/` subtrees (the
+    # latest checkpoint is downloaded and resumed from). Note that when resuming, the default behavior
+    # is to use the configuration from the checkpoint, regardless of what's provided with the training
+    # command at the time of resumption (CLI `--*` flags still override).
     resume: bool = False
+    # `seed` is used for training (eg: model initialization, dataset shuffling)
+    # AND for the evaluation environments.
     seed: int | None = 1000
+    # Set to True to use deterministic cuDNN algorithms for reproducibility.
+    # This disables cudnn.benchmark and may reduce training speed by ~10-20 percent.
     cudnn_deterministic: bool = False
+    # Number of workers for the dataloader.
     num_workers: int = 4
     batch_size: int = 8
     prefetch_factor: int = 4
     persistent_workers: bool = True
+    # DataLoader worker start method. "spawn" is safer than "fork" with
+    # non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
+    # worker-startup time per run since workers re-import modules instead
+    # of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
+    # when appropriate, or set it to `null` to use Python's platform default.
     dataloader_multiprocessing_context: str | None = "spawn"
     steps: int = 100_000
+    # Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
     env_eval_freq: int = 20_000
     log_freq: int = 200
+    # Compute eval loss on held-out episodes every N steps (0 = disabled). Requires eval_split > 0.
     eval_steps: int = 0
+    # Cap on total eval samples, split uniformly across tasks (0 = use all held-out data).
     max_eval_samples: int = 0
     tolerance_s: float = 1e-4
     save_checkpoint: bool = True
+    # Checkpoint is saved every `save_freq` training iterations and after the last training step.
+    # A non-positive value disables periodic saving, keeping only the final checkpoint.
     save_freq: int = 20_000
+    # Model-artifact format inside checkpoints; non-default values require a sharded run.
     checkpoint_format: CheckpointFormat = CheckpointFormat.SAFETENSORS
     use_policy_training_preset: bool = True
     optimizer: OptimizerConfig | None = None
     scheduler: LRSchedulerConfig | None = None
+    # Process topology: dp_replicate / dp_shard (HSDP) and context-parallel degree placeholders.
     parallelism: ParallelismConfig = field(default_factory=ParallelismConfig)
+    # Execution runtime handed to the Accelerator: mixed precision, gradient accumulation,
+    # FSDP/DDP tuning knobs, compile & activation-checkpointing placeholders.
     accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig)
     eval: EvalConfig = field(default_factory=EvalConfig)
     wandb: WandBConfig = field(default_factory=WandBConfig)
     peft: PeftConfig | None = None
 
+    # Where to run training (local default, or an HF Jobs flavor). See JobConfig.
     job: JobConfig = field(default_factory=JobConfig)
+    # Push each saved checkpoint to the Hub (policy.repo_id) as it is written, not
+    # just the final model (useful to monitor progress mid-run). Optional; the
+    # final model is pushed regardless. Works the same locally and remotely.
     save_checkpoint_to_hub: bool = False
 
+    # Sample weighting configuration (e.g., for RA-BC training)
     sample_weighting: SampleWeightingConfig | None = None
 
+    # Rename map for the observation to override the image and state keys
     rename_map: dict[str, str] = field(default_factory=dict)
     checkpoint_path: Path | None = field(init=False, default=None)
 
@@ -306,21 +262,6 @@ 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
@@ -448,7 +389,6 @@ 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:
@@ -469,35 +409,6 @@ 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,18 +18,6 @@ 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"
@@ -39,28 +27,11 @@ 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"
@@ -70,30 +41,11 @@ 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,33 +84,18 @@ DEPTH_ENCODER_INFO_FIELD_NAMES: frozenset[str] = frozenset({"depth_min", "depth_
 
 @dataclass
 class VideoEncoderConfig:
-    """Video encoder configuration.
+    """Video encoder configuration."""
 
-    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
+    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.
     # TODO(CarolinePascal): add torchcodec support + find a way to unify the
     # two backends (encoding and decoding).
-    video_backend: str = "pyav"
+    video_backend: str = "pyav"  # Encoding backend. Only "pyav" is currently supported.
+    # Extra codec options merged last, e.g. {"tune": "film"}.
     extra_options: dict[str, Any] = field(default_factory=dict)
 
     # Source-data channel count this encoder is expected to handle. ``None``
@@ -119,7 +104,6 @@ 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":
@@ -128,7 +112,9 @@ 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] = {}
 
@@ -161,7 +147,6 @@ 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.
         """
@@ -226,7 +211,6 @@ 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)
 
@@ -318,10 +302,9 @@ class DepthEncoderConfig(VideoEncoderConfig):
 
     @classmethod
     def _kwargs_from_video_info(cls, video_info: dict | None) -> dict[str, Any]:
-        """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.
+        """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.
         """
         kwargs = super()._kwargs_from_video_info(video_info)
         video_info = video_info or {}
@@ -345,8 +328,8 @@ def encoder_config_from_video_info(video_info: dict | None) -> VideoEncoderConfi
     otherwise.
 
     Args:
-        video_info (`dict | None`): A feature's ``info`` dict as persisted in ``info.json``, or ``None``
-            (treated as an empty dict).
+        video_info: 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/optim/factory.py

@@ -25,14 +25,20 @@ from lerobot.policies import PreTrainedPolicy
 def make_optimizer_and_scheduler(
     cfg: TrainPipelineConfig, policy: PreTrainedPolicy
 ) -> tuple[Optimizer, LRScheduler | None]:
-    """Generates the optimizer and scheduler based on configs.
+    """Build the optimizer and, if configured, the learning rate scheduler for training a policy.
 
     Args:
-        cfg (TrainPipelineConfig): The training config that contains optimizer and scheduler configs
-        policy (PreTrainedPolicy): The policy config from which parameters and presets must be taken from.
+        cfg (`TrainPipelineConfig`):
+            The training config, whose `optimizer` and `scheduler` fields are built.
+        policy (`PreTrainedPolicy`):
+            The policy being trained; its parameters (or optimizer-preset groups, if
+            `cfg.use_policy_training_preset` is `True`) are passed to the optimizer.
 
     Returns:
-        tuple[Optimizer, LRScheduler | None]: The couple (Optimizer, Scheduler). Scheduler can be `None`.
+        `tuple[Optimizer, LRScheduler | None]`: The built optimizer, and scheduler if one was configured.
+
+    Raises:
+        ValueError: If `cfg.optimizer` is `None`.
     """
     params = policy.get_optim_params() if cfg.use_policy_training_preset else policy.parameters()
     if cfg.optimizer is None:

src/lerobot/optim/optimizers.py

@@ -44,12 +44,32 @@ OptimizerParams = (
 
 @dataclass
 class OptimizerConfig(draccus.ChoiceRegistry, abc.ABC):
+    """Base configuration shared by every optimizer.
+
+    Concrete optimizers subclass this and register themselves with
+    `@OptimizerConfig.register_subclass("name")`, which is what makes `--optimizer.type=name` work on the
+    command line.
+
+    Args:
+        lr (`float`):
+            Learning rate.
+        weight_decay (`float`):
+            Weight decay (L2 penalty) applied by the optimizer.
+        grad_clip_norm (`float`):
+            Maximum gradient norm; gradients are clipped to this value before each optimizer step.
+    """
+
     lr: float
     weight_decay: float
     grad_clip_norm: float
 
     @property
     def type(self) -> str:
+        """Return the registered name this config was registered under.
+
+        Returns:
+            `str`: The name passed to `@OptimizerConfig.register_subclass`, e.g. `"adam"`.
+        """
         return self.get_choice_name(self.__class__)
 
     @property
@@ -59,12 +79,16 @@ class OptimizerConfig(draccus.ChoiceRegistry, abc.ABC):
 
     @classmethod
     def default_choice_name(cls) -> str | None:
+        """Return the registered name used when `--optimizer.type` is not specified.
+
+        Returns:
+            `str | None`: `"adam"`.
+        """
         return "adam"
 
     @abc.abstractmethod
     def build(self, params: OptimizerParams) -> torch.optim.Optimizer | dict[str, torch.optim.Optimizer]:
-        """
-        Build the optimizer. It can be a single optimizer or a dictionary of optimizers.
+        """Build the optimizer. It can be a single optimizer or a dictionary of optimizers.
 
         NOTE: Multiple optimizers are useful when you have different models to optimize.
         For example, you can have one optimizer for the policy and another one for the value function
@@ -89,6 +113,21 @@ class OptimizerConfig(draccus.ChoiceRegistry, abc.ABC):
 @OptimizerConfig.register_subclass("adam")
 @dataclass
 class AdamConfig(OptimizerConfig):
+    """Configuration for [`torch.optim.Adam`](https://docs.pytorch.org/docs/stable/generated/torch.optim.Adam.html).
+
+    Args:
+        lr (`float`, *optional*, defaults to 0.001):
+            Learning rate.
+        weight_decay (`float`, *optional*, defaults to 0.0):
+            Weight decay (L2 penalty).
+        grad_clip_norm (`float`, *optional*, defaults to 10.0):
+            Maximum gradient norm; gradients are clipped to this value before each optimizer step.
+        betas (`tuple[float, float]`, *optional*, defaults to `(0.9, 0.999)`):
+            Coefficients used for computing running averages of the gradient and its square.
+        eps (`float`, *optional*, defaults to 1e-08):
+            Term added to the denominator to improve numerical stability.
+    """
+
     lr: float = 1e-3
     betas: tuple[float, float] = (0.9, 0.999)
     eps: float = 1e-8
@@ -96,6 +135,15 @@ class AdamConfig(OptimizerConfig):
     grad_clip_norm: float = 10.0
 
     def build(self, params: OptimizerParams) -> torch.optim.Optimizer:
+        """Build a [`torch.optim.Adam`](https://docs.pytorch.org/docs/stable/generated/torch.optim.Adam.html) instance from this config.
+
+        Args:
+            params (`OptimizerParams`):
+                Parameters to optimize, as accepted by `torch.optim.Adam`.
+
+        Returns:
+            `torch.optim.Optimizer`: The built optimizer.
+        """
         kwargs = asdict(self)
         kwargs.pop("grad_clip_norm")
         return torch.optim.Adam(params, **kwargs)
@@ -104,6 +152,21 @@ class AdamConfig(OptimizerConfig):
 @OptimizerConfig.register_subclass("adamw")
 @dataclass
 class AdamWConfig(OptimizerConfig):
+    """Configuration for [`torch.optim.AdamW`](https://docs.pytorch.org/docs/stable/generated/torch.optim.AdamW.html).
+
+    Args:
+        lr (`float`, *optional*, defaults to 0.001):
+            Learning rate.
+        weight_decay (`float`, *optional*, defaults to 0.01):
+            Weight decay, applied decoupled from the gradient update as in the AdamW paper.
+        grad_clip_norm (`float`, *optional*, defaults to 10.0):
+            Maximum gradient norm; gradients are clipped to this value before each optimizer step.
+        betas (`tuple[float, float]`, *optional*, defaults to `(0.9, 0.999)`):
+            Coefficients used for computing running averages of the gradient and its square.
+        eps (`float`, *optional*, defaults to 1e-08):
+            Term added to the denominator to improve numerical stability.
+    """
+
     lr: float = 1e-3
     betas: tuple[float, float] = (0.9, 0.999)
     eps: float = 1e-8
@@ -111,6 +174,15 @@ class AdamWConfig(OptimizerConfig):
     grad_clip_norm: float = 10.0
 
     def build(self, params: OptimizerParams) -> torch.optim.Optimizer:
+        """Build a [`torch.optim.AdamW`](https://docs.pytorch.org/docs/stable/generated/torch.optim.AdamW.html) instance from this config.
+
+        Args:
+            params (`OptimizerParams`):
+                Parameters to optimize, as accepted by `torch.optim.AdamW`.
+
+        Returns:
+            `torch.optim.Optimizer`: The built optimizer.
+        """
         kwargs = asdict(self)
         kwargs.pop("grad_clip_norm")
         return torch.optim.AdamW(params, **kwargs)
@@ -119,6 +191,23 @@ class AdamWConfig(OptimizerConfig):
 @OptimizerConfig.register_subclass("sgd")
 @dataclass
 class SGDConfig(OptimizerConfig):
+    """Configuration for [`torch.optim.SGD`](https://docs.pytorch.org/docs/stable/generated/torch.optim.SGD.html).
+
+    Args:
+        lr (`float`, *optional*, defaults to 0.001):
+            Learning rate.
+        weight_decay (`float`, *optional*, defaults to 0.0):
+            Weight decay (L2 penalty).
+        grad_clip_norm (`float`, *optional*, defaults to 10.0):
+            Maximum gradient norm; gradients are clipped to this value before each optimizer step.
+        momentum (`float`, *optional*, defaults to 0.0):
+            Momentum factor.
+        dampening (`float`, *optional*, defaults to 0.0):
+            Dampening for momentum.
+        nesterov (`bool`, *optional*, defaults to `False`):
+            Whether to enable Nesterov momentum.
+    """
+
     lr: float = 1e-3
     momentum: float = 0.0
     dampening: float = 0.0
@@ -127,6 +216,15 @@ class SGDConfig(OptimizerConfig):
     grad_clip_norm: float = 10.0
 
     def build(self, params: OptimizerParams) -> torch.optim.Optimizer:
+        """Build a [`torch.optim.SGD`](https://docs.pytorch.org/docs/stable/generated/torch.optim.SGD.html) instance from this config.
+
+        Args:
+            params (`OptimizerParams`):
+                Parameters to optimize, as accepted by `torch.optim.SGD`.
+
+        Returns:
+            `torch.optim.Optimizer`: The built optimizer.
+        """
         kwargs = asdict(self)
         kwargs.pop("grad_clip_norm")
         return torch.optim.SGD(params, **kwargs)
@@ -168,8 +266,7 @@ class XVLAAdamWConfig(OptimizerConfig):
     soft_prompt_warmup_lr_scale: float | None = None  # If set, start soft-prompts at this scale (e.g., 0.01)
 
     def build(self, params: OptimizerParams) -> torch.optim.Optimizer:
-        """
-        Build AdamW optimizer with differential learning rates.
+        """Build AdamW optimizer with differential learning rates.
 
         Args:
             params: Must be a dict[str, Parameter] from dict(model.named_parameters())
@@ -239,10 +336,14 @@ class MultiAdamConfig(OptimizerConfig):
     This creates a dictionary of Adam optimizers, each with its own hyperparameters.
 
     Args:
-        lr: Default learning rate (used if not specified for a group)
-        weight_decay: Default weight decay (used if not specified for a group)
-        optimizer_groups: Dictionary mapping parameter group names to their hyperparameters
-        grad_clip_norm: Gradient clipping norm
+        lr (`float`, *optional*, defaults to 0.001):
+            Default learning rate, used for a group unless overridden in `optimizer_groups`.
+        weight_decay (`float`, *optional*, defaults to 0.0):
+            Default weight decay, used for a group unless overridden in `optimizer_groups`.
+        grad_clip_norm (`float`, *optional*, defaults to 10.0):
+            Maximum gradient norm; gradients are clipped to this value before each optimizer step.
+        optimizer_groups (`dict[str, dict[str, Any]]`, *optional*):
+            Per-group hyperparameter overrides (`lr`, `betas`, `eps`, `weight_decay`), keyed by group name.
     """
 
     lr: float = 1e-3
@@ -252,6 +353,7 @@ class MultiAdamConfig(OptimizerConfig):
 
     @property
     def builds_multiple_optimizers(self) -> bool:
+        """`bool`: Always `True`; `build()` returns a dict of optimizers, one per parameter group."""
         return True
 
     def build(self, params: OptimizerParams) -> dict[str, torch.optim.Optimizer]:
@@ -296,8 +398,10 @@ def save_optimizer_state(
     """Save optimizer state to disk (non-sharded runs; sharded runs use the DCP channel).
 
     Args:
-        optimizer: Either a single optimizer or a dictionary of optimizers.
-        save_dir: Directory to save the optimizer state.
+        optimizer (`torch.optim.Optimizer | dict[str, torch.optim.Optimizer]`):
+            Either a single optimizer or a dictionary of optimizers.
+        save_dir (`Path`):
+            Directory to save the optimizer state.
     """
     if isinstance(optimizer, dict):
         # Handle dictionary of optimizers
@@ -325,11 +429,14 @@ def load_optimizer_state(
     """Load optimizer state from disk.
 
     Args:
-        optimizer: Either a single optimizer or a dictionary of optimizers.
-        save_dir: Directory to load the optimizer state from.
+        optimizer (`torch.optim.Optimizer | dict[str, torch.optim.Optimizer]`):
+            Either a single optimizer or a dictionary of optimizers.
+        save_dir (`Path`):
+            Directory to load the optimizer state from.
 
     Returns:
-        The updated optimizer(s) with loaded state.
+        `torch.optim.Optimizer | dict[str, torch.optim.Optimizer]`: The updated optimizer(s) with loaded
+        state.
     """
     if isinstance(optimizer, dict):
         # Handle dictionary of optimizers

src/lerobot/optim/schedulers.py

@@ -36,24 +36,64 @@ else:
 
 @dataclass
 class LRSchedulerConfig(draccus.ChoiceRegistry, abc.ABC):
+    """Base configuration shared by every learning rate scheduler.
+
+    Concrete schedulers subclass this and register themselves with
+    `@LRSchedulerConfig.register_subclass("name")`, which is what makes `--scheduler.type=name` work on the
+    command line.
+
+    Args:
+        num_warmup_steps (`int | None`):
+            Number of steps over which the learning rate ramps up from 0 before the scheduler's own
+            behavior takes over. `None` disables warmup.
+    """
+
     num_warmup_steps: int | None
 
     @property
     def type(self) -> str:
+        """Return the registered name this config was registered under.
+
+        Returns:
+            `str`: The name passed to `@LRSchedulerConfig.register_subclass`, e.g. `"diffuser"`.
+        """
         return self.get_choice_name(self.__class__)
 
     @abc.abstractmethod
     def build(self, optimizer: Optimizer, num_training_steps: int) -> LRScheduler | None:
+        """Build the scheduler for a given optimizer and training length.
+
+        Args:
+            optimizer (`Optimizer`):
+                The optimizer whose learning rate the scheduler will adjust.
+            num_training_steps (`int`):
+                Total number of training steps, used to compute decay/annealing schedules.
+
+        Returns:
+            `LRScheduler | None`: The built scheduler.
+        """
         raise NotImplementedError
 
 
 @LRSchedulerConfig.register_subclass("diffuser")
 @dataclass
 class DiffuserSchedulerConfig(LRSchedulerConfig):
+    """A [`diffusers`](https://huggingface.co/docs/diffusers) learning rate schedule.
+
+    Args:
+        num_warmup_steps (`int`, *optional*):
+            Number of steps over which the learning rate ramps up from 0. `None` disables warmup.
+        name (`str`, *optional*, defaults to `"cosine"`):
+            Name of the `diffusers` schedule to build, e.g. `"cosine"`, `"linear"`, `"constant"`. See
+            [`diffusers.optimization.get_scheduler`](https://huggingface.co/docs/diffusers/api/schedulers/overview)
+            for the full list.
+    """
+
     name: str = "cosine"
     num_warmup_steps: int | None = None
 
     def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
+        """See [`~optim.schedulers.LRSchedulerConfig.build`]. Delegates to `diffusers.optimization.get_scheduler`."""
         require_package("diffusers", extra="diffusion")
 
         kwargs = {**asdict(self), "num_training_steps": num_training_steps, "optimizer": optimizer}
@@ -63,12 +103,31 @@ class DiffuserSchedulerConfig(LRSchedulerConfig):
 @LRSchedulerConfig.register_subclass("vqbet")
 @dataclass
 class VQBeTSchedulerConfig(LRSchedulerConfig):
+    """Used to train VQ-BeT: constant LR during VQ-VAE pretraining, then warmup and cosine decay.
+
+    Args:
+        num_warmup_steps (`int`):
+            Number of steps over which the learning rate ramps up from 0, counted from the end of VQ-VAE
+            pretraining.
+        num_vqvae_training_steps (`int`):
+            Number of initial steps spent pretraining the VQ-VAE, during which the LR stays at its peak.
+        num_cycles (`float`, *optional*, defaults to 0.5):
+            Number of cosine cycles in the decay phase; 0.5 decays smoothly to 0 by the end of training.
+    """
+
     num_warmup_steps: int
     num_vqvae_training_steps: int
     num_cycles: float = 0.5
 
     def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
+        """See [`~optim.schedulers.LRSchedulerConfig.build`].
+
+        Holds the LR at its peak during VQ-VAE pretraining, then applies linear warmup followed by cosine
+        decay for the remaining steps.
+        """
+
         def lr_lambda(current_step):
+            """Return the LR multiplier for `current_step`, per the VQ-BeT schedule."""
             if current_step < self.num_vqvae_training_steps:
                 return float(1)
             else:
@@ -90,14 +149,20 @@ class ConstantWithWarmupSchedulerConfig(LRSchedulerConfig):
 
     Mirrors the ``warmup_constant_lambda`` used by LingBot-VA (upstream ``wan_va/train.py``):
     the LR ramps linearly from 0 to the peak over ``num_warmup_steps`` steps, then stays flat.
+
+    Args:
+        num_warmup_steps (`int`, *optional*, defaults to 1000):
+            Number of steps over which the learning rate ramps up from 0 to its peak.
     """
 
     num_warmup_steps: int = 1000
 
     def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
+        """See [`~optim.schedulers.LRSchedulerConfig.build`]."""
         warmup_steps = self.num_warmup_steps or 0
 
         def lr_lambda(current_step):
+            """Return the LR multiplier for `current_step`: linear ramp, then constant `1.0`."""
             if current_step < warmup_steps:
                 return float(current_step) / float(max(1, warmup_steps))
             return 1.0
@@ -111,12 +176,19 @@ class CosineAnnealingWithWarmupSchedulerConfig(LRSchedulerConfig):
     """Linear warmup followed by cosine annealing from the peak LR to zero.
 
     Used by EVO1; the annealing phase always spans the remaining training steps.
+
+    Args:
+        num_warmup_steps (`int`):
+            Number of steps over which the learning rate ramps up from 0 to its peak.
     """
 
     num_warmup_steps: int
 
     def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
+        """See [`~optim.schedulers.LRSchedulerConfig.build`]."""
+
         def lr_lambda(current_step: int) -> float:
+            """Return the LR multiplier for `current_step`: linear warmup, then cosine annealing to 0."""
             if current_step < self.num_warmup_steps:
                 return current_step / max(1, self.num_warmup_steps)
             progress = (current_step - self.num_warmup_steps) / max(
@@ -134,6 +206,18 @@ class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):
 
     Automatically scales warmup and decay steps if num_training_steps < num_decay_steps.
     This ensures the learning rate schedule completes properly even with shorter training runs.
+
+    Args:
+        num_warmup_steps (`int`):
+            Number of steps over which the learning rate ramps up from `peak_lr / (num_warmup_steps + 1)`
+            to `peak_lr`.
+        num_decay_steps (`int`):
+            Number of steps over which the learning rate decays from `peak_lr` to `decay_lr`. Scaled down
+            automatically if `num_training_steps` is shorter than this.
+        peak_lr (`float`):
+            Learning rate reached at the end of warmup.
+        decay_lr (`float`):
+            Learning rate reached at the end of decay.
     """
 
     num_warmup_steps: int
@@ -142,6 +226,11 @@ class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):
     decay_lr: float
 
     def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
+        """See [`~optim.schedulers.LRSchedulerConfig.build`].
+
+        If `num_training_steps` is shorter than `num_decay_steps`, scales `num_warmup_steps` and
+        `num_decay_steps` down proportionally so the schedule still completes.
+        """
         # Auto-scale scheduler parameters if training steps are shorter than configured decay steps
         actual_warmup_steps = self.num_warmup_steps
         actual_decay_steps = self.num_decay_steps
@@ -161,13 +250,17 @@ class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):
             )
 
         def lr_lambda(current_step):
+            """Return the LR multiplier for `current_step`: linear warmup, then cosine decay."""
+
             def linear_warmup_schedule(current_step):
+                """Return the LR multiplier during warmup, ramping from `1 / (warmup + 1)` to 1."""
                 if current_step <= 0:
                     return 1 / (actual_warmup_steps + 1)
                 frac = 1 - current_step / actual_warmup_steps
                 return (1 / (actual_warmup_steps + 1) - 1) * frac + 1
 
             def cosine_decay_schedule(current_step):
+                """Return the LR multiplier during decay, from 1 down to `decay_lr / peak_lr`."""
                 step = min(current_step, actual_decay_steps)
                 cosine_decay = 0.5 * (1 + math.cos(math.pi * step / actual_decay_steps))
                 alpha = self.decay_lr / self.peak_lr
@@ -183,11 +276,30 @@ class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):
 
 
 def save_scheduler_state(scheduler: LRScheduler, save_dir: Path) -> None:
+    """Save a scheduler's state to disk.
+
+    Args:
+        scheduler (`LRScheduler`):
+            The scheduler whose state to save.
+        save_dir (`Path`):
+            Directory to save the scheduler state.
+    """
     state_dict = scheduler.state_dict()
     write_json(state_dict, save_dir / SCHEDULER_STATE)
 
 
 def load_scheduler_state(scheduler: LRScheduler, save_dir: Path) -> LRScheduler:
+    """Load a scheduler's state from disk.
+
+    Args:
+        scheduler (`LRScheduler`):
+            The scheduler to load state into.
+        save_dir (`Path`):
+            Directory to load the scheduler state from.
+
+    Returns:
+        `LRScheduler`: The same scheduler, with its state loaded.
+    """
     state_dict = deserialize_json_into_object(save_dir / SCHEDULER_STATE, scheduler.state_dict())
     scheduler.load_state_dict(state_dict)
     return scheduler

utils/check_docstrings.py

@@ -60,7 +60,7 @@ PATH_TO_LEROBOT = PATH_TO_REPO / "src" / "lerobot"
 # Modules whose public objects are checked. Add a module here once its docstrings follow the standard.
 MODULES_TO_CHECK = [
     "lerobot.robots",
-    "lerobot.configs",
+    "lerobot.optim",
 ]
 
 # Objects that do not yet follow the standard, so the check can be green from day one. Removing an entry