From 1fe58f2d3afe0e7c46e86fee03de2e4122fbe9a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ahmet=20Faruk=20G=C3=9CM=C3=9C=C5=9ETA=C5=9E?= Date: Thu, 30 Jul 2026 18:56:56 +0300 Subject: [PATCH] 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 --- src/lerobot/common/train_utils.py | 11 +++++++++++ src/lerobot/configs/train.py | 1 + src/lerobot/scripts/lerobot_train.py | 3 ++- tests/utils/test_train_utils.py | 13 +++++++++++++ 4 files changed, 27 insertions(+), 1 deletion(-) diff --git a/src/lerobot/common/train_utils.py b/src/lerobot/common/train_utils.py index b26196f14..a300cc76c 100644 --- a/src/lerobot/common/train_utils.py +++ b/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: diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index 8318b56ac..1240b749c 100644 --- a/src/lerobot/configs/train.py +++ b/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 diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 4c0e12eaf..178a68cfa 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/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 diff --git a/tests/utils/test_train_utils.py b/tests/utils/test_train_utils.py index ccd769bd0..3098c080d 100644 --- a/tests/utils/test_train_utils.py +++ b/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)