Defer pxr-importing usd_prim_tree import in pass-2 resolver

The pass-2 resolver imported usd_prim_tree at module top, which pulls in
pxr before SimulationApp is created and broke unit tests that import the
generation agent. Import load_usd_prim_tree lazily inside infer() and
guard UsdPrimRecord under TYPE_CHECKING, and patch the source module in
the affected tests.

Signed-off-by: Qian Lin <qianl@nvidia.com>

Author: qianl-nv

isaaclab_arena/agentic_environment_generation/object_reference_prim_resolver.py

@@ -8,7 +8,7 @@
 from __future__ import annotations
 
 import json
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from pydantic import BaseModel, Field, ValidationError, ValidationInfo, model_validator
 
@@ -18,7 +18,9 @@
 from isaaclab_arena.environments.arena_env_graph_spec import ArenaEnvGraphSpec
 from isaaclab_arena.environments.arena_env_graph_types import ObjectReferenceSpec
 from isaaclab_arena.utils.asset_usd import resolve_asset_usd_path
-from isaaclab_arena.utils.usd_prim_tree import UsdPrimRecord, load_usd_prim_tree
+
+if TYPE_CHECKING:
+    from isaaclab_arena.utils.usd_prim_tree import UsdPrimRecord
 
 
 class ResolvedObjectReferences(BaseModel):
@@ -147,6 +149,9 @@ def infer(
             The spec with resolved prim paths on success, otherwise ``None`` (with
             error lines appended to ``traces``).
         """
+        # Defer pxr import until call time to avoid conflict with SimulationApp.
+        from isaaclab_arena.utils.usd_prim_tree import load_usd_prim_tree
+
         usd_path = resolve_asset_usd_path(spec.background)
         prim_tree = load_usd_prim_tree(usd_path)
         data = self._query_backend.run_json(

isaaclab_arena/tests/test_environment_generation_agent.py

@@ -101,7 +101,7 @@ def test_builds_catalogues_from_singleton_registries_when_none(self, agent):
         mock_build_relations.assert_called_once_with()
         mock_build_tasks.assert_called_once_with()
 
-    @patch("isaaclab_arena.agentic_environment_generation.object_reference_prim_resolver.load_usd_prim_tree")
+    @patch("isaaclab_arena.utils.usd_prim_tree.load_usd_prim_tree")
     @patch("isaaclab_arena.agentic_environment_generation.object_reference_prim_resolver.resolve_asset_usd_path")
     def test_two_pass_generate_spec_resolves_object_references(self, mock_resolve_usd, mock_load_tree, agent):
         mock_resolve_usd.return_value = "/tmp/scene.usd"
@@ -120,7 +120,7 @@ def test_two_pass_generate_spec_resolves_object_references(self, mock_resolve_us
         assert agent.client.chat.completions.create.call_count == 2
         assert result.object_references
 
-    @patch("isaaclab_arena.agentic_environment_generation.object_reference_prim_resolver.load_usd_prim_tree")
+    @patch("isaaclab_arena.utils.usd_prim_tree.load_usd_prim_tree")
     @patch("isaaclab_arena.agentic_environment_generation.object_reference_prim_resolver.resolve_asset_usd_path")
     def test_two_pass_generate_spec_returns_dict_on_pass2_failure(self, mock_resolve_usd, mock_load_tree, agent):
         mock_resolve_usd.return_value = "/tmp/scene.usd"

isaaclab_arena/tests/test_object_reference_prim_resolver.py

@@ -23,7 +23,7 @@
 )
 
 
-@patch("isaaclab_arena.agentic_environment_generation.object_reference_prim_resolver.load_usd_prim_tree")
+@patch("isaaclab_arena.utils.usd_prim_tree.load_usd_prim_tree")
 @patch("isaaclab_arena.agentic_environment_generation.object_reference_prim_resolver.resolve_asset_usd_path")
 def test_object_reference_prim_resolver_infer_merges_llm_output(mock_resolve_usd, mock_load_tree):
     mock_resolve_usd.return_value = "/tmp/scene.usd"
@@ -79,7 +79,7 @@ def test_object_reference_prim_resolver_infer_merges_llm_output(mock_resolve_usd
         ),
     ],
 )
-@patch("isaaclab_arena.agentic_environment_generation.object_reference_prim_resolver.load_usd_prim_tree")
+@patch("isaaclab_arena.utils.usd_prim_tree.load_usd_prim_tree")
 @patch("isaaclab_arena.agentic_environment_generation.object_reference_prim_resolver.resolve_asset_usd_path")
 def test_object_reference_prim_resolver_infer_records_invalid_llm_output(
     mock_resolve_usd,