diff --git a/isaaclab_arena/environments/arena_env_graph_spec.py b/isaaclab_arena/environments/arena_env_graph_spec.py index 0dc3d6db6..03221af86 100644 --- a/isaaclab_arena/environments/arena_env_graph_spec.py +++ b/isaaclab_arena/environments/arena_env_graph_spec.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Self -from pydantic import BaseModel, Field, SerializeAsAny, field_validator, model_validator +from pydantic import BaseModel, Field, SerializeAsAny, ValidationInfo, field_validator, model_validator from isaaclab_arena.assets.registries import TaskRegistry from isaaclab_arena.environments.arena_env_graph_types import ( @@ -158,8 +158,10 @@ class ArenaEnvInitialGraphSpec(ArenaEnvGraphSpecBase): initial_state_spec: ArenaEnvGraphStateSpec @model_validator(mode="after") - def validate(self) -> Self: + def validate(self, info: ValidationInfo) -> Self: """Check unique IDs, constraint references, and spatial constraint shapes.""" + if info.context and info.context.get("skip_registry"): + return self assert_unique_ids(self.nodes, [], [self.initial_state_spec]) assert_constraint_references(self.nodes, [self.initial_state_spec]) assert_spatial_constraint_shapes([self.initial_state_spec]) diff --git a/isaaclab_arena/environments/arena_env_graph_types.py b/isaaclab_arena/environments/arena_env_graph_types.py index 740f6a642..38e121175 100644 --- a/isaaclab_arena/environments/arena_env_graph_types.py +++ b/isaaclab_arena/environments/arena_env_graph_types.py @@ -16,7 +16,7 @@ from enum import Enum from typing import Any -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, Field, ValidationInfo, field_validator, model_validator from isaaclab_arena.assets.object_type import ObjectType from isaaclab_arena.assets.registries import ObjectRelationLibraryRegistry, TaskRegistry @@ -115,7 +115,9 @@ class TaskSpec(BaseModel): @field_validator("kind") @classmethod - def _validate_registered_task_type(cls, value: str) -> str: + def _validate_registered_task_type(cls, value: str, info: ValidationInfo) -> str: + if info.context and info.context.get("skip_registry"): + return value registry = TaskRegistry() assert registry.is_registered(value), f"Unknown task kind '{value}'" return value @@ -169,7 +171,9 @@ class SpatialRelationSpec(BaseModel): ) @model_validator(mode="after") - def _validate_kind_and_arity(self) -> SpatialRelationSpec: + def _validate_kind_and_arity(self, info: ValidationInfo) -> SpatialRelationSpec: + if info.context and info.context.get("skip_registry"): + return self registry = ObjectRelationLibraryRegistry() assert registry.is_registered(self.kind), f"Unknown relation kind '{self.kind}'" relation_cls = registry.get_object_relation_by_name(self.kind) diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/dashboard.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/dashboard.py index 39a135dd0..ee08bd731 100644 --- a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/dashboard.py +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/dashboard.py @@ -17,9 +17,10 @@ from isaaclab_arena_examples.agentic_environment_generation.review_gui.render.styles import DASHBOARD_CSS -def render_dashboard_html(spec: ArenaEnvInitialGraphSpec) -> str: +def render_dashboard_html(spec: ArenaEnvInitialGraphSpec, thumbnails: dict[str, bytes] | None = None) -> str: """Render the self-contained review dashboard HTML for ``spec``.""" initial_state = spec.initial_state_spec + thumbnails = thumbnails or {} return f""" @@ -36,7 +37,7 @@ def render_dashboard_html(spec: ArenaEnvInitialGraphSpec) -> str:

Nodes

-
{render_node_cards(spec)}
+
{render_node_cards(spec, thumbnails)}

Spatial graph (initial state: {html_lib.escape(initial_state.id)})

diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/panels.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/panels.py index 6ab00637d..6bc253c9e 100644 --- a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/panels.py +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/panels.py @@ -10,9 +10,7 @@ from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvInitialGraphSpec from isaaclab_arena.environments.arena_env_graph_types import ArenaEnvGraphNodeSpec, ArenaEnvGraphStateSpec -from isaaclab_arena_examples.agentic_environment_generation.review_gui.render.thumbnails import ( - render_placeholder_thumbnail, -) +from isaaclab_arena_examples.agentic_environment_generation.review_gui.render.thumbnails import render_node_thumbnail def render_unary_constraints(state: ArenaEnvGraphStateSpec) -> str: @@ -63,16 +61,17 @@ def render_tasks_table(spec: ArenaEnvInitialGraphSpec) -> str: ) -def render_node_cards(spec: ArenaEnvInitialGraphSpec) -> str: +def render_node_cards(spec: ArenaEnvInitialGraphSpec, thumbnails: dict[str, bytes] | None = None) -> str: """Render one card per graph node for the dashboard nodes panel.""" - return "\n".join(render_node_card(node) for node in spec.nodes) + thumbnails = thumbnails or {} + return "\n".join(render_node_card(node, thumbnails.get(node.id)) for node in spec.nodes) -def render_node_card(node: ArenaEnvGraphNodeSpec) -> str: - """Render a single node card with placeholder thumbnail and YAML dump.""" +def render_node_card(node: ArenaEnvGraphNodeSpec, png_bytes: bytes | None = None) -> str: + """Render a single node card with USD snapshot or placeholder thumbnail and YAML dump.""" node_dict = node.model_dump(mode="json", exclude_none=True) node_yaml = yaml.safe_dump(node_dict, sort_keys=False).rstrip() - thumb = render_placeholder_thumbnail(node) + thumb = render_node_thumbnail(node, png_bytes) return f"""
{thumb}
diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/thumbnails.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/thumbnails.py index f2742ee8f..4b2a5eb23 100644 --- a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/thumbnails.py +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/thumbnails.py @@ -5,14 +5,22 @@ from __future__ import annotations +import base64 import html as html_lib from isaaclab_arena.environments.arena_env_graph_types import ArenaEnvGraphNodeSpec -# TODO(qianl): Replace placeholder thumbnails with sim-rendered snapshots . -def render_placeholder_thumbnail(node: ArenaEnvGraphNodeSpec) -> str: - """Per-node placeholder thumbnail — two-letter initial.""" +def render_node_thumbnail(node: ArenaEnvGraphNodeSpec, png_bytes: bytes | None = None) -> str: + """Per-node thumbnail: USD capture if available, else two-letter placeholder.""" + if png_bytes: + b64 = base64.b64encode(png_bytes).decode("ascii") + return ( + '
' + f'{html_lib.escape(node.name)} thumbnail' + f'{html_lib.escape(node.name)}' + "
" + ) initial = (node.name[:2] if node.name else "?").upper() return f"""
{html_lib.escape(initial)} diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/simapp_sidecar.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/simapp_sidecar.py new file mode 100644 index 000000000..2ba7a7d65 --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/simapp_sidecar.py @@ -0,0 +1,285 @@ +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Long-lived ``SimulationApp`` host process for the live review editor. + +Boots Kit's ``SimulationApp`` once (with ``--viz kit``) on *its own* main thread and serves +validation and thumbnail-render requests over a newline-delimited JSON-RPC +pipe on stdin/stdout. The parent (``streamlit_ui.py`` running inside +Streamlit) spawns exactly one of these and reuses it for the entire server +lifetime via +:class:`isaaclab_arena_examples.agentic_environment_generation.review_gui.simapp_sidecar_client.SimAppSidecar`. + +Why a sidecar and not an in-process ``SimulationApp``: + +* ``signal.signal`` only works in the main thread. Streamlit's + ``ScriptRunner`` runs the script in a worker thread, so SimApp's signal + setup raises ``ValueError("signal only works in main thread …")``. +* ``omni.usd.UsdContext`` is process-singleton AND can't tolerate cross- + thread driving from Streamlit reruns — driving it from worker threads + triggers ``[Error] [omni.usd] UsdContext busy`` and the open_stage call + fails. A dedicated process with serialized request handling avoids both. + +Protocol (newline-delimited JSON over stdin/stdout): + + Ready handshake (sent by sidecar on boot before reading any request): + {"ready": true} # SimApp boot succeeded + {"ready": false, "error": "..."} # boot failed; sidecar exits + + Requests: + {"cmd": "ping"} + → {"ok": true} + + {"cmd": "validate_spec", "yaml_text": "..."} + → {"ok": true, "spec_dict": {...}} + (full :class:`ArenaEnvInitialGraphSpec` validation including registry + lookups — runs in the sidecar where registries are already warm) + + {"cmd": "render_spec", "yaml_text": "..."} + → {"ok": true, "paths": {"node_id": "/abs/path/to.png", ...}, + "errors": [{"node_id": "...", "error": "..."}]} + (paths are absolute filesystem paths on the disk cache. The PNGs + themselves stay on disk — the parent reads them itself.) + + {"cmd": "build_catalogues"} + → {"ok": true, "asset_catalogue": {...}, "relation_catalogue": {...}, + "task_catalogue": {...}} + (registry vocabulary for :meth:`EnvironmentGenerationAgent.fetch_intent_from_prompt`) + + {"cmd": "compile_intent", "intent_dict": {...}} + → {"ok": true, "spec_dict": {...}, "has_resolution_errors": bool, + "trace": [{"stage": "...", "query": "...", ...}]} + (validates :class:`EnvironmentIntentSpec` and compiles to initial graph spec) + + {"cmd": "shutdown"} + → {"ok": true} # sidecar exits cleanly after replying + + Parent EOF on stdin (parent process died) triggers the same graceful + shutdown as the explicit "shutdown" cmd. + +stdout multiplexing: + +Kit writes a lot to stdout (warnings, replicator startup, etc.) and that +would corrupt the JSON channel the parent reads. We dup the original +stdout fd before touching Kit, then redirect Kit's stdout to stderr — +JSON replies go out through the saved fd; everything else from Kit +appears on the user's terminal via inherited stderr. +""" + +from __future__ import annotations + +import contextlib +import json +import os +import signal +import sys +import traceback +import yaml +from pathlib import Path +from typing import Any + +_JSON_FD = os.dup(1) +os.dup2(2, 1) +sys.stdout = sys.stderr + + +def _send(payload: dict[str, Any]) -> None: + """Write one JSON line to the parent on the saved stdout fd.""" + data = (json.dumps(payload) + "\n").encode("utf-8") + os.write(_JSON_FD, data) + + +def _install_signal_handlers() -> None: + def _exit(signum, _frame): + raise SystemExit(0) + + signal.signal(signal.SIGTERM, _exit) + signal.signal(signal.SIGINT, _exit) + + +def _serve() -> int: + """Boot SimApp, hand-shake with the parent, then service requests.""" + _install_signal_handlers() + + try: + from isaaclab_arena_examples.agentic_environment_generation.review_gui.thumbnail_render import ( # noqa: PLC0415 + _launch_simulation_app, + ) + except Exception as exc: + _send({"ready": False, "error": f"import failed: {exc}", "traceback": traceback.format_exc()}) + return 1 + + app = _launch_simulation_app() + if app is None: + _send({"ready": False, "error": "SimulationApp launch returned None"}) + return 1 + + from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvInitialGraphSpec # noqa: PLC0415 + from isaaclab_arena_examples.agentic_environment_generation.review_gui.thumbnail_render import ( # noqa: PLC0415 + _render_thumbnails_with_app, + ) + + _send({"ready": True}) + + try: + for raw_line in sys.stdin: + line = raw_line.strip() + if not line: + continue + try: + req = json.loads(line) + except json.JSONDecodeError as exc: + _send({"ok": False, "error": f"bad json: {exc}"}) + continue + + cmd = req.get("cmd") + if cmd == "shutdown": + _send({"ok": True}) + return 0 + + if cmd == "ping": + _send({"ok": True}) + continue + + if cmd == "validate_spec": + _send(_handle_validate_spec(req, ArenaEnvInitialGraphSpec)) + continue + + if cmd == "render_spec": + _send(_handle_render_spec(app, req, _render_thumbnails_with_app, ArenaEnvInitialGraphSpec)) + continue + + if cmd == "build_catalogues": + _send(_handle_build_catalogues()) + continue + + if cmd == "compile_intent": + _send(_handle_compile_intent(req)) + continue + + _send({"ok": False, "error": f"unknown cmd: {cmd!r}"}) + + return 0 + finally: + with contextlib.suppress(Exception): + app.close() + + +def _handle_validate_spec(req: dict[str, Any], spec_cls) -> dict[str, Any]: + """Parse and fully validate YAML as an initial graph spec.""" + yaml_text = req.get("yaml_text") + if not isinstance(yaml_text, str): + return {"ok": False, "error": "validate_spec requires string 'yaml_text'"} + + try: + raw = yaml.safe_load(yaml_text) + except Exception as exc: + return {"ok": False, "error": f"yaml parse failed: {exc}", "traceback": traceback.format_exc()} + + if raw is None: + return {"ok": False, "error": "YAML is empty"} + if not isinstance(raw, dict): + return {"ok": False, "error": f"expected mapping, got {type(raw).__name__}"} + + try: + spec = spec_cls.model_validate(raw) + except Exception as exc: + return {"ok": False, "error": f"spec validation failed: {exc}", "traceback": traceback.format_exc()} + + return {"ok": True, "spec_dict": spec.to_dict()} + + +def _handle_build_catalogues() -> dict[str, Any]: + """Return asset/relation/task catalogues for the env-generation agent.""" + from dataclasses import asdict # noqa: PLC0415 + + from isaaclab_arena.agentic_environment_generation.environment_generation_agent import ( # noqa: PLC0415 + build_asset_catalogue, + build_relation_catalogue, + build_task_catalogue, + ) + + try: + asset_catalogue = build_asset_catalogue() + relation_catalogue = build_relation_catalogue() + task_catalogue = build_task_catalogue() + except Exception as exc: + return {"ok": False, "error": f"catalogue build failed: {exc}", "traceback": traceback.format_exc()} + + return { + "ok": True, + "asset_catalogue": asdict(asset_catalogue), + "relation_catalogue": { + "relations": [asdict(entry) for entry in relation_catalogue.relations], + }, + "task_catalogue": { + "tasks": [asdict(entry) for entry in task_catalogue.tasks], + }, + } + + +def _handle_compile_intent(req: dict[str, Any]) -> dict[str, Any]: + """Validate an EnvironmentIntentSpec and compile it to an initial graph spec.""" + from dataclasses import asdict # noqa: PLC0415 + + from isaaclab_arena.agentic_environment_generation.environment_intent_spec import ( # noqa: PLC0415 + EnvironmentIntentSpec, + ) + from isaaclab_arena.agentic_environment_generation.intent_compiler import IntentCompiler # noqa: PLC0415 + + intent_dict = req.get("intent_dict") + if not isinstance(intent_dict, dict): + return {"ok": False, "error": "compile_intent requires mapping 'intent_dict'"} + + try: + intent = EnvironmentIntentSpec.model_validate(intent_dict) + compiler = IntentCompiler() + spec = compiler.compile(intent) + except Exception as exc: + return {"ok": False, "error": f"intent compile failed: {exc}", "traceback": traceback.format_exc()} + + return { + "ok": True, + "spec_dict": spec.to_dict(), + "has_resolution_errors": compiler.has_resolution_errors, + "trace": [asdict(event) for event in compiler.trace], + "reasoning": intent.reasoning, + } + + +def _handle_render_spec( + app, + req: dict[str, Any], + render_fn, + spec_cls, +) -> dict[str, Any]: + """Parse the spec, run thumbnail rendering, marshal the response.""" + yaml_text = req.get("yaml_text") + if not isinstance(yaml_text, str): + return {"ok": False, "error": "render_spec requires string 'yaml_text'"} + + try: + spec = spec_cls.from_dict(yaml.safe_load(yaml_text)) + except Exception as exc: + return {"ok": False, "error": f"spec parse failed: {exc}", "traceback": traceback.format_exc()} + + try: + paths: dict[str, Path] = render_fn(app, spec) + except Exception as exc: + return {"ok": False, "error": f"render failed: {exc}", "traceback": traceback.format_exc()} + + return { + "ok": True, + "paths": {node_id: str(p) for node_id, p in paths.items()}, + "errors": [], + } + + +def main() -> int: + return _serve() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/simapp_sidecar_client.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/simapp_sidecar_client.py new file mode 100644 index 000000000..d54a01462 --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/simapp_sidecar_client.py @@ -0,0 +1,200 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""JSON-RPC client for the persistent :mod:`isaaclab_arena_examples.agentic_environment_generation.review_gui.simapp_sidecar` process.""" + +from __future__ import annotations + +import contextlib +import json +import os +import subprocess +import sys +import threading +import yaml +from pathlib import Path +from typing import Any + +from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvInitialGraphSpec + + +class SimAppSidecarError(RuntimeError): + """Raised when the SimApp sidecar process can't fulfil a request.""" + + +class SimAppSidecar: + """Long-lived Kit/SimApp host process exposed as a validation and render service.""" + + def __init__(self, *, boot_timeout_s: float = 180.0, shutdown_timeout_s: float = 10.0) -> None: + self._proc: subprocess.Popen | None = None + self._lock = threading.Lock() + self._boot_timeout_s = boot_timeout_s + self._shutdown_timeout_s = shutdown_timeout_s + + def start(self) -> None: + """Spawn the sidecar process and wait for its ``{"ready": true}`` handshake.""" + if self._proc is not None and self._proc.poll() is None: + return + + cmd = [ + sys.executable, + "-m", + "isaaclab_arena_examples.agentic_environment_generation.review_gui.simapp_sidecar", + ] + self._proc = subprocess.Popen( + cmd, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=None, + text=True, + bufsize=1, + env=os.environ.copy(), + ) + + line = self._readline_or_die() + try: + msg = json.loads(line) + except json.JSONDecodeError as exc: + self._terminate() + raise SimAppSidecarError(f"Sidecar emitted non-JSON handshake: {line!r}") from exc + + if not msg.get("ready"): + self._terminate() + raise SimAppSidecarError( + f"Sidecar boot failed: {msg.get('error', 'unknown error')}\n{msg.get('traceback', '')}" + ) + + def is_alive(self) -> bool: + return self._proc is not None and self._proc.poll() is None + + def close(self) -> None: + """Send ``shutdown``, then terminate/kill if the process doesn't exit.""" + proc = self._proc + if proc is None: + return + self._proc = None + + if proc.poll() is None: + with contextlib.suppress(Exception): + proc.stdin.write(json.dumps({"cmd": "shutdown"}) + "\n") + proc.stdin.flush() + try: + proc.wait(timeout=self._shutdown_timeout_s) + except subprocess.TimeoutExpired: + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + with contextlib.suppress(Exception): + proc.wait(timeout=5) + + with contextlib.suppress(Exception): + if proc.stdin: + proc.stdin.close() + with contextlib.suppress(Exception): + if proc.stdout: + proc.stdout.close() + + def validate_yaml_text(self, yaml_text: str) -> dict[str, Any]: + """Run full spec validation (including registry lookups) in the sidecar.""" + if not self.is_alive(): + raise SimAppSidecarError("SimApp sidecar is not running — start it first") + + with self._lock: + return self._request({"cmd": "validate_spec", "yaml_text": yaml_text}) + + def build_catalogues(self) -> dict[str, Any]: + """Build asset/relation/task catalogues from warm registries in the sidecar.""" + if not self.is_alive(): + raise SimAppSidecarError("SimApp sidecar is not running — start it first") + + with self._lock: + return self._request({"cmd": "build_catalogues"}) + + def compile_intent(self, intent_dict: dict[str, Any]) -> dict[str, Any]: + """Validate and compile an EnvironmentIntentSpec in the sidecar.""" + if not self.is_alive(): + raise SimAppSidecarError("SimApp sidecar is not running — start it first") + + with self._lock: + return self._request({"cmd": "compile_intent", "intent_dict": intent_dict}) + + def render_spec(self, spec: ArenaEnvInitialGraphSpec) -> dict[str, bytes]: + """Ask the sidecar to render thumbnails for ``spec``.""" + if not self.is_alive(): + raise SimAppSidecarError("SimApp sidecar is not running — start it first") + + yaml_text = yaml.safe_dump(spec.to_dict(), sort_keys=False) + + with self._lock: + response = self._request({"cmd": "render_spec", "yaml_text": yaml_text}) + + if not response.get("ok"): + raise SimAppSidecarError( + f"sidecar render failed: {response.get('error', 'unknown')}\n{response.get('traceback', '')}" + ) + + paths: dict[str, str] = response.get("paths", {}) or {} + results: dict[str, bytes] = {} + for node_id, path_str in paths.items(): + path = Path(path_str) + if path.exists() and path.stat().st_size > 0: + results[node_id] = path.read_bytes() + else: + print( + f"[review_gui] sidecar reported {node_id} -> {path_str} but file is missing.", + file=sys.stderr, + ) + return results + + def ping(self) -> bool: + """Cheap liveness check round-trip — returns True on a healthy reply.""" + if not self.is_alive(): + return False + with self._lock: + try: + response = self._request({"cmd": "ping"}) + except SimAppSidecarError: + return False + return bool(response.get("ok")) + + def _request(self, payload: dict[str, Any]) -> dict[str, Any]: + assert self._proc is not None and self._proc.stdin is not None and self._proc.stdout is not None + line = json.dumps(payload) + "\n" + try: + self._proc.stdin.write(line) + self._proc.stdin.flush() + except BrokenPipeError as exc: + raise SimAppSidecarError("sidecar pipe closed unexpectedly") from exc + + reply_line = self._readline_or_die() + try: + return json.loads(reply_line) + except json.JSONDecodeError as exc: + raise SimAppSidecarError(f"sidecar replied with non-JSON: {reply_line!r}") from exc + + def _readline_or_die(self) -> str: + assert self._proc is not None and self._proc.stdout is not None + line = self._proc.stdout.readline() + if line == "": + exit_code = self._proc.poll() + raise SimAppSidecarError( + f"sidecar exited prematurely (exit code: {exit_code}). " + "See its stderr output above for the underlying cause." + ) + return line + + def _terminate(self) -> None: + if self._proc is None: + return + with contextlib.suppress(Exception): + self._proc.terminate() + try: + self._proc.wait(timeout=5) + except subprocess.TimeoutExpired: + with contextlib.suppress(Exception): + self._proc.kill() + self._proc = None diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py index c3b5b0a59..1a680cc7d 100644 --- a/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py @@ -14,6 +14,10 @@ /isaac-sim/python.sh -m isaaclab_arena_examples.agentic_environment_generation.gui_runner \\ --yaml isaaclab_arena/tests/test_data/pick_and_place_maple_table_init_env_graph.yaml +Registry lookups (task kinds, relation kinds, asset USD paths) run in a persistent +SimApp sidecar so YAML re-validation stays fast after the first ~30s sidecar boot. +Node thumbnails are rendered live by the same sidecar. + Natural-language generation calls the LLM from Streamlit (``NV_API_KEY``) and compiles the returned intent in-process with :class:`IntentCompiler`. """ @@ -21,6 +25,7 @@ from __future__ import annotations import argparse +import atexit import traceback import yaml from dataclasses import asdict, dataclass @@ -42,9 +47,15 @@ from isaaclab_arena.agentic_environment_generation.intent_compiler import IntentCompiler from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvInitialGraphSpec from isaaclab_arena_examples.agentic_environment_generation.review_gui.render.dashboard import render_dashboard_html +from isaaclab_arena_examples.agentic_environment_generation.review_gui.simapp_sidecar_client import ( + SimAppSidecar, + SimAppSidecarError, +) _IFRAME_HEIGHT_PX = 1100 +_SKIP_REGISTRY_CONTEXT: dict[str, Any] = {"skip_registry": True} + _DEFAULT_GENERATION_PROMPT = ( "franka pick up avocado from the maple table and place it into a bowl on the table. " "there are other veggies on the table as distractor" @@ -59,6 +70,61 @@ _DEFAULT_SAVE_PATH = "isaaclab_arena_environments/llm_generated/generated_spec.yaml" +@st.cache_resource(show_spinner="Booting Isaac Sim sidecar (≈30s first run, cached afterwards)…") +def _get_simapp_sidecar() -> SimAppSidecar | None: + """Spawn the SimApp sidecar once per Streamlit server process.""" + sidecar = SimAppSidecar() + try: + sidecar.start() + except SimAppSidecarError as exc: + print(f"[review_gui] SimApp sidecar failed to start: {exc}", flush=True) + return None + + atexit.register(sidecar.close) + return sidecar + + +def _ensure_sidecar() -> SimAppSidecar | None: + """Return a healthy sidecar, re-spawning if the cached one died.""" + sidecar = _get_simapp_sidecar() + if sidecar is not None and sidecar.is_alive(): + return sidecar + if sidecar is not None: + sidecar.close() + _get_simapp_sidecar.clear() + return _get_simapp_sidecar() + + +def _spec_from_sidecar_dict(spec_dict: dict[str, Any]) -> ArenaEnvInitialGraphSpec: + """Rebuild a validated spec locally without registry imports.""" + return ArenaEnvInitialGraphSpec.model_validate(spec_dict, context=_SKIP_REGISTRY_CONTEXT) + + +def _render_with_thumbnails(spec: ArenaEnvInitialGraphSpec) -> str: + """Render review HTML, asking the sidecar for live USD thumbnails when available.""" + sidecar = _ensure_sidecar() + if sidecar is None: + st.warning( + "Isaac Sim sidecar is unavailable — showing placeholder thumbnails. " + "Check the terminal where you launched the server for the underlying error.", + icon="⚠️", + ) + return render_dashboard_html(spec) + + try: + thumbnails = sidecar.render_spec(spec) + except SimAppSidecarError as exc: + st.error( + f"Sidecar render failed; showing placeholder thumbnails.\n\n```\n{exc}\n```", + icon="🛑", + ) + with st.spinner("Resetting the SimApp sidecar…"): + _get_simapp_sidecar.clear() + return render_dashboard_html(spec) + + return render_dashboard_html(spec, thumbnails=thumbnails if thumbnails else None) + + @dataclass class ValidationResult: """Outcome of parsing and validating YAML text as an initial graph spec.""" @@ -92,7 +158,7 @@ def _get_catalogue_bundle() -> CatalogueBundle: def validate_yaml_text(text: str) -> ValidationResult: - """Parse ``text`` as YAML and validate it as an :class:`ArenaEnvInitialGraphSpec`.""" + """Parse YAML and validate via the SimApp sidecar (registry lookups run there).""" cached_text = st.session_state.get("_validation_text") cached_result = st.session_state.get("_validation_result") if cached_text == text and isinstance(cached_result, ValidationResult): @@ -103,15 +169,41 @@ def validate_yaml_text(text: str) -> ValidationResult: else: try: raw = yaml.safe_load(text) + except Exception: + result = ValidationResult(spec=None, error=traceback.format_exc()) + else: if raw is None: result = ValidationResult(spec=None, error="YAML is empty") elif not isinstance(raw, dict): result = ValidationResult(spec=None, error=f"Expected mapping, got {type(raw).__name__}") else: - spec = ArenaEnvInitialGraphSpec.from_dict(raw) - result = ValidationResult(spec=spec, error=None) - except Exception: - result = ValidationResult(spec=None, error=traceback.format_exc()) + sidecar = _ensure_sidecar() + if sidecar is None: + result = ValidationResult( + spec=None, + error=( + "SimApp sidecar is unavailable — cannot validate registry entries. " + "Check the terminal where you launched the server." + ), + ) + else: + try: + response = sidecar.validate_yaml_text(text) + except SimAppSidecarError as exc: + _get_simapp_sidecar.clear() + result = ValidationResult(spec=None, error=str(exc)) + else: + if not response.get("ok"): + err = response.get("error", "validation failed") + tb = response.get("traceback", "") + message = f"{err}\n\n{tb}" if tb else str(err) + result = ValidationResult(spec=None, error=message) + else: + try: + spec = _spec_from_sidecar_dict(response["spec_dict"]) + result = ValidationResult(spec=spec, error=None) + except Exception: + result = ValidationResult(spec=None, error=traceback.format_exc()) st.session_state["_validation_text"] = text st.session_state["_validation_result"] = result @@ -160,7 +252,7 @@ def _apply_generated_yaml(yaml_text: str, *, spec: ArenaEnvInitialGraphSpec | No st.session_state["editor_version"] = st.session_state.get("editor_version", 0) + 1 st.session_state["last_rendered_text"] = yaml_text if spec is not None: - st.session_state["rendered_html"] = render_dashboard_html(spec) + st.session_state["rendered_html"] = _render_with_thumbnails(spec) st.session_state["_validation_text"] = yaml_text st.session_state["_validation_result"] = ValidationResult(spec=spec, error=None) else: @@ -360,7 +452,7 @@ def render_editor_panel(yaml_path: Path | None) -> ValidationResult: if edited_since_render: if validation.is_valid: with st.spinner("Rendering visualization…"): - st.session_state["rendered_html"] = render_dashboard_html(validation.spec) + st.session_state["rendered_html"] = _render_with_thumbnails(validation.spec) else: st.session_state["rendered_html"] = _BROKEN_PLACEHOLDER_HTML st.session_state["last_rendered_text"] = st.session_state["edited_text"] diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/thumbnail_render.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/thumbnail_render.py new file mode 100644 index 000000000..f63f17d4e --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/thumbnail_render.py @@ -0,0 +1,360 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""USD viewport thumbnail rendering for the review GUI SimApp sidecar.""" + +from __future__ import annotations + +import argparse +import hashlib +import sys +from pathlib import Path + +from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvInitialGraphSpec + +THUMBNAIL_CACHE_DIR = Path(__file__).resolve().parents[3] / ".cache" / "llm_env_gen_thumbnails" + + +def _render_thumbnails_with_app(app, spec: ArenaEnvInitialGraphSpec) -> dict[str, Path]: + """Resolve each node's USD via ``AssetRegistry``, render cache-misses, return PNG paths. + + ``app`` must already be a booted ``SimulationApp``. The caller owns the + lifecycle (Kit may turn ``app.close()`` into ``os._exit(0)`` — that's why + the sidecar holds the only reference and closes it inside its ``finally``). + + Returns ``{node.id: png_path}`` for nodes whose asset USD could be located + *and* whose PNG exists on disk (either from the persistent cache under + ``THUMBNAIL_CACHE_DIR`` or freshly rendered into the cache by + :func:`_capture_usd_thumbnails`). Missing entries fall through to the + placeholder in :func:`_render_node_thumbnail`, so a partial failure (one + bad asset) never breaks the rest of the page. + + We return ``Path`` rather than ``bytes`` so the sidecar protocol can ship + just the filenames over its stdin/stdout pipe (a few hundred bytes of JSON + instead of multiple MB of base64 PNG data). The parent reads the bytes + itself off the shared filesystem cache. + + Ordering matters: ``SimulationApp`` MUST be launched before any + ``AssetRegistry`` access, because ``ensure_assets_registered()`` imports + isaaclab asset modules which transitively load ``pxr``. ``pxr`` loaded + before ``AppLauncher`` puts Kit's extension manager into an unrecoverable + state ("extension class wrapper for base class ... has not been created + yet"). This is the same root cause we fixed for the pytest suite. + """ + asset_paths = _resolve_node_usd_paths(spec) + if not asset_paths: + print("[thumbnail_render] no asset USD paths resolved; skipping thumbnail rendering.", file=sys.stderr) + return {} + + THUMBNAIL_CACHE_DIR.mkdir(parents=True, exist_ok=True) + + # Split into cache-hits vs to-render. Cache key is sha1(usd_path) so + # the same USD across multiple envs / nodes hits the same PNG. + resolved: dict[str, Path] = {} + to_render: dict[str, tuple[str, Path]] = {} + for node_id, usd_path in asset_paths.items(): + cache_path = THUMBNAIL_CACHE_DIR / f"{_usd_cache_key(usd_path)}.png" + if cache_path.exists() and cache_path.stat().st_size > 0: + resolved[node_id] = cache_path + else: + to_render[node_id] = (usd_path, cache_path) + + if to_render: + print( + f"[thumbnail_render] rendering {len(to_render)} new thumbnail(s) " + f"(reusing {len(resolved)} from cache at {THUMBNAIL_CACHE_DIR})...", + file=sys.stderr, + ) + # ``_capture_usd_thumbnails`` still returns ``{node_id: bytes}``, but + # we only use it as a presence signal here — the same call also wrote + # the PNG to ``cache_path`` as a side effect, which is what we return. + captured = _capture_usd_thumbnails(app, to_render) + for node_id, (_usd_path, cache_path) in to_render.items(): + if node_id in captured and cache_path.exists() and cache_path.stat().st_size > 0: + resolved[node_id] = cache_path + else: + print(f"[thumbnail_render] all {len(resolved)} thumbnail(s) served from cache.", file=sys.stderr) + + return resolved + + +def _sidecar_launch_args() -> argparse.Namespace: + """AppLauncher args for the review GUI sidecar (Kit UI + viewport capture).""" + return argparse.Namespace(visualizer=["kit"], enable_cameras=True, livestream=-1) + + +def _launch_simulation_app(): + """Boot Isaac Sim's ``SimulationApp`` with the Kit visualizer, or ``None`` on failure. + + Kept as a tiny helper so the call site can lazy-import inside this + function — module-level import of ``simulation_app`` would drag Kit + into every invocation, including ``--help``. + """ + try: + # Lazy-import: keeps the default ``review_graph`` invocation Kit-free. + from isaaclab_arena.utils.isaaclab_utils.simulation_app import get_app_launcher # noqa: PLC0415 + + return get_app_launcher(_sidecar_launch_args()).app + except Exception as exc: + print(f"[thumbnail_render] SimulationApp launch failed: {exc}", file=sys.stderr) + return None + + +def _resolve_node_usd_paths(spec: ArenaEnvInitialGraphSpec) -> dict[str, str]: + """Map ``node.id → usd_path`` via :class:`AssetRegistry`, skipping unresolvable nodes. + + Tries two lookup strategies in order: + + 1. Class-attribute ``cls.usd_path`` — the convention every ``LibraryObject`` + subclass in ``object_library.py`` follows. No instantiation, cheap. + + 2. ``cls().scene_config.robot.spawn.usd_path`` — the convention every + :class:`EmbodimentBase` subclass uses. Requires instantiating the + embodiment because the Franka embodiments populate ``scene_config.robot`` + inside ``__init__`` rather than as a class default. Embodiment + ``__init__`` is light (no Kit / sim required) — it only constructs + configclass objects. + + This function MUST be called only after ``SimulationApp`` has booted — see + the docstring of :func:`_render_thumbnails_with_app` for why. + """ + try: + from isaaclab_arena.assets.registries import AssetRegistry # noqa: PLC0415 + except Exception as exc: + print(f"[thumbnail_render] AssetRegistry import failed: {exc}", file=sys.stderr) + return {} + + registry = AssetRegistry() + paths: dict[str, str] = {} + for node in spec.nodes: + try: + if not registry.is_registered(node.name): + print(f"[thumbnail_render] {node.id}: asset '{node.name}' not registered, skipping.", file=sys.stderr) + continue + cls = registry.get_asset_by_name(node.name) + usd_path = _extract_usd_path(cls) + if not usd_path: + print(f"[thumbnail_render] {node.id}: '{node.name}' has no usd_path, skipping.", file=sys.stderr) + continue + paths[node.id] = usd_path + except Exception as exc: + print(f"[thumbnail_render] {node.id}: lookup failed for '{node.name}': {exc}", file=sys.stderr) + return paths + + +def _extract_usd_path(cls) -> str | None: + """Return the asset's root USD path, or ``None`` if not extractable. + + See :func:`_resolve_node_usd_paths` for the two strategies tried in order. + """ + # Strategy 1: ``LibraryObject`` convention. + usd_path = getattr(cls, "usd_path", None) + if usd_path: + return usd_path + + # Strategy 2: ``EmbodimentBase`` convention. Walk + # ``instance.scene_config.robot.spawn.usd_path``. We instantiate with no + # args; every embodiment ``__init__`` defaults all parameters. + # NoEmbodiment legitimately has no robot — its instance.scene_config + # exists but ``.robot`` is absent / None, so the getattr chain returns + # None and we silently fall through. + try: + instance = cls() + except Exception: + return None + scene_config = getattr(instance, "scene_config", None) + robot = getattr(scene_config, "robot", None) if scene_config is not None else None + spawn = getattr(robot, "spawn", None) if robot is not None else None + return getattr(spawn, "usd_path", None) if spawn is not None else None + + +def _usd_cache_key(usd_path: str) -> str: + return hashlib.sha1(usd_path.encode("utf-8")).hexdigest()[:16] + + +def _capture_usd_thumbnails(app, to_render: dict[str, tuple[str, Path]]) -> dict[str, bytes]: + """Capture all queued USDs under one already-booted ``SimulationApp``. + + Deduplicates by ``usd_path`` so the same USD shared by multiple nodes is + only rendered once and the bytes are fanned back out. + """ + out: dict[str, bytes] = {} + + path_to_node_ids: dict[str, list[str]] = {} + path_to_cache: dict[str, Path] = {} + for node_id, (usd_path, cache_path) in to_render.items(): + path_to_node_ids.setdefault(usd_path, []).append(node_id) + path_to_cache[usd_path] = cache_path + + for usd_path, node_ids in path_to_node_ids.items(): + cache_path = path_to_cache[usd_path] + try: + png_bytes = _render_one_usd(app, usd_path, cache_path) + except Exception as exc: + print(f"[thumbnail_render] render failed for {usd_path}: {exc}", file=sys.stderr) + continue + if png_bytes: + for node_id in node_ids: + out[node_id] = png_bytes + + return out + + +def _render_one_usd(app, usd_path: str, cache_path: Path) -> bytes | None: + """Open ``usd_path`` directly as the stage, frame the camera, capture PNG. + + Opening the USD as the stage root (rather than ``new_stage`` + reference + wrapper) is what makes viewport capture actually produce a file in + headless mode — Kit's viewport machinery binds to the just-opened stage + cleanly, whereas a referenced sub-stage left the render product empty in + every test we tried. The trade-off is that we lose isolation between + captures (each call replaces the stage), but Kit handles that fine + because we call ``open_stage`` again on the next asset. + """ + import omni.usd # noqa: PLC0415 + from omni.kit.viewport.utility import ( # noqa: PLC0415 + capture_viewport_to_file, + frame_viewport_prims, + get_active_viewport, + ) + from pxr import Sdf # noqa: PLC0415 + + ctx = omni.usd.get_context() + if not ctx.open_stage(usd_path): + print(f"[thumbnail_render] open_stage failed: {usd_path}", file=sys.stderr) + return None + stage = ctx.get_stage() + + # Wait for textures / payloads / Nucleus fetches to settle before framing. + _wait_for_stage_load(app, ctx) + + # Standalone object USDs (avocado, bowl, ...) ship no lights, so a viewport + # capture renders them as a near-black silhouette against the dark skybox + # — that's the "blank thumbnail" symptom. Complete scene USDs (maple table) + # already include their own lighting, so this is a no-op for them. + _ensure_default_lighting(stage) + + # Use the default prim if present, otherwise the pseudo-root, for framing. + target_prim = stage.GetDefaultPrim() + if not target_prim or not target_prim.IsValid(): + target_prim = stage.GetPrimAtPath(Sdf.Path("/")) + + viewport = get_active_viewport() + + # Use Kit's own ``frame_viewport_prims`` (the "F"-key equivalent / ``FramePrimsCommand``) + # so we go through the viewport camera controller. Manually editing the + # ``/OmniverseKit_Persp`` xform op directly worked sometimes but Kit's + # camera controller treats /OmniverseKit_Persp as an internal state and + # silently overrode our edits for small assets — that's why avocado / bowl + # captured as tiny specks even with the right math. Letting Kit do the + # framing is both correct and avoids us re-implementing the math. + framed = frame_viewport_prims(viewport, prims=[str(target_prim.GetPath())]) + if not framed: + print(f"[thumbnail_render] warning: frame_viewport_prims failed for {usd_path}", file=sys.stderr) + + # Settle Hydra after camera change so the captured frame matches the new pose. + for _ in range(30): + app.update() + + cache_path.parent.mkdir(parents=True, exist_ok=True) + capture_obj = capture_viewport_to_file(viewport, str(cache_path)) + + _wait_for_capture(app, capture_obj, cache_path, max_updates=600) + + if cache_path.exists() and cache_path.stat().st_size > 0: + return cache_path.read_bytes() + print(f"[thumbnail_render] capture produced no file: {cache_path}", file=sys.stderr) + return None + + +def _wait_for_stage_load(app, usd_context, max_updates: int = 600) -> None: + """Pump frames until ``usd_context.get_stage_loading_status()`` reports nothing pending. + + Returns after stage load completes or after the budget is exhausted. We + also need a few extra frames after the count goes to zero so material + binding / texture upload finishes — they don't show up in the load count. + """ + settled = 0 + for _ in range(max_updates): + app.update() + try: + _msg, loading_count, loaded_count = usd_context.get_stage_loading_status() + except Exception: + return + if loading_count == 0 and loaded_count == 0: + settled += 1 + if settled > 15: + return + else: + settled = 0 + + +def _wait_for_capture(app, capture_obj, cache_path: Path, max_updates: int = 600) -> None: + """Pump ``app.update()`` until the capture PNG lands on disk (or we time out). + + Kit's capture future is fulfilled inside its async loop during + ``app.update()``, but future completion doesn't always coincide with the + file being flushed — checking the file directly is the most reliable + completion signal. We also keep the future-based fast path so a + successful capture doesn't have to wait for the file system to settle. + """ + if capture_obj is None: + for _ in range(max_updates): + app.update() + return + + future = ( + getattr(capture_obj, "_Capture__future", None) + or getattr(capture_obj, "_RenderCapture__future", None) + or getattr(capture_obj, "future", None) + ) + + for _ in range(max_updates): + app.update() + if cache_path.exists() and cache_path.stat().st_size > 0: + return + if future is not None and future.done(): + # Future is done but file might still be flushing — give it a few frames. + for _ in range(15): + app.update() + if cache_path.exists() and cache_path.stat().st_size > 0: + return + return + + +def _ensure_default_lighting(stage) -> None: + """Add a dome + key distant light if the stage has none. + + Without this, standalone object USDs (which don't ship their own lights) + render as a near-black silhouette. We skip the addition if any + ``UsdLuxLight``-derived prim already exists on the stage to avoid + double-lighting scenes like the maple table that bake in their own rig. + """ + from pxr import Gf, Sdf, UsdGeom, UsdLux # noqa: PLC0415 + + for prim in stage.Traverse(): + if ( + prim.HasAPI(UsdLux.LightAPI) + or prim.IsA(UsdLux.BoundableLightBase) + or prim.IsA(UsdLux.NonboundableLightBase) + ): + return + + # Soft hemispherical fill so the asset is visible from any angle, plus a + # weak directional key for shape definition. Intensities are tuned for + # OmniPBR / RTX defaults; tweak if asset libraries adopt darker materials. + dome = UsdLux.DomeLight.Define(stage, Sdf.Path("/_ReviewDomeLight")) + dome.CreateIntensityAttr(800.0) + dome.CreateColorAttr(Gf.Vec3f(1.0, 1.0, 1.0)) + + key = UsdLux.DistantLight.Define(stage, Sdf.Path("/_ReviewKeyLight")) + key.CreateIntensityAttr(2500.0) + key.CreateAngleAttr(2.0) + # Aim the key roughly from the camera's 3/4 angle so the lit side faces + # the viewport. + key_xformable = UsdGeom.Xformable(key.GetPrim()) + key_xformable.ClearXformOpOrder() + rot = key_xformable.AddRotateXYZOp() + rot.Set(Gf.Vec3f(-45.0, 30.0, 0.0))