diff --git a/docs/source/api/lab/isaaclab.sim.views.rst b/docs/source/api/lab/isaaclab.sim.views.rst index 3a5f9bdecfe9..e06c4e54a246 100644 --- a/docs/source/api/lab/isaaclab.sim.views.rst +++ b/docs/source/api/lab/isaaclab.sim.views.rst @@ -7,11 +7,27 @@ .. autosummary:: - XformPrimView + BaseFrameView + UsdFrameView + FrameView -XForm Prim View +Base Frame View --------------- -.. autoclass:: XformPrimView +.. autoclass:: BaseFrameView + :members: + :show-inheritance: + +USD Frame View +-------------- + +.. autoclass:: UsdFrameView + :members: + :show-inheritance: + +Frame View +---------- + +.. autoclass:: FrameView :members: :show-inheritance: diff --git a/docs/source/migration/migrating_to_isaaclab_3-0.rst b/docs/source/migration/migrating_to_isaaclab_3-0.rst index 6855a4fda28a..37306fcbae39 100644 --- a/docs/source/migration/migrating_to_isaaclab_3-0.rst +++ b/docs/source/migration/migrating_to_isaaclab_3-0.rst @@ -98,6 +98,49 @@ The following classes have been moved to ``isaaclab_physx``: installation steps are required. +Renaming of ``XformPrimView`` to ``FrameView`` +----------------------------------------------- + +Isaac Lab's ``XformPrimView`` and related classes have been renamed to ``FrameView`` to +better reflect their purpose and avoid confusion with Isaac Sim's ``XFormPrim`` class +hierarchy. The old ``XformPrimView`` name is kept as a deprecated alias. + +The rename applies across all backends: + +.. list-table:: + :header-rows: 1 + :widths: 50 50 + + * - Isaac Lab 2.x + - Isaac Lab 3.0 + * - ``BaseXformPrimView`` + - :class:`~isaaclab.sim.views.BaseFrameView` + * - ``UsdXformPrimView`` + - :class:`~isaaclab.sim.views.UsdFrameView` + * - ``XformPrimView`` + - :class:`~isaaclab.sim.views.FrameView` + * - ``FabricXformPrimView`` + - :class:`~isaaclab_physx.sim.views.FabricFrameView` + * - ``NewtonSiteXformPrimView`` + - :class:`~isaaclab_newton.sim.views.NewtonSiteFrameView` + +For most users, the only change needed is updating imports: + +.. code-block:: python + + # Before + from isaaclab.sim.views import XformPrimView + + # After + from isaaclab.sim.views import FrameView + +The :class:`~isaaclab.sim.views.FrameView` factory automatically dispatches to the correct +backend (:class:`~isaaclab_physx.sim.views.FabricFrameView` for PhysX, +:class:`~isaaclab_newton.sim.views.NewtonSiteFrameView` for Newton) based on the active +physics backend. The deprecated ``XformPrimView`` alias continues to work but will be +removed in a future release. + + Unchanged Imports ----------------- diff --git a/docs/source/overview/core-concepts/scene_data_providers.rst b/docs/source/overview/core-concepts/scene_data_providers.rst index 684dfcefcbef..46244317ba2e 100644 --- a/docs/source/overview/core-concepts/scene_data_providers.rst +++ b/docs/source/overview/core-concepts/scene_data_providers.rst @@ -55,7 +55,7 @@ Newton-based visualizers (Newton, Rerun, Viser) require a Newton model/state to The sync pipeline: 1. Reads transforms from PhysX ``RigidBodyView`` (fast tensor API) -2. Falls back to ``XformPrimView`` for bodies not covered by the rigid body view +2. Falls back to :class:`~isaaclab.sim.views.FrameView` for bodies not covered by the rigid body view 3. Converts and writes merged poses into the Newton state via Warp kernels Newton Scene Data Provider diff --git a/docs/source/tutorials/04_sensors/add_sensors_on_robot.rst b/docs/source/tutorials/04_sensors/add_sensors_on_robot.rst index 3d9f40667b62..85383a876ed4 100644 --- a/docs/source/tutorials/04_sensors/add_sensors_on_robot.rst +++ b/docs/source/tutorials/04_sensors/add_sensors_on_robot.rst @@ -93,11 +93,11 @@ Height scanner The height-scanner is implemented as a virtual sensor using the NVIDIA Warp ray-casting kernels. Through the :class:`sensors.RayCasterCfg`, we can specify the pattern of rays to cast and the -meshes against which to cast the rays. Since they are virtual sensors, there is no corresponding -prim created in the scene for them. Instead they are attached to a prim in the scene, which is -used to specify the location of the sensor. +meshes against which to cast the rays. By default, :attr:`~sensors.RayCasterCfg.spawn` creates +a plain USD Xform at :attr:`~sensors.RayCasterCfg.prim_path` to serve as the sensor's +attachment frame, similar to how :class:`sensors.CameraCfg` spawns a Camera prim. -For this tutorial, the ray-cast based height scanner is attached to the base frame of the robot. +For this tutorial, the ray-cast based height scanner is attached under the base frame of the robot. The pattern of rays is specified using the :attr:`~sensors.RayCasterCfg.pattern` attribute. For a uniform grid pattern, we specify the pattern using :class:`~sensors.patterns.GridPatternCfg`. Since we only care about the height information, we do not need to consider the roll and pitch diff --git a/scripts/benchmarks/benchmark_view_comparison.py b/scripts/benchmarks/benchmark_view_comparison.py index 8f2b60c49077..a637f687803e 100644 --- a/scripts/benchmarks/benchmark_view_comparison.py +++ b/scripts/benchmarks/benchmark_view_comparison.py @@ -3,54 +3,50 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Benchmark script comparing XformPrimView vs PhysX RigidBodyView for transform operations. +"""Benchmark script comparing FrameView backends and PhysX RigidBodyView. -This script tests the performance of batched transform operations using: +Compares batched transform operation performance across: -- Isaac Lab's XformPrimView (USD-based) -- Isaac Lab's XformPrimView (Fabric-based) -- PhysX RigidBodyView (PhysX tensors-based, as used in RigidObject) - -Note: - XformPrimView operates on USD attributes directly (useful for non-physics prims), - or on Fabric attributes when Fabric is enabled. - while RigidBodyView requires rigid body physics components and operates on PhysX tensors. - This benchmark helps understand the performance trade-offs between the two approaches. +- **USD** (baseline): Isaac Lab's FrameView via USD XformCache +- **Fabric**: Isaac Lab's FrameView via Fabric GPU arrays +- **Newton**: Isaac Lab's Newton FrameView via Warp site kernels +- **PhysX**: PhysX RigidBodyView via PhysX tensor API (reference) Usage: - # Basic benchmark + # All backends ./isaaclab.sh -p scripts/benchmarks/benchmark_view_comparison.py --num_envs 1024 --device cuda:0 --headless - # With profiling enabled (for snakeviz visualization) - ./isaaclab.sh -p scripts/benchmarks/benchmark_view_comparison.py --num_envs 1024 --profile --headless + # Select specific backends + ./isaaclab.sh -p scripts/benchmarks/benchmark_view_comparison.py --backends usd fabric newton --headless - # Then visualize with snakeviz: - snakeviz profile_results/xform_view_benchmark.prof - snakeviz profile_results/physx_view_benchmark.prof + # With profiling + ./isaaclab.sh -p scripts/benchmarks/benchmark_view_comparison.py --num_envs 1024 --profile --headless """ from __future__ import annotations -"""Launch Isaac Sim Simulator first.""" - import argparse from isaaclab.app import AppLauncher -# parse the arguments -args_cli = argparse.Namespace() - -parser = argparse.ArgumentParser(description="Benchmark XformPrimView vs PhysX RigidBodyView performance.") +parser = argparse.ArgumentParser(description="Benchmark FrameView backends and PhysX RigidBodyView.") parser.add_argument("--num_envs", type=int, default=1000, help="Number of environments to simulate.") parser.add_argument("--num_iterations", type=int, default=50, help="Number of iterations for each test.") +parser.add_argument( + "--backends", + nargs="+", + default=["usd", "fabric", "newton", "physx"], + choices=["usd", "fabric", "newton", "physx"], + help="Backends to benchmark. Default: all four.", +) parser.add_argument( "--profile", action="store_true", help="Enable profiling with cProfile. Results saved as .prof files for snakeviz visualization.", ) parser.add_argument( - "--profile-dir", + "--profile_dir", type=str, default="./profile_results", help="Directory to save profile results. Default: ./profile_results", @@ -59,7 +55,6 @@ AppLauncher.add_app_launcher_args(parser) args_cli = parser.parse_args() -# launch omniverse app app_launcher = AppLauncher(args_cli) simulation_app = app_launcher.app @@ -69,40 +64,40 @@ import time import torch +import warp as wp + +from pxr import Gf import isaaclab.sim as sim_utils -from isaaclab.sim.views import XformPrimView +from isaaclab.sim.views import FrameView + +try: + from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + from isaaclab_newton.sim.views import NewtonSiteFrameView + + HAS_NEWTON = True +except ImportError: + HAS_NEWTON = False + + +# ------------------------------------------------------------------ +# Benchmark functions +# ------------------------------------------------------------------ @torch.no_grad() -def benchmark_view(view_type: str, num_iterations: int) -> tuple[dict[str, float], dict[str, torch.Tensor]]: - """Benchmark the specified view class. - - Args: - view_type: Type of view to benchmark ("xform", "xform_fabric", or "physx"). - num_iterations: Number of iterations to run. - - Returns: - A tuple of (timing_results, computed_results) where: - - timing_results: Dictionary containing timing results for various operations - - computed_results: Dictionary containing the computed values for validation - """ +def benchmark_usd_or_fabric(view_type: str, num_iterations: int) -> dict[str, float]: + """Benchmark USD or Fabric FrameView.""" timing_results = {} - computed_results = {} - # Setup scene print(" Setting up scene") - # Clear stage sim_utils.create_new_stage() - # Create simulation context start_time = time.perf_counter() - sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device, use_fabric=(view_type == "xform_fabric")) + sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device, use_fabric=(view_type == "fabric")) sim = sim_utils.SimulationContext(sim_cfg) stage = sim_utils.get_current_stage() + print(f" SimulationContext: {time.perf_counter() - start_time:.4f}s") - print(f" Time taken to create simulation context: {time.perf_counter() - start_time:.4f} seconds") - - # create a rigid object object_cfg = sim_utils.ConeCfg( radius=0.15, height=0.5, @@ -111,222 +106,223 @@ def benchmark_view(view_type: str, num_iterations: int) -> tuple[dict[str, float collision_props=sim_utils.CollisionPropertiesCfg(), visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0)), ) - # Create prims for i in range(args_cli.num_envs): sim_utils.create_prim(f"/World/Env_{i}", "Xform", stage=stage, translation=(i * 2.0, 0.0, 0.0)) object_cfg.func(f"/World/Env_{i}/Object", object_cfg, translation=(0.0, 0.0, 1.0)) + prim = stage.DefinePrim(f"/World/Env_{i}/Object/Sensor", "Xform") + sim_utils.standardize_xform_ops(prim) + prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(0.1, 0.0, 0.05)) + prim.GetAttribute("xformOp:orient").Set(Gf.Quatd(1.0, 0.0, 0.0, 0.0)) - # Play simulation sim.reset() - # Pattern to match all prims - pattern = "/World/Env_.*/Object" if view_type == "xform" else "/World/Env_*/Object" - print(f" Pattern: {pattern}") + pattern = "/World/Env_.*/Object/Sensor" - # Create view based on type start_time = time.perf_counter() - if view_type == "xform": - view = XformPrimView(pattern, device=args_cli.device, validate_xform_ops=False) - num_prims = view.count - view_name = "XformPrimView (USD)" - elif view_type == "xform_fabric": - if "cuda" not in args_cli.device: - raise ValueError("Fabric backend requires CUDA. Please use --device cuda:0 for this benchmark.") - view = XformPrimView(pattern, device=args_cli.device, validate_xform_ops=False) - num_prims = view.count - view_name = "XformPrimView (Fabric)" - else: # physx - physics_sim_view = sim.physics_manager.get_physics_sim_view() - view = physics_sim_view.create_rigid_body_view(pattern) - num_prims = view.count - view_name = "PhysX RigidBodyView" + if view_type == "fabric" and "cuda" not in args_cli.device: + raise ValueError("Fabric backend requires CUDA.") + view = FrameView(pattern, device=args_cli.device, validate_xform_ops=False) + num_prims = view.count timing_results["init"] = time.perf_counter() - start_time - # prepare indices for benchmarking - all_indices = torch.arange(num_prims, device=args_cli.device) - - print(f" {view_name} managing {num_prims} prims") - - # Fabric is write-first: initialize it to match USD before benchmarking reads. - if view_type == "xform_fabric" and num_prims > 0: - init_positions = torch.zeros((num_prims, 3), dtype=torch.float32, device=args_cli.device) - init_positions[:, 0] = 2.0 * torch.arange(num_prims, device=args_cli.device, dtype=torch.float32) - init_positions[:, 2] = 1.0 - init_orientations = torch.tensor( - [[1.0, 0.0, 0.0, 0.0]] * num_prims, dtype=torch.float32, device=args_cli.device + + print(f" FrameView ({view_type.upper()}) managing {num_prims} prims") + + positions, orientations = view.get_world_poses() + + _run_pose_benchmarks(view, num_prims, num_iterations, timing_results, positions, orientations) + + sim.clear_instance() + return timing_results + + +@torch.no_grad() +def benchmark_newton(num_iterations: int) -> dict[str, float]: + """Benchmark Newton FrameView.""" + from isaaclab.assets import RigidObjectCfg + from isaaclab.scene import InteractiveScene, InteractiveSceneCfg + from isaaclab.sim import SimulationCfg, build_simulation_context + from isaaclab.utils import configclass + + timing_results = {} + + @configclass + class _SceneCfg(InteractiveSceneCfg): + cube: RigidObjectCfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube", + spawn=sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), ) - view.set_world_poses(init_positions, init_orientations) - # Benchmark get_world_poses + print(" Setting up Newton scene") + newton_cfg = SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg()), device=args_cli.device) start_time = time.perf_counter() - for _ in range(num_iterations): - if view_type in ("xform", "xform_fabric"): - positions, orientations = view.get_world_poses() - else: # physx - transforms = view.get_transforms() - positions = transforms[:, :3] - orientations = transforms[:, 3:7] - timing_results["get_world_poses"] = (time.perf_counter() - start_time) / num_iterations + ctx = build_simulation_context(device=args_cli.device, sim_cfg=newton_cfg, add_ground_plane=True) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None + InteractiveScene(_SceneCfg(num_envs=args_cli.num_envs, env_spacing=2.0)) - # Store initial world poses - computed_results["initial_world_positions"] = positions.clone() - computed_results["initial_world_orientations"] = orientations.clone() + stage = sim_utils.get_current_stage() + for i in range(args_cli.num_envs): + prim = stage.DefinePrim(f"/World/envs/env_{i}/Cube/Sensor", "Xform") + sim_utils.standardize_xform_ops(prim) + prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(0.1, 0.0, 0.05)) + prim.GetAttribute("xformOp:orient").Set(Gf.Quatd(1.0, 0.0, 0.0, 0.0)) + + sim.reset() + print(f" Newton scene setup: {time.perf_counter() - start_time:.4f}s") - # Benchmark set_world_poses - new_positions = positions.clone() - new_positions[:, 2] += 0.5 start_time = time.perf_counter() - for _ in range(num_iterations): - if view_type in ("xform", "xform_fabric"): - view.set_world_poses(new_positions, orientations) - else: # physx - new_transforms = torch.cat([new_positions, orientations], dim=-1) - view.set_transforms(new_transforms, indices=all_indices) - timing_results["set_world_poses"] = (time.perf_counter() - start_time) / num_iterations + view = NewtonSiteFrameView("/World/envs/env_.*/Cube/Sensor", device=args_cli.device) + num_prims = view.count + timing_results["init"] = time.perf_counter() - start_time - # Get world poses after setting to verify - if view_type in ("xform", "xform_fabric"): - positions_after_set, orientations_after_set = view.get_world_poses() - else: # physx - transforms_after = view.get_transforms() - positions_after_set = transforms_after[:, :3] - orientations_after_set = transforms_after[:, 3:7] - computed_results["world_positions_after_set"] = positions_after_set.clone() - computed_results["world_orientations_after_set"] = orientations_after_set.clone() - - # close simulation - sim.clear_instance() + print(f" Newton FrameView managing {num_prims} prims") - return timing_results, computed_results + positions, orientations = view.get_world_poses() + _run_pose_benchmarks(view, num_prims, num_iterations, timing_results, positions, orientations) -def compare_results( - results_dict: dict[str, dict[str, torch.Tensor]], tolerance: float = 1e-4 -) -> dict[str, dict[str, dict[str, float]]]: - """Compare computed results across implementations. + ctx.__exit__(None, None, None) + return timing_results - Args: - results_dict: Dictionary mapping implementation names to their computed values. - tolerance: Tolerance for numerical comparison. - Returns: - Nested dictionary: {comparison_pair: {metric: {stats}}} - """ - comparison_stats = {} - impl_names = list(results_dict.keys()) +@torch.no_grad() +def benchmark_physx(num_iterations: int) -> dict[str, float]: + """Benchmark PhysX RigidBodyView.""" + timing_results = {} - # Compare each pair of implementations - for i, impl1 in enumerate(impl_names): - for impl2 in impl_names[i + 1 :]: - pair_key = f"{impl1}_vs_{impl2}" - comparison_stats[pair_key] = {} + print(" Setting up scene") + sim_utils.create_new_stage() + start_time = time.perf_counter() + sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device, use_fabric=False) + sim = sim_utils.SimulationContext(sim_cfg) + stage = sim_utils.get_current_stage() + print(f" SimulationContext: {time.perf_counter() - start_time:.4f}s") - computed1 = results_dict[impl1] - computed2 = results_dict[impl2] + object_cfg = sim_utils.ConeCfg( + radius=0.15, + height=0.5, + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0)), + ) + for i in range(args_cli.num_envs): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", stage=stage, translation=(i * 2.0, 0.0, 0.0)) + object_cfg.func(f"/World/Env_{i}/Object", object_cfg, translation=(0.0, 0.0, 1.0)) - for key in computed1.keys(): - if key not in computed2: - continue + sim.reset() - val1 = computed1[key] - val2 = computed2[key] + pattern = "/World/Env_*/Object" + start_time = time.perf_counter() + physics_sim_view = sim.physics_manager.get_physics_sim_view() + view = physics_sim_view.create_rigid_body_view(pattern) + num_prims = view.count + timing_results["init"] = time.perf_counter() - start_time - # Skip zero tensors (not applicable tests) - if torch.all(val1 == 0) or torch.all(val2 == 0): - continue + print(f" PhysX RigidBodyView managing {num_prims} prims") - # Compute differences - diff = torch.abs(val1 - val2) - max_diff = torch.max(diff).item() - mean_diff = torch.mean(diff).item() + all_indices = wp.from_torch(torch.arange(num_prims, dtype=torch.int32, device=args_cli.device)) - # Check if within tolerance - all_close = torch.allclose(val1, val2, atol=tolerance, rtol=0) + transforms = view.get_transforms() + transforms_t = wp.to_torch(transforms) if isinstance(transforms, wp.array) else transforms + positions_t = transforms_t[:, :3] + orientations_t = transforms_t[:, 3:7] - comparison_stats[pair_key][key] = { - "max_diff": max_diff, - "mean_diff": mean_diff, - "all_close": all_close, - } + start_time = time.perf_counter() + for _ in range(num_iterations): + transforms = view.get_transforms() + timing_results["get_world_poses"] = (time.perf_counter() - start_time) / num_iterations - return comparison_stats + new_positions = positions_t.clone() + new_positions[:, 2] += 0.5 + expected_positions = new_positions.clone() + new_transforms = wp.from_torch(torch.cat([new_positions, orientations_t], dim=-1).contiguous()) + start_time = time.perf_counter() + for _ in range(num_iterations): + view.set_transforms(new_transforms, indices=all_indices) + timing_results["set_world_poses"] = (time.perf_counter() - start_time) / num_iterations + transforms_after = view.get_transforms() + ta = wp.to_torch(transforms_after) if isinstance(transforms_after, wp.array) else transforms_after + pos_ok = torch.allclose(ta[:, :3], expected_positions, atol=1e-4, rtol=0) + quat_ok = torch.allclose(ta[:, 3:7], orientations_t, atol=1e-4, rtol=0) + if pos_ok and quat_ok: + print(" Round-trip verification: PASS") + else: + pos_diff = (ta[:, :3] - expected_positions).abs().max().item() + quat_diff = (ta[:, 3:7] - orientations_t).abs().max().item() + print(f" Round-trip verification: FAIL (pos max_diff={pos_diff:.6e}, quat max_diff={quat_diff:.6e})") -def print_comparison_results(comparison_stats: dict[str, dict[str, dict[str, float]]], tolerance: float): - """Print comparison results. + sim.clear_instance() + return timing_results + + +def _run_pose_benchmarks( + view, + num_prims: int, + num_iterations: int, + timing_results: dict, + positions: wp.array, + orientations: wp.array, +): + """Shared benchmark loop for get/set world poses on any FrameView.""" + start_time = time.perf_counter() + for _ in range(num_iterations): + view.get_world_poses() + timing_results["get_world_poses"] = (time.perf_counter() - start_time) / num_iterations - Args: - comparison_stats: Nested dictionary containing comparison statistics. - tolerance: Tolerance used for comparison. - """ - for pair_key, pair_stats in comparison_stats.items(): - if not pair_stats: # Skip if no comparable results - continue + new_positions = wp.clone(positions) + new_positions_t = wp.to_torch(new_positions) + new_positions_t[:, 2] += 0.5 + expected_positions = new_positions_t.clone() - # Format the pair key for display - impl1, impl2 = pair_key.split("_vs_") - display_impl1 = impl1.replace("_", " ").title() - display_impl2 = impl2.replace("_", " ").title() - comparison_title = f"{display_impl1} vs {display_impl2}" - - # Check if all results match - all_match = all(stats["all_close"] for stats in pair_stats.values()) - - if all_match: - # Compact output when everything matches - print("\n" + "=" * 100) - print(f"RESULT COMPARISON: {comparison_title}") - print("=" * 100) - print(f"✓ All computed values match within tolerance ({tolerance})") - print("=" * 100) - else: - # Detailed output when there are mismatches - print("\n" + "=" * 100) - print(f"RESULT COMPARISON: {comparison_title}") - print("=" * 100) - print(f"{'Computed Value':<40} {'Max Diff':<15} {'Mean Diff':<15} {'Match':<10}") - print("-" * 100) - - for key, stats in pair_stats.items(): - # Format the key for display - display_key = key.replace("_", " ").title() - match_str = "✓ Yes" if stats["all_close"] else "✗ No" - - print(f"{display_key:<40} {stats['max_diff']:<15.6e} {stats['mean_diff']:<15.6e} {match_str:<10}") - - print("=" * 100) - print(f"\n✗ Some results differ beyond tolerance ({tolerance})") - print(f" This may indicate implementation differences between {display_impl1} and {display_impl2}") + start_time = time.perf_counter() + for _ in range(num_iterations): + view.set_world_poses(new_positions, orientations) + timing_results["set_world_poses"] = (time.perf_counter() - start_time) / num_iterations + + ret_pos, ret_quat = view.get_world_poses() + ret_pos_t = wp.to_torch(ret_pos) + ret_quat_t = wp.to_torch(ret_quat) + ori_t = wp.to_torch(orientations) + + pos_ok = torch.allclose(ret_pos_t, expected_positions, atol=1e-4, rtol=0) + quat_ok = torch.allclose(ret_quat_t, ori_t, atol=1e-4, rtol=0) + if pos_ok and quat_ok: + print(" Round-trip verification: PASS") + else: + pos_diff = (ret_pos_t - expected_positions).abs().max().item() + quat_diff = (ret_quat_t - ori_t).abs().max().item() + print(f" Round-trip verification: FAIL (pos max_diff={pos_diff:.6e}, quat max_diff={quat_diff:.6e})") - print() + +# ------------------------------------------------------------------ +# Reporting +# ------------------------------------------------------------------ def print_results(results_dict: dict[str, dict[str, float]], num_prims: int, num_iterations: int): - """Print benchmark results in a formatted table. - - Args: - results_dict: Dictionary mapping implementation names to their timing results. - num_prims: Number of prims tested. - num_iterations: Number of iterations run. - """ - print("\n" + "=" * 100) + """Print benchmark results in a formatted table.""" + print("\n" + "=" * 120) print(f"BENCHMARK RESULTS: {num_prims} prims, {num_iterations} iterations") - print("=" * 100) + print("=" * 120) impl_names = list(results_dict.keys()) - # Format names for display - display_names = [name.replace("_", " ").title() for name in impl_names] - - # Calculate column width - col_width = 20 + display_names = {n: n.replace("_", " ").title() for n in impl_names} + col_width = 22 - # Print header - header = f"{'Operation':<30}" - for display_name in display_names: - header += f" {display_name + ' (ms)':<{col_width}}" + header = f"{'Operation':<25}" + for name in impl_names: + header += f" {display_names[name] + ' (ms)':>{col_width}}" print(header) - print("-" * 100) + print("-" * 120) - # Print each operation operations = [ ("Initialization", "init"), ("Get World Poses", "get_world_poses"), @@ -334,168 +330,117 @@ def print_results(results_dict: dict[str, dict[str, float]], num_prims: int, num ] for op_name, op_key in operations: - row = f"{op_name:<30}" - for impl_name in impl_names: - impl_time = results_dict[impl_name].get(op_key, 0) * 1000 # Convert to ms - row += f" {impl_time:>{col_width - 1}.4f}" + row = f"{op_name:<25}" + for name in impl_names: + val = results_dict[name].get(op_key, 0) * 1000 + row += f" {val:>{col_width}.4f}" print(row) - print("=" * 100) - - # Calculate and print total time (excluding N/A operations) - total_row = f"{'Total Time':<30}" - for impl_name in impl_names: - if impl_name == "physx_view": - # Exclude local pose operations for PhysX - total_time = ( - results_dict[impl_name].get("init", 0) * 1000 - + results_dict[impl_name].get("get_world_poses", 0) * 1000 - + results_dict[impl_name].get("set_world_poses", 0) * 1000 - ) - else: - total_time = sum(results_dict[impl_name].values()) * 1000 - total_row += f" {total_time:>{col_width - 1}.4f}" - print(f"\n{total_row}") - - # Calculate speedups relative to XformPrimView (USD baseline) - if "xform_view" in impl_names: - print("\n" + "=" * 100) - print("SPEEDUP vs XformPrimView (USD)") - print("=" * 100) - print(f"{'Operation':<30}", end="") - for impl_name, display_name in zip(impl_names, display_names): - if impl_name != "xform_view": - print(f" {display_name + ' Speedup':<{col_width}}", end="") - print() - print("-" * 100) - - xform_results = results_dict["xform_view"] + print("=" * 120) + + total_row = f"{'Total':<25}" + for name in impl_names: + total = sum(results_dict[name].values()) * 1000 + total_row += f" {total:>{col_width}.4f}" + print(total_row) + + baseline = "usd" + if baseline in results_dict and len(impl_names) > 1: + print("\n" + "=" * 120) + print(f"SPEEDUP vs {display_names[baseline]}") + print("=" * 120) + header = f"{'Operation':<25}" + for name in impl_names: + if name != baseline: + header += f" {display_names[name]:>{col_width}}" + print(header) + print("-" * 120) + + base = results_dict[baseline] for op_name, op_key in operations: - print(f"{op_name:<30}", end="") - xform_time = xform_results.get(op_key, 0) - for impl_name, display_name in zip(impl_names, display_names): - if impl_name != "xform_view": - impl_time = results_dict[impl_name].get(op_key, 0) - if xform_time > 0 and impl_time > 0: - speedup = impl_time / xform_time - print(f" {speedup:>{col_width - 1}.2f}x", end="") + row = f"{op_name:<25}" + base_t = base.get(op_key, 0) + for name in impl_names: + if name != baseline: + impl_t = results_dict[name].get(op_key, 0) + if base_t > 0 and impl_t > 0: + row += f" {base_t / impl_t:>{col_width}.2f}x" else: - print(f" {'N/A':>{col_width}}", end="") - print() + row += f" {'N/A':>{col_width}}" + print(row) + print("=" * 120) - # Overall speedup (only world pose operations) - print("=" * 100) - print(f"{'Overall Speedup (World Ops)':<30}", end="") - total_xform = ( - xform_results.get("init", 0) - + xform_results.get("get_world_poses", 0) - + xform_results.get("set_world_poses", 0) - ) - for impl_name, display_name in zip(impl_names, display_names): - if impl_name != "xform_view": - total_impl = ( - results_dict[impl_name].get("init", 0) - + results_dict[impl_name].get("get_world_poses", 0) - + results_dict[impl_name].get("set_world_poses", 0) - ) - if total_xform > 0 and total_impl > 0: - overall_speedup = total_impl / total_xform - print(f" {overall_speedup:>{col_width - 1}.2f}x", end="") - else: - print(f" {'N/A':>{col_width}}", end="") - print() - - print("\n" + "=" * 100) print("\nNotes:") print(" - Times are averaged over all iterations") - print(" - Speedup = (Implementation time) / (XformPrimView USD time)") - print(" - Speedup > 1.0 means USD XformPrimView is faster") - print(" - Speedup < 1.0 means the implementation is faster than USD") - print(" - PhysX View requires rigid body physics components") - print(" - XformPrimView works with any Xform prim (physics or non-physics)") - print(" - PhysX View does not support local pose operations directly") + print(" - Speedup > 1.0 means faster than USD baseline") + print(" - PhysX RigidBodyView requires rigid body physics; FrameView works with any Xformable prim") print() +# ------------------------------------------------------------------ +# Main +# ------------------------------------------------------------------ + + def main(): - """Main benchmark function.""" - print("=" * 100) - print("View Comparison Benchmark - XformPrimView vs PhysX RigidBodyView") - print("=" * 100) - print("Configuration:") - print(f" Number of environments: {args_cli.num_envs}") - print(f" Iterations per test: {args_cli.num_iterations}") - print(f" Device: {args_cli.device}") - print(f" Profiling: {'Enabled' if args_cli.profile else 'Disabled'}") - if args_cli.profile: - print(f" Profile directory: {args_cli.profile_dir}") + print("=" * 120) + print("FrameView Benchmark: USD vs Fabric vs Newton vs PhysX") + print("=" * 120) + print(f" Environments: {args_cli.num_envs}") + print(f" Iterations: {args_cli.num_iterations}") + print(f" Device: {args_cli.device}") + print(f" Backends: {', '.join(args_cli.backends)}") print() - # Create profile directory if profiling is enabled if args_cli.profile: import os os.makedirs(args_cli.profile_dir, exist_ok=True) - # Dictionary to store all results - all_timing_results = {} - all_computed_results = {} + all_timing = {} profile_files = {} - # Implementations to benchmark - implementations = [ - ("xform_view", "XformPrimView (USD)", "xform"), - ("xform_fabric_view", "XformPrimView (Fabric)", "xform_fabric"), - ("physx_view", "PhysX RigidBodyView", "physx"), - ] + dispatch = { + "usd": ("usd", "FrameView (USD)", lambda n: benchmark_usd_or_fabric("usd", n)), + "fabric": ("fabric", "FrameView (Fabric)", lambda n: benchmark_usd_or_fabric("fabric", n)), + "newton": ("newton", "FrameView (Newton)", lambda n: benchmark_newton(n)), + "physx": ("physx", "PhysX RigidBodyView", lambda n: benchmark_physx(n)), + } - # Benchmark each implementation - for impl_key, impl_name, view_type in implementations: - print(f"Benchmarking {impl_name}...") + for backend in args_cli.backends: + if backend == "newton" and not HAS_NEWTON: + print(f"Skipping {backend}: isaaclab_newton not installed") + continue + + key, display_name, bench_fn = dispatch[backend] + print(f"Benchmarking {display_name}...") if args_cli.profile: profiler = cProfile.Profile() profiler.enable() - timing, computed = benchmark_view(view_type=view_type, num_iterations=args_cli.num_iterations) + timing = bench_fn(args_cli.num_iterations) if args_cli.profile: profiler.disable() - profile_file = f"{args_cli.profile_dir}/{impl_key}_benchmark.prof" - profiler.dump_stats(profile_file) - profile_files[impl_key] = profile_file - print(f" Profile saved to: {profile_file}") - - all_timing_results[impl_key] = timing - all_computed_results[impl_key] = computed + pf = f"{args_cli.profile_dir}/{key}_benchmark.prof" + profiler.dump_stats(pf) + profile_files[key] = pf + print(f" Profile saved to: {pf}") - print(" Done!") - print() + all_timing[key] = timing + print(" Done!\n") - # Print timing results - print_results(all_timing_results, args_cli.num_envs, args_cli.num_iterations) + print_results(all_timing, args_cli.num_envs, args_cli.num_iterations) - # Compare computed results - print("\nComparing computed results across implementations...") - comparison_stats = compare_results(all_computed_results, tolerance=1e-4) - print_comparison_results(comparison_stats, tolerance=1e-4) - - # Print profiling instructions if enabled if args_cli.profile: print("\n" + "=" * 100) print("PROFILING RESULTS") print("=" * 100) - print("Profile files have been saved. To visualize with snakeviz, run:") - for impl_key, profile_file in profile_files.items(): - impl_display = impl_key.replace("_", " ").title() - print(f" # {impl_display}") - print(f" snakeviz {profile_file}") - print("\nAlternatively, use pstats to analyze in terminal:") - print(" python -m pstats ") - print("=" * 100) + for key, pf in profile_files.items(): + print(f" snakeviz {pf}") print() - # Clean up sim_utils.SimulationContext.clear_instance() diff --git a/scripts/benchmarks/benchmark_xform_prim_view.py b/scripts/benchmarks/benchmark_xform_prim_view.py index e76796e20271..b682c03f71fc 100644 --- a/scripts/benchmarks/benchmark_xform_prim_view.py +++ b/scripts/benchmarks/benchmark_xform_prim_view.py @@ -3,59 +3,35 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Benchmark script comparing XformPrimView implementations across different APIs. +"""Benchmark script comparing FrameView implementations across backends. -This script tests the performance of batched transform operations using: -- Isaac Lab's XformPrimView implementation with USD backend -- Isaac Lab's XformPrimView implementation with Fabric backend -- Isaac Sim's XformPrimView implementation (legacy) -- Isaac Sim Experimental's XformPrim implementation (latest) +Compares batched transform operation performance across: +- Isaac Lab FrameView (USD backend) -- baseline +- Isaac Lab FrameView (Fabric backend) +- Isaac Lab FrameView (Newton backend) Usage: - # Basic benchmark (all APIs) ./isaaclab.sh -p scripts/benchmarks/benchmark_xform_prim_view.py --num_envs 1024 --device cuda:0 --headless - # With profiling enabled (for snakeviz visualization) + # With profiling ./isaaclab.sh -p scripts/benchmarks/benchmark_xform_prim_view.py --num_envs 1024 --profile --headless - - # Then visualize with snakeviz: - snakeviz profile_results/isaaclab_usd_benchmark.prof - snakeviz profile_results/isaaclab_fabric_benchmark.prof - snakeviz profile_results/isaacsim_benchmark.prof - snakeviz profile_results/isaacsim_exp_benchmark.prof """ from __future__ import annotations -"""Launch Isaac Sim Simulator first.""" - import argparse from isaaclab.app import AppLauncher -# parse the arguments -args_cli = argparse.Namespace() - -parser = argparse.ArgumentParser(description="This script can help you benchmark the performance of XformPrimView.") - +parser = argparse.ArgumentParser(description="Benchmark FrameView performance across backends.") parser.add_argument("--num_envs", type=int, default=100, help="Number of environments to simulate.") parser.add_argument("--num_iterations", type=int, default=50, help="Number of iterations for each test.") -parser.add_argument( - "--profile", - action="store_true", - help="Enable profiling with cProfile. Results saved as .prof files for snakeviz visualization.", -) -parser.add_argument( - "--profile-dir", - type=str, - default="./profile_results", - help="Directory to save profile results. Default: ./profile_results", -) +parser.add_argument("--profile", action="store_true", help="Enable cProfile profiling.") +parser.add_argument("--profile_dir", type=str, default="./profile_results", help="Directory for .prof files.") AppLauncher.add_app_launcher_args(parser) args_cli = parser.parse_args() -# launch omniverse app app_launcher = AppLauncher(args_cli) simulation_app = app_launcher.app @@ -66,402 +42,231 @@ from typing import Literal import torch +import warp as wp +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.sim.views import NewtonSiteFrameView +from isaaclab_physx.sim.views import FabricFrameView -from isaacsim.core.prims import XFormPrim as IsaacSimXformPrimView -from isaacsim.core.utils.extensions import enable_extension - -# compare against latest Isaac Sim implementation -enable_extension("isaacsim.core.experimental.prims") -from isaacsim.core.experimental.prims import XformPrim as IsaacSimExperimentalXformPrimView +from pxr import Gf import isaaclab.sim as sim_utils -from isaaclab.sim.views import XformPrimView as IsaacLabXformPrimView +from isaaclab.assets import RigidObjectCfg +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.sim.views import UsdFrameView +from isaaclab.utils import configclass + + +@configclass +class _NewtonSceneCfg(InteractiveSceneCfg): + cube: RigidObjectCfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Object", + spawn=sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), + ) + + +# ------------------------------------------------------------------ +# Benchmark +# ------------------------------------------------------------------ @torch.no_grad() -def benchmark_xform_prim_view( # noqa: C901 - api: Literal["isaaclab-usd", "isaaclab-fabric", "isaacsim-usd", "isaacsim-fabric", "isaacsim-exp"], +def benchmark_frame_view( # noqa: C901 + api: Literal["isaaclab-usd", "isaaclab-fabric", "isaaclab-newton-site"], num_iterations: int, ) -> tuple[dict[str, float], dict[str, torch.Tensor]]: - """Benchmark the Xform view class from Isaac Lab, Isaac Sim, or Isaac Sim Experimental. - - Args: - api: Which API to benchmark: - - "isaaclab-usd": Isaac Lab XformPrimView with USD backend - - "isaaclab-fabric": Isaac Lab XformPrimView with Fabric backend - - "isaacsim-usd": Isaac Sim legacy XformPrimView with USD (usd=True) - - "isaacsim-fabric": Isaac Sim legacy XformPrimView with Fabric (usd=False) - - "isaacsim-exp": Isaac Sim Experimental XformPrim - num_iterations: Number of iterations to run. - - Returns: - A tuple of (timing_results, computed_results) where: - - timing_results: Dictionary containing timing results for various operations - - computed_results: Dictionary containing the computed values for validation - """ - timing_results = {} - computed_results = {} - - # Setup scene + """Benchmark get/set world/local poses for the given FrameView backend.""" + timing_results: dict[str, float] = {} + computed_results: dict[str, torch.Tensor] = {} + device = args_cli.device + num_envs = args_cli.num_envs + + # -- Scene setup (backend-specific) -------------------------------- + print(" Setting up scene") - # Clear stage - sim_utils.create_new_stage() - # Create simulation context - start_time = time.perf_counter() - sim_cfg = sim_utils.SimulationCfg( - dt=0.01, - device=args_cli.device, - use_fabric=api in ("isaaclab-fabric", "isaacsim-fabric"), - ) - sim = sim_utils.SimulationContext(sim_cfg) - stage = sim_utils.get_current_stage() - - print(f" Time taken to create simulation context: {time.perf_counter() - start_time} seconds") - - # Create prims - prim_paths = [] - for i in range(args_cli.num_envs): - sim_utils.create_prim(f"/World/Env_{i}", "Xform", stage=stage, translation=(i * 2.0, 0.0, 1.0)) - sim_utils.create_prim(f"/World/Env_{i}/Object", "Xform", stage=stage, translation=(0.0, 0.0, 0.0)) - prim_paths.append(f"/World/Env_{i}/Object") - # Play simulation - sim.reset() - - # Pattern to match all prims - pattern = "/World/Env_.*/Object" - print(f" Pattern: {pattern}") - - # Create view - start_time = time.perf_counter() - if api == "isaaclab-usd" or api == "isaaclab-fabric": - xform_view = IsaacLabXformPrimView(pattern, device=args_cli.device, validate_xform_ops=False) - elif api == "isaacsim-usd": - xform_view = IsaacSimXformPrimView(pattern, reset_xform_properties=False, usd=True) - elif api == "isaacsim-fabric": - xform_view = IsaacSimXformPrimView(pattern, reset_xform_properties=False, usd=False) - elif api == "isaacsim-exp": - xform_view = IsaacSimExperimentalXformPrimView(pattern) + cleanup = None + + if api == "isaaclab-newton-site": + newton_cfg = SimulationCfg(device=device, physics=NewtonCfg(solver_cfg=MJWarpSolverCfg())) + ctx = build_simulation_context(device=device, sim_cfg=newton_cfg, add_ground_plane=True) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None + InteractiveScene(_NewtonSceneCfg(num_envs=num_envs, env_spacing=2.0)) + + stage = sim_utils.get_current_stage() + for i in range(num_envs): + prim = stage.DefinePrim(f"/World/envs/env_{i}/Object/Sensor", "Xform") + sim_utils.standardize_xform_ops(prim) + prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(0.1, 0.0, 0.05)) + prim.GetAttribute("xformOp:orient").Set(Gf.Quatd(1.0, 0.0, 0.0, 0.0)) + + sim.reset() + + start_time = time.perf_counter() + xform_view = NewtonSiteFrameView("/World/envs/env_.*/Object/Sensor", device=device) + timing_results["init"] = time.perf_counter() - start_time + cleanup = lambda: ctx.__exit__(None, None, None) # noqa: E731 + else: - raise ValueError(f"Invalid API: {api}") - timing_results["init"] = time.perf_counter() - start_time - - if api in ("isaaclab-usd", "isaaclab-fabric", "isaacsim-usd", "isaacsim-fabric"): - num_prims = xform_view.count - elif api == "isaacsim-exp": - num_prims = len(xform_view.prims) - print(f" XformView managing {num_prims} prims") - - # Benchmark get_world_poses - # Warmup call to initialize Fabric (if needed) - excluded from timing - positions, orientations = xform_view.get_world_poses() - - # Now time the actual iterations (steady-state performance) - start_time = time.perf_counter() - for _ in range(num_iterations): - positions, orientations = xform_view.get_world_poses() - - # Ensure tensors are torch tensors (do this AFTER timing) - if not isinstance(positions, torch.Tensor): - positions = torch.tensor(positions, dtype=torch.float32) - if not isinstance(orientations, torch.Tensor): - orientations = torch.tensor(orientations, dtype=torch.float32) - - timing_results["get_world_poses"] = (time.perf_counter() - start_time) / num_iterations - - # Store initial world poses - computed_results["initial_world_positions"] = positions.clone() - computed_results["initial_world_orientations"] = orientations.clone() - - # Benchmark set_world_poses - new_positions = positions.clone() - new_positions[:, 2] += 0.1 - start_time = time.perf_counter() - for _ in range(num_iterations): - if api in ("isaaclab-usd", "isaaclab-fabric", "isaacsim-usd", "isaacsim-fabric"): + sim_utils.create_new_stage() + start_time = time.perf_counter() + use_fabric = api == "isaaclab-fabric" + sim = sim_utils.SimulationContext(sim_utils.SimulationCfg(dt=0.01, device=device, use_fabric=use_fabric)) + stage = sim_utils.get_current_stage() + + for i in range(num_envs): + sim_utils.create_prim(f"/World/Env_{i}", "Xform", stage=stage, translation=(i * 2.0, 0.0, 1.0)) + sim_utils.create_prim(f"/World/Env_{i}/Object", "Xform", stage=stage, translation=(0.0, 0.0, 0.0)) + + sim.reset() + + pattern = "/World/Env_.*/Object" + start_time = time.perf_counter() + ViewClass = FabricFrameView if use_fabric else UsdFrameView + xform_view = ViewClass(pattern, device=device, validate_xform_ops=False) + timing_results["init"] = time.perf_counter() - start_time + cleanup = lambda: sim.clear_instance() # noqa: E731 + + num_prims = xform_view.count + print(f" {api} managing {num_prims} prims") + + is_newton = api == "isaaclab-newton-site" + + def to_torch(a): + return wp.to_torch(a) if isinstance(a, wp.array) else a + + try: + # -- Warmup -------------------------------------------------------- + xform_view.get_world_poses() + + # -- get_world_poses ----------------------------------------------- + if is_newton: + torch.cuda.synchronize() + start_time = time.perf_counter() + for _ in range(num_iterations): + positions, orientations = xform_view.get_world_poses() + if is_newton: + torch.cuda.synchronize() + timing_results["get_world_poses"] = (time.perf_counter() - start_time) / num_iterations + + positions_t = to_torch(positions) + orientations_t = to_torch(orientations) + computed_results["initial_world_positions"] = positions_t.clone() + computed_results["initial_world_orientations"] = orientations_t.clone() + + # -- set_world_poses ----------------------------------------------- + if is_newton: + new_positions = wp.clone(positions) + wp.to_torch(new_positions)[:, 2] += 0.1 + else: + new_positions = positions_t.clone() + new_positions[:, 2] += 0.1 + + if is_newton: + torch.cuda.synchronize() + start_time = time.perf_counter() + for _ in range(num_iterations): xform_view.set_world_poses(new_positions, orientations) - elif api == "isaacsim-exp": - xform_view.set_world_poses(new_positions.cpu().numpy(), orientations.cpu().numpy()) - timing_results["set_world_poses"] = (time.perf_counter() - start_time) / num_iterations - - # Get world poses after setting to verify - positions_after_set, orientations_after_set = xform_view.get_world_poses() - if not isinstance(positions_after_set, torch.Tensor): - positions_after_set = torch.tensor(positions_after_set, dtype=torch.float32) - if not isinstance(orientations_after_set, torch.Tensor): - orientations_after_set = torch.tensor(orientations_after_set, dtype=torch.float32) - computed_results["world_positions_after_set"] = positions_after_set.clone() - computed_results["world_orientations_after_set"] = orientations_after_set.clone() - - # Benchmark get_local_poses - # Warmup call (though local poses use USD, so minimal overhead) - translations, orientations_local = xform_view.get_local_poses() - - # Now time the actual iterations - start_time = time.perf_counter() - for _ in range(num_iterations): - translations, orientations_local = xform_view.get_local_poses() - # Ensure tensors are torch tensors (do this AFTER timing) - if not isinstance(translations, torch.Tensor): - translations = torch.tensor(translations, dtype=torch.float32, device=args_cli.device) - if not isinstance(orientations_local, torch.Tensor): - orientations_local = torch.tensor(orientations_local, dtype=torch.float32, device=args_cli.device) - - timing_results["get_local_poses"] = (time.perf_counter() - start_time) / num_iterations - - # Store initial local poses - computed_results["initial_local_translations"] = translations.clone() - computed_results["initial_local_orientations"] = orientations_local.clone() - - # Benchmark set_local_poses - new_translations = translations.clone() - new_translations[:, 2] += 0.1 - start_time = time.perf_counter() - for _ in range(num_iterations): - if api in ("isaaclab-usd", "isaaclab-fabric", "isaacsim-usd", "isaacsim-fabric"): + if is_newton: + torch.cuda.synchronize() + timing_results["set_world_poses"] = (time.perf_counter() - start_time) / num_iterations + + pa, oa = xform_view.get_world_poses() + computed_results["world_positions_after_set"] = to_torch(pa).clone() + computed_results["world_orientations_after_set"] = to_torch(oa).clone() + + # -- get_local_poses ----------------------------------------------- + if is_newton: + torch.cuda.synchronize() + start_time = time.perf_counter() + for _ in range(num_iterations): + translations, orientations_local = xform_view.get_local_poses() + if is_newton: + torch.cuda.synchronize() + timing_results["get_local_poses"] = (time.perf_counter() - start_time) / num_iterations + + translations_t = to_torch(translations) + orientations_local_t = to_torch(orientations_local) + computed_results["initial_local_translations"] = translations_t.clone() + computed_results["initial_local_orientations"] = orientations_local_t.clone() + + # -- set_local_poses ----------------------------------------------- + if is_newton: + new_translations = wp.clone(translations) + wp.to_torch(new_translations)[:, 2] += 0.1 + else: + new_translations = translations_t.clone() + new_translations[:, 2] += 0.1 + + if is_newton: + torch.cuda.synchronize() + start_time = time.perf_counter() + for _ in range(num_iterations): xform_view.set_local_poses(new_translations, orientations_local) - elif api == "isaacsim-exp": - xform_view.set_local_poses(new_translations.cpu().numpy(), orientations_local.cpu().numpy()) - timing_results["set_local_poses"] = (time.perf_counter() - start_time) / num_iterations - - # Get local poses after setting to verify - translations_after_set, orientations_local_after_set = xform_view.get_local_poses() - if not isinstance(translations_after_set, torch.Tensor): - translations_after_set = torch.tensor(translations_after_set, dtype=torch.float32) - if not isinstance(orientations_local_after_set, torch.Tensor): - orientations_local_after_set = torch.tensor(orientations_local_after_set, dtype=torch.float32) - computed_results["local_translations_after_set"] = translations_after_set.clone() - computed_results["local_orientations_after_set"] = orientations_local_after_set.clone() - - # Benchmark combined get operation - # Warmup call (Fabric should already be initialized by now, but for consistency) - positions, orientations = xform_view.get_world_poses() - translations, local_orientations = xform_view.get_local_poses() - - # Now time the actual iterations - start_time = time.perf_counter() - for _ in range(num_iterations): - positions, orientations = xform_view.get_world_poses() - translations, local_orientations = xform_view.get_local_poses() - timing_results["get_both"] = (time.perf_counter() - start_time) / num_iterations - - # Benchmark interleaved set/get (realistic workflow pattern) - # Pre-convert tensors for experimental API to avoid conversion overhead in loop - if api == "isaacsim-exp": - new_positions_np = new_positions.cpu().numpy() - orientations_np = orientations - - # Warmup - if api in ("isaaclab-usd", "isaaclab-fabric", "isaacsim-usd", "isaacsim-fabric"): - xform_view.set_world_poses(new_positions, orientations) - positions, orientations = xform_view.get_world_poses() - elif api == "isaacsim-exp": - xform_view.set_world_poses(new_positions_np, orientations_np) - positions, orientations = xform_view.get_world_poses() - positions = torch.tensor(positions, dtype=torch.float32) - orientations = torch.tensor(orientations, dtype=torch.float32) - - # Now time the actual interleaved iterations - start_time = time.perf_counter() - for _ in range(num_iterations): - # Write then immediately read (common pattern: set pose, verify/query result) - if api in ("isaaclab-usd", "isaaclab-fabric", "isaacsim-usd", "isaacsim-fabric"): + if is_newton: + torch.cuda.synchronize() + timing_results["set_local_poses"] = (time.perf_counter() - start_time) / num_iterations + + ta, ola = xform_view.get_local_poses() + computed_results["local_translations_after_set"] = to_torch(ta).clone() + computed_results["local_orientations_after_set"] = to_torch(ola).clone() + + # -- get_both (world + local) -------------------------------------- + if is_newton: + torch.cuda.synchronize() + start_time = time.perf_counter() + for _ in range(num_iterations): + xform_view.get_world_poses() + xform_view.get_local_poses() + if is_newton: + torch.cuda.synchronize() + timing_results["get_both"] = (time.perf_counter() - start_time) / num_iterations + + # -- interleaved set -> get ---------------------------------------- + if is_newton: + torch.cuda.synchronize() + start_time = time.perf_counter() + for _ in range(num_iterations): xform_view.set_world_poses(new_positions, orientations) - positions, orientations = xform_view.get_world_poses() - elif api == "isaacsim-exp": - xform_view.set_world_poses(new_positions_np, orientations_np) - positions, orientations = xform_view.get_world_poses() + xform_view.get_world_poses() + if is_newton: + torch.cuda.synchronize() + timing_results["interleaved_world_set_get"] = (time.perf_counter() - start_time) / num_iterations - timing_results["interleaved_world_set_get"] = (time.perf_counter() - start_time) / num_iterations - - # close simulation - sim.clear_instance() + finally: + if cleanup: + cleanup() return timing_results, computed_results -def compare_results( - results_dict: dict[str, dict[str, torch.Tensor]], tolerance: float = 1e-4 -) -> dict[str, dict[str, dict[str, float]]]: - """Compare computed results across multiple implementations. - - Only compares implementations using the same data path: - - USD implementations (isaaclab-usd, isaacsim-usd) are compared with each other - - Fabric implementations (isaaclab-fabric, isaacsim-fabric) are compared with each other - - This is because Fabric is designed for write-first workflows and may not match - USD reads on initialization. - - Args: - results_dict: Dictionary mapping API names to their computed values. - tolerance: Tolerance for numerical comparison. - - Returns: - Nested dictionary: {comparison_pair: {metric: {stats}}}, e.g., - {"isaaclab-usd_vs_isaacsim-usd": {"initial_world_positions": {"max_diff": 0.001, ...}}} - """ - comparison_stats = {} - - # Group APIs by their data path (USD vs Fabric) - usd_apis = [api for api in results_dict.keys() if "usd" in api and "fabric" not in api] - fabric_apis = [api for api in results_dict.keys() if "fabric" in api] - - # Compare within USD group - for i, api1 in enumerate(usd_apis): - for api2 in usd_apis[i + 1 :]: - pair_key = f"{api1}_vs_{api2}" - comparison_stats[pair_key] = {} - - computed1 = results_dict[api1] - computed2 = results_dict[api2] - - for key in computed1.keys(): - if key not in computed2: - print(f" Warning: Key '{key}' not found in {api2} results") - continue - - val1 = computed1[key] - val2 = computed2[key] - - # Compute differences - diff = torch.abs(val1 - val2) - max_diff = torch.max(diff).item() - mean_diff = torch.mean(diff).item() - - # Check if within tolerance - all_close = torch.allclose(val1, val2, atol=tolerance, rtol=0) - - comparison_stats[pair_key][key] = { - "max_diff": max_diff, - "mean_diff": mean_diff, - "all_close": all_close, - } - - # Compare within Fabric group - for i, api1 in enumerate(fabric_apis): - for api2 in fabric_apis[i + 1 :]: - pair_key = f"{api1}_vs_{api2}" - comparison_stats[pair_key] = {} - - computed1 = results_dict[api1] - computed2 = results_dict[api2] - - for key in computed1.keys(): - if key not in computed2: - print(f" Warning: Key '{key}' not found in {api2} results") - continue - - val1 = computed1[key] - val2 = computed2[key] - - # Compute differences - diff = torch.abs(val1 - val2) - max_diff = torch.max(diff).item() - mean_diff = torch.mean(diff).item() - - # Check if within tolerance - all_close = torch.allclose(val1, val2, atol=tolerance, rtol=0) - - comparison_stats[pair_key][key] = { - "max_diff": max_diff, - "mean_diff": mean_diff, - "all_close": all_close, - } - - return comparison_stats - - -def print_comparison_results(comparison_stats: dict[str, dict[str, dict[str, float]]], tolerance: float): - """Print comparison results across implementations. - - Args: - comparison_stats: Nested dictionary containing comparison statistics for each API pair. - tolerance: Tolerance used for comparison. - """ - if not comparison_stats: - print("\n" + "=" * 100) - print("RESULT COMPARISON") - print("=" * 100) - print("ℹ️ No comparisons performed.") - print(" USD and Fabric implementations are not compared because Fabric uses a") - print(" write-first workflow and may not match USD reads on initialization.") - print("=" * 100) - print() - return - - for pair_key, pair_stats in comparison_stats.items(): - # Format the pair key for display (e.g., "isaaclab_vs_isaacsim" -> "Isaac Lab vs Isaac Sim") - api1, api2 = pair_key.split("_vs_") - display_api1 = api1.replace("-", " ").title() - display_api2 = api2.replace("-", " ").title() - comparison_title = f"{display_api1} vs {display_api2}" - - # Check if all results match - all_match = all(stats["all_close"] for stats in pair_stats.values()) - - if all_match: - # Compact output when everything matches - print("\n" + "=" * 100) - print(f"RESULT COMPARISON: {comparison_title}") - print("=" * 100) - print(f"✓ All computed values match within tolerance ({tolerance})") - print("=" * 100) - else: - # Detailed output when there are mismatches - print("\n" + "=" * 100) - print(f"RESULT COMPARISON: {comparison_title}") - print("=" * 100) - print(f"{'Computed Value':<40} {'Max Diff':<15} {'Mean Diff':<15} {'Match':<10}") - print("-" * 100) - - for key, stats in pair_stats.items(): - # Format the key for display - display_key = key.replace("_", " ").title() - match_str = "✓ Yes" if stats["all_close"] else "✗ No" - - print(f"{display_key:<40} {stats['max_diff']:<15.6e} {stats['mean_diff']:<15.6e} {match_str:<10}") - - print("=" * 100) - print(f"\n✗ Some results differ beyond tolerance ({tolerance})") - - # Special note for Isaac Sim Fabric local pose bug - if "isaacsim-fabric" in pair_key and any("local_translations_after_set" in k for k in pair_stats.keys()): - if not pair_stats.get("local_translations_after_set", {}).get("all_close", True): - print("\n ⚠️ Known Issue: Isaac Sim Fabric has a bug where get_local_poses() returns stale") - print(" values after set_local_poses(). Isaac Lab Fabric correctly returns updated values.") - print(" This is a correctness issue in Isaac Sim's implementation, not Isaac Lab's.") - else: - print(f" This may indicate implementation differences between {display_api1} and {display_api2}") - - print() +# ------------------------------------------------------------------ +# Reporting +# ------------------------------------------------------------------ def print_results(results_dict: dict[str, dict[str, float]], num_prims: int, num_iterations: int): - """Print benchmark results in a formatted table. - - Args: - results_dict: Dictionary mapping API names to their timing results. - num_prims: Number of prims tested. - num_iterations: Number of iterations run. - """ - print("\n" + "=" * 100) + """Print benchmark results in a formatted table.""" + print("\n" + "=" * 120) print(f"BENCHMARK RESULTS: {num_prims} prims, {num_iterations} iterations") - print("=" * 100) + print("=" * 120) api_names = list(results_dict.keys()) - # Format API names for display - display_names = [name.replace("-", " ").replace("_", " ").title() for name in api_names] - - # Calculate column width based on number of APIs - col_width = 20 + display_names = [name.replace("-", " ").title() for name in api_names] + col_width = 22 - # Print header - header = f"{'Operation':<25}" - for display_name in display_names: - header += f" {display_name + ' (ms)':<{col_width}}" + header = f"{'Operation':<28}" + for dn in display_names: + header += f" {dn + ' (ms)':>{col_width}}" print(header) - print("-" * 100) + print("-" * 120) - # Print each operation operations = [ ("Initialization", "init"), ("Get World Poses", "get_world_poses"), @@ -469,159 +274,119 @@ def print_results(results_dict: dict[str, dict[str, float]], num_prims: int, num ("Get Local Poses", "get_local_poses"), ("Set Local Poses", "set_local_poses"), ("Get Both (World+Local)", "get_both"), - ("Interleaved World Set→Get", "interleaved_world_set_get"), + ("Interleaved World Set->Get", "interleaved_world_set_get"), ] for op_name, op_key in operations: - row = f"{op_name:<25}" - for api_name in api_names: - api_time = results_dict[api_name].get(op_key, 0) * 1000 # Convert to ms - row += f" {api_time:>{col_width - 1}.4f}" + row = f"{op_name:<28}" + for name in api_names: + val = results_dict[name].get(op_key, 0) * 1000 + row += f" {val:>{col_width}.4f}" print(row) - print("=" * 100) + print("=" * 120) - # Calculate and print total time - total_row = f"{'Total Time':<25}" - for api_name in api_names: - total_time = sum(results_dict[api_name].values()) * 1000 - total_row += f" {total_time:>{col_width - 1}.4f}" + total_row = f"{'Total':<28}" + for name in api_names: + total_row += f" {sum(results_dict[name].values()) * 1000:>{col_width}.4f}" print(f"\n{total_row}") - # Calculate speedups relative to Isaac Lab USD (baseline) - if "isaaclab-usd" in api_names: - print("\n" + "=" * 100) - print("SPEEDUP vs Isaac Lab USD (Baseline)") - print("=" * 100) - print(f"{'Operation':<25}", end="") - for api_name, display_name in zip(api_names, display_names): - if api_name != "isaaclab-usd": - print(f" {display_name:<{col_width}}", end="") - print() - print("-" * 100) - - isaaclab_usd_results = results_dict["isaaclab-usd"] + baseline = "isaaclab-usd" + if baseline in results_dict and len(api_names) > 1: + print("\n" + "=" * 120) + print(f"SPEEDUP vs {baseline.replace('-', ' ').title()}") + print("=" * 120) + header = f"{'Operation':<28}" + for name in api_names: + if name != baseline: + header += f" {name.replace('-', ' ').title():>{col_width}}" + print(header) + print("-" * 120) + + base = results_dict[baseline] for op_name, op_key in operations: - print(f"{op_name:<25}", end="") - isaaclab_usd_time = isaaclab_usd_results.get(op_key, 0) - for api_name, display_name in zip(api_names, display_names): - if api_name != "isaaclab-usd": - api_time = results_dict[api_name].get(op_key, 0) - if isaaclab_usd_time > 0 and api_time > 0: - speedup = isaaclab_usd_time / api_time - print(f" {speedup:>{col_width - 1}.2f}x", end="") + row = f"{op_name:<28}" + base_t = base.get(op_key, 0) + for name in api_names: + if name != baseline: + impl_t = results_dict[name].get(op_key, 0) + if base_t > 0 and impl_t > 0: + row += f" {base_t / impl_t:>{col_width}.2f}x" else: - print(f" {'N/A':>{col_width}}", end="") - print() - - # Overall speedup - print("=" * 100) - print(f"{'Overall Speedup':<25}", end="") - total_isaaclab_usd = sum(isaaclab_usd_results.values()) - for api_name, display_name in zip(api_names, display_names): - if api_name != "isaaclab-usd": - total_api = sum(results_dict[api_name].values()) - if total_isaaclab_usd > 0 and total_api > 0: - overall_speedup = total_isaaclab_usd / total_api - print(f" {overall_speedup:>{col_width - 1}.2f}x", end="") + row += f" {'N/A':>{col_width}}" + print(row) + + print("=" * 120) + print(f"{'Overall':>28}", end="") + total_base = sum(base.values()) + for name in api_names: + if name != baseline: + total_impl = sum(results_dict[name].values()) + if total_base > 0 and total_impl > 0: + print(f" {total_base / total_impl:>{col_width}.2f}x", end="") else: print(f" {'N/A':>{col_width}}", end="") print() - print("\n" + "=" * 100) + print("\n" + "=" * 120) print("\nNotes:") print(" - Times are averaged over all iterations") - print(" - Speedup = (Isaac Lab USD time) / (Other API time)") - print(" - Speedup > 1.0 means the other API is faster than Isaac Lab USD") - print(" - Speedup < 1.0 means the other API is slower than Isaac Lab USD") + print(" - Speedup > 1.0 means faster than USD baseline") print() def main(): - """Main benchmark function.""" - print("=" * 100) - print("XformPrimView Benchmark - Comparing Multiple APIs") - print("=" * 100) - print("Configuration:") - print(f" Number of environments: {args_cli.num_envs}") - print(f" Iterations per test: {args_cli.num_iterations}") - print(f" Device: {args_cli.device}") - print(f" Profiling: {'Enabled' if args_cli.profile else 'Disabled'}") - if args_cli.profile: - print(f" Profile directory: {args_cli.profile_dir}") + print("=" * 120) + print("FrameView Benchmark") + print("=" * 120) + print(f" Environments: {args_cli.num_envs}") + print(f" Iterations: {args_cli.num_iterations}") + print(f" Device: {args_cli.device}") print() - # Create profile directory if profiling is enabled if args_cli.profile: import os os.makedirs(args_cli.profile_dir, exist_ok=True) - # Dictionary to store all results - all_timing_results = {} - all_computed_results = {} + all_timing = {} + all_computed = {} profile_files = {} - # APIs to benchmark - apis_to_test = [ - ("isaaclab-usd", "Isaac Lab XformPrimView (USD)"), - ("isaaclab-fabric", "Isaac Lab XformPrimView (Fabric)"), - ("isaacsim-usd", "Isaac Sim XformPrimView (USD)"), - ("isaacsim-fabric", "Isaac Sim XformPrimView (Fabric)"), - ("isaacsim-exp", "Isaac Sim Experimental XformPrim"), + apis = [ + ("isaaclab-usd", "Isaac Lab FrameView (USD)"), + ("isaaclab-fabric", "Isaac Lab FrameView (Fabric)"), + ("isaaclab-newton-site", "Isaac Lab FrameView (Newton Site)"), ] - # Benchmark each API - for api_key, api_name in apis_to_test: + for api_key, api_name in apis: print(f"Benchmarking {api_name}...") if args_cli.profile: profiler = cProfile.Profile() profiler.enable() - # Cast api_key to Literal type for type checker - timing, computed = benchmark_xform_prim_view( - api=api_key, # type: ignore[arg-type] - num_iterations=args_cli.num_iterations, - ) + timing, computed = benchmark_frame_view(api=api_key, num_iterations=args_cli.num_iterations) if args_cli.profile: profiler.disable() - profile_file = f"{args_cli.profile_dir}/{api_key.replace('-', '_')}_benchmark.prof" - profiler.dump_stats(profile_file) - profile_files[api_key] = profile_file - print(f" Profile saved to: {profile_file}") - - all_timing_results[api_key] = timing - all_computed_results[api_key] = computed - - print(" Done!") - print() + pf = f"{args_cli.profile_dir}/{api_key.replace('-', '_')}_benchmark.prof" + profiler.dump_stats(pf) + profile_files[api_key] = pf + print(f" Profile saved to: {pf}") - # Print timing results - print_results(all_timing_results, args_cli.num_envs, args_cli.num_iterations) + all_timing[api_key] = timing + all_computed[api_key] = computed + print(" Done!\n") - # Compare computed results - print("\nComparing computed results across APIs...") - comparison_stats = compare_results(all_computed_results, tolerance=1e-6) - print_comparison_results(comparison_stats, tolerance=1e-4) + print_results(all_timing, args_cli.num_envs, args_cli.num_iterations) - # Print profiling instructions if enabled if args_cli.profile: - print("\n" + "=" * 100) - print("PROFILING RESULTS") - print("=" * 100) - print("Profile files have been saved. To visualize with snakeviz, run:") - for api_key, profile_file in profile_files.items(): - api_display = api_key.replace("-", " ").title() - print(f" # {api_display}") - print(f" snakeviz {profile_file}") - print("\nAlternatively, use pstats to analyze in terminal:") - print(" python -m pstats ") - print("=" * 100) + print("\nProfile files:") + for key, pf in profile_files.items(): + print(f" snakeviz {pf}") print() - # Clean up sim_utils.SimulationContext.clear_instance() diff --git a/scripts/demos/sensors/raycaster_sensor.py b/scripts/demos/sensors/raycaster_sensor.py index 4f758274b61c..dd0b454ad636 100644 --- a/scripts/demos/sensors/raycaster_sensor.py +++ b/scripts/demos/sensors/raycaster_sensor.py @@ -62,7 +62,7 @@ class RaycasterSensorSceneCfg(InteractiveSceneCfg): robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") ray_caster = RayCasterCfg( - prim_path="{ENV_REGEX_NS}/Robot/base/lidar_cage", + prim_path="{ENV_REGEX_NS}/Robot/base", update_period=1 / 60, offset=RayCasterCfg.OffsetCfg(pos=(0, 0, 0.5)), mesh_prim_paths=["/World/Ground"], diff --git a/source/isaaclab/config/extension.toml b/source/isaaclab/config/extension.toml index 7092d52f3ff7..86024cf07e6a 100644 --- a/source/isaaclab/config/extension.toml +++ b/source/isaaclab/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "4.6.10" +version = "4.6.11" # Description title = "Isaac Lab framework for Robot Learning" diff --git a/source/isaaclab/docs/CHANGELOG.rst b/source/isaaclab/docs/CHANGELOG.rst index 0251fcb03ca2..489be9f6aa15 100644 --- a/source/isaaclab/docs/CHANGELOG.rst +++ b/source/isaaclab/docs/CHANGELOG.rst @@ -1,6 +1,40 @@ Changelog --------- +4.6.11 (2026-04-22) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Changed :class:`~isaaclab.sensors.RayCaster` to spawn its own non-physics Xform prim via + the new :attr:`~isaaclab.sensors.RayCasterCfg.spawn` attribute. ``prim_path`` should now + point to a child under the parent link (e.g. ``{ENV_REGEX_NS}/Robot/base/raycaster``). +* Renamed :class:`~isaaclab.sim.views.XformPrimView` to :class:`~isaaclab.sim.views.FrameView`, + ``BaseXformPrimView`` to :class:`~isaaclab.sim.views.BaseFrameView`, + and ``UsdXformPrimView`` to :class:`~isaaclab.sim.views.UsdFrameView`. + ``XformPrimView`` is kept as a deprecated alias. +* Moved :class:`~isaaclab.sensors.RayCasterCfg` offset into the spawned prim's local transform + instead of applying it at runtime. The :class:`~isaaclab.sim.views.FrameView` world pose now + includes the offset directly. +* Unified sensor prim path resolution in :class:`~isaaclab.sensors.SensorBase`. When + ``prim_path`` points at a physics body and a spawner is configured, a child prim is + automatically created underneath. + +Deprecated +^^^^^^^^^^ + +* Deprecated passing a ``prim_path`` with ``ArticulationRootAPI`` or ``RigidBodyAPI`` to + :class:`~isaaclab.sensors.RayCasterCfg`. The path is automatically extended with + ``/raycaster``; users should migrate to the child-path convention. + +Removed +^^^^^^^ + +* Removed :attr:`~isaaclab.sensors.RayCasterCfg.attach_yaw_only` (deprecated since 2.1.1). + Use ``ray_alignment="yaw"`` or ``ray_alignment="base"`` instead. + + 4.6.10 (2026-04-22) ~~~~~~~~~~~~~~~~~~~ @@ -381,6 +415,21 @@ Added Added ^^^^^ + +* Added :class:`~isaaclab.sim.views.BaseXformPrimView` abstract base class that defines + the common interface for backend-specific ``XformPrimView`` implementations. +* Added :class:`~isaaclab.sim.views.XformPrimView` factory to instantiate the correct + backend-specific ``XformPrimView`` based on the active simulation backend. + +Changed +^^^^^^^ + +* Refactored :class:`~isaaclab.sim.views.XformPrimView` to delegate backend-specific + logic to :class:`~isaaclab_physx.sim.views.FabricXformPrimView` and + :class:`~isaaclab_newton.sim.views.NewtonSiteXformPrimView`. The public API is + unchanged; use :class:`~isaaclab.sim.views.XformPrimView` for backend-aware + instantiation. + * Added release version to :class:`~isaaclab.test.benchmark.recorders.VersionInfoRecorder` output. diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index a4ac4424f485..da6f5eff75d7 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -32,7 +32,7 @@ from isaaclab.sensors import ContactSensorCfg, FrameTransformerCfg, SensorBase, SensorBaseCfg from isaaclab.sim import SimulationContext from isaaclab.sim.utils.stage import get_current_stage, get_current_stage_id -from isaaclab.sim.views import XformPrimView +from isaaclab.sim.views import FrameView from isaaclab.terrains import TerrainImporter, TerrainImporterCfg # Note: This is a temporary import for the VisuoTactileSensorCfg class. @@ -403,11 +403,11 @@ def surface_grippers(self) -> dict[str, SurfaceGripper]: return self._surface_grippers @property - def extras(self) -> dict[str, XformPrimView]: + def extras(self) -> dict[str, FrameView]: """A dictionary of miscellaneous simulation objects that neither inherit from assets nor sensors. The keys are the names of the miscellaneous objects, and the values are the - :class:`~isaaclab.sim.views.XformPrimView` instances of the corresponding prims. + :class:`~isaaclab.sim.views.FrameView` instances of the corresponding prims. As an example, lights or other props in the scene that do not have any attributes or properties that you want to alter at runtime can be added to this dictionary. @@ -833,7 +833,7 @@ def _add_entities_from_cfg(self): # noqa: C901 ) # store xform prim view corresponding to this asset # all prims in the scene are Xform prims (i.e. have a transform component) - self._extras[asset_name] = XformPrimView(asset_cfg.prim_path, device=self.device, stage=self.stage) + self._extras[asset_name] = FrameView(asset_cfg.prim_path, device=self.device, stage=self.stage) else: raise ValueError(f"Unknown asset config type for {asset_name}: {asset_cfg}") diff --git a/source/isaaclab/isaaclab/sensors/__init__.py b/source/isaaclab/isaaclab/sensors/__init__.py index 1128b11c1202..717fc4a7163c 100644 --- a/source/isaaclab/isaaclab/sensors/__init__.py +++ b/source/isaaclab/isaaclab/sensors/__init__.py @@ -26,7 +26,7 @@ +---------------------+---------------------------+---------------------------------------------------------------+ | Contact Sensor | /World/robot/feet_* | Leaf is available and checks if the schema exists | +---------------------+---------------------------+---------------------------------------------------------------+ -| Ray Caster | /World/robot/base | Leaf exists and is a physics body (Articulation / Rigid Body) | +| Ray Caster | /World/robot/base/raycast | ``spawn`` creates an Xform leaf; else the leaf must exist | +---------------------+---------------------------+---------------------------------------------------------------+ | Frame Transformer | /World/robot/base | Leaf exists and is a physics body (Articulation / Rigid Body) | +---------------------+---------------------------+---------------------------------------------------------------+ diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index eb588489f729..22c96af1779e 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -16,11 +16,10 @@ from pxr import Sdf, UsdGeom -import isaaclab.sim as sim_utils import isaaclab.utils.sensors as sensor_utils from isaaclab.app.settings_manager import get_settings_manager from isaaclab.renderers import BaseRenderer, Renderer -from isaaclab.sim.views import XformPrimView +from isaaclab.sim.views import FrameView from isaaclab.utils import has_kit, to_camel_case from isaaclab.utils.math import ( convert_camera_frame_orientation_convention, @@ -121,36 +120,15 @@ def __init__(self, cfg: CameraCfg): settings = get_settings_manager() settings.set_bool("/isaaclab/render/rtx_sensors", True) - # spawn the asset - if self.cfg.spawn is not None: - # Use spawn_path when set (points to template location for scene-cloned sensors). - # This allows the camera to be spawned inside the asset template (e.g. inside - # proto_asset_0) before clone_environments replicates it to all env paths. - spawn_target = ( - self.cfg.spawn.spawn_path - if getattr(self.cfg.spawn, "spawn_path", None) is not None - else self.cfg.prim_path - ) - # compute the rotation offset - rot = torch.tensor(self.cfg.offset.rot, dtype=torch.float32, device="cpu").unsqueeze(0) - rot_offset = convert_camera_frame_orientation_convention( - rot, origin=self.cfg.offset.convention, target="opengl" - ) - rot_offset = rot_offset.squeeze(0).cpu().numpy() - # ensure vertical aperture is set, otherwise replace with default for squared pixels - if self.cfg.spawn.vertical_aperture is None: - self.cfg.spawn.vertical_aperture = self.cfg.spawn.horizontal_aperture * self.cfg.height / self.cfg.width - self.cfg.spawn.func(spawn_target, self.cfg.spawn, translation=self.cfg.offset.pos, orientation=rot_offset) - # check that spawn was successful; use spawn_path if set (template location) since env - # paths are not yet populated at init time — they are filled in by clone_environments. - check_path = ( - self.cfg.spawn.spawn_path - if self.cfg.spawn is not None and getattr(self.cfg.spawn, "spawn_path", None) is not None - else self.cfg.prim_path + # Compute camera orientation (convention conversion) and spawn + rot = torch.tensor(self.cfg.offset.rot, dtype=torch.float32, device="cpu").unsqueeze(0) + rot_offset = convert_camera_frame_orientation_convention( + rot, origin=self.cfg.offset.convention, target="opengl" ) - matching_prims = sim_utils.find_matching_prims(check_path) - if len(matching_prims) == 0: - raise RuntimeError(f"Could not find prim with path {check_path}.") + rot_offset = rot_offset.squeeze(0).cpu().numpy() + if self.cfg.spawn is not None and self.cfg.spawn.vertical_aperture is None: + self.cfg.spawn.vertical_aperture = self.cfg.spawn.horizontal_aperture * self.cfg.height / self.cfg.width + self._resolve_and_spawn("camera", translation=self.cfg.offset.pos, orientation=rot_offset) # UsdGeom Camera prim for the sensor self._sensor_prims: list[UsdGeom.Camera] = list() @@ -335,8 +313,16 @@ def set_world_poses( elif not isinstance(orientations, torch.Tensor): orientations = torch.tensor(orientations, device=self._device) orientations = convert_camera_frame_orientation_convention(orientations, origin=convention, target="opengl") - # set the pose - self._view.set_world_poses(positions, orientations, env_ids) + # convert torch tensors to warp arrays for the view + pos_wp = wp.from_torch(positions.contiguous()) if positions is not None else None + ori_wp = wp.from_torch(orientations.contiguous()) if orientations is not None else None + if env_ids is not None: + if not isinstance(env_ids, torch.Tensor): + env_ids = torch.tensor(env_ids, dtype=torch.int32, device=self._device) + idx_wp = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) + else: + idx_wp = None + self._view.set_world_poses(pos_wp, ori_wp, idx_wp) def set_world_poses_from_view( self, eyes: torch.Tensor, targets: torch.Tensor, env_ids: Sequence[int] | None = None @@ -359,7 +345,10 @@ def set_world_poses_from_view( up_axis = UsdGeom.GetStageUpAxis(self.stage) # set camera poses using the view orientations = quat_from_matrix(create_rotation_matrix_from_view(eyes, targets, up_axis, device=self._device)) - self._view.set_world_poses(eyes, orientations, env_ids) + if not isinstance(env_ids, torch.Tensor): + env_ids = torch.tensor(env_ids, dtype=torch.int32, device=self._device) + idx_wp = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) + self._view.set_world_poses(wp.from_torch(eyes.contiguous()), wp.from_torch(orientations.contiguous()), idx_wp) """ Operations @@ -418,9 +407,7 @@ def _initialize_impl(self): # Create a view for the sensor with Fabric enabled for fast pose queries. # TODO: remove sync_usd_on_fabric_write=True once the GPU Fabric sync bug is fixed. - self._view = XformPrimView( - self.cfg.prim_path, device=self._device, stage=self.stage, sync_usd_on_fabric_write=True - ) + self._view = FrameView(self.cfg.prim_path, device=self._device, stage=self.stage, sync_usd_on_fabric_write=True) # Check that sizes are correct if self._view.count != self._num_envs: raise RuntimeError( @@ -612,11 +599,14 @@ def _update_poses(self, env_ids: Sequence[int]): if len(self._sensor_prims) == 0: raise RuntimeError("Camera prim is None. Please call 'sim.play()' first.") - # get the poses from the view - poses, quat = self._view.get_world_poses(env_ids) - self._data.pos_w[env_ids] = poses + # get the poses from the view (returns wp.array, convert to torch) + if env_ids is not None and not isinstance(env_ids, torch.Tensor): + env_ids = torch.tensor(env_ids, dtype=torch.int32, device=self._device) + indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None + pos_wp, quat_wp = self._view.get_world_poses(indices) + self._data.pos_w[env_ids] = wp.to_torch(pos_wp) self._data.quat_w_world[env_ids] = convert_camera_frame_orientation_convention( - quat, origin="opengl", target="world" + wp.to_torch(quat_wp), origin="opengl", target="world" ) # notify renderer of updated poses (guarded in case called before initialization completes) if self._render_data is not None: diff --git a/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py b/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py index 1b5070cfd214..efd8e1f304c1 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py @@ -85,7 +85,7 @@ class OffsetCfg: """Whether to update the latest camera pose when fetching the camera's data. Defaults to False. If True, the latest camera pose is updated in the camera's data which will slow down performance - due to the use of :class:`XformPrimView`. + due to the use of :class:`FrameView`. If False, the pose of the camera during initialization is returned. """ diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py index 002456b6e101..9f3d692b13cd 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py @@ -14,18 +14,15 @@ import trimesh import warp as wp -import omni.physics.tensors.impl.api as physx - import isaaclab.sim as sim_utils -from isaaclab.sim.views import XformPrimView -from isaaclab.utils.math import combine_frame_transforms, matrix_from_quat +from isaaclab.sim.views import BaseFrameView, FrameView +from isaaclab.utils.math import matrix_from_quat from isaaclab.utils.mesh import PRIMITIVE_MESH_TYPES, create_trimesh_from_geom_mesh, create_trimesh_from_geom_shape from isaaclab.utils.warp import convert_to_warp_mesh from isaaclab.utils.warp import kernels as warp_kernels from .kernels import fill_float2d_masked_kernel, fill_vec3_inf_kernel from .multi_mesh_ray_caster_data import MultiMeshRayCasterData -from .ray_cast_utils import obtain_world_pose_from_view from .ray_caster import RayCaster if TYPE_CHECKING: @@ -86,12 +83,7 @@ class MultiMeshRayCaster(RayCaster): cfg: MultiMeshRayCasterCfg """The configuration parameters.""" - mesh_offsets: ClassVar[dict[str, tuple[torch.Tensor, torch.Tensor]]] = {} - """Per-mesh position and orientation offsets relative to their physics views, shared across instances. - - Keys are prim path expressions; values are ``(pos_offset, ori_offset)`` tuples.""" - - mesh_views: ClassVar[dict[str, XformPrimView | physx.ArticulationView | physx.RigidBodyView]] = {} + mesh_views: ClassVar[dict[str, BaseFrameView]] = {} """A dictionary to store mesh views for raycasting, shared across all instances. The keys correspond to the prim path for the mesh views, and values are the corresponding view objects. @@ -281,8 +273,8 @@ def _initialize_warp_meshes(self): mesh_idx += n_meshes_per_env if target_cfg.track_mesh_transforms: - MultiMeshRayCaster.mesh_views[target_prim_path], MultiMeshRayCaster.mesh_offsets[target_prim_path] = ( - self._obtain_trackable_prim_view(target_prim_path) + MultiMeshRayCaster.mesh_views[target_prim_path] = FrameView( + target_prim_path, device=self._device, stage=self.stage ) if all([target_cfg.prim_expr not in multi_mesh_ids for target_cfg in self._raycast_targets_cfg]): @@ -357,19 +349,12 @@ def _update_mesh_transforms(self) -> None: mesh_idx += self._num_meshes_per_env[target_cfg.prim_expr] continue - pos_w, ori_w = obtain_world_pose_from_view(view, None) + # update position of the target meshes + pos_wp, quat_wp = view.get_world_poses(None) + pos_w, ori_w = wp.to_torch(pos_wp), wp.to_torch(quat_wp) pos_w = pos_w.squeeze(0) if len(pos_w.shape) == 3 else pos_w ori_w = ori_w.squeeze(0) if len(ori_w.shape) == 3 else ori_w - if target_cfg.prim_expr in MultiMeshRayCaster.mesh_offsets: - pos_offset, ori_offset = MultiMeshRayCaster.mesh_offsets[target_cfg.prim_expr] - pos_w, ori_w = combine_frame_transforms( - pos_w, - ori_w, - pos_offset.expand(pos_w.shape[0], -1), - ori_offset.expand(ori_w.shape[0], -1), - ) - count = view.count if count != 1: count = count // self._num_envs @@ -434,7 +419,6 @@ def _invalidate_initialize_callback(self, event): def __del__(self): super().__del__() if RayCaster._instance_count == 0: - MultiMeshRayCaster.mesh_offsets.clear() MultiMeshRayCaster.mesh_views.clear() diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py index d5e084abb32e..f184c28b20e2 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py @@ -21,7 +21,6 @@ ) from .multi_mesh_ray_caster import MultiMeshRayCaster from .multi_mesh_ray_caster_camera_data import MultiMeshRayCasterCameraData -from .ray_cast_utils import obtain_world_pose_from_view from .ray_caster_camera import RayCasterCamera if TYPE_CHECKING: @@ -178,7 +177,9 @@ def _update_ray_infos(self, env_mask: wp.array): return # Compute camera world poses by composing view pose with sensor offset - pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids) + indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) + pos_wp, quat_wp = self._view.get_world_poses(indices) + pos_w, quat_w = wp.to_torch(pos_wp), wp.to_torch(quat_wp) pos_w, quat_w = math_utils.combine_frame_transforms( pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] ) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_cast_utils.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_cast_utils.py deleted file mode 100644 index ac503b28bf52..000000000000 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_cast_utils.py +++ /dev/null @@ -1,49 +0,0 @@ -# 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 - -"""Utility functions for ray-cast sensors.""" - -from __future__ import annotations - -import torch -import warp as wp - -import omni.physics.tensors.impl.api as physx - -from isaaclab.sim.views import XformPrimView - - -def obtain_world_pose_from_view( - physx_view: XformPrimView | physx.ArticulationView | physx.RigidBodyView, - env_ids: torch.Tensor, - clone: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - """Get the world poses of the prim referenced by the prim view. - - Args: - physx_view: The prim view to get the world poses from. - env_ids: The environment ids of the prims to get the world poses for. - clone: Whether to clone the returned tensors (default: False). - - Returns: - A tuple containing the world positions and orientations of the prims. - Orientation is in (x, y, z, w) format. - - Raises: - NotImplementedError: If the prim view is not of the supported type. - """ - if isinstance(physx_view, XformPrimView): - pos_w, quat_w = physx_view.get_world_poses(env_ids) - elif isinstance(physx_view, physx.ArticulationView): - pos_w, quat_w = wp.to_torch(physx_view.get_root_transforms())[env_ids].split([3, 4], dim=-1) - elif isinstance(physx_view, physx.RigidBodyView): - pos_w, quat_w = wp.to_torch(physx_view.get_transforms())[env_ids].split([3, 4], dim=-1) - else: - raise NotImplementedError(f"Cannot get world poses for prim view of type '{type(physx_view)}'.") - - if clone: - return pos_w.clone(), quat_w.clone() - else: - return pos_w, quat_w diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py index 89cc9aaa674f..1e23ee00c1b6 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py @@ -13,13 +13,12 @@ import torch import warp as wp -import omni.physics.tensors.impl.api as physx -from pxr import Gf, Usd, UsdGeom, UsdPhysics +from pxr import Gf, Usd, UsdGeom import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils from isaaclab.markers import VisualizationMarkers -from isaaclab.sim.views import XformPrimView +from isaaclab.sim.views import FrameView from isaaclab.terrains.trimesh.utils import make_plane from isaaclab.utils.warp import convert_to_warp_mesh from isaaclab.utils.warp.kernels import raycast_mesh_masked_kernel @@ -74,6 +73,8 @@ def __init__(self, cfg: RayCasterCfg): """ RayCaster._instance_count += 1 super().__init__(cfg) + # Resolve physics-body paths and spawn the sensor Xform child if needed. + self._resolve_and_spawn("raycaster") self._data = RayCasterData() def __str__(self) -> str: @@ -135,42 +136,20 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None def _initialize_impl(self): super()._initialize_impl() - # obtain global simulation view - self._physics_sim_view = sim_utils.SimulationContext.instance().physics_manager.get_physics_sim_view() - prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) - if prim is None: - available_prims = ",".join([str(p.GetPath()) for p in sim_utils.get_current_stage().Traverse()]) - raise RuntimeError( - f"Failed to find a prim at path expression: {self.cfg.prim_path}. Available prims: {available_prims}" - ) - - self._view, self._offset = self._obtain_trackable_prim_view(self.cfg.prim_path) - - # Convert offsets to warp (zero-copy from existing torch tensors). - # Store the contiguous tensors explicitly so they are not garbage-collected while - # the wp.array views (_offset_pos_wp / _offset_quat_wp) are alive. If the tensor - # returned by .contiguous() is a temporary copy (non-contiguous input), the warp - # view would otherwise point to freed memory once GC reclaims it. - self._offset_pos_contiguous = self._offset[0].contiguous() - self._offset_quat_contiguous = self._offset[1].contiguous() - self._offset_pos_wp = wp.from_torch(self._offset_pos_contiguous, dtype=wp.vec3f) + # Build a FrameView over the sensor prim paths. The FrameView tracks the spawned + # (non-physics) Xform directly, so no physics-body redirect or offset resolution + # is needed at runtime — the world pose returned already includes any offset + # baked into the prim's local transform. + self._view = FrameView(self.cfg.prim_path, device=self._device, stage=self.stage) + + # Per-env identity offsets (kept for kernel ABI compatibility): the sensor frame is + # already the FrameView's tracked prim, so no additional view-to-sensor offset applies. + self._offset_pos_wp = wp.zeros(self._view.count, dtype=wp.vec3f, device=self._device) + identity_quat = torch.zeros(self._view.count, 4, device=self._device) + identity_quat[:, 3] = 1.0 + self._offset_quat_contiguous = identity_quat.contiguous() self._offset_quat_wp = wp.from_torch(self._offset_quat_contiguous, dtype=wp.quatf) - # Handle deprecated attach_yaw_only at init time - if self.cfg.attach_yaw_only is not None: - msg = ( - "Raycaster attribute 'attach_yaw_only' property will be deprecated in a future release." - " Please use the parameter 'ray_alignment' instead." - ) - if self.cfg.attach_yaw_only: - self.cfg.ray_alignment = "yaw" - msg += " Setting ray_alignment to 'yaw'." - else: - self.cfg.ray_alignment = "base" - msg += " Setting ray_alignment to 'base'." - logger.warning(msg) - self.cfg.attach_yaw_only = None - # Resolve alignment mode to integer constant for kernel dispatch alignment_map = {"world": 0, "yaw": 1, "base": 2} if self.cfg.ray_alignment not in alignment_map: @@ -278,23 +257,18 @@ def _initialize_rays_impl(self): self._dummy_ray_normal = wp.empty((1, 1), dtype=wp.vec3f, device=self._device) def _get_view_transforms_wp(self) -> wp.array: - """Get world transforms from the physics view as a warp array. + """Get world transforms from the frame view as a warp array of ``wp.transformf``. Returns: - Warp array of ``wp.transformf`` with shape (num_envs,). + Warp array of ``wp.transformf`` with shape ``(num_envs,)``. Layout is + ``(tx, ty, tz, qx, qy, qz, qw)`` per element, matching the quaternion + convention returned by :class:`~isaaclab.sim.views.FrameView`. """ - if isinstance(self._view, XformPrimView): - # XformPrimView.get_world_poses() returns quaternions in (x, y, z, w) convention, - # which matches the wp.transformf layout (translation then xyzw quaternion). - pos_w, quat_w = self._view.get_world_poses() - poses = torch.cat([pos_w, quat_w], dim=-1).contiguous() - return wp.from_torch(poses).view(wp.transformf) - elif isinstance(self._view, physx.ArticulationView): - return self._view.get_root_transforms().view(wp.transformf) - elif isinstance(self._view, physx.RigidBodyView): - return self._view.get_transforms().view(wp.transformf) - else: - raise NotImplementedError(f"Cannot get transforms for view type '{type(self._view)}'.") + pos_wp, quat_wp = self._view.get_world_poses() + pos_torch = wp.to_torch(pos_wp).reshape(-1, 3) + quat_torch = wp.to_torch(quat_wp).reshape(-1, 4) + poses = torch.cat([pos_torch, quat_torch], dim=-1).contiguous() + return wp.from_torch(poses).view(wp.transformf) def _update_ray_infos(self, env_mask: wp.array): """Updates sensor poses and ray world-frame buffers via a single warp kernel.""" @@ -385,86 +359,6 @@ def _debug_vis_callback(self, event): self.ray_visualizer.visualize(viz_points) - """ - Internal Helpers. - """ - - def _obtain_trackable_prim_view( - self, target_prim_path: str - ) -> tuple[XformPrimView | physx.ArticulationView | physx.RigidBodyView, tuple[torch.Tensor, torch.Tensor]]: - """Obtain a prim view that can be used to track the pose of the target prim. - - The target prim path is a regex expression that matches one or more mesh prims. While we can track its - pose directly using XFormPrim, this is not efficient and can be slow. Instead, we create a prim view - using the physics simulation view, which provides a more efficient way to track the pose of the mesh prims. - - The function additionally resolves the relative pose between the mesh and its corresponding physics prim. - This is especially useful if the mesh is not directly parented to the physics prim. - - Args: - target_prim_path: The target prim path to obtain the prim view for. - - Returns: - A tuple containing: - - - An XFormPrim or a physics prim view (ArticulationView or RigidBodyView). - - A tuple containing the positions and orientations of the mesh prims in the physics prim frame. - - """ - - mesh_prim = sim_utils.find_first_matching_prim(target_prim_path) - current_prim = mesh_prim - current_path_expr = target_prim_path - - prim_view = None - - while prim_view is None: - if current_prim.HasAPI(UsdPhysics.ArticulationRootAPI): - prim_view = self._physics_sim_view.create_articulation_view(current_path_expr.replace(".*", "*")) - logger.info(f"Created articulation view for mesh prim at path: {target_prim_path}") - break - - if current_prim.HasAPI(UsdPhysics.RigidBodyAPI): - prim_view = self._physics_sim_view.create_rigid_body_view(current_path_expr.replace(".*", "*")) - logger.info(f"Created rigid body view for mesh prim at path: {target_prim_path}") - break - - new_root_prim = current_prim.GetParent() - current_path_expr = current_path_expr.rsplit("/", 1)[0] - if not new_root_prim.IsValid(): - prim_view = XformPrimView(target_prim_path, device=self._device, stage=self.stage) - current_path_expr = target_prim_path - logger.warning( - f"The prim at path {target_prim_path} which is used for raycasting is not a physics prim." - " Defaulting to XFormPrim. \n The pose of the mesh will most likely not" - " be updated correctly when running in headless mode and position lookups will be much slower. \n" - " If possible, ensure that the mesh or its parent is a physics prim (rigid body or articulation)." - ) - break - - current_prim = new_root_prim - - # obtain the relative transforms between target prim and the view prims - mesh_prims = sim_utils.find_matching_prims(target_prim_path) - view_prims = sim_utils.find_matching_prims(current_path_expr) - if len(mesh_prims) != len(view_prims): - raise RuntimeError( - f"The number of mesh prims ({len(mesh_prims)}) does not match the number of physics prims" - f" ({len(view_prims)})Please specify the correct mesh and physics prim paths more" - " specifically in your target expressions." - ) - positions = [] - quaternions = [] - for mesh_prim, view_prim in zip(mesh_prims, view_prims): - pos, orientation = sim_utils.resolve_prim_pose(mesh_prim, view_prim) - positions.append(torch.tensor(pos, dtype=torch.float32, device=self.device)) - quaternions.append(torch.tensor(orientation, dtype=torch.float32, device=self.device)) - - positions = torch.stack(positions).to(device=self.device, dtype=torch.float32) - quaternions = torch.stack(quaternions).to(device=self.device, dtype=torch.float32) - - return prim_view, (positions, quaternions) - """ Internal simulation callbacks. """ diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py index b9f53ea8c491..17bb1e601980 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py @@ -27,7 +27,6 @@ fill_vec3_inf_kernel, update_ray_caster_kernel, ) -from .ray_cast_utils import obtain_world_pose_from_view from .ray_caster import RayCaster if TYPE_CHECKING: @@ -168,9 +167,13 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) elif env_ids is None or isinstance(env_ids, slice): env_ids = self._ALL_INDICES + if not isinstance(env_ids, torch.Tensor): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self._device) # reset the data # note: this recomputation is useful if one performs events such as randomizations on the camera poses. - pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids, clone=True) + indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None + pos_wp, quat_wp = self._view.get_world_poses(indices) + pos_w, quat_w = wp.to_torch(pos_wp).clone(), wp.to_torch(quat_wp).clone() pos_w, quat_w = math_utils.combine_frame_transforms( pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] ) @@ -214,7 +217,9 @@ def set_world_poses( env_ids = self._ALL_INDICES # get current positions - pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids) + indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None + pos_wp, quat_wp = self._view.get_world_poses(indices) + pos_w, quat_w = wp.to_torch(pos_wp), wp.to_torch(quat_wp) if positions is not None: # transform to camera frame pos_offset_world_frame = positions - pos_w @@ -227,7 +232,8 @@ def set_world_poses( self._offset_quat[env_ids] = math_utils.quat_mul(math_utils.quat_inv(quat_w), quat_w_set) # update the data - pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids, clone=True) + pos_wp2, quat_wp2 = self._view.get_world_poses(indices) + pos_w, quat_w = wp.to_torch(pos_wp2).clone(), wp.to_torch(quat_wp2).clone() pos_w, quat_w = math_utils.combine_frame_transforms( pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] ) @@ -574,21 +580,22 @@ def _compute_view_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tenso """Obtains the pose of the view the camera is attached to in the world frame. .. deprecated v2.3.1: - This function will be removed in a future release in favor of implementation - :meth:`obtain_world_pose_from_view`. + This function will be removed in a future release. Call + ``self._view.get_world_poses(indices)`` directly instead. Returns: A tuple of the position (in meters) and quaternion (x, y, z, w). """ - # deprecation logger.warning( - "The function '_compute_view_world_poses' will be deprecated in favor of the util method" - " 'obtain_world_pose_from_view'. Please use 'obtain_world_pose_from_view' instead...." + "The function '_compute_view_world_poses' is deprecated." + " Call 'self._view.get_world_poses(indices)' directly instead." ) - return obtain_world_pose_from_view(self._view, env_ids, clone=True) + indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None + pos_wp, quat_wp = self._view.get_world_poses(indices) + return wp.to_torch(pos_wp).clone(), wp.to_torch(quat_wp).clone() def _compute_camera_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]: """Computes the pose of the camera in the world frame. @@ -600,7 +607,9 @@ def _compute_camera_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Ten .. code-block:: python - pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids, clone=True) + indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) + pos_wp, quat_wp = self._view.get_world_poses(indices) + pos_w, quat_w = wp.to_torch(pos_wp).clone(), wp.to_torch(quat_wp).clone() pos_w, quat_w = math_utils.combine_frame_transforms( pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] ) @@ -608,14 +617,12 @@ def _compute_camera_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Ten Returns: A tuple of the position (in meters) and quaternion (x, y, z, w) in "world" convention. """ - - # deprecation logger.warning( - "The function '_compute_camera_world_poses' will be deprecated in favor of the combination of methods" - " 'obtain_world_pose_from_view' and 'math_utils.combine_frame_transforms'. Please use" - " 'obtain_world_pose_from_view' and 'math_utils.combine_frame_transforms' instead...." + "The function '_compute_camera_world_poses' is deprecated." + " Call 'self._view.get_world_poses(indices)' and 'math_utils.combine_frame_transforms' directly instead." ) - # get the pose of the view the camera is attached to - pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids, clone=True) + indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None + pos_wp, quat_wp = self._view.get_world_poses(indices) + pos_w, quat_w = wp.to_torch(pos_wp).clone(), wp.to_torch(quat_wp).clone() return math_utils.combine_frame_transforms(pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids]) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_cfg.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_cfg.py index 9c3ef14091c5..3e862e389c1e 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_cfg.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_cfg.py @@ -12,6 +12,7 @@ from isaaclab.markers import VisualizationMarkersCfg from isaaclab.markers.config import RAY_CASTER_MARKER_CFG +from isaaclab.sim.spawners.sensors.sensors_cfg import SensorFrameCfg from isaaclab.utils import configclass from ..sensor_base_cfg import SensorBaseCfg @@ -36,6 +37,21 @@ class OffsetCfg: class_type: type[RayCaster] | str = "{DIR}.ray_caster:RayCaster" + spawn: SensorFrameCfg | None = SensorFrameCfg() + """Spawn configuration for the sensor Xform prim. + + A plain USD Xform is created at :attr:`prim_path` before initialization, matching the + pattern used by :class:`~isaaclab.sensors.camera.camera_cfg.CameraCfg` (which spawns a + Camera prim). The :attr:`prim_path` can be either: + + - A **new** child path under a parent link (e.g. ``{ENV_REGEX_NS}/Robot/base``). + - A **physics body** path (e.g. ``{ENV_REGEX_NS}/Robot/base``). In this case, the sensor + will automatically create a child Xform at ``{prim_path}``. + + If ``None``, the prim at :attr:`prim_path` must already exist on the USD stage and must + **not** be a physics body. + """ + mesh_prim_paths: list[str] = MISSING """The list of mesh primitive paths to ray cast against. @@ -47,22 +63,6 @@ class OffsetCfg: offset: OffsetCfg = OffsetCfg() """The offset pose of the sensor's frame from the sensor's parent frame. Defaults to identity.""" - attach_yaw_only: bool | None = None - """Whether the rays' starting positions and directions only track the yaw orientation. - Defaults to None, which doesn't raise a warning of deprecated usage. - - This is useful for ray-casting height maps, where only yaw rotation is needed. - - .. deprecated:: 2.1.1 - - This attribute is deprecated and will be removed in the future. Please use - :attr:`ray_alignment` instead. - - To get the same behavior as setting this parameter to ``True`` or ``False``, set - :attr:`ray_alignment` to ``"yaw"`` or "base" respectively. - - """ - ray_alignment: Literal["base", "yaw", "world"] = "base" """Specify in what frame the rays are projected onto the ground. Default is "base". diff --git a/source/isaaclab/isaaclab/sensors/sensor_base.py b/source/isaaclab/isaaclab/sensors/sensor_base.py index 4a9fb91786e5..3b15d8a0171e 100644 --- a/source/isaaclab/isaaclab/sensors/sensor_base.py +++ b/source/isaaclab/isaaclab/sensors/sensor_base.py @@ -12,6 +12,7 @@ from __future__ import annotations import inspect +import logging import re import weakref from abc import ABC, abstractmethod @@ -29,6 +30,8 @@ if TYPE_CHECKING: from .sensor_base_cfg import SensorBaseCfg +logger = logging.getLogger(__name__) + class SensorBase(ABC): """The base class for implementing a sensor. @@ -386,3 +389,71 @@ def _resolve_indices_and_mask( self._reset_mask.zero_() self._reset_mask_torch[env_ids] = True return self._reset_mask + + def _resolve_and_spawn(self, sensor_name: str, **spawn_kwargs) -> None: + """Resolve physics-body prim paths and spawn the sensor prim if needed. + + Behavior matrix (``spawn`` refers to ``cfg.spawn``): + + +----------------+------------------+--------------------------------------------+ + | ``spawn`` | ``prim_path`` | Action | + +================+==================+============================================+ + | not ``None`` | physics body | Append ``/``, spawn child. | + +----------------+------------------+--------------------------------------------+ + | not ``None`` | non-physics prim | Use existing prim, skip spawn. | + | | (already exists) | | + +----------------+------------------+--------------------------------------------+ + | not ``None`` | does not exist | Spawn prim at ``prim_path``. | + +----------------+------------------+--------------------------------------------+ + | ``None`` | physics body | Raise ``ValueError``. | + +----------------+------------------+--------------------------------------------+ + | ``None`` | non-physics prim | Use as-is (no spawn). | + +----------------+------------------+--------------------------------------------+ + + Args: + sensor_name: Short identifier (e.g. ``"raycaster"``, ``"camera"``). + **spawn_kwargs: Extra keyword arguments forwarded to ``cfg.spawn.func`` + (e.g. ``translation``, ``orientation``). + + Raises: + ValueError: If ``spawn`` is ``None`` and ``prim_path`` is a physics body. + RuntimeError: If the prim does not exist after the spawn attempt. + """ + from pxr import UsdPhysics # noqa: PLC0415 + + spawn = getattr(self.cfg, "spawn", None) + has_spawn = spawn is not None + + # Determine the path to probe for physics-body redirect + spawn_path = (getattr(spawn, "spawn_path", None) or self.cfg.prim_path) if has_spawn else None + probe_path = spawn_path if spawn_path is not None else self.cfg.prim_path + + prim = sim_utils.find_first_matching_prim(probe_path) + if prim is not None and prim.IsValid(): + is_physics = prim.HasAPI(UsdPhysics.ArticulationRootAPI) or prim.HasAPI(UsdPhysics.RigidBodyAPI) + if is_physics: + if not has_spawn: + raise ValueError( + f"Sensor prim_path '{self.cfg.prim_path}' resolves to a physics body but" + f" no spawner is configured (spawn=None). Either set spawn or point" + f" prim_path at a non-physics child (e.g. '{self.cfg.prim_path}/{sensor_name}')." + ) + logger.info( + f"Sensor prim_path '{self.cfg.prim_path}' points at a physics body." + f" Redirecting to '{self.cfg.prim_path}/{sensor_name}'." + ) + self.cfg.prim_path = f"{self.cfg.prim_path}/{sensor_name}" + if getattr(spawn, "spawn_path", None) is not None: + spawn.spawn_path = f"{spawn.spawn_path}/{sensor_name}" + + if not has_spawn: + return + + spawn_target = getattr(spawn, "spawn_path", None) or self.cfg.prim_path + prim = sim_utils.find_first_matching_prim(spawn_target) + if prim is None or not prim.IsValid(): + spawn.func(spawn_target, spawn, **spawn_kwargs) + + check_path = getattr(spawn, "spawn_path", None) or self.cfg.prim_path + if len(sim_utils.find_matching_prims(check_path)) == 0: + raise RuntimeError(f"Could not find prim with path {check_path!r}.") diff --git a/source/isaaclab/isaaclab/sim/__init__.pyi b/source/isaaclab/isaaclab/sim/__init__.pyi index aa1816845242..a718ccdcb989 100644 --- a/source/isaaclab/isaaclab/sim/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/__init__.pyi @@ -90,8 +90,10 @@ __all__ = [ "MeshSphereCfg", "MeshSquareCfg", "spawn_camera", + "spawn_sensor_frame", "FisheyeCameraCfg", "PinholeCameraCfg", + "SensorFrameCfg", "spawn_capsule", "spawn_cone", "spawn_cuboid", @@ -160,6 +162,10 @@ __all__ = [ "resolve_prim_pose", "resolve_prim_scale", "convert_world_pose_to_local", + "BaseFrameView", + "UsdFrameView", + "FrameView", + # Deprecated alias "XformPrimView", ] @@ -252,8 +258,10 @@ from .spawners import ( MeshSphereCfg, MeshSquareCfg, spawn_camera, + spawn_sensor_frame, FisheyeCameraCfg, PinholeCameraCfg, + SensorFrameCfg, spawn_capsule, spawn_cone, spawn_cuboid, @@ -325,4 +333,5 @@ from .utils import ( resolve_prim_scale, convert_world_pose_to_local, ) -from .views import XformPrimView +from .views import BaseFrameView, UsdFrameView, FrameView +from .views import XformPrimView # deprecated alias diff --git a/source/isaaclab/isaaclab/sim/spawners/__init__.pyi b/source/isaaclab/isaaclab/sim/spawners/__init__.pyi index ba8f6d3d7b69..dae1a432b47e 100644 --- a/source/isaaclab/isaaclab/sim/spawners/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/spawners/__init__.pyi @@ -46,8 +46,10 @@ __all__ = [ "MeshSphereCfg", "MeshSquareCfg", "spawn_camera", + "spawn_sensor_frame", "FisheyeCameraCfg", "PinholeCameraCfg", + "SensorFrameCfg", "spawn_capsule", "spawn_cone", "spawn_cuboid", @@ -113,7 +115,7 @@ from .meshes import ( MeshSquareCfg, MeshSphereCfg, ) -from .sensors import spawn_camera, FisheyeCameraCfg, PinholeCameraCfg +from .sensors import spawn_camera, spawn_sensor_frame, FisheyeCameraCfg, PinholeCameraCfg, SensorFrameCfg from .shapes import ( spawn_capsule, spawn_cone, diff --git a/source/isaaclab/isaaclab/sim/spawners/sensors/__init__.pyi b/source/isaaclab/isaaclab/sim/spawners/sensors/__init__.pyi index 46d24596932c..ba5b96a44c7d 100644 --- a/source/isaaclab/isaaclab/sim/spawners/sensors/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/spawners/sensors/__init__.pyi @@ -5,9 +5,11 @@ __all__ = [ "spawn_camera", + "spawn_sensor_frame", "FisheyeCameraCfg", "PinholeCameraCfg", + "SensorFrameCfg", ] -from .sensors import spawn_camera -from .sensors_cfg import FisheyeCameraCfg, PinholeCameraCfg +from .sensors import spawn_camera, spawn_sensor_frame +from .sensors_cfg import FisheyeCameraCfg, PinholeCameraCfg, SensorFrameCfg diff --git a/source/isaaclab/isaaclab/sim/spawners/sensors/sensors.py b/source/isaaclab/isaaclab/sim/spawners/sensors/sensors.py index 4eb70005e487..db68d21d8a90 100644 --- a/source/isaaclab/isaaclab/sim/spawners/sensors/sensors.py +++ b/source/isaaclab/isaaclab/sim/spawners/sensors/sensors.py @@ -145,3 +145,46 @@ def spawn_camera( prim.GetAttribute(prim_prop_name).Set(param_value) # return the prim return prim + + +@clone +def spawn_sensor_frame( + prim_path: str, + cfg: sensors_cfg.SensorFrameCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Create a plain USD Xform prim as a sensor attachment frame. + + .. note:: + This function is decorated with :func:`clone` that resolves prim path into list of paths + if the input prim path is a regex pattern. + + Args: + prim_path: The prim path or pattern to spawn the asset at. + cfg: The configuration instance. + translation: Local translation (x, y, z) [m] w.r.t. the parent prim. Defaults to None + (origin). + orientation: Local orientation as quaternion (x, y, z, w) w.r.t. the parent prim. + Defaults to None (identity). + **kwargs: Additional keyword arguments, like ``clone_in_fabric``. + + Returns: + The created USD prim. + + Raises: + ValueError: If a prim already exists at the given path. + """ + stage = get_current_stage() + if not stage.GetPrimAtPath(prim_path).IsValid(): + prim = create_prim( + prim_path, + "Xform", + translation=translation, + orientation=orientation, + stage=stage, + ) + else: + raise ValueError(f"A prim already exists at path: '{prim_path}'.") + return prim diff --git a/source/isaaclab/isaaclab/sim/spawners/sensors/sensors_cfg.py b/source/isaaclab/isaaclab/sim/spawners/sensors/sensors_cfg.py index 56c9102cf1f1..1ad9dcd73bf2 100644 --- a/source/isaaclab/isaaclab/sim/spawners/sensors/sensors_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/sensors/sensors_cfg.py @@ -222,3 +222,15 @@ class FisheyeCameraCfg(PinholeCameraCfg): fisheye_polynomial_f: float = 0.0 """Sixth component of fisheye polynomial. Defaults to 0.0.""" + + +@configclass +class SensorFrameCfg(SpawnerCfg): + """Spawns a plain USD Xform as a sensor attachment frame. + + The spawned prim carries no rigid body or collision API. It serves as a + non-physics child under a link so that :class:`~isaaclab.sim.views.FrameView` + can track it on all backends (including Newton, which rejects physics body prims). + """ + + func: Callable | str = "{DIR}.sensors:spawn_sensor_frame" diff --git a/source/isaaclab/isaaclab/sim/views/__init__.pyi b/source/isaaclab/isaaclab/sim/views/__init__.pyi index a666958e4387..d578f85d6ada 100644 --- a/source/isaaclab/isaaclab/sim/views/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/views/__init__.pyi @@ -4,7 +4,15 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "BaseFrameView", + "UsdFrameView", + "FrameView", + # Deprecated alias "XformPrimView", ] +from .base_frame_view import BaseFrameView +from .usd_frame_view import UsdFrameView +from .frame_view import FrameView +# Deprecated alias from .xform_prim_view import XformPrimView diff --git a/source/isaaclab/isaaclab/sim/views/base_frame_view.py b/source/isaaclab/isaaclab/sim/views/base_frame_view.py new file mode 100644 index 000000000000..fc59c2ed83ab --- /dev/null +++ b/source/isaaclab/isaaclab/sim/views/base_frame_view.py @@ -0,0 +1,108 @@ +# 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 + +"""Abstract base class for batched prim transform views.""" + +from __future__ import annotations + +import abc + +import warp as wp + + +class BaseFrameView(abc.ABC): + """Abstract interface for reading and writing world-space transforms of multiple prims. + + Backend-specific implementations (USD/Fabric, Newton GPU state, etc.) subclass + this to provide efficient batched pose queries. The factory + :class:`~isaaclab.sim.views.FrameView` selects the correct + implementation at runtime based on the active physics backend. + + All getters return ``wp.array``. Setters accept ``wp.array``. + """ + + @property + @abc.abstractmethod + def count(self) -> int: + """Number of prims in this view.""" + ... + + @abc.abstractmethod + def get_world_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + """Get world-space positions and orientations for prims in the view. + + Args: + indices: Subset of prims to query. ``None`` means all prims. + + Returns: + A tuple ``(positions (M, 3), orientations (M, 4))`` as ``wp.array``. + """ + ... + + @abc.abstractmethod + def set_world_poses( + self, + positions: wp.array | None = None, + orientations: wp.array | None = None, + indices: wp.array | None = None, + ) -> None: + """Set world-space positions and/or orientations for prims in the view. + + Args: + positions: World-space positions ``(M, 3)``. ``None`` leaves positions unchanged. + orientations: World-space quaternions ``(M, 4)``. ``None`` leaves orientations unchanged. + indices: Subset of prims to update. ``None`` means all prims. + """ + ... + + @abc.abstractmethod + def get_local_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + """Get local-space positions and orientations for prims in the view. + + Args: + indices: Subset of prims to query. ``None`` means all prims. + + Returns: + A tuple ``(translations (M, 3), orientations (M, 4))`` as ``wp.array``. + """ + ... + + @abc.abstractmethod + def set_local_poses( + self, + translations: wp.array | None = None, + orientations: wp.array | None = None, + indices: wp.array | None = None, + ) -> None: + """Set local-space translations and/or orientations for prims in the view. + + Args: + translations: Local-space translations ``(M, 3)``. ``None`` leaves translations unchanged. + orientations: Local-space quaternions ``(M, 4)``. ``None`` leaves orientations unchanged. + indices: Subset of prims to update. ``None`` means all prims. + """ + ... + + @abc.abstractmethod + def get_scales(self, indices: wp.array | None = None) -> wp.array: + """Get scales for prims in the view. + + Args: + indices: Subset of prims to query. ``None`` means all prims. + + Returns: + A ``wp.array`` of shape ``(M, 3)``. + """ + ... + + @abc.abstractmethod + def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set scales for prims in the view. + + Args: + scales: Scales ``(M, 3)`` as ``wp.array``. + indices: Subset of prims to update. ``None`` means all prims. + """ + ... diff --git a/source/isaaclab/isaaclab/sim/views/frame_view.py b/source/isaaclab/isaaclab/sim/views/frame_view.py new file mode 100644 index 000000000000..ea9d5bfbeea9 --- /dev/null +++ b/source/isaaclab/isaaclab/sim/views/frame_view.py @@ -0,0 +1,48 @@ +# 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 + +"""Backend-dispatching FrameView. + +``FrameView(path, device=...)`` automatically selects the right backend: +- PhysX: :class:`~isaaclab_physx.sim.views.FabricFrameView` +- Newton: :class:`~isaaclab_newton.sim.views.NewtonSiteFrameView` +""" + +from __future__ import annotations + +from isaaclab.utils.backend_utils import FactoryBase + +from .base_frame_view import BaseFrameView + + +class FrameView(FactoryBase, BaseFrameView): + """FrameView that dispatches to the active physics backend. + + Callers use ``FrameView(prim_path, device=device)`` and get the + correct implementation automatically: + + - **PhysX / no backend**: :class:`~isaaclab_physx.sim.views.FabricFrameView` + (Fabric GPU acceleration with USD fallback). + - **Newton**: :class:`~isaaclab_newton.sim.views.NewtonSiteFrameView` + (GPU-resident site-based transforms). + """ + + _backend_class_names = {"physx": "FabricFrameView", "newton": "NewtonSiteFrameView"} + + @classmethod + def _get_backend(cls, *args, **kwargs) -> str: + from isaaclab.sim.simulation_context import SimulationContext # noqa: PLC0415 + + ctx = SimulationContext.instance() + if ctx is None: + return "physx" + manager_name = ctx.physics_manager.__name__.lower() + if "newton" in manager_name: + return "newton" + return "physx" + + def __new__(cls, *args, **kwargs) -> BaseFrameView: + """Create a new FrameView for the active physics backend.""" + return super().__new__(cls, *args, **kwargs) diff --git a/source/isaaclab/isaaclab/sim/views/usd_frame_view.py b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py new file mode 100644 index 000000000000..4421fa5391ea --- /dev/null +++ b/source/isaaclab/isaaclab/sim/views/usd_frame_view.py @@ -0,0 +1,359 @@ +# 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 + +from __future__ import annotations + +import logging + +import numpy as np +import torch +import warp as wp + +from pxr import Gf, Sdf, Usd, UsdGeom, Vt + +import isaaclab.sim as sim_utils + +from .base_frame_view import BaseFrameView + +logger = logging.getLogger(__name__) + + +class UsdFrameView(BaseFrameView): + """Batched interface for reading and writing transforms of multiple USD prims. + + Provides batch operations for getting and setting poses (position and orientation) + of multiple prims at once via USD's ``XformCache``. + + The class supports both world-space and local-space pose operations: + + - **World poses**: Positions and orientations in the global world frame + - **Local poses**: Positions and orientations relative to each prim's parent + + For GPU-accelerated Fabric operations, use the PhysX backend variant + obtained via :class:`~isaaclab.sim.views.FrameView`. + + All getters return ``wp.array``. Setters accept ``wp.array``. + + .. note:: + **Transform Requirements:** + + All prims in the view must be Xformable and have standardized transform operations: + ``[translate, orient, scale]``. Non-standard prims will raise a ValueError during + initialization if :attr:`validate_xform_ops` is True. Please use the function + :func:`isaaclab.sim.utils.standardize_xform_ops` to prepare prims before using this view. + + .. warning:: + This class operates at the USD default time code. Any animation or time-sampled data + will not be affected by write operations. For animated transforms, you need to handle + time-sampled keyframes separately. + """ + + def __init__( + self, + prim_path: str, + device: str = "cpu", + validate_xform_ops: bool = True, + stage: Usd.Stage | None = None, + **kwargs, + ): + """Initialize the view with matching prims. + + Args: + prim_path: USD prim path pattern to match prims. Supports wildcards (``*``) and + regex patterns (e.g., ``"/World/Env_.*/Robot"``). See + :func:`isaaclab.sim.utils.find_matching_prims` for pattern syntax. + device: Device to place arrays on. Can be ``"cpu"`` or CUDA devices like + ``"cuda:0"``. Defaults to ``"cpu"``. + validate_xform_ops: Whether to validate that the prims have standard xform operations. + Defaults to True. + stage: USD stage to search for prims. Defaults to None, in which case the current + active stage from the simulation context is used. + **kwargs: Additional keyword arguments (ignored). Allows forward-compatible + construction when callers pass backend-specific options like + ``sync_usd_on_fabric_write``. + + Raises: + ValueError: If any matched prim is not Xformable or doesn't have standardized + transform operations (translate, orient, scale in that order). + """ + self._prim_path = prim_path + self._device = device + + stage = sim_utils.get_current_stage() if stage is None else stage + self._prims: list[Usd.Prim] = sim_utils.find_matching_prims(prim_path, stage=stage) + + if validate_xform_ops: + for prim in self._prims: + sim_utils.standardize_xform_ops(prim) + if not sim_utils.validate_standard_xform_ops(prim): + raise ValueError( + f"Prim at path '{prim.GetPath().pathString}' is not a xformable prim with standard transform" + f" operations [translate, orient, scale]. Received type: '{prim.GetTypeName()}'." + " Use sim_utils.standardize_xform_ops() to prepare the prim." + ) + + self._ALL_INDICES = list(range(len(self._prims))) + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def count(self) -> int: + """Number of prims in this view.""" + return len(self._prims) + + @property + def device(self) -> str: + """Device where arrays are allocated (cpu or cuda).""" + return self._device + + @property + def prims(self) -> list[Usd.Prim]: + """List of USD prims being managed by this view.""" + return self._prims + + @property + def prim_paths(self) -> list[str]: + """List of prim paths (as strings) for all prims being managed by this view. + + The conversion is performed lazily on first access and cached. + """ + if not hasattr(self, "_prim_paths"): + self._prim_paths = [prim.GetPath().pathString for prim in self._prims] + return self._prim_paths + + # ------------------------------------------------------------------ + # Setters + # ------------------------------------------------------------------ + + def set_world_poses( + self, + positions: wp.array | None = None, + orientations: wp.array | None = None, + indices: wp.array | None = None, + ): + """Set world-space poses for prims in the view. + + Converts the desired world pose to local-space relative to each prim's + parent before writing to USD xform ops. + + Args: + positions: World-space positions of shape ``(M, 3)``. + orientations: World-space quaternions ``(w, x, y, z)`` of shape ``(M, 4)``. + indices: Indices of prims to set poses for. Defaults to None (all prims). + """ + indices_list = self._resolve_indices(indices) + + positions_array = Vt.Vec3dArray.FromNumpy(self._to_numpy(positions)) if positions is not None else None + orientations_array = Vt.QuatdArray.FromNumpy(self._to_numpy(orientations)) if orientations is not None else None + + xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + + with Sdf.ChangeBlock(): + for idx, prim_idx in enumerate(indices_list): + prim = self._prims[prim_idx] + parent_prim = prim.GetParent() + + world_pos = positions_array[idx] if positions_array is not None else None + world_quat = orientations_array[idx] if orientations_array is not None else None + + if parent_prim.IsValid() and parent_prim.GetPath() != Sdf.Path.absoluteRootPath: + if positions_array is None or orientations_array is None: + prim_tf = xform_cache.GetLocalToWorldTransform(prim) + prim_tf.Orthonormalize() + if world_pos is not None: + prim_tf.SetTranslateOnly(world_pos) + if world_quat is not None: + prim_tf.SetRotateOnly(world_quat) + else: + prim_tf = Gf.Matrix4d() + prim_tf.SetTranslateOnly(world_pos) + prim_tf.SetRotateOnly(world_quat) + + parent_world_tf = xform_cache.GetLocalToWorldTransform(parent_prim) + local_tf = prim_tf * parent_world_tf.GetInverse() + local_pos = local_tf.ExtractTranslation() + local_quat = local_tf.ExtractRotationQuat() + else: + # Root-level prim: world == local + local_pos = world_pos + local_quat = world_quat + + if local_pos is not None: + prim.GetAttribute("xformOp:translate").Set(local_pos) + if local_quat is not None: + prim.GetAttribute("xformOp:orient").Set(local_quat) + + def set_local_poses( + self, + translations: wp.array | None = None, + orientations: wp.array | None = None, + indices: wp.array | None = None, + ): + """Set local-space poses for prims in the view. + + Args: + translations: Local-space translations of shape ``(M, 3)``. + orientations: Local-space quaternions ``(w, x, y, z)`` of shape ``(M, 4)``. + indices: Indices of prims to set poses for. Defaults to None (all prims). + """ + indices_list = self._resolve_indices(indices) + + translations_array = Vt.Vec3dArray.FromNumpy(self._to_numpy(translations)) if translations is not None else None + orientations_array = Vt.QuatdArray.FromNumpy(self._to_numpy(orientations)) if orientations is not None else None + + with Sdf.ChangeBlock(): + for idx, prim_idx in enumerate(indices_list): + prim = self._prims[prim_idx] + if translations_array is not None: + prim.GetAttribute("xformOp:translate").Set(translations_array[idx]) + if orientations_array is not None: + prim.GetAttribute("xformOp:orient").Set(orientations_array[idx]) + + def set_scales(self, scales: wp.array, indices: wp.array | None = None): + """Set scales for prims in the view. + + Args: + scales: Scales of shape ``(M, 3)``. + indices: Indices of prims to set scales for. Defaults to None (all prims). + """ + indices_list = self._resolve_indices(indices) + scales_array = Vt.Vec3dArray.FromNumpy(self._to_numpy(scales)) + + with Sdf.ChangeBlock(): + for idx, prim_idx in enumerate(indices_list): + prim = self._prims[prim_idx] + prim.GetAttribute("xformOp:scale").Set(scales_array[idx]) + + def set_visibility(self, visibility: torch.Tensor, indices: wp.array | None = None): + """Set visibility for prims in the view. + + Args: + visibility: Visibility as a boolean tensor of shape ``(M,)``. + indices: Indices of prims to set visibility for. Defaults to None (all prims). + """ + indices_list = self._resolve_indices(indices) + + if visibility.shape != (len(indices_list),): + raise ValueError(f"Expected visibility shape ({len(indices_list)},), got {visibility.shape}.") + + with Sdf.ChangeBlock(): + for idx, prim_idx in enumerate(indices_list): + imageable = UsdGeom.Imageable(self._prims[prim_idx]) + if visibility[idx]: + imageable.MakeVisible() + else: + imageable.MakeInvisible() + + # ------------------------------------------------------------------ + # Getters + # ------------------------------------------------------------------ + + def get_world_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + """Get world-space poses for prims in the view. + + Args: + indices: Indices of prims to get poses for. Defaults to None (all prims). + + Returns: + A tuple of ``(positions, orientations)`` as ``wp.array``. + """ + indices_list = self._resolve_indices(indices) + + positions = Vt.Vec3dArray(len(indices_list)) + orientations = Vt.QuatdArray(len(indices_list)) + xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + + for idx, prim_idx in enumerate(indices_list): + prim = self._prims[prim_idx] + prim_tf = xform_cache.GetLocalToWorldTransform(prim) + prim_tf.Orthonormalize() + positions[idx] = prim_tf.ExtractTranslation() + orientations[idx] = prim_tf.ExtractRotationQuat() + + return ( + wp.array(np.array(positions, dtype=np.float32), dtype=wp.float32, device=self._device), + wp.array(np.array(orientations, dtype=np.float32), dtype=wp.float32, device=self._device), + ) + + def get_local_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + """Get local-space poses for prims in the view. + + Args: + indices: Indices of prims to get poses for. Defaults to None (all prims). + + Returns: + A tuple of ``(translations, orientations)`` as ``wp.array``. + """ + indices_list = self._resolve_indices(indices) + + translations = Vt.Vec3dArray(len(indices_list)) + orientations = Vt.QuatdArray(len(indices_list)) + xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + + for idx, prim_idx in enumerate(indices_list): + prim = self._prims[prim_idx] + prim_tf = xform_cache.GetLocalTransformation(prim)[0] + prim_tf.Orthonormalize() + translations[idx] = prim_tf.ExtractTranslation() + orientations[idx] = prim_tf.ExtractRotationQuat() + + return ( + wp.array(np.array(translations, dtype=np.float32), dtype=wp.float32, device=self._device), + wp.array(np.array(orientations, dtype=np.float32), dtype=wp.float32, device=self._device), + ) + + def get_scales(self, indices: wp.array | None = None) -> wp.array: + """Get scales for prims in the view. + + Args: + indices: Indices of prims to get scales for. Defaults to None (all prims). + + Returns: + A ``wp.array`` of shape ``(M, 3)``. + """ + indices_list = self._resolve_indices(indices) + + scales = Vt.Vec3dArray(len(indices_list)) + for idx, prim_idx in enumerate(indices_list): + prim = self._prims[prim_idx] + scales[idx] = prim.GetAttribute("xformOp:scale").Get() + + return wp.array(np.array(scales, dtype=np.float32), dtype=wp.float32, device=self._device) + + def get_visibility(self, indices: wp.array | None = None) -> torch.Tensor: + """Get visibility for prims in the view. + + Args: + indices: Indices of prims to get visibility for. Defaults to None (all prims). + + Returns: + A tensor of shape ``(M,)`` containing the visibility of each prim (bool). + """ + indices_list = self._resolve_indices(indices) + + visibility = torch.zeros(len(indices_list), dtype=torch.bool, device=self._device) + for idx, prim_idx in enumerate(indices_list): + imageable = UsdGeom.Imageable(self._prims[prim_idx]) + visibility[idx] = imageable.ComputeVisibility() != UsdGeom.Tokens.invisible + return visibility + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _resolve_indices(self, indices: wp.array | None): + """Resolve warp indices to an iterable of ints for per-prim USD operations.""" + if indices is None or indices == slice(None): + return self._ALL_INDICES + return indices.numpy() + + @staticmethod + def _to_numpy(data: wp.array | torch.Tensor) -> np.ndarray: + """Convert a ``wp.array`` or ``torch.Tensor`` to a numpy array on CPU.""" + if isinstance(data, wp.array): + return data.numpy() + return data.cpu().numpy() diff --git a/source/isaaclab/isaaclab/sim/views/xform_prim_view.py b/source/isaaclab/isaaclab/sim/views/xform_prim_view.py index 211994a7226b..ce480fa65594 100644 --- a/source/isaaclab/isaaclab/sim/views/xform_prim_view.py +++ b/source/isaaclab/isaaclab/sim/views/xform_prim_view.py @@ -3,1138 +3,8 @@ # # SPDX-License-Identifier: BSD-3-Clause -from __future__ import annotations +"""Backward-compatibility alias: ``XformPrimView`` -> :class:`FrameView`.""" -import logging -from collections.abc import Sequence +from .frame_view import FrameView -import numpy as np -import torch -import warp as wp - -from pxr import Gf, Sdf, Usd, UsdGeom, Vt - -import isaaclab.sim as sim_utils -from isaaclab.app.settings_manager import SettingsManager -from isaaclab.utils.warp import fabric as fabric_utils - -logger = logging.getLogger(__name__) - - -class XformPrimView: - """Optimized batched interface for reading and writing transforms of multiple USD prims. - - This class provides efficient batch operations for getting and setting poses (position and orientation) - of multiple prims at once using torch tensors. It is designed for scenarios where you need to manipulate - many prims simultaneously, such as in multi-agent simulations or large-scale procedural generation. - - The class supports both world-space and local-space pose operations: - - - **World poses**: Positions and orientations in the global world frame - - **Local poses**: Positions and orientations relative to each prim's parent - - When Fabric is enabled, the class leverages NVIDIA's Fabric API for GPU-accelerated batch operations: - - - Uses `omni:fabric:worldMatrix` and `omni:fabric:localMatrix` attributes for all Boundable prims - - Performs batch matrix decomposition/composition using Warp kernels on GPU - - Achieves performance comparable to Isaac Sim's XFormPrim implementation - - Works for both physics-enabled and non-physics prims (cameras, meshes, etc.). - Note: renderers typically consume USD-authored camera transforms. - - .. warning:: - **Fabric requires CUDA**: Fabric is only supported with on CUDA devices. - Warp's CPU backend for fabric-array writes has known issues, so attempting to use - Fabric with CPU device (``device="cpu"``) will raise a ValueError at initialization. - - .. note:: - **Fabric Support:** - - When Fabric is enabled, this view ensures prims have the required Fabric hierarchy - attributes (``omni:fabric:localMatrix`` and ``omni:fabric:worldMatrix``). On first Fabric - read, USD-authored transforms initialize Fabric state. Fabric writes can optionally - be mirrored back to USD via :attr:`sync_usd_on_fabric_write`. - - For more information, see the `Fabric Hierarchy documentation`_. - - .. _Fabric Hierarchy documentation: https://docs.omniverse.nvidia.com/kit/docs/usdrt/latest/docs/fabric_hierarchy.html - - .. note:: - **Performance Considerations:** - - * Tensor operations are performed on the specified device (CPU/CUDA) - * USD write operations use ``Sdf.ChangeBlock`` for batched updates - * Fabric operations use GPU-accelerated Warp kernels for maximum performance - * For maximum performance, minimize get/set operations within tight loops - - .. note:: - **Transform Requirements:** - - All prims in the view must be Xformable and have standardized transform operations: - ``[translate, orient, scale]``. Non-standard prims will raise a ValueError during - initialization if :attr:`validate_xform_ops` is True. Please use the function - :func:`isaaclab.sim.utils.standardize_xform_ops` to prepare prims before using this view. - - .. warning:: - This class operates at the USD default time code. Any animation or time-sampled data - will not be affected by write operations. For animated transforms, you need to handle - time-sampled keyframes separately. - """ - - def __init__( - self, - prim_path: str, - device: str = "cpu", - validate_xform_ops: bool = True, - sync_usd_on_fabric_write: bool = False, - stage: Usd.Stage | None = None, - ): - """Initialize the view with matching prims. - - This method searches the USD stage for all prims matching the provided path pattern, - validates that they are Xformable with standard transform operations, and stores - references for efficient batch operations. - - We generally recommend to validate the xform operations, as it ensures that the prims are in a consistent state - and have the standard transform operations (translate, orient, scale in that order). - However, if you are sure that the prims are in a consistent state, you can set this to False to improve - performance. This can save around 45-50% of the time taken to initialize the view. - - Args: - prim_path: USD prim path pattern to match prims. Supports wildcards (``*``) and - regex patterns (e.g., ``"/World/Env_.*/Robot"``). See - :func:`isaaclab.sim.utils.find_matching_prims` for pattern syntax. - device: Device to place the tensors on. Can be ``"cpu"`` or CUDA devices like - ``"cuda:0"``. Defaults to ``"cpu"``. - validate_xform_ops: Whether to validate that the prims have standard xform operations. - Defaults to True. - sync_usd_on_fabric_write: Whether to mirror Fabric transform writes back to USD. - When True, transform updates are synchronized to USD so that USD data readers (e.g., rendering - cameras) can observe these changes. Defaults to False for better performance. - stage: USD stage to search for prims. Defaults to None, in which case the current active stage - from the simulation context is used. - - Raises: - ValueError: If any matched prim is not Xformable or doesn't have standardized - transform operations (translate, orient, scale in that order). - """ - # Store configuration - self._prim_path = prim_path - self._device = device - - # Find and validate matching prims - stage = sim_utils.get_current_stage() if stage is None else stage - self._prims: list[Usd.Prim] = sim_utils.find_matching_prims(prim_path, stage=stage) - - # Validate all prims have standard xform operations - if validate_xform_ops: - for prim in self._prims: - sim_utils.standardize_xform_ops(prim) - if not sim_utils.validate_standard_xform_ops(prim): - raise ValueError( - f"Prim at path '{prim.GetPath().pathString}' is not a xformable prim with standard transform" - f" operations [translate, orient, scale]. Received type: '{prim.GetTypeName()}'." - " Use sim_utils.standardize_xform_ops() to prepare the prim." - ) - - # Determine if Fabric is supported on the device - settings = SettingsManager.instance() - self._use_fabric = bool(settings.get("/physics/fabricEnabled", False)) - - # Check for unsupported Fabric + CPU combination - if self._use_fabric and self._device == "cpu": - logger.warning( - "Fabric mode with Warp fabric-array operations is not supported on CPU devices. " - "While Fabric itself can run on both CPU and GPU, our batch Warp kernels for " - "fabric-array operations require CUDA and are not reliable on the CPU backend. " - "To ensure stability, Fabric is being disabled and execution will fall back " - "to standard USD operations on the CPU. This may impact performance." - ) - self._use_fabric = False - - # Check for unsupported Fabric + non-primary CUDA device combination. - # USDRT SelectPrims and Warp fabric arrays only support cuda:0 internally. - # When running on cuda:1 or higher, SelectPrims raises a C++ error regardless of - # the device argument, because USDRT uses the active CUDA context (which is cuda:1). - if self._use_fabric and self._device not in ("cuda", "cuda:0"): - logger.warning( - f"Fabric mode is not supported on device '{self._device}'. " - "USDRT SelectPrims and Warp fabric arrays only support cuda:0. " - "Falling back to standard USD operations. This may impact performance." - ) - self._use_fabric = False - - # Create indices buffer - # Since we iterate over the indices, we need to use range instead of torch tensor - self._ALL_INDICES = list(range(len(self._prims))) - - # Some prims (e.g., Cameras) require USD-authored transforms for rendering. - # When enabled, mirror Fabric pose writes to USD for those prims. - self._sync_usd_on_fabric_write = sync_usd_on_fabric_write - - # Fabric batch infrastructure (initialized lazily on first use) - self._fabric_initialized = False - self._fabric_usd_sync_done = False - self._fabric_selection = None - self._fabric_to_view: wp.array | None = None - self._view_to_fabric: wp.array | None = None - self._default_view_indices: wp.array | None = None - self._fabric_hierarchy = None - # Create a valid USD attribute name: namespace:name - # Use "isaaclab" namespace to identify our custom attributes - self._view_index_attr = f"isaaclab:view_index:{abs(hash(self))}" - - """ - Properties. - """ - - @property - def count(self) -> int: - """Number of prims in this view.""" - return len(self._prims) - - @property - def device(self) -> str: - """Device where tensors are allocated (cpu or cuda).""" - return self._device - - @property - def prims(self) -> list[Usd.Prim]: - """List of USD prims being managed by this view.""" - return self._prims - - @property - def prim_paths(self) -> list[str]: - """List of prim paths (as strings) for all prims being managed by this view. - - This property converts each prim to its path string representation. The conversion is - performed lazily on first access and cached for subsequent accesses. - - Note: - For most use cases, prefer using :attr:`prims` directly as it provides direct access - to the USD prim objects without the conversion overhead. This property is mainly useful - for logging, debugging, or when string paths are explicitly required. - """ - # we cache it the first time it is accessed. - # we don't compute it in constructor because it is expensive and we don't need it most of the time. - # users should usually deal with prims directly as they typically need to access the prims directly. - if not hasattr(self, "_prim_paths"): - self._prim_paths = [prim.GetPath().pathString for prim in self._prims] - return self._prim_paths - - """ - Operations - Setters. - """ - - def set_world_poses( - self, - positions: torch.Tensor | None = None, - orientations: torch.Tensor | None = None, - indices: Sequence[int] | None = None, - ): - """Set world-space poses for prims in the view. - - This method sets the position and/or orientation of each prim in world space. - - - When Fabric is enabled, the function writes directly to Fabric's ``omni:fabric:worldMatrix`` - attribute using GPU-accelerated batch operations. - - When Fabric is disabled, the function converts to local space and writes to USD's ``xformOp:translate`` - and ``xformOp:orient`` attributes. - - Args: - positions: World-space positions as a tensor of shape (M, 3) where M is the number of prims - to set (either all prims if indices is None, or the number of indices provided). - Defaults to None, in which case positions are not modified. - orientations: World-space orientations as quaternions (w, x, y, z) with shape (M, 4). - Defaults to None, in which case orientations are not modified. - indices: Indices of prims to set poses for. Defaults to None, in which case poses are set - for all prims in the view. - - Raises: - ValueError: If positions shape is not (M, 3) or orientations shape is not (M, 4). - ValueError: If the number of poses doesn't match the number of indices provided. - """ - if self._use_fabric: - self._set_world_poses_fabric(positions, orientations, indices) - else: - self._set_world_poses_usd(positions, orientations, indices) - - def set_local_poses( - self, - translations: torch.Tensor | None = None, - orientations: torch.Tensor | None = None, - indices: Sequence[int] | None = None, - ): - """Set local-space poses for prims in the view. - - This method sets the position and/or orientation of each prim in local space (relative to - their parent prims). - - The function writes directly to USD's ``xformOp:translate`` and ``xformOp:orient`` attributes. - - Note: - Even in Fabric mode, local pose operations use USD. This behavior is based on Isaac Sim's design - where Fabric is only used for world pose operations. - - Rationale: - - Local pose writes need correct parent-child hierarchy relationships - - USD maintains these relationships correctly and efficiently - - Fabric is optimized for world pose operations, not local hierarchies - - Args: - translations: Local-space translations as a tensor of shape (M, 3) where M is the number of prims - to set (either all prims if indices is None, or the number of indices provided). - Defaults to None, in which case translations are not modified. - orientations: Local-space orientations as quaternions (w, x, y, z) with shape (M, 4). - Defaults to None, in which case orientations are not modified. - indices: Indices of prims to set poses for. Defaults to None, in which case poses are set - for all prims in the view. - - Raises: - ValueError: If translations shape is not (M, 3) or orientations shape is not (M, 4). - ValueError: If the number of poses doesn't match the number of indices provided. - """ - if self._use_fabric: - self._set_local_poses_fabric(translations, orientations, indices) - else: - self._set_local_poses_usd(translations, orientations, indices) - - def set_scales(self, scales: torch.Tensor, indices: Sequence[int] | None = None): - """Set scales for prims in the view. - - This method sets the scale of each prim in the view. - - - When Fabric is enabled, the function updates scales in Fabric matrices using GPU-accelerated batch operations. - - When Fabric is disabled, the function writes to USD's ``xformOp:scale`` attributes. - - Args: - scales: Scales as a tensor of shape (M, 3) where M is the number of prims - to set (either all prims if indices is None, or the number of indices provided). - indices: Indices of prims to set scales for. Defaults to None, in which case scales are set - for all prims in the view. - - Raises: - ValueError: If scales shape is not (M, 3). - """ - if self._use_fabric: - self._set_scales_fabric(scales, indices) - else: - self._set_scales_usd(scales, indices) - - def set_visibility(self, visibility: torch.Tensor, indices: Sequence[int] | None = None): - """Set visibility for prims in the view. - - This method sets the visibility of each prim in the view. - - Args: - visibility: Visibility as a boolean tensor of shape (M,) where M is the - number of prims to set (either all prims if indices is None, or the number of indices provided). - indices: Indices of prims to set visibility for. Defaults to None, in which case visibility is set - for all prims in the view. - - Raises: - ValueError: If visibility shape is not (M,). - """ - # Resolve indices - if indices is None or indices == slice(None): - indices_list = self._ALL_INDICES - else: - indices_list = indices.tolist() if isinstance(indices, torch.Tensor) else list(indices) - - # Validate inputs - if visibility.shape != (len(indices_list),): - raise ValueError(f"Expected visibility shape ({len(indices_list)},), got {visibility.shape}.") - - # Set visibility for each prim - with Sdf.ChangeBlock(): - for idx, prim_idx in enumerate(indices_list): - # Convert prim to imageable - imageable = UsdGeom.Imageable(self._prims[prim_idx]) - # Set visibility - if visibility[idx]: - imageable.MakeVisible() - else: - imageable.MakeInvisible() - - """ - Operations - Getters. - """ - - def get_world_poses(self, indices: Sequence[int] | None = None) -> tuple[torch.Tensor, torch.Tensor]: - """Get world-space poses for prims in the view. - - This method retrieves the position and orientation of each prim in world space by computing - the full transform hierarchy from the prim to the world root. - - - When Fabric is enabled, the function uses Fabric batch operations with Warp kernels. - - When Fabric is disabled, the function uses USD XformCache. - - Note: - Scale and skew are ignored. The returned poses contain only translation and rotation. - - Args: - indices: Indices of prims to get poses for. Defaults to None, in which case poses are retrieved - for all prims in the view. - - Returns: - A tuple of (positions, orientations) where: - - - positions: Torch tensor of shape (M, 3) containing world-space positions (x, y, z), - where M is the number of prims queried. - - orientations: Torch tensor of shape (M, 4) containing world-space quaternions (w, x, y, z) - """ - if self._use_fabric: - return self._get_world_poses_fabric(indices) - else: - return self._get_world_poses_usd(indices) - - def get_local_poses(self, indices: Sequence[int] | None = None) -> tuple[torch.Tensor, torch.Tensor]: - """Get local-space poses for prims in the view. - - This method retrieves the position and orientation of each prim in local space (relative to - their parent prims). It reads directly from USD's ``xformOp:translate`` and ``xformOp:orient`` attributes. - - Note: - Even in Fabric mode, local pose operations use USD. This behavior is based on Isaac Sim's design - where Fabric is only used for world pose operations. - - Rationale: - - Local pose reads need correct parent-child hierarchy relationships - - USD maintains these relationships correctly and efficiently - - Fabric is optimized for world pose operations, not local hierarchies - - Note: - Scale is ignored. The returned poses contain only translation and rotation. - - Args: - indices: Indices of prims to get poses for. Defaults to None, in which case poses are retrieved - for all prims in the view. - - Returns: - A tuple of (translations, orientations) where: - - - translations: Torch tensor of shape (M, 3) containing local-space translations (x, y, z), - where M is the number of prims queried. - - orientations: Torch tensor of shape (M, 4) containing local-space quaternions (w, x, y, z) - """ - if self._use_fabric: - return self._get_local_poses_fabric(indices) - else: - return self._get_local_poses_usd(indices) - - def get_scales(self, indices: Sequence[int] | None = None) -> torch.Tensor: - """Get scales for prims in the view. - - This method retrieves the scale of each prim in the view. - - - When Fabric is enabled, the function extracts scales from Fabric matrices using batch operations with - Warp kernels. - - When Fabric is disabled, the function reads from USD's ``xformOp:scale`` attributes. - - Args: - indices: Indices of prims to get scales for. Defaults to None, in which case scales are retrieved - for all prims in the view. - - Returns: - A tensor of shape (M, 3) containing the scales of each prim, where M is the number of prims queried. - """ - if self._use_fabric: - return self._get_scales_fabric(indices) - else: - return self._get_scales_usd(indices) - - def get_visibility(self, indices: Sequence[int] | None = None) -> torch.Tensor: - """Get visibility for prims in the view. - - This method retrieves the visibility of each prim in the view. - - Args: - indices: Indices of prims to get visibility for. Defaults to None, in which case visibility is retrieved - for all prims in the view. - - Returns: - A tensor of shape (M,) containing the visibility of each prim, where M is the number of prims queried. - The tensor is of type bool. - """ - # Resolve indices - if indices is None or indices == slice(None): - indices_list = self._ALL_INDICES - else: - # Convert to list if it is a tensor array - indices_list = indices.tolist() if isinstance(indices, torch.Tensor) else list(indices) - - # Create buffers - visibility = torch.zeros(len(indices_list), dtype=torch.bool, device=self._device) - - for idx, prim_idx in enumerate(indices_list): - # Get prim - imageable = UsdGeom.Imageable(self._prims[prim_idx]) - # Get visibility - visibility[idx] = imageable.ComputeVisibility() != UsdGeom.Tokens.invisible - - return visibility - - """ - Internal Functions - USD. - """ - - def _set_world_poses_usd( - self, - positions: torch.Tensor | None = None, - orientations: torch.Tensor | None = None, - indices: Sequence[int] | None = None, - ): - """Set world poses to USD.""" - # Resolve indices - if indices is None or indices == slice(None): - indices_list = self._ALL_INDICES - else: - # Convert to list if it is a tensor array - indices_list = indices.tolist() if isinstance(indices, torch.Tensor) else list(indices) - - # Validate inputs - if positions is not None: - if positions.shape != (len(indices_list), 3): - raise ValueError( - f"Expected positions shape ({len(indices_list)}, 3), got {positions.shape}. " - "Number of positions must match the number of prims in the view." - ) - positions_array = Vt.Vec3dArray.FromNumpy(positions.cpu().numpy()) - else: - positions_array = None - if orientations is not None: - if orientations.shape != (len(indices_list), 4): - raise ValueError( - f"Expected orientations shape ({len(indices_list)}, 4), got {orientations.shape}. " - "Number of orientations must match the number of prims in the view." - ) - # Vt expects quaternions in xyzw order - orientations_array = Vt.QuatdArray.FromNumpy(orientations.cpu().numpy()) - else: - orientations_array = None - - # Create xform cache instance - xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) - - # Set poses for each prim - # We use Sdf.ChangeBlock to minimize notification overhead. - with Sdf.ChangeBlock(): - for idx, prim_idx in enumerate(indices_list): - # Get prim - prim = self._prims[prim_idx] - # Get parent prim for local space conversion - parent_prim = prim.GetParent() - - # Determine what to set - world_pos = positions_array[idx] if positions_array is not None else None - world_quat = orientations_array[idx] if orientations_array is not None else None - - # Convert world pose to local if we have a valid parent - # Note: We don't use :func:`isaaclab.sim.utils.transforms.convert_world_pose_to_local` - # here since it isn't optimized for batch operations. - if parent_prim.IsValid() and parent_prim.GetPath() != Sdf.Path.absoluteRootPath: - # Get current world pose if we're only setting one component - if positions_array is None or orientations_array is None: - # get prim xform - prim_tf = xform_cache.GetLocalToWorldTransform(prim) - # sanitize quaternion - # this is needed, otherwise the quaternion might be non-normalized - prim_tf.Orthonormalize() - # populate desired world transform - if world_pos is not None: - prim_tf.SetTranslateOnly(world_pos) - if world_quat is not None: - prim_tf.SetRotateOnly(world_quat) - else: - # Both position and orientation are provided, create new transform - prim_tf = Gf.Matrix4d() - prim_tf.SetTranslateOnly(world_pos) - prim_tf.SetRotateOnly(world_quat) - - # Convert to local space - parent_world_tf = xform_cache.GetLocalToWorldTransform(parent_prim) - local_tf = prim_tf * parent_world_tf.GetInverse() - local_pos = local_tf.ExtractTranslation() - local_quat = local_tf.ExtractRotationQuat() - else: - # No parent or parent is root, world == local - local_pos = world_pos - local_quat = world_quat - - # Get or create the standard transform operations - if local_pos is not None: - prim.GetAttribute("xformOp:translate").Set(local_pos) - if local_quat is not None: - prim.GetAttribute("xformOp:orient").Set(local_quat) - - def _set_local_poses_usd( - self, - translations: torch.Tensor | None = None, - orientations: torch.Tensor | None = None, - indices: Sequence[int] | None = None, - ): - """Set local poses to USD.""" - # Resolve indices - if indices is None or indices == slice(None): - indices_list = self._ALL_INDICES - else: - indices_list = indices.tolist() if isinstance(indices, torch.Tensor) else list(indices) - - # Validate inputs - if translations is not None: - if translations.shape != (len(indices_list), 3): - raise ValueError(f"Expected translations shape ({len(indices_list)}, 3), got {translations.shape}.") - translations_array = Vt.Vec3dArray.FromNumpy(translations.cpu().numpy()) - else: - translations_array = None - if orientations is not None: - if orientations.shape != (len(indices_list), 4): - raise ValueError(f"Expected orientations shape ({len(indices_list)}, 4), got {orientations.shape}.") - orientations_array = Vt.QuatdArray.FromNumpy(orientations.cpu().numpy()) - else: - orientations_array = None - - # Set local poses - with Sdf.ChangeBlock(): - for idx, prim_idx in enumerate(indices_list): - prim = self._prims[prim_idx] - if translations_array is not None: - prim.GetAttribute("xformOp:translate").Set(translations_array[idx]) - if orientations_array is not None: - prim.GetAttribute("xformOp:orient").Set(orientations_array[idx]) - - def _set_scales_usd(self, scales: torch.Tensor, indices: Sequence[int] | None = None): - """Set scales to USD.""" - # Resolve indices - if indices is None or indices == slice(None): - indices_list = self._ALL_INDICES - else: - indices_list = indices.tolist() if isinstance(indices, torch.Tensor) else list(indices) - - # Validate inputs - if scales.shape != (len(indices_list), 3): - raise ValueError(f"Expected scales shape ({len(indices_list)}, 3), got {scales.shape}.") - - scales_array = Vt.Vec3dArray.FromNumpy(scales.cpu().numpy()) - # Set scales for each prim - with Sdf.ChangeBlock(): - for idx, prim_idx in enumerate(indices_list): - prim = self._prims[prim_idx] - prim.GetAttribute("xformOp:scale").Set(scales_array[idx]) - - def _get_world_poses_usd(self, indices: Sequence[int] | None = None) -> tuple[torch.Tensor, torch.Tensor]: - """Get world poses from USD.""" - # Resolve indices - if indices is None or indices == slice(None): - indices_list = self._ALL_INDICES - else: - # Convert to list if it is a tensor array - indices_list = indices.tolist() if isinstance(indices, torch.Tensor) else list(indices) - - # Create buffers - positions = Vt.Vec3dArray(len(indices_list)) - orientations = Vt.QuatdArray(len(indices_list)) - # Create xform cache instance - xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) - - # Note: We don't use :func:`isaaclab.sim.utils.transforms.resolve_prim_pose` - # here since it isn't optimized for batch operations. - for idx, prim_idx in enumerate(indices_list): - # Get prim - prim = self._prims[prim_idx] - # get prim xform - prim_tf = xform_cache.GetLocalToWorldTransform(prim) - # sanitize quaternion - # this is needed, otherwise the quaternion might be non-normalized - prim_tf.Orthonormalize() - # extract position and orientation - positions[idx] = prim_tf.ExtractTranslation() - orientations[idx] = prim_tf.ExtractRotationQuat() - - # move to torch tensors - positions = torch.tensor(np.array(positions), dtype=torch.float32, device=self._device) - orientations = torch.tensor(np.array(orientations), dtype=torch.float32, device=self._device) - return positions, orientations # type: ignore - - def _get_local_poses_usd(self, indices: Sequence[int] | None = None) -> tuple[torch.Tensor, torch.Tensor]: - """Get local poses from USD.""" - # Resolve indices - if indices is None or indices == slice(None): - indices_list = self._ALL_INDICES - else: - indices_list = indices.tolist() if isinstance(indices, torch.Tensor) else list(indices) - - # Create buffers - translations = Vt.Vec3dArray(len(indices_list)) - orientations = Vt.QuatdArray(len(indices_list)) - - # Create a fresh XformCache to avoid stale cached values - xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) - - for idx, prim_idx in enumerate(indices_list): - prim = self._prims[prim_idx] - prim_tf = xform_cache.GetLocalTransformation(prim)[0] - prim_tf.Orthonormalize() - translations[idx] = prim_tf.ExtractTranslation() - orientations[idx] = prim_tf.ExtractRotationQuat() - - translations = torch.tensor(np.array(translations), dtype=torch.float32, device=self._device) - orientations = torch.tensor(np.array(orientations), dtype=torch.float32, device=self._device) - return translations, orientations # type: ignore - - def _get_scales_usd(self, indices: Sequence[int] | None = None) -> torch.Tensor: - """Get scales from USD.""" - # Resolve indices - if indices is None or indices == slice(None): - indices_list = self._ALL_INDICES - else: - indices_list = indices.tolist() if isinstance(indices, torch.Tensor) else list(indices) - - # Create buffers - scales = Vt.Vec3dArray(len(indices_list)) - - for idx, prim_idx in enumerate(indices_list): - prim = self._prims[prim_idx] - scales[idx] = prim.GetAttribute("xformOp:scale").Get() - - # Convert to tensor - return torch.tensor(np.array(scales), dtype=torch.float32, device=self._device) - - """ - Internal Functions - Fabric. - """ - - def _set_world_poses_fabric( - self, - positions: torch.Tensor | None = None, - orientations: torch.Tensor | None = None, - indices: Sequence[int] | None = None, - ): - """Set world poses using Fabric GPU batch operations. - - Writes directly to Fabric's ``omni:fabric:worldMatrix`` attribute using Warp kernels. - Changes are propagated through Fabric's hierarchy system but remain GPU-resident. - - For workflows mixing Fabric world pose writes with USD local pose queries, note - that local poses read from USD's xformOp:* attributes, which may not immediately - reflect Fabric changes. For best performance and consistency, use Fabric methods - exclusively (get_world_poses/set_world_poses with Fabric enabled). - """ - # Lazy initialization - if not self._fabric_initialized: - self._initialize_fabric() - - # Resolve indices (treat slice(None) as None for consistency with USD path) - indices_wp = self._resolve_indices_wp(indices) - - count = indices_wp.shape[0] - - # Convert torch to warp (if provided), use dummy arrays for None to avoid Warp kernel issues - if positions is not None: - positions_wp = wp.from_torch(positions) - else: - positions_wp = wp.zeros((0, 3), dtype=wp.float32).to(self._device) - - if orientations is not None: - orientations_wp = wp.from_torch(orientations) - else: - orientations_wp = wp.zeros((0, 4), dtype=wp.float32).to(self._device) - - # Dummy array for scales (not modifying) - scales_wp = wp.zeros((0, 3), dtype=wp.float32).to(self._device) - - # Use cached fabricarray for world matrices - world_matrices = self._fabric_world_matrices - - # Batch compose matrices with a single kernel launch - # Note: world_matrices is a fabricarray on fabric_device, so we must launch on fabric_device - wp.launch( - kernel=fabric_utils.compose_fabric_transformation_matrix_from_warp_arrays, - dim=count, - inputs=[ - world_matrices, - positions_wp, - orientations_wp, - scales_wp, # dummy array instead of None - False, # broadcast_positions - False, # broadcast_orientations - False, # broadcast_scales - indices_wp, - self._view_to_fabric, - ], - device=self._fabric_device, - ) - - # Synchronize to ensure kernel completes - wp.synchronize() - - # Update world transforms within Fabric hierarchy - self._fabric_hierarchy.update_world_xforms() - # Fabric now has authoritative data; skip future USD syncs - self._fabric_usd_sync_done = True - # Mirror to USD for renderer-facing prims when enabled. - if self._sync_usd_on_fabric_write: - self._set_world_poses_usd(positions, orientations, indices) - - # Fabric writes are GPU-resident; local pose operations still use USD. - - def _set_local_poses_fabric( - self, - translations: torch.Tensor | None = None, - orientations: torch.Tensor | None = None, - indices: Sequence[int] | None = None, - ): - """Set local poses using USD (matches Isaac Sim's design). - - Note: Even in Fabric mode, local pose operations use USD. - This is Isaac Sim's design: the ``usd=False`` parameter only affects world poses. - - Rationale: - - Local pose writes need correct parent-child hierarchy relationships - - USD maintains these relationships correctly and efficiently - - Fabric is optimized for world pose operations, not local hierarchies - """ - self._set_local_poses_usd(translations, orientations, indices) - - def _set_scales_fabric(self, scales: torch.Tensor, indices: Sequence[int] | None = None): - """Set scales using Fabric GPU batch operations.""" - # Lazy initialization - if not self._fabric_initialized: - self._initialize_fabric() - - # Resolve indices (treat slice(None) as None for consistency with USD path) - indices_wp = self._resolve_indices_wp(indices) - - count = indices_wp.shape[0] - - # Convert torch to warp - scales_wp = wp.from_torch(scales) - - # Dummy arrays for positions and orientations (not modifying) - positions_wp = wp.zeros((0, 3), dtype=wp.float32).to(self._device) - orientations_wp = wp.zeros((0, 4), dtype=wp.float32).to(self._device) - - # Use cached fabricarray for world matrices - world_matrices = self._fabric_world_matrices - - # Batch compose matrices on GPU with a single kernel launch - # Note: world_matrices is a fabricarray on fabric_device, so we must launch on fabric_device - wp.launch( - kernel=fabric_utils.compose_fabric_transformation_matrix_from_warp_arrays, - dim=count, - inputs=[ - world_matrices, - positions_wp, # dummy array instead of None - orientations_wp, # dummy array instead of None - scales_wp, - False, # broadcast_positions - False, # broadcast_orientations - False, # broadcast_scales - indices_wp, - self._view_to_fabric, - ], - device=self._fabric_device, - ) - - # Synchronize to ensure kernel completes before syncing - wp.synchronize() - - # Update world transforms to propagate changes - self._fabric_hierarchy.update_world_xforms() - # Fabric now has authoritative data; skip future USD syncs - self._fabric_usd_sync_done = True - # Mirror to USD for renderer-facing prims when enabled. - if self._sync_usd_on_fabric_write: - self._set_scales_usd(scales, indices) - - def _get_world_poses_fabric(self, indices: Sequence[int] | None = None) -> tuple[torch.Tensor, torch.Tensor]: - """Get world poses from Fabric using GPU batch operations.""" - # Lazy initialization of Fabric infrastructure - if not self._fabric_initialized: - self._initialize_fabric() - # Sync once from USD to ensure reads see the latest authored transforms - if not self._fabric_usd_sync_done: - self._sync_fabric_from_usd_once() - - # Resolve indices (treat slice(None) as None for consistency with USD path) - indices_wp = self._resolve_indices_wp(indices) - - count = indices_wp.shape[0] - - # Use pre-allocated buffers for full reads, allocate only for partial reads - use_cached_buffers = indices is None or indices == slice(None) - if use_cached_buffers: - # Full read: Use cached buffers (zero allocation overhead!) - positions_wp = self._fabric_positions_buffer - orientations_wp = self._fabric_orientations_buffer - scales_wp = self._fabric_dummy_buffer - else: - # Partial read: Need to allocate buffers of appropriate size - positions_wp = wp.zeros((count, 3), dtype=wp.float32).to(self._device) - orientations_wp = wp.zeros((count, 4), dtype=wp.float32).to(self._device) - scales_wp = self._fabric_dummy_buffer # Always use dummy for scales - - # Use cached fabricarray for world matrices - # This eliminates the 0.06-0.30ms variability from creating fabricarray each call - world_matrices = self._fabric_world_matrices - - # Launch GPU kernel to decompose matrices in parallel - # Note: world_matrices is a fabricarray on fabric_device, so we must launch on fabric_device - wp.launch( - kernel=fabric_utils.decompose_fabric_transformation_matrix_to_warp_arrays, - dim=count, - inputs=[ - world_matrices, - positions_wp, - orientations_wp, - scales_wp, # dummy array instead of None - indices_wp, - self._view_to_fabric, - ], - device=self._fabric_device, - ) - - # Return tensors: zero-copy for cached buffers, conversion for partial reads - if use_cached_buffers: - # Zero-copy! The Warp kernel wrote directly into the PyTorch tensors - # We just need to synchronize to ensure the kernel is done - wp.synchronize() - return self._fabric_positions_torch, self._fabric_orientations_torch - else: - # Partial read: Need to convert from Warp to torch - positions = wp.to_torch(positions_wp) - orientations = wp.to_torch(orientations_wp) - return positions, orientations - - def _get_local_poses_fabric(self, indices: Sequence[int] | None = None) -> tuple[torch.Tensor, torch.Tensor]: - """Get local poses using USD (matches Isaac Sim's design). - - Note: - Even in Fabric mode, local pose operations use USD's XformCache. - This is Isaac Sim's design: the ``usd=False`` parameter only affects world poses. - - Rationale: - - Local pose computation requires parent transforms which may not be in the view - - USD's XformCache provides efficient hierarchy-aware local transform queries - - Fabric is optimized for world pose operations, not local hierarchies - """ - return self._get_local_poses_usd(indices) - - def _get_scales_fabric(self, indices: Sequence[int] | None = None) -> torch.Tensor: - """Get scales from Fabric using GPU batch operations.""" - # Lazy initialization - if not self._fabric_initialized: - self._initialize_fabric() - # Sync once from USD to ensure reads see the latest authored transforms - if not self._fabric_usd_sync_done: - self._sync_fabric_from_usd_once() - - # Resolve indices (treat slice(None) as None for consistency with USD path) - indices_wp = self._resolve_indices_wp(indices) - - count = indices_wp.shape[0] - - # Use pre-allocated buffers for full reads, allocate only for partial reads - use_cached_buffers = indices is None or indices == slice(None) - if use_cached_buffers: - # Full read: Use cached buffers (zero allocation overhead!) - scales_wp = self._fabric_scales_buffer - else: - # Partial read: Need to allocate buffer of appropriate size - scales_wp = wp.zeros((count, 3), dtype=wp.float32).to(self._device) - - # Always use dummy buffers for positions and orientations (not needed for scales) - positions_wp = self._fabric_dummy_buffer - orientations_wp = self._fabric_dummy_buffer - - # Use cached fabricarray for world matrices - world_matrices = self._fabric_world_matrices - - # Launch GPU kernel to decompose matrices in parallel - # Note: world_matrices is a fabricarray on fabric_device, so we must launch on fabric_device - wp.launch( - kernel=fabric_utils.decompose_fabric_transformation_matrix_to_warp_arrays, - dim=count, - inputs=[ - world_matrices, - positions_wp, # dummy array instead of None - orientations_wp, # dummy array instead of None - scales_wp, - indices_wp, - self._view_to_fabric, - ], - device=self._fabric_device, - ) - - # Return tensor: zero-copy for cached buffers, conversion for partial reads - if use_cached_buffers: - # Zero-copy! The Warp kernel wrote directly into the PyTorch tensor - wp.synchronize() - return self._fabric_scales_torch - else: - # Partial read: Need to convert from Warp to torch - return wp.to_torch(scales_wp) - - """ - Internal Functions - Initialization. - """ - - def _initialize_fabric(self) -> None: - """Initialize Fabric batch infrastructure for GPU-accelerated pose queries. - - This method ensures all prims have the required Fabric hierarchy attributes - (``omni:fabric:localMatrix`` and ``omni:fabric:worldMatrix``) and creates the necessary - infrastructure for batch GPU operations using Warp. - - Based on the Fabric Hierarchy documentation, when Fabric Scene Delegate is enabled, - all boundable prims should have these attributes. This method ensures they exist - and are properly synchronized with USD. - """ - import usdrt - from usdrt import Rt - - # Get USDRT (Fabric) stage - stage_id = sim_utils.get_current_stage_id() - fabric_stage = usdrt.Usd.Stage.Attach(stage_id) - - # Step 1: Ensure all prims have Fabric hierarchy attributes - # According to the documentation, these attributes are created automatically - # when Fabric Scene Delegate is enabled, but we ensure they exist - for i in range(self.count): - rt_prim = fabric_stage.GetPrimAtPath(self.prim_paths[i]) - rt_xformable = Rt.Xformable(rt_prim) - - # Create Fabric hierarchy world matrix attribute if it doesn't exist - has_attr = ( - rt_xformable.HasFabricHierarchyWorldMatrixAttr() - if hasattr(rt_xformable, "HasFabricHierarchyWorldMatrixAttr") - else False - ) - if not has_attr: - rt_xformable.CreateFabricHierarchyWorldMatrixAttr() - - # Best-effort USD->Fabric sync; authoritative initialization happens on first read. - rt_xformable.SetWorldXformFromUsd() - - # Create view index attribute for batch operations - rt_prim.CreateAttribute(self._view_index_attr, usdrt.Sdf.ValueTypeNames.UInt, custom=True) - rt_prim.GetAttribute(self._view_index_attr).Set(i) - - # After syncing all prims, update the Fabric hierarchy to ensure world matrices are computed - self._fabric_hierarchy = usdrt.hierarchy.IFabricHierarchy().get_fabric_hierarchy( - fabric_stage.GetFabricId(), fabric_stage.GetStageIdAsStageId() - ) - self._fabric_hierarchy.update_world_xforms() - - # Step 2: Create index arrays for batch operations - self._default_view_indices = wp.zeros((self.count,), dtype=wp.uint32).to(self._device) - wp.launch( - kernel=fabric_utils.arange_k, - dim=self.count, - inputs=[self._default_view_indices], - device=self._device, - ) - wp.synchronize() # Ensure indices are ready - - # Step 3: Create Fabric selection with attribute filtering - # SelectPrims expects device format like "cuda:0" not "cuda" - # - # KNOWN ISSUE: SelectPrims may return prims in a different order than self._prims - # (which comes from USD's find_matching_prims). We create a bidirectional mapping - # (_view_to_fabric and _fabric_to_view) to handle this ordering difference. - # This works correctly for full-view operations but partial indexing still has issues. - # - # NOTE: SelectPrims only supports "cuda:0" regardless of which GPU the simulation - # is running on. In multi-GPU setups, we must use "cuda:0" for SelectPrims even if - # the simulation device is "cuda:1" or higher. - fabric_device = self._device - if self._device == "cuda": - logger.warning("Fabric device is not specified, defaulting to 'cuda:0'.") - fabric_device = "cuda:0" - elif self._device.startswith("cuda:"): - # SelectPrims only supports cuda:0, so we always use cuda:0 for SelectPrims - # even if the simulation is running on a different GPU - if self._device != "cuda:0": - logger.debug( - f"SelectPrims only supports cuda:0. Using cuda:0 for SelectPrims " - f"even though simulation device is {self._device}." - ) - fabric_device = "cuda:0" - - self._fabric_selection = fabric_stage.SelectPrims( - require_attrs=[ - (usdrt.Sdf.ValueTypeNames.UInt, self._view_index_attr, usdrt.Usd.Access.Read), - (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.ReadWrite), - ], - device=fabric_device, - ) - - # Step 4: Create bidirectional mapping between view and fabric indices - # Note: fabric_to_view is tied to fabric_device (cuda:0) because it's created from SelectPrims. - # view_to_fabric must also be on fabric_device since it's always used with fabricarrays in kernels. - self._view_to_fabric = wp.zeros((self.count,), dtype=wp.uint32).to(fabric_device) - self._fabric_to_view = wp.fabricarray(self._fabric_selection, self._view_index_attr) - - wp.launch( - kernel=fabric_utils.set_view_to_fabric_array, - dim=self._fabric_to_view.shape[0], - inputs=[self._fabric_to_view, self._view_to_fabric], - device=fabric_device, - ) - # Synchronize to ensure mapping is ready before any operations - wp.synchronize() - - # Pre-allocate reusable output buffers for read operations - self._fabric_positions_torch = torch.zeros((self.count, 3), dtype=torch.float32, device=self._device) - self._fabric_orientations_torch = torch.zeros((self.count, 4), dtype=torch.float32, device=self._device) - self._fabric_scales_torch = torch.zeros((self.count, 3), dtype=torch.float32, device=self._device) - - # Create Warp views of the PyTorch tensors - self._fabric_positions_buffer = wp.from_torch(self._fabric_positions_torch, dtype=wp.float32) - self._fabric_orientations_buffer = wp.from_torch(self._fabric_orientations_torch, dtype=wp.float32) - self._fabric_scales_buffer = wp.from_torch(self._fabric_scales_torch, dtype=wp.float32) - - # Dummy array for unused outputs (always empty) - self._fabric_dummy_buffer = wp.zeros((0, 3), dtype=wp.float32).to(self._device) - - # Cache fabricarray for world matrices to avoid recreation overhead - # Refs: https://docs.omniverse.nvidia.com/kit/docs/usdrt/latest/docs/usdrt_prim_selection.html - # https://docs.omniverse.nvidia.com/kit/docs/usdrt/latest/docs/scenegraph_use.html - self._fabric_world_matrices = wp.fabricarray(self._fabric_selection, "omni:fabric:worldMatrix") - - # Cache Fabric stage to avoid expensive get_current_stage() calls - self._fabric_stage = fabric_stage - - # Store fabric_device for use in kernel launches that involve fabricarrays - self._fabric_device = fabric_device - - self._fabric_initialized = True - # Force a one-time USD->Fabric sync on first read to pick up any USD edits - # made after the view was constructed. - self._fabric_usd_sync_done = False - - def _sync_fabric_from_usd_once(self) -> None: - """Sync Fabric world matrices from USD once, on the first read.""" - # Ensure Fabric is initialized - if not self._fabric_initialized: - self._initialize_fabric() - - # Read authoritative transforms from USD and write once into Fabric. - positions_usd, orientations_usd = self._get_world_poses_usd() - scales_usd = self._get_scales_usd() - - prev_sync = self._sync_usd_on_fabric_write - self._sync_usd_on_fabric_write = False - self._set_world_poses_fabric(positions_usd, orientations_usd) - self._set_scales_fabric(scales_usd) - self._sync_usd_on_fabric_write = prev_sync - - self._fabric_usd_sync_done = True - - def _resolve_indices_wp(self, indices: Sequence[int] | None) -> wp.array: - """Resolve view indices as a Warp array.""" - if indices is None or indices == slice(None): - if self._default_view_indices is None: - raise RuntimeError("Fabric indices are not initialized.") - return self._default_view_indices - indices_list = indices.tolist() if isinstance(indices, torch.Tensor) else list(indices) - return wp.array(indices_list, dtype=wp.uint32).to(self._device) +XformPrimView = FrameView diff --git a/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py b/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py index 03221e1ce366..4824d968284f 100644 --- a/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py +++ b/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py @@ -142,7 +142,7 @@ def main(): prim_path="/World/envs/env_.*/ball", mesh_prim_paths=mesh_targets, pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=(1.6, 1.0)), - attach_yaw_only=True, + ray_alignment="yaw", debug_vis=not args_cli.headless, ) ray_caster = MultiMeshRayCaster(cfg=ray_caster_cfg) diff --git a/source/isaaclab/test/sensors/test_ray_caster.py b/source/isaaclab/test/sensors/test_ray_caster.py index 5dffb5ccd015..4e29b25ce351 100644 --- a/source/isaaclab/test/sensors/test_ray_caster.py +++ b/source/isaaclab/test/sensors/test_ray_caster.py @@ -243,6 +243,84 @@ def test_raycast_random_cube(raycast_setup): torch.testing.assert_close(ray_face_id, ray_face_id_m) +## +# RayCaster sensor-level tests +## + + +def test_raycaster_offset_does_not_affect_pos_w(): + """Verify that cfg.offset.pos shifts ray starts but NOT data.pos_w. + + data.pos_w must reflect the parent body position so that downstream + observations like height_scan (pos_w_z - hit_z - 0.5) produce values + relative to the body, not relative to the offset sensor frame. + + Regression test: previously the offset was baked into the FrameView's + Xform local transform, causing data.pos_w to include the 20m offset + and breaking height-scan observations during training. + """ + import isaaclab.sim as sim_utils + from isaaclab.sensors.ray_caster import RayCaster, RayCasterCfg, patterns + from isaaclab.terrains.trimesh.utils import make_plane + from isaaclab.terrains.utils import create_prim_from_mesh + + sim_utils.create_new_stage() + + # ground plane at z=0 + mesh = make_plane(size=(100, 100), height=0.0, center_zero=True) + create_prim_from_mesh("/World/ground", mesh) + + # parent body at known position + body_pos = (0.0, 0.0, 0.6) + sim_utils.create_prim("/World/Robot", "Xform", translation=body_pos) + + # large z-offset to make the regression obvious + offset_z = 20.0 + cfg = RayCasterCfg( + prim_path="/World/Robot", + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, offset_z)), + mesh_prim_paths=["/World/ground"], + pattern_cfg=patterns.GridPatternCfg(resolution=0.5, size=[1.0, 1.0]), + ray_alignment="yaw", + ) + + dt = 0.01 + sim = sim_utils.SimulationContext(sim_utils.SimulationCfg(dt=dt)) + + sensor = RayCaster(cfg) + sim.reset() + sensor.update(dt) + + # data.pos_w / data.ray_hits_w are wp.array after the ray caster warp-backend + # migration (PR #4967); convert to torch views for indexing. + pos_w = wp.to_torch(sensor.data.pos_w)[0].cpu() + + # pos_w.z should be near the body height, NOT body_height + offset + assert abs(pos_w[2].item() - body_pos[2]) < 1.0, ( + f"data.pos_w.z = {pos_w[2].item():.2f}, expected near body height {body_pos[2]}." + f" If pos_w.z ≈ {body_pos[2] + offset_z}, the offset was incorrectly baked into the FrameView." + ) + + # ray_hits should be near z=0 (ground plane) + hits_z = wp.to_torch(sensor.data.ray_hits_w)[0, :, 2].cpu() + valid = hits_z[~torch.isinf(hits_z)] + if len(valid) > 0: + assert valid.abs().max().item() < 2.0, ( + f"Ray hits z range [{valid.min().item():.2f}, {valid.max().item():.2f}] — expected near ground (z≈0)." + ) + + # height_scan observation: pos_w_z - hit_z - 0.5 should be small, not ~20 + if len(valid) > 0: + height_obs = pos_w[2].item() - valid.mean().item() - 0.5 + assert abs(height_obs) < 5.0, ( + f"height_scan observation = {height_obs:.2f}, expected near 0." + f" If ≈{offset_z}, the offset leaked into data.pos_w." + ) + + sim.stop() + sim.clear_instance() + + # --------------------------------------------------------------------------- # Tests for raycast_mesh_masked_kernel (new kernel in utils/warp/kernels.py) # --------------------------------------------------------------------------- diff --git a/source/isaaclab/test/sim/frame_view_contract_utils.py b/source/isaaclab/test/sim/frame_view_contract_utils.py new file mode 100644 index 000000000000..37734fee4322 --- /dev/null +++ b/source/isaaclab/test/sim/frame_view_contract_utils.py @@ -0,0 +1,359 @@ +# 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 + +"""Shared FrameView contract tests. + +This module defines the invariants that **every** FrameView backend +(USD, Fabric, Newton) must satisfy. Backend test files import these tests +via ``from frame_view_contract_utils import *`` and provide a +``view_factory`` pytest fixture that builds the backend-specific scene. + +The factory signature is:: + + def view_factory() -> Callable[[int, str], ViewBundle]: ... + +Where ``ViewBundle`` is a :class:`NamedTuple`:: + + class ViewBundle(NamedTuple): + view: BaseFrameView + get_parent_pos: Callable[[int, str], torch.Tensor] + set_parent_pos: Callable[[torch.Tensor, int], None] + teardown: Callable[[], None] + +- ``view``: The FrameView under test. Must track child prims at + :data:`CHILD_OFFSET` under parent prims/bodies. +- ``get_parent_pos(n, device)``: Read the parent prim/body positions. +- ``set_parent_pos(positions, n)``: Write the parent prim/body positions. +- ``teardown()``: Cleanup (close context, clear stage, etc.). + +Tolerance policy: + - Indexed reads (exact copy): ``atol=0`` + - Composition / decomposition through float32 transforms: ``atol=ATOL`` + - Parent position identity checks (should be untouched): ``atol=0`` +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import NamedTuple + +import pytest +import torch +import warp as wp + +CHILD_OFFSET = (0.1, 0.0, 0.05) +"""Local offset of the child prim from its parent, shared by all backend fixtures.""" + +ATOL = 1e-5 +"""Default absolute tolerance for float32 transform composition.""" + + +class ViewBundle(NamedTuple): + """Return type of the ``view_factory`` fixture.""" + + view: object + get_parent_pos: Callable + set_parent_pos: Callable + teardown: Callable + + +def _t(a): + """Convert wp.array to torch.Tensor (pass-through for Tensor).""" + return wp.to_torch(a) if isinstance(a, wp.array) else a + + +def _wp_vec3f(data, device="cpu"): + return wp.array([wp.vec3f(*row) for row in data], dtype=wp.vec3f, device=device) + + +def _wp_vec4f(data, device="cpu"): + return wp.array([wp.vec4f(*row) for row in data], dtype=wp.vec4f, device=device) + + +# ================================================================== +# Contract: Getters +# ================================================================== + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_world_pose_equals_parent_plus_offset(device, view_factory): + """world_pose == parent_pos + local offset (identity parent orientation).""" + bundle = view_factory(num_envs=4, device=device) + try: + child_pos = _t(bundle.view.get_world_poses()[0]) + parent_pos = bundle.get_parent_pos(4, device) + offset = torch.tensor(CHILD_OFFSET, device=device) + + torch.testing.assert_close(child_pos, parent_pos + offset.unsqueeze(0), atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_local_pose_equals_structural_offset(device, view_factory): + """local_pose == the authored offset (0.1, 0, 0.05) for every prim.""" + bundle = view_factory(num_envs=4, device=device) + try: + local_pos, local_quat = bundle.view.get_local_poses() + expected_pos = torch.tensor(CHILD_OFFSET, device=device).expand(4, -1) + expected_quat = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device).expand(4, -1) + + torch.testing.assert_close(_t(local_pos), expected_pos, atol=ATOL, rtol=0) + torch.testing.assert_close(_t(local_quat), expected_quat, atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_local_differs_from_world(device, view_factory): + """local != world when parent is not at the origin. + + Asserts |world - local| > 0.5 to catch any implementation that returns + world as local. The parent is offset from the origin so the z-component + alone provides > 0.5 difference. + """ + bundle = view_factory(num_envs=2, device=device) + try: + world_pos = _t(bundle.view.get_world_poses()[0]) + local_pos = _t(bundle.view.get_local_poses()[0]) + + diff = (world_pos - local_pos).abs().max().item() + assert diff > 0.5, ( + f"Expected |world - local| > 0.5, got {diff:.4f}. world={world_pos.tolist()}, local={local_pos.tolist()}" + ) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_local_stable_after_parent_move(device, view_factory): + """Moving the parent changes world but NOT local.""" + bundle = view_factory(num_envs=2, device=device) + try: + local_before = _t(bundle.view.get_local_poses()[0]).clone() + bundle.set_parent_pos(torch.tensor([[99.0, 0.0, 0.0], [0.0, 99.0, 0.0]], device=device), 2) + local_after = _t(bundle.view.get_local_poses()[0]) + + torch.testing.assert_close(local_after, local_before, atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_world_tracks_parent_move(device, view_factory): + """Moving the parent shifts world poses by the same amount.""" + bundle = view_factory(num_envs=2, device=device) + try: + new_parent_pos = torch.tensor([[5.0, 0.0, 0.0], [0.0, 5.0, 0.0]], device=device) + bundle.set_parent_pos(new_parent_pos, 2) + + child_pos = _t(bundle.view.get_world_poses()[0]) + offset = torch.tensor(CHILD_OFFSET, device=device) + + torch.testing.assert_close(child_pos, new_parent_pos + offset.unsqueeze(0), atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_indexed_get_returns_correct_subset(device, view_factory): + """Indexed get (out-of-order) returns exact copies for both world and local.""" + bundle = view_factory(num_envs=5, device=device) + try: + all_world = _t(bundle.view.get_world_poses()[0]) + all_local = _t(bundle.view.get_local_poses()[0]) + + indices_list = [4, 1, 3] + indices = wp.array(indices_list, dtype=wp.int32, device=device) + sub_world = _t(bundle.view.get_world_poses(indices)[0]) + sub_local = _t(bundle.view.get_local_poses(indices)[0]) + + for out_i, view_i in enumerate(indices_list): + torch.testing.assert_close(sub_world[out_i], all_world[view_i], atol=0, rtol=0) + torch.testing.assert_close(sub_local[out_i], all_local[view_i], atol=0, rtol=0) + finally: + bundle.teardown() + + +# ================================================================== +# Contract: Setters +# ================================================================== + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_world_roundtrip(device, view_factory): + """set_world_poses -> get_world_poses returns the same values.""" + bundle = view_factory(num_envs=2, device=device) + try: + new_pos = _wp_vec3f([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], device=device) + new_quat = _wp_vec4f([[0.0, 0.0, 0.7071068, 0.7071068], [0.0, 0.0, 0.0, 1.0]], device=device) + bundle.view.set_world_poses(new_pos, new_quat) + + ret_pos, ret_quat = bundle.view.get_world_poses() + torch.testing.assert_close(_t(ret_pos), _t(new_pos), atol=ATOL, rtol=0) + torch.testing.assert_close(_t(ret_quat), _t(new_quat), atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_local_roundtrip(device, view_factory): + """set_local_poses -> get_local_poses returns the same values.""" + bundle = view_factory(num_envs=2, device=device) + try: + new_pos = _wp_vec3f([[0.5, 0.3, 0.1], [0.2, 0.7, 0.4]], device=device) + new_quat = _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device) + bundle.view.set_local_poses(new_pos, new_quat) + + ret_pos, ret_quat = bundle.view.get_local_poses() + torch.testing.assert_close(_t(ret_pos), _t(new_pos), atol=ATOL, rtol=0) + torch.testing.assert_close(_t(ret_quat), _t(new_quat), atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_world_does_not_move_parent(device, view_factory): + """set_world_poses must not modify the parent prim/body position.""" + bundle = view_factory(num_envs=2, device=device) + try: + parent_before = bundle.get_parent_pos(2, device).clone() + bundle.view.set_world_poses( + _wp_vec3f([[99.0, 99.0, 99.0], [88.0, 88.0, 88.0]], device=device), + _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), + ) + parent_after = bundle.get_parent_pos(2, device) + + torch.testing.assert_close(parent_after, parent_before, atol=0, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_local_does_not_move_parent(device, view_factory): + """set_local_poses must not modify the parent prim/body position.""" + bundle = view_factory(num_envs=2, device=device) + try: + parent_before = bundle.get_parent_pos(2, device).clone() + bundle.view.set_local_poses( + _wp_vec3f([[0.5, 0.5, 0.5], [1.0, 1.0, 1.0]], device=device), + _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), + ) + parent_after = bundle.get_parent_pos(2, device) + + torch.testing.assert_close(parent_after, parent_before, atol=0, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_world_updates_local(device, view_factory): + """After set_world_poses, get_local_poses reflects the new offset. + + Uses non-axis-aligned offsets to catch coordinate swap bugs. + """ + bundle = view_factory(num_envs=2, device=device) + try: + parent_pos = bundle.get_parent_pos(2, device) + desired_offset = torch.tensor([[0.3, 0.7, 0.2], [0.8, 0.1, 0.6]], device=device) + new_world = parent_pos + desired_offset + + bundle.view.set_world_poses( + _wp_vec3f(new_world.tolist(), device=device), + _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), + ) + + local_pos = _t(bundle.view.get_local_poses()[0]) + torch.testing.assert_close(local_pos, desired_offset, atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_local_updates_world(device, view_factory): + """After set_local_poses, get_world_poses == parent + new_local. + + Uses non-axis-aligned offsets to catch coordinate swap bugs. + """ + bundle = view_factory(num_envs=2, device=device) + try: + parent_pos = bundle.get_parent_pos(2, device) + new_offset = torch.tensor([[0.4, 0.9, 0.15], [0.6, 0.2, 0.85]], device=device) + bundle.view.set_local_poses( + _wp_vec3f(new_offset.tolist(), device=device), + _wp_vec4f([[0.0, 0.0, 0.0, 1.0]] * 2, device=device), + ) + + world_pos = _t(bundle.view.get_world_poses()[0]) + torch.testing.assert_close(world_pos, parent_pos + new_offset, atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_world_partial_position_only(device, view_factory): + """Setting only positions: new positions written, orientations preserved.""" + bundle = view_factory(num_envs=2, device=device) + try: + _, orig_quat = bundle.view.get_world_poses() + new_pos = _wp_vec3f([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], device=device) + bundle.view.set_world_poses(positions=new_pos) + + ret_pos, ret_quat = bundle.view.get_world_poses() + torch.testing.assert_close(_t(ret_pos), _t(new_pos), atol=ATOL, rtol=0) + torch.testing.assert_close(_t(ret_quat), _t(orig_quat), atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_world_partial_orientation_only(device, view_factory): + """Setting only orientations: new orientations written, positions preserved.""" + bundle = view_factory(num_envs=2, device=device) + try: + orig_pos, _ = bundle.view.get_world_poses() + new_quat = _wp_vec4f([[0.0, 0.0, 0.7071068, 0.7071068], [0.7071068, 0.0, 0.0, 0.7071068]], device=device) + bundle.view.set_world_poses(orientations=new_quat) + + ret_pos, ret_quat = bundle.view.get_world_poses() + torch.testing.assert_close(_t(ret_pos), _t(orig_pos), atol=ATOL, rtol=0) + torch.testing.assert_close(_t(ret_quat), _t(new_quat), atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_local_partial_position_only(device, view_factory): + """Setting only local translations: new translations written, orientations preserved.""" + bundle = view_factory(num_envs=2, device=device) + try: + _, orig_quat = bundle.view.get_local_poses() + new_pos = _wp_vec3f([[0.2, 0.3, 0.4], [0.5, 0.6, 0.7]], device=device) + bundle.view.set_local_poses(translations=new_pos) + + ret_pos, ret_quat = bundle.view.get_local_poses() + torch.testing.assert_close(_t(ret_pos), _t(new_pos), atol=ATOL, rtol=0) + torch.testing.assert_close(_t(ret_quat), _t(orig_quat), atol=ATOL, rtol=0) + finally: + bundle.teardown() + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_set_world_indexed_only_affects_subset(device, view_factory): + """Indexed set_world_poses writes requested indices, leaves others untouched.""" + bundle = view_factory(num_envs=4, device=device) + try: + orig_pos = _t(bundle.view.get_world_poses()[0]).clone() + indices = wp.array([1, 3], dtype=wp.int32, device=device) + new_pos = _wp_vec3f([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], device=device) + bundle.view.set_world_poses(positions=new_pos, indices=indices) + + updated = _t(bundle.view.get_world_poses()[0]) + torch.testing.assert_close(updated[0], orig_pos[0], atol=0, rtol=0) + torch.testing.assert_close(updated[2], orig_pos[2], atol=0, rtol=0) + torch.testing.assert_close(updated[1], _t(new_pos)[0], atol=ATOL, rtol=0) + torch.testing.assert_close(updated[3], _t(new_pos)[1], atol=ATOL, rtol=0) + finally: + bundle.teardown() diff --git a/source/isaaclab/test/sim/test_views_xform_prim.py b/source/isaaclab/test/sim/test_views_xform_prim.py index 3de7a0b357a2..2b40705f732a 100644 --- a/source/isaaclab/test/sim/test_views_xform_prim.py +++ b/source/isaaclab/test/sim/test_views_xform_prim.py @@ -3,1513 +3,294 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Launch Isaac Sim Simulator first.""" +"""USD backend tests for FrameView. + +Imports the shared contract tests and provides the USD-specific +``view_factory`` fixture. Also includes USD-only tests for visibility, +prim ordering, xformOp standardization, and Isaac Sim comparison. +""" from isaaclab.app import AppLauncher -# launch omniverse app simulation_app = AppLauncher(headless=True).app -"""Rest everything follows.""" - import pytest # noqa: E402 import torch # noqa: E402 +import warp as wp # noqa: E402 + +from pxr import Gf, UsdGeom # noqa: E402 try: from isaacsim.core.prims import XFormPrim as _IsaacSimXformPrimView except (ModuleNotFoundError, ImportError): _IsaacSimXformPrimView = None +from frame_view_contract_utils import * # noqa: F401, F403, E402 +from frame_view_contract_utils import CHILD_OFFSET, ViewBundle # noqa: E402 import isaaclab.sim as sim_utils # noqa: E402 -from isaaclab.sim.views import XformPrimView as XformPrimView # noqa: E402 +from isaaclab.sim.views import UsdFrameView as FrameView # noqa: E402 from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR # noqa: E402 +PARENT_POS = (0.0, 0.0, 1.0) + @pytest.fixture(autouse=True) def test_setup_teardown(): - """Create a blank new stage for each test.""" - # Setup: Create a new stage sim_utils.create_new_stage() sim_utils.update_stage() - - # Yield for the test yield - - # Teardown: Clear stage after each test sim_utils.clear_stage() sim_utils.SimulationContext.clear_instance() -""" -Helper functions. -""" - - -def _prepare_indices(index_type, target_indices, num_prims, device): - """Helper function to prepare indices based on type.""" - if index_type == "list": - return target_indices, target_indices - elif index_type == "torch_tensor": - return torch.tensor(target_indices, dtype=torch.int64, device=device), target_indices - elif index_type == "slice_none": - return slice(None), list(range(num_prims)) - else: - raise ValueError(f"Unknown index type: {index_type}") - - -def _skip_if_backend_unavailable(backend: str, device: str): - """Skip tests when the requested backend is unavailable.""" - if device.startswith("cuda") and not torch.cuda.is_available(): - pytest.skip("CUDA not available") - if backend == "fabric" and device == "cpu": - pytest.skip("Warp fabricarray operations on CPU have known issues") - - -def _prim_type_for_backend(backend: str) -> str: - """Return a prim type that is compatible with the backend.""" - return "Camera" if backend == "fabric" else "Xform" - - -def _create_view(pattern: str, device: str, backend: str) -> XformPrimView: - """Create an XformPrimView for the requested backend.""" - if backend == "fabric": - sim_utils.SimulationContext(sim_utils.SimulationCfg(dt=0.01, device=device, use_fabric=True)) - return XformPrimView(pattern, device=device) - - -""" -Tests - Initialization. -""" - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_xform_prim_view_initialization_single_prim(device): - """Test XformPrimView initialization with a single prim.""" - # check if CUDA is available - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - # Create a single xform prim - stage = sim_utils.get_current_stage() - sim_utils.create_prim("/World/Object", "Xform", translation=(1.0, 2.0, 3.0), stage=stage) - - # Create view - view = XformPrimView("/World/Object", device=device) - - # Verify properties - assert view.count == 1 - assert view.prim_paths == ["/World/Object"] - assert view.device == device - assert len(view.prims) == 1 - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_xform_prim_view_initialization_multiple_prims(device): - """Test XformPrimView initialization with multiple prims using pattern matching.""" - # check if CUDA is available - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - # Create multiple prims - num_prims = 10 - stage = sim_utils.get_current_stage() - for i in range(num_prims): - sim_utils.create_prim(f"/World/Env_{i}/Object", "Xform", translation=(i * 2.0, 0.0, 1.0), stage=stage) - - # Create view with pattern - view = XformPrimView("/World/Env_.*/Object", device=device) - - # Verify properties - assert view.count == num_prims - assert view.device == device - assert len(view.prims) == num_prims - assert view.prim_paths == [f"/World/Env_{i}/Object" for i in range(num_prims)] - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_xform_prim_view_initialization_multiple_prims_order(device): - """Test XformPrimView initialization with multiple prims using pattern matching with multiple objects per prim. - - This test validates that XformPrimView respects USD stage traversal order, which is based on - creation order (depth-first search), NOT alphabetical/lexical sorting. This is an important - edge case that ensures deterministic prim ordering that matches USD's internal representation. - - The test creates prims in a deliberately non-alphabetical order (1, 0, A, a, 2) and verifies - that they are retrieved in creation order, not sorted order (0, 1, 2, A, a). - """ - # check if CUDA is available - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - # Create multiple prims - num_prims = 10 - stage = sim_utils.get_current_stage() - - # NOTE: Prims are created in a specific order to test that XformPrimView respects - # USD stage traversal order (DFS based on creation order), NOT alphabetical/lexical order. - # This is an important edge case: children under the same parent are returned in the - # order they were created, not sorted by name. - - # First batch: Create Object_1, Object_0, Object_A for each environment - # (intentionally non-alphabetical: 1, 0, A instead of 0, 1, A) - for i in range(num_prims): - sim_utils.create_prim(f"/World/Env_{i}/Object_1", "Xform", translation=(i * 2.0, -2.0, 1.0), stage=stage) - sim_utils.create_prim(f"/World/Env_{i}/Object_0", "Xform", translation=(i * 2.0, 2.0, 1.0), stage=stage) - sim_utils.create_prim(f"/World/Env_{i}/Object_A", "Xform", translation=(i * 2.0, 0.0, -1.0), stage=stage) - - # Second batch: Create Object_a, Object_2 for each environment - # (created after the first batch to verify traversal is depth-first per environment) - for i in range(num_prims): - sim_utils.create_prim(f"/World/Env_{i}/Object_a", "Xform", translation=(i * 2.0, 2.0, -1.0), stage=stage) - sim_utils.create_prim(f"/World/Env_{i}/Object_2", "Xform", translation=(i * 2.0, 2.0, 1.0), stage=stage) - - # Create view with pattern - view = XformPrimView("/World/Env_.*/Object_.*", device=device) - - # Expected ordering: DFS traversal by environment, with children in creation order - # For each Env_i, we expect: Object_1, Object_0, Object_A, Object_a, Object_2 - # (matches creation order, NOT alphabetical: would be 0, 1, 2, A, a if sorted) - expected_prim_paths_ordering = [] - for i in range(num_prims): - expected_prim_paths_ordering.append(f"/World/Env_{i}/Object_1") - expected_prim_paths_ordering.append(f"/World/Env_{i}/Object_0") - expected_prim_paths_ordering.append(f"/World/Env_{i}/Object_A") - expected_prim_paths_ordering.append(f"/World/Env_{i}/Object_a") - expected_prim_paths_ordering.append(f"/World/Env_{i}/Object_2") - - # Verify properties - assert view.count == num_prims * 5 - assert view.device == device - assert len(view.prims) == num_prims * 5 - assert view.prim_paths == expected_prim_paths_ordering - - # Additional validation: Verify ordering is NOT alphabetical - # If it were alphabetical, Object_0 would come before Object_1 - alphabetical_order = [] - for i in range(num_prims): - alphabetical_order.append(f"/World/Env_{i}/Object_0") - alphabetical_order.append(f"/World/Env_{i}/Object_1") - alphabetical_order.append(f"/World/Env_{i}/Object_2") - alphabetical_order.append(f"/World/Env_{i}/Object_A") - alphabetical_order.append(f"/World/Env_{i}/Object_a") - - assert view.prim_paths != alphabetical_order, ( - "Prim paths should follow creation order, not alphabetical order. " - "This test validates that USD stage traversal respects creation order." - ) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_xform_prim_view_standardizes_transform_op(device): - """Test that XformPrimView standardizes a prim with xformOp:transform to translate/orient/scale.""" - from pxr import Gf, UsdGeom - - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - expected_pos = (3.0, -1.0, 0.5) - matrix = Gf.Matrix4d(1.0) - matrix.SetTranslateOnly(Gf.Vec3d(*expected_pos)) - - stage = sim_utils.get_current_stage() - prim = stage.DefinePrim("/World/TransformPrim", "Xform") - UsdGeom.Xformable(prim).AddTransformOp().Set(matrix) - - view = XformPrimView("/World/TransformPrim", device=device) - - assert view.count == 1 - assert sim_utils.validate_standard_xform_ops(view.prims[0]) - - xformable = UsdGeom.Xformable(view.prims[0]) - ordered_ops = xformable.GetOrderedXformOps() - op_names = [op.GetOpName() for op in ordered_ops] - assert op_names == ["xformOp:translate", "xformOp:orient", "xformOp:scale"] - - assert ordered_ops[0].Get() == Gf.Vec3d(*expected_pos) - assert ordered_ops[1].Get() == Gf.Quatd(1.0, 0.0, 0.0, 0.0) - assert ordered_ops[2].Get() == Gf.Vec3d(1.0, 1.0, 1.0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_xform_prim_view_initialization_empty_pattern(device): - """Test XformPrimView initialization with pattern that matches no prims.""" - # check if CUDA is available - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - sim_utils.create_new_stage() - - # Create view with pattern that matches nothing - view = XformPrimView("/World/NonExistent_.*", device=device) - - # Should have zero count - assert view.count == 0 - assert len(view.prims) == 0 - - -""" -Tests - Getters. -""" - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_get_world_poses(device, backend): - """Test getting world poses from XformPrimView.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create prims with known world poses - expected_positions = [(1.0, 2.0, 3.0), (4.0, 5.0, 6.0), (7.0, 8.0, 9.0)] - expected_orientations = [(0.0, 0.0, 0.0, 1.0), (0.0, 0.0, 0.7071068, 0.7071068), (0.7071068, 0.0, 0.0, 0.7071068)] - - for i, (pos, quat) in enumerate(zip(expected_positions, expected_orientations)): - sim_utils.create_prim(f"/World/Object_{i}", prim_type, translation=pos, orientation=quat, stage=stage) - - # Create view - view = _create_view("/World/Object_.*", device=device, backend=backend) - - # Convert expected values to tensors - expected_positions_tensor = torch.tensor(expected_positions, dtype=torch.float32, device=device) - expected_orientations_tensor = torch.tensor(expected_orientations, dtype=torch.float32, device=device) - - # Get world poses - positions, orientations = view.get_world_poses() - - # Verify shapes - assert positions.shape == (3, 3) - assert orientations.shape == (3, 4) - - # Verify positions - torch.testing.assert_close(positions, expected_positions_tensor, atol=1e-5, rtol=0) - - # Verify orientations (allow for quaternion sign ambiguity) - try: - torch.testing.assert_close(orientations, expected_orientations_tensor, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(orientations, -expected_orientations_tensor, atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_get_local_poses(device, backend): - """Test getting local poses from XformPrimView.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create parent and child prims - sim_utils.create_prim("/World/Parent", "Xform", translation=(10.0, 0.0, 0.0), stage=stage) - - # Children with different local poses - expected_local_positions = [(1.0, 0.0, 0.0), (0.0, 2.0, 0.0), (0.0, 0.0, 3.0)] - expected_local_orientations = [ - (0.0, 0.0, 0.0, 1.0), - (0.0, 0.0, 0.7071068, 0.7071068), - (0.7071068, 0.0, 0.0, 0.7071068), - ] - - for i, (pos, quat) in enumerate(zip(expected_local_positions, expected_local_orientations)): - sim_utils.create_prim(f"/World/Parent/Child_{i}", prim_type, translation=pos, orientation=quat, stage=stage) - - # Create view - view = _create_view("/World/Parent/Child_.*", device=device, backend=backend) - - # Get local poses - translations, orientations = view.get_local_poses() - - # Verify shapes - assert translations.shape == (3, 3) - assert orientations.shape == (3, 4) - - # Convert expected values to tensors - expected_translations_tensor = torch.tensor(expected_local_positions, dtype=torch.float32, device=device) - expected_orientations_tensor = torch.tensor(expected_local_orientations, dtype=torch.float32, device=device) - - # Verify translations - torch.testing.assert_close(translations, expected_translations_tensor, atol=1e-5, rtol=0) - - # Verify orientations (allow for quaternion sign ambiguity) - try: - torch.testing.assert_close(orientations, expected_orientations_tensor, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(orientations, -expected_orientations_tensor, atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_get_scales(device, backend): - """Test getting scales from XformPrimView.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create prims with different scales - expected_scales = [(1.0, 1.0, 1.0), (2.0, 2.0, 2.0), (1.0, 2.0, 3.0)] - - for i, scale in enumerate(expected_scales): - sim_utils.create_prim(f"/World/Object_{i}", prim_type, scale=scale, stage=stage) - - # Create view - view = _create_view("/World/Object_.*", device=device, backend=backend) - - expected_scales_tensor = torch.tensor(expected_scales, dtype=torch.float32, device=device) - - # Get scales - scales = view.get_scales() - - # Verify shape and values - assert scales.shape == (3, 3) - torch.testing.assert_close(scales, expected_scales_tensor, atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_get_visibility(device): - """Test getting visibility when all prims are visible.""" - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - stage = sim_utils.get_current_stage() - - # Create prims (default is visible) - num_prims = 5 - for i in range(num_prims): - sim_utils.create_prim(f"/World/Object_{i}", "Xform", translation=(float(i), 0.0, 0.0), stage=stage) - - # Create view - view = XformPrimView("/World/Object_.*", device=device) - - # Get visibility - visibility = view.get_visibility() - - # Verify shape and values - assert visibility.shape == (num_prims,) - assert visibility.dtype == torch.bool - assert torch.all(visibility), "All prims should be visible by default" - - -""" -Tests - Setters. -""" - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_set_world_poses(device, backend): - """Test setting world poses in XformPrimView.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create prims - num_prims = 5 - for i in range(num_prims): - sim_utils.create_prim(f"/World/Object_{i}", prim_type, translation=(0.0, 0.0, 0.0), stage=stage) - - # Create view - view = _create_view("/World/Object_.*", device=device, backend=backend) - - # Set new world poses - new_positions = torch.tensor( - [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]], device=device - ) - new_orientations = torch.tensor( - [ - [0.0, 0.0, 0.0, 1.0], - [0.0, 0.0, 0.7071068, 0.7071068], - [0.7071068, 0.0, 0.0, 0.7071068], - [0.3826834, 0.0, 0.0, 0.9238795], - [0.0, 0.7071068, 0.0, 0.7071068], - ], - device=device, - ) - - view.set_world_poses(new_positions, new_orientations) - - # Get the poses back - retrieved_positions, retrieved_orientations = view.get_world_poses() - - # Verify they match - torch.testing.assert_close(retrieved_positions, new_positions, atol=1e-5, rtol=0) - # Check quaternions (allow sign flip) - try: - torch.testing.assert_close(retrieved_orientations, new_orientations, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(retrieved_orientations, -new_orientations, atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_set_world_poses_only_positions(device, backend): - """Test setting only positions, leaving orientations unchanged.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create prims with specific orientations - initial_quat = (0.0, 0.0, 0.7071068, 0.7071068) # 90 deg around Z - for i in range(3): - sim_utils.create_prim( - f"/World/Object_{i}", prim_type, translation=(0.0, 0.0, 0.0), orientation=initial_quat, stage=stage - ) - - # Create view - view = _create_view("/World/Object_.*", device=device, backend=backend) - - # Get initial orientations - _, initial_orientations = view.get_world_poses() - - # Set only positions - new_positions = torch.tensor([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]], device=device) - view.set_world_poses(positions=new_positions, orientations=None) - - # Get poses back - retrieved_positions, retrieved_orientations = view.get_world_poses() - - # Positions should be updated - torch.testing.assert_close(retrieved_positions, new_positions, atol=1e-5, rtol=0) - - # Orientations should be unchanged - try: - torch.testing.assert_close(retrieved_orientations, initial_orientations, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(retrieved_orientations, -initial_orientations, atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_set_world_poses_only_orientations(device, backend): - """Test setting only orientations, leaving positions unchanged.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create prims with specific positions - for i in range(3): - sim_utils.create_prim(f"/World/Object_{i}", prim_type, translation=(float(i), 0.0, 0.0), stage=stage) - - # Create view - view = _create_view("/World/Object_.*", device=device, backend=backend) - - # Get initial positions - initial_positions, _ = view.get_world_poses() - - # Set only orientations - new_orientations = torch.tensor( - [[0.0, 0.0, 0.7071068, 0.7071068], [0.7071068, 0.0, 0.0, 0.7071068], [0.3826834, 0.0, 0.0, 0.9238795]], - device=device, - ) - view.set_world_poses(positions=None, orientations=new_orientations) - - # Get poses back - retrieved_positions, retrieved_orientations = view.get_world_poses() - - # Positions should be unchanged - torch.testing.assert_close(retrieved_positions, initial_positions, atol=1e-5, rtol=0) - - # Orientations should be updated - try: - torch.testing.assert_close(retrieved_orientations, new_orientations, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(retrieved_orientations, -new_orientations, atol=1e-5, rtol=0) - +# ------------------------------------------------------------------ +# Contract fixture +# ------------------------------------------------------------------ -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_set_world_poses_with_hierarchy(device, backend): - """Test setting world poses correctly handles parent transformations.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - child_prim_type = _prim_type_for_backend(backend) - - # Create parent prims - for i in range(3): - parent_pos = (i * 10.0, 0.0, 0.0) - parent_quat = (0.0, 0.0, 0.7071068, 0.7071068) # 90 deg around Z - sim_utils.create_prim( - f"/World/Parent_{i}", "Xform", translation=parent_pos, orientation=parent_quat, stage=stage - ) - # Create child prims - sim_utils.create_prim(f"/World/Parent_{i}/Child", child_prim_type, translation=(0.0, 0.0, 0.0), stage=stage) - - # Create view for children - view = _create_view("/World/Parent_.*/Child", device=device, backend=backend) - - # Set world poses for children - desired_world_positions = torch.tensor([[5.0, 5.0, 0.0], [15.0, 5.0, 0.0], [25.0, 5.0, 0.0]], device=device) - desired_world_orientations = torch.tensor( - [[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0]], device=device - ) - - view.set_world_poses(desired_world_positions, desired_world_orientations) - - # Get world poses back - retrieved_positions, retrieved_orientations = view.get_world_poses() - - # Should match desired world poses - torch.testing.assert_close(retrieved_positions, desired_world_positions, atol=1e-4, rtol=0) - try: - torch.testing.assert_close(retrieved_orientations, desired_world_orientations, atol=1e-4, rtol=0) - except AssertionError: - torch.testing.assert_close(retrieved_orientations, -desired_world_orientations, atol=1e-4, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_set_local_poses(device, backend): - """Test setting local poses in XformPrimView.""" - _skip_if_backend_unavailable(backend, device) +def _get_parent_positions(num_envs, device="cpu"): + """Read parent Xform positions from USD.""" stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create parent - sim_utils.create_prim("/World/Parent", "Xform", translation=(5.0, 5.0, 5.0), stage=stage) - - # Create children - num_prims = 4 - for i in range(num_prims): - sim_utils.create_prim(f"/World/Parent/Child_{i}", prim_type, translation=(0.0, 0.0, 0.0), stage=stage) - - # Create view - view = _create_view("/World/Parent/Child_.*", device=device, backend=backend) - - # Set new local poses - new_translations = torch.tensor([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0], [4.0, 4.0, 4.0]], device=device) - new_orientations = torch.tensor( - [ - [0.0, 0.0, 0.0, 1.0], - [0.0, 0.0, 0.7071068, 0.7071068], - [0.7071068, 0.0, 0.0, 0.7071068], - [0.3826834, 0.0, 0.0, 0.9238795], - ], - device=device, - ) - - view.set_local_poses(new_translations, new_orientations) + xform_cache = UsdGeom.XformCache() + positions = [] + for i in range(num_envs): + prim = stage.GetPrimAtPath(f"/World/Parent_{i}") + tf = xform_cache.GetLocalToWorldTransform(prim) + t = tf.ExtractTranslation() + positions.append([float(t[0]), float(t[1]), float(t[2])]) + return torch.tensor(positions, dtype=torch.float32, device=device) - # Get local poses back - retrieved_translations, retrieved_orientations = view.get_local_poses() - # Verify they match - torch.testing.assert_close(retrieved_translations, new_translations, atol=1e-5, rtol=0) - try: - torch.testing.assert_close(retrieved_orientations, new_orientations, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(retrieved_orientations, -new_orientations, atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_set_local_poses_only_translations(device, backend): - """Test setting only local translations.""" - _skip_if_backend_unavailable(backend, device) +def _set_parent_positions(positions, num_envs): + """Write parent Xform positions to USD.""" + from pxr import Sdf # noqa: PLC0415 stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create parent and children with specific orientations - sim_utils.create_prim("/World/Parent", "Xform", translation=(0.0, 0.0, 0.0), stage=stage) - initial_quat = (0.0, 0.0, 0.7071068, 0.7071068) - - for i in range(3): - sim_utils.create_prim( - f"/World/Parent/Child_{i}", - prim_type, - translation=(0.0, 0.0, 0.0), - orientation=initial_quat, - stage=stage, + with Sdf.ChangeBlock(): + for i in range(num_envs): + prim = stage.GetPrimAtPath(f"/World/Parent_{i}") + pos = positions[i] + prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(float(pos[0]), float(pos[1]), float(pos[2]))) + + +@pytest.fixture +def view_factory(): + """USD factory: parent Xform at PARENT_POS + child Xform at CHILD_OFFSET.""" + + def factory(num_envs: int, device: str) -> ViewBundle: + stage = sim_utils.get_current_stage() + for i in range(num_envs): + sim_utils.create_prim(f"/World/Parent_{i}", "Xform", translation=PARENT_POS, stage=stage) + sim_utils.create_prim(f"/World/Parent_{i}/Child", "Xform", translation=CHILD_OFFSET, stage=stage) + + view = FrameView("/World/Parent_.*/Child", device=device) + return ViewBundle( + view=view, + get_parent_pos=_get_parent_positions, + set_parent_pos=_set_parent_positions, + teardown=lambda: None, ) - # Create view - view = _create_view("/World/Parent/Child_.*", device=device, backend=backend) - - # Get initial orientations - _, initial_orientations = view.get_local_poses() + return factory - # Set only translations - new_translations = torch.tensor([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]], device=device) - view.set_local_poses(translations=new_translations, orientations=None) - # Get poses back - retrieved_translations, retrieved_orientations = view.get_local_poses() - - # Translations should be updated - torch.testing.assert_close(retrieved_translations, new_translations, atol=1e-5, rtol=0) - - # Orientations should be unchanged - try: - torch.testing.assert_close(retrieved_orientations, initial_orientations, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(retrieved_orientations, -initial_orientations, atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_set_scales(device, backend): - """Test setting scales in XformPrimView.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create prims - num_prims = 5 - for i in range(num_prims): - sim_utils.create_prim(f"/World/Object_{i}", prim_type, scale=(1.0, 1.0, 1.0), stage=stage) - - # Create view - view = _create_view("/World/Object_.*", device=device, backend=backend) - - # Set new scales - new_scales = torch.tensor( - [[2.0, 2.0, 2.0], [1.0, 2.0, 3.0], [0.5, 0.5, 0.5], [3.0, 1.0, 2.0], [1.5, 1.5, 1.5]], device=device - ) - - view.set_scales(new_scales) - - # Get scales back - retrieved_scales = view.get_scales() - - # Verify they match - torch.testing.assert_close(retrieved_scales, new_scales, atol=1e-5, rtol=0) +# ================================================================== +# USD-only: Visibility +# ================================================================== @pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_set_visibility(device): +def test_visibility_toggle(device): """Test toggling visibility multiple times.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") stage = sim_utils.get_current_stage() - - # Create prims num_prims = 3 for i in range(num_prims): sim_utils.create_prim(f"/World/Object_{i}", "Xform", stage=stage) - # Create view - view = XformPrimView("/World/Object_.*", device=device) + view = FrameView("/World/Object_.*", device=device) - # Initial state: all visible - visibility = view.get_visibility() - assert torch.all(visibility), "All should be visible initially" + assert torch.all(view.get_visibility()) - # Make all invisible view.set_visibility(torch.zeros(num_prims, dtype=torch.bool, device=device)) - visibility = view.get_visibility() - assert not torch.any(visibility), "All should be invisible" + assert not torch.any(view.get_visibility()) - # Make all visible again view.set_visibility(torch.ones(num_prims, dtype=torch.bool, device=device)) - visibility = view.get_visibility() - assert torch.all(visibility), "All should be visible again" - - # Toggle individual prims - view.set_visibility(torch.tensor([False], dtype=torch.bool, device=device), indices=[1]) - visibility = view.get_visibility() - assert visibility[0] and not visibility[1] and visibility[2], "Only middle prim should be invisible" - - -""" -Tests - Index Handling. -""" - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("index_type", ["list", "torch_tensor", "slice_none"]) -@pytest.mark.parametrize("method", ["world_poses", "local_poses", "scales", "visibility"]) -def test_index_types_get_methods(device, index_type, method): - """Test that getter methods work with different index types.""" - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - stage = sim_utils.get_current_stage() - - # Create prims based on method type - num_prims = 10 - if method == "local_poses": - # Create parent and children for local poses - sim_utils.create_prim("/World/Parent", "Xform", translation=(10.0, 0.0, 0.0), stage=stage) - for i in range(num_prims): - sim_utils.create_prim( - f"/World/Parent/Child_{i}", "Xform", translation=(float(i), float(i) * 0.5, 0.0), stage=stage - ) - view = XformPrimView("/World/Parent/Child_.*", device=device) - elif method == "scales": - # Create prims with different scales - for i in range(num_prims): - scale = (1.0 + i * 0.5, 1.0 + i * 0.3, 1.0 + i * 0.2) - sim_utils.create_prim(f"/World/Object_{i}", "Xform", scale=scale, stage=stage) - view = XformPrimView("/World/Object_.*", device=device) - else: # world_poses - # Create prims with different positions - for i in range(num_prims): - sim_utils.create_prim(f"/World/Object_{i}", "Xform", translation=(float(i), 0.0, 0.0), stage=stage) - view = XformPrimView("/World/Object_.*", device=device) - - # Get all data as reference - if method == "world_poses": - all_data1, all_data2 = view.get_world_poses() - elif method == "local_poses": - all_data1, all_data2 = view.get_local_poses() - elif method == "scales": - all_data1 = view.get_scales() - all_data2 = None - else: # visibility - all_data1 = view.get_visibility() - all_data2 = None - - # Prepare indices - target_indices_base = [2, 5, 7] - indices, target_indices = _prepare_indices(index_type, target_indices_base, num_prims, device) - - # Get subset - if method == "world_poses": - subset_data1, subset_data2 = view.get_world_poses(indices=indices) # type: ignore[arg-type] - elif method == "local_poses": - subset_data1, subset_data2 = view.get_local_poses(indices=indices) # type: ignore[arg-type] - elif method == "scales": - subset_data1 = view.get_scales(indices=indices) # type: ignore[arg-type] - subset_data2 = None - else: # visibility - subset_data1 = view.get_visibility(indices=indices) # type: ignore[arg-type] - subset_data2 = None - - # Verify shapes - expected_count = len(target_indices) - if method == "visibility": - assert subset_data1.shape == (expected_count,) - else: - assert subset_data1.shape == (expected_count, 3) - if subset_data2 is not None: - assert subset_data2.shape == (expected_count, 4) - - # Verify values - target_indices_tensor = torch.tensor(target_indices, dtype=torch.int64, device=device) - torch.testing.assert_close(subset_data1, all_data1[target_indices_tensor], atol=1e-5, rtol=0) - if subset_data2 is not None and all_data2 is not None: - torch.testing.assert_close(subset_data2, all_data2[target_indices_tensor], atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("index_type", ["list", "torch_tensor", "slice_none"]) -@pytest.mark.parametrize("method", ["world_poses", "local_poses", "scales", "visibility"]) -def test_index_types_set_methods(device, index_type, method): - """Test that setter methods work with different index types.""" - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - stage = sim_utils.get_current_stage() - - # Create prims based on method type - num_prims = 10 - if method == "local_poses": - # Create parent and children for local poses - sim_utils.create_prim("/World/Parent", "Xform", translation=(5.0, 5.0, 0.0), stage=stage) - for i in range(num_prims): - sim_utils.create_prim(f"/World/Parent/Child_{i}", "Xform", translation=(float(i), 0.0, 0.0), stage=stage) - view = XformPrimView("/World/Parent/Child_.*", device=device) - else: # world_poses or scales - for i in range(num_prims): - sim_utils.create_prim(f"/World/Object_{i}", "Xform", translation=(0.0, 0.0, 0.0), stage=stage) - view = XformPrimView("/World/Object_.*", device=device) - - # Get initial data - if method == "world_poses": - initial_data1, initial_data2 = view.get_world_poses() - elif method == "local_poses": - initial_data1, initial_data2 = view.get_local_poses() - elif method == "scales": - initial_data1 = view.get_scales() - initial_data2 = None - else: # visibility - initial_data1 = view.get_visibility() - initial_data2 = None - - # Prepare indices - target_indices_base = [2, 5, 7] - indices, target_indices = _prepare_indices(index_type, target_indices_base, num_prims, device) - - # Prepare new data - num_to_set = len(target_indices) - if method in ["world_poses", "local_poses"]: - new_data1 = torch.randn(num_to_set, 3, device=device) * 10.0 - new_data2 = torch.tensor([[0.0, 0.0, 0.0, 1.0]] * num_to_set, dtype=torch.float32, device=device) - elif method == "scales": - new_data1 = torch.rand(num_to_set, 3, device=device) * 2.0 + 0.5 - new_data2 = None - else: # visibility - # Set to False to test change (default is True) - new_data1 = torch.zeros(num_to_set, dtype=torch.bool, device=device) - new_data2 = None - - # Set data - if method == "world_poses": - view.set_world_poses(positions=new_data1, orientations=new_data2, indices=indices) # type: ignore[arg-type] - elif method == "local_poses": - view.set_local_poses(translations=new_data1, orientations=new_data2, indices=indices) # type: ignore[arg-type] - elif method == "scales": - view.set_scales(scales=new_data1, indices=indices) # type: ignore[arg-type] - else: # visibility - view.set_visibility(visibility=new_data1, indices=indices) # type: ignore[arg-type] - - # Get all data after update - if method == "world_poses": - updated_data1, updated_data2 = view.get_world_poses() - elif method == "local_poses": - updated_data1, updated_data2 = view.get_local_poses() - elif method == "scales": - updated_data1 = view.get_scales() - updated_data2 = None - else: # visibility - updated_data1 = view.get_visibility() - updated_data2 = None - - # Verify that specified indices were updated - for i, target_idx in enumerate(target_indices): - torch.testing.assert_close(updated_data1[target_idx], new_data1[i], atol=1e-5, rtol=0) - if new_data2 is not None and updated_data2 is not None: - try: - torch.testing.assert_close(updated_data2[target_idx], new_data2[i], atol=1e-5, rtol=0) - except AssertionError: - # Account for quaternion sign ambiguity - torch.testing.assert_close(updated_data2[target_idx], -new_data2[i], atol=1e-5, rtol=0) - - # Verify that other indices were NOT updated (only for non-slice(None) cases) - if index_type != "slice_none": - for i in range(num_prims): - if i not in target_indices: - torch.testing.assert_close(updated_data1[i], initial_data1[i], atol=1e-5, rtol=0) - if initial_data2 is not None and updated_data2 is not None: - try: - torch.testing.assert_close(updated_data2[i], initial_data2[i], atol=1e-5, rtol=0) - except AssertionError: - # Account for quaternion sign ambiguity - torch.testing.assert_close(updated_data2[i], -initial_data2[i], atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_indices_single_element(device, backend): - """Test with a single index.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create prims - num_prims = 5 - for i in range(num_prims): - sim_utils.create_prim(f"/World/Object_{i}", prim_type, translation=(float(i), 0.0, 0.0), stage=stage) - - # Create view - view = _create_view("/World/Object_.*", device=device, backend=backend) - - # Test with single index - indices = [3] - positions, orientations = view.get_world_poses(indices=indices) - - # Verify shapes - assert positions.shape == (1, 3) - assert orientations.shape == (1, 4) - - # Set pose for single index - new_position = torch.tensor([[100.0, 200.0, 300.0]], device=device) - view.set_world_poses(positions=new_position, indices=indices) - - # Verify it was set - retrieved_positions, _ = view.get_world_poses(indices=indices) - torch.testing.assert_close(retrieved_positions, new_position, atol=1e-5, rtol=0) - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_indices_out_of_order(device, backend): - """Test with indices provided in non-sequential order.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) + assert torch.all(view.get_visibility()) - # Create prims - num_prims = 10 - for i in range(num_prims): - sim_utils.create_prim(f"/World/Object_{i}", prim_type, translation=(0.0, 0.0, 0.0), stage=stage) - - # Create view - view = _create_view("/World/Object_.*", device=device, backend=backend) - - # Use out-of-order indices - indices = [7, 2, 9, 0, 5] - new_positions = torch.tensor( - [[7.0, 0.0, 0.0], [2.0, 0.0, 0.0], [9.0, 0.0, 0.0], [0.0, 0.0, 0.0], [5.0, 0.0, 0.0]], device=device + view.set_visibility( + torch.tensor([False], dtype=torch.bool, device=device), indices=wp.array([1], dtype=wp.int32, device=device) ) - - # Set poses with out-of-order indices - view.set_world_poses(positions=new_positions, indices=indices) - - # Get all poses - all_positions, _ = view.get_world_poses() - - # Verify each index got the correct value - expected_x_values = [0.0, 0.0, 2.0, 0.0, 0.0, 5.0, 0.0, 7.0, 0.0, 9.0] - for i in range(num_prims): - assert abs(all_positions[i, 0].item() - expected_x_values[i]) < 1e-5 - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("backend", ["usd", "fabric"]) -def test_indices_with_only_positions_or_orientations(device, backend): - """Test indices work correctly when setting only positions or only orientations.""" - _skip_if_backend_unavailable(backend, device) - - stage = sim_utils.get_current_stage() - prim_type = _prim_type_for_backend(backend) - - # Create prims - num_prims = 5 - for i in range(num_prims): - sim_utils.create_prim( - f"/World/Object_{i}", - prim_type, - translation=(0.0, 0.0, 0.0), - orientation=(0.0, 0.0, 0.0, 1.0), - stage=stage, - ) - - # Create view - view = _create_view("/World/Object_.*", device=device, backend=backend) - - # Get initial poses - initial_positions, initial_orientations = view.get_world_poses() - - # Set only positions for specific indices - indices = [1, 3] - new_positions = torch.tensor([[10.0, 0.0, 0.0], [30.0, 0.0, 0.0]], device=device) - view.set_world_poses(positions=new_positions, orientations=None, indices=indices) - - # Get updated poses - updated_positions, updated_orientations = view.get_world_poses() - - # Verify positions updated for indices 1 and 3, others unchanged - torch.testing.assert_close(updated_positions[1], new_positions[0], atol=1e-5, rtol=0) - torch.testing.assert_close(updated_positions[3], new_positions[1], atol=1e-5, rtol=0) - torch.testing.assert_close(updated_positions[0], initial_positions[0], atol=1e-5, rtol=0) - - # Verify all orientations unchanged - try: - torch.testing.assert_close(updated_orientations, initial_orientations, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(updated_orientations, -initial_orientations, atol=1e-5, rtol=0) - - # Now set only orientations for different indices - indices2 = [0, 4] - new_orientations = torch.tensor([[0.0, 0.0, 0.7071068, 0.7071068], [0.7071068, 0.0, 0.0, 0.7071068]], device=device) - view.set_world_poses(positions=None, orientations=new_orientations, indices=indices2) - - # Get final poses - final_positions, final_orientations = view.get_world_poses() - - # Verify positions unchanged from previous step - torch.testing.assert_close(final_positions, updated_positions, atol=1e-5, rtol=0) - - # Verify orientations updated for indices 0 and 4 - try: - torch.testing.assert_close(final_orientations[0], new_orientations[0], atol=1e-5, rtol=0) - torch.testing.assert_close(final_orientations[4], new_orientations[1], atol=1e-5, rtol=0) - except AssertionError: - # Account for quaternion sign ambiguity - torch.testing.assert_close(final_orientations[0], -new_orientations[0], atol=1e-5, rtol=0) - torch.testing.assert_close(final_orientations[4], -new_orientations[1], atol=1e-5, rtol=0) + vis = view.get_visibility() + assert vis[0] and not vis[1] and vis[2] @pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_index_type_none_equivalent_to_all(device): - """Test that indices=None is equivalent to getting/setting all prims.""" +def test_visibility_parent_inheritance(device): + """Making a parent invisible hides all children.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") stage = sim_utils.get_current_stage() + sim_utils.create_prim("/World/Parent", "Xform", stage=stage) + for i in range(4): + sim_utils.create_prim(f"/World/Parent/Child_{i}", "Xform", stage=stage) - # Create prims - num_prims = 6 - for i in range(num_prims): - sim_utils.create_prim(f"/World/Object_{i}", "Xform", translation=(float(i), 0.0, 0.0), stage=stage) - - # Create view - view = XformPrimView("/World/Object_.*", device=device) - - # Get poses with indices=None - pos_none, quat_none = view.get_world_poses(indices=None) - - # Get poses with no argument (default) - pos_default, quat_default = view.get_world_poses() - - # Get poses with slice(None) - pos_slice, quat_slice = view.get_world_poses(indices=slice(None)) # type: ignore[arg-type] - - # All should be equivalent - torch.testing.assert_close(pos_none, pos_default, atol=1e-10, rtol=0) - torch.testing.assert_close(quat_none, quat_default, atol=1e-10, rtol=0) - torch.testing.assert_close(pos_none, pos_slice, atol=1e-10, rtol=0) - torch.testing.assert_close(quat_none, quat_slice, atol=1e-10, rtol=0) - - # Test the same for set operations - new_positions = torch.randn(num_prims, 3, device=device) * 10.0 - new_orientations = torch.tensor([[0.0, 0.0, 0.0, 1.0]] * num_prims, dtype=torch.float32, device=device) - - # Set with indices=None - view.set_world_poses(positions=new_positions, orientations=new_orientations, indices=None) - pos_after_none, quat_after_none = view.get_world_poses() - - # Reset - view.set_world_poses(positions=torch.zeros(num_prims, 3, device=device), indices=None) + parent_view = FrameView("/World/Parent", device=device) + children_view = FrameView("/World/Parent/Child_.*", device=device) - # Set with slice(None) - view.set_world_poses( - positions=new_positions, - orientations=new_orientations, - indices=slice(None), # type: ignore[arg-type] - ) - pos_after_slice, quat_after_slice = view.get_world_poses() + parent_view.set_visibility(torch.tensor([False], dtype=torch.bool, device=device)) + assert not torch.any(children_view.get_visibility()) - # Should be equivalent - torch.testing.assert_close(pos_after_none, pos_after_slice, atol=1e-5, rtol=0) - torch.testing.assert_close(quat_after_none, quat_after_slice, atol=1e-5, rtol=0) + parent_view.set_visibility(torch.tensor([True], dtype=torch.bool, device=device)) + assert torch.all(children_view.get_visibility()) -""" -Tests - Integration. -""" +# ================================================================== +# USD-only: Prim ordering +# ================================================================== @pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_with_franka_robots(device): - """Test XformPrimView with real Franka robot USD assets.""" +def test_prim_ordering_follows_creation_order(device): + """Prims are returned in USD creation order (DFS), not alphabetical.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") stage = sim_utils.get_current_stage() + num_envs = 3 + for i in range(num_envs): + sim_utils.create_prim(f"/World/Env_{i}/Object_1", "Xform", stage=stage) + sim_utils.create_prim(f"/World/Env_{i}/Object_0", "Xform", stage=stage) + sim_utils.create_prim(f"/World/Env_{i}/Object_A", "Xform", stage=stage) - # Load Franka robot assets - franka_usd_path = f"{ISAAC_NUCLEUS_DIR}/Robots/FrankaRobotics/FrankaPanda/franka.usd" - - # Add two Franka robots to the stage - sim_utils.create_prim("/World/Franka_1", "Xform", usd_path=franka_usd_path, stage=stage) - sim_utils.create_prim("/World/Franka_2", "Xform", usd_path=franka_usd_path, stage=stage) - - # Create view for both Frankas - frankas_view = XformPrimView("/World/Franka_.*", device=device) - - # Verify count - assert frankas_view.count == 2 + view = FrameView("/World/Env_.*/Object_.*", device=device) + expected = [] + for i in range(num_envs): + expected += [f"/World/Env_{i}/Object_1", f"/World/Env_{i}/Object_0", f"/World/Env_{i}/Object_A"] - # Get initial world poses (should be at origin) - initial_positions, initial_orientations = frankas_view.get_world_poses() + assert view.prim_paths == expected - # Verify initial positions are at origin - expected_initial_positions = torch.zeros(2, 3, device=device) - torch.testing.assert_close(initial_positions, expected_initial_positions, atol=1e-5, rtol=0) - # Verify initial orientations are identity - expected_initial_orientations = torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0]], device=device) - try: - torch.testing.assert_close(initial_orientations, expected_initial_orientations, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(initial_orientations, -expected_initial_orientations, atol=1e-5, rtol=0) - - # Set new world poses - new_positions = torch.tensor([[10.0, 10.0, 0.0], [-40.0, -40.0, 0.0]], device=device) - # 90° rotation around Z axis for first, -90° for second - new_orientations = torch.tensor( - [[0.0, 0.0, 0.7071068, 0.7071068], [0.0, 0.0, -0.7071068, 0.7071068]], device=device - ) - - frankas_view.set_world_poses(positions=new_positions, orientations=new_orientations) - - # Get poses back and verify - retrieved_positions, retrieved_orientations = frankas_view.get_world_poses() - - torch.testing.assert_close(retrieved_positions, new_positions, atol=1e-5, rtol=0) - try: - torch.testing.assert_close(retrieved_orientations, new_orientations, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(retrieved_orientations, -new_orientations, atol=1e-5, rtol=0) +# ================================================================== +# USD-only: xformOp standardization +# ================================================================== @pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_with_nested_targets(device): - """Test with nested frame/target structure similar to Isaac Sim tests.""" +def test_standardize_transform_op(device): + """FrameView standardizes a prim with xformOp:transform to translate/orient/scale.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") - stage = sim_utils.get_current_stage() - - # Create frames and targets - for i in range(1, 4): - sim_utils.create_prim(f"/World/Frame_{i}", "Xform", stage=stage) - sim_utils.create_prim(f"/World/Frame_{i}/Target", "Xform", stage=stage) - - # Create views - frames_view = XformPrimView("/World/Frame_.*", device=device) - targets_view = XformPrimView("/World/Frame_.*/Target", device=device) - - assert frames_view.count == 3 - assert targets_view.count == 3 + expected_pos = (3.0, -1.0, 0.5) + matrix = Gf.Matrix4d(1.0) + matrix.SetTranslateOnly(Gf.Vec3d(*expected_pos)) - # Set local poses for frames - frame_translations = torch.tensor([[0.0, 0.0, 0.0], [0.0, 10.0, 5.0], [0.0, 3.0, 5.0]], device=device) - frames_view.set_local_poses(translations=frame_translations) + stage = sim_utils.get_current_stage() + prim = stage.DefinePrim("/World/TransformPrim", "Xform") + UsdGeom.Xformable(prim).AddTransformOp().Set(matrix) - # Set local poses for targets - target_translations = torch.tensor([[0.0, 20.0, 10.0], [0.0, 30.0, 20.0], [0.0, 50.0, 10.0]], device=device) - targets_view.set_local_poses(translations=target_translations) + view = FrameView("/World/TransformPrim", device=device) + assert sim_utils.validate_standard_xform_ops(view.prims[0]) - # Get world poses of targets - world_positions, _ = targets_view.get_world_poses() + ordered_ops = UsdGeom.Xformable(view.prims[0]).GetOrderedXformOps() + op_names = [op.GetOpName() for op in ordered_ops] + assert op_names == ["xformOp:translate", "xformOp:orient", "xformOp:scale"] + assert ordered_ops[0].Get() == Gf.Vec3d(*expected_pos) - # Expected world positions are frame_translation + target_translation - expected_positions = torch.tensor([[0.0, 20.0, 10.0], [0.0, 40.0, 25.0], [0.0, 53.0, 15.0]], device=device) - torch.testing.assert_close(world_positions, expected_positions, atol=1e-5, rtol=0) +# ================================================================== +# USD-only: Nested hierarchy (frame + target) +# ================================================================== @pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_visibility_with_hierarchy(device): - """Test visibility with parent-child hierarchy and inheritance.""" +def test_nested_hierarchy_world_poses(device): + """World pose of nested child == sum of parent + child translations.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") stage = sim_utils.get_current_stage() + frame_positions = [(0.0, 0.0, 0.0), (0.0, 10.0, 5.0), (0.0, 3.0, 5.0)] + target_positions = [(0.0, 20.0, 10.0), (0.0, 30.0, 20.0), (0.0, 50.0, 10.0)] - # Create parent and children - sim_utils.create_prim("/World/Parent", "Xform", stage=stage) - - num_children = 4 - for i in range(num_children): - sim_utils.create_prim(f"/World/Parent/Child_{i}", "Xform", stage=stage) - - # Create views for both parent and children - parent_view = XformPrimView("/World/Parent", device=device) - children_view = XformPrimView("/World/Parent/Child_.*", device=device) - - # Verify parent and all children are visible initially - parent_visibility = parent_view.get_visibility() - children_visibility = children_view.get_visibility() - assert parent_visibility[0], "Parent should be visible initially" - assert torch.all(children_visibility), "All children should be visible initially" - - # Make some children invisible directly - new_visibility = torch.tensor([True, False, True, False], dtype=torch.bool, device=device) - children_view.set_visibility(new_visibility) - - # Verify the visibility changes - retrieved_visibility = children_view.get_visibility() - torch.testing.assert_close(retrieved_visibility, new_visibility) - - # Make all children visible again - children_view.set_visibility(torch.ones(num_children, dtype=torch.bool, device=device)) - all_visible = children_view.get_visibility() - assert torch.all(all_visible), "All children should be visible again" - - # Now test parent visibility inheritance: - # Make parent invisible - parent_view.set_visibility(torch.tensor([False], dtype=torch.bool, device=device)) - - # Verify parent is invisible - parent_visibility = parent_view.get_visibility() - assert not parent_visibility[0], "Parent should be invisible" + for i in range(3): + sim_utils.create_prim(f"/World/Frame_{i}", "Xform", translation=frame_positions[i], stage=stage) + sim_utils.create_prim(f"/World/Frame_{i}/Target", "Xform", translation=target_positions[i], stage=stage) - # Verify children are also invisible (due to parent being invisible) - children_visibility = children_view.get_visibility() - assert not torch.any(children_visibility), "All children should be invisible when parent is invisible" + frames_view = FrameView("/World/Frame_.*", device=device) + targets_view = FrameView("/World/Frame_.*/Target", device=device) - # Make parent visible again - parent_view.set_visibility(torch.tensor([True], dtype=torch.bool, device=device)) + frames_view.set_local_poses(translations=torch.tensor(frame_positions, device=device)) + targets_view.set_local_poses(translations=torch.tensor(target_positions, device=device)) - # Verify parent is visible - parent_visibility = parent_view.get_visibility() - assert parent_visibility[0], "Parent should be visible again" - - # Verify children are also visible again - children_visibility = children_view.get_visibility() - assert torch.all(children_visibility), "All children should be visible again when parent is visible" + world_pos = wp.to_torch(targets_view.get_world_poses()[0]) + expected = torch.tensor( + [[f[j] + t[j] for j in range(3)] for f, t in zip(frame_positions, target_positions)], + device=device, + ) + torch.testing.assert_close(world_pos, expected, atol=1e-5, rtol=0) -""" -Tests - Comparison with Isaac Sim Implementation. -""" +# ================================================================== +# USD-only: Comparison with Isaac Sim +# ================================================================== def test_compare_get_world_poses_with_isaacsim(): """Compare get_world_poses with Isaac Sim's implementation.""" - stage = sim_utils.get_current_stage() - - # Check if Isaac Sim is available if _IsaacSimXformPrimView is None: pytest.skip("Isaac Sim is not available") - # Create prims with various poses + stage = sim_utils.get_current_stage() num_prims = 10 for i in range(num_prims): pos = (i * 2.0, i * 0.5, i * 1.5) - # Vary orientations - if i % 3 == 0: - quat = (0.0, 0.0, 0.0, 1.0) # Identity - elif i % 3 == 1: - quat = (0.0, 0.0, 0.7071068, 0.7071068) # 90 deg around Z - else: - quat = (0.7071068, 0.0, 0.0, 0.7071068) # 90 deg around X + quat = (0.0, 0.0, 0.0, 1.0) if i % 2 == 0 else (0.0, 0.0, 0.7071068, 0.7071068) sim_utils.create_prim(f"/World/Env_{i}/Object", "Xform", translation=pos, orientation=quat, stage=stage) pattern = "/World/Env_.*/Object" - - # Create both views - isaaclab_view = XformPrimView(pattern, device="cpu") + isaaclab_view = FrameView(pattern, device="cpu") isaacsim_view = _IsaacSimXformPrimView(pattern, reset_xform_properties=False) - # Get world poses from both - isaaclab_pos, isaaclab_quat = isaaclab_view.get_world_poses() # xyzw - isaacsim_pos, isaacsim_quat = isaacsim_view.get_world_poses() # wxyz - - # Convert Isaac Sim results to torch tensors if needed + isaaclab_pos = wp.to_torch(isaaclab_view.get_world_poses()[0]) + isaacsim_pos, isaacsim_quat = isaacsim_view.get_world_poses() if not isinstance(isaacsim_pos, torch.Tensor): isaacsim_pos = torch.tensor(isaacsim_pos, dtype=torch.float32) - if not isinstance(isaacsim_quat, torch.Tensor): - isaacsim_quat = torch.tensor(isaacsim_quat, dtype=torch.float32).roll(-1, dims=1) - # Compare results torch.testing.assert_close(isaaclab_pos, isaacsim_pos, atol=1e-5, rtol=0) - # Compare quaternions (account for sign ambiguity) - try: - torch.testing.assert_close(isaaclab_quat, isaacsim_quat, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(isaaclab_quat, -isaacsim_quat, atol=1e-5, rtol=0) - - -def test_compare_set_world_poses_with_isaacsim(): - """Compare set_world_poses with Isaac Sim's implementation.""" - stage = sim_utils.get_current_stage() - - # Check if Isaac Sim is available - if _IsaacSimXformPrimView is None: - pytest.skip("Isaac Sim is not available") - - # Create prims - num_prims = 8 - for i in range(num_prims): - sim_utils.create_prim(f"/World/Env_{i}/Object", "Xform", translation=(0.0, 0.0, 0.0), stage=stage) - - pattern = "/World/Env_.*/Object" - - # Create both views - isaaclab_view = XformPrimView(pattern, device="cpu") - isaacsim_view = _IsaacSimXformPrimView(pattern, reset_xform_properties=False) - - # Generate new poses - new_positions = torch.randn(num_prims, 3) * 10.0 - new_orientations = torch.tensor([[0.0, 0.0, 0.0, 1.0]] * num_prims, dtype=torch.float32) - - # Set poses using both implementations - isaaclab_view.set_world_poses(new_positions.clone(), new_orientations.clone()) # xyzw - isaacsim_view.set_world_poses(new_positions.clone(), new_orientations.clone().roll(1, dims=1)) # wxyz - - # Get poses back from both - isaaclab_pos, isaaclab_quat = isaaclab_view.get_world_poses() # xyzw - isaacsim_pos, isaacsim_quat = isaacsim_view.get_world_poses() # wxyz - - # Convert Isaac Sim results to torch tensors if needed - if not isinstance(isaacsim_pos, torch.Tensor): - isaacsim_pos = torch.tensor(isaacsim_pos, dtype=torch.float32) - if not isinstance(isaacsim_quat, torch.Tensor): - isaacsim_quat = torch.tensor(isaacsim_quat, dtype=torch.float32).roll(-1, dims=1) - - # Compare results - both implementations should produce the same world poses - torch.testing.assert_close(isaaclab_pos, isaacsim_pos, atol=1e-4, rtol=0) - try: - torch.testing.assert_close(isaaclab_quat, isaacsim_quat, atol=1e-4, rtol=0) - except AssertionError: - torch.testing.assert_close(isaaclab_quat, -isaacsim_quat, atol=1e-4, rtol=0) - - -def test_compare_get_local_poses_with_isaacsim(): - """Compare get_local_poses with Isaac Sim's implementation.""" - stage = sim_utils.get_current_stage() - - # Check if Isaac Sim is available - if _IsaacSimXformPrimView is None: - pytest.skip("Isaac Sim is not available") - - # Create hierarchical prims - num_prims = 5 - for i in range(num_prims): - # Create parent - sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 5.0, 0.0, 0.0), stage=stage) - # Create child with local pose - local_pos = (1.0, float(i), 0.0) - local_quat = (0.0, 0.0, 0.0, 1.0) if i % 2 == 0 else (0.0, 0.0, 0.7071068, 0.7071068) - sim_utils.create_prim( - f"/World/Env_{i}/Object", "Xform", translation=local_pos, orientation=local_quat, stage=stage - ) - - pattern = "/World/Env_.*/Object" - - # Create both views - isaaclab_view = XformPrimView(pattern, device="cpu") - isaacsim_view = _IsaacSimXformPrimView(pattern, reset_xform_properties=False) - - # Get local poses from both - isaaclab_trans, isaaclab_quat = isaaclab_view.get_local_poses() - isaacsim_trans, isaacsim_quat = isaacsim_view.get_local_poses() - - # Convert Isaac Sim results to torch tensors if needed - if not isinstance(isaacsim_trans, torch.Tensor): - isaacsim_trans = torch.tensor(isaacsim_trans, dtype=torch.float32) - if not isinstance(isaacsim_quat, torch.Tensor): - isaacsim_quat = torch.tensor(isaacsim_quat, dtype=torch.float32).roll(-1, dims=1) - - # Compare results - torch.testing.assert_close(isaaclab_trans, isaacsim_trans, atol=1e-5, rtol=0) - try: - torch.testing.assert_close(isaaclab_quat, isaacsim_quat, atol=1e-5, rtol=0) - except AssertionError: - torch.testing.assert_close(isaaclab_quat, -isaacsim_quat, atol=1e-5, rtol=0) - - -def test_compare_set_local_poses_with_isaacsim(): - """Compare set_local_poses with Isaac Sim's implementation.""" - stage = sim_utils.get_current_stage() - - # Check if Isaac Sim is available - if _IsaacSimXformPrimView is None: - pytest.skip("Isaac Sim is not available") - - # Create hierarchical prims - num_prims = 6 - for i in range(num_prims): - sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0.0, 0.0), stage=stage) - sim_utils.create_prim(f"/World/Env_{i}/Object", "Xform", translation=(0.0, 0.0, 0.0), stage=stage) - - pattern = "/World/Env_.*/Object" - - # Create both views - isaaclab_view = XformPrimView(pattern, device="cpu") - isaacsim_view = _IsaacSimXformPrimView(pattern, reset_xform_properties=False) - - # Generate new local poses - new_translations = torch.randn(num_prims, 3) * 5.0 - new_orientations = torch.tensor( - [[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.7071068, 0.7071068]] * (num_prims // 2), dtype=torch.float32 - ) - # Set local poses using both implementations - isaaclab_view.set_local_poses(new_translations.clone(), new_orientations.clone()) - isaacsim_view.set_local_poses(new_translations.clone(), new_orientations.clone().roll(1, dims=1)) - - # Get local poses back from both - isaaclab_trans, isaaclab_quat = isaaclab_view.get_local_poses() - isaacsim_trans, isaacsim_quat = isaacsim_view.get_local_poses() - - # Convert Isaac Sim results to torch tensors if needed - if not isinstance(isaacsim_trans, torch.Tensor): - isaacsim_trans = torch.tensor(isaacsim_trans, dtype=torch.float32) - if not isinstance(isaacsim_quat, torch.Tensor): - isaacsim_quat = torch.tensor(isaacsim_quat, dtype=torch.float32).roll(-1, dims=1) - - # Compare results - torch.testing.assert_close(isaaclab_trans, isaacsim_trans, atol=1e-4, rtol=0) - try: - torch.testing.assert_close(isaaclab_quat, isaacsim_quat, atol=1e-4, rtol=0) - except AssertionError: - torch.testing.assert_close(isaaclab_quat, -isaacsim_quat, atol=1e-4, rtol=0) - - -""" -Tests - Fabric Operations. -""" - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_fabric_initialization(device): - """Test XformPrimView initialization with Fabric enabled.""" - _skip_if_backend_unavailable("fabric", device) - - stage = sim_utils.get_current_stage() - - # Create camera prims (Boundable prims that support Fabric) - num_prims = 5 - for i in range(num_prims): - sim_utils.create_prim(f"/World/Cam_{i}", "Camera", translation=(i * 1.0, 0.0, 1.0), stage=stage) - - # Create view with Fabric enabled - view = _create_view("/World/Cam_.*", device=device, backend="fabric") - - # Verify properties - assert view.count == num_prims - assert view.device == device - assert len(view.prims) == num_prims +# ================================================================== +# USD-only: Franka integration +# ================================================================== @pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_fabric_usd_consistency(device): - """Test that Fabric round-trip (write→read) is consistent, matching Isaac Sim's design. - - Note: This does NOT test Fabric vs USD reads on initialization, as Fabric is designed - for write-first workflows. Instead, it tests that: - 1. Fabric write→read round-trip works correctly - 2. This matches Isaac Sim's Fabric behavior - """ - _skip_if_backend_unavailable("fabric", device) +def test_with_franka_robots(device): + """Verify FrameView works with real Franka robot USD assets.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") stage = sim_utils.get_current_stage() + franka_usd_path = f"{ISAAC_NUCLEUS_DIR}/Robots/FrankaRobotics/FrankaPanda/franka.usd" - # Create prims - num_prims = 5 - for i in range(num_prims): - sim_utils.create_prim( - f"/World/Cam_{i}", - "Camera", - translation=(i * 1.0, 2.0, 3.0), - orientation=(0.0, 0.0, 0.7071068, 0.7071068), - stage=stage, - ) - - # Create Fabric view - view_fabric = _create_view("/World/Cam_.*", device=device, backend="fabric") - - # Test Fabric write→read round-trip (Isaac Sim's intended workflow) - # Initialize Fabric state by WRITING first - init_positions = torch.zeros((num_prims, 3), dtype=torch.float32, device=device) - init_positions[:, 0] = torch.arange(num_prims, dtype=torch.float32, device=device) - init_positions[:, 1] = 2.0 - init_positions[:, 2] = 3.0 - init_orientations = torch.tensor([[0.0, 0.0, 0.7071068, 0.7071068]] * num_prims, dtype=torch.float32, device=device) - - view_fabric.set_world_poses(init_positions, init_orientations) + sim_utils.create_prim("/World/Franka_1", "Xform", usd_path=franka_usd_path, stage=stage) + sim_utils.create_prim("/World/Franka_2", "Xform", usd_path=franka_usd_path, stage=stage) - # Read back from Fabric (should match what we wrote) - pos_fabric, quat_fabric = view_fabric.get_world_poses() - torch.testing.assert_close(pos_fabric, init_positions, atol=1e-4, rtol=0) - torch.testing.assert_close(quat_fabric, init_orientations, atol=1e-4, rtol=0) + view = FrameView("/World/Franka_.*", device=device) + assert view.count == 2 - # Test another round-trip with different values - new_positions = torch.rand((num_prims, 3), dtype=torch.float32, device=device) * 10.0 - new_orientations = torch.tensor([[0.0, 0.0, 0.0, 1.0]] * num_prims, dtype=torch.float32, device=device) + positions = wp.to_torch(view.get_world_poses()[0]) + torch.testing.assert_close(positions, torch.zeros(2, 3, device=device), atol=1e-5, rtol=0) - view_fabric.set_world_poses(new_positions, new_orientations) + new_pos = torch.tensor([[10.0, 10.0, 0.0], [-40.0, -40.0, 0.0]], device=device) + new_quat = torch.tensor([[0.0, 0.0, 0.7071068, 0.7071068], [0.0, 0.0, -0.7071068, 0.7071068]], device=device) + view.set_world_poses(positions=new_pos, orientations=new_quat) - # Read back from Fabric (should match) - pos_fabric_after, quat_fabric_after = view_fabric.get_world_poses() - torch.testing.assert_close(pos_fabric_after, new_positions, atol=1e-4, rtol=0) - torch.testing.assert_close(quat_fabric_after, new_orientations, atol=1e-4, rtol=0) + ret_pos = wp.to_torch(view.get_world_poses()[0]) + torch.testing.assert_close(ret_pos, new_pos, atol=1e-5, rtol=0) diff --git a/source/isaaclab/test/terrains/check_terrain_importer.py b/source/isaaclab/test/terrains/check_terrain_importer.py index 519f84fc2743..c024d33bb5f1 100644 --- a/source/isaaclab/test/terrains/check_terrain_importer.py +++ b/source/isaaclab/test/terrains/check_terrain_importer.py @@ -153,8 +153,8 @@ def main(): physics_scene_path, "/World/collisions", prim_paths=envs_prim_paths, global_paths=["/World/ground"] ) - # Set ball positions over terrain origins using XformPrimView (before simulation starts) - xform_view = sim_utils.XformPrimView("/World/envs/env_.*/ball") + # Set ball positions over terrain origins using FrameView (before simulation starts) + xform_view = sim_utils.FrameView("/World/envs/env_.*/ball") # cache initial state of the balls ball_initial_positions = terrain_importer.env_origins.clone() ball_initial_positions[:, 2] += 5.0 diff --git a/source/isaaclab/test/terrains/test_terrain_importer.py b/source/isaaclab/test/terrains/test_terrain_importer.py index 3951296e978a..8842c6df673b 100644 --- a/source/isaaclab/test/terrains/test_terrain_importer.py +++ b/source/isaaclab/test/terrains/test_terrain_importer.py @@ -327,8 +327,8 @@ def _populate_scene(sim: SimulationContext, num_balls: int = 2048, geom_sphere: ) # Set ball positions over terrain origins - # Create a view over all the balls using Isaac Lab's XformPrimView - ball_view = sim_utils.XformPrimView("/World/envs/env_.*/ball") + # Create a view over all the balls using Isaac Lab's FrameView + ball_view = sim_utils.FrameView("/World/envs/env_.*/ball") # cache initial state of the balls ball_initial_positions = terrain_importer.env_origins.clone() ball_initial_positions[:, 2] += 5.0 diff --git a/source/isaaclab_contrib/test/sensors/test_visuotactile_sensor.py b/source/isaaclab_contrib/test/sensors/test_visuotactile_sensor.py index 88a2249bafa0..5ea9a373fb3b 100644 --- a/source/isaaclab_contrib/test/sensors/test_visuotactile_sensor.py +++ b/source/isaaclab_contrib/test/sensors/test_visuotactile_sensor.py @@ -307,7 +307,8 @@ def test_sensor_cam_set_wrong_prim(setup_tactile_cam): sim.reset() robot.update(dt) sensor.update(dt) - assert "Could not find prim with path" in str(excinfo.value) + err_msg = str(excinfo.value) + assert "Could not find prim with path" in err_msg or "does not match the number of environments" in err_msg @pytest.mark.isaacsim_ci diff --git a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py index 79cf307c9cb3..2a05b5f70096 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py +++ b/source/isaaclab_mimic/isaaclab_mimic/locomanipulation_sdg/scene_utils.py @@ -10,7 +10,7 @@ import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.sim.views import XformPrimView +from isaaclab.sim.views import FrameView from .occupancy_map_utils import OccupancyMap, intersect_occupancy_maps from .transform_utils import transform_mul @@ -101,19 +101,21 @@ def __init__(self, scene, entity_name: str): self.scene = scene self.entity_name = entity_name - def _get_xform_view(self) -> XformPrimView: - """Return the XformPrimView for this asset, refreshing it if prims were not yet cloned.""" + def _get_xform_view(self) -> FrameView: + """Return the FrameView for this asset, refreshing it if prims were not yet cloned.""" xform_prim = self.scene[self.entity_name] if xform_prim.count == 0: # The view was created before environment cloning; rebuild it now that prims exist. - xform_prim = XformPrimView(xform_prim._prim_path, device=xform_prim.device) + xform_prim = FrameView(xform_prim._prim_path, device=xform_prim.device) self.scene.extras[self.entity_name] = xform_prim return xform_prim def get_pose(self): """Get the 3D pose of the entity.""" xform_prim = self._get_xform_view() - position, orientation = xform_prim.get_world_poses() + pos_wp, ori_wp = xform_prim.get_world_poses() + position = wp.to_torch(pos_wp) + orientation = wp.to_torch(ori_wp) pose = torch.cat([position, orientation], dim=-1) return pose @@ -122,7 +124,7 @@ def set_pose(self, pose: torch.Tensor): xform_prim = self._get_xform_view() position = pose[..., :3] orientation = pose[..., 3:] - xform_prim.set_world_poses(position, orientation, None) + xform_prim.set_world_poses(wp.from_torch(position.contiguous()), wp.from_torch(orientation.contiguous()), None) class RelativePose(HasPose): diff --git a/source/isaaclab_newton/config/extension.toml b/source/isaaclab_newton/config/extension.toml index 3639311a8a91..8a95afe963a0 100644 --- a/source/isaaclab_newton/config/extension.toml +++ b/source/isaaclab_newton/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.5.19" +version = "0.5.20" # Description title = "Newton simulation interfaces for IsaacLab core package" diff --git a/source/isaaclab_newton/docs/CHANGELOG.rst b/source/isaaclab_newton/docs/CHANGELOG.rst index f8f670baf455..3e6cdd8cd115 100644 --- a/source/isaaclab_newton/docs/CHANGELOG.rst +++ b/source/isaaclab_newton/docs/CHANGELOG.rst @@ -1,6 +1,22 @@ Changelog --------- +0.5.20 (2026-04-22) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :class:`~isaaclab_newton.sim.views.XformPrimView` providing the Newton + backend implementation for xform prim views. + +Changed +^^^^^^^ + +* Renamed :class:`~isaaclab_newton.sim.views.NewtonSiteXformPrimView` to + :class:`~isaaclab_newton.sim.views.NewtonSiteFrameView`. Old name is kept as a deprecated alias. + + 0.5.19 (2026-04-22) ~~~~~~~~~~~~~~~~~~~ @@ -129,10 +145,6 @@ Fixed so articulation write methods trigger ``eval_fk`` before the next ``collide()``. - -0.5.9 (2026-03-16) -~~~~~~~~~~~~~~~~~~ - Fixed ^^^^^ diff --git a/source/isaaclab_newton/isaaclab_newton/sim/__init__.py b/source/isaaclab_newton/isaaclab_newton/sim/__init__.py new file mode 100644 index 000000000000..b4646aabbd0a --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sim/__init__.py @@ -0,0 +1,10 @@ +# 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 simulation utilities.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_newton/isaaclab_newton/sim/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/sim/__init__.pyi new file mode 100644 index 000000000000..aac4c8327ccb --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sim/__init__.pyi @@ -0,0 +1,10 @@ +# 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 + +__all__ = [ + "views", +] + +from . import views diff --git a/source/isaaclab_newton/isaaclab_newton/sim/views/__init__.py b/source/isaaclab_newton/isaaclab_newton/sim/views/__init__.py new file mode 100644 index 000000000000..44e8303bcedf --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sim/views/__init__.py @@ -0,0 +1,10 @@ +# 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 simulation views.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_newton/isaaclab_newton/sim/views/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/sim/views/__init__.pyi new file mode 100644 index 000000000000..433dfc1e8b6a --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sim/views/__init__.pyi @@ -0,0 +1,10 @@ +# 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 + +__all__ = [ + "NewtonSiteFrameView", +] + +from .newton_site_frame_view import NewtonSiteFrameView diff --git a/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py b/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py new file mode 100644 index 000000000000..e4f2285cb528 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py @@ -0,0 +1,939 @@ +# 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 — Warp-native, GPU-resident pose queries.""" + +from __future__ import annotations + +import logging + +import warp as wp + +from pxr import Gf, Usd, UsdGeom + +import isaaclab.sim as sim_utils +from isaaclab.physics import PhysicsEvent +from isaaclab.sim.views.base_frame_view import BaseFrameView + +from isaaclab_newton.physics.newton_manager import NewtonManager + +logger = logging.getLogger(__name__) + +WORLD_BODY_INDEX = -1 + + +# ------------------------------------------------------------------ +# Warp kernels +# ------------------------------------------------------------------ + + +@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), + out_pos: wp.array(dtype=wp.vec3f), + out_quat: wp.array(dtype=wp.vec4f), +): + """Compute world-space transforms for every site in the view. + + For each site *i*, computes ``world = body_q[site_body[i]] * site_local[i]`` + and splits the result into position and quaternion outputs. When + ``site_body[i] == -1`` the site is world-attached and ``site_local[i]`` is + returned directly. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + A value of ``-1`` indicates a world-attached site. + site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. + out_pos: Output world positions [m], shape ``[num_sites]``. + out_quat: Output world orientations as ``(qx, qy, qz, qw)``, shape ``[num_sites]``. + """ + i = wp.tid() + bid = site_body[i] + if bid == -1: + world = site_local[i] + else: + world = wp.transform_multiply(body_q[bid], site_local[i]) + 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 _compute_site_world_transforms_indexed( + 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), +): + """Indexed variant of :func:`_compute_site_world_transforms`. + + Only computes world transforms for the subset of sites selected by + ``indices``. Thread *i* reads ``indices[i]`` to obtain the site index, + then writes the result to ``out_pos[i]`` / ``out_quat[i]``. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. + indices: Site indices to query, shape ``[M]``. + out_pos: Output world positions [m], shape ``[M]``. + out_quat: Output world orientations as ``(qx, qy, qz, qw)``, shape ``[M]``. + """ + i = wp.tid() + si = indices[i] + bid = site_body[si] + if bid == -1: + 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_scales( + shape_scale: wp.array(dtype=wp.vec3f), + shape_body: wp.array(dtype=wp.int32), + site_body: wp.array(dtype=wp.int32), + num_shapes: wp.int32, + out_scales: wp.array(dtype=wp.vec3f), +): + """Gather per-site scales from collision shapes on the same body. + + For each site *i*, linearly scans all shapes to find the first one whose + ``shape_body`` matches ``site_body[i]`` and copies its scale. Falls back + to ``(1, 1, 1)`` if no shape is found on that body. + + Args: + shape_scale: Per-shape scale vectors from the Newton model, shape ``[num_shapes]``. + shape_body: Per-shape parent body index, shape ``[num_shapes]``. + site_body: Per-site body index, shape ``[num_sites]``. + num_shapes: Total number of shapes in the model. + out_scales: Output scale per site, shape ``[num_sites]``. + """ + i = wp.tid() + bid = site_body[i] + 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 _gather_scales_indexed( + 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), +): + """Indexed variant of :func:`_gather_scales`. + + Args: + shape_scale: Per-shape scale vectors from the Newton model, shape ``[num_shapes]``. + shape_body: Per-shape parent body index, shape ``[num_shapes]``. + site_body: Per-site body index, shape ``[num_sites]``. + indices: Site indices to query, shape ``[M]``. + num_shapes: Total number of shapes in the model. + out_scales: Output scale per queried site, shape ``[M]``. + """ + 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_scales( + site_body: 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 per-site scales to all collision shapes on the same body. + + For each site *i*, writes ``new_scales[i]`` to every shape whose + ``shape_body`` matches ``site_body[i]``. Multiple shapes on the same + body all receive the same scale. + + Args: + site_body: Per-site body index, shape ``[num_sites]``. + new_scales: New scale to apply per site, shape ``[num_sites]``. + shape_body: Per-shape parent body index, shape ``[num_shapes]``. + num_shapes: Total number of shapes in the model. + shape_scale: Per-shape scale vectors to write into (modified in-place), + shape ``[num_shapes]``. + """ + i = wp.tid() + bid = site_body[i] + for s in range(num_shapes): + if shape_body[s] == bid: + shape_scale[s] = new_scales[i] + + +@wp.kernel +def _scatter_scales_indexed( + 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), +): + """Indexed variant of :func:`_scatter_scales`. + + Args: + site_body: Per-site body index, shape ``[num_sites]``. + indices: Site indices to update, shape ``[M]``. + new_scales: New scale to apply per selected site, shape ``[M]``. + shape_body: Per-shape parent body index, shape ``[num_shapes]``. + num_shapes: Total number of shapes in the model. + shape_scale: Per-shape scale vectors to write into (modified in-place), + shape ``[num_shapes]``. + """ + 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] + + +# ------------------------------------------------------------------ +# World-pose site_local write kernels +# ------------------------------------------------------------------ + + +@wp.kernel +def _write_site_local_from_world_poses( + body_q: wp.array(dtype=wp.transformf), + site_body: 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 site local offsets so that the sites reach desired world poses. + + For each site *i*, computes + ``site_local[i] = inv(body_q[site_body[i]]) * desired_world`` so that + a subsequent ``body_q[bid] * site_local[i]`` yields the requested world + pose. For world-attached sites (``site_body[i] == -1``) the desired world + transform is written directly into ``site_local[i]``. + + Does **not** modify ``body_q``. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + world_pos: Desired world positions [m], shape ``[num_sites]``. + world_quat: Desired world orientations as ``(qx, qy, qz, qw)``, shape ``[num_sites]``. + site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. + """ + i = wp.tid() + 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[i] + if bid == -1: + site_local[i] = desired_world + else: + site_local[i] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) + + +@wp.kernel +def _write_site_local_from_world_poses_indexed( + 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), +): + """Indexed variant of :func:`_write_site_local_from_world_poses`. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + indices: Site indices to update, shape ``[M]``. + world_pos: Desired world positions [m], shape ``[M]``. + world_quat: Desired world orientations as ``(qx, qy, qz, qw)``, shape ``[M]``. + site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. + """ + 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 == -1: + site_local[si] = desired_world + else: + site_local[si] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) + + +# ------------------------------------------------------------------ +# Local-pose Warp kernels +# ------------------------------------------------------------------ + + +@wp.kernel +def _compute_site_local_transforms( + body_q: wp.array(dtype=wp.transformf), + site_body: wp.array(dtype=wp.int32), + site_local: wp.array(dtype=wp.transformf), + parent_site_body: wp.array(dtype=wp.int32), + parent_site_local: wp.array(dtype=wp.transformf), + out_pos: wp.array(dtype=wp.vec3f), + out_quat: wp.array(dtype=wp.vec4f), +): + """Compute parent-relative transforms for every site in the view. + + For each site *i*, computes the world pose of both the site and its USD + parent, then returns ``inv(parent_world) * prim_world``. When + ``site_body[i] == -1`` the site is world-attached and ``site_local[i]`` + is used as the world transform directly. The same convention applies to + the parent arrays. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. + parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. + parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. + out_pos: Output parent-relative positions [m], shape ``[num_sites]``. + out_quat: Output parent-relative orientations as ``(qx, qy, qz, qw)``, + shape ``[num_sites]``. + """ + i = wp.tid() + prim_bid = site_body[i] + if prim_bid == -1: + prim_world = site_local[i] + else: + prim_world = wp.transform_multiply(body_q[prim_bid], site_local[i]) + + parent_bid = parent_site_body[i] + if parent_bid == -1: + parent_world = parent_site_local[i] + else: + parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[i]) + + local_tf = wp.transform_multiply(wp.transform_inverse(parent_world), prim_world) + 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 _compute_site_local_transforms_indexed( + body_q: wp.array(dtype=wp.transformf), + site_body: wp.array(dtype=wp.int32), + site_local: wp.array(dtype=wp.transformf), + parent_site_body: wp.array(dtype=wp.int32), + parent_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), +): + """Indexed variant of :func:`_compute_site_local_transforms`. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. + parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. + parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. + indices: Site indices to query, shape ``[M]``. + out_pos: Output parent-relative positions [m], shape ``[M]``. + out_quat: Output parent-relative orientations as ``(qx, qy, qz, qw)``, + shape ``[M]``. + """ + i = wp.tid() + si = indices[i] + prim_bid = site_body[si] + if prim_bid == -1: + prim_world = site_local[si] + else: + prim_world = wp.transform_multiply(body_q[prim_bid], site_local[si]) + + parent_bid = parent_site_body[si] + if parent_bid == -1: + parent_world = parent_site_local[si] + else: + parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[si]) + + local_tf = wp.transform_multiply(wp.transform_inverse(parent_world), prim_world) + 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_local_poses( + body_q: wp.array(dtype=wp.transformf), + site_body: wp.array(dtype=wp.int32), + parent_site_body: wp.array(dtype=wp.int32), + parent_site_local: wp.array(dtype=wp.transformf), + local_pos: wp.array(dtype=wp.vec3f), + local_quat: wp.array(dtype=wp.vec4f), + site_local: wp.array(dtype=wp.transformf), +): + """Update site local offsets so that sites reach desired parent-relative poses. + + For each site *i*, reconstructs the desired world pose as + ``parent_world * desired_local``, then solves for the body-relative offset: + ``site_local[i] = inv(body_q[bid]) * desired_world``. For world-attached + sites (``site_body[i] == -1``) the world transform is written directly. + + Does **not** modify ``body_q``. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. + parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. + local_pos: Desired parent-relative positions [m], shape ``[num_sites]``. + local_quat: Desired parent-relative orientations as ``(qx, qy, qz, qw)``, + shape ``[num_sites]``. + site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. + """ + i = wp.tid() + parent_bid = parent_site_body[i] + if parent_bid == -1: + parent_world = parent_site_local[i] + else: + parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[i]) + + l_pos = local_pos[i] + l_q = local_quat[i] + local_tf = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) + desired_world = wp.transform_multiply(parent_world, local_tf) + + bid = site_body[i] + if bid == -1: + site_local[i] = desired_world + else: + site_local[i] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) + + +@wp.kernel +def _write_site_local_from_local_poses_indexed( + body_q: wp.array(dtype=wp.transformf), + site_body: wp.array(dtype=wp.int32), + parent_site_body: wp.array(dtype=wp.int32), + parent_site_local: wp.array(dtype=wp.transformf), + 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), +): + """Indexed variant of :func:`_write_site_local_from_local_poses`. + + Args: + body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. + site_body: Per-site body index (flat model-level), shape ``[num_sites]``. + parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. + parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. + indices: Site indices to update, shape ``[M]``. + local_pos: Desired parent-relative positions [m], shape ``[M]``. + local_quat: Desired parent-relative orientations as ``(qx, qy, qz, qw)``, + shape ``[M]``. + site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. + """ + i = wp.tid() + si = indices[i] + parent_bid = parent_site_body[si] + if parent_bid == -1: + parent_world = parent_site_local[si] + else: + parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[si]) + + l_pos = local_pos[i] + l_q = local_quat[i] + local_tf = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) + desired_world = wp.transform_multiply(parent_world, local_tf) + + bid = site_body[si] + if bid == -1: + site_local[si] = desired_world + else: + site_local[si] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) + + +# ------------------------------------------------------------------ +# View class +# ------------------------------------------------------------------ + + +class NewtonSiteFrameView(BaseFrameView): + """Batched prim view for non-physics prims tracked as sites on Newton bodies. + + Each matched USD prim must be a **non-physics** prim (camera, sensor, + Xform marker, etc.) that sits as a child of a Newton rigid body in the + USD hierarchy. The prim path must **not** resolve directly to a physics + body or collision shape -- those are owned by Newton and should be + accessed through :class:`~isaaclab_newton.assets.Articulation` or + :class:`~isaaclab_newton.assets.RigidObject` instead. + + At init time each prim is resolved to a ``(body_index, site_local)`` + pair via ancestor walk: the nearest ancestor that appears in + ``model.body_label`` becomes the attachment body, and the relative USD + transform becomes the site offset. If no body ancestor exists the prim + is attached to the world frame (``body_index = -1``). + + World poses are computed on GPU as + ``body_q[body_index] * site_local`` via a Warp kernel. Both + ``set_world_poses`` and ``set_local_poses`` update ``site_local`` -- + neither touches ``body_q``. + + All getters return ``wp.array``. Setters accept ``wp.array``. + + Raises: + ValueError: If any matched prim resolves to a Newton physics body + or collision shape. + """ + + def __init__(self, prim_path: str, device: str = "cpu", stage: Usd.Stage | None = None, **kwargs): + """Initialize the Newton site-based frame view. + + Resolves all USD prims matching ``prim_path`` and, for each one, walks + the USD ancestor hierarchy to find the nearest Newton rigid body. The + relative transform between the prim and its ancestor body becomes the + site's local offset. + + If the Newton model is already finalized the view initializes + immediately; otherwise initialization is deferred to a + :attr:`PhysicsEvent.PHYSICS_READY` callback. + + Args: + prim_path: USD prim path pattern (may contain regex). + device: Warp device for GPU arrays (e.g. ``"cuda:0"``). + stage: USD stage to search. Defaults to the current stage. + **kwargs: Unused; accepted for interface compatibility with other + :class:`~isaaclab.sim.views.BaseFrameView` backends. + """ + self._prim_path = prim_path + self._device = device + + stage = sim_utils.get_current_stage() if stage is None else stage + self._prims: list[Usd.Prim] = sim_utils.find_matching_prims(prim_path, stage=stage) + + model = NewtonManager.get_model() + if model is not None: + self._initialize_impl(model) + else: + self._physics_ready_handle = NewtonManager.register_callback( + self._on_physics_ready, PhysicsEvent.PHYSICS_READY, name=f"site_view_{prim_path}" + ) + + def _on_physics_ready(self, _event) -> None: + """Callback invoked when the Newton model becomes available.""" + self._initialize_impl(NewtonManager.get_model()) + + def _initialize_impl(self, model) -> None: + """Resolve USD prims to Newton body indices and allocate GPU buffers.""" + body_labels = list(model.body_label) + body_label_set = set(body_labels) + body_label_to_idx = {path: idx for idx, path in enumerate(body_labels)} + shape_label_set = set(model.shape_label) + + xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + + site_bodies: list[int] = [] + site_locals: list[list[float]] = [] + parent_bodies: list[int] = [] + parent_locals: list[list[float]] = [] + + identity_xform = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] + resolve_cache: dict[str, tuple[int, list[float]]] = {} + + for prim in self._prims: + pp = prim.GetPath().pathString + if pp in body_label_set: + raise ValueError( + f"FrameView prim '{pp}' is a Newton physics body. " + "FrameView should only be used for non-physics prims (cameras, sensors, Xform markers). " + "Use Articulation or RigidObject APIs to control physics bodies." + ) + if pp in shape_label_set: + raise ValueError( + f"FrameView prim '{pp}' is a Newton collision shape. " + "FrameView should only be used for non-physics prims (cameras, sensors, Xform markers). " + "Use Articulation or RigidObject APIs to control collision shapes." + ) + + body_idx, local_xform = self._resolve_ancestor_body(prim, body_label_to_idx, xform_cache) + site_bodies.append(body_idx) + site_locals.append(local_xform) + + parent = prim.GetParent() + if not parent or not parent.IsValid() or parent.GetPath().pathString == "/": + parent_bodies.append(WORLD_BODY_INDEX) + parent_locals.append(identity_xform) + else: + parent_path = parent.GetPath().pathString + if parent_path in resolve_cache: + pb_idx, pb_local = resolve_cache[parent_path] + elif parent_path in body_label_to_idx: + pb_idx = body_label_to_idx[parent_path] + pb_local = identity_xform + resolve_cache[parent_path] = (pb_idx, pb_local) + else: + pb_idx, pb_local = self._resolve_ancestor_body(parent, body_label_to_idx, xform_cache) + resolve_cache[parent_path] = (pb_idx, pb_local) + parent_bodies.append(pb_idx) + parent_locals.append(pb_local) + + 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._parent_site_body = wp.array(parent_bodies, dtype=wp.int32, device=device) + self._parent_site_local = wp.array( + [wp.transform(*x) for x in parent_locals], + dtype=wp.transformf, + 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) + + @staticmethod + def _resolve_ancestor_body( + prim: Usd.Prim, + body_label_to_idx: dict[str, int], + xform_cache: UsdGeom.XformCache, + ) -> tuple[int, list[float]]: + """Walk USD ancestors to find the nearest Newton body and compute the relative local transform. + + Args: + prim: The USD prim to resolve. + body_label_to_idx: Dict mapping body prim paths to their Newton body indices. + xform_cache: USD xform cache for efficient transform lookups. + + Returns: + A tuple ``(body_index, local_xform_7)`` where *local_xform_7* is + ``[tx, ty, tz, qx, qy, qz, qw]``. If no body ancestor exists, + ``body_index`` is :data:`WORLD_BODY_INDEX` and the local transform + is the prim's world transform. + """ + prim_world_tf = xform_cache.GetLocalToWorldTransform(prim) + prim_world_tf.Orthonormalize() + + ancestor = prim.GetParent() + while ancestor and ancestor.IsValid() and ancestor.GetPath().pathString != "/": + ancestor_path = ancestor.GetPath().pathString + body_idx = body_label_to_idx.get(ancestor_path) + if body_idx is not None: + ancestor_world_tf = xform_cache.GetLocalToWorldTransform(ancestor) + ancestor_world_tf.Orthonormalize() + local_tf = prim_world_tf * ancestor_world_tf.GetInverse() + return body_idx, _gf_matrix_to_xform7(local_tf) + ancestor = ancestor.GetParent() + + return WORLD_BODY_INDEX, _gf_matrix_to_xform7(prim_world_tf) + + @property + def prims(self) -> list: + """List of USD prims being managed by this view.""" + return self._prims + + @property + def count(self) -> int: + """Number of prims in this view.""" + return len(self._prims) + + # ------------------------------------------------------------------ + # World poses + # ------------------------------------------------------------------ + + def get_world_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + """Get world-space positions and orientations. + + Args: + indices: Subset of sites to query. ``None`` means all sites. + + Returns: + A tuple ``(positions, orientations)`` as ``wp.array`` of shapes + ``(M, 3)`` and ``(M, 4)`` respectively. + """ + state = NewtonManager.get_state_0() + + if indices is not None: + n = len(indices) + pos_buf = wp.zeros(n, dtype=wp.vec3f, device=self._device) + quat_buf = wp.zeros(n, dtype=wp.vec4f, device=self._device) + wp.launch( + _compute_site_world_transforms_indexed, + dim=n, + inputs=[state.body_q, self._site_body, self._site_local, indices], + outputs=[pos_buf, quat_buf], + device=self._device, + ) + return pos_buf, quat_buf + + wp.launch( + _compute_site_world_transforms, + dim=self.count, + inputs=[state.body_q, self._site_body, self._site_local], + outputs=[self._pos_buf, self._quat_buf], + device=self._device, + ) + return self._pos_buf, self._quat_buf + + def set_world_poses( + self, + positions: wp.array | None = None, + orientations: wp.array | None = None, + indices: wp.array | None = None, + ) -> None: + """Set world-space positions and/or orientations. + + Updates the internal ``site_local`` offsets so that + ``body_q[body] * new_site_local`` yields the desired world pose. + Does **not** modify ``body_q``. + + Args: + positions: Desired world positions ``(M, 3)``. ``None`` leaves + positions unchanged. + orientations: Desired world quaternions ``(M, 4)`` as + ``(qx, qy, qz, qw)``. ``None`` leaves orientations unchanged. + indices: Subset of sites to update. ``None`` means all sites. + """ + if positions is None and orientations is None: + return + + state = NewtonManager.get_state_0() + + if positions is None or orientations is None: + cur_pos, cur_quat = self.get_world_poses(indices) + if positions is None: + positions = cur_pos + if orientations is None: + orientations = cur_quat + + if indices is not None: + wp.launch( + _write_site_local_from_world_poses_indexed, + dim=len(indices), + inputs=[state.body_q, self._site_body, indices, positions, orientations, self._site_local], + device=self._device, + ) + else: + wp.launch( + _write_site_local_from_world_poses, + dim=self.count, + inputs=[state.body_q, self._site_body, positions, orientations, self._site_local], + device=self._device, + ) + + # ------------------------------------------------------------------ + # Local poses (parent-relative) + # ------------------------------------------------------------------ + + def get_local_poses(self, indices: wp.array | None = None) -> tuple[wp.array, wp.array]: + """Get parent-relative positions and orientations. + + Computes ``inv(parent_world) * prim_world`` for each site. + + Args: + indices: Subset of sites to query. ``None`` means all sites. + + Returns: + A tuple ``(translations, orientations)`` as ``wp.array`` of shapes + ``(M, 3)`` and ``(M, 4)`` respectively. + """ + state = NewtonManager.get_state_0() + + if indices is not None: + n = len(indices) + pos_buf = wp.zeros(n, dtype=wp.vec3f, device=self._device) + quat_buf = wp.zeros(n, dtype=wp.vec4f, device=self._device) + wp.launch( + _compute_site_local_transforms_indexed, + dim=n, + inputs=[ + state.body_q, + self._site_body, + self._site_local, + self._parent_site_body, + self._parent_site_local, + indices, + ], + outputs=[pos_buf, quat_buf], + device=self._device, + ) + return pos_buf, quat_buf + + wp.launch( + _compute_site_local_transforms, + dim=self.count, + inputs=[ + state.body_q, + self._site_body, + self._site_local, + self._parent_site_body, + self._parent_site_local, + ], + outputs=[self._local_pos_buf, self._local_quat_buf], + device=self._device, + ) + return self._local_pos_buf, self._local_quat_buf + + def set_local_poses( + self, + translations: wp.array | None = None, + orientations: wp.array | None = None, + indices: wp.array | None = None, + ) -> None: + """Set parent-relative translations and/or orientations. + + Updates the internal ``site_local`` offsets so that + ``inv(parent_world) * (body_q[bid] * site_local)`` yields the desired + local pose. Does **not** modify ``body_q``. + + Args: + translations: Desired parent-relative translations ``(M, 3)``. + ``None`` leaves translations unchanged. + orientations: Desired parent-relative quaternions ``(M, 4)`` as + ``(qx, qy, qz, qw)``. ``None`` leaves orientations unchanged. + indices: Subset of sites to update. ``None`` means all sites. + """ + if translations is None and orientations is None: + return + + state = NewtonManager.get_state_0() + + if translations is None or orientations is None: + cur_pos, cur_quat = self.get_local_poses(indices) + if translations is None: + translations = cur_pos + if orientations is None: + orientations = cur_quat + + if indices is not None: + wp.launch( + _write_site_local_from_local_poses_indexed, + dim=len(indices), + inputs=[ + state.body_q, + self._site_body, + self._parent_site_body, + self._parent_site_local, + indices, + translations, + orientations, + self._site_local, + ], + device=self._device, + ) + else: + wp.launch( + _write_site_local_from_local_poses, + dim=self.count, + inputs=[ + state.body_q, + self._site_body, + self._parent_site_body, + self._parent_site_local, + translations, + orientations, + self._site_local, + ], + device=self._device, + ) + + # ------------------------------------------------------------------ + # Scales + # ------------------------------------------------------------------ + + def get_scales(self, indices: wp.array | None = None) -> wp.array: + """Get per-site scales by reading from the first collision shape on the same body. + + Args: + indices: Subset of sites to query. ``None`` means all sites. + + Returns: + A ``wp.array`` of shape ``(M, 3)``. + """ + model = NewtonManager.get_model() + num_shapes = model.shape_count + + if indices is not None: + n = len(indices) + out = wp.zeros(n, dtype=wp.vec3f, device=self._device) + wp.launch( + _gather_scales_indexed, + dim=n, + inputs=[model.shape_scale, model.shape_body, self._site_body, indices, num_shapes], + outputs=[out], + device=self._device, + ) + else: + out = wp.zeros(self.count, dtype=wp.vec3f, device=self._device) + wp.launch( + _gather_scales, + dim=self.count, + inputs=[model.shape_scale, model.shape_body, self._site_body, num_shapes], + outputs=[out], + device=self._device, + ) + return out + + def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: + """Set per-site scales by writing to all collision shapes on the same body. + + Args: + scales: New scales ``(M, 3)`` as ``wp.array``. + indices: Subset of sites to update. ``None`` means all sites. + """ + model = NewtonManager.get_model() + num_shapes = model.shape_count + + if indices is not None: + wp.launch( + _scatter_scales_indexed, + dim=len(indices), + inputs=[self._site_body, indices, scales, model.shape_body, num_shapes, model.shape_scale], + device=self._device, + ) + else: + wp.launch( + _scatter_scales, + dim=self.count, + inputs=[self._site_body, scales, model.shape_body, num_shapes, model.shape_scale], + device=self._device, + ) + + +def _gf_matrix_to_xform7(mat: Gf.Matrix4d) -> list[float]: + """Convert a ``Gf.Matrix4d`` to ``[tx, ty, tz, qx, qy, qz, qw]``.""" + t = mat.ExtractTranslation() + q = mat.ExtractRotationQuat() + imag = q.GetImaginary() + return [float(t[0]), float(t[1]), float(t[2]), float(imag[0]), float(imag[1]), float(imag[2]), float(q.GetReal())] diff --git a/source/isaaclab_newton/test/sim/__init__.py b/source/isaaclab_newton/test/sim/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_newton/test/sim/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py b/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py new file mode 100644 index 000000000000..9785b6d62e2e --- /dev/null +++ b/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py @@ -0,0 +1,198 @@ +# 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 backend tests for FrameView. + +Imports the shared contract tests and provides the Newton-specific +``view_factory`` fixture. Also includes Newton-only guard tests and +the world-attached prim edge case. +""" + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "isaaclab" / "test" / "sim")) + +import pytest +import torch +import warp as wp +from frame_view_contract_utils import * # noqa: F401, F403 — import all contract tests +from frame_view_contract_utils import CHILD_OFFSET, ViewBundle, _wp_vec3f, _wp_vec4f +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg +from isaaclab_newton.physics.newton_manager import NewtonManager +from isaaclab_newton.sim.views import NewtonSiteFrameView as FrameView + +from pxr import Gf + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObjectCfg +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.utils import configclass + +NEWTON_SIM_CFG = SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg())) +WORLD_MARKER_POS = (5.0, 3.0, 1.0) + + +@configclass +class _SceneCfg(InteractiveSceneCfg): + cube: RigidObjectCfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube", + spawn=sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 1.0)), + ) + + +def _sim_context(device, num_envs=4): + NEWTON_SIM_CFG.device = device + return build_simulation_context(device=device, sim_cfg=NEWTON_SIM_CFG, add_ground_plane=True) + + +def _get_body_positions(num_envs, device="cpu"): + model = NewtonManager.get_model() + body_labels = list(model.body_label) + body_q_t = wp.to_torch(NewtonManager.get_state_0().body_q) + return torch.stack([body_q_t[body_labels.index(f"/World/envs/env_{i}/Cube"), :3] for i in range(num_envs)]) + + +def _set_body_positions(positions, num_envs): + model = NewtonManager.get_model() + body_labels = list(model.body_label) + body_q_t = wp.to_torch(NewtonManager.get_state_0().body_q) + for i in range(num_envs): + body_q_t[body_labels.index(f"/World/envs/env_{i}/Cube"), :3] = positions[i] + + +# ------------------------------------------------------------------ +# Contract fixture +# ------------------------------------------------------------------ + + +@pytest.fixture +def view_factory(): + """Newton factory: CameraMount child Xform at CHILD_OFFSET under each Cube body.""" + + def factory(num_envs: int, device: str) -> ViewBundle: + ctx = _sim_context(device, num_envs=num_envs) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None + InteractiveScene(_SceneCfg(num_envs=num_envs, env_spacing=2.0)) + + stage = sim_utils.get_current_stage() + for i in range(num_envs): + prim = stage.DefinePrim(f"/World/envs/env_{i}/Cube/CameraMount", "Xform") + sim_utils.standardize_xform_ops(prim) + prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(*CHILD_OFFSET)) + prim.GetAttribute("xformOp:orient").Set(Gf.Quatd(1.0, 0.0, 0.0, 0.0)) + + sim.reset() + view = FrameView("/World/envs/env_.*/Cube/CameraMount", device=device) + + return ViewBundle( + view=view, + get_parent_pos=_get_body_positions, + set_parent_pos=_set_body_positions, + teardown=lambda: ctx.__exit__(None, None, None), + ) + + return factory + + +# ================================================================== +# Newton-only: guard tests +# ================================================================== + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_reject_body_path(device): + """FrameView rejects prim paths that resolve to a Newton physics body.""" + ctx = _sim_context(device, num_envs=2) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None + InteractiveScene(_SceneCfg(num_envs=2, env_spacing=2.0)) + sim.reset() + + with pytest.raises(ValueError, match="physics body"): + FrameView("/World/envs/env_.*/Cube", device=device) + ctx.__exit__(None, None, None) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_reject_shape_path(device): + """FrameView rejects prim paths that resolve to a Newton collision shape.""" + ctx = _sim_context(device, num_envs=2) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None + InteractiveScene(_SceneCfg(num_envs=2, env_spacing=2.0)) + sim.reset() + + shape_labels = list(NewtonManager.get_model().shape_label) + if not shape_labels: + pytest.skip("No shapes in model") + + with pytest.raises(ValueError, match="collision shape"): + FrameView(shape_labels[0], device=device) + ctx.__exit__(None, None, None) + + +# ================================================================== +# Newton edge case: world-attached prim (body=-1) +# ================================================================== + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_world_attached_returns_initial_pose(device): + """A world-rooted Xform returns its USD-authored position.""" + ctx = _sim_context(device, num_envs=2) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None + InteractiveScene(_SceneCfg(num_envs=2, env_spacing=2.0)) + + stage = sim_utils.get_current_stage() + prim = stage.DefinePrim("/World/StaticMarker", "Xform") + sim_utils.standardize_xform_ops(prim) + prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(*WORLD_MARKER_POS)) + prim.GetAttribute("xformOp:orient").Set(Gf.Quatd(1.0, 0.0, 0.0, 0.0)) + + sim.reset() + view = FrameView("/World/StaticMarker", device=device) + + pos = wp.to_torch(view.get_world_poses()[0]) + expected = torch.tensor([list(WORLD_MARKER_POS)], device=device) + torch.testing.assert_close(pos, expected, atol=1e-5, rtol=0) + ctx.__exit__(None, None, None) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_world_attached_set_world_roundtrip(device): + """A world-attached prim can be repositioned via set_world_poses.""" + ctx = _sim_context(device, num_envs=2) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None + InteractiveScene(_SceneCfg(num_envs=2, env_spacing=2.0)) + + stage = sim_utils.get_current_stage() + prim = stage.DefinePrim("/World/StaticMarker", "Xform") + sim_utils.standardize_xform_ops(prim) + prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(*WORLD_MARKER_POS)) + prim.GetAttribute("xformOp:orient").Set(Gf.Quatd(1.0, 0.0, 0.0, 0.0)) + + sim.reset() + view = FrameView("/World/StaticMarker", device=device) + + new_pos = _wp_vec3f([[10.0, 20.0, 30.0]], device=device) + new_quat = _wp_vec4f([[0.0, 0.0, 0.0, 1.0]], device=device) + view.set_world_poses(new_pos, new_quat) + + ret_pos, ret_quat = view.get_world_poses() + torch.testing.assert_close(wp.to_torch(ret_pos), wp.to_torch(new_pos), atol=1e-5, rtol=0) + torch.testing.assert_close(wp.to_torch(ret_quat), wp.to_torch(new_quat), atol=1e-5, rtol=0) + ctx.__exit__(None, None, None) diff --git a/source/isaaclab_physx/config/extension.toml b/source/isaaclab_physx/config/extension.toml index d05b38808199..f9368e59a0f5 100644 --- a/source/isaaclab_physx/config/extension.toml +++ b/source/isaaclab_physx/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.5.20" +version = "0.5.21" # Description title = "PhysX simulation interfaces for IsaacLab core package" diff --git a/source/isaaclab_physx/docs/CHANGELOG.rst b/source/isaaclab_physx/docs/CHANGELOG.rst index d722acb0687a..a343bf0367fd 100644 --- a/source/isaaclab_physx/docs/CHANGELOG.rst +++ b/source/isaaclab_physx/docs/CHANGELOG.rst @@ -1,6 +1,22 @@ Changelog --------- +0.5.21 (2026-04-22) +~~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +* Added :class:`~isaaclab_physx.sim.views.XformPrimView` providing the PhysX/Fabric + backend implementation for xform prim views. + +Changed +^^^^^^^ + +* Renamed :class:`~isaaclab_physx.sim.views.FabricXformPrimView` to + :class:`~isaaclab_physx.sim.views.FabricFrameView`. Old name is kept as a deprecated alias. + + 0.5.20 (2026-04-21) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py index ceb6089dc2ce..766137753dfe 100644 --- a/source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py +++ b/source/isaaclab_physx/isaaclab_physx/scene_data_providers/physx_scene_data_provider.py @@ -54,7 +54,7 @@ class PhysxSceneDataProvider(BaseSceneDataProvider): """Scene data provider for Omni PhysX backend. Supports: - - body poses via PhysX tensor views, with XformPrimView fallback + - body poses via PhysX tensor views, with FrameView fallback - camera poses & intrinsics - USD stage handles - Newton model/state handles @@ -560,13 +560,13 @@ def _apply_view_poses(self, view: Any, view_key: str, positions: Any, orientatio return count def _apply_xform_poses(self, positions: Any, orientations: Any, covered: Any, xform_mask: Any) -> int: - """Fill remaining poses using XformPrimView (USD fallback). + """Fill remaining poses using FrameView (USD fallback). This is slower but more robust when PhysX views don't cover all bodies. """ import torch - from isaaclab.sim.views import XformPrimView + from isaaclab.sim.views import FrameView uncovered = torch.where(~covered)[0].cpu().tolist() if not uncovered: @@ -578,14 +578,14 @@ def _apply_xform_poses(self, positions: Any, orientations: Any, covered: Any, xf path = self._rigid_body_paths[idx] try: if path not in self._xform_views: - self._xform_views[path] = XformPrimView( + self._xform_views[path] = FrameView( path, device=self._device, stage=self._stage, validate_xform_ops=False ) - pos, quat = self._xform_views[path].get_world_poses() - if pos is not None and quat is not None: - positions[idx] = pos.to(device=self._device, dtype=torch.float32).squeeze() - orientations[idx] = quat.to(device=self._device, dtype=torch.float32).squeeze() + pos_wp, quat_wp = self._xform_views[path].get_world_poses() + if pos_wp is not None and quat_wp is not None: + positions[idx] = wp.to_torch(pos_wp).to(device=self._device, dtype=torch.float32).squeeze() + orientations[idx] = wp.to_torch(quat_wp).to(device=self._device, dtype=torch.float32).squeeze() covered[idx] = True xform_mask[idx] = True count += 1 @@ -605,7 +605,7 @@ def _apply_xform_poses(self, positions: Any, orientations: Any, covered: Any, xf def _convert_xform_quats(self, orientations: Any, xform_mask: Any) -> Any: """Return quaternions in xyzw convention. - PhysX views, XformPrimView, and resolve_prim_pose() in Isaac Lab all use xyzw. + PhysX views, FrameView, and resolve_prim_pose() in Isaac Lab all use xyzw. Keeping this helper as a no-op preserves a single conversion point if conventions ever diverge again. """ diff --git a/source/isaaclab_physx/isaaclab_physx/sim/__init__.pyi b/source/isaaclab_physx/isaaclab_physx/sim/__init__.pyi index c75cccf3f04a..abc8d0087afd 100644 --- a/source/isaaclab_physx/isaaclab_physx/sim/__init__.pyi +++ b/source/isaaclab_physx/isaaclab_physx/sim/__init__.pyi @@ -11,6 +11,7 @@ __all__ = [ "spawn_deformable_body_material", "DeformableBodyMaterialCfg", "SurfaceDeformableBodyMaterialCfg", + "views", ] from .schemas import ( @@ -24,3 +25,4 @@ from .spawners import ( DeformableBodyMaterialCfg, SurfaceDeformableBodyMaterialCfg, ) +from . import views diff --git a/source/isaaclab_physx/isaaclab_physx/sim/views/__init__.py b/source/isaaclab_physx/isaaclab_physx/sim/views/__init__.py new file mode 100644 index 000000000000..85c69b44a24f --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/sim/views/__init__.py @@ -0,0 +1,10 @@ +# 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 + +"""PhysX simulation views.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_physx/isaaclab_physx/sim/views/__init__.pyi b/source/isaaclab_physx/isaaclab_physx/sim/views/__init__.pyi new file mode 100644 index 000000000000..789d62af9d14 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/sim/views/__init__.pyi @@ -0,0 +1,10 @@ +# 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 + +__all__ = [ + "FabricFrameView", +] + +from .fabric_frame_view import FabricFrameView diff --git a/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py b/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py new file mode 100644 index 000000000000..87adad2238c4 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/sim/views/fabric_frame_view.py @@ -0,0 +1,403 @@ +# 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 + +"""PhysX FrameView with Fabric GPU acceleration.""" + +from __future__ import annotations + +import logging + +import torch +import warp as wp + +from pxr import Usd + +import isaaclab.sim as sim_utils +from isaaclab.app.settings_manager import SettingsManager +from isaaclab.sim.views.base_frame_view import BaseFrameView +from isaaclab.sim.views.usd_frame_view import UsdFrameView +from isaaclab.utils.warp import fabric as fabric_utils + +logger = logging.getLogger(__name__) + + +def _to_float32_2d(a: wp.array | torch.Tensor) -> wp.array | torch.Tensor: + """Ensure array is compatible with Fabric kernels (2-D float32). + + For ``wp.array`` with vec dtypes (``vec3f``, ``vec4f``), uses + :meth:`wp.array.view` for zero-copy reinterpretation. + ``torch.Tensor`` and already-correct 2-D float32 arrays pass through. + """ + if not isinstance(a, wp.array): + return a + if a.shape[0] == 0: + return a + if a.ndim == 2 and a.dtype == wp.float32: + return a + return a.view(dtype=wp.float32) + + +class FabricFrameView(BaseFrameView): + """FrameView with Fabric GPU acceleration for the PhysX backend. + + Uses composition: holds a :class:`UsdFrameView` internally for USD + fallback and non-accelerated operations (local poses, visibility, scales + when Fabric is disabled). + + When Fabric is enabled, world-pose and scale operations use GPU-accelerated + Warp kernels operating on ``omni:fabric:worldMatrix``. All other operations + delegate to the internal USD view. + + All getters return ``wp.array``. Setters accept ``wp.array``. + """ + + def __init__( + self, + prim_path: str, + device: str = "cpu", + validate_xform_ops: bool = True, + sync_usd_on_fabric_write: bool = False, + stage: Usd.Stage | None = None, + ): + self._usd_view = UsdFrameView(prim_path, device=device, validate_xform_ops=validate_xform_ops, stage=stage) + self._device = device + self._sync_usd_on_fabric_write = sync_usd_on_fabric_write + + settings = SettingsManager.instance() + self._use_fabric = bool(settings.get("/physics/fabricEnabled", False)) + + if self._use_fabric and self._device == "cpu": + logger.warning( + "Fabric mode with Warp fabric-array operations is not supported on CPU devices. " + "Falling back to standard USD operations on the CPU. This may impact performance." + ) + self._use_fabric = False + + if self._use_fabric and self._device not in ("cuda", "cuda:0"): + logger.warning( + f"Fabric mode is not supported on device '{self._device}'. " + "USDRT SelectPrims and Warp fabric arrays only support cuda:0. " + "Falling back to standard USD operations. This may impact performance." + ) + self._use_fabric = False + + self._fabric_initialized = False + self._fabric_usd_sync_done = False + self._fabric_selection = None + self._fabric_to_view: wp.array | None = None + self._view_to_fabric: wp.array | None = None + self._default_view_indices: wp.array | None = None + self._fabric_hierarchy = None + self._view_index_attr = f"isaaclab:view_index:{abs(hash(self))}" + + # ------------------------------------------------------------------ + # Delegated properties + # ------------------------------------------------------------------ + + @property + def count(self) -> int: + return self._usd_view.count + + @property + def device(self) -> str: + return self._device + + @property + def prims(self) -> list: + return self._usd_view.prims + + @property + def prim_paths(self) -> list[str]: + return self._usd_view.prim_paths + + # ------------------------------------------------------------------ + # Delegated operations (USD-only) + # ------------------------------------------------------------------ + + def get_visibility(self, indices=None): + return self._usd_view.get_visibility(indices) + + def set_visibility(self, visibility, indices=None): + self._usd_view.set_visibility(visibility, indices) + + # ------------------------------------------------------------------ + # World poses — Fabric-accelerated or USD fallback + # ------------------------------------------------------------------ + + def set_world_poses(self, positions=None, orientations=None, indices=None): + if not self._use_fabric: + self._usd_view.set_world_poses(positions, orientations, indices) + return + + if not self._fabric_initialized: + self._initialize_fabric() + + indices_wp = self._resolve_indices_wp(indices) + count = indices_wp.shape[0] + + dummy = wp.zeros((0, 3), dtype=wp.float32, device=self._device) + positions_wp = _to_float32_2d(positions) if positions is not None else dummy + orientations_wp = ( + _to_float32_2d(orientations) + if orientations is not None + else wp.zeros((0, 4), dtype=wp.float32, device=self._device) + ) + + wp.launch( + kernel=fabric_utils.compose_fabric_transformation_matrix_from_warp_arrays, + dim=count, + inputs=[ + self._fabric_world_matrices, + positions_wp, + orientations_wp, + dummy, + False, + False, + False, + indices_wp, + self._view_to_fabric, + ], + device=self._fabric_device, + ) + wp.synchronize() + + self._fabric_hierarchy.update_world_xforms() + self._fabric_usd_sync_done = True + if self._sync_usd_on_fabric_write: + self._usd_view.set_world_poses(positions, orientations, indices) + + def get_world_poses(self, indices=None): + if not self._use_fabric: + return self._usd_view.get_world_poses(indices) + + if not self._fabric_initialized: + self._initialize_fabric() + if not self._fabric_usd_sync_done: + self._sync_fabric_from_usd_once() + + indices_wp = self._resolve_indices_wp(indices) + count = indices_wp.shape[0] + + use_cached = indices is None or indices == slice(None) + if use_cached: + positions_wp = self._fabric_positions_buf + orientations_wp = self._fabric_orientations_buf + else: + positions_wp = wp.zeros((count, 3), dtype=wp.float32, device=self._device) + orientations_wp = wp.zeros((count, 4), dtype=wp.float32, device=self._device) + + wp.launch( + kernel=fabric_utils.decompose_fabric_transformation_matrix_to_warp_arrays, + dim=count, + inputs=[ + self._fabric_world_matrices, + positions_wp, + orientations_wp, + self._fabric_dummy_buffer, + indices_wp, + self._view_to_fabric, + ], + device=self._fabric_device, + ) + + if use_cached: + wp.synchronize() + return positions_wp, orientations_wp + + # ------------------------------------------------------------------ + # Local poses — USD fallback (Fabric only accelerates world poses) + # ------------------------------------------------------------------ + + def set_local_poses(self, translations=None, orientations=None, indices=None): + self._usd_view.set_local_poses(translations, orientations, indices) + + def get_local_poses(self, indices=None): + return self._usd_view.get_local_poses(indices) + + # ------------------------------------------------------------------ + # Scales — Fabric-accelerated or USD fallback + # ------------------------------------------------------------------ + + def set_scales(self, scales, indices=None): + if not self._use_fabric: + self._usd_view.set_scales(scales, indices) + return + + if not self._fabric_initialized: + self._initialize_fabric() + + indices_wp = self._resolve_indices_wp(indices) + count = indices_wp.shape[0] + + dummy3 = wp.zeros((0, 3), dtype=wp.float32, device=self._device) + dummy4 = wp.zeros((0, 4), dtype=wp.float32, device=self._device) + scales_wp = _to_float32_2d(scales) + + wp.launch( + kernel=fabric_utils.compose_fabric_transformation_matrix_from_warp_arrays, + dim=count, + inputs=[ + self._fabric_world_matrices, + dummy3, + dummy4, + scales_wp, + False, + False, + False, + indices_wp, + self._view_to_fabric, + ], + device=self._fabric_device, + ) + wp.synchronize() + + self._fabric_hierarchy.update_world_xforms() + self._fabric_usd_sync_done = True + if self._sync_usd_on_fabric_write: + self._usd_view.set_scales(scales, indices) + + def get_scales(self, indices=None): + if not self._use_fabric: + return self._usd_view.get_scales(indices) + + if not self._fabric_initialized: + self._initialize_fabric() + if not self._fabric_usd_sync_done: + self._sync_fabric_from_usd_once() + + indices_wp = self._resolve_indices_wp(indices) + count = indices_wp.shape[0] + + use_cached = indices is None or indices == slice(None) + if use_cached: + scales_wp = self._fabric_scales_buf + else: + scales_wp = wp.zeros((count, 3), dtype=wp.float32, device=self._device) + + wp.launch( + kernel=fabric_utils.decompose_fabric_transformation_matrix_to_warp_arrays, + dim=count, + inputs=[ + self._fabric_world_matrices, + self._fabric_dummy_buffer, + self._fabric_dummy_buffer, + scales_wp, + indices_wp, + self._view_to_fabric, + ], + device=self._fabric_device, + ) + + if use_cached: + wp.synchronize() + return scales_wp + + # ------------------------------------------------------------------ + # Internal — Fabric initialization + # ------------------------------------------------------------------ + + def _initialize_fabric(self) -> None: + """Initialize Fabric batch infrastructure for GPU-accelerated pose queries.""" + import usdrt # noqa: PLC0415 + from usdrt import Rt # noqa: PLC0415 + + stage_id = sim_utils.get_current_stage_id() + fabric_stage = usdrt.Usd.Stage.Attach(stage_id) + + for i in range(self.count): + rt_prim = fabric_stage.GetPrimAtPath(self.prim_paths[i]) + rt_xformable = Rt.Xformable(rt_prim) + + has_attr = ( + rt_xformable.HasFabricHierarchyWorldMatrixAttr() + if hasattr(rt_xformable, "HasFabricHierarchyWorldMatrixAttr") + else False + ) + if not has_attr: + rt_xformable.CreateFabricHierarchyWorldMatrixAttr() + + rt_xformable.SetWorldXformFromUsd() + + rt_prim.CreateAttribute(self._view_index_attr, usdrt.Sdf.ValueTypeNames.UInt, custom=True) + rt_prim.GetAttribute(self._view_index_attr).Set(i) + + self._fabric_hierarchy = usdrt.hierarchy.IFabricHierarchy().get_fabric_hierarchy( + fabric_stage.GetFabricId(), fabric_stage.GetStageIdAsStageId() + ) + self._fabric_hierarchy.update_world_xforms() + + self._default_view_indices = wp.zeros((self.count,), dtype=wp.uint32, device=self._device) + wp.launch( + kernel=fabric_utils.arange_k, dim=self.count, inputs=[self._default_view_indices], device=self._device + ) + wp.synchronize() + + fabric_device = self._device + if self._device == "cuda": + logger.warning("Fabric device is not specified, defaulting to 'cuda:0'.") + fabric_device = "cuda:0" + elif self._device.startswith("cuda:"): + if self._device != "cuda:0": + logger.debug( + f"SelectPrims only supports cuda:0. Using cuda:0 for SelectPrims " + f"even though simulation device is {self._device}." + ) + fabric_device = "cuda:0" + + self._fabric_selection = fabric_stage.SelectPrims( + require_attrs=[ + (usdrt.Sdf.ValueTypeNames.UInt, self._view_index_attr, usdrt.Usd.Access.Read), + (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.ReadWrite), + ], + device=fabric_device, + ) + + self._view_to_fabric = wp.zeros((self.count,), dtype=wp.uint32, device=fabric_device) + self._fabric_to_view = wp.fabricarray(self._fabric_selection, self._view_index_attr) + + wp.launch( + kernel=fabric_utils.set_view_to_fabric_array, + dim=self._fabric_to_view.shape[0], + inputs=[self._fabric_to_view, self._view_to_fabric], + device=fabric_device, + ) + wp.synchronize() + + self._fabric_positions_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) + self._fabric_orientations_buf = wp.zeros((self.count, 4), dtype=wp.float32, device=self._device) + self._fabric_scales_buf = wp.zeros((self.count, 3), dtype=wp.float32, device=self._device) + self._fabric_dummy_buffer = wp.zeros((0, 3), dtype=wp.float32, device=self._device) + self._fabric_world_matrices = wp.fabricarray(self._fabric_selection, "omni:fabric:worldMatrix") + self._fabric_stage = fabric_stage + self._fabric_device = fabric_device + + self._fabric_initialized = True + self._fabric_usd_sync_done = False + + def _sync_fabric_from_usd_once(self) -> None: + """Sync Fabric world matrices from USD once, on the first read.""" + if not self._fabric_initialized: + self._initialize_fabric() + + positions_usd, orientations_usd = self._usd_view.get_world_poses() + scales_usd = self._usd_view.get_scales() + + prev_sync = self._sync_usd_on_fabric_write + self._sync_usd_on_fabric_write = False + self.set_world_poses(positions_usd, orientations_usd) + self.set_scales(scales_usd) + self._sync_usd_on_fabric_write = prev_sync + + self._fabric_usd_sync_done = True + + def _resolve_indices_wp(self, indices: wp.array | None) -> wp.array: + """Resolve view indices as a Warp uint32 array.""" + if indices is None or indices == slice(None): + if self._default_view_indices is None: + raise RuntimeError("Fabric indices are not initialized.") + return self._default_view_indices + if indices.dtype != wp.uint32: + return wp.array(indices.numpy().astype("uint32"), dtype=wp.uint32, device=self._device) + return indices diff --git a/source/isaaclab_physx/test/sim/__init__.py b/source/isaaclab_physx/test/sim/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_physx/test/sim/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/source/isaaclab_physx/test/sim/test_views_xform_prim_fabric.py b/source/isaaclab_physx/test/sim/test_views_xform_prim_fabric.py new file mode 100644 index 000000000000..0bc77ccf7223 --- /dev/null +++ b/source/isaaclab_physx/test/sim/test_views_xform_prim_fabric.py @@ -0,0 +1,105 @@ +# 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 + +"""PhysX Fabric backend tests for FrameView. + +Imports the shared contract tests and provides the Fabric-specific +``view_factory`` fixture (SimulationContext with use_fabric=True, +Camera prim type for Fabric SelectPrims compatibility). +""" + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "isaaclab" / "test" / "sim")) + +from isaaclab.app import AppLauncher + +simulation_app = AppLauncher(headless=True).app + +import pytest # noqa: E402 +import torch # noqa: E402 +from frame_view_contract_utils import * # noqa: F401, F403, E402 +from frame_view_contract_utils import CHILD_OFFSET, ViewBundle # noqa: E402 +from isaaclab_physx.sim.views import FabricFrameView as FrameView # noqa: E402 + +from pxr import Gf, UsdGeom # noqa: E402 + +import isaaclab.sim as sim_utils # noqa: E402 + +PARENT_POS = (0.0, 0.0, 1.0) + + +@pytest.fixture(autouse=True) +def test_setup_teardown(): + sim_utils.create_new_stage() + sim_utils.update_stage() + yield + sim_utils.clear_stage() + sim_utils.SimulationContext.clear_instance() + + +def _skip_if_unavailable(device: str): + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + if device == "cpu": + pytest.skip("Warp fabricarray operations on CPU have known issues") + + +# ------------------------------------------------------------------ +# Parent position helpers (via USD xformOps) +# ------------------------------------------------------------------ + + +def _get_parent_positions(num_envs, device="cpu"): + stage = sim_utils.get_current_stage() + xform_cache = UsdGeom.XformCache() + positions = [] + for i in range(num_envs): + prim = stage.GetPrimAtPath(f"/World/Parent_{i}") + tf = xform_cache.GetLocalToWorldTransform(prim) + t = tf.ExtractTranslation() + positions.append([float(t[0]), float(t[1]), float(t[2])]) + return torch.tensor(positions, dtype=torch.float32, device=device) + + +def _set_parent_positions(positions, num_envs): + from pxr import Sdf # noqa: PLC0415 + + stage = sim_utils.get_current_stage() + with Sdf.ChangeBlock(): + for i in range(num_envs): + prim = stage.GetPrimAtPath(f"/World/Parent_{i}") + pos = positions[i] + prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(float(pos[0]), float(pos[1]), float(pos[2]))) + + +# ------------------------------------------------------------------ +# Contract fixture +# ------------------------------------------------------------------ + + +@pytest.fixture +def view_factory(): + """Fabric factory: Camera child at CHILD_OFFSET under parent Xforms, with Fabric enabled.""" + + def factory(num_envs: int, device: str) -> ViewBundle: + _skip_if_unavailable(device) + + stage = sim_utils.get_current_stage() + for i in range(num_envs): + sim_utils.create_prim(f"/World/Parent_{i}", "Xform", translation=PARENT_POS, stage=stage) + sim_utils.create_prim(f"/World/Parent_{i}/Child", "Camera", translation=CHILD_OFFSET, stage=stage) + + sim_utils.SimulationContext(sim_utils.SimulationCfg(dt=0.01, device=device, use_fabric=True)) + view = FrameView("/World/Parent_.*/Child", device=device, sync_usd_on_fabric_write=True) + return ViewBundle( + view=view, + get_parent_pos=_get_parent_positions, + set_parent_pos=_set_parent_positions, + teardown=lambda: None, + ) + + return factory diff --git a/source/isaaclab_tasks/config/extension.toml b/source/isaaclab_tasks/config/extension.toml index bc0640c4cbd6..81334689ea9f 100644 --- a/source/isaaclab_tasks/config/extension.toml +++ b/source/isaaclab_tasks/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "1.5.23" +version = "1.5.24" # Description title = "Isaac Lab Environments" diff --git a/source/isaaclab_tasks/docs/CHANGELOG.rst b/source/isaaclab_tasks/docs/CHANGELOG.rst index b0cf36c9369b..c81768405791 100644 --- a/source/isaaclab_tasks/docs/CHANGELOG.rst +++ b/source/isaaclab_tasks/docs/CHANGELOG.rst @@ -1,6 +1,20 @@ Changelog --------- +1.5.24 (2026-04-22) +~~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Updated locomotion :class:`~isaaclab.sensors.ray_caster.ray_caster_cfg.RayCasterCfg` + height-scanner defaults to spawn a ``raycaster`` Xform child under the robot attachment link + (using :class:`~isaaclab.sim.spawners.sensors.sensors_cfg.RayCasterXformCfg`) so the sensor + works with Newton site-based :class:`~isaaclab.sim.views.FrameView` tracking. +* Updated all sensor configurations to use :class:`~isaaclab.sim.views.FrameView` instead of + the deprecated ``XformPrimView``. + + 1.5.23 (2026-04-21) ~~~~~~~~~~~~~~~~~~~ @@ -10,6 +24,7 @@ Fixed * Refreshed Newton Warp renderer golden images for Dexsuite Kuka-Allegro environment case in ``test_rendering_correctness`` because Newton Warp renderer honors visibility of prims now. + 1.5.22 (2026-04-20) ~~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/terminations.py index 2b87dc69df76..8e530a7d71e0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/terminations.py @@ -51,7 +51,7 @@ def task_done_pick_place_table_frame( env: The RL environment instance. task_link_name: Name of the right wrist link on the robot. object_cfg: Configuration for the object entity. - table_cfg: Configuration for the destination table entity (must be an XformPrimView). + table_cfg: Configuration for the destination table entity (must be a FrameView). right_wrist_max_x: Maximum x position of the right wrist in table frame for task completion. min_x: Minimum x position of the object relative to the table for task completion. max_x: Maximum x position of the object relative to the table for task completion. diff --git a/source/isaaclab_teleop/test/test_oxr_device.py b/source/isaaclab_teleop/test/test_oxr_device.py index 2d2ed8444969..1663a56612d9 100644 --- a/source/isaaclab_teleop/test/test_oxr_device.py +++ b/source/isaaclab_teleop/test/test_oxr_device.py @@ -179,12 +179,12 @@ def test_xr_anchor(empty_env, mock_xrcore): device = OpenXRDevice(OpenXRDeviceCfg(xr_cfg=env_cfg.xr)) # Check that the xr anchor prim is created with the correct pose - xr_anchor_view = sim_utils.XformPrimView("/World/XRAnchor") + xr_anchor_view = sim_utils.FrameView("/World/XRAnchor") assert xr_anchor_view.count == 1 position, orientation = xr_anchor_view.get_world_poses() np.testing.assert_almost_equal(position.numpy(), [[1, 2, 3]]) - # XformPrimView returns quaternion in xyzw format, identity is [0, 0, 0, 1] + # FrameView returns quaternion in xyzw format, identity is [0, 0, 0, 1] np.testing.assert_almost_equal(orientation.numpy(), [[0, 0, 0, 1]]) # Check that xr anchor mode and custom anchor are set correctly @@ -202,7 +202,7 @@ def test_xr_anchor_default(empty_env, mock_xrcore): device = OpenXRDevice(OpenXRDeviceCfg()) # Check that the xr anchor prim is created with the correct default pose - xr_anchor_view = sim_utils.XformPrimView("/World/XRAnchor") + xr_anchor_view = sim_utils.FrameView("/World/XRAnchor") assert xr_anchor_view.count == 1 position, orientation = xr_anchor_view.get_world_poses() @@ -225,7 +225,7 @@ def test_xr_anchor_multiple_devices(empty_env, mock_xrcore): device_2 = OpenXRDevice(OpenXRDeviceCfg()) # Check that the xr anchor prim is created with the correct default pose - xr_anchor_view = sim_utils.XformPrimView("/World/XRAnchor") + xr_anchor_view = sim_utils.FrameView("/World/XRAnchor") assert xr_anchor_view.count == 1 position, orientation = xr_anchor_view.get_world_poses()