diff --git a/isaaclab_arena_examples/agentic_environment_generation/gui_runner.py b/isaaclab_arena_examples/agentic_environment_generation/gui_runner.py new file mode 100644 index 000000000..8bf2ae3c5 --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/gui_runner.py @@ -0,0 +1,91 @@ +# 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 + +"""CLI launcher for the ArenaEnvInitialGraphSpec live editor. + +Spawns Streamlit with :mod:`~isaaclab_arena_examples.agentic_environment_generation.review_gui.streamlit_ui`. + +Usage: + # Default — prompt-only (empty editor until you generate or paste YAML): + /isaac-sim/python.sh -m isaaclab_arena_examples.agentic_environment_generation.gui_runner + + # Open an existing spec: + /isaac-sim/python.sh -m isaaclab_arena_examples.agentic_environment_generation.gui_runner \\ + --yaml isaaclab_arena/tests/test_data/pick_and_place_maple_table_init_env_graph.yaml + + # Custom port: + /isaac-sim/python.sh -m isaaclab_arena_examples.agentic_environment_generation.gui_runner \\ + --yaml --port 8600 +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path + +_REVIEW_GUI_DIR = Path(__file__).resolve().parent / "review_gui" + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--yaml", + type=Path, + default=None, + help="Optional ArenaEnvInitialGraphSpec YAML to open in the editor.", + ) + parser.add_argument( + "--port", + type=int, + default=8501, + help="Streamlit server port (default: 8501).", + ) + args = parser.parse_args() + serve_live_editor(args.yaml, port=args.port) + + +def serve_live_editor(yaml_path: Path | None, port: int = 8501) -> None: + """Spawn ``streamlit run streamlit_ui.py`` and wait.""" + app_path = _REVIEW_GUI_DIR / "streamlit_ui.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), + "--browser.gatherUsageStats", + "false", + "--server.fileWatcherType", + "none", + "--", + ] + if yaml_path is not None: + cmd.extend(["--yaml", str(yaml_path.resolve())]) + + print(f"[review_gui] launching Streamlit live editor: {' '.join(cmd)}", file=sys.stderr) + try: + subprocess.run(cmd, env=os.environ.copy(), check=True) + except FileNotFoundError as exc: + 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: + pass + + +if __name__ == "__main__": + main() diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/__init__.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/__init__.py new file mode 100644 index 000000000..821e672b4 --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""Live editor and HTML dashboard for :class:`~isaaclab_arena.environments.arena_env_graph_spec.ArenaEnvInitialGraphSpec`.""" diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/__init__.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/__init__.py new file mode 100644 index 000000000..2fd363882 --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""HTML rendering backend for the initial-graph review dashboard.""" diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/dashboard.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/dashboard.py new file mode 100644 index 000000000..39a135dd0 --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/dashboard.py @@ -0,0 +1,60 @@ +# 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 + +from __future__ import annotations + +import html as html_lib + +from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvInitialGraphSpec +from isaaclab_arena_examples.agentic_environment_generation.review_gui.render.mermaid_graph import render_mermaid_graph +from isaaclab_arena_examples.agentic_environment_generation.review_gui.render.panels import ( + render_node_cards, + render_tasks_table, + render_unary_constraints, +) +from isaaclab_arena_examples.agentic_environment_generation.review_gui.render.styles import DASHBOARD_CSS + + +def render_dashboard_html(spec: ArenaEnvInitialGraphSpec) -> str: + """Render the self-contained review dashboard HTML for ``spec``.""" + initial_state = spec.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 · initial state: {html_lib.escape(initial_state.id)}

+
+
+
+

Nodes

+
{render_node_cards(spec)}
+
+
+

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

+
+
+
{render_mermaid_graph(spec, initial_state)}
+
+ +
+
+
+

Tasks

