refactor(rewards): add get_optim_params() with vision_encoder_lr_multiplier for differential LR

src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py

@@ -77,6 +77,7 @@ class DistributionalVFConfig(RewardModelConfig):
     freeze_vision_encoder: bool = False
     freeze_language_model: bool = False
     stop_gradient_to_vlm: bool = False
+    vision_encoder_lr_multiplier: float = 0.5
 
     # Normalization
     normalization_mapping: dict[str, NormalizationMode] = field(
@@ -87,16 +88,16 @@ class DistributionalVFConfig(RewardModelConfig):
 
     def get_optimizer_preset(self) -> AdamWConfig:
         return AdamWConfig(
-            lr=3e-4,
+            lr=1e-4,
-            weight_decay=1e-4,
+            weight_decay=1e-5,
             grad_clip_norm=1.0,
         )
 
     def get_scheduler_preset(self) -> CosineDecayWithWarmupSchedulerConfig:
         return CosineDecayWithWarmupSchedulerConfig(
             num_warmup_steps=500,
-            num_decay_steps=50000,
+            num_decay_steps=30000,
-            peak_lr=3e-4,
+            peak_lr=1e-4,
             decay_lr=1e-6,
         )
 

src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py

@@ -194,6 +194,25 @@ class DistributionalVFRewardModel(PreTrainedRewardModel):
             self.gemma3.eval()
         return self
 
+    def get_optim_params(self) -> list[dict]:
+        """Optimizer param groups with per-component learning rates."""
+        vision_params = []
+        other_params = []
+
+        for name, param in self.named_parameters():
+            if not param.requires_grad:
+                continue
+            if name.startswith("vision_encoder"):
+                vision_params.append(param)
+            else:
+                other_params.append(param)
+
+        base_lr = self.config.get_optimizer_preset().lr
+        return [
+            {"params": other_params},
+            {"params": vision_params, "lr": base_lr * self.config.vision_encoder_lr_multiplier},
+        ]
+
     def embed_image(self, image: Tensor) -> Tensor:
         """Embed images: SigLIP2 → projection → [B, num_patches, gemma3_hidden].