From 2ed6009c21da81b2c3b43064c34d2114c4c2e1c9 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Tue, 26 May 2026 15:52:09 +0800 Subject: [PATCH 01/10] Add graph review script --- isaaclab_arena/llm_env_gen/review_graph.py | 361 +++++++++++++++++++++ 1 file changed, 361 insertions(+) create mode 100644 isaaclab_arena/llm_env_gen/review_graph.py diff --git a/isaaclab_arena/llm_env_gen/review_graph.py b/isaaclab_arena/llm_env_gen/review_graph.py new file mode 100644 index 000000000..91f269f98 --- /dev/null +++ b/isaaclab_arena/llm_env_gen/review_graph.py @@ -0,0 +1,361 @@ +# 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 + +"""Render an ArenaEnvGraphSpec YAML to a self-contained HTML review page. + +Three panels (dark dashboard style): + * Top-left — graph diagram (mermaid.js, CDN-loaded) of the initial-state + spatial constraints. Anchor nodes are highlighted; constraints without + a child (is_anchor / position_limits / at_pose / ...) are listed below + the graph rather than rendered as self-loops. + * Bottom-left — task table (id, type, initial/success state ids, task_args). + * Right — node card grid: type badge, asset name, and the per-node YAML + stanza. ``_render_node_thumbnail`` is the single integration point for + a future real USD-snapshot renderer (e.g. ``pxr.UsdAppUtils.FrameRecorder`` + / ``usdrecord``); v1 emits a styled placeholder so the script stays + lightweight and runs outside the Isaac Sim Docker container. + +Usage: + # Default: writes .html alongside the input file. + python -m isaaclab_arena.llm_env_gen.review_graph \\ + --yaml isaaclab_arena/tests/test_data/pick_and_place_maple_table_env_graph.yaml + + # Explicit output path: + python -m isaaclab_arena.llm_env_gen.review_graph \\ + --yaml isaaclab_arena_environments/llm_generated/_proposal.yaml \\ + --out /tmp/review.html +""" + +from __future__ import annotations + +import argparse +import html as html_lib +import re +import webbrowser +import yaml +from dataclasses import asdict +from pathlib import Path + +from isaaclab_arena.environments.arena_env_graph_spec import ( + ArenaEnvGraphNodeSpec, + ArenaEnvGraphSpec, + ArenaEnvGraphStateSpec, + _yaml_dict_factory, +) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--yaml", type=Path, required=True, help="Path to an ArenaEnvGraphSpec YAML file.") + parser.add_argument( + "--out", + type=Path, + default=None, + help="Output HTML path. Defaults to .html next to the input.", + ) + parser.add_argument("--open", action="store_true", help="Open the resulting HTML in the default browser.") + args = parser.parse_args() + + spec = ArenaEnvGraphSpec.from_yaml(args.yaml) + out_path = args.out or args.yaml.with_suffix(".html") + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(_render_html(spec), encoding="utf-8") + print(f"Wrote {out_path}") + if args.open: + webbrowser.open(out_path.resolve().as_uri()) + + +# --------------------------------------------------------------------------- +# Top-level HTML +# --------------------------------------------------------------------------- + + +def _render_html(spec: ArenaEnvGraphSpec) -> str: + initial_state = _pick_initial_state(spec) + return f""" + + + +{html_lib.escape(spec.env_name)} — graph review + + + + +
+

{html_lib.escape(spec.env_name)}

+

{len(spec.nodes)} nodes · {len(spec.tasks)} tasks · {len(spec.state_specs)} state specs

+
+
+
+

Spatial graph (initial state: { + html_lib.escape(initial_state.id if initial_state else "") + })

+ 
{_render_mermaid(spec, initial_state)}
+ {_render_unary_constraints(initial_state)} +
+
+

Tasks

+ {_render_tasks_table(spec)} +
+
+

Nodes

