Antoiner/feat/torch array

[AntoineRichard]

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context.
List any dependencies that are required for this change.

Fixes # (issue)

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (existing functionality will not work without user modification)
• Documentation update

Screenshots

Please attach before and after screenshots of the change if applicable.

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 updated the changelog and the corresponding version in the extension's config/extension.toml file
• I have added my name to the CONTRIBUTORS.md or my name already exists there

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR introduces TorchArray, a warp-first dual-access array wrapper that replaces raw wp.array returns from all data properties across the framework. The design provides explicit .torch and .warp accessors with a deprecation bridge for backward compatibility. The migration is comprehensive — 143 files touching articulation, rigid object, sensor, MDP, and task environment code across both PhysX and Newton backends.

The core TorchArray class is clean and well-tested. The Newton pinning optimization is clever but has a correctness gap in the rebind path for sliced properties. The PhysX backend takes a simpler (allocation-per-access) approach that sidesteps this class of bug entirely.

Design Assessment

Design is sound. The warp-first wrapper with cached zero-copy torch views is the right approach for this codebase. The deprecation bridge via __torch_function__ + operator overloads provides a clean migration path. The Newton pinning optimization is well-motivated (avoiding per-access allocation for stable pointers), and the FrameTransformer revert shows good judgment about when pinning isn't safe.

Two design asymmetries worth noting:

• PhysX vs Newton allocation strategy: PhysX creates a new TorchArray per property access (~80 allocations/step), while Newton pins them. Both are valid but the inconsistency may confuse contributors. Consider documenting this difference.
• _deprecation_warned is a class-level singleton: Once any TorchArray triggers the warning, all instances go silent. This is reasonable for one-time migration warnings but could mask issues if different subsystems migrate at different rates.

Findings

🔴 Critical: Newton sliced wp.array buffers not invalidated on sim rebind — source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py:1433-1450

When _create_simulation_bindings() is called after a sim reset (via the PHYSICS_READY callback), the code correctly invalidates the pinned TorchArray wrappers (_root_link_pos_w_ta = None, etc.) but does not invalidate the intermediate sliced wp.array buffers (_root_link_pos_w, _root_link_quat_w, etc.).

For contiguous transforms, _get_pos_from_transform() creates a strided wp.array view via wp.array(ptr=transform.ptr, ...). After a sim reset:

1. The transform sim binding is rebound to a new pointer ✅
2. The TorchArray wrapper is invalidated ✅
3. But self._root_link_pos_w still holds the old strided wp.array with ptr pointing to the old (freed) transform memory ❌

On next property access, _get_pos_from_transform(self._root_link_pos_w, self.root_link_pose_w.warp) sees source is not None and transform.is_contiguous, so it returns the stale source unchanged — serving data from freed memory.

Fix: Also reset the underlying wp.array sliced buffers alongside the TorchArray wrappers:

# Invalidate lazy sliced TorchArrays AND their backing wp.arrays
# so they are re-created from the new sim bindings on next access.
self._root_link_pos_w_ta = None
self._root_link_pos_w = None
self._root_link_quat_w_ta = None
self._root_link_quat_w = None
self._root_link_lin_vel_w_ta = None
self._root_link_lin_vel_w = None
self._root_link_ang_vel_w_ta = None
self._root_link_ang_vel_w = None
# ... same pattern for root_com_*, body_link_*, body_com_*

This same pattern should be verified in RigidObjectData and RigidObjectCollectionData if they also have sim rebind callbacks.

🟡 Warning: No __getitem__ on TorchArray — direct indexing will fail — source/isaaclab/isaaclab/utils/warp/torch_array.py

TorchArray does not implement __getitem__, so ta[0], ta[:, 2], etc. will raise TypeError: 'TorchArray' object is not subscriptable. While the consumer code in this PR consistently uses .torch before indexing, this is a footgun for anyone writing new code against the TorchArray API — especially since the old wp.array supported indexing.

The __torch_function__ protocol only intercepts calls to torch.* functions, not Python's [] operator. Consider adding:

def __getitem__(self, key):
    self._warn_implicit()
    return self.torch[key]

This would maintain the deprecation bridge's backward compatibility promise more completely.

🟡 Warning: PhysX root_com_vel_w reassigns .data but TorchArray wraps it per-access — source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py:613

self._root_com_vel_w.data = self._root_view.get_root_velocities().view(wp.spatial_vectorf)

Some PhysX properties replace the TimestampedBuffer.data wp.array reference entirely each step. Since PhysX creates a new TorchArray per access, this works correctly. But if someone later applies the Newton pinning optimization to PhysX without understanding this pattern, pinned TorchArrays would serve stale data. A brief comment noting why PhysX uses per-access TorchArray allocation (and why pinning would be unsafe here) would prevent future bugs.

🔵 Suggestion: Add __init__ type guard for wp_array parameter — source/isaaclab/isaaclab/utils/warp/torch_array.py:54-61

def __init__(self, wp_array: wp.array) -> None:
    if not isinstance(wp_array, wp.array):
        raise TypeError(f"TorchArray expects a wp.array, got {type(wp_array).__name__}")
    self._warp = wp_array
    self._torch_cache: torch.Tensor | None = None

This would catch accidental wrapping of non-warp objects (e.g., passing a torch.Tensor or None) at construction time rather than producing confusing errors on .torch or .warp access.

🔵 Suggestion: Missing rebind() test — source/isaaclab/test/utils/warp/test_torch_array.py

