admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-26 11:16:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(train): streamline reward model training logic

Browse Source
This commit is contained in:
Khalil Meftah
2026-04-21 13:56:18 +02:00
parent 25df7fdd67
commit 001439f02f
1 changed files with 13 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/lerobot_train.py
+13 -12
View File
@@ -193,8 +193,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
# Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting. # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting.
# Force the device to be CPU when policy.device is set to CPU. # Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training).
force_cpu = cfg.policy.device == "cpu" force_cpu = cfg.trainable_config.device == "cpu"
accelerator = Accelerator( accelerator = Accelerator(
step_scheduler_with_optimizer=False, step_scheduler_with_optimizer=False,
kwargs_handlers=[ddp_kwargs], kwargs_handlers=[ddp_kwargs],
@@ -295,7 +295,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path: if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path:
processor_kwargs["dataset_stats"] = dataset.meta.stats processor_kwargs["dataset_stats"] = dataset.meta.stats
if cfg.is_reward_model_training and cfg.reward_model.type == "sarm": if cfg.is_reward_model_training:
processor_kwargs["dataset_meta"] = dataset.meta processor_kwargs["dataset_meta"] = dataset.meta
if not cfg.is_reward_model_training and processor_pretrained_path is not None: if not cfg.is_reward_model_training and processor_pretrained_path is not None:
@@ -376,13 +376,13 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})") logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
# create dataloader for offline training # create dataloader for offline training
if hasattr(cfg.policy, "drop_n_last_frames"): if hasattr(active_cfg, "drop_n_last_frames"):
shuffle = False shuffle = False
sampler = EpisodeAwareSampler( sampler = EpisodeAwareSampler(
dataset.meta.episodes["dataset_from_index"], dataset.meta.episodes["dataset_from_index"],
dataset.meta.episodes["dataset_to_index"], dataset.meta.episodes["dataset_to_index"],
episode_indices_to_use=dataset.episodes, episode_indices_to_use=dataset.episodes,
drop_n_last_frames=cfg.policy.drop_n_last_frames, drop_n_last_frames=active_cfg.drop_n_last_frames,
shuffle=True, shuffle=True,
) )
else: else:
@@ -559,14 +559,15 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process: if is_main_process:
logging.info("End of training") logging.info("End of training")
if cfg.policy.push_to_hub: if getattr(active_cfg, "push_to_hub", False):
unwrapped_policy = accelerator.unwrap_model(policy) unwrapped_model = accelerator.unwrap_model(policy)
if cfg.policy.use_peft: # PEFT only applies when training a policy — reward models use the plain path.
unwrapped_policy.push_model_to_hub(cfg, peft_model=unwrapped_policy) if not cfg.is_reward_model_training and cfg.policy.use_peft:
unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model)
else: else:
unwrapped_policy.push_model_to_hub(cfg) unwrapped_model.push_model_to_hub(cfg)
preprocessor.push_to_hub(cfg.policy.repo_id) preprocessor.push_to_hub(active_cfg.repo_id)
postprocessor.push_to_hub(cfg.policy.repo_id) postprocessor.push_to_hub(active_cfg.repo_id)
# Properly clean up the distributed process group # Properly clean up the distributed process group
accelerator.wait_for_everyone() accelerator.wait_for_everyone()