Propagate Newton shape colors before cloning (isaac-sim#6194)

# Description

This PR integrates @ooctipus's proposed fix for newton shape color
replacement to occur pre-clone. Original PR:
rilei-nvidia#1

Replace the post-finalize `replace_newton_shape_colors` pass on the
Newton model with `replace_newton_builder_shape_colors`, applied to each
ModelBuilder before clone replication. This ensures all cloned
environments automatically inherit the correct USD material and
displayColor values. This not only improves performance but also
addresses the regression that @huidongc discovered last week where only
env_0 (the clone source) had the right color for the robot arm, and the
other cloned envs got the fallback colors

- Add replace_newton_builder_shape_colors to newton_model_utils.  
- Removed replace_newton_shape_colors and any un-used functionality
- Call it in NewtonManager (flat, env, and proto-source builders),
  newton_clone_utils, and replicate
- Update tests to exercise the new replace_newton_builder_shape_colors
- Update dexsuite_kuka_hetero golden images for correctness 

# Validation

Validated for correctness in the Newton visualizer with (all robot arms
should have the same color):
```bash
./isaaclab.sh train --rl_library rsl_rl \
    --task Isaac-Lift-KukaAllegro-Camera \
    presets=newton_mjwarp,newton_renderer,rgb64,single_camera \
    --seed 42 \
    --num_envs 4 \
    --max_iterations 2 \
    --visualizer newton
```
<img width="1917" height="1159" alt="image"
src="https://github.com/user-attachments/assets/fabf9ef4-60f6-4e29-a11a-e266e77d7f61"
/>


Rendering test golden images for the dexsuite_kuka_hetero kitless suite
have also been updated to reflect fix to the regression
```bash
./isaaclab.sh -p -m pytest -v -rA source/isaaclab_tasks/test/core/test_rendering_dexsuite_kuka_hetero_kitless.py::test_rendering_dexsuite_kuka_hetero_kitless[newton-newton_warp-rgb]
```
<img width="742" height="1046" alt="image"
src="https://github.com/user-attachments/assets/6e357b06-5cf2-458f-886b-4b283952c94f"
/>


# Performance

Testing on a AMD Ryzen Threadripper PRO 7965WX 24-Cores (48) @ 5.36 GHz

For testing the performance impact I used the benchmark script:
```bash
./isaaclab.sh -p scripts/benchmarks/benchmark_non_rl.py \
    --task Isaac-Lift-KukaAllegro-Camera \
    presets=newton_mjwarp,ovrtx_renderer,rgb64,single_camera \
    --seed 42 \
    --num_envs "8000" \
    --benchmark_backend json \
    --output_path "hdc/benchmarks/20260611/heterogeneous_8000"
```

The time of replace_newton_shape_colors has been moved to calls to the
new replace_newton_builder_shape_colors pre-clone:
| # Env | replace_newton_shape_colors (Before) |
replace_newton_builder_shape_colors (After) |
| ------ | ----- | ----- |
| 4000 | 2 s | 0.0045~ s|  
| 8000 | 4 s | 0.0045~ s|  

With @ooctipus's change the time for shape color propagation is not only
"free" now, but constant to increasing # of environments.

## Checklist

- [x] I have read and understood the [contribution
guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html)
- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

Author: rilei-nvidia

source/isaaclab/changelog.d/new-replace-newton-builder-colors.rst

@@ -0,0 +1,7 @@
+Added
+^^^^^
+
+* Added :func:`~isaaclab.sim.utils.newton_model_utils.replace_newton_builder_shape_colors` to
+  propagate USD material and ``displayColor`` values into a Newton ``ModelBuilder``'s shape colors
+  before clone replication, so cloned environments inherit correct colors without a separate
+  post-finalize pass.

source/isaaclab/isaaclab/sim/utils/newton_model_utils.py

@@ -13,7 +13,6 @@
 from __future__ import annotations
 
 import logging
-import os
 import warnings
 from typing import Any
 
@@ -22,7 +21,7 @@
 
 from pxr import Usd, UsdGeom, UsdShade
 
-__all__ = ["replace_newton_shape_colors"]
+__all__ = ["replace_newton_builder_shape_colors"]
 
 logger = logging.getLogger(__name__)
 
@@ -36,39 +35,15 @@
 _UNBOUND_DEFAULT_FALLBACK_GRAY = (0.18, 0.18, 0.18)
 
 
-@wp.func
-def _linear_channel_to_srgb_warp(c: float) -> float:
-    """Per-channel sRGB OETF on device: linear ``[0, 1]`` to sRGB-encoded ``[0, 1]``."""
+def _linear_channel_to_srgb(c: float) -> float:
+    """Per-channel sRGB OETF on host: linear ``[0, 1]`` to sRGB-encoded ``[0, 1]``."""
     if c <= 0.0:
         return 0.0
     if c <= 0.0031308:
         return 12.92 * c
     if c >= 1.0:
         return 1.0
-    return 1.055 * wp.pow(c, 1.0 / 2.4) - 0.055
-
-
-@wp.func
-def _linear_rgb_to_srgb_warp(linear_rgb: wp.vec3) -> wp.vec3:
-    """Apply sRGB OETF per channel: linear RGB ``[0, 1]`` to sRGB-encoded ``[0, 1]``."""
-    return wp.vec3(
-        _linear_channel_to_srgb_warp(linear_rgb[0]),
-        _linear_channel_to_srgb_warp(linear_rgb[1]),
-        _linear_channel_to_srgb_warp(linear_rgb[2]),
-    )
-
-
-@wp.kernel
-def _scatter_shape_color_rows_kernel(
-    shape_colors: wp.array(dtype=wp.vec3),  # type: ignore
-    row_indices: wp.array(dtype=wp.int32),  # type: ignore
-    row_colors: wp.array(dtype=wp.vec3),  # type: ignore
-):
-    """Write per-row sRGB colors into ``shape_colors``."""
-    tid = wp.tid()
-    index = row_indices[tid]
-    color = row_colors[tid]
-    shape_colors[index] = _linear_rgb_to_srgb_warp(color)
+    return 1.055 * (c ** (1.0 / 2.4)) - 0.055
 
 
 def _canonical_prim_lookup_key(prim: Usd.Prim) -> str:
@@ -203,7 +178,7 @@ def _resolve_shape_color(
     """Resolve replacement linear RGB for one prim path (sRGB encoding is applied in the scatter kernel).
 
     Returns:
-        Linear RGB to pass to :func:`_scatter_shape_color_rows_kernel`, or ``None`` to leave the row unchanged.
+        Linear RGB to pass, or ``None`` to leave the row unchanged.
     """
     shape_prim = stage.GetPrimAtPath(prim_path)
     if not shape_prim.IsValid():
@@ -231,122 +206,68 @@ def _resolve_shape_color(
     return material_color
 
 
-def replace_newton_shape_colors(model: Any, stage: Usd.Stage | None = None) -> int:
-    """Align Newton visualization colors with the USD stage.
+def replace_newton_builder_shape_colors(builder: Any, stage: Usd.Stage) -> int:
+    """Align a Newton ``ModelBuilder``'s shape colors with the USD stage before clone replication.
+
+    Overwrites entries in ``builder.shape_color`` so that colors match the authored USD data:
 
-    Newton assigns a per-shape palette to ``shape_color``. This overwrites those rows so rendering matches authored USD
-    data where supported:
+    - **No bound material**: use authored ``primvars:displayColor`` (treated as linear RGB), or a
+      neutral 18% linear gray if ``displayColor`` is not authored.
+    - **OmniPBR**: use ``diffuse_color_constant`` × ``diffuse_tint`` (linear RGB, with MDL defaults
+      when inputs are not authored).
+    - **Other materials**: leave the existing Newton color for that shape unchanged.
+    - **Guide purpose** prims (``UsdGeom.Tokens.guide``): leave unchanged so guide visualization
+      stays on the Newton palette.
 
-    - **No bound material**: use authored ``primvars:displayColor`` (treated as linear RGB), or a neutral 18% linear
-      gray if ``displayColor`` is not authored; linear values are encoded to sRGB in the scatter kernel.
-    - **OmniPBR**: use ``diffuse_color_constant`` times ``diffuse_tint`` (linear RGB, with MDL defaults when inputs are
-      not authored), encoded to sRGB in the scatter kernel.
-    - **Other materials**: leave the existing Newton color for that shape.
-    - **Guide purpose** prims (``UsdGeom.Tokens.guide``): leave unchanged so guide visualization stays on the palette.
+    Linear RGB values are encoded to sRGB before being written into ``builder.shape_color``.
 
     Args:
-        model: Object with ``shape_label`` (``list`` of USD prim paths) and ``shape_color`` (``wp.array`` of
-            ``wp.vec3``), typically a finalized Newton model.
-        stage: USD stage to read from. If ``None``, uses :func:`~isaaclab.sim.utils.stage.get_current_stage`.
+        builder: Object with ``shape_label`` (``list`` of USD prim paths) and ``shape_color``
+            (``list`` of ``wp.vec3``), typically a Newton ``ModelBuilder`` before finalization.
+        stage: USD stage to read material and primvar data from.
 
     Returns:
         Number of shapes that had their colors replaced.
-
-    Note:
-        Set ``ISAACLAB_REPLACE_NEWTON_SHAPE_COLORS`` to ``0``, ``false``, ``off``, or ``no`` to skip this pass
-        entirely (returns ``0``).
-
-        This pass exists only while Isaac Lab and Isaac Sim content still relies on NVIDIA-specific MDL and OmniPBR
-        materials; after migration to neutral USD materials that Newton can consume directly, this path is expected to
-        be deprecated and removed.
-
-        Wall time for USD resolution and the GPU scatter is measured with :class:`~isaaclab.utils.timer.Timer`, which
-        may print a timing summary when the timer is enabled.
     """
-    env_val = os.getenv("ISAACLAB_REPLACE_NEWTON_SHAPE_COLORS")
-    if env_val is not None and env_val.strip().lower() in ["false", "0", "off", "no"]:
-        logger.debug("Newton shape color replacement is disabled")
-        return 0
-
     warnings.warn(
         "Newton shape color replacement is enabled; this workaround will be deprecated in a future release.",
         FutureWarning,
         stacklevel=2,
     )
 
     # Use duck typing to avoid introducing hard dependencies on newton.
-    shape_labels = getattr(model, "shape_label", None)
-    shape_colors = getattr(model, "shape_color", None)
+    shape_labels = getattr(builder, "shape_label", None)
+    shape_colors = getattr(builder, "shape_color", None)
 
     if not isinstance(shape_labels, list):
         logger.debug("shape_label must be a list, got %s", type(shape_labels))
         return 0
 
-    if not isinstance(shape_colors, wp.array):
-        logger.debug("shape_color must be a Warp array, got %s", type(shape_colors))
-        return 0
-
-    num_shapes = len(shape_labels)
-    if num_shapes == 0:
-        logger.debug("Found empty list of shape labels")
+    if not isinstance(shape_colors, list):
+        logger.debug("shape_color must be a list, got %s", type(shape_colors))
         return 0
 
-    if num_shapes != len(shape_colors):
-        logger.debug("Mismatching length of shape_labels and shape_colors: %d != %d", num_shapes, len(shape_colors))
-        return 0
+    if len(shape_labels) != len(shape_colors):
+        raise ValueError(
+            f"Mismatching length of shape_label and shape_color: {len(shape_labels)} != {len(shape_colors)}"
+        )
 
     from isaaclab.utils.timer import Timer
 
-    num_color_updates = 0
-
-    with Timer(f"[INFO]: Time taken for replace_newton_shape_colors for {num_shapes} shapes", enable=False):
-        if stage is None:
-            from .stage import get_current_stage
-
-            stage = get_current_stage()
-
-        shape_keys: list[str] = []
-
-        for label in shape_labels:
-            prim = stage.GetPrimAtPath(label)
-            shape_keys.append(_canonical_prim_lookup_key(prim) if prim.IsValid() else label)
-
-        # shape_keys must stay the same length as shape labels, to guarantee the correctness of
-        # shape indices that will be used in the scatter kernel.
-        assert num_shapes == len(shape_keys)
-
-        resolved_color_cache: dict[str, tuple[float, float, float] | None] = {}
+    with Timer(
+        f"[INFO]: Time taken for replace_newton_builder_shape_colors for {len(shape_labels)} shapes", enable=False
+    ):
+        num_color_updates = 0
         material_color_cache: dict[str, tuple[float, float, float] | None] = {}
-
-        unique_keys = dict.fromkeys(shape_keys)
-        for key in unique_keys:
-            color = _resolve_shape_color(stage, key, material_color_cache)
-            resolved_color_cache[key] = color
-
-        # Prepare the indices and colors for the scatter kernel:
-        # - Indices point to the slots in the shape_colors array that should be updated
-        # - Colors are the new values to write into the slots
-        indices_np = np.empty(num_shapes, dtype=np.int32)
-        colors_np = np.empty((num_shapes, 3), dtype=np.float32)
-
-        for i, shape_key in enumerate(shape_keys):
-            if rgb := resolved_color_cache.get(shape_key):
-                indices_np[num_color_updates] = i
-                colors_np[num_color_updates] = rgb
+        for i, label in enumerate(shape_labels):
+            rgb = _resolve_shape_color(stage, label, material_color_cache)
+            if rgb is not None:
+                shape_colors[i] = wp.vec3(
+                    _linear_channel_to_srgb(rgb[0]),
+                    _linear_channel_to_srgb(rgb[1]),
+                    _linear_channel_to_srgb(rgb[2]),
+                )
                 num_color_updates += 1
 
-        # If there are any color updates, launch the scatter kernel to update the shape_colors array.
-        if num_color_updates != 0:
-            indices_wp = wp.from_numpy(indices_np[:num_color_updates], dtype=wp.int32, device=shape_colors.device)
-            colors_wp = wp.from_numpy(colors_np[:num_color_updates], dtype=wp.vec3, device=shape_colors.device)
-
-            wp.launch(
-                kernel=_scatter_shape_color_rows_kernel,
-                dim=num_color_updates,
-                inputs=[shape_colors, indices_wp, colors_wp],
-                device=shape_colors.device,
-            )
-
-        logger.debug("Replaced colors for %d / %d shapes", num_color_updates, num_shapes)
-
-    return num_color_updates
+        logger.debug("Replaced builder colors for %d / %d shapes", num_color_updates, len(shape_labels))
+        return num_color_updates

source/isaaclab/test/sim/test_newton_manager_visualization_state.py

@@ -100,7 +100,6 @@ def test_ensure_visualization_model_builds_from_stage_when_backend_is_physx(monk
     monkeypatch.setattr(nm.PhysicsManager, "_sim", None, raising=False)
     _set_sim_context(monkeypatch, nm)
     monkeypatch.setattr(nm.PhysicsManager, "_device", "cpu", raising=False)
-    monkeypatch.setattr(nm, "replace_newton_shape_colors", lambda model, *a, **kw: 0)
 
     finalize_calls: list[str] = []
 
@@ -155,7 +154,6 @@ def test_ensure_visualization_model_populates_num_envs_when_backend_is_physx(mon
     monkeypatch.setattr(nm.PhysicsManager, "_sim", None, raising=False)
     _set_sim_context(monkeypatch, nm)
     monkeypatch.setattr(nm.PhysicsManager, "_device", "cpu", raising=False)
-    monkeypatch.setattr(nm, "replace_newton_shape_colors", lambda model, *a, **kw: 0)
 
     class _FakeBuilder:
         body_count = 3