The test suite thoroughly covers basic functionality, structured types, convenience properties, and the deprecation bridge — good coverage. However, rebind() is a critical method (it's the mechanism that prevents stale data after sim resets) and has no direct test. Consider adding:

def test_rebind_invalidates_cache(self, device):
    from isaaclab.utils.warp.torch_array import TorchArray
    arr1 = wp.zeros(10, dtype=wp.float32, device=device)
    ta = TorchArray(arr1)
    t1 = ta.torch  # Cache the tensor
    arr2 = wp.ones(10, dtype=wp.float32, device=device)
    ta.rebind(arr2)
    t2 = ta.torch
    assert t1 is not t2  # Cache was invalidated
    assert ta.warp is arr2  # Warp points to new array
    assert t2[0].item() == 1.0  # New data is served

Test Coverage

• Core TorchArray class: Well tested — basic ops, structured types, deprecation bridge, operator forwarding. ✅
• Missing: rebind() method test, __getitem__ behavior (would currently fail), behavior when wrapping a 0-dimensional array.
• Integration: Existing PhysX/Newton articulation/sensor tests updated to use .torch accessor. The test updates appear mechanical and correct.
• Coverage adequate: Yes (for a feature PR), with the rebind() gap noted above.

CI Status

Only the labeler check has run (pass). No test or lint CI results available yet — recommend waiting for full CI before merge.

Verdict

Significant concerns

The core TorchArray design is solid and the migration is comprehensive, but the Newton rebind bug (stale sliced wp.array pointers after sim reset) is a correctness issue that could cause silent data corruption in RL training loops. This should be addressed before merge. The other findings are lower priority but would improve robustness.

[isaaclab-review-bot]

🔴 Critical: Sliced wp.array buffers not invalidated on rebind

This block invalidates the TorchArray wrappers (_root_link_pos_w_ta = None) but not the underlying sliced wp.array buffers (_root_link_pos_w). After sim reset, contiguous strided views created by _get_pos_from_transform() still hold ptr into the old (freed) transform memory. Since source is not None, the helper returns the stale buffer without re-creating it from the new transform.

Fix: Also set the wp.array buffers to None:

Suggested change

	self._root_link_pos_w_ta = None
	# Invalidate lazy sliced TorchArrays AND their backing wp.arrays
	# so they are re-created from the new sim bindings on next access.
	self._root_link_pos_w_ta = None
	self._root_link_pos_w = None

(Same pattern for all 16 sliced properties: _root_link_quat_w, _root_link_lin_vel_w, _root_link_ang_vel_w, _root_com_*, _body_link_*, _body_com_*)

[isaaclab-review-bot]

🟡 Warning: No __getitem__ support

TorchArray doesn't implement __getitem__, so ta[0] or ta[:, 2] will raise TypeError. The __torch_function__ protocol only intercepts torch.* function calls, not Python's [] operator.

Consider adding a subscript forwarder for the deprecation bridge:

def __getitem__(self, key):
    self._warn_implicit()
    return self.torch[key]

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot — Follow-up

New commits (f0f4c82→15d8fe2a): adds missing invalidations in articulation_data.py and registers a PHYSICS_READY callback in rigid_object_collection.py. The articulation_data fix looks complete — all lazy-sliced buffers are now invalidated. One new issue found in the collection callback registration.

[isaaclab-review-bot]

🟡 Warning: _physics_ready_handle not deregistered on cleanup

This registers a PHYSICS_READY callback but the class does not override _clear_callbacks() to deregister it. The base AssetBase._clear_callbacks() only cleans up _initialize_handle, _invalidate_initialize_handle, _prim_deletion_handle, and _debug_vis_handle — it has no knowledge of this new handle.

The sibling class Articulation (articulation.py L3223-3228) correctly overrides _clear_callbacks() for exactly this reason:

def _clear_callbacks(self) -> None:
    super()._clear_callbacks()
    if hasattr(self, "_physics_ready_handle") and self._physics_ready_handle is not None:
        self._physics_ready_handle.deregister()
        self._physics_ready_handle = None

Without this, when the rigid object collection is destroyed (e.g., scene teardown), the callback will persist and fire into a dead self._data reference, likely causing a crash or silent corruption on the next physics reset.

[StafaH]

Might be interesting to also add the automatic unwarp for warp kernels as well? Something like below. Might needs some solid tests (cuda vs cpu kernels, make sure they get passed in, no functional or performance difference between passing array vs array.warp.

   # ------------------------------------------------------------------
    # Warp kernel interop
    # ------------------------------------------------------------------
    #
    # Expose the array-interface protocols from the underlying warp array so that
    # a TorchArray can be passed directly as an input to :func:`warp.launch`.
    # Warp's ``pack_arg`` routine looks for these attributes when the runtime
    # type does not match the expected :class:`warp.array` and builds an
    # ``array_t`` descriptor from the interface (keeping ``self`` alive via its
    # ``owner`` reference). Delegating to ``self._warp`` reuses the cached
    # interface dict that warp already maintains per array, so there is no extra
    # allocation on each launch.
    #
    # AttributeError is propagated (not caught) when the device doesn't match
    # the protocol — warp's ``pack_arg`` relies on the ``AttributeError`` to
    # fall through to the next candidate protocol.

    @property
    def __cuda_array_interface__(self):
        """Delegate the CUDA array interface to the underlying warp array.

        Raises:
            AttributeError: If the underlying warp array is not on a CUDA device.
        """
        return self._warp.__cuda_array_interface__

    @property
    def __array_interface__(self):
        """Delegate the NumPy array interface to the underlying warp array.

        Raises:
            AttributeError: If the underlying warp array is not on a CPU device.
        """
        return self._warp.__array_interface__

[hougantc-nvda]

Profiling the physx side, it doesn't appear the caching in working as intended.
All these TorchArray constructors are causing wp.to_torch to be called on every call to these functions since init sets self._torch_cache: to None
Including the overhead of TorchArray it's actually a slight bit slower.