File size: 4,064 Bytes
c2c0566 d9123b7 c2c0566 d9123b7 c2c0566 d9123b7 1b5c78c d9123b7 c2c0566 d9123b7 c2c0566 d9123b7 c2c0566 1b5c78c d9123b7 c2c0566 d9123b7 c2c0566 1b5c78c c2c0566 1b5c78c c2c0566 1b5c78c c2c0566 1b5c78c c2c0566 0240e37 1b5c78c c2c0566 0240e37 d9123b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """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()
|