diff --git a/AGENT_GUIDE.md b/AGENT_GUIDE.md index 03b270dce..005ce6f1f 100644 --- a/AGENT_GUIDE.md +++ b/AGENT_GUIDE.md @@ -321,10 +321,11 @@ SmolVLA ships with `freeze_vision_encoder=True`. Unfreezing usually **improves p ```bash lerobot-train ... --policy.type=smolvla \ - --policy.freeze_vision_encoder=false \ - --policy.train_expert_only=false + --policy.fine_tune_vision_encoder=true ``` +This selectively trains the vision encoder and connector while leaving the language model frozen. Their learning rate defaults to `0.1 × optimizer_lr`; adjust it with `--policy.vision_encoder_lr_multiplier` if needed. + ### 7.7 Signals to stop / keep going - Train loss plateaus → stop, save a Hub checkpoint. diff --git a/docs/source/smolvla.mdx b/docs/source/smolvla.mdx index e28270c9b..7deb16714 100644 --- a/docs/source/smolvla.mdx +++ b/docs/source/smolvla.mdx @@ -70,6 +70,19 @@ cd lerobot && lerobot-train \ GPU allows it, as long as loading times remain short. +For tasks that require adapting visual features, such as distinguishing new colors or shapes, selectively +fine-tune the vision encoder and its connector: + +```bash +lerobot-train ... \ + --policy.path=lerobot/smolvla_base \ + --policy.fine_tune_vision_encoder=true +``` + +This keeps the language model frozen with the default `train_expert_only=true` setting and trains the vision +path at `0.1` times the main learning rate by default. Fine-tuning the vision encoder increases memory use and +can reduce the model's general visual knowledge, so enable it only when the frozen encoder is insufficient. + Fine-tuning is an art. For a complete overview of the options for finetuning, run ```bash diff --git a/src/lerobot/policies/smolvla/configuration_smolvla.py b/src/lerobot/policies/smolvla/configuration_smolvla.py index f810c8e2f..fd35dfc62 100644 --- a/src/lerobot/policies/smolvla/configuration_smolvla.py +++ b/src/lerobot/policies/smolvla/configuration_smolvla.py @@ -67,7 +67,8 @@ class SmolVLAConfig(PreTrainedConfig): # Finetuning settings freeze_vision_encoder: bool = True - fine_tune_vision_encoder: bool = False + fine_tune_vision_encoder: bool = False # Fine-tune vision + connector; takes priority over freezing. + vision_encoder_lr_multiplier: float = 0.1 train_expert_only: bool = True train_state_proj: bool = True @@ -111,6 +112,12 @@ class SmolVLAConfig(PreTrainedConfig): super().__post_init__() """Input validation (not exhaustive).""" + if self.fine_tune_vision_encoder: + self.freeze_vision_encoder = False + if self.vision_encoder_lr_multiplier <= 0: + raise ValueError( + f"`vision_encoder_lr_multiplier` must be positive, got {self.vision_encoder_lr_multiplier}." + ) if self.n_action_steps > self.chunk_size: raise ValueError( f"The chunk size is the upper bound for the number of action steps per model invocation. Got " diff --git a/src/lerobot/policies/smolvla/modeling_smolvla.py b/src/lerobot/policies/smolvla/modeling_smolvla.py index 7b42c2857..a7c0c0219 100644 --- a/src/lerobot/policies/smolvla/modeling_smolvla.py +++ b/src/lerobot/policies/smolvla/modeling_smolvla.py @@ -162,14 +162,6 @@ 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): @@ -194,8 +186,27 @@ class SmolVLAPolicy(PreTrainedPolicy): if model_value is not None: model_value.rtc_processor = self.rtc_processor - def get_optim_params(self) -> dict: - return self.parameters() + def get_optim_params(self): + if not self.config.fine_tune_vision_encoder: + return self.parameters() + + vision_params = [] + other_params = [] + for name, param in self.named_parameters(): + if not param.requires_grad: + continue + if ".vision_model." in name or ".connector." in name: + vision_params.append(param) + else: + other_params.append(param) + + return [ + {"params": other_params}, + { + "params": vision_params, + "lr": self.config.optimizer_lr * self.config.vision_encoder_lr_multiplier, + }, + ] def _get_action_chunk( self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs] @@ -501,6 +512,7 @@ class VLAFlowMatching(nn.Module): self.vlm_with_expert = SmolVLMWithExpertModel( model_id=self.config.vlm_model_name, freeze_vision_encoder=self.config.freeze_vision_encoder, + fine_tune_vision_encoder=self.config.fine_tune_vision_encoder, train_expert_only=self.config.train_expert_only, load_vlm_weights=self.config.load_vlm_weights, attention_mode=self.config.attention_mode, diff --git a/src/lerobot/policies/smolvla/smolvlm_with_expert.py b/src/lerobot/policies/smolvla/smolvlm_with_expert.py index e4cbd3d52..62a30a417 100644 --- a/src/lerobot/policies/smolvla/smolvlm_with_expert.py +++ b/src/lerobot/policies/smolvla/smolvlm_with_expert.py @@ -78,6 +78,7 @@ class SmolVLMWithExpertModel(nn.Module): load_vlm_weights: bool = True, train_expert_only: bool = True, freeze_vision_encoder: bool = False, + fine_tune_vision_encoder: bool = False, attention_mode: str = "self_attn", num_expert_layers: int = -1, num_vlm_layers: int = -1, @@ -141,6 +142,7 @@ class SmolVLMWithExpertModel(nn.Module): self.num_key_value_heads = self.config.text_config.num_key_value_heads self.freeze_vision_encoder = freeze_vision_encoder + self.fine_tune_vision_encoder = fine_tune_vision_encoder self.train_expert_only = train_expert_only self.attention_mode = attention_mode self.expert_hidden_size = lm_expert_config.hidden_size @@ -150,10 +152,6 @@ class SmolVLMWithExpertModel(nn.Module): return self.vlm.model def set_requires_grad(self): - if self.freeze_vision_encoder: - self.get_vlm_model().vision_model.eval() - for params in self.get_vlm_model().vision_model.parameters(): - params.requires_grad = False if self.train_expert_only: self.vlm.eval() for params in self.vlm.parameters(): @@ -176,6 +174,18 @@ class SmolVLMWithExpertModel(nn.Module): for name, params in self.vlm.named_parameters(): if any(k in name for k in frozen_layers): params.requires_grad = False + + if self.freeze_vision_encoder: + self.get_vlm_model().vision_model.eval() + for params in self.get_vlm_model().vision_model.parameters(): + params.requires_grad = False + + if self.fine_tune_vision_encoder: + for params in self.get_vlm_model().vision_model.parameters(): + params.requires_grad = True + for params in self.get_vlm_model().connector.parameters(): + params.requires_grad = True + # To avoid unused params issue with distributed training for name, params in self.lm_expert.named_parameters(): if "lm_head" in name: @@ -184,11 +194,15 @@ class SmolVLMWithExpertModel(nn.Module): def train(self, mode: bool = True): super().train(mode) + if self.train_expert_only: + self.vlm.eval() + if self.freeze_vision_encoder: self.get_vlm_model().vision_model.eval() - if self.train_expert_only: - self.vlm.eval() + if self.fine_tune_vision_encoder: + self.get_vlm_model().vision_model.train(mode) + self.get_vlm_model().connector.train(mode) def embed_image(self, image: torch.Tensor): patch_attention_mask = None