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)