Capture Newton scene-query tasks into per-task CUDA graphs

Every registered scene-query task ran inside a wp.capture_if conditional body.
CUDA forbids memory-allocation nodes there, and wp.Mesh.refit() allocates scratch
on every call, so tasks with deformable geometry failed to capture and all sensor
work fell back to eager execution.

Capture the shared BVH refit and each task into their own top-level graph, where
allocation nodes are legal. This preserves the previous per-task selection
semantics and drops the device flag array and its per-call host-to-device copy.

Author: mmichelis

source/isaaclab_newton/changelog.d/fix-sensor-graph-allocating-task.rst

@@ -0,0 +1,7 @@
+Fixed
+^^^^^
+
+* Fixed Newton sensor CUDA graph capture failing with ``RuntimeError: Conditional body graph contains an
+  unsupported operation (memory allocation)`` on tasks with deformable geometry, such as
+  ``Isaac-Lift-Cloth-Franka-Camera``. Scene-query tasks are now captured into one graph each, because Warp
+  forbids memory allocation inside the conditional body that previously held them.

source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py

@@ -477,9 +477,8 @@ class NewtonManager(PhysicsManager):
 
     # Newton scene-query scheduling and graph execution.
     _sensor_tasks: dict[str, Callable[[], None]] = {}
-    _sensor_graph: wp.Graph | None = None
-    _sensor_flags: wp.array | None = None
-    _sensor_flags_host: np.ndarray | None = None
+    _sensor_refit_graph: wp.Graph | None = None
+    _sensor_task_graphs: dict[str, wp.Graph] = {}
     _sensor_state: State | None = None
     _sensor_state_dirty: bool = True
     _sensor_graph_capture_failed: bool = False
@@ -1161,7 +1160,7 @@ def clear(cls):
         NewtonManager._graph = None
         NewtonManager._graph_capture_pending = False
         NewtonManager._sensor_tasks = {}
-        NewtonManager._invalidate_sensor_graph()
+        NewtonManager._invalidate_sensor_graphs()
         NewtonManager._sensor_state = None
         NewtonManager._sensor_state_dirty = True
         NewtonManager._sensor_graph_capture_failed = False
@@ -2698,13 +2697,13 @@ def _register_sensor_task(cls, name: str, update_fn: Callable[[], None]) -> None
         cls._sensor_tasks[name] = update_fn
         cls._sensor_state = state
         cls._sensor_state_dirty = True
-        cls._invalidate_sensor_graph()
+        cls._invalidate_sensor_graphs()
 
     @classmethod
     def _unregister_sensor_task(cls, name: str) -> None:
         """Remove a scene-query task, ignoring unknown names."""
         if cls._sensor_tasks.pop(name, None) is not None:
-            cls._invalidate_sensor_graph()
+            cls._invalidate_sensor_graphs()
 
     @classmethod
     def _update_sensor_tasks(cls, *names: str) -> None:
@@ -2718,29 +2717,24 @@ def _update_sensor_tasks(cls, *names: str) -> None:
         if state is not cls._sensor_state:
             cls._sensor_state = state
             cls._sensor_state_dirty = True
-            cls._invalidate_sensor_graph()
+            cls._invalidate_sensor_graphs()
         cfg = PhysicsManager._cfg
         use_cuda_graph = bool(getattr(cfg, "use_cuda_graph", False)) and "cuda" in str(PhysicsManager._device)
-        if use_cuda_graph and cls._sensor_graph is None and not cls._sensor_graph_capture_failed:
-            cls._capture_sensor_graph()
-        if cls._sensor_graph is None:
+        if use_cuda_graph and cls._sensor_refit_graph is None and not cls._sensor_graph_capture_failed:
+            cls._capture_sensor_graphs()
+        if cls._sensor_refit_graph is None:
             if cls._sensor_state_dirty:
                 cls._refit_sensor_bvh()
                 cls._sensor_state_dirty = False
             for name in names:
                 cls._sensor_tasks[name]()
             return
 
-        assert cls._sensor_flags_host is not None
-        assert cls._sensor_flags is not None
-        cls._sensor_flags_host.fill(0)
-        cls._sensor_flags_host[0] = int(cls._sensor_state_dirty)
-        task_names = tuple(cls._sensor_tasks)
+        if cls._sensor_state_dirty:
+            wp.capture_launch(cls._sensor_refit_graph)
+            cls._sensor_state_dirty = False
         for name in names:
