From 582e9536764b2135effb6c3da0062948b1d935d6 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Wed, 8 Jul 2026 12:08:17 +0200 Subject: [PATCH] fix(train): fix VF scheduler config and reward model hub push Add missing peak_lr and decay_lr to DistributionalVFConfig scheduler preset. Fix push_model_to_hub call for reward models. --- .../configuration_distributional_value_function.py | 2 ++ src/lerobot/scripts/lerobot_train.py | 2 ++ 2 files changed, 4 insertions(+) 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 1a1f33c34..894d8eb2a 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 @@ -98,6 +98,8 @@ class DistributionalVFConfig(RewardModelConfig): return CosineDecayWithWarmupSchedulerConfig( num_warmup_steps=500, num_decay_steps=50000, + peak_lr=3e-4, + decay_lr=1e-6, ) def validate_features(self) -> None: diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 6e8458523..d8d3dc803 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -741,6 +741,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # PEFT only applies when training a policy — reward models use the plain path. if not cfg.is_reward_model_training and cfg.policy.use_peft: unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model, dataset_meta=dataset.meta) + elif cfg.is_reward_model_training: + unwrapped_model.push_model_to_hub(cfg) else: unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict, dataset_meta=dataset.meta) preprocessor.push_to_hub(active_cfg.repo_id)