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

@@ -447,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/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
+        cfg (`TrainPipelineConfig`):
-        policy (PreTrainedPolicy): The policy config from which parameters and presets must be taken from.
+            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)
+        lr (`float`, *optional*, defaults to 0.001):
-        weight_decay: Default weight decay (used if not specified for a group)
+            Default learning rate, used for a group unless overridden in `optimizer_groups`.
-        optimizer_groups: Dictionary mapping parameter group names to their hyperparameters
+        weight_decay (`float`, *optional*, defaults to 0.0):
-        grad_clip_norm: Gradient clipping norm
+            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.
+        optimizer (`torch.optim.Optimizer | dict[str, torch.optim.Optimizer]`):
-        save_dir: Directory to save the optimizer state.
+            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.
+        optimizer (`torch.optim.Optimizer | dict[str, torch.optim.Optimizer]`):
-        save_dir: Directory to load the optimizer state from.
+            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,6 +60,7 @@ PATH_TO_LEROBOT = PATH_TO_REPO / "src" / "lerobot"
 # Modules whose public objects are checked. Add a module here once its docstrings follow the standard.
 MODULES_TO_CHECK = [
     "lerobot.robots",
+    "lerobot.optim",
 ]
 
 # Objects that do not yet follow the standard, so the check can be green from day one. Removing an entry