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

fix

Browse Source
This commit is contained in:
Pepijn
2025-09-24 12:01:03 +02:00
parent bab60cf02f
commit 5489d8073d
3 changed files with 127 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/configs/train.py
-4
View File
@@ -63,10 +63,6 @@ class TrainPipelineConfig(HubMixin):
scheduler: LRSchedulerConfig | None = None scheduler: LRSchedulerConfig | None = None
eval: EvalConfig = field(default_factory=EvalConfig) eval: EvalConfig = field(default_factory=EvalConfig)
wandb: WandBConfig = field(default_factory=WandBConfig) wandb: WandBConfig = field(default_factory=WandBConfig)
# Accelerate configuration for multi-GPU training
use_accelerate: bool = False
gradient_accumulation_steps: int = 1
mixed_precision: str = "no" # Options: "no", "fp16", "bf16"
def __post_init__(self): def __post_init__(self):
self.checkpoint_path = None self.checkpoint_path = None
src/lerobot/optim/schedulers.py
+5 -5
View File
@@ -92,20 +92,20 @@ class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):
def lr_lambda(current_step): def lr_lambda(current_step):
def linear_warmup_schedule(current_step): def linear_warmup_schedule(current_step):
if current_step <= 0: if current_step <= 0:
return 0.1 # Start at 10% of peak LR return 0.1 # Start at 10% of peak LR instead of 0.1%
if current_step >= self.num_warmup_steps: if current_step >= self.num_warmup_steps:
return 1.0 # Reach 100% at end of warmup return 1.0 # Reach 100% at end of warmup
# Linear interpolation from 0.1 to 1.0 # Linear interpolation from 10% to 100% of peak LR
return 0.1 + 0.9 * (current_step / self.num_warmup_steps) return 0.1 + 0.9 * (current_step / self.num_warmup_steps)
def cosine_decay_schedule(current_step): def cosine_decay_schedule(current_step):
# Steps since warmup ended (this was the bug!)
decay_step = current_step - self.num_warmup_steps decay_step = current_step - self.num_warmup_steps
decay_step = min(decay_step, self.num_decay_steps) decay_step = max(0, min(decay_step, self.num_decay_steps))
cosine_decay = 0.5 * (1 + math.cos(math.pi * decay_step / self.num_decay_steps)) cosine_decay = 0.5 * (1 + math.cos(math.pi * decay_step / self.num_decay_steps))
alpha = self.decay_lr / self.peak_lr alpha = self.decay_lr / self.peak_lr
return (1 - alpha) * cosine_decay + alpha decayed = (1 - alpha) * cosine_decay + alpha
return decayed
if current_step < self.num_warmup_steps: if current_step < self.num_warmup_steps:
return linear_warmup_schedule(current_step) return linear_warmup_schedule(current_step)
src/lerobot/scripts/train.py
+122 -209
View File
@@ -91,49 +91,42 @@ def update_policy(
- A dictionary of outputs from the policy's forward pass, for logging purposes. - A dictionary of outputs from the policy's forward pass, for logging purposes.
""" """
start_time = time.perf_counter() start_time = time.perf_counter()
device = get_device_from_parameters(policy) device = get_device_from_parameters(policy) if accelerator is None else accelerator.device
policy.train() policy.train()
# Handle mixed precision differently for accelerate vs non-accelerate
if accelerator is not None: if accelerator is not None:
# Accelerate handles mixed precision internally # Use accelerate's autocast and backward
with accelerator.autocast() if use_amp else nullcontext(): with accelerator.autocast():
loss, output_dict = policy.forward(batch) loss, output_dict = policy.forward(batch)
# Use accelerator's backward method
accelerator.backward(loss) accelerator.backward(loss)
else:
# Original behavior for non-accelerate
with torch.autocast(device_type=device.type) if use_amp else nullcontext():
loss, output_dict = policy.forward(batch)
# TODO(rcadene): policy.unnormalize_outputs(out_dict)
grad_scaler.scale(loss).backward()
if accelerator is not None: # Use accelerate's gradient clipping
# Accelerate handles gradient scaling internally if grad_clip_norm > 0:
grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm) grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
if grad_norm is None: if grad_norm is None:
grad_norm = 0.0 grad_norm = 0.0
else:
grad_norm = torch.tensor(0.0, device=device)
with lock if lock is not None else nullcontext(): with lock if lock is not None else nullcontext():
optimizer.step() optimizer.step()
optimizer.zero_grad() optimizer.zero_grad()
else: else:
# Original gradient handling for non-accelerate # Original single-GPU path
# Unscale the gradient of the optimizer's assigned params in-place **prior to gradient clipping**. with torch.autocast(device_type=device.type) if use_amp else nullcontext():
grad_scaler.unscale_(optimizer) loss, output_dict = policy.forward(batch)
grad_scaler.scale(loss).backward()
grad_scaler.unscale_(optimizer)
grad_norm = torch.nn.utils.clip_grad_norm_( grad_norm = torch.nn.utils.clip_grad_norm_(
policy.parameters(), policy.parameters(),
grad_clip_norm, grad_clip_norm,
error_if_nonfinite=False, error_if_nonfinite=False,
) )
# Optimizer's gradients are already unscaled, so scaler.step does not unscale them,
# although it still skips optimizer.step() if the gradients contain infs or NaNs.
with lock if lock is not None else nullcontext(): with lock if lock is not None else nullcontext():
grad_scaler.step(optimizer) grad_scaler.step(optimizer)
# Updates the scale for next iteration.
grad_scaler.update() grad_scaler.update()
optimizer.zero_grad() optimizer.zero_grad()
# Step through pytorch scheduler at every batch instead of epoch # Step through pytorch scheduler at every batch instead of epoch
@@ -168,16 +161,11 @@ def train(cfg: TrainPipelineConfig):
cfg: A `TrainPipelineConfig` object containing all training configurations. cfg: A `TrainPipelineConfig` object containing all training configurations.
""" """
cfg.validate() cfg.validate()
# Only log config on main process when using accelerate
# For now we don't know if we're using accelerate yet, so we'll log this always
# and fix the duplicate later if needed
logging.info(pformat(cfg.to_dict())) logging.info(pformat(cfg.to_dict()))
# Initialize Accelerate if requested # Initialize Accelerate if enabled
accelerator = None accelerator = None
if cfg.use_accelerate: if cfg.use_accelerate:
# Configure DDP to handle unused parameters
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
accelerator = Accelerator( accelerator = Accelerator(
gradient_accumulation_steps=cfg.gradient_accumulation_steps, gradient_accumulation_steps=cfg.gradient_accumulation_steps,
@@ -191,19 +179,17 @@ def train(cfg: TrainPipelineConfig):
) )
logging.info(f"Training on {accelerator.num_processes} processes") logging.info(f"Training on {accelerator.num_processes} processes")
else: else:
# Check device is available (original behavior)
device = get_safe_torch_device(cfg.policy.device, log=True) device = get_safe_torch_device(cfg.policy.device, log=True)
# Only create wandb logger on main process # Only create wandb logger on main process when using accelerate
if cfg.wandb.enable and cfg.wandb.project: if accelerator is None or accelerator.is_main_process:
if accelerator is None or accelerator.is_main_process: if cfg.wandb.enable and cfg.wandb.project:
wandb_logger = WandBLogger(cfg) wandb_logger = WandBLogger(cfg)
else: else:
wandb_logger = None wandb_logger = None
logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
else: else:
wandb_logger = None wandb_logger = None
if accelerator is None or accelerator.is_main_process:
logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
if cfg.seed is not None: if cfg.seed is not None:
set_seed(cfg.seed) set_seed(cfg.seed)
@@ -211,8 +197,7 @@ def train(cfg: TrainPipelineConfig):
torch.backends.cudnn.benchmark = True torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cuda.matmul.allow_tf32 = True
if accelerator is None or accelerator.is_main_process: logging.info("Creating dataset")
logging.info("Creating dataset")
dataset = make_dataset(cfg) dataset = make_dataset(cfg)
# Create environment used for evaluating checkpoints during training on simulation data. # Create environment used for evaluating checkpoints during training on simulation data.
@@ -220,12 +205,10 @@ def train(cfg: TrainPipelineConfig):
# using the eval.py instead, with gym_dora environment and dora-rs. # using the eval.py instead, with gym_dora environment and dora-rs.
eval_env = None eval_env = None
if cfg.eval_freq > 0 and cfg.env is not None: if cfg.eval_freq > 0 and cfg.env is not None:
if accelerator is None or accelerator.is_main_process: logging.info("Creating env")
logging.info("Creating env")
eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs) eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
if accelerator is None or accelerator.is_main_process: logging.info("Creating policy")
logging.info("Creating policy")
policy = make_policy( policy = make_policy(
cfg=cfg.policy, cfg=cfg.policy,
ds_meta=dataset.meta, ds_meta=dataset.meta,
@@ -246,44 +229,43 @@ def train(cfg: TrainPipelineConfig):
if accelerator is None or accelerator.is_main_process: if accelerator is None or accelerator.is_main_process:
logging.info("Creating optimizer and scheduler") logging.info("Creating optimizer and scheduler")
optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp) # Scale scheduler parameters for multi-GPU training
if accelerator is not None and accelerator.num_processes > 1:
# With more GPUs, we process data faster, so scheduler should adapt faster
original_warmup_steps = cfg.scheduler.num_warmup_steps if cfg.scheduler else 0
original_decay_steps = cfg.scheduler.num_decay_steps if cfg.scheduler else 0
if cfg.scheduler is not None:
cfg.scheduler.num_warmup_steps = max(
1, cfg.scheduler.num_warmup_steps // accelerator.num_processes
)
cfg.scheduler.num_decay_steps = max(1, cfg.scheduler.num_decay_steps // accelerator.num_processes)
if accelerator.is_main_process:
logging.info(f"Scaled scheduler for {accelerator.num_processes} GPUs:")
logging.info(f" Warmup steps: {original_warmup_steps} → {cfg.scheduler.num_warmup_steps}")
logging.info(f" Decay steps: {original_decay_steps} → {cfg.scheduler.num_decay_steps}")
optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp and accelerator is None)
step = 0 # number of policy updates (forward + backward + optim) step = 0 # number of policy updates (forward + backward + optim)
if cfg.resume: if cfg.resume:
if accelerator is not None:
# Load accelerate-specific state if available
accelerate_state_path = cfg.checkpoint_path / "accelerate_state"
if accelerate_state_path.exists():
accelerator.load_state(str(accelerate_state_path))
if accelerator.is_main_process:
logging.info("Loaded Accelerate state from checkpoint")
step, optimizer, lr_scheduler = load_training_state(cfg.checkpoint_path, optimizer, lr_scheduler) step, optimizer, lr_scheduler = load_training_state(cfg.checkpoint_path, optimizer, lr_scheduler)
num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad) num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad)
num_total_params = sum(p.numel() for p in policy.parameters()) num_total_params = sum(p.numel() for p in policy.parameters())
# Only log setup info on main process logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
if accelerator is None or accelerator.is_main_process: if cfg.env is not None:
logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}") logging.info(f"{cfg.env.task=}")
if cfg.env is not None: logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})")
logging.info(f"{cfg.env.task=}") logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})")
logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})") logging.info(f"{dataset.num_episodes=}")
logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})") logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
logging.info(f"{dataset.num_episodes=}") logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
# Log batch size and learning rate info
if accelerator is not None:
logging.info(f"Per-GPU batch size: {cfg.batch_size}")
logging.info(f"Effective batch size (total): {cfg.batch_size * accelerator.num_processes}")
else:
logging.info(f"Batch size: {cfg.batch_size}")
logging.info(f"Learning rate: {optimizer.param_groups[0]['lr']:.2e}")
# create dataloader for offline training # create dataloader for offline training
if hasattr(cfg.policy, "drop_n_last_frames"): if hasattr(cfg.policy, "drop_n_last_frames"):
@@ -318,7 +300,6 @@ def train(cfg: TrainPipelineConfig):
logging.info("Policy, optimizer, dataloader, and scheduler prepared with Accelerate") logging.info("Policy, optimizer, dataloader, and scheduler prepared with Accelerate")
dl_iter = cycle(dataloader) dl_iter = cycle(dataloader)
policy.train() policy.train()
train_metrics = { train_metrics = {
@@ -336,45 +317,24 @@ def train(cfg: TrainPipelineConfig):
effective_batch_size, dataset.num_frames, dataset.num_episodes, train_metrics, initial_step=step effective_batch_size, dataset.num_frames, dataset.num_episodes, train_metrics, initial_step=step
) )
if accelerator is None or accelerator.is_main_process: logging.info("Start offline training on a fixed dataset")
logging.info("Start offline training on a fixed dataset")
for _ in range(step, cfg.steps): for _ in range(step, cfg.steps):
# Handle gradient accumulation start_time = time.perf_counter()
if accelerator is not None: batch = next(dl_iter)
with accelerator.accumulate(policy): batch = preprocessor(batch)
start_time = time.perf_counter() train_tracker.dataloading_s = time.perf_counter() - start_time
batch = next(dl_iter)
batch = preprocessor(batch)
train_tracker.dataloading_s = time.perf_counter() - start_time
train_tracker, output_dict = update_policy( train_tracker, output_dict = update_policy(
train_tracker, train_tracker,
policy, policy,
batch, batch,
optimizer, optimizer,
cfg.optimizer.grad_clip_norm, cfg.optimizer.grad_clip_norm,
grad_scaler=grad_scaler, grad_scaler=grad_scaler,
lr_scheduler=lr_scheduler, lr_scheduler=lr_scheduler,
use_amp=cfg.policy.use_amp, use_amp=cfg.policy.use_amp,
accelerator=accelerator, accelerator=accelerator,
) )
else:
start_time = time.perf_counter()
batch = next(dl_iter)
batch = preprocessor(batch)
train_tracker.dataloading_s = time.perf_counter() - start_time
train_tracker, output_dict = update_policy(
train_tracker,
policy,
batch,
optimizer,
cfg.optimizer.grad_clip_norm,
grad_scaler=grad_scaler,
lr_scheduler=lr_scheduler,
use_amp=cfg.policy.use_amp,
accelerator=accelerator,
)
# Note: eval and checkpoint happens *after* the `step`th training update has completed, so we # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
# increment `step` here. # increment `step` here.
@@ -385,121 +345,74 @@ def train(cfg: TrainPipelineConfig):
is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0 is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
if is_log_step: if is_log_step:
# Only log training metrics on main process logging.info(train_tracker)
if accelerator is None or accelerator.is_main_process: if wandb_logger:
logging.info(train_tracker) wandb_log_dict = train_tracker.to_dict()
if wandb_logger: if output_dict:
wandb_log_dict = train_tracker.to_dict() wandb_log_dict.update(output_dict)
if output_dict: wandb_logger.log_dict(wandb_log_dict, step)
wandb_log_dict.update(output_dict)
wandb_logger.log_dict(wandb_log_dict, step)
train_tracker.reset_averages() train_tracker.reset_averages()
if cfg.save_checkpoint and is_saving_step: if cfg.save_checkpoint and is_saving_step:
if accelerator is None or accelerator.is_main_process: logging.info(f"Checkpoint policy after step {step}")
logging.info(f"Checkpoint policy after step {step}")
checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step) checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)
save_checkpoint(
if accelerator is not None: checkpoint_dir, step, cfg, policy, optimizer, lr_scheduler, preprocessor, postprocessor
# Use accelerate's checkpointing - only saves on main process )
accelerator.wait_for_everyone() # Synchronize all processes update_last_checkpoint(checkpoint_dir)
if accelerator.is_main_process: if wandb_logger:
# Use unwrapped model for saving wandb_logger.log_policy(checkpoint_dir)
unwrapped_policy = accelerator.unwrap_model(policy)
save_checkpoint(
checkpoint_dir,
step,
cfg,
unwrapped_policy,
optimizer,
lr_scheduler,
preprocessor,
postprocessor,
)
update_last_checkpoint(checkpoint_dir)
if wandb_logger:
wandb_logger.log_policy(checkpoint_dir)
# Save accelerate-specific state
accelerator.save_state(checkpoint_dir / "accelerate_state")
else:
# Original behavior for non-accelerate
save_checkpoint(
checkpoint_dir, step, cfg, policy, optimizer, lr_scheduler, preprocessor, postprocessor
)
update_last_checkpoint(checkpoint_dir)
if wandb_logger:
wandb_logger.log_policy(checkpoint_dir)
if cfg.env and is_eval_step: if cfg.env and is_eval_step:
# Only evaluate on main process when using accelerate step_id = get_step_identifier(step, cfg.steps)
if accelerator is None or accelerator.is_main_process: logging.info(f"Eval policy at step {step}")
step_id = get_step_identifier(step, cfg.steps) with (
logging.info(f"Eval policy at step {step}") torch.no_grad(),
torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
# Use unwrapped model for evaluation if using accelerate ):
eval_policy = accelerator.unwrap_model(policy) if accelerator is not None else policy eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env
with ( policy=policy,
torch.no_grad(), preprocessor=preprocessor,
torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(), postprocessor=postprocessor,
): n_episodes=cfg.eval.n_episodes,
eval_info = eval_policy_all( videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
envs=eval_env, # dict[suite][task_id] -> vec_env max_episodes_rendered=4,
policy=eval_policy, start_seed=cfg.seed,
preprocessor=preprocessor, max_parallel_tasks=cfg.env.max_parallel_tasks,
postprocessor=postprocessor,
n_episodes=cfg.eval.n_episodes,
videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
max_episodes_rendered=4,
start_seed=cfg.seed,
max_parallel_tasks=cfg.env.max_parallel_tasks,
)
# overall metrics (suite-agnostic)
aggregated = eval_info["overall"]
# optional: per-suite logging
for suite, suite_info in eval_info.items():
logging.info("Suite %s aggregated: %s", suite, suite_info)
# meters/tracker
eval_metrics = {
"avg_sum_reward": AverageMeter("∑rwrd", ":.3f"),
"pc_success": AverageMeter("success", ":.1f"),
"eval_s": AverageMeter("eval_s", ":.3f"),
}
eval_tracker = MetricsTracker(
cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step
) )
eval_tracker.eval_s = aggregated.pop("eval_s") # overall metrics (suite-agnostic)
eval_tracker.avg_sum_reward = aggregated.pop("avg_sum_reward") aggregated = eval_info["overall"]
eval_tracker.pc_success = aggregated.pop("pc_success")
if wandb_logger:
wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
wandb_logger.log_video(eval_info["overall"]["video_paths"][0], step, mode="eval")
# Synchronize all processes after evaluation # optional: per-suite logging
if accelerator is not None: for suite, suite_info in eval_info.items():
accelerator.wait_for_everyone() logging.info("Suite %s aggregated: %s", suite, suite_info)
# meters/tracker
eval_metrics = {
"avg_sum_reward": AverageMeter("∑rwrd", ":.3f"),
"pc_success": AverageMeter("success", ":.1f"),
"eval_s": AverageMeter("eval_s", ":.3f"),
}
eval_tracker = MetricsTracker(
cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step
)
eval_tracker.eval_s = aggregated.pop("eval_s")
eval_tracker.avg_sum_reward = aggregated.pop("avg_sum_reward")
eval_tracker.pc_success = aggregated.pop("pc_success")
if wandb_logger:
wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
wandb_logger.log_video(eval_info["overall"]["video_paths"][0], step, mode="eval")
if eval_env: if eval_env:
close_envs(eval_env) close_envs(eval_env)
logging.info("End of training")
if accelerator is None or accelerator.is_main_process:
logging.info("End of training")
# Synchronize all processes before finishing
if accelerator is not None:
accelerator.wait_for_everyone()
if cfg.policy.push_to_hub: if cfg.policy.push_to_hub:
# Only push to hub from main process when using accelerate policy.push_model_to_hub(cfg)
if accelerator is None or accelerator.is_main_process: preprocessor.push_to_hub(cfg.policy.repo_id)
# Use unwrapped model for hub pushing if using accelerate postprocessor.push_to_hub(cfg.policy.repo_id)
hub_policy = accelerator.unwrap_model(policy) if accelerator is not None else policy
hub_policy.push_model_to_hub(cfg)
preprocessor.push_to_hub(cfg.policy.repo_id)
postprocessor.push_to_hub(cfg.policy.repo_id)
def main(): def main():