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>
This commit is contained in:
Kartik
2026-01-02 20:36:24 +01:00
committed by GitHub
parent 9701b9c273
commit fc296548cb
10 changed files with 704 additions and 15 deletions
+2
View File
@@ -59,6 +59,8 @@
title: Environments from the Hub title: Environments from the Hub
- local: envhub_leisaac - local: envhub_leisaac
title: Control & Train Robots in Sim (LeIsaac) title: Control & Train Robots in Sim (LeIsaac)
- local: envhub_isaaclab_arena
title: NVIDIA IsaacLab Arena Environments
- local: libero - local: libero
title: Using Libero title: Using Libero
- local: metaworld - local: metaworld
+474
View File
@@ -0,0 +1,474 @@
# NVIDIA IsaacLab Arena & LeRobot
LeRobot EnvHub now supports **GPU-accelerated simulation** with IsaacLab Arena for policy evaluation at scale.
Train and evaluate imitation learning policies with high-fidelity simulation — all integrated into the LeRobot ecosystem.
<img
src="https://huggingface.co/nvidia/isaaclab-arena-envs/resolve/main/assets/Gr1OpenMicrowaveEnvironment.png"
alt="IsaacLab Arena - GR1 Microwave Environment"
style={{ maxWidth: "100%", borderRadius: "8px", marginBottom: "1rem" }}
/>
[IsaacLab Arena](https://github.com/isaac-sim/IsaacLab-Arena) integrates with NVIDIA IsaacLab to provide:
- 🤖 **Humanoid embodiments**: GR1, G1, Galileo with various configurations
- 🎯 **Manipulation & loco-manipulation tasks**: Microwave opening, pick-and-place, button pressing
- ⚡ **GPU-accelerated rollouts**: Parallel environment execution on NVIDIA GPUs
- 🖼️ **RTX Rendering**: Evaluate vision-based policies with realistic rendering, reflections and refractions.
- 📦 **LeRobot-compatible datasets**: Ready for training with GR00T Nx, PI0, SmolVLA, ACT, Diffusion policies
- 🔄 **EnvHub integration**: Load environments from HuggingFace Hub with one line
## Installation
### Prerequisites
Hardware requirements are shared with Isaac Sim, and are detailed in [Isaac Sim Requirements](https://docs.isaacsim.omniverse.nvidia.com/5.1.0/installation/requirements.html).
- NVIDIA GPU with CUDA support
- NVIDIA driver compatible with IsaacSim 5.1.0
- Linux (Ubuntu 22.04 / 24.04)
### Setup
```bash
# 1. Create conda environment
conda create -y -n lerobot-arena python=3.11
conda activate lerobot-arena
conda install -y -c conda-forge ffmpeg=7.1.1
# 2. Install Isaac Sim 5.1.0
pip install "isaacsim[all,extscache]==5.1.0" --extra-index-url https://pypi.nvidia.com
# Accept NVIDIA EULA (required)
export ACCEPT_EULA=Y
export PRIVACY_CONSENT=Y
# 3. Install IsaacLab 2.3.0
git clone https://github.com/isaac-sim/IsaacLab.git
cd IsaacLab
git checkout v2.3.0
./isaaclab.sh -i
cd ..
# 4. Install IsaacLab Arena
git clone https://github.com/isaac-sim/IsaacLab-Arena.git
cd IsaacLab-Arena
git checkout release/0.1.1
pip install -e .
cd ..
# 5. Install LeRobot
git clone https://github.com/huggingface/lerobot.git
cd lerobot
pip install -e .
cd ..
# 6. Install additional dependencies
pip install onnxruntime==1.23.2 lightwheel-sdk==1.0.1 vuer[all]==0.0.70 qpsolvers==4.8.1
pip install numpy==1.26.0 # Isaac Sim 5.1 depends on numpy==1.26.0, this will be fixed in next release
```
## Evaluating Policies
### Pre-trained Policies
The following trained policies are available:
| Policy | Architecture | Task | Link |
| :-------------------------- | :----------- | :------------ | :----------------------------------------------------------------------- |
| pi05-arena-gr1-microwave | PI0.5 | GR1 Microwave | [HuggingFace](https://huggingface.co/nvidia/pi05-arena-gr1-microwave) |
| smolvla-arena-gr1-microwave | SmolVLA | GR1 Microwave | [HuggingFace](https://huggingface.co/nvidia/smolvla-arena-gr1-microwave) |
### Evaluate SmolVLA
```bash
pip install -e ".[smolvla]"
pip install numpy==1.26.0 # revert to numpy version is 1.26
```
```bash
lerobot-eval \
--policy.path=nvidia/smolvla-arena-gr1-microwave \
--env.type=isaaclab_arena \
--env.hub_path=nvidia/isaaclab-arena-envs \
--rename_map='{"observation.images.robot_pov_cam_rgb": "observation.images.robot_pov_cam"}' \
--policy.device=cuda \
--env.environment=gr1_microwave \
--env.embodiment=gr1_pink \
--env.object=mustard_bottle \
--env.headless=false \
--env.enable_cameras=true \
--env.video=true \
--env.video_length=10 \
--env.video_interval=15 \
--env.state_keys=robot_joint_pos \
--env.camera_keys=robot_pov_cam_rgb \
--trust_remote_code=True \
--eval.batch_size=1
```
### Evaluate PI0.5
```bash
pip install -e ".[pi]"
pip install numpy==1.26.0 # revert to numpy version is 1.26
```
<Tip>PI0.5 requires disabling torch compile for evaluation:</Tip>
```bash
TORCH_COMPILE_DISABLE=1 TORCHINDUCTOR_DISABLE=1 lerobot-eval \
--policy.path=nvidia/pi05-arena-gr1-microwave \
--env.type=isaaclab_arena \
--env.hub_path=nvidia/isaaclab-arena-envs \
--rename_map='{"observation.images.robot_pov_cam_rgb": "observation.images.robot_pov_cam"}' \
--policy.device=cuda \
--env.environment=gr1_microwave \
--env.embodiment=gr1_pink \
--env.object=mustard_bottle \
--env.headless=false \
--env.enable_cameras=true \
--env.video=true \
--env.video_length 15 \
--env.video_interval 15 \
--env.state_keys=robot_joint_pos \
--env.camera_keys=robot_pov_cam_rgb \
--trust_remote_code=True \
--eval.batch_size=1
```
<Tip>
To change the number of parallel environments, use the ```--eval.batch_size```
flag.
</Tip>
### What to Expect
During evaluation, you will see a progress bar showing the running success rate:
```
Stepping through eval batches: 8%|██████▍ | 4/50 [00:45<08:06, 10.58s/it, running_success_rate=25.0%]
```
### Video Recording
To enable video recording during evaluation, add the following flags to your command:
```bash
--env.video=true \
--env.video_length=15 \
--env.video_interval=15
```
For more details on video recording, see the [IsaacLab Recording Documentation](https://isaac-sim.github.io/IsaacLab/main/source/how-to/record_video.html).
<Tip>
When running headless with `--env.headless=true`, you must also enable cameras explicitly for camera enabled environments:
```bash
--env.headless=true --env.enable_cameras=true
```
</Tip>
### Output Directory
Evaluation videos are saved to the output directory with the following structure:
```
outputs/eval/<date>/<timestamp>_<env>_<policy>/videos/<task>_<env_id>/eval_episode_<n>.mp4
```
For example:
```
outputs/eval/2026-01-02/14-38-01_isaaclab_arena_smolvla/videos/gr1_microwave_0/eval_episode_0.mp4
```
## Training Policies
To learn more about training policies with LeRobot, please refer to training documentation:
- [SmolVLA](./smolvla)
- [Pi0.5](./pi05)
- [GR00T N1.5](./groot)
Sample IsaacLab Arena datasets are available on HuggingFace Hub for experimentation:
| Dataset | Description | Frames |
| :-------------------------------------------------------------------------------------------------------- | :------------------------- | :----- |
| [Arena-GR1-Manipulation-Task](https://huggingface.co/datasets/nvidia/Arena-GR1-Manipulation-Task-v3) | GR1 microwave manipulation | ~4K |
| [Arena-G1-Loco-Manipulation-Task](https://huggingface.co/datasets/nvidia/Arena-G1-Loco-Manipulation-Task) | G1 loco-manipulation | ~4K |
## Environment Configuration
### Full Configuration Options
```python
from lerobot.envs.configs import IsaaclabArenaEnv
config = IsaaclabArenaEnv(
# Environment selection
environment="gr1_microwave", # Task environment
embodiment="gr1_pink", # Robot embodiment
object="power_drill", # Object to manipulate
# Simulation settings
episode_length=300, # Max steps per episode
headless=True, # Run without GUI
device="cuda:0", # GPU device
seed=42, # Random seed
# Observation configuration
state_keys="robot_joint_pos", # State observation keys (comma-separated)
camera_keys="robot_pov_cam_rgb", # Camera observation keys (comma-separated)
state_dim=54, # Expected state dimension
action_dim=36, # Expected action dimension
camera_height=512, # Camera image height
camera_width=512, # Camera image width
enable_cameras=True, # Enable camera observations
# Video recording
video=False, # Enable video recording
video_length=100, # Frames per video
video_interval=200, # Steps between recordings
# Advanced
mimic=False, # Enable mimic mode
teleop_device=None, # Teleoperation device
disable_fabric=False, # Disable fabric optimization
enable_pinocchio=True, # Enable Pinocchio for IK
)
```
### Using Environment Hub directly for advanced usage
Create a file called `test_env_load_arena.py` or [download from the EnvHub](https://huggingface.co/nvidia/isaaclab-arena-envs/blob/main/tests/test_env_load_arena.py):
```python
import logging
from dataclasses import asdict
from pprint import pformat
import torch
import tqdm
from lerobot.configs import parser
from lerobot.configs.eval import EvalPipelineConfig
@parser.wrap()
def main(cfg: EvalPipelineConfig):
"""Run zero action rollout for IsaacLab Arena environment."""
logging.info(pformat(asdict(cfg)))
from lerobot.envs.factory import make_env
env_dict = make_env(
cfg.env,
n_envs=cfg.env.num_envs,
trust_remote_code=True,
)
env = next(iter(env_dict.values()))[0]
env.reset()
for _ in tqdm.tqdm(range(cfg.env.episode_length)):
with torch.inference_mode():
actions = env.action_space.sample()
obs, rewards, terminated, truncated, info = env.step(actions)
if terminated.any() or truncated.any():
obs, info = env.reset()
env.close()
if __name__ == "__main__":
main()
```
Run with:
```bash
python test_env_load_arena.py \
--env.environment=g1_locomanip_pnp \
--env.embodiment=gr1_pink \
--env.object=cracker_box \
--env.num_envs=4 \
--env.enable_cameras=true \
--env.seed=1000 \
--env.video=true \
--env.video_length=10 \
--env.video_interval=15 \
--env.headless=false \
--env.hub_path=nvidia/isaaclab-arena-envs \
--env.type=isaaclab_arena
```
## Troubleshooting
### CUDA out of memory
Reduce `batch_size` or use a GPU with more VRAM:
```bash
--eval.batch_size=1
```
### EULA not accepted
Set environment variables before running:
```bash
export ACCEPT_EULA=Y
export PRIVACY_CONSENT=Y
```
### Video recording not working
Enable cameras when running headless:
```bash
--env.video=true --env.enable_cameras=true --env.headless=true
```
### Policy output dimension mismatch
E.g. ensure `action_dim` matches your policy:
```bash
--env.action_dim=36
```
### libGLU.so.1 Errors during Isaac Sim initialization
Ensure you have the following dependencies installed, this is likely to happen on headless machines.
```bash
sudo apt update && sudo apt install -y libglu1-mesa libxt6
```
## See Also
- [EnvHub Documentation](./envhub.mdx) - General EnvHub usage
- [IsaacLab Arena GitHub](https://github.com/isaac-sim/IsaacLab-Arena)
- [IsaacLab Documentation](https://isaac-sim.github.io/IsaacLab/)
## LightWheel LW-BenchHub
[LightWheel AI](https://www.lightwheel.ai) are bringing `Lightwheel-Libero-Tasks` and `Lightwheel-RoboCasa-Tasks` with 268 tasks to the LeRobot ecosystem.
LW-BenchHub collects and generates large-scale datasets via teleoperation that comply with the LeRobot specification, enabling out-of-the-box training and evaluation workflows.
With the unified interface provided by EnvHub, developers can quickly build end-to-end experimental pipelines.
### Install
Assuming you followed the [Installation](#installation) steps, you can install LW-BenchHub with:
```bash
conda install pinocchio -c conda-forge -y
git clone https://github.com/LightwheelAI/lw_benchhub
cd lw_benchhub
pip install -e .
```
For more detailed instructions, please refer to the [LW-BenchHub Documentation](https://docs.lightwheel.net/lw_benchhub/usage/Installation).
### Lightwheel Tasks Dataset
LW-BenchHub datasets are available on HuggingFace Hub:
| Dataset | Description | Tasks | Frames |
| :------------------------------------------------------------------------------------------------------------ | :---------------------- | :---- | :----- |
| [Lightwheel-Tasks-X7S](https://huggingface.co/datasets/LightwheelAI/Lightwheel-Tasks-X7S) | X7S LIBERO and RoboCasa | 117 | ~10.3M |
| [Lightwheel-Tasks-Double-Piper](https://huggingface.co/datasets/LightwheelAI/Lightwheel-Tasks-Double-Piper) | Double-Piper LIBERO | 130 | ~6.0M |
| [Lightwheel-Tasks-G1-Controller](https://huggingface.co/datasets/LightwheelAI/Lightwheel-Tasks-G1-Controller) | G1-Controller LIBERO | 62 | ~2.7M |
| [Lightwheel-Tasks-G1-WBC](https://huggingface.co/datasets/LightwheelAI/Lightwheel-Tasks-G1-WBC) | G1-WBC RoboCasa | 32 | ~1.5M |
For training policies, refer to the [Training Policies](#training-policies) section.
### Evaluating Policies
#### Pre-trained Policies
The following trained policies are available:
| Policy | Architecture | Task | Layout | Robot | Link |
| :----------------------- | :----------- | :----------------------------- | :--------- | :-------------- | :------------------------------------------------------------------------------------ |
| smolvla-double-piper-pnp | SmolVLA | L90K1PutTheBlackBowlOnThePlate | libero-1-1 | DoublePiper-Abs | [HuggingFace](https://huggingface.co/LightwheelAI/smolvla-double-piper-pnp/tree/main) |
#### Evaluate SmolVLA
```bash
lerobot-eval \
--policy.path=LightwheelAI/smolvla-double-piper-pnp \
--env.type=isaaclab_arena \
--rename_map='{"observation.images.left_hand_camera_rgb": "observation.images.left_hand", "observation.images.right_hand_camera_rgb": "observation.images.right_hand", "observation.images.first_person_camera_rgb": "observation.images.first_person"}' \
--env.hub_path=LightwheelAI/lw_benchhub_env \
--env.kwargs='{"config_path": "configs/envhub/example.yml"}' \
--trust_remote_code=true \
--env.state_keys=joint_pos \
--env.action_dim=12 \
--env.camera_keys=left_hand_camera_rgb,right_hand_camera_rgb,first_person_camera_rgb \
--policy.device=cuda \
--eval.batch_size=10 \
--eval.n_episodes=100
```
### Environment Configuration
Evaluation can be quickly launched by modifying the `robot`, `task`, and `layout` settings in the configuration file.
#### Full Configuration Options
```yml
# =========================
# Basic Settings
# =========================
disable_fabric: false
device: cuda:0
sensitivity: 1.0
step_hz: 50
enable_cameras: true
execute_mode: eval
episode_length_s: 20.0 # Episode length in seconds, increase if episodes timeout during eval
# =========================
# Robot Settings
# =========================
robot: DoublePiper-Abs # Robot type, DoublePiper-Abs, X7S-Abs, G1-Controller or G1-Controller-DecoupledWBC
robot_scale: 1.0
# =========================
# Task & Scene Settings
# =========================
task: L90K1PutTheBlackBowlOnThePlate # Task name
scene_backend: robocasa
task_backend: robocasa
debug_assets: null
layout: libero-1-1 # Layout and style ID
sources:
- objaverse
- lightwheel
- aigen_objs
object_projects: []
usd_simplify: false
seed: 42
# =========================
# Object Placement Retry Settings
# =========================
max_scene_retry: 4
max_object_placement_retry: 3
resample_objects_placement_on_reset: true
resample_robot_placement_on_reset: true
# =========================
# Replay Configuration Settings
# =========================
replay_cfgs:
add_camera_to_observation: true
render_resolution: [640, 480]
```
### See Also
- [LW-BenchHub GitHub](https://github.com/LightwheelAI/LW-BenchHub)
- [LW-BenchHub Documentation](https://docs.lightwheel.net/lw_benchhub/)
+2
View File
@@ -38,6 +38,8 @@ class EvalPipelineConfig:
seed: int | None = 1000 seed: int | None = 1000
# Rename map for the observation to override the image and state keys # Rename map for the observation to override the image and state keys
rename_map: dict[str, str] = field(default_factory=dict) 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: def __post_init__(self) -> None:
# HACK: We parse again the cli args here to get the pretrained path if there was one. # HACK: We parse again the cli args here to get the pretrained path if there was one.
+1 -1
View File
@@ -12,4 +12,4 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
from .configs import AlohaEnv, EnvConfig, PushtEnv # noqa: F401 from .configs import AlohaEnv, EnvConfig, HubEnvConfig, PushtEnv # noqa: F401
+85 -1
View File
@@ -13,7 +13,7 @@
# limitations under the License. # limitations under the License.
import abc import abc
from dataclasses import dataclass, field from dataclasses import dataclass, field, fields
from typing import Any from typing import Any
import draccus import draccus
@@ -68,6 +68,22 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
raise NotImplementedError() 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") @EnvConfig.register_subclass("aloha")
@dataclass @dataclass
class AlohaEnv(EnvConfig): class AlohaEnv(EnvConfig):
@@ -368,3 +384,71 @@ class MetaworldEnv(EnvConfig):
"obs_type": self.obs_type, "obs_type": self.obs_type,
"render_mode": self.render_mode, "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 {}
+40 -5
View File
@@ -20,11 +20,11 @@ import gymnasium as gym
from gymnasium.envs.registration import registry as gym_registry from gymnasium.envs.registration import registry as gym_registry
from lerobot.configs.policies import PreTrainedConfig 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.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.policies.xvla.configuration_xvla import XVLAConfig
from lerobot.processor import ProcessorStep 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 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: if isinstance(env_cfg, LiberoEnv) or "libero" in env_cfg.type:
preprocessor_steps.append(LiberoProcessorStep()) 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) preprocessor = PolicyProcessorPipeline(steps=preprocessor_steps)
postprocessor = PolicyProcessorPipeline(steps=postprocessor_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. 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. 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`. Default False must be set to True to import/exec hub `env.py`.
Raises: Raises:
ValueError: if n_envs < 1 ValueError: if n_envs < 1
ModuleNotFoundError: If the requested env package is not installed 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") # if user passed a hub id string (e.g., "username/repo", "username/repo@main:env.py")
# simplified: only support hub-provided `make_env` # simplified: only support hub-provided `make_env`
# TODO: (jadechoghari): deprecate string API and remove this check
if isinstance(cfg, str): 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 # _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 # import and surface clear import errors
module = _import_hub_module(local_file, repo_id) module = _import_hub_module(local_file, repo_id)
# call the hub-provided make_env # 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}} # normalize the return into {suite: {task_id: vec_env}}
return _normalize_hub_result(raw_result) 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: if n_envs < 1:
raise ValueError("`n_envs` must be at least 1") raise ValueError("`n_envs` must be at least 1")
+16 -4
View File
@@ -46,7 +46,7 @@ def _convert_nested_dict(d):
def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Tensor]: 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. """Convert environment observation to LeRobot format observation.
Args: Args:
observation: Dictionary of observation batches from a Gym vector environment. 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: if "robot_state" in observations:
return_observations[f"{OBS_STR}.robot_state"] = _convert_nested_dict(observations["robot_state"]) 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 return return_observations
def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]: 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) # (need to also refactor preprocess_observation and externalize normalization from policies)
policy_features = {} policy_features = {}
for key, ft in env_cfg.features.items(): 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 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. 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)`." 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 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]]: def _normalize_hub_result(result: Any) -> dict[str, dict[int, gym.vector.VectorEnv]]:
+76 -1
View File
@@ -18,7 +18,7 @@ from dataclasses import dataclass
import torch import torch
from lerobot.configs.types import PipelineFeatureType, PolicyFeature 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 from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
@@ -152,3 +152,78 @@ class LiberoProcessorStep(ObservationProcessorStep):
result[mask] = axis * angle.unsqueeze(1) result[mask] = axis * angle.unsqueeze(1)
return result 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)
+6 -1
View File
@@ -509,7 +509,12 @@ def eval_main(cfg: EvalPipelineConfig):
logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}") logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
logging.info("Making environment.") 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.") logging.info("Making policy.")
+2 -2
View File
@@ -17,7 +17,7 @@
import json import json
import tempfile import tempfile
from collections.abc import Callable from collections.abc import Callable
from dataclasses import dataclass from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@@ -1884,7 +1884,7 @@ class FeatureContractAddStep(ProcessorStep):
"""Adds a PolicyFeature""" """Adds a PolicyFeature"""
key: str = "a" key: str = "a"
value: PolicyFeature = PolicyFeature(type=FeatureType.STATE, shape=(1,)) value: PolicyFeature = field(default_factory=lambda: PolicyFeature(type=FeatureType.STATE, shape=(1,)))
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition return transition