test_isaaclab_arena_simple.py

#!/usr/bin/env python
"""
快速测试 IsaacLab Arena 环境
运行随机动作，验证环境是否正常工作
"""

import os
os.environ["ACCEPT_EULA"] = "Y"
os.environ["PRIVACY_CONSENT"] = "Y"

# 设置代理（如果需要）
# 如果环境变量未设置，使用本地代理
if "http_proxy" not in os.environ and "HTTPS_PROXY" not in os.environ:
    os.environ["http_proxy"] = "http://127.0.0.1:7890"
    os.environ["https_proxy"] = "http://127.0.0.1:7890"
    os.environ["HTTP_PROXY"] = "http://127.0.0.1:7890"
    os.environ["HTTPS_PROXY"] = "http://127.0.0.1:7890"

import torch
from lerobot.envs.factory import make_env
from lerobot.envs.configs import IsaaclabArenaEnv
import tqdm

def main():
    print("=" * 60)
    print("IsaacLab Arena 快速测试")
    print("=" * 60)
    
    # 创建环境配置
    env_cfg = IsaaclabArenaEnv(
        hub_path="nvidia/isaaclab-arena-envs",
        environment="gr1_microwave",
        embodiment="gr1_pink",
        object="mustard_bottle",
        headless=False,  # 设置为 True 则无 GUI
        enable_cameras=True,
        state_keys="robot_joint_pos",
        camera_keys="robot_pov_cam_rgb",
        state_dim=54,
        action_dim=36,
        device="cuda:0",
        episode_length=100,
    )
    
    print(f"\n环境配置:")
    print(f"  - Environment: {env_cfg.environment}")
    print(f"  - Embodiment: {env_cfg.embodiment}")
    print(f"  - Object: {env_cfg.object}")
    print(f"  - Headless: {env_cfg.headless}")
    
    # 创建环境
    print("\n正在加载环境（这可能需要 30-60 秒）...")
    envs_dict = make_env(env_cfg, n_envs=1, trust_remote_code=True)
    suite_name = next(iter(envs_dict))
    env = envs_dict[suite_name][0]
    
    print(f"✓ 环境加载成功: {suite_name}")
    print(f"  - 动作空间: {env.action_space}")
    print(f"  - 观察空间键: {list(env.observation_space.keys())[:5]}...")
    
    # 运行随机动作
    print("\n开始运行随机动作（100 步）...")
    obs, info = env.reset()
    
    for step in tqdm.tqdm(range(1000)):
        action = env.action_space.sample()
        obs, reward, terminated, truncated, info = env.step(action)
        
        if terminated.any() or truncated.any():
            obs, info = env.reset()
    
    env.close()
    print("\n✓ 测试完成！环境工作正常。")

if __name__ == "__main__":
    main()