feat(policiy): fine tune vision encoder refactor completness

Fixes #1774

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.

docs/source/smolvla.mdx

@@ -70,6 +70,19 @@ cd lerobot && lerobot-train \
   GPU allows it, as long as loading times remain short.
 </Tip>
 
+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

src/lerobot/policies/smolvla/configuration_smolvla.py

@@ -67,6 +67,8 @@ class SmolVLAConfig(PreTrainedConfig):
 
     # Finetuning settings
     freeze_vision_encoder: bool = True
+    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
 
@@ -110,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 "

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -186,9 +186,28 @@ class SmolVLAPolicy(PreTrainedPolicy):
             if model_value is not None:
                 model_value.rtc_processor = self.rtc_processor
 
-    def get_optim_params(self) -> dict:
+    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]
     ) -> Tensor:
@@ -493,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,

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