feat(envs): Add NVIDIA IsaacLab-Arena Lerobot (#2699)

* adding Isaaclab Arena from collab * adding into lerobot-eval * minor modification * added bash script for env setup * setups * fix applauncher not getting the arguments * data conversion, train and eval smolvla * fixed imports * clean-up * added test suits & clean up - wip * fixed video recording * clean-up * hub integration working * clean-up * added kwargs * Revert "added kwargs" This reverts commit 9b445356385d0707655cf04d02be058b25138119. * added kwargs * clean-up * cleaned unused function * added logging * docs * cleaned up IsaaclabArenaEnv * clean-up * clean-up * clean up * added tests * minor clean-up * fix: support for state based envs * feat(envs): Add NVIDIA IsaacLab Arena integration with LeRobot for policy evaluation at scale * feat(envs): Add IsaacLab Arena integration for policy evaluation Integrate NVIDIA IsaacLab Arena with LeRobot to enable GPU-accelerated simulation through the EnvHub infrastructure. This enables: - Training imitation learning policies (PI0, SmolVLA, etc.) - Evaluating trained policies in with IsaacLab Arena The implementation adds: - IsaaclabArenaEnv config with Arena-specific parameters - IsaaclabArenaProcessorStep for observation processing - Hub loading from nvkartik/isaaclab-arena-envs repository - Video recording support Available environments include GR1 microwave manipulation, Galileo pick-and-place, G1 loco-manipulation, and button pressing tasks. Datasets: nvkartik/Arena-GR1-Manipulation-Task Policies: nvkartik/pi05-arena-gr1-microwave, nvkartik/smolvla-arena-gr1-microwave * added isaaclab arena wrapper and corresponding tests * added error handling * renamed wrapper file: isaaclab_arena to isaaclab * added extra kwarg changes * adjustments for hub envs * correct class name in test file * fixed parsing of env_kwargs * tested end to end * removed unused code * refactor design * shifted IsaacLab to hub * removed IsaacLab tests * docs: Add LW-BenchHub evaluation instructions * docs: Add LW-BenchHub evaluation instructions * docs diet * minor edits to texts * IL Arena commit hash * update links * minor edits * fix numpy version after install of lerobot * links update * valideated on vanilla brev * docs: Add LW-BenchHub evaluation instructions * remove kwargs from all make_env calls * remove kwargs from all make_env calls * fix LW table and indentations * remove environment list from docs * docs: Update lw-benchhub eval config in envhub docs * removing kwargs * removed extra line * ensure pinocchio install for lightwheel + add lightwheel website link * remove env_kwargs * no default empty value for hub_path * not using assert method * remove env_processor defaults * revert and adding default "" value for hub_path * pinning down packages versions * explicit None value for hub_path * Update src/lerobot/configs/eval.py Co-authored-by: Jade Choghari <chogharijade@gmail.com> Signed-off-by: Lior Ben Horin <liorbenhorin@gmail.com> * corrected formatting * corrected job_name var in config * updated docs and namespace * updated namespace * updated docs * updated docs * added hardware requirements * updated docs --------- Signed-off-by: Lior Ben Horin <liorbenhorin@gmail.com> Co-authored-by: lbenhorin <lbenhorin@nvidia.com> Co-authored-by: Lior Ben Horin <liorbenhorin@gmail.com> Co-authored-by: Jade Choghari <chogharijade@gmail.com> Co-authored-by: Steven Palma <imstevenpmwork@ieee.org> Co-authored-by: tianheng.wu <tianheng.wu@lightwheel.ai>
2026-05-16 17:20:05 +00:00 · 2026-01-02 20:36:24 +01:00
parent 9701b9c273
commit fc296548cb
10 changed files with 704 additions and 15 deletions
@@ -38,6 +38,8 @@ class EvalPipelineConfig:
    seed: int | None = 1000
    # Rename map for the observation to override the image and state keys
    rename_map: dict[str, str] = field(default_factory=dict)
+    # Explicit consent to execute remote code from the Hub (required for hub environments).
+    trust_remote_code: bool = False

    def __post_init__(self) -> None:
        # HACK: We parse again the cli args here to get the pretrained path if there was one.
@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-from .configs import AlohaEnv, EnvConfig, PushtEnv  # noqa: F401
+from .configs import AlohaEnv, EnvConfig, HubEnvConfig, PushtEnv  # noqa: F401
@@ -13,7 +13,7 @@
 # limitations under the License.

 import abc
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, fields
 from typing import Any

 import draccus
@@ -68,6 +68,22 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
        raise NotImplementedError()


+@dataclass
+class HubEnvConfig(EnvConfig):
+    """Base class for environments that delegate creation to a hub-hosted make_env.
+
+    Hub environments download and execute remote code from the HF Hub.
+    The hub_path points to a repository containing an env.py with a make_env function.
+    """
+
+    hub_path: str | None = None  # required: e.g., "username/repo" or "username/repo@branch:file.py"
+
+    @property
+    def gym_kwargs(self) -> dict:
+        # Not used for hub environments - the hub's make_env handles everything
+        return {}
+
+
@EnvConfig.register_subclass("aloha")
@dataclass
 class AlohaEnv(EnvConfig):
@@ -368,3 +384,71 @@ class MetaworldEnv(EnvConfig):
            "obs_type": self.obs_type,
            "render_mode": self.render_mode,
        }
