scripts/smoke_hud.py

"""HUD integration smoke: Hub → make_env → ``hud.environment.robot.wrap`` → step.

This is the loop-owning HUD path (traces stream when HUD_API_KEY is set).
For the served path, run::

    python -m hud.environment.server hud_env.py --port 8765

Usage (isaac conda, free GPU)::

    ACCEPT_EULA=Y PRIVACY_CONSENT=Y OMNI_KIT_ACCEPT_EULA=YES \\
      ASSEMBLY_BENCH_HUB=lukasskellijs/env_assembly_bench \\
      python scripts/smoke_hud.py --steps 3
"""

from __future__ import annotations

import argparse
import importlib.util
import logging
import os
import sys
import tempfile
from pathlib import Path
from types import SimpleNamespace

ROOT = Path(__file__).resolve().parents[1]


def _hub_make_env(variant: str, num_envs: int):
    from huggingface_hub import snapshot_download

    hub = os.environ.get("ASSEMBLY_BENCH_HUB", "lukasskellijs/env_assembly_bench")
    root = snapshot_download(repo_id=hub)
    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]

    env_py = Path(root) / "env.py"
    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)

    cfg = SimpleNamespace(
        environment="assembly_bench",
        embodiment="droid_abs_joint_pos_softmimic",
        object=None,
        mimic=False,
        teleop_device=None,
        seed=0,
        device="cuda:0",
        disable_fabric=False,
        enable_cameras=True,
        headless=True,
        enable_pinocchio=False,
        episode_length=50,
        state_dim=15,
        action_dim=8,
        camera_height=720,
        camera_width=1280,
        video=False,
        video_length=10,
        video_interval=15,
        state_keys="joint_pos,gripper_pos,eef_pos,eef_quat",
        camera_keys="front_cam_rgb,wrist_camera_rgb",
        task=None,
        variant=variant,
        reward="none",
        hdr="asm_machine_shop",
        light_intensity=1500.0,
    )
    suites = module.make_env(n_envs=num_envs, use_async_envs=False, cfg=cfg)
    env = next(iter(next(iter(suites.values())).values()))
    return env, Path(root)


def main() -> None:
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
    p = argparse.ArgumentParser()
    p.add_argument("--variant", default="peg_round_M1_loose")
    p.add_argument("--steps", type=int, default=3)
    p.add_argument("--num-envs", type=int, default=1)
    args = p.parse_args()

    from hud.environment.robot import wrap

    logging.info("loading EnvHub package + building Arena env...")
    raw, hub_root = _hub_make_env(args.variant, args.num_envs)
    # Writable contract path (cwd may be read-only in containers).
    contract = Path(tempfile.gettempdir()) / "env_assembly_bench_hud_contract.json"
    if not contract.is_file():
        bundled = hub_root / "contract.json"
        if bundled.is_file():
            contract.write_text(bundled.read_text())
    # HUD wrap: contract probe + optional telemetry job for the rollout.
    env = wrap(raw, job="envhub-assembly-smoke", contract=str(contract))
    obs, info = env.reset(seed=0)
    print("HUD reset ok", flush=True)
    logging.info("HUD reset ok; obs type=%s", type(obs))
    for i in range(args.steps):
        action = env.action_space.sample()
        obs, reward, terminated, truncated, info = env.step(action)
        print(f"HUD step {i}", flush=True)
        logging.info("HUD step %d reward=%s", i, reward)
        if getattr(terminated, "any", lambda: bool(terminated))() or getattr(
            truncated, "any", lambda: bool(truncated)
        )():
            obs, info = env.reset()
    # Kit's SimulationApp.close() terminates the process — log success first.
    print("HUD smoke ok", flush=True)
    logging.info("HUD smoke ok")
    env.close()


if __name__ == "__main__":
    main()