envs/so101_pick_orange.py

import os
from pathlib import Path
from huggingface_hub import snapshot_download


def make_env(n_envs: int = 1, use_async_envs: bool = False):
    """
    Create vectorized environments for leisaac task.

    Args:
        n_envs: Number of parallel environments
        use_async_envs: Whether to use AsyncVectorEnv or SyncVectorEnv
            (not use now)

    Returns:
        ManagerBasedRLEnv or DirectRLEnv inherit from gym.Env implemented in IsaacLab
    """

    # download assets from huggingface first
    snapshot_download(repo_id='LightwheelAI/leisaac_env', revision=None, cache_dir=None)

    # set environment variable for leisaac asset
    assets_root = Path(__file__).parent.parent / 'assets'
    os.environ['LEISAAC_ASSETS_ROOT'] = str(assets_root)

    from isaaclab.app import AppLauncher
    _ = AppLauncher({"enable_cameras": True})

    from leisaac.utils.envhub_utils import export_env

    return export_env(
        'LeIsaac-SO101-PickOrange-v0', device="cuda:0", num_envs=n_envs
    )