Add SimReady search to agentic environment generation.

isaaclab_arena/agentic_environment_generation/environment_generation_agent.py

@@ -12,6 +12,15 @@
 
 from isaaclab_arena.agentic_environment_generation.inference_backend import InferenceBackend
 from isaaclab_arena.agentic_environment_generation.prim_path_inference import PrimPathInference
+from isaaclab_arena.agentic_environment_generation.prompt_normalization import (
+    PromptNormalizationInference,
+    format_normalized_prompt_block,
+)
+from isaaclab_arena.agentic_environment_generation.simready_asset_search import (
+    SimReadyCandidateCatalogue,
+    SimReadySearchConfig,
+    search_simready_objects,
+)
 from isaaclab_arena.agentic_environment_generation.spec_inference import SpecInference
 from isaaclab_arena.agentic_environment_generation.spec_validation import required_task_init_param_names
 from isaaclab_arena.assets.registries import AssetRegistry, ObjectRelationLibraryRegistry, TaskRegistry
@@ -34,6 +43,9 @@ def __init__(
         temperature: float = 0.2,
         max_tokens: int = 4096,
         max_retries: int = 3,
+        *,
+        enable_simready_search: bool = False,
+        simready_config: SimReadySearchConfig | None = None,
     ):
         """Configure the OpenAI-compatible client and validate the model.
 
@@ -51,6 +63,8 @@ def __init__(
             max_retries: Number of additional attempts after a recoverable failure
                 (network errors, timeouts, empty responses, malformed JSON). Each
                 retry is a fresh API call.
+            enable_simready_search: When ``True``, run SimReady search on normalized object phrases.
+            simready_config: Optional SimReady search configuration.
         """
         inference_backend = InferenceBackend(
             api_key=api_key,
@@ -60,8 +74,11 @@ def __init__(
             max_tokens=max_tokens,
             max_retries=max_retries,
         )
+        self.prompt_normalization = PromptNormalizationInference(inference_backend)
         self.spec_inference = SpecInference(inference_backend)
         self.prim_path_inference = PrimPathInference(inference_backend)
+        self.enable_simready_search = enable_simready_search
+        self.simready_config = simready_config or SimReadySearchConfig(enabled=enable_simready_search)
         self._traces: list[str] = []
 
     @property
@@ -75,6 +92,8 @@ def generate_spec(
         asset_catalog: AssetCatalogue | None = None,
         relation_catalog: RelationCatalogue | None = None,
         task_catalog: TaskCatalogue | None = None,
+        *,
+        enable_simready_search: bool | None = None,
     ) -> tuple[ArenaEnvGraphSpec | None, dict[str, Any] | None]:
         """Call the model with user prompt and return the parsed ArenaEnvGraphSpec.
 
@@ -86,13 +105,40 @@ def generate_spec(
                 from the live ``ObjectRelationLibraryRegistry``.
             task_catalog: Pre-built task vocabulary. When ``None``, built from
                 ``TaskRegistry`` tasks marked ``@agent_ready``.
+            enable_simready_search: Override the agent-level SimReady search flag.
 
         Returns:
             A ``(spec, data)`` tuple. On success, ``spec`` is validated and
             ``data`` is None. On failure, ``spec`` is None and ``data`` is the corresponding JSON dict.
             When validation fails, ``agent.traces`` holds the diagnostic trace.
         """
         self._traces = []
+        normalized = self.prompt_normalization.infer(prompt, self._traces)
+        if normalized is None:
+            return None, None
+
+        normalized_block = format_normalized_prompt_block(normalized)
+        self._traces.append(normalized_block)
+
+        use_simready = self.enable_simready_search if enable_simready_search is None else enable_simready_search
+        simready_catalog: SimReadyCandidateCatalogue | None = None
+        if use_simready:
+            simready_config = SimReadySearchConfig(
+                enabled=True,
+                source=self.simready_config.source,
+                s3_url=self.simready_config.s3_url,
+                service_url=self.simready_config.service_url,
+                project_config_path=self.simready_config.project_config_path,
+                indexed_path=self.simready_config.indexed_path,
+                indexed_directory_type=self.simready_config.indexed_directory_type,
+                max_results_per_object=self.simready_config.max_results_per_object,
+                use_service_fallback=self.simready_config.use_service_fallback,
+            )
+            simready_catalog = search_simready_objects(normalized.objects, simready_config, self._traces)
+            simready_block = simready_catalog.to_catalog_string()
+            if simready_block:
+                self._traces.append(simready_block)
+
         asset_catalog = asset_catalog or build_asset_catalogue()
         relation_catalog = relation_catalog or build_relation_catalogue()
         task_catalog = task_catalog or build_task_catalogue()
@@ -102,6 +148,8 @@ def generate_spec(
             asset_catalog=asset_catalog,
             relation_catalog=relation_catalog,
             task_catalog=task_catalog,
+            normalized_prompt_block=normalized_block,
+            simready_candidate_catalog=simready_catalog,
         )
         if spec is None:
             return None, data

isaaclab_arena/agentic_environment_generation/prompt_normalization.py

@@ -0,0 +1,132 @@
+# 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
+
+"""LLM inference to normalize user prompts before spec generation."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field, ValidationError
+
+from isaaclab_arena.agentic_environment_generation.inference_backend import (
+    InferenceBackend,
+    StructuredOutputRequest,
+    build_strict_schema,
+)
+from isaaclab_arena.agentic_environment_generation.spec_validation import format_validation_error
+
+
+class NormalizedPromptDescriptions(BaseModel):
+    """Structured descriptions for each ArenaEnvGraphSpec section."""
+
+    env_name: str = Field(
+        min_length=1,
+        description="Short snake_case label summarizing the scene and tasks.",
+    )
+    embodiment: str = Field(
+        min_length=1,
+        description="Robot/embodiment intent in plain language (family, control mode, cameras).",
+    )
+    background: str = Field(
+        min_length=1,
+        description="Static scene/background intent in plain language.",
+    )
+    object_references: str = Field(
+        default="",
+        description=(
+            "Optional named sub-parts inside a multi-prim background (e.g. counter top, "
+            "fridge door) that should become object_reference nodes; empty when none are needed. "
+            "Do not list the background itself (e.g. 'maple table') here."
+        ),
+    )
+    objects: list[str] = Field(
+        default_factory=list,
+        description=(
+            "One short search phrase per manipulable object or distractor, e.g. 'red hammer' or 'ceramic bowl'."
+        ),
+    )
+    relations: str = Field(
+        default="",
+        description="Spatial layout intent in plain language (on/next_to/anchor).",
+    )
+    task: str = Field(
+        min_length=1,
+        description="Overall manipulation task intent in plain language.",
+    )
+
+
+def format_normalized_prompt_block(normalized: NormalizedPromptDescriptions) -> str:
+    """Format normalized descriptions for downstream LLM prompts."""
+    object_lines = "\n".join(f"- {phrase}" for phrase in normalized.objects) or "- (none)"
+    refs = normalized.object_references.strip() or "(none)"
+    relations = normalized.relations.strip() or "(none)"
+    return (
+        "NORMALIZED PROMPT:\n"
+        f"env_name: {normalized.env_name}\n"
+        f"embodiment: {normalized.embodiment}\n"
+        f"background: {normalized.background}\n"
+        f"object_references: {refs}\n"
+        f"objects:\n{object_lines}\n"
+        f"relations: {relations}\n"
+        f"task: {normalized.task}"
+    )
+
+
+class PromptNormalizationInference:
+    """Natural-language prompt -> normalized section descriptions."""
+
+    def __init__(self, inference_backend: InferenceBackend):
+        """Wire prompt normalization to a structured-output backend.
+
+        Args:
+            inference_backend: Shared LLM client for JSON-schema completion requests.
+        """
+        self._inference_backend = inference_backend
+        self._schema = build_strict_schema(NormalizedPromptDescriptions)
+
+    def infer(self, prompt: str, traces: list[str]) -> NormalizedPromptDescriptions | None:
+        """Normalize a user prompt into section descriptions for later inference passes.
+
+        Args:
+            prompt: End-user environment description.
+            traces: Accumulator for validation error lines, extended in place on failure.
+
+        Returns:
+            Validated normalized descriptions on success, otherwise ``None``.
+        """
+        data = self._inference_backend.run_json(
+            StructuredOutputRequest(
+                schema_name="NormalizedPromptDescriptions",
+                schema=self._schema,
+                system=self._system_prompt(),
+                user=f"USER PROMPT:\n{prompt.strip()}",
+                retry_label="normalize_prompt",
+            )
+        )
+        try:
+            return NormalizedPromptDescriptions.model_validate(data)
+        except ValidationError as exc:
+            traces.extend(format_validation_error(exc))
+            return None
+
+    @staticmethod
+    def _system_prompt() -> str:
+        return """\
+You normalize robot environment-generation prompts into structured descriptions.
+Do not choose registry asset names. Describe intent only.
+
+GUIDANCE:
+- ``env_name`` should be short snake_case summarizing the scene and task.
+- ``objects`` must list one short search phrase per distinct manipulable object or distractor.
+  Keep phrases compact and visual, e.g. "red hammer", "blue bowl", "spring clamp".
+- Put the scene/table/room name in ``background`` only (e.g. "maple table", "kitchen").
+  Naming the background is not an ``object_references`` cue — the background asset itself is
+  the resting surface, so leave ``object_references`` empty for prompts like
+  "on the maple table" or "from the maple table".
+- Leave ``object_references`` empty unless the prompt names a distinct sub-part or appliance
+  *inside* a multi-prim background (e.g. "counter top", "fridge door", "microwave door").
+  Do not invent a table-surface reference for a table background.
+- ``relations`` should capture placement intent (what sits on what, next_to sides, anchor surface).
+- ``task`` should summarize the robot's goal in one or two sentences.
+"""