diff --git a/docs/source/api/optim.mdx b/docs/source/api/optim.mdx new file mode 100644 index 000000000..678e38886 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index a350cd785..992b47dbf 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/lerobot/optim/factory.py b/src/lerobot/optim/factory.py index ce519e0b2..0faee6d98 100644 --- a/src/lerobot/optim/factory.py +++ b/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: diff --git a/src/lerobot/optim/optimizers.py b/src/lerobot/optim/optimizers.py index 820579c06..f2fdb60a0 100644 --- a/src/lerobot/optim/optimizers.py +++ b/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 diff --git a/src/lerobot/optim/schedulers.py b/src/lerobot/optim/schedulers.py index 2a80f74fb..58cedf9d7 100644 --- a/src/lerobot/optim/schedulers.py +++ b/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 diff --git a/utils/check_docstrings.py b/utils/check_docstrings.py index 254c7c872..57678d410 100644 --- a/utils/check_docstrings.py +++ b/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