-            cls._sensor_flags_host[1 + task_names.index(name)] = 1
-        cls._sensor_flags.assign(cls._sensor_flags_host)
-        wp.capture_launch(cls._sensor_graph)
-        cls._sensor_state_dirty = False
+            wp.capture_launch(cls._sensor_task_graphs[name])
 
     @classmethod
     def _mark_sensor_state_dirty(cls) -> None:
@@ -2757,7 +2751,7 @@ def _mark_sensor_state_dirty(cls) -> None:
             return
         if NewtonManager._state_0 is not NewtonManager._sensor_state:
             NewtonManager._sensor_state = NewtonManager._state_0
-            NewtonManager._invalidate_sensor_graph()
+            NewtonManager._invalidate_sensor_graphs()
         NewtonManager._sensor_state_dirty = True
 
     @classmethod
@@ -2781,49 +2775,61 @@ def _refit_sensor_bvh(cls) -> None:
             cls._model.bvh_refit_particles(cls._sensor_state)
 
     @classmethod
-    def _invalidate_sensor_graph(cls) -> None:
-        """Discard captured scene-query graph resources."""
-        cls._sensor_graph = None
-        cls._sensor_flags = None
-        cls._sensor_flags_host = None
+    def _invalidate_sensor_graphs(cls) -> None:
+        """Discard the captured scene-query graphs."""
+        cls._sensor_refit_graph = None
+        cls._sensor_task_graphs = {}
         cls._sensor_graph_capture_failed = False
 
     @classmethod
-    def _capture_sensor_graph(cls) -> None:
-        """Capture BVH refit and scene-query tasks into a conditional graph."""
-        with wp.ScopedDevice(PhysicsManager._device):
+    def _capture_sensor_graphs(cls) -> None:
+        """Capture the BVH refit and each scene-query task into its own graph.
+
+        Each one gets a standalone top-level graph rather than a ``wp.capture_if``
+        conditional body of one shared graph: Warp rejects memory allocation inside a
+        conditional body, and ``wp.Mesh.refit`` allocates scratch on every call, so
+        deformable geometry in the tiled-camera render path would fail to capture.
+        """
+        device = PhysicsManager._device
+        with wp.ScopedDevice(device):
             cls._refit_sensor_bvh()
             for update_fn in cls._sensor_tasks.values():
                 update_fn()
 
-        cls._sensor_flags = wp.zeros(1 + len(cls._sensor_tasks), dtype=wp.int32, device=PhysicsManager._device)
-        cls._sensor_flags_host = np.zeros(1 + len(cls._sensor_tasks), dtype=np.int32)
-        update_fns = tuple(cls._sensor_tasks.values())
+        refit_graph = cls._capture_sensor_graph(device, cls._refit_sensor_bvh)
+        failed = None if refit_graph is not None else "bvh refit"
+        task_graphs: dict[str, wp.Graph] = {}
+        if failed is None:
+            for name, update_fn in cls._sensor_tasks.items():
+                graph = cls._capture_sensor_graph(device, update_fn)
+                if graph is None:
+                    failed = name
+                    break
+                task_graphs[name] = graph
+
+        if failed is not None:
+            cls._invalidate_sensor_graphs()
+            # Latch after invalidating: _invalidate_sensor_graphs() clears the flag.
+            cls._sensor_graph_capture_failed = True
+            logger.warning("Newton sensor graph capture failed for '%s'; falling back to eager execution.", failed)
+            return
 
-        def pipeline() -> None:
-            assert cls._sensor_flags is not None
-            wp.capture_if(cls._sensor_flags[0:1], cls._refit_sensor_bvh)
-            for index, update_fn in enumerate(update_fns):
-                wp.capture_if(cls._sensor_flags[index + 1 : index + 2], update_fn)
+        cls._sensor_refit_graph = refit_graph
+        cls._sensor_task_graphs = task_graphs
+        logger.info("Captured Newton sensor graphs for %d task(s).", len(task_graphs))
 
-        device = PhysicsManager._device
+    @classmethod
+    def _capture_sensor_graph(cls, device: str, capture_target: Callable[[], None]) -> wp.Graph | None:
+        """Capture one scene-query callable into a standalone graph, or ``None`` on failure."""
         if cls._usdrt_stage is not None:
