Skip to content

Accelerate Newton transform sync for Isaac RTX #7553

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ooctipus wants to merge 1 commit into
base: develop
Choose a base branch
from
+77 −100
Open

Accelerate Newton transform sync for Isaac RTX #7553

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions source/isaaclab_newton/changelog.d/accelerate-fabric-transform-sync.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
Fixed
^^^^^

* Accelerated Newton-to-Isaac RTX transform synchronization by reusing Fabric bindings and GPU hierarchy propagation.
173 changes: 73 additions & 100 deletions source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)

@isaaclab-review-bot isaaclab-review-bot Bot Sep 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Warning · Implementation — CUDA-only stream sync breaks CPU device

wp.synchronize_stream resolves its argument via get_device(...).stream, which raises for non-CUDA devices, while this path still binds and launches on PhysicsManager._device. On a CPU device this now throws into the broad except, _transforms_dirty stays set, and every render frame logs a traceback with no transform update, whereas the removed update_world_xforms() fallback previously handled it. wp.synchronize_device on line 762 already provides the ordering.

if not cls._fabric_hierarchy.update_world_xforms_gpu(not topology_changed):

@isaaclab-review-bot isaaclab-review-bot Bot Sep 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Warning · Design Architecture — Retry after failed propagation compounds local writes

Local matrices are written and synchronized before this call. When update_world_xforms_gpu returns false or raises, the handler leaves _transforms_dirty true so pre_render() retries, but the kernel derives the next local from the already-mutated local combined with the still-stale world matrix, so each retry reapplies the delta and corrupts poses. Restore a fallback propagation or make the local write idempotent before allowing a retry.

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")

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()
Expand Down
Loading