diff --git a/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py b/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py index acf2b7b6e..2902f8216 100644 --- a/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py +++ b/src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py @@ -66,6 +66,9 @@ class PI05OpenPIConfig(PreTrainedConfig): compile_mode: str = "max-autotune" # Torch compile mode device: str | None = None # Device to use for the model (None = auto-detect) + # Pretrained model loading + pretrained_path: str | None = None # Path or repo_id to load pretrained weights from + # Optimizer settings: see openpi `AdamW` and optimizer_lr: float = 2.5e-5 # see openpi `CosineDecaySchedule: peak_lr` optimizer_betas: tuple[float, float] = (0.9, 0.95) diff --git a/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py b/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py index 455640213..fb2a4ea49 100644 --- a/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py +++ b/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py @@ -875,8 +875,8 @@ class PI05OpenPIPolicy(PreTrainedPolicy): if pretrained_name_or_path is None: raise ValueError("pretrained_name_or_path is required") - # Create default config - config = cls.config_class() + # Use provided config if available, otherwise create default config + config = kwargs.get("config", cls.config_class()) # Initialize model without loading weights # Check if dataset_stats were provided in kwargs diff --git a/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py b/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py index 973638fae..52e587abe 100644 --- a/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py @@ -63,6 +63,9 @@ class PI0OpenPIConfig(PreTrainedConfig): compile_mode: str = "max-autotune" # Torch compile mode device: str | None = None # Device to use for the model (None = auto-detect) + # Pretrained model loading + pretrained_path: str | None = None # Path or repo_id to load pretrained weights from + # Optimizer settings: see openpi `AdamW` and optimizer_lr: float = 2.5e-5 # see openpi `CosineDecaySchedule: peak_lr` optimizer_betas: tuple[float, float] = (0.9, 0.95) diff --git a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py index 3f82cfbcf..4bb1c83d9 100644 --- a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py @@ -894,8 +894,8 @@ class PI0OpenPIPolicy(PreTrainedPolicy): if pretrained_name_or_path is None: raise ValueError("pretrained_name_or_path is required") - # Create default config - config = cls.config_class() + # Use provided config if available, otherwise create default config + config = kwargs.get("config", cls.config_class()) # Initialize model without loading weights # Check if dataset_stats were provided in kwargs