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 commanded joint position relative to the current measured one
before issuing it on the bus. before issuing it on the bus.
""" """
import importlib # noqa: PLC0415
import json # noqa: PLC0415 import json # noqa: PLC0415
import pkgutil # noqa: PLC0415
import lerobot.robots as _robots_pkg # noqa: PLC0415
from lerobot.robots import ( # noqa: PLC0415 from lerobot.robots import ( # noqa: PLC0415
RobotConfig, RobotConfig,
make_robot_from_config, make_robot_from_config,
) )
# ``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) 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] = {} kwargs: dict[str, Any] = {}
if robot_port: if robot_port:
kwargs["port"] = robot_port kwargs["port"] = robot_port