Fix OVRTX renderer device mismatch on multi-GPU

CONTRIBUTORS.md

@@ -80,6 +80,7 @@ Guidelines for modifications:
 * Emily Sturman
 * Emmanuel Ferdman
 * Fabian Jenelten
+* Fatima Anes
 * Felipe Mohr
 * Felix Yu
 * Frank Lai

source/isaaclab_ov/changelog.d/fix-ovrtx-device-mismatch.rst

@@ -0,0 +1,7 @@
+Fixed
+^^^^^
+
+* Fixed :class:`OVRTXRenderer` crash on multi-GPU systems when ``sim.device``
+  is not ``cuda:0``. All Warp kernel launches, buffer allocations, and OVRTX
+  ``binding.map()`` calls now use the device from :class:`CameraRenderSpec`
+  instead of hardcoded defaults.

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer.py

@@ -12,7 +12,7 @@
   updates camera/object transforms (using kernels), steps the renderer, then extracts
   tiles from the tiled framebuffer (kernels).
 
-- **ovrtx_renderer_kernels.py**: Warp GPU kernels and DEVICE constant.
+- **ovrtx_renderer_kernels.py**: Warp GPU kernels for OVRTX rendering pipeline.
 
 - **ovrtx_usd.py**: USD helpers for OVRTX: render var config, camera injection, etc.
 """
@@ -47,7 +47,6 @@
 
 from .ovrtx_renderer_cfg import OVRTXRendererCfg
 from .ovrtx_renderer_kernels import (
-    DEVICE,
     create_camera_transforms_kernel,
     extract_all_depth_tiles_kernel,
     extract_all_depth_tiles_kernel_legacy,
@@ -147,8 +146,15 @@ def supported_output_types(self) -> dict[RenderBufferKind, RenderBufferSpec]:
             RenderBufferKind.DISTANCE_TO_CAMERA: RenderBufferSpec(1, torch.float32),
         }
 
+    @property
+    def _device_id(self) -> int:
+        """CUDA device index extracted from ``self._device`` for OVRTX ``binding.map()`` calls."""
+        parts = self._device.split(":")
+        return int(parts[1]) if len(parts) > 1 else 0
+
     def __init__(self, cfg: OVRTXRendererCfg):
         self.cfg = cfg
+        self._device = "cuda:0"  # default; overridden by create_render_data(spec)
         self._usd_handles = []
         self._render_product_paths = []
         self._camera_binding = None
@@ -355,7 +361,7 @@ def _setup_object_bindings(self):
 
             if self._object_binding is not None:
                 logger.info("Object binding created successfully")
-                self._object_newton_indices = wp.array(newton_indices, dtype=wp.int32, device=DEVICE)
+                self._object_newton_indices = wp.array(newton_indices, dtype=wp.int32, device=self._device)
             else:
                 logger.warning("Object binding is None")
         except ImportError:
@@ -369,9 +375,10 @@ def create_render_data(self, spec: CameraRenderSpec) -> OVRTXRenderData:
         Performs OVRTX initialization (stage export, USD load, bindings) on first call,
         matching the interface of Isaac RTX and Newton Warp which need no separate initialize().
         """
+        self._device = spec.device
         if not self._initialized_scene:
             self._initialize_from_spec(spec)
