diff --git a/scripts/play_interactive.py b/scripts/play_interactive.py index 675c2d239..310197ca1 100644 --- a/scripts/play_interactive.py +++ b/scripts/play_interactive.py @@ -14,9 +14,7 @@ """ import argparse -import importlib import os -import pkgutil import sys import time from pathlib import Path @@ -29,24 +27,7 @@ ROOT_DIR = Path(__file__).parent.parent sys.path.append(str(ROOT_DIR)) - -def ensure_registries(): - for pkg_name in ( - "unilab.envs.locomotion", - "unilab.envs.manipulation", - "unilab.envs.motion_tracking", - ): - try: - package = importlib.import_module(pkg_name) - if hasattr(package, "__path__"): - for _, name, _ in pkgutil.walk_packages(package.__path__, package.__name__ + "."): - try: - importlib.import_module(name) - except Exception: - pass - except ImportError: - pass - +from unilab.training import ensure_registries ensure_registries() diff --git a/src/unilab/utils/algo_utils.py b/src/unilab/utils/algo_utils.py index 679af2493..cd3193dfd 100644 --- a/src/unilab/utils/algo_utils.py +++ b/src/unilab/utils/algo_utils.py @@ -1,26 +1,74 @@ """Common utilities for RL algorithms.""" +from __future__ import annotations + import importlib +import logging import pkgutil +from typing import Sequence + +logger = logging.getLogger(__name__) + +# Default packages to scan for env registration +_DEFAULT_REGISTRY_PACKAGES = ( + "unilab.envs.locomotion", + "unilab.envs.manipulation", + "unilab.envs.motion_tracking", +) + +def ensure_registries( + packages: Sequence[str] | None = None, + *, + optional_packages: Sequence[str] | None = None, + fail_on_error: bool = True, +) -> None: + """Import all env modules so they are registered. -def ensure_registries(): - """Import all env modules so they are registered.""" - for pkg_name in ( - "unilab.envs.locomotion", - "unilab.envs.manipulation", - "unilab.envs.motion_tracking", - ): + Args: + packages: List of package names to scan for env modules. + Defaults to standard unilab env packages. + optional_packages: List of optional package names that may not be present. + Import failures for these are logged as warnings, not raised. + fail_on_error: If True (default), raise exceptions for non-optional packages. + If False, log warnings instead of raising. + + Raises: + ImportError: If a non-optional package fails to import and fail_on_error is True. + Exception: If a module within a package fails to import and fail_on_error is True. + """ + pkgs = list(packages) if packages is not None else list(_DEFAULT_REGISTRY_PACKAGES) + optional = set(optional_packages) if optional_packages else set() + + for pkg_name in pkgs: + is_optional = pkg_name in optional try: package = importlib.import_module(pkg_name) - if hasattr(package, "__path__"): - for _, name, _ in pkgutil.walk_packages(package.__path__, package.__name__ + "."): - try: - importlib.import_module(name) - except Exception: - pass - except ImportError: - pass + except ImportError as e: + if is_optional: + logger.warning("Optional registry package not found: %s (%s)", pkg_name, e) + elif fail_on_error: + raise ImportError( + f"Failed to import registry package '{pkg_name}'. " + f"Add to optional_packages if this is expected to be absent." + ) from e + else: + logger.warning("Registry package not found: %s (%s)", pkg_name, e) + continue + + if not hasattr(package, "__path__"): + continue + + for _, name, _ in pkgutil.walk_packages(package.__path__, package.__name__ + "."): + try: + importlib.import_module(name) + except Exception as e: + if fail_on_error and not is_optional: + raise RuntimeError( + f"Failed to import env module '{name}'. " + f"Fix the import error or add '{pkg_name}' to optional_packages." + ) from e + logger.warning("Failed to import env module '%s': %s", name, e) def build_actor( diff --git a/tests/utils/test_algo_utils.py b/tests/utils/test_algo_utils.py index e1e9b423d..6a5fb520d 100644 --- a/tests/utils/test_algo_utils.py +++ b/tests/utils/test_algo_utils.py @@ -11,13 +11,60 @@ class TestEnsureRegistries: """Tests for ensure_registries.""" def test_runs_without_error(self) -> None: + """Default call should work with standard packages.""" ensure_registries() def test_is_idempotent(self) -> None: + """Multiple calls should be safe.""" ensure_registries() ensure_registries() ensure_registries() + def test_fail_on_error_default_raises(self) -> None: + """By default, invalid packages should raise ImportError.""" + with pytest.raises(ImportError): + ensure_registries(["nonexistent_package_12345"]) + + def test_fail_on_error_false_logs_warning(self, caplog) -> None: + """With fail_on_error=False, invalid packages log warning.""" + import logging + + with caplog.at_level(logging.WARNING): + ensure_registries(["nonexistent_package_12345"], fail_on_error=False) + + assert "nonexistent_package_12345" in caplog.text + + def test_optional_package_logs_warning(self, caplog) -> None: + """Optional packages that fail to import log warning instead of raising.""" + import logging + + with caplog.at_level(logging.WARNING): + ensure_registries( + ["nonexistent_package_12345"], + optional_packages=["nonexistent_package_12345"], + ) + + assert "nonexistent_package_12345" in caplog.text + assert "Optional" in caplog.text + + def test_mixed_optional_and_required(self, caplog) -> None: + """Mix of optional (fails gracefully) and required (works).""" + import logging + + # Use real package + fake optional package + with caplog.at_level(logging.WARNING): + ensure_registries( + ["unilab.envs.locomotion", "nonexistent_optional_12345"], + optional_packages=["nonexistent_optional_12345"], + ) + + # Should warn about optional but not raise + assert "nonexistent_optional_12345" in caplog.text + + def test_empty_packages_list(self) -> None: + """Empty packages list should be a no-op.""" + ensure_registries([]) + class TestBuildActor: """Tests for build_actor."""