+
{_render_node_cards(spec)}
+
+
+ + + +""" + + +def _pick_initial_state(spec: ArenaEnvGraphSpec) -> ArenaEnvGraphStateSpec | None: + """Pick the state spec that tasks point at as their initial state. + + Falls back to the first state spec in the list. Returns ``None`` only if + there are no state specs at all. + """ + if spec.tasks: + target_id = spec.tasks[0].initial_state_spec_id + for s in spec.state_specs: + if s.id == target_id: + return s + return spec.state_specs[0] if spec.state_specs else None + + +# --------------------------------------------------------------------------- +# Mermaid graph rendering +# --------------------------------------------------------------------------- + + +def _render_mermaid(spec: ArenaEnvGraphSpec, state: ArenaEnvGraphStateSpec | None) -> str: + """Emit a left-to-right mermaid graph of the spatial constraints with children. + + Constraints without a child (is_anchor / position_limits / at_pose / ...) + are dropped here and surfaced separately by :func:`_render_unary_constraints`. + """ + lines = ["graph LR"] + if state is None: + lines.append(" empty[no state spec]") + return "\n".join(lines) + + anchor_ids: set[str] = set() + edge_nodes: set[str] = set() + + for c in state.spatial_constraints: + kind = c.type.value + if kind == "is_anchor": + anchor_ids.add(c.parent) + elif c.child is not None: + lines.append( + f" {_mermaid_id(c.child)}[{_mermaid_label(c.child)}]" + f" -->|{kind}| " + f"{_mermaid_id(c.parent)}[{_mermaid_label(c.parent)}]" + ) + edge_nodes.add(c.child) + edge_nodes.add(c.parent) + + # Include every node from the spec so disconnected ones still appear. + for node in spec.nodes: + if node.id not in edge_nodes: + lines.append(f" {_mermaid_id(node.id)}[{_mermaid_label(node.id)}]") + + # Anchor highlight. + for anchor_id in anchor_ids: + lines.append(f" style {_mermaid_id(anchor_id)} fill:#3a7d44,color:#fff,stroke:#7fd17f,stroke-width:2px") + + # Color nodes by type for quick visual scanning. + type_palette = { + "background": ("#3a4f7a", "#7aa0d8"), + "embodiment": ("#7a3a3a", "#d87a7a"), + "object": ("#7a6b3a", "#d8c47a"), + "object_reference": ("#6b3a7a", "#c47ad8"), + "lighting": ("#3a7a7a", "#7ad8d8"), + } + for node in spec.nodes: + if node.id in anchor_ids: + continue # anchor style wins + fill, stroke = type_palette.get(node.type.value, ("#3a3d44", "#888")) + lines.append(f" style {_mermaid_id(node.id)} fill:{fill},color:#fff,stroke:{stroke}") + + return "\n".join(lines) + + +_MERMAID_ID_SAFE = re.compile(r"[^A-Za-z0-9_]") + + +def _mermaid_id(s: str) -> str: + """Mermaid node identifiers must be alphanumeric / underscore.""" + return _MERMAID_ID_SAFE.sub("_", s) + + +def _mermaid_label(s: str) -> str: + """Escape mermaid-significant characters inside node labels.""" + return s.replace('"', """).replace("|", "|") + + +def _render_unary_constraints(state: ArenaEnvGraphStateSpec | None) -> str: + """List constraints without a child below the graph (anchors, position_limits, ...).""" + if state is None: + return "" + rows = [] + for c in state.spatial_constraints: + if c.child is not None: + continue + params = ( + f' {html_lib.escape(yaml.safe_dump(c.params, default_flow_style=True).rstrip())}' + if c.params + else "" + ) + rows.append( + f'
  • {html_lib.escape(c.type.value)}' + f" on {html_lib.escape(c.parent)}{params}
    • " + ) + if not rows: + return "" + return ( + f'
    Unary constraints ({len(rows)})' + f'
      {"".join(rows)}
    ' + ) + + +# --------------------------------------------------------------------------- +# Tasks panel +# --------------------------------------------------------------------------- + + +def _render_tasks_table(spec: ArenaEnvGraphSpec) -> str: + if not spec.tasks: + return "

    No tasks defined.

    " + rows = [] + for t in spec.tasks: + task_args_str = yaml.safe_dump(t.task_args, sort_keys=False).rstrip() if t.task_args else "(empty)" + rows.append( + "" + f"{html_lib.escape(t.id)}" + f'{html_lib.escape(t.type)}' + f"{html_lib.escape(t.initial_state_spec_id)}" + f"{html_lib.escape(t.success_state_spec_id) or 'unset'}" + f"
    {html_lib.escape(task_args_str)}
    " + "" + ) + return ( + "" + "" + f"{''.join(rows)}" + "
    idtypeinitialsuccesstask_args
    " + ) + + +# --------------------------------------------------------------------------- +# Node cards +# --------------------------------------------------------------------------- + + +def _render_node_cards(spec: ArenaEnvGraphSpec) -> str: + return "\n".join(_render_one_node_card(node) for node in spec.nodes) + + +def _render_one_node_card(node: ArenaEnvGraphNodeSpec) -> str: + node_dict = asdict(node, dict_factory=_yaml_dict_factory) + node_yaml = yaml.safe_dump(node_dict, sort_keys=False).rstrip() + thumb = _render_node_thumbnail(node) + return f"""
    + {thumb} +
    +
    {html_lib.escape(node.id)}
    + {html_lib.escape(node.type.value)} +
    + 
    {html_lib.escape(node_yaml)}
    +
    """ + + +def _render_node_thumbnail(node: ArenaEnvGraphNodeSpec) -> str: + """Single integration point for per-node preview rendering. + + Currently emits a styled placeholder. To wire in real USD snapshots, look + up the asset's USD path via ``AssetRegistry.get_asset_by_name(node.name)``, + render a PNG with ``pxr.UsdAppUtils.FrameRecorder`` (or shell out to the + ``usdrecord`` CLI), and return an ```` + instead. The rest of the layout doesn't need to change. + """ + initial = (node.name[:2] if node.name else "?").upper() + return f"""
    + {html_lib.escape(initial)} + {html_lib.escape(node.name)} +
    """ + + +# --------------------------------------------------------------------------- +# Styling +# --------------------------------------------------------------------------- + +_CSS = """ +:root { + --bg: #15181d; + --bg-elev: #1d2128; + --bg-elev2: #262b34; + --border: #2f343d; + --fg: #e4e6eb; + --fg-muted: #8a9099; + --accent: #7fd17f; +} +* { box-sizing: border-box; } +body { margin: 0; padding: 24px; font: 14px/1.5 -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; + background: var(--bg); color: var(--fg); } +header { margin-bottom: 16px; } +header h1 { margin: 0; font-size: 28px; font-weight: 700; } +header .sub { margin: 4px 0 0; color: var(--fg-muted); font-size: 13px; } +main { display: grid; grid-template-columns: 1fr 1fr; grid-template-rows: auto auto; + grid-template-areas: "graph nodes" "tasks nodes"; gap: 16px; } +.graph-panel { grid-area: graph; } +.tasks-panel { grid-area: tasks; } +.nodes-panel { grid-area: nodes; } +.panel { background: var(--bg-elev); border: 1px solid var(--border); border-radius: 8px; padding: 16px; } +.panel h2 { margin: 0 0 12px; font-size: 16px; font-weight: 600; letter-spacing: 0.02em; } +.panel h2 .muted { color: var(--fg-muted); font-weight: 400; font-size: 13px; } +code { font-family: ui-monospace, 'SF Mono', Menlo, monospace; font-size: 12px; + background: var(--bg-elev2); padding: 1px 6px; border-radius: 4px; } +pre { font-family: ui-monospace, 'SF Mono', Menlo, monospace; font-size: 12px; + background: var(--bg-elev2); padding: 10px 12px; border-radius: 6px; margin: 0; + white-space: pre-wrap; word-break: break-word; } +.muted { color: var(--fg-muted); } +.badge { display: inline-block; padding: 2px 8px; border-radius: 999px; font-size: 11px; + font-weight: 600; letter-spacing: 0.03em; background: var(--bg-elev2); color: var(--fg); } +.badge.type-background { background: #3a4f7a; } +.badge.type-embodiment { background: #7a3a3a; } +.badge.type-object { background: #7a6b3a; } +.badge.type-object_reference { background: #6b3a7a; } +.badge.type-lighting { background: #3a7a7a; } +.badge.type-is_anchor { background: #3a7d44; } +.badge.type-position_limits, .badge.type-at_pose, .badge.type-at_position { background: #6b3a7a; } +.badge.type-task { background: #2f343d; border: 1px solid #4a5; color: var(--accent); } +.mermaid { background: var(--bg-elev2); padding: 8px; border-radius: 6px; min-height: 220px; + display: flex; align-items: center; justify-content: center; } +.unary { margin-top: 12px; } +.unary summary { cursor: pointer; color: var(--fg-muted); font-size: 13px; padding: 4px 0; } +.unary ul { margin: 8px 0 0; padding-left: 20px; list-style: disc; color: var(--fg); } +.unary li { padding: 3px 0; } +table.tasks { width: 100%; border-collapse: collapse; } +table.tasks th, table.tasks td { text-align: left; padding: 8px 10px; border-bottom: 1px solid var(--border); + vertical-align: top; font-size: 12px; } +table.tasks th { color: var(--fg-muted); font-weight: 600; text-transform: uppercase; letter-spacing: 0.05em; } +table.tasks pre { padding: 6px 8px; font-size: 11px; } +.node-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(220px, 1fr)); gap: 12px; } +.node-card { background: var(--bg-elev2); border: 1px solid var(--border); border-radius: 8px; + padding: 12px; display: flex; flex-direction: column; gap: 10px; } +.node-card .thumb { aspect-ratio: 1 / 1; background: linear-gradient(135deg, #2a2f37, #1c2026); + border-radius: 6px; display: flex; flex-direction: column; + align-items: center; justify-content: center; color: var(--fg-muted); position: relative; } +.thumb-initial { font-size: 36px; font-weight: 700; color: var(--fg); opacity: 0.6; + font-family: ui-monospace, monospace; } +.thumb-name { font-size: 10px; margin-top: 6px; padding: 0 8px; text-align: center; word-break: break-word; } +.node-meta { display: flex; align-items: center; justify-content: space-between; gap: 8px; } +.node-id { font-family: ui-monospace, monospace; font-size: 13px; font-weight: 600; word-break: break-all; } +.node-yaml { font-size: 11px; } +""" + + +if __name__ == "__main__": + main() From 1d0a8346a00cf0cbab0b426c64b802d2858d5dbe Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Tue, 26 May 2026 17:39:16 +0800 Subject: [PATCH 02/10] Add snapshot to html reviewer --- isaaclab_arena/llm_env_gen/review_graph.py | 411 +++++++++++++++++++-- 1 file changed, 385 insertions(+), 26 deletions(-) diff --git a/isaaclab_arena/llm_env_gen/review_graph.py b/isaaclab_arena/llm_env_gen/review_graph.py index 91f269f98..63b237882 100644 --- a/isaaclab_arena/llm_env_gen/review_graph.py +++ b/isaaclab_arena/llm_env_gen/review_graph.py @@ -12,27 +12,40 @@ the graph rather than rendered as self-loops. * Bottom-left — task table (id, type, initial/success state ids, task_args). * Right — node card grid: type badge, asset name, and the per-node YAML - stanza. ``_render_node_thumbnail`` is the single integration point for - a future real USD-snapshot renderer (e.g. ``pxr.UsdAppUtils.FrameRecorder`` - / ``usdrecord``); v1 emits a styled placeholder so the script stays - lightweight and runs outside the Isaac Sim Docker container. + stanza. With ``--render-thumbnails``, the per-node thumbnail is a real + USD viewport capture (cached on disk and inlined as base64); otherwise + a styled placeholder keeps the script lightweight. Usage: - # Default: writes .html alongside the input file. + # Default: writes .html alongside the input file. Lightweight. python -m isaaclab_arena.llm_env_gen.review_graph \\ --yaml isaaclab_arena/tests/test_data/pick_and_place_maple_table_env_graph.yaml - # Explicit output path: - python -m isaaclab_arena.llm_env_gen.review_graph \\ + # With real per-node USD snapshots (boots Isaac Sim once, ~30s): + /isaac-sim/python.sh -m isaaclab_arena.llm_env_gen.review_graph \\ --yaml isaaclab_arena_environments/llm_generated/_proposal.yaml \\ - --out /tmp/review.html + --render-thumbnails --open + +Note on USD rendering: + ``pxr.UsdAppUtils.FrameRecorder`` and the ``usdrecord`` CLI are NOT + available inside the Isaac Sim container (Kit ships ``UsdAppUtils.py`` + but strips out ``libusd_usdAppUtils.so``, and ``usdrecord`` is omitted + entirely). The Kit-equivalent path used here is: + ``omni.usd`` to open the stage + ``omni.kit.viewport.utility`` to + capture the active viewport. Kit transparently uses cached Nucleus + thumbnails when opening ``omniverse://`` URIs, so we don't need a + separate Nucleus-HTTPS probe path. """ from __future__ import annotations import argparse +import base64 +import contextlib +import hashlib import html as html_lib import re +import sys import webbrowser import yaml from dataclasses import asdict @@ -45,6 +58,12 @@ _yaml_dict_factory, ) +# Disk cache for rendered thumbnails. Keyed by sha1(usd_path) so identical +# USDs across envs reuse the same PNG. Survives across runs to avoid the +# ~30s SimulationApp boot when nothing changed. +_THUMBNAIL_CACHE_DIR = Path(__file__).resolve().parents[2] / ".cache" / "llm_env_gen_thumbnails" +_THUMBNAIL_SIZE = 256 + def main() -> None: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) @@ -56,15 +75,41 @@ def main() -> None: help="Output HTML path. Defaults to .html next to the input.", ) parser.add_argument("--open", action="store_true", help="Open the resulting HTML in the default browser.") + parser.add_argument( + "--render-thumbnails", + action="store_true", + help=( + "Boot Isaac Sim once and capture per-node USD viewport thumbnails " + "(cached under .cache/llm_env_gen_thumbnails/). Slow first run " + "(~30s SimulationApp boot + ~2s per unique USD); subsequent runs " + "reuse cached PNGs. Must run inside the Isaac Sim container." + ), + ) args = parser.parse_args() spec = ArenaEnvGraphSpec.from_yaml(args.yaml) out_path = args.out or args.yaml.with_suffix(".html") out_path.parent.mkdir(parents=True, exist_ok=True) - out_path.write_text(_render_html(spec), encoding="utf-8") - print(f"Wrote {out_path}") - if args.open: - webbrowser.open(out_path.resolve().as_uri()) + + # Important: when --render-thumbnails is set, we keep SimulationApp open + # across the HTML write. Calling ``app.close()`` first can ``os._exit(0)`` + # (Kit's normal shutdown behavior) and silently drop the write_text below. + app = None + try: + thumbnails: dict[str, bytes] = {} + if args.render_thumbnails: + app = _launch_simulation_app() + if app is not None: + thumbnails = _render_thumbnails_with_app(app, spec) + + out_path.write_text(_render_html(spec, thumbnails), encoding="utf-8") + print(f"Wrote {out_path}") + if args.open: + webbrowser.open(out_path.resolve().as_uri()) + finally: + if app is not None: + with contextlib.suppress(Exception): + app.close() # --------------------------------------------------------------------------- @@ -72,8 +117,9 @@ def main() -> None: # --------------------------------------------------------------------------- -def _render_html(spec: ArenaEnvGraphSpec) -> str: +def _render_html(spec: ArenaEnvGraphSpec, thumbnails: dict[str, bytes] | None = None) -> str: initial_state = _pick_initial_state(spec) + thumbnails = thumbnails or {} return f""" @@ -101,7 +147,7 @@ def _render_html(spec: ArenaEnvGraphSpec) -> str:

    Nodes

    -
    {_render_node_cards(spec)}
    +
    {_render_node_cards(spec, thumbnails)}
    @@ -253,14 +299,14 @@ def _render_tasks_table(spec: ArenaEnvGraphSpec) -> str: # --------------------------------------------------------------------------- -def _render_node_cards(spec: ArenaEnvGraphSpec) -> str: - return "\n".join(_render_one_node_card(node) for node in spec.nodes) +def _render_node_cards(spec: ArenaEnvGraphSpec, thumbnails: dict[str, bytes]) -> str: + return "\n".join(_render_one_node_card(node, thumbnails.get(node.id)) for node in spec.nodes) -def _render_one_node_card(node: ArenaEnvGraphNodeSpec) -> str: +def _render_one_node_card(node: ArenaEnvGraphNodeSpec, png_bytes: bytes | None) -> str: node_dict = asdict(node, dict_factory=_yaml_dict_factory) node_yaml = yaml.safe_dump(node_dict, sort_keys=False).rstrip() - thumb = _render_node_thumbnail(node) + thumb = _render_node_thumbnail(node, png_bytes) return f"""
    {thumb}
    @@ -271,15 +317,26 @@ def _render_one_node_card(node: ArenaEnvGraphNodeSpec) -> str:
    """ -def _render_node_thumbnail(node: ArenaEnvGraphNodeSpec) -> str: - """Single integration point for per-node preview rendering. +def _render_node_thumbnail(node: ArenaEnvGraphNodeSpec, png_bytes: bytes | None = None) -> str: + """Per-node thumbnail: real USD viewport capture if rendered, else placeholder. + + When ``png_bytes`` is provided (i.e. ``--render-thumbnails`` ran and the + asset was successfully captured by :func:`_render_thumbnails_for_spec`), + inline the PNG as a ``data:image/png;base64,...`` URI so the resulting + HTML is fully self-contained — no sidecar files to keep next to the page. - Currently emits a styled placeholder. To wire in real USD snapshots, look - up the asset's USD path via ``AssetRegistry.get_asset_by_name(node.name)``, - render a PNG with ``pxr.UsdAppUtils.FrameRecorder`` (or shell out to the - ``usdrecord`` CLI), and return an ```` - instead. The rest of the layout doesn't need to change. + Otherwise fall back to the lightweight two-letter placeholder card, so + a default ``python -m ... review_graph --yaml ...`` invocation still + produces a useful page without booting Isaac Sim. """ + 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)} @@ -287,6 +344,302 @@ def _render_node_thumbnail(node: ArenaEnvGraphNodeSpec) -> str:
    """ +# --------------------------------------------------------------------------- +# USD viewport capture (opt-in via --render-thumbnails) +# --------------------------------------------------------------------------- + + +def _render_thumbnails_with_app(app, spec: ArenaEnvGraphSpec) -> dict[str, bytes]: + """Resolve each node's USD via ``AssetRegistry``, render or read cache. + + ``app`` must already be a booted ``SimulationApp``. The caller owns the + lifecycle so the HTML write can happen before ``app.close()`` (which Kit + may turn into ``os._exit(0)``). + + Returns ``{node.id: png_bytes}`` for nodes whose asset USD could be + located *and* rendered. 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. + + 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("[review_graph] 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. + rendered: dict[str, bytes] = {} + 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: + rendered[node_id] = cache_path.read_bytes() + else: + to_render[node_id] = (usd_path, cache_path) + + if to_render: + print( + f"[review_graph] rendering {len(to_render)} new thumbnail(s) " + f"(reusing {len(rendered)} from cache at {_THUMBNAIL_CACHE_DIR})...", + file=sys.stderr, + ) + rendered.update(_capture_usd_thumbnails(app, to_render)) + else: + print(f"[review_graph] all {len(rendered)} thumbnail(s) served from cache.", file=sys.stderr) + + return rendered + + +def _launch_simulation_app(): + """Boot Isaac Sim's ``SimulationApp`` for headless viewport capture, 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 + + sim_args = argparse.Namespace(headless=True, enable_cameras=True, hide_ui=True, livestream=-1) + return get_app_launcher(sim_args).app + except Exception as exc: + print(f"[review_graph] SimulationApp launch failed: {exc}", file=sys.stderr) + return None + + +def _resolve_node_usd_paths(spec: ArenaEnvGraphSpec) -> dict[str, str]: + """Map ``node.id → usd_path`` via :class:`AssetRegistry`, skipping unresolvable nodes. + + ``usd_path`` is read as a class attribute (the convention used by every + ``LibraryObject`` subclass in ``object_library.py``); we never instantiate + the asset class. This function MUST be called only after ``SimulationApp`` + has booted — see the docstring of :func:`_render_thumbnails_for_spec` for + why. + """ + try: + from isaaclab_arena.assets.registries import AssetRegistry # noqa: PLC0415 + except Exception as exc: + print(f"[review_graph] 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"[review_graph] {node.id}: asset '{node.name}' not registered, skipping.", file=sys.stderr) + continue + cls = registry.get_asset_by_name(node.name) + usd_path = getattr(cls, "usd_path", None) + if not usd_path: + print(f"[review_graph] {node.id}: '{node.name}' has no usd_path, skipping.", file=sys.stderr) + continue + paths[node.id] = usd_path + except Exception as exc: + print(f"[review_graph] {node.id}: lookup failed for '{node.name}': {exc}", file=sys.stderr) + return paths + + +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"[review_graph] 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"[review_graph] 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"[review_graph] 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"[review_graph] 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)) + + # --------------------------------------------------------------------------- # Styling # --------------------------------------------------------------------------- @@ -347,7 +700,13 @@ def _render_node_thumbnail(node: ArenaEnvGraphNodeSpec) -> str: padding: 12px; display: flex; flex-direction: column; gap: 10px; } .node-card .thumb { aspect-ratio: 1 / 1; background: linear-gradient(135deg, #2a2f37, #1c2026); border-radius: 6px; display: flex; flex-direction: column; - align-items: center; justify-content: center; color: var(--fg-muted); position: relative; } + align-items: center; justify-content: center; color: var(--fg-muted); + position: relative; overflow: hidden; } +.node-card .thumb-rendered { background: #0e1115; } +.node-card .thumb-rendered img { width: 100%; height: 100%; object-fit: contain; display: block; } +.node-card .thumb-rendered .thumb-name { position: absolute; bottom: 0; left: 0; right: 0; + padding: 4px 6px; background: rgba(15, 17, 21, 0.78); + color: var(--fg); margin: 0; } .thumb-initial { font-size: 36px; font-weight: 700; color: var(--fg); opacity: 0.6; font-family: ui-monospace, monospace; } .thumb-name { font-size: 10px; margin-top: 6px; padding: 0 8px; text-align: center; word-break: break-word; } From d9f7c04193b599fde2f7cf1184ea79fec44d1e5f Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Tue, 26 May 2026 18:05:21 +0800 Subject: [PATCH 03/10] Add robot snapshot --- isaaclab_arena/llm_env_gen/review_graph.py | 47 +++++++++++++++++++--- 1 file changed, 41 insertions(+), 6 deletions(-) diff --git a/isaaclab_arena/llm_env_gen/review_graph.py b/isaaclab_arena/llm_env_gen/review_graph.py index 63b237882..59fe3899f 100644 --- a/isaaclab_arena/llm_env_gen/review_graph.py +++ b/isaaclab_arena/llm_env_gen/review_graph.py @@ -420,11 +420,20 @@ def _launch_simulation_app(): def _resolve_node_usd_paths(spec: ArenaEnvGraphSpec) -> dict[str, str]: """Map ``node.id → usd_path`` via :class:`AssetRegistry`, skipping unresolvable nodes. - ``usd_path`` is read as a class attribute (the convention used by every - ``LibraryObject`` subclass in ``object_library.py``); we never instantiate - the asset class. This function MUST be called only after ``SimulationApp`` - has booted — see the docstring of :func:`_render_thumbnails_for_spec` for - why. + 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 @@ -440,7 +449,7 @@ def _resolve_node_usd_paths(spec: ArenaEnvGraphSpec) -> dict[str, str]: print(f"[review_graph] {node.id}: asset '{node.name}' not registered, skipping.", file=sys.stderr) continue cls = registry.get_asset_by_name(node.name) - usd_path = getattr(cls, "usd_path", None) + usd_path = _extract_usd_path(cls) if not usd_path: print(f"[review_graph] {node.id}: '{node.name}' has no usd_path, skipping.", file=sys.stderr) continue @@ -450,6 +459,32 @@ def _resolve_node_usd_paths(spec: ArenaEnvGraphSpec) -> dict[str, str]: 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] From ec34903beff8446de5beb06483b4af44b6b99ba2 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Thu, 28 May 2026 23:50:17 +0800 Subject: [PATCH 04/10] Move script folder --- .../{llm_env_gen => environments/agentic_env_gen}/review_graph.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename isaaclab_arena/{llm_env_gen => environments/agentic_env_gen}/review_graph.py (100%) diff --git a/isaaclab_arena/llm_env_gen/review_graph.py b/isaaclab_arena/environments/agentic_env_gen/review_graph.py similarity index 100% rename from isaaclab_arena/llm_env_gen/review_graph.py rename to isaaclab_arena/environments/agentic_env_gen/review_graph.py From b445b450c13bdd10705e0e9ac19b8b2ee1b286a1 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Tue, 26 May 2026 18:16:06 +0800 Subject: [PATCH 05/10] Move assert_unique_ids + assert_references_exist to post_init so it's always called when loading from yaml/dict --- .../environments/arena_env_graph_spec.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/isaaclab_arena/environments/arena_env_graph_spec.py b/isaaclab_arena/environments/arena_env_graph_spec.py index 2f69f97d6..a62ab5d93 100644 --- a/isaaclab_arena/environments/arena_env_graph_spec.py +++ b/isaaclab_arena/environments/arena_env_graph_spec.py @@ -149,6 +149,16 @@ class ArenaEnvGraphSpec: tasks: list[ArenaEnvGraphTaskSpec] = field(default_factory=list) state_specs: list[ArenaEnvGraphStateSpec] = field(default_factory=list) + def __post_init__(self) -> None: + # Enforce graph invariants on EVERY construction path (YAML parse, direct + # dataclass instantiation, programmatic build, ...). Centralizing here means + # downstream consumers — including ``nodes_by_id`` / ``tasks_by_id`` / + # ``state_specs_by_id``, which collapse duplicates silently in their dict + # comprehensions — can rely on globally-unique ids and valid references + # without re-validating. + assert_unique_ids(self.nodes, self.tasks, self.state_specs) + assert_references_exist(self.nodes, self.tasks, self.state_specs) + @classmethod def from_yaml(cls, path: str | Path) -> "ArenaEnvGraphSpec": with Path(path).open("r", encoding="utf-8") as f: @@ -157,18 +167,11 @@ def from_yaml(cls, path: str | Path) -> "ArenaEnvGraphSpec": @classmethod def from_dict(cls, data: dict[str, Any]) -> "ArenaEnvGraphSpec": data = as_dict(data, "Env graph spec") - nodes = parse_list(data, "nodes", _parse_node) - tasks = parse_list(data, "tasks", _parse_task) - state_specs = parse_list(data, "state_specs", _parse_state_spec) - - assert_unique_ids(nodes, tasks, state_specs) - assert_references_exist(nodes, tasks, state_specs) - return cls( env_name=required_str(data, "env_name"), - nodes=nodes, - tasks=tasks, - state_specs=state_specs, + nodes=parse_list(data, "nodes", _parse_node), + tasks=parse_list(data, "tasks", _parse_task), + state_specs=parse_list(data, "state_specs", _parse_state_spec), ) @property From 61c032adcc185ce1a6804d58ea7abd2ab37f4fd9 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Thu, 28 May 2026 00:45:42 +0800 Subject: [PATCH 06/10] Add ArenaEnvGraphSpec YAML-serialization --- .../environments/arena_env_graph_spec.py | 43 ++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/isaaclab_arena/environments/arena_env_graph_spec.py b/isaaclab_arena/environments/arena_env_graph_spec.py index a62ab5d93..bf80433f4 100644 --- a/isaaclab_arena/environments/arena_env_graph_spec.py +++ b/isaaclab_arena/environments/arena_env_graph_spec.py @@ -4,7 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 import yaml -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass, field from enum import Enum from pathlib import Path from typing import Any @@ -174,6 +174,27 @@ def from_dict(cls, data: dict[str, Any]) -> "ArenaEnvGraphSpec": state_specs=parse_list(data, "state_specs", _parse_state_spec), ) + def to_dict(self) -> dict[str, Any]: + """Return a YAML/JSON-serializable dict. + + Output shape round-trips through :meth:`from_dict` / :meth:`from_yaml`: + enums become their ``.value`` strings and ``None`` / empty-dict fields + are omitted so the optional-field parsers fall back to their defaults. + """ + return asdict(self, dict_factory=_yaml_dict_factory) + + def to_yaml(self, path: str | Path) -> Path: + """Write this spec to ``path`` as YAML. Creates parent dirs as needed. + + Returns the resolved :class:`Path` written. Symmetric with + :meth:`from_yaml`. + """ + out_path = Path(path) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w", encoding="utf-8") as f: + yaml.safe_dump(self.to_dict(), f, sort_keys=False) + return out_path + @property def nodes_by_id(self) -> dict[str, ArenaEnvGraphNodeSpec]: return {node.id: node for node in self.nodes} @@ -257,3 +278,23 @@ def _parse_task(data: Any) -> ArenaEnvGraphTaskSpec: success_state_spec_id=required_str(data, "success_state_spec_id"), task_args=optional_dict(data, "task_args"), ) + + +def _yaml_dict_factory(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + """``dataclasses.asdict`` hook used by :meth:`ArenaEnvGraphSpec.to_dict`. + + Two responsibilities: + * convert :class:`Enum` field values to their ``.value`` strings so + ``yaml.safe_dump`` can serialize them, and + * drop ``None`` / empty-dict fields so the emitted YAML stays clean + and ``optional_str`` / ``optional_dict`` parsers pick up defaults + instead of seeing redundant keys. + """ + out: dict[str, Any] = {} + for key, value in pairs: + if isinstance(value, Enum): + value = value.value + if value is None or (isinstance(value, dict) and not value): + continue + out[key] = value + return out From 65cd65ca33c5038d242a34bd73da278c4f4781e4 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Fri, 29 May 2026 10:45:13 +0800 Subject: [PATCH 07/10] Add interative editing --- .../agentic_env_gen/review_app.py | 426 ++++++++++++++++++ .../agentic_env_gen/review_graph.py | 190 ++++++-- setup.py | 9 + 3 files changed, 577 insertions(+), 48 deletions(-) create mode 100644 isaaclab_arena/environments/agentic_env_gen/review_app.py diff --git a/isaaclab_arena/environments/agentic_env_gen/review_app.py b/isaaclab_arena/environments/agentic_env_gen/review_app.py new file mode 100644 index 000000000..3649a1e0d --- /dev/null +++ b/isaaclab_arena/environments/agentic_env_gen/review_app.py @@ -0,0 +1,426 @@ +# 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 + +"""Streamlit-backed live editor for ArenaEnvGraphSpec YAMLs. + +Wraps :func:`isaaclab_arena.environments.agentic_env_gen.review_graph.render_html_for_spec` +in a two-pane Streamlit page so the user can edit the YAML in the browser and +re-render the dark-dashboard visualization on demand. ``SimulationApp`` is +booted once via ``@st.cache_resource`` and reused for every thumbnail render +(disk-cache hit when possible, live USD viewport capture when not). + +Launch (always via the wrapper in review_graph.py — handles streamlit flags): + /isaac-sim/python.sh -m isaaclab_arena.environments.agentic_env_gen.review_graph \\ + --yaml path/to/spec.yaml + +Design: + * Left pane — ``streamlit-ace`` YAML editor + validation badge + action + buttons. Validation runs on every rerun (i.e. after each editor blur). + * Right pane — sandboxed iframe with the rendered review HTML. Updates only + when the user clicks "Regenerate", so editing never causes mid-typing + re-renders. + * Thumbnails — real USD viewport captures. Booted ``SimulationApp`` lives + inside an ``@st.cache_resource`` so its ~30s startup is paid once per + server lifetime. PNGs are cached on disk under + ``.cache/llm_env_gen_thumbnails/`` and survive across runs. +""" + +from __future__ import annotations + +import argparse +import traceback +import yaml +from dataclasses import dataclass +from pathlib import Path + +import streamlit as st + +from isaaclab_arena.environments.agentic_env_gen.review_graph import ( + launch_simulation_app, + render_html_for_spec, + render_thumbnails_for_spec, +) +from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvGraphSpec + +# Visualization iframe height. Tuned so the graph + tasks + node grid all +# fit without an outer Streamlit scrollbar swallowing the inner one. +_IFRAME_HEIGHT_PX = 1100 + + +# --------------------------------------------------------------------------- +# SimulationApp lifecycle +# --------------------------------------------------------------------------- + + +def _patch_signal_for_non_main_thread() -> None: + """No-op ``signal.signal`` for non-main threads so Kit's bootstrap doesn't raise. + + Kit's ``SimulationApp.__init__`` installs SIGINT/SIGTERM handlers via + ``signal.signal``, which raises ``ValueError("signal only works in main + thread of the main interpreter")`` when called outside the interpreter's + main thread. Streamlit's ``ScriptRunner`` runs the user script in a + worker thread per session — so the SimApp boot would crash there even + though it would succeed under a plain ``python -m`` invocation. + + We swallow *only* the "main thread" ValueError and let everything else + propagate. The Streamlit main thread keeps its own SIGINT handler (it's + installed before our script ever runs), so Ctrl-C still tears the server + down cleanly. The patch is idempotent — repeated calls re-wrap the + already-wrapped function harmlessly. + """ + import signal # noqa: PLC0415 + + if getattr(signal.signal, "_arena_patched", False): + return + + _original_signal = signal.signal + + def _safe_signal(signum, handler): + try: + return _original_signal(signum, handler) + except ValueError as exc: + if "main thread" in str(exc): + return None + raise + + _safe_signal._arena_patched = True # type: ignore[attr-defined] + signal.signal = _safe_signal # type: ignore[assignment] + + +@st.cache_resource(show_spinner="Booting Isaac Sim (≈30s on first run, cached afterwards)…") +def _get_simulation_app(): + """Boot ``SimulationApp`` exactly once per Streamlit server process. + + ``@st.cache_resource`` is exactly the right primitive here: the returned + handle is shared across reruns and across browser sessions, and Streamlit + won't try to pickle it (unlike ``@st.cache_data``). If the boot fails we + return ``None`` and the app degrades to placeholder thumbnails — the + review page still works, just with the styled two-letter cards instead + of real USD captures. + + The signal patch must happen *before* the SimApp launch (which happens + on the Streamlit worker thread); see ``_patch_signal_for_non_main_thread``. + """ + _patch_signal_for_non_main_thread() + return launch_simulation_app() + + +def _render_with_thumbnails(spec: ArenaEnvGraphSpec) -> str: + """Render the review HTML for ``spec`` with thumbnails resolved through + the cached ``SimulationApp``. + + Cache-aware in two layers: + * The disk cache under ``.cache/llm_env_gen_thumbnails/`` survives + across runs; ``render_thumbnails_for_spec`` reads it directly. + * Within a run, ``@st.cache_resource`` ensures we never re-boot Kit. + + If ``SimulationApp`` is unavailable (boot failed) we fall back to + placeholder thumbnails so the user still gets a usable page. + """ + app = _get_simulation_app() + if app is None: + st.warning( + "Isaac Sim is unavailable — falling back to placeholder thumbnails. " + "Check the terminal where you launched the server for the boot error.", + icon="⚠️", + ) + return render_html_for_spec(spec, thumbnails=None) + thumbnails = render_thumbnails_for_spec(app, spec) + return render_html_for_spec(spec, thumbnails=thumbnails) + + +# --------------------------------------------------------------------------- +# Args + session-state init +# --------------------------------------------------------------------------- + + +def _parse_args() -> argparse.Namespace: + """Pull ``--yaml`` from ``sys.argv`` (post ``--`` from the streamlit CLI). + + Streamlit forwards anything after ``--`` on its command line into the + script's ``sys.argv``. We use a tolerant parser so reruns (which keep + argv intact) never abort the app. + """ + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--yaml", type=Path, required=True) + args, _unknown = parser.parse_known_args() + return args + + +@dataclass +class _ValidationResult: + """Outcome of validating editor text against ``ArenaEnvGraphSpec``.""" + + spec: ArenaEnvGraphSpec | None + error: str | None # human-readable, multi-line; None iff spec is not None + + @property + def is_valid(self) -> bool: + return self.spec is not None + + +def _validate_yaml_text(text: str) -> _ValidationResult: + """Two-stage validation: yaml.safe_load → ArenaEnvGraphSpec.from_dict. + + Returns a populated error string on the first failing stage so the UI can + render exactly one red banner with the most actionable message. + """ + # Stage 1: YAML parse. PyYAML's ``problem_mark`` is the most useful + # location info we get for syntax errors — surface it explicitly. + try: + loaded = yaml.safe_load(text) + except yaml.YAMLError as exc: + mark = getattr(exc, "problem_mark", None) + loc = f" (line {mark.line + 1}, column {mark.column + 1})" if mark is not None else "" + return _ValidationResult(spec=None, error=f"YAML parse error{loc}:\n{exc}") + + if not isinstance(loaded, dict): + return _ValidationResult( + spec=None, + error=f"Top-level YAML must be a mapping, got {type(loaded).__name__}.", + ) + + # Stage 2: schema validation via the existing dataclass parser. + try: + spec = ArenaEnvGraphSpec.from_dict(loaded) + except Exception as exc: + # ``from_dict`` raises a mix of AssertionError / ValueError / KeyError + # depending on which graph_spec_utils check trips. Keep the type so + # users can tell schema errors apart from YAML errors. + tb = traceback.format_exception_only(type(exc), exc) + return _ValidationResult(spec=None, error="".join(tb).rstrip()) + + return _ValidationResult(spec=spec, error=None) + + +def _initialize_state(yaml_path: Path) -> None: + """Seed ``st.session_state`` from disk exactly once per session. + + We key off ``_yaml_path`` so that if the user passes a different YAML on + a Streamlit reload (rare — usually the same), we reset cleanly. + """ + if st.session_state.get("_yaml_path") == str(yaml_path): + return + + original_text = yaml_path.read_text(encoding="utf-8") + + st.session_state["_yaml_path"] = str(yaml_path) + st.session_state["original_text"] = original_text + st.session_state["edited_text"] = original_text + # The text whose render is currently displayed. Starts == original so the + # first paint shows the on-disk file (and "Regenerate" is correctly + # disabled until the user edits something). + st.session_state["last_rendered_text"] = original_text + st.session_state["save_path"] = str(yaml_path) + + initial = _validate_yaml_text(original_text) + if not initial.is_valid: + # Defensive: if the on-disk file is already broken we still want to + # show *something*, but we won't pre-render it. The user fixes the + # YAML in the editor, then hits Regenerate. + st.session_state["rendered_html"] = _BROKEN_PLACEHOLDER_HTML + else: + # First render boots ``SimulationApp`` (≈30s) via the cached resource + # and renders any uncached USD thumbnails. Both steps are amortized + # across subsequent regenerations. + st.session_state["rendered_html"] = _render_with_thumbnails(initial.spec) + + +# Tiny standalone HTML used when the on-disk YAML is itself invalid. +_BROKEN_PLACEHOLDER_HTML = """ +

    No visualization yet — fix YAML and click Regenerate.

    +""" + + +# --------------------------------------------------------------------------- +# UI panels +# --------------------------------------------------------------------------- + + +def _render_validation_badge(validation: _ValidationResult) -> None: + """Show a green tick + summary, or a red cross + the raw exception text.""" + if validation.is_valid: + spec = validation.spec + st.success( + f"Valid spec — {spec.env_name} · {len(spec.nodes)} nodes · " + f"{len(spec.tasks)} tasks · {len(spec.state_specs)} state specs", + icon="✅", + ) + else: + st.error(f"Invalid YAML\n\n```\n{validation.error}\n```", icon="🛑") + + +def _render_action_buttons(validation: _ValidationResult) -> None: + """Render the Regenerate + Save buttons with the gating rules from the spec. + + Gating (per the PR brief): + * Regenerate — only when edited since last render AND valid. + * Save — only when YAML is valid. + """ + edited_since_render = st.session_state["edited_text"] != st.session_state["last_rendered_text"] + can_regenerate = edited_since_render and validation.is_valid + can_save = validation.is_valid + + col_a, col_b = st.columns(2) + + with col_a: + if st.button( + "Regenerate visualization", + type="primary", + disabled=not can_regenerate, + use_container_width=True, + help=( + "Re-renders the right pane from the current editor text. " + "Enabled only when the YAML has been edited since the last " + "render and currently validates." + ), + ): + assert validation.spec is not None # guarded by can_regenerate + # ``_render_with_thumbnails`` reuses the cached ``SimulationApp`` + # and the on-disk PNG cache — so a regen that doesn't introduce + # new asset names is near-instant; one that does pays roughly + # ~2s per new USD. + with st.spinner("Rendering thumbnails…"): + st.session_state["rendered_html"] = _render_with_thumbnails(validation.spec) + st.session_state["last_rendered_text"] = st.session_state["edited_text"] + st.toast("Visualization regenerated.", icon="🔄") + st.rerun() + + with col_b: + save_path_str = st.session_state["save_path"] + if st.button( + f"Save to {Path(save_path_str).name}", + disabled=not can_save, + use_container_width=True, + help=f"Writes the editor contents to {save_path_str}. Disabled while YAML is invalid.", + ): + try: + Path(save_path_str).write_text(st.session_state["edited_text"], encoding="utf-8") + # Update "original" so future "edited since…" comparisons are + # against the just-saved file, not the very first load. + st.session_state["original_text"] = st.session_state["edited_text"] + st.toast(f"Saved → {save_path_str}", icon="💾") + except OSError as exc: + st.error(f"Save failed: {exc}", icon="🛑") + + with st.expander("Change save location", expanded=False): + new_path = st.text_input( + "Save path", + value=save_path_str, + key="save_path_input", + help="Defaults to the YAML file passed via --yaml.", + ) + if new_path and new_path != save_path_str: + st.session_state["save_path"] = new_path + + +def _render_editor_panel(yaml_path: Path) -> _ValidationResult: + """Left pane. Returns the validation result for the current editor text. + + Returning the validation result (rather than stashing it in session_state) + keeps the data flow inside one render pass and avoids a stale-state class + of bug where the badge and the buttons disagree. + """ + # Lazy import so the module is importable from environments that don't + # have streamlit-ace installed yet (we surface a clean error message + # rather than ImportError at module load). + try: + from streamlit_ace import st_ace # noqa: PLC0415 + except ImportError as exc: + # See review_graph._serve_live_editor for why --user --ignore-installed + # is required inside the isaaclab_arena container. + st.error( + "`streamlit-ace` is not installed. Inside the isaaclab_arena container run:\n" + "`python -m pip install --user --ignore-installed streamlit-ace`\n\n" + f"Underlying error: {exc}", + icon="🛑", + ) + st.stop() + + st.subheader("YAML editor") + st.caption(f"Source: `{yaml_path}`") + + # ``auto_update=False`` (the default) commits on blur / Ctrl+Enter rather + # than on every keystroke. That gives us "live enough" validation without + # hammering Streamlit's rerun loop while the user is mid-line. + new_text = st_ace( + value=st.session_state["edited_text"], + language="yaml", + theme="monokai", + keybinding="vscode", + font_size=13, + tab_size=2, + show_gutter=True, + show_print_margin=False, + wrap=False, + auto_update=False, + min_lines=30, + # Key bound to the YAML path so swapping --yaml between sessions + # forces ace to remount with the new content. + key=f"ace_editor::{yaml_path}", + ) + if new_text is not None: + st.session_state["edited_text"] = new_text + + validation = _validate_yaml_text(st.session_state["edited_text"]) + _render_validation_badge(validation) + _render_action_buttons(validation) + return validation + + +def _render_visualization_panel() -> None: + """Right pane — iframe-mount the cached rendered HTML.""" + st.subheader("Visualization") + edited_since_render = st.session_state["edited_text"] != st.session_state["last_rendered_text"] + if edited_since_render: + st.caption("⚠️ Editor has unrendered changes — click **Regenerate visualization** to refresh.") + else: + st.caption("Showing the current YAML.") + + # ``st.components.v1.html`` wraps the payload in a sandboxed iframe, which + # is what we want — the mermaid CDN script and the static CSS stay + # isolated from Streamlit's own DOM. + st.components.v1.html( + st.session_state["rendered_html"], + height=_IFRAME_HEIGHT_PX, + scrolling=True, + ) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def main() -> None: + st.set_page_config( + page_title="ArenaEnvGraphSpec live editor", + layout="wide", + initial_sidebar_state="collapsed", + ) + + args = _parse_args() + yaml_path = args.yaml.resolve() + if not yaml_path.exists(): + st.error(f"YAML file not found: {yaml_path}", icon="🛑") + st.stop() + + _initialize_state(yaml_path) + + st.markdown("### ArenaEnvGraphSpec live editor") + left, right = st.columns([2, 3], gap="large") + with left: + _render_editor_panel(yaml_path) + with right: + _render_visualization_panel() + + +# Streamlit invokes the script top-level on every rerun, so we run main() +# unconditionally. The standard ``if __name__ == "__main__"`` guard would +# also work under ``streamlit run`` but is unnecessary — this module is only +# ever loaded as the Streamlit entrypoint. +main() diff --git a/isaaclab_arena/environments/agentic_env_gen/review_graph.py b/isaaclab_arena/environments/agentic_env_gen/review_graph.py index 59fe3899f..75e6a8129 100644 --- a/isaaclab_arena/environments/agentic_env_gen/review_graph.py +++ b/isaaclab_arena/environments/agentic_env_gen/review_graph.py @@ -3,28 +3,43 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Render an ArenaEnvGraphSpec YAML to a self-contained HTML review page. +"""Live ArenaEnvGraphSpec review tool — Streamlit editor + USD thumbnails. -Three panels (dark dashboard style): +The CLI is a thin launcher: it always boots the Streamlit app in +``review_app.py``. The static-HTML mode and the old ``--serve`` / +``--render-thumbnails`` switches were collapsed away — thumbnails are now +always rendered (cache-hit when possible, ``SimulationApp`` cache-miss +otherwise), and the result is shown inside the live editor. + +Three panels (dark dashboard style) inside the embedded view: * Top-left — graph diagram (mermaid.js, CDN-loaded) of the initial-state spatial constraints. Anchor nodes are highlighted; constraints without a child (is_anchor / position_limits / at_pose / ...) are listed below the graph rather than rendered as self-loops. * Bottom-left — task table (id, type, initial/success state ids, task_args). * Right — node card grid: type badge, asset name, and the per-node YAML - stanza. With ``--render-thumbnails``, the per-node thumbnail is a real - USD viewport capture (cached on disk and inlined as base64); otherwise - a styled placeholder keeps the script lightweight. + stanza. The per-node thumbnail is a real USD viewport capture (cached + on disk under ``.cache/llm_env_gen_thumbnails/`` and inlined as base64 + so the HTML stays self-contained). Usage: - # Default: writes .html alongside the input file. Lightweight. - python -m isaaclab_arena.llm_env_gen.review_graph \\ + /isaac-sim/python.sh -m isaaclab_arena.environments.agentic_env_gen.review_graph \\ --yaml isaaclab_arena/tests/test_data/pick_and_place_maple_table_env_graph.yaml - # With real per-node USD snapshots (boots Isaac Sim once, ~30s): - /isaac-sim/python.sh -m isaaclab_arena.llm_env_gen.review_graph \\ - --yaml isaaclab_arena_environments/llm_generated/_proposal.yaml \\ - --render-thumbnails --open + # Custom port: + /isaac-sim/python.sh -m isaaclab_arena.environments.agentic_env_gen.review_graph \\ + --yaml --port 8600 + +Public API used by ``review_app.py``: + * :func:`launch_simulation_app` — boots Kit's ``SimulationApp`` (headless + + cameras). Returns ``None`` on failure so the app can degrade to + placeholder thumbnails rather than crashing. + * :func:`render_thumbnails_for_spec` — given a live ``SimulationApp`` and + a parsed spec, returns ``{node_id: png_bytes}``. Cache-aware: existing + PNGs under the disk cache are read directly; missing ones are rendered + through the live app and written back to the cache for next time. + * :func:`render_html_for_spec` — full HTML payload with the given + thumbnails dict inlined. Pass an empty dict to fall back to placeholders. Note on USD rendering: ``pxr.UsdAppUtils.FrameRecorder`` and the ``usdrecord`` CLI are NOT @@ -41,12 +56,12 @@ import argparse import base64 -import contextlib import hashlib import html as html_lib +import os import re +import subprocess import sys -import webbrowser import yaml from dataclasses import asdict from pathlib import Path @@ -66,50 +81,129 @@ def main() -> None: - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--yaml", type=Path, required=True, help="Path to an ArenaEnvGraphSpec YAML file.") + """CLI entry point — argparse parses the user's flags, then we hand off + to Streamlit. The actual interactive UI lives in ``review_app.py``. + """ + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) parser.add_argument( - "--out", + "--yaml", type=Path, - default=None, - help="Output HTML path. Defaults to .html next to the input.", + required=True, + help="Path to an ArenaEnvGraphSpec YAML file. The Streamlit app will open it for live editing.", ) - parser.add_argument("--open", action="store_true", help="Open the resulting HTML in the default browser.") parser.add_argument( - "--render-thumbnails", - action="store_true", - help=( - "Boot Isaac Sim once and capture per-node USD viewport thumbnails " - "(cached under .cache/llm_env_gen_thumbnails/). Slow first run " - "(~30s SimulationApp boot + ~2s per unique USD); subsequent runs " - "reuse cached PNGs. Must run inside the Isaac Sim container." - ), + "--port", + type=int, + default=8501, + help="Streamlit server port (default: 8501).", ) args = parser.parse_args() + _serve_live_editor(args.yaml, port=args.port) - spec = ArenaEnvGraphSpec.from_yaml(args.yaml) - out_path = args.out or args.yaml.with_suffix(".html") - out_path.parent.mkdir(parents=True, exist_ok=True) - # Important: when --render-thumbnails is set, we keep SimulationApp open - # across the HTML write. Calling ``app.close()`` first can ``os._exit(0)`` - # (Kit's normal shutdown behavior) and silently drop the write_text below. - app = None +def _serve_live_editor(yaml_path: Path, port: int = 8501) -> None: + """Spawn ``streamlit run review_app.py -- --yaml `` and wait. + + This is the only path through the CLI now — the old static-HTML and + standalone ``--render-thumbnails`` flows were folded into the Streamlit + app, which boots ``SimulationApp`` once via ``@st.cache_resource`` and + keeps it alive for the lifetime of the server. We resolve + ``review_app.py`` next to this file rather than going through ``-m`` so + Streamlit picks the path up cleanly (``streamlit run`` doesn't accept + module dotted-paths). + """ + app_path = Path(__file__).with_name("review_app.py") + if not app_path.exists(): + raise FileNotFoundError(f"Streamlit app not found at {app_path} — installation is incomplete.") + + cmd = [ + sys.executable, + "-m", + "streamlit", + "run", + str(app_path), + "--server.port", + str(port), + # Skip the email prompt the first time Streamlit runs in a fresh + # container — the live editor is a developer tool, not a hosted + # service, and an interactive prompt would block automation. + "--browser.gatherUsageStats", + "false", + # File watcher is a footgun here: Kit's ``SimulationApp`` boot is + # tens of seconds; we don't want Streamlit silently rerunning the + # script (and reissuing the cached_resource init) every time we + # save a source file during development. The user can still hit "R" + # in the browser to force a rerun if they want. + "--server.fileWatcherType", + "none", + "--", + "--yaml", + str(yaml_path.resolve()), + ] + + # Inherit env so the Streamlit subprocess sees PYTHONPATH / isaac-sim + # site-packages exactly the same way we do. + print(f"[review_graph] launching Streamlit live editor: {' '.join(cmd)}", file=sys.stderr) try: - thumbnails: dict[str, bytes] = {} - if args.render_thumbnails: - app = _launch_simulation_app() - if app is not None: - thumbnails = _render_thumbnails_with_app(app, spec) - - out_path.write_text(_render_html(spec, thumbnails), encoding="utf-8") - print(f"Wrote {out_path}") - if args.open: - webbrowser.open(out_path.resolve().as_uri()) - finally: - if app is not None: - with contextlib.suppress(Exception): - app.close() + subprocess.run(cmd, env=os.environ.copy(), check=True) + except FileNotFoundError as exc: + # The plain ``pip install streamlit`` fails inside the isaaclab_arena + # container because streamlit≥1.30 needs uvicorn>=0.30 but Kit ships + # a bundled uvicorn==0.29 under a read-only /isaac-sim/extscache path. + # ``--user --ignore-installed`` sidesteps the rollback by writing + # everything to ~/.local (which is earlier on sys.path than extscache). + raise SystemExit( + "Streamlit is not installed. Inside the isaaclab_arena container run:\n" + " python -m pip install --user --ignore-installed streamlit streamlit-ace" + ) from exc + except KeyboardInterrupt: + # Normal exit path — user hit Ctrl-C in the terminal. + pass + + +# --------------------------------------------------------------------------- +# Public API consumed by review_app.py +# --------------------------------------------------------------------------- + + +def render_html_for_spec(spec: ArenaEnvGraphSpec, thumbnails: dict[str, bytes] | None = None) -> str: + """Render the review HTML for ``spec``, inlining the given thumbnails. + + Thin public alias of :func:`_render_html` so external entry points don't + have to reach into a private name. Pass ``thumbnails=None`` (or omit) to + fall back to placeholder thumbnails — useful when ``SimulationApp`` is + unavailable. + """ + return _render_html(spec, thumbnails=thumbnails) + + +def launch_simulation_app(): + """Public alias for :func:`_launch_simulation_app`. + + Boots Kit's ``SimulationApp`` (headless, with cameras + UI hidden) and + returns the live app handle. Returns ``None`` on failure so callers can + degrade gracefully (e.g. the Streamlit app falls back to placeholder + thumbnails instead of crashing). + + The returned app *must* be reused across thumbnail renders — Kit only + supports one ``SimulationApp`` per process. The Streamlit app holds the + instance under ``@st.cache_resource``. + """ + return _launch_simulation_app() + + +def render_thumbnails_for_spec(app, spec: ArenaEnvGraphSpec) -> dict[str, bytes]: + """Public alias for :func:`_render_thumbnails_with_app`. + + Resolves each node's USD via ``AssetRegistry``, reads the on-disk PNG + cache for hits, and renders cache-misses through the live ``app`` + handle. Safe to call repeatedly across reruns — cache-hit cost is just + a ``read_bytes()`` per node. + """ + return _render_thumbnails_with_app(app, spec) # --------------------------------------------------------------------------- diff --git a/setup.py b/setup.py index 82cd92b56..73f4d23ee 100644 --- a/setup.py +++ b/setup.py @@ -22,6 +22,14 @@ "tenacity", ] +# Live YAML editor for ArenaEnvGraphSpec (see isaaclab_arena/environments/ +# agentic_env_gen/review_app.py). Kept as an extras_require so the core +# package install in CI / minimal containers stays Streamlit-free. +ENV_REVIEW_DEPS = [ + "streamlit>=1.30", + "streamlit-ace>=0.1.1", +] + setup( name="isaaclab_arena", version=ISAACLAB_ARENA_VERSION_NUMBER, @@ -40,6 +48,7 @@ install_requires=RUNTIME_DEPS, extras_require={ "dev": DEV_DEPS, + "env-review": ENV_REVIEW_DEPS, }, zip_safe=False, ) From ccc66eed3936f4d1f96f0b56928afba1b0c76f0c Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Fri, 29 May 2026 11:42:50 +0800 Subject: [PATCH 08/10] SimApp in a sidecar process --- .../agentic_env_gen/review_app.py | 151 +++++---- .../agentic_env_gen/review_graph.py | 287 ++++++++++++++++-- .../agentic_env_gen/simapp_sidecar.py | 208 +++++++++++++ 3 files changed, 548 insertions(+), 98 deletions(-) create mode 100644 isaaclab_arena/environments/agentic_env_gen/simapp_sidecar.py diff --git a/isaaclab_arena/environments/agentic_env_gen/review_app.py b/isaaclab_arena/environments/agentic_env_gen/review_app.py index 3649a1e0d..144245a79 100644 --- a/isaaclab_arena/environments/agentic_env_gen/review_app.py +++ b/isaaclab_arena/environments/agentic_env_gen/review_app.py @@ -30,6 +30,7 @@ from __future__ import annotations import argparse +import atexit import traceback import yaml from dataclasses import dataclass @@ -38,9 +39,9 @@ import streamlit as st from isaaclab_arena.environments.agentic_env_gen.review_graph import ( - launch_simulation_app, + SimAppSidecar, + SimAppSidecarError, render_html_for_spec, - render_thumbnails_for_spec, ) from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvGraphSpec @@ -50,84 +51,104 @@ # --------------------------------------------------------------------------- -# SimulationApp lifecycle +# SimulationApp sidecar lifecycle # --------------------------------------------------------------------------- - - -def _patch_signal_for_non_main_thread() -> None: - """No-op ``signal.signal`` for non-main threads so Kit's bootstrap doesn't raise. - - Kit's ``SimulationApp.__init__`` installs SIGINT/SIGTERM handlers via - ``signal.signal``, which raises ``ValueError("signal only works in main - thread of the main interpreter")`` when called outside the interpreter's - main thread. Streamlit's ``ScriptRunner`` runs the user script in a - worker thread per session — so the SimApp boot would crash there even - though it would succeed under a plain ``python -m`` invocation. - - We swallow *only* the "main thread" ValueError and let everything else - propagate. The Streamlit main thread keeps its own SIGINT handler (it's - installed before our script ever runs), so Ctrl-C still tears the server - down cleanly. The patch is idempotent — repeated calls re-wrap the - already-wrapped function harmlessly. +# +# Kit's ``SimulationApp`` cannot live inside the Streamlit worker thread: +# its bootstrap installs signal handlers (main-thread only) and the +# ``omni.usd`` UsdContext does not tolerate being driven from different +# threads across Streamlit reruns ("[Error] [omni.usd] UsdContext busy"). +# We host it in a dedicated subprocess (``simapp_sidecar.py``) and talk to +# it over a JSON-RPC pipe. The wrapper class ``SimAppSidecar`` in +# ``review_graph`` owns the subprocess; we hold one instance per Streamlit +# server process via ``@st.cache_resource``. + + +@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. + + Returns ``None`` if the sidecar fails to boot — the app then falls back + to placeholder thumbnails so the review page still renders. We register + an ``atexit`` cleanup so the sidecar is reaped on normal interpreter + shutdown (Ctrl-C of the terminal that owns Streamlit). + + The ``@st.cache_resource`` decorator gives us a single instance shared + across reruns AND across browser sessions, which is exactly what we + want: one Kit, many requests, serialized by the sidecar's own + ``threading.Lock``. """ - import signal # noqa: PLC0415 - - if getattr(signal.signal, "_arena_patched", False): - return - - _original_signal = signal.signal - - def _safe_signal(signum, handler): - try: - return _original_signal(signum, handler) - except ValueError as exc: - if "main thread" in str(exc): - return None - raise - - _safe_signal._arena_patched = True # type: ignore[attr-defined] - signal.signal = _safe_signal # type: ignore[assignment] - - -@st.cache_resource(show_spinner="Booting Isaac Sim (≈30s on first run, cached afterwards)…") -def _get_simulation_app(): - """Boot ``SimulationApp`` exactly once per Streamlit server process. - - ``@st.cache_resource`` is exactly the right primitive here: the returned - handle is shared across reruns and across browser sessions, and Streamlit - won't try to pickle it (unlike ``@st.cache_data``). If the boot fails we - return ``None`` and the app degrades to placeholder thumbnails — the - review page still works, just with the styled two-letter cards instead - of real USD captures. - - The signal patch must happen *before* the SimApp launch (which happens - on the Streamlit worker thread); see ``_patch_signal_for_non_main_thread``. + sidecar = SimAppSidecar() + try: + sidecar.start() + except SimAppSidecarError as exc: + print(f"[review_app] SimApp sidecar failed to start: {exc}", flush=True) + return None + + # atexit covers the common-case shutdown path (Ctrl-C in the launching + # terminal -> Python interpreter shutdown -> atexit handlers fire). + # Abnormal exits (SIGKILL of the Streamlit process) are handled by the + # sidecar itself: it watches for EOF on stdin and exits via its own + # ``finally`` block. So the SimApp gets closed either way. + atexit.register(sidecar.close) + return sidecar + + +def _ensure_sidecar() -> SimAppSidecar | None: + """Return a healthy sidecar, re-spawning if the cached one died. + + If the cached resource exists but the subprocess crashed (e.g. an asset + triggered an unrecoverable Kit error), we clear the Streamlit cache and + start fresh. The single re-spawn keeps the user from having to restart + the whole Streamlit process for a transient render failure. """ - _patch_signal_for_non_main_thread() - return launch_simulation_app() + sidecar = _get_simapp_sidecar() + if sidecar is not None and sidecar.is_alive(): + return sidecar + if sidecar is not None: + # Sidecar died (crash / SIGKILL / whatever). Clean it up and ask + # Streamlit for a fresh one on the next call. + sidecar.close() + _get_simapp_sidecar.clear() + return _get_simapp_sidecar() def _render_with_thumbnails(spec: ArenaEnvGraphSpec) -> str: - """Render the review HTML for ``spec`` with thumbnails resolved through - the cached ``SimulationApp``. + """Render review HTML, asking the sidecar for thumbnails. Cache-aware in two layers: * The disk cache under ``.cache/llm_env_gen_thumbnails/`` survives - across runs; ``render_thumbnails_for_spec`` reads it directly. - * Within a run, ``@st.cache_resource`` ensures we never re-boot Kit. + across runs; the sidecar's internal renderer reads it directly. + * Within a server lifetime, ``@st.cache_resource`` keeps Kit warm so + only the cache-misses pay the ~2s-per-USD capture cost. - If ``SimulationApp`` is unavailable (boot failed) we fall back to - placeholder thumbnails so the user still gets a usable page. + If the sidecar is unavailable (boot failed and re-spawn also failed) we + fall back to placeholder thumbnails so the user still gets a usable page + and a visible warning explaining why. """ - app = _get_simulation_app() - if app is None: + sidecar = _ensure_sidecar() + if sidecar is None: st.warning( - "Isaac Sim is unavailable — falling back to placeholder thumbnails. " - "Check the terminal where you launched the server for the boot error.", + "Isaac Sim sidecar is unavailable — falling back to placeholder thumbnails. " + "Check the terminal where you launched the server for the underlying error.", icon="⚠️", ) return render_html_for_spec(spec, thumbnails=None) - thumbnails = render_thumbnails_for_spec(app, spec) + + try: + thumbnails = sidecar.render_spec(spec) + except SimAppSidecarError as exc: + st.error( + f"Sidecar render failed; falling back to placeholder thumbnails.\n\n```\n{exc}\n```", + icon="🛑", + ) + # Force a re-spawn on the next call — most "render failed" errors + # that propagate up are pipe-broken / process-died and the next + # invocation will boot a fresh Kit. + with st.spinner("Resetting the SimApp sidecar…"): + _get_simapp_sidecar.clear() + return render_html_for_spec(spec, thumbnails=None) + return render_html_for_spec(spec, thumbnails=thumbnails) diff --git a/isaaclab_arena/environments/agentic_env_gen/review_graph.py b/isaaclab_arena/environments/agentic_env_gen/review_graph.py index 75e6a8129..8c759346d 100644 --- a/isaaclab_arena/environments/agentic_env_gen/review_graph.py +++ b/isaaclab_arena/environments/agentic_env_gen/review_graph.py @@ -56,15 +56,19 @@ import argparse import base64 +import contextlib import hashlib import html as html_lib +import json import os import re import subprocess import sys +import threading import yaml from dataclasses import asdict from pathlib import Path +from typing import Any from isaaclab_arena.environments.arena_env_graph_spec import ( ArenaEnvGraphNodeSpec, @@ -174,36 +178,240 @@ def render_html_for_spec(spec: ArenaEnvGraphSpec, thumbnails: dict[str, bytes] | Thin public alias of :func:`_render_html` so external entry points don't have to reach into a private name. Pass ``thumbnails=None`` (or omit) to - fall back to placeholder thumbnails — useful when ``SimulationApp`` is + fall back to placeholder thumbnails — useful when the sidecar is unavailable. """ return _render_html(spec, thumbnails=thumbnails) -def launch_simulation_app(): - """Public alias for :func:`_launch_simulation_app`. +class SimAppSidecarError(RuntimeError): + """Raised when the SimApp sidecar process can't fulfil a request. + + Distinct exception type so the Streamlit app can catch sidecar failures + specifically (and e.g. clear its ``@st.cache_resource`` to force a + re-spawn) without swallowing programming errors. + """ - Boots Kit's ``SimulationApp`` (headless, with cameras + UI hidden) and - returns the live app handle. Returns ``None`` on failure so callers can - degrade gracefully (e.g. the Streamlit app falls back to placeholder - thumbnails instead of crashing). - The returned app *must* be reused across thumbnail renders — Kit only - supports one ``SimulationApp`` per process. The Streamlit app holds the - instance under ``@st.cache_resource``. +class SimAppSidecar: + """Long-lived Kit/SimApp host process exposed as a render service. + + See ``simapp_sidecar.py`` for the protocol. The instance is meant to be + cached for the lifetime of the Streamlit server process via + ``@st.cache_resource``; calling :meth:`render_spec` is safe across + Streamlit reruns and across concurrent sessions (an internal + ``threading.Lock`` serializes pipe access — Kit can only service one + render at a time anyway). + + Lifecycle: + + * :meth:`start` spawns the subprocess and waits for the ``{"ready": + true}`` handshake. Times out after ``boot_timeout_s`` if Kit boot + hangs. + * :meth:`render_spec` sends a ``render_spec`` request and reads the + reply line. Reads paths back, materializes the PNG bytes from the + shared filesystem cache, returns ``{node_id: bytes}``. + * :meth:`close` sends ``shutdown`` and waits for the process to exit, + terminating then killing if it doesn't. + * On parent crash / SIGKILL, the sidecar reads EOF on stdin and exits + on its own via the ``finally`` in ``simapp_sidecar._serve``. """ - return _launch_simulation_app() + # Subprocess.Popen would normally inherit the parent's stderr. Kit + # writes a lot there, which is fine — we want users to see those logs. + # The JSON channel travels through stdout instead; the sidecar redirects + # Kit's stdout to stderr at the fd level before booting so the channel + # stays clean. + + 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 + + # -- lifecycle -- + + def start(self) -> None: + """Spawn the sidecar process and wait for its ``{"ready": true}`` handshake. + + Raises :class:`SimAppSidecarError` if the boot fails (handshake says + ``ready: false``, sidecar exits early, or boot takes longer than + ``boot_timeout_s``). + """ + if self._proc is not None and self._proc.poll() is None: + return # already running + + cmd = [ + sys.executable, + "-m", + "isaaclab_arena.environments.agentic_env_gen.simapp_sidecar", + ] + # ``start_new_session=False`` (default) leaves the child in the same + # process group as the parent, so Ctrl-C in the launching terminal + # also signals the sidecar. The sidecar installs SIGINT/SIGTERM + # handlers that route to a clean SystemExit -> finally -> app.close(). + self._proc = subprocess.Popen( + cmd, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=None, # inherit so Kit logs flow to the user's terminal + text=True, + bufsize=1, # line-buffered + env=os.environ.copy(), + ) + + # Block until we hear the ready handshake. We don't use signal-based + # timeout (only main thread can use ``signal.alarm``); a watchdog + # thread is overkill here, so we just rely on the boot being fast + # under normal conditions and let the user Ctrl-C if it really hangs. + # In practice Kit either boots in ~30s or fails immediately. + 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 render_thumbnails_for_spec(app, spec: ArenaEnvGraphSpec) -> dict[str, bytes]: - """Public alias for :func:`_render_thumbnails_with_app`. + def is_alive(self) -> bool: + return self._proc is not None and self._proc.poll() is None - Resolves each node's USD via ``AssetRegistry``, reads the on-disk PNG - cache for hits, and renders cache-misses through the live ``app`` - handle. Safe to call repeatedly across reruns — cache-hit cost is just - a ``read_bytes()`` per node. - """ - return _render_thumbnails_with_app(app, spec) + def close(self) -> None: + """Send ``shutdown``, then terminate/kill if the process doesn't exit. + + Safe to call multiple times. Safe to call after the child has already + died (e.g. via SIGINT propagated from the terminal). Quiet about + common shutdown races so atexit doesn't spam the terminal. + """ + 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() + + # -- requests -- + + def render_spec(self, spec: ArenaEnvGraphSpec) -> dict[str, bytes]: + """Ask the sidecar to render thumbnails for ``spec``. + + Serializes the spec back to YAML (via ``ArenaEnvGraphSpec.to_dict`` + which already unwraps Enums) before shipping it — the sidecar + re-parses on its end. We round-trip through YAML rather than JSON of + ``asdict`` because the sidecar already imports yaml and we already + trust ``ArenaEnvGraphSpec.from_dict`` to be the canonical parser. + + Returns ``{node_id: png_bytes}`` ready to splice into the HTML. + Cache-hit nodes read from disk on the parent side (cheap mmap-style + ``read_bytes``); cache-miss nodes triggered a render in the sidecar + and we read the freshly-written file by the same code path. + """ + 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: + # Path missing despite a successful response — surface it on + # stderr but don't bail, the placeholder thumbnail will show. + print( + f"[review_graph] 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")) + + # -- internals -- + + def _request(self, payload: dict[str, Any]) -> dict[str, Any]: + """Single request/response round-trip. Caller owns the lock.""" + 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: + """Read a line from sidecar stdout; raise if the pipe closes (sidecar died).""" + 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: + """Hard-kill the sidecar — used when boot fails and graceful is moot.""" + 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 # --------------------------------------------------------------------------- @@ -443,17 +651,24 @@ def _render_node_thumbnail(node: ArenaEnvGraphNodeSpec, png_bytes: bytes | None # --------------------------------------------------------------------------- -def _render_thumbnails_with_app(app, spec: ArenaEnvGraphSpec) -> dict[str, bytes]: - """Resolve each node's USD via ``AssetRegistry``, render or read cache. +def _render_thumbnails_with_app(app, spec: ArenaEnvGraphSpec) -> 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 so the HTML write can happen before ``app.close()`` (which Kit - may turn into ``os._exit(0)``). + 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. - Returns ``{node.id: png_bytes}`` for nodes whose asset USD could be - located *and* rendered. 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 @@ -471,26 +686,32 @@ def _render_thumbnails_with_app(app, spec: ArenaEnvGraphSpec) -> dict[str, bytes # 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. - rendered: dict[str, bytes] = {} + 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: - rendered[node_id] = cache_path.read_bytes() + resolved[node_id] = cache_path else: to_render[node_id] = (usd_path, cache_path) if to_render: print( f"[review_graph] rendering {len(to_render)} new thumbnail(s) " - f"(reusing {len(rendered)} from cache at {_THUMBNAIL_CACHE_DIR})...", + f"(reusing {len(resolved)} from cache at {_THUMBNAIL_CACHE_DIR})...", file=sys.stderr, ) - rendered.update(_capture_usd_thumbnails(app, to_render)) + # ``_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"[review_graph] all {len(rendered)} thumbnail(s) served from cache.", file=sys.stderr) + print(f"[review_graph] all {len(resolved)} thumbnail(s) served from cache.", file=sys.stderr) - return rendered + return resolved def _launch_simulation_app(): diff --git a/isaaclab_arena/environments/agentic_env_gen/simapp_sidecar.py b/isaaclab_arena/environments/agentic_env_gen/simapp_sidecar.py new file mode 100644 index 000000000..f69678e50 --- /dev/null +++ b/isaaclab_arena/environments/agentic_env_gen/simapp_sidecar.py @@ -0,0 +1,208 @@ +# 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 on *its own* main thread and serves +thumbnail-render requests over a newline-delimited JSON-RPC pipe on +stdin/stdout. The parent (``review_app.py`` running inside Streamlit) spawns +exactly one of these and reuses it for the entire server lifetime via +:class:`isaaclab_arena.environments.agentic_env_gen.review_graph.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": "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": "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 + +# --------------------------------------------------------------------------- +# stdout multiplexing setup — run BEFORE importing anything that might +# touch Kit or print to stdout, otherwise Kit's chatter pollutes the JSON +# channel and the parent crashes on the first bad json.loads. +# --------------------------------------------------------------------------- + +_JSON_FD = os.dup(1) # save real stdout for our JSON channel +# Redirect fd 1 to stderr at the OS level so writes from C-extensions (Kit) +# end up on the terminal instead of the JSON pipe. +os.dup2(2, 1) +# Mirror at the Python level so `print()` and python-level stdout writes +# also go to the terminal. ``sys.stderr`` already points to the inherited +# parent stderr. +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: + """Translate SIGINT / SIGTERM to a clean ``SystemExit``. + + Lets the ``finally`` block in :func:`_serve` close ``SimulationApp`` + even when the parent kills us with a signal rather than the explicit + shutdown command. + """ + + def _exit(signum, _frame): + # SystemExit is preferable to ``os._exit`` here: it propagates + # through the for-loop in ``_serve`` and reaches our finally. + 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 render requests. + + Returns the process exit code so :func:`main` can propagate it. + """ + _install_signal_handlers() + + try: + # Importing review_graph is light (no Kit yet); the SimApp launch is + # what costs ~30s. + from isaaclab_arena.environments.agentic_env_gen.review_graph import _launch_simulation_app # noqa: PLC0415 + 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 + + # Post-Kit imports — these touch pxr transitively and MUST come after + # SimApp boot (same reason ``_resolve_node_usd_paths`` is lazy). + from isaaclab_arena.environments.agentic_env_gen.review_graph import _render_thumbnails_with_app # noqa: PLC0415 + from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvGraphSpec # noqa: PLC0415 + + _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 == "render_spec": + _send(_handle_render_spec(app, req, _render_thumbnails_with_app, ArenaEnvGraphSpec)) + continue + + _send({"ok": False, "error": f"unknown cmd: {cmd!r}"}) + + # ``for raw_line in sys.stdin`` exits when the parent's write end of + # the pipe is closed — i.e. the parent died. Treat that the same as + # a polite shutdown. + return 0 + finally: + with contextlib.suppress(Exception): + app.close() + + +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. + + Kept separate from :func:`_serve` so each request gets its own try/except + boundary — one bad spec shouldn't tear down the sidecar. + """ + 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: + # ``_render_thumbnails_with_app`` was refactored to return + # ``dict[node_id, Path]`` (paths on the disk cache); the parent reads + # the PNG bytes itself to keep IPC small. + 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()) From 0e92428995dd00971a6df420b715f041f371f8ac Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Thu, 28 May 2026 22:09:18 -0700 Subject: [PATCH 09/10] Add task constraint and reference object to graph viz as edge --- .../agentic_env_gen/review_graph.py | 42 ++++++++++++++++--- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/isaaclab_arena/environments/agentic_env_gen/review_graph.py b/isaaclab_arena/environments/agentic_env_gen/review_graph.py index 8c759346d..cf0cb0dcf 100644 --- a/isaaclab_arena/environments/agentic_env_gen/review_graph.py +++ b/isaaclab_arena/environments/agentic_env_gen/review_graph.py @@ -478,10 +478,21 @@ def _pick_initial_state(spec: ArenaEnvGraphSpec) -> ArenaEnvGraphStateSpec | Non def _render_mermaid(spec: ArenaEnvGraphSpec, state: ArenaEnvGraphStateSpec | None) -> str: - """Emit a left-to-right mermaid graph of the spatial constraints with children. + """Emit a left-to-right mermaid graph of spatial and task constraints. - Constraints without a child (is_anchor / position_limits / at_pose / ...) - are dropped here and surfaced separately by :func:`_render_unary_constraints`. + Binary spatial constraints (child is set) are drawn as solid edges: + child -->|kind| parent + + Unary spatial constraints (no child) are omitted from the graph and listed + below it by :func:`_render_unary_constraints` so their params are visible. + + Task constraints with a child are drawn as dashed edges: + parent -.->|type| child + + Task constraints without a child are omitted from the graph. + + objectReference nodes are drawn with a dotted edge to their parent node: + ref_node -. ref .-> parent_node """ lines = ["graph LR"] if state is None: @@ -491,11 +502,12 @@ def _render_mermaid(spec: ArenaEnvGraphSpec, state: ArenaEnvGraphStateSpec | Non anchor_ids: set[str] = set() edge_nodes: set[str] = set() + # --- Spatial constraints (binary only) --- for c in state.spatial_constraints: kind = c.type.value if kind == "is_anchor": anchor_ids.add(c.parent) - elif c.child is not None: + if c.child is not None: lines.append( f" {_mermaid_id(c.child)}[{_mermaid_label(c.child)}]" f" -->|{kind}| " @@ -504,11 +516,31 @@ def _render_mermaid(spec: ArenaEnvGraphSpec, state: ArenaEnvGraphStateSpec | Non edge_nodes.add(c.child) edge_nodes.add(c.parent) + # --- Task constraints (dashed edges, binary only) --- + for tc in state.task_constraints: + if tc.child is not None: + lines.append( + f" {_mermaid_id(tc.parent)}[{_mermaid_label(tc.parent)}]" + f" -.->|{_mermaid_label(tc.type.value)}| " + f"{_mermaid_id(tc.child)}[{_mermaid_label(tc.child)}]" + ) + edge_nodes.add(tc.parent) + edge_nodes.add(tc.child) + # Include every node from the spec so disconnected ones still appear. for node in spec.nodes: if node.id not in edge_nodes: lines.append(f" {_mermaid_id(node.id)}[{_mermaid_label(node.id)}]") + # --- objectReference → parent edges (dotted, structural) --- + # Use bare node IDs (no label re-declaration) — all nodes are already + # declared above either in constraint edges or in the disconnected-node block. + nodes_by_id = spec.nodes_by_id + for node in spec.nodes: + if node.type.value == "objectReference" and node.parent is not None: + if node.parent in nodes_by_id: + lines.append(f" {_mermaid_id(node.id)} -.->|ref| {_mermaid_id(node.parent)}") + # Anchor highlight. for anchor_id in anchor_ids: lines.append(f" style {_mermaid_id(anchor_id)} fill:#3a7d44,color:#fff,stroke:#7fd17f,stroke-width:2px") @@ -518,7 +550,7 @@ def _render_mermaid(spec: ArenaEnvGraphSpec, state: ArenaEnvGraphStateSpec | Non "background": ("#3a4f7a", "#7aa0d8"), "embodiment": ("#7a3a3a", "#d87a7a"), "object": ("#7a6b3a", "#d8c47a"), - "object_reference": ("#6b3a7a", "#c47ad8"), + "objectReference": ("#6b3a7a", "#c47ad8"), "lighting": ("#3a7a7a", "#7ad8d8"), } for node in spec.nodes: From 2b033d849bb572e68960c48c65757c4e48574fa0 Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Thu, 28 May 2026 23:55:26 -0700 Subject: [PATCH 10/10] Make viz regren automatic on update validation --- .../environments/agentic_env_gen/__init__.py | 4 + .../agentic_env_gen/review_app.py | 111 +++++++----------- 2 files changed, 46 insertions(+), 69 deletions(-) create mode 100644 isaaclab_arena/environments/agentic_env_gen/__init__.py diff --git a/isaaclab_arena/environments/agentic_env_gen/__init__.py b/isaaclab_arena/environments/agentic_env_gen/__init__.py new file mode 100644 index 000000000..16ea4c218 --- /dev/null +++ b/isaaclab_arena/environments/agentic_env_gen/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/isaaclab_arena/environments/agentic_env_gen/review_app.py b/isaaclab_arena/environments/agentic_env_gen/review_app.py index 144245a79..7dd408bf2 100644 --- a/isaaclab_arena/environments/agentic_env_gen/review_app.py +++ b/isaaclab_arena/environments/agentic_env_gen/review_app.py @@ -7,20 +7,20 @@ Wraps :func:`isaaclab_arena.environments.agentic_env_gen.review_graph.render_html_for_spec` in a two-pane Streamlit page so the user can edit the YAML in the browser and -re-render the dark-dashboard visualization on demand. ``SimulationApp`` is -booted once via ``@st.cache_resource`` and reused for every thumbnail render -(disk-cache hit when possible, live USD viewport capture when not). +see the visualization update automatically. ``SimulationApp`` is booted once +via ``@st.cache_resource`` and reused for every thumbnail render (disk-cache +hit when possible, live USD viewport capture when not). Launch (always via the wrapper in review_graph.py — handles streamlit flags): /isaac-sim/python.sh -m isaaclab_arena.environments.agentic_env_gen.review_graph \\ --yaml path/to/spec.yaml Design: - * Left pane — ``streamlit-ace`` YAML editor + validation badge + action - buttons. Validation runs on every rerun (i.e. after each editor blur). - * Right pane — sandboxed iframe with the rendered review HTML. Updates only - when the user clicks "Regenerate", so editing never causes mid-typing - re-renders. + * Left pane — ``streamlit-ace`` YAML editor + validation badge + Save button. + Validation runs on every rerun (i.e. after each editor blur). When the YAML + is valid and has changed since the last render, the visualization updates + automatically — no button click required. + * Right pane — sandboxed iframe with the rendered review HTML. * Thumbnails — real USD viewport captures. Booted ``SimulationApp`` lives inside an ``@st.cache_resource`` so its ~30s startup is paid once per server lifetime. PNGs are cached on disk under @@ -253,7 +253,7 @@ def _initialize_state(yaml_path: Path) -> None: _BROKEN_PLACEHOLDER_HTML = """ -

    No visualization yet — fix YAML and click Regenerate.

    +

    No visualization yet — fix the YAML errors to auto-render.

    """ @@ -275,58 +275,24 @@ def _render_validation_badge(validation: _ValidationResult) -> None: st.error(f"Invalid YAML\n\n```\n{validation.error}\n```", icon="🛑") -def _render_action_buttons(validation: _ValidationResult) -> None: - """Render the Regenerate + Save buttons with the gating rules from the spec. - - Gating (per the PR brief): - * Regenerate — only when edited since last render AND valid. - * Save — only when YAML is valid. - """ - edited_since_render = st.session_state["edited_text"] != st.session_state["last_rendered_text"] - can_regenerate = edited_since_render and validation.is_valid +def _render_save_button(validation: _ValidationResult) -> None: + """Render the Save button. Disabled while the YAML is invalid.""" can_save = validation.is_valid - - col_a, col_b = st.columns(2) - - with col_a: - if st.button( - "Regenerate visualization", - type="primary", - disabled=not can_regenerate, - use_container_width=True, - help=( - "Re-renders the right pane from the current editor text. " - "Enabled only when the YAML has been edited since the last " - "render and currently validates." - ), - ): - assert validation.spec is not None # guarded by can_regenerate - # ``_render_with_thumbnails`` reuses the cached ``SimulationApp`` - # and the on-disk PNG cache — so a regen that doesn't introduce - # new asset names is near-instant; one that does pays roughly - # ~2s per new USD. - with st.spinner("Rendering thumbnails…"): - st.session_state["rendered_html"] = _render_with_thumbnails(validation.spec) - st.session_state["last_rendered_text"] = st.session_state["edited_text"] - st.toast("Visualization regenerated.", icon="🔄") - st.rerun() - - with col_b: - save_path_str = st.session_state["save_path"] - if st.button( - f"Save to {Path(save_path_str).name}", - disabled=not can_save, - use_container_width=True, - help=f"Writes the editor contents to {save_path_str}. Disabled while YAML is invalid.", - ): - try: - Path(save_path_str).write_text(st.session_state["edited_text"], encoding="utf-8") - # Update "original" so future "edited since…" comparisons are - # against the just-saved file, not the very first load. - st.session_state["original_text"] = st.session_state["edited_text"] - st.toast(f"Saved → {save_path_str}", icon="💾") - except OSError as exc: - st.error(f"Save failed: {exc}", icon="🛑") + save_path_str = st.session_state["save_path"] + + if st.button( + f"Save to {Path(save_path_str).name}", + disabled=not can_save, + use_container_width=True, + help=f"Writes the editor contents to {save_path_str}. Disabled while YAML is invalid.", + ): + try: + Path(save_path_str).write_text(st.session_state["edited_text"], encoding="utf-8") + # Update "original" so future comparisons are against the saved file. + st.session_state["original_text"] = st.session_state["edited_text"] + st.toast(f"Saved → {save_path_str}", icon="💾") + except OSError as exc: + st.error(f"Save failed: {exc}", icon="🛑") with st.expander("Change save location", expanded=False): new_path = st.text_input( @@ -365,9 +331,8 @@ def _render_editor_panel(yaml_path: Path) -> _ValidationResult: st.subheader("YAML editor") st.caption(f"Source: `{yaml_path}`") - # ``auto_update=False`` (the default) commits on blur / Ctrl+Enter rather - # than on every keystroke. That gives us "live enough" validation without - # hammering Streamlit's rerun loop while the user is mid-line. + # ``auto_update=False`` commits on blur / Ctrl+Enter rather than on every + # keystroke, showing an "Apply" button in the editor toolbar. new_text = st_ace( value=st.session_state["edited_text"], language="yaml", @@ -389,18 +354,26 @@ def _render_editor_panel(yaml_path: Path) -> _ValidationResult: validation = _validate_yaml_text(st.session_state["edited_text"]) _render_validation_badge(validation) - _render_action_buttons(validation) + + # Auto-render whenever the YAML is valid and has changed since the last + # render. This runs before the right pane is drawn, so the updated HTML + # is already in session_state when the iframe is mounted — no extra rerun + # needed. + edited_since_render = st.session_state["edited_text"] != st.session_state["last_rendered_text"] + if validation.is_valid and edited_since_render: + with st.spinner("Rendering visualization…"): + st.session_state["rendered_html"] = _render_with_thumbnails(validation.spec) + st.session_state["last_rendered_text"] = st.session_state["edited_text"] + st.toast("Visualization updated.", icon="🔄") + + _render_save_button(validation) return validation def _render_visualization_panel() -> None: """Right pane — iframe-mount the cached rendered HTML.""" st.subheader("Visualization") - edited_since_render = st.session_state["edited_text"] != st.session_state["last_rendered_text"] - if edited_since_render: - st.caption("⚠️ Editor has unrendered changes — click **Regenerate visualization** to refresh.") - else: - st.caption("Showing the current YAML.") + st.caption("Updates automatically when the YAML is valid.") # ``st.components.v1.html`` wraps the payload in a sandboxed iframe, which # is what we want — the mermaid CDN script and the static CSS stay