hud_env.py

"""HUD environment that loads assembly_bench from EnvHub (Hub download path).

Standard HUD robot flow: ``env.gym(make_env)`` derives the contract/capability
and serves over the robot WebSocket. ``make_env`` pulls this repo from the Hub
(the same way LeRobot EnvHub does), then builds the Isaac Lab Arena sim.

Serve (isaac conda, GPU)::

    OMNI_KIT_ACCEPT_EULA=YES ACCEPT_EULA=Y PRIVACY_CONSENT=Y \\
      python -m hud.environment.server hud_env.py --port 8765

Scene config is the factory signature (``variant`` / ``num_envs`` / …); episodic
args (seed) go through ``sim.reset``.
"""

from __future__ import annotations

import importlib.util
import os
import sys
from pathlib import Path
from types import SimpleNamespace

from hud import Environment

HUB = os.environ.get("ASSEMBLY_BENCH_HUB", "lukasskellijs/env_assembly_bench")

# Authored contract next to this file (absolute — cwd-relative would silently drift).
_CONTRACT = Path(__file__).absolute().parent / "contract.json"


def make_env(
    variant: str = "peg_round_M1_loose",
    num_envs: int = int(os.environ.get("ASSEMBLY_NUM_ENVS", "1")),
    embodiment: str = os.environ.get("ASSEMBLY_EMBODIMENT", "droid_abs_joint_pos_softmimic"),
    reward: str = os.environ.get("ASSEMBLY_REWARD", "none"),
):
    """Download EnvHub package (if needed) and return the vectorized Arena env."""
    from huggingface_hub import snapshot_download

    # Fresh-enough Hub snapshot; local checkout of this repo also works when
    # ASSEMBLY_BENCH_HUB points at a path... snapshot always hits the Hub id.
    root = snapshot_download(repo_id=HUB)
    env_py = Path(root) / "env.py"
    if not env_py.is_file():
        raise FileNotFoundError(f"Hub repo {HUB} missing env.py (got {root})")

    # Prefer Hub package over any editable/local assembly_bench install.
    if root in sys.path:
        sys.path.remove(root)
    sys.path.insert(0, root)
    for name in list(sys.modules):
        if name == "assembly_bench" or name.startswith("assembly_bench."):
            del sys.modules[name]

    spec = importlib.util.spec_from_file_location("env_assembly_bench_hub_env", env_py)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)

    # IsaaclabArenaEnv-shaped config (no lerobot import required in the isaac env).
    cfg = SimpleNamespace(
        environment="assembly_bench",
        embodiment=embodiment,
        object=None,
        mimic=False,
        teleop_device=None,
        seed=42,
        device=os.environ.get("ASSEMBLY_DEVICE", "cuda:0"),
        disable_fabric=False,
        enable_cameras=True,
        headless=True,
        enable_pinocchio=False,
        episode_length=300,  # hub make_env replaces 300 with variant duration
        state_dim=15,
        action_dim=8,
        camera_height=720,
        camera_width=1280,
        video=False,
        video_length=100,
        video_interval=200,
        state_keys="joint_pos,gripper_pos,eef_pos,eef_quat",
        camera_keys="front_cam_rgb,wrist_camera_rgb",
        task=None,  # filled from variants.instruction
        variant=variant,
        reward=reward,
        hdr="asm_machine_shop",
        light_intensity=1500.0,
    )
    suites = module.make_env(n_envs=num_envs, use_async_envs=False, cfg=cfg)
    # {suite: {task_id: VectorEnv}} → the live vectorized env
    return next(iter(next(iter(suites.values())).values()))


env = Environment(name="assembly-bench")
sim = env.gym(
    make_env,
    contract=_CONTRACT,
    variant="peg_round_M1_loose",
)


@env.template(id="assembly")
async def assembly(
    variant: str = "peg_round_M1_loose",
    seed: int = 0,
    num_envs: int | None = None,
    embodiment: str | None = None,
    reward: str | None = None,
):
    """One assembly episode. Env-defining args rebuild the scene when they change."""
    reset_kw: dict = {"variant": variant, "seed": seed}
    if num_envs is not None:
        reset_kw["num_envs"] = num_envs
    if embodiment is not None:
        reset_kw["embodiment"] = embodiment
    if reward is not None:
        reset_kw["reward"] = reward
    ep = await sim.reset(**reset_kw)
    yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}}
    yield await sim.result(token=ep["token"])