-            cls._sensor_graph = cls._capture_relaxed_graph(device, capture_target=pipeline)
-        else:
-            try:
-                with wp.ScopedCapture(device=device) as capture:
-                    pipeline()
-                cls._sensor_graph = capture.graph
-            except Exception:
-                logger.exception("[NewtonManager] sensor CUDA graph capture failed")
-                cls._sensor_graph = None
-        if cls._sensor_graph is None:
-            cls._sensor_flags = None
-            cls._sensor_flags_host = None
-            cls._sensor_graph_capture_failed = True
-            logger.warning("Newton sensor graph capture failed; falling back to eager execution.")
-        else:
-            logger.info("Captured Newton sensor graph with %d task(s).", len(cls._sensor_tasks))
+            return cls._capture_relaxed_graph(device, capture_target=capture_target)
+        try:
+            with _paused_gc(), wp.ScopedCapture(device=device) as capture:
+                capture_target()
+            return capture.graph
+        except Exception:
+            logger.exception("[NewtonManager] sensor CUDA graph capture failed")
+            return None
 
     @classmethod
     def get_num_envs(cls) -> int:

source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py

@@ -338,9 +338,8 @@ def get_state(cls):
     monkeypatch.setattr(NewtonManager, "_sensor_tasks", {}, raising=False)
     monkeypatch.setattr(NewtonManager, "_sensor_state", None, raising=False)
     monkeypatch.setattr(NewtonManager, "_sensor_state_dirty", True, raising=False)
-    monkeypatch.setattr(NewtonManager, "_sensor_graph", None, raising=False)
-    monkeypatch.setattr(NewtonManager, "_sensor_flags", None, raising=False)
-    monkeypatch.setattr(NewtonManager, "_sensor_flags_host", None, raising=False)
+    monkeypatch.setattr(NewtonManager, "_sensor_refit_graph", None, raising=False)
+    monkeypatch.setattr(NewtonManager, "_sensor_task_graphs", {}, raising=False)
     monkeypatch.setattr(NewtonManager, "_sensor_graph_capture_failed", False, raising=False)
     monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=False), raising=False)
 
@@ -350,6 +349,38 @@ def get_state(cls):
     assert status["rendered"]
 
 
+@pytest.mark.skipif(not wp.get_cuda_device_count(), reason="CUDA is unavailable")
+def test_sensor_graph_captures_allocating_task(monkeypatch):
+    """A sensor task that allocates scratch is still graph-capturable.
+
+    ``wp.Mesh.refit`` allocates scratch natively on every call, which Warp rejects
+    inside a ``wp.capture_if`` conditional body. Deformable geometry reaches it through
+    the tiled-camera render task, so each task must be captured as its own graph.
+    """
+
+    points = wp.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=wp.vec3f, device="cuda:0")
+    mesh = wp.Mesh(points, wp.array([0, 1, 2], dtype=wp.int32, device="cuda:0"))
+
+    state = object()
+    monkeypatch.setattr(NewtonManager, "get_state", classmethod(lambda cls: state))
+    monkeypatch.setattr(NewtonManager, "_model", None, raising=False)
+    monkeypatch.setattr(NewtonManager, "_sensor_tasks", {"mesh_refit": mesh.refit}, raising=False)
+    monkeypatch.setattr(NewtonManager, "_sensor_state", state, raising=False)
+    monkeypatch.setattr(NewtonManager, "_sensor_state_dirty", True, raising=False)
+    monkeypatch.setattr(NewtonManager, "_sensor_refit_graph", None, raising=False)
+    monkeypatch.setattr(NewtonManager, "_sensor_task_graphs", {}, raising=False)
+    monkeypatch.setattr(NewtonManager, "_sensor_graph_capture_failed", False, raising=False)
+    monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False)
+    monkeypatch.setattr(PhysicsManager, "_device", "cuda:0", raising=False)
+    monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False)
+
+    NewtonManager._update_sensor_tasks("mesh_refit")
+
+    assert not NewtonManager._sensor_graph_capture_failed
+    assert NewtonManager._sensor_refit_graph is not None
+    assert "mesh_refit" in NewtonManager._sensor_task_graphs
+
+
 def test_sensor_bvh_shape_flags_are_fixed_before_builder_creation(monkeypatch):
     """Builder finalization includes collision-only shapes without a later BVH rebuild."""
     import newton

source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py

@@ -252,4 +252,4 @@ def test_renderer_and_raycast_share_newton_manager_graph(sim):
     task_names = sorted(NewtonManager._sensor_tasks)
     assert any(name.startswith("newton_raycast:") for name in task_names)
     assert any(name.startswith("newton_warp_render:") for name in task_names)
-    assert (NewtonManager._sensor_graph is not None) == sim.cfg.physics.use_cuda_graph
+    assert bool(NewtonManager._sensor_task_graphs) == sim.cfg.physics.use_cuda_graph