Fix Newton visualization ClonePlan handling

[nblauch]

Description

Newton's PhysX-backed shadow visualization model ignored the active ClonePlan and built every rendered world by importing env_0 once, then replicating that prototype across all worlds. That made Newton RGB/depth visualization homogeneous even when the PhysX scene used heterogeneous cloned sources, for example different object prototypes per environment.

This change builds the Newton shadow visualization model from the active ClonePlan when available:

• caches one Newton ModelBuilder per ClonePlan source prototype
• assembles each render world from the source rows selected by clone_mask
• rewrites copied Newton labels from source paths to destination env paths so PhysX scene-data transform sync continues to map poses correctly
• keeps the existing env_0 prototype fallback for scenes without a ClonePlan

Fixes # (not filed)

Type of change

• Bug fix (non-breaking change which fixes an issue)

Screenshots

N/A. This was validated with a downstream Dextrah Newton RGB distillation video smoke; env 0/1/2 rendered different object geometries after the fix.

Checklist

• I have read and understood the contribution guidelines
• I have run the pre-commit checks with ./isaaclab.sh --format
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• I have added a changelog fragment under source/<pkg>/changelog.d/ for every touched package (do not edit CHANGELOG.rst or bump extension.toml — CI handles that)
• I have added my name to the CONTRIBUTORS.md or my name already exists there

Validation run:

• python -m py_compile source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
• python -m pytest source/isaaclab_newton/test/physics/test_newton_visualization_clone_plan.py -q
• ./isaaclab.sh --format
• downstream Dextrah Newton distillation smoke with video recording enabled

[isaaclab-review-bot]

Review Summary

This PR fixes a legitimate bug where Newton's shadow visualization model ignored the active ClonePlan, building every render world from env_0 and producing homogeneous visualization regardless of the PhysX scene's heterogeneous cloning. The fix is well-architected with clean decomposition, defensive error handling, and a solid test suite.

Verdict: Looks good overall. A few suggestions below, none blocking.

---

Strengths

• Clean decomposition — Each helper method (_visualization_world_transform, _rewrite_visualization_label_prefix, _clone_plan_selects_visualization_source, etc.) has a single responsibility and is independently testable.
• Source builder caching — _get_clone_plan_visualization_source_builder correctly memoizes the expensive add_usd call per source root, avoiding redundant USD parsing across N environments.
• Robust fallback chain — Missing prims, invalid transforms, and malformed clone_mask shapes all produce warnings and graceful degradation rather than crashes.
• Improved label rewriting — The new startswith-based prefix matching in _rewrite_visualization_label_prefix is more correct than the old str.replace(..., 1) approach which could match non-prefix occurrences.

---

Suggestions

1. GPU clone_mask may cause per-element device sync (performance, non-blocking)

File: newton_manager.py, _clone_plan_selects_visualization_source

return bool(clone_mask[source_index, env_id].item())

ClonePlan.clone_mask may reside on a CUDA device (see ClonePlan.from_env_0 which accepts a device parameter). Calling .item() inside a loop over N_envs × N_sources would trigger a host↔device sync per iteration. For 4096 envs × 4 sources, that's ~16K synchronization points.

Suggestion: Move the mask to CPU once before the env loop in _add_clone_plan_visualization_worlds:

clone_mask = clone_plan.clone_mask.cpu() if clone_plan.clone_mask.is_cuda else clone_plan.clone_mask

Since this runs at model-build time (not per-frame), it's not critical, but it avoids a surprising latency cliff when scaling env count.

2. Consider a regression test for the refactored _add_env0_prototype_visualization_worlds (test coverage, non-blocking)

The env_0 prototype replication logic was moved from inline code into a standalone method and the label rewrite was changed from str.replace to the new prefix-based helper. While the behavior is equivalent for standard cases, a dedicated test exercising this fallback path (e.g., when SimulationContext.instance() returns None) would guard against future regressions.

3. _visualization_destination_for_env — clarify the env_name heuristic (documentation, non-blocking)

File: newton_manager.py, lines ~1708-1715

