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:
Pepijn
2026-05-12 14:31:58 +02:00
parent c661d81409
commit e0fa957569
@@ -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