+ {render_tasks_table(spec)} +
+
+ + + +""" diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/mermaid_graph.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/mermaid_graph.py new file mode 100644 index 000000000..7ec9c3fb0 --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/mermaid_graph.py @@ -0,0 +1,96 @@ +# 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 + +from __future__ import annotations + +import re + +from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvInitialGraphSpec +from isaaclab_arena.environments.arena_env_graph_types import ArenaEnvGraphStateSpec + +_MERMAID_ID_SAFE = re.compile(r"[^A-Za-z0-9_]") + + +def render_mermaid_graph(spec: ArenaEnvInitialGraphSpec, state: ArenaEnvGraphStateSpec) -> str: + """Emit a left-to-right mermaid graph of spatial and task constraints. + + Binary spatial constraints (reference is set) are drawn as solid edges: + subject -->|kind| reference + + Unary spatial constraints (no reference) are omitted from the graph and + listed to its right by :func:`render_unary_constraints` so their params are + visible. + + Task constraints with a child are drawn as dashed edges: + parent -.->|type| child + + object_reference nodes are drawn with a dotted edge to their parent node: + ref_node -. ref .-> parent_node + """ + lines = ["graph LR"] + + anchor_ids: set[str] = set() + edge_nodes: set[str] = set() + + for constraint in state.spatial_constraints: + kind = constraint.kind + if kind == "is_anchor": + anchor_ids.add(constraint.subject) + if constraint.reference is not None: + lines.append( + f" {_mermaid_id(constraint.subject)}[{_mermaid_label(constraint.subject)}]" + f" -->|{kind}| " + f"{_mermaid_id(constraint.reference)}[{_mermaid_label(constraint.reference)}]" + ) + edge_nodes.add(constraint.subject) + edge_nodes.add(constraint.reference) + + for task_constraint in state.task_constraints: + if task_constraint.child is not None: + lines.append( + f" {_mermaid_id(task_constraint.parent)}[{_mermaid_label(task_constraint.parent)}]" + f" -.->|{_mermaid_label(task_constraint.type.value)}| " + f"{_mermaid_id(task_constraint.child)}[{_mermaid_label(task_constraint.child)}]" + ) + edge_nodes.add(task_constraint.parent) + edge_nodes.add(task_constraint.child) + + for node in spec.nodes: + if node.id not in edge_nodes: + lines.append(f" {_mermaid_id(node.id)}[{_mermaid_label(node.id)}]") + + nodes_by_id = spec.nodes_by_id + for node in spec.nodes: + if node.type.value == "object_reference" 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)}") + + for anchor_id in anchor_ids: + lines.append(f" style {_mermaid_id(anchor_id)} fill:#3a7d44,color:#fff,stroke:#7fd17f,stroke-width:2px") + + 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 + 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) + + +def _mermaid_id(value: str) -> str: + """Mermaid node identifiers must be alphanumeric / underscore.""" + return _MERMAID_ID_SAFE.sub("_", value) + + +def _mermaid_label(value: str) -> str: + """Escape mermaid-significant characters inside node labels.""" + return value.replace('"', """).replace("|", "|") diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/panels.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/panels.py new file mode 100644 index 000000000..6ab00637d --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/panels.py @@ -0,0 +1,83 @@ +# 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 + +from __future__ import annotations + +import html as html_lib +import yaml + +from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvInitialGraphSpec +from isaaclab_arena.environments.arena_env_graph_types import ArenaEnvGraphNodeSpec, ArenaEnvGraphStateSpec +from isaaclab_arena_examples.agentic_environment_generation.review_gui.render.thumbnails import ( + render_placeholder_thumbnail, +) + + +def render_unary_constraints(state: ArenaEnvGraphStateSpec) -> str: + """List constraints without a reference beside the spatial graph.""" + rows = [] + for constraint in state.spatial_constraints: + if constraint.reference is not None: + continue + params = ( + " {html_lib.escape(yaml.safe_dump(constraint.params, default_flow_style=True).rstrip())}' + if constraint.params + else "" + ) + rows.append( + f'
  • {html_lib.escape(constraint.kind)}' + f" on {html_lib.escape(constraint.subject)}{params}
  • " + ) + if not rows: + return '

    No unary constraints.

    ' + return ( + f'

    Unary constraints ({len(rows)})

    ' + f'' + ) + + +def render_tasks_table(spec: ArenaEnvInitialGraphSpec) -> str: + """Render task rows as an HTML table for the dashboard tasks panel.""" + if not spec.tasks: + return "

    No tasks defined.

    " + rows = [] + for index, task in enumerate(spec.tasks): + params_str = yaml.safe_dump(task.params, sort_keys=False).rstrip() if task.params else "(empty)" + description = html_lib.escape(task.description or "") + rows.append( + "" + f"{index}" + f'{html_lib.escape(task.kind)}' + f"{description}" + f"
    {html_lib.escape(params_str)}
    " + "" + ) + return ( + "" + "" + f"{''.join(rows)}" + "
    #kinddescriptionparams
    " + ) + + +def render_node_cards(spec: ArenaEnvInitialGraphSpec) -> str: + """Render one card per graph node for the dashboard nodes panel.""" + return "\n".join(render_node_card(node) for node in spec.nodes) + + +def render_node_card(node: ArenaEnvGraphNodeSpec) -> str: + """Render a single node card with placeholder thumbnail and YAML dump.""" + node_dict = node.model_dump(mode="json", exclude_none=True) + node_yaml = yaml.safe_dump(node_dict, sort_keys=False).rstrip() + thumb = render_placeholder_thumbnail(node) + return f"""
    + {thumb} +
    +
    {html_lib.escape(node.id)}
    + {html_lib.escape(node.type.value)} +
    +
    {html_lib.escape(node_yaml)}
    +
    """ diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/styles.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/styles.py new file mode 100644 index 000000000..ece136de4 --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/styles.py @@ -0,0 +1,77 @@ +# 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 + +"""Dark-theme CSS for the review GUI embedded HTML dashboard.""" + +DASHBOARD_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: flex; flex-direction: column; gap: 16px; } +.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; } +.graph-row { display: grid; grid-template-columns: minmax(0, 1fr) minmax(220px, 300px); gap: 16px; align-items: start; } +.graph-mermaid { min-width: 0; } +.graph-unary { background: var(--bg-elev2); border: 1px solid var(--border); border-radius: 6px; padding: 12px; } +.unary-heading { margin: 0 0 10px; font-size: 13px; font-weight: 600; letter-spacing: 0.02em; } +.unary-heading .muted { font-weight: 400; } +.unary-list { margin: 0; padding-left: 18px; list-style: disc; color: var(--fg); } +.unary-list li { padding: 4px 0; font-size: 12px; } +.unary-empty { margin: 0; font-size: 12px; } +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; margin: 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(3, minmax(0, 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; 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; } +.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; } +""" diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/thumbnails.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/thumbnails.py new file mode 100644 index 000000000..f2742ee8f --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/render/thumbnails.py @@ -0,0 +1,20 @@ +# 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 + +from __future__ import annotations + +import html as html_lib + +from isaaclab_arena.environments.arena_env_graph_types import ArenaEnvGraphNodeSpec + + +# TODO(qianl): Replace placeholder thumbnails with sim-rendered snapshots . +def render_placeholder_thumbnail(node: ArenaEnvGraphNodeSpec) -> str: + """Per-node placeholder thumbnail — two-letter initial.""" + initial = (node.name[:2] if node.name else "?").upper() + return f"""
    + {html_lib.escape(initial)} + {html_lib.escape(node.name)} +
    """ diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py new file mode 100644 index 000000000..c3b5b0a59 --- /dev/null +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py @@ -0,0 +1,454 @@ +# 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 + +"""Streamlit UI for the initial-graph live editor. + +Launch via :mod:`~isaaclab_arena_examples.agentic_environment_generation.gui_runner`: + + # Default — prompt-only (empty editor until you generate or paste YAML): + /isaac-sim/python.sh -m isaaclab_arena_examples.agentic_environment_generation.gui_runner + + # Open an existing spec: + /isaac-sim/python.sh -m isaaclab_arena_examples.agentic_environment_generation.gui_runner \\ + --yaml isaaclab_arena/tests/test_data/pick_and_place_maple_table_init_env_graph.yaml + +Natural-language generation calls the LLM from Streamlit (``NV_API_KEY``) and +compiles the returned intent in-process with :class:`IntentCompiler`. +""" + +from __future__ import annotations + +import argparse +import traceback +import yaml +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import streamlit as st + +from isaaclab_arena.agentic_environment_generation.asset_matcher import ASSET_ERROR_STAGES +from isaaclab_arena.agentic_environment_generation.environment_generation_agent import ( + AssetCatalogue, + EnvironmentGenerationAgent, + RelationCatalogue, + TaskCatalogue, + build_asset_catalogue, + build_relation_catalogue, + build_task_catalogue, +) +from isaaclab_arena.agentic_environment_generation.intent_compiler import IntentCompiler +from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvInitialGraphSpec +from isaaclab_arena_examples.agentic_environment_generation.review_gui.render.dashboard import render_dashboard_html + +_IFRAME_HEIGHT_PX = 1100 + +_DEFAULT_GENERATION_PROMPT = ( + "franka pick up avocado from the maple table and place it into a bowl on the table. " + "there are other veggies on the table as distractor" +) + +_BROKEN_PLACEHOLDER_HTML = """ +

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

    +""" + +_DEFAULT_SAVE_PATH = "isaaclab_arena_environments/llm_generated/generated_spec.yaml" + + +@dataclass +class ValidationResult: + """Outcome of parsing and validating YAML text as an initial graph spec.""" + + spec: ArenaEnvInitialGraphSpec | None + error: str | None + + @property + def is_valid(self) -> bool: + """True when ``spec`` parsed successfully.""" + return self.spec is not None + + +@dataclass +class CatalogueBundle: + """Asset/relation/task vocabulary for the env-generation agent.""" + + asset_catalogue: AssetCatalogue + relation_catalogue: RelationCatalogue + task_catalogue: TaskCatalogue + + +@st.cache_resource(show_spinner="Building asset catalogues (first run)…") +def _get_catalogue_bundle() -> CatalogueBundle: + """Build and cache registry-backed catalogues for LLM prompt assembly.""" + return CatalogueBundle( + asset_catalogue=build_asset_catalogue(), + relation_catalogue=build_relation_catalogue(), + task_catalogue=build_task_catalogue(), + ) + + +def validate_yaml_text(text: str) -> ValidationResult: + """Parse ``text`` as YAML and validate it as an :class:`ArenaEnvInitialGraphSpec`.""" + cached_text = st.session_state.get("_validation_text") + cached_result = st.session_state.get("_validation_result") + if cached_text == text and isinstance(cached_result, ValidationResult): + return cached_result + + if not text.strip(): + result = ValidationResult(spec=None, error=None) + else: + try: + raw = yaml.safe_load(text) + if raw is None: + result = ValidationResult(spec=None, error="YAML is empty") + elif not isinstance(raw, dict): + result = ValidationResult(spec=None, error=f"Expected mapping, got {type(raw).__name__}") + else: + spec = ArenaEnvInitialGraphSpec.from_dict(raw) + result = ValidationResult(spec=spec, error=None) + except Exception: + result = ValidationResult(spec=None, error=traceback.format_exc()) + + st.session_state["_validation_text"] = text + st.session_state["_validation_result"] = result + return result + + +def _get_generation_agent() -> EnvironmentGenerationAgent | None: + """Lazy-init the LLM agent when ``NV_API_KEY`` is available.""" + if st.session_state.get("generation_agent_error"): + return None + agent = st.session_state.get("generation_agent") + if agent is not None: + return agent + try: + agent = EnvironmentGenerationAgent() + except AssertionError as exc: + st.session_state["generation_agent_error"] = str(exc) + return None + except Exception as exc: + st.session_state["generation_agent_error"] = f"{type(exc).__name__}: {exc}" + return None + st.session_state["generation_agent"] = agent + st.session_state.pop("generation_agent_error", None) + return agent + + +def _format_trace_lines(trace: list[dict[str, Any]], *, errors_only: bool = False) -> str: + """Format intent-compiler trace events as fixed-width log lines.""" + error_stages = ASSET_ERROR_STAGES | IntentCompiler._ERROR_TRACE_STAGES + lines: list[str] = [] + for event in trace: + stage = event.get("stage", "") + if errors_only and stage not in error_stages: + continue + chosen = event.get("chosen") + chosen_str = chosen if chosen is not None else "" + note = event.get("note") or "" + note_str = f" [{note}]" if note else "" + lines.append(f"{stage:34s} {event.get('query', ''):24s} -> {chosen_str}{note_str}") + return "\n".join(lines) + + +def _apply_generated_yaml(yaml_text: str, *, spec: ArenaEnvInitialGraphSpec | None = None) -> None: + """Push compiled spec YAML into the editor and sync the dashboard preview.""" + st.session_state["edited_text"] = yaml_text + st.session_state["editor_version"] = st.session_state.get("editor_version", 0) + 1 + st.session_state["last_rendered_text"] = yaml_text + if spec is not None: + st.session_state["rendered_html"] = render_dashboard_html(spec) + st.session_state["_validation_text"] = yaml_text + st.session_state["_validation_result"] = ValidationResult(spec=spec, error=None) + else: + st.session_state["rendered_html"] = "" + st.session_state.pop("_validation_text", None) + st.session_state.pop("_validation_result", None) + + +def run_generation_pipeline(prompt: str) -> tuple[bool, str]: + """Call the LLM, compile intent in-process, and load YAML into the editor.""" + prompt = prompt.strip() + if not prompt: + return False, "Enter a prompt describing the environment." + + agent = _get_generation_agent() + if agent is None: + err = st.session_state.get( + "generation_agent_error", + "Set NV_API_KEY in the environment before generating specs.", + ) + return False, err + + try: + catalogues = _get_catalogue_bundle() + except Exception: + return False, traceback.format_exc() + + try: + intent, _raw = agent.generate_spec( + prompt, + asset_catalog=catalogues.asset_catalogue, + relation_catalog=catalogues.relation_catalogue, + task_catalog=catalogues.task_catalogue, + ) + except Exception: + return False, traceback.format_exc() + + try: + compiler = IntentCompiler() + spec = compiler.compile(intent) + yaml_text = yaml.safe_dump(spec.to_dict(), sort_keys=False) + trace = [asdict(event) for event in compiler.trace] + has_resolution_errors = compiler.has_resolution_errors + reasoning = intent.reasoning + except Exception: + return False, traceback.format_exc() + + _apply_generated_yaml(yaml_text, spec=spec) + + if reasoning: + st.session_state["last_generation_reasoning"] = reasoning + if trace: + st.session_state["last_generation_trace"] = trace + + if has_resolution_errors: + error_trace = _format_trace_lines(trace, errors_only=True) + return ( + True, + ( + "Spec generated with resolution warnings — review the trace below and edit the YAML as needed.\n\n" + f"{error_trace}" + ), + ) + + return True, "Spec generated and loaded into the YAML editor." + + +def parse_args() -> argparse.Namespace: + """Parse Streamlit CLI args forwarded after ``--`` by :mod:`gui_runner`.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--yaml", + type=Path, + default=None, + help="Optional path to an ArenaEnvInitialGraphSpec YAML to open in the editor.", + ) + return parser.parse_args() + + +def initialize_state(yaml_path: Path | None) -> None: + """Seed ``st.session_state`` from disk exactly once per session.""" + session_key = str(yaml_path.resolve()) if yaml_path is not None else "" + if st.session_state.get("_yaml_path") == session_key: + return + + st.session_state["_yaml_path"] = session_key + st.session_state.setdefault("generation_prompt", _DEFAULT_GENERATION_PROMPT) + st.session_state.setdefault("editor_version", 0) + st.session_state.pop("_validation_text", None) + st.session_state.pop("_validation_result", None) + + if yaml_path is None: + st.session_state["original_text"] = "" + st.session_state["edited_text"] = "" + st.session_state["last_rendered_text"] = "" + st.session_state["rendered_html"] = "" + st.session_state["save_path"] = _DEFAULT_SAVE_PATH + return + + original_text = yaml_path.read_text(encoding="utf-8") + + st.session_state["original_text"] = original_text + st.session_state["edited_text"] = original_text + st.session_state["last_rendered_text"] = "" + st.session_state["rendered_html"] = "" + st.session_state["save_path"] = str(yaml_path) + + +def render_validation_badge(validation: ValidationResult) -> None: + """Show a success or error badge for the current editor YAML.""" + if validation.spec is None and validation.error is None: + return + if validation.is_valid: + spec = validation.spec + st.success( + f"Valid spec — {spec.env_name} · {len(spec.nodes)} nodes · " + f"{len(spec.tasks)} tasks · initial state: {spec.initial_state_spec.id}", + icon="✅", + ) + else: + st.error(f"Invalid YAML\n\n```\n{validation.error}\n```", icon="🛑") + + +def render_save_button(validation: ValidationResult) -> None: + """Render save controls and optional save-path editor.""" + can_save = validation.is_valid + save_path_str = st.session_state["save_path"] + save_label = f"Save to {Path(save_path_str).name}" if save_path_str else "Save YAML" + + if st.button( + save_label, + disabled=not can_save, + use_container_width=True, + help=f"Writes the editor contents to {save_path_str}. Disabled while YAML is invalid.", + ): + try: + out_path = Path(save_path_str) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(st.session_state["edited_text"], encoding="utf-8") + 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, or a generated-spec path when none was given.", + ) + if new_path and new_path != save_path_str: + st.session_state["save_path"] = new_path + + +def render_editor_panel(yaml_path: Path | None) -> ValidationResult: + """Render the ACE YAML editor and refresh the dashboard when text changes.""" + try: + from streamlit_ace import st_ace # noqa: PLC0415 + except ImportError as exc: + 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") + if yaml_path is not None: + st.caption(f"Source: `{yaml_path}`") + else: + st.caption("No file loaded — generate a spec or paste YAML.") + + editor_key = str(yaml_path) if yaml_path is not None else "new" + 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=f"ace_editor::{editor_key}::{st.session_state.get('editor_version', 0)}", + ) + 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) + + edited_since_render = st.session_state["edited_text"] != st.session_state["last_rendered_text"] + if edited_since_render: + if validation.is_valid: + with st.spinner("Rendering visualization…"): + st.session_state["rendered_html"] = render_dashboard_html(validation.spec) + else: + st.session_state["rendered_html"] = _BROKEN_PLACEHOLDER_HTML + st.session_state["last_rendered_text"] = st.session_state["edited_text"] + if validation.is_valid: + st.toast("Visualization updated.", icon="🔄") + + render_save_button(validation) + return validation + + +def render_generation_panel() -> None: + """Prompt input and generate-spec controls (top of the left column).""" + st.subheader("Generate from prompt") + st.caption("Calls the env-generation agent (LLM) then compiles intent in-process.") + + prompt = st.text_area( + "Prompt", + value=st.session_state.get("generation_prompt", _DEFAULT_GENERATION_PROMPT), + height=120, + placeholder="Describe the robot task, scene, objects, and distractors…", + ) + st.session_state["generation_prompt"] = prompt + + agent_error = st.session_state.get("generation_agent_error") + if agent_error: + st.info(f"LLM agent unavailable: {agent_error}", icon="ℹ️") + + if st.button("Generate & compile", type="primary", use_container_width=True): + with st.spinner("Generating spec (LLM call + intent compile)…"): + ok, message = run_generation_pipeline(st.session_state["generation_prompt"]) + if ok: + if "resolution warnings" in message: + st.warning(message, icon="⚠️") + else: + st.success(message, icon="✅") + st.rerun() + else: + st.error(f"Generation failed\n\n```\n{message}\n```", icon="🛑") + + reasoning = st.session_state.get("last_generation_reasoning") + if reasoning: + with st.expander("Agent reasoning (last run)", expanded=False): + st.markdown(reasoning) + + trace = st.session_state.get("last_generation_trace") + if trace: + with st.expander("Resolution trace (last run)", expanded=False): + st.code(_format_trace_lines(trace), language=None) + + +def render_visualization_panel() -> None: + """Embed the rendered dashboard HTML in the right-hand column.""" + st.subheader("Visualization") + if not st.session_state.get("last_rendered_text", "").strip(): + st.caption("Generate or enter valid YAML to see the visualization.") + return + + st.caption("Updates automatically when the YAML is valid.") + st.components.v1.html( + st.session_state["rendered_html"], + height=_IFRAME_HEIGHT_PX, + scrolling=True, + ) + + +def main() -> None: + """Build the two-column Streamlit layout for generation, editing, and preview.""" + st.set_page_config( + page_title="ArenaEnvInitialGraphSpec live editor", + layout="wide", + initial_sidebar_state="collapsed", + ) + + args = parse_args() + yaml_path = args.yaml.resolve() if args.yaml is not None else None + if yaml_path is not None and not yaml_path.exists(): + st.error(f"YAML file not found: {yaml_path}", icon="🛑") + st.stop() + + initialize_state(yaml_path) + + st.markdown("### ArenaEnvInitialGraphSpec live editor") + left, right = st.columns([2, 3], gap="large") + with left: + render_generation_panel() + render_editor_panel(yaml_path) + with right: + render_visualization_panel() + + +main() diff --git a/setup.py b/setup.py index 582f669ec..c411f629b 100644 --- a/setup.py +++ b/setup.py @@ -26,6 +26,8 @@ "jupyter", "debugpy", "tenacity", + "streamlit>=1.30", + "streamlit-ace>=0.1.1", ] setup(