diff --git a/src/lerobot/policies/smolvla/configuration_smolvla.py b/src/lerobot/policies/smolvla/configuration_smolvla.py index 6d5288db3..f810c8e2f 100644 --- a/src/lerobot/policies/smolvla/configuration_smolvla.py +++ b/src/lerobot/policies/smolvla/configuration_smolvla.py @@ -67,6 +67,7 @@ class SmolVLAConfig(PreTrainedConfig): # Finetuning settings freeze_vision_encoder: bool = True + fine_tune_vision_encoder: bool = False train_expert_only: bool = True train_state_proj: bool = True diff --git a/src/lerobot/policies/smolvla/modeling_smolvla.py b/src/lerobot/policies/smolvla/modeling_smolvla.py index 7f54d03f0..7b42c2857 100644 --- a/src/lerobot/policies/smolvla/modeling_smolvla.py +++ b/src/lerobot/policies/smolvla/modeling_smolvla.py @@ -162,6 +162,14 @@ class SmolVLAPolicy(PreTrainedPolicy): self.config = config self.init_rtc_processor() self.model = VLAFlowMatching(config, rtc_processor=self.rtc_processor) + + if self.config.fine_tune_vision_encoder: + self.model.vlm_with_expert.freeze_vision_encoder = False + for params in self.model.vlm_with_expert.get_vlm_model().vision_model.parameters(): + params.requires_grad = True + for params in self.model.vlm_with_expert.get_vlm_model().connector.parameters(): + params.requires_grad = True + self.reset() def reset(self):