Add torch support to collection reshape helper

[AntoineRichard]

Description

Adds torch.Tensor input support to RigidObjectCollection.reshape_data_to_view_3d() for the PhysX, Newton, and OVPhysX backends. Torch inputs now reshape from instance-major (num_instances, num_bodies, data_dim) to body-major (num_bodies * num_instances, data_dim), preserve dtype, return a contiguous torch tensor, and honor the requested output device.

The existing Warp array paths remain unchanged after the torch type check where the helper already existed. Newton now exposes the same helper for backend API consistency.

Fixes #5893

Type of change

• New feature (non-breaking change which adds functionality)

Screenshots

N/A

Test Plan

• PYTHONPATH=/home/antoiner/Documents/IsaacLab_Newton/.worktrees/fix-issue-5893/source/isaaclab_newton:/home/antoiner/Documents/IsaacLab_Newton/.worktrees/fix-issue-5893/source/isaaclab_ovphysx:/home/antoiner/Documents/IsaacLab_Newton/.worktrees/fix-issue-5893/source/isaaclab_physx:/home/antoiner/Documents/IsaacLab_Newton/.worktrees/fix-issue-5893/source/isaaclab:$PYTHONPATH ./isaaclab.sh -p -m pytest source/isaaclab/test/assets/test_rigid_object_collection_iface.py::TestCollectionViewReshape -q
  • Result: 21 passed
• git diff --check
  • Result: passed
• ./isaaclab.sh --format
  • Result: ruff, format, whitespace, spelling, RST, and other non-LFS hooks passed; check-git-lfs-pointers failed because git-lfs is not installed in this local environment.

Checklist

• I have read and understood the contribution guidelines
• I have run the pre-commit checks with ./isaaclab.sh --format (ran the exact command; blocked only by missing local git-lfs)
• I have made corresponding changes to the documentation (N/A: no user-facing docs page change needed for this helper-only API update)
• My changes generate no new warnings (targeted tests emit existing deprecation warnings unrelated to this change)
• 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

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR adds torch.Tensor input support to RigidObjectCollection.reshape_data_to_view_3d() in both PhysX and OVPhysX backends, fixing #5893. The implementation dispatches on isinstance(data, torch.Tensor), performs the instance-major → body-major transpose, and honors the target device. Tests cover dtype preservation, device transfer, and Warp/torch output equivalence.

Design Assessment

Architecture: The approach is sound — type dispatch via isinstance before the existing Warp path preserves backward compatibility cleanly. Implementing in both backend classes mirrors the existing pattern.

Reshape correctness: The reshape logic data.transpose(0, 1).reshape(N*M, D).to(device).contiguous() correctly converts from (num_instances, num_bodies, data_dim) to (num_bodies * num_instances, data_dim) in body-major order.

API consistency: The type annotations (wp.array | torch.Tensor) and return type annotations are correctly updated. The docstrings are informative.

Findings

🟡 Minor Issues

1. .to(device).contiguous() ordering — source/isaaclab_physx/.../rigid_object_collection.py:L1147 and source/isaaclab_ovphysx/.../rigid_object_collection.py:L1234

   After .reshape() on a transposed tensor, the tensor is likely non-contiguous. Calling .to(device) on a non-contiguous tensor may trigger an implicit copy anyway, but calling .contiguous() after .to() means the contiguity enforcement happens on the target device. This is actually fine for correctness (and may even be preferable for avoiding a double copy), but it is worth noting that if device differs from the source, there could be a brief non-contiguous tensor on the target device before the final .contiguous() call. The current order is acceptable.

2. PhysX device default changed semantics — source/isaaclab_physx/.../rigid_object_collection.py:L1122

   The PhysX signature changes from device: str = "cpu" to device: str | None = "cpu". This correctly allows None for torch inputs (defaulting to data.device), but for the Warp path further down, device is still used as a string directly (no None check before passing to wp.clone). This is safe only because the Warp path is reached exclusively when data is NOT a torch tensor, and in that case device defaults to "cpu" (the callers would never pass None explicitly for Warp data). However, if a user explicitly passes device=None with a Warp array, the code would pass None to wp.clone, which may raise. Consider adding a guard:

   if device is None:
       device = str(data.device)

   at the top of the Warp path (similar to what OVPhysX already does).

3. Duplicate logic across backends — Both PhysX and OVPhysX implement the identical torch path. This is a minor DRY concern but consistent with the existing codebase pattern where each backend has its own implementation. Not blocking.

🟢 No Issues

• Type dispatch (isinstance(data, torch.Tensor)) is correct and robust.
• Device placement is honored for torch tensors.
• Dtype preservation is verified by tests.
• The Warp code path is untouched and unaffected.
• Changelog fragments follow project conventions.

Test Coverage

Strong. The PR adds 14 test cases (parametrized across backends, devices, and dtypes):

• ✅ test_reshape_data_to_view_3d_accepts_torch_tensor — validates shape, dtype, device, contiguity, and correctness of the reorder against permute(1,0,2).reshape()
• ✅ test_reshape_data_to_view_3d_moves_torch_tensor_to_requested_device — cross-device (CUDA→CPU) transfer
• ✅ test_reshape_data_to_view_3d_keeps_warp_array_behavior — regression test ensuring Warp path unchanged and outputs match torch

Suggestions:

• Consider adding a test that explicitly passes device=None with a torch tensor to verify the default-device fallback path.
• A test with a 1D data_dim=1 edge case would strengthen coverage.

Verdict

