Fix policy.path not working in YAML config files (#3145)

* fix(config): support policy.path in YAML config files

policy.path was only handled via CLI args (filtered from sys.argv before
draccus, then retrieved in validate()). When specified in YAML, draccus
would crash because 'path' is not a valid field on PreTrainedConfig.

Extract path fields from the YAML/JSON config before draccus processes
it, store them in a module-level dict, and fall back to it in
get_path_arg() when the CLI doesn't have the path.

Fixes #2957

* fix(parser): preserve YAML policy overrides when loading from pretrained

When policy.path is set in YAML, validate() was calling from_pretrained
with only CLI overrides, discarding any YAML policy fields (e.g. lr,
batch_size) that draccus had already parsed. Fix by capturing the
remaining YAML fields as CLI-style args in _config_yaml_overrides and
merging them into the overrides passed to from_pretrained in train.py,
eval.py, and lerobot_record.py (CLI args still take precedence).

Also fix the NamedTemporaryFile SIM115 ruff warning and add types-PyYAML
to the mypy pre-commit hook.

* fix(parser): serialize bool/None values correctly in YAML policy overrides

Bool values from YAML configs (e.g. push_to_hub: true) were passed as
Python "True"/"False" strings instead of lowercase "true"/"false" that
draccus expects. Also skip None values to avoid passing "None" strings.

* revert: remove types-PyYAML from .pre-commit-config.yaml

* chore: fix quality check caused by untyped YAML import

Co-authored-by: masato-ka <jp6uzv@gmail.com>
Signed-off-by: Khalil Meftah <khalil.meftah@huggingface.co>

---------

Signed-off-by: Khalil Meftah <khalil.meftah@huggingface.co>
Co-authored-by: Khalil Meftah <khalil.meftah@huggingface.co>
Co-authored-by: masato-ka <jp6uzv@gmail.com>

src/lerobot/configs/eval.py

@@ -46,8 +46,11 @@ class EvalPipelineConfig:
         # HACK: We parse again the cli args here to get the pretrained path if there was one.
         policy_path = parser.get_path_arg("policy")
         if policy_path:
-            cli_overrides = parser.get_cli_overrides("policy")
-            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
+            yaml_overrides = parser.get_yaml_overrides("policy")
+            cli_overrides = parser.get_cli_overrides("policy") or []
+            self.policy = PreTrainedConfig.from_pretrained(
+                policy_path, cli_overrides=yaml_overrides + cli_overrides
+            )
             self.policy.pretrained_path = Path(policy_path)
 
         else:

src/lerobot/configs/parser.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 import importlib
 import inspect
+import json
 import pkgutil
 import sys
+import tempfile
 from argparse import ArgumentError
 from collections.abc import Callable, Iterable, Sequence
 from functools import wraps
@@ -24,6 +26,7 @@ from types import ModuleType
 from typing import Any, TypeVar, cast
 
 import draccus
+import yaml  # type: ignore[import-untyped]
 
 from lerobot.utils.utils import has_method
 
@@ -32,6 +35,29 @@ F = TypeVar("F", bound=Callable[..., object])
 PATH_KEY = "path"
 PLUGIN_DISCOVERY_SUFFIX = "discover_packages_path"
 
+# Storage for path args extracted from YAML/JSON config files, so that
+# get_path_arg() can find them even when they weren't passed via CLI.
+_config_path_args: dict[str, str] = {}
+
+# Storage for non-path YAML overrides so validate() can pass them to from_pretrained.
+_config_yaml_overrides: dict[str, list[str]] = {}
+
+
+def _flatten_to_cli_args(d: dict, prefix: str = "") -> list[str]:
+    """Recursively flatten a nested dict to CLI-style args (e.g. {"lr": 1e-4} -> ["--lr=0.0001"])."""
+    args = []
+    for key, value in d.items():
+        if key in (PATH_KEY, draccus.CHOICE_TYPE_KEY):
+            continue
+        full_key = f"{prefix}.{key}" if prefix else key
+        if isinstance(value, bool):
+            value = str(value).lower()
+        if isinstance(value, dict):
+            args.extend(_flatten_to_cli_args(value, full_key))
+        elif value is not None and not isinstance(value, list):
+            args.append(f"--{full_key}={value}")
+    return args
+
 
 def get_cli_overrides(field_name: str, args: Sequence[str] | None = None) -> list[str] | None:
     """Parses arguments from cli at a given nested attribute level.
@@ -145,7 +171,14 @@ def load_plugin(plugin_path: str) -> None:
 
 
 def get_path_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
-    return parse_arg(f"{field_name}.{PATH_KEY}", args)
+    result = parse_arg(f"{field_name}.{PATH_KEY}", args)
+    if result is None:
+        result = _config_path_args.get(field_name)
+    return result
+
+
+def get_yaml_overrides(field_name: str) -> list[str]:
+    return _config_yaml_overrides.get(field_name, [])
 
 
 def get_type_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
@@ -192,6 +225,52 @@ def filter_path_args(fields_to_filter: str | list[str], args: Sequence[str] | No
     return filtered_args
 
 
+def extract_path_fields_from_config(config_path: str, path_fields: list[str]) -> str:
+    """Extract `path` fields from a YAML/JSON config before draccus processes it.
+
+    When a user specifies e.g. ``policy.path: lerobot/smolvla_base`` in a YAML config,
+    draccus will fail because ``path`` is not a valid field on policy config classes.
+    This function extracts those path values, stores them in ``_config_path_args`` for
+    later retrieval by ``get_path_arg()``, and returns a cleaned temp config file path.
+    """
+    config_file = Path(config_path)
+    suffix = config_file.suffix.lower()
+
+    if suffix in (".yaml", ".yml"):
+        with open(config_file) as f:
+            config_data = yaml.safe_load(f)
+    elif suffix == ".json":
+        with open(config_file) as f:
+            config_data = json.load(f)
+    else:
+        return config_path
+
+    if not isinstance(config_data, dict):
+        return config_path
+
+    modified = False
+    for field in path_fields:
+        if field in config_data and isinstance(config_data[field], dict) and PATH_KEY in config_data[field]:
+            _config_path_args[field] = str(config_data[field].pop(PATH_KEY))
+            remaining = config_data[field]
+            if remaining:
+                _config_yaml_overrides[field] = _flatten_to_cli_args(remaining)
+            else:
+                del config_data[field]
+            modified = True
+
+    if not modified:
+        return config_path
+
+    # Write cleaned config to a temp file
+    with tempfile.NamedTemporaryFile(mode="w", suffix=suffix, delete=False) as tmp:
+        if suffix in (".yaml", ".yml"):
+            yaml.dump(config_data, tmp, default_flow_style=False)
+        else:
+            json.dump(config_data, tmp, indent=2)
+    return tmp.name
+
+
 def wrap(config_path: Path | None = None) -> Callable[[F], F]:
     """
     HACK: Similar to draccus.wrap but does three additional things:
@@ -225,6 +304,9 @@ def wrap(config_path: Path | None = None) -> Callable[[F], F]:
                 if has_method(argtype, "__get_path_fields__"):
                     path_fields = argtype.__get_path_fields__()
                     cli_args = filter_path_args(path_fields, cli_args)
+                    # Also extract path fields from the YAML/JSON config file
+                    if config_path_cli:
+                        config_path_cli = extract_path_fields_from_config(config_path_cli, path_fields)
                 if has_method(argtype, "from_pretrained") and config_path_cli:
                     cli_args = filter_arg("config_path", cli_args)
                     cfg = argtype.from_pretrained(config_path_cli, cli_args=cli_args)

src/lerobot/configs/train.py

@@ -144,8 +144,11 @@ class TrainPipelineConfig(HubMixin):
             )
             self.reward_model.pretrained_path = str(Path(reward_model_path))
         elif policy_path:
-            cli_overrides = parser.get_cli_overrides("policy")
-            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
+            yaml_overrides = parser.get_yaml_overrides("policy")
+            cli_overrides = parser.get_cli_overrides("policy") or []
+            self.policy = PreTrainedConfig.from_pretrained(
+                policy_path, cli_overrides=yaml_overrides + cli_overrides
+            )
             self.policy.pretrained_path = Path(policy_path)
         elif self.resume:
             config_path = parser.parse_arg("config_path")