moved isaaclab to envhub

env.py

@@ -19,25 +19,47 @@ from __future__ import annotations
 
 import argparse
 import importlib
+import importlib.util
 import logging
 import os
+import sys
 from pathlib import Path
 from typing import Any
 
 import gymnasium as gym
-import numpy as np
-import torch
 import yaml
 
 from lerobot.envs.configs import EnvConfig
-from lerobot.envs.isaaclab import IsaacLabEnvWrapper, cleanup_isaaclab
-from lerobot.utils.errors import IsaacLabArenaCameraKeyError, IsaacLabArenaStateKeyError
 
 # Hub constants for downloading additional files
 HUB_REPO_ID = "nvkartik/isaaclab-arena-envs"
 DEFAULT_CONFIG_FILE = "configs/config.yaml"
 
 
+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_name = filename.replace(".py", "")
+    spec = importlib.util.spec_from_file_location(module_name, local_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+try:
+    from .isaaclab import IsaacLabEnvWrapper, cleanup_isaaclab
+    from .errors import IsaacLabArenaCameraKeyError, IsaacLabArenaStateKeyError
+except ImportError:
+    _errors = _download_and_import("errors.py")
+    _isaaclab = _download_and_import("isaaclab.py")
+    IsaacLabEnvWrapper = _isaaclab.IsaacLabEnvWrapper
+    cleanup_isaaclab = _isaaclab.cleanup_isaaclab
+    IsaacLabArenaCameraKeyError = _errors.IsaacLabArenaCameraKeyError
+    IsaacLabArenaStateKeyError = _errors.IsaacLabArenaStateKeyError
+
+
 def validate_config(
     env,
     state_keys: tuple[str, ...],
@@ -262,6 +284,7 @@ def make_env(
     n_envs: int = 1,
     use_async_envs: bool = False,
     cfg: EnvConfig | None = None,
+    cfg: EnvConfig | None = None,
     **kwargs: Any,
 ) -> dict[str, dict[int, gym.vector.VectorEnv]]:
     """Create IsaacLab Arena environments (EnvHub-compatible API).
@@ -296,6 +319,42 @@ def make_env(
         raise ValueError("`n_envs` must be at least 1")
 
     # Build config from cfg attributes directly (hub path) or YAML (local dev)
+    if cfg is not None and hasattr(cfg, 'environment'):
+        # Extract config directly from EnvConfig attributes
+        config = {
+            "environment": cfg.environment,
+            "embodiment": cfg.embodiment,
+            "object": cfg.object,
+            "mimic": cfg.mimic,
+            "teleop_device": cfg.teleop_device,
+            "seed": cfg.seed,
+            "device": cfg.device,
+            "disable_fabric": cfg.disable_fabric,
+            "enable_cameras": cfg.enable_cameras,
+            "headless": cfg.headless,
+            "enable_pinocchio": cfg.enable_pinocchio,
+            "episode_length": cfg.episode_length,
+            "state_dim": cfg.state_dim,
+            "action_dim": cfg.action_dim,
+            "camera_height": cfg.camera_height,
+            "camera_width": cfg.camera_width,
+            "video": cfg.video,
+            "video_length": cfg.video_length,
+            "video_interval": cfg.video_interval,
+            "state_keys": cfg.state_keys,
+            "camera_keys": cfg.camera_keys or "",  # Pass empty string for no cameras
+            "task": cfg.task,
+        }
+    else:
+        # Fallback to YAML for local development
+        yaml_path = _find_config_file(kwargs.get("config_path"))
+        if yaml_path is not None:
+            logging.info(f"Loading config from: {yaml_path}")
+            with open(yaml_path) as f:
+                config = yaml.safe_load(f) or {}
+        else:
+            config = {}
+    # Build config from cfg attributes directly (hub path) or YAML (local dev)
     if cfg is not None and hasattr(cfg, 'environment'):
         # Extract config directly from EnvConfig attributes
         config = {

errors.py

@@ -0,0 +1,30 @@
+class IsaacLabArenaError(RuntimeError):
+    """Base exception for IsaacLab Arena environment errors."""
+
+    def __init__(self, message: str = "IsaacLab Arena error"):
+        self.message = message
+        super().__init__(self.message)
+
+
+class IsaacLabArenaConfigError(IsaacLabArenaError):
+    """Exception raised for invalid environment configuration."""
+
+    def __init__(self, invalid: list, available: list, key_type: str = "keys"):
+        msg = f"Invalid {key_type}: {invalid}. Available: {sorted(available)}"
+        super().__init__(msg)
+        self.invalid = invalid
+        self.available = available
+
+
+class IsaacLabArenaCameraKeyError(IsaacLabArenaConfigError):
+    """Exception raised when camera_keys don't match available cameras."""
+
+    def __init__(self, invalid: list, available: list):
+        super().__init__(invalid, available, "camera_keys")
+
+
+class IsaacLabArenaStateKeyError(IsaacLabArenaConfigError):
+    """Exception raised when state_keys don't match available state terms."""
+
+    def __init__(self, invalid: list, available: list):
+        super().__init__(invalid, available, "state_keys")

isaaclab.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+import atexit
+import logging
+import os
+import signal
+from contextlib import suppress
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+import torch
+
+from .errors import IsaacLabArenaError
+
+
+def cleanup_isaaclab(env, simulation_app) -> None:
+    """Cleanup IsaacLab env and simulation app resources."""
+    # Ignore signals during cleanup to prevent interruption
+    old_sigint = signal.signal(signal.SIGINT, signal.SIG_IGN)
+    old_sigterm = signal.signal(signal.SIGTERM, signal.SIG_IGN)
+    try:
+        with suppress(Exception):
+            if env is not None:
+                env.close()
+        with suppress(Exception):
+            if simulation_app is not None:
+                simulation_app.app.close()
+    finally:
+        # Restore signal handlers
+        signal.signal(signal.SIGINT, old_sigint)
+        signal.signal(signal.SIGTERM, old_sigterm)
+
+
+class IsaacLabEnvWrapper(gym.vector.AsyncVectorEnv):
+    """Wrapper adapting IsaacLab batched GPU env to AsyncVectorEnv.
+    IsaacLab handles vectorization internally on GPU. We inherit from
+    AsyncVectorEnv for compatibility with LeRobot."""
+
+    metadata = {"render_modes": ["rgb_array"], "render_fps": 30}
+    _cleanup_in_progress = False  # Class-level flag for re-entrant protection
+
+    def __init__(
+        self,
+        env,
+        episode_length: int = 500,
+        task: str | None = None,
+        render_mode: str | None = "rgb_array",
+        simulation_app=None,
+    ):
+        self._env = env
+        self._num_envs = env.num_envs
+        self._episode_length = episode_length
+        self._closed = False
+        self.render_mode = render_mode
+        self._simulation_app = simulation_app
+
+        self.observation_space = env.observation_space
+        self.action_space = env.action_space
+        self.single_observation_space = env.observation_space
+        self.single_action_space = env.action_space
+        self.task = task
+
+        if hasattr(env, "metadata") and env.metadata:
+            self.metadata = {**self.metadata, **env.metadata}
+
+        # Register cleanup handlers
+        atexit.register(self._cleanup)
+        signal.signal(signal.SIGINT, self._signal_handler)
+        signal.signal(signal.SIGTERM, self._signal_handler)
+
+    def _signal_handler(self, signum, frame):
+        if IsaacLabEnvWrapper._cleanup_in_progress:
+            return  # Prevent re-entrant cleanup
+        IsaacLabEnvWrapper._cleanup_in_progress = True
+        logging.info(f"Received signal {signum}, cleaning up...")
+        self._cleanup()
+        # Exit without raising to avoid propagating through callbacks
+        os._exit(0)
+
+    def _check_closed(self):
+        if self._closed:
+            raise IsaacLabArenaError()
+
+    @property
+    def unwrapped(self):
+        return self
+
+    @property
+    def num_envs(self) -> int:
+        return self._num_envs
+
+    @property
+    def _max_episode_steps(self) -> int:
+        return self._episode_length
+
+    @property
+    def device(self) -> str:
+        return getattr(self._env, "device", "cpu")
+
+    def reset(
+        self,
+        *,
+        seed: int | list[int] | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        self._check_closed()
+        if isinstance(seed, (list, tuple, range)):
+            seed = seed[0] if len(seed) > 0 else None
+
+        obs, info = self._env.reset(seed=seed, options=options)
+
+        if "final_info" not in info:
+            zeros = np.zeros(self._num_envs, dtype=bool)
+            info["final_info"] = {"is_success": zeros}
+
+        return obs, info
+
+    def step(
+        self, actions: np.ndarray | torch.Tensor
+    ) -> tuple[dict, np.ndarray, np.ndarray, np.ndarray, dict]:
+        self._check_closed()
+        if isinstance(actions, np.ndarray):
+            actions = torch.from_numpy(actions).to(self._env.device)
+
+        obs, reward, terminated, truncated, info = self._env.step(actions)
+
+        # Convert to numpy for gym compatibility
+        reward = reward.cpu().numpy().astype(np.float32)
+        terminated = terminated.cpu().numpy().astype(bool)
+        truncated = truncated.cpu().numpy().astype(bool)
+
+        is_success = self._get_success(terminated, truncated)
+        info["final_info"] = {"is_success": is_success}
+
+        return obs, reward, terminated, truncated, info
+
+    def _get_success(self, terminated: np.ndarray, truncated: np.ndarray) -> np.ndarray:
+        is_success = np.zeros(self._num_envs, dtype=bool)
+
+        if not hasattr(self._env, "termination_manager"):
+            return is_success & (terminated | truncated)
+
+        term_manager = self._env.termination_manager
+        if not hasattr(term_manager, "get_term"):
+            return is_success & (terminated | truncated)
+
+        success_tensor = term_manager.get_term("success")
+        if success_tensor is None:
+            return is_success & (terminated | truncated)
+
+        is_success = success_tensor.cpu().numpy().astype(bool)
+
+        return is_success & (terminated | truncated)
+
+    def call(self, method_name: str, *args, **kwargs) -> list[Any]:
+        if method_name == "_max_episode_steps":
+            return [self._episode_length] * self._num_envs
+        if method_name == "task":
+            return [self.task] * self._num_envs
+        if method_name == "render":
+            return self.render_all()
+
+        if hasattr(self._env, method_name):
+            attr = getattr(self._env, method_name)
+            result = attr(*args, **kwargs) if callable(attr) else attr
+            if isinstance(result, list):
+                return result
+            return [result] * self._num_envs
+
+        raise AttributeError(f"IsaacLab-Arena has no method/attribute '{method_name}'")
+
+    def render_all(self) -> list[np.ndarray]:
+        self._check_closed()
+        frames = self.render()
+        if frames is None:
+            placeholder = np.zeros((480, 640, 3), dtype=np.uint8)
+            return [placeholder] * self._num_envs
+
+        if frames.ndim == 4:
+            return [frames[i] for i in range(min(len(frames), self._num_envs))]
+
+        return [np.zeros((480, 640, 3), dtype=np.uint8)] * self._num_envs
+
+    def render(self) -> np.ndarray | None:
+        """Render all environments and return list of frames."""
+        self._check_closed()
+        if self.render_mode != "rgb_array":
+            return None
+
+        frames = self._env.render() if hasattr(self._env, "render") else None
+        if frames is None:
+            return None
+
+        if isinstance(frames, torch.Tensor):
+            frames = frames.cpu().numpy()
+
+        return frames[0] if frames.ndim == 4 else frames
+
+    def _cleanup(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        IsaacLabEnvWrapper._cleanup_in_progress = True
+        logging.info("Cleaning up IsaacLab Arena environment...")
+        cleanup_isaaclab(self._env, self._simulation_app)
+
+    def close(self) -> None:
+        self._cleanup()
+
+    @property
+    def envs(self) -> list[IsaacLabEnvWrapper]:
+        return [self] * self._num_envs
+
+    def __del__(self):
+        self._cleanup()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._cleanup()
+        return False