Minor fixes needed — The core implementation is correct and well-tested. The only actionable concern is the potential device=None passthrough to wp.clone in the PhysX backend if a user explicitly sets device=None with a Warp array (unlikely but possible). A one-line guard would close this gap.

---

Update (139d07d): New commit adds the Newton backend implementation of reshape_data_to_view_3d with torch support, a changelog fragment, and expands test parametrization to include the Newton backend.

• ✅ The Newton implementation correctly includes the device=None guard for the Warp path (if device is None: device = str(data.device)).
• The PhysX backend still lacks this guard on its Warp path (default "cpu" mitigates practical risk, but an explicit None could still propagate). Not re-posting since it's minor and unlikely to be hit in practice.
• Overall: LGTM. Clean implementation consistent across all three backends.

[greptile-apps]

Greptile Summary

This PR extends RigidObjectCollection.reshape_data_to_view_3d() in both the PhysX and OVPhysX backends to accept torch.Tensor inputs in addition to the existing Warp array path. The new torch branch transposes axes 0 and 1 (instance → body major), reshapes to (num_bodies * num_instances, data_dim), honours the requested output device, and returns a contiguous tensor.

• Core logic: data.transpose(0, 1).reshape(...).to(device).contiguous() correctly produces body-major layout matching the Warp strided-view path; verified by the accompanying test suite (14 tests covering dtype preservation, device movement, and Warp regression).
• Type annotation gap: In both backends the device parameter is typed str | None, but the torch branch assigns data.device (a torch.device object) to it when device=None; this diverges from the declared type and will be flagged by static analysers.
• Changelog and skip files: Correctly added for both isaaclab_physx and isaaclab_ovphysx; the isaaclab package correctly uses a .skip file since only test code changed there.

Confidence Score: 4/5

Safe to merge; the torch reshape logic is correct and well-tested, with only a minor type annotation inaccuracy.

The reshape arithmetic (transpose(0,1).reshape(B*N, D)) produces the correct body-major layout and is validated against the Warp path in tests. The one concrete inaccuracy is that device = data.device assigns a torch.device object to a variable declared str | None, which will surface as a type-checker warning in both backends but has no runtime impact.

Both backend rigid_object_collection.py files share the same device type annotation gap in the new torch branch.

Important Files Changed

Filename	Overview
source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py	Adds torch.Tensor branch to reshape_data_to_view_3d; logic is correct but device = data.device assigns a torch.device to a str
source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py	Identical torch branch added; same type annotation inaccuracy for device parameter.
source/isaaclab/test/assets/test_rigid_object_collection_iface.py	New TestCollectionViewReshape class covers torch dtype preservation, device movement, and Warp path regression; tests are well-structured and thorough.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["reshape_data_to_view_3d(data, data_dim, device)"] --> B{isinstance data torch.Tensor?}
    B -- Yes --> C{device is None?}
    C -- Yes --> D["device = data.device"]
    C -- No --> E["use supplied device"]
    D --> F["data.transpose(0,1)\n.reshape(num_bodies*num_instances, data_dim)\n.to(device).contiguous()"]
    E --> F
    F --> G["return torch.Tensor\n(body-major, contiguous)"]
    B -- No --> H{device is None?}
    H -- Yes --> I["device = str(data.device)"]
    H -- No --> J["use supplied device (str)"]
    I --> K["wp.array strided view\n(B, N, data_dim)"]
    J --> K
    K --> L["wp.clone → reshape\n(num_bodies*num_instances, data_dim)"]
    L --> M["return wp.array\n(body-major, contiguous)"]

Loading

_{Reviews (1): Last reviewed commit: "Add torch support to collection reshape ..." | Re-trigger Greptile}

[greptile-apps]

The device parameter is annotated as str | None, but when device=None and a torch.Tensor is passed, the local variable is immediately reassigned to data.device, which is a torch.device object — not a str. Static type checkers (mypy/pyright) will flag the re-assignment as a type error. Using str(data.device) keeps the annotation correct and is consistent with how the warp path normalises the device string below.

Suggested change

	if isinstance(data, torch.Tensor):
	if device is None:
	device = data.device
	return data.transpose(0, 1).reshape(self.num_bodies * self.num_instances, data_dim).to(device).contiguous()
	if isinstance(data, torch.Tensor):
	if device is None:
	device = str(data.device)
	return data.transpose(0, 1).reshape(self.num_bodies * self.num_instances, data_dim).to(device).contiguous()

[kellyguo11]

seems like a valid point. there's some inconsistency between the check on line 1235 and 1239

[isaaclab-review-bot]

Thanks for confirming! The author addressed this in the Newton backend commit — adding a device is not None guard before the Warp wp.clone(..., device=device) call. The OVPhysX backend here could benefit from the same treatment for consistency.

[isaaclab-review-bot]

Thanks for confirming! The fix is straightforward—changing line 1235 from:

device = data.device

to:

device = str(data.device)

This keeps the type annotation (str | None) accurate and maintains consistency with the Warp path below. The .to(device) call on torch tensors accepts both str and torch.device, so this change is safe.

[greptile-apps]

Same type annotation issue as in the ovphysx backend: device is typed as str | None but is reassigned to data.device (torch.device) when it is None and the input is a torch tensor. Use str(data.device) to stay within the declared type.

Suggested change

	if isinstance(data, torch.Tensor):
	if device is None:
	device = data.device
	return data.transpose(0, 1).reshape(self.num_bodies * self.num_instances, data_dim).to(device).contiguous()
	if isinstance(data, torch.Tensor):
	if device is None:
	device = str(data.device)
	return data.transpose(0, 1).reshape(self.num_bodies * self.num_instances, data_dim).to(device).contiguous()