diff --git a/isaaclab_arena/agentic_environment_generation/environment_generation_agent.py b/isaaclab_arena/agentic_environment_generation/environment_generation_agent.py index 4f4ab4a51a..0930e95a9e 100644 --- a/isaaclab_arena/agentic_environment_generation/environment_generation_agent.py +++ b/isaaclab_arena/agentic_environment_generation/environment_generation_agent.py @@ -21,10 +21,8 @@ ) from isaaclab_arena.agentic_environment_generation.missing_object_inference import MissingObjectInference from isaaclab_arena.agentic_environment_generation.prim_path_inference import PrimPathInference -from isaaclab_arena.agentic_environment_generation.simready_asset_search import ( - SimReadySearchConfig, - search_simready_objects, -) +from isaaclab_arena.agentic_environment_generation.simready_asset_search import search_simready_objects +from isaaclab_arena.agentic_environment_generation.simready_search_interface import SimReadySearchConfig from isaaclab_arena.agentic_environment_generation.spec_inference import SpecInference from isaaclab_arena.assets.simready_constants import SIMREADY_USD_OBJECT_REGISTRY_NAME from isaaclab_arena.environment_spec.arena_env_graph_spec import ArenaEnvGraphSpec @@ -211,13 +209,14 @@ def _extend_catalogue_with_simready(self, prompt: str, asset_catalog: AssetCatal return asset_catalog extended = replace(asset_catalog, objects=list(asset_catalog.objects)) for candidate in search_result.candidates: - asset_cls = register_searched_simready_object(candidate.search_phrase, candidate.usd_path, candidate.tags) + usd_path = candidate.source_url or candidate.usd_path + asset_cls = register_searched_simready_object(candidate.search_phrase, usd_path, candidate.tags) tags = [tag for tag in asset_cls.tags if tag != "object"] extended.objects.append({ "name": asset_cls.name, "tags": tags, "object_type": asset_cls.object_type.value, }) - self._simready_usd_paths[asset_cls.name] = candidate.usd_path + self._simready_usd_paths[asset_cls.name] = usd_path _logger.info("catalogued %r for %r: %s", asset_cls.name, candidate.search_phrase, candidate.usd_path) return extended diff --git a/isaaclab_arena/agentic_environment_generation/simready_asset_search.py b/isaaclab_arena/agentic_environment_generation/simready_asset_search.py index 5a7e2e7f77..127b11ad37 100644 --- a/isaaclab_arena/agentic_environment_generation/simready_asset_search.py +++ b/isaaclab_arena/agentic_environment_generation/simready_asset_search.py @@ -3,75 +3,20 @@ # # SPDX-License-Identifier: Apache-2.0 -"""SimReady asset search for agentic environment generation.""" - +"""Coordinate selected SimReady backends and merge their search results.""" from __future__ import annotations -import asyncio import logging -import re -from dataclasses import dataclass, field -from enum import Enum -from typing import TYPE_CHECKING, Any -from isaaclab_arena.assets.simready_constants import ( - DEFAULT_SIMREADY_SERVICE_URL, - ISAAC_SIMREADY_GA_S3_URL, - SIMREADY_PHYSICS_VARIANTS, - SIMREADY_USD_OBJECT_REGISTRY_NAME, +from isaaclab_arena.agentic_environment_generation.simready_search_interface import ( + SimReadyCandidateCatalogue, + SimReadySearchBackend, + SimReadySearchConfig, + SimReadySourceKind, ) -from isaaclab_arena.utils.usd.rigid_bodies import read_asset_rigid_body_paths - -if TYPE_CHECKING: - from simready.search import AssetLibrary _logger = logging.getLogger(__name__) -MIN_INSPECTED_MATCHES_PER_PHRASE = 5 -"""How many hits per phrase to inspect for rigid bodies, unless more results than this were asked for.""" - - -class SimReadySourceKind(str, Enum): - """Which SimReady backend to search.""" - - ISAAC_SIM_GA = "isaac-sim-ga" - S3 = "s3" - SERVICE = "service" - - -@dataclass -class SimReadySearchConfig: - """Where to search for SimReady objects and how many hits to keep.""" - - source: SimReadySourceKind = SimReadySourceKind.ISAAC_SIM_GA - s3_url: str = ISAAC_SIMREADY_GA_S3_URL - service_url: str = DEFAULT_SIMREADY_SERVICE_URL - max_results_per_object: int = 1 - - -@dataclass(frozen=True) -class SimReadyObjectCandidate: - """One SimReady search hit exposed to spec inference.""" - - search_phrase: str - usd_path: str - tags: tuple[str, ...] = ("sim-ready",) - relevance_score: float | None = None - - @property - def registry_name(self) -> str: - return SIMREADY_USD_OBJECT_REGISTRY_NAME - - -@dataclass -class SimReadyCandidateCatalogue: - """The SimReady hits found for one prompt, handed to spec inference.""" - - candidates: list[SimReadyObjectCandidate] = field(default_factory=list) - - unmatched_phrases: list[str] = field(default_factory=list) - """Objects the search found no usable asset for, so they are left out of the spec.""" - def search_simready_objects(object_phrases: list[str], config: SimReadySearchConfig) -> SimReadyCandidateCatalogue: """Search SimReady for an asset for each object phrase. A hit that cannot be spawned as a rigid object is turned down, and the next hit is tried in @@ -88,162 +33,33 @@ def search_simready_objects(object_phrases: list[str], config: SimReadySearchCon if not phrases: return SimReadyCandidateCatalogue() - library = _configure_asset_library(config) - if library is None: - return SimReadyCandidateCatalogue() - - candidates: list[SimReadyObjectCandidate] = [] - unmatched_phrases: list[str] = [] - for phrase in phrases: + candidates = [] + seen = set() + for source in dict.fromkeys(config.sources): try: - hits = _search_phrase(library, phrase, max_results=config.max_results_per_object) + backend = _create_search_backend(source, config) + result = backend.search(phrases, config.max_results_per_object) except Exception as exc: - _logger.warning("simready search failed for %r: %s", phrase, exc) - hits = [] - if hits: - candidates.extend(hits) - else: - _logger.info("simready search found no usable asset for %r", phrase) - unmatched_phrases.append(phrase) - - return SimReadyCandidateCatalogue(candidates=candidates, unmatched_phrases=unmatched_phrases) - - -def simready_search_config_from_cli( - source: str, - s3_url: str | None, - service_url: str | None, - max_results_per_object: int, -) -> SimReadySearchConfig: - """Build a search configuration from CLI or GUI arguments, filling in the default URLs.""" - return SimReadySearchConfig( - source=SimReadySourceKind(source), - s3_url=s3_url or ISAAC_SIMREADY_GA_S3_URL, - service_url=service_url or DEFAULT_SIMREADY_SERVICE_URL, - max_results_per_object=max_results_per_object, - ) - - -_CAMEL_CASE_BOUNDARY = re.compile(r"(?<=[a-z0-9])(?=[A-Z])") -_NON_ALPHANUMERIC = re.compile(r"[^a-z0-9]+") - - -def _configure_asset_library(config: SimReadySearchConfig) -> Any | None: - """Build the SimReady asset library for the configured source, or None if the source is unknown.""" - if config.source not in tuple(SimReadySourceKind): - _logger.error("unknown simready source: %s", config.source) - return None - - # Imported here rather than at module scope only to keep the AWS stack out of the import path - # of every caller. - from simready.search import AssetLibrary - - library = AssetLibrary() - if config.source in (SimReadySourceKind.ISAAC_SIM_GA, SimReadySourceKind.S3): - asyncio.run(library.add_s3_source(config.s3_url or ISAAC_SIMREADY_GA_S3_URL)) - else: - library.add_service_source(config.service_url or DEFAULT_SIMREADY_SERVICE_URL) - return library - - -def _search_phrase(library: AssetLibrary, phrase: str, *, max_results: int) -> list[SimReadyObjectCandidate]: - matches = _keep_whole_word_matches(library.search(include_any=_phrase_path_filters(phrase)), phrase) - - # Inspect a few more hits than are wanted back, so a rejected one can be replaced rather than - # costing the phrase a result. Past that, give up and let the agent fall back to the Arena - # asset registry: every further hit inspected is another asset downloaded. - ranked = matches[: max(MIN_INSPECTED_MATCHES_PER_PHRASE, max_results)] - if len(ranked) < len(matches): - _logger.info("simready search checked the best %d of %d hits for %r", len(ranked), len(matches), phrase) - - candidates: list[SimReadyObjectCandidate] = [] - for match in ranked: - usd_path = str(match.asset_path) - is_valid, rejection_reason = _is_valid_isaaclab_rigidbody(usd_path) - if is_valid: - candidates.append(_build_simready_object_candidate_from_match(match, phrase)) - if len(candidates) >= max_results: - break - else: - _logger.info("simready rejected %s for %r: %s", usd_path, phrase, rejection_reason) - return candidates - - -def _phrase_path_filters(phrase: str) -> list[Any]: - """Build one path filter per word of the phrase, so a hit only has to match one word.""" - from simready.search import SearchFilterPathContains - - return [SearchFilterPathContains(word) for word in _phrase_words(phrase)] - - -def _keep_whole_word_matches(matches: list[Any], phrase: str) -> list[Any]: - """Keep only the hits named after the object being asked for, best match first.""" - words = _phrase_words(phrase) - assert words, "a search phrase needs at least one word" - object_word = words[-1] - kept = [match for match in matches if _is_word_in_path(object_word, _split_path_into_words(str(match.asset_path)))] - # The sort is stable, so hits that tie on both counts keep the order the search gave them. - kept.sort( - key=lambda match: ( - getattr(match, "relevance_score", None) or 0.0, - _count_matching_words(phrase, str(match.asset_path)), - ), - reverse=True, - ) - return kept - - -def _is_valid_isaaclab_rigidbody(usd_path: str) -> tuple[bool, str]: - """Say whether a SimReady asset can be spawned as an Isaac Lab rigid object. - - Args: - usd_path: The asset to look at, local or remote. - - Returns: - Whether the asset is usable, and a phrase naming the problem when it is not, such as - "it has no rigid body". The phrase is empty when the asset is usable. - """ - try: - rigid_body_paths = read_asset_rigid_body_paths(usd_path, SIMREADY_PHYSICS_VARIANTS) - except Exception as exc: - return False, f"its USD could not be read: {exc}" - # Only one RigidBodyAPI is allowed on a USD asset. - if len(rigid_body_paths) == 1: - return True, "" - if not rigid_body_paths: - return False, "it has no rigid body" - return False, f"it has {len(rigid_body_paths)} rigid bodies" - - -def _build_simready_object_candidate_from_match(match: Any, phrase: str) -> SimReadyObjectCandidate: - return SimReadyObjectCandidate( - search_phrase=phrase, - usd_path=str(match.asset_path), - tags=tuple(dict.fromkeys(("sim-ready", *_phrase_words(phrase)))), - relevance_score=getattr(match, "relevance_score", None), + _logger.warning("simready backend %s failed: %s", source, exc) + continue + for candidate in result: + key = (candidate.search_phrase, candidate.usd_path) + if key not in seen: + seen.add(key) + candidates.append(candidate) + matched = {candidate.search_phrase for candidate in candidates} + return SimReadyCandidateCatalogue( + candidates=candidates, + unmatched_phrases=[phrase for phrase in phrases if phrase not in matched], ) -def _phrase_words(phrase: str) -> list[str]: - """Split a search phrase into its lowercase words, in order.""" - return [word for word in _NON_ALPHANUMERIC.split(phrase.lower()) if word] - - -def _split_path_into_words(asset_path: str) -> frozenset[str]: - """Split an asset path into lowercase words, also splitting camelCase.""" - # e.g. trashcan_wheeled_green_a01_01.usd -> trashcan wheeled green a01 01 - spaced = _CAMEL_CASE_BOUNDARY.sub(" ", asset_path) - return frozenset(word for word in _NON_ALPHANUMERIC.split(spaced.lower()) if word) - - -def _is_word_in_path(word: str, path_words: frozenset[str]) -> bool: - """True if the word is one of the path words. A trailing "s" on either side is ignored.""" - if word in path_words or f"{word}s" in path_words: - return True - return word.endswith("s") and word[:-1] in path_words +def _create_search_backend(source: SimReadySourceKind, config: SimReadySearchConfig) -> SimReadySearchBackend: + """Construct the implementation for one selected source.""" + if source == SimReadySourceKind.CENTRAL: + from isaaclab_arena.agentic_environment_generation.simready_central_assets import CentralSearchBackend + return CentralSearchBackend() + from isaaclab_arena.agentic_environment_generation.simready_ga_assets import AssetLibrarySearchBackend -def _count_matching_words(phrase: str, asset_path: str) -> int: - """Count how many words of the phrase appear as whole words in the asset path.""" - path_words = _split_path_into_words(asset_path) - return sum(1 for word in _phrase_words(phrase) if _is_word_in_path(word, path_words)) + return AssetLibrarySearchBackend(source, s3_url=config.s3_url, service_url=config.service_url) diff --git a/isaaclab_arena/agentic_environment_generation/simready_central_assets.py b/isaaclab_arena/agentic_environment_generation/simready_central_assets.py new file mode 100644 index 0000000000..0283cfd3e6 --- /dev/null +++ b/isaaclab_arena/agentic_environment_generation/simready_central_assets.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 + +"""Search SimReady Central through its specialized REST API.""" +from __future__ import annotations + +import logging + +import requests + +from isaaclab_arena.agentic_environment_generation.simready_search_interface import ( + MIN_INSPECTED_MATCHES_PER_PHRASE, + SimReadyObjectCandidate, +) +from isaaclab_arena.assets.simready_central_asset_loader import download_and_wrap_central_asset + +CENTRAL_SEARCH_URL = "https://nvidia-simready-central.hf.space/search_hybrid" +_logger = logging.getLogger(__name__) + + +class CentralSearchBackend: + """Search Central and validate downloaded single-rigid-body packages.""" + + def search(self, phrases: list[str], max_results: int) -> list[SimReadyObjectCandidate]: + """Search each phrase independently so one request failure does not discard other hits.""" + candidates = [] + for phrase in phrases: + try: + hits = search_central_objects(phrase, max_results=max_results) + if not hits: + _logger.info("Central found no usable asset for %r", phrase) + candidates.extend(hits) + except Exception as exc: + _logger.warning("Central search failed for %r: %s", phrase, exc) + return candidates + + +def search_central_objects(phrase: str, *, max_results: int) -> list[SimReadyObjectCandidate]: + """Search Central and retain only locally validated single-body props.""" + assert max_results > 0, "max_results must be positive" + response = requests.post( + CENTRAL_SEARCH_URL, + headers={"x-usdsearch-storage-backend": "usdsearch-dev"}, + json={ + "hybrid_text_query": phrase, + "file_extension_include": "usd*", + "filter_by_properties_numeric": "simready_metadata_rigid_body_count=1", + "return_usd_properties": True, + "limit": max(MIN_INSPECTED_MATCHES_PER_PHRASE, max_results), + }, + timeout=30, + ) + response.raise_for_status() + candidates = [] + for hit in response.json()["hits"]: + try: + source = hit["source"] + # Fail closed even if a backend ignores the numeric filter. + counts = [ + prop.get("value_numeric", prop.get("value")) + for prop in source.get("usd_properties", []) + if prop.get("name") == "simready_metadata_rigid_body_count" + ] + if not counts or any(float(count) != 1 for count in counts): + continue + usd_path = download_and_wrap_central_asset(source["path"]) + candidates.append( + SimReadyObjectCandidate( + search_phrase=phrase, + usd_path=usd_path, + tags=("sim-ready", "single-rigid-body"), + relevance_score=hit.get("rrf_score"), + source_url=source["path"].replace("hf://", "https://", 1), + ) + ) + except Exception as exc: + _logger.info("Skipping Central asset %s: %s", hit.get("id"), exc) + continue + if len(candidates) >= max_results: + break + return candidates diff --git a/isaaclab_arena/agentic_environment_generation/simready_ga_assets.py b/isaaclab_arena/agentic_environment_generation/simready_ga_assets.py new file mode 100644 index 0000000000..cc3bd5beaa --- /dev/null +++ b/isaaclab_arena/agentic_environment_generation/simready_ga_assets.py @@ -0,0 +1,184 @@ +# 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 + +"""Search GA, S3, and service assets through the simready-search AssetLibrary API.""" +from __future__ import annotations + +import asyncio +import logging +import re +from typing import TYPE_CHECKING, Any + +from isaaclab_arena.agentic_environment_generation.simready_search_interface import ( + MIN_INSPECTED_MATCHES_PER_PHRASE, + SimReadyObjectCandidate, + SimReadySourceKind, +) +from isaaclab_arena.assets.simready_constants import ( + DEFAULT_SIMREADY_SERVICE_URL, + ISAAC_SIMREADY_GA_S3_URL, + SIMREADY_PHYSICS_VARIANTS, +) +from isaaclab_arena.utils.usd.rigid_bodies import read_asset_rigid_body_paths + +if TYPE_CHECKING: + from simready.search import AssetLibrary +_logger = logging.getLogger(__name__) + + +class AssetLibrarySearchBackend: + """Search a GA, S3, or service source through AssetLibrary.""" + + def __init__(self, source: SimReadySourceKind, *, s3_url: str, service_url: str): + self.library = _create_asset_library_for_source(source, s3_url=s3_url, service_url=service_url) + + def search(self, phrases: list[str], max_results: int) -> list[SimReadyObjectCandidate]: + """Search each phrase and retain single-rigid-body candidates.""" + if self.library is None: + return [] + candidates = [] + for phrase in phrases: + try: + # Query the simready-search API through its AssetLibrary backend. + hits = _search_library_rigid_body_candidates(self.library, phrase, max_results=max_results) + if not hits: + _logger.info("simready found no usable asset for %r", phrase) + candidates.extend(hits) + except Exception as exc: + _logger.warning("simready search failed for %r: %s", phrase, exc) + return candidates + + +_CAMEL_CASE_BOUNDARY = re.compile(r"(?<=[a-z0-9])(?=[A-Z])") +_NON_ALPHANUMERIC = re.compile(r"[^a-z0-9]+") + + +def _create_asset_library_for_source(source: SimReadySourceKind, *, s3_url: str, service_url: str) -> Any | None: + """Build the SimReady asset library for the configured source, or None if the source is unknown.""" + if source not in (SimReadySourceKind.ISAAC_SIM_GA, SimReadySourceKind.S3, SimReadySourceKind.SERVICE): + _logger.error("unknown simready source: %s", source) + return None + + # Imported here rather than at module scope only to keep the AWS stack out of the import path + # of every caller. + from simready.search import AssetLibrary + + library = AssetLibrary() + if source in (SimReadySourceKind.ISAAC_SIM_GA, SimReadySourceKind.S3): + url = ISAAC_SIMREADY_GA_S3_URL if source == SimReadySourceKind.ISAAC_SIM_GA else s3_url + asyncio.run(library.add_s3_source(url or ISAAC_SIMREADY_GA_S3_URL)) + else: + library.add_service_source(service_url or DEFAULT_SIMREADY_SERVICE_URL) + return library + + +def _search_library_rigid_body_candidates( + library: AssetLibrary, phrase: str, *, max_results: int +) -> list[SimReadyObjectCandidate]: + """Search the library by phrase and retain ranked assets with exactly one rigid body.""" + matches = _filter_and_rank_path_matches(library.search(include_any=_build_path_filters_from_phrase(phrase)), phrase) + + # Inspect a few more hits than are wanted back, so a rejected one can be replaced rather than + # costing the phrase a result. Past that, give up and let the agent fall back to the Arena + # asset registry: every further hit inspected is another asset downloaded. + ranked = matches[: max(MIN_INSPECTED_MATCHES_PER_PHRASE, max_results)] + if len(ranked) < len(matches): + _logger.info("simready search checked the best %d of %d hits for %r", len(ranked), len(matches), phrase) + + candidates: list[SimReadyObjectCandidate] = [] + for match in ranked: + usd_path = str(match.asset_path) + is_valid, rejection_reason = _has_single_rigid_body(usd_path) + if is_valid: + candidates.append(_build_simready_object_candidate_from_match(match, phrase)) + if len(candidates) >= max_results: + break + else: + _logger.info("simready rejected %s for %r: %s", usd_path, phrase, rejection_reason) + return candidates + + +def _build_path_filters_from_phrase(phrase: str) -> list[Any]: + """Build one path filter per word of the phrase, so a hit only has to match one word.""" + from simready.search import SearchFilterPathContains + + return [SearchFilterPathContains(word) for word in _split_phrase_into_words(phrase)] + + +def _filter_and_rank_path_matches(matches: list[Any], phrase: str) -> list[Any]: + """Keep only the hits named after the object being asked for, best match first.""" + words = _split_phrase_into_words(phrase) + assert words, "a search phrase needs at least one word" + object_word = words[-1] + kept = [ + match + for match in matches + if _path_words_match_singular_or_plural(object_word, _split_path_into_words(str(match.asset_path))) + ] + # The sort is stable, so hits that tie on both counts keep the order the search gave them. + kept.sort( + key=lambda match: ( + getattr(match, "relevance_score", None) or 0.0, + _count_phrase_words_in_path(phrase, str(match.asset_path)), + ), + reverse=True, + ) + return kept + + +def _has_single_rigid_body(usd_path: str) -> tuple[bool, str]: + """Check that a readable USD contains exactly one rigid body with the SimReady variants selected. + + Args: + usd_path: The asset to look at, local or remote. + + Returns: + Whether exactly one rigid body was found, and a reason when the check fails, such as + "it has no rigid body". The phrase is empty when the asset is usable. + """ + try: + rigid_body_paths = read_asset_rigid_body_paths(usd_path, SIMREADY_PHYSICS_VARIANTS) + except Exception as exc: + return False, f"its USD could not be read: {exc}" + # Only one RigidBodyAPI is allowed on a USD asset. + if len(rigid_body_paths) == 1: + return True, "" + if not rigid_body_paths: + return False, "it has no rigid body" + return False, f"it has {len(rigid_body_paths)} rigid bodies" + + +def _build_simready_object_candidate_from_match(match: Any, phrase: str) -> SimReadyObjectCandidate: + return SimReadyObjectCandidate( + search_phrase=phrase, + usd_path=str(match.asset_path), + tags=tuple(dict.fromkeys(("sim-ready", *_split_phrase_into_words(phrase)))), + relevance_score=getattr(match, "relevance_score", None), + ) + + +def _split_phrase_into_words(phrase: str) -> list[str]: + """Split a search phrase into its lowercase words, in order.""" + return [word for word in _NON_ALPHANUMERIC.split(phrase.lower()) if word] + + +def _split_path_into_words(asset_path: str) -> frozenset[str]: + """Split an asset path into lowercase words, also splitting camelCase.""" + # e.g. trashcan_wheeled_green_a01_01.usd -> trashcan wheeled green a01 01 + spaced = _CAMEL_CASE_BOUNDARY.sub(" ", asset_path) + return frozenset(word for word in _NON_ALPHANUMERIC.split(spaced.lower()) if word) + + +def _path_words_match_singular_or_plural(word: str, path_words: frozenset[str]) -> bool: + """True if the word is one of the path words. A trailing "s" on either side is ignored.""" + if word in path_words or f"{word}s" in path_words: + return True + return word.endswith("s") and word[:-1] in path_words + + +def _count_phrase_words_in_path(phrase: str, asset_path: str) -> int: + """Count how many words of the phrase appear as whole words in the asset path.""" + path_words = _split_path_into_words(asset_path) + return sum(1 for word in _split_phrase_into_words(phrase) if _path_words_match_singular_or_plural(word, path_words)) diff --git a/isaaclab_arena/agentic_environment_generation/simready_search_interface.py b/isaaclab_arena/agentic_environment_generation/simready_search_interface.py new file mode 100644 index 0000000000..b119281d5d --- /dev/null +++ b/isaaclab_arena/agentic_environment_generation/simready_search_interface.py @@ -0,0 +1,87 @@ +# 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 + +"""Shared types and protocol for SimReady search backends.""" +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Protocol + +from isaaclab_arena.assets.simready_constants import ( + DEFAULT_SIMREADY_SERVICE_URL, + ISAAC_SIMREADY_GA_S3_URL, + SIMREADY_USD_OBJECT_REGISTRY_NAME, +) + +MIN_INSPECTED_MATCHES_PER_PHRASE = 5 +"""How many hits per phrase to inspect for rigid bodies, unless more results than this were asked for.""" + + +class SimReadySourceKind(str, Enum): + """Which SimReady backend to search.""" + + ISAAC_SIM_GA = "isaac-sim-ga" + S3 = "s3" + SERVICE = "service" + CENTRAL = "simready-central" + + +@dataclass +class SimReadySearchConfig: + """Where to search for SimReady objects and how many hits to keep.""" + + sources: list[SimReadySourceKind] = field(default_factory=lambda: [SimReadySourceKind.ISAAC_SIM_GA]) + s3_url: str = ISAAC_SIMREADY_GA_S3_URL + service_url: str = DEFAULT_SIMREADY_SERVICE_URL + max_results_per_object: int = 1 + + +@dataclass(frozen=True) +class SimReadyObjectCandidate: + """One SimReady search hit exposed to spec inference.""" + + search_phrase: str + usd_path: str + tags: tuple[str, ...] = ("sim-ready",) + relevance_score: float | None = None + source_url: str | None = None + + @property + def registry_name(self) -> str: + return SIMREADY_USD_OBJECT_REGISTRY_NAME + + +@dataclass +class SimReadyCandidateCatalogue: + """The SimReady hits found for one prompt, handed to spec inference.""" + + candidates: list[SimReadyObjectCandidate] = field(default_factory=list) + + unmatched_phrases: list[str] = field(default_factory=list) + """Objects the search found no usable asset for, so they are left out of the spec.""" + + +class SimReadySearchBackend(Protocol): + """Return validated candidates for object phrases from one backend.""" + + def search(self, phrases: list[str], max_results: int) -> list[SimReadyObjectCandidate]: + """Search each phrase, retaining at most max_results candidates per phrase.""" + ... + + +def simready_search_config_from_cli( + sources: list[str], + s3_url: str | None, + service_url: str | None, + max_results_per_object: int, +) -> SimReadySearchConfig: + """Build a search configuration from CLI or GUI arguments, filling in the default URLs.""" + return SimReadySearchConfig( + sources=[SimReadySourceKind(source) for source in dict.fromkeys(sources)], + s3_url=s3_url or ISAAC_SIMREADY_GA_S3_URL, + service_url=service_url or DEFAULT_SIMREADY_SERVICE_URL, + max_results_per_object=max_results_per_object, + ) diff --git a/isaaclab_arena/assets/simready_central_asset_loader.py b/isaaclab_arena/assets/simready_central_asset_loader.py new file mode 100644 index 0000000000..3cd8f9b8bb --- /dev/null +++ b/isaaclab_arena/assets/simready_central_asset_loader.py @@ -0,0 +1,98 @@ +# 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 + +"""Download Central asset packages and create validated PhysX USD wrappers.""" +from __future__ import annotations + +import hashlib +import re +from pathlib import Path, PurePosixPath +from urllib.parse import urlsplit + + +def is_central_asset_url(url: str) -> bool: + """Identify Hugging Face dataset URLs handled by the Central asset loader.""" + parsed = urlsplit(url) + return ( + parsed.scheme in ("https", "hf") and parsed.netloc == "huggingface.co" and parsed.path.startswith("/datasets/") + ) + + +def _parse_central_asset_url(url: str) -> tuple[str, str, str, str]: + """Parse a revision-pinned Central dataset URL and its asset package folder.""" + assert is_central_asset_url(url), "expected a Hugging Face dataset URL" + parsed = urlsplit(url) + assert not parsed.query and not parsed.fragment, "expected a plain asset URL" + match = re.fullmatch(r"/datasets/([^/]+/[^/]+)/(?:resolve|blob)/([0-9a-f]{40})/(.+)", parsed.path) + assert match, "expected a pinned Hugging Face dataset URL" + repo_id, revision, filename = match.groups() + path = PurePosixPath(filename) + assert not path.is_absolute() and ".." not in path.parts, "invalid asset path" + assert not any(char in filename for char in "*?[]\\"), "invalid asset path" + assert path.suffix in (".usd", ".usda", ".usdc", ".usdz"), "expected USD asset" + # Central packages place the entrypoint in /simready_usd/. Fetch its + # siblings too (physics layers, materials, textures), never the entire dataset. + assert "simready_usd" in path.parts, "unsupported Central package layout" + folder = PurePosixPath(*path.parts[: path.parts.index("simready_usd")]) + assert str(folder) != ".", "asset must have a package folder" + return repo_id, revision, filename, str(folder) + + +def _download_central_asset_package_to_cache(url: str) -> str: + """Download the pinned asset package and return its cached USD entrypoint.""" + from concurrent.futures import ThreadPoolExecutor + + from huggingface_hub import HfApi, hf_hub_download + from huggingface_hub.hf_api import RepoFile + + repo_id, revision, filename, folder = _parse_central_asset_url(url) + # List only this package: a dataset-wide snapshot scan is expensive for large libraries. + files = [ + entry.path + for entry in HfApi().list_repo_tree( + repo_id=repo_id, repo_type="dataset", revision=revision, path_in_repo=folder, recursive=True + ) + if isinstance(entry, RepoFile) + ] + assert filename in files, "asset package is missing its USD entrypoint" + + def download_file(path: str) -> str: + return hf_hub_download(repo_id=repo_id, repo_type="dataset", revision=revision, filename=path) + + with ThreadPoolExecutor(max_workers=8) as executor: + cached_paths = dict(zip(files, executor.map(download_file, files), strict=True)) + return cached_paths[filename] + + +def download_and_wrap_central_asset(url: str) -> str: + """Download a Central package, validate its rigid body, and write a cached PhysX USD wrapper.""" + from pxr import Usd, UsdGeom, UsdPhysics + + from isaaclab_arena.utils.usd.rigid_bodies import apply_usd_variant_selections + + asset_path = _download_central_asset_package_to_cache(url) + asset_stage = Usd.Stage.Open(asset_path) + assert asset_stage and asset_stage.GetDefaultPrim(), "asset requires a default prim" + stage = Usd.Stage.CreateInMemory() + UsdGeom.SetStageMetersPerUnit(stage, UsdGeom.GetStageMetersPerUnit(asset_stage)) + UsdGeom.SetStageUpAxis(stage, UsdGeom.GetStageUpAxis(asset_stage)) + root = stage.DefinePrim("/Asset") + root.GetReferences().AddReference(asset_path) + stage.SetDefaultPrim(root) + # TODO(xinjieyao, 20261022): Support Newton. + apply_usd_variant_selections(stage, {"Physics": "physics", "PhysX": "Enabled"}) + assert not stage.GetCompositionErrors(), "USD has unresolved composition errors" + bodies = [] + for prim in Usd.PrimRange.Stage(stage, Usd.TraverseInstanceProxies()): + assert not prim.HasAPI(UsdPhysics.ArticulationRootAPI), "asset has an articulation root" + if prim.HasAPI(UsdPhysics.RigidBodyAPI): + bodies.append(prim) + assert len(bodies) == 1, f"expected one rigid body, found {len(bodies)}" + assert UsdPhysics.RigidBodyAPI(bodies[0]).GetRigidBodyEnabledAttr().Get(), "rigid body is disabled" + cache = Path.home() / ".cache" / "isaaclab_arena" / "simready-central" + cache.mkdir(parents=True, exist_ok=True) + wrapper = cache / f"{hashlib.sha256(url.encode()).hexdigest()}.usda" + stage.GetRootLayer().Export(str(wrapper)) + return str(wrapper) diff --git a/isaaclab_arena/assets/simready_object_library.py b/isaaclab_arena/assets/simready_object_library.py index 1bbb214381..09e541e0cf 100644 --- a/isaaclab_arena/assets/simready_object_library.py +++ b/isaaclab_arena/assets/simready_object_library.py @@ -7,6 +7,7 @@ from __future__ import annotations +import hashlib import re from collections.abc import Sequence from typing import Any @@ -46,6 +47,13 @@ def __init__( **kwargs: Any, ): assert usd_path, "simready_usd_object requires params.usd_path" + from isaaclab_arena.assets.simready_central_asset_loader import ( + download_and_wrap_central_asset, + is_central_asset_url, + ) + + if is_central_asset_url(usd_path): + usd_path = download_and_wrap_central_asset(usd_path) spawn_cfg_addon = dict(kwargs.pop("spawn_cfg_addon", None) or {}) spawn_cfg_addon.setdefault("variants", dict(SIMREADY_PHYSICS_VARIANTS)) super().__init__( @@ -61,11 +69,12 @@ def __init__( ) -def simready_search_registry_name(search_phrase: str) -> str: - """Return the catalogue name used to register the SimReady asset found for a search phrase.""" +def simready_search_registry_name(search_phrase: str, usd_path: str) -> str: + """Return a stable catalogue name for a phrase and asset path.""" slug = _NON_IDENTIFIER_CHARACTERS.sub("_", search_phrase.strip().lower()).strip("_") assert slug, f"search phrase {search_phrase!r} has no characters usable in a registry name" - return f"{SIMREADY_SEARCH_REGISTRY_PREFIX}{slug}" + digest = hashlib.sha256(usd_path.encode()).hexdigest()[:16] + return f"{SIMREADY_SEARCH_REGISTRY_PREFIX}{slug}_{digest}" def register_searched_simready_object( @@ -76,14 +85,14 @@ def register_searched_simready_object( """Register a searched SimReady asset so it can be picked like any other catalogue object. Args: - search_phrase: Phrase the asset was found for; it determines the catalogue name. + search_phrase: Phrase the asset was found for; used with its path to name the catalogue entry. usd_path: USD path or URL of the found asset. tags: Extra catalogue tags to expose alongside the SimReady ones. Returns: The registered asset class. """ - registry_name = simready_search_registry_name(search_phrase) + registry_name = simready_search_registry_name(search_phrase, usd_path) registry = AssetRegistry() if registry.is_registered(registry_name, ensure_loaded=False): return registry.get_asset_by_name(registry_name) diff --git a/isaaclab_arena/tests/test_environment_generation_agent.py b/isaaclab_arena/tests/test_environment_generation_agent.py index d5243a4026..f18d70d336 100644 --- a/isaaclab_arena/tests/test_environment_generation_agent.py +++ b/isaaclab_arena/tests/test_environment_generation_agent.py @@ -18,7 +18,7 @@ build_task_catalogue, ) from isaaclab_arena.agentic_environment_generation.environment_generation_agent import EnvironmentGenerationAgent -from isaaclab_arena.agentic_environment_generation.simready_asset_search import ( +from isaaclab_arena.agentic_environment_generation.simready_search_interface import ( SimReadyCandidateCatalogue, SimReadyObjectCandidate, ) @@ -215,6 +215,8 @@ def _generate_with_simready(self, agent_obj, asset_catalog=None): @patch("isaaclab_arena.agentic_environment_generation.environment_generation_agent.search_simready_objects") def test_generate_spec_catalogues_what_the_search_finds_before_inferring_the_spec(self, mock_search, agent): + from isaaclab_arena.assets.simready_object_library import simready_search_registry_name + agent_obj, client = agent mock_search.return_value = SimReadyCandidateCatalogue( candidates=[ @@ -236,12 +238,16 @@ def test_generate_spec_catalogues_what_the_search_finds_before_inferring_the_spe # The found asset reaches spec inference as an ordinary OBJECTS entry, so the model picks # it by name and the spec needs no usd_path of its own. spec_user_msg = client.chat.completions.create.call_args.kwargs["messages"][1]["content"] - assert "simready_green_trash_can" in spec_user_msg + assert simready_search_registry_name("green trash can", "s3://bucket/trash_can.usd") in spec_user_msg assert "s3://bucket/trash_can.usd" not in spec_user_msg - assert AssetRegistry().is_registered("simready_green_trash_can") + assert AssetRegistry().is_registered( + simready_search_registry_name("green trash can", "s3://bucket/trash_can.usd") + ) @patch("isaaclab_arena.agentic_environment_generation.environment_generation_agent.search_simready_objects") def test_a_searched_simready_object_reaches_the_spec_with_its_usd_path(self, mock_search, agent): + from isaaclab_arena.assets.simready_object_library import simready_search_registry_name + agent_obj, client = agent mock_search.return_value = SimReadyCandidateCatalogue( candidates=[ @@ -255,7 +261,7 @@ def test_a_searched_simready_object_reaches_the_spec_with_its_usd_path(self, moc spec_dict = minimal_spec_dict() spec_dict["objects"].append({ "id": "trash_can", - "registry_name": "simready_green_trash_can", + "registry_name": simready_search_registry_name("green trash can", "s3://bucket/trash_can.usd"), "params": {}, }) client.chat.completions.create.side_effect = [ @@ -272,14 +278,58 @@ def test_a_searched_simready_object_reaches_the_spec_with_its_usd_path(self, moc # Tags in params would be a duplicate keyword argument at build time. assert all("tags" not in asset["params"] for asset in spec.to_dict()["objects"]) + @patch("isaaclab_arena.agentic_environment_generation.environment_generation_agent.search_simready_objects") + def test_multiple_candidates_export_distinct_portable_paths(self, mock_search, agent): + from isaaclab_arena.assets.simready_object_library import simready_search_registry_name + + agent_obj, client = agent + central_url = ( + "https://huggingface.co/datasets/nvidia/simready-assets/resolve/" + "a" * 40 + "/can/simready_usd/can.usd" + ) + paths = ["s3://bucket/can.usd", central_url] + mock_search.return_value = SimReadyCandidateCatalogue( + candidates=[ + SimReadyObjectCandidate(search_phrase="can", usd_path=paths[0], tags=("ga-can",)), + SimReadyObjectCandidate( + search_phrase="can", + usd_path="/other-machine/cache/can.usda", + source_url=central_url, + tags=("central-can",), + ), + ] + ) + spec_dict = minimal_spec_dict() + for index, path in enumerate(paths): + spec_dict["objects"].append( + {"id": f"can_{index}", "registry_name": simready_search_registry_name("can", path), "params": {}} + ) + client.chat.completions.create.side_effect = [ + self._missing_objects_response("can"), + chat_response(content=json.dumps(spec_dict)), + ] + spec, _ = self._generate_with_simready(agent_obj) + assert isinstance(spec, ArenaEnvGraphSpec) + exported = spec.to_dict() + selected = [obj for obj in exported["objects"] if obj["id"].startswith("can_")] + assert [obj["params"]["usd_path"] for obj in selected] == paths + assert all(obj["registry_name"] == SIMREADY_USD_OBJECT_REGISTRY_NAME for obj in selected) + assert "/other-machine/" not in json.dumps(exported) + for path, tag in zip(paths, ["ga-can", "central-can"], strict=True): + cls = AssetRegistry().get_asset_by_name(simready_search_registry_name("can", path)) + assert tag in cls.tags + @patch("isaaclab_arena.agentic_environment_generation.environment_generation_agent.search_simready_objects") def test_a_searched_simready_object_in_an_object_set_is_rejected(self, mock_search, agent): + from isaaclab_arena.assets.simready_object_library import simready_search_registry_name + agent_obj, client = agent mock_search.return_value = SimReadyCandidateCatalogue( candidates=[SimReadyObjectCandidate(search_phrase="green trash can", usd_path="s3://bucket/trash_can.usd")] ) spec_dict = minimal_spec_dict() - spec_dict["object_sets"] = [{"id": "bins", "members": ["simready_green_trash_can"]}] + spec_dict["object_sets"] = [ + {"id": "bins", "members": [simready_search_registry_name("green trash can", "s3://bucket/trash_can.usd")]} + ] client.chat.completions.create.side_effect = [ self._missing_objects_response("green trash can"), chat_response(content=json.dumps(spec_dict)), diff --git a/isaaclab_arena/tests/test_simready_asset_search.py b/isaaclab_arena/tests/test_simready_asset_search.py index 8cea746cfd..494794190a 100644 --- a/isaaclab_arena/tests/test_simready_asset_search.py +++ b/isaaclab_arena/tests/test_simready_asset_search.py @@ -8,22 +8,26 @@ from __future__ import annotations import logging -from unittest.mock import patch +from unittest.mock import AsyncMock, patch -from isaaclab_arena.agentic_environment_generation.simready_asset_search import ( +import pytest + +from isaaclab_arena.agentic_environment_generation.simready_asset_search import search_simready_objects +from isaaclab_arena.agentic_environment_generation.simready_ga_assets import ( + _count_phrase_words_in_path, + _filter_and_rank_path_matches, + _has_single_rigid_body, + _split_path_into_words, +) +from isaaclab_arena.agentic_environment_generation.simready_search_interface import ( SimReadyObjectCandidate, SimReadySearchConfig, SimReadySourceKind, - _count_matching_words, - _is_valid_isaaclab_rigidbody, - _keep_whole_word_matches, - _split_path_into_words, - search_simready_objects, simready_search_config_from_cli, ) from isaaclab_arena.assets.simready_constants import SIMREADY_USD_OBJECT_REGISTRY_NAME -SEARCH_LOGGER = "isaaclab_arena.agentic_environment_generation.simready_asset_search" +SEARCH_LOGGER = "isaaclab_arena.agentic_environment_generation.simready_ga_assets" """Where the search reports its progress, now that the agent's traces carry only errors.""" CABINET_PATH = "SimReady/Residential/Kitchen/Cabinets/Cabinet_D01/sm_fixture_cabinet_d01_01.usd" @@ -31,15 +35,11 @@ TRASH_CAN_PATH = "SimReady/Residential/Garage/sm_trashCan_wheeled_green_a01_01.usd" PLAIN_TRASH_CAN_PATH = "SimReady/Residential/Kitchen/sm_trashCan_a01_01.usd" -RIGID_BODY_CHECK_TARGET = ( - "isaaclab_arena.agentic_environment_generation.simready_asset_search._is_valid_isaaclab_rigidbody" -) -# Patched where it is looked up, not where it is defined: simready_asset_search binds the name at +RIGID_BODY_CHECK_TARGET = "isaaclab_arena.agentic_environment_generation.simready_ga_assets._has_single_rigid_body" +# Patched where it is looked up, not where it is defined: simready_ga_assets binds the name at # import time, so patching the defining module would leave that binding untouched. -RIGID_BODY_PATHS_TARGET = ( - "isaaclab_arena.agentic_environment_generation.simready_asset_search.read_asset_rigid_body_paths" -) -PATH_FILTERS_TARGET = "isaaclab_arena.agentic_environment_generation.simready_asset_search._phrase_path_filters" +RIGID_BODY_PATHS_TARGET = "isaaclab_arena.agentic_environment_generation.simready_ga_assets.read_asset_rigid_body_paths" +PATH_FILTERS_TARGET = "isaaclab_arena.agentic_environment_generation.simready_ga_assets._build_path_filters_from_phrase" # Stubbed wherever the fake library stands in for a real one: building a path filter needs the # simready package, which only ships with Isaac Sim, and the fake library ignores the filters. @@ -65,66 +65,66 @@ def test_split_path_into_words_splits_camel_case_and_separators(): assert "trashcan" not in words -def test_count_matching_words_ignores_substring_hits(): +def test_count_phrase_words_in_path_ignores_substring_hits(): # "bin" only appears inside "Cabinets", which is how a cabinet got picked for a grey bin. - assert _count_matching_words("grey bin", CABINET_PATH) == 0 - assert _count_matching_words("kitchen cabinet", CABINET_PATH) == 2 + assert _count_phrase_words_in_path("grey bin", CABINET_PATH) == 0 + assert _count_phrase_words_in_path("kitchen cabinet", CABINET_PATH) == 2 -def test_count_matching_words_tolerates_plurals(): - assert _count_matching_words("cabinet", CABINET_PATH) == 1 - assert _count_matching_words("cabinets", CABINET_PATH) == 1 +def test_count_phrase_words_in_path_tolerates_plurals(): + assert _count_phrase_words_in_path("cabinet", CABINET_PATH) == 1 + assert _count_phrase_words_in_path("cabinets", CABINET_PATH) == 1 -def test_keep_whole_word_matches_filters_and_ranks(): +def test_filter_and_rank_path_matches_filters_and_ranks(): cabinet = _FakeMatch(CABINET_PATH) trash_can = _FakeMatch(TRASH_CAN_PATH) - assert _keep_whole_word_matches([cabinet], "grey bin") == [] + assert _filter_and_rank_path_matches([cabinet], "grey bin") == [] # "green trash can" matches three words in the trash can and none in the cabinet. - assert _keep_whole_word_matches([cabinet, trash_can], "green trash can") == [trash_can] + assert _filter_and_rank_path_matches([cabinet, trash_can], "green trash can") == [trash_can] -def test_keep_whole_word_matches_orders_by_word_overlap(): +def test_filter_and_rank_path_matches_orders_by_word_overlap(): green_can = _FakeMatch(TRASH_CAN_PATH) plain_can = _FakeMatch(PLAIN_TRASH_CAN_PATH) - ranked = _keep_whole_word_matches([plain_can, green_can], "green trash can") + ranked = _filter_and_rank_path_matches([plain_can, green_can], "green trash can") assert ranked == [green_can, plain_can] -def test_keep_whole_word_matches_rejects_an_asset_that_only_shares_a_describing_word(): +def test_filter_and_rank_path_matches_rejects_an_asset_that_only_shares_a_describing_word(): grey_cabinet = _FakeMatch(GREY_CABINET_PATH) # A grey cabinet is not a grey bin: the colour matches but the object itself does not, and # handing back the cabinet is how a bin became a cabinet in a generated environment. - assert _keep_whole_word_matches([grey_cabinet], "grey bin") == [] - assert _keep_whole_word_matches([grey_cabinet], "grey cabinet") == [grey_cabinet] + assert _filter_and_rank_path_matches([grey_cabinet], "grey bin") == [] + assert _filter_and_rank_path_matches([grey_cabinet], "grey cabinet") == [grey_cabinet] -def test_is_valid_isaaclab_rigidbody_accepts_a_single_rigid_body(): +def test_has_single_rigid_body_accepts_a_single_rigid_body(): with patch(RIGID_BODY_PATHS_TARGET, return_value=["/Asset/bottle"]): - assert _is_valid_isaaclab_rigidbody("s3://bucket/bottle.usd") == (True, "") + assert _has_single_rigid_body("s3://bucket/bottle.usd") == (True, "") -def test_is_valid_isaaclab_rigidbody_turns_down_several_rigid_bodies(): +def test_has_single_rigid_body_turns_down_several_rigid_bodies(): # A bottle and its cap are two bodies even though a fixed joint holds them together. with patch(RIGID_BODY_PATHS_TARGET, return_value=["/Asset/bottle", "/Asset/cap"]): - assert _is_valid_isaaclab_rigidbody("s3://bucket/bottle.usd") == (False, "it has 2 rigid bodies") + assert _has_single_rigid_body("s3://bucket/bottle.usd") == (False, "it has 2 rigid bodies") -def test_is_valid_isaaclab_rigidbody_turns_down_an_asset_without_physics(): +def test_has_single_rigid_body_turns_down_an_asset_without_physics(): with patch(RIGID_BODY_PATHS_TARGET, return_value=[]): - assert _is_valid_isaaclab_rigidbody("s3://bucket/decoration.usd") == (False, "it has no rigid body") + assert _has_single_rigid_body("s3://bucket/decoration.usd") == (False, "it has no rigid body") -def test_is_valid_isaaclab_rigidbody_turns_down_an_unreadable_asset(): +def test_has_single_rigid_body_turns_down_an_unreadable_asset(): with patch(RIGID_BODY_PATHS_TARGET, side_effect=FileNotFoundError("no such file")): - is_valid, reason = _is_valid_isaaclab_rigidbody("s3://bucket/missing.usd") + is_valid, reason = _has_single_rigid_body("s3://bucket/missing.usd") assert not is_valid assert "could not be read" in reason @patch(PATH_FILTERS_TARGET, return_value=[]) @patch( - "isaaclab_arena.agentic_environment_generation.simready_asset_search._configure_asset_library", + "isaaclab_arena.agentic_environment_generation.simready_ga_assets._create_asset_library_for_source", ) def test_search_falls_back_to_the_next_hit_when_one_is_rejected(mock_configure, mock_path_filters, caplog): mock_configure.return_value = _FakeLibrary([_FakeMatch(TRASH_CAN_PATH), _FakeMatch(PLAIN_TRASH_CAN_PATH)]) @@ -140,7 +140,7 @@ def test_search_falls_back_to_the_next_hit_when_one_is_rejected(mock_configure, @patch(PATH_FILTERS_TARGET, return_value=[]) @patch( - "isaaclab_arena.agentic_environment_generation.simready_asset_search._configure_asset_library", + "isaaclab_arena.agentic_environment_generation.simready_ga_assets._create_asset_library_for_source", ) def test_search_reports_a_phrase_with_only_rejected_hits_as_unmatched(mock_configure, mock_path_filters, caplog): mock_configure.return_value = _FakeLibrary([_FakeMatch(CABINET_PATH)]) @@ -154,47 +154,49 @@ def test_search_reports_a_phrase_with_only_rejected_hits_as_unmatched(mock_confi def test_simready_search_config_from_cli_defaults(): config = simready_search_config_from_cli( - source=SimReadySourceKind.ISAAC_SIM_GA.value, + sources=[SimReadySourceKind.ISAAC_SIM_GA.value], s3_url=None, service_url=None, max_results_per_object=2, ) - assert config.source == SimReadySourceKind.ISAAC_SIM_GA + assert config.sources == [SimReadySourceKind.ISAAC_SIM_GA] assert config.max_results_per_object == 2 @patch( - "isaaclab_arena.agentic_environment_generation.simready_asset_search._search_phrase", + "isaaclab_arena.agentic_environment_generation.simready_ga_assets._search_library_rigid_body_candidates", ) @patch( - "isaaclab_arena.agentic_environment_generation.simready_asset_search._configure_asset_library", + "isaaclab_arena.agentic_environment_generation.simready_ga_assets._create_asset_library_for_source", ) -def test_search_simready_objects_returns_candidates(mock_configure, mock_search_phrase): +def test_search_simready_objects_returns_candidates(mock_configure, mock_search_library_rigid_body_candidates): mock_configure.return_value = _FakeLibrary() - mock_search_phrase.return_value = [ + mock_search_library_rigid_body_candidates.return_value = [ SimReadyObjectCandidate( search_phrase="ceramic bowl", usd_path="s3://bucket/bowl.usd", ) ] - catalog = search_simready_objects(["ceramic bowl"], SimReadySearchConfig(source=SimReadySourceKind.ISAAC_SIM_GA)) + catalog = search_simready_objects(["ceramic bowl"], SimReadySearchConfig(sources=[SimReadySourceKind.ISAAC_SIM_GA])) assert len(catalog.candidates) == 1 assert catalog.candidates[0].usd_path == "s3://bucket/bowl.usd" assert catalog.candidates[0].registry_name == SIMREADY_USD_OBJECT_REGISTRY_NAME -def test_configure_asset_library_records_an_unknown_source(caplog): - from isaaclab_arena.agentic_environment_generation.simready_asset_search import _configure_asset_library +def test_create_asset_library_for_source_records_an_unknown_source(caplog): + from isaaclab_arena.agentic_environment_generation.simready_ga_assets import _create_asset_library_for_source config = SimReadySearchConfig() - config.source = "not-a-source" with caplog.at_level(logging.ERROR, logger=SEARCH_LOGGER): - assert _configure_asset_library(config) is None + assert ( + _create_asset_library_for_source("not-a-source", s3_url=config.s3_url, service_url=config.service_url) + is None + ) assert "unknown simready source" in caplog.text @patch( - "isaaclab_arena.agentic_environment_generation.simready_asset_search._configure_asset_library", + "isaaclab_arena.agentic_environment_generation.simready_ga_assets._create_asset_library_for_source", ) def test_search_simready_objects_returns_empty_when_library_unavailable(mock_configure): mock_configure.return_value = None @@ -203,16 +205,121 @@ def test_search_simready_objects_returns_empty_when_library_unavailable(mock_con @patch( - "isaaclab_arena.agentic_environment_generation.simready_asset_search._search_phrase", + "isaaclab_arena.agentic_environment_generation.simready_ga_assets._search_library_rigid_body_candidates", ) @patch( - "isaaclab_arena.agentic_environment_generation.simready_asset_search._configure_asset_library", + "isaaclab_arena.agentic_environment_generation.simready_ga_assets._create_asset_library_for_source", ) -def test_search_simready_objects_records_no_matches(mock_configure, mock_search_phrase, caplog): +def test_search_simready_objects_records_no_matches(mock_configure, mock_search_library_rigid_body_candidates, caplog): mock_configure.return_value = _FakeLibrary() - mock_search_phrase.return_value = [] + mock_search_library_rigid_body_candidates.return_value = [] with caplog.at_level(logging.INFO, logger=SEARCH_LOGGER): catalog = search_simready_objects(["missing object"], SimReadySearchConfig()) assert catalog.candidates == [] assert catalog.unmatched_phrases == ["missing object"] assert "no usable asset" in caplog.text + + +def test_multiple_sources_merge_and_deduplicate_candidates(): + shared = SimReadyObjectCandidate(search_phrase="egg", usd_path="shared.usd") + central = SimReadyObjectCandidate(search_phrase="egg", usd_path="central.usd") + with ( + patch(f"{SEARCH_LOGGER}._create_asset_library_for_source", return_value=object()), + patch(f"{SEARCH_LOGGER}._search_library_rigid_body_candidates", side_effect=[[shared], []]), + patch( + "isaaclab_arena.agentic_environment_generation.simready_central_assets.search_central_objects", + side_effect=[[shared, central], []], + ), + ): + result = search_simready_objects( + ["egg", "missing"], + SimReadySearchConfig(sources=[SimReadySourceKind.ISAAC_SIM_GA, SimReadySourceKind.CENTRAL]), + ) + assert result.candidates == [shared, central] + assert result.unmatched_phrases == ["missing"] + + +def test_multiple_sources_keep_central_when_library_initialization_fails(): + central = SimReadyObjectCandidate(search_phrase="egg", usd_path="central.usd") + with ( + patch(f"{SEARCH_LOGGER}._create_asset_library_for_source", side_effect=RuntimeError("unavailable")), + patch( + "isaaclab_arena.agentic_environment_generation.simready_central_assets.search_central_objects", + return_value=[central], + ), + ): + result = search_simready_objects( + ["egg"], SimReadySearchConfig(sources=[SimReadySourceKind.ISAAC_SIM_GA, SimReadySourceKind.CENTRAL]) + ) + assert result.candidates == [central] + assert not result.unmatched_phrases + + +def test_multiple_sources_keep_library_when_central_fails(): + candidate = SimReadyObjectCandidate(search_phrase="egg", usd_path="library.usd") + with ( + patch(f"{SEARCH_LOGGER}._create_asset_library_for_source", return_value=object()), + patch(f"{SEARCH_LOGGER}._search_library_rigid_body_candidates", return_value=[candidate]), + patch( + "isaaclab_arena.agentic_environment_generation.simready_central_assets.search_central_objects", + side_effect=RuntimeError("unavailable"), + ), + ): + result = search_simready_objects( + ["egg"], SimReadySearchConfig(sources=[SimReadySourceKind.ISAAC_SIM_GA, SimReadySourceKind.CENTRAL]) + ) + assert result.candidates == [candidate] + assert not result.unmatched_phrases + + +def test_source_list_searches_selected_backends_once(): + config = simready_search_config_from_cli( + sources=["service", "simready-central", "service"], + s3_url=None, + service_url="http://asset-service", + max_results_per_object=2, + ) + assert config.sources == [SimReadySourceKind.SERVICE, SimReadySourceKind.CENTRAL] + with ( + patch(f"{SEARCH_LOGGER}._create_asset_library_for_source", return_value=object()) as configure, + patch(f"{SEARCH_LOGGER}._search_library_rigid_body_candidates", return_value=[]) as library_search, + patch( + "isaaclab_arena.agentic_environment_generation.simready_central_assets.search_central_objects", + return_value=[], + ) as central_search, + ): + result = search_simready_objects(["egg"], config) + configure.assert_called_once_with(SimReadySourceKind.SERVICE, s3_url=config.s3_url, service_url=config.service_url) + library_search.assert_called_once() + central_search.assert_called_once_with("egg", max_results=2) + assert result.unmatched_phrases == ["egg"] + + +@pytest.mark.parametrize("source", [SimReadySourceKind.ISAAC_SIM_GA, SimReadySourceKind.CENTRAL]) +def test_backend_keeps_other_phrase_results_when_one_phrase_fails(source): + candidate = SimReadyObjectCandidate(search_phrase="wrench", usd_path="wrench.usd") + target = ( + f"{SEARCH_LOGGER}._search_library_rigid_body_candidates" + if source == SimReadySourceKind.ISAAC_SIM_GA + else "isaaclab_arena.agentic_environment_generation.simready_central_assets.search_central_objects" + ) + with ( + patch(f"{SEARCH_LOGGER}._create_asset_library_for_source", return_value=object()), + patch(target, side_effect=[RuntimeError("unavailable"), [candidate]]) as search, + ): + result = search_simready_objects(["missing", "wrench"], SimReadySearchConfig(sources=[source])) + assert search.call_count == 2 + assert result.candidates == [candidate] + assert result.unmatched_phrases == ["missing"] + + +@pytest.mark.parametrize("source", [SimReadySourceKind.ISAAC_SIM_GA, SimReadySourceKind.S3]) +def test_ga_ignores_custom_s3_url(source): + from isaaclab_arena.agentic_environment_generation.simready_ga_assets import _create_asset_library_for_source + from isaaclab_arena.assets.simready_constants import ISAAC_SIMREADY_GA_S3_URL + + with patch("simready.search.AssetLibrary") as library: + library.return_value.add_s3_source = AsyncMock() + _create_asset_library_for_source(source, s3_url="s3://custom/assets", service_url="") + expected = ISAAC_SIMREADY_GA_S3_URL if source == SimReadySourceKind.ISAAC_SIM_GA else "s3://custom/assets" + library.return_value.add_s3_source.assert_awaited_once_with(expected) diff --git a/isaaclab_arena/tests/test_simready_central.py b/isaaclab_arena/tests/test_simready_central.py new file mode 100644 index 0000000000..ae085fb2c4 --- /dev/null +++ b/isaaclab_arena/tests/test_simready_central.py @@ -0,0 +1,164 @@ +# 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 + +"""Central search filtering and pinned asset preparation tests.""" + +from unittest.mock import Mock, patch + +import pytest + +from isaaclab_arena.agentic_environment_generation.simready_asset_search import search_simready_objects +from isaaclab_arena.agentic_environment_generation.simready_central_assets import search_central_objects +from isaaclab_arena.agentic_environment_generation.simready_search_interface import ( + SimReadySearchConfig, + SimReadySourceKind, +) +from isaaclab_arena.assets.simready_central_asset_loader import _parse_central_asset_url, is_central_asset_url + +MODULE = "isaaclab_arena.agentic_environment_generation.simready_central_assets" +LOADER = "isaaclab_arena.assets.simready_central_asset_loader" +REVISION = "a" * 40 +URL = f"hf://huggingface.co/datasets/nvidia/simready-assets/resolve/{REVISION}/pallet/simready_usd/pallet.usd" + + +def _hit(count, path=URL): + properties = [] if count is None else [{"name": "simready_metadata_rigid_body_count", "value_numeric": count}] + return {"source": {"path": path, "usd_properties": properties}, "rrf_score": 0.5} + + +def test_central_requires_single_body_metadata_and_skips_invalid_usd(): + response = Mock() + response.json.return_value = {"hits": [_hit(None), _hit(2), _hit(1), _hit(1)]} + with ( + patch(f"{MODULE}.requests.post", return_value=response) as post, + patch( + f"{MODULE}.download_and_wrap_central_asset", side_effect=[ValueError("articulation"), "/cache/pallet.usda"] + ) as prepare, + ): + candidates = search_central_objects("pallet", max_results=1) + assert [candidate.usd_path for candidate in candidates] == ["/cache/pallet.usda"] + assert prepare.call_count == 2 + assert post.call_args.kwargs["json"]["filter_by_properties_numeric"] == "simready_metadata_rigid_body_count=1" + assert post.call_args.kwargs["timeout"] == 30 + + +def test_central_failure_leaves_phrase_unmatched(): + with patch(f"{MODULE}.requests.post", side_effect=TimeoutError("offline")): + catalogue = search_simready_objects(["pallet"], SimReadySearchConfig(sources=[SimReadySourceKind.CENTRAL])) + assert catalogue.candidates == [] + assert catalogue.unmatched_phrases == ["pallet"] + + +def test_central_package_location_preserves_revision_and_folder(): + assert _parse_central_asset_url(URL) == ( + "nvidia/simready-assets", + REVISION, + "pallet/simready_usd/pallet.usd", + "pallet", + ) + + +@pytest.mark.parametrize( + "url", + [ + URL.replace(REVISION, "main"), + URL.replace("pallet/simready_usd", "../simready_usd"), + URL.replace("pallet/simready_usd/", ""), + ], +) +def test_central_rejects_unpinned_or_unbounded_packages(url): + with pytest.raises(AssertionError): + _parse_central_asset_url(url) + + +def _test_central_usd_validation(simulation_app): + import tempfile + from pathlib import Path + + from pxr import Usd, UsdGeom, UsdPhysics + + from isaaclab_arena.assets.simready_central_asset_loader import download_and_wrap_central_asset + + with tempfile.TemporaryDirectory() as directory: + folder = Path(directory) + asset = folder / "pallet/simready_usd/pallet.usd" + asset.parent.mkdir(parents=True) + stage = Usd.Stage.CreateNew(str(asset)) + root = UsdGeom.Xform.Define(stage, "/Pallet").GetPrim() + stage.SetDefaultPrim(root) + body = UsdPhysics.RigidBodyAPI.Apply(root) + body.CreateRigidBodyEnabledAttr(True) + stage.GetRootLayer().Save() + with ( + patch(f"{LOADER}._download_central_asset_package_to_cache", return_value=str(asset)), + patch(f"{LOADER}.Path.home", return_value=folder), + ): + wrapper = download_and_wrap_central_asset(URL) + assert Usd.Stage.Open(wrapper).GetDefaultPrim().HasAPI(UsdPhysics.RigidBodyAPI) + body.GetRigidBodyEnabledAttr().Set(False) + stage.GetRootLayer().Save() + with pytest.raises(AssertionError, match="disabled"): + download_and_wrap_central_asset(URL) + body.GetRigidBodyEnabledAttr().Set(True) + UsdPhysics.ArticulationRootAPI.Apply(root) + stage.GetRootLayer().Save() + with pytest.raises(AssertionError, match="articulation root"): + download_and_wrap_central_asset(URL) + root.RemoveAPI(UsdPhysics.ArticulationRootAPI) + UsdPhysics.RigidBodyAPI.Apply(UsdGeom.Xform.Define(stage, "/Pallet/Other").GetPrim()) + stage.GetRootLayer().Save() + with pytest.raises(AssertionError, match="found 2"): + download_and_wrap_central_asset(URL) + return True + + +def test_central_usd_validation(): + from isaaclab_arena.tests.utils.persistent_simulation_app import run_function_with_persistent_simulation_app + + assert run_function_with_persistent_simulation_app(_test_central_usd_validation) + + +def test_central_download_lists_only_the_selected_package(): + from huggingface_hub.hf_api import RepoFile + + from isaaclab_arena.assets.simready_central_asset_loader import _download_central_asset_package_to_cache + + filenames = ["pallet/simready_usd/pallet.usd", "pallet/textures/color.png"] + entries = [RepoFile(path=path, size=1, oid="a" * 40) for path in filenames] + with ( + patch("huggingface_hub.HfApi.list_repo_tree", return_value=iter(entries)) as list_tree, + patch("huggingface_hub.hf_hub_download", side_effect=lambda **kw: "/cache/" + kw["filename"]) as download, + ): + assert _download_central_asset_package_to_cache(URL) == "/cache/" + filenames[0] + list_tree.assert_called_once_with( + repo_id="nvidia/simready-assets", + repo_type="dataset", + revision=REVISION, + path_in_repo="pallet", + recursive=True, + ) + assert {call.kwargs["filename"] for call in download.call_args_list} == set(filenames) + assert all(call.kwargs["revision"] == REVISION for call in download.call_args_list) + + +@pytest.mark.parametrize("scheme", ["https", "hf"]) +def test_central_url_recognition_and_parsing_share_supported_schemes(scheme): + url = URL.replace("hf://", f"{scheme}://") + assert is_central_asset_url(url) + assert _parse_central_asset_url(url) == _parse_central_asset_url(URL) + + +@pytest.mark.parametrize( + "url", + [ + "/cache/pallet.usda", + "s3://bucket/pallet.usd", + URL.replace("huggingface.co", "huggingface.co.example.com"), + URL.replace("/datasets/", "/models/"), + URL.replace("hf://", "http://"), + ], +) +def test_central_url_recognition_leaves_other_asset_paths_unchanged(url): + assert not is_central_asset_url(url) diff --git a/isaaclab_arena/tests/test_simready_object_library.py b/isaaclab_arena/tests/test_simready_object_library.py index a2f624ecb2..1b7b4fe88e 100644 --- a/isaaclab_arena/tests/test_simready_object_library.py +++ b/isaaclab_arena/tests/test_simready_object_library.py @@ -74,8 +74,8 @@ def test_simready_usd_object_enables_physics_variant_by_default(tmp_path): def test_simready_search_registry_name_derives_an_identifier_from_the_phrase(): from isaaclab_arena.assets.simready_object_library import simready_search_registry_name - assert simready_search_registry_name("green trash can") == "simready_green_trash_can" - assert simready_search_registry_name(" Chrome/Watering-Can ") == "simready_chrome_watering_can" + assert simready_search_registry_name("green trash can", "a.usd").startswith("simready_green_trash_can_") + assert simready_search_registry_name(" Chrome/Watering-Can ", "a.usd").startswith("simready_chrome_watering_can_") def test_register_searched_simready_object_makes_it_a_catalogue_entry(): @@ -86,8 +86,8 @@ def test_register_searched_simready_object_makes_it_a_catalogue_entry(): asset_cls = register_searched_simready_object( "green trash can", "s3://bucket/trash_can.usd", ("sim-ready", "green") ) - assert asset_cls.name == "simready_green_trash_can" - assert AssetRegistry().get_asset_by_name("simready_green_trash_can") is asset_cls + assert asset_cls.name.startswith("simready_green_trash_can_") + assert AssetRegistry().get_asset_by_name(asset_cls.name) is asset_cls # Rigid, so the entry can also be an object_set member; the search only accepts one rigid body. assert asset_cls.object_type is ObjectType.RIGID assert "green" in asset_cls.tags @@ -97,16 +97,20 @@ def test_a_searched_asset_resolves_its_usd_path_without_spec_params(): from isaaclab_arena.assets.simready_object_library import register_searched_simready_object ensure_assets_registered() - register_searched_simready_object("chrome watering can", "s3://bucket/watering_can.usd") + asset_cls = register_searched_simready_object("chrome watering can", "s3://bucket/watering_can.usd") # The path is baked into the registered class, so a spec naming it carries no params at all. - spec = AssetSpec(id="watering_can", registry_name="simready_chrome_watering_can") + spec = AssetSpec(id="watering_can", registry_name=asset_cls.name) assert spec.resolve_usd_path() == "s3://bucket/watering_can.usd" -def test_registering_the_same_phrase_twice_keeps_the_first_asset(): +def test_registering_distinct_assets_for_the_same_phrase_preserves_both(): from isaaclab_arena.assets.simready_object_library import register_searched_simready_object ensure_assets_registered() first = register_searched_simready_object("steel kettle", "s3://bucket/kettle_a.usd") second = register_searched_simready_object("steel kettle", "s3://bucket/kettle_b.usd") - assert second is first + assert second is not first + assert first.name != second.name + assert AssetSpec(id="first", registry_name=first.name).resolve_usd_path() == "s3://bucket/kettle_a.usd" + assert AssetSpec(id="second", registry_name=second.name).resolve_usd_path() == "s3://bucket/kettle_b.usd" + assert register_searched_simready_object("steel kettle", "s3://bucket/kettle_a.usd") is first diff --git a/isaaclab_arena_environments/maple_table_top/tabletop_pick_place_spray_gun_stand_bearing_wrench.yaml b/isaaclab_arena_environments/maple_table_top/tabletop_pick_place_spray_gun_stand_bearing_wrench.yaml new file mode 100644 index 0000000000..ba30fdd830 --- /dev/null +++ b/isaaclab_arena_environments/maple_table_top/tabletop_pick_place_spray_gun_stand_bearing_wrench.yaml @@ -0,0 +1,88 @@ +# 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 + +# Assets: NVIDIA SimReady Central (nvidia/simready-assets), CC-BY-4.0. +env_name: tabletop_pick_place_spray_gun_stand_bearing_wrench +embodiment: + id: franka + registry_name: franka_ik + params: {} +background: + id: maple_table + registry_name: maple_table_robolab + params: {} +objects: +- id: spray_gun + registry_name: simready_usd_object + params: + usd_path: https://huggingface.co/datasets/nvidia/simready-assets/resolve/653d91e182011f357cf1644013cc94ddf194d349/paint_supplies_spraygun_b01/simready_usd/sm_paint_supplies_spraygun_b01_01.usd +- id: spray_gun_stand + registry_name: simready_usd_object + params: + usd_path: https://huggingface.co/datasets/nvidia/simready-assets/resolve/653d91e182011f357cf1644013cc94ddf194d349/serviceTool_sprayGunStand_benchMount_a01/simready_usd/sm_serviceTool_sprayGunStand_benchMount_a01_01.usd +- id: ball_bearing + registry_name: simready_usd_object + params: + usd_path: https://huggingface.co/datasets/nvidia/simready-assets/resolve/653d91e182011f357cf1644013cc94ddf194d349/consumable_bearing_ball_a01/simready_usd/sm_consumable_bearing_ball_a01_01.usd +- id: combination_wrench + registry_name: simready_usd_object + params: + usd_path: https://huggingface.co/datasets/nvidia/simready-assets/resolve/653d91e182011f357cf1644013cc94ddf194d349/handTool_wrench_combination_a01/simready_usd/sm_handTool_wrench_combination_a01_01.usd +- id: paint_jar + registry_name: simready_usd_object + params: + usd_path: https://huggingface.co/datasets/nvidia/simready-assets/resolve/653d91e182011f357cf1644013cc94ddf194d349/paint_container_jar_orange_a01/simready_usd/sm_paint_container_jar_orange_a01_01.usd +- id: paint_brush + registry_name: simready_usd_object + params: + usd_path: https://huggingface.co/datasets/nvidia/simready-assets/resolve/653d91e182011f357cf1644013cc94ddf194d349/paint_supplies_brush_a08/simready_usd/sm_paint_supplies_brush_a08_01.usd +- id: paint_tube + registry_name: simready_usd_object + params: + usd_path: https://huggingface.co/datasets/nvidia/simready-assets/resolve/653d91e182011f357cf1644013cc94ddf194d349/paint_container_tube_d02/simready_usd/sm_paint_container_tube_d02_01.usd +relations: +- kind: is_anchor + subject: maple_table + params: {} +- kind: 'on' + subject: spray_gun + reference: maple_table + params: {} +- kind: 'on' + subject: spray_gun_stand + reference: maple_table + params: {} +- kind: 'on' + subject: ball_bearing + reference: maple_table + params: {} +- kind: 'on' + subject: combination_wrench + reference: maple_table + params: {} +- kind: 'on' + subject: paint_jar + reference: maple_table + params: {} +- kind: 'on' + subject: paint_brush + reference: maple_table + params: {} +- kind: 'on' + subject: paint_tube + reference: maple_table + params: {} +task: + composition: atomic + description: Franka picks up the spray gun from the maple table and places it on + the bench stand, leaving the paint jar, brush, paint tube, ball bearing, and + combination wrench as distractors. + subtasks: + - kind: PickAndPlaceTask + params: + pick_up_object: spray_gun + destination_location: spray_gun_stand + placement_consecutive_steps: 10 + background_scene: maple_table diff --git a/isaaclab_arena_examples/agentic_environment_generation/cli_runner.py b/isaaclab_arena_examples/agentic_environment_generation/cli_runner.py index ce1958689b..c8ef4eff32 100644 --- a/isaaclab_arena_examples/agentic_environment_generation/cli_runner.py +++ b/isaaclab_arena_examples/agentic_environment_generation/cli_runner.py @@ -56,7 +56,7 @@ def add_agentic_env_gen_runner_cli_args(parser: argparse.ArgumentParser) -> None: - from isaaclab_arena.agentic_environment_generation.simready_asset_search import SimReadySourceKind + from isaaclab_arena.agentic_environment_generation.simready_search_interface import SimReadySourceKind group = parser.add_argument_group("Agentic Environment Generation Runner") group.add_argument( @@ -119,11 +119,13 @@ def add_agentic_env_gen_runner_cli_args(parser: argparse.ArgumentParser) -> None help="Search SimReady for objects the Arena asset catalog does not cover.", ) group.add_argument( + "--simready_sources", "--simready_source", + nargs="+", type=str, choices=tuple(kind.value for kind in SimReadySourceKind), - default="isaac-sim-ga", - help="SimReady search backend (default: isaac-sim-ga Isaac Sim 6.0 GA props).", + default=["isaac-sim-ga"], + help="SimReady backends to search. Result limits apply per backend.", ) group.add_argument( "--simready_s3_url", @@ -153,7 +155,7 @@ def resolve_env_spec(args_cli: argparse.Namespace) -> Path | None: build_task_catalogue, ) from isaaclab_arena.agentic_environment_generation.environment_generation_agent import EnvironmentGenerationAgent - from isaaclab_arena.agentic_environment_generation.simready_asset_search import simready_search_config_from_cli + from isaaclab_arena.agentic_environment_generation.simready_search_interface import simready_search_config_from_cli # The generation passes log what they searched for and found; this is a console tool, so show # it. A no-op once something else owns the root logger, such as Kit in the modes that start it. @@ -165,7 +167,7 @@ def resolve_env_spec(args_cli: argparse.Namespace) -> Path | None: task_catalog = build_task_catalogue() simready_config = simready_search_config_from_cli( - source=args_cli.simready_source, + sources=args_cli.simready_sources, s3_url=args_cli.simready_s3_url, service_url=args_cli.simready_service_url, max_results_per_object=args_cli.simready_max_results_per_object, @@ -355,7 +357,6 @@ def main() -> int: parser = get_isaaclab_arena_cli_parser() add_agentic_env_gen_runner_cli_args(parser) args_cli = parser.parse_args() - if args_cli.mode == "schema": print_schema() return 0 diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/generation_panel.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/generation_panel.py index d06945c9a5..f359e7491d 100644 --- a/isaaclab_arena_examples/agentic_environment_generation/review_gui/generation_panel.py +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/generation_panel.py @@ -22,7 +22,7 @@ build_task_catalogue, ) from isaaclab_arena.agentic_environment_generation.environment_generation_agent import EnvironmentGenerationAgent -from isaaclab_arena.agentic_environment_generation.simready_asset_search import ( +from isaaclab_arena.agentic_environment_generation.simready_search_interface import ( SimReadySearchConfig, SimReadySourceKind, simready_search_config_from_cli, @@ -63,12 +63,13 @@ def get_catalogue_bundle() -> CatalogueBundle: def _simready_config_from_session() -> SimReadySearchConfig: """Read the SimReady search settings from Streamlit session state.""" - return simready_search_config_from_cli( - source=st.session_state.get("simready_source", SimReadySourceKind.ISAAC_SIM_GA.value), + config = simready_search_config_from_cli( + sources=st.session_state.get("simready_sources", [SimReadySourceKind.ISAAC_SIM_GA.value]), s3_url=st.session_state.get("simready_s3_url") or None, service_url=st.session_state.get("simready_service_url") or None, max_results_per_object=int(st.session_state.get("simready_max_results_per_object", 1)), ) + return config def available_inference_endpoints() -> list[str]: @@ -111,7 +112,7 @@ def _generation_agent_cache_key( "generation_agent", endpoint_name, simready_enabled, - simready_config.source.value, + tuple(source.value for source in simready_config.sources), simready_config.s3_url, simready_config.service_url, simready_config.max_results_per_object, @@ -318,13 +319,14 @@ def render_generation_panel() -> None: st.session_state["enable_simready_search"] = st.checkbox( "Enable SimReady search", value=bool(st.session_state.get("enable_simready_search", False)), - help="Search Isaac Sim GA SimReady props for objects the Arena asset catalog does not cover.", + help="Search SimReady for objects the Arena asset catalog does not cover.", ) source_options = [kind.value for kind in SimReadySourceKind] - st.session_state["simready_source"] = st.selectbox( - "SimReady source", + st.session_state["simready_sources"] = st.multiselect( + "SimReady sources", options=source_options, - index=source_options.index(st.session_state.get("simready_source", SimReadySourceKind.ISAAC_SIM_GA.value)), + help="Search each selected backend; the result limit applies to each backend.", + default=st.session_state.get("simready_sources", [SimReadySourceKind.ISAAC_SIM_GA.value]), ) st.session_state["simready_max_results_per_object"] = st.number_input( "Max results per object", @@ -332,25 +334,29 @@ def render_generation_panel() -> None: max_value=10, value=int(st.session_state.get("simready_max_results_per_object", 1)), ) - # Each source reads exactly one location, so only that one is offered. Showing both - # invites filling in a field the search then ignores. - source = st.session_state["simready_source"] - if source == SimReadySourceKind.S3.value: + sources = st.session_state["simready_sources"] + if SimReadySourceKind.S3.value in sources: st.session_state["simready_s3_url"] = st.text_input( "S3 URL", value=st.session_state.get("simready_s3_url", ""), placeholder="Leave empty for the Isaac Sim 6.0 GA SimReady bucket", ) - elif source == SimReadySourceKind.SERVICE.value: + if SimReadySourceKind.CENTRAL.value in sources: + st.caption( + "Searches Hugging Face SimReady Central; accepts only single rigid bodies without articulations." + ) + if SimReadySourceKind.SERVICE.value in sources: st.session_state["simready_service_url"] = st.text_input( "Service URL", value=st.session_state.get("simready_service_url", ""), placeholder=DEFAULT_SIMREADY_SERVICE_URL, ) - else: + if SimReadySourceKind.ISAAC_SIM_GA.value in sources: st.caption("Searches the Isaac Sim 6.0 GA SimReady library — no further configuration needed.") - generate_disabled = not available + generate_disabled = not available or ( + st.session_state.get("enable_simready_search", False) and not st.session_state.get("simready_sources") + ) if st.button("Generate spec", type="primary", width="stretch", disabled=generate_disabled): with st.spinner("Generating spec (LLM call)…"): ok, message = run_generation_pipeline(st.session_state["generation_prompt"]) diff --git a/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py b/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py index 1281acd349..b693fbc337 100644 --- a/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py +++ b/isaaclab_arena_examples/agentic_environment_generation/review_gui/streamlit_ui.py @@ -68,7 +68,6 @@ def initialize_state(yaml_path: Path | None, out_dir: Path) -> None: st.session_state["out_dir"] = str(out_dir.resolve()) 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"] = "" diff --git a/isaaclab_arena_examples/tests/test_review_gui.py b/isaaclab_arena_examples/tests/test_review_gui.py index 763ba8880f..a4b65e7b7e 100644 --- a/isaaclab_arena_examples/tests/test_review_gui.py +++ b/isaaclab_arena_examples/tests/test_review_gui.py @@ -13,7 +13,7 @@ import pytest -from isaaclab_arena.agentic_environment_generation.simready_asset_search import SimReadySearchConfig +from isaaclab_arena.agentic_environment_generation.simready_search_interface import SimReadySearchConfig from isaaclab_arena.agentic_environment_generation.spec_io import env_graph_spec_path, write_env_graph_spec from isaaclab_arena.assets.object_type import ObjectType from isaaclab_arena.environment_spec.arena_env_graph_spec import ArenaEnvGraphSpec diff --git a/pyproject.toml b/pyproject.toml index 752f0fb778..de86d5cb68 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,7 @@ dev = [ # its own: it adds one wheel, since Isaac Sim already brings everything it depends on, and an # extra only bought a way to install the review GUI with its asset search missing. "simready-search>=2026.4.2", + "huggingface-hub>=0.36", ] # PEP 735 dependency groups for the native uv install. pip (used by the Docker diff --git a/uv.lock b/uv.lock index 95ce179f2b..2729adccbf 100644 --- a/uv.lock +++ b/uv.lock @@ -1573,6 +1573,7 @@ dependencies = [ [package.optional-dependencies] dev = [ { name = "debugpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "huggingface-hub", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "jupyter", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "simready-search", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "streamlit", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -1613,6 +1614,7 @@ openpi = [ requires-dist = [ { name = "debugpy", marker = "extra == 'dev'" }, { name = "decorator", specifier = ">=4.0.2,<5" }, + { name = "huggingface-hub", marker = "extra == 'dev'", specifier = ">=0.36" }, { name = "imageio" }, { name = "imageio-ffmpeg" }, { name = "isaacteleop", extras = ["retargeters", "ui", "cloudxr"], specifier = "==1.4.126rc1", index = "https://pypi.org/simple" },