env.py

"""Assembly bench EnvHub entry — LeRobot ``make_env`` for Isaac Lab Arena.

Docs: https://huggingface.co/docs/lerobot/en/envhub_isaaclab_arena

Variant selection (Arena ``--task``) is *not* ``cfg.task`` — that field is the
natural-language prompt on ``IsaaclabArenaEnv``. Pass the variant via kwargs:

    --env.kwargs='{"variant": "peg_round_M1_loose", "reward": "none"}'
"""

from __future__ import annotations

import importlib.util
import logging
import os
import shutil
import sys
from pathlib import Path
from typing import Any

import gymnasium as gym
import yaml

from huggingface_hub import hf_hub_download, snapshot_download

# Hub id for sibling downloads + full-repo snapshot (assets + package).
HUB_REPO_ID = os.environ.get("ASSEMBLY_BENCH_HUB", "lukasskellijs/env_assembly_bench")
EXAMPLE_ENVS = "example_envs.yaml"

# Default proprio / cameras match assembly_bench/contract.json (DROID 8-D).
_DEFAULT_STATE_KEYS = "joint_pos,gripper_pos,eef_pos,eef_quat"
_DEFAULT_CAMERA_KEYS = "front_cam_rgb,wrist_camera_rgb"
_DEFAULT_STATE_DIM = 15  # 7+1+3+4
_DEFAULT_ACTION_DIM = 8
_DEFAULT_VARIANT = "peg_round_M1_loose"
_DEFAULT_EMBODIMENT = "droid_abs_joint_pos_softmimic"
_MICROWAVE_PROMPT = "Reach out to the microwave and open it."  # IsaaclabArenaEnv default

_REPO_ROOT: str | None = None


def _materialize_hub_tree(root: Path) -> Path:
    """Copy a HF snapshot to a real directory (dereference blob symlinks).

    Hub stores files as symlinks into ``blobs/<hash>`` (no ``.usd`` suffix). USD /
    PhysX open the resolved path and then hang or fail on the extension-less blob.
    ``symlinks=False`` copies content under the original names.
    """
    sample = next((root / "assembly_bench" / "assets").rglob("*.usd"), None)
    if sample is None or not sample.is_symlink():
        return root

    cache = Path(
        os.environ.get(
            "ASSEMBLY_BENCH_MATERIALIZE_DIR",
            Path.home() / ".cache" / "env_assembly_bench" / "materialized",
        )
    )
    dest = cache / root.name
    marker = dest / ".materialized"
    if marker.is_file():
        return dest

    logging.info("Materializing Hub snapshot (deref blob symlinks) → %s", dest)
    if dest.exists():
        shutil.rmtree(dest)
    shutil.copytree(root, dest, symlinks=False)
    marker.write_text("ok\n", encoding="utf-8")
    return dest


def _ensure_repo_root() -> str:
    """Put this EnvHub checkout/snapshot on sys.path so ``assembly_bench`` imports."""
    global _REPO_ROOT
    if _REPO_ROOT is not None:
        return _REPO_ROOT

    # absolute() keeps the Hub snapshot path; resolve() follows blob symlinks.
    here = Path(__file__).absolute().parent
    if (here / "assembly_bench" / "environments" / "assembly" / "assembly.py").is_file():
        root = here
    else:
        # Loaded as a lone hub file — pull the full repo (USDs + package).
        root = Path(snapshot_download(repo_id=HUB_REPO_ID))

    root = _materialize_hub_tree(root)
    _REPO_ROOT = str(root)
    # Materialized tree must win over a still-on-path HF snapshot (blob symlinks).
    if _REPO_ROOT in sys.path:
        sys.path.remove(_REPO_ROOT)
    sys.path.insert(0, _REPO_ROOT)
    for name in list(sys.modules):
        if name == "assembly_bench" or name.startswith("assembly_bench."):
            del sys.modules[name]
    return _REPO_ROOT


def _download_hub_file(filename: str) -> str:
    root = _ensure_repo_root()
    local = Path(root) / filename
    if local.is_file():
        return str(local)
    return hf_hub_download(repo_id=HUB_REPO_ID, filename=filename)