env_name = env_path.rsplit("/", 1)[-1]
if env_name != str(env_id):
    named_destination_root = destination_template.format(env_name)

This fallback handles the case where env prims aren't named env_<id> but use custom names. A brief inline comment explaining when this can occur (e.g., custom scene layouts or renamed env prims) would help future readers understand why the extra lookup is needed.

4. Type annotation for clone_plan parameter (style, non-blocking)

In _add_clone_plan_visualization_worlds, the clone_plan parameter is untyped:

clone_plan,

Adding clone_plan: ClonePlan would improve IDE support and self-documentation. The import can be guarded under TYPE_CHECKING to avoid circular dependencies since the runtime from isaaclab.sim import SimulationContext is already conditional.

---

Test Analysis

The test suite is solid:

• ✅ Heterogeneous clone_mask with correct source selection per env
• ✅ Label rewriting from source paths to destination env paths
• ✅ Trailing slash handling in prefix rewrite
• ✅ Graceful degradation when clone_mask row count is insufficient

One note on test_visualization_builder_falls_back_when_clone_mask_source_row_is_missing: the test correctly verifies that envs 1 and 2 fall back to direct add_usd with root_path=env_path (the full env root, not the Object subtree). This is intentional — when the ClonePlan cannot determine which source to use, importing the entire env subtree is the safest fallback.

---

Changelog

The fragment accurately describes the fix. No issues.

---

Reviewed commit: 856bb84

[isaaclab-review-bot]

The expected values here seem inconsistent with the test description. When the clone_mask is malformed (only 1 row for 2 sources), the code logs a warning and falls back to direct env import via builder.add_usd(stage, root_path=env_path). This would import the full env path, not the /Object subpath.

Is this the intended fallback behavior? The env_1 and env_2 assertions expect /World/envs/env_1 and /World/envs/env_2 respectively, while env_0 expects /World/envs/env_0/Object. This asymmetry suggests either:

1. The fallback is intentionally different (imports the entire env vs the cloned subtree), or
2. There might be a bug in the test expectations

Please clarify the intended semantics.

[isaaclab-review-bot]

Consider adding explicit type annotation for clone_mask. Based on usage with .shape and .item(), it appears to be a torch.Tensor, but the type isn't documented. A type hint like clone_mask: torch.Tensor would improve readability.

[isaaclab-review-bot]

The string prefix matching here could produce unexpected results if prim paths have overlapping prefixes (e.g., /World/envs/env_0 and /World/envs/env_0_extended). The trailing slash stripping helps prevent full-path mismatches, but a brief comment explaining why this approach is safe for USD prim paths would be helpful for future maintainers.

[greptile-apps]

Greptile Summary

This PR fixes Newton's PhysX-backed shadow visualization model to respect the active ClonePlan when building per-env Newton worlds, replacing the old env_0-only prototype approach with a per-source-builder path that correctly handles heterogeneous environments.

• Extracts build_source_builders, replicate_builder_mapping, and rename_builder_labels into a shared newton_clone_utils module and rewires both the physics-replication path and the new visualization path to use them.
• Adds build_visualization_builder_from_stage_envs in a new visualization_builder.py module and calls it from _ensure_visualization_model; however, the call site passes get_clone_plan() which returns ClonePlan | None and the function has no None guard, breaking all scenes that do not publish a ClonePlan.
• Removes the try/except safety net that previously wrapped the visualization model build, so any construction error now propagates uncaught through get_model() / get_state_0() / update_visualization_state().

Confidence Score: 3/5

Not safe to merge as-is: any scene without a published ClonePlan will hit an uncaught AttributeError in _ensure_visualization_model, which is called from get_model(), get_state_0(), and update_visualization_state() on every render frame.

The visualization path now unconditionally dereferences get_clone_plan() without a None guard. That method returns None for any scene that does not use the cloner, making the visualization model build fail with an uncaught AttributeError on every frame for those users. The removed try/except wrapper compounds this because the error propagates rather than being logged. The core ClonePlan logic and the physics-replication refactor look correct; it is specifically the call site in _ensure_visualization_model and the missing fallback that need fixing before merge.