-        return OVRTXRenderData(spec, DEVICE)
+        return OVRTXRenderData(spec, self._device)
 
     # Map torch dtypes to their warp counterparts for zero-copy wrapping.
     _TORCH_TO_WP_DTYPE: dict[torch.dtype, Any] = {
@@ -427,13 +434,13 @@ def update_transforms(self) -> None:
             if body_q is None:
                 return
 
-            with self._object_binding.map(device=Device.CUDA, device_id=0) as attr_mapping:
+            with self._object_binding.map(device=Device.CUDA, device_id=self._device_id) as attr_mapping:
                 ovrtx_transforms = wp.from_dlpack(attr_mapping.tensor, dtype=wp.mat44d)
                 wp.launch(
                     kernel=sync_newton_transforms_kernel,
                     dim=len(self._object_newton_indices),
                     inputs=[ovrtx_transforms, self._object_newton_indices, body_q],
-                    device=DEVICE,
+                    device=self._device,
                 )
         except Exception as e:
             logger.warning("Failed to update object transforms: %s", e)
@@ -450,15 +457,15 @@ def update_camera(
         camera_quats_opengl = convert_camera_frame_orientation_convention(orientations, origin="world", target="opengl")
         camera_positions_wp = wp.from_torch(positions.contiguous(), dtype=wp.vec3)
         camera_orientations_wp = wp.from_torch(camera_quats_opengl.contiguous(), dtype=wp.quatf)
-        camera_transforms = wp.zeros(num_envs, dtype=wp.mat44d, device=DEVICE)
+        camera_transforms = wp.zeros(num_envs, dtype=wp.mat44d, device=self._device)
         wp.launch(
             kernel=create_camera_transforms_kernel,
             dim=num_envs,
             inputs=[camera_positions_wp, camera_orientations_wp, camera_transforms],
-            device=DEVICE,
+            device=self._device,
         )
         if self._camera_binding is not None:
-            with self._camera_binding.map(device=Device.CUDA, device_id=0) as attr_mapping:
+            with self._camera_binding.map(device=Device.CUDA, device_id=self._device_id) as attr_mapping:
                 wp_transforms_view = wp.from_dlpack(attr_mapping.tensor, dtype=wp.mat44d)
                 wp.copy(wp_transforms_view, camera_transforms)
 
@@ -480,7 +487,7 @@ def read_output(
     def _generate_random_colors_from_ids(self, input_ids: wp.array) -> wp.array:
         """Generate pseudo-random colors from semantic IDs."""
         if self._output_semantic_color_buffer is None or self._output_semantic_color_buffer.shape != input_ids.shape:
-            self._output_semantic_color_buffer = wp.zeros(shape=input_ids.shape, dtype=wp.uint32, device=DEVICE)
+            self._output_semantic_color_buffer = wp.zeros(shape=input_ids.shape, dtype=wp.uint32, device=self._device)
 
         output_colors = self._output_semantic_color_buffer
 
@@ -492,7 +499,7 @@ def _generate_random_colors_from_ids(self, input_ids: wp.array) -> wp.array:
             ),
             dim=input_ids.shape,
             inputs=[input_ids, output_colors],
-            device=DEVICE,
+            device=self._device,
         )
 
         return output_colors
@@ -522,7 +529,7 @@ def _extract_rgba_tiles(
                 render_data.height,
                 num_channels,
             ],
-            device=DEVICE,
+            device=self._device,
         )
 
     def _extract_depth_tiles(
@@ -543,7 +550,7 @@ def _extract_depth_tiles(
                         render_data.width,
                         render_data.height,
                     ],
-                    device=DEVICE,
+                    device=self._device,
                 )
 
     def _process_render_frame(self, render_data: OVRTXRenderData, frame, output_buffers: dict) -> None:

source/isaaclab_ov/isaaclab_ov/renderers/ovrtx_renderer_kernels.py

@@ -3,12 +3,10 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""Warp kernels and device constant for OVRTX renderer."""
+"""Warp kernels for OVRTX rendering pipeline."""
 
 import warp as wp
 
-DEVICE = "cuda:0"
-
 
 @wp.kernel
 def create_camera_transforms_kernel(

source/isaaclab_ov/test/test_ovrtx_renderer_kernels.py

@@ -11,14 +11,15 @@
 import pytest
 import warp as wp
 from isaaclab_ov.renderers.ovrtx_renderer_kernels import (
-    DEVICE,
     extract_all_depth_tiles_kernel,
     extract_all_depth_tiles_kernel_legacy,
     extract_all_rgba_tiles_kernel,
     generate_random_colors_from_ids_kernel,
     generate_random_colors_from_ids_kernel_legacy,
 )
 
+DEVICE = "cuda:0"
+
 
 def _color_hash(seed: int) -> int:
     h = seed