mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-21 19:49:49 +00:00
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:
@@ -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
|
||||||
|
|||||||
@@ -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/)
|
||||||
@@ -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.
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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 {}
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|
||||||
|
|||||||
@@ -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]]:
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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.")
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user