newton_manager.py (call site lacks None guard and lost its exception handler) and visualization_builder.py / newton_clone_utils.py (both assume a non-None ClonePlan with fully-populated source rows).

Important Files Changed

Filename	Overview
source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py	Replaces the old env_0-prototype visualization builder with a ClonePlan-aware call, but drops the try/except safety net and has no None guard for get_clone_plan(), breaking all non-ClonePlan scenes at runtime.
source/isaaclab_newton/isaaclab_newton/physics/visualization_builder.py	New module that builds the shadow visualization model from a ClonePlan; correct when a plan is present, but has no guard for a None argument passed from the call site.
source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py	New shared utilities for building and replicating Newton source builders; source_world_indices comprehension regresses from the prior None-safe guard and will raise IndexError for all-False mapping rows.
source/isaaclab_newton/isaaclab_newton/cloner/replicate.py	Refactored to delegate core builder construction to newton_clone_utils; logic is equivalent to the prior implementation; _deformable_registry guard removed safely because the attribute is now always declared.
source/isaaclab/isaaclab/cloner/cloner_utils.py	Adds replace_path_prefix helper with correct boundary-safe prefix replacement; no issues.
source/isaaclab_newton/test/cloner/test_rename_builder_labels.py	Tests updated and extended with TestVisualizationClonePlan; the new test covers the happy path but does not exercise the None-ClonePlan fallback that the PR description claims to preserve.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["_ensure_visualization_model()"] --> B{backend is Newton?}
    B -- yes --> C[return no-op]
    B -- no --> D[get USD stage]
    D --> E{stage is None?}
    E -- yes --> F[log error, return]
    E -- no --> G[collect env_paths from /World/envs]
    G --> H{env_paths empty?}
    H -- yes --> I[log error, return]
    H -- no --> J["get_clone_plan() -> ClonePlan | None"]
    J --> K["build_visualization_builder_from_stage_envs(stage, env_paths, clone_plan)"]
    K -- clone_plan is None --> L["AttributeError: NoneType has no attribute sources"]
    K -- clone_plan valid --> M["build_source_builders()"]
    M --> N["replicate_builder_mapping()"]
    N --> O["rename_builder_labels()"]
    O --> P{body_count == 0?}
    P -- yes --> Q[log error, return]
    P -- no --> R["builder.finalize() wrapped in try/except"]
    R --> S[model and state set]
    style L fill:#ff6b6b,color:#fff

Loading

_{Reviews (2): Last reviewed commit: "Fix Newton visualization clone-plan sour..." | Re-trigger Greptile}

[greptile-apps]

env_id used as clone_mask column index, not sequential env position

env_id is the integer extracted from the prim name (e.g., 5 for env_5) and is used directly as the column index into clone_mask. After env_paths.sort(key=lambda x: x[0]), the zero-based position of each entry in the sorted list is the correct column index, not env_id. For the standard IsaacLab case (envs named env_0, env_1, … contiguously), both coincide. But if any run uses non-contiguous numbering (e.g., env_0, env_2, env_4), the code would read clone_mask[:, 4] for env_4 instead of clone_mask[:, 2], silently selecting wrong sources or logging spurious "env id exceeds clone_mask column count" warnings. The same env_id is forwarded to _visualization_destination_for_env for the format() call, which is correct — only the mask indexing needs the sequential index.

[greptile-apps]

begin_world / end_world not paired with try/finally

builder.begin_world() is called unconditionally at the top of the env loop, but builder.end_world() sits at the bottom of the loop body. If the fallback builder.add_usd(...) (line 1819) raises — for example because schema_resolvers fails to parse an env prim — end_world() is never called, potentially leaving the builder in a half-open world state for all subsequent envs. The outer try/except in _ensure_visualization_model will catch the exception and abort model building, so the visible impact is limited, but a try/finally around the loop body would make this robust.

[isaaclab-review-bot]

Review Summary

This PR fixes a legitimate bug where Newton's shadow visualization model ignored the active ClonePlan, building every render world from env_0 and producing homogeneous visualization regardless of the PhysX scene's heterogeneous cloning. The fix is well-architected with clean decomposition, defensive error handling, and a solid test suite.

