[Fix] Defer pxr loading in isaaclab.sim, .assets, .scene, and .sim/utils

[hujc7]

1. Summary

• Closes NVBug 6121889: the direct Cartpole env (Warp backend) + --visualizer kit crashes during SimulationApp.startup with a TfNotice / UsdAPISchemaBase / GfVec3f converter cascade.
• Root cause: module-import-time from pxr import … in IsaacLab library code. Env-cfg parse runs before AppLauncher and transitively pulls pxr into sys.modules, so Kit's USD-binding registration fails against an already partially-initialized pxr.
• Fix: defer pxr into the function bodies that use it across the 9 modules on the env-cfg parse path. 0 pxr/omni/carb/isaacsim modules now leak across all 197 Isaac-* env-cfg parses (the Cartpole path previously pulled in 63 pxr/omni/isaacsim modules).
• No public API change: function bodies are otherwise unchanged; the diff is import placement only.

2. Mechanism

2.1 Per-file pattern

Each module-top from pxr import … is removed. The runtime import moves into each function that uses USD; type hints stay under TYPE_CHECKING so static checkers and IDE autocomplete keep resolving Usd.Stage, Sdf.Path, etc. against the real classes.

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pxr import Sdf, Usd  # noqa: F401


def some_fn(prim: Usd.Prim) -> Sdf.Path:
    from pxr import Sdf, Usd  # noqa: PLC0415
    ...

# noqa: PLC0415 matches the function-local import sites already present across the repo; the TYPE_CHECKING guard matches the 200+ files already using if TYPE_CHECKING: from X import Y.

2.2 Affected modules

Module	pxr names deferred to function bodies
isaaclab.sim.simulation_context	Gf, UsdGeom, UsdPhysics, UsdUtils
isaaclab.assets.asset_base	Usd (type-only — no runtime pxr use)
isaaclab.scene.interactive_scene	Sdf
isaaclab.sim.spawners.from_files.from_files	Gf, Sdf, UsdGeom
isaaclab.sim.utils.prims	Sdf, Usd, UsdGeom, UsdPhysics, UsdShade, UsdUtils
isaaclab.sim.utils.queries	Sdf, Usd, UsdPhysics
isaaclab.sim.utils.semantics	Usd, UsdGeom, UsdSemantics
isaaclab.sim.utils.stage	Sdf, Usd, UsdUtils
isaaclab.sim.utils.transforms	Gf, Sdf, Usd, UsdGeom

3. Verification

• source/isaaclab_tasks/test/test_env_cfg_no_forbidden_imports.py (existing integration guard): all 197 Isaac-* tasks parse their env-cfg with 0 pxr/omni/carb/isaacsim/scipy imports.
• Static check: every runtime pxr symbol use in the 9 modules is covered by a function-local from pxr import …; type-only uses remain under TYPE_CHECKING (no module-scope pxr reference remains).
• ./isaaclab.sh -f clean.

4. Test plan

• ./isaaclab.sh -f clean.
• test_env_cfg_no_forbidden_imports.py — 197 tasks, no backend leak.
• Direct Cartpole (Warp) + --visualizer kit: env-cfg parse loads 0 pxr modules; SimulationApp.startup proceeds without the TfNotice / GfVec3f cascade.
• PhysX backend tests unaffected — function bodies unchanged, pxr still imports on first use.

5. Out of scope / follow-ups

~32 more files in isaaclab/, isaaclab_newton/, isaaclab_physx/ still from pxr import … at module top. They are not on the Cartpole-Warp env-cfg parse path so they don't block this fix, but should be converted the same way as the bug class will resurface for new envs that touch them:

• Tier A — likely on env-cfg chains (~14 files): cloner/cloner_utils, sim/spawners/{lights,materials,meshes,shapes,sensors,wrappers}, sim/schemas/schemas, sim/utils/{legacy,newton_model_utils}, sim/views/usd_frame_view, isaaclab_newton/{assets,cloner,sensors,sim/views}.
• Tier B — PhysX equivalents (~13 files): isaaclab_physx/*.
• Tier C — feature-gated, narrow trigger (~7 files): sensors/{camera,ray_caster}, envs/mdp/actions/task_space_actions, envs/ui/base_env_window, ui/xr_widgets, sim/converters/mesh_converter, utils/mesh.

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR defers pxr loading in asset base classes to runtime, unblocking workflows that import isaaclab.assets before launching Kit. The implementation correctly leverages the existing lazy-loading infrastructure in isaaclab.sim to achieve deferred loading without introducing new complexity.

Design Assessment

Design is sound. The approach is minimal and surgical:

1. Moving from pxr import Usd under TYPE_CHECKING preserves type hints while deferring the actual import
2. Using sim_utils.SimulationContext.instance() instead of direct imports leverages the existing lazy_loader infrastructure
3. Deferring the SimulationContext import to __init__ in base_articulation.py is appropriate since that's the only call site

The PR description is exemplary — it clearly explains the mechanism, provides verification steps, and acknowledges follow-up work for sensors.

Findings

🔵 Suggestion: Consider documenting the lazy-loading pattern — source/isaaclab/isaaclab/assets/asset_base.py

The TYPE_CHECKING pattern for deferring pxr imports could benefit from a brief inline comment explaining why it's structured this way. Future contributors unfamiliar with the import ordering constraints might inadvertently revert to direct imports.

if TYPE_CHECKING:
    # pxr is imported under TYPE_CHECKING to defer loading until runtime.
    # This allows `from isaaclab.assets import Articulation` before AppLauncher.
    from pxr import Usd

Test Coverage

• Bug fix: This is a fix for import-time side effects
• Regression test: The PR description mentions manual verification steps, and references test_env_cfg_no_forbidden_imports.py as a potential extension. The existing articulation tests should catch any runtime regressions.
• Gaps: Consider adding a test that explicitly validates from isaaclab.assets import Articulation doesn't trigger pxr loading (as mentioned in "Out of scope / follow-ups" section 4)

CI Status

Several checks still in progress (Build Base Docker Image, Installation Tests). Core checks (pre-commit, changelog, Build Wheel) have passed.

Verdict

No issues found

Clean implementation with proper lazy-loading patterns. The changes are minimal, correct, and well-documented. The only suggestion is adding a clarifying comment for maintainability.

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Update: This review reflects the expanded scope (now 10 files, ~155 LOC) since the initial review on commit 80b86a01.

Summary

This PR addresses NVBug 6121889 by deferring pxr (USD) imports that were loading at module-import time, which broke Kit's USD binding registration when AppLauncher started after env-cfg parsing. The fix extends lazy-loading to isaaclab.sim, .assets, .scene, and .sim/utils modules.

Architecture & Design ✅

The approach of using lazy_export() with .pyi stubs and TYPE_CHECKING guards is well-designed and maintains type-safety while achieving the deferred import behavior. The pattern is consistent across all modified modules.

Implementation Quality

Positive aspects:

• Clean, consistent pattern applied across all affected modules
• Type hints preserved via TYPE_CHECKING blocks
• Proper noqa: F401 annotations to suppress unused import warnings from type-only imports

No Blocking Issues Found

The implementation appears correct and follows best practices for deferred imports in Python.

---

Update (2573abe):
The author has introduced a cleaner architecture with the new lazy_imports() utility function:

1. New _LazyModule proxy class — Creates transparent proxies that import on first attribute access and self-replace in caller globals
2. lazy_imports(package, names) API — One-liner that binds multiple lazy submodules; returns None (side-effect-only API keeps call sites clean)
3. Comprehensive test suite — 12 pytest tests covering: no-preload guarantees, per-name isolation, self-replacement semantics, error handling for missing submodules, caller isolation between modules, and multi-call composition
4. Enhanced changelog — Now documents both the new utility (Added section) and provides detailed technical context for the fix

This is a significant improvement over the previous commit's ad-hoc approach. The lazy_imports() + _LazyModule pattern is well-documented, properly encapsulated, and thoroughly tested. Nice work! 👍

---

Update (28514a4):
Minor cleanup commit — moves TYPE_CHECKING imports and lazy_imports() calls to proper positions in several isaaclab.sim.utils.* modules:

• queries.py, semantics.py, stage.py, transforms.py: Added missing TYPE_CHECKING import, ensured lazy_imports() is called before any code that might reference the lazy-bound names
• interactive_scene.py: Moved the TYPE_CHECKING block below lazy_imports() call (correct ordering)

No functional changes — just import ordering fixes to ensure the lazy-import pattern works correctly. LGTM! ✅

---

Update (a82e275):
This commit improves the _LazyModule proxy class with better robustness and documentation:

1. Documentation enhancement — Added a note explaining that attribute assignment on the proxy is blocked by __slots__ before first access. Suggests either triggering first access first, or patching sys.modules[fully_qualified_name] directly for testing scenarios.

2. Defensive __getattr__ fix — The self-replace logic now uses sys.modules.get() instead of direct indexing, gracefully handling cases where the owner module was evicted from sys.modules (e.g., via importlib.reload or test cleanup with sys.modules.pop()). The proxy continues to work correctly without the native-lookup speedup.

3. New test case — test_lazy_module_owner_evicted_before_first_access() validates that the proxy doesn't crash with KeyError when the owner module is removed before first attribute access.

4. Docstring ordering — Swapped the lazy_imports() call before the TYPE_CHECKING block in the example, reflecting best practice.

Solid defensive improvement for edge cases. LGTM! ✅

---

Update (701cc34):
Significant simplification! The author removed the lazy_imports() / _LazyModule infrastructure in favor of direct local imports inside function bodies:

Changes:

• ❌ Removed lazy_imports() function and _LazyModule class (~100 LOC)
• ❌ Removed test_module.py test file (~265 LOC)
• ✅ All pxr imports now use explicit from pxr import X # noqa: PLC0415 inside each function that needs them
• ✅ Changelog updated to reflect the simpler approach

Assessment:
This is a good architectural pivot. The local-import pattern is:

• Simpler — no custom proxy classes to understand
• Explicit — every function shows exactly what it needs
• Standard Python — PLC0415 (import-outside-toplevel) is a well-known pattern for deferred imports

Trade-off: slightly more import statements scattered across functions vs. maintaining custom lazy-loading machinery. For this use case (deferring a handful of pxr imports), the explicit approach wins on maintainability.

LGTM! ✅

---

Update (dd576b0):
Minimal-diff cleanup! This commit reverts incidental import reroutes that were no longer needed:

Changes in asset_base.py:

• Now imports SimulationContext and get_current_stage directly at module top, instead of accessing them via sim_utils.SimulationContext / sim_utils.get_current_stage()
• This is cleaner — the module directly expresses its dependencies rather than going through the lazy namespace indirection

Changes in simulation_context.py:

• Uses stage_utils.close_stage() directly instead of sim_utils.close_stage()
• Uses create_new_stage directly from the explicit import, not through sim_utils
• The build_simulation_context() context manager now correctly uses the locally-imported create_new_stage function

Rationale (from commit message):

The pxr fix only requires deferring the foreign from pxr import ... statements; it does not require changing how these two modules import their own isaaclab dependencies.

This is exactly right. The earlier commits had rerouted isaaclab-internal imports through the lazy sim_utils namespace as a workaround during development. Now that pxr is deferred everywhere, those reroutes are unnecessary overhead. Reverting them keeps the diff scoped to the actual bug fix.

LGTM! ✅ Nice attention to minimal, focused changes.

---

Update (3ca9899):
Bug fix for name shadowing in build_simulation_context():

• In simulation_context.py, the function parameter create_new_stage: bool was shadowing the imported create_new_stage function
• Changed call from create_new_stage() to sim_utils.create_new_stage() to access the function through the namespace
• Added clarifying comment explaining the shadowing issue

This is a necessary fix — without it, build_simulation_context(create_new_stage=True) would fail because Python would try to call the boolean parameter instead of the function.

LGTM! ✅

---

Update (1b53736):
Changelog wording refinement only. The .rst file now uses more concise phrasing:

• "Each affected module keeps its type hints under TYPE_CHECKING and defers the runtime import" → "Type hints stay under TYPE_CHECKING; where pxr is used at runtime, the import is deferred into the function body"

No code changes, just documentation polish. LGTM! ✅

[greptile-apps]

Greptile Summary

This PR fixes a crash triggered when env-cfg parsing (which runs before AppLauncher) transitively loaded pxr into sys.modules, causing Kit's USD-binding registration to fail. The fix introduces lazy_imports() — a side-effect helper that installs _LazyModule proxies in the caller's module globals, deferring the actual from pxr import … until first attribute access inside a function body.

• New helper (isaaclab.utils.module.lazy_imports): a proxy-based lazy-import mechanism with self-replacement on first access, complementing the existing lazy_export() for package __init__ files; 13 unit tests including a direct regression test verifying pxr.* stays out of sys.modules after the call.
• Nine leaf modules converted (sim/simulation_context, assets/asset_base, scene/interactive_scene, sim/spawners/from_files, sim/utils/{prims,queries,semantics,stage,transforms}): each replaces from pxr import … with lazy_imports("pxr", […]) plus a TYPE_CHECKING guard, leaving all function bodies unchanged.
• simulation_context.py uses inline method-level from pxr import … (inside __init__ and _init_usd_physics_scene) rather than lazy_imports, consistent with the singleton guard pattern and correct for the same reason.

Confidence Score: 4/5

Safe to merge — the change is purely additive (new helper, proxy wrappers) with no functional changes to any function bodies or business logic.

The core mechanism is well-tested and the root cause fix is correct. The findings are edge-case robustness gaps in the new _LazyModule helper — a confusing KeyError if an owner module is evicted from sys.modules before first proxy access, a reversed docstring example, and an undocumented write-to-proxy limitation — none of which affect the primary use path.

source/isaaclab/isaaclab/utils/module.py — the new _LazyModule.__getattr__ implementation and lazy_imports docstring example are the areas worth a second look.

Important Files Changed

Filename	Overview
source/isaaclab/isaaclab/utils/module.py	Introduces _LazyModule proxy class and lazy_imports() helper. Core logic is correct; minor issues with docstring example ordering (reversed vs. all call sites) and KeyError when owner module is evicted from sys.modules before first proxy access.
source/isaaclab/isaaclab/sim/simulation_context.py	Moves pxr imports (UsdUtils, Gf, UsdGeom, UsdPhysics) to inline method-level imports instead of lazy_imports; inlines the stage_utils import into __init__. Functionally correct — pxr is never loaded at module import time.
source/isaaclab/isaaclab/assets/asset_base.py	Removes direct from pxr import Usd and eager SimulationContext/get_current_stage imports; replaces with sim_utils.SimulationContext and sim_utils.get_current_stage() through the lazy isaaclab.sim namespace — correct and clean.
source/isaaclab/isaaclab/scene/interactive_scene.py	Replaces from pxr import Sdf with lazy_imports("pxr", ["Sdf"]) and adds the corresponding TYPE_CHECKING guard. No functional changes.
source/isaaclab/isaaclab/sim/utils/stage.py	Replaces from pxr import Sdf, Usd, UsdUtils with lazy_imports pattern and adds TYPE_CHECKING guard. Correct; _context = threading.local() is unaffected.
source/isaaclab/isaaclab/sim/utils/prims.py	Replaces from pxr import Sdf, Usd, UsdGeom, UsdPhysics, UsdShade, UsdUtils with lazy_imports pattern; adds TYPE_CHECKING guard. Uniform and correct.
source/isaaclab/test/utils/test_module.py	13 well-structured pytest tests covering construction, first-access swap, per-name isolation, caller isolation, empty-list noop, unknown-module error, and real-semantics checks. Includes a direct regression test (test_lazy_imports_does_not_preload_pxr).

Sequence Diagram

sequenceDiagram
    participant EC as env-cfg parse
    participant IA as isaaclab.assets
    participant AB as asset_base.py
    participant SU as isaaclab.sim (lazy_export)
    participant SG as sim/utils/stage.py
    participant LI as lazy_imports proxy
    participant PXR as pxr (USD)
    participant AL as AppLauncher / SimulationApp

    Note over EC,AL: BEFORE this PR — crash path
    EC->>IA: from isaaclab.assets import Articulation
    IA->>AB: import asset_base
    AB->>PXR: from pxr import Usd  ← module-top import
    PXR-->>AB: pxr loaded into sys.modules
    EC->>AL: AppLauncher.__init__()
    AL->>AL: SimulationApp.startup()
    AL-->>AL: ❌ pxr already initialised — TfNotice / GfVec3f cascade

    Note over EC,AL: AFTER this PR — fixed path
    EC->>IA: from isaaclab.assets import Articulation
    IA->>AB: import asset_base
    AB->>SU: import isaaclab.sim as sim_utils (lazy_export, no pxr)
    AB->>SG: stage.py lazy_imports proxy bound (pxr NOT loaded)
    SG->>LI: _LazyModule stored in module globals
    EC->>AL: AppLauncher.__init__()
    AL->>AL: SimulationApp.startup() ✅ pxr not in sys.modules
    AL->>PXR: Kit registers USD bindings cleanly
    AL-->>AL: SimulationContext.__init__() called
    AL->>PXR: from pxr import UsdUtils  ← first actual pxr load
    LI->>PXR: _LazyModule.__getattr__ → importlib.import_module
    PXR-->>LI: real module returned, proxy self-replaced in globals

Loading

_{Reviews (1): Last reviewed commit: "Extend pxr lazy-load to scene, spawners,..." | Re-trigger Greptile}

[greptile-apps]

KeyError instead of AttributeError when owner module is evicted from sys.modules

If the owner module is removed from sys.modules (e.g. during a importlib.reload cycle or a test that does sys.modules.pop(name)) between proxy construction and first attribute access, sys.modules[self._owner_mod_name] raises a bare KeyError. Callers see a confusing KeyError: '__some_internal_module__' instead of an AttributeError or the real ModuleNotFoundError. The self-replacement optimization can safely be skipped in this case — the proxy will still return the correct attribute from the real module on every call via __getattr__.

Using sys.modules.get() instead avoids the error without changing observable behaviour for the common path.

[greptile-apps]

lazy_imports docstring example reverses the convention used at every call site

The docstring shows if TYPE_CHECKING: placed before the lazy_imports() call, but all 9 call sites in this PR place it after (and the PR description itself shows the same "after" order). A developer copying from the docstring as a template would produce code with the opposite ordering.

Suggested change

	Example:
	::
	from isaaclab.utils.module import lazy_imports
	if TYPE_CHECKING:
	from pxr import Sdf, Usd, UsdGeom # noqa: F401
	lazy_imports("pxr", ["Sdf", "Usd", "UsdGeom"])
	Example:
	::
	from isaaclab.utils.module import lazy_imports
	lazy_imports("pxr", ["Sdf", "Usd", "UsdGeom"])
	if TYPE_CHECKING:
	from pxr import Sdf, Usd, UsdGeom # noqa: F401

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

_LazyModule silently drops __setattr__ calls before first access

__slots__ prevents arbitrary attribute writes on the proxy, so any code that tries to assign to a lazy-module name before the real module is loaded — e.g., a mock that patches module.Sdf — will raise AttributeError: '_LazyModule' object has no attribute 'X' rather than transparently forwarding the write to the real module. The current use patterns (read-only accesses inside function bodies) are not affected, but this is an undocumented limitation worth calling out in the class docstring for future adopters.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[kellyguo11]

what's the main reason why we can't import pxr before these modules?

[hujc7]

what's the main reason why we can't import pxr before these modules?

Kit registers a specialized pxr such that, if pxr is imported before appLaunch, the later kit registration will have problems. The crashing behavior could be different but it seems the root cause is the same.