+
+
+@EnvConfig.register_subclass("isaaclab_arena")
+@dataclass
+class IsaaclabArenaEnv(HubEnvConfig):
+    hub_path: str = "nvidia/isaaclab-arena-envs"
+    episode_length: int = 300
+    num_envs: int = 1
+    embodiment: str | None = "gr1_pink"
+    object: str | None = "power_drill"
+    mimic: bool = False
+    teleop_device: str | None = None
+    seed: int | None = 42
+    device: str | None = "cuda:0"
+    disable_fabric: bool = False
+    enable_cameras: bool = False
+    headless: bool = False
+    enable_pinocchio: bool = True
+    environment: str | None = "gr1_microwave"
+    task: str | None = "Reach out to the microwave and open it."
+    state_dim: int = 54
+    action_dim: int = 36
+    camera_height: int = 512
+    camera_width: int = 512
+    video: bool = False
+    video_length: int = 100
+    video_interval: int = 200
+    # Comma-separated keys, e.g., "robot_joint_pos,left_eef_pos"
+    state_keys: str = "robot_joint_pos"
+    # Comma-separated keys, e.g., "robot_pov_cam_rgb,front_cam_rgb"
+    # Set to None or "" for environments without cameras
+    camera_keys: str | None = None
+    features: dict[str, PolicyFeature] = field(default_factory=dict)
+    features_map: dict[str, str] = field(default_factory=dict)
+    kwargs: dict | None = None
+
+    def __post_init__(self):
+        if self.kwargs:
+            # dynamically convert kwargs to fields in the dataclass
+            # NOTE! the new fields will not bee seen by the dataclass repr
+            field_names = {f.name for f in fields(self)}
+            for key, value in self.kwargs.items():
+                if key not in field_names and key != "kwargs":
+                    setattr(self, key, value)
+            self.kwargs = None
+
+        # Set action feature
+        self.features[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(self.action_dim,))
+        self.features_map[ACTION] = ACTION
+
+        # Set state feature
+        self.features[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(self.state_dim,))
+        self.features_map[OBS_STATE] = OBS_STATE
+
+        # Add camera features for each camera key
+        if self.enable_cameras and self.camera_keys:
+            for cam_key in self.camera_keys.split(","):
+                cam_key = cam_key.strip()
+                if cam_key:
+                    self.features[cam_key] = PolicyFeature(
+                        type=FeatureType.VISUAL,
+                        shape=(self.camera_height, self.camera_width, 3),
+                    )
+                    self.features_map[cam_key] = f"{OBS_IMAGES}.{cam_key}"
+
+    @property
+    def gym_kwargs(self) -> dict:
+        return {}
@@ -20,11 +20,11 @@ import gymnasium as gym
 from gymnasium.envs.registration import registry as gym_registry

 from lerobot.configs.policies import PreTrainedConfig
-from lerobot.envs.configs import AlohaEnv, EnvConfig, LiberoEnv, PushtEnv
+from lerobot.envs.configs import AlohaEnv, EnvConfig, HubEnvConfig, IsaaclabArenaEnv, LiberoEnv, PushtEnv
 from lerobot.envs.utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result
 from lerobot.policies.xvla.configuration_xvla import XVLAConfig
 from lerobot.processor import ProcessorStep
-from lerobot.processor.env_processor import LiberoProcessorStep
+from lerobot.processor.env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
 from lerobot.processor.pipeline import PolicyProcessorPipeline