def _download_and_import(filename: str):
    """Load a sibling module from the hub snapshot (absolute imports)."""
    local_path = _download_hub_file(filename)
    module_dir = os.path.dirname(local_path)
    if module_dir not in sys.path:
        sys.path.insert(0, module_dir)

    module_name = filename.replace(".py", "")
    with open(local_path) as f:
        content = f.read()
    content = content.replace("from .errors import", "from errors import")
    content = content.replace("from .isaaclab_env_wrapper import", "from isaaclab_env_wrapper import")

    spec = importlib.util.spec_from_file_location(module_name, local_path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    exec(compile(content, local_path, "exec"), module.__dict__)  # noqa: S102
    return module


try:
    from .errors import IsaacLabArenaCameraKeyError, IsaacLabArenaStateKeyError
    from .isaaclab_env_wrapper import IsaacLabEnvWrapper, cleanup_isaaclab
except ImportError:
    _errors = _download_and_import("errors.py")
    _wrapper = _download_and_import("isaaclab_env_wrapper.py")
    IsaacLabEnvWrapper = _wrapper.IsaacLabEnvWrapper
    cleanup_isaaclab = _wrapper.cleanup_isaaclab
    IsaacLabArenaCameraKeyError = _errors.IsaacLabArenaCameraKeyError
    IsaacLabArenaStateKeyError = _errors.IsaacLabArenaStateKeyError


def validate_config(
    env,
    state_keys: tuple[str, ...],
    camera_keys: tuple[str, ...],
    cfg_state_dim: int,
    cfg_action_dim: int,
) -> None:
    """Check observation keys and dims against the live Isaac Lab managers."""
    obs_manager = env.observation_manager
    active_terms = obs_manager.active_terms
    policy_terms = set(active_terms.get("policy", []))
    camera_terms = set(active_terms.get("camera_obs", []))

    missing_state = [k for k in state_keys if k not in policy_terms]
    if missing_state:
        raise IsaacLabArenaStateKeyError(missing_state, policy_terms)

    missing_cam = [k for k in camera_keys if k not in camera_terms]
    if missing_cam:
        raise IsaacLabArenaCameraKeyError(missing_cam, camera_terms)

    env_action_dim = env.action_space.shape[-1]
    if cfg_action_dim != env_action_dim:
        raise ValueError(f"action_dim mismatch: config={cfg_action_dim}, env={env_action_dim}")

    policy_dims = obs_manager.group_obs_dim.get("policy", [])
    policy_names = active_terms.get("policy", [])
    term_dims = dict(zip(policy_names, policy_dims, strict=False))

    expected_state_dim = 0
    for key in state_keys:
        if key in term_dims:
            shape = term_dims[key]
            dim = 1
            for s in shape if isinstance(shape, (tuple, list)) else [shape]:
                dim *= s
            expected_state_dim += dim

    if cfg_state_dim != expected_state_dim:
        raise ValueError(
            f"state_dim mismatch: config={cfg_state_dim}, "
            f"computed={expected_state_dim}. "
            f"Term dims: {term_dims}"
        )

    logging.info(f"Validated: state_keys={state_keys}, camera_keys={camera_keys}")


def resolve_environment_alias(environment: str) -> str:
    """Map ``--env.environment`` alias → dotted class path."""
    with open(_download_hub_file(EXAMPLE_ENVS)) as f:
        envs_mapping = yaml.safe_load(f)
    return f"{envs_mapping['repo']['base_module']}.{envs_mapping['repo']['envs'][environment]}"


def _cfg_get(cfg: Any, name: str, default: Any = None) -> Any:
    if cfg is None:
        return default
    return getattr(cfg, name, default)


def _create_isaaclab_env(config: dict, n_envs: int) -> dict[str, dict[int, gym.vector.VectorEnv]]:
    """Boot Kit, build AssemblyBench via Arena, wrap for EnvHub."""
    _ensure_repo_root()

    from isaaclab.app import AppLauncher

    if config.get("enable_pinocchio", False):
        import pinocchio  # noqa: F401

    config = {**config, "num_envs": n_envs}

    # Only Kit-relevant fields — extra Arena/EnvHub keys confuse nothing, but a
    # minimal Namespace keeps AppLauncher resolution predictable.
    logging.info("Launching IsaacLab simulation app...")
    app_launcher = AppLauncher(
        headless=bool(config.get("headless", True)),
        enable_cameras=bool(config.get("enable_cameras", True)),
        device=config.get("device") or "cuda:0",
    )

    # Package imports need the hub root on path (assets resolve from scene.py).
    from assembly_bench.environments.assembly.assembly import make_assembly_env
    from assembly_bench.environments.assembly.variants import VARIANTS

    variant = config.get("variant") or _DEFAULT_VARIANT
    if variant not in VARIANTS:
        raise ValueError(f"Unknown variant {variant!r}. Choices: {sorted(VARIANTS)}")
    # Confirm the EnvHub alias maps to this package (discoverability / typos).
    resolve_environment_alias(config.get("environment") or "assembly_bench")

    # Shared factory: Arena CLI defaults + reward/hdr wiring.
    raw_env = make_assembly_env(
        task=variant,
        num_envs=n_envs,
        embodiment=config.get("embodiment") or _DEFAULT_EMBODIMENT,
        hdr=config.get("hdr") or "asm_machine_shop",
        light_intensity=float(config.get("light_intensity") or 1500.0),
        reward=config.get("reward") or "none",
    )
    # Peel Gymnasium wrappers (OrderEnforcing, …) — managers live on the Isaac env.
    isaac_env = raw_env.unwrapped if hasattr(raw_env, "unwrapped") else raw_env

    # make_assembly_env already registers with render_mode="rgb_array"; Gymnasium's
    # OrderEnforcing wrapper exposes render_mode as a read-only property, so don't
    # assign on raw_env. IsaacLabEnvWrapper owns the EnvHub-facing value.
    render_mode = "rgb_array" if config.get("enable_cameras", False) else None

    state_keys = tuple(k.strip() for k in (config.get("state_keys") or "").split(",") if k.strip())
    camera_keys = tuple(k.strip() for k in (config.get("camera_keys") or "").split(",") if k.strip())
    try:
        validate_config(
            isaac_env,
            state_keys,
            camera_keys,
            config.get("state_dim", _DEFAULT_STATE_DIM),
            config.get("action_dim", _DEFAULT_ACTION_DIM),
        )
    except (IsaacLabArenaCameraKeyError, IsaacLabArenaStateKeyError, ValueError):
        cleanup_isaaclab(raw_env, app_launcher)
        raise
    except Exception as e:
        logging.error(f"Validation failed with unexpected error: {type(e).__name__}: {e}")
        cleanup_isaaclab(raw_env, app_launcher)
        raise

    # NL prompt: cfg.task unless still the Arena microwave default → use variant text.
    task_prompt = config.get("task")
    if not task_prompt or task_prompt == _MICROWAVE_PROMPT:
        task_prompt = VARIANTS[variant].instruction

    # 15 Hz control; fall back to the variant's authored duration (seconds).
    episode_length = config.get("episode_length")
    if episode_length in (None, 300):
        episode_length = int(VARIANTS[variant].episode_length_s * 15)

    environment = config.get("environment") or "assembly_bench"
    wrapped_env = IsaacLabEnvWrapper(
        isaac_env,
        episode_length=int(episode_length),
        task=task_prompt,
        render_mode=render_mode,
        simulation_app=app_launcher,
    )
    # Keep the Gymnasium outer wrapper alive: its __del__/close() would otherwise
    # tear down the Isaac env as soon as this frame returns (exit 0 mid-smoke).
    wrapped_env._gym_root = raw_env
    # OrderEnforcing.close() would close the Isaac env on GC; neutralize it.
    raw_env.close = lambda *a, **k: None  # type: ignore[method-assign]
    logging.info(
        f"Created: {environment} variant={variant} n_envs={wrapped_env.num_envs} "
        f"render_mode={render_mode}"
    )
    return {environment: {0: wrapped_env}}


def make_env(
    n_envs: int = 1,
    use_async_envs: bool = False,  # noqa: ARG001 — GPU batching; AsyncVectorEnv unused
    cfg: Any | None = None,
) -> dict[str, dict[int, gym.vector.VectorEnv]]:
    """EnvHub entry point. ``cfg`` is typically ``IsaaclabArenaEnv`` (or SimpleNamespace)."""
    if n_envs < 1:
        raise ValueError("`n_envs` must be at least 1")

    environment = _cfg_get(cfg, "environment", "assembly_bench")
    if environment is None:
        raise ValueError("No 'environment' specified (expected 'assembly_bench').")

    # IsaaclabArenaEnv.__post_init__ promotes kwargs (variant/reward/hdr) onto cfg.
    variant = _cfg_get(cfg, "variant", _DEFAULT_VARIANT)
    reward = _cfg_get(cfg, "reward", "none")
    hdr = _cfg_get(cfg, "hdr", "asm_machine_shop")
    light_intensity = _cfg_get(cfg, "light_intensity", 1500.0)

    config = {
        "environment": environment,
        "embodiment": _cfg_get(cfg, "embodiment", _DEFAULT_EMBODIMENT),
        "object": _cfg_get(cfg, "object", None),
        "mimic": _cfg_get(cfg, "mimic", False),
        "teleop_device": _cfg_get(cfg, "teleop_device", None),
        "seed": _cfg_get(cfg, "seed", 42),
        "device": _cfg_get(cfg, "device", "cuda:0"),
        "disable_fabric": _cfg_get(cfg, "disable_fabric", False),
        "enable_cameras": _cfg_get(cfg, "enable_cameras", True),
        "headless": _cfg_get(cfg, "headless", True),
        "enable_pinocchio": _cfg_get(cfg, "enable_pinocchio", False),
        "episode_length": _cfg_get(cfg, "episode_length", None),
        "state_dim": _cfg_get(cfg, "state_dim", _DEFAULT_STATE_DIM),
        "action_dim": _cfg_get(cfg, "action_dim", _DEFAULT_ACTION_DIM),
        "camera_height": _cfg_get(cfg, "camera_height", 720),
        "camera_width": _cfg_get(cfg, "camera_width", 1280),
        "video": _cfg_get(cfg, "video", False),
        "video_length": _cfg_get(cfg, "video_length", 100),
        "video_interval": _cfg_get(cfg, "video_interval", 200),
        "state_keys": _cfg_get(cfg, "state_keys", _DEFAULT_STATE_KEYS),
        "camera_keys": _cfg_get(cfg, "camera_keys", _DEFAULT_CAMERA_KEYS) or "",
        "task": _cfg_get(cfg, "task", None),
        "variant": variant,
        "reward": reward,
        "hdr": hdr,
        "light_intensity": light_intensity,
    }

    logging.info(
        f"EnvHub make_env: environment={environment}, variant={variant}, n_envs={n_envs}, "
        f"headless={config['headless']}, enable_cameras={config['enable_cameras']}"
    )
    return _create_isaaclab_env(config, n_envs)