fix(train): guard checkpoint saving against save_freq=0 (#4112)

is_saving_step was the only step-frequency check without a `> 0` guard,
so save_freq=0 raised ZeroDivisionError from `step % 0` on the first
step. Route the decision through a should_save_checkpoint helper that
treats a non-positive save_freq as "save only the final checkpoint",
matching how log_freq/eval_freq handle non-positive values.

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

src/lerobot/common/train_utils.py

@@ -52,6 +52,17 @@ def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Pa
     return output_dir / CHECKPOINTS_DIR / step_identifier
 
 
+def should_save_checkpoint(step: int, save_freq: int, total_steps: int) -> bool:
+    """Whether a checkpoint should be saved at ``step``.
+
+    A checkpoint is saved every ``save_freq`` steps and always after the final step. A
+    non-positive ``save_freq`` disables periodic saving (only the final checkpoint is
+    written), mirroring how ``log_freq``/``eval_freq`` treat non-positive values and
+    avoiding a ``ZeroDivisionError`` from ``step % 0``.
+    """
+    return (save_freq > 0 and step % save_freq == 0) or step == total_steps
+
+
 def save_training_step(
     step: int, save_dir: Path, num_processes: int | None = None, batch_size: int | None = None
 ) -> None:

src/lerobot/configs/train.py

@@ -119,6 +119,7 @@ class TrainPipelineConfig(HubMixin):
     tolerance_s: float = 1e-4
     save_checkpoint: bool = True
     # Checkpoint is saved every `save_freq` training iterations and after the last training step.
+    # A non-positive value disables periodic saving, keeping only the final checkpoint.
     save_freq: int = 20_000
     use_policy_training_preset: bool = True
     optimizer: OptimizerConfig | None = None

src/lerobot/scripts/lerobot_train.py

@@ -45,6 +45,7 @@ from lerobot.common.train_utils import (
     load_training_state,
     push_checkpoint_to_hub,
     save_checkpoint,
+    should_save_checkpoint,
     update_last_checkpoint,
 )
 from lerobot.common.wandb_utils import WandBLogger
@@ -616,7 +617,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             progbar.update(1)
         train_tracker.step()
         is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0
-        is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
+        is_saving_step = should_save_checkpoint(step, cfg.save_freq, cfg.steps)
         is_env_eval_step = cfg.env_eval_freq > 0 and step % cfg.env_eval_freq == 0
         is_eval_step = cfg.eval_steps > 0 and eval_dataloader is not None and step % cfg.eval_steps == 0
 

tests/utils/test_train_utils.py

@@ -30,6 +30,7 @@ from lerobot.common.train_utils import (
     save_checkpoint,
     save_training_state,
     save_training_step,
+    should_save_checkpoint,
     update_last_checkpoint,
 )
 from lerobot.utils.constants import (
@@ -50,6 +51,18 @@ def test_get_step_identifier():
     assert get_step_identifier(456789, 1_000_000) == "0456789"
 
 
+def test_should_save_checkpoint():
+    # Periodic checkpoints land on multiples of save_freq.
+    assert should_save_checkpoint(10, save_freq=10, total_steps=100) is True
+    assert should_save_checkpoint(5, save_freq=10, total_steps=100) is False
+    # The final step always saves, even when it is not a multiple of save_freq.
+    assert should_save_checkpoint(100, save_freq=30, total_steps=100) is True
+    # save_freq <= 0 disables periodic saving without raising ZeroDivisionError.
+    assert should_save_checkpoint(1, save_freq=0, total_steps=100) is False
+    assert should_save_checkpoint(100, save_freq=0, total_steps=100) is True
+    assert should_save_checkpoint(1, save_freq=-1, total_steps=100) is False
+
+
 def test_get_step_checkpoint_dir():
     output_dir = Path("/checkpoints")
     step_dir = get_step_checkpoint_dir(output_dir, 1000, 5)