fix

src/lerobot/configs/train.py

@@ -63,10 +63,6 @@ class TrainPipelineConfig(HubMixin):
     scheduler: LRSchedulerConfig | None = None
     eval: EvalConfig = field(default_factory=EvalConfig)
     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):
         self.checkpoint_path = None

src/lerobot/optim/schedulers.py

@@ -92,20 +92,20 @@ class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):
         def lr_lambda(current_step):
             def linear_warmup_schedule(current_step):
                 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:
                     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)
 
             def cosine_decay_schedule(current_step):
-                # Steps since warmup ended (this was the bug!)
                 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))
                 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:
                 return linear_warmup_schedule(current_step)

src/lerobot/scripts/train.py

@@ -91,49 +91,42 @@ def update_policy(
         - A dictionary of outputs from the policy's forward pass, for logging purposes.
     """
     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()
 
-    # Handle mixed precision differently for accelerate vs non-accelerate
     if accelerator is not None:
-        # Accelerate handles mixed precision internally
-        with accelerator.autocast() if use_amp else nullcontext():
+        # Use accelerate's autocast and backward
+        with accelerator.autocast():
             loss, output_dict = policy.forward(batch)
-        # Use accelerator's backward method
         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:
-        # Accelerate handles gradient scaling internally
-        grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
-        if grad_norm is None:
-            grad_norm = 0.0
+        # Use accelerate's gradient clipping
+        if grad_clip_norm > 0:
+            grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
+            if grad_norm is None:
+                grad_norm = 0.0
+        else:
+            grad_norm = torch.tensor(0.0, device=device)
+
         with lock if lock is not None else nullcontext():
             optimizer.step()
         optimizer.zero_grad()
     else:
-        # Original gradient handling for non-accelerate
-        # Unscale the gradient of the optimizer's assigned params in-place **prior to gradient clipping**.
-        grad_scaler.unscale_(optimizer)
+        # Original single-GPU path
+        with torch.autocast(device_type=device.type) if use_amp else nullcontext():
+            loss, output_dict = policy.forward(batch)
+        grad_scaler.scale(loss).backward()
 
+        grad_scaler.unscale_(optimizer)
         grad_norm = torch.nn.utils.clip_grad_norm_(
             policy.parameters(),
             grad_clip_norm,
             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():
             grad_scaler.step(optimizer)
-        # Updates the scale for next iteration.
         grad_scaler.update()
-
         optimizer.zero_grad()
 
     # 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.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()))
 
-    # Initialize Accelerate if requested
+    # Initialize Accelerate if enabled
     accelerator = None
     if cfg.use_accelerate:
-        # Configure DDP to handle unused parameters
         ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
         accelerator = Accelerator(
             gradient_accumulation_steps=cfg.gradient_accumulation_steps,
@@ -191,19 +179,17 @@ def train(cfg: TrainPipelineConfig):
             )
             logging.info(f"Training on {accelerator.num_processes} processes")
     else:
-        # Check device is available (original behavior)
         device = get_safe_torch_device(cfg.policy.device, log=True)
 
-    # Only create wandb logger on main process
-    if cfg.wandb.enable and cfg.wandb.project:
-        if accelerator is None or accelerator.is_main_process:
+    # Only create wandb logger on main process when using accelerate
+    if accelerator is None or accelerator.is_main_process:
+        if cfg.wandb.enable and cfg.wandb.project:
             wandb_logger = WandBLogger(cfg)
         else:
             wandb_logger = None
+            logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
     else:
         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:
         set_seed(cfg.seed)
@@ -211,8 +197,7 @@ def train(cfg: TrainPipelineConfig):
     torch.backends.cudnn.benchmark = 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)
 
     # 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.
     eval_env = 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)
 
-    if accelerator is None or accelerator.is_main_process:
-        logging.info("Creating policy")
+    logging.info("Creating policy")
     policy = make_policy(
         cfg=cfg.policy,
         ds_meta=dataset.meta,
@@ -246,44 +229,43 @@ def train(cfg: TrainPipelineConfig):
 
     if accelerator is None or accelerator.is_main_process:
         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)
 
     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)
 
     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())
 
-    # Only log setup info on main process
-    if accelerator is None or accelerator.is_main_process:
-        logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
-        if cfg.env is not None:
-            logging.info(f"{cfg.env.task=}")
-        logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})")
-        logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})")
-        logging.info(f"{dataset.num_episodes=}")
-        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}")
+    logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
+    if cfg.env is not None:
+        logging.info(f"{cfg.env.task=}")
+    logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})")
+    logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})")
+    logging.info(f"{dataset.num_episodes=}")
+    logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
+    logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
 
     # create dataloader for offline training
     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")
 
     dl_iter = cycle(dataloader)
-
     policy.train()
 
     train_metrics = {
@@ -336,45 +317,24 @@ def train(cfg: TrainPipelineConfig):
         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):
-        # Handle gradient accumulation
-        if accelerator is not None:
-            with accelerator.accumulate(policy):
-                start_time = time.perf_counter()
-                batch = next(dl_iter)
-                batch = preprocessor(batch)
-                train_tracker.dataloading_s = time.perf_counter() - start_time
+        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,
-                )
-        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,
-            )
+        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
         # increment `step` here.
@@ -385,121 +345,74 @@ def train(cfg: TrainPipelineConfig):
         is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
 
         if is_log_step:
-            # Only log training metrics on main process
-            if accelerator is None or accelerator.is_main_process:
-                logging.info(train_tracker)
-                if wandb_logger:
-                    wandb_log_dict = train_tracker.to_dict()
-                    if output_dict:
-                        wandb_log_dict.update(output_dict)
-                    wandb_logger.log_dict(wandb_log_dict, step)
+            logging.info(train_tracker)
+            if wandb_logger:
+                wandb_log_dict = train_tracker.to_dict()
+                if output_dict:
+                    wandb_log_dict.update(output_dict)
+                wandb_logger.log_dict(wandb_log_dict, step)
             train_tracker.reset_averages()
 
         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)
-
-            if accelerator is not None:
-                # Use accelerate's checkpointing - only saves on main process
-                accelerator.wait_for_everyone()  # Synchronize all processes
-                if accelerator.is_main_process:
-                    # Use unwrapped model for saving
-                    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)
+            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:
-            # Only evaluate on main process when using accelerate
-            if accelerator is None or accelerator.is_main_process:
-                step_id = get_step_identifier(step, cfg.steps)
-                logging.info(f"Eval policy at step {step}")
-
-                # Use unwrapped model for evaluation if using accelerate
-                eval_policy = accelerator.unwrap_model(policy) if accelerator is not None else policy
-
-                with (
-                    torch.no_grad(),
-                    torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
-                ):
-                    eval_info = eval_policy_all(
-                        envs=eval_env,  # dict[suite][task_id] -> vec_env
-                        policy=eval_policy,
-                        preprocessor=preprocessor,
-                        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
+            step_id = get_step_identifier(step, cfg.steps)
+            logging.info(f"Eval policy at step {step}")
+            with (
+                torch.no_grad(),
+                torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
+            ):
+                eval_info = eval_policy_all(
+                    envs=eval_env,  # dict[suite][task_id] -> vec_env
+                    policy=policy,
+                    preprocessor=preprocessor,
+                    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,
                 )
-                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")
+            # overall metrics (suite-agnostic)
+            aggregated = eval_info["overall"]
 
-            # Synchronize all processes after evaluation
-            if accelerator is not None:
-                accelerator.wait_for_everyone()
+            # 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")
+            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:
         close_envs(eval_env)
-
-    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()
+    logging.info("End of training")
 
     if cfg.policy.push_to_hub:
-        # Only push to hub from main process when using accelerate
-        if accelerator is None or accelerator.is_main_process:
-            # Use unwrapped model for hub pushing if using accelerate
-            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)
+        policy.push_model_to_hub(cfg)
+        preprocessor.push_to_hub(cfg.policy.repo_id)
+        postprocessor.push_to_hub(cfg.policy.repo_id)
 
 
 def main():