diff --git a/src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py b/src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py index 59d0d94d7..31bf16e8d 100644 --- a/src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py +++ b/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, - weight_decay=1e-4, + lr=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, - peak_lr=3e-4, + num_decay_steps=30000, + peak_lr=1e-4, decay_lr=1e-6, ) diff --git a/src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py b/src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py index 519de30ea..5d1c610b3 100644 --- a/src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py +++ b/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].