@@ -73,6 +73,26 @@ def make_env_pre_post_processors(
    if isinstance(env_cfg, LiberoEnv) or "libero" in env_cfg.type:
        preprocessor_steps.append(LiberoProcessorStep())

+    # For Isaaclab Arena environments, add the IsaaclabArenaProcessorStep
+    if isinstance(env_cfg, IsaaclabArenaEnv) or "isaaclab_arena" in env_cfg.type:
+        # Parse comma-separated keys (handle None for state-based policies)
+        if env_cfg.state_keys:
+            state_keys = tuple(k.strip() for k in env_cfg.state_keys.split(",") if k.strip())
+        else:
+            state_keys = ()
+        if env_cfg.camera_keys:
+            camera_keys = tuple(k.strip() for k in env_cfg.camera_keys.split(",") if k.strip())
+        else:
+            camera_keys = ()
+        if not state_keys and not camera_keys:
+            raise ValueError("At least one of state_keys or camera_keys must be specified.")
+        preprocessor_steps.append(
+            IsaaclabArenaProcessorStep(
+                state_keys=state_keys,
+                camera_keys=camera_keys,
+            )
+        )
+
    preprocessor = PolicyProcessorPipeline(steps=preprocessor_steps)
    postprocessor = PolicyProcessorPipeline(steps=postprocessor_steps)

@@ -98,7 +118,6 @@ def make_env(
        hub_cache_dir (str | None): Optional cache path for downloaded hub files.
        trust_remote_code (bool): **Explicit consent** to execute remote code from the Hub.
            Default False — must be set to True to import/exec hub `env.py`.
-
    Raises:
        ValueError: if n_envs < 1
        ModuleNotFoundError: If the requested env package is not installed
@@ -112,19 +131,35 @@ def make_env(
    """
    # if user passed a hub id string (e.g., "username/repo", "username/repo@main:env.py")
    # simplified: only support hub-provided `make_env`
+    # TODO: (jadechoghari): deprecate string API and remove this check
    if isinstance(cfg, str):
+        hub_path: str | None = cfg
+    elif isinstance(cfg, HubEnvConfig):
+        hub_path = cfg.hub_path
+    else:
+        hub_path = None
+
+    # If hub_path is set, download and call hub-provided `make_env`
+    if hub_path:
        # _download_hub_file will raise the same RuntimeError if trust_remote_code is False
-        repo_id, file_path, local_file, revision = _download_hub_file(cfg, trust_remote_code, hub_cache_dir)
+        repo_id, file_path, local_file, revision = _download_hub_file(
+            hub_path, trust_remote_code, hub_cache_dir
+        )

        # import and surface clear import errors
        module = _import_hub_module(local_file, repo_id)

        # call the hub-provided make_env
-        raw_result = _call_make_env(module, n_envs=n_envs, use_async_envs=use_async_envs)
+        env_cfg = None if isinstance(cfg, str) else cfg
+        raw_result = _call_make_env(module, n_envs=n_envs, use_async_envs=use_async_envs, cfg=env_cfg)

        # normalize the return into {suite: {task_id: vec_env}}
        return _normalize_hub_result(raw_result)

+    # At this point, cfg must be an EnvConfig (not a string) since hub_path would have been set otherwise
+    if isinstance(cfg, str):
+        raise TypeError("cfg should be an EnvConfig at this point")
+
    if n_envs < 1:
        raise ValueError("`n_envs` must be at least 1")

@@ -46,7 +46,7 @@ def _convert_nested_dict(d):


 def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Tensor]:
-    # TODO(aliberts, rcadene): refactor this to use features from the environment (no hardcoding)
+    # TODO(jadechoghari, imstevenpmwork): refactor this to use features from the environment (no hardcoding)
    """Convert environment observation to LeRobot format observation.
    Args:
        observation: Dictionary of observation batches from a Gym vector environment.
@@ -98,11 +98,19 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten

    if "robot_state" in observations:
        return_observations[f"{OBS_STR}.robot_state"] = _convert_nested_dict(observations["robot_state"])
+
+    # Handle IsaacLab Arena format: observations have 'policy' and 'camera_obs' keys
+    if "policy" in observations:
+        return_observations[f"{OBS_STR}.policy"] = observations["policy"]
+
+    if "camera_obs" in observations:
+        return_observations[f"{OBS_STR}.camera_obs"] = observations["camera_obs"]
+
    return return_observations


 def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
-    # TODO(aliberts, rcadene): remove this hardcoding of keys and just use the nested keys as is
+    # TODO(jadechoghari, imstevenpmwork): remove this hardcoding of keys and just use the nested keys as is
    # (need to also refactor preprocess_observation and externalize normalization from policies)
    policy_features = {}
    for key, ft in env_cfg.features.items():
@@ -302,7 +310,7 @@ def _import_hub_module(local_file: str, repo_id: str) -> Any:
    return module


-def _call_make_env(module: Any, n_envs: int, use_async_envs: bool) -> Any:
+def _call_make_env(module: Any, n_envs: int, use_async_envs: bool, cfg: EnvConfig | None) -> Any:
    """
    Ensure module exposes make_env and call it.
    """
@@ -311,7 +319,11 @@ def _call_make_env(module: Any, n_envs: int, use_async_envs: bool) -> Any:
            f"The hub module {getattr(module, '__name__', 'hub_module')} must expose `make_env(n_envs=int, use_async_envs=bool)`."
        )
    entry_fn = module.make_env
-    return entry_fn(n_envs=n_envs, use_async_envs=use_async_envs)
+    # Only pass cfg if it's not None (i.e., when an EnvConfig was provided, not a string hub ID)
+    if cfg is not None:
+        return entry_fn(n_envs=n_envs, use_async_envs=use_async_envs, cfg=cfg)
+    else:
+        return entry_fn(n_envs=n_envs, use_async_envs=use_async_envs)


 def _normalize_hub_result(result: Any) -> dict[str, dict[int, gym.vector.VectorEnv]]:
@@ -18,7 +18,7 @@ from dataclasses import dataclass
 import torch

 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.utils.constants import OBS_IMAGES, OBS_STATE
+from lerobot.utils.constants import OBS_IMAGES, OBS_STATE, OBS_STR

 from .pipeline import ObservationProcessorStep, ProcessorStepRegistry

@@ -152,3 +152,78 @@ class LiberoProcessorStep(ObservationProcessorStep):
            result[mask] = axis * angle.unsqueeze(1)

        return result
+
+
+@dataclass
+@ProcessorStepRegistry.register(name="isaaclab_arena_processor")
+class IsaaclabArenaProcessorStep(ObservationProcessorStep):
+    """
+    Processes IsaacLab Arena observations into LeRobot format.
+
+    **State Processing:**
+    - Extracts state components from obs["policy"] based on `state_keys`.
+    - Concatenates into a flat vector mapped to "observation.state".
+
+    **Image Processing:**
+    - Extracts images from obs["camera_obs"] based on `camera_keys`.
+    - Converts from (B, H, W, C) uint8 to (B, C, H, W) float32 [0, 1].
+    - Maps to "observation.images.<camera_name>".
+    """
+
+    # Configurable from IsaacLabEnv config / cli args: --env.state_keys="robot_joint_pos,left_eef_pos"
+    state_keys: tuple[str, ...]
+
+    # Configurable from IsaacLabEnv config / cli args: --env.camera_keys="robot_pov_cam_rgb"
+    camera_keys: tuple[str, ...]
+
+    def _process_observation(self, observation):
+        """
+        Processes both image and policy state observations from IsaacLab Arena.
+        """
+        processed_obs = {}
+
+        if f"{OBS_STR}.camera_obs" in observation:
+            camera_obs = observation[f"{OBS_STR}.camera_obs"]
+
+            for cam_name, img in camera_obs.items():
+                if cam_name not in self.camera_keys:
+                    continue
+
+                img = img.permute(0, 3, 1, 2).contiguous()
+                if img.dtype == torch.uint8:
+                    img = img.float() / 255.0
+                elif img.dtype != torch.float32:
+                    img = img.float()
+
+                processed_obs[f"{OBS_IMAGES}.{cam_name}"] = img
+
+        # Process policy state -> observation.state
+        if f"{OBS_STR}.policy" in observation:
+            policy_obs = observation[f"{OBS_STR}.policy"]
+
+            # Collect state components in order
+            state_components = []
+            for key in self.state_keys:
+                if key in policy_obs:
+                    component = policy_obs[key]
+                    # Flatten extra dims: (B, N, M) -> (B, N*M)
+                    if component.dim() > 2:
+                        batch_size = component.shape[0]
+                        component = component.view(batch_size, -1)
+                    state_components.append(component)
+
+            if state_components:
+                state = torch.cat(state_components, dim=-1)
+                state = state.float()
+                processed_obs[OBS_STATE] = state
+
+        return processed_obs
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """Not used for policy evaluation."""
+        return features
+
+    def observation(self, observation):
+        return self._process_observation(observation)
@@ -509,7 +509,12 @@ def eval_main(cfg: EvalPipelineConfig):
    logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")

    logging.info("Making environment.")
-    envs = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
+    envs = make_env(
+        cfg.env,
+        n_envs=cfg.eval.batch_size,
+        use_async_envs=cfg.eval.use_async_envs,
+        trust_remote_code=cfg.trust_remote_code,
+    )

    logging.info("Making policy.")