diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index e92247188..8318b56ac 100644 --- a/src/lerobot/configs/train.py +++ b/src/lerobot/configs/train.py @@ -194,7 +194,11 @@ class TrainPipelineConfig(HubMixin): ) if Path(config_path).resolve().exists(): - policy_dir = Path(config_path).parent + # `config_path` may point at the checkpoint's train_config.json or at its + # pretrained_model/ directory (both documented above) — resolve either to + # the pretrained_model/ directory. + config_path_obj = Path(config_path) + policy_dir = config_path_obj.parent if config_path_obj.is_file() else config_path_obj self.checkpoint_path = policy_dir.parent elif self.job.is_remote: return diff --git a/tests/configs/test_resume_from_hub.py b/tests/configs/test_resume_from_hub.py index 2cc9fd7ae..d5affc2bc 100644 --- a/tests/configs/test_resume_from_hub.py +++ b/tests/configs/test_resume_from_hub.py @@ -66,3 +66,27 @@ def test_from_pretrained_raises_when_no_root_config_and_no_checkpoints(monkeypat with pytest.raises(FileNotFoundError, match="train_config.json not found"): TrainPipelineConfig.from_pretrained("user/empty-repo") + + +@pytest.mark.parametrize("pass_dir", [False, True]) +def test_resolve_resume_checkpoint_accepts_file_or_pretrained_model_dir(tmp_path, monkeypatch, pass_dir): + """`--config_path` may point at the checkpoint's train_config.json or at its + pretrained_model/ directory; both must resolve `policy.pretrained_path` to the + pretrained_model/ directory (regression test for the directory case, which + previously resolved one level too high and failed on model.safetensors).""" + pretrained_dir = tmp_path / "checkpoints" / "000002" / "pretrained_model" + pretrained_dir.mkdir(parents=True) + (pretrained_dir / "train_config.json").touch() + target = pretrained_dir if pass_dir else pretrained_dir / "train_config.json" + + from lerobot.policies.act.configuration_act import ACTConfig + + cfg = tc.draccus.parse(TrainPipelineConfig, args=["--dataset.repo_id", "u/d"]) + cfg.policy = ACTConfig() + cfg.resume = True + monkeypatch.setattr(tc.parser, "parse_arg", lambda name: str(target) if name == "config_path" else None) + + cfg._resolve_resume_checkpoint() + + assert cfg.policy.pretrained_path == pretrained_dir + assert cfg.checkpoint_path == pretrained_dir.parent