refactor: unify registry bootstrap and remove silent exception swallowing

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()
 

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(

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."""