Use selective ManiSkill tabletop registration for CTT rollout

workspace/scripts/eval_ctt_generated_rollout.py

@@ -2,12 +2,16 @@
 from __future__ import annotations
 
 import argparse
+import importlib
+import importlib.machinery
 import json
 import math
+import os
 import pickle
 import shutil
 import subprocess
 import sys
+import types
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
@@ -155,7 +159,7 @@ def main(argv: list[str] | None = None) -> int:
     _append_log(log_path, "importing gymnasium/mani_skill")
     try:
         import gymnasium as gym
-        import mani_skill  # noqa: F401 - importing registers environments
+        _register_required_maniskill_envs(log_path)
     except ImportError as exc:  # pragma: no cover - exercised in the Apptainer env
         raise ImportError(
             "CTT measured rollout requires gymnasium, mani_skill, numpy, and torch. "
@@ -954,6 +958,70 @@ def _uses_single_env_cpu_backend(sim_backend: Any) -> bool:
     return value in {"cpu", "physx_cpu"} or value.endswith("_cpu")
 
 
+def _register_required_maniskill_envs(log_path: Path | None = None) -> None:
+    """Register only the tabletop ManiSkill tasks used by the CIL diagnostic.
+
+    ManiSkill 3.0.1 imports all task families from its top-level package,
+    including digital-twin Bridge tasks that import cv2. On the current HPC
+    container that full import can hang before our five tabletop tasks are
+    registered. A lightweight package stub lets Python resolve ManiSkill
+    submodules without executing the top-level eager task import.
+    """
+
+    package = sys.modules.get("mani_skill")
+    if package is None:
+        spec = importlib.machinery.PathFinder.find_spec("mani_skill")
+        if spec is None or not spec.submodule_search_locations:
+            raise ImportError("Could not locate installed mani_skill package")
+        package_dir = Path(list(spec.submodule_search_locations)[0]).resolve()
+        package = types.ModuleType("mani_skill")
+        package.__file__ = str(package_dir / "__init__.py")
+        package.__path__ = [str(package_dir)]
+        package.__package__ = "mani_skill"
+        package.__version__ = "3.0.1"
+        package.PACKAGE_DIR = package_dir
+        package.PACKAGE_ASSET_DIR = package_dir / "assets"
+        asset_root = Path(
+            os.getenv("MS_ASSET_DIR", os.path.join(os.path.expanduser("~"), ".maniskill"))
+        )
+        package.ASSET_DIR = asset_root / "data"
+        package.DEMO_DIR = asset_root / "demos"
+        package.format_path = lambda p: p.format(
+            PACKAGE_DIR=package.PACKAGE_DIR,
+            PACKAGE_ASSET_DIR=package.PACKAGE_ASSET_DIR,
+            ASSET_DIR=package.ASSET_DIR,
+        )
+        sys.modules["mani_skill"] = package
+        _install_package_stub("mani_skill.envs", package_dir / "envs")
+        _install_package_stub("mani_skill.envs.tasks", package_dir / "envs" / "tasks")
+        _install_package_stub(
+            "mani_skill.envs.tasks.tabletop",
+            package_dir / "envs" / "tasks" / "tabletop",
+        )
+        package.logger = importlib.import_module("mani_skill.utils.logging_utils").logger
+        _append_log(log_path, "installed lightweight mani_skill tabletop import stub")
+
+    modules = [
+        "mani_skill.envs.tasks.tabletop.pick_cube",
+        "mani_skill.envs.tasks.tabletop.pull_cube",
+        "mani_skill.envs.tasks.tabletop.push_cube",
+        "mani_skill.envs.tasks.tabletop.stack_cube",
+        "mani_skill.envs.tasks.tabletop.lift_peg_upright",
+    ]
+    for module_name in modules:
+        importlib.import_module(module_name)
+    _append_log(log_path, "registered tabletop ManiSkill tasks selectively")
+
+
+def _install_package_stub(name: str, path: Path) -> None:
+    if name in sys.modules:
+        return
+    module = types.ModuleType(name)
+    module.__path__ = [str(path)]
+    module.__package__ = name
+    sys.modules[name] = module
+
+
 def _validate_indexes(
     source_path: Path,
     source_index: dict[str, Any],