fix(finetune): use LeRobot policy path selector (#310)

Author: rylinjames

scripts/modal_act_from_scratch_smoke.py

@@ -12,7 +12,7 @@
 confirm:
   1. The tether install pulls all the new code (config + run + cli).
   2. _build_lerobot_command produces a valid lerobot-train invocation
-     with --policy.type=act, no --policy.pretrained_path, --policy.chunk_size=N.
+     with --policy.type=act, no pretrained policy path, --policy.chunk_size=N.
   3. lerobot-train accepts the args and starts training.
   4. A checkpoint actually lands at the expected path.
 
@@ -166,7 +166,12 @@ def act_smoke(
     # Sanity check the construction matches Phase 4 expectations:
     cmd_str = " ".join(cmd)
     expected_in_cmd = ["--policy.type=act", "--policy.chunk_size=31"]
-    forbidden_in_cmd = ["--policy.pretrained_path", "--peft.method_type"]
+    forbidden_in_cmd = [
+        "--policy.path",
+        "--policy.pretrained_path",
+        "--policy.pretrained_model_path",
+        "--peft.method_type",
+    ]
     construction_ok = (
         all(s in cmd_str for s in expected_in_cmd)
         and not any(s in cmd_str for s in forbidden_in_cmd)

src/tether/finetune/cli.py

@@ -28,8 +28,9 @@ def finetune_command(
     policy: str = typer.Option(
         "auto",
         "--policy",
-        help="Policy class. 'auto' (default) infers from --base. Set explicitly "
-             "(e.g. 'act') for from-scratch training. Per ADR 2026-05-06.",
+        help="Policy class. 'auto' (default) loads the config from --base. "
+             "Set explicitly (e.g. 'act') for from-scratch training. "
+             "Per ADR 2026-05-06.",
     ),
     chunk_size: int = typer.Option(
         50,

src/tether/finetune/config.py

@@ -43,10 +43,10 @@ class FinetuneConfig:
 
     policy: str = "auto"
     """Policy class to train.
-      - 'auto' (default): infer from `base` (e.g. lerobot/smolvla_base → smolvla).
+      - 'auto' (default): load policy config + weights from `base`.
       - 'act' / 'diffusion' / 'pi0' / etc.: explicit lerobot policy.type.
-    Used by from-scratch training paths where `base` is empty or a non-pretrained
-    sentinel; lerobot-train requires --policy.type either way.
+    Explicit types are used by from-scratch training paths where `base` is empty.
+    Pretrained paths use LeRobot's --policy.path and must not also pass policy.type.
 
     ACT-from-scratch (vendored from auto_soarm 2026-05-06): set policy='act',
     mode='full', leave base='' (or set to ''). Recipe defaults:

src/tether/finetune/run.py

@@ -81,15 +81,11 @@ def _validate_config(cfg: FinetuneConfig) -> list[str]:
 
 
 def _infer_policy_type(base: str) -> str:
-    """Derive lerobot's policy registry name from the base model id.
+    """Derive LeRobot's policy registry name from a model id.
 
-    lerobot 0.5.1 registers policies by short name (smolvla, pi0, pi05,
-    act, diffusion, vqbet, ...) and requires `--policy.type=<name>` at
-    CLI time. The full HF model id (e.g. 'lerobot/smolvla_base') goes
-    to `--policy.pretrained_model_path=...` separately.
-
-    Falls back to raising a clear error for unrecognized bases rather
-    than guessing. Customers can override via extra_lerobot_args={"policy.type": "..."}.
+    This metadata helper is retained for callers that need a short policy
+    name. It is deliberately not used to select a pretrained policy on the
+    LeRobot 0.5.1 CLI: ``--policy.path`` loads the type from the checkpoint.
     """
     base_lower = base.lower()
     if "smolvla" in base_lower:
@@ -102,9 +98,8 @@ def _infer_policy_type(base: str) -> str:
         # lerobot 0.5.1 can't load N1.6 per prior Step-3 finding; v0.6 work.
         return "gr00t_n1_5"
     raise ValueError(
-        f"Could not infer --policy.type from base={base!r}. "
-        f"Supported in v0.3: lerobot/smolvla_base. For other bases, "
-        f"pass policy.type explicitly via extra_lerobot_args."
+        f"Could not infer policy type from base={base!r}. "
+        "Supported identifiers include smolvla, pi0, pi05, and gr00t."
     )
 
 
@@ -121,25 +116,42 @@ def _build_lerobot_command(cfg: FinetuneConfig) -> list[str]:
     customers can reproduce manually.
 
     Key arg names per lerobot 0.5.1:
-      --policy.pretrained_model_path   (NOT --policy.path)
+      --policy.path                    (pretrained checkpoint)
+      --policy.type                    (from-scratch policy only)
       --optimizer.lr                   (NOT --policy.optimizer_lr)
       --peft.method_type=lora          (enables PEFT)
       --peft.r                         (LoRA rank)
 
+    LeRobot's parser rejects using ``--policy.path`` and ``--policy.type``
+    together. A pretrained command therefore selects the policy only by
+    checkpoint path; LeRobot loads the policy type from that checkpoint's
+    config. From-scratch commands select the policy only by type.
+
     cfg.precision is intentionally NOT passed through — lerobot 0.5.1
     doesn't expose a top-level precision flag; it's baked into the
     policy config. v0.5 will add per-policy precision overrides.
     """
-    # draccus requires `policy.type` to select which PreTrainedConfig
-    # subclass to decode into. Use explicit cfg.policy when set, else
-    # infer from the base-model id.
     explicit_policy = getattr(cfg, "policy", "auto")
-    if explicit_policy != "auto":
-        policy_type = explicit_policy
-    else:
-        policy_type = _infer_policy_type(cfg.base)
     is_from_scratch = (explicit_policy != "auto" and cfg.mode == "full")
 
+    # These select the policy and must be emitted exactly once by this
+    # builder. In particular, LeRobot 0.5.1 rejects policy.path + policy.type,
+    # while the older pretrained_* spellings are not valid selection flags.
+    reserved_policy_args = {
+        "policy.path",
+        "policy.type",
+        "policy.pretrained_path",
+        "policy.pretrained_model_path",
+    }
+    conflicting_args = reserved_policy_args.intersection(cfg.extra_lerobot_args)
+    if conflicting_args:
+        names = ", ".join(sorted(conflicting_args))
+        raise ValueError(
+            f"extra_lerobot_args cannot override policy selection ({names}); "
+            "use FinetuneConfig.base for pretrained policies or "
+            "FinetuneConfig.policy for from-scratch policies"
+        )
+
     # lerobot-train wants to OWN its output_dir (errors if pre-existing
     # and resume=False). We keep cfg.output as the tether orchestration
     # root, and give lerobot a subdirectory it creates fresh on each run.
@@ -153,7 +165,6 @@ def _build_lerobot_command(cfg: FinetuneConfig) -> list[str]:
 
     cmd = [
         "lerobot-train",
-        f"--policy.type={policy_type}",
         f"--policy.repo_id={repo_id}",
         f"--policy.push_to_hub=false",
         f"--dataset.repo_id={cfg.dataset}",
@@ -163,9 +174,11 @@ def _build_lerobot_command(cfg: FinetuneConfig) -> list[str]:
         f"--optimizer.lr={cfg.learning_rate}",
         f"--seed={cfg.seed}",
     ]
-    if not is_from_scratch:
-        # Pretrained-base path: pass the HF id / local checkpoint to load weights from.
-        cmd.append(f"--policy.pretrained_path={cfg.base}")
+    if is_from_scratch:
+        cmd.append(f"--policy.type={explicit_policy}")
+    else:
+        # LeRobot 0.5.1 loads both config and weights through policy.path.
+        cmd.append(f"--policy.path={cfg.base}")
     if is_from_scratch and cfg.chunk_size:
         # ACT (and similar chunked policies) need chunk_size; pretrained bases
         # bake this in. auto_soarm convention (per its train.py): set

tests/test_finetune.py

@@ -124,8 +124,10 @@ def test_basic_command_shape(self, tmp_path):
         # Schema is pinned to lerobot 0.5.1. If lerobot renames flags
         # upstream, this test catches it.
         joined = " ".join(cmd)
-        assert "--policy.type=smolvla" in joined
-        assert "--policy.pretrained_path=lerobot/smolvla_base" in joined
+        assert cmd.count("--policy.path=lerobot/smolvla_base") == 1
+        assert "--policy.type=" not in joined
+        assert "--policy.pretrained_path=" not in joined
+        assert "--policy.pretrained_model_path=" not in joined
         assert "--policy.repo_id=" in joined
         assert "--policy.push_to_hub=false" in joined
         assert "--dataset.repo_id=lerobot/libero" in joined
@@ -139,16 +141,58 @@ def test_basic_command_shape(self, tmp_path):
 
     def test_policy_type_inference(self):
         from tether.finetune.run import _infer_policy_type
+
         assert _infer_policy_type("lerobot/smolvla_base") == "smolvla"
         assert _infer_policy_type("lerobot/pi0_base") == "pi0"
         assert _infer_policy_type("lerobot/pi05_base") == "pi05"
         assert _infer_policy_type("nvidia/GR00T-N1.6-3B") == "gr00t_n1_5"
 
     def test_policy_type_unknown_rejected(self):
         from tether.finetune.run import _infer_policy_type
+
         with pytest.raises(ValueError, match="Could not infer"):
             _infer_policy_type("some-random/unknown-model")
 
+    def test_from_scratch_selects_type_without_pretrained_path(self, tmp_path):
+        cfg = FinetuneConfig(
+            base="",
+            dataset="lerobot/pusht",
+            output=tmp_path,
+            mode="full",
+            policy="act",
+            chunk_size=31,
+        )
+
+        cmd = _build_lerobot_command(cfg)
+        joined = " ".join(cmd)
+
+        assert cmd.count("--policy.type=act") == 1
+        assert "--policy.path=" not in joined
+        assert "--policy.pretrained_path=" not in joined
+        assert "--policy.pretrained_model_path=" not in joined
+
+    @pytest.mark.parametrize(
+        "arg_name",
+        [
+            "policy.path",
+            "policy.type",
+            "policy.pretrained_path",
+            "policy.pretrained_model_path",
+        ],
+    )
+    def test_policy_selection_cannot_be_duplicated_by_extra_args(
+        self, tmp_path, arg_name
+    ):
+        cfg = FinetuneConfig(
+            base="lerobot/smolvla_base",
+            dataset="lerobot/libero",
+            output=tmp_path,
+            extra_lerobot_args={arg_name: "other/value"},
+        )
+
+        with pytest.raises(ValueError, match="cannot override policy selection"):
+            _build_lerobot_command(cfg)
+
     def test_extra_args_pass_through(self, tmp_path):
         cfg = FinetuneConfig(
             base="lerobot/smolvla_base",