Verdict: Looks good overall. A few suggestions below, none blocking.

---

Strengths

• Clean decomposition — Each helper method (_visualization_world_transform, _rewrite_visualization_label_prefix, _clone_plan_selects_visualization_source, etc.) has a single responsibility and is independently testable.
• Source builder caching — _get_clone_plan_visualization_source_builder correctly memoizes the expensive add_usd call per source root, avoiding redundant USD parsing across N environments.
• Robust fallback chain — Missing prims, invalid transforms, and malformed clone_mask shapes all produce warnings and graceful degradation rather than crashes.
• Improved label rewriting — The new startswith-based prefix matching in _rewrite_visualization_label_prefix is more correct than the old str.replace(..., 1) approach which could match non-prefix occurrences.

---

Suggestions

1. GPU clone_mask may cause per-element device sync (medium — performance)

File: newton_manager.py, _clone_plan_selects_visualization_source

return bool(clone_mask[source_index, env_id].item())

ClonePlan.clone_mask may reside on a CUDA device. Calling .item() inside a loop over N_envs × N_sources would trigger a host↔device sync per iteration. For 4096 envs × 4 sources, that's ~16K synchronization points.

Suggestion: Move the mask to CPU once before the env loop in _add_clone_plan_visualization_worlds:

clone_mask = clone_plan.clone_mask
if hasattr(clone_mask, 'is_cuda') and clone_mask.is_cuda:
    clone_mask = clone_mask.cpu()

Since this runs at model-build time (not per-frame), it's not critical, but it avoids a surprising latency cliff when scaling env count.

2. Consider a regression test for the refactored _add_env0_prototype_visualization_worlds (low — test coverage)

The env_0 prototype replication logic was moved from inline code into a standalone method and the label rewrite was changed from str.replace to the new prefix-based helper. While the behavior is equivalent for standard cases, a dedicated test exercising this fallback path (e.g., when SimulationContext.instance() returns None) would guard against future regressions in the non-ClonePlan code path.

3. _visualization_destination_for_env returns potentially non-existent path (low — robustness)

File: newton_manager.py, _visualization_destination_for_env

return destination_root  # Returns path even if prim doesn't exist

When neither destination_template.format(env_id) nor destination_template.format(env_name) resolves to a valid prim, the method returns the unvalidated destination_root. This is caught downstream by _visualization_relative_transform returning None (which triggers a skip + warning), so it doesn't cause a crash — but the resulting warning message could be misleading since it says the "destination prim is missing" without indicating that the lookup exhausted fallback heuristics.

A brief debug-level log here (before the final return) would make debugging destination resolution failures easier.

4. Type annotation for clone_plan parameter (nit — style)

In _add_clone_plan_visualization_worlds, the clone_plan parameter is untyped. Adding clone_plan: ClonePlan (with a TYPE_CHECKING-guarded import) would improve IDE support and self-documentation.

---

Test Analysis

The test suite is solid:

• ✅ Heterogeneous clone_mask with correct source selection per env
• ✅ Label rewriting from source paths to destination env paths
• ✅ Trailing slash handling in prefix rewrite
• ✅ Graceful degradation when clone_mask row count is insufficient

One note on test_visualization_builder_falls_back_when_clone_mask_source_row_is_missing: the fallback for envs 1 and 2 correctly imports the full env path (not the /Object subtree), since when the ClonePlan cannot determine which source to use, importing the entire env subtree is the safest fallback. The asymmetry between env_0's expected value (/World/envs/env_0/Object) and envs 1-2 (/World/envs/env_1, /World/envs/env_2) accurately reflects the two different code paths being exercised.

---

Changelog

The fragment accurately describes the fix. No issues.

---

Reviewed commit: 856bb84

---

Update (2b9b99b): All previously raised concerns have been addressed in this revision:

1. ✅ GPU clone_mask sync — Now converted to CPU upfront via getattr(clone_mask, "is_cuda", False).
2. ✅ env_id vs sequential index for mask column — Fixed by using env_index from enumerate(env_paths) instead of env_id. Test updated with non-contiguous IDs (0, 2, 4) to validate.
3. ✅ begin_world/end_world try/finally — Both clone-plan and env0-prototype paths now use try/finally.
4. ✅ Type annotation for clone_plan — Added clone_plan: ClonePlan with TYPE_CHECKING import.
5. ✅ Prefix matching comment — Clarifying comment added to _rewrite_visualization_label_prefix.
6. ✅ Destination fallback helper — Removed entirely; logic inlined with direct warning at failure point.

The refactoring also consolidates several small helper methods into the main loop, reducing indirection while maintaining the same defensive behavior. No new issues found. LGTM 👍

Update (0a28e22): The PR has been substantially refactored — the visualization model construction now uses an inline approach rather than multiple helper methods. Key observations on the new changes:

Previously raised concerns — status:

1. ✅ GPU clone_mask sync — Still addressed; .detach().cpu() is called upfront.
2. ✅ env_id vs sequential index — Resolved differently via env_id_to_col mapping from clone_plan.env_ids.
3. ⚠️ begin_world/end_world try/finally — Not present in the new inline visualization path. Low risk since failures here abort the entire model build, but still worth noting.

New changes in replicate.py:

• _build_newton_builder_from_mapping now accepts a destinations parameter and computes per-source-world transforms (source_world_indices) instead of a single env_0 delta. This is a correctness fix for heterogeneous clone plans where different sources have different world transforms. Looks correct.

Unrelated addition (lambda security fix in string.py):

• Adds AST validation (_validate_lambda_expression) blocking Call, Attribute, NamedExpr nodes and dunder names in lambda strings before eval. Also restricts eval to empty __builtins__. Good hardening with solid test coverage. No issues.

No new blocking issues found. LGTM 👍

Update (50020c8): Further refactoring simplifies the visualization model builder significantly. Changes reviewed:

replicate.py: Extracted _rewrite_label_prefix as a standalone module-level function (previously inline logic in _rename_builder_labels). Clean, reusable, well-documented with path-boundary semantics.

newton_manager.py: Major cleanup of _build_visualization_model_from_stage:

• _world_transform helper closure eliminates repeated transform extraction boilerplate.
• Heterogeneous ClonePlan path reuses _rewrite_label_prefix from replicate.py — correct and DRY.
• Homogeneous fallback now also uses _rewrite_label_prefix instead of str.replace(..., 1), fixing the prefix-boundary issue.

Previous concerns — status:

1. ✅ GPU clone_mask sync — Still addressed via .detach().cpu() upfront.
2. ✅ env_id vs sequential index — Still resolved via env_id_to_col mapping.
3. ✅ Prefix matching — Now uses well-documented _rewrite_label_prefix with explicit boundary check.
4. ⚠️ begin_world/end_world try/finally — Still not present. Low risk (build-time only, failures abort entirely).

No new issues found. LGTM 👍

Update (7bb6bd6): Major architectural simplification — the visualization builder is now fully delegated to a new build_newton_visualization_builder_from_stage function in replicate.py, which reuses _build_newton_builder_from_mapping with apply_manager_hooks=False. This eliminates ~80 lines of duplicated logic from newton_manager.py.

Key changes:

• New replace_path_prefix utility in cloner_utils.py — clean extraction of path-boundary-aware prefix replacement, now shared by _rename_builder_labels.
• _build_newton_builder_from_mapping gains apply_manager_hooks flag to skip NewtonManager site injection and builder hooks for visualization-only builds. Clean conditional gating throughout.
• _rename_builder_labels simplified to use replace_path_prefix instead of inline boundary logic.

Previous concerns — status:

1. ✅ GPU clone_mask sync — Still addressed via .detach().cpu() in build_newton_visualization_builder_from_stage.
2. ✅ env_id vs sequential index — Correctly handled via clone_plan.env_ids tensor passed directly to _build_newton_builder_from_mapping.
3. ✅ Prefix matching — Now factored into replace_path_prefix with proper path-boundary semantics.
4. ✅ begin_world/end_world try/finally — No longer applicable; the visualization path delegates to _build_newton_builder_from_mapping which manages world lifecycle internally.

