# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""Newton-backed FrameView using Newton body labels and injected sites."""
from __future__ import annotations
import contextlib
import logging
import re
import sys
import numpy as np
import warp as wp
from newton import ShapeFlags
from pxr import UsdPhysics
import isaaclab.sim as sim_utils
from isaaclab import cloner
from isaaclab.physics import PhysicsEvent
from isaaclab.sim.views.base_frame_view import BaseFrameView
from isaaclab.sim.views.fabric_xform_selection import FabricXformSelection
from isaaclab.sim.views.xform_space_writer import FrameViewLocalSpaceWriter, FrameViewWorldSpaceWriter
from isaaclab.utils.string import resolve_matching_names
from isaaclab.utils.warp import ProxyArray
from isaaclab.utils.warp import fabric as fabric_utils
from isaaclab_newton.physics.newton_manager import NewtonManager
logger = logging.getLogger(__name__)
WORLD_BODY_INDEX = -1
# Regex metacharacters that mark a body pattern as a genuine expression rather than a literal
# USD path. Patterns free of these can be resolved via an exact dict lookup instead of scanning
# every body label with a compiled regex.
_REGEX_TOKENS = frozenset(".*[]()+?|\\^$")
# One resolved site registration: (body_patterns, local transform, xform scale, per_world, env_ids,
# destination prim paths). Prim paths follow the spec's expansion order.
_Strs = tuple[str, ...]
_SiteSpec = tuple[_Strs | None, wp.transform, tuple[float, float, float], bool, tuple[int, ...] | None, _Strs | None]
def _extend_prim_paths(collected: list[str] | None, paths: tuple[str, ...] | None, expected: int) -> list[str] | None:
"""Append one spec's destination paths, or return ``None`` to give up on mirroring this view."""
if collected is None or paths is None or len(paths) != expected:
return None
collected.extend(paths)
return collected
@wp.kernel
def _compute_site_world_transforms(
body_q: wp.array(dtype=wp.transformf),
site_body: wp.array(dtype=wp.int32),
site_local: wp.array(dtype=wp.transformf),
indices: wp.array(dtype=wp.int32),
out_pos: wp.array(dtype=wp.vec3f),
out_quat: wp.array(dtype=wp.vec4f),
):
"""Compute world-space transforms for selected sites."""
i = wp.tid()
si = indices[i]
bid = site_body[si]
if bid == WORLD_BODY_INDEX:
world = site_local[si]
else:
world = wp.transform_multiply(body_q[bid], site_local[si])
out_pos[i] = wp.transform_get_translation(world)
q = wp.transform_get_rotation(world)
out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3])
@wp.kernel
def _gather_site_local_transforms(
site_local: wp.array(dtype=wp.transformf),
indices: wp.array(dtype=wp.int32),
out_pos: wp.array(dtype=wp.vec3f),
out_quat: wp.array(dtype=wp.vec4f),
):
"""Gather local transforms for selected sites."""
i = wp.tid()
si = indices[i]
local_tf = site_local[si]
out_pos[i] = wp.transform_get_translation(local_tf)
q = wp.transform_get_rotation(local_tf)
out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3])
@wp.kernel
def _write_site_local_from_world_poses(
body_q: wp.array(dtype=wp.transformf),
site_body: wp.array(dtype=wp.int32),
indices: wp.array(dtype=wp.int32),
world_pos: wp.array(dtype=wp.vec3f),
world_quat: wp.array(dtype=wp.vec4f),
site_local: wp.array(dtype=wp.transformf),
):
"""Update local offsets so selected sites reach desired world poses."""
i = wp.tid()
si = indices[i]
w_pos = world_pos[i]
w_q = world_quat[i]
desired_world = wp.transform(w_pos, wp.quatf(w_q[0], w_q[1], w_q[2], w_q[3]))
bid = site_body[si]
if bid == WORLD_BODY_INDEX:
site_local[si] = desired_world
else:
site_local[si] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world)
@wp.kernel
def _write_site_local_from_local_poses(
indices: wp.array(dtype=wp.int32),
local_pos: wp.array(dtype=wp.vec3f),
local_quat: wp.array(dtype=wp.vec4f),
site_local: wp.array(dtype=wp.transformf),
):
"""Update local offsets for selected sites."""
i = wp.tid()
si = indices[i]
l_pos = local_pos[i]
l_q = local_quat[i]
site_local[si] = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3]))
@wp.kernel(enable_backward=False)
def _gather_mirrored_site_poses(
site_positions: wp.array(dtype=wp.vec3f),
site_orientations: wp.array(dtype=wp.vec4f),
site_indices: wp.array(dtype=wp.int32),
out_positions: wp.array(dtype=wp.float32, ndim=2),
out_orientations: wp.array(dtype=wp.float32, ndim=2),
):
"""Gather the mirrored sites' world poses into the flat layout the Fabric kernels expect."""
i = wp.tid()
site = site_indices[i]
position = site_positions[site]
orientation = site_orientations[site]
for axis in range(3):
out_positions[i, axis] = position[axis]
for component in range(4):
out_orientations[i, component] = orientation[component]
@wp.kernel
def _gather_shape_scales(
shape_scale: wp.array(dtype=wp.vec3f),
shape_body: wp.array(dtype=wp.int32),
site_body: wp.array(dtype=wp.int32),
indices: wp.array(dtype=wp.int32),
num_shapes: wp.int32,
out_scales: wp.array(dtype=wp.vec3f),
):
"""Gather legacy per-site geometry scales from collision shapes on the same body."""
i = wp.tid()
si = indices[i]
bid = site_body[si]
found = int(0)
for s in range(num_shapes):
if shape_body[s] == bid and found == 0:
out_scales[i] = shape_scale[s]
found = 1
if found == 0:
out_scales[i] = wp.vec3f(1.0, 1.0, 1.0)
@wp.kernel
def _scatter_shape_scales(
site_body: wp.array(dtype=wp.int32),
indices: wp.array(dtype=wp.int32),
new_scales: wp.array(dtype=wp.vec3f),
shape_body: wp.array(dtype=wp.int32),
num_shapes: wp.int32,
shape_scale: wp.array(dtype=wp.vec3f),
):
"""Scatter legacy per-site geometry scales to collision shapes on the same body."""
i = wp.tid()
si = indices[i]
bid = site_body[si]
for s in range(num_shapes):
if shape_body[s] == bid:
shape_scale[s] = new_scales[i]
@wp.kernel
def _gather_xform_scales(
site_xform_scale: wp.array(dtype=wp.vec3f),
indices: wp.array(dtype=wp.int32),
out_scales: wp.array(dtype=wp.vec3f),
):
"""Gather per-site xform scales."""
i = wp.tid()
out_scales[i] = site_xform_scale[indices[i]]
@wp.kernel
def _scatter_xform_scales(
indices: wp.array(dtype=wp.int32),
new_scales: wp.array(dtype=wp.vec3f),
site_xform_scale: wp.array(dtype=wp.vec3f),
):
"""Scatter per-site xform scales."""
i = wp.tid()
site_xform_scale[indices[i]] = new_scales[i]
[docs]
class NewtonSiteFrameView(BaseFrameView):
"""Batched Newton site view for non-physics frames.
The public construction contract matches the generic :class:`FrameView`:
callers provide a prim expression and the backend resolves the source prim
into Newton body-local or world-local sites.
"""
[docs]
def __init__(
self,
prim_path: str | list[str],
device: str = "cpu",
validate_xform_ops: bool = True,
stage: object | None = None,
**kwargs,
):
"""Initialize the Newton site frame view.
Args:
prim_path: User-facing frame path pattern, or list of patterns.
device: Warp device for GPU arrays.
validate_xform_ops: Whether to validate source USD xform ops.
stage: USD stage that contains the source prims.
**kwargs: Unused.
"""
del kwargs
self._prim_paths = [prim_path] if isinstance(prim_path, str) else list(prim_path)
self._prim_path = prim_path if isinstance(prim_path, str) else ", ".join(self._prim_paths)
self._device = device
self._prims = []
stage = sim_utils.get_current_stage() if stage is None else stage
self._site_specs = self._resolve_site_specs(stage, validate_xform_ops)
self._site_labels: list[str] = []
self._site_label_scales: list[tuple[float, float, float]] = []
# Destination prim paths per label, in expansion order; ``None`` when that is not yet known.
self._site_label_prim_paths: list[tuple[str, ...] | None] = []
self._site_prim_paths: list[str] | None = None
# Fabric mirror state, built on the first write (see :meth:`_mirror_to_fabric`).
self._fabric_sel: FabricXformSelection | None = None
self._mirror_disabled = False
# Set only on the pre-model path below; released again by :meth:`close`.
self._physics_ready_handle = None
self._site_body: wp.array | None = None
self._site_local: wp.array | None = None
self._site_xform_scale: wp.array | None = None
self._site_indices: wp.array | None = None
self._pos_buf: wp.array | None = None
self._quat_buf: wp.array | None = None
self._local_pos_buf: wp.array | None = None
self._local_quat_buf: wp.array | None = None
self._scale_buf: wp.array | None = None
self._pos_ta: ProxyArray | None = None
self._quat_ta: ProxyArray | None = None
self._local_pos_ta: ProxyArray | None = None
self._local_quat_ta: ProxyArray | None = None
self._scale_ta: ProxyArray | None = None
self._count = 0
model = NewtonManager.get_model()
if model is not None:
self._initialize_from_specs(model)
else:
for body_patterns, xform, scale, per_world, _env_ids, spec_paths in self._site_specs:
if body_patterns is None:
self._site_labels.append(NewtonManager.cl_register_site(None, xform, per_world=per_world))
self._site_label_scales.append(scale)
self._site_label_prim_paths.append(spec_paths)
else:
for body_pattern in body_patterns:
self._site_labels.append(NewtonManager.cl_register_site(body_pattern, xform))
self._site_label_scales.append(scale)
self._site_label_prim_paths.append(spec_paths)
self._physics_ready_handle = NewtonManager.register_callback(
self._on_physics_ready, PhysicsEvent.PHYSICS_READY, name=f"site_view_{self._prim_path}"
)
def _resolve_site_specs(self, stage, validate_xform_ops: bool) -> list[_SiteSpec]:
"""Resolve source prims into Newton site registration specs."""
plan = sim_utils.SimulationContext.instance().get_clone_plan()
groups = ()
if plan is not None:
sources = cloner.path.get_asset_prototype_paths(plan)
templates, starts, worlds, world_starts = cloner.path.get_world_prototype_asset_templates(
plan, include_world_indices=True
)
groups = np.flatnonzero(np.diff(world_starts[1:])) + 1
model = NewtonManager.get_model()
body_labels = list(model.body_label) if model is not None else ()
shape_labels = list(model.shape_label) if model is not None else ()
shape_flags = None
before_physics = model is None
specs: list[_SiteSpec] = []
for path_expr in self._prim_paths:
if resolve_matching_names(path_expr, body_labels, raise_when_no_match=False)[1]:
raise ValueError(
f"FrameView prim '{path_expr}' is a Newton physics body. "
"FrameView should only be used for non-physics frames."
)
shape_indices, _ = resolve_matching_names(path_expr, shape_labels, raise_when_no_match=False)
if shape_indices:
if shape_flags is None:
shape_flags = model.shape_flags.numpy()
collision_flags = int(ShapeFlags.COLLIDE_SHAPES | ShapeFlags.COLLIDE_PARTICLES)
if np.any(shape_flags[shape_indices] & collision_flags):
raise ValueError(
f"FrameView prim '{path_expr}' matches a Newton collision shape. "
"FrameView should only be used for non-physics frames."
)
# Resolve frame expressions through the closest declared asset, then inspect only its prototype subtree.
matches = []
for group in groups:
for index in range(*starts[group : group + 2]):
if (matched := cloner.path.match(path_expr, templates[index])) is not None:
matches.append((group, index, matched))
if matches:
suffix = min((matched.suffix for _, _, matched in matches), key=len)
for group, index, matched in matches:
if matched.suffix != suffix:
continue
env_ids = worlds[world_starts[group] : world_starts[group + 1]]
env_ids = env_ids[
resolve_matching_names(matched.instance, env_ids.astype(str), raise_when_no_match=False)[0]
]
if not len(env_ids):
continue
root, template = sources[plan.topology.world_prototypes[index]], templates[index]
pattern = re.compile(root + suffix)
prims = sim_utils.get_all_matching_child_prims(
root, lambda prim: pattern.fullmatch(prim.GetPath().pathString) is not None, stage=stage
)
if not prims:
raise RuntimeError(f"FrameView '{path_expr}' could not resolve source prim '{root + suffix}'.")
ids = tuple(map(int, env_ids))
source_args = validate_xform_ops, root, template, ids, before_physics, stage
specs.extend(self._resolve_source_prim(prim, *source_args) for prim in prims)
continue
prims = sim_utils.find_matching_prims(path_expr, stage)
if not prims:
raise RuntimeError(f"FrameView '{path_expr}' could not resolve a source prim.")
source_args = validate_xform_ops, None, None, None, before_physics, stage
specs.extend(self._resolve_source_prim(prim, *source_args) for prim in prims)
return specs
def _resolve_source_prim(
self,
prim,
validate_xform_ops: bool,
source_root: str | None,
destination_template: str | None,
env_ids: tuple[int, ...] | None,
use_clone_body_pattern: bool,
stage,
) -> _SiteSpec:
"""Resolve one source prim into body patterns, local frame, xform scale, and destination paths."""
prim_path = prim.GetPath().pathString
dest_paths = (prim_path,)
if source_root is not None and destination_template is not None and env_ids is not None:
suffix = cloner.path.relative_to(prim_path, source_root)
dest_paths = tuple(destination_template.format(env_id) + suffix for env_id in env_ids)
if prim.HasAPI(UsdPhysics.CollisionAPI):
raise ValueError(
f"FrameView prim '{prim_path}' is a Newton collision shape. "
"FrameView should only be used for non-physics frames."
)
if prim.HasAPI(UsdPhysics.RigidBodyAPI) or prim.HasAPI(UsdPhysics.ArticulationRootAPI):
raise ValueError(
f"FrameView prim '{prim_path}' is a Newton physics body. "
"FrameView should only be used for non-physics frames."
)
if validate_xform_ops:
sim_utils.standardize_xform_ops(prim)
if not sim_utils.validate_standard_xform_ops(prim):
raise ValueError(f"FrameView prim '{prim_path}' does not have standard xform ops.")
scale_attr = prim.GetAttribute("xformOp:scale")
scale = (
tuple(float(v) for v in scale_attr.Get())
if scale_attr and scale_attr.HasAuthoredValue()
else (1.0, 1.0, 1.0)
)
body_prim = prim.GetParent()
while body_prim and body_prim.IsValid():
if body_prim.HasAPI(UsdPhysics.RigidBodyAPI) or body_prim.HasAPI(UsdPhysics.ArticulationRootAPI):
pos, quat = sim_utils.resolve_prim_pose(prim, body_prim)
body_path = body_prim.GetPath().pathString
spec_tail = (wp.transform(pos, quat), scale, False, env_ids, dest_paths)
if source_root is not None and destination_template is not None:
assert env_ids is not None
worlds = (".*",) if use_clone_body_pattern else env_ids
suffix = cloner.path.relative_to(body_path, source_root)
# A separately declared frame can be cloned below a body outside its source root.
if suffix is None and source_root.startswith(body_path + "/"):
suffix = source_root[len(body_path) :]
body_patterns = []
for env_id in worlds:
destination_root = destination_template.format(env_id)
if not destination_root.endswith(suffix):
raise RuntimeError(
f"FrameView destination root '{destination_root}' does not end with '{suffix}'."
)
body_patterns.append(destination_root[: -len(suffix)])
return (tuple(body_patterns), *spec_tail)
if suffix is None:
raise RuntimeError(f"FrameView source body '{body_path}' is not under '{source_root}'.")
body_patterns = tuple(destination_template.format(env_id) + suffix for env_id in worlds)
else:
body_patterns = (body_path,)
return (body_patterns, *spec_tail)
body_prim = body_prim.GetParent()
ref_path = source_root
if source_root is not None and destination_template is not None:
matched = cloner.path.match(source_root, destination_template.partition("{}")[0] + "{}")
if matched is not None:
ref_path = source_root.removesuffix(matched.suffix)
ref_prim = stage.GetPrimAtPath(ref_path) if ref_path is not None else None
pos, quat = sim_utils.resolve_prim_pose(prim, ref_prim if ref_prim and ref_prim.IsValid() else None)
return None, wp.transform(pos, quat), scale, source_root is not None, env_ids, dest_paths
def _on_physics_ready(self, _event) -> None:
"""Callback invoked when the Newton model becomes available."""
self._initialize_from_site_map(NewtonManager.get_model())
def _initialize_from_site_map(self, model) -> None:
"""Initialize arrays from injected Newton sites."""
site_map = NewtonManager._cl_site_index_map
body_t = wp.to_torch(model.shape_body)
xform_t = wp.to_torch(model.shape_transform)
site_bodies: list[int] = []
site_locals: list[list[float]] = []
site_scales: list[tuple[float, float, float]] = []
site_prim_paths: list[str] | None = []
for site_label, scale, label_paths in zip(
self._site_labels, self._site_label_scales, self._site_label_prim_paths, strict=True
):
global_idx, per_world = site_map[site_label]
site_indices = (
[global_idx] if per_world is None else [site_idx for sites in per_world for site_idx in sites]
)
for site_idx in site_indices:
site_bodies.append(int(body_t[site_idx].item()))
site_locals.append([float(v) for v in xform_t[site_idx].tolist()])
site_scales.append(scale)
site_prim_paths = _extend_prim_paths(site_prim_paths, label_paths, len(site_indices))
self._create_buffers(site_bodies, site_locals, site_scales, site_prim_paths)
def _initialize_from_specs(self, model) -> None:
"""Initialize arrays directly from resolved specs and Newton body labels."""
body_labels = list(model.body_label)
# Exact label -> index map, built once. Replicated frames expand to one concrete
# body path per environment, so matching each against every label via regex is
# ``O(num_envs * num_bodies)`` (quadratic in ``num_envs``). Fast-pathing literal
# paths through this map keeps the common per-environment case linear; genuine
# regex patterns (e.g. the cloned ``.*`` pattern) still fall back to a full scan.
label_to_index = {label: idx for idx, label in enumerate(body_labels)}
site_bodies: list[int] = []
site_locals: list[list[float]] = []
site_scales: list[tuple[float, float, float]] = []
site_prim_paths: list[str] | None = []
for body_patterns, xform, scale, per_world, env_ids, spec_paths in self._site_specs:
if body_patterns is None:
if per_world:
if NewtonManager._world_xforms is None:
raise RuntimeError(f"FrameView '{self._prim_path}' needs Newton cloned-world transforms.")
world_ids = range(len(NewtonManager._world_xforms)) if env_ids is None else env_ids
for world_id in world_ids:
world_xform = NewtonManager._world_xforms[world_id]
site_bodies.append(WORLD_BODY_INDEX)
site_locals.append([float(v) for v in wp.transform_multiply(world_xform, xform)])
site_scales.append(scale)
site_prim_paths = _extend_prim_paths(site_prim_paths, spec_paths, len(world_ids))
else:
site_bodies.append(WORLD_BODY_INDEX)
site_locals.append([float(v) for v in xform])
site_scales.append(scale)
site_prim_paths = _extend_prim_paths(site_prim_paths, spec_paths, 1)
continue
for index, body_pattern in enumerate(body_patterns):
exact_index = label_to_index.get(body_pattern) if _REGEX_TOKENS.isdisjoint(body_pattern) else None
if exact_index is not None:
matched_indices = [exact_index]
else:
matched_indices, _ = resolve_matching_names(body_pattern, body_labels, raise_when_no_match=False)
if not matched_indices:
raise ValueError(
f"FrameView '{self._prim_path}' body pattern '{body_pattern}' matched no Newton bodies."
)
for body_idx in matched_indices:
site_bodies.append(body_idx)
site_locals.append([float(v) for v in xform])
site_scales.append(scale)
spec_path = None if spec_paths is None else (spec_paths[index],)
site_prim_paths = _extend_prim_paths(site_prim_paths, spec_path, len(matched_indices))
self._create_buffers(site_bodies, site_locals, site_scales, site_prim_paths)
def _create_buffers(
self,
site_bodies: list[int],
site_locals: list[list[float]],
site_scales: list[tuple[float, float, float]],
site_prim_paths: list[str] | None = None,
) -> None:
"""Allocate view buffers from body indices, local transforms, and destination prim paths."""
self._count = len(site_bodies)
paired = site_prim_paths is not None and len(site_prim_paths) == self._count
self._site_prim_paths = site_prim_paths if paired else None
if not paired and self._count:
logger.warning(
f"FrameView '{self._prim_path}' could not pair its sites with destination prims; pose writes"
" update Newton state but are not visible to the renderer."
)
device = self._device
self._site_body = wp.array(site_bodies, dtype=wp.int32, device=device)
self._site_local = wp.array([wp.transform(*x) for x in site_locals], dtype=wp.transformf, device=device)
self._site_xform_scale = wp.array([wp.vec3f(*scale) for scale in site_scales], dtype=wp.vec3f, device=device)
self._site_indices = wp.array(list(range(self._count)), dtype=wp.int32, device=device)
self._pos_buf = wp.zeros(self._count, dtype=wp.vec3f, device=device)
self._quat_buf = wp.zeros(self._count, dtype=wp.vec4f, device=device)
self._local_pos_buf = wp.zeros(self._count, dtype=wp.vec3f, device=device)
self._local_quat_buf = wp.zeros(self._count, dtype=wp.vec4f, device=device)
self._scale_buf = wp.zeros(self._count, dtype=wp.vec3f, device=device)
self._pos_ta = ProxyArray(self._pos_buf)
self._quat_ta = ProxyArray(self._quat_buf)
self._local_pos_ta = ProxyArray(self._local_pos_buf)
self._local_quat_ta = ProxyArray(self._local_quat_buf)
self._scale_ta = ProxyArray(self._site_xform_scale)
def _mirror_to_fabric(self) -> None:
"""Stamp the current site world poses onto the Fabric transforms the renderer reads.
The local matrix is written too, else Newton's body sync forward-propagates ``parent * local``
over the world matrix for frames under a body.
"""
if self._fabric_sel is None and not self._initialize_fabric_mirror():
return
# Bodies sync at render cadence, so after a ``render=False`` step the local derivation below
# would read a stale parent. No-op when clean.
sim = sim_utils.SimulationContext.instance()
from isaaclab_physx.renderers.fabric import FabricBackendCfg # noqa: PLC0415 - requires Kit
fabric = sim.get_or_create_backend(FabricBackendCfg(stage=sim.stage, device=sim.device))
fabric.update_transforms(sim.get_scene_data_provider())
count = self._fabric_sel.count
pos_ta, quat_ta = self._get_world_poses_impl(None)
def launch(kernel, inputs, outputs=()):
wp.launch(kernel, dim=count, inputs=inputs, outputs=outputs, device=self._device)
world_ifa, local_ifa = self._fabric_sel.child_ifas()
parent_ifa = self._fabric_sel.parent_world_ifa()
view_indices = self._fabric_sel.view_indices
mirrored = [self._mirror_positions, self._mirror_orientations]
launch(_gather_mirrored_site_poses, [pos_ta.warp, quat_ta.warp, self._mirror_site_indices], mirrored)
# ``False`` x3: no broadcasting, every site has its own pose. The empty scale array keeps each
# matrix's accumulated-parent scale, as the body sync does for USD scale.
compose = [world_ifa, *mirrored, self._mirror_empty_scales, False, False, False, view_indices]
launch(fabric_utils.compose_indexed_fabric_transforms, compose)
launch(fabric_utils.update_indexed_local_matrix_from_world, [world_ifa, parent_ifa, local_ifa, view_indices])
def _initialize_fabric_mirror(self) -> bool:
"""Build the Fabric selection backing :meth:`_mirror_to_fabric` (a ``False`` result is sticky).
Seeding from USD is off: Newton never writes poses back, so it would reset the prim to its spawn
pose. Coverage can be partial -- Newton clones physics without USD, so a site can outlive its prim.
"""
if self._mirror_disabled or self._site_prim_paths is None or self._count == 0:
self._mirror_disabled = True
return False
try:
selection = FabricXformSelection(
self._site_prim_paths,
self._device,
owner=type(self).__name__,
seed_from_usd=False,
skip_missing_prims=True,
)
except ImportError:
# No Fabric runtime (kitless run): the site state is still correct, nothing consumes it.
self._mirror_disabled = True
logger.info("Fabric runtime unavailable; Newton site poses will not be mirrored to prims.")
return False
count = selection.count
if count == 0:
self._mirror_disabled = True
return False
if count != len(self._site_prim_paths):
logger.info(
f"FrameView '{self._prim_path}' mirrors {count} of {len(self._site_prim_paths)} sites to Fabric;"
" the rest have no prim on the stage (physics-only clones) and nothing to render."
)
# Refreshing the read-write selection is what notifies the renderer; the read-only one would
# land in Fabric but keep showing the old pose.
selection.read_write = True
self._fabric_sel = selection
self._mirror_site_indices = wp.array(selection.kept_indices, dtype=wp.int32, device=self._device)
self._mirror_positions = wp.empty((count, 3), dtype=wp.float32, device=self._device)
self._mirror_orientations = wp.empty((count, 4), dtype=wp.float32, device=self._device)
self._mirror_empty_scales = wp.zeros((0, 0), dtype=wp.float32, device=self._device)
return True
def close(self) -> None:
"""Release the Fabric attributes and the model-ready callback authored by this view."""
handle = self._physics_ready_handle
self._physics_ready_handle = None # cleared first so a repeat close() cannot deregister twice
if handle is not None:
handle.deregister()
if self._fabric_sel is not None:
self._fabric_sel.close()
self._fabric_sel = None
self._mirror_disabled = True
def __del__(self, _sys=sys):
"""Best-effort cleanup when the view is collected without :meth:`close`.
``sys`` is a default argument so it survives module teardown; nothing runs during finalization,
when the tags die with Fabric anyway.
"""
if _sys.is_finalizing() or _sys.meta_path is None:
return
with contextlib.suppress(Exception): # never propagate from __del__
self.close()
@property
def prims(self) -> list:
"""List of USD prims being managed by this view.
Newton site views do not retain USD prim handles.
"""
return self._prims
@property
def count(self) -> int:
"""Number of frames in this view."""
return self._count
@property
def device(self) -> str:
"""Device where arrays are allocated."""
return self._device
# ------------------------------------------------------------------
# Writer factory hooks (pass-through; Newton has no separate Fabric storage)
# ------------------------------------------------------------------
def _make_world_space_writer(self) -> FrameViewWorldSpaceWriter:
return _NewtonWorldSpaceWriter(self)
def _make_local_space_writer(self) -> FrameViewLocalSpaceWriter:
return _NewtonLocalSpaceWriter(self)
# ------------------------------------------------------------------
# Backend hooks
# ------------------------------------------------------------------
def _get_world_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]:
"""Get world-space positions and orientations."""
state = NewtonManager.get_state_0()
site_indices = self._site_indices if indices is None else indices
n = self.count if indices is None else len(indices)
pos_buf = self._pos_buf if indices is None else wp.zeros(n, dtype=wp.vec3f, device=self._device)
quat_buf = self._quat_buf if indices is None else wp.zeros(n, dtype=wp.vec4f, device=self._device)
wp.launch(
_compute_site_world_transforms,
dim=n,
inputs=[state.body_q, self._site_body, self._site_local, site_indices],
outputs=[pos_buf, quat_buf],
device=self._device,
)
if indices is None:
return self._pos_ta, self._quat_ta
return ProxyArray(pos_buf), ProxyArray(quat_buf)
def _apply_world_pose_write(
self,
positions: wp.array | None = None,
orientations: wp.array | None = None,
indices: wp.array | None = None,
) -> None:
"""Set world-space positions and/or orientations."""
if positions is None and orientations is None:
return
state = NewtonManager.get_state_0()
if positions is None or orientations is None:
cur_pos_ta, cur_quat_ta = self._get_world_poses_impl(indices)
if positions is None:
positions = cur_pos_ta.warp
if orientations is None:
orientations = cur_quat_ta.warp
site_indices = self._site_indices if indices is None else indices
n = self.count if indices is None else len(indices)
wp.launch(
_write_site_local_from_world_poses,
dim=n,
inputs=[state.body_q, self._site_body, site_indices, positions, orientations, self._site_local],
device=self._device,
)
def _get_local_poses_impl(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]:
"""Get body-local positions and orientations."""
site_indices = self._site_indices if indices is None else indices
n = self.count if indices is None else len(indices)
pos_buf = self._local_pos_buf if indices is None else wp.zeros(n, dtype=wp.vec3f, device=self._device)
quat_buf = self._local_quat_buf if indices is None else wp.zeros(n, dtype=wp.vec4f, device=self._device)
wp.launch(
_gather_site_local_transforms,
dim=n,
inputs=[self._site_local, site_indices],
outputs=[pos_buf, quat_buf],
device=self._device,
)
if indices is None:
return self._local_pos_ta, self._local_quat_ta
return ProxyArray(pos_buf), ProxyArray(quat_buf)
def _apply_local_pose_write(
self,
translations: wp.array | None = None,
orientations: wp.array | None = None,
indices: wp.array | None = None,
) -> None:
"""Set body-local translations and/or orientations."""
if translations is None and orientations is None:
return
if translations is None or orientations is None:
cur_pos_ta, cur_quat_ta = self._get_local_poses_impl(indices)
if translations is None:
translations = cur_pos_ta.warp
if orientations is None:
orientations = cur_quat_ta.warp
site_indices = self._site_indices if indices is None else indices
n = self.count if indices is None else len(indices)
wp.launch(
_write_site_local_from_local_poses,
dim=n,
inputs=[site_indices, translations, orientations, self._site_local],
device=self._device,
)
# ------------------------------------------------------------------
# Scales
# ------------------------------------------------------------------
def _get_world_scales_impl(self, indices: wp.array | None = None) -> ProxyArray:
"""Get per-site world xform scales.
These are transform scales, matching the USD FrameView scale API. They
are intentionally separate from Newton collision shape geometry sizes.
"""
if indices is None:
return self._scale_ta
n = len(indices)
out = wp.zeros(n, dtype=wp.vec3f, device=self._device)
wp.launch(
_gather_xform_scales,
dim=n,
inputs=[self._site_xform_scale, indices],
outputs=[out],
device=self._device,
)
return ProxyArray(out)
def _get_local_scales_impl(self, indices: wp.array | None = None) -> ProxyArray:
"""Get per-site local xform scales.
These are transform scales, matching the USD FrameView scale API. They
are intentionally separate from Newton collision shape geometry sizes.
"""
return self._get_world_scales_impl(indices)
def _apply_world_scale_write(self, scales: wp.array, indices: wp.array | None = None) -> None:
"""Set per-site world xform scales.
These update transform scale state only; use deprecated ``set_scales`` if
legacy Newton collision shape geometry-scale behavior is required.
"""
if indices is None:
indices = self._site_indices
n = self.count if indices is self._site_indices else len(indices)
wp.launch(
_scatter_xform_scales,
dim=n,
inputs=[indices, scales, self._site_xform_scale],
device=self._device,
)
def _apply_local_scale_write(self, scales: wp.array, indices: wp.array | None = None) -> None:
"""Set per-site local xform scales.
These update transform scale state only; use deprecated ``set_scales`` if
legacy Newton collision shape geometry-scale behavior is required.
"""
self._apply_world_scale_write(scales, indices)
def _get_legacy_shape_scales(self, indices: wp.array | None = None) -> ProxyArray:
"""Get Newton legacy geometry scales from collision shapes."""
model = NewtonManager.get_model()
num_shapes = model.shape_count
site_indices = self._site_indices if indices is None else indices
n = self.count if indices is None else len(indices)
out = wp.zeros(n, dtype=wp.vec3f, device=self._device)
wp.launch(
_gather_shape_scales,
dim=n,
inputs=[model.shape_scale, model.shape_body, self._site_body, site_indices, num_shapes],
outputs=[out],
device=self._device,
)
return ProxyArray(out)
def _set_legacy_shape_scales(self, scales: wp.array, indices: wp.array | None = None) -> None:
"""Set Newton legacy geometry scales on collision shapes."""
model = NewtonManager.get_model()
num_shapes = model.shape_count
site_indices = self._site_indices if indices is None else indices
n = self.count if indices is None else len(indices)
wp.launch(
_scatter_shape_scales,
dim=n,
inputs=[self._site_body, site_indices, scales, model.shape_body, num_shapes, model.shape_scale],
device=self._device,
)
def _get_scales_impl(self, indices: wp.array | None = None) -> ProxyArray:
"""Newton legacy: get_scales returns collision shape geometry scales."""
return self._get_legacy_shape_scales(indices)
def _set_scales_impl(self, scales: wp.array, indices: wp.array | None = None) -> None:
"""Newton legacy: deprecated set_scales writes collision shape geometry scales.
Newton's legacy ``set_scales`` path is *not* routed through the
:class:`FrameViewSpaceWriterBase` API because it targets a different state
(collision-shape geometry sizes) than the transform-scale state that
the writer's :meth:`~FrameViewSpaceWriterBase.set_scales` operates on.
"""
self._set_legacy_shape_scales(scales, indices)
# ----------------------------------------------------------------------
# Pass-through writer classes
# ----------------------------------------------------------------------
class _NewtonWriterMixin:
"""Mirrors the scope's pose writes onto Fabric on exit.
A full-view sync, so it runs only when a pose changed -- and on the exception path too, since the
Newton write is already committed. A mirror failing mid-unwind is logged, never masking the original.
"""
def _enter_impl(self) -> None:
self._wrote_poses = False
def _exit_impl(self, exc_type, exc_val, exc_tb) -> None:
if not self._wrote_poses:
return
try:
self._view._mirror_to_fabric() # type: ignore[attr-defined]
except Exception as mirror_exc: # noqa: BLE001 -- see the docstring
if exc_type is None:
raise
logger.error("Fabric mirror failed while a writer scope unwound: %s", mirror_exc)
class _NewtonWorldSpaceWriter(_NewtonWriterMixin, FrameViewWorldSpaceWriter):
"""Newton world-space writer: pass-through to backend ``_apply_*`` hooks."""
def set_poses(self, positions=None, orientations=None, indices=None) -> None:
if positions is None and orientations is None:
return
self._wrote_poses = True
self._view._apply_world_pose_write(positions, orientations, indices) # type: ignore[attr-defined]
def set_scales(self, scales, indices=None) -> None:
self._view._apply_world_scale_write(scales, indices) # type: ignore[attr-defined]
def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]:
return self._view._get_world_poses_impl(indices) # type: ignore[attr-defined]
def get_scales(self, indices=None) -> ProxyArray:
return self._view._get_world_scales_impl(indices) # type: ignore[attr-defined]
class _NewtonLocalSpaceWriter(_NewtonWriterMixin, FrameViewLocalSpaceWriter):
"""Newton local-space writer: pass-through to backend ``_apply_*`` hooks."""
def set_poses(self, positions=None, orientations=None, indices=None) -> None:
if positions is None and orientations is None:
return
self._wrote_poses = True
self._view._apply_local_pose_write(positions, orientations, indices) # type: ignore[attr-defined]
def set_scales(self, scales, indices=None) -> None:
self._view._apply_local_scale_write(scales, indices) # type: ignore[attr-defined]
def get_poses(self, indices=None) -> tuple[ProxyArray, ProxyArray]:
return self._view._get_local_poses_impl(indices) # type: ignore[attr-defined]
def get_scales(self, indices=None) -> ProxyArray:
return self._view._get_local_scales_impl(indices) # type: ignore[attr-defined]
