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)