No new issues found. The consolidation is a clear improvement — single code path for both sim and visualization model building, differentiated only by the apply_manager_hooks and simplify_meshes flags. LGTM 👍

Update (271247f): Branch squashed/rebased with additional architectural improvements over 7bb6bd6f:

Key refactoring in replicate.py:

• Extracted _build_source_prototypes — builds prototype builders for each source path, now shared between sim and visualization paths.
• Extracted _replicate_builder_mapping — handles the per-env world replication loop with hooks, decoupled from site injection.
• build_newton_visualization_builder_from_stage now directly calls these helpers rather than _build_newton_builder_from_mapping, giving cleaner separation between sim (needs manager hooks, site injection) and visualization (just needs prototypes + replication).

Key changes in newton_manager.py:

• Added class-level registries (_deformable_registry, _per_world_builder_hooks, _post_replicate_hooks) as explicit lists, replacing the previous hasattr() pattern. Cleaner and more explicit.
• _build_visualization_model_from_stage method removed; logic moved inline into _ensure_visualization_model with direct delegation to build_newton_visualization_builder_from_stage.

Previous concerns — status:

1. ✅ GPU clone_mask sync — Still addressed via .detach().cpu().
2. ✅ env_id vs sequential index — Still correct via clone_plan.env_ids.
3. ✅ Prefix matching — replace_path_prefix utility from cloner_utils.py.
4. ✅ World lifecycle — Handled in _replicate_builder_mapping.

The decomposition into _build_source_prototypes and _replicate_builder_mapping improves testability and allows the visualization path to skip site injection cleanly. Test coverage remains solid with TestVisualizationClonePlan. No new issues. LGTM 👍

Update (4a4d54b): Final refactoring pass — the PR now cleanly decomposes into dedicated modules:

• builder_utils.py — extracted build_source_builders and replicate_builder_mapping (shared by sim and visualization paths)
• label_utils.py — moved _rename_builder_labels and _BUILTIN_LABEL_TYPES out of replicate.py
• visualization_builder.py — standalone build_visualization_builder_from_stage_envs function

Variable naming consistently updated from proto/protos to source_builder/source_builders throughout, including _cl_inject_sites docstrings. The hasattr() checks for deformable hooks replaced with explicit class-level list attributes. Test imports updated to match new module locations.

All previously raised concerns remain resolved. No new issues. LGTM 👍

Update (5679765): Module consolidation — builder_utils.py and label_utils.py merged into a single newton_clone_utils.py module under isaaclab_newton.cloner. replace_path_prefix moved to the shared isaaclab.cloner.cloner_utils module. All imports updated accordingly. visualization_builder.py now uses resolve_prim_pose from isaaclab.sim.utils.transforms instead of manual Xform extraction. Test imports updated. Logic unchanged from previous revision. No new issues. LGTM 👍

Update (0dd04d7): Branch rebased onto latest develop (incorporates upstream a4a08ac6 — safe lambda callable fix). No functional changes to this PR's code. All previous concerns remain resolved. LGTM 👍

Update (4b7cf73): Single commit since last review — appears to be a squash/rebase of the previous force-push. The code is structurally identical to the revision reviewed in the incremental review above. All previously raised concerns remain resolved. The only outstanding nit from the most recent review (return type annotation list[tuple[str, str, int]] → list[tuple[str, int]] in _rename_builder_labels) is minor and non-blocking. LGTM 👍

Update (4decd4b): New commit promotes internal helper to public API:

• _rename_builder_labels → rename_builder_labels (underscore prefix removed)
• Test file updated with matching function name
• Minor refactoring of deformable ignore path logic (explicit loop instead of nested comprehension, functionally equivalent)

This is a clean API exposure change. No functional differences, no new issues. LGTM 👍

[isaaclab-review-bot]

Performance (non-blocking): clone_mask[source_index, env_id].item() is called per (source, env) pair. If the mask is on GPU, this triggers a device→host sync each time. Consider converting to CPU upfront:

clone_mask = clone_plan.clone_mask
if hasattr(clone_mask, 'is_cuda') and clone_mask.is_cuda:
    clone_mask = clone_mask.cpu()

[isaaclab-review-bot]

Nit: This final return destination_root returns a path that may not correspond to a valid prim. It works because callers handle the None from _visualization_relative_transform, but a logger.debug(...) here documenting the exhausted-fallback case would help future debugging.

[nblauch]

Pushed an additive simplification commit (2b9b99b) instead of rewriting history. This folds the one-off ClonePlan helpers back into the ClonePlan assembly path while preserving the reusable transform/label helpers, moves clone_mask to CPU once, indexes clone_mask by sequential env position rather than env id, pairs begin_world/end_world with try/finally, and updates the regression test to use non-contiguous env ids.

[isaaclab-review-bot]

🔄 Incremental Review (New Commits)

Reviewing changes from 0dd04d7 → 77d12a6

This is a substantial improvement over the previous version. The refactoring extracts reusable utilities into newton_clone_utils.py and creates a dedicated visualization_builder.py module — much cleaner architecture.

✅ What Looks Good

1. Clean module extraction: Moving build_source_builders(), replicate_builder_mapping(), and _rename_builder_labels() into newton_clone_utils.py enables reuse between the live-sim replication path and the visualization builder.

