From 86f851d6d8168eeea24a1ebad7d0ae4ebc2279d4 Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Thu, 3 Sep 2026 18:19:45 -0700 Subject: [PATCH] Accelerate Newton Fabric transform sync --- .../accelerate-fabric-transform-sync.rst | 4 + .../isaaclab_newton/physics/newton_manager.py | 173 ++++++++---------- 2 files changed, 77 insertions(+), 100 deletions(-) create mode 100644 source/isaaclab_newton/changelog.d/accelerate-fabric-transform-sync.rst diff --git a/source/isaaclab_newton/changelog.d/accelerate-fabric-transform-sync.rst b/source/isaaclab_newton/changelog.d/accelerate-fabric-transform-sync.rst new file mode 100644 index 000000000000..f671d6b4f990 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/accelerate-fabric-transform-sync.rst @@ -0,0 +1,4 @@ +Fixed +^^^^^ + +* Accelerated Newton-to-Isaac RTX transform synchronization by reusing Fabric bindings and GPU hierarchy propagation. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 24363b5c1147..e9eceac774c4 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -164,12 +164,13 @@ def _capture_fabric_scales( @wp.kernel(enable_backward=False) def _set_fabric_transforms( - fabric_transforms: wp.fabricarray(dtype=wp.mat44d), + fabric_world_transforms: wp.fabricarray(dtype=wp.mat44d), + fabric_local_transforms: wp.fabricarray(dtype=wp.mat44d), newton_indices: wp.fabricarray(dtype=wp.uint32), newton_body_q: wp.array(ndim=1, dtype=wp.transformf), body_scales: wp.array(dtype=wp.vec3f), ): - """Write Newton body poses to Fabric world matrices with their initialized scale. + """Write Newton body poses to Fabric local matrices with their initialized scale. For each Fabric prim at thread ``i``, reads the Newton body transform at ``newton_body_q[newton_indices[i]]`` and combines its translation and rotation @@ -181,15 +182,12 @@ def _set_fabric_transforms( idx = int(newton_indices[i]) transform = newton_body_q[idx] scale = body_scales[idx] - fabric_transforms[i] = wp.mat44d( - wp.transpose( - wp.transform_compose( - wp.transform_get_translation(transform), - wp.transform_get_rotation(transform), - scale, - ) + world_transform = wp.transpose( + wp.mat44d( + wp.transform_compose(wp.transform_get_translation(transform), wp.transform_get_rotation(transform), scale) ) ) + fabric_local_transforms[i] = world_transform * wp.inverse(fabric_world_transforms[i]) * fabric_local_transforms[i] @wp.kernel(enable_backward=False) @@ -488,9 +486,13 @@ class NewtonManager(PhysicsManager): # USD/Fabric sync _newton_stage_path = None _usdrt_stage = None + _fabric_hierarchy = None _newton_index_attr = "newton:index" - # Body-indexed world scales captured before Newton first overwrites Fabric transforms. + # Body-indexed world scales captured before Newton first updates Fabric transforms. _fabric_body_scales: wp.array | None = None + # Persistent query and arrays: selection, world matrix, local matrix, and Newton index. + _fabric_body_sync: tuple[Any, Any, Any, Any] | None = None + _fabric_body_binding_count: int = 0 _clone_physics_only = False _transforms_dirty: bool = False _transforms_may_change_on_graph_replay: bool = False @@ -504,9 +506,6 @@ class NewtonManager(PhysicsManager): _newton_particle_count_attr = "newton:particleCount" _particle_visual_prims: dict[str, _ParticleVisualPrim] = {} - # Cached after the first fabric sync that probes IFabricHierarchy GPU APIs. - _use_fabric_gpu_hierarchy: bool | None = None - # Set to True after sync_transforms_to_usd() successfully writes body positions for # the first time in each simulation session. Reset to False in clear(). Polled by # test drain helpers to know when the GPU has propagated the newton:index Fabric @@ -687,15 +686,8 @@ def sync_transforms_to_usd(cls) -> None: Uses ``wp.fabricarray`` directly (no ``isaacsim.physics.newton`` extension needed). On the first successful sync, a Warp kernel captures each initialized Fabric world scale by Newton body index. The pose kernel then combines that scale with - ``state_0.body_q[newton_index[i]]`` and writes the corresponding ``mat44d`` to - ``omni:fabric:worldMatrix`` for each prim. - - When ``IFabricHierarchy.update_world_xforms_gpu_with_options`` is - available the method mirrors PhysX's ``DirectGpuHelper`` pattern: pause - Fabric change tracking, write transforms, resume tracking, then run the - GPU hierarchy update with ``RIGID_BODY | FORCE_UPDATE`` so Newton-authored - world matrices stay authoritative on rigid-body prims. Otherwise it - falls back to the CPU ``update_world_xforms()`` path. + ``state_0.body_q[newton_index[i]]``, derives the corresponding local matrix, + and lets Fabric propagate the hierarchy on the GPU. """ if cls._usdrt_stage is None or cls._model is None or cls._state_0 is None: return @@ -704,95 +696,73 @@ def sync_transforms_to_usd(cls) -> None: try: import usdrt - fabric_hierarchy = None - gpu_opts_cls = None - if hasattr(usdrt, "hierarchy"): - fabric_hierarchy = usdrt.hierarchy.IFabricHierarchy().get_fabric_hierarchy( - cls._usdrt_stage.GetFabricId(), cls._usdrt_stage.GetStageIdAsStageId() - ) - gpu_opts_cls = getattr(usdrt.hierarchy, "FabricHierarchyGpuUpdateOptions", None) - - if cls._use_fabric_gpu_hierarchy is None and hasattr(usdrt, "hierarchy"): - # Probe the pybind class once so a transient null hierarchy handle does - # not permanently disable the GPU path for the session. - NewtonManager._use_fabric_gpu_hierarchy = gpu_opts_cls is not None and hasattr( - usdrt.hierarchy.IFabricHierarchy, "update_world_xforms_gpu_with_options" - ) - if cls._use_fabric_gpu_hierarchy: - logger.info("Fabric GPU transform hierarchy enabled via IFabricHierarchy") - else: - logger.info("Fabric GPU transform hierarchy unavailable; falling back to update_world_xforms()") - - use_gpu_hierarchy = bool( - cls._use_fabric_gpu_hierarchy and fabric_hierarchy is not None and gpu_opts_cls is not None - ) - - # Pause hierarchy change tracking BEFORE SelectPrims. - # SelectPrims with ReadWrite access calls getAttributeArrayGpu - # internally, which marks Fabric buffers dirty. If tracking is - # still active at that point the hierarchy records the change and - # Kit's updateWorldXforms will do an expensive connectivity - # rebuild every frame. PhysX avoids this via ScopedUSDRT which - # pauses tracking before any Fabric writes. - if use_gpu_hierarchy: - fabric_hierarchy.track_world_xform_changes(False) - fabric_hierarchy.track_local_xform_changes(False) - - try: + topology_changed = False + fabric_sync = cls._fabric_body_sync + if fabric_sync is None: selection = cls._usdrt_stage.SelectPrims( require_attrs=[ - (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.ReadWrite), + (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.Read), + ( + usdrt.Sdf.ValueTypeNames.Matrix4d, + "omni:fabric:localMatrix", + usdrt.Usd.Access.ReadWrite, + ), (usdrt.Sdf.ValueTypeNames.UInt, cls._newton_index_attr, usdrt.Usd.Access.Read), ], device=str(PhysicsManager._device), ) - if selection.GetCount() == 0: - # The newton:index attribute is written CPU-side by start_simulation() but - # GPU propagation is deferred. Keep _transforms_dirty=True so the next - # pre_render() retries once initialize_solver() has completed (FK delegate - # bound) and body_q holds valid values. - if cls._eval_fk is _eval_fk_unbound: + else: + selection, fabric_world_transforms, fabric_local_transforms, newton_indices = fabric_sync + topology_changed = selection.PrepareForReuse() + + if fabric_sync is None or topology_changed: + if selection.GetCount() != cls._fabric_body_binding_count: + NewtonManager._fabric_body_sync = None + # CPU-authored newton:index attributes may not have reached Fabric yet. + if selection.GetCount() == 0 and cls._eval_fk is _eval_fk_unbound: NewtonManager._transforms_dirty = False return - - fabric_transforms = wp.fabricarray(selection, "omni:fabric:worldMatrix") + fabric_world_transforms = wp.fabricarray(selection, "omni:fabric:worldMatrix") + fabric_local_transforms = wp.fabricarray(selection, "omni:fabric:localMatrix") newton_indices = wp.fabricarray(selection, cls._newton_index_attr) - if cls._fabric_body_scales is None: - NewtonManager._fabric_body_scales = wp.empty( - cls._model.body_count, - dtype=wp.vec3f, - device=PhysicsManager._device, - ) - wp.launch( - _capture_fabric_scales, - dim=newton_indices.shape[0], - inputs=[fabric_transforms, newton_indices, cls._fabric_body_scales], - device=PhysicsManager._device, - ) + NewtonManager._fabric_body_sync = ( + selection, + fabric_world_transforms, + fabric_local_transforms, + newton_indices, + ) + if cls._fabric_body_scales is None: + NewtonManager._fabric_body_scales = wp.empty( + cls._model.body_count, + dtype=wp.vec3f, + device=PhysicsManager._device, + ) wp.launch( - _set_fabric_transforms, + _capture_fabric_scales, dim=newton_indices.shape[0], - inputs=[fabric_transforms, newton_indices, cls._state_0.body_q, cls._fabric_body_scales], + inputs=[fabric_world_transforms, newton_indices, cls._fabric_body_scales], device=PhysicsManager._device, ) - wp.synchronize_device(PhysicsManager._device) + wp.launch( + _set_fabric_transforms, + dim=newton_indices.shape[0], + inputs=[ + fabric_world_transforms, + fabric_local_transforms, + newton_indices, + cls._state_0.body_q, + cls._fabric_body_scales, + ], + device=PhysicsManager._device, + ) - NewtonManager._newton_fabric_ready = True - NewtonManager._transforms_dirty = False + wp.synchronize_stream(PhysicsManager._device) + if not cls._fabric_hierarchy.update_world_xforms_gpu(not topology_changed): + raise RuntimeError("Fabric GPU hierarchy propagation failed.") + wp.synchronize_device(PhysicsManager._device) - if use_gpu_hierarchy: - # RIGID_BODY: inverse-propagate on PhysicsRigidBodyAPI buckets - # (keep Newton world matrices, derive local). FORCE_UPDATE: - # bypass the change-listener dirty check after tracking pause. - fabric_hierarchy.update_world_xforms_gpu_with_options( - gpu_opts_cls.RIGID_BODY | gpu_opts_cls.FORCE_UPDATE - ) - elif fabric_hierarchy is not None: - fabric_hierarchy.update_world_xforms() - finally: - if use_gpu_hierarchy: - fabric_hierarchy.track_world_xform_changes(True) - fabric_hierarchy.track_local_xform_changes(True) + NewtonManager._newton_fabric_ready = True + NewtonManager._transforms_dirty = False except Exception: logger.exception("[NewtonManager] sync_transforms_to_usd FAILED") @@ -1123,7 +1093,6 @@ def clear(cls): NewtonManager._visualization_stop_callback = None if callback is not None: callback.deregister() - NewtonManager._use_fabric_gpu_hierarchy = None NewtonManager._newton_fabric_ready = False NewtonManager._builder = None NewtonManager._model = None @@ -1168,7 +1137,10 @@ def clear(cls): NewtonManager._sensor_bvh_shape_flags = ShapeFlags.VISIBLE NewtonManager._newton_stage_path = None NewtonManager._usdrt_stage = None + NewtonManager._fabric_hierarchy = None NewtonManager._fabric_body_scales = None + NewtonManager._fabric_body_sync = None + NewtonManager._fabric_body_binding_count = 0 NewtonManager._transforms_dirty = False NewtonManager._transforms_may_change_on_graph_replay = False NewtonManager._particles_dirty = False @@ -1656,10 +1628,13 @@ def start_simulation(cls) -> None: if body_bindings is None: # Non-replicated Newton stages do not pass through NewtonReplicateContext. body_bindings = [(body_path, i) for i, body_path in enumerate(body_paths)] + NewtonManager._fabric_body_sync = None + NewtonManager._fabric_body_binding_count = len(body_bindings) fabric_hierarchy = usdrt.hierarchy.IFabricHierarchy().get_fabric_hierarchy( cls._usdrt_stage.GetFabricId(), cls._usdrt_stage.GetStageIdAsStageId() ) + NewtonManager._fabric_hierarchy = fabric_hierarchy NewtonManager._initialize_fabric_body_prims(cls._usdrt_stage, fabric_hierarchy, usdrt, body_bindings) NewtonManager._initialize_fabric_cable_prims(cls._usdrt_stage, fabric_hierarchy, usdrt) @@ -1690,9 +1665,7 @@ def _initialize_fabric_body_prims(stage, fabric_hierarchy, usdrt, body_bindings: prim.CreateAttribute(NewtonManager._newton_index_attr, usdrt.Sdf.ValueTypeNames.UInt, custom=True) prim.GetAttribute(NewtonManager._newton_index_attr).Set(body_index) - # Tag with PhysicsRigidBodyAPI so FabricHierarchyGpuUpdateOptions.RIGID_BODY - # applies Inverse propagation (preserves Newton's world transforms and derives - # local) instead of Forward. + # Generated Xforms need the same rigid-body identity as authored body prims. prim.AddAppliedSchema("PhysicsRigidBodyAPI") fabric_hierarchy.update_world_xforms()