Refactor: Move PEFT config from training script to policy level (#2806)

* move peft config from `lerobot_train` to policy level

* Update src/lerobot/scripts/lerobot_train.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Michel Aractingi <michel.aractingi@huggingface.co>

* copilot response

* Change the polciy function to return targets rather than peft config.`_get_default_peft_targets()` override in PI0, PI0.5, SmolVLA

* remove none check when building config dict

---------

Signed-off-by: Michel Aractingi <michel.aractingi@huggingface.co>

src/lerobot/policies/pretrained.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import abc
 import builtins
+import dataclasses
 import logging
 import os
 from importlib.resources import files
@@ -265,3 +266,166 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         card = ModelCard.from_template(card_data, template_str=template_card)
         card.validate()
         return card
+
+    def wrap_with_peft(
+        self,
+        peft_config=None,
+        peft_cli_overrides: dict | None = None,
+    ) -> "PreTrainedPolicy":
+        """
+        Wrap this policy with PEFT adapters for parameter-efficient fine-tuning.
+
+        This method is the single entry point for PEFT integration. Subclasses should
+        override `_get_default_peft_targets()` to provide default target modules, and
+        `_validate_peft_config()` for policy-specific validation.
+
+        Args:
+            peft_config: Optional PEFT adapter configuration (e.g., LoraConfig).
+                If provided, used directly (with CLI overrides applied).
+            peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
+                These are merged with policy defaults to build the final config.
+        """
+        from peft import get_peft_model
+
+        # If user provided a complete config, use it directly (with overrides)
+        if peft_config is not None:
+            final_config = peft_config
+            if peft_cli_overrides:
+                final_config = self._apply_peft_cli_overrides(final_config, peft_cli_overrides)
+        else:
+            # Build config from defaults + CLI overrides
+            final_config = self._build_peft_config(peft_cli_overrides or {})
+
+        # Validate the configuration
+        self._validate_peft_config(final_config)
+
+        # Freeze base parameters, only adapter params will be trained
+        for p in self.parameters():
+            p.requires_grad_(False)
+
+        # Store pretrained path for PEFT's base_model_name_or_path
+        if self.config.pretrained_path:
+            self.name_or_path = str(self.config.pretrained_path)
+
+        # Wrap with PEFT
+        peft_model = get_peft_model(self, final_config)
+
+        # Mark config as using PEFT for proper loading later
+        peft_model.config.use_peft = True
+
+        logging.info(f"Wrapped {self.name} with PEFT ({type(final_config).__name__})")
+        return peft_model
+
+    def _get_default_peft_targets(self) -> dict[str, any] | None:
+        """
+        Return default PEFT target modules for this policy.
+
+        Override this in subclasses to provide policy-specific defaults. These defaults
+        are PEFT-method agnostic - they only specify which modules to target.
+
+        """
+        return None
+
+    def _validate_peft_config(self, peft_config) -> None:
+        """
+        Validate the PEFT configuration for this policy.
+
+        Override this in subclasses to add policy-specific validation or warnings.
+        The default implementation checks that a pretrained_path exists.
+
+        Args:
+            peft_config: The PEFT configuration to validate.
+
+        Raises:
+            ValueError: If the configuration is invalid.
+        """
+        if not self.config.pretrained_path:
+            raise ValueError(
+                "Training from scratch using PEFT is unlikely to yield good results. "
+                "Supply a `policy.pretrained_path` to fine-tune an existing model."
+            )
+
+    def _preprocess_peft_cli_overrides(self, cli_overrides: dict, peft_method_type) -> dict:
+        """
+        Preprocess CLI overrides: rename keys and handle method-specific init_type.
+
+        Args:
+            cli_overrides: Dict of CLI options (will be copied, not mutated).
+            peft_method_type: The PeftType enum value for the PEFT method.
+
+        Returns:
+            Preprocessed dict with renamed keys and init_type mapped to method-specific key.
+        """
+        from peft import PeftType
+
+        cli_overrides = cli_overrides.copy()
+
+        # Handle the full_training_modules -> modules_to_save rename
+        if "full_training_modules" in cli_overrides:
+            cli_overrides["modules_to_save"] = cli_overrides.pop("full_training_modules")
+
+        # Remove method_type as it's handled separately
+        cli_overrides.pop("method_type", None)
+
+        # Handle init_type specially based on PEFT method
+        init_type = cli_overrides.pop("init_type", None)
+        if init_type is not None:
+            if peft_method_type == PeftType.LORA:
+                cli_overrides["init_lora_weights"] = init_type
+            elif peft_method_type == PeftType.MISS:
+                cli_overrides["init_weights"] = init_type
+            else:
+                raise ValueError(f"Init type '{init_type}' unknown for PEFT method {peft_method_type}.")
+
+        return cli_overrides
+
+    def _build_peft_config(self, cli_overrides: dict):
+        """Build a PEFT config from policy defaults and CLI overrides."""
+        from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
+
+        # Determine PEFT method type (default to LORA)
+        method_type_str = cli_overrides.get("method_type") or "lora"
+        peft_method_type = PeftType[method_type_str.upper()]
+        peft_config_cls = PEFT_TYPE_TO_CONFIG_MAPPING[peft_method_type]
+
+        # Preprocess CLI overrides
+        cli_overrides = self._preprocess_peft_cli_overrides(cli_overrides, peft_method_type)
+
+        # Start with policy defaults, apply CLI overrides
+        config_dict = dict(self._get_default_peft_targets() or {})
+        for key, value in cli_overrides.items():
+            if value is not None:
+                config_dict[key] = value
+
+        # Ensure we have target_modules
+        if not config_dict.get("target_modules"):
+            raise ValueError(
+                f"Policy '{self.name}' does not define default target_modules. "
+                "Please pass --peft.target_modules explicitly."
+            )
+
+        return peft_config_cls(**config_dict)
+
+    def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
+        """Apply CLI overrides to an existing PEFT config."""
+        from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
+
+        # Get method type from existing config or CLI override
+        method_type_str = cli_overrides.get("method_type")
+        if method_type_str:
+            peft_method_type = PeftType[method_type_str.upper()]
+            peft_config_cls = PEFT_TYPE_TO_CONFIG_MAPPING[peft_method_type]
+        else:
+            peft_method_type = PeftType(peft_config.peft_type)
+            peft_config_cls = type(peft_config)
+
+        # Preprocess CLI overrides
+        cli_overrides = self._preprocess_peft_cli_overrides(cli_overrides, peft_method_type)
+
+        # Start with existing config, apply CLI overrides
+        config_dict = {k: v for k, v in dataclasses.asdict(peft_config).items() if not k.startswith("_")}
+        for key, value in cli_overrides.items():
+            if value is not None:
+                config_dict[key] = value
+
+        return peft_config_cls(**config_dict)