mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-27 06:29:47 +00:00
fix(smolvla2): eagerly import robot submodules before get_choice_class
``RobotConfig._choice_registry`` is populated as a side-effect of each robot's ``@RobotConfig.register_subclass`` decorator running, and those decorators only fire when the corresponding ``lerobot.robots.<name>`` module is imported. The package's ``__init__.py`` doesn't import them — instead ``make_robot_from_config`` does it lazily in its big if/elif chain. ``_build_robot`` jumped the gun: called ``RobotConfig.get_choice_class (robot_type)`` before any robot module had been imported, so the registry was empty and every ``--robot.type=<X>`` produced ``KeyError: 'X'`` (e.g. ``KeyError: 'omx_follower'``). Walk ``lerobot.robots``'s submodules via ``pkgutil.iter_modules`` and ``importlib.import_module`` each one before the lookup. ~200ms on the first invocation, negligible for an autonomous run. On a real ``KeyError`` (typo / unsupported robot), raise a clean ``ValueError`` listing the registry's available choices instead of a bare KeyError. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -452,14 +452,41 @@ def _build_robot(
|
||||
commanded joint position relative to the current measured one
|
||||
before issuing it on the bus.
|
||||
"""
|
||||
import importlib # noqa: PLC0415
|
||||
import json # noqa: PLC0415
|
||||
import pkgutil # noqa: PLC0415
|
||||
|
||||
import lerobot.robots as _robots_pkg # noqa: PLC0415
|
||||
from lerobot.robots import ( # noqa: PLC0415
|
||||
RobotConfig,
|
||||
make_robot_from_config,
|
||||
)
|
||||
|
||||
cls = RobotConfig.get_choice_class(robot_type)
|
||||
# ``RobotConfig._choice_registry`` is populated lazily — each robot's
|
||||
# ``config_<name>.py`` calls ``@RobotConfig.register_subclass`` at
|
||||
# import time. ``lerobot.robots/__init__.py`` doesn't import the
|
||||
# individual robot packages, so ``get_choice_class(robot_type)``
|
||||
# raises ``KeyError`` until at least one robot module has been
|
||||
# imported. Mirror what ``make_robot_from_config`` does internally:
|
||||
# walk the robots package's submodules and import each so the
|
||||
# decorator side-effect runs. Slow only on the first call (~200ms
|
||||
# for ~10 dataclass modules); negligible for an autonomous run that
|
||||
# then loops at ctrl_hz for minutes.
|
||||
for _modinfo in pkgutil.iter_modules(_robots_pkg.__path__):
|
||||
if _modinfo.name.startswith("_"):
|
||||
continue
|
||||
try:
|
||||
importlib.import_module(f"lerobot.robots.{_modinfo.name}")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("could not import lerobot.robots.%s: %s", _modinfo.name, exc)
|
||||
|
||||
try:
|
||||
cls = RobotConfig.get_choice_class(robot_type)
|
||||
except KeyError as exc:
|
||||
available = sorted(RobotConfig._choice_registry.keys())
|
||||
raise ValueError(
|
||||
f"Unknown robot type {robot_type!r}. Available choices: {available}"
|
||||
) from exc
|
||||
kwargs: dict[str, Any] = {}
|
||||
if robot_port:
|
||||
kwargs["port"] = robot_port
|
||||
|
||||
Reference in New Issue
Block a user