fix(config): accept pretrained_model dir for --config_path on resume (#4023)

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
This commit is contained in:
Jash Shah
2026-07-30 04:48:27 -07:00
committed by GitHub
parent 0cef9cd197
commit 40a5e70352
2 changed files with 29 additions and 1 deletions
+24
View File
@@ -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