update env.py

env.py

@@ -1,15 +1,55 @@
+from __future__ import annotations
 import logging
 from typing import Any
 
+import importlib
+import importlib.util
+import logging
 import os
 import sys
 
-# Ensure the repository root is on sys.path
-_REPO_ROOT = os.path.dirname(os.path.abspath(__file__))
-if _REPO_ROOT not in sys.path:
-    sys.path.insert(0, _REPO_ROOT)
 
-from isaaclab_env_wrapper import IsaacLabEnvWrapper
+# Hub constants for downloading additional files
+HUB_REPO_ID = "LightwheelAI/lw_benchhub_env"
+
+
+def _download_and_import(filename: str):
+    """Download file from Hub and import as module."""
+    from huggingface_hub import hf_hub_download
+
+    local_path = hf_hub_download(repo_id=HUB_REPO_ID, filename=filename)
+    module_dir = os.path.dirname(local_path)
+
+    # Add directory to sys.path so modules can import each other
+    if module_dir not in sys.path:
+        sys.path.insert(0, module_dir)
+
+    module_name = filename.replace(".py", "")
+
+    # Read file content and replace relative imports with absolute imports
+    with open(local_path) as f:
+        content = f.read()
+
+    # Replace relative imports (e.g., "from .errors" -> "from errors")
+    # content = content.replace("from .errors import", "from errors import")
+    content = content.replace("from .isaaclab_env_wrapper import", "from isaaclab_env_wrapper import")
+
+    # Create module and execute modified content
+    module = importlib.util.module_from_spec(importlib.util.spec_from_file_location(module_name, local_path))
+    sys.modules[module_name] = module
+
+    # Compile and execute the modified code
+    code = compile(content, local_path, "exec")
+    exec(code, module.__dict__)  # noqa: S102
+
+    return module
+
+
+try:
+    from .isaaclab_env_wrapper import IsaacLabEnvWrapper
+except ImportError:
+    _isaaclab_wrapper = _download_and_import("isaaclab_env_wrapper.py")
+    IsaacLabEnvWrapper = _isaaclab_wrapper.IsaacLabEnvWrapper
 
 
 def make_env(n_envs: int = 1, use_async_envs: bool = False, **kwargs: Any):