2. New replace_path_prefix() utility: The boundary-safe path rewriting in cloner_utils.py is well-designed — it correctly handles edge cases like partial prefix matches (e.g., /Sources/protoA shouldn't match /Sources/protoAB).

3. Test coverage for heterogeneous ClonePlans: The new TestVisualizationClonePlan test class validates the core fix — ensuring visualization builds the correct source geometry for each world rather than always using env_0.

4. Proper changelog entries: Both the .skip file for the internal utility extraction and the .rst changelog for the actual fix are well-written.

5. Consistent naming: Renaming proto_* → source_* throughout improves clarity since these builders aren't true "prototypes" in the heterogeneous case.

📝 Suggestions

1. _rename_builder_labels return type mismatch (newton_clone_utils.py:120): The function signature declares list[tuple[str, str, int]] but actually returns list[tuple[str, int]]. This was present in the original code but should be fixed.

2. Consider adding docstring to replace_path_prefix: While the implementation is clear, a brief example in the docstring would help future readers:

   # >>> replace_path_prefix("/World/envs/env_0/Robot", "/World/envs/env_0", "/World/envs/env_1")
   # '/World/envs/env_1/Robot'

3. Hook list initialization: NewtonManager._deformable_registry, _per_world_builder_hooks, and _post_replicate_hooks are now always defined (good!), but the hasattr() guards were removed from replicate.py without ensuring other code paths initialize these. The class-level defaults handle this, but worth double-checking nothing else expects the old behavior.

🔍 Minor Observations

• The _FakeVisualizationModelBuilder test helper is comprehensive and should make future test maintenance easier.
• The visualization builder now correctly ignores both /World/envs and the source paths when loading stage-level geometry, preventing duplicate geometry.

---

Overall: This is a well-executed refactor that both fixes the original bug and improves code organization. The test coverage gives confidence in the heterogeneous ClonePlan handling. Nice work on the cleanup! 🎯

---

Update (commit 4decd4b): Reviewed incremental changes from 77d12a6 → 4decd4bf.

This commit completes the refactoring by extracting the visualization builder into its own module (visualization_builder.py) that properly uses the ClonePlan API. The old _build_visualization_model_from_stage method is removed entirely.

Previous inline comments status:

• ✅ env_id as column index bug — Fixed. The new visualization_builder.py uses clone_plan.env_ids passed through replicate_builder_mapping() which correctly iterates by sequential column index.
• ✅ begin_world/end_world not paired with try/finally — Addressed. The per-env loop is now handled by the shared replicate_builder_mapping() utility (same pattern as the live-sim path).
• ✅ GPU→CPU sync per item — Fixed. clone_plan.clone_mask.detach().cpu() is called upfront in visualization_builder.py.
• ✅ Fallback return destination_root logging — No longer applicable; the old code path is entirely removed.
• ⚠️ Return type annotation (rename_builder_labels declares list[tuple[str, str, int]] but returns list[tuple[str, int]]) — Still present in newton_clone_utils.py:120. Low priority but worth fixing.

No new issues found in the incremental diff. The architecture is clean and well-tested.

---

Update (9f5c8d0): Reviewed incremental changes from b1f67249 → 9f5c8d02.

✅ Return type annotation fixed — rename_builder_labels now correctly declares -> list[tuple[str, int]] instead of the incorrect 3-tuple annotation. This was the last remaining item from previous reviews.

The changes in this commit match the previously reviewed version. No new issues found.

---

Update (e18301a): Force-push squashed all prior commits into a single clean commit. All previously identified issues have been addressed in this final version (return type annotation fix ✅, architecture cleanup ✅). No new issues found. LGTM.

[isaaclab-review-bot]

Re-review for commit b1f6724

This is a follow-up review after new commits were pushed. The changes look solid overall.

Summary

This PR fixes Newton visualization ClonePlan handling for heterogeneous scenes by:

1. Extracting reusable utilities (newton_clone_utils.py) — build_source_builders, replicate_builder_mapping, rename_builder_labels, and replace_path_prefix are now shared between the cloner and visualization paths

2. Adding a dedicated visualization builder (visualization_builder.py) — build_visualization_builder_from_stage_envs() constructs the Newton shadow model using the active ClonePlan instead of always replicating from env_0

3. Improved test coverage — The new TestVisualizationClonePlan test class verifies that heterogeneous ClonePlans use the correct source builders per world and properly rewrite labels

Code Quality

✅ Good separation of concerns with the new utility module
✅ Clean refactoring of _rename_builder_labels → rename_builder_labels (public API)
✅ Consistent renaming from proto_* terminology to source_* for clarity
✅ The _BUILTIN_LABEL_TYPES constant is now properly exposed from the utils module
✅ Test coverage includes a fake builder that tracks geometry sources per world

Minor Observations

1. In newton_clone_utils.py line 75, the return type annotation shows list[tuple[str, str, int]] in the docstring but the actual return is list[tuple[str, int]] (2-tuple, not 3-tuple). This appears to be a copy-paste artifact from the old _rename_builder_labels function.

2. The replace_path_prefix utility in cloner_utils.py is a nice addition that handles the boundary-safe prefix replacement cleanly.

Overall this is a well-structured fix that properly handles heterogeneous ClonePlan sources for Newton visualization. The refactoring to share code between the cloner and visualization paths reduces duplication and makes the codebase more maintainable.

[ooctipus]

@greptile-apps review again

[isaaclab-review-bot]

🔄 Update: Reviewed commit 18b32ab7

Incremental review of changes from b1f67249 → 18b32ab7.

This commit squashes/rewrites the PR history into a single clean commit. All previously raised concerns have been addressed:

Issue	Status
Return type annotation (rename_builder_labels → list[tuple[str, int]])	✅ Fixed
GPU→CPU sync per item	✅ Fixed in visualization_builder.py
Clean module organization	✅ newton_clone_utils.py and visualization_builder.py properly extracted

No new issues found in this incremental diff. Ready to merge! 🚀

[greptile-apps]

Missing None guard crashes all non-ClonePlan scenes

PhysicsManager._sim.get_clone_plan() returns ClonePlan | None — it is None until the scene replicates, and is always None for scenes that do not use the cloner at all. build_visualization_builder_from_stage_envs immediately accesses clone_plan.sources with no guard, so every such call raises AttributeError. This removes the unconditional env_0-prototype path the PR description says is preserved.

Additionally, the old try/except Exception block that surrounded _build_visualization_model_from_stage is gone. The uncaught AttributeError now propagates through get_model() / get_state_0() / update_visualization_state() — the latter is called every render frame — rather than being logged and returning gracefully. The function needs a None guard with an env_0 fallback (or at minimum an early-return with a log) before calling build_visualization